diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e3129c2..7df6703 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/.gitignore b/.gitignore index 8079cec..4581abf 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +node_modules .hypothesis .vscode *cache* diff --git a/merkly/mtree.py b/merkly/mtree.py index ba80e9a..43a985e 100644 --- a/merkly/mtree.py +++ b/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: diff --git a/merkly/utils.py b/merkly/utils.py index 2dec063..d7eb232 100644 --- a/merkly/utils.py +++ b/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) diff --git a/test/errors/test_errors.py b/test/errors/test_errors.py new file mode 100644 index 0000000..b584e3b --- /dev/null +++ b/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, + ) diff --git a/test/merkletreejs/merkle_proof/merkle_proof_test.js b/test/merkletreejs/merkle_proof/merkle_proof_test.js new file mode 100644 index 0000000..4d2e44c --- /dev/null +++ b/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()) })) diff --git a/test/merkletreejs/merkle_proof/merkle_proof_test.py b/test/merkletreejs/merkle_proof/merkle_proof_test.py new file mode 100644 index 0000000..eaa3957 --- /dev/null +++ b/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)})) diff --git a/test/merkletreejs/merkle_root/merkle_root_test.js b/test/merkletreejs/merkle_root/merkle_root_test.js new file mode 100644 index 0000000..5f3e2ad --- /dev/null +++ b/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 })); diff --git a/test/merkletreejs/merkle_root/merkle_root_test.py b/test/merkletreejs/merkle_root/merkle_root_test.py new file mode 100644 index 0000000..cb69112 --- /dev/null +++ b/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})) diff --git a/test/merkletreejs/package.json b/test/merkletreejs/package.json new file mode 100644 index 0000000..120e486 --- /dev/null +++ b/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" + } +} \ No newline at end of file diff --git a/test/merkletreejs/test_merkle_proof_compatibility.py b/test/merkletreejs/test_merkle_proof_compatibility.py new file mode 100644 index 0000000..ad756dc --- /dev/null +++ b/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 diff --git a/test/merkletreejs/test_merkle_root_compatibility.py b/test/merkletreejs/test_merkle_root_compatibility.py new file mode 100644 index 0000000..44a69a6 --- /dev/null +++ b/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