From 24323d3ce305f13580d5baa88ac623f0bc3c6bc0 Mon Sep 17 00:00:00 2001
From: Yondon Fu <yondon.fu@gmail.com>
Date: Mon, 7 Aug 2017 11:23:52 -0400
Subject: [PATCH] MerkleTree util class hashes elements

---
 test/MerkleProof.js        | 16 ++++++++--------
 test/helpers/merkleTree.js | 32 ++++++++++++++------------------
 2 files changed, 22 insertions(+), 26 deletions(-)

diff --git a/test/MerkleProof.js b/test/MerkleProof.js
index eef07b2de..7f449b431 100644
--- a/test/MerkleProof.js
+++ b/test/MerkleProof.js
@@ -1,7 +1,7 @@
 var MerkleProof = artifacts.require("./MerkleProof.sol");
 
-import { sha3 } from "ethereumjs-util";
 import MerkleTree from "./helpers/merkleTree.js";
+import { sha3, bufferToHex } from "ethereumjs-util";
 
 contract('MerkleProof', function(accounts) {
   let merkleProof;
@@ -12,28 +12,28 @@ contract('MerkleProof', function(accounts) {
 
   describe("verifyProof", function() {
     it("should return true for a valid Merkle proof", async function() {
-      const elements = ["a", "b", "c", "d"].map(el => sha3(el));
+      const elements = ["a", "b", "c", "d"];
       const merkleTree = new MerkleTree(elements);
 
       const root = merkleTree.getHexRoot();
 
       const proof = merkleTree.getHexProof(elements[0]);
 
-      const leaf = merkleTree.bufToHex(elements[0]);
+      const leaf = bufferToHex(sha3(elements[0]));
 
       const result = await merkleProof.verifyProof(proof, root, leaf);
       assert.isOk(result, "verifyProof did not return true for a valid proof");
     });
 
     it("should return false for an invalid Merkle proof", async function() {
-      const correctElements = ["a", "b", "c"].map(el => sha3(el));
+      const correctElements = ["a", "b", "c"]
       const correctMerkleTree = new MerkleTree(correctElements);
 
       const correctRoot = correctMerkleTree.getHexRoot();
 
-      const correctLeaf = correctMerkleTree.bufToHex(correctElements[0]);
+      const correctLeaf = bufferToHex(sha3(correctElements[0]));
 
-      const badElements = ["d", "e", "f"].map(el => sha3(el))
+      const badElements = ["d", "e", "f"]
       const badMerkleTree = new MerkleTree(badElements)
 
       const badProof = badMerkleTree.getHexProof(badElements[0])
@@ -43,7 +43,7 @@ contract('MerkleProof', function(accounts) {
     });
 
     it("should return false for a Merkle proof of invalid length", async function() {
-      const elements = ["a", "b", "c"].map(el => sha3(el));
+      const elements = ["a", "b", "c"]
       const merkleTree = new MerkleTree(elements);
 
       const root = merkleTree.getHexRoot();
@@ -51,7 +51,7 @@ contract('MerkleProof', function(accounts) {
       const proof = merkleTree.getHexProof(elements[0]);
       const badProof = proof.slice(0, proof.length - 5);
 
-      const leaf = merkleTree.bufToHex(elements[0]);
+      const leaf = bufferToHex(sha3(elements[0]));
 
       const result = await merkleProof.verifyProof(badProof, root, leaf);
       assert.isNotOk(result, "verifyProof did not return false for proof of invalid length");
diff --git a/test/helpers/merkleTree.js b/test/helpers/merkleTree.js
index dc00eb12f..0de95eb15 100644
--- a/test/helpers/merkleTree.js
+++ b/test/helpers/merkleTree.js
@@ -1,14 +1,9 @@
-import { sha3 } from "ethereumjs-util";
+import { sha3, bufferToHex } from "ethereumjs-util";
 
 export default class MerkleTree {
   constructor(elements) {
-    // Filter empty strings
-    this.elements = elements.filter(el => el);
-
-    // Check if elements are 32 byte buffers
-    if (this.elements.some(el => el.length !== 32 || !Buffer.isBuffer(el))) {
-      throw new Error("Elements must be 32 byte buffers");
-    }
+    // Filter empty strings and hash elements
+    this.elements = elements.filter(el => el).map(el => sha3(el));
 
     // Deduplicate elements
     this.elements = this.bufDedup(this.elements);
@@ -58,7 +53,7 @@ export default class MerkleTree {
   }
 
   getHexRoot() {
-    return this.bufToHex(this.getRoot());
+    return bufferToHex(this.getRoot());
   }
 
   getProof(el) {
@@ -98,8 +93,17 @@ export default class MerkleTree {
   }
 
   bufIndexOf(el, arr) {
+    let hash;
+
+    // Convert element to 32 byte hash if it is not one already
+    if (el.length !== 32 || !Buffer.isBuffer(el)) {
+      hash = sha3(el);
+    } else {
+      hash = el;
+    }
+
     for (let i = 0; i < arr.length; i++) {
-      if (el.equals(arr[i])) {
+      if (hash.equals(arr[i])) {
         return i;
       }
     }
@@ -113,14 +117,6 @@ export default class MerkleTree {
     });
   }
 
-  bufToHex(el) {
-    if (!Buffer.isBuffer(el)) {
-      throw new Error("Element is not a buffer");
-    }
-
-    return "0x" + el.toString("hex");
-  }
-
   bufArrToHex(arr) {
     if (arr.some(el => !Buffer.isBuffer(el))) {
       throw new Error("Array is not an array of buffers");