chore: Use prefix-based routing for IBCv2 contracts

Author: kulikthebird

app/app.go

@@ -582,8 +582,6 @@ func NewWasmApp(
 		panic(fmt.Sprintf("error while reading wasm config: %s", err))
 	}
 
-	ibcRouterV2 := ibcapi.NewRouter()
-
 	// The last arguments can contain custom message handlers, and custom query handlers,
 	// if we want to allow any custom callbacks
 	app.WasmKeeper = wasmkeeper.NewKeeper(
@@ -604,7 +602,6 @@ func NewWasmApp(
 		wasmtypes.VMConfig{},
 		wasmkeeper.BuiltInCapabilities(),
 		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
-		ibcRouterV2,
 		wasmOpts...,
 	)
 
@@ -646,8 +643,11 @@ func NewWasmApp(
 		AddRoute(icahosttypes.SubModuleName, icaHostStack)
 	app.IBCKeeper.SetRouter(ibcRouter)
 
+	ibcRouterV2 := ibcapi.NewRouter()
 	ibcRouterV2 = ibcRouterV2.
-		AddRoute(ibctransfertypes.PortID, transferv2.NewIBCModule(app.TransferKeeper))
+		AddRoute(ibctransfertypes.PortID, transferv2.NewIBCModule(app.TransferKeeper)).
+		AddRoute(wasmkeeper.PortIDPrefixV2, wasmkeeper.NewIBC2Handler(app.WasmKeeper))
+
 	app.IBCKeeper.SetRouterV2(ibcRouterV2)
 
 	clientKeeper := app.IBCKeeper.ClientKeeper
@@ -875,10 +875,6 @@ func NewWasmApp(
 		if err := app.WasmKeeper.InitializePinnedCodes(ctx); err != nil {
 			panic(fmt.Sprintf("failed initialize pinned codes %s", err))
 		}
-
-		// TODO: Remove the registration of each contract in https://github.com/CosmWasm/wasmd/issues/2278
-		//       and add IBCv2 port ID prefix before calling `SetRouterV2` during WasmKeeper creation.
-		app.WasmKeeper.RegisterContractsInIbc2Router(ctx)
 	}
 
 	return app

x/wasm/keeper/ibc.go

@@ -25,17 +25,17 @@ func ContractFromPortID(portID string) (sdk.AccAddress, error) {
 }
 
 // The port prefix refers to "CosmWasm over IBC v2" and ensures packets are routed to the right entry points
-const portIDPrefixV2 = "wasm2"
+const PortIDPrefixV2 = "wasm2"
 
 func PortIDForContractV2(addr sdk.AccAddress) string {
 	blockchainPrefix := sdk.GetConfig().GetBech32AccountAddrPrefix()
-	return portIDPrefixV2 + strings.TrimPrefix(addr.String(), blockchainPrefix)
+	return PortIDPrefixV2 + strings.TrimPrefix(addr.String(), blockchainPrefix)
 }
 
 func ContractFromPortID2(portID string) (sdk.AccAddress, error) {
-	if !strings.HasPrefix(portID, portIDPrefixV2) {
+	if !strings.HasPrefix(portID, PortIDPrefixV2) {
 		return nil, errorsmod.Wrapf(types.ErrInvalid, "without prefix")
 	}
 	blockchainPrefix := sdk.GetConfig().GetBech32AccountAddrPrefix()
-	return sdk.AccAddressFromBech32(blockchainPrefix + portID[len(portIDPrefixV2):])
+	return sdk.AccAddressFromBech32(blockchainPrefix + portID[len(PortIDPrefixV2):])
 }

x/wasm/keeper/keeper.go

@@ -15,7 +15,6 @@ import (
 	wasmvm "github.com/CosmWasm/wasmvm/v3"
 	wasmvmtypes "github.com/CosmWasm/wasmvm/v3/types"
 	channeltypes "github.com/cosmos/ibc-go/v10/modules/core/04-channel/types"
-	ibcapi "github.com/cosmos/ibc-go/v10/modules/core/api"
 
 	"cosmossdk.io/collections"
 	corestoretypes "cosmossdk.io/core/store"
@@ -111,12 +110,6 @@ type Keeper struct {
 
 	// wasmLimits contains the limits sent to wasmvm on init
 	wasmLimits wasmvmtypes.WasmLimits
-
-	ibcRouterV2 *ibcapi.Router
-}
-
-func (k Keeper) GetIBCRouterV2() *ibcapi.Router {
-	return k.ibcRouterV2
 }
 
 func (k Keeper) getUploadAccessConfig(ctx context.Context) types.AccessConfig {
@@ -379,13 +372,7 @@ func (k Keeper) instantiate(
 		contractInfo.IBCPortID = ibcPort
 	}
 
-	ibc2Port := PortIDForContractV2(contractAddress)
-	contractInfo.IBC2PortID = ibc2Port
-
-	// TODO: Remove the registration of each contract in https://github.com/CosmWasm/wasmd/issues/2278
-	if !k.ibcRouterV2.HasRoute(ibc2Port) {
-		k.ibcRouterV2.AddRoute(ibc2Port, NewIBC2Handler(k))
-	}
+	contractInfo.IBC2PortID = PortIDForContractV2(contractAddress)
 
 	// store contract before dispatch so that contract could be called back
 	historyEntry := contractInfo.InitialHistory(initMsg)
@@ -554,6 +541,8 @@ func (k Keeper) migrate(
 	}
 	k.mustStoreContractInfo(ctx, contractAddress, contractInfo)
 
+	contractInfo.IBC2PortID = PortIDForContractV2(contractAddress)
+
 	sdkCtx.EventManager().EmitEvent(sdk.NewEvent(
 		types.EventTypeMigrate,
 		sdk.NewAttribute(types.AttributeKeyCodeID, strconv.FormatUint(newCodeID, 10)),
@@ -1092,12 +1081,6 @@ func (k Keeper) importContractState(ctx context.Context, contractAddress sdk.Acc
 		prefixStore.Set(model.Key, model.Value)
 	}
 
-	// TODO: Remove the registration of each contract in https://github.com/CosmWasm/wasmd/issues/2278
-	contractPortIBC2ID := PortIDForContractV2(contractAddress)
-	if !k.ibcRouterV2.HasRoute(contractPortIBC2ID) {
-		k.ibcRouterV2.AddRoute(contractPortIBC2ID, NewIBC2Handler(k))
-	}
-
 	return nil
 }
 
@@ -1139,24 +1122,6 @@ func (k Keeper) IterateCodeInfos(ctx context.Context, cb func(uint64, types.Code
 	}
 }
 
-// TODO: Remove the registration of each contract in https://github.com/CosmWasm/wasmd/issues/2278 (remove this method)
-func (k Keeper) IterateContractAddresses(ctx context.Context, cb func(sdk.AccAddress)) {
-	prefixStore := prefix.NewStore(runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx)), types.ContractKeyPrefix)
-	iter := prefixStore.Iterator(nil, nil)
-	defer iter.Close()
-
-	for ; iter.Valid(); iter.Next() {
-		cb(iter.Key())
-	}
-}
-
-// TODO: Remove the registration of each contract in https://github.com/CosmWasm/wasmd/issues/2278 (remove this method)
-func (k Keeper) RegisterContractsInIbc2Router(ctx context.Context) {
-	k.IterateContractAddresses(ctx, func(contractAddr sdk.AccAddress) {
-		k.ibcRouterV2.AddRoute(PortIDForContractV2(contractAddr), NewIBC2Handler(k))
-	})
-}
-
 func (k Keeper) GetByteCode(ctx context.Context, codeID uint64) ([]byte, error) {
 	store := k.storeService.OpenKVStore(ctx)
 	var codeInfo types.CodeInfo

x/wasm/keeper/keeper_cgo.go

@@ -7,7 +7,6 @@ import (
 
 	wasmvm "github.com/CosmWasm/wasmvm/v3"
 	wasmvmtypes "github.com/CosmWasm/wasmvm/v3/types"
-	ibcapi "github.com/cosmos/ibc-go/v10/modules/core/api"
 
 	"cosmossdk.io/collections"
 	corestoretypes "cosmossdk.io/core/store"
@@ -37,7 +36,6 @@ func NewKeeper(
 	vmConfig types.VMConfig,
 	availableCapabilities []string,
 	authority string,
-	ibcRouterV2 *ibcapi.Router,
 	opts ...Option,
 ) Keeper {
 	sb := collections.NewSchemaBuilder(storeService)
@@ -57,9 +55,8 @@ func NewKeeper(
 		propagateGovAuthorization: map[types.AuthorizationPolicyAction]struct{}{
 			types.AuthZActionInstantiate: {},
 		},
-		authority:   authority,
-		wasmLimits:  vmConfig.WasmLimits,
-		ibcRouterV2: ibcRouterV2,
+		authority:  authority,
+		wasmLimits: vmConfig.WasmLimits,
 	}
 	keeper.messenger = NewDefaultMessageHandler(keeper, router, ics4Wrapper, channelKeeperV2, bankKeeper, cdc, portSource)
 	keeper.wasmVMQueryHandler = DefaultQueryPlugins(bankKeeper, stakingKeeper, distrKeeper, channelKeeper, keeper)

x/wasm/keeper/test_common.go

@@ -17,7 +17,6 @@ import (
 	"github.com/cosmos/ibc-go/v10/modules/apps/transfer"
 	ibctransfertypes "github.com/cosmos/ibc-go/v10/modules/apps/transfer/types"
 	ibc "github.com/cosmos/ibc-go/v10/modules/core"
-	ibcapi "github.com/cosmos/ibc-go/v10/modules/core/api"
 	ibcexported "github.com/cosmos/ibc-go/v10/modules/core/exported"
 	ibckeeper "github.com/cosmos/ibc-go/v10/modules/core/keeper"
 	"github.com/stretchr/testify/require"
@@ -389,7 +388,6 @@ func createTestInput(
 		vmConfig,
 		availableCapabilities,
 		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
-		ibcapi.NewRouter(),
 		opts...,
 	)
 	require.NoError(t, keeper.SetParams(ctx, types.DefaultParams()))