Do not panic on nil proof when handling votes (#186)

If not would receive finality vote with `nil` Proof, it would panic. 

Pr:
- fixes that
- add test case

CHANGELOG.md

@@ -45,6 +45,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ### Bug fixes
 
 * [#154](https://github.com/babylonlabs-io/babylon/pull/154) Fix "edit-finality-provider" cmd argument index
+* [#186](https://github.com/babylonlabs-io/babylon/pull/186) Do not panic on `nil`
+Proof when handling finality votes
 
 ### Improvements
 

x/finality/keeper/msg_server.go

@@ -49,9 +49,10 @@ func (ms msgServer) UpdateParams(goCtx context.Context, req *types.MsgUpdatePara
 func (ms msgServer) AddFinalitySig(goCtx context.Context, req *types.MsgAddFinalitySig) (*types.MsgAddFinalitySigResponse, error) {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), types.MetricsKeyAddFinalitySig)
 
-	if req.FpBtcPk == nil {
-		return nil, types.ErrInvalidFinalitySig.Wrap("empty finality provider BTC PK")
+	if err := req.ValidateBasic(); err != nil {
+		return nil, err
 	}
+
 	fpPK := req.FpBtcPk
 
 	ctx := sdk.UnwrapSDKContext(goCtx)
@@ -91,10 +92,6 @@ func (ms msgServer) AddFinalitySig(goCtx context.Context, req *types.MsgAddFinal
 		return nil, types.ErrInvalidFinalitySig.Wrapf("the finality provider %s does not have voting power at height %d", fpPK.MarshalHex(), req.BlockHeight)
 	}
 
-	// ensure the finality provider has not cast the same vote yet
-	if req.FinalitySig == nil {
-		return nil, types.ErrInvalidFinalitySig.Wrap("empty finality signature")
-	}
 	existingSig, err := ms.GetSig(ctx, req.BlockHeight, fpPK)
 	if err == nil && existingSig.Equals(req.FinalitySig) {
 		ms.Logger(ctx).Debug("Received duplicated finality vote", "block height", req.BlockHeight, "finality provider", req.FpBtcPk)

x/finality/keeper/msg_server_test.go

@@ -378,3 +378,69 @@ func TestVoteForConflictingHashShouldRetrieveEvidenceAndSlash(t *testing.T) {
 	require.Equal(t, msg.FinalitySig.MustMarshal(),
 		sig.MustMarshal())
 }
+
+func TestDoNotPanicOnNilProof(t *testing.T) {
+	r := rand.New(rand.NewSource(time.Now().Unix()))
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	bsKeeper := types.NewMockBTCStakingKeeper(ctrl)
+	cKeeper := types.NewMockCheckpointingKeeper(ctrl)
+	iKeeper := types.NewMockIncentiveKeeper(ctrl)
+	iKeeper.EXPECT().IndexRefundableMsg(gomock.Any(), gomock.Any()).AnyTimes()
+	fKeeper, ctx := keepertest.FinalityKeeper(t, bsKeeper, iKeeper, cKeeper)
+	ms := keeper.NewMsgServerImpl(*fKeeper)
+
+	// create and register a random finality provider
+	btcSK, btcPK, err := datagen.GenRandomBTCKeyPair(r)
+	require.NoError(t, err)
+	fp, err := datagen.GenRandomFinalityProviderWithBTCSK(r, btcSK)
+	require.NoError(t, err)
+	fpBTCPK := bbn.NewBIP340PubKeyFromBTCPK(btcPK)
+	fpBTCPKBytes := fpBTCPK.MustMarshal()
+	require.NoError(t, err)
+	bsKeeper.EXPECT().HasFinalityProvider(gomock.Any(), gomock.Eq(fpBTCPKBytes)).Return(true).AnyTimes()
+
+	// set committed epoch num
+	committedEpochNum := datagen.GenRandomEpochNum(r) + 1
+	cKeeper.EXPECT().GetEpoch(gomock.Any()).Return(&epochingtypes.Epoch{EpochNumber: committedEpochNum}).AnyTimes()
+
+	// commit some public randomness
+	startHeight := uint64(0)
+	numPubRand := uint64(200)
+	randListInfo, msgCommitPubRandList, err := datagen.GenRandomMsgCommitPubRandList(r, btcSK, startHeight, numPubRand)
+	require.NoError(t, err)
+	_, err = ms.CommitPubRandList(ctx, msgCommitPubRandList)
+	require.NoError(t, err)
+
+	// generate a vote
+	blockHeight := startHeight + uint64(1)
+	blockAppHash := datagen.GenRandomByteArray(r, 32)
+	signer := datagen.GenRandomAccount().Address
+	msg, err := datagen.NewMsgAddFinalitySig(
+		signer,
+		btcSK,
+		startHeight,
+		blockHeight,
+		randListInfo,
+		blockAppHash,
+	)
+	require.NoError(t, err)
+
+	// Not panic on empty proof
+	msg.Proof = nil
+	// Case 3: successful if the finality provider has voting power and has not casted this vote yet
+	// index this block first
+	ctx = ctx.WithHeaderInfo(header.Info{Height: int64(blockHeight), AppHash: blockAppHash})
+	fKeeper.IndexBlock(ctx)
+	bsKeeper.EXPECT().GetFinalityProvider(gomock.Any(), gomock.Eq(fpBTCPKBytes)).Return(fp, nil).AnyTimes()
+	// mock voting power
+	bsKeeper.EXPECT().GetVotingPower(gomock.Any(), gomock.Eq(fpBTCPKBytes), gomock.Eq(blockHeight)).Return(uint64(1)).AnyTimes()
+	// set the committed epoch finalized for the rest of the cases
+	lastFinalizedEpoch := datagen.GenRandomEpochNum(r) + committedEpochNum
+	cKeeper.EXPECT().GetLastFinalizedEpoch(gomock.Any()).Return(lastFinalizedEpoch).AnyTimes()
+
+	// add vote and it should work
+	_, err = ms.AddFinalitySig(ctx, msg)
+	require.Error(t, err)
+}

x/finality/types/msg.go

@@ -4,6 +4,7 @@ import (
 	fmt "fmt"
 
 	"github.com/babylonlabs-io/babylon/crypto/eots"
+	bbn "github.com/babylonlabs-io/babylon/types"
 	"github.com/cometbft/cometbft/crypto/merkle"
 	"github.com/cometbft/cometbft/crypto/tmhash"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -20,6 +21,30 @@ func (m *MsgAddFinalitySig) MsgToSign() []byte {
 	return msgToSignForVote(m.BlockHeight, m.BlockAppHash)
 }
 
+func (m *MsgAddFinalitySig) ValidateBasic() error {
+	if m.FpBtcPk.Size() != bbn.BIP340PubKeyLen {
+		return ErrInvalidFinalitySig.Wrapf("invalid finality provider BTC public key length: got %d, want %d", m.FpBtcPk.Size(), bbn.BIP340PubKeyLen)
+	}
+
+	if m.PubRand.Size() != bbn.SchnorrPubRandLen {
+		return ErrInvalidFinalitySig.Wrapf("invalind public randomness length: got %d, want %d", m.PubRand.Size(), bbn.SchnorrPubRandLen)
+	}
+
+	if m.Proof == nil {
+		return ErrInvalidFinalitySig.Wrap("empty inclusion proof")
+	}
+
+	if m.FinalitySig.Size() != bbn.SchnorrEOTSSigLen {
+		return ErrInvalidFinalitySig.Wrapf("invalid finality signature length: got %d, want %d", m.FinalitySig.Size(), bbn.BIP340SignatureLen)
+	}
+
+	if len(m.BlockAppHash) != tmhash.Size {
+		return ErrInvalidFinalitySig.Wrapf("invalid block app hash length: got %d, want %d", len(m.BlockAppHash), tmhash.Size)
+	}
+
+	return nil
+}
+
 // VerifyFinalitySig verifies the finality signature message w.r.t. the
 // public randomness commitment. The verification includes
 // - verifying the proof of inclusion of the given public randomness