From a8ff617773542b815136f63f5e8c01b221fe0ff6 Mon Sep 17 00:00:00 2001
From: Callum Waters <cmwaters19@gmail.com>
Date: Thu, 21 Oct 2021 13:40:43 +0200
Subject: [PATCH] state: add height assertion to rollback function (#7143)

---
 internal/state/rollback.go      | 17 +++++++
 internal/state/rollback_test.go | 79 ++++++++++++++-------------------
 2 files changed, 50 insertions(+), 46 deletions(-)

diff --git a/internal/state/rollback.go b/internal/state/rollback.go
index 6e13da0e2..e78957b02 100644
--- a/internal/state/rollback.go
+++ b/internal/state/rollback.go
@@ -19,6 +19,23 @@ func Rollback(bs BlockStore, ss Store) (int64, []byte, error) {
 		return -1, nil, errors.New("no state found")
 	}
 
+	height := bs.Height()
+
+	// NOTE: persistence of state and blocks don't happen atomically. Therefore it is possible that
+	// when the user stopped the node the state wasn't updated but the blockstore was. In this situation
+	// we don't need to rollback any state and can just return early
+	if height == invalidState.LastBlockHeight+1 {
+		return invalidState.LastBlockHeight, invalidState.AppHash, nil
+	}
+
+	// If the state store isn't one below nor equal to the blockstore height than this violates the
+	// invariant
+	if height != invalidState.LastBlockHeight {
+		return -1, nil, fmt.Errorf("statestore height (%d) is not one below or equal to blockstore height (%d)",
+			invalidState.LastBlockHeight, height)
+	}
+
+	// state store height is equal to blockstore height. We're good to proceed with rolling back state
 	rollbackHeight := invalidState.LastBlockHeight
 	rollbackBlock := bs.LoadBlockMeta(rollbackHeight)
 	if rollbackBlock == nil {
diff --git a/internal/state/rollback_test.go b/internal/state/rollback_test.go
index ae5c8ee84..e782b4d89 100644
--- a/internal/state/rollback_test.go
+++ b/internal/state/rollback_test.go
@@ -14,42 +14,14 @@ import (
 )
 
 func TestRollback(t *testing.T) {
-	stateStore := state.NewStore(dbm.NewMemDB())
-	blockStore := &mocks.BlockStore{}
 	var (
 		height     int64  = 100
 		appVersion uint64 = 10
 	)
-
-	valSet, _ := factory.RandValidatorSet(5, 10)
-
-	params := types.DefaultConsensusParams()
-	params.Version.AppVersion = appVersion
-	newParams := types.DefaultConsensusParams()
-	newParams.Block.MaxBytes = 10000
-
-	initialState := state.State{
-		Version: state.Version{
-			Consensus: version.Consensus{
-				Block: version.BlockProtocol,
-				App:   10,
-			},
-			Software: version.TMVersion,
-		},
-		ChainID:                          factory.DefaultTestChainID,
-		InitialHeight:                    10,
-		LastBlockID:                      factory.MakeBlockID(),
-		AppHash:                          factory.RandomHash(),
-		LastResultsHash:                  factory.RandomHash(),
-		LastBlockHeight:                  height,
-		LastValidators:                   valSet,
-		Validators:                       valSet.CopyIncrementProposerPriority(1),
-		NextValidators:                   valSet.CopyIncrementProposerPriority(2),
-		LastHeightValidatorsChanged:      height + 1,
-		ConsensusParams:                  *params,
-		LastHeightConsensusParamsChanged: height + 1,
-	}
-	require.NoError(t, stateStore.Bootstrap(initialState))
+	blockStore := &mocks.BlockStore{}
+	stateStore := setupStateStore(t, height)
+	initialState, err := stateStore.Load()
+	require.NoError(t, err)
 
 	height++
 	block := &types.BlockMeta{
@@ -61,9 +33,13 @@ func TestRollback(t *testing.T) {
 		},
 	}
 	blockStore.On("LoadBlockMeta", height).Return(block)
+	blockStore.On("Height").Return(height)
 
+	// perform the rollback over a version bump
 	appVersion++
+	newParams := types.DefaultConsensusParams()
 	newParams.Version.AppVersion = appVersion
+	newParams.Block.MaxBytes = 1000
 	nextState := initialState.Copy()
 	nextState.LastBlockHeight = height
 	nextState.Version.Consensus.App = appVersion
@@ -102,19 +78,34 @@ func TestRollbackNoState(t *testing.T) {
 }
 
 func TestRollbackNoBlocks(t *testing.T) {
-	stateStore := state.NewStore(dbm.NewMemDB())
+	const height = int64(100)
+	stateStore := setupStateStore(t, height)
 	blockStore := &mocks.BlockStore{}
-	var (
-		height     int64  = 100
-		appVersion uint64 = 10
-	)
+	blockStore.On("Height").Return(height)
+	blockStore.On("LoadBlockMeta", height).Return(nil)
+
+	_, _, err := state.Rollback(blockStore, stateStore)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "block at height 100 not found")
+}
 
+func TestRollbackDifferentStateHeight(t *testing.T) {
+	const height = int64(100)
+	stateStore := setupStateStore(t, height)
+	blockStore := &mocks.BlockStore{}
+	blockStore.On("Height").Return(height + 2)
+
+	_, _, err := state.Rollback(blockStore, stateStore)
+	require.Error(t, err)
+	require.Equal(t, err.Error(), "statestore height (100) is not one below or equal to blockstore height (102)")
+}
+
+func setupStateStore(t *testing.T, height int64) state.Store {
+	stateStore := state.NewStore(dbm.NewMemDB())
 	valSet, _ := factory.RandValidatorSet(5, 10)
 
 	params := types.DefaultConsensusParams()
-	params.Version.AppVersion = appVersion
-	newParams := types.DefaultConsensusParams()
-	newParams.Block.MaxBytes = 10000
+	params.Version.AppVersion = 10
 
 	initialState := state.State{
 		Version: state.Version{
@@ -137,10 +128,6 @@ func TestRollbackNoBlocks(t *testing.T) {
 		ConsensusParams:                  *params,
 		LastHeightConsensusParamsChanged: height + 1,
 	}
-	require.NoError(t, stateStore.Save(initialState))
-	blockStore.On("LoadBlockMeta", height).Return(nil)
-
-	_, _, err := state.Rollback(blockStore, stateStore)
-	require.Error(t, err)
-	require.Contains(t, err.Error(), "block at height 100 not found")
+	require.NoError(t, stateStore.Bootstrap(initialState))
+	return stateStore
 }