consensus/cometbft: Fail ImmutableState creation if version is missing

.changelog/6040.bugfix.md

@@ -0,0 +1,8 @@
+go/consensus/cometbft: Fail ImmutableState creation if version is missing
+
+Previously, when an `ImmutableState` was requested for a block version that
+didn't exist, the function would silently default to the latest available
+block. This could lead to inconsistencies since clients might receive state
+for a different block than expected. With this change, calls to create
+an `ImmutableState` for a missing version now explicitly fail with a
+"version not found" error, ensuring that such cases are handled properly.

go/consensus/cometbft/abci/state.go

@@ -278,28 +278,39 @@ func (s *applicationState) GetEpoch(ctx context.Context, blockHeight int64) (bea
 	return s.timeSource.GetEpoch(ctx, blockHeight)
 }
 
-func (s *applicationState) GetCurrentEpoch(ctx context.Context) (beacon.EpochTime, error) {
+func (s *applicationState) getCurrentEpoch(ctx context.Context, inCommitPhase bool) (beacon.EpochTime, error) {
 	blockHeight := s.BlockHeight()
 	if blockHeight == 0 {
 		return beacon.EpochInvalid, nil
 	}
+	latestStateHeight := blockHeight
+	// If we are not in the commit phase, then the latest state height is the next height.
+	// Otherwise, it is the current height.
+	if !inCommitPhase {
+		latestStateHeight++
+	}
+
 	// Check if there is an epoch transition scheduled for the current height. This should be taken
 	// into account when GetCurrentEpoch is called before the time keeping app does the transition.
-	future, err := s.timeSource.GetFutureEpoch(ctx, blockHeight+1)
+	future, err := s.timeSource.GetFutureEpoch(ctx, latestStateHeight)
 	if err != nil {
-		return beacon.EpochInvalid, fmt.Errorf("failed to get future epoch for height %d: %w", blockHeight+1, err)
+		return beacon.EpochInvalid, fmt.Errorf("failed to get future epoch for height %d: %w", latestStateHeight, err)
 	}
 	if future != nil && future.Height == blockHeight+1 {
 		return future.Epoch, nil
 	}
 
-	currentEpoch, err := s.timeSource.GetEpoch(ctx, blockHeight+1)
+	currentEpoch, err := s.timeSource.GetEpoch(ctx, latestStateHeight)
 	if err != nil {
-		return beacon.EpochInvalid, fmt.Errorf("failed to get epoch for height %d: %w", blockHeight+1, err)
+		return beacon.EpochInvalid, fmt.Errorf("failed to get epoch for height %d: %w", latestStateHeight, err)
 	}
 	return currentEpoch, nil
 }
 
+func (s *applicationState) GetCurrentEpoch(ctx context.Context) (beacon.EpochTime, error) {
+	return s.getCurrentEpoch(ctx, false)
+}
+
 func (s *applicationState) EpochChanged(ctx *api.Context) (bool, beacon.EpochTime) {
 	blockHeight := s.BlockHeight()
 	if blockHeight == 0 {

go/consensus/cometbft/abci/upgrade.go

@@ -17,7 +17,7 @@ func (mux *abciMux) maybeHaltForUpgrade() {
 	}
 
 	// Determine the epoch that will become the current epoch in the next block.
-	epoch, err := mux.state.GetCurrentEpoch(context.Background())
+	epoch, err := mux.state.getCurrentEpoch(context.Background(), true)
 	if err != nil {
 		mux.logger.Warn("failed to get current epoch",
 			"err", err,

go/consensus/cometbft/api/state.go

@@ -294,7 +294,10 @@ func NewImmutableState(ctx context.Context, state ApplicationQueryState, version
 	if state.BlockHeight() == 0 {
 		return nil, consensus.ErrNoCommittedBlocks
 	}
-	if version <= 0 || version > state.BlockHeight() {
+	if version > state.BlockHeight() {
+		return nil, consensus.ErrVersionNotFound
+	}
+	if version <= 0 {
 		version = state.BlockHeight()
 	}
 

go/consensus/cometbft/apps/staking/slashing.go

@@ -69,7 +69,7 @@ func onEvidenceByzantineConsensus(
 	// validator from being scheduled in the next epoch.
 	if penalty.FreezeInterval > 0 {
 		var epoch beacon.EpochTime
-		epoch, err = ctx.AppState().GetEpoch(context.Background(), ctx.BlockHeight()+1)
+		epoch, err = ctx.AppState().GetEpoch(context.Background(), ctx.BlockHeight())
 		if err != nil {
 			return err
 		}