trie: merkle proofs

trie/proof.go

@@ -0,0 +1,182 @@
+package trie
+
+import (
+	"bytes"
+
+	"github.com/ethereum/go-ethereum/common"
+	//	"github.com/ethereum/go-ethereum/rlp"
+)
+
+/*
+A merkle proof for the trie consists of the RLP data for all nodes
+on the path from the root to the node of interest
+*/
+
+type TrieProof struct {
+	Key        []byte
+	Value      []byte
+	InnerNodes []ProofNode // ShortNode or FullNode
+	RootHash   []byte
+}
+
+func (proof *TrieProof) RlpData() interface{} {
+	return []interface{}{proof.Key, proof.Value, proof.InnerNodes, proof.RootHash}
+}
+
+// Prove a byte array is in the trie
+func (trie *Trie) Prove(key []byte) *TrieProof {
+	if trie == nil {
+		return nil
+	}
+	k := CompactHexDecode(key) // nibbles
+
+	rootHash := trie.Hash()
+	proof := &TrieProof{
+		Key:      key,
+		RootHash: rootHash,
+	}
+	// recursively appends nodes on the path from root to k to proof.InnerNodes
+	if exists := trie.constructProof(trie.root, k, proof); !exists {
+		return nil
+	}
+	return proof
+}
+
+// Verify the proof actually proves the given key byte-array is in the trie
+func (proof *TrieProof) Verify(key, value, rootHash []byte) bool {
+	if !bytes.Equal(key, proof.Key) {
+		return false
+	}
+	if !bytes.Equal(value, proof.Value) {
+		return false
+	}
+	if !bytes.Equal(rootHash, proof.RootHash) {
+		return false
+	}
+	// build up proof nodes into proper nodes in a dumy trie
+	trie := New(nil, nil)
+	nextNode := Node(NewValueNode(trie, value))
+	for _, proofNode := range proof.InnerNodes {
+		nextNode = linkProofNodes(trie, proofNode, nextNode)
+		if nextNode == nil {
+			return false
+		}
+	}
+	trie.root = nextNode
+	finalHash := trie.Hash()
+	return bytes.Equal(proof.RootHash, finalHash)
+}
+
+//--------------------------------------------------------------------------------
+
+// Proof node contains some rlp encoded data that can be decoded into a Node
+type ProofNode struct {
+	Key   []byte   // bytes
+	Nodes [][]byte `rlp:"nil"` // empty (ShortNode) or FullNode
+}
+
+func (proof ProofNode) RlpData() interface{} {
+	var t []interface{}
+	if proof.Nodes != nil {
+		t = make([]interface{}, 17)
+		for i, n := range proof.Nodes {
+			t[i] = n
+		}
+	}
+	return []interface{}{proof.Key, t}
+}
+
+// Create ShortNodefrom or FullNode. FullNode's should only have 17 entries in byte array
+func (proof ProofNode) TrieNode(trie *Trie) Node {
+	if len(proof.Nodes) == 0 {
+		return &ShortNode{trie: trie, key: proof.Key}
+	} else {
+		fullNode := NewFullNode(trie)
+		if len(proof.Nodes) != len(fullNode.nodes) {
+			return nil
+		}
+		for i, hash := range proof.Nodes {
+			fullNode.nodes[i] = trie.mknode(common.NewValueFromBytes(hash))
+		}
+		return fullNode
+	}
+}
+
+// Create proof node from full node by rlp encoding all children
+func FullNodeToProof(node *FullNode, key byte) ProofNode {
+	proofNode := ProofNode{Key: []byte{key}, Nodes: make([][]byte, 17)}
+	for i, n := range node.nodes {
+		if n == nil || i == int(key) { // don't store the node for the branch we're proving
+			proofNode.Nodes[i] = common.Encode("")
+		} else {
+			proofNode.Nodes[i] = common.Encode(n)
+		}
+	}
+	return proofNode
+}
+
+func ShortNodeToProof(node *ShortNode) ProofNode {
+	return ProofNode{Key: node.key, Nodes: nil}
+}
+
+//--------------------------------------------------------------------------------
+
+// key should be nibbles
+func (trie *Trie) constructProof(node Node, key []byte, proof *TrieProof) (exists bool) {
+	if node == nil {
+		return false
+	}
+	switch n := node.(type) {
+	case *ValueNode:
+		// NOTE: we should have made sure the key matches by now
+		// getting here == much success, very proof
+		proof.Value = n.Val()
+	case *HashNode:
+		// resolve the hash node and call constructProof
+		if exists := trie.constructProof(trie.trans(n), key, proof); !exists {
+			return false
+		}
+	case *ShortNode:
+		// chew off some key, constructProof on the value
+		if !bytes.HasPrefix(key, n.Key()) {
+			return false
+		}
+		k := key[len(n.Key()):]
+		if exists := trie.constructProof(n.Value(), k, proof); !exists {
+			return false
+		}
+		proof.InnerNodes = append(proof.InnerNodes, ShortNodeToProof(n))
+	case *FullNode:
+		// pick the right branch and carry on
+		if exists := trie.constructProof(n.branch(key[0]), key[1:], proof); !exists {
+			return false
+		}
+		proof.InnerNodes = append(proof.InnerNodes, FullNodeToProof(n, key[0]))
+	default:
+		panic("unknown node type")
+	}
+	return true
+}
+
+// fill in the nodes as we climb up the tree.
+//returns nil if full node has an empty key or too many branches
+func linkProofNodes(trie *Trie, proofNode ProofNode, node Node) Node {
+	nextNode := proofNode.TrieNode(trie)
+	if nextNode == nil {
+		return nil
+	}
+	switch nextNode := nextNode.(type) {
+	case *ShortNode:
+		nextNode.value = node
+		return nextNode
+	case *FullNode:
+		if len(proofNode.Key) == 0 {
+			return nil
+		}
+		nextNode.nodes[proofNode.Key[0]] = node
+		return nextNode
+	default:
+		panic("invalid proof node")
+	}
+	return nil
+}

trie/trie.go

@@ -239,7 +239,7 @@ func (self *Trie) get(node Node, key []byte) Node {
 		k := node.Key()
 		cnode := node.Value()
 
-		if len(key) >= len(k) && bytes.Equal(k, key[:len(k)]) {
+		if len(key) >= len(k) && bytes.HasPrefix(key, k) {
 			return self.get(cnode, key[len(k):])
 		}
 
@@ -373,7 +373,7 @@ func (self *Trie) store(node Node) interface{} {
 	data := common.Encode(node)
 	if len(data) >= 32 {
 		key := crypto.Sha3(data)
-		if node.Dirty() {
+		if node.Dirty() && self.cache != nil {
 			//fmt.Println("save", node)
 			//fmt.Println()
 			self.cache.Put(key, data)

trie/trie_test.go

@@ -18,11 +18,14 @@ package trie
 
 import (
 	"bytes"
+	crand "crypto/rand"
 	"fmt"
+	mrand "math/rand"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/rlp"
 )
 
 type Db map[string][]byte
@@ -352,3 +355,184 @@ func TestSecureDelete(t *testing.T) {
 		t.Errorf("expected %x got %x", exp, hash)
 	}
 }
+
+//------------------------------------------------------------------------------------
+// proof tests (and helpers)
+
+func randBytes(n int) []byte {
+	r := make([]byte, n)
+	crand.Read(r)
+	return r
+}
+
+func randInt(m int) int {
+	return int(mrand.Int31n(int32(m)))
+}
+
+// genuinely new byte
+func newByte(c byte) byte {
+	c2 := byte(randInt(255))
+	if c == c2 {
+		return newByte(c)
+	}
+	return c2
+}
+
+// genuinely change a byte
+func mutateBytes(b []byte) []byte {
+	b2 := make([]byte, len(b))
+	copy(b2, b)
+	b = b2
+
+	// Mutate a single byte
+	r := randInt(len(b))
+	c := b[r]
+	if c == byte(128) || c == byte(32) { // indeterminacy in rlp?
+		return mutateBytes(b)
+	}
+	d := newByte(c)
+	b[r] = d
+	return b
+}
+
+func makeTrieForProofs(n int) (*Trie, map[string]*kv) {
+	trie := NewEmpty()
+	vals := make(map[string]*kv)
+
+	for i := byte(0); i < 100; i++ {
+		value := &kv{common.LeftPadBytes([]byte{i}, 32), []byte{i}, false}
+		value2 := &kv{common.LeftPadBytes([]byte{i + 10}, 32), []byte{i}, false}
+		trie.Update(value.k, value.v)
+		trie.Update(value2.k, value2.v)
+		vals[string(value.k)] = value
+		vals[string(value2.k)] = value2
+	}
+
+	for i := 0; i < n; i++ {
+		value := &kv{randBytes(32), randBytes(20), false}
+		trie.Update(value.k, value.v)
+		vals[string(value.k)] = value
+	}
+
+	return trie, vals
+}
+
+func TestProof(t *testing.T) {
+	trie, vals := makeTrieForProofs(100)
+	// prove things are in the tree
+	for _, kvv := range vals {
+		proof := trie.Prove(kvv.k)
+		if proof == nil {
+			t.Fatalf("Failed to find key %X while constructing proof", kvv.k)
+		}
+		proven := proof.Verify(kvv.k, kvv.v, trie.Hash())
+		if !proven {
+			t.Fatalf("failed to prove key %X", kvv.k)
+		}
+	}
+}
+
+func testBadProof(t *testing.T, trie *Trie, kvv *kv, originalProofBytes []byte) {
+	proofBytes := mutateBytes(originalProofBytes)
+
+	proof2 := new(TrieProof)
+	if err := rlp.Decode(bytes.NewBuffer(proofBytes), proof2); err == nil {
+		proven := proof2.Verify(kvv.k, kvv.v, trie.Hash())
+		if proven {
+			t.Fatalf("expected proof to fail for %X", kvv.k)
+		}
+	} else {
+		// if we failed to decode, we mutated the rlp too badly.
+		// try again
+		testBadProof(t, trie, kvv, originalProofBytes)
+	}
+}
+
+func TestBadProof(t *testing.T) {
+	trie, vals := makeTrieForProofs(100)
+	for _, kvv := range vals {
+		proof := trie.Prove(kvv.k)
+		proven := proof.Verify(kvv.k, kvv.v, trie.Hash())
+		if !proven {
+			t.Fatalf("expected proof not to fail for %X", kvv.k)
+		}
+		proofBytes := common.Encode(proof)
+		testBadProof(t, trie, kvv, proofBytes)
+	}
+}
+
+func compareProofs(proof, proof2 *TrieProof) error {
+	if !bytes.Equal(proof.Key, proof2.Key) {
+		return fmt.Errorf("codec error: keys are not same. got %X, expected %X\n", proof2.Key, proof.Key)
+	}
+	if !bytes.Equal(proof.Value, proof2.Value) {
+		return fmt.Errorf("codec error: values are not same. got %X, expected %X\n", proof2.Value, proof.Value)
+	}
+	if !bytes.Equal(proof.RootHash, proof2.RootHash) {
+		return fmt.Errorf("codec error: root hashes are not same. got %X, expected %X\n", proof2.RootHash, proof.RootHash)
+	}
+	if len(proof.InnerNodes) != len(proof2.InnerNodes) {
+		return fmt.Errorf("codec error: wrong number of inner nodes. got %d, expected %d\n", len(proof2.InnerNodes), len(proof.InnerNodes))
+	}
+
+	for i, in := range proof.InnerNodes {
+		in2 := proof2.InnerNodes[i]
+		if !bytes.Equal(in.Key, in2.Key) {
+			return fmt.Errorf("codec error: inner keys for node %d are not same. got %X, expected %X\n", i, in2.Key, in.Key)
+		}
+		for j, n := range in.Nodes {
+			if !bytes.Equal(n, in2.Nodes[j]) {
+				return fmt.Errorf("codec error: inner nodes %d (%d) are not same. got %X, expected %X\n", i, j, in2.Nodes[j], n)
+			}
+		}
+	}
+	return nil
+}
+
+func TestProofCodec(t *testing.T) {
+	trie, vals := makeTrieForProofs(100)
+	for _, kvv := range vals {
+		proof := trie.Prove(kvv.k)
+		proofBytes := common.Encode(proof)
+
+		proof2 := new(TrieProof)
+		if err := rlp.Decode(bytes.NewBuffer(proofBytes), proof2); err != nil {
+			t.Fatalf("error decoding proof bytes: %v", err)
+		}
+
+		if err := compareProofs(proof, proof2); err != nil {
+			t.Fatal(err)
+		}
+
+		proven := proof2.Verify(kvv.k, kvv.v, trie.Hash())
+		if !proven {
+			t.Fatalf("failed to prove key %X", kvv.k)
+		}
+	}
+}
+
+func BenchmarkProof(b *testing.B) {
+	trie, vals := makeTrieForProofs(100)
+
+	proofs := make([]*TrieProof, len(vals))
+
+	i := 0
+	for _, kvv := range vals {
+		proof := trie.Prove(kvv.k)
+		if proof == nil {
+			b.Fatalf("Failed to find key %X while constructing proof", kvv.k)
+		}
+		proofs[i] = proof
+		i += 1
+	}
+
+	N := len(vals)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		im := i % N
+		v := proofs[im].Verify(proofs[im].Key, proofs[im].Value, proofs[im].RootHash)
+		if !v {
+			b.Fatalf("failed to prove key %X", proofs[im].Key)
+		}
+	}
+}