feat(lib/trie): Implement verify_proof function

lib/trie/database.go

@@ -18,13 +18,17 @@ package trie
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 
 	"github.com/ChainSafe/gossamer/lib/common"
 
 	"github.com/ChainSafe/chaindb"
 )
 
+// ErrIncompleteProof indicates the proof slice is empty
+var ErrIncompleteProof = errors.New("incomplete proof")
+
 // Store stores each trie node in the database, where the key is the hash of the encoded node and the value is the encoded node.
 // Generally, this will only be used for the genesis trie.
 func (t *Trie) Store(db chaindb.Database) error {
@@ -73,6 +77,72 @@ func (t *Trie) store(db chaindb.Batch, curr node) error {
 	return nil
 }
 
+// LoadFromProof create a partial trie based on the proof slice, as it only contains nodes that are in the proof afaik.
+func (t *Trie) LoadFromProof(proof [][]byte, root []byte) error {
+	mappedNodes := make(map[string]node)
+
+	// map all the proofs hash -> decoded node
+	// and takes the loop to indentify the root node
+	for _, rawNode := range proof {
+		var (
+			decNode node
+			err     error
+		)
+
+		if decNode, err = decodeBytes(rawNode); err != nil {
+			return err
+		}
+
+		decNode.setDirty(false)
+		decNode.setEncodingAndHash(rawNode, nil)
+
+		_, computedRoot, err := decNode.encodeAndHash()
+		if err != nil {
+			return err
+		}
+
+		mappedNodes[common.BytesToHex(computedRoot)] = decNode
+
+		if bytes.Equal(computedRoot, root) {
+			t.root = decNode
+		}
+	}
+
+	if len(mappedNodes) == 0 {
+		return ErrIncompleteProof
+	}
+
+	t.loadProof(mappedNodes, t.root)
+	return nil
+}
+
+// loadFromProof is a recursive function that will create all the trie paths based
+// on the mapped proofs slice starting by the root
+func (t *Trie) loadProof(proof map[string]node, curr node) {
+	var (
+		c  *branch
+		ok bool
+	)
+
+	if c, ok = curr.(*branch); ok {
+		return
+	}
+
+	for i, child := range c.children {
+		if child == nil {
+			continue
+		}
+
+		proofNode, ok := proof[common.BytesToHex(child.getHash())]
+		if !ok {
+			continue
+		}
+
+		c.children[i] = proofNode
+		t.loadProof(proof, proofNode)
+	}
+}
+
 // Load reconstructs the trie from the database from the given root hash. Used when restarting the node to load the current state trie.
 func (t *Trie) Load(db chaindb.Database, root common.Hash) error {
 	if root == EmptyHash {

lib/trie/lookup.go

@@ -2,86 +2,53 @@ package trie
 
 import (
 	"bytes"
-	"errors"
-
-	"github.com/ChainSafe/chaindb"
 )
 
-var (
-	// ErrProofNodeNotFound when a needed proof node is not in the database
-	ErrProofNodeNotFound = errors.New("cannot find a trie node in the database")
-)
+func findAndRecord(t *Trie, key []byte, recorder *recorder) []byte {
+	l, err := find(t.root, key, recorder)
+	if l == nil || err != nil {
+		return nil
+	}
 
-// lookup struct holds the state root and database reference
-// used to retrieve trie information from database
-type lookup struct {
-	// root to start the lookup
-	root []byte
-	db   chaindb.Database
+	return l.value
 }
 
-// newLookup returns a Lookup to helps the proof generator
-func newLookup(rootHash []byte, db chaindb.Database) *lookup {
-	lk := &lookup{db: db}
-	lk.root = make([]byte, len(rootHash))
-	copy(lk.root, rootHash)
+func find(parent node, key []byte, recorder *recorder) (*leaf, error) {
+	enc, hash, err := parent.encodeAndHash()
+	if err != nil {
+		return nil, err
+	}
 
-	return lk
-}
+	recorder.record(hash, enc)
 
-// find will return the desired value or nil if key cannot be found and will record visited nodes
-func (l *lookup) find(key []byte, recorder *recorder) ([]byte, error) {
-	partial := key
-	hash := l.root
+	switch p := parent.(type) {
+	case *branch:
+		length := lenCommonPrefix(p.key, key)
 
-	for {
-		nodeData, err := l.db.Get(hash)
-		if err != nil {
-			return nil, ErrProofNodeNotFound
+		// found the value at this node
+		if bytes.Equal(p.key, key) || len(key) == 0 {
+			return &leaf{key: p.key, value: p.value, dirty: false}, nil
 		}
 
-		nodeHash := make([]byte, len(hash))
-		copy(nodeHash, hash)
-
-		recorder.record(nodeHash, nodeData)
+		// did not find value
+		if bytes.Equal(p.key[:length], key) && len(key) < len(p.key) {
+			return nil, nil
+		}
 
-		decoded, err := decodeBytes(nodeData)
+		return find(p.children[key[length]], key[length+1:], recorder)
+	case *leaf:
+		enc, hash, err := p.encodeAndHash()
 		if err != nil {
 			return nil, err
 		}
 
-		switch currNode := decoded.(type) {
-		case nil:
-			return nil, nil
-
-		case *leaf:
-			if bytes.Equal(currNode.key, partial) {
-				return currNode.value, nil
-			}
-			return nil, nil
-
-		case *branch:
-			switch len(partial) {
-			case 0:
-				return currNode.value, nil
-			default:
-				if !bytes.HasPrefix(partial, currNode.key) {
-					return nil, nil
-				}
-
-				if bytes.Equal(partial, currNode.key) {
-					return currNode.value, nil
-				}
-
-				length := lenCommonPrefix(currNode.key, partial)
-				switch child := currNode.children[partial[length]].(type) {
-				case nil:
-					return nil, nil
-				default:
-					partial = partial[length+1:]
-					copy(hash, child.getHash())
-				}
-			}
+		recorder.record(hash, enc)
+		if bytes.Equal(p.key, key) {
+			return p, nil
 		}
+	default:
+		return nil, nil
 	}
+
+	return nil, nil
 }

lib/trie/proof.go

@@ -17,6 +17,7 @@
 package trie
 
 import (
+	"bytes"
 	"errors"
 
 	"github.com/ChainSafe/chaindb"
@@ -26,23 +27,29 @@ import (
 var (
 	// ErrEmptyTrieRoot occurs when trying to craft a prove with an empty trie root
 	ErrEmptyTrieRoot = errors.New("provided trie must have a root")
+
+	// ErrValueMismatch indicates that a returned verify proof value doesnt match with the expected value on items array
+	ErrValueMismatch = errors.New("expected value not found in the trie")
+
+	// ErrDuplicateKeys not allowed to verify proof with duplicate keys
+	ErrDuplicateKeys = errors.New("duplicate keys on verify proof")
 )
 
 // GenerateProof receive the keys to proof, the trie root and a reference to database
 // will
 func GenerateProof(root []byte, keys [][]byte, db chaindb.Database) ([][]byte, error) {
 	trackedProofs := make(map[string][]byte)
 
+	proofTrie := NewEmptyTrie()
+	if err := proofTrie.Load(db, common.BytesToHash(root)); err != nil {
+		return nil, err
+	}
+
 	for _, k := range keys {
 		nk := keyToNibbles(k)
 
-		lookup := newLookup(root, db)
 		recorder := new(recorder)
-
-		_, err := lookup.find(nk, recorder)
-		if err != nil {
-			return nil, err
-		}
+		findAndRecord(proofTrie, nk, recorder)
 
 		for !recorder.isEmpty() {
 			recNode := recorder.next()
@@ -54,10 +61,36 @@ func GenerateProof(root []byte, keys [][]byte, db chaindb.Database) ([][]byte, e
 	}
 
 	proofs := make([][]byte, 0)
-
 	for _, p := range trackedProofs {
 		proofs = append(proofs, p)
 	}
 
 	return proofs, nil
 }
+
+// Pair holds the key and value to check while verifying the proof
+type Pair struct{ Key, Value []byte }
+
+// VerifyProof ensure a given key is inside a proof by creating a proof trie based on the proof slice
+// this function ignores the order of proofs
+func VerifyProof(proof [][]byte, root []byte, items []Pair) (bool, error) {
+	for i := 1; i < len(items); i++ {
+		if bytes.Equal(items[i-1].Key, items[i].Key) {
+			return false, ErrDuplicateKeys
+		}
+	}
+
+	proofTrie := NewEmptyTrie()
+	if err := proofTrie.LoadFromProof(proof, root); err != nil {
+		return false, err
+	}
+
+	for _, i := range items {
+		recValue := proofTrie.Get(i.Key)
+		if !bytes.Equal(i.Value, recValue) {
+			return false, ErrValueMismatch
+		}
+	}
+
+	return true, nil
+}