feat: mekletreejs compatibility

.github/workflows/ci.yml

@@ -33,6 +33,14 @@ jobs:
           pip install pytest
           pip install -e .
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+      - name: Set up Node.js
+        uses: actions/setup-node@v2
+        with:
+          node-version: "18"
+      - name: Install Node.js dependencies
+        run: |
+          cd test/merkletreejs
+          yarn install
       - name: Test with pytest
         run: |
           pytest -m "not benchmark" -vv

.gitignore

@@ -1,3 +1,4 @@
+node_modules
 .hypothesis
 .vscode
 *cache*

merkly/mtree.py

@@ -40,23 +40,25 @@ def __init__(
         self.short_leafs: List[str] = self.short(self.leafs)
 
     def __hash_leafs(self, leafs: List[str]) -> List[str]:
-        return list(map(lambda x: self.hash_function(x.encode(), b""), leafs))
+        return list(map(lambda x: self.hash_function(x.encode(), bytes()), leafs))
 
     def __repr__(self) -> str:
         return f"""MerkleTree(\nraw_leafs: {self.raw_leafs}\nleafs: {self.leafs}\nshort_leafs: {self.short(self.leafs)})"""
 
     def short(self, data: List[str]) -> List[str]:
-        return [f"{x[:4]}..." for x in data]
+        return [x[:2] for x in data]
 
     @property
     def root(self) -> bytes:
         return self.make_root(self.leafs)
 
     def proof(self, raw_leaf: str) -> List[Node]:
-        return self.make_proof(self.leafs, [], self.hash_function(raw_leaf, ""))
+        return self.make_proof(
+            self.leafs, [], self.hash_function(raw_leaf.encode(), bytes())
+        )
 
-    def verify(self, proof: List[str], raw_leaf: str) -> bool:
-        full_proof = [self.hash_function(raw_leaf, "")]
+    def verify(self, proof: List[bytes], raw_leaf: str) -> bool:
+        full_proof = [self.hash_function(raw_leaf.encode(), bytes())]
         full_proof.extend(proof)
 
         def concat_nodes(left: Node, right: Node) -> Node:
@@ -78,18 +80,22 @@ def concat_nodes(left: Node, right: Node) -> Node:
 
         return reduce(concat_nodes, full_proof).data == self.root
 
-    def make_root(self, leafs: List[str]) -> List[str]:
-        if len(leafs) == 1:
-            return leafs
+    def make_root(self, leafs: List[bytes]) -> List[str]:
+        while len(leafs) > 1:
+            next_level = []
+            for i in range(0, len(leafs) - 1, 2):
+                next_level.append(self.hash_function(leafs[i], leafs[i + 1]))
 
-        return self.make_root(
-            [
-                self.hash_function(pair[0], pair[1]) if len(pair) > 1 else pair[0]
-                for pair in slice_in_pairs(leafs)
-            ]
-        )
+            if len(leafs) % 2 == 1:
+                next_level.append(leafs[-1])
+
+            leafs = next_level
 
-    def make_proof(self, leafs: List[str], proof: List[Node], leaf: str) -> List[Node]:
+        return leafs[0]
+
+    def make_proof(
+        self, leafs: List[bytes], proof: List[Node], leaf: bytes
+    ) -> List[Node]:
         """
         # Make a proof
 
@@ -126,14 +132,14 @@ def make_proof(self, leafs: List[str], proof: List[Node], leaf: str) -> List[Nod
         left, right = half(leafs)
 
         if index < len(leafs) / 2:
-            proof.append(Node(data=self.make_root(right)[0], side=Side.RIGHT))
+            proof.append(Node(data=self.make_root(right), side=Side.RIGHT))
             return self.make_proof(left, proof, leaf)
         else:
-            proof.append(Node(data=self.make_root(left)[0], side=Side.LEFT))
+            proof.append(Node(data=self.make_root(left), side=Side.LEFT))
             return self.make_proof(right, proof, leaf)
 
     def mix_tree(
-        self, leaves: List[str], proof: List[Node], leaf_index: int
+        self, leaves: List[bytes], proof: List[Node], leaf_index: int
     ) -> List[Node]:
         if len(leaves) == 1:
             return proof
@@ -148,7 +154,7 @@ def mix_tree(
 
         return self.mix_tree(self.up_layer(leaves), proof, leaf_index // 2)
 
-    def up_layer(self, leaves: List[str]) -> List[str]:
+    def up_layer(self, leaves: List[bytes]) -> List[bytes]:
         new_layer = []
         for pair in slice_in_pairs(leaves):
             if len(pair) == 1:

merkly/utils.py

@@ -19,7 +19,7 @@ class InvalidHashFunctionError(Exception):
     """Exception raised for invalid hash function."""
 
     def __init__(self) -> None:
-        self.message = "Must type of: (str) -> str"
+        self.message = "Must type of: (bytes, bytes) -> bytes"
         super().__init__(self.message)
 
 
@@ -89,7 +89,9 @@ def slice_in_pairs(list_item: list):
 
 
 def validate_leafs(leafs: List[str]):
-    if len(leafs) < 2:
+    size = len(leafs)
+
+    if size < 2:
         raise Exception("Invalid size, need > 2")
 
     a = isinstance(leafs, List)

test/errors/test_errors.py

@@ -0,0 +1,27 @@
+from merkly.utils import InvalidHashFunctionError
+from merkly.mtree import MerkleTree
+from pytest import raises
+
+
+def test_make_proof_value_error():
+    leafs = ["a", "b", "c", "d", "e", "f", "g", "h"]
+    tree = MerkleTree(leafs)
+
+    invalid_leaf = "invalid"
+    with raises(ValueError) as error:
+        tree.make_proof(leafs, [], invalid_leaf)
+
+    assert (
+        str(error.value) == f"Leaf: {invalid_leaf} does not exist in the tree: {leafs}"
+    )
+
+
+def test_invalid_hash_function_error():
+    def invalid_hash_function(data):
+        return 123
+
+    with raises(InvalidHashFunctionError):
+        MerkleTree(
+            ["a", "b", "c", "d"],
+            invalid_hash_function,
+        )

test/merkletreejs/merkle_proof/merkle_proof_test.js

@@ -0,0 +1,9 @@
+const { MerkleTree } = require('merkletreejs');
+const SHA256 = require('crypto-js/sha256');
+
+const leaves = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'].map(x => SHA256(x));
+const tree = new MerkleTree(leaves, SHA256);
+const leaf = SHA256('a');
+const proof = tree.getProof(leaf).map(node => ({ data: node.data.toString('hex'), position: node.position }));
+
+console.log(JSON.stringify({ proof, isValid: tree.verify(proof, leaf, tree.getRoot()) }))

test/merkletreejs/merkle_proof/merkle_proof_test.py

@@ -0,0 +1,19 @@
+from merkly.mtree import MerkleTree
+import hashlib
+import json
+
+
+def sha256(x, y):
+    data = x + y
+    return hashlib.sha256(data).digest()
+
+
+leaves = ["a", "b", "c", "d", "e", "f", "g", "h"]
+tree = MerkleTree(leaves, sha256)
+leaf = "a"
+proof = tree.proof(leaf)
+formatted_proof = [
+    {"data": node.data.hex(), "position": node.side.name.lower()} for node in proof
+]
+
+print(json.dumps({"proof": formatted_proof, "isValid": tree.verify(proof, leaf)}))

test/merkletreejs/merkle_root/merkle_root_test.js

@@ -0,0 +1,8 @@
+const { MerkleTree } = require('merkletreejs');
+const SHA256 = require('crypto-js/sha256');
+
+const leaves = ['a', 'b', 'c', 'd'].map(SHA256);
+const tree = new MerkleTree(leaves, SHA256, {});
+const root = tree.getRoot().toString('hex');
+
+console.log(JSON.stringify({ root }));

test/merkletreejs/merkle_root/merkle_root_test.py

@@ -0,0 +1,15 @@
+from merkly.mtree import MerkleTree
+import hashlib
+import json
+
+
+def sha256(x, y):
+    data = x + y
+    return hashlib.sha256(data).digest()
+
+
+leaves = ["a", "b", "c", "d"]
+tree = MerkleTree(leaves, sha256)
+root = tree.root.hex()
+
+print(json.dumps({"root": root}))

test/merkletreejs/package.json

@@ -0,0 +1,10 @@
+{
+    "name": "merkletreejs",
+    "version": "1.0.0",
+    "main": "index.js",
+    "license": "MIT",
+    "dependencies": {
+        "merkletreejs": "^0.3.11",
+        "web3": "^4.3.0"
+    }
+}

test/merkletreejs/test_merkle_proof_compatibility.py

@@ -0,0 +1,30 @@
+from pytest import mark
+import subprocess
+import json
+
+
+@mark.merkletreejs
+def test_merkle_proof_compatibility_between_merkletreejs_and_merkly():
+    result = subprocess.run(["yarn"], check=False)
+
+    assert result.returncode == 0, result.stderr
+
+    result_js = subprocess.run(
+        ["node", "./test/merkletreejs/merkle_proof/merkle_proof_test.js"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    assert result_js.returncode == 0, result_js.stderr
+    data_js = json.loads(result_js.stdout)
+
+    result_py = subprocess.run(
+        ["python", "./test/merkletreejs/merkle_proof/merkle_proof_test.py"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    assert result_py.returncode == 0, result_py.stderr
+    data_py = json.loads(result_py.stdout)
+
+    assert data_js == data_py

test/merkletreejs/test_merkle_root_compatibility.py

@@ -0,0 +1,30 @@
+from pytest import mark
+import subprocess
+import json
+
+
+@mark.merkletreejs
+def test_merkle_root_compatibility_between_merkletreejs_and_merkly():
+    result = subprocess.run(["yarn"], check=False)
+
+    assert result.returncode == 0, result.stderr
+
+    result_js = subprocess.run(
+        ["node", "./test/merkletreejs/merkle_root/merkle_root_test.js"],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert result_js.returncode == 0, result_js.stderr
+    merkle_root_js = json.loads(result_js.stdout)
+
+    result_py = subprocess.run(
+        ["python", "./test/merkletreejs/merkle_root/merkle_root_test.py"],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert result_py.returncode == 0, result_py.stderr
+    merkle_root_py = json.loads(result_py.stdout)
+
+    assert merkle_root_js == merkle_root_py