diff options
Diffstat (limited to 'trie')
| -rw-r--r-- | trie/secure_trie.go | 33 | ||||
| -rw-r--r-- | trie/secure_trie_test.go | 71 |
2 files changed, 94 insertions, 10 deletions
diff --git a/trie/secure_trie.go b/trie/secure_trie.go index efe875bc8..2a8b57214 100644 --- a/trie/secure_trie.go +++ b/trie/secure_trie.go @@ -24,6 +24,8 @@ import ( var secureKeyPrefix = []byte("secure-key-") +const secureKeyLength = 11 + 32 // Length of the above prefix + 32byte hash + // SecureTrie wraps a trie with key hashing. In a secure trie, all // access operations hash the key using keccak256. This prevents // calling code from creating long chains of nodes that @@ -35,10 +37,11 @@ var secureKeyPrefix = []byte("secure-key-") // // SecureTrie is not safe for concurrent use. type SecureTrie struct { - trie Trie - hashKeyBuf []byte - secKeyBuf [200]byte - secKeyCache map[string][]byte + trie Trie + hashKeyBuf [secureKeyLength]byte + secKeyBuf [200]byte + secKeyCache map[string][]byte + secKeyCacheOwner *SecureTrie // Pointer to self, replace the key cache on mismatch } // NewSecure creates a trie with an existing root node from db. @@ -56,8 +59,7 @@ func NewSecure(root common.Hash, db Database) (*SecureTrie, error) { return nil, err } return &SecureTrie{ - trie: *trie, - secKeyCache: make(map[string][]byte), + trie: *trie, }, nil } @@ -104,7 +106,7 @@ func (t *SecureTrie) TryUpdate(key, value []byte) error { if err != nil { return err } - t.secKeyCache[string(hk)] = common.CopyBytes(key) + t.getSecKeyCache()[string(hk)] = common.CopyBytes(key) return nil } @@ -119,14 +121,14 @@ func (t *SecureTrie) Delete(key []byte) { // If a node was not found in the database, a MissingNodeError is returned. func (t *SecureTrie) TryDelete(key []byte) error { hk := t.hashKey(key) - delete(t.secKeyCache, string(hk)) + delete(t.getSecKeyCache(), string(hk)) return t.trie.TryDelete(hk) } // GetKey returns the sha3 preimage of a hashed key that was // previously used to store a value. func (t *SecureTrie) GetKey(shaKey []byte) []byte { - if key, ok := t.secKeyCache[string(shaKey)]; ok { + if key, ok := t.getSecKeyCache()[string(shaKey)]; ok { return key } key, _ := t.trie.db.Get(t.secKey(shaKey)) @@ -165,7 +167,7 @@ func (t *SecureTrie) NodeIterator() *NodeIterator { // the trie's database. Calling code must ensure that the changes made to db are // written back to the trie's attached database before using the trie. func (t *SecureTrie) CommitTo(db DatabaseWriter) (root common.Hash, err error) { - if len(t.secKeyCache) > 0 { + if len(t.getSecKeyCache()) > 0 { for hk, key := range t.secKeyCache { if err := db.Put(t.secKey([]byte(hk)), key); err != nil { return common.Hash{}, err @@ -196,3 +198,14 @@ func (t *SecureTrie) hashKey(key []byte) []byte { returnHasherToPool(h) return buf } + +// getSecKeyCache returns the current secure key cache, creating a new one if +// ownership changed (i.e. the current secure trie is a copy of another owning +// the actual cache). +func (t *SecureTrie) getSecKeyCache() map[string][]byte { + if t != t.secKeyCacheOwner { + t.secKeyCacheOwner = t + t.secKeyCache = make(map[string][]byte) + } + return t.secKeyCache +} diff --git a/trie/secure_trie_test.go b/trie/secure_trie_test.go index 0be5b3d15..3171b8c31 100644 --- a/trie/secure_trie_test.go +++ b/trie/secure_trie_test.go @@ -18,6 +18,8 @@ package trie import ( "bytes" + "runtime" + "sync" "testing" "github.com/ethereum/go-ethereum/common" @@ -31,6 +33,37 @@ func newEmptySecure() *SecureTrie { return trie } +// makeTestSecureTrie creates a large enough secure trie for testing. +func makeTestSecureTrie() (ethdb.Database, *SecureTrie, map[string][]byte) { + // Create an empty trie + db, _ := ethdb.NewMemDatabase() + trie, _ := NewSecure(common.Hash{}, db) + + // Fill it with some arbitrary data + content := make(map[string][]byte) + for i := byte(0); i < 255; i++ { + // Map the same data under multiple keys + key, val := common.LeftPadBytes([]byte{1, i}, 32), []byte{i} + content[string(key)] = val + trie.Update(key, val) + + key, val = common.LeftPadBytes([]byte{2, i}, 32), []byte{i} + content[string(key)] = val + trie.Update(key, val) + + // Add some other data to inflate th trie + for j := byte(3); j < 13; j++ { + key, val = common.LeftPadBytes([]byte{j, i}, 32), []byte{j, i} + content[string(key)] = val + trie.Update(key, val) + } + } + trie.Commit() + + // Return the generated trie + return db, trie, content +} + func TestSecureDelete(t *testing.T) { trie := newEmptySecure() vals := []struct{ k, v string }{ @@ -72,3 +105,41 @@ func TestSecureGetKey(t *testing.T) { t.Errorf("GetKey returned %q, want %q", k, key) } } + +func TestSecureTrieConcurrency(t *testing.T) { + // Create an initial trie and copy if for concurrent access + _, trie, _ := makeTestSecureTrie() + + threads := runtime.NumCPU() + tries := make([]*SecureTrie, threads) + for i := 0; i < threads; i++ { + cpy := *trie + tries[i] = &cpy + } + // Start a batch of goroutines interactng with the trie + pend := new(sync.WaitGroup) + pend.Add(threads) + for i := 0; i < threads; i++ { + go func(index int) { + defer pend.Done() + + for j := byte(0); j < 255; j++ { + // Map the same data under multiple keys + key, val := common.LeftPadBytes([]byte{byte(index), 1, j}, 32), []byte{j} + tries[index].Update(key, val) + + key, val = common.LeftPadBytes([]byte{byte(index), 2, j}, 32), []byte{j} + tries[index].Update(key, val) + + // Add some other data to inflate the trie + for k := byte(3); k < 13; k++ { + key, val = common.LeftPadBytes([]byte{byte(index), k, j}, 32), []byte{k, j} + tries[index].Update(key, val) + } + } + tries[index].Commit() + }(i) + } + // Wait for all threads to finish + pend.Wait() +} |