go/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/beacon/tests/tester.go

@@ -35,6 +35,10 @@ func BeaconImplementationTests(t *testing.T, backend api.SetableBackend) {
 	latestEpoch, err := backend.GetEpoch(context.Background(), consensus.HeightLatest)
 	require.NoError(err, "GetEpoch")
 
+	// Querying epoch for a non-existing height should fail.
+	_, err = backend.GetEpoch(context.Background(), 100000000000)
+	require.ErrorIs(err, consensus.ErrVersionNotFound, "GetEpoch should fail for non-existing height")
+
 	var lastHeight int64
 	for epoch := api.EpochTime(0); epoch <= latestEpoch; epoch++ {
 		height, err := backend.GetEpochBlock(context.Background(), epoch)

go/consensus/cometbft/abci/state.go

@@ -283,17 +283,25 @@ func (s *applicationState) GetCurrentEpoch(ctx context.Context) (beacon.EpochTim
 	if blockHeight == 0 {
 		return beacon.EpochInvalid, nil
 	}
+
+	latestHeight := blockHeight
+	if abciCtx := api.FromCtx(ctx); abciCtx != nil {
+		// If request was made from an ABCI application context, then use blockHeight + 1, to fetch
+		// the epoch at current (future) height. See cometbft/api.NewImmutableState for details.
+		latestHeight++
+	}
+
 	// 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, latestHeight)
 	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", latestHeight, 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, latestHeight)
 	if err != nil {
 		return beacon.EpochInvalid, fmt.Errorf("failed to get epoch for height %d: %w", blockHeight+1, err)
 	}
@@ -305,8 +313,14 @@ func (s *applicationState) EpochChanged(ctx *api.Context) (bool, beacon.EpochTim
 	if blockHeight == 0 {
 		return false, beacon.EpochInvalid
 	}
+	latestHeight := blockHeight
+	if abciCtx := api.FromCtx(ctx); abciCtx != nil {
+		// If request was made from an ABCI application context, then use blockHeight + 1, to fetch
+		// the epoch at current (future) height. See cometbft/api.NewImmutableState for details.
+		latestHeight++
+	}
 
-	currentEpoch, err := s.timeSource.GetEpoch(ctx, blockHeight+1)
+	currentEpoch, err := s.timeSource.GetEpoch(ctx, latestHeight)
 	if err != nil {
 		s.logger.Error("EpochChanged: failed to get current epoch",
 			"err", err,
@@ -320,7 +334,7 @@ func (s *applicationState) EpochChanged(ctx *api.Context) (bool, beacon.EpochTim
 		return false, currentEpoch
 	}
 
-	previousEpoch, err := s.timeSource.GetEpoch(ctx, blockHeight)
+	previousEpoch, err := s.timeSource.GetEpoch(ctx, latestHeight-1)
 	if err != nil {
 		s.logger.Error("EpochChanged: failed to get previous 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,9 @@ 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)
+		// XXX: This should maybe use abciCtx and ctx.BlockHeight() + 1, but was
+		// kept like this to avoid potential consensus breaking changes at this time.
+		epoch, err = ctx.AppState().GetEpoch(context.Background(), ctx.BlockHeight())
 		if err != nil {
 			return err
 		}