From 5211f5103a25ec30eea4dd1ccb11fedfff01469c Mon Sep 17 00:00:00 2001
From: KonradStaniec <konrad.staniec@gmail.com>
Date: Thu, 2 Jan 2025 09:23:16 +0100
Subject: [PATCH] fix block tally in finality module

---
 x/finality/keeper/tallying.go      |  3 +-
 x/finality/keeper/tallying_test.go | 58 ++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/x/finality/keeper/tallying.go b/x/finality/keeper/tallying.go
index 3a11c82d..d4ff91ce 100644
--- a/x/finality/keeper/tallying.go
+++ b/x/finality/keeper/tallying.go
@@ -39,6 +39,7 @@ func (k Keeper) TallyBlocks(ctx context.Context) {
 	// - has finality providers, finalised: impossible to happen, panic
 	// - does not have finality providers, finalised: impossible to happen, panic
 	// After this for loop, the blocks since earliest activated height are either finalised or non-finalisable
+finalizationLoop:
 	for i := startHeight; i <= uint64(sdkCtx.HeaderInfo().Height); i++ {
 		ib, err := k.GetBlock(ctx, i)
 		if err != nil {
@@ -58,7 +59,7 @@ func (k Keeper) TallyBlocks(ctx context.Context) {
 			} else {
 				// if not, then this block and all subsequent blocks should not be finalised
 				// thus, we need to break here
-				break
+				break finalizationLoop
 			}
 		case fpSet == nil && !ib.Finalized:
 			// does not have finality providers, non-finalised: not finalisable,
diff --git a/x/finality/keeper/tallying_test.go b/x/finality/keeper/tallying_test.go
index 6e2392d8..560d4d29 100644
--- a/x/finality/keeper/tallying_test.go
+++ b/x/finality/keeper/tallying_test.go
@@ -167,3 +167,61 @@ func giveNoQCToHeight(r *rand.Rand, ctx sdk.Context, fKeeper *keeper.Keeper, hei
 
 	return nil
 }
+
+func FuzzConsecutiveFinalization(f *testing.F) {
+	datagen.AddRandomSeedsToFuzzer(f, 10)
+
+	f.Fuzz(func(t *testing.T, seed int64) {
+		r := rand.New(rand.NewSource(seed))
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+
+		bsKeeper := types.NewMockBTCStakingKeeper(ctrl)
+		iKeeper := types.NewMockIncentiveKeeper(ctrl)
+		cKeeper := types.NewMockCheckpointingKeeper(ctrl)
+		fKeeper, ctx := keepertest.FinalityKeeper(t, bsKeeper, iKeeper, cKeeper)
+
+		// activate BTC staking protocol at a random height
+		activatedHeight := datagen.RandomInt(r, 10) + 1
+		numBlockToInspect := uint64(30)
+		// There will be a block in between activatedHeight and activatedHeight + numBlockToInspect
+		// that woud not get necessary votes to be finalised
+		firstNonFinalizedBlock := activatedHeight + 1 + datagen.RandomInt(r, 20)
+
+		for i := activatedHeight; i < activatedHeight+numBlockToInspect; i++ {
+			// index blocks
+			fKeeper.SetBlock(ctx, &types.IndexedBlock{
+				Height:    i,
+				AppHash:   datagen.GenRandomByteArray(r, 32),
+				Finalized: false,
+			})
+
+			if i == firstNonFinalizedBlock {
+				// this block does not have QC
+				err := giveNoQCToHeight(r, ctx, fKeeper, i)
+				require.NoError(t, err)
+			} else {
+				// this block has QC
+				err := giveQCToHeight(r, ctx, fKeeper, i)
+				require.NoError(t, err)
+			}
+		}
+
+		ctx = datagen.WithCtxHeight(ctx, activatedHeight+numBlockToInspect-1)
+		fKeeper.TallyBlocks(ctx)
+
+		// all blocks up to firstNonFinalizedBlock must be finalised
+		for i := activatedHeight; i < firstNonFinalizedBlock; i++ {
+			ib, err := fKeeper.GetBlock(ctx, i)
+			require.NoError(t, err)
+			require.True(t, ib.Finalized)
+		}
+
+		// all blocks from the firstNonFinalizedBlock must not be finalised
+		for i := firstNonFinalizedBlock; i < activatedHeight+numBlockToInspect; i++ {
+			ib, err := fKeeper.GetBlock(ctx, i)
+			require.NoError(t, err)
+			require.False(t, ib.Finalized)
+		}
+	})
+}