diff --git a/trie/proof.go b/trie/proof.go new file mode 100644 index 0000000000..4e68605393 --- /dev/null +++ b/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 +} diff --git a/trie/trie.go b/trie/trie.go index e08fb3bf6e..71e014ca3a 100644 --- a/trie/trie.go +++ b/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) diff --git a/trie/trie_test.go b/trie/trie_test.go index ae4e5efe40..a90f65e86c 100644 --- a/trie/trie_test.go +++ b/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) + } + } +}