Generic Solidity Merkle-tree library, includes Merkle-root computatio…

…n function

* MerkleTrees library contains a generic Merkle-tree implementation, that may be configured with aribitrary leaf/node-pair hash functions
* LegacyMerkleTrees library contains a Merkle tree implementation that uses plain keccak256 for both leaves and nodes.
* MerkleProof library behaviour is unchanged, except that it now accepts the leaf data directly instead of the leaf hash. Otherwise it keeps the same interface and tests, but internally it forwards calls to LegacyMerkleTrees library
* SecondPreimageResistantMerkleTrees contains a Merkle-tree implementation that uses different hash functions for leaves and node-pairs

contracts/cryptography/LegacyMerkleTrees.sol

@@ -0,0 +1,54 @@
+pragma solidity ^0.4.24;
+
+import "./MerkleTrees.sol";
+
+
+/**
+ * @title LegacyMerkleTree library
+ * @dev Computation of Merkle root and verification of Merkle proof based on
+ * https://github.com/ameensol/merkle-tree-solidity/blob/master/src/MerkleProof.sol
+ */
+library LegacyMerkleTrees {
+
+  using MerkleTrees for MerkleTrees.TreeConfig;
+
+  function _keccak256Leaf(bytes leafDataBlock)
+    internal
+    pure
+    returns (bytes32)
+  {
+    return keccak256(leafDataBlock);
+  }
+
+  function _keccak256SortNodePair(bytes32 h1, bytes32 h2)
+    internal
+    pure
+    returns (bytes32)
+  {
+    return keccak256(
+      h1 < h2 ? abi.encodePacked(h1, h2) : abi.encodePacked(h2, h1)
+    );
+  }
+
+  function _config() internal pure returns (MerkleTrees.TreeConfig config) {
+    config._hashLeafData = _keccak256Leaf;
+    config._hashNodePair = _keccak256SortNodePair;
+  }
+
+  function newTree(uint256 size)
+    internal
+    pure
+    returns (MerkleTrees.Tree tree)
+  {
+    return _config().newTree(size);
+  }
+
+  function verifyProof(bytes32 root, bytes leafDataBlock, bytes32[] proof)
+    internal
+    pure
+    returns (bool)
+  {
+    return _config().verifyProof(root, leafDataBlock, proof);
+  }
+
+}

contracts/cryptography/MerkleProof.sol

@@ -1,5 +1,8 @@
 pragma solidity ^0.4.24;
 
+import "./LegacyMerkleTrees.sol";
+
+
 /**
  * @title MerkleProof
  * @dev Merkle proof verification based on
@@ -11,32 +14,17 @@ library MerkleProof {
    * and each pair of pre-images are sorted.
    * @param proof Merkle proof containing sibling hashes on the branch from the leaf to the root of the Merkle tree
    * @param root Merkle root
-   * @param leaf Leaf of Merkle tree
+   * @param leafDataBlock Leaf data block of Merkle tree, before hashing
    */
   function verify(
     bytes32[] proof,
     bytes32 root,
-    bytes32 leaf
+    bytes leafDataBlock
   )
     internal
     pure
     returns (bool)
   {
-    bytes32 computedHash = leaf;
-
-    for (uint256 i = 0; i < proof.length; i++) {
-      bytes32 proofElement = proof[i];
-
-      if (computedHash < proofElement) {
-        // Hash(current computed hash + current element of the proof)
-        computedHash = keccak256(abi.encodePacked(computedHash, proofElement));
-      } else {
-        // Hash(current element of the proof + current computed hash)
-        computedHash = keccak256(abi.encodePacked(proofElement, computedHash));
-      }
-    }
-
-    // Check if the computed hash (root) is equal to the provided root
-    return computedHash == root;
+    return LegacyMerkleTrees.verifyProof(root, leafDataBlock, proof);
   }
 }

contracts/cryptography/MerkleTrees.sol

@@ -0,0 +1,151 @@
+pragma solidity ^0.4.24;
+
+/**
+ * @title Generic Merkle-tree library
+ * @author Edward Grech (@dwardu)
+ * @notice Functions to compute a Merkle-root and verify a Merkle-proof.
+ * @dev The Merkle-tree implementation in this library may be configured
+ * with arbitrary leaf/node-pair hash functions.
+ */
+library MerkleTrees {
+
+  using MerkleTrees for TreeConfig;
+
+  /**
+   * @dev Holds references to the leaf/node hash functions
+   * to be used to construct a Merkle tree.
+   * By choosing the two functions to be different, it is possible
+   * to guard against second pre-image attacks.
+   * See https://flawed.net.nz/2018/02/21/attacking-merkle-trees-with-a-second-preimage-attack/
+   * Give a TreeConfig, it is then possible to build a Merkle tree
+   * to compute its root, or to verify a Merkle proof for a leaf in
+   * a tree with a specific root.
+   */
+  struct TreeConfig {
+    function(bytes memory) internal pure returns (bytes32) _hashLeafData;
+    function(bytes32, bytes32) internal pure returns (bytes32) _hashNodePair;
+  }
+
+  using MerkleTrees for Tree;
+
+  struct Tree {
+    /**
+     * @dev Holds references to the leaf/node hash functions
+     * to be used for this tree.
+     */
+    TreeConfig _config;
+
+    /**
+     * @dev Used as a guard against computing the root more than once, as
+     * the first computation destroys the original leaf nodes.
+     */
+    bytes32[] _nodes;
+
+    /**
+     * @dev Used as a guard against computing the root more than once, as
+     * the first computation overwrites the original leaf nodes.
+     */
+    bool _wasRootComputed;
+  }
+
+  /**
+   * @notice The first step towards computing a Merkle root is to allocate
+   * a new tree by calling this function.
+   */
+  function newTree(TreeConfig self, uint256 size)
+    internal
+    pure
+    returns (Tree tree)
+  {
+    assert(1 <= size);
+    tree._config = self;
+    tree._nodes = new bytes32[](size);
+    tree._wasRootComputed = false;
+  }
+
+  /**
+   * @notice After allocating a tree, this function should be called
+   * with each leaf data block.
+   */
+  function setLeafDataBlock(Tree self, uint256 index, bytes leafDataBlock)
+    internal
+    pure
+  {
+    self._nodes[index] = self._config._hashLeafData(leafDataBlock);
+  }
+
+  /**
+   * @notice Once the tree has been allocated and the leaves set,
+   * this function is called to compute the Merkle root.
+   * @return The Merkle root.
+   */
+  function computeRoot(Tree self)
+    internal
+    pure
+    returns (bytes32 root)
+  {
+    assert(!self._wasRootComputed);
+
+    uint256 nCurr = self._nodes.length;
+
+    while (1 < nCurr) {
+
+      // We pair and hash sibling elements in the current layer starting from
+      // the left to the right, and store the hashes in the next layer.
+      // If nCurr is odd, then the right-most element in current layer will
+      // remain unpaired - we do not account for it in `nNext` right now, as
+      // `nCurr / 2` rounds down, but we will account for it later.
+      uint256 nNext = nCurr / 2;
+
+      // Loop over all paired sibling elements
+      for (uint256 iNext = 0; iNext < nNext; iNext++) {
+        uint256 iCurr = iNext * 2;
+        self._nodes[iNext] = self._config._hashNodePair(
+          self._nodes[iCurr],
+          self._nodes[iCurr + 1]
+        );
+      }
+
+      // If the right-most element remained unpaired, promote it to the
+      // end of the next layer, and increment nNext to account for it.
+      if (nCurr % 2 == 1) {
+        self._nodes[++nNext - 1] = self._nodes[nCurr - 1];
+      }
+
+      nCurr = nNext;
+    }
+
+    self._wasRootComputed = true;
+
+    return self._nodes[0];
+  }
+
+  /**
+   * @notice Verifies a Merkle proof proving the existence of a leaf data
+   * block in a Merkle tree.
+   * @param self The Merkle tree configuration instance.
+   * @param root The root of the Merkle tree to verify the proof against.
+   * @param leafDataBlock The leaf data block (unhashed) whose exitence to verify.
+   * @param proof Merkle proof containing sibling hashes on the branch from
+   * the leaf to the root of the Merkle tree.
+   */
+  function verifyProof(
+    TreeConfig self,
+    bytes32 root,
+    bytes leafDataBlock,
+    bytes32[] proof
+  )
+    internal
+    pure
+    returns (bool)
+  {
+    bytes32 computedHash = self._hashLeafData(leafDataBlock);
+
+    for (uint256 i = 0; i < proof.length; i++) {
+      computedHash = self._hashNodePair(computedHash, proof[i]);
+    }
+
+    return computedHash == root;
+  }
+
+}

contracts/cryptography/SecondPreimageResistantMerkleTrees.sol

@@ -0,0 +1,63 @@
+pragma solidity ^0.4.24;
+
+import "./MerkleTrees.sol";
+
+/**
+ * @title SecondPreimageResistantMerkleTrees library
+ * @dev Computation of Merkle root and verification of Merkle proof based on
+ * a Merkle tree implementation that uses different hash functions for leaf
+ * data blocks and node hash pairs, in order to give
+ * second preimage resistance.
+ * The implementation is similar to that described in the
+ * Certificate Transparency RFC
+ * (see https://tools.ietf.org/html/rfc6962#section-2.1), in that it
+ * prepends 0x00 and 0x01 to leaves and nodes respectively, but differs
+ * in that it uses keccak256 instead of sha256, and sorts
+ * node hash pairs before hashing.
+ */
+library SecondPreimageResistantMerkleTrees {
+
+  using MerkleTrees for MerkleTrees.TreeConfig;
+
+  function _keccak256Prepend00Leaf(bytes leafDataBlock)
+    internal
+    pure
+    returns (bytes32)
+  {
+    return keccak256(abi.encodePacked(bytes1(0x00), leafDataBlock));
+  }
+
+  function _keccak256Prepend01SortNodePair(bytes32 h1, bytes32 h2)
+    internal
+    pure
+    returns (bytes32)
+  {
+    return keccak256(
+      h1 < h2
+        ? abi.encodePacked(bytes1(0x01), h1, h2)
+        : abi.encodePacked(bytes1(0x01), h2, h1)
+    );
+  }
+
+  function _config() internal pure returns (MerkleTrees.TreeConfig config) {
+    config._hashLeafData = _keccak256Prepend00Leaf;
+    config._hashNodePair = _keccak256Prepend01SortNodePair;
+  }
+
+  function newTree(uint256 size)
+    internal
+    pure
+    returns (MerkleTrees.Tree tree)
+  {
+    return _config().newTree(size);
+  }
+
+  function verifyProof(bytes32 root, bytes leafDataBlock, bytes32[] proof)
+    internal
+    pure
+    returns (bool)
+  {
+    return _config().verifyProof(root, leafDataBlock, proof);
+  }
+
+}

contracts/mocks/MerkleProofWrapper.sol

@@ -7,12 +7,12 @@ contract MerkleProofWrapper {
   function verify(
     bytes32[] proof,
     bytes32 root,
-    bytes32 leaf
+    bytes leafDataBlock
   )
     public
     pure
     returns (bool)
   {
-    return MerkleProof.verify(proof, root, leaf);
+    return MerkleProof.verify(proof, root, leafDataBlock);
   }
 }

test/cryptography/MerkleProof.test.js

@@ -1,10 +1,9 @@
 const { MerkleTree } = require('../helpers/merkleTree.js');
-const { sha3, bufferToHex } = require('ethereumjs-util');
+const { toBuffer, bufferToHex } = require('ethereumjs-util');
 
 const MerkleProofWrapper = artifacts.require('MerkleProofWrapper');
 
-require('chai')
-  .should();
+require('chai').should();
 
 contract('MerkleProof', function () {
   beforeEach(async function () {
@@ -20,7 +19,7 @@ contract('MerkleProof', function () {
 
       const proof = merkleTree.getHexProof(elements[0]);
 
-      const leaf = bufferToHex(sha3(elements[0]));
+      const leaf = bufferToHex(toBuffer(elements[0]));
 
       (await this.merkleProof.verify(proof, root, leaf)).should.equal(true);
     });
@@ -31,7 +30,7 @@ contract('MerkleProof', function () {
 
       const correctRoot = correctMerkleTree.getHexRoot();
 
-      const correctLeaf = bufferToHex(sha3(correctElements[0]));
+      const correctLeaf = bufferToHex(toBuffer(correctElements[0]));
 
       const badElements = ['d', 'e', 'f'];
       const badMerkleTree = new MerkleTree(badElements);
@@ -50,7 +49,7 @@ contract('MerkleProof', function () {
       const proof = merkleTree.getHexProof(elements[0]);
       const badProof = proof.slice(0, proof.length - 5);
 
-      const leaf = bufferToHex(sha3(elements[0]));
+      const leaf = bufferToHex(toBuffer(elements[0]));
 
       (await this.merkleProof.verify(badProof, root, leaf)).should.equal(false);
     });