Self review.

dot/core/interface.go

@@ -50,9 +50,8 @@ type BlockState interface {
 	SubChain(start, end common.Hash) ([]common.Hash, error)
 	GetBlockBody(hash common.Hash) (*types.Body, error)
 	HandleRuntimeChanges(newState *rtstorage.TrieState, in runtime.Instance, bHash common.Hash) error
-	GetRuntime(*common.Hash) (runtime.Instance, bool)
+	GetRuntime(*common.Hash) (runtime.Instance, error)
 	StoreRuntime(common.Hash, runtime.Instance)
-	GetAllBlocks() []common.Hash
 }
 
 // StorageState interface for storage state methods

dot/core/messages.go

@@ -19,7 +19,6 @@ package core
 import (
 	"github.com/ChainSafe/gossamer/dot/network"
 	"github.com/ChainSafe/gossamer/dot/types"
-	"github.com/ChainSafe/gossamer/lib/blocktree"
 	"github.com/ChainSafe/gossamer/lib/transaction"
 )
 
@@ -33,9 +32,9 @@ func (s *Service) HandleTransactionMessage(msg *network.TransactionMessage) (boo
 	txs := msg.Extrinsics
 	var toPropagate []types.Extrinsic
 
-	rt, ok := s.blockState.GetRuntime(nil)
-	if !ok {
-		return false, blocktree.ErrFailedToGetRuntime
+	rt, err := s.blockState.GetRuntime(nil)
+	if err != nil {
+		return false, err
 	}
 
 	for _, tx := range txs {

dot/core/messages_test.go

@@ -150,8 +150,8 @@ func TestService_HandleTransactionMessage(t *testing.T) {
 	header, err := types.NewHeader(genHash, common.Hash{}, common.Hash{}, big.NewInt(1), types.NewEmptyDigest())
 	require.NoError(t, err)
 
-	rt, ok := s.blockState.GetRuntime(nil)
-	require.True(t, ok)
+	rt, err := s.blockState.GetRuntime(nil)
+	require.NoError(t, err)
 
 	// initialise block header
 	err = rt.InitializeBlock(header)

dot/core/service.go

@@ -199,7 +199,7 @@ func (s *Service) handleBlock(block *types.Block, state *rtstorage.TrieState) er
 	}
 
 	// store block in database
-	if err := s.blockState.AddBlock(block); err != nil {
+	if err = s.blockState.AddBlock(block); err != nil {
 		if err == blocktree.ErrParentNotFound && block.Header.Number.Cmp(big.NewInt(0)) != 0 {
 			return err
 		} else if err == blocktree.ErrBlockExists || block.Header.Number.Cmp(big.NewInt(0)) == 0 {
@@ -214,9 +214,9 @@ func (s *Service) handleBlock(block *types.Block, state *rtstorage.TrieState) er
 	// handle consensus digests
 	s.digestHandler.HandleDigests(block.Header)
 
-	rt, ok := s.blockState.GetRuntime(&block.Header.ParentHash)
-	if !ok {
-		return blocktree.ErrFailedToGetRuntime
+	rt, err := s.blockState.GetRuntime(&block.Header.ParentHash)
+	if err != nil {
+		return err
 	}
 
 	// check for runtime changes
@@ -262,12 +262,12 @@ func (s *Service) handleCodeSubstitution(hash common.Hash) error {
 		return ErrEmptyRuntimeCode
 	}
 
-	rt, ok := s.blockState.GetRuntime(&hash)
-	if !ok {
-		return blocktree.ErrFailedToGetRuntime
+	rt, err := s.blockState.GetRuntime(&hash)
+	if err != nil {
+		return err
 	}
 
-	err := rt.UpdateRuntimeCode(code)
+	err = rt.UpdateRuntimeCode(code)
 	if err != nil {
 		return err
 	}
@@ -352,19 +352,19 @@ func (s *Service) handleChainReorg(prev, curr common.Hash) error {
 		subchain = subchain[1:]
 	}
 
+	// Check transaction validation on the best block.
+	rt, err := s.blockState.GetRuntime(nil)
+	if err != nil {
+		return err
+	}
+
 	// for each block in the previous chain, re-add its extrinsics back into the pool
 	for _, hash := range subchain {
 		body, err := s.blockState.GetBlockBody(hash)
 		if err != nil {
 			continue
 		}
 
-		rt, ok := s.blockState.GetRuntime(&hash)
-		if !ok {
-			logger.Debug("failed to get runtime instance", "block", hash)
-			continue
-		}
-
 		exts, err := body.AsExtrinsics()
 		if err != nil {
 			continue
@@ -478,9 +478,9 @@ func (s *Service) GetRuntimeVersion(bhash *common.Hash) (runtime.Version, error)
 		return nil, err
 	}
 
-	rt, ok := s.blockState.GetRuntime(bhash)
-	if !ok {
-		return nil, blocktree.ErrFailedToGetRuntime
+	rt, err := s.blockState.GetRuntime(bhash)
+	if err != nil {
+		return nil, err
 	}
 
 	rt.SetContextStorage(ts)
@@ -499,9 +499,9 @@ func (s *Service) HandleSubmittedExtrinsic(ext types.Extrinsic) error {
 		return err
 	}
 
-	rt, ok := s.blockState.GetRuntime(nil)
-	if !ok {
-		return blocktree.ErrFailedToGetRuntime
+	rt, err := s.blockState.GetRuntime(nil)
+	if err != nil {
+		return err
 	}
 
 	rt.SetContextStorage(ts)
@@ -543,9 +543,9 @@ func (s *Service) GetMetadata(bhash *common.Hash) ([]byte, error) {
 		return nil, err
 	}
 
-	rt, ok := s.blockState.GetRuntime(bhash)
-	if !ok {
-		return nil, blocktree.ErrFailedToGetRuntime
+	rt, err := s.blockState.GetRuntime(bhash)
+	if err != nil {
+		return nil, err
 	}
 
 	rt.SetContextStorage(ts)

dot/core/service_test.go

@@ -50,9 +50,9 @@ func addTestBlocksToStateWithParent(t *testing.T, previousHash common.Hash, dept
 	require.NoError(t, err)
 	previousNum := prevHeader.Number
 
-	headers := []*types.Header{}
-	rt, ok := blockState.GetRuntime(nil)
-	require.True(t, ok)
+	var headers []*types.Header
+	rt, err := blockState.GetRuntime(nil)
+	require.NoError(t, err)
 
 	for i := 1; i <= depth; i++ {
 		block := &types.Block{
@@ -195,8 +195,8 @@ func TestHandleChainReorg_WithReorg_Trans(t *testing.T) {
 	parent, err := bs.BestBlockHeader()
 	require.NoError(t, err)
 
-	rt, ok := s.blockState.GetRuntime(nil)
-	require.True(t, ok)
+	rt, err := s.blockState.GetRuntime(nil)
+	require.NoError(t, err)
 
 	block1 := sync.BuildBlock(t, rt, parent, nil)
 	bs.StoreRuntime(block1.Header.Hash(), rt)
@@ -285,8 +285,8 @@ func TestHandleChainReorg_WithReorg_Transactions(t *testing.T) {
 	require.NoError(t, err)
 
 	bhash := s.blockState.BestBlockHash()
-	rt, ok := s.blockState.GetRuntime(&bhash)
-	require.True(t, ok)
+	rt, err := s.blockState.GetRuntime(&bhash)
+	require.NoError(t, err)
 
 	validity, err := rt.ValidateTransaction(tx)
 	require.NoError(t, err)
@@ -438,8 +438,8 @@ func TestMaintainTransactionPool_BlockWithExtrinsics(t *testing.T) {
 
 func TestService_GetRuntimeVersion(t *testing.T) {
 	s := NewTestService(t, nil)
-	rt, ok := s.blockState.GetRuntime(nil)
-	require.True(t, ok)
+	rt, err := s.blockState.GetRuntime(nil)
+	require.NoError(t, err)
 
 	rtExpected, err := rt.Version()
 	require.NoError(t, err)
@@ -458,10 +458,10 @@ func TestService_HandleSubmittedExtrinsic(t *testing.T) {
 	header, err := types.NewHeader(block.Header.Hash(), common.Hash{}, common.Hash{}, big.NewInt(block.Header.Number.Int64()+1), types.NewEmptyDigest())
 	require.NoError(t, err)
 
-	rt, ok := s.blockState.GetRuntime(nil)
-	require.True(t, ok)
+	rt, err := s.blockState.GetRuntime(nil)
+	require.NoError(t, err)
 
-	//initialise block header
+	// Initialise block header
 	err = rt.InitializeBlock(header)
 	require.NoError(t, err)
 
@@ -507,8 +507,8 @@ func TestService_HandleRuntimeChanges(t *testing.T) {
 	ts, err := s.storageState.TrieState(nil) // Pass genesis root
 	require.NoError(t, err)
 
-	parentRt, ok := s.blockState.GetRuntime(&hash)
-	require.True(t, ok)
+	parentRt, err := s.blockState.GetRuntime(&hash)
+	require.NoError(t, err)
 
 	v, err := parentRt.Version()
 	require.NoError(t, err)
@@ -529,15 +529,15 @@ func TestService_HandleRuntimeChanges(t *testing.T) {
 	require.NoError(t, err)
 
 	// bhash1 runtime should not be updated
-	rt, ok := s.blockState.GetRuntime(&bhash1)
-	require.True(t, ok)
+	rt, err := s.blockState.GetRuntime(&bhash1)
+	require.NoError(t, err)
 
 	v, err = rt.Version()
 	require.NoError(t, err)
 	require.Equal(t, v.SpecVersion(), currSpecVersion)
 
-	rt, ok = s.blockState.GetRuntime(&rtUpdateBhash)
-	require.True(t, ok)
+	rt, err = s.blockState.GetRuntime(&rtUpdateBhash)
+	require.NoError(t, err)
 
 	v, err = rt.Version()
 	require.NoError(t, err)
@@ -555,8 +555,8 @@ func TestService_HandleCodeSubstitutes(t *testing.T) {
 		blockHash: common.BytesToHex(testRuntime),
 	}
 
-	rt, ok := s.blockState.GetRuntime(nil)
-	require.True(t, ok)
+	rt, err := s.blockState.GetRuntime(nil)
+	require.NoError(t, err)
 
 	s.blockState.StoreRuntime(blockHash, rt)
 
@@ -566,23 +566,42 @@ func TestService_HandleCodeSubstitutes(t *testing.T) {
 	require.Equal(t, blockHash, codSub)
 }
 
-//func TestService_HandleRuntimeChangesAfterCodeSubstitutes(t *testing.T) {
-//	s := NewTestService(t, nil)
-//	codeHashBefore := s.codeHash
-//	blockHash := common.MustHexToHash("0x86aa36a140dfc449c30dbce16ce0fea33d5c3786766baa764e33f336841b9e29") // hash for known test code substitution
-//
-//	err := s.handleCodeSubstitution(blockHash)
-//	require.NoError(t, err)
-//	require.Equal(t, codeHashBefore, s.codeHash) // codeHash should remain unchanged after code substitute
-//
-//	testRuntime, err := ioutil.ReadFile(runtime.POLKADOT_RUNTIME_FP)
-//	require.NoError(t, err)
-//
-//	ts, err := s.storageState.TrieState(nil)
-//	require.NoError(t, err)
-//
-//	ts.Set(common.CodeKey, testRuntime)
-//	err = s.handleRuntimeChanges(ts)
-//	require.NoError(t, err)
-//	require.NotEqualf(t, codeHashBefore, s.codeHash, "expected different code hash after runtime update") // codeHash should change after runtime change
-//}
+func TestService_HandleRuntimeChangesAfterCodeSubstitutes(t *testing.T) {
+	s := NewTestService(t, nil)
+
+	parentRt, err := s.blockState.GetRuntime(nil)
+	require.NoError(t, err)
+
+	codeHashBefore := parentRt.GetCodeHash()
+	blockHash := common.MustHexToHash("0x86aa36a140dfc449c30dbce16ce0fea33d5c3786766baa764e33f336841b9e29") // hash for known test code substitution
+
+	newBlock := &types.Block{
+		Header: &types.Header{
+			ParentHash: blockHash,
+			Number:     big.NewInt(1),
+			Digest:     types.Digest{utils.NewMockDigestItem(1)}},
+		Body: types.NewBody([]byte("Updated Runtime")),
+	}
+
+	err = s.handleCodeSubstitution(blockHash)
+	require.NoError(t, err)
+	require.Equal(t, codeHashBefore, parentRt.GetCodeHash()) // codeHash should remain unchanged after code substitute
+
+	testRuntime, err := ioutil.ReadFile(runtime.POLKADOT_RUNTIME_FP)
+	require.NoError(t, err)
+
+	ts, err := s.storageState.TrieState(nil)
+	require.NoError(t, err)
+
+	ts.Set(common.CodeKey, testRuntime)
+	rtUpdateBhash := newBlock.Header.Hash()
+
+	// update runtime for new block
+	err = s.blockState.HandleRuntimeChanges(ts, parentRt, rtUpdateBhash)
+	require.NoError(t, err)
+
+	rt, err := s.blockState.GetRuntime(&rtUpdateBhash)
+	require.NoError(t, err)
+
+	require.NotEqualf(t, codeHashBefore, rt.GetCodeHash(), "expected different code hash after runtime update") // codeHash should change after runtime change
+}

dot/core/test_helpers.go

@@ -123,13 +123,10 @@ func NewTestService(t *testing.T, cfg *Config) *Service {
 		rtCfg.CodeHash, err = cfg.StorageState.LoadCodeHash(nil)
 		require.NoError(t, err)
 
-		rt, err := wasmer.NewRuntimeFromGenesis(gen, rtCfg)
+		cfg.Runtime, err = wasmer.NewRuntimeFromGenesis(gen, rtCfg)
 		require.NoError(t, err)
-
-		cfg.BlockState.StoreRuntime(cfg.BlockState.BestBlockHash(), rt)
-	} else {
-		cfg.BlockState.StoreRuntime(cfg.BlockState.BestBlockHash(), cfg.Runtime)
 	}
+	cfg.BlockState.StoreRuntime(cfg.BlockState.BestBlockHash(), cfg.Runtime)
 
 	if cfg.Network == nil {
 		config := &network.Config{

dot/node.go

@@ -34,6 +34,7 @@ import (
 	"github.com/ChainSafe/gossamer/dot/state/pruner"
 	"github.com/ChainSafe/gossamer/dot/telemetry"
 	"github.com/ChainSafe/gossamer/dot/types"
+	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/genesis"
 	"github.com/ChainSafe/gossamer/lib/keystore"
 	"github.com/ChainSafe/gossamer/lib/runtime"
@@ -451,14 +452,20 @@ func (n *Node) Stop() {
 func loadRuntime(cfg *Config, stateSrvc *state.Service, ks *keystore.GlobalKeystore, net *network.Service) error {
 	blocks := stateSrvc.Block.GetAllBlocks()
 	runtimeCode := make(map[string]runtime.Instance)
-	for _, hash := range blocks {
-		code, err := stateSrvc.Storage.GetStorageByBlockHash(hash, []byte(":code"))
+	for i := range blocks {
+		hash := &blocks[i]
+		code, err := stateSrvc.Storage.GetStorageByBlockHash(*hash, []byte(":code"))
 		if err != nil {
 			return err
 		}
 
-		if rt, ok := runtimeCode[string(code)]; ok {
-			stateSrvc.Block.StoreRuntime(hash, rt)
+		codeHash, err := common.Blake2bHash(code)
+		if err != nil {
+			return err
+		}
+
+		if rt, ok := runtimeCode[codeHash.String()]; ok {
+			stateSrvc.Block.StoreRuntime(*hash, rt)
 			continue
 		}
 
@@ -467,7 +474,7 @@ func loadRuntime(cfg *Config, stateSrvc *state.Service, ks *keystore.GlobalKeyst
 			return err
 		}
 
-		runtimeCode[string(code)] = rt
+		runtimeCode[codeHash.String()] = rt
 	}
 
 	return nil

dot/rpc/modules/state_test.go

@@ -444,8 +444,8 @@ func setupStateModule(t *testing.T) (*StateModule, *common.Hash, *common.Hash) {
 	err = chain.Block.AddBlock(b)
 	require.NoError(t, err)
 
-	rt, ok := chain.Block.GetRuntime(&b.Header.ParentHash)
-	require.True(t, ok)
+	rt, err := chain.Block.GetRuntime(&b.Header.ParentHash)
+	require.NoError(t, err)
 
 	chain.Block.StoreRuntime(b.Header.Hash(), rt)
 

dot/rpc/websocket_test.go

@@ -17,7 +17,6 @@ package rpc
 
 import (
 	"flag"
-	"fmt"
 	"log"
 	"net/url"
 	"testing"
@@ -92,8 +91,6 @@ func TestHTTPServer_ServeHTTP(t *testing.T) {
 
 		_, message, err := c.ReadMessage()
 		require.Nil(t, err)
-		fmt.Println("message", string(message))
-		fmt.Println("item.expected", string(item.expected))
 		require.Equal(t, item.expected, message)
 	}
 }