diff --git a/starknet/src/authenticators/eth_sig.cairo b/starknet/src/authenticators/eth_sig.cairo index 0bc90c2f..ae841170 100644 --- a/starknet/src/authenticators/eth_sig.cairo +++ b/starknet/src/authenticators/eth_sig.cairo @@ -47,7 +47,7 @@ mod EthSigAuthenticator { use clone::Clone; use sx::space::space::{ISpaceDispatcher, ISpaceDispatcherTrait}; use sx::types::{Strategy, IndexedStrategy, Choice, UserAddress}; - use sx::utils::{signatures, LegacyHashEthAddress}; + use sx::utils::{signatures, legacy_hash::LegacyHashEthAddress}; #[storage] struct Storage { diff --git a/starknet/src/space/space.cairo b/starknet/src/space/space.cairo index 0cc7b795..f88eb4ef 100644 --- a/starknet/src/space/space.cairo +++ b/starknet/src/space/space.cairo @@ -75,6 +75,7 @@ mod Space { NoUpdateStrategy, NoUpdateArray }; use sx::utils::bits::BitSetter; + use sx::utils::legacy_hash::LegacyHashChoice; use sx::external::ownable::Ownable; #[storage] diff --git a/starknet/src/tests.cairo b/starknet/src/tests.cairo index ff157d78..996a672c 100644 --- a/starknet/src/tests.cairo +++ b/starknet/src/tests.cairo @@ -1,5 +1,6 @@ -mod test_space; +mod test_merkle_whitelist; mod test_factory; +mod test_space; mod test_upgrade; mod mocks; diff --git a/starknet/src/tests/test_merkle_whitelist.cairo b/starknet/src/tests/test_merkle_whitelist.cairo new file mode 100644 index 00000000..a7124049 --- /dev/null +++ b/starknet/src/tests/test_merkle_whitelist.cairo @@ -0,0 +1,472 @@ +#[cfg(test)] +mod merkle_utils { + use array::{ArrayTrait, SpanTrait}; + use clone::Clone; + use traits::Into; + use result::ResultTrait; + use option::OptionTrait; + use hash::LegacyHash; + use sx::utils::merkle::{Leaf, Hash}; + use starknet::contract_address_try_from_felt252; + use sx::types::UserAddress; + + impl SpanIntoArray, impl TDrop: Drop> of Into, Array> { + fn into(self: Span) -> Array { + let mut self = self; + let mut output = ArrayTrait::::new(); + loop { + match self.pop_front() { + Option::Some(val) => output.append(val.clone()), + Option::None => { + break; + } + }; + }; + output + } + } + + // Generates the proof for the given `index` in the `merkle_data`. + fn generate_proof(mut merkle_data: Span, mut index: usize) -> Array { + let mut proof = ArrayTrait::new(); + + loop { + if merkle_data.len() == 1 { + break; + } + + if merkle_data.len() % 2 != 0 { + let mut cpy = merkle_data.into(); + cpy.append(0_felt252); // append 0 because of odd length + merkle_data = cpy.span(); + } + + let next_level = get_next_level(merkle_data); + + let mut index_parent = 0_usize; + let mut i = 0_usize; + loop { + if i == merkle_data.len() { + break; + } + if i == index { + index_parent = i / 2; + if i % 2 == 0 { + proof.append(*merkle_data.at(index + 1)); + } else { + proof.append(*merkle_data.at(index - 1)); + } + } + i += 1; + }; + merkle_data = next_level.span(); + index = index_parent; + }; + proof + } + + // Generates the merkle root from the + fn generate_merkle_root(mut merkle_data: Span) -> felt252 { + if merkle_data.len() == 1 { + return *merkle_data.pop_front().unwrap(); + } + + if merkle_data.len() % 2 != 0 { + let mut cpy = merkle_data.into(); + cpy.append(0_felt252); // append 0 because of odd length + merkle_data = cpy.span(); + } + + let next_level = get_next_level(merkle_data); + generate_merkle_root(next_level.span()) + } + + fn get_next_level(mut merkle_data: Span) -> Array { + let mut next_level = ArrayTrait::::new(); + loop { + match merkle_data.pop_front() { + Option::Some(a) => { + match merkle_data.pop_front() { + Option::Some(b) => { + // compare + let a_: u256 = (*a).into(); + let b_: u256 = (*b).into(); + if a_ > b_ { + let node = LegacyHash::hash(*a, *b); + next_level.append(node); + } else { + let node = LegacyHash::hash(*b, *a); + next_level.append(node); + } + }, + Option::None => panic_with_felt252('Incorrect array length'), + } + }, + Option::None => { + break; + } + }; + }; + next_level + } + + // Generates the `merkle_data` from the members. + // The `merkle_data` corresponds to the hashes leaves of the members. + fn generate_merkle_data(members: Span) -> Array { + let mut members_ = members; + let mut output = ArrayTrait::::new(); + loop { + match members_.pop_front() { + Option::Some(leaf) => { + output.append(leaf.hash()); + }, + Option::None => { + break; + }, + }; + }; + output + } + + // Generates n members with voting power 1, 2, 3, and + // address 1, 2, 3, ... + // Even members will be Ethereum addresses and odd members will be Starknet addresses. + fn generate_n_members(n: usize) -> Array { + let mut members = ArrayTrait::::new(); + let mut i = 1_usize; + loop { + if i >= n + 1 { + break; + } + let mut address = UserAddress::Custom(0); + if i % 2 == 0 { + address = UserAddress::Ethereum(starknet::EthAddress { address: i.into() }); + } else { + address = + UserAddress::Starknet(contract_address_try_from_felt252(i.into()).unwrap()); + } + members.append(Leaf { address: address, voting_power: i.into() }); + i += 1; + }; + members + } +} +#[cfg(test)] +mod assert_valid_proof { + use sx::tests::setup::setup::setup::{setup, deploy}; + use array::{ArrayTrait, SpanTrait}; + use option::OptionTrait; + use sx::utils::merkle::{Leaf, assert_valid_proof, Hash}; + use starknet::{contract_address_const, contract_address_try_from_felt252}; + use clone::Clone; + use traits::Into; + use hash::LegacyHash; + use serde::Serde; + use super::merkle_utils::{ + generate_n_members, generate_merkle_data, generate_merkle_root, generate_proof + }; + use sx::types::UserAddress; + + // Generates the proof and verifies the proof for every member in `members`. + fn verify_all_members(members: Span) { + let merkle_data = generate_merkle_data(members); + let root = generate_merkle_root(merkle_data.span()); + let mut index = 0; + loop { + let proof = generate_proof(merkle_data.span(), index); + if index == members.len() { + break; + } + assert_valid_proof(root, *members.at(index), proof.span()); + index += 1; + } + } + + // Replaces the first element of `arr` with `value`. + fn replace_first_element, impl TCopy: Copy>( + mut arr: Span, value: T + ) -> Array { + let mut output = ArrayTrait::new(); + output.append(value); + + arr.pop_front(); // remove first element + loop { + match arr.pop_front() { + Option::Some(v) => output.append(*v), + Option::None => { + break; + }, + }; + }; + output + } + + #[test] + #[available_gas(10000000)] + fn one_member() { + let mut members = generate_n_members(1); + verify_all_members(members.span()); + } + + #[test] + #[available_gas(10000000)] + fn two_members() { + let members = generate_n_members(2); + verify_all_members(members.span()); + } + + #[test] + #[available_gas(10000000)] + fn three_members() { + let members = generate_n_members(3); + verify_all_members(members.span()); + } + + #[test] + #[available_gas(10000000)] + fn four_members() { + let members = generate_n_members(4); + verify_all_members(members.span()); + } + + #[test] + #[available_gas(1000000000)] + fn one_hundred_members() { + let members = generate_n_members(100); + verify_all_members(members.span()); + } + + #[test] + #[available_gas(1000000000)] + fn one_hundred_and_one_members() { + let members = generate_n_members(101); + verify_all_members(members.span()); + } + + #[test] + #[available_gas(1000000000)] + #[should_panic(expected: ('Merkle: Invalid proof', ))] + fn no_leaf() { + let root = 0; + let leaf = Leaf { + address: UserAddress::Starknet(contract_address_const::<0>()), voting_power: 0 + }; + let proof = ArrayTrait::new(); + assert_valid_proof(root, leaf, proof.span()); + } + + #[test] + #[available_gas(10000000)] + #[should_panic(expected: ('Merkle: Invalid proof', ))] + fn invalid_extra_node() { + let mut members = ArrayTrait::new(); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<5>()), voting_power: 5 + } + ); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<4>()), voting_power: 4 + } + ); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<3>()), voting_power: 3 + } + ); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<2>()), voting_power: 2 + } + ); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<1>()), voting_power: 1 + } + ); + let merkle_data = generate_merkle_data(members.span()); + + let root = generate_merkle_root(merkle_data.span()); + let index = 2; + let mut proof = generate_proof(merkle_data.span(), index); + proof.append(0x1337); // Adding a useless node + assert_valid_proof(root, *members.at(index), proof.span()); + } + + + #[test] + #[available_gas(10000000)] + #[should_panic(expected: ('Merkle: Invalid proof', ))] + fn invalid_proof() { + let mut members = ArrayTrait::new(); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<5>()), voting_power: 5 + } + ); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<4>()), voting_power: 4 + } + ); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<3>()), voting_power: 3 + } + ); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<2>()), voting_power: 2 + } + ); + members + .append( + Leaf { + address: UserAddress::Starknet(contract_address_const::<1>()), voting_power: 1 + } + ); + let merkle_data = generate_merkle_data(members.span()); + + let root = generate_merkle_root(merkle_data.span()); + let index = 2; + let proof = generate_proof(merkle_data.span(), index); + let fake_proof = replace_first_element(proof.span(), 0x1337); + + assert_valid_proof(root, *members.at(index), fake_proof.span()); + } +} + +#[cfg(test)] +mod merkle_whitelist_voting_power { + use array::ArrayTrait; + use sx::utils::merkle::Leaf; + use super::merkle_utils::{ + generate_merkle_root, generate_n_members, generate_merkle_data, generate_proof + }; + use sx::voting_strategies::merkle_whitelist::{MerkleWhitelistVotingStrategy}; + use sx::interfaces::IVotingStrategy; + use starknet::syscalls::deploy_syscall; + use starknet::SyscallResult; + use result::ResultTrait; + use option::OptionTrait; + use traits::TryInto; + use sx::interfaces::{IVotingStrategyDispatcher, IVotingStrategyDispatcherTrait}; + use serde::Serde; + use starknet::contract_address_const; + use sx::types::UserAddress; + + #[test] + #[available_gas(1000000000)] + fn one_hundred_members() { + let members = generate_n_members(20); + + let (contract, _) = deploy_syscall( + MerkleWhitelistVotingStrategy::TEST_CLASS_HASH.try_into().unwrap(), + 0, + array::ArrayTrait::::new().span(), + false, + ) + .unwrap(); + let voting_strategy = IVotingStrategyDispatcher { contract_address: contract }; + let timestamp = 0x1234; + let index = 2; + let leaf = *members.at(index); + let voter = leaf.address; + + let merkle_data = generate_merkle_data(members.span()); + let root = generate_merkle_root(merkle_data.span()); + let proof = generate_proof(merkle_data.span(), index); + + let mut params = ArrayTrait::::new(); + root.serialize(ref params); + + let mut user_params = ArrayTrait::::new(); + leaf.serialize(ref user_params); + proof.serialize(ref user_params); + + voting_strategy.get_voting_power(timestamp, voter, params, user_params); + } + + #[test] + #[available_gas(1000000000)] + #[should_panic(expected: ('Merkle: Invalid proof', 'ENTRYPOINT_FAILED'))] + fn lying_voting_power() { + let members = generate_n_members(20); + + let (contract, _) = deploy_syscall( + MerkleWhitelistVotingStrategy::TEST_CLASS_HASH.try_into().unwrap(), + 0, + array::ArrayTrait::::new().span(), + false, + ) + .unwrap(); + let voting_strategy = IVotingStrategyDispatcher { contract_address: contract }; + let timestamp = 0x1234; + let index = 2; + let leaf = *members.at(index); + let voter = leaf.address; + + let merkle_data = generate_merkle_data(members.span()); + let root = generate_merkle_root(merkle_data.span()); + let proof = generate_proof(merkle_data.span(), index); + + let mut params = ArrayTrait::::new(); + root.serialize(ref params); + + let mut user_params = ArrayTrait::::new(); + let fake_leaf = Leaf { + address: leaf.address, voting_power: leaf.voting_power + 1, + }; // lying about voting power here + fake_leaf.serialize(ref user_params); + proof.serialize(ref user_params); + + voting_strategy.get_voting_power(timestamp, voter, params, user_params); + } + + #[test] + #[available_gas(1000000000)] + #[should_panic(expected: ('Merkle: Invalid proof', 'ENTRYPOINT_FAILED'))] + fn lying_address_power() { + let members = generate_n_members(20); + + let (contract, _) = deploy_syscall( + MerkleWhitelistVotingStrategy::TEST_CLASS_HASH.try_into().unwrap(), + 0, + array::ArrayTrait::::new().span(), + false, + ) + .unwrap(); + let voting_strategy = IVotingStrategyDispatcher { contract_address: contract }; + let timestamp = 0x1234; + let index = 2; + let leaf = *members.at(index); + let voter = leaf.address; + + let merkle_data = generate_merkle_data(members.span()); + let root = generate_merkle_root(merkle_data.span()); + let proof = generate_proof(merkle_data.span(), index); + + let mut params = ArrayTrait::::new(); + root.serialize(ref params); + + let mut user_params = ArrayTrait::::new(); + let fake_leaf = Leaf { + address: UserAddress::Starknet(contract_address_const::<0x1337>()), + voting_power: leaf.voting_power, + }; // lying about address here + fake_leaf.serialize(ref user_params); + proof.serialize(ref user_params); + + voting_strategy.get_voting_power(timestamp, voter, params, user_params); + } +} diff --git a/starknet/src/types/choice.cairo b/starknet/src/types/choice.cairo index 0da101b1..8e3450d9 100644 --- a/starknet/src/types/choice.cairo +++ b/starknet/src/types/choice.cairo @@ -2,7 +2,6 @@ use serde::Serde; use traits::Into; use hash::LegacyHash; - #[derive(Copy, Drop, Serde)] enum Choice { Against: (), @@ -25,9 +24,3 @@ impl ChoiceIntoU256 of Into { ChoiceIntoU8::into(self).into() } } - -impl LegacyHashChoice of LegacyHash { - fn hash(state: felt252, value: Choice) -> felt252 { - LegacyHash::hash(state, ChoiceIntoU8::into(value)) - } -} diff --git a/starknet/src/types/strategy.cairo b/starknet/src/types/strategy.cairo index 0e8e9474..79fab588 100644 --- a/starknet/src/types/strategy.cairo +++ b/starknet/src/types/strategy.cairo @@ -5,7 +5,6 @@ use option::OptionTrait; use clone::Clone; use result::ResultTrait; use traits::TryInto; -use hash::LegacyHash; use starknet::{StorageBaseAddress, Store, SyscallResult}; #[derive(Clone, Drop, Option, Serde, starknet::Store)] diff --git a/starknet/src/types/user_address.cairo b/starknet/src/types/user_address.cairo index a1778e84..1066eb96 100644 --- a/starknet/src/types/user_address.cairo +++ b/starknet/src/types/user_address.cairo @@ -1,8 +1,8 @@ use starknet::{ContractAddress, EthAddress}; use traits::{PartialEq, TryInto, Into}; -use hash::LegacyHash; use serde::Serde; use array::ArrayTrait; +use sx::utils::legacy_hash::LegacyHashUserAddress; #[derive(Copy, Drop, Serde, PartialEq, starknet::Store)] enum UserAddress { @@ -14,16 +14,6 @@ enum UserAddress { Custom: u256 } -impl LegacyHashUserAddress of LegacyHash { - fn hash(state: felt252, value: UserAddress) -> felt252 { - match value { - UserAddress::Starknet(address) => LegacyHash::::hash(state, address.into()), - UserAddress::Ethereum(address) => LegacyHash::::hash(state, address.into()), - UserAddress::Custom(address) => LegacyHash::::hash(state, address), - } - } -} - trait UserAddressTrait { fn to_starknet_address(self: UserAddress) -> ContractAddress; fn to_ethereum_address(self: UserAddress) -> EthAddress; diff --git a/starknet/src/utils.cairo b/starknet/src/utils.cairo index b2717f7e..c8dc1bd0 100644 --- a/starknet/src/utils.cairo +++ b/starknet/src/utils.cairo @@ -5,10 +5,10 @@ mod constants; mod felt_arr_to_uint_arr; use felt_arr_to_uint_arr::Felt252ArrayIntoU256Array; -mod legacy_hash_eth_address; -use legacy_hash_eth_address::LegacyHashEthAddress; +mod legacy_hash; mod math; +mod merkle; mod single_slot_proof; diff --git a/starknet/src/utils/legacy_hash.cairo b/starknet/src/utils/legacy_hash.cairo new file mode 100644 index 00000000..6d4eacf6 --- /dev/null +++ b/starknet/src/utils/legacy_hash.cairo @@ -0,0 +1,47 @@ +use hash::LegacyHash; +use traits::Into; +use starknet::EthAddress; +use sx::types::{Choice, UserAddress}; +use array::{ArrayTrait, SpanTrait}; + +impl LegacyHashChoice of LegacyHash { + fn hash(state: felt252, value: Choice) -> felt252 { + let choice: u8 = value.into(); + LegacyHash::hash(state, choice) + } +} + +impl LegacyHashEthAddress of LegacyHash { + fn hash(state: felt252, value: EthAddress) -> felt252 { + LegacyHash::::hash(state, value.into()) + } +} + +impl LegacyHashSpan of LegacyHash> { + fn hash(mut state: felt252, mut value: Span) -> felt252 { + let len = value.len(); + loop { + match value.pop_front() { + Option::Some(current) => { + state = LegacyHash::hash(state, *current); + }, + Option::None => { + break; + }, + }; + }; + LegacyHash::hash( + state, len + ) // append the length to conform to computeHashOnElements in starknet.js + } +} + +impl LegacyHashUserAddress of LegacyHash { + fn hash(state: felt252, value: UserAddress) -> felt252 { + match value { + UserAddress::Starknet(address) => LegacyHash::::hash(state, address.into()), + UserAddress::Ethereum(address) => LegacyHash::::hash(state, address.into()), + UserAddress::Custom(address) => LegacyHash::::hash(state, address), + } + } +} diff --git a/starknet/src/utils/legacy_hash_eth_address.cairo b/starknet/src/utils/legacy_hash_eth_address.cairo deleted file mode 100644 index 43bc0b05..00000000 --- a/starknet/src/utils/legacy_hash_eth_address.cairo +++ /dev/null @@ -1,9 +0,0 @@ -use hash::LegacyHash; -use traits::Into; -use starknet::EthAddress; - -impl LegacyHashEthAddress of LegacyHash { - fn hash(state: felt252, value: EthAddress) -> felt252 { - LegacyHash::::hash(state, value.into()) - } -} diff --git a/starknet/src/utils/merkle.cairo b/starknet/src/utils/merkle.cairo new file mode 100644 index 00000000..3132df8f --- /dev/null +++ b/starknet/src/utils/merkle.cairo @@ -0,0 +1,56 @@ +use core::traits::Into; +use array::{ArrayTrait, Span, SpanTrait}; +use option::OptionTrait; +use serde::Serde; +use sx::types::UserAddress; +use clone::Clone; +use hash::{LegacyHash}; +use debug::PrintTrait; +use sx::utils::legacy_hash::LegacyHashSpan; + +/// Leaf struct for the merkle tree +#[derive(Copy, Clone, Drop, Serde)] +struct Leaf { + address: UserAddress, + voting_power: u256, +} + +trait Hash { + fn hash(self: @T) -> felt252; +} + +impl HashSerde> of Hash { + fn hash(self: @T) -> felt252 { + let mut serialized = ArrayTrait::new(); + Serde::::serialize(self, ref serialized); + let hashed = LegacyHash::hash(0, serialized.span()); + hashed + } +} + +/// Asserts that the given proof is valid for the given leaf and root. +fn assert_valid_proof(root: felt252, leaf: Leaf, proof: Span) { + let leaf_node = leaf.hash(); + let computed_root = _compute_merkle_root(leaf_node, proof); + assert(computed_root == root, 'Merkle: Invalid proof'); +} + +/// Internal helper function that computes the merkle root, given a leaf node and a proof. +fn _compute_merkle_root(mut current: felt252, proof: Span) -> felt252 { + let mut proof = proof; + loop { + match proof.pop_front() { + Option::Some(val) => { + let p_u256: u256 = (*val).into(); // Needed for type annotation + if current.into() >= p_u256 { + current = LegacyHash::hash(current, *val); + } else { + current = LegacyHash::hash(*val, current); + }; + }, + Option::None => { + break current; + }, + }; + } +} diff --git a/starknet/src/voting_strategies.cairo b/starknet/src/voting_strategies.cairo index e5266c7b..1093dffb 100644 --- a/starknet/src/voting_strategies.cairo +++ b/starknet/src/voting_strategies.cairo @@ -1,3 +1,5 @@ mod vanilla; mod eth_balance_of; + +mod merkle_whitelist; diff --git a/starknet/src/voting_strategies/merkle_whitelist.cairo b/starknet/src/voting_strategies/merkle_whitelist.cairo new file mode 100644 index 00000000..fb4a6f4c --- /dev/null +++ b/starknet/src/voting_strategies/merkle_whitelist.cairo @@ -0,0 +1,39 @@ +#[starknet::contract] +mod MerkleWhitelistVotingStrategy { + use sx::interfaces::IVotingStrategy; + use serde::Serde; + use sx::types::UserAddress; + use array::{ArrayTrait, Span, SpanTrait}; + use option::OptionTrait; + use sx::utils::merkle::{assert_valid_proof, Leaf}; + use debug::PrintTrait; + + const LEAF_SIZE: usize = 4; // Serde::::serialize().len() + + #[storage] + struct Storage {} + + #[external(v0)] + impl MerkleWhitelistImpl of IVotingStrategy { + fn get_voting_power( + self: @ContractState, + block_number: u32, + voter: UserAddress, + params: Array, // [root] + user_params: Array, // [Serde(leaf), Serde(proofs)] + ) -> u256 { + let cache = user_params.span(); // cache + + let mut leaf_raw = cache.slice(0, LEAF_SIZE); + let leaf = Serde::::deserialize(ref leaf_raw).unwrap(); + + let mut proofs_raw = cache.slice(LEAF_SIZE, cache.len() - LEAF_SIZE); + let proofs = Serde::>::deserialize(ref proofs_raw).unwrap(); + + let root = *params.at(0); // no need to deserialize because it's a simple value + + assert_valid_proof(root, leaf, proofs.span()); + leaf.voting_power + } + } +}