light, trie: use NodeSet and NodeList for Merkle proofs

les/handler.go

@@ -35,6 +35,7 @@ import (
 	"github.com/ethereum/go-ethereum/eth/downloader"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/event"
+	"github.com/ethereum/go-ethereum/light"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/discover"
@@ -690,9 +691,10 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 						}
 					}
 					if tr != nil {
-						proof := tr.Prove(req.Key)
+						var proof light.NodeList
+						tr.Prove(req.Key, 0, &proof)
 						proofs = append(proofs, proof)
-						bytes += len(proof)
+						bytes += proof.DataSize()
 					}
 				}
 			}
@@ -751,9 +753,10 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 					if tr, _ := trie.New(root, pm.chainDb); tr != nil {
 						var encNumber [8]byte
 						binary.BigEndian.PutUint64(encNumber[:], req.BlockNum)
-						proof := tr.Prove(encNumber[:])
+						var proof light.NodeList
+						tr.Prove(encNumber[:], 0, &proof)
 						proofs = append(proofs, ChtResp{Header: header, Proof: proof})
-						bytes += len(proof) + estHeaderRlpSize
+						bytes += proof.DataSize() + estHeaderRlpSize
 					}
 				}
 			}

les/handler_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/eth/downloader"
 	"github.com/ethereum/go-ethereum/ethdb"
+	"github.com/ethereum/go-ethereum/light"
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
@@ -331,7 +332,8 @@ func testGetProofs(t *testing.T, protocol int) {
 			}
 			proofreqs = append(proofreqs, req)
 
-			proof := trie.Prove(crypto.Keccak256(acc[:]))
+			var proof light.NodeList
+			trie.Prove(crypto.Keccak256(acc[:]), 0, &proof)
 			proofs = append(proofs, proof)
 		}
 	}

les/odr_requests.go

@@ -220,7 +220,7 @@ func (r *TrieRequest) Validate(db ethdb.Database, msg *Msg) error {
 		return errMultipleEntries
 	}
 	// Verify the proof and store if checks out
-	if _, err := trie.VerifyProof(r.Id.Root, r.Key, proofs[0]); err != nil {
+	if _, err, _ := trie.VerifyProof(r.Id.Root, r.Key, light.NodeList(proofs[0]).NodeSet()); err != nil {
 		return fmt.Errorf("merkle proof verification failed: %v", err)
 	}
 	r.Proof = proofs[0]
@@ -336,7 +336,7 @@ func (r *ChtRequest) Validate(db ethdb.Database, msg *Msg) error {
 	var encNumber [8]byte
 	binary.BigEndian.PutUint64(encNumber[:], r.BlockNum)
 
-	value, err := trie.VerifyProof(r.ChtRoot, encNumber[:], proof.Proof)
+	value, err, _ := trie.VerifyProof(r.ChtRoot, encNumber[:], light.NodeList(proof.Proof).NodeSet())
 	if err != nil {
 		return err
 	}

light/nodeset.go

@@ -0,0 +1,173 @@
+// Copyright 2014 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package light
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/ethereum/go-ethereum/trie"
+)
+
+// NodeSet stores a set of trie nodes. It implements trie.Database and can also
+// act as a cache for another trie.Database.
+type NodeSet struct {
+	db                                map[string][]byte
+	dataSize                          int
+	lock                              sync.RWMutex
+	fallback                          trie.Database
+	copyFromFallback, writeToFallback bool
+}
+
+// NewNodeSet creates an empty node set
+func NewNodeSet() *NodeSet {
+	return &NodeSet{
+		db: make(map[string][]byte),
+	}
+}
+
+// SetFallback will add a fallback database, making this node set a cache for the backing database.
+// If copyFromFallback is true, it keeps any node it fetches from the fallback database.
+// If writeToFallback is true, it writes stored nodes to the fallback database too.
+func (db *NodeSet) SetFallback(fallback trie.Database, copyFromFallback, writeToFallback bool) {
+	db.lock.Lock()
+	defer db.lock.Unlock()
+
+	db.fallback = fallback
+	db.copyFromFallback = copyFromFallback
+	db.writeToFallback = writeToFallback
+}
+
+// ReadCache returns a new read cache (copyFromFallback=true) for this node set
+func (db *NodeSet) ReadCache() *NodeSet {
+	cdb := NewNodeSet()
+	cdb.SetFallback(db, true, false)
+	return cdb
+}
+
+// Put stores a new node in the set
+func (db *NodeSet) Put(key []byte, value []byte) error {
+	db.lock.Lock()
+	defer db.lock.Unlock()
+
+	if _, ok := db.db[string(key)]; !ok {
+		db.db[string(key)] = common.CopyBytes(value)
+		db.dataSize += len(value)
+		if db.writeToFallback && db.fallback != nil {
+			db.fallback.Put(key, value)
+		}
+	}
+	return nil
+}
+
+// Get returns a stored node
+func (db *NodeSet) Get(key []byte) ([]byte, error) {
+	db.lock.RLock()
+	defer db.lock.RUnlock()
+
+	if entry, ok := db.db[string(key)]; ok {
+		return entry, nil
+	}
+	if db.fallback != nil {
+		value, err := db.fallback.Get(key)
+		if db.copyFromFallback && err == nil {
+			db.db[string(key)] = value
+			db.dataSize += len(value)
+		}
+		return value, err
+	}
+	return nil, errors.New("not found")
+}
+
+// Has returns true if the node set contains the given key
+func (db *NodeSet) Has(key []byte) (bool, error) {
+	_, err := db.Get(key)
+	return err == nil, nil
+}
+
+// KeyCount returns the number of nodes in the set
+func (db *NodeSet) KeyCount() int {
+	db.lock.RLock()
+	defer db.lock.RUnlock()
+
+	return len(db.db)
+}
+
+// DataSize returns the aggregated data size of nodes in the set
+func (db *NodeSet) DataSize() int {
+	db.lock.RLock()
+	defer db.lock.RUnlock()
+
+	return db.dataSize
+}
+
+// NodeList converts the node set to a NodeList
+func (db *NodeSet) NodeList() NodeList {
+	db.lock.RLock()
+	defer db.lock.RUnlock()
+
+	var values NodeList
+	for _, value := range db.db {
+		values = append(values, value)
+	}
+	return values
+}
+
+// Store writes the contents of the set to the given database
+func (db *NodeSet) Store(target trie.Database) {
+	db.lock.RLock()
+	defer db.lock.RUnlock()
+
+	for key, value := range db.db {
+		target.Put([]byte(key), value)
+	}
+}
+
+// NodeList stores an ordered list of trie nodes. It implements trie.DatabaseWriter.
+type NodeList []rlp.RawValue
+
+// Store writes the contents of the list to the given database
+func (n NodeList) Store(db trie.Database) {
+	for _, node := range n {
+		db.Put(crypto.Keccak256(node), node)
+	}
+}
+
+// NodeSet converts the node list to a NodeSet
+func (n NodeList) NodeSet() *NodeSet {
+	db := NewNodeSet()
+	n.Store(db)
+	return db
+}
+
+// Put stores a new node at the end of the list
+func (n *NodeList) Put(key []byte, value []byte) error {
+	*n = append(*n, value)
+	return nil
+}
+
+// DataSize returns the aggregated data size of nodes in the list
+func (n NodeList) DataSize() int {
+	var size int
+	for _, node := range n {
+		size += len(node)
+	}
+	return size
+}

light/odr_test.go

@@ -77,7 +77,9 @@ func (odr *testOdr) Retrieve(ctx context.Context, req OdrRequest) error {
 		req.Receipts = core.GetBlockReceipts(odr.sdb, req.Hash, core.GetBlockNumber(odr.sdb, req.Hash))
 	case *TrieRequest:
 		t, _ := trie.New(req.Id.Root, odr.sdb)
-		req.Proof = t.Prove(req.Key)
+		nodes := NewNodeSet()
+		t.Prove(req.Key, 0, nodes)
+		req.Proof = nodes.NodeList()
 	case *CodeRequest:
 		req.Data, _ = odr.sdb.Get(req.Hash[:])
 	}

trie/proof.go

@@ -18,11 +18,10 @@ package trie
 
 import (
 	"bytes"
-	"errors"
 	"fmt"
 
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/crypto/sha3"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rlp"
 )
@@ -36,7 +35,7 @@ import (
 // contains all nodes of the longest existing prefix of the key
 // (at least the root node), ending with the node that proves the
 // absence of the key.
-func (t *Trie) Prove(key []byte) []rlp.RawValue {
+func (t *Trie) Prove(key []byte, fromLevel uint, proofDb DatabaseWriter) error {
 	// Collect all nodes on the path to key.
 	key = keybytesToHex(key)
 	nodes := []node{}
@@ -61,67 +60,63 @@ func (t *Trie) Prove(key []byte) []rlp.RawValue {
 			tn, err = t.resolveHash(n, nil)
 			if err != nil {
 				log.Error(fmt.Sprintf("Unhandled trie error: %v", err))
-				return nil
+				return err
 			}
 		default:
 			panic(fmt.Sprintf("%T: invalid node: %v", tn, tn))
 		}
 	}
 	hasher := newHasher(0, 0)
-	proof := make([]rlp.RawValue, 0, len(nodes))
 	for i, n := range nodes {
 		// Don't bother checking for errors here since hasher panics
 		// if encoding doesn't work and we're not writing to any database.
 		n, _, _ = hasher.hashChildren(n, nil)
 		hn, _ := hasher.store(n, nil, false)
-		if _, ok := hn.(hashNode); ok || i == 0 {
+		if hash, ok := hn.(hashNode); ok || i == 0 {
 			// If the node's database encoding is a hash (or is the
 			// root node), it becomes a proof element.
-			enc, _ := rlp.EncodeToBytes(n)
+			if fromLevel > 0 {
-			proof = append(proof, enc)
+				fromLevel--
+			} else {
+				enc, _ := rlp.EncodeToBytes(n)
+				if !ok {
+					hash = crypto.Keccak256(enc)
+				}
+				proofDb.Put(hash, enc)
+			}
 		}
 	}
-	return proof
+	return nil
 }
 
 // VerifyProof checks merkle proofs. The given proof must contain the
 // value for key in a trie with the given root hash. VerifyProof
 // returns an error if the proof contains invalid trie nodes or the
 // wrong value.
-func VerifyProof(rootHash common.Hash, key []byte, proof []rlp.RawValue) (value []byte, err error) {
+func VerifyProof(rootHash common.Hash, key []byte, proofDb DatabaseReader) (value []byte, err error, nodes int) {
 	key = keybytesToHex(key)
-	sha := sha3.NewKeccak256()
+	wantHash := rootHash[:]
-	wantHash := rootHash.Bytes()
+	for i := 0; ; i++ {
-	for i, buf := range proof {
+		buf, _ := proofDb.Get(wantHash)
-		sha.Reset()
+		if buf == nil {
-		sha.Write(buf)
+			return nil, fmt.Errorf("proof node %d (hash %064x) missing", i, wantHash[:]), i
-		if !bytes.Equal(sha.Sum(nil), wantHash) {
-			return nil, fmt.Errorf("bad proof node %d: hash mismatch", i)
 		}
 		n, err := decodeNode(wantHash, buf, 0)
 		if err != nil {
-			return nil, fmt.Errorf("bad proof node %d: %v", i, err)
+			return nil, fmt.Errorf("bad proof node %d: %v", i, err), i
 		}
 		keyrest, cld := get(n, key)
 		switch cld := cld.(type) {
 		case nil:
-			if i != len(proof)-1 {
+			// The trie doesn't contain the key.
-				return nil, fmt.Errorf("key mismatch at proof node %d", i)
+			return nil, nil, i
-			} else {
-				// The trie doesn't contain the key.
-				return nil, nil
-			}
 		case hashNode:
 			key = keyrest
 			wantHash = cld
 		case valueNode:
-			if i != len(proof)-1 {
+			return cld, nil, i + 1
-				return nil, errors.New("additional nodes at end of proof")
-			}
-			return cld, nil
 		}
 	}
-	return nil, errors.New("unexpected end of proof")
 }
 
 func get(tn node, key []byte) ([]byte, node) {

trie/proof_test.go

@@ -19,12 +19,13 @@ package trie
 import (
 	"bytes"
 	crand "crypto/rand"
+	"errors"
 	mrand "math/rand"
 	"testing"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/ethereum/go-ethereum/crypto"
 )
 
 func init() {
@@ -35,13 +36,13 @@ func TestProof(t *testing.T) {
 	trie, vals := randomTrie(500)
 	root := trie.Hash()
 	for _, kv := range vals {
-		proof := trie.Prove(kv.k)
+		proofs := newTestProofDb()
-		if proof == nil {
+		if trie.Prove(kv.k, 0, proofs) != nil {
 			t.Fatalf("missing key %x while constructing proof", kv.k)
 		}
-		val, err := VerifyProof(root, kv.k, proof)
+		val, err, _ := VerifyProof(root, kv.k, proofs)
 		if err != nil {
-			t.Fatalf("VerifyProof error for key %x: %v\nraw proof: %x", kv.k, err, proof)
+			t.Fatalf("VerifyProof error for key %x: %v\nraw proof: %v", kv.k, err, proofs)
 		}
 		if !bytes.Equal(val, kv.v) {
 			t.Fatalf("VerifyProof returned wrong value for key %x: got %x, want %x", kv.k, val, kv.v)
@@ -52,16 +53,14 @@ func TestProof(t *testing.T) {
 func TestOneElementProof(t *testing.T) {
 	trie := new(Trie)
 	updateString(trie, "k", "v")
-	proof := trie.Prove([]byte("k"))
+	proofs := newTestProofDb()
-	if proof == nil {
+	trie.Prove([]byte("k"), 0, proofs)
-		t.Fatal("nil proof")
+	if len(proofs.db) != 1 {
-	}
-	if len(proof) != 1 {
 		t.Error("proof should have one element")
 	}
-	val, err := VerifyProof(trie.Hash(), []byte("k"), proof)
+	val, err, _ := VerifyProof(trie.Hash(), []byte("k"), proofs)
 	if err != nil {
-		t.Fatalf("VerifyProof error: %v\nraw proof: %x", err, proof)
+		t.Fatalf("VerifyProof error: %v\nraw proof: %v", err, proofs.db)
 	}
 	if !bytes.Equal(val, []byte("v")) {
 		t.Fatalf("VerifyProof returned wrong value: got %x, want 'k'", val)
@@ -72,12 +71,22 @@ func TestVerifyBadProof(t *testing.T) {
 	trie, vals := randomTrie(800)
 	root := trie.Hash()
 	for _, kv := range vals {
-		proof := trie.Prove(kv.k)
+		proofs := newTestProofDb()
-		if proof == nil {
+		trie.Prove(kv.k, 0, proofs)
-			t.Fatal("nil proof")
+		if len(proofs.db) == 0 {
+			t.Fatal("zero length proof")
 		}
-		mutateByte(proof[mrand.Intn(len(proof))])
+		idx := mrand.Intn(len(proofs.db))
-		if _, err := VerifyProof(root, kv.k, proof); err == nil {
+		for key, node := range proofs.db {
+			if idx == 0 {
+				delete(proofs.db, key)
+				mutateByte(node)
+				proofs.Put(crypto.Keccak256(node), node)
+				break
+			}
+			idx--
+		}
+		if _, err, _ := VerifyProof(root, kv.k, proofs); err == nil {
 			t.Fatalf("expected proof to fail for key %x", kv.k)
 		}
 	}
@@ -104,8 +113,9 @@ func BenchmarkProve(b *testing.B) {
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		kv := vals[keys[i%len(keys)]]
-		if trie.Prove(kv.k) == nil {
+		proofs := newTestProofDb()
-			b.Fatalf("nil proof for %x", kv.k)
+		if trie.Prove(kv.k, 0, proofs); len(proofs.db) == 0 {
+			b.Fatalf("zero length proof for %x", kv.k)
 		}
 	}
 }
@@ -114,16 +124,18 @@ func BenchmarkVerifyProof(b *testing.B) {
 	trie, vals := randomTrie(100)
 	root := trie.Hash()
 	var keys []string
-	var proofs [][]rlp.RawValue
+	var proofs []*testProofDb
 	for k := range vals {
 		keys = append(keys, k)
-		proofs = append(proofs, trie.Prove([]byte(k)))
+		proof := newTestProofDb()
+		trie.Prove([]byte(k), 0, proof)
+		proofs = append(proofs, proof)
 	}
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		im := i % len(keys)
-		if _, err := VerifyProof(root, []byte(keys[im]), proofs[im]); err != nil {
+		if _, err, _ := VerifyProof(root, []byte(keys[im]), proofs[im]); err != nil {
 			b.Fatalf("key %x: %v", keys[im], err)
 		}
 	}
@@ -153,3 +165,30 @@ func randBytes(n int) []byte {
 	crand.Read(r)
 	return r
 }
+
+type testProofDb struct {
+	db map[string][]byte
+}
+
+func newTestProofDb() *testProofDb {
+	return &testProofDb{
+		db: make(map[string][]byte),
+	}
+}
+
+func (db *testProofDb) Put(key []byte, value []byte) error {
+	db.db[string(key)] = common.CopyBytes(value)
+	return nil
+}
+
+func (db *testProofDb) Get(key []byte) ([]byte, error) {
+	if entry, ok := db.db[string(key)]; ok {
+		return entry, nil
+	}
+	return nil, errors.New("not found")
+}
+
+func (db *testProofDb) Has(key []byte) (bool, error) {
+	_, err := db.Get(key)
+	return err == nil, nil
+}