add hash function parameter

Author: michael-spengler

src/merkle-tree.spec.ts

@@ -1,10 +1,10 @@
 import { assertEquals } from "https://deno.land/std@0.86.0/testing/asserts.ts"
 import { MerkleTree, Helper } from "../mod.ts"
 
-const exampleArray = ["dog", "horse", "cow", "chicken", "rabbit", "bird", "bee", "me"]
-
 Deno.test("should return valid as the investigated entry is in array", async () => {
 
+    const exampleArray = ["dog", "horse", "cow", "chicken", "rabbit", "bird", "bee", "me"]
+
     for (const investigatedEntry of exampleArray) {
         const merkleTree = new MerkleTree(exampleArray)
 
@@ -20,6 +20,8 @@ Deno.test("should return valid as the investigated entry is in array", async ()
 
 Deno.test("should return invalid as the investigated entry is not in array", async () => {
 
+    const exampleArray = ["dog", "horse", "cow", "chicken", "rabbit", "bird", "bee", "me"]
+
     const merkleTree = new MerkleTree(exampleArray)
     const investigatedEntry = "an entry which is not in the array"
     const proof = merkleTree.getProofElements(exampleArray.indexOf(investigatedEntry))
@@ -30,3 +32,25 @@ Deno.test("should return invalid as the investigated entry is not in array", asy
     assertEquals(isValid, false)
 
 })
+
+
+Deno.test("should be able to utilize a custom hash function", async () => {
+
+    const exampleArray = [2, 4, 8, 1]
+
+    const merkleTree = new MerkleTree(exampleArray, (x: number) => x * 2 % 10)
+    const rootHash = merkleTree.getRootHash()
+
+    assertEquals(rootHash, 8)
+
+    const proof = merkleTree.getProofElements(exampleArray.indexOf(1))
+    console.log(proof)
+
+    const investigatedEntry = 1
+    const investigatedEntryHashed = investigatedEntry * 2 % 10
+    const isValid = merkleTree.verify(proof, investigatedEntryHashed, rootHash, exampleArray.indexOf(investigatedEntry))
+
+    assertEquals(isValid, true)
+
+
+})

src/merkle-tree.ts

@@ -10,26 +10,37 @@ export class MerkleTree {
 
     private hashes: IMerkleHashes[] = []
 
+    private hashFunction: any | undefined
 
-    public constructor(array: any[]) {
+    public constructor(array: any[], hashFunction?: any) {
 
         this.checkInput(array)
 
+        this.hashFunction = hashFunction
+
         this.generateTree(array)
 
     }
 
 
-    public verify(proof: string[], leaf: string, rootHash: string, index: number) {
+    public verify(proof: string[] | number[], leaf: string | number, rootHash: string | number, index: number) {
 
         let hash = leaf
 
         for (let i = 0; i < proof.length; i++) {
 
             if (index % 2 === 0) {
-                hash = Helper.sha256(`${hash}${proof[i]}`)
+                if (this.hashFunction === undefined) {
+                    hash = Helper.sha256(`${hash}${proof[i]}`)
+                } else {
+                    hash = this.hashFunction(`${hash}${proof[i]}`)
+                }
             } else {
-                hash = Helper.sha256(`${proof[i]}${hash}`)
+                if (this.hashFunction === undefined) {
+                    hash = Helper.sha256(`${proof[i]}${hash}`)
+                } else {
+                    hash = this.hashFunction(`${proof[i]}${hash}`)
+                }
             }
 
             index = Math.floor(index / 2)
@@ -95,7 +106,6 @@ export class MerkleTree {
 
         while (itemsOnThisLevel.length > 1) {
             itemsOnThisLevel = this.getHashesForLevel(level, itemsOnThisLevel)
-
             this.hashes.push({ level, hashes: itemsOnThisLevel })
 
             level++
@@ -110,9 +120,20 @@ export class MerkleTree {
 
         for (let i = 0; i <= array.length; i++) {
             if (level === 0) {
-                hashesOnThisLevel.push(Helper.sha256(array[i]))
+                let hash
+                if (this.hashFunction === undefined) {
+                    hash = Helper.sha256(array[i])
+                } else {
+                    hash = this.hashFunction(array[i])
+                }
+                hashesOnThisLevel.push(hash)
             } else if (i % 2 === 0 && i > 0) {
-                const hash: string = Helper.sha256(`${array[i - 2]}${array[i - 1]}`)
+                let hash: string = Helper.sha256(`${array[i - 2]}${array[i - 1]}`)
+                if (this.hashFunction === undefined) {
+                    hash = Helper.sha256(`${array[i - 2]}${array[i - 1]}`)
+                } else {
+                    hash = this.hashFunction(`${array[i - 2]}${array[i - 1]}`)
+                }
                 hashesOnThisLevel.push(hash)
             }
         }