core, les, light, trie: correctly use the memdb for tries

core/state/database.go

@@ -21,6 +21,7 @@ import (
 	"sync"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/trie"
 	lru "github.com/hashicorp/golang-lru"
 )
@@ -67,6 +68,7 @@ type Trie interface {
 	Hash() common.Hash
 	NodeIterator(startKey []byte) trie.NodeIterator
 	GetKey([]byte) []byte // TODO(fjl): remove this when SecureTrie is removed
+	Prove(key []byte, fromLevel uint, proofDb ethdb.Putter) error
 }
 
 // NewDatabase creates a backing store for state. The returned database is safe for
@@ -173,3 +175,7 @@ func (m cachedTrie) Commit(onleaf trie.LeafCallback) (common.Hash, error) {
 	}
 	return root, err
 }
+
+func (m cachedTrie) Prove(key []byte, fromLevel uint, proofDb ethdb.Putter) error {
+	return m.SecureTrie.Prove(key, fromLevel, proofDb)
+}

core/state/statedb.go

@@ -243,6 +243,11 @@ func (self *StateDB) GetState(a common.Address, b common.Hash) common.Hash {
 	return common.Hash{}
 }
 
+// Database retrieves the low level database supporting the lower level trie ops.
+func (self *StateDB) Database() Database {
+	return self.db
+}
+
 // StorageTrie returns the storage trie of an account.
 // The return value is a copy and is nil for non-existent accounts.
 func (self *StateDB) StorageTrie(a common.Address) Trie {

les/handler.go

@@ -18,7 +18,6 @@
 package les
 
 import (
-	"bytes"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -78,6 +77,7 @@ type BlockChain interface {
 	GetHeaderByHash(hash common.Hash) *types.Header
 	CurrentHeader() *types.Header
 	GetTd(hash common.Hash, number uint64) *big.Int
+	State() (*state.StateDB, error)
 	InsertHeaderChain(chain []*types.Header, checkFreq int) (int, error)
 	Rollback(chain []common.Hash)
 	GetHeaderByNumber(number uint64) *types.Header
@@ -579,18 +579,20 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 		for _, req := range req.Reqs {
 			// Retrieve the requested state entry, stopping if enough was found
 			if header := core.GetHeader(pm.chainDb, req.BHash, core.GetBlockNumber(pm.chainDb, req.BHash)); header != nil {
-				if trie, _ := trie.New(header.Root, trie.NewDatabase(pm.chainDb)); trie != nil {
+				statedb, err := pm.blockchain.State()
-					sdata := trie.Get(req.AccKey)
+				if err != nil {
-					var acc state.Account
+					continue
-					if err := rlp.DecodeBytes(sdata, &acc); err == nil {
+				}
-						entry, _ := pm.chainDb.Get(acc.CodeHash)
+				account, err := pm.getAccount(statedb, header.Root, common.BytesToHash(req.AccKey))
-						if bytes+len(entry) >= softResponseLimit {
+				if err != nil {
+					continue
+				}
+				code, _ := statedb.Database().TrieDB().Node(common.BytesToHash(account.CodeHash))
+
+				data = append(data, code)
+				if bytes += len(code); bytes >= softResponseLimit {
 					break
 				}
-						data = append(data, entry)
-						bytes += len(entry)
-					}
-				}
 			}
 		}
 		bv, rcost := p.fcClient.RequestProcessed(costs.baseCost + uint64(reqCnt)*costs.reqCost)
@@ -701,25 +703,29 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 			return errResp(ErrRequestRejected, "")
 		}
 		for _, req := range req.Reqs {
-			if bytes >= softResponseLimit {
-				break
-			}
 			// Retrieve the requested state entry, stopping if enough was found
 			if header := core.GetHeader(pm.chainDb, req.BHash, core.GetBlockNumber(pm.chainDb, req.BHash)); header != nil {
-				if tr, _ := trie.New(header.Root, trie.NewDatabase(pm.chainDb)); tr != nil {
+				statedb, err := pm.blockchain.State()
+				if err != nil {
+					continue
+				}
+				var trie state.Trie
 				if len(req.AccKey) > 0 {
-						sdata := tr.Get(req.AccKey)
+					account, err := pm.getAccount(statedb, header.Root, common.BytesToHash(req.AccKey))
-						tr = nil
+					if err != nil {
-						var acc state.Account
+						continue
-						if err := rlp.DecodeBytes(sdata, &acc); err == nil {
-							tr, _ = trie.New(acc.Root, trie.NewDatabase(pm.chainDb))
 					}
+					trie, _ = statedb.Database().OpenStorageTrie(common.BytesToHash(req.AccKey), account.Root)
+				} else {
+					trie, _ = statedb.Database().OpenTrie(header.Root)
 				}
-					if tr != nil {
+				if trie != nil {
 					var proof light.NodeList
-						tr.Prove(req.Key, 0, &proof)
+					trie.Prove(req.Key, 0, &proof)
+
 					proofs = append(proofs, proof)
-						bytes += proof.DataSize()
+					if bytes += proof.DataSize(); bytes >= softResponseLimit {
+						break
 					}
 				}
 			}
@@ -741,8 +747,8 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 		// Gather state data until the fetch or network limits is reached
 		var (
 			lastBHash common.Hash
-			lastAccKey []byte
+			statedb   *state.StateDB
-			tr, str    *trie.Trie
+			root      common.Hash
 		)
 		reqCnt := len(req.Reqs)
 		if reject(uint64(reqCnt), MaxProofsFetch) {
@@ -752,36 +758,37 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 		nodes := light.NewNodeSet()
 
 		for _, req := range req.Reqs {
+			// Look up the state belonging to the request
+			if statedb == nil || req.BHash != lastBHash {
+				statedb, root, lastBHash = nil, common.Hash{}, req.BHash
+
+				if header := core.GetHeader(pm.chainDb, req.BHash, core.GetBlockNumber(pm.chainDb, req.BHash)); header != nil {
+					statedb, _ = pm.blockchain.State()
+					root = header.Root
+				}
+			}
+			if statedb == nil {
+				continue
+			}
+			// Pull the account or storage trie of the request
+			var trie state.Trie
+			if len(req.AccKey) > 0 {
+				account, err := pm.getAccount(statedb, root, common.BytesToHash(req.AccKey))
+				if err != nil {
+					continue
+				}
+				trie, _ = statedb.Database().OpenStorageTrie(common.BytesToHash(req.AccKey), account.Root)
+			} else {
+				trie, _ = statedb.Database().OpenTrie(root)
+			}
+			if trie == nil {
+				continue
+			}
+			// Prove the user's request from the account or stroage trie
+			trie.Prove(req.Key, req.FromLevel, nodes)
 			if nodes.DataSize() >= softResponseLimit {
 				break
 			}
-			if tr == nil || req.BHash != lastBHash {
-				if header := core.GetHeader(pm.chainDb, req.BHash, core.GetBlockNumber(pm.chainDb, req.BHash)); header != nil {
-					tr, _ = trie.New(header.Root, trie.NewDatabase(pm.chainDb))
-				} else {
-					tr = nil
-				}
-				lastBHash = req.BHash
-				str = nil
-			}
-			if tr != nil {
-				if len(req.AccKey) > 0 {
-					if str == nil || !bytes.Equal(req.AccKey, lastAccKey) {
-						sdata := tr.Get(req.AccKey)
-						str = nil
-						var acc state.Account
-						if err := rlp.DecodeBytes(sdata, &acc); err == nil {
-							str, _ = trie.New(acc.Root, trie.NewDatabase(pm.chainDb))
-						}
-						lastAccKey = common.CopyBytes(req.AccKey)
-					}
-					if str != nil {
-						str.Prove(req.Key, req.FromLevel, nodes)
-					}
-				} else {
-					tr.Prove(req.Key, req.FromLevel, nodes)
-				}
-			}
 		}
 		proofs := nodes.NodeList()
 		bv, rcost := p.fcClient.RequestProcessed(costs.baseCost + uint64(reqCnt)*costs.reqCost)
@@ -849,23 +856,29 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 		if reject(uint64(reqCnt), MaxHelperTrieProofsFetch) {
 			return errResp(ErrRequestRejected, "")
 		}
-		trieDb := ethdb.NewTable(pm.chainDb, light.ChtTablePrefix)
 		for _, req := range req.Reqs {
-			if bytes >= softResponseLimit {
-				break
-			}
-
 			if header := pm.blockchain.GetHeaderByNumber(req.BlockNum); header != nil {
 				sectionHead := core.GetCanonicalHash(pm.chainDb, req.ChtNum*light.ChtV1Frequency-1)
 				if root := light.GetChtRoot(pm.chainDb, req.ChtNum-1, sectionHead); root != (common.Hash{}) {
-					if tr, _ := trie.New(root, trie.NewDatabase(trieDb)); tr != nil {
+					statedb, err := pm.blockchain.State()
+					if err != nil {
+						continue
+					}
+					trie, err := statedb.Database().OpenTrie(root)
+					if err != nil {
+						continue
+					}
 					var encNumber [8]byte
 					binary.BigEndian.PutUint64(encNumber[:], req.BlockNum)
+
 					var proof light.NodeList
-						tr.Prove(encNumber[:], 0, &proof)
+					trie.Prove(encNumber[:], 0, &proof)
+
 					proofs = append(proofs, ChtResp{Header: header, Proof: proof})
-						bytes += proof.DataSize() + estHeaderRlpSize
+					if bytes += proof.DataSize() + estHeaderRlpSize; bytes >= softResponseLimit {
+						break
 					}
+
 				}
 			}
 		}
@@ -897,25 +910,21 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 			lastIdx  uint64
 			lastType uint
 			root     common.Hash
-			tr       *trie.Trie
+			statedb  *state.StateDB
+			trie     state.Trie
 		)
 
 		nodes := light.NewNodeSet()
 
 		for _, req := range req.Reqs {
-			if nodes.DataSize()+auxBytes >= softResponseLimit {
+			if trie == nil || req.HelperTrieType != lastType || req.TrieIdx != lastIdx {
-				break
+				statedb, trie, lastType, lastIdx = nil, nil, req.HelperTrieType, req.TrieIdx
-			}
+
-			if tr == nil || req.HelperTrieType != lastType || req.TrieIdx != lastIdx {
+				if root, _ = pm.getHelperTrie(req.HelperTrieType, req.TrieIdx); root != (common.Hash{}) {
-				var prefix string
+					if statedb, _ = pm.blockchain.State(); statedb != nil {
-				root, prefix = pm.getHelperTrie(req.HelperTrieType, req.TrieIdx)
+						trie, _ = statedb.Database().OpenTrie(root)
-				if root != (common.Hash{}) {
-					if t, err := trie.New(root, trie.NewDatabase(ethdb.NewTable(pm.chainDb, prefix))); err == nil {
-						tr = t
 					}
 				}
-				lastType = req.HelperTrieType
-				lastIdx = req.TrieIdx
 			}
 			if req.AuxReq == auxRoot {
 				var data []byte
@@ -925,8 +934,8 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 				auxData = append(auxData, data)
 				auxBytes += len(data)
 			} else {
-				if tr != nil {
+				if trie != nil {
-					tr.Prove(req.Key, req.FromLevel, nodes)
+					trie.Prove(req.Key, req.FromLevel, nodes)
 				}
 				if req.AuxReq != 0 {
 					data := pm.getHelperTrieAuxData(req)
@@ -934,6 +943,9 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 					auxBytes += len(data)
 				}
 			}
+			if nodes.DataSize()+auxBytes >= softResponseLimit {
+				break
+			}
 		}
 		proofs := nodes.NodeList()
 		bv, rcost := p.fcClient.RequestProcessed(costs.baseCost + uint64(reqCnt)*costs.reqCost)
@@ -1090,6 +1102,23 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 	return nil
 }
 
+// getAccount retrieves an account from the state based at root.
+func (pm *ProtocolManager) getAccount(statedb *state.StateDB, root, hash common.Hash) (state.Account, error) {
+	trie, err := trie.New(root, statedb.Database().TrieDB())
+	if err != nil {
+		return state.Account{}, err
+	}
+	blob, err := trie.TryGet(hash[:])
+	if err != nil {
+		return state.Account{}, err
+	}
+	var account state.Account
+	if err = rlp.DecodeBytes(blob, &account); err != nil {
+		return state.Account{}, err
+	}
+	return account, nil
+}
+
 // getHelperTrie returns the post-processed trie root for the given trie ID and section index
 func (pm *ProtocolManager) getHelperTrie(id uint, idx uint64) (common.Hash, string) {
 	switch id {

les/odr_test.go

@@ -102,7 +102,6 @@ func odrAccounts(ctx context.Context, db ethdb.Database, config *params.ChainCon
 			res = append(res, rlp...)
 		}
 	}
-
 	return res
 }
 

light/lightchain.go

@@ -18,6 +18,7 @@ package light
 
 import (
 	"context"
+	"errors"
 	"math/big"
 	"sync"
 	"sync/atomic"
@@ -26,6 +27,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/consensus"
 	"github.com/ethereum/go-ethereum/core"
+	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/event"
@@ -212,6 +214,11 @@ func (bc *LightChain) Genesis() *types.Block {
 	return bc.genesisBlock
 }
 
+// State returns a new mutable state based on the current HEAD block.
+func (bc *LightChain) State() (*state.StateDB, error) {
+	return nil, errors.New("not implemented, needs client/server interface split")
+}
+
 // GetBody retrieves a block body (transactions and uncles) from the database
 // or ODR service by hash, caching it if found.
 func (self *LightChain) GetBody(ctx context.Context, hash common.Hash) (*types.Body, error) {

light/trie.go

@@ -18,12 +18,14 @@ package light
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/trie"
 )
 
@@ -139,6 +141,10 @@ func (t *odrTrie) GetKey(sha []byte) []byte {
 	return nil
 }
 
+func (t *odrTrie) Prove(key []byte, fromLevel uint, proofDb ethdb.Putter) error {
+	return errors.New("not implemented, needs client/server interface split")
+}
+
 // do tries and retries to execute a function until it returns with no error or
 // an error type other than MissingNodeError
 func (t *odrTrie) do(key []byte, fn func() error) error {

trie/proof.go

@@ -88,6 +88,17 @@ func (t *Trie) Prove(key []byte, fromLevel uint, proofDb ethdb.Putter) error {
 	return nil
 }
 
+// Prove constructs a merkle proof for key. The result contains all encoded nodes
+// on the path to the value at key. The value itself is also included in the last
+// node and can be retrieved by verifying the proof.
+//
+// If the trie does not contain a value for key, the returned proof 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 *SecureTrie) Prove(key []byte, fromLevel uint, proofDb ethdb.Putter) error {
+	return t.trie.Prove(key, fromLevel, proofDb)
+}
+
 // 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.