From e5b144d94843e08e3708ec7057f287ade70d6a7a Mon Sep 17 00:00:00 2001
From: Runchao Han <me@runchao.rocks>
Date: Thu, 27 Jun 2024 11:48:28 +1000
Subject: [PATCH] fix

---
 x/babylon/keeper/wasm.go | 41 ++++++++++++++++++++++++++++++----------
 1 file changed, 31 insertions(+), 10 deletions(-)

diff --git a/x/babylon/keeper/wasm.go b/x/babylon/keeper/wasm.go
index 3a39d29b..294d159a 100644
--- a/x/babylon/keeper/wasm.go
+++ b/x/babylon/keeper/wasm.go
@@ -10,17 +10,40 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
-// SendBeginBlockMsg sends a BeginBlock sudo message to the BTC staking contract via sudo
-func (k Keeper) SendBeginBlockMsg(c context.Context) error {
-	ctx := sdk.UnwrapSDKContext(c)
-
+func (k Keeper) getBTCStakingContractAddr(ctx sdk.Context) sdk.AccAddress {
 	// get address of the BTC staking contract
 	addrStr := k.GetParams(ctx).BtcStakingContractAddress
 	if len(addrStr) == 0 {
 		// the BTC staking contract address is not set yet, skip sending BeginBlockMsg
 		return nil
 	}
-	addr := sdk.MustAccAddressFromBech32(addrStr)
+	addr, err := sdk.AccAddressFromBech32(addrStr)
+	if err != nil {
+		// Although this is a programming error so we should panic, we emit
+		// a warning message to minimise the impact on the consumer chain's operation
+		k.Logger(ctx).Warn("the BTC staking contract address is malformed", "contract", addrStr, "error", err)
+		return nil
+	}
+	if k.wasm.HasContractInfo(ctx, addr) {
+		// NOTE: it's possible that the default contract address does not correspond to
+		// any contract. We emit a warning message rather than panic to minimise the
+		// impact on the consumer chain's operation
+		k.Logger(ctx).Warn("the BTC staking contract address is not on-chain", "contract", addrStr)
+		return nil
+	}
+
+	return addr
+}
+
+// SendBeginBlockMsg sends a BeginBlock sudo message to the BTC staking contract via sudo
+func (k Keeper) SendBeginBlockMsg(c context.Context) error {
+	ctx := sdk.UnwrapSDKContext(c)
+
+	// try to get and parse BTC staking contract
+	addr := k.getBTCStakingContractAddr(ctx)
+	if addr == nil {
+		return nil
+	}
 
 	// construct the sudo message
 	headerInfo := ctx.HeaderInfo()
@@ -39,13 +62,11 @@ func (k Keeper) SendBeginBlockMsg(c context.Context) error {
 func (k Keeper) SendEndBlockMsg(c context.Context) error {
 	ctx := sdk.UnwrapSDKContext(c)
 
-	// get address of the BTC staking contract
-	addrStr := k.GetParams(ctx).BtcStakingContractAddress
-	if len(addrStr) == 0 {
-		// the BTC staking contract address is not set yet, skip sending EndBlockMsg
+	// try to get and parse BTC staking contract
+	addr := k.getBTCStakingContractAddr(ctx)
+	if addr == nil {
 		return nil
 	}
-	addr := sdk.MustAccAddressFromBech32(addrStr)
 
 	// construct the sudo message
 	headerInfo := ctx.HeaderInfo()