1 files changed, 40 insertions, 30 deletions

diff --git a/trie/secure_trie.go b/trie/secure_trie.go
index 1d027c102..efe875bc8 100644
--- a/trie/secure_trie.go
+++ b/trie/secure_trie.go
@@ -17,10 +17,7 @@
 package trie
 
 import (
-	"hash"
-
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/crypto/sha3"
 	"github.com/ethereum/go-ethereum/logger"
 	"github.com/ethereum/go-ethereum/logger/glog"
 )
@@ -38,11 +35,9 @@ var secureKeyPrefix = []byte("secure-key-")
 //
 // SecureTrie is not safe for concurrent use.
 type SecureTrie struct {
-	*Trie
-
-	hash        hash.Hash
+	trie        Trie
 	hashKeyBuf  []byte
-	secKeyBuf   []byte
+	secKeyBuf   [200]byte
 	secKeyCache map[string][]byte
 }
 
@@ -61,7 +56,7 @@ func NewSecure(root common.Hash, db Database) (*SecureTrie, error) {
 		return nil, err
 	}
 	return &SecureTrie{
-		Trie:        trie,
+		trie:        *trie,
 		secKeyCache: make(map[string][]byte),
 	}, nil
 }
@@ -80,7 +75,7 @@ func (t *SecureTrie) Get(key []byte) []byte {
 // The value bytes must not be modified by the caller.
 // If a node was not found in the database, a MissingNodeError is returned.
 func (t *SecureTrie) TryGet(key []byte) ([]byte, error) {
-	return t.Trie.TryGet(t.hashKey(key))
+	return t.trie.TryGet(t.hashKey(key))
 }
 
 // Update associates key with value in the trie. Subsequent calls to
@@ -105,7 +100,7 @@ func (t *SecureTrie) Update(key, value []byte) {
 // If a node was not found in the database, a MissingNodeError is returned.
 func (t *SecureTrie) TryUpdate(key, value []byte) error {
 	hk := t.hashKey(key)
-	err := t.Trie.TryUpdate(hk, value)
+	err := t.trie.TryUpdate(hk, value)
 	if err != nil {
 		return err
 	}
@@ -125,7 +120,7 @@ func (t *SecureTrie) Delete(key []byte) {
 func (t *SecureTrie) TryDelete(key []byte) error {
 	hk := t.hashKey(key)
 	delete(t.secKeyCache, string(hk))
-	return t.Trie.TryDelete(hk)
+	return t.trie.TryDelete(hk)
 }
 
 // GetKey returns the sha3 preimage of a hashed key that was
@@ -134,7 +129,7 @@ func (t *SecureTrie) GetKey(shaKey []byte) []byte {
 	if key, ok := t.secKeyCache[string(shaKey)]; ok {
 		return key
 	}
-	key, _ := t.Trie.db.Get(t.secKey(shaKey))
+	key, _ := t.trie.db.Get(t.secKey(shaKey))
 	return key
 }
 
@@ -144,7 +139,23 @@ func (t *SecureTrie) GetKey(shaKey []byte) []byte {
 // Committing flushes nodes from memory. Subsequent Get calls will load nodes
 // from the database.
 func (t *SecureTrie) Commit() (root common.Hash, err error) {
-	return t.CommitTo(t.db)
+	return t.CommitTo(t.trie.db)
+}
+
+func (t *SecureTrie) Hash() common.Hash {
+	return t.trie.Hash()
+}
+
+func (t *SecureTrie) Root() []byte {
+	return t.trie.Root()
+}
+
+func (t *SecureTrie) Iterator() *Iterator {
+	return t.trie.Iterator()
+}
+
+func (t *SecureTrie) NodeIterator() *NodeIterator {
+	return NewNodeIterator(&t.trie)
 }
 
 // CommitTo writes all nodes and the secure hash pre-images to the given database.
@@ -162,27 +173,26 @@ func (t *SecureTrie) CommitTo(db DatabaseWriter) (root common.Hash, err error) {
 		}
 		t.secKeyCache = make(map[string][]byte)
 	}
-	n, clean, err := t.hashRoot(db)
-	if err != nil {
-		return (common.Hash{}), err
-	}
-	t.root = clean
-	return common.BytesToHash(n.(hashNode)), nil
+	return t.trie.CommitTo(db)
 }
 
+// secKey returns the database key for the preimage of key, as an ephemeral buffer.
+// The caller must not hold onto the return value because it will become
+// invalid on the next call to hashKey or secKey.
 func (t *SecureTrie) secKey(key []byte) []byte {
-	t.secKeyBuf = append(t.secKeyBuf[:0], secureKeyPrefix...)
-	t.secKeyBuf = append(t.secKeyBuf, key...)
-	return t.secKeyBuf
+	buf := append(t.secKeyBuf[:0], secureKeyPrefix...)
+	buf = append(buf, key...)
+	return buf
 }
 
+// hashKey returns the hash of key as an ephemeral buffer.
+// The caller must not hold onto the return value because it will become
+// invalid on the next call to hashKey or secKey.
 func (t *SecureTrie) hashKey(key []byte) []byte {
-	if t.hash == nil {
-		t.hash = sha3.NewKeccak256()
-		t.hashKeyBuf = make([]byte, 32)
-	}
-	t.hash.Reset()
-	t.hash.Write(key)
-	t.hashKeyBuf = t.hash.Sum(t.hashKeyBuf[:0])
-	return t.hashKeyBuf
+	h := newHasher()
+	h.sha.Reset()
+	h.sha.Write(key)
+	buf := h.sha.Sum(t.hashKeyBuf[:0])
+	returnHasherToPool(h)
+	return buf
 }