Fix: Monitor/fix reported checkpoint BTC height query bugs

app/app.go

@@ -508,7 +508,7 @@ func NewBabylonApp(
 			privSigner.ClientCtx,
 		)
 	app.CheckpointingKeeper = *checkpointingKeeper.SetHooks(
-		checkpointingtypes.NewMultiCheckpointingHooks(app.EpochingKeeper.Hooks(), app.ZoneConciergeKeeper.Hooks()),
+		checkpointingtypes.NewMultiCheckpointingHooks(app.EpochingKeeper.Hooks(), app.ZoneConciergeKeeper.Hooks(), app.MonitorKeeper.Hooks()),
 	)
 	app.ZoneConciergeKeeper.SetCheckpointingKeeper(app.CheckpointingKeeper)
 

x/checkpointing/abci.go

@@ -26,7 +26,7 @@ func BeginBlocker(ctx sdk.Context, k keeper.Keeper, req abci.RequestBeginBlock)
 	if epoch.IsFirstBlock(ctx) {
 		err := k.InitValidatorBLSSet(ctx)
 		if err != nil {
-			panic(fmt.Errorf("failed to store validator BLS set"))
+			panic(fmt.Errorf("failed to store validator BLS set: %w", err))
 		}
 	}
 	if epoch.IsSecondBlock(ctx) {

x/checkpointing/keeper/keeper.go

@@ -3,7 +3,6 @@ package keeper
 import (
 	"errors"
 	"fmt"
-
 	txformat "github.com/babylonchain/babylon/btctxformatter"
 
 	"github.com/babylonchain/babylon/crypto/bls12381"
@@ -210,6 +209,11 @@ func (k Keeper) verifyCkptBytes(ctx sdk.Context, rawCheckpoint *txformat.RawBtcC
 
 	// can skip the checks if it is identical with the local checkpoint that is not accumulating
 	if ckptWithMeta.Ckpt.Equal(ckpt) && ckptWithMeta.Status != types.Accumulating {
+		// record verified checkpoint
+		err = k.AfterRawCheckpointBlsSigVerified(ctx, ckpt)
+		if err != nil {
+			return nil, fmt.Errorf("failed to record verified checkpoint of epoch %d for monitoring: %w", ckpt.EpochNum, err)
+		}
 		return ckptWithMeta, nil
 	}
 
@@ -244,7 +248,7 @@ func (k Keeper) verifyCkptBytes(ctx sdk.Context, rawCheckpoint *txformat.RawBtcC
 	// record verified checkpoint
 	err = k.AfterRawCheckpointBlsSigVerified(ctx, ckpt)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to record verified checkpoint of epoch %d for monitoring: %w", ckpt.EpochNum, err)
 	}
 
 	// now the checkpoint's multi-sig is valid, if the lastcommithash is the
@@ -270,6 +274,10 @@ func (k Keeper) verifyCkptBytes(ctx sdk.Context, rawCheckpoint *txformat.RawBtcC
 	return nil, types.ErrConflictingCheckpoint
 }
 
+func (k *Keeper) SetEpochingKeeper(ek types.EpochingKeeper) {
+	k.epochingKeeper = ek
+}
+
 // SetCheckpointSubmitted sets the status of a checkpoint to SUBMITTED,
 // and records the associated state update in lifecycle
 func (k Keeper) SetCheckpointSubmitted(ctx sdk.Context, epoch uint64) {

x/checkpointing/types/hooks.go

@@ -49,7 +49,9 @@ func (h MultiCheckpointingHooks) AfterRawCheckpointFinalized(ctx sdk.Context, ep
 
 func (h MultiCheckpointingHooks) AfterRawCheckpointBlsSigVerified(ctx sdk.Context, ckpt *RawCheckpoint) error {
 	for i := range h {
-		return h[i].AfterRawCheckpointBlsSigVerified(ctx, ckpt)
+		if err := h[i].AfterRawCheckpointBlsSigVerified(ctx, ckpt); err != nil {
+			return err
+		}
 	}
 	return nil
 }

x/monitor/keeper/grpc_query_params.go

@@ -9,6 +9,13 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+// Querier is used as Keeper will have duplicate methods if used directly, and gRPC names take precedence over keeper
+type Querier struct {
+	Keeper
+}
+
+var _ types.QueryServer = Querier{}
+
 func (k Keeper) Params(c context.Context, req *types.QueryParamsRequest) (*types.QueryParamsResponse, error) {
 	if req == nil {
 		return nil, status.Error(codes.InvalidArgument, "invalid request")

x/monitor/keeper/grpc_query_test.go

@@ -0,0 +1,158 @@
+package keeper_test
+
+import (
+	"github.com/babylonchain/babylon/btctxformatter"
+	"github.com/babylonchain/babylon/testutil/datagen"
+	"github.com/babylonchain/babylon/testutil/mocks"
+	btclightclienttypes "github.com/babylonchain/babylon/x/btclightclient/types"
+	ckpttypes "github.com/babylonchain/babylon/x/checkpointing/types"
+	"github.com/babylonchain/babylon/x/epoching/testepoching"
+	types2 "github.com/babylonchain/babylon/x/epoching/types"
+	monitorkeeper "github.com/babylonchain/babylon/x/monitor/keeper"
+	"github.com/babylonchain/babylon/x/monitor/types"
+	"github.com/cosmos/cosmos-sdk/baseapp"
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/require"
+	"math/rand"
+	"testing"
+)
+
+func FuzzQueryEndedEpochBtcHeight(f *testing.F) {
+	datagen.AddRandomSeedsToFuzzer(f, 10)
+	f.Fuzz(func(t *testing.T, seed int64) {
+		rand.Seed(seed)
+		// a genesis validator is generated for setup
+		helper := testepoching.NewHelper(t)
+		lck := helper.App.BTCLightClientKeeper
+		mk := helper.App.MonitorKeeper
+		ek := helper.EpochingKeeper
+		querier := monitorkeeper.Querier{Keeper: mk}
+		queryHelper := baseapp.NewQueryServerTestHelper(helper.Ctx, helper.App.InterfaceRegistry())
+		types.RegisterQueryServer(queryHelper, querier)
+		queryClient := types.NewQueryClient(queryHelper)
+
+		// BeginBlock of block 1, and thus entering epoch 1
+		ctx := helper.BeginBlock()
+		epoch := ek.GetEpoch(ctx)
+		require.Equal(t, uint64(1), epoch.EpochNumber)
+
+		// Insert header tree
+		tree := datagen.NewBTCHeaderTree()
+		root := lck.GetBaseBTCHeader(ctx)
+		tree.Add(root, nil)
+		tree.GenRandomBTCHeaderTree(1, 10, root, func(header *btclightclienttypes.BTCHeaderInfo) bool {
+			err := lck.InsertHeader(ctx, header.Header)
+			require.NoError(t, err)
+			return true
+		})
+
+		// EndBlock of block 1
+		ctx = helper.EndBlock()
+
+		// go to BeginBlock of block 11, and thus entering epoch 2
+		for i := uint64(0); i < ek.GetParams(ctx).EpochInterval; i++ {
+			ctx = helper.GenAndApplyEmptyBlock()
+		}
+		epoch = ek.GetEpoch(ctx)
+		require.Equal(t, uint64(2), epoch.EpochNumber)
+
+		// query epoch 0 ended BTC light client height, should return base height
+		req := types.QueryEndedEpochBtcHeightRequest{
+			EpochNum: 0,
+		}
+		resp, err := queryClient.EndedEpochBtcHeight(ctx, &req)
+		require.NoError(t, err)
+		require.Equal(t, lck.GetBaseBTCHeader(ctx).Height, resp.BtcLightClientHeight)
+
+		// query epoch 1 ended BTC light client height, should return tip height
+		req = types.QueryEndedEpochBtcHeightRequest{
+			EpochNum: 1,
+		}
+		resp, err = queryClient.EndedEpochBtcHeight(ctx, &req)
+		require.NoError(t, err)
+		require.Equal(t, lck.GetTipInfo(ctx).Height, resp.BtcLightClientHeight)
+	})
+}
+
+func FuzzQueryReportedCheckpointBtcHeight(f *testing.F) {
+	datagen.AddRandomSeedsToFuzzer(f, 10)
+	f.Fuzz(func(t *testing.T, seed int64) {
+		rand.Seed(seed)
+		// a genesis validator is generated for setup
+		helper := testepoching.NewHelper(t)
+		ctl := gomock.NewController(t)
+		defer ctl.Finish()
+		lck := helper.App.BTCLightClientKeeper
+		mk := helper.App.MonitorKeeper
+		ek := helper.EpochingKeeper
+		ck := helper.App.CheckpointingKeeper
+		mockEk := mocks.NewMockEpochingKeeper(ctl)
+		ck.SetEpochingKeeper(mockEk)
+		querier := monitorkeeper.Querier{Keeper: mk}
+		queryHelper := baseapp.NewQueryServerTestHelper(helper.Ctx, helper.App.InterfaceRegistry())
+		types.RegisterQueryServer(queryHelper, querier)
+		queryClient := types.NewQueryClient(queryHelper)
+
+		// BeginBlock of block 1, and thus entering epoch 1
+		ctx := helper.BeginBlock()
+		epoch := ek.GetEpoch(ctx)
+		require.Equal(t, uint64(1), epoch.EpochNumber)
+
+		// Insert header tree
+		tree := datagen.NewBTCHeaderTree()
+		root := lck.GetBaseBTCHeader(ctx)
+		tree.Add(root, nil)
+		tree.GenRandomBTCHeaderTree(1, 10, root, func(header *btclightclienttypes.BTCHeaderInfo) bool {
+			err := lck.InsertHeader(ctx, header.Header)
+			require.NoError(t, err)
+			return true
+		})
+
+		// Add checkpoint
+		valBlsSet, privKeys := datagen.GenerateValidatorSetWithBLSPrivKeys(int(datagen.RandomIntOtherThan(0, 10)))
+		valSet := make([]types2.Validator, len(valBlsSet.ValSet))
+		for i, val := range valBlsSet.ValSet {
+			valSet[i] = types2.Validator{
+				Addr:  []byte(val.ValidatorAddress),
+				Power: int64(val.VotingPower),
+			}
+			err := ck.CreateRegistration(ctx, val.BlsPubKey, []byte(val.ValidatorAddress))
+			require.NoError(t, err)
+		}
+		mockCkptWithMeta := &ckpttypes.RawCheckpointWithMeta{Ckpt: datagen.GenerateLegitimateRawCheckpoint(privKeys)}
+		mockEk.EXPECT().GetValidatorSet(gomock.Any(), gomock.Eq(mockCkptWithMeta.Ckpt.EpochNum)).Return(valSet).AnyTimes()
+		// make sure voting power is always sufficient
+		mockEk.EXPECT().GetTotalVotingPower(gomock.Any(), gomock.Eq(mockCkptWithMeta.Ckpt.EpochNum)).Return(int64(0)).AnyTimes()
+		err := ck.AddRawCheckpoint(
+			ctx,
+			mockCkptWithMeta,
+		)
+		require.NoError(t, err)
+
+		// Verify checkpoint
+		btcCkpt := btctxformatter.RawBtcCheckpoint{
+			Epoch:            mockCkptWithMeta.Ckpt.EpochNum,
+			LastCommitHash:   *mockCkptWithMeta.Ckpt.LastCommitHash,
+			BitMap:           mockCkptWithMeta.Ckpt.Bitmap,
+			SubmitterAddress: datagen.GenRandomByteArray(btctxformatter.AddressLength),
+			BlsSig:           *mockCkptWithMeta.Ckpt.BlsMultiSig,
+		}
+		err = ck.VerifyCheckpoint(ctx, btcCkpt)
+		require.NoError(t, err)
+
+		// query reported checkpoint BTC light client height
+		req := types.QueryReportedCheckpointBtcHeightRequest{
+			CkptHash: mockCkptWithMeta.Ckpt.HashStr(),
+		}
+		resp, err := queryClient.ReportedCheckpointBtcHeight(ctx, &req)
+		require.NoError(t, err)
+		require.Equal(t, lck.GetTipInfo(ctx).Height, resp.BtcLightClientHeight)
+
+		// query not reported checkpoint BTC light client height, should expect an ErrCheckpointNotReported
+		req = types.QueryReportedCheckpointBtcHeightRequest{
+			CkptHash: datagen.GenRandomHexStr(32),
+		}
+		_, err = queryClient.ReportedCheckpointBtcHeight(ctx, &req)
+		require.ErrorIs(t, err, types.ErrCheckpointNotReported)
+	})
+}

x/monitor/keeper/hooks.go

@@ -16,8 +16,7 @@ type Hooks struct {
 	k Keeper
 }
 
-var _ HandledHooks = Hooks{}
-
+// Create new distribution hooks
 func (k Keeper) Hooks() Hooks { return Hooks{k} }
 
 func (h Hooks) AfterEpochBegins(ctx sdk.Context, epoch uint64) {}

x/monitor/keeper/keeper.go

@@ -65,18 +65,18 @@ func (k Keeper) updateBtcLightClientHeightForCheckpoint(ctx sdk.Context, ckpt *c
 	store := ctx.KVStore(k.storeKey)
 	ckptHashStr := ckpt.HashStr()
 
+	storeKey, err := types.GetCheckpointReportedLightClientHeightKey(ckptHashStr)
+	if err != nil {
+		return err
+	}
+
 	// if the checkpoint exists, meaning an earlier checkpoint with a lower btc height is already recorded
 	// we should keep the lower btc height in the store
-	if store.Has([]byte(ckptHashStr)) {
+	if store.Has(storeKey) {
 		k.Logger(ctx).With("module", fmt.Sprintf("checkpoint %s is already recorded", ckptHashStr))
 		return nil
 	}
 
-	storeKey, err := types.GetCheckpointReportedLightClientHeightKey(ckptHashStr)
-	if err != nil {
-		return err
-	}
-
 	currentTipHeight := k.btcLightClientKeeper.GetTipInfo(ctx).Height
 	store.Set(storeKey, sdk.Uint64ToBigEndian(currentTipHeight))
 
@@ -98,6 +98,10 @@ func (k Keeper) removeCheckpointRecord(ctx sdk.Context, ckpt *ckpttypes.RawCheck
 }
 
 func (k Keeper) LightclientHeightAtEpochEnd(ctx sdk.Context, epoch uint64) (uint64, error) {
+	if epoch == 0 {
+		return k.btcLightClientKeeper.GetBaseBTCHeader(ctx).Height, nil
+	}
+
 	store := ctx.KVStore(k.storeKey)
 
 	btcHeightBytes := store.Get(types.GetEpochEndLightClientHeightKey(epoch))

x/monitor/types/expected_keepers.go

@@ -20,4 +20,5 @@ type BankKeeper interface {
 
 type BTCLightClientKeeper interface {
 	GetTipInfo(ctx sdk.Context) *lc.BTCHeaderInfo
+	GetBaseBTCHeader(ctx sdk.Context) *lc.BTCHeaderInfo
 }