Merge 8f28a7c799 into edaea69817

trie/encoding.go

@@ -16,13 +16,20 @@
 
 package trie
 
+// This file implements a codec to convert between byte arrays
+// and nibble (half-byte) arrays, where the latter are used for traversing the trie.
+// Hence the "compact" encoded form for a key is as the byte array, which
+// is (roughly) half the length of the decoded nibble array.
+// Nibble arrays for keys that represent leaf nodes have a terminator flag (numerical 16) appended to the end.
+// The compact encoded form uses Hex Prefix (HP) encoding, to encode the terminator status
+// and whether the key length is even or odd in the first two bytes (note the original description uses only
+// one byte, but that's inconvenient)
+
+// Encode a slice of nibbles into a HP byte array
 func CompactEncode(hexSlice []byte) []byte {
 	terminator := 0
 	if hexSlice[len(hexSlice)-1] == 16 {
 		terminator = 1
-	}
-
-	if terminator == 1 {
 		hexSlice = hexSlice[:len(hexSlice)-1]
 	}
 
@@ -42,12 +49,16 @@ func CompactEncode(hexSlice []byte) []byte {
 	return buf
 }
 
-func CompactDecode(str []byte) []byte {
+// Decode a HP encoded byte array into a nibble array
-	base := CompactHexDecode(str)
+// with terminator flag if applicable
+func CompactDecode(key []byte) []byte {
+	base := CompactHexDecode(key) // appends the terminator flag by default
+	if base[0] < 2 {
+		// remove the terminator flag if its not in the HP
 		base = base[:len(base)-1]
-	if base[0] >= 2 {
-		base = append(base, 16)
 	}
+
+	// HP tells us if key length is even or odd
 	if base[0]%2 == 1 {
 		base = base[1:]
 	} else {
@@ -57,10 +68,13 @@ func CompactDecode(str []byte) []byte {
 	return base
 }
 
-func CompactHexDecode(str []byte) []byte {
+// Decode a byte array into a nibble array.
-	l := len(str)*2 + 1
+// Assumes the key coressponds to a terminator node (ie appends 16)
+// CompactHexDecode is called immediately by the Get/Update/Remove functions.
+func CompactHexDecode(key []byte) []byte {
+	l := len(key)*2 + 1
 	var nibbles = make([]byte, l)
-	for i, b := range str {
+	for i, b := range key {
 		nibbles[i*2] = b / 16
 		nibbles[i*2+1] = b % 16
 	}
@@ -68,12 +82,13 @@ func CompactHexDecode(str []byte) []byte {
 	return nibbles
 }
 
-func DecodeCompact(key []byte) []byte {
+// This is really a compact encoding of nibbles
-	l := len(key) / 2
+// without hex-prefix
+func DecodeCompact(nibbles []byte) []byte {
+	l := len(nibbles) / 2
 	var res = make([]byte, l)
 	for i := 0; i < l; i++ {
-		v1, v0 := key[2*i], key[2*i+1]
+		res[i] = 16*nibbles[2*i] + nibbles[2*i+1]
-		res[i] = v1*16 + v0
 	}
 	return res
 }

trie/fullnode.go

@@ -16,6 +16,9 @@
 
 package trie
 
+// FullNode represents the main node type for a radix tree.
+// The first 16 children are branches for each possible next letter in the key.
+// The final slot is for a terminating node (a ValueNode), whose key ends at the FullNode.
 type FullNode struct {
 	trie  *Trie
 	nodes [17]Node

trie/hashnode.go

@@ -18,6 +18,10 @@ package trie
 
 import "github.com/ethereum/go-ethereum/common"
 
+// HashNode represents a node that is too big to be a ShortNode.
+// The actual node of interest (probably a FullNode) is rlp encoded and hashed,
+// yielding the key of the HashNode. The key can be used to fetch the actual
+// node from the database
 type HashNode struct {
 	key   []byte
 	trie  *Trie

trie/node.go

@@ -25,7 +25,7 @@ type Node interface {
 	Copy(*Trie) Node // All nodes, for now, return them self
 	Dirty() bool
 	fstring(string) string
-	Hash() interface{}
+	Hash() interface{} // only really a hash if size(node) > 32
 	RlpData() interface{}
 	setDirty(dirty bool)
 }

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/shortnode.go

@@ -18,9 +18,12 @@ package trie
 
 import "github.com/ethereum/go-ethereum/common"
 
+// ShortNode holds the nibble key for a ValueNode, HashNode, or FullNode.
+// Note a ShortNode should never hold another ShortNode, as they can just
+// be combined into one ShortNode
 type ShortNode struct {
 	trie  *Trie
-	key   []byte
+	key   []byte // hex-prefixed bytes
 	value Node
 	dirty bool
 }
@@ -30,7 +33,6 @@ func NewShortNode(t *Trie, key []byte, value Node) *ShortNode {
 }
 func (self *ShortNode) Value() Node {
 	self.value = self.trie.trans(self.value)
-
 	return self.value
 }
 func (self *ShortNode) Dirty() bool { return self.dirty }
@@ -48,6 +50,7 @@ func (self *ShortNode) Hash() interface{} {
 	return self.trie.store(self)
 }
 
+// returns nibbles
 func (self *ShortNode) Key() []byte {
 	return CompactDecode(self.key)
 }

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):])
 		}
 
@@ -358,6 +358,7 @@ func (self *Trie) mknode(value *common.Value) Node {
 	return NewValueNode(self, value.Bytes())
 }
 
+// resolve HashNodes by fetching from the db
 func (self *Trie) trans(node Node) Node {
 	switch node := node.(type) {
 	case *HashNode:
@@ -372,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)
+		}
+	}
+}

trie/valuenode.go

@@ -18,6 +18,8 @@ package trie
 
 import "github.com/ethereum/go-ethereum/common"
 
+// ValueNode represents a leaf node, a terminal point in the tree.
+// It is implied that if you found the ValueNode, then you know its key
 type ValueNode struct {
 	trie  *Trie
 	data  []byte