From 45e24073d2c2c2d0b0ac92646a82afc0127b26ab Mon Sep 17 00:00:00 2001 From: Luke Murphy Date: Thu, 1 Aug 2019 13:06:32 +0200 Subject: [PATCH] Another pass: rough shape of tree and tests It isn't working yet, the tests don't pass. But it is on the way! --- merkle_tree_stream/generator.py | 88 +++++++++++++++++++++++---------- setup.cfg | 3 ++ test/conftest.py | 22 +++++++++ test/test_generator.py | 66 +++++++++++++++++++++++++ 4 files changed, 154 insertions(+), 25 deletions(-) create mode 100644 test/conftest.py create mode 100644 test/test_generator.py diff --git a/merkle_tree_stream/generator.py b/merkle_tree_stream/generator.py index d29201b..20035e2 100644 --- a/merkle_tree_stream/generator.py +++ b/merkle_tree_stream/generator.py @@ -5,11 +5,10 @@ from typing import Any, Callable, List, Optional import attr from flat_tree import FlatTreeAccessor -__all__ = ['MerkleTreeGenerator', 'MerkleTreeNode'] - - Hash = str +__all__ = ['MerkleTreeGenerator', 'MerkleTreeNode'] + flat_tree = FlatTreeAccessor() @@ -17,17 +16,17 @@ flat_tree = FlatTreeAccessor() class MerkleTreeNode: """A node in a merkle tree. - :param index: TODO - :param parent: TODO - :param size: TODO - :param data: TODO - :param hash: TODO + :param index: The index of node + :param parent: The parent of the node + :param size: The size of the data + :param data: The data of the node + :param hash: The hash of the data """ index: int - parent: Optional[int] + parent: int size: int - data: bytes + data: Optional[bytes] hash: Optional[str] = None def __attrs_post_init__(self) -> Any: @@ -39,30 +38,69 @@ class MerkleTreeNode: class MerkleTreeGenerator: """A stream that generates a merkle tree based on the incoming data. - :param leaf: TODO - :param parent: TODO - :param roots: TODO - :param blocks: TODO + :param leaf: The leaf hash generation function + :param parent: The parent hash generation function + :param roots: The tree roots """ - leaf: Callable[[bytes], Hash] - parent: Callable[[bytes], Hash] - blocks: int - roots: Optional[List[MerkleTreeNode]] = attr.Factory(list) + leaf: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash] + parent: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash] + roots: List[MerkleTreeNode] = attr.Factory(list) - def next(self, data: bytes) -> List[MerkleTreeNode]: - """Further generate the treem based on the incoming data. + def next( + self, data: bytes, nodes: Optional[List[MerkleTreeNode]] = None + ) -> List[MerkleTreeNode]: + """Further generate the tree based on the incoming data. :param data: Incoming data + :param nodes: Pre-existing nodes """ - pass + nodes = nodes or [] + + index = 2 * (self.blocks + 1) + + leaf_node = MerkleTreeNode( + index=index, + parent=flat_tree.parent(index), + hash=None, + data=data, + size=len(data), + ) + + leaf_node.hash = self.leaf(leaf_node, self.roots) + + self.roots.append(leaf_node) + nodes.append(leaf_node) + + while len(self.roots) > 1: + left = self.roots[len(self.roots) - 2] + right = self.roots[len(self.roots) - 1] + + if left.parent != right.parent: + break + + self.roots.pop() + + new_node = MerkleTreeNode( + index=left.parent, + parent=flat_tree.parent(left.parent), + hash=self.parent(left, [right]), + size=left.size + right.size, + data=None, + ) + + self.roots[len(self.roots) - 1] = new_node + + nodes.append(new_node) + + return nodes def __attrs_post_init__(self) -> Any: """Initialise parent and block defaults.""" + index = self.roots[len(self.roots) - 1].index + right_span = flat_tree.right_span(index) + self.blocks = (1 + (right_span / 2)) if self.roots else 0 + for root in self.roots: if not root.parent: root.parent = flat_tree.parent(root.index) - - # https://github.com/mafintosh/merkle-tree-stream/blob/master/generator.js#L14 - # self.roots[self.roots.length] ... - # self.blocks = (1 + (flat_tree.right_span(...) / 2)) if self.roots else 0 diff --git a/setup.cfg b/setup.cfg index 0583280..c3a210a 100644 --- a/setup.cfg +++ b/setup.cfg @@ -45,6 +45,9 @@ package_dir = = . packages = find: zip_safe = False +install_requires = + attrs >= 19.1.0, < 20.0 + flat-tree == 0.0.1a3 # TODO(decentral1se): use bounds when 0.0.1 lands [options.packages.find] where = . diff --git a/test/conftest.py b/test/conftest.py new file mode 100644 index 0000000..8c40afc --- /dev/null +++ b/test/conftest.py @@ -0,0 +1,22 @@ +import hashlib + +import pytest + + +@pytest.fixture +def leaf(): + def _leaf(node): + return hashlib.sha256(leaf.data).hexdigest() + + return _leaf + + +@pytest.fixture +def parent(): + def _parent(left, right): + sha256 = hashlib.sha256() + sha256.update(left) + sha256.update(right) + return sha256.hexdigest() + + return _parent diff --git a/test/test_generator.py b/test/test_generator.py new file mode 100644 index 0000000..7ea3a91 --- /dev/null +++ b/test/test_generator.py @@ -0,0 +1,66 @@ +"""Generator test module.""" + +import hashlib + +from merkle_tree_stream import MerkleTreeGenerator, MerkleTreeNode + + +def test_hashes(leaf, parent): + stream = MerkleTreeGenerator(leaf=leaf, parent=parent) + + stream.next(b'a') + + first_node = ( + MerkleTreeNode( + index=0, + parent=1, + hash=hashlib.sha256(b'a').hexdigest(), + size=1, + data=b'a', + ), + ) + + stream.next(b'b') + + second_node = ( + MerkleTreeNode( + index=2, + parent=1, + hash=hashlib.sha256(b'b').hexdigest(), + size=1, + data=b'a', + ), + ) + + stream.next(b'c') + + third = hashlib.sha256(b'a') + third.update(b'b') + third_hash = third.hexdigest() + + third_node = ( + MerkleTreeNode(index=1, parent=3, hash=third_hash, size=2, data=b'a'), + ) + + assert stream.nodes == [first_node, second_node, third_node] + + +def test_single_root(leaf, parent): + stream = MerkleTreeGenerator(leaf=leaf, parent=parent) + + stream.next('a') + stream.next('b') + stream.next('c') + stream.next('d') + + assert stream.roots.length == 1 + + +def multiple_roots(leaf, parent): + stream = MerkleTreeGenerator(leaf=leaf, parent=parent) + + stream.next('a') + stream.next('b') + stream.next('c') + + assert stream.roots.length > 1