From 85c3b96f82125eb39246cd0063728c7df6c5e3c3 Mon Sep 17 00:00:00 2001
From: Ekaterina Pavlova <ekt@morphbits.io>
Date: Fri, 22 Nov 2024 11:33:05 +0300
Subject: [PATCH] core: fix restoring chain with `StateRootInHeader = true`

Close #3597

Signed-off-by: Ekaterina Pavlova <ekt@morphbits.io>
---
 pkg/core/blockchain_core_test.go | 71 ++++++++++++++++++++++++++++++++
 pkg/core/stateroot/module.go     |  8 ++--
 2 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/pkg/core/blockchain_core_test.go b/pkg/core/blockchain_core_test.go
index b2a10f21a..1b0c4c00d 100644
--- a/pkg/core/blockchain_core_test.go
+++ b/pkg/core/blockchain_core_test.go
@@ -12,9 +12,11 @@ import (
 	"github.com/nspcc-dev/neo-go/internal/testchain"
 	"github.com/nspcc-dev/neo-go/pkg/config"
 	"github.com/nspcc-dev/neo-go/pkg/core/block"
+	"github.com/nspcc-dev/neo-go/pkg/core/chaindump"
 	"github.com/nspcc-dev/neo-go/pkg/core/state"
 	"github.com/nspcc-dev/neo-go/pkg/core/storage"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
+	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/trigger"
 	"github.com/nspcc-dev/neo-go/pkg/util"
@@ -414,3 +416,72 @@ func TestNewBlockchain_InitHardforks(t *testing.T) {
 		}, bc.GetConfig().Hardforks)
 	})
 }
+
+type nopCloserStorage struct {
+	storage.Store
+}
+
+func (nopCloserStorage) Close() error {
+	return nil
+}
+
+func TestBlockchainRestoreStateRootInHeader(t *testing.T) {
+	bc := newTestChainWithCustomCfg(t, func(c *config.Config) {
+		c.ProtocolConfiguration.StateRootInHeader = true
+	})
+	require.NoError(t, bc.AddBlock(bc.newBlock()))
+	require.NoError(t, bc.AddBlock(bc.newBlock()))
+	require.NoError(t, bc.AddBlock(bc.newBlock()))
+
+	w := io.NewBufBinWriter()
+	require.NoError(t, chaindump.Dump(bc, w.BinWriter, 0, 4))
+	require.NoError(t, w.Err)
+
+	data := w.Bytes()
+	fcfg := func(c *config.Config) { c.ProtocolConfiguration.StateRootInHeader = true }
+	initChain := func(st storage.Store) *Blockchain {
+		chain := initTestChain(t, st, fcfg)
+		go chain.Run()
+		return chain
+	}
+
+	t.Run("empty MemoryStore", func(t *testing.T) {
+		bc := initChain(storage.NewMemoryStore())
+		r := io.NewBinReaderFromBuf(data)
+		require.NoError(t, chaindump.Restore(bc, r, 0, 4, nil))
+		bc.Close()
+	})
+	t.Run("not empty MemoryStore, one block", func(t *testing.T) {
+		st := nopCloserStorage{Store: storage.NewMemoryStore()}
+
+		bc := initChain(st)
+		r := io.NewBinReaderFromBuf(data)
+		require.NoError(t, chaindump.Restore(bc, r, 0, 1, nil))
+		expected := bc.stateRoot.CurrentLocalStateRoot()
+		bc.Close()
+
+		bc = initChain(st)
+		actual := bc.stateRoot.CurrentLocalStateRoot()
+		require.Equal(t, expected, actual)
+		r = io.NewBinReaderFromBuf(data)
+		require.NoError(t, chaindump.Restore(bc, r, 1, 1, nil))
+		bc.Close()
+	})
+	t.Run("not empty MemoryStore, multiple blocks", func(t *testing.T) {
+		st := nopCloserStorage{Store: storage.NewMemoryStore()}
+		{
+			bc := initChain(st)
+			r := io.NewBinReaderFromBuf(data)
+			require.NoError(t, chaindump.Restore(bc, r, 0, 2, nil))
+			expected := bc.stateRoot.CurrentLocalStateRoot()
+			bc.Close()
+
+			bc = initChain(st)
+			actual := bc.stateRoot.CurrentLocalStateRoot()
+			require.Equal(t, expected, actual)
+			r = io.NewBinReaderFromBuf(data)
+			require.NoError(t, chaindump.Restore(bc, r, 2, 1, nil))
+			bc.Close()
+		}
+	})
+}
diff --git a/pkg/core/stateroot/module.go b/pkg/core/stateroot/module.go
index 3b109b4a6..e06c8a328 100644
--- a/pkg/core/stateroot/module.go
+++ b/pkg/core/stateroot/module.go
@@ -183,13 +183,15 @@ func (s *Module) Init(height uint32) error {
 		updateStateHeightMetric(h)
 	}
 
-	if height == 0 {
+	if s.mpt == nil {
 		s.mpt = mpt.NewTrie(nil, s.mode, s.Store)
-		s.currentLocal.Store(util.Uint256{})
-		return nil
 	}
 	r, err := s.getStateRoot(makeStateRootKey(height))
 	if err != nil {
+		if height == 0 {
+			s.currentLocal.Store(util.Uint256{})
+			return nil
+		}
 		return err
 	}
 	s.currentLocal.Store(r.Root)