trie: add merkle proof functions

1 files changed, 139 insertions, 0 deletions

diff --git a/trie/proof_test.go b/trie/proof_test.go
new file mode 100644
index 000000000..6b5bef05c
--- /dev/null
+++ b/trie/proof_test.go
@@ -0,0 +1,139 @@
+package trie
+
+import (
+	"bytes"
+	crand "crypto/rand"
+	mrand "math/rand"
+	"testing"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/rlp"
+)
+
+func init() {
+	mrand.Seed(time.Now().Unix())
+}
+
+func TestProof(t *testing.T) {
+	trie, vals := randomTrie(500)
+	root := trie.Hash()
+	for _, kv := range vals {
+		proof := trie.Prove(kv.k)
+		if proof == nil {
+			t.Fatalf("missing key %x while constructing proof", kv.k)
+		}
+		val, err := VerifyProof(root, kv.k, proof)
+		if err != nil {
+			t.Fatalf("VerifyProof error for key %x: %v\nraw proof: %x", kv.k, err, proof)
+		}
+		if !bytes.Equal(val, kv.v) {
+			t.Fatalf("VerifyProof returned wrong value for key %x: got %x, want %x", kv.k, val, kv.v)
+		}
+	}
+}
+
+func TestOneElementProof(t *testing.T) {
+	trie := new(Trie)
+	updateString(trie, "k", "v")
+	proof := trie.Prove([]byte("k"))
+	if proof == nil {
+		t.Fatal("nil proof")
+	}
+	if len(proof) != 1 {
+		t.Error("proof should have one element")
+	}
+	val, err := VerifyProof(trie.Hash(), []byte("k"), proof)
+	if err != nil {
+		t.Fatalf("VerifyProof error: %v\nraw proof: %x", err, proof)
+	}
+	if !bytes.Equal(val, []byte("v")) {
+		t.Fatalf("VerifyProof returned wrong value: got %x, want 'k'", val)
+	}
+}
+
+func TestVerifyBadProof(t *testing.T) {
+	trie, vals := randomTrie(800)
+	root := trie.Hash()
+	for _, kv := range vals {
+		proof := trie.Prove(kv.k)
+		if proof == nil {
+			t.Fatal("nil proof")
+		}
+		mutateByte(proof[mrand.Intn(len(proof))])
+		if _, err := VerifyProof(root, kv.k, proof); err == nil {
+			t.Fatalf("expected proof to fail for key %x", kv.k)
+		}
+	}
+}
+
+// mutateByte changes one byte in b.
+func mutateByte(b []byte) {
+	for r := mrand.Intn(len(b)); ; {
+		new := byte(mrand.Intn(255))
+		if new != b[r] {
+			b[r] = new
+			break
+		}
+	}
+}
+
+func BenchmarkProve(b *testing.B) {
+	trie, vals := randomTrie(100)
+	var keys []string
+	for k := range vals {
+		keys = append(keys, k)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		kv := vals[keys[i%len(keys)]]
+		if trie.Prove(kv.k) == nil {
+			b.Fatalf("nil proof for %x", kv.k)
+		}
+	}
+}
+
+func BenchmarkVerifyProof(b *testing.B) {
+	trie, vals := randomTrie(100)
+	root := trie.Hash()
+	var keys []string
+	var proofs [][]rlp.RawValue
+	for k := range vals {
+		keys = append(keys, k)
+		proofs = append(proofs, trie.Prove([]byte(k)))
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		im := i % len(keys)
+		if _, err := VerifyProof(root, []byte(keys[im]), proofs[im]); err != nil {
+			b.Fatalf("key %x: error", keys[im], err)
+		}
+	}
+}
+
+func randomTrie(n int) (*Trie, map[string]*kv) {
+	trie := new(Trie)
+	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 randBytes(n int) []byte {
+	r := make([]byte, n)
+	crand.Read(r)
+	return r
+}