generating and verifying state proofs (#364)

* generating and verifying state proofs

* verify correct value during proof verification

* add proof and verification to state

* exploring proof generation for arbitrary roots

Author: lovesh

state/pruning_state.py

@@ -8,6 +8,7 @@
 from state.util.fast_rlp import encode_optimized as rlp_encode, \
     decode_optimized as rlp_decode
 from state.util.utils import to_string, isHex
+from storage.kv_in_memory import KeyValueStorageInMemory
 from storage.kv_store import KeyValueStorage
 
 
@@ -86,6 +87,15 @@ def revertToHead(self, headHash=None):
             head = BLANK_NODE
         self._trie.replace_root_hash(self._trie.root_node, head)
 
+    # Proofs are always generated over committed state
+    def generate_state_proof(self, key: bytes, root=None, serialize=False):
+        return self._trie.generate_state_proof(key, root, serialize)
+
+    @staticmethod
+    def verify_state_proof(root, key, value, proof_nodes, serialized=False):
+        return Trie.verify_spv_proof(root, key, rlp_encode([value]),
+                                     proof_nodes, serialized)
+
     @property
     def as_dict(self):
         d = self._trie.to_dict()

state/test/test_state_proof_verification.py

@@ -0,0 +1,57 @@
+import pytest
+
+from state.pruning_state import PruningState
+from state.state import State
+from storage.kv_in_memory import KeyValueStorageInMemory
+from storage.kv_store_leveldb import KeyValueStorageLeveldb
+
+
+@pytest.yield_fixture(params=['memory', 'leveldb'])
+def state(request, tmpdir_factory) -> State:
+    if request.param == 'memory':
+        db = KeyValueStorageInMemory()
+    if request.param == 'leveldb':
+        db = KeyValueStorageLeveldb(tmpdir_factory.mktemp('').strpath,
+                                    'some_db')
+    state = PruningState(db)
+    yield state
+    state.close()
+
+
+def test_state_proof_and_verification(state):
+    state.set(b'k1', b'v1')
+    state.set(b'k2', b'v2')
+    state.set(b'k3', b'v3')
+    state.set(b'k4', b'v4')
+
+    p1 = state.generate_state_proof(b'k1')
+    p2 = state.generate_state_proof(b'k2')
+    p3 = state.generate_state_proof(b'k3')
+    p4 = state.generate_state_proof(b'k4')
+
+    # Verify correct proofs and values
+    assert PruningState.verify_state_proof(state.headHash, b'k1', b'v1', p1)
+    assert PruningState.verify_state_proof(state.headHash, b'k2', b'v2', p2)
+    assert PruningState.verify_state_proof(state.headHash, b'k3', b'v3', p3)
+    assert PruningState.verify_state_proof(state.headHash, b'k4', b'v4', p4)
+
+    # Incorrect proof
+    assert PruningState.verify_state_proof(state.headHash, b'k3', b'v3', p4)
+
+    # Correct proof but incorrect value
+    assert not PruningState.verify_state_proof(state.headHash, b'k2', b'v1', p2)
+    assert not PruningState.verify_state_proof(state.headHash, b'k4', b'v2', p4)
+
+
+def test_state_proof_and_verification_serialized(state):
+    data = {k.encode(): v.encode() for k, v in
+            [('k1', 'v1'), ('k2', 'v2'), ('k35', 'v55'), ('k70', 'v99')]}
+
+    for k, v in data.items():
+        state.set(k, v)
+
+    proofs = {k: state.generate_state_proof(k, serialize=True) for k in data}
+
+    for k, v in data.items():
+        assert PruningState.verify_state_proof(state.headHash, k, v,
+                                               proofs[k], serialized=True)

state/test/trie/test_proof.py

@@ -0,0 +1,215 @@
+from copy import deepcopy
+from random import random, randint, choice
+
+import pytest
+
+from plenum.common.util import randomString
+from state.db.persistent_db import PersistentDB
+from state.trie.pruning_trie import Trie, rlp_encode
+from storage.kv_in_memory import KeyValueStorageInMemory
+
+
+def gen_test_data(num_keys, max_key_size=64, max_val_size=256):
+    return {randomString(max_key_size).encode():
+            rlp_encode([randomString(max_val_size)]) for i in range(num_keys)}
+
+
+def test_verify_proof():
+    node_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+    client_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+
+    node_trie.update('k1'.encode(), rlp_encode(['v1']))
+    node_trie.update('k2'.encode(), rlp_encode(['v2']))
+
+    root_hash_0 = node_trie.root_hash
+    p0 = node_trie.produce_spv_proof('k2'.encode())
+    p0.append(deepcopy(node_trie.root_node))
+    p00 = deepcopy(p0)
+    assert client_trie.verify_spv_proof(root_hash_0, 'k2'.encode(),
+                                        rlp_encode(['v2']), p0)
+    assert p00 == p0
+
+    node_trie.update('k3'.encode(), rlp_encode(['v3']))
+    node_trie.update('k4'.encode(), rlp_encode(['v4']))
+    node_trie.update('x1'.encode(), rlp_encode(['y1']))
+    node_trie.update('x2'.encode(), rlp_encode(['y2']))
+    root_hash_1 = node_trie.root_hash
+
+    # Generate 1 proof and then verify that proof
+    p1 = node_trie.produce_spv_proof('k1'.encode())
+    p1.append(node_trie.root_node)
+    assert client_trie.verify_spv_proof(root_hash_1, 'k1'.encode(),
+                                        rlp_encode(['v1']), p1)
+
+    p2 = node_trie.produce_spv_proof('x2'.encode())
+    p2.append(node_trie.root_node)
+    assert client_trie.verify_spv_proof(root_hash_1, 'x2'.encode(),
+                                        rlp_encode(['y2']), p2)
+
+    # Generate more than 1 proof and then verify all proofs
+
+    p3 = node_trie.produce_spv_proof('k3'.encode())
+    p3.append(node_trie.root_node)
+
+    p4 = node_trie.produce_spv_proof('x1'.encode())
+    p4.append(node_trie.root_node)
+
+    assert client_trie.verify_spv_proof(root_hash_1, 'k3'.encode(),
+                                        rlp_encode(['v3']), p3)
+    assert client_trie.verify_spv_proof(root_hash_1, 'x1'.encode(),
+                                        rlp_encode(['y1']), p4)
+
+    # Proof is correct but value is different
+    assert not client_trie.verify_spv_proof(root_hash_1, 'x1'.encode(),
+                                            rlp_encode(['y99']), p4)
+
+    # Verify same proof again
+    assert client_trie.verify_spv_proof(root_hash_1, 'k3'.encode(),
+                                        rlp_encode(['v3']), p3)
+
+    assert p00 == p0
+    assert client_trie.verify_spv_proof(root_hash_0, 'k2'.encode(),
+                                        rlp_encode(['v2']), p0)
+
+
+def test_verify_proof_generated_using_helper():
+    node_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+    client_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+
+    node_trie.update('k1'.encode(), rlp_encode(['v1']))
+    node_trie.update('k2'.encode(), rlp_encode(['v2']))
+
+    root_hash_0 = node_trie.root_hash
+    p0 = node_trie.generate_state_proof('k2'.encode())
+    assert client_trie.verify_spv_proof(root_hash_0, 'k2'.encode(),
+                                        rlp_encode(['v2']), p0)
+
+    node_trie.update('k3'.encode(), rlp_encode(['v3']))
+    node_trie.update('k4'.encode(), rlp_encode(['v4']))
+    node_trie.update('x1'.encode(), rlp_encode(['y1']))
+    node_trie.update('x2'.encode(), rlp_encode(['y2']))
+    root_hash_1 = node_trie.root_hash
+
+    # Generate 1 proof and then verify that proof
+    p1 = node_trie.generate_state_proof('k1'.encode())
+    assert client_trie.verify_spv_proof(root_hash_1, 'k1'.encode(),
+                                        rlp_encode(['v1']), p1)
+
+    p2 = node_trie.generate_state_proof('x2'.encode())
+    assert client_trie.verify_spv_proof(root_hash_1, 'x2'.encode(),
+                                        rlp_encode(['y2']), p2)
+
+    # Generate more than 1 proof and then verify all proofs
+
+    p3 = node_trie.generate_state_proof('k3'.encode())
+
+    p4 = node_trie.generate_state_proof('x1'.encode())
+
+    assert client_trie.verify_spv_proof(root_hash_1, 'k3'.encode(),
+                                        rlp_encode(['v3']), p3)
+    assert client_trie.verify_spv_proof(root_hash_1, 'x1'.encode(),
+                                        rlp_encode(['y1']), p4)
+
+    # Proof is correct but value is different
+    assert not client_trie.verify_spv_proof(root_hash_1, 'x1'.encode(),
+                                            rlp_encode(['y99']), p4)
+
+    # Verify same proof again
+    assert client_trie.verify_spv_proof(root_hash_1, 'k3'.encode(),
+                                        rlp_encode(['v3']), p3)
+
+    assert client_trie.verify_spv_proof(root_hash_0, 'k2'.encode(),
+                                        rlp_encode(['v2']), p0)
+
+    # Proof generated using non-existent key fails verification
+    p5 = node_trie.generate_state_proof('x909'.encode())
+    assert not client_trie.verify_spv_proof(root_hash_1, 'x909'.encode(),
+                                            rlp_encode(['y909']), p5)
+
+
+def test_verify_proof_random_data():
+    """
+    Add some key value pairs in trie. Generate and verify proof for them.
+    :return:
+    """
+    num_keys = 100
+    test_data = gen_test_data(num_keys)
+    partitions = 4
+    partition_size = num_keys // partitions
+    keys = [list(list(test_data.keys())[i:i + partition_size])
+            for i in range(0, len(test_data), partition_size)]
+
+    node_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+    client_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+    root_hashes = []
+    proofs = []
+    for i in range(0, partitions):
+        for k in keys[i]:
+            node_trie.update(k, test_data[k])
+
+        root_hashes.append(node_trie.root_hash)
+        proofs.append({k: node_trie.generate_state_proof(k) for k in keys[i]})
+        assert all([client_trie.verify_spv_proof(root_hashes[i],
+                                                 k, test_data[k],
+                                                 proofs[i][k]) for k in keys[i]])
+
+    # Pick any keys from any partition and verify the already generated proof
+    for _ in range(400):
+        p = randint(0, partitions - 1)
+        key = choice(keys[p])
+        assert client_trie.verify_spv_proof(root_hashes[p], key,
+                                            test_data[key], proofs[p][key])
+
+    # Pick any key randomly, generate new proof corresponding to current root
+    # and verify proof
+    all_keys = [k for i in keys for k in i]
+    root_hash = node_trie.root_hash
+    for _ in range(400):
+        key = choice(all_keys)
+        proof = node_trie.generate_state_proof(key)
+        assert client_trie.verify_spv_proof(root_hash, key,
+                                            test_data[key], proof)
+
+
+def test_proof_serialize_deserialize():
+    node_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+    client_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+    keys = {k.encode(): [rlp_encode([v]), ] for k, v in
+            [('k1', 'v1'), ('k2', 'v2'), ('k35', 'v55'), ('k70', 'v99')]}
+
+    for k, v in keys.items():
+        node_trie.update(k, v[0])
+
+    for k in keys:
+        keys[k].append(node_trie.generate_state_proof(k, serialize=True))
+
+    for k in keys:
+        prf = keys[k][1]
+        assert isinstance(prf, bytes)
+        assert client_trie.verify_spv_proof(node_trie.root_hash, k, keys[k][0],
+                                            prf, serialized=True)
+
+
+@pytest.mark.skip(reason='Need to check if need/possible to build proof for an '
+                         'arbitrary old root')
+def test_proof_specific_root():
+    node_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+    client_trie = Trie(PersistentDB(KeyValueStorageInMemory()))
+
+    node_trie.update('k1'.encode(), rlp_encode(['v1']))
+    node_trie.update('k2'.encode(), rlp_encode(['v2']))
+    node_trie.update('x3'.encode(), rlp_encode(['v3']))
+
+    root_hash_0 = node_trie.root_hash
+    root_node_0 = node_trie.root_node
+
+    node_trie.update('x4'.encode(), rlp_encode(['v5']))
+    node_trie.update('y99'.encode(), rlp_encode(['v6']))
+    node_trie.update('x5'.encode(), rlp_encode(['v7']))
+
+    # root_hash_1 = node_trie.root_hash
+    # root_node_1 = node_trie.root_node
+
+    k, v = 'k1'.encode(), rlp_encode(['v1'])
+    old_root_proof = node_trie.generate_state_proof(k, root=root_node_0)
+    assert client_trie.verify_spv_proof(root_hash_0, k, v, old_root_proof)

state/trie/pruning_trie.py

@@ -6,11 +6,12 @@
 import rlp
 from rlp.utils import decode_hex, encode_hex, ascii_chr, str_to_bytes
 from state.db.db import BaseDB
-from state.util.fast_rlp import encode_optimized
+from state.util.fast_rlp import encode_optimized, decode_optimized
 from state.util.utils import is_string, to_string, sha3, sha3rlp, encode_int
 from storage.kv_in_memory import KeyValueStorageInMemory
 
 rlp_encode = encode_optimized
+rlp_decode = decode_optimized
 
 bin_to_nibbles_cache = {}
 
@@ -63,12 +64,13 @@ def __init__(self):
         self.nodes = []
         self.exempt = []
 
-    def push(self, mode, nodes=[]):
+    def push(self, mode, nodes=None):
         global proving
         proving = True
         self.mode.append(mode)
         self.exempt.append(set())
         if mode == VERIFYING:
+            nodes = nodes or []
             self.nodes.append(set([rlp_encode(x) for x in nodes]))
         else:
             self.nodes.append(set())
@@ -968,9 +970,10 @@ def root_hash_valid(self):
             return True
         return self.root_hash in self._db
 
-    def produce_spv_proof(self, key):
+    def produce_spv_proof(self, key, root=None):
+        root = root or self.root_node
         proof.push(RECORDING)
-        self.get(key)
+        self.get_at(root, key)
         o = proof.get_nodelist()
         proof.pop()
         return o
@@ -984,25 +987,45 @@ def get_at(self, root_node, key):
         """
         return self._get(root_node, bin_to_nibbles(to_string(key)))
 
+    def generate_state_proof(self, key, root=None, serialize=False):
+        # NOTE: The method `produce_spv_proof` is not deliberately modified
+        root = root or self.root_node
+        pf = self.produce_spv_proof(key, root)
+        pf.append(copy.deepcopy(root))
+        return pf if not serialize else self.serialize_proof(pf)
+
     @staticmethod
-    def verify_spv_proof(root, key, proof_nodes):
+    def verify_spv_proof(root, key, value, proof_nodes, serialized=False):
+        # NOTE: `root` is a derivative of the last element of `proof_nodes`
+        # but it's important to keep `root` as a separate as signed root
+        # hashes will be published.
+        if serialized:
+            proof_nodes = Trie.deserialize_proof(proof_nodes)
         proof.push(VERIFYING, proof_nodes)
-        t = Trie(KeyValueStorageInMemory())
+        new_trie = Trie(KeyValueStorageInMemory())
 
-        for i, node in enumerate(proof_nodes):
+        for node in proof_nodes:
             R = rlp_encode(node)
             H = sha3(R)
-            t._db.put(H, R)
+            new_trie._db.put(H, R)
         try:
-            t.root_hash = root
-            t.get(key)
+            new_trie.root_hash = root
+            v = new_trie.get(key)
             proof.pop()
-            return True
+            return v == value
         except Exception as e:
             print(e)
             proof.pop()
             return False
 
+    @staticmethod
+    def serialize_proof(proof_nodes):
+        return rlp_encode(proof_nodes)
+
+    @staticmethod
+    def deserialize_proof(ser_proof):
+        return rlp_decode(ser_proof)
+
 
 if __name__ == "__main__":