github.com/tacshi/go-ethereum@v0.0.0-20230616113857-84a434e20921/trie/secure_trie_test.go (about)

     1  // Copyright 2015 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package trie
    18  
    19  import (
    20  	"bytes"
    21  	"fmt"
    22  	"runtime"
    23  	"sync"
    24  	"testing"
    25  
    26  	"github.com/tacshi/go-ethereum/common"
    27  	"github.com/tacshi/go-ethereum/core/rawdb"
    28  	"github.com/tacshi/go-ethereum/crypto"
    29  )
    30  
    31  func newEmptySecure() *StateTrie {
    32  	trie, _ := NewStateTrie(TrieID(common.Hash{}), NewDatabase(rawdb.NewMemoryDatabase()))
    33  	return trie
    34  }
    35  
    36  // makeTestStateTrie creates a large enough secure trie for testing.
    37  func makeTestStateTrie() (*Database, *StateTrie, map[string][]byte) {
    38  	// Create an empty trie
    39  	triedb := NewDatabase(rawdb.NewMemoryDatabase())
    40  	trie, _ := NewStateTrie(TrieID(common.Hash{}), triedb)
    41  
    42  	// Fill it with some arbitrary data
    43  	content := make(map[string][]byte)
    44  	for i := byte(0); i < 255; i++ {
    45  		// Map the same data under multiple keys
    46  		key, val := common.LeftPadBytes([]byte{1, i}, 32), []byte{i}
    47  		content[string(key)] = val
    48  		trie.Update(key, val)
    49  
    50  		key, val = common.LeftPadBytes([]byte{2, i}, 32), []byte{i}
    51  		content[string(key)] = val
    52  		trie.Update(key, val)
    53  
    54  		// Add some other data to inflate the trie
    55  		for j := byte(3); j < 13; j++ {
    56  			key, val = common.LeftPadBytes([]byte{j, i}, 32), []byte{j, i}
    57  			content[string(key)] = val
    58  			trie.Update(key, val)
    59  		}
    60  	}
    61  	root, nodes := trie.Commit(false)
    62  	if err := triedb.Update(NewWithNodeSet(nodes)); err != nil {
    63  		panic(fmt.Errorf("failed to commit db %v", err))
    64  	}
    65  	// Re-create the trie based on the new state
    66  	trie, _ = NewStateTrie(TrieID(root), triedb)
    67  	return triedb, trie, content
    68  }
    69  
    70  func TestSecureDelete(t *testing.T) {
    71  	trie := newEmptySecure()
    72  	vals := []struct{ k, v string }{
    73  		{"do", "verb"},
    74  		{"ether", "wookiedoo"},
    75  		{"horse", "stallion"},
    76  		{"shaman", "horse"},
    77  		{"doge", "coin"},
    78  		{"ether", ""},
    79  		{"dog", "puppy"},
    80  		{"shaman", ""},
    81  	}
    82  	for _, val := range vals {
    83  		if val.v != "" {
    84  			trie.Update([]byte(val.k), []byte(val.v))
    85  		} else {
    86  			trie.Delete([]byte(val.k))
    87  		}
    88  	}
    89  	hash := trie.Hash()
    90  	exp := common.HexToHash("29b235a58c3c25ab83010c327d5932bcf05324b7d6b1185e650798034783ca9d")
    91  	if hash != exp {
    92  		t.Errorf("expected %x got %x", exp, hash)
    93  	}
    94  }
    95  
    96  func TestSecureGetKey(t *testing.T) {
    97  	trie := newEmptySecure()
    98  	trie.Update([]byte("foo"), []byte("bar"))
    99  
   100  	key := []byte("foo")
   101  	value := []byte("bar")
   102  	seckey := crypto.Keccak256(key)
   103  
   104  	if !bytes.Equal(trie.Get(key), value) {
   105  		t.Errorf("Get did not return bar")
   106  	}
   107  	if k := trie.GetKey(seckey); !bytes.Equal(k, key) {
   108  		t.Errorf("GetKey returned %q, want %q", k, key)
   109  	}
   110  }
   111  
   112  func TestStateTrieConcurrency(t *testing.T) {
   113  	// Create an initial trie and copy if for concurrent access
   114  	_, trie, _ := makeTestStateTrie()
   115  
   116  	threads := runtime.NumCPU()
   117  	tries := make([]*StateTrie, threads)
   118  	for i := 0; i < threads; i++ {
   119  		tries[i] = trie.Copy()
   120  	}
   121  	// Start a batch of goroutines interacting with the trie
   122  	pend := new(sync.WaitGroup)
   123  	pend.Add(threads)
   124  	for i := 0; i < threads; i++ {
   125  		go func(index int) {
   126  			defer pend.Done()
   127  
   128  			for j := byte(0); j < 255; j++ {
   129  				// Map the same data under multiple keys
   130  				key, val := common.LeftPadBytes([]byte{byte(index), 1, j}, 32), []byte{j}
   131  				tries[index].Update(key, val)
   132  
   133  				key, val = common.LeftPadBytes([]byte{byte(index), 2, j}, 32), []byte{j}
   134  				tries[index].Update(key, val)
   135  
   136  				// Add some other data to inflate the trie
   137  				for k := byte(3); k < 13; k++ {
   138  					key, val = common.LeftPadBytes([]byte{byte(index), k, j}, 32), []byte{k, j}
   139  					tries[index].Update(key, val)
   140  				}
   141  			}
   142  			tries[index].Commit(false)
   143  		}(i)
   144  	}
   145  	// Wait for all threads to finish
   146  	pend.Wait()
   147  }