fix(lib/trie): use MustBeHashed for V1 trie nodes with larger stora…

…ge values (#3739)
ChainSafe · Feb 28, 2024 · f5e48a9 · f5e48a9
1 parent 75ef976
commit f5e48a9
Show file tree

Hide file tree

Showing 62 changed files with 578 additions and 358 deletions.
diff --git a/dot/core/service.go b/dot/core/service.go
@@ -247,7 +247,7 @@ func (s *Service) handleBlock(block *types.Block, state *rtstorage.TrieState) er
 	}
 
 	logger.Debugf("imported block %s and stored state trie with root %s",
-		block.Header.Hash(), state.MustRoot(trie.NoMaxInlineValueSize))
+		block.Header.Hash(), state.MustRoot())
 
 	parentRuntimeInstance, err := s.blockState.GetRuntime(block.Header.ParentHash)
 	if err != nil {

diff --git a/dot/core/service_test.go b/dot/core/service_test.go
@@ -126,8 +126,6 @@ func generateExtrinsic(t *testing.T) (extrinsic, externalExtrinsic types.Extrins
 func Test_Service_StorageRoot(t *testing.T) {
 	t.Parallel()
 
-	ts := rtstorage.NewTrieState(trie.NewEmptyTrie())
-
 	tests := []struct {
 		name          string
 		service       *Service
@@ -152,7 +150,7 @@ func Test_Service_StorageRoot(t *testing.T) {
 			service: &Service{},
 			exp: common.Hash{0x3, 0x17, 0xa, 0x2e, 0x75, 0x97, 0xb7, 0xb7, 0xe3, 0xd8, 0x4c, 0x5, 0x39, 0x1d, 0x13, 0x9a,
 				0x62, 0xb1, 0x57, 0xe7, 0x87, 0x86, 0xd8, 0xc0, 0x82, 0xf2, 0x9d, 0xcf, 0x4c, 0x11, 0x13, 0x14},
-			retTrieState:  ts,
+			retTrieState:  rtstorage.NewTrieState(trie.NewEmptyTrie()),
 			trieStateCall: true,
 			stateVersion:  0,
 		},
@@ -161,7 +159,7 @@ func Test_Service_StorageRoot(t *testing.T) {
 			service: &Service{},
 			exp: common.Hash{0x3, 0x17, 0xa, 0x2e, 0x75, 0x97, 0xb7, 0xb7, 0xe3, 0xd8, 0x4c, 0x5, 0x39, 0x1d, 0x13, 0x9a,
 				0x62, 0xb1, 0x57, 0xe7, 0x87, 0x86, 0xd8, 0xc0, 0x82, 0xf2, 0x9d, 0xcf, 0x4c, 0x11, 0x13, 0x14},
-			retTrieState:  ts,
+			retTrieState:  rtstorage.NewTrieState(trie.NewEmptyTrie()),
 			trieStateCall: true,
 			stateVersion:  1,
 		},
@@ -282,6 +280,7 @@ func Test_Service_handleCodeSubstitution(t *testing.T) {
 			errWrapped: io.ErrUnexpectedEOF,
 		},
 		"store_code_substitution_block_hash_error": {
+			trieState: rtstorage.NewTrieState(trie.NewEmptyTrie()),
 			serviceBuilder: func(ctrl *gomock.Controller) *Service {
 				storedRuntime := NewMockInstance(ctrl)
 				storedRuntime.EXPECT().Keystore().Return(nil)
@@ -339,6 +338,7 @@ func Test_Service_handleCodeSubstitution(t *testing.T) {
 				}
 			},
 			blockHash: common.Hash{0x01},
+			trieState: rtstorage.NewTrieState(trie.NewEmptyTrie()),
 		},
 	}
 

diff --git a/dot/import.go b/dot/import.go
@@ -21,7 +21,7 @@ import (
 
 // ImportState imports the state in the given files to the database with the given path.
 func ImportState(basepath, stateFP, headerFP string, stateTrieVersion trie.TrieLayout, firstSlot uint64) error {
-	tr, err := newTrieFromPairs(stateFP)
+	tr, err := newTrieFromPairs(stateFP, trie.V0)
 	if err != nil {
 		return err
 	}
@@ -41,7 +41,7 @@ func ImportState(basepath, stateFP, headerFP string, stateTrieVersion trie.TrieL
 	return srv.Import(header, tr, stateTrieVersion, firstSlot)
 }
 
-func newTrieFromPairs(filename string) (*trie.Trie, error) {
+func newTrieFromPairs(filename string, version trie.TrieLayout) (*trie.Trie, error) {
 	data, err := os.ReadFile(filepath.Clean(filename))
 	if err != nil {
 		return nil, err
@@ -62,7 +62,7 @@ func newTrieFromPairs(filename string) (*trie.Trie, error) {
 		entries[pairArr[0].(string)] = pairArr[1].(string)
 	}
 
-	tr, err := trie.LoadFromMap(entries)
+	tr, err := trie.LoadFromMap(entries, version)
 	if err != nil {
 		return nil, err
 	}

diff --git a/dot/import_integration_test.go b/dot/import_integration_test.go
@@ -52,7 +52,7 @@ func Test_newTrieFromPairs(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			got, err := newTrieFromPairs(tt.filename)
+			got, err := newTrieFromPairs(tt.filename, tt.stateVersion)
 			if tt.err != nil {
 				assert.EqualError(t, err, tt.err.Error())
 			} else {

diff --git a/dot/node_integration_test.go b/dot/node_integration_test.go
@@ -385,7 +385,7 @@ func TestInitNode_LoadStorageRoot(t *testing.T) {
 	node, err := NewNode(config, ks)
 	require.NoError(t, err)
 
-	expected, err := trie.LoadFromMap(gen.GenesisFields().Raw["top"])
+	expected, err := trie.LoadFromMap(gen.GenesisFields().Raw["top"], trie.V0)
 	require.NoError(t, err)
 
 	expectedRoot, err := trie.V0.Hash(&expected) // Since we are using a runtime with state trie V0

diff --git a/dot/rpc/modules/childstate_integration_test.go b/dot/rpc/modules/childstate_integration_test.go
@@ -255,7 +255,7 @@ func setupChildStateStorage(t *testing.T) (*ChildStateModule, common.Hash) {
 	err = tr.SetChild([]byte(":child_storage_key"), childTr)
 	require.NoError(t, err)
 
-	stateRoot, err := tr.Root(trie.NoMaxInlineValueSize)
+	stateRoot, err := tr.Root()
 	require.NoError(t, err)
 
 	bb, err := st.Block.BestBlock()

diff --git a/dot/rpc/modules/childstate_test.go b/dot/rpc/modules/childstate_test.go
@@ -35,7 +35,7 @@ func createTestTrieState(t *testing.T) (*trie.Trie, common.Hash) {
 	err := tr.SetChild([]byte(":child_storage_key"), childTr)
 	require.NoError(t, err)
 
-	stateRoot, err := tr.Root(trie.NoMaxInlineValueSize)
+	stateRoot, err := tr.Root()
 	require.NoError(t, err)
 
 	return tr.Trie(), stateRoot

diff --git a/dot/rpc/modules/dev_integration_test.go b/dot/rpc/modules/dev_integration_test.go
@@ -57,7 +57,7 @@ func newBABEService(t *testing.T) *babe.Service {
 
 	bs, es := newState(t)
 	tt := trie.NewEmptyTrie()
-	rt := wazero_runtime.NewTestInstanceWithTrie(t, runtime.WESTEND_RUNTIME_v0929, tt)
+	rt := wazero_runtime.NewTestInstance(t, runtime.WESTEND_RUNTIME_v0929, wazero_runtime.TestWithTrie(tt))
 	bs.StoreRuntime(bs.GenesisHash(), rt)
 	tt.Put(
 		common.MustHexToBytes("0x886726f904d8372fdabb7707870c2fad"),

diff --git a/dot/rpc/modules/helpers_test.go b/dot/rpc/modules/helpers_test.go
@@ -35,7 +35,7 @@ func newWestendLocalGenesisWithTrieAndHeader(t *testing.T) (
 	require.NoError(t, err)
 
 	parentHash := common.NewHash([]byte{0})
-	stateRoot := genesisTrie.MustHash(trie.NoMaxInlineValueSize)
+	stateRoot := genesisTrie.MustHash()
 	extrinsicRoot := trie.EmptyHash
 	const number = 0
 	digest := types.NewDigest()

diff --git a/dot/rpc/modules/state_integration_test.go b/dot/rpc/modules/state_integration_test.go
@@ -17,7 +17,6 @@ import (
 	"github.com/ChainSafe/gossamer/dot/rpc/modules/mocks"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
-	"github.com/ChainSafe/gossamer/pkg/trie"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
@@ -575,7 +574,7 @@ func setupStateModule(t *testing.T) (*StateModule, *common.Hash, *common.Hash) {
 	err = ts.SetChildStorage([]byte(`:child1`), []byte(`:key1`), []byte(`:childValue1`))
 	require.NoError(t, err)
 
-	sr1, err := ts.Root(trie.NoMaxInlineValueSize)
+	sr1, err := ts.Root()
 	require.NoError(t, err)
 	err = chain.Storage.StoreTrie(ts, nil)
 	require.NoError(t, err)

diff --git a/dot/rpc/modules/system_integration_test.go b/dot/rpc/modules/system_integration_test.go
@@ -330,7 +330,7 @@ func setupSystemModule(t *testing.T) *SystemModule {
 		Header: types.Header{
 			Number:     3,
 			ParentHash: chain.Block.BestBlockHash(),
-			StateRoot:  ts.MustRoot(trie.NoMaxInlineValueSize),
+			StateRoot:  ts.MustRoot(),
 			Digest:     digest,
 		},
 		Body: types.Body{},
@@ -354,7 +354,7 @@ func setupSystemModule(t *testing.T) *SystemModule {
 func newCoreService(t *testing.T, srvc *state.Service) *core.Service {
 	// setup service
 	tt := trie.NewEmptyTrie()
-	rt := wazero_runtime.NewTestInstanceWithTrie(t, runtime.WESTEND_RUNTIME_v0929, tt)
+	rt := wazero_runtime.NewTestInstance(t, runtime.WESTEND_RUNTIME_v0929, wazero_runtime.TestWithTrie(tt))
 	ks := keystore.NewGlobalKeystore()
 	t.Cleanup(func() {
 		rt.Stop()

diff --git a/dot/state/base_test.go b/dot/state/base_test.go
@@ -29,15 +29,15 @@ func TestTrie_StoreAndLoadFromDB(t *testing.T) {
 	err := tt.WriteDirty(db)
 	require.NoError(t, err)
 
-	encroot, err := tt.Hash(trie.NoMaxInlineValueSize)
+	encroot, err := tt.Hash()
 	require.NoError(t, err)
 
-	expected := tt.MustHash(trie.NoMaxInlineValueSize)
+	expected := tt.MustHash()
 
 	tt = trie.NewEmptyTrie()
 	err = tt.Load(db, encroot)
 	require.NoError(t, err)
-	require.Equal(t, expected, tt.MustHash(trie.NoMaxInlineValueSize))
+	require.Equal(t, expected, tt.MustHash())
 }
 
 func TestStoreAndLoadGenesisData(t *testing.T) {

diff --git a/dot/state/block.go b/dot/state/block.go
@@ -880,7 +880,7 @@ func (bs *BlockState) HandleRuntimeChanges(newState *rtstorage.TrieState,
 		return fmt.Errorf("failed to update code substituted block hash: %w", err)
 	}
 
-	newVersion, err := parentRuntimeInstance.Version()
+	newVersion, err := instance.Version()
 	if err != nil {
 		return err
 	}

diff --git a/dot/state/helpers_test.go b/dot/state/helpers_test.go
@@ -102,7 +102,7 @@ func newWestendDevGenesisWithTrieAndHeader(t *testing.T) (
 	require.NoError(t, err)
 
 	parentHash := common.NewHash([]byte{0})
-	stateRoot := genesisTrie.MustHash(trie.NoMaxInlineValueSize)
+	stateRoot := genesisTrie.MustHash()
 	extrinsicRoot := trie.EmptyHash
 	const number = 0
 	digest := types.NewDigest()

diff --git a/dot/state/service_integration_test.go b/dot/state/service_integration_test.go
@@ -90,7 +90,7 @@ func TestService_Initialise(t *testing.T) {
 	require.NoError(t, err)
 
 	genesisHeaderPtr := types.NewHeader(common.NewHash([]byte{77}),
-		genTrie.MustHash(trie.NoMaxInlineValueSize), trie.EmptyHash, 0, nil)
+		genTrie.MustHash(), trie.EmptyHash, 0, nil)
 
 	err = state.Initialise(&genData, genesisHeaderPtr, genTrieCopy)
 	require.NoError(t, err)
@@ -287,7 +287,7 @@ func TestService_PruneStorage(t *testing.T) {
 		copiedTrie := trieState.Trie().DeepCopy()
 
 		var rootHash common.Hash
-		rootHash, err = copiedTrie.Hash(trie.NoMaxInlineValueSize)
+		rootHash, err = copiedTrie.Hash()
 		require.NoError(t, err)
 
 		prunedArr = append(prunedArr, prunedBlock{hash: block.Header.StateRoot, dbKey: rootHash[:]})
@@ -400,7 +400,7 @@ func TestService_Import(t *testing.T) {
 	require.NoError(t, err)
 	header := &types.Header{
 		Number:    77,
-		StateRoot: tr.MustHash(trie.NoMaxInlineValueSize),
+		StateRoot: tr.MustHash(),
 		Digest:    digest,
 	}
 
@@ -440,7 +440,7 @@ func generateBlockWithRandomTrie(t *testing.T, serv *Service,
 	err = trieState.Put(key, value)
 	require.NoError(t, err)
 
-	trieStateRoot, err := trieState.Root(trie.NoMaxInlineValueSize)
+	trieStateRoot, err := trieState.Root()
 	require.NoError(t, err)
 
 	if parent == nil {

diff --git a/dot/state/storage.go b/dot/state/storage.go
@@ -59,8 +59,7 @@ func NewStorageState(db database.Database, blockState *BlockState,
 
 // StoreTrie stores the given trie in the StorageState and writes it to the database
 func (s *StorageState) StoreTrie(ts *rtstorage.TrieState, header *types.Header) error {
-	root := ts.MustRoot(trie.NoMaxInlineValueSize)
-
+	root := ts.MustRoot()
 	s.tries.softSet(root, ts.Trie())
 
 	if header != nil {
@@ -92,7 +91,7 @@ func (s *StorageState) TrieState(root *common.Hash) (*rtstorage.TrieState, error
 	if root == nil {
 		sr, err := s.blockState.BestBlockStateRoot()
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("while getting best block state root: %w", err)
 		}
 		root = &sr
 	}
@@ -102,11 +101,11 @@ func (s *StorageState) TrieState(root *common.Hash) (*rtstorage.TrieState, error
 		var err error
 		t, err = s.LoadFromDB(*root)
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("while loading from database: %w", err)
 		}
 
 		s.tries.softSet(*root, t)
-	} else if t.MustHash(trie.NoMaxInlineValueSize) != *root {
+	} else if t.MustHash() != *root {
 		panic("trie does not have expected root")
 	}
 
@@ -125,7 +124,7 @@ func (s *StorageState) LoadFromDB(root common.Hash) (*trie.Trie, error) {
 		return nil, err
 	}
 
-	s.tries.softSet(t.MustHash(trie.NoMaxInlineValueSize), t)
+	s.tries.softSet(t.MustHash(), t)
 	return t, nil
 }
 

diff --git a/dot/state/storage_test.go b/dot/state/storage_test.go
@@ -35,7 +35,7 @@ func TestStorage_StoreAndLoadTrie(t *testing.T) {
 	ts, err := storage.TrieState(&trie.EmptyHash)
 	require.NoError(t, err)
 
-	root, err := ts.Root(trie.NoMaxInlineValueSize)
+	root, err := ts.Root()
 	require.NoError(t, err)
 	err = storage.StoreTrie(ts, nil)
 	require.NoError(t, err)
@@ -59,7 +59,7 @@ func TestStorage_GetStorageByBlockHash(t *testing.T) {
 	value := []byte("testvalue")
 	ts.Put(key, value)
 
-	root, err := ts.Root(trie.NoMaxInlineValueSize)
+	root, err := ts.Root()
 	require.NoError(t, err)
 	err = storage.StoreTrie(ts, nil)
 	require.NoError(t, err)
@@ -91,7 +91,7 @@ func TestStorage_TrieState(t *testing.T) {
 	require.NoError(t, err)
 	ts.Put([]byte("noot"), []byte("washere"))
 
-	root, err := ts.Root(trie.NoMaxInlineValueSize)
+	root, err := ts.Root()
 	require.NoError(t, err)
 	err = storage.StoreTrie(ts, nil)
 	require.NoError(t, err)
@@ -102,7 +102,7 @@ func TestStorage_TrieState(t *testing.T) {
 	storage.blockState.tries.delete(root)
 	ts3, err := storage.TrieState(&root)
 	require.NoError(t, err)
-	require.Equal(t, ts.Trie().MustHash(trie.NoMaxInlineValueSize), ts3.Trie().MustHash(trie.NoMaxInlineValueSize))
+	require.Equal(t, ts.Trie().MustHash(), ts3.Trie().MustHash())
 }
 
 func TestStorage_LoadFromDB(t *testing.T) {
@@ -125,7 +125,7 @@ func TestStorage_LoadFromDB(t *testing.T) {
 		ts.Put(kv.key, kv.value)
 	}
 
-	root, err := ts.Root(trie.NoMaxInlineValueSize)
+	root, err := ts.Root()
 	require.NoError(t, err)
 
 	// Write trie to disk.
@@ -205,13 +205,13 @@ func TestGetStorageChildAndGetStorageFromChild(t *testing.T) {
 
 	trieState := runtime.NewTrieState(&genTrie)
 
-	header := types.NewHeader(blockState.GenesisHash(), trieState.MustRoot(trie.NoMaxInlineValueSize),
+	header := types.NewHeader(blockState.GenesisHash(), trieState.MustRoot(),
 		common.Hash{}, 1, types.NewDigest())
 
 	err = storage.StoreTrie(trieState, header)
 	require.NoError(t, err)
 
-	rootHash, err := genTrie.Hash(trie.NoMaxInlineValueSize)
+	rootHash, err := genTrie.Hash()
 	require.NoError(t, err)
 
 	_, err = storage.GetStorageChild(&rootHash, []byte("keyToChild"))

diff --git a/dot/state/tries.go b/dot/state/tries.go
@@ -60,7 +60,7 @@ func (t *Tries) SetEmptyTrie() {
 
 // SetTrie sets the trie at its root hash in the tries map.
 func (t *Tries) SetTrie(tr *trie.Trie) {
-	t.softSet(tr.MustHash(trie.NoMaxInlineValueSize), tr)
+	t.softSet(tr.MustHash(), tr)
 }
 
 // softSet sets the given trie at the given root hash

diff --git a/dot/state/tries_test.go b/dot/state/tries_test.go
@@ -59,7 +59,7 @@ func Test_Tries_SetTrie(t *testing.T) {
 
 	expectedTries := &Tries{
 		rootToTrie: map[common.Hash]*trie.Trie{
-			tr.MustHash(trie.NoMaxInlineValueSize): tr,
+			tr.MustHash(): tr,
 		},
 		triesGauge:    triesGauge,
 		setCounter:    setCounter,

diff --git a/dot/sync/chain_sync.go b/dot/sync/chain_sync.go
@@ -24,7 +24,6 @@ import (
 	"github.com/ChainSafe/gossamer/internal/database"
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/common/variadic"
-	"github.com/ChainSafe/gossamer/pkg/trie"
 )
 
 var _ ChainSync = (*chainSync)(nil)
@@ -872,7 +871,7 @@ func (cs *chainSync) handleBlock(block *types.Block, announceImportedBlock bool)
 		return err
 	}
 
-	root := ts.MustRoot(trie.NoMaxInlineValueSize)
+	root := ts.MustRoot()
 	if !bytes.Equal(parent.StateRoot[:], root[:]) {
 		panic("parent state root does not match snapshot state root")
 	}