github.com/klaytn/klaytn@v1.12.1/storage/statedb/secure_trie_test.go (about)

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