tests: Add TestStateReexec

node/cn/state_accessor.go

@@ -116,7 +116,7 @@ func (cn *CN) stateAtBlock(block *types.Block, reexec uint64, base *state.StateD
 	for current.NumberU64() < origin {
 		// Print progress logs if long enough time elapsed
 		if report && time.Since(logged) > 8*time.Second {
-			logger.Info("Regenerating historical state", "block", current.NumberU64()+1, "target", origin, "remaining", origin-block.NumberU64()-1, "elapsed", time.Since(start))
+			logger.Info("Regenerating historical state", "block", current.NumberU64()+1, "target", origin, "remaining", origin-current.NumberU64()-1, "elapsed", time.Since(start))
 			logged = time.Now()
 		}
 		// Quit the state regeneration if time limit exceeds

tests/state_reexec_test.go

@@ -0,0 +1,163 @@
+package tests
+
+import (
+	"context"
+	"math/big"
+	"testing"
+
+	"github.com/kaiachain/kaia/accounts/abi/bind"
+	"github.com/kaiachain/kaia/accounts/abi/bind/backends"
+	"github.com/kaiachain/kaia/blockchain"
+	"github.com/kaiachain/kaia/blockchain/system"
+	"github.com/kaiachain/kaia/common"
+	"github.com/kaiachain/kaia/consensus/istanbul"
+	testcontract "github.com/kaiachain/kaia/contracts/contracts/testing/reward"
+	"github.com/kaiachain/kaia/log"
+	"github.com/kaiachain/kaia/params"
+	"github.com/kaiachain/kaia/reward"
+	"github.com/kaiachain/kaia/storage/database"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// Test State Regeneration (reexecution) after pruning state trie nodes.
+// This test ensures that the state regeneration yields the exact same state as the block's stateRoot.
+// Post-Kaia engine.Finalize() relies on the state trie to calculate rewards, so the state regeneration
+// can be interfered. This test ensures that the state regeneration is robust against such interference.
+func TestStateReexec(t *testing.T) {
+	log.EnableLogForTest(log.LvlCrit, log.LvlWarn)
+
+	// Test parameters
+	var (
+		numNodes = 1
+		forkNum  = big.NewInt(4)
+		nodeId   = bind.NewKeyedTransactor(deriveTestAccount(0)).From
+		owner    = bind.NewKeyedTransactor(deriveTestAccount(5))
+
+		config = testStateReexec_config(forkNum)
+		alloc  = testStateReexec_alloc(t, owner, nodeId)
+	)
+
+	// Start the chain
+	ctx, err := newBlockchainTestContext(&blockchainTestOverrides{
+		numNodes:    numNodes,
+		numAccounts: 8,
+		config:      config,
+		alloc:       alloc,
+	})
+	require.Nil(t, err)
+	ctx.Subscribe(t, func(ev *blockchain.ChainEvent) {
+		b := ev.Block
+		t.Logf("block[%3d] stateRoot=%x", b.NumberU64(), b.Header().Root)
+	})
+	ctx.Start()
+	defer ctx.Cleanup()
+
+	ctx.WaitBlock(t, 6)
+
+	// Clear staking cache to force GetStakingInfo post-Kaia to utilize the state trie.
+	reward.PurgeStakingInfoCache()
+
+	// Delete state roots to force historical state regeneration
+	testStateReexec_prune(t, ctx.nodes[0], []uint64{2, 3, 4, 5})
+	testStateReexec_run(t, ctx.nodes[0], 3) // pre-kaia
+
+	testStateReexec_prune(t, ctx.nodes[0], []uint64{2, 3, 4, 5})
+	testStateReexec_run(t, ctx.nodes[0], 5) // post-kaia
+}
+
+func testStateReexec_config(forkNum *big.Int) *params.ChainConfig {
+	config := blockchainTestChainConfig.Copy()
+	config.LondonCompatibleBlock = common.Big0
+	config.IstanbulCompatibleBlock = common.Big0
+	config.EthTxTypeCompatibleBlock = common.Big0
+	config.MagmaCompatibleBlock = common.Big0
+	config.KoreCompatibleBlock = common.Big0
+	config.ShanghaiCompatibleBlock = common.Big0
+	config.CancunCompatibleBlock = common.Big0
+	config.KaiaCompatibleBlock = forkNum
+
+	// Use WeightedRandom to test reward distribution based on StakingInfo
+	config.Istanbul.ProposerPolicy = uint64(istanbul.WeightedRandom)
+	// Set the reward ratio so that reward distribution is different from the 'all to proposer' fallback.
+	// If the GetStakingInfo() fails during state regen, the regenerated state would just give all
+	// rewards to the proposer, deviating from the actual historical state.
+	config.Governance.Reward.Ratio = "34/54/12"
+	return config
+}
+
+// Create a genesis state with an AddressBookMock
+func testStateReexec_alloc(t *testing.T, owner *bind.TransactOpts, nodeId common.Address) blockchain.GenesisAlloc {
+	// Create a simulated state with the mock contract populated.
+	var (
+		abookAddr   = system.AddressBookAddr
+		abookCode   = common.FromHex(testcontract.AddressBookMockBinRuntime)
+		stakingAddr = common.HexToAddress("0x1000")
+		rewardAddr  = common.HexToAddress("0x2000")
+		fund1Addr   = common.HexToAddress("0xa000")
+		fund2Addr   = common.HexToAddress("0xb000")
+
+		alloc = blockchain.GenesisAlloc{
+			owner.From:             {Balance: big.NewInt(params.KAIA)},
+			system.AddressBookAddr: {Balance: common.Big0, Code: abookCode},
+		}
+		db          = database.NewMemoryDBManager()
+		backend     = backends.NewSimulatedBackendWithDatabase(db, alloc, &params.ChainConfig{})
+		contract, _ = testcontract.NewAddressBookMockTransactor(abookAddr, backend)
+	)
+	_, err := contract.ConstructContract(owner, []common.Address{owner.From}, common.Big1)
+	backend.Commit()
+	require.Nil(t, err)
+
+	_, err = contract.RegisterCnStakingContract(owner, nodeId, stakingAddr, rewardAddr)
+	require.Nil(t, err)
+	_, err = contract.UpdatePocContract(owner, fund1Addr, common.Big1)
+	require.Nil(t, err)
+	_, err = contract.UpdateKirContract(owner, fund2Addr, common.Big1)
+	backend.Commit()
+	require.Nil(t, err)
+
+	_, err = contract.ActivateAddressBook(owner)
+	backend.Commit()
+	require.Nil(t, err)
+
+	// Copy contract storage from the simulated state to the genesis account.
+	abookStorage := make(map[common.Hash]common.Hash)
+	stateDB, _ := backend.BlockChain().State()
+	stateDB.ForEachStorage(abookAddr, func(key common.Hash, value common.Hash) bool {
+		abookStorage[key] = value
+		return true
+	})
+	return blockchain.GenesisAlloc{
+		abookAddr: {
+			Balance: common.Big0,
+			Code:    abookCode,
+			Storage: abookStorage,
+		},
+		stakingAddr: {
+			Balance: new(big.Int).Mul(big.NewInt(params.KAIA), big.NewInt(5_000_000)),
+		},
+	}
+}
+
+func testStateReexec_prune(t *testing.T, node *blockchainTestNode, nums []uint64) {
+	db := node.cn.ChainDB()
+
+	for _, num := range nums {
+		block := node.cn.BlockChain().GetBlockByNumber(num)
+		root := block.Header().Root
+		db.DeleteTrieNode(root.ExtendZero())
+	}
+}
+
+func testStateReexec_run(t *testing.T, node *blockchainTestNode, num uint64) {
+	block := node.cn.BlockChain().GetBlockByNumber(num)
+
+	state, err := node.cn.APIBackend.StateAtBlock(context.Background(), block, 10, nil, false, false)
+	require.Nil(t, err)
+	root, err := state.Commit(false)
+	require.Nil(t, err)
+
+	// Regenerated state must match the stored block's stateRoot
+	assert.Equal(t, block.Header().Root, root)
+}

tests/testutil_blockchain_test.go

@@ -56,6 +56,8 @@ type blockchainTestContext struct {
 	config       *params.ChainConfig
 	genesis      *blockchain.Genesis
 
+	savedStakingManager *reward.StakingManager
+
 	workspace string
 	nodes     []*blockchainTestNode
 }
@@ -261,6 +263,18 @@ func (ctx *blockchainTestContext) forEachNode(f func(*blockchainTestNode) error)
 }
 
 func (ctx *blockchainTestContext) Start() error {
+	// TODO: make StakingManager not singleton OR recreate new in cn.New()
+	// Manually re-wire StakingManager with the new blockchain.
+	// Because StakingManager is a singleton, it has to be part of one node. Here we use the first node.
+	ctx.savedStakingManager = reward.GetStakingManager()
+	reward.SetTestStakingManagerWithChain(
+		ctx.nodes[0].cn.BlockChain(),
+		ctx.nodes[0].cn.Governance(),
+		ctx.nodes[0].cn.ChainDB(),
+	)
+	reward.StakingManagerUnsubscribe()
+	reward.StakingManagerSubscribe() // re-subscribe to the new blockchain
+
 	return ctx.forEachNode(func(tn *blockchainTestNode) error {
 		return tn.cn.StartMining(false)
 	})
@@ -280,6 +294,7 @@ func (ctx *blockchainTestContext) Stop() error {
 	// other tests can use StakingManager as if it's fresh.
 	reward.PurgeStakingInfoCache()
 	blockchain.ClearMigrationPrerequisites()
+	reward.SetTestStakingManager(ctx.savedStakingManager)
 	return nil
 }