POC: feature gate for zone conceirge module

app/e2e_include_upgrades.go → app/e2e_code.go

@@ -1,13 +1,18 @@
 //go:build e2e
 
+// This file contains code specific to end-to-end (e2e) testing for the Babylon application.
+// It includes the signet upgrade and enables integration.
+
 package app
 
 import (
 	"github.com/babylonlabs-io/babylon/app/upgrades/signetlaunch"
+	zctypes "github.com/babylonlabs-io/babylon/x/zoneconcierge/types"
 )
 
 // init is used to include signet upgrade used for e2e testing
 // this file should be removed once the upgrade testing with signet ends.
 func init() {
 	Upgrades = append(Upgrades, signetlaunch.Upgrade)
+	zctypes.EnableIntegration = true
 }

app/keepers/keepers.go

@@ -614,8 +614,11 @@ func (ak *AppKeepers) InitKeepers(
 	// Create static IBC router, add ibc-transfer module route, then set and seal it
 	ibcRouter := porttypes.NewRouter().
 		AddRoute(ibctransfertypes.ModuleName, transferStack).
-		AddRoute(zctypes.ModuleName, zoneConciergeStack).
 		AddRoute(wasmtypes.ModuleName, wasmStack)
+	// Add zoneconcierge module to IBC router if integration is enabled
+	if zctypes.EnableIntegration {
+		ibcRouter.AddRoute(zctypes.ModuleName, zoneConciergeStack)
+	}
 
 	// Setting Router will finalize all routes by sealing router
 	// No more routes can be added

x/zoneconcierge/keeper/grpc_query.go

@@ -11,22 +11,37 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-var _ types.QueryServer = Keeper{}
+// 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{}
 
 const maxQueryChainsInfoLimit = 100
 
-func (k Keeper) Params(c context.Context, req *types.QueryParamsRequest) (*types.QueryParamsResponse, error) {
+func validateRequest(req sdk.Msg) error {
+	if !types.EnableIntegration {
+		return types.ErrIntegrationDisabled
+	}
 	if req == nil {
-		return nil, status.Error(codes.InvalidArgument, "invalid request")
+		return status.Error(codes.InvalidArgument, "invalid request")
+	}
+	return nil
+}
+
+func (k Keeper) Params(c context.Context, req *types.QueryParamsRequest) (*types.QueryParamsResponse, error) {
+	if err := validateRequest(req); err != nil {
+		return nil, err
 	}
 	ctx := sdk.UnwrapSDKContext(c)
 
 	return &types.QueryParamsResponse{Params: k.GetParams(ctx)}, nil
 }
 
 func (k Keeper) ChainList(c context.Context, req *types.QueryChainListRequest) (*types.QueryChainListResponse, error) {
-	if req == nil {
-		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	if err := validateRequest(req); err != nil {
+		return nil, err
 	}
 
 	ctx := sdk.UnwrapSDKContext(c)
@@ -51,8 +66,8 @@ func (k Keeper) ChainList(c context.Context, req *types.QueryChainListRequest) (
 
 // ChainsInfo returns the latest info for a given list of chains
 func (k Keeper) ChainsInfo(c context.Context, req *types.QueryChainsInfoRequest) (*types.QueryChainsInfoResponse, error) {
-	if req == nil {
-		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	if err := validateRequest(req); err != nil {
+		return nil, err
 	}
 
 	// return if no chain IDs are provided
@@ -87,8 +102,8 @@ func (k Keeper) ChainsInfo(c context.Context, req *types.QueryChainsInfoRequest)
 
 // Header returns the header and fork headers at a given height
 func (k Keeper) Header(c context.Context, req *types.QueryHeaderRequest) (*types.QueryHeaderResponse, error) {
-	if req == nil {
-		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	if err := validateRequest(req); err != nil {
+		return nil, err
 	}
 
 	if len(req.ConsumerId) == 0 {
@@ -112,8 +127,8 @@ func (k Keeper) Header(c context.Context, req *types.QueryHeaderRequest) (*types
 
 // EpochChainsInfo returns the latest info for list of chains in a given epoch
 func (k Keeper) EpochChainsInfo(c context.Context, req *types.QueryEpochChainsInfoRequest) (*types.QueryEpochChainsInfoResponse, error) {
-	if req == nil {
-		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	if err := validateRequest(req); err != nil {
+		return nil, err
 	}
 
 	// return if no chain IDs are provided
@@ -160,8 +175,8 @@ func (k Keeper) EpochChainsInfo(c context.Context, req *types.QueryEpochChainsIn
 
 // ListHeaders returns all headers of a chain with given ID, with pagination support
 func (k Keeper) ListHeaders(c context.Context, req *types.QueryListHeadersRequest) (*types.QueryListHeadersResponse, error) {
-	if req == nil {
-		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	if err := validateRequest(req); err != nil {
+		return nil, err
 	}
 
 	if len(req.ConsumerId) == 0 {
@@ -192,8 +207,8 @@ func (k Keeper) ListHeaders(c context.Context, req *types.QueryListHeadersReques
 // ListEpochHeaders returns all headers of a chain with given ID
 // TODO: support pagination in this RPC
 func (k Keeper) ListEpochHeaders(c context.Context, req *types.QueryListEpochHeadersRequest) (*types.QueryListEpochHeadersResponse, error) {
-	if req == nil {
-		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	if err := validateRequest(req); err != nil {
+		return nil, err
 	}
 
 	if len(req.ConsumerId) == 0 {
@@ -215,8 +230,8 @@ func (k Keeper) ListEpochHeaders(c context.Context, req *types.QueryListEpochHea
 
 // FinalizedChainsInfo returns the finalized info for a given list of chains
 func (k Keeper) FinalizedChainsInfo(c context.Context, req *types.QueryFinalizedChainsInfoRequest) (*types.QueryFinalizedChainsInfoResponse, error) {
-	if req == nil {
-		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	if err := validateRequest(req); err != nil {
+		return nil, err
 	}
 
 	// return if no chain IDs are provided
@@ -305,8 +320,8 @@ func (k Keeper) FinalizedChainsInfo(c context.Context, req *types.QueryFinalized
 }
 
 func (k Keeper) FinalizedChainInfoUntilHeight(c context.Context, req *types.QueryFinalizedChainInfoUntilHeightRequest) (*types.QueryFinalizedChainInfoUntilHeightResponse, error) {
-	if req == nil {
-		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	if err := validateRequest(req); err != nil {
+		return nil, err
 	}
 
 	if len(req.ConsumerId) == 0 {

x/zoneconcierge/keeper/header_handler.go

@@ -11,6 +11,11 @@ import (
 
 // HandleHeaderWithValidCommit handles a CZ header with a valid QC
 func (k Keeper) HandleHeaderWithValidCommit(ctx context.Context, txHash []byte, header *types.HeaderInfo, isOnFork bool) {
+	// if integration is not enabled, do not process headers
+	if !types.EnableIntegration {
+		return
+	}
+
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	babylonHeader := sdkCtx.HeaderInfo()
 	indexedHeader := types.IndexedHeader{

x/zoneconcierge/keeper/hooks.go

@@ -21,6 +21,12 @@ var _ epochingtypes.EpochingHooks = Hooks{}
 func (k Keeper) Hooks() Hooks { return Hooks{k} }
 
 func (h Hooks) AfterEpochEnds(ctx context.Context, epoch uint64) {
+	// if integration is not enabled, do not trigger hooks
+	// NOTE: might not be needed, since the post handler is disabled
+	if !types.EnableIntegration {
+		return
+	}
+
 	// upon an epoch has ended, index the current chain info for each CZ
 	// TODO: do this together when epoch is sealed?
 	for _, consumerID := range h.k.GetAllConsumerIDs(ctx) {
@@ -29,6 +35,12 @@ func (h Hooks) AfterEpochEnds(ctx context.Context, epoch uint64) {
 }
 
 func (h Hooks) AfterRawCheckpointSealed(ctx context.Context, epoch uint64) error {
+	// if integration is not enabled, do not trigger hooks
+	// NOTE: might not be needed, since the proofs do not depend on consumer chains
+	if !types.EnableIntegration {
+		return nil
+	}
+
 	// upon a raw checkpoint is sealed, index the current chain info for each consumer,
 	// and generate/save the proof that the epoch is sealed
 	h.k.recordEpochChainInfoProofs(ctx, epoch)
@@ -38,6 +50,11 @@ func (h Hooks) AfterRawCheckpointSealed(ctx context.Context, epoch uint64) error
 
 // AfterRawCheckpointFinalized is triggered upon an epoch has been finalised
 func (h Hooks) AfterRawCheckpointFinalized(ctx context.Context, epoch uint64) error {
+	// if integration is not enabled, do not trigger hooks
+	if !types.EnableIntegration {
+		return nil
+	}
+
 	headersToBroadcast := h.k.getHeadersToBroadcast(ctx)
 
 	// send BTC timestamp to all open channels with ZoneConcierge

x/zoneconcierge/keeper/keeper_test.go

@@ -3,13 +3,22 @@ package keeper_test
 import (
 	"context"
 	"math/rand"
-
-	ibctmtypes "github.com/cosmos/ibc-go/v8/modules/light-clients/07-tendermint"
+	"testing"
 
 	"github.com/babylonlabs-io/babylon/testutil/datagen"
+	testhelper "github.com/babylonlabs-io/babylon/testutil/helper"
 	zckeeper "github.com/babylonlabs-io/babylon/x/zoneconcierge/keeper"
+	"github.com/babylonlabs-io/babylon/x/zoneconcierge/types"
+	zctypes "github.com/babylonlabs-io/babylon/x/zoneconcierge/types"
+	"github.com/cosmos/cosmos-sdk/baseapp"
+	ibctmtypes "github.com/cosmos/ibc-go/v8/modules/light-clients/07-tendermint"
+	"github.com/stretchr/testify/require"
 )
 
+func init() {
+	zctypes.EnableIntegration = true
+}
+
 // SimulateNewHeaders generates a non-zero number of canonical headers
 func SimulateNewHeaders(ctx context.Context, r *rand.Rand, k *zckeeper.Keeper, consumerID string, startHeight uint64, numHeaders uint64) []*ibctmtypes.Header {
 	headers := []*ibctmtypes.Header{}
@@ -41,3 +50,94 @@ func SimulateNewHeadersAndForks(ctx context.Context, r *rand.Rand, k *zckeeper.K
 	}
 	return headers, forkHeaders
 }
+
+func FuzzFeatureGate(f *testing.F) {
+	datagen.AddRandomSeedsToFuzzer(f, 10)
+
+	f.Fuzz(func(t *testing.T, seed int64) {
+		r := rand.New(rand.NewSource(seed))
+
+		// Save the original value of EnableIntegration
+		originalEnableIntegration := zctypes.EnableIntegration
+		// Restore the original value after the test
+		defer func() {
+			zctypes.EnableIntegration = originalEnableIntegration
+		}()
+		// Set EnableIntegration to a random value
+		currentEnableIntegration := datagen.OneInN(r, 2)
+		zctypes.EnableIntegration = currentEnableIntegration
+
+		helper := testhelper.NewHelper(t)
+		zcKeeper := helper.App.ZoneConciergeKeeper
+		ctx := helper.Ctx
+
+		// Create a random header and consumer ID
+		header := datagen.GenRandomIBCTMHeader(r, 1)
+		consumerID := datagen.GenRandomHexStr(r, 30)
+		headerInfo := datagen.NewZCHeaderInfo(header, consumerID)
+
+		/*
+			Ensure PostHandler is feature gated
+		*/
+		// Call HandleHeaderWithValidCommit
+		zcKeeper.HandleHeaderWithValidCommit(ctx, datagen.GenRandomByteArray(r, 32), headerInfo, false)
+		// Check that the header was not stored
+		_, err := zcKeeper.GetHeader(ctx, consumerID, 1)
+		if currentEnableIntegration {
+			require.NoError(t, err, "Header should be stored when EnableIntegration is true")
+		} else {
+			require.Error(t, err, "Header should not be stored when EnableIntegration is false")
+		}
+
+		/*
+			Ensure GRPC query server is feature gated
+		*/
+		// Get zone concierge query client
+		zcQueryHelper := baseapp.NewQueryServerTestHelper(ctx, helper.App.InterfaceRegistry())
+		querier := zckeeper.Querier{Keeper: zcKeeper}
+		types.RegisterQueryServer(zcQueryHelper, querier)
+
+		queryClient := zctypes.NewQueryClient(zcQueryHelper)
+
+		// Test GetParams query
+		paramsReq := &zctypes.QueryParamsRequest{}
+		_, err = queryClient.Params(ctx, paramsReq)
+		if currentEnableIntegration {
+			require.NoError(t, err, "Params query should work when EnableIntegration is true")
+		} else {
+			require.Error(t, err, "Params query should be blocked when EnableIntegration is false")
+			require.ErrorIs(t, err, types.ErrIntegrationDisabled)
+		}
+
+		/*
+			Ensure msg server is feature gated
+		*/
+		msgSrvr := zckeeper.NewMsgServerImpl(zcKeeper)
+		msgReq := &zctypes.MsgUpdateParams{
+			Authority: helper.App.GovKeeper.GetGovernanceAccount(ctx).GetAddress().String(),
+			Params:    zctypes.DefaultParams(),
+		}
+		_, err = msgSrvr.UpdateParams(ctx, msgReq)
+		if currentEnableIntegration {
+			require.NoError(t, err, "MsgUpdateParams should work when EnableIntegration is true")
+		} else {
+			require.Error(t, err, "MsgUpdateParams should be blocked when EnableIntegration is false")
+			require.ErrorIs(t, err, types.ErrIntegrationDisabled)
+		}
+
+		/*
+			Ensure IBC route does not contain zone concierge
+		*/
+		// Get the IBC keeper
+		ibcKeeper := helper.App.IBCKeeper
+		// Get the IBC router
+		router := ibcKeeper.Router
+		// Ensure the zone concierge module is not in the router
+		_, found := router.GetRoute(zctypes.ModuleName)
+		if currentEnableIntegration {
+			require.True(t, found, "Zone concierge module should be in the IBC router when EnableIntegration is true")
+		} else {
+			require.False(t, found, "Zone concierge module should not be in the IBC router when EnableIntegration is false")
+		}
+	})
+}

x/zoneconcierge/keeper/msg_server.go

@@ -23,6 +23,9 @@ var _ types.MsgServer = msgServer{}
 
 // UpdateParams updates the params
 func (ms msgServer) UpdateParams(goCtx context.Context, req *types.MsgUpdateParams) (*types.MsgUpdateParamsResponse, error) {
+	if !types.EnableIntegration {
+		return nil, types.ErrIntegrationDisabled
+	}
 	if ms.authority != req.Authority {
 		return nil, errorsmod.Wrapf(govtypes.ErrInvalidSigner, "invalid authority; expected %s, got %s", ms.authority, req.Authority)
 	}

x/zoneconcierge/module.go

@@ -2,10 +2,11 @@ package zoneconcierge
 
 import (
 	"context"
-	"cosmossdk.io/core/appmodule"
 	"encoding/json"
 	"fmt"
 
+	"cosmossdk.io/core/appmodule"
+
 	"github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"github.com/spf13/cobra"
 

x/zoneconcierge/module_ibc.go

@@ -34,6 +34,10 @@ func (im IBCModule) OnChanOpenInit(
 	counterparty channeltypes.Counterparty,
 	version string,
 ) (string, error) {
+	if !types.EnableIntegration {
+		return "", types.ErrIntegrationDisabled
+	}
+
 	// the IBC channel has to be ordered
 	if order != channeltypes.ORDERED {
 		return "", errorsmod.Wrapf(channeltypes.ErrInvalidChannelOrdering, "expected %s channel, got %s ", channeltypes.ORDERED, order)
@@ -69,6 +73,10 @@ func (im IBCModule) OnChanOpenTry(
 	counterparty channeltypes.Counterparty,
 	counterpartyVersion string,
 ) (string, error) {
+	if !types.EnableIntegration {
+		return "", types.ErrIntegrationDisabled
+	}
+
 	// the IBC channel has to be ordered
 	if order != channeltypes.ORDERED {
 		return "", errorsmod.Wrapf(channeltypes.ErrInvalidChannelOrdering, "expected %s channel, got %s ", channeltypes.ORDERED, order)
@@ -107,6 +115,10 @@ func (im IBCModule) OnChanOpenAck(
 	_,
 	counterpartyVersion string,
 ) error {
+	if !types.EnableIntegration {
+		return types.ErrIntegrationDisabled
+	}
+
 	// check version consistency
 	if counterpartyVersion != types.Version {
 		return errorsmod.Wrapf(types.ErrInvalidVersion, "invalid counterparty version: %s, expected %s", counterpartyVersion, types.Version)
@@ -121,6 +133,10 @@ func (im IBCModule) OnChanOpenConfirm(
 	portID,
 	channelID string,
 ) error {
+	if !types.EnableIntegration {
+		return types.ErrIntegrationDisabled
+	}
+
 	return nil
 }