Add merkle whitelist voting strategy

contracts/starknet/TestContracts/Test_Merkle.cairo

@@ -0,0 +1,13 @@
+%lang starknet
+from starkware.cairo.common.alloc import alloc
+from starkware.cairo.common.cairo_builtins import HashBuiltin
+from contracts.starknet.lib.merkle import Merkle
+
+@view
+func testAssertValidLeaf{range_check_ptr, pedersen_ptr: HashBuiltin*}(
+    root: felt, leaf_len: felt, leaf: felt*, proof_len: felt, proof: felt*
+) {
+    alloc_locals;
+    Merkle.assert_valid_leaf(root, leaf_len, leaf, proof_len, proof);
+    return ();
+}

contracts/starknet/VotingStrategies/MerkleWhitelist.cairo

@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: MIT
+
+%lang starknet
+
+from starkware.cairo.common.uint256 import Uint256
+from starkware.cairo.common.cairo_builtins import HashBuiltin
+from starkware.cairo.common.memcpy import memcpy
+from starkware.cairo.common.alloc import alloc
+from starkware.cairo.common.math import assert_nn_le
+
+from contracts.starknet.lib.general_address import Address
+from contracts.starknet.lib.merkle import Merkle
+
+//
+// @title Merkle Whitelist Voting Strategy
+// @author SnapshotLabs
+// @notice Contract to allow a merkle tree based whitelist to be used to compute voting power for each user
+//
+
+// @dev Returns the voting power for a user obtained from the whitelist
+// @param timestamp The snapshot timestamp (not used)
+// @param voter_address The address of the user
+// @param params Configuration parameter array that is the same for every voter in the proposal. Should be as follows:
+//      params[0] = The merkle root of the whitelist data, should be computed off-chain
+// @param user_params Array containing the leaf and merkle proof data. Should be as follows:
+//      user_params[0] = address of the whitelisted user
+//      user_params[1] = Low 128 bits of the voting power of the user
+//      user_params[2] = High 128 bits of the voting power of the user
+// @return voting_power The voting power of the user as a Uint256
+@view
+func getVotingPower{syscall_ptr: felt*, pedersen_ptr: HashBuiltin*, range_check_ptr: felt}(
+    timestamp: felt,
+    voter_address: Address,
+    params_len: felt,
+    params: felt*,
+    user_params_len: felt,
+    user_params: felt*,
+) -> (voting_power: Uint256) {
+    alloc_locals;
+
+    with_attr error_message("MerkleWhitelist: Invalid parameters supplied") {
+        assert_nn_le(3, user_params_len);
+    }
+
+    // Extracting leaf data from user params array
+    let (leaf: felt*) = alloc();
+    let leaf_len = 3;
+    memcpy(leaf, user_params, leaf_len);
+
+    // Checking that the leaf corresponds to the voter's address
+    with_attr error_message("MerkleWhitelist: Invalid proof supplied") {
+        // The address resides at the beginning of the leaf data array
+        assert leaf[0] = voter_address.value;
+    }
+
+    // Extracting proof from user params array
+    let (proof: felt*) = alloc();
+    let proof_len = user_params_len - leaf_len;
+    memcpy(proof, user_params + leaf_len, proof_len);
+
+    // Extracting merkle root from params array
+    let merkle_root = params[0];
+
+    // Checking the merkle proof
+    Merkle.assert_valid_leaf(merkle_root, leaf_len, leaf, proof_len, proof);
+
+    // Extract voting power from leaf and cast to Uint256
+    let voting_power = Uint256(leaf[1], leaf[2]);
+
+    return (voting_power,);
+}

contracts/starknet/lib/merkle.cairo

@@ -0,0 +1,54 @@
+// SPDX-License-Identifier: MIT
+
+%lang starknet
+
+from starkware.cairo.common.cairo_builtins import HashBuiltin
+from starkware.cairo.common.math_cmp import is_le_felt
+from starkware.cairo.common.hash import hash2
+
+from contracts.starknet.lib.array_utils import ArrayUtils
+
+//
+// @title Merkle Proof Library
+// @author SnapshotLabs
+// @notice A library to to verify merkle proofs
+//
+
+namespace Merkle {
+    // @dev Asserts a given leaf is a member of the set with the specified root by verifing a proof
+    // @param root The merkle root of the data
+    // @param leaf The leaf data array
+    // @param proof The proof
+    func assert_valid_leaf{pedersen_ptr: HashBuiltin*, range_check_ptr}(
+        root: felt, leaf_len: felt, leaf: felt*, proof_len: felt, proof: felt*
+    ) {
+        let (leaf_node) = ArrayUtils.hash(leaf_len, leaf);
+        let (computed_root) = _compute_merkle_root(leaf_node, proof_len, proof);
+        with_attr error_message("Merkle: Invalid proof") {
+            assert root = computed_root;
+        }
+        return ();
+    }
+}
+
+func _compute_merkle_root{pedersen_ptr: HashBuiltin*, range_check_ptr}(
+    curr: felt, proof_len: felt, proof: felt*
+) -> (root: felt) {
+    alloc_locals;
+
+    if (proof_len == 0) {
+        return (curr,);
+    }
+
+    let le = is_le_felt(curr, proof[0]);
+    if (le == 1) {
+        let (n) = hash2{hash_ptr=pedersen_ptr}(curr, proof[0]);
+        tempvar node = n;
+    } else {
+        let (n) = hash2{hash_ptr=pedersen_ptr}(proof[0], curr);
+        tempvar node = n;
+    }
+
+    let (root) = _compute_merkle_root(node, proof_len - 1, &proof[1]);
+    return (root,);
+}

package.json

@@ -31,10 +31,12 @@
     "@gnosis.pm/safe-contracts": "^1.3.0",
     "@gnosis.pm/zodiac": "^1.1.5",
     "@openzeppelin/contracts": "^4.7.3",
-    "@snapshot-labs/sx": "0.1.0-beta.12",
     "@shardlabs/starknet-hardhat-plugin": "0.6.7",
-    "starknet": "^4.4.2",
+    "@snapshot-labs/sx": "0.1.0-beta.12",
     "concurrently": "^7.4.0",
+    "keccak256": "^1.0.6",
+    "merkletreejs": "^0.3.1",
+    "starknet": "^4.4.2",
     "wait-on": "^6.0.1"
   },
   "devDependencies": {

test/shared/merkle.ts

@@ -0,0 +1,60 @@
+import { pedersen } from 'starknet/dist/utils/hash';
+
+export class MerkleTree {
+  root: string;
+  constructor(values: string[]) {
+    this.root = MerkleTree.generateMerkleRoot(values);
+  }
+
+  getProof(values: string[], index: number): string[] {
+    return MerkleTree.getProofHelper(values, index, []);
+  }
+
+  static generateMerkleRoot(values: string[]): string {
+    if (values.length == 1) {
+      return values[0];
+    }
+    if (values.length % 2 != 0) {
+      values.push('0x0');
+    }
+    const nextLevel = MerkleTree.getNextLevel(values);
+    return MerkleTree.generateMerkleRoot(nextLevel);
+  }
+
+  static getNextLevel(level: string[]): string[] {
+    const nextLevel = [];
+    for (let i = 0; i < level.length; i += 2) {
+      let node = '0x0';
+      if (BigInt(level[i]) < BigInt(level[i + 1])) {
+        node = pedersen([level[i], level[i + 1]]);
+      } else {
+        node = pedersen([level[i + 1], level[i]]);
+      }
+      nextLevel.push(node);
+    }
+    return nextLevel;
+  }
+
+  static getProofHelper(level: string[], index: number, proof: string[]): string[] {
+    if (level.length == 1) {
+      return proof;
+    }
+    if (level.length % 2 != 0) {
+      level.push('0x0');
+    }
+    const nextLevel = MerkleTree.getNextLevel(level);
+    let indexParent = 0;
+
+    for (let i = 0; i < level.length; i++) {
+      if (i == index) {
+        indexParent = Math.floor(i / 2);
+        if (i % 2 == 0) {
+          proof.push(level[index + 1]);
+        } else {
+          proof.push(level[index - 1]);
+        }
+      }
+    }
+    return MerkleTree.getProofHelper(nextLevel, indexParent, proof);
+  }
+}

test/starknet/MerkleProof.test.ts

@@ -0,0 +1,165 @@
+import { StarknetContract } from 'hardhat/types/runtime';
+import { expect } from 'chai';
+import { ethers, starknet } from 'hardhat';
+import { computeHashOnElements } from 'starknet/dist/utils/hash';
+import { MerkleTree } from '../shared/merkle';
+
+describe('Merkle:', () => {
+  let testMerkle: StarknetContract;
+
+  before(async function () {
+    this.timeout(800000);
+    const testMerkleFactory = await starknet.getContractFactory(
+      './contracts/starknet/TestContracts/Test_Merkle.cairo'
+    );
+    testMerkle = await testMerkleFactory.deploy();
+  });
+
+  it('The library should handle a tree with one leaf', async () => {
+    const values = [];
+    // Generating random data for the merkle tree
+    for (let i = 0; i < 1; i++) {
+      values.push([
+        ethers.Wallet.createRandom().address,
+        ethers.utils.hexlify(ethers.utils.randomBytes(1)),
+      ]);
+    }
+
+    // computing the hash of each address value pair, and sorting
+    const merkleData = values
+      .map((v, i) => [computeHashOnElements(v), v[0], v[1]])
+      .sort(function (a, b) {
+        if (a > b) return 1;
+        if (a < b) return -1;
+        return 0;
+      })
+      .map((x, i) => [x[0], x[1], x[2], i]);
+    const leaves = merkleData.map((x) => x[0].toString());
+    const tree = new MerkleTree(leaves);
+    // Picking random leaf to prove
+    const address = values[Math.floor(Math.random() * 1)][0];
+    const leafData = merkleData.find((leaf) => leaf[1] == address)!;
+
+    await testMerkle.call('testAssertValidLeaf', {
+      root: tree.root,
+      leaf: [leafData[1], leafData[2]],
+      proof: tree.getProof(leaves, Number(leafData[3])),
+    });
+
+    it('The library should handle a tree with two leaves', async () => {
+      const values = [];
+      // Generating random data for the merkle tree
+      for (let i = 0; i < 2; i++) {
+        values.push([
+          ethers.Wallet.createRandom().address,
+          ethers.utils.hexlify(ethers.utils.randomBytes(1)),
+        ]);
+      }
+
+      // computing the hash of each address value pair, and sorting
+      const merkleData = values
+        .map((v, i) => [computeHashOnElements(v), v[0], v[1]])
+        .sort(function (a, b) {
+          if (a > b) return 1;
+          if (a < b) return -1;
+          return 0;
+        })
+        .map((x, i) => [x[0], x[1], x[2], i]);
+      const leaves = merkleData.map((x) => x[0].toString());
+      const tree = new MerkleTree(leaves);
+      // Picking random leaf to prove
+      const address = values[Math.floor(Math.random() * 2)][0];
+      const leafData = merkleData.find((leaf) => leaf[1] == address)!;
+
+      await testMerkle.call('testAssertValidLeaf', {
+        root: tree.root,
+        leaf: [leafData[1], leafData[2]],
+        proof: tree.getProof(leaves, Number(leafData[3])),
+      });
+    });
+
+    it('The library should verify a merkle proof for a leaf in a large tree', async () => {
+      const values = [];
+      // Generating random data for the merkle tree
+      for (let i = 0; i < 1000; i++) {
+        values.push([
+          ethers.Wallet.createRandom().address,
+          ethers.utils.hexlify(ethers.utils.randomBytes(1)),
+        ]);
+      }
+
+      // computing the hash of each address value pair, and sorting
+      const merkleData = values
+        .map((v, i) => [computeHashOnElements(v), v[0], v[1]])
+        .sort(function (a, b) {
+          if (a > b) return 1;
+          if (a < b) return -1;
+          return 0;
+        })
+        .map((x, i) => [x[0], x[1], x[2], i]);
+      const leaves = merkleData.map((x) => x[0].toString());
+      const tree = new MerkleTree(leaves);
+
+      // Picking random leaf to prove
+      const address = values[Math.floor(Math.random() * 99)][0];
+      const leafData = merkleData.find((leaf) => leaf[1] == address)!;
+
+      await testMerkle.call('testAssertValidLeaf', {
+        root: tree.root,
+        leaf: [leafData[1], leafData[2]],
+        proof: tree.getProof(leaves, Number(leafData[3])),
+      });
+    }).timeout(600000);
+
+    it('The library should fail to verify if an invalid proof is supplied', async () => {
+      const values = [];
+      // Generating random data for the merkle tree
+      for (let i = 0; i < 100; i++) {
+        values.push([
+          ethers.Wallet.createRandom().address,
+          ethers.utils.hexlify(ethers.utils.randomBytes(1)),
+        ]);
+      }
+
+      // computing the hash of each address value pair, and sorting
+      const merkleData = values
+        .map((v, i) => [computeHashOnElements(v), v[0], v[1]])
+        .sort(function (a, b) {
+          if (a > b) return 1;
+          if (a < b) return -1;
+          return 0;
+        })
+        .map((x, i) => [x[0], x[1], x[2], i]);
+      const leaves = merkleData.map((x) => x[0].toString());
+      const tree = new MerkleTree(leaves);
+
+      // Picking random leaf to prove
+      const address = values[Math.floor(Math.random() * 99)][0];
+      const leafData = merkleData.find((leaf) => leaf[1] == address)!;
+
+      const corruptedProof = tree.getProof(leaves, Number(leafData[3]));
+      corruptedProof[0] = ethers.utils.hexlify(ethers.utils.randomBytes(4));
+      try {
+        await testMerkle.call('testAssertValidLeaf', {
+          root: tree.root,
+          leaf: [leafData[1], leafData[2]],
+          proof: corruptedProof,
+        });
+      } catch (error: any) {
+        expect(error.message).to.contain('Merkle: Invalid proof');
+      }
+    }).timeout(600000);
+
+    it('The library should handle a tree with no leaves', async () => {
+      try {
+        await testMerkle.call('testAssertValidLeaf', {
+          root: '0x0',
+          leaf: ['0x0', '0x0'],
+          proof: '0x0',
+        });
+      } catch (error: any) {
+        expect(error.message).to.contain('Merkle: Invalid proof');
+      }
+    });
+  }).timeout(600000);
+});