trie: implement stacktrie-based verification

trie: implement proof iteration + stacktrie.insert_hash

trie/stackproof.go

@@ -0,0 +1,224 @@
+package trie
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/ethdb"
+)
+
+// nodeToStNode converts from `node` to `*stNode`.
+func nodeToStNode(n node, key []byte) *stNode {
+	st := new(stNode)
+	switch n := n.(type) {
+	case *shortNode:
+		st.typ = extNode
+		st.key = append([]byte{}, n.Key...)
+	case *fullNode:
+		st.typ = branchNode
+		idx := int(key[0])
+		for i := 0; i < idx; i++ {
+			sibling := n.Children[i]
+			if sibling == nil {
+				continue
+			}
+			siblingNode := new(stNode)
+			siblingNode.typ = hashedNode
+
+			if hash, ok := sibling.(hashNode); ok {
+				siblingNode.val = []byte(hash)
+			} else {
+				// This happens is the sibling is small enough (<32B) to be inlined,
+				// in which case the rlp-encoded node is embedded instead of the hash
+				short := sibling.(*shortNode)
+				short.Key = hexToCompact(short.Key)
+				siblingNode.val = nodeToBytes(short)
+			}
+			st.children[i] = siblingNode
+		}
+	default:
+		panic(fmt.Sprintf("%T", n))
+	}
+	return st
+}
+
+func resolveFromProof(proofDb ethdb.KeyValueReader, hash common.Hash) (node, error) {
+	data, _ := proofDb.Get(hash[:])
+	if data == nil {
+		return nil, fmt.Errorf("proof node (hash %064x) missing", hash)
+	}
+	n, err := decodeNode(data[:], data)
+	if err != nil {
+		return nil, fmt.Errorf("bad proof node: %v", err)
+	}
+	return n, err
+}
+
+// newStackTrieFromProof creates a new stacktrie, and initialises it from the given
+// proof. It does so by starting at the given root, traverses along the given
+// key, and, one by one, converts the nodes into stacktrie elements.
+//
+// OBS: The resulting stacktrie instance is not guaranteed to be structurally
+// identical to a stacktrie which is initialized from scratch by feeding the
+// corresponding elements!
+// A proof-initialized (PI) stack-trie has some implicit prescient knowledge! Therefore,
+// a PI can have already expanded a shortnode into shortnode+fullnode, which a non-PI
+// will do only later.
+//
+// However, the two guarantees that PI gives are:
+// - Identical hash,
+// - Identical commit-sequence of nodes.
+//
+// OBS 2: The element in proof should _not_ be added again during value-filling.
+// OBS 3: Proofs-of-abscence have not been fully tested. TODO @holiman
+func newStackTrieFromProof(rootHash common.Hash, key []byte, proofDb ethdb.KeyValueReader, writeFn NodeWriteFunc) (*StackTrie, error) {
+	var (
+		err               error
+		child, parent     node
+		stChild, stParent *stNode
+		keyrest           []byte
+		stack             = NewStackTrie(writeFn)
+	)
+	key = keybytesToHex(key)
+	// First we need to resolve the root node from the proof.
+	if parent, err = resolveFromProof(proofDb, rootHash); err != nil {
+		return nil, err
+	}
+	stParent = nodeToStNode(parent, key)
+	stack.root = stParent
+	// Now we pursue the given key downwards, and populate the stacktrie too
+	for {
+		keyrest, child = get(parent, key, false)
+		switch cld := child.(type) {
+		case nil:
+			return nil, errors.New("no node at given path")
+		case hashNode:
+			child, err = resolveFromProof(proofDb, common.BytesToHash(cld))
+			if err != nil {
+				return nil, err
+			}
+		case valueNode:
+			// The value node goes right into the child
+			stParent.val = common.CopyBytes(cld)
+			stParent.typ = leafNode
+			// remove the terminator
+			stParent.key = stParent.key[:len(stParent.key)-1]
+			return stack, nil
+		case *shortNode:
+			// In the case of small leaves, we might end up here with a fullnode
+			// whose child is an embedded *shortNode.
+		default:
+			// we don't expect fullnodes
+			panic(fmt.Sprintf("got %T", cld))
+		}
+		stChild = nodeToStNode(child, keyrest) // convert to stacktrie equivalent
+		// Link the parent and child.
+		switch pnode := parent.(type) {
+		case *shortNode:
+			stParent.children[0] = stChild
+		case *fullNode:
+			stParent.children[key[0]] = stChild
+		default:
+			panic(fmt.Sprintf("%T: invalid node: %v", pnode, pnode))
+		}
+		key = keyrest
+		parent = child
+		stParent = stChild
+	}
+}
+
+func (st *stNode) dumpTrie(lvl int) {
+	var indent []byte
+	for i := 0; i < lvl; i++ {
+		indent = append(indent, ' ')
+	}
+	switch st.typ {
+	case branchNode:
+		fmt.Printf("\n%s FN (key='%#x')", string(indent), st.key)
+
+		for i := 0; i < 16; i++ {
+			if st.children[i] == nil {
+				continue
+			}
+			fmt.Printf("\n%s %#x. ", string(indent), i)
+			st.children[i].dumpTrie(lvl + 1)
+		}
+		fmt.Println("")
+	case extNode:
+		fmt.Printf("%s: sn('%#x')", string(indent), st.key)
+		st.children[0].dumpTrie(lvl + 1)
+	case leafNode:
+		fmt.Printf("%s: leaf('%#x'): %x ", string(indent), st.key, st.val)
+	case hashedNode:
+		fmt.Printf("hash: %#x %x", st.val, st.key)
+	default:
+		fmt.Printf("Foo: %d ? ", st.typ)
+	}
+}
+
+type hashPath struct {
+	path []byte
+	hash []byte
+}
+
+// iterateProof iterates through a proof, starting at the root given by rootHash, and follows the path.
+// Along the way, the hashes/paths are collected and delivered.
+// If 'ascending' is true, the paths will be on the left side of the proof going down,
+// If 'ascending' is false, the paths will be on the right side of the proof going up.
+func iterateProof(rootHash common.Hash, path []byte, ascending bool, proof ethdb.KeyValueReader) ([]*hashPath, error) {
+	path = keybytesToHex(path)
+	var (
+		position = 0
+		n, _     = resolveFromProof(proof, rootHash)
+		paths    []*hashPath
+	)
+	if n == nil {
+		return nil, fmt.Errorf("proof node (hash %064x) missing", rootHash)
+	}
+	for {
+		//fmt.Printf("At position %x\n", path[:position])
+		switch typ := n.(type) {
+		case *shortNode:
+			n = typ.Val
+			position += len(typ.Key)
+		case *fullNode:
+			i, delta := 0, 1 // Start at zero, iterate upwards
+			if !ascending {
+				i, delta = len(typ.Children)-1, -1 // Start at max, iterate down
+			}
+			for ; byte(i) != path[position]; i += delta {
+				if typ.Children[i] == nil {
+					continue
+				}
+				currentPath := append([]byte{}, path[:position]...)
+				currentPath = append(currentPath, byte(i))
+				if hn, ok := typ.Children[i].(hashNode); ok {
+					paths = append(paths, &hashPath{currentPath, []byte(hn)})
+				} else {
+					// This happens is the sibling is small enough (<32B) to be inlined,
+					// in which case the rlp-encoded node is embedded instead of the hash
+					short := typ.Children[i].(*shortNode)
+					short.Key = hexToCompact(short.Key)
+					data := nodeToBytes(short)
+					paths = append(paths, &hashPath{currentPath, data})
+				}
+				//fmt.Printf("%d. node at (typ %T) : %x\n", i, typ.Children[i], currentPath)
+			}
+			n = typ.Children[path[position]]
+			position++
+		default:
+			break
+		}
+		if position == len(path) {
+			break
+		}
+		if hn, ok := n.(hashNode); ok {
+			n, _ = resolveFromProof(proof, common.Hash(hn))
+			if n == nil {
+				return nil, fmt.Errorf("proof node (hash %064x) missing", rootHash)
+			}
+		}
+	}
+	return paths, nil
+}

trie/stackproof_test.go

@@ -0,0 +1,144 @@
+package trie
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/ethdb/memorydb"
+	"golang.org/x/crypto/sha3"
+	"golang.org/x/exp/slices"
+)
+
+func trieWithSmallValues() (*Trie, map[string]*kv) {
+	trie := NewEmpty(NewDatabase(rawdb.NewMemoryDatabase(), nil))
+	vals := make(map[string]*kv)
+	// This loop creates a few dense nodes with small leafs: hence will
+	// cause embedded nodes.
+	for i := byte(0); i < 100; i++ {
+		value := &kv{common.LeftPadBytes([]byte{i}, 32), []byte{i}, false}
+		trie.MustUpdate(value.k, value.v)
+		vals[string(value.k)] = value
+	}
+	return trie, vals
+}
+
+func TestStRangeProofLeftside(t *testing.T) {
+	trie, vals := randomTrie(4096)
+	testStRangeProofLeftside(t, trie, vals)
+}
+
+func TestStRangeProofLeftsideSmallValues(t *testing.T) {
+	trie, vals := trieWithSmallValues()
+	testStRangeProofLeftside(t, trie, vals)
+}
+
+func testStRangeProofLeftside(t *testing.T, trie *Trie, vals map[string]*kv) {
+	var (
+		want    = trie.Hash()
+		entries []*kv
+	)
+	for _, kv := range vals {
+		entries = append(entries, kv)
+	}
+	slices.SortFunc(entries, (*kv).cmp)
+	for start := 10; start < len(vals); start *= 2 {
+		// Set write-fn on both stacktries, to compare outputs
+		var (
+			haveSponge = &spongeDb{sponge: sha3.NewLegacyKeccak256(), id: "have"}
+			wantSponge = &spongeDb{sponge: sha3.NewLegacyKeccak256(), id: "want"}
+			proof      = memorydb.New()
+		)
+		// Provide the proof for the first entry
+		if err := trie.Prove(entries[start].k, proof); err != nil {
+			t.Fatalf("Failed to prove the first node %v", err)
+		}
+		// Initiate the stacktrie with the proof
+		stTrie, err := newStackTrieFromProof(trie.Hash(), entries[start].k, proof, func(owner common.Hash, path []byte, hash common.Hash, blob []byte) {
+			rawdb.WriteTrieNode(haveSponge, owner, path, hash, blob, "path")
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		// Initiate a reference stacktrie without proof (filling manually)
+		refTrie := NewStackTrie(nil)
+		for i := 0; i <= start; i++ { // do prefill
+			k, v := common.CopyBytes(entries[i].k), common.CopyBytes(entries[i].v)
+			refTrie.Update(k, v)
+		}
+		refTrie.writeFn = func(owner common.Hash, path []byte, hash common.Hash, blob []byte) {
+			rawdb.WriteTrieNode(wantSponge, owner, path, hash, blob, "path")
+		}
+		// Feed the remaining values into them both
+		for i := start + 1; i < len(vals); i++ {
+			stTrie.Update(entries[i].k, common.CopyBytes(entries[i].v))
+			refTrie.Update(entries[i].k, common.CopyBytes(entries[i].v))
+		}
+		// Verify the final trie hash
+		if have := stTrie.Hash(); have != want {
+			t.Fatalf("wrong hash, have %x want %x\n", have, want)
+		}
+		if have := refTrie.Hash(); have != want {
+			t.Fatalf("wrong hash, have %x want %x\n", have, want)
+		}
+		// Verify the sequence of committed nodes
+		if have, want := haveSponge.sponge.Sum(nil), wantSponge.sponge.Sum(nil); !bytes.Equal(have, want) {
+			// Show the journal
+			t.Logf("Want:")
+			for i, v := range wantSponge.journal {
+				t.Logf("op %d: %v", i, v)
+			}
+			t.Logf("Have:")
+			for i, v := range haveSponge.journal {
+				t.Logf("op %d: %v", i, v)
+			}
+			t.Errorf("proof from %d: disk write sequence wrong:\nhave %x want %x\n", start, have, want)
+		}
+	}
+}
+
+func TestStackInsertHash(t *testing.T) {
+	trie, vals := randomTrie(4096)
+	testStackInsertHash(t, trie, vals)
+}
+
+func testStackInsertHash(t *testing.T, trie *Trie, vals map[string]*kv) {
+	var (
+		entries []*kv
+		want    = trie.Hash()
+	)
+	for _, kv := range vals {
+		entries = append(entries, kv)
+	}
+	slices.SortFunc(entries, (*kv).cmp)
+	for start := 10; start < len(vals); start *= 2 {
+		var (
+			proof = memorydb.New()
+		)
+		// Provide the proof for the first entry
+		if err := trie.Prove(entries[start].k, proof); err != nil {
+			t.Fatalf("Failed to prove the first node %v", err)
+		}
+		// Now we have a proof: use it to initiate the stacktrie
+		stTrie, err := newStackTrieFromProof(trie.Hash(), entries[start].k, proof, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		// Obtain the hashes
+		hps, err := iterateProof(trie.Hash(), entries[start].k, false, proof)
+		if err != nil {
+			t.Fatal(err)
+		}
+		slices.Reverse(hps)
+		// Insert into stacktrie
+		for _, hp := range hps {
+			//fmt.Printf("%d. Adding hash/val %x: %x\n", i, hp.path, hp.hash)
+			stTrie.insert(stTrie.root, hp.path, hp.hash[:], nil, newHashed)
+		}
+		// Verify the final trie hash
+		if have := stTrie.Hash(); have != want {
+			t.Fatalf("wrong hash, have %x want %x\n", have, want)
+		}
+	}
+}

trie/stacktrie.go

@@ -59,7 +59,7 @@ func (t *StackTrie) Update(key, value []byte) error {
 	if len(value) == 0 {
 		panic("deletion not supported")
 	}
-	t.insert(t.root, k[:len(k)-1], value, nil)
+	t.insert(t.root, k[:len(k)-1], value, nil, newLeaf)
 	return nil
 }
 
@@ -84,6 +84,16 @@ type stNode struct {
 	children [16]*stNode // list of children (for branch and exts)
 }
 
+// newHashed constructs a hashed-node with provided value. The key
+// will be ignored, thus safe to modify afterwards, but
+// value is not safe to modify afterwards.
+func newHashed(key, val []byte) *stNode {
+	st := stPool.Get().(*stNode)
+	st.typ = hashedNode
+	st.val = val
+	return st
+}
+
 // newLeaf constructs a leaf node with provided node key and value. The key
 // will be deep-copied in the function and safe to modify afterwards, but
 // value is not.
@@ -138,7 +148,7 @@ func (n *stNode) getDiffIndex(key []byte) int {
 
 // Helper function to that inserts a (key, value) pair into
 // the trie.
-func (t *StackTrie) insert(st *stNode, key, value []byte, prefix []byte) {
+func (t *StackTrie) insert(st *stNode, key, value []byte, prefix []byte, leafCtor func([]byte, []byte) *stNode) {
 	switch st.typ {
 	case branchNode: /* Branch */
 		idx := int(key[0])
@@ -155,9 +165,9 @@ func (t *StackTrie) insert(st *stNode, key, value []byte, prefix []byte) {
 
 		// Add new child
 		if st.children[idx] == nil {
-			st.children[idx] = newLeaf(key[1:], value)
+			st.children[idx] = leafCtor(key[1:], value)
 		} else {
-			t.insert(st.children[idx], key[1:], value, append(prefix, key[0]))
+			t.insert(st.children[idx], key[1:], value, append(prefix, key[0]), leafCtor)
 		}
 
 	case extNode: /* Ext */
@@ -172,7 +182,7 @@ func (t *StackTrie) insert(st *stNode, key, value []byte, prefix []byte) {
 		if diffidx == len(st.key) {
 			// Ext key and key segment are identical, recurse into
 			// the child node.
-			t.insert(st.children[0], key[diffidx:], value, append(prefix, key[:diffidx]...))
+			t.insert(st.children[0], key[diffidx:], value, append(prefix, key[:diffidx]...), leafCtor)
 			return
 		}
 		// Save the original part. Depending if the break is
@@ -211,7 +221,7 @@ func (t *StackTrie) insert(st *stNode, key, value []byte, prefix []byte) {
 			p = st.children[0]
 		}
 		// Create a leaf for the inserted part
-		o := newLeaf(key[diffidx+1:], value)
+		o := leafCtor(key[diffidx+1:], value)
 
 		// Insert both child leaves where they belong:
 		origIdx := st.key[diffidx]
@@ -260,7 +270,7 @@ func (t *StackTrie) insert(st *stNode, key, value []byte, prefix []byte) {
 		t.hash(p.children[origIdx], append(prefix, st.key[:diffidx+1]...))
 
 		newIdx := key[diffidx]
-		p.children[newIdx] = newLeaf(key[diffidx+1:], value)
+		p.children[newIdx] = leafCtor(key[diffidx+1:], value)
 
 		// Finally, cut off the key part that has been passed
 		// over to the children.