Improve tests and verification robustness

hippiepug/tree.py

@@ -7,6 +7,7 @@
 from warnings import warn
 
 from .struct import TreeNode, TreeLeaf
+from .store import IntegrityValidationError
 from .pack import encode, decode
 
 
@@ -76,30 +77,48 @@ def _get_inclusion_proof(self, lookup_key):
 
             try:
                 if isinstance(current_node, TreeNode):
+
                     if current_node.left_hash:
-                        left_child = self._get_node_by_hash(current_node.left_hash)
+                        left_child = self._get_node_by_hash(
+                                current_node.left_hash)
                     if current_node.right_hash:
-                        right_child = self._get_node_by_hash(current_node.right_hash)
+                        right_child = self._get_node_by_hash(
+                                current_node.right_hash)
 
+                    integrity_issue_msg = ('Not enough nodes to validate '
+                                           'inclusion.')
                     if lookup_key < current_node.pivot_prefix:
                         current_node = left_child
                         if right_child is not None:
                             closure_nodes.append(right_child)
+
+                        # If right hash is specified, but right node is not
+                        # found, can not validate inclusion.
+                        elif current_node.right_hash is not None:
+                            raise ValueError(integrity_issue_msg)
+
                     else:
                         current_node = right_child
                         if left_child is not None:
                             closure_nodes.append(left_child)
 
-                # Stop when leaf found.
+                        # Similarly, if left hash is specified, but right
+                        # node is not found, can not validate inclusion.
+                        elif current_node.left_hash is not None:
+                            raise ValueError(integrity_issue_msg)
+
+                # Stop when leaf is found.
                 elif isinstance(current_node, TreeLeaf):
                     break
+
+                # Not a tree node.
                 else:
                     raise TypeError('Invalid node type.')
 
+            # If something happened, likely a node was malformed.
             except Exception as e:
                 warn('Exception occured when handling node %s: %s' % (
                     current_node, e))
-                # Also stop when something is broken.
                 break
 
         return path_nodes, closure_nodes
@@ -115,14 +134,13 @@ def get_value_by_lookup_key(self, lookup_key, return_proof=False):
         path, closure = self._get_inclusion_proof(lookup_key)
         result = None
         if path:
-            try:
-                if path[-1].lookup_key == lookup_key:
-                    serialized_payload = self.object_store.get(
-                            path[-1].payload_hash)
-                    result = serialized_payload
-            except Exception as e:
-                warn('Exception occured when handling node %s: %s' % (
-                    path[-1], e))
+            maybe_leaf = path[-1]
+            if maybe_leaf is not None and (
+                    hasattr(maybe_leaf, 'lookup_key')) and (
+                        maybe_leaf.lookup_key == lookup_key):
+                serialized_payload = self.object_store.get(
+                        path[-1].payload_hash)
+                result = serialized_payload
 
         if return_proof:
             return result, (path, closure)
@@ -264,13 +282,14 @@ def __repr__(self):
                     self=self)
 
 
-def verify_tree_inclusion_proof(store, root, lookup_key, payload, proof):
+# TODO: Verify correctness of the tree root!
+def verify_tree_inclusion_proof(store, root, lookup_key, value, proof):
     """Verify inclusion proof for a tree.
 
     :param store: Object store
     :param head: Tree root
     :param lookup_key: Lookup key
-    :param payload: Payload hash associated with the lookup key
+    :param value: Value associated with the lookup key
     :param proof: Inclusion proof
     :type proof: tuple, containing list of encoded path nodes, and encoded
                  closure nodes.
@@ -280,7 +299,7 @@ def verify_tree_inclusion_proof(store, root, lookup_key, payload, proof):
     path, closure = proof
     for node in path + closure:
         store.add(encode(node))
-    store.add(payload)
+    store.add(value)
     verifier_tree = Tree(store, root=root)
     retrieved_payload = verifier_tree.get_value_by_lookup_key(lookup_key)
-    return retrieved_payload == payload
+    return retrieved_payload == value

tests/test_chain.py

@@ -188,7 +188,7 @@ def test_chain_iterator(chain_and_hashes):
         assert block_hash == expected_block_hash
 
 
-def test_chain_proof(object_store):
+def test_chain_inclusion_proof(object_store):
     """Check returned inclusion proof."""
     chain1 = Chain(object_store)
     for i in range(10):
@@ -207,39 +207,59 @@ def test_chain_proof(object_store):
     assert chain2.get_block_by_index(2).payload == 'Block 2'
 
 
-def test_chain_proof_verification_util(object_store):
-    """Check verification utility."""
-    chain1 = Chain(object_store)
-    for i in range(10):
-        block_builder = BlockBuilder(chain1)
-        block_builder.payload='Block %i' % i
-        block_builder.commit()
+def test_chain_proof_verif(chain_and_hashes):
+    """Check proof verification utility."""
+    chain, hashes = chain_and_hashes
+
+    for i in range(len(hashes)):
+        result, proof = chain.get_block_by_index(i, return_proof=True)
+        store = chain.object_store.__class__()
+        assert verify_chain_inclusion_proof(
+                store, chain.head, result, proof)
 
-    # Check proof of inclusion.
-    store = chain1.object_store.__class__()
-    result, proof = chain1.get_block_by_index(2, return_proof=True)
-    assert verify_chain_inclusion_proof(store, chain1.head, result, proof)
 
-    # Check proof of inclusion detects bad block.
-    store = chain1.object_store.__class__()
+def test_chain_proof_verif_bad_block(chain_and_hashes):
+    """Check proof of inclusion detects bad block."""
+    chain, hashes = chain_and_hashes
+    store = chain.object_store.__class__()
     bad_block = ChainBlock('non-existent')
+    result, proof = chain.get_block_by_index(0, return_proof=True)
+
     assert not verify_chain_inclusion_proof(
-            store, chain1.head, bad_block, proof)
+            store, chain.head, bad_block, proof)
 
-    # Check proof of inclusion detects bad block with the same index.
-    store = chain1.object_store.__class__()
+
+def test_chain_proof_verif_bad_block_same_index(chain_and_hashes):
+    """Check proof of inclusion detects bad block with the same index."""
+    chain, hashes = chain_and_hashes
+    store = chain.object_store.__class__()
+    result, proof = chain.get_block_by_index(0, return_proof=True)
     bad_block = ChainBlock('non-existent', index=result.index)
+
     assert not verify_chain_inclusion_proof(
-            store, chain1.head, bad_block, proof)
+            store, chain.head, bad_block, proof)
 
-    # Check proof of inclusion fails when it's insufficient
-    store = chain1.object_store.__class__()
+
+def test_chain_proof_verif_insufficient(chain_and_hashes):
+    """Check proof of inclusion fails when it's insufficient."""
+    chain, hashes = chain_and_hashes
+    store = chain.object_store.__class__()
+    result, proof = chain.get_block_by_index(0, return_proof=True)
     assert not verify_chain_inclusion_proof(
-            store, chain1.head, result, proof=[])
+            store, chain.head, result, proof=[])
 
-    # Check proof of inclusion fails when it's malformed
-    store = chain1.object_store.__class__()
+
+def test_chain_proof_verif_malformed(chain_and_hashes):
+    """Check proof of inclusion fails when it's malformed."""
+    chain, hashes = chain_and_hashes
+    if chain.head_block.index == 0:
+        return
+
+    store = chain.object_store.__class__()
+    result, proof = chain.get_block_by_index(0, return_proof=True)
     proof[0].fingers = None
     bad_head = store.hash_object(encode(proof[0]))
-    assert not verify_chain_inclusion_proof(
-            store, bad_head, result, proof)
+
+    with pytest.warns(UserWarning, match='Exception occured'):
+        assert not verify_chain_inclusion_proof(
+                store, bad_head, result, proof)

tests/test_tree.py

@@ -10,6 +10,8 @@
 from hippiepug.pack import encode
 
 
+LOOKUP_KEYS = ['AB', 'AC', 'ZZZ', 'Z']
+
 # Test tree:
 #     /ZZZ-|
 #   ZZ
@@ -20,10 +22,8 @@
 @pytest.fixture
 def populated_tree(object_store):
     builder = TreeBuilder(object_store)
-    builder['AB']  = b'AB value'
-    builder['AC']  = b'AC value'
-    builder['ZZZ'] = b'ZZZ value'
-    builder['Z']   = b'Z value'
+    for lookup_key in LOOKUP_KEYS:
+        builder[lookup_key] = b'%s value' % lookup_key.encode('utf-8')
     return builder.commit()
 
 
@@ -100,7 +100,7 @@ def test_tree_get_node_by_hash_fails_if_not_node(populated_tree):
 
 
 def test_tree_inclusion_proof(populated_tree):
-    """Check tree (non-)inclusion proof."""
+    """Manually check tree (non-)inclusion proof."""
 
     # Inclusion in the right subtree.
     _, (path, closure) = populated_tree.get_value_by_lookup_key(
@@ -139,22 +139,71 @@ def test_tree_inclusion_proof(populated_tree):
     assert closure[1].lookup_key == 'Z'
 
 
-@pytest.mark.parametrize('lookup_key', ['AB', 'AC', 'ZZZ', 'Z'])
-def test_tree_proof_verification_util(populated_tree, lookup_key):
+@pytest.mark.parametrize('lookup_key', LOOKUP_KEYS)
+def test_tree_proof_verif(populated_tree, lookup_key):
+    """Check proof of inclusion verification."""
     root = populated_tree.root
     payload, proof = populated_tree.get_value_by_lookup_key(lookup_key,
             return_proof=True)
-
-    # Check proof of inclusion.
     store = populated_tree.object_store.__class__()
-    assert verify_tree_inclusion_proof(store, root, lookup_key, payload, proof)
+    assert verify_tree_inclusion_proof(
+            store, root, lookup_key, payload, proof)
 
-    # Check proof of inclusion fails when lookup_key in the leaf is different
+
+@pytest.mark.parametrize('lookup_key', LOOKUP_KEYS)
+def test_tree_proof_verif_fails_when_leaf_different(
+        populated_tree, lookup_key):
+    """Check proof fails when lookup_key in the leaf is different."""
+    root = populated_tree.root
+    payload, proof = populated_tree.get_value_by_lookup_key(lookup_key,
+            return_proof=True)
     store = populated_tree.object_store.__class__()
     proof[0][-1].lookup_key = 'hacked'
-    assert not verify_tree_inclusion_proof(store, root, lookup_key, payload, proof)
 
-    # Check proof of inclusion fails when payload is different
+    assert not verify_tree_inclusion_proof(
+            store, root, lookup_key, payload, proof)
+
+
+@pytest.mark.parametrize('lookup_key', LOOKUP_KEYS)
+def test_tree_proof_verif_fails_when_payload_different(
+        populated_tree, lookup_key):
+    """Check proof of inclusion fails when payload is different."""
+    root = populated_tree.root
     store = populated_tree.object_store.__class__()
+    payload, proof = populated_tree.get_value_by_lookup_key(lookup_key,
+            return_proof=True)
     payload = b'non-existent'
-    assert not verify_tree_inclusion_proof(store, root, lookup_key, payload, proof)
+
+    assert not verify_tree_inclusion_proof(
+            store, root, lookup_key, payload, proof)
+
+
+@pytest.mark.parametrize('lookup_key', LOOKUP_KEYS)
+def test_tree_proof_verif_fails_when_lacks_closure(
+        populated_tree, lookup_key):
+    """Check proof of inclusion fails when it lacks closure."""
+    root = populated_tree.root
+    payload, proof = populated_tree.get_value_by_lookup_key(lookup_key,
+            return_proof=True)
+    store = populated_tree.object_store.__class__()
+    path, closure = proof
+    bad_proof = (path, [])
+
+    with pytest.warns(UserWarning, match='Not enough nodes'):
+        assert not verify_tree_inclusion_proof(
+                store, root, lookup_key, payload, bad_proof)
+
+
+@pytest.mark.parametrize('lookup_key', LOOKUP_KEYS)
+def test_tree_proof_verif_fails_when_path_is_bad(
+        populated_tree, lookup_key):
+    """Check proof of inclusion fails when path is bad."""
+    root = populated_tree.root
+    store = populated_tree.object_store.__class__()
+    payload, proof = populated_tree.get_value_by_lookup_key(lookup_key,
+            return_proof=True)
+    path, closure = proof
+    bad_proof = (path[:2], closure)
+
+    assert not verify_tree_inclusion_proof(
+            store, root, lookup_key, payload, bad_proof)