From 9d45caa7b49d4c1ce4e0996605bd6d02f82769a2 Mon Sep 17 00:00:00 2001
From: shotaronowhere <shawtarohgn@gmail.com>
Date: Thu, 12 May 2022 09:44:41 +0100
Subject: [PATCH 1/2] fix: sha256 leaves to avoid 2nd pre-image attack

---
 contracts/src/bridge/merkle/MerkleProof.sol       | 2 +-
 contracts/src/bridge/merkle/MerkleTreeHistory.sol | 7 +++++--
 contracts/test/bridge/merkle/MerkleTree.ts        | 2 +-
 contracts/test/bridge/merkle/index.ts             | 2 +-
 4 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/contracts/src/bridge/merkle/MerkleProof.sol b/contracts/src/bridge/merkle/MerkleProof.sol
index cdfa3835e..c01647408 100644
--- a/contracts/src/bridge/merkle/MerkleProof.sol
+++ b/contracts/src/bridge/merkle/MerkleProof.sol
@@ -39,7 +39,7 @@ contract MerkleProof {
         bytes memory data,
         bytes32 merkleRoot
     ) public pure returns (bool) {
-        return validateProof(proof, keccak256(data), merkleRoot);
+        return validateProof(proof, sha256(data), merkleRoot);
     }
 
     /** @dev Calculates merkle root from proof.
diff --git a/contracts/src/bridge/merkle/MerkleTreeHistory.sol b/contracts/src/bridge/merkle/MerkleTreeHistory.sol
index 6d7059e24..1d3f6eda9 100644
--- a/contracts/src/bridge/merkle/MerkleTreeHistory.sol
+++ b/contracts/src/bridge/merkle/MerkleTreeHistory.sol
@@ -40,7 +40,10 @@ contract MerkleTreeHistory {
      *  @param data The data to insert in the merkle tree.
      */
     function append(bytes memory data) public {
-        bytes32 leaf = keccak256(data);
+        // Differentiate leaves from interior nodes with different
+        // hash functions to prevent 2nd order pre-image attack.
+        // https://flawed.net.nz/2018/02/21/attacking-merkle-trees-with-a-second-preimage-attack/
+        bytes32 leaf = sha256(data);
         count += 1;
         uint256 size = count;
         uint256 hashBitField = (size ^ (size - 1)) & size;
@@ -98,8 +101,8 @@ contract MerkleTreeHistory {
         uint256 height = 0;
         bool isFirstHash = true;
         while (size > 0) {
-            // avoid redundant calculation
             if ((size & 1) == 1) {
+                // avoid redundant calculation
                 if (isFirstHash) {
                     node = branch[height];
                     isFirstHash = false;
diff --git a/contracts/test/bridge/merkle/MerkleTree.ts b/contracts/test/bridge/merkle/MerkleTree.ts
index 1bc941eca..2bb1756ed 100644
--- a/contracts/test/bridge/merkle/MerkleTree.ts
+++ b/contracts/test/bridge/merkle/MerkleTree.ts
@@ -31,7 +31,7 @@ export class MerkleTree {
    * @return node The `sha3` (A.K.A. `keccak256`) hash of `first, ...params` as a 32-byte hex string.
    */
   public static makeLeafNode(data: string): string {
-    const result = ethers.utils.keccak256(data);
+    const result = ethers.utils.sha256(data);
 
     if (!result) {
       throw new Error("Leaf node must not be empty");
diff --git a/contracts/test/bridge/merkle/index.ts b/contracts/test/bridge/merkle/index.ts
index c6af5edd6..e766c4c31 100644
--- a/contracts/test/bridge/merkle/index.ts
+++ b/contracts/test/bridge/merkle/index.ts
@@ -61,7 +61,7 @@ describe("Merkle", function () {
     });
   it("Should correctly verify all nodes in the tree", async () => {
     for (var message of data)  {
-      const leaf = ethers.utils.keccak256(message);
+      const leaf = ethers.utils.sha256(message);
       proof = mt.getHexProof(leaf);
       const validation = await merkleProof.validateProof(proof, message,rootOnChain);
       expect(validation).equal(true);

From 04e4f30b409db98d2339bc82c05068b54a9e03e4 Mon Sep 17 00:00:00 2001
From: shotaronowhere <shawtarohgn@gmail.com>
Date: Fri, 13 May 2022 00:35:40 +0100
Subject: [PATCH 2/2] feat: gas savings, compile w optimization

---
 .../src/bridge/merkle/MerkleTreeHistory.sol     | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/contracts/src/bridge/merkle/MerkleTreeHistory.sol b/contracts/src/bridge/merkle/MerkleTreeHistory.sol
index 1d3f6eda9..01f6260b0 100644
--- a/contracts/src/bridge/merkle/MerkleTreeHistory.sol
+++ b/contracts/src/bridge/merkle/MerkleTreeHistory.sol
@@ -44,31 +44,30 @@ contract MerkleTreeHistory {
         // hash functions to prevent 2nd order pre-image attack.
         // https://flawed.net.nz/2018/02/21/attacking-merkle-trees-with-a-second-preimage-attack/
         bytes32 leaf = sha256(data);
-        count += 1;
-        uint256 size = count;
+        uint256 size = count + 1;
+        count = size;
         uint256 hashBitField = (size ^ (size - 1)) & size;
-
-        for (uint256 height = 0; height < 64; height++) {
-            if ((hashBitField & 1) == 1) {
-                branch[height] = leaf;
-                return;
-            }
+        uint256 height;
+        while ((hashBitField & 1) == 0) {
             bytes32 node = branch[height];
-            // effecient hash
             if (node > leaf)
                 assembly {
+                    // effecient hash
                     mstore(0x00, leaf)
                     mstore(0x20, node)
                     leaf := keccak256(0x00, 0x40)
                 }
             else
                 assembly {
+                    // effecient hash
                     mstore(0x00, node)
                     mstore(0x20, leaf)
                     leaf := keccak256(0x00, 0x40)
                 }
             hashBitField /= 2;
+            height = height + 1;
         }
+        branch[height] = leaf;
     }
 
     /** @dev Saves the merkle root state in history and resets.