Feat: Whitelist tokenfactory hooks

proto/osmosis/tokenfactory/params.proto

@@ -0,0 +1,39 @@
+syntax = "proto3";
+package osmosis.tokenfactory;
+
+import "cosmos/base/v1beta1/coin.proto";
+import "cosmos_proto/cosmos.proto";
+import "gogoproto/gogo.proto";
+
+option go_package = "github.com/neutron-org/neutron/v4/x/tokenfactory/types";
+
+message HookWhitelist {
+  uint64 code_id = 1 [(gogoproto.customname) = "CodeID"];
+  string denom_creator = 2;
+}
+
+// Params defines the parameters for the tokenfactory module.
+message Params {
+  // DenomCreationFee defines the fee to be charged on the creation of a new
+  // denom. The fee is drawn from the MsgCreateDenom's sender account, and
+  // transferred to the community pool.
+  repeated cosmos.base.v1beta1.Coin denom_creation_fee = 1 [
+    (gogoproto.castrepeated) = "github.com/cosmos/cosmos-sdk/types.Coins",
+    (gogoproto.moretags) = "yaml:\"denom_creation_fee\"",
+    (gogoproto.nullable) = false
+  ];
+
+  // DenomCreationGasConsume defines the gas cost for creating a new denom.
+  // This is intended as a spam deterrence mechanism.
+  //
+  // See: https://github.com/CosmWasm/token-factory/issues/11
+  uint64 denom_creation_gas_consume = 2 [
+    (gogoproto.moretags) = "yaml:\"denom_creation_gas_consume\"",
+    (gogoproto.nullable) = true
+  ];
+
+  // FeeCollectorAddress is the address where fees collected from denom creation
+  // are sent to
+  string fee_collector_address = 3;
+  repeated HookWhitelist whitelisted_hooks = 4;
+}

proto/osmosis/tokenfactory/v1beta1/genesis.proto

@@ -2,8 +2,8 @@ syntax = "proto3";
 package osmosis.tokenfactory.v1beta1;
 
 import "gogoproto/gogo.proto";
+import "osmosis/tokenfactory/params.proto";
 import "osmosis/tokenfactory/v1beta1/authorityMetadata.proto";
-import "osmosis/tokenfactory/v1beta1/params.proto";
 
 option go_package = "github.com/neutron-org/neutron/v4/x/tokenfactory/types";
 

proto/osmosis/tokenfactory/v1beta1/params.proto

@@ -5,7 +5,7 @@ import "cosmos/base/v1beta1/coin.proto";
 import "cosmos_proto/cosmos.proto";
 import "gogoproto/gogo.proto";
 
-option go_package = "github.com/neutron-org/neutron/v4/x/tokenfactory/types";
+option go_package = "github.com/neutron-org/neutron/v4/x/tokenfactory/types/v1beta1";
 
 // Params defines the parameters for the tokenfactory module.
 message Params {

proto/osmosis/tokenfactory/v1beta1/query.proto

@@ -3,8 +3,8 @@ package osmosis.tokenfactory.v1beta1;
 
 import "gogoproto/gogo.proto";
 import "google/api/annotations.proto";
+import "osmosis/tokenfactory/params.proto";
 import "osmosis/tokenfactory/v1beta1/authorityMetadata.proto";
-import "osmosis/tokenfactory/v1beta1/params.proto";
 
 option go_package = "github.com/neutron-org/neutron/v4/x/tokenfactory/types";
 

proto/osmosis/tokenfactory/v1beta1/tx.proto

@@ -7,7 +7,7 @@ import "cosmos/base/v1beta1/coin.proto";
 import "cosmos/msg/v1/msg.proto";
 import "cosmos_proto/cosmos.proto";
 import "gogoproto/gogo.proto";
-import "osmosis/tokenfactory/v1beta1/params.proto";
+import "osmosis/tokenfactory/params.proto";
 
 option go_package = "github.com/neutron-org/neutron/v4/x/tokenfactory/types";
 

wasmbinding/test/custom_message_test.go

@@ -76,6 +76,7 @@ func (suite *CustomMessengerTestSuite) SetupTest() {
 		sdk.NewCoins(sdk.NewInt64Coin(params.DefaultDenom, 10_000_000)),
 		0,
 		FeeCollectorAddress,
+		tokenfactorytypes.DefaultWhitelistedHooks,
 	))
 	suite.Require().NoError(err)
 

wasmbinding/test/custom_query_test.go

@@ -255,6 +255,7 @@ func (suite *CustomQuerierTestSuite) TestDenomAdmin() {
 		sdk.NewCoins(sdk.NewInt64Coin(params.DefaultDenom, 10_000_000)),
 		0,
 		FeeCollectorAddress,
+		tokenfactorytypes.DefaultWhitelistedHooks,
 	))
 	suite.Require().NoError(err)
 

x/tokenfactory/keeper/before_send.go

@@ -99,6 +99,11 @@ func (k Keeper) callBeforeSendListener(ctx context.Context, from, to sdk.AccAddr
 				return err
 			}
 
+			// Do not invoke hook if denom is not whitelisted
+			if err := k.AssertIsHookWhitelisted(c, coin.Denom, cwAddr); err != nil {
+				return nil
+			}
+
 			var msgBz []byte
 
 			// get msgBz, either BlockBeforeSend or TrackBeforeSend

x/tokenfactory/keeper/before_send_test.go

@@ -2,7 +2,6 @@ package keeper_test
 
 import (
 	"encoding/json"
-	"fmt"
 	"os"
 
 	"cosmossdk.io/math"
@@ -12,62 +11,107 @@ import (
 	"github.com/neutron-org/neutron/v4/x/tokenfactory/types"
 )
 
-func (suite *KeeperTestSuite) TestTrackBeforeSendWasm() {
-	for _, tc := range []struct {
-		name     string
-		wasmFile string
-	}{
-		{
-			name: "Test bank hook tracking contract ",
-			// https://github.com/neutron-org/neutron-dev-contracts/tree/feat/balance-tracker-contract/contracts/balance-tracker
-			wasmFile: "./testdata/balance_tracker.wasm",
+func (suite *KeeperTestSuite) initBalanceTrackContract(denom string) (sdk.AccAddress, uint64, string) {
+	senderAddress := suite.ChainA.SenderAccounts[0].SenderAccount.GetAddress()
+	suite.TopUpWallet(suite.ChainA.GetContext(), senderAddress, suite.TestAccs[0])
+
+	// https://github.com/neutron-org/neutron-dev-contracts/tree/feat/balance-tracker-contract/contracts/balance-tracker
+	wasmFile := "./testdata/balance_tracker.wasm"
+
+	// load wasm file
+	wasmCode, err := os.ReadFile(wasmFile)
+	suite.Require().NoError(err)
+
+	// create new denom
+	res, err := suite.msgServer.CreateDenom(suite.ChainA.GetContext(), types.NewMsgCreateDenom(suite.TestAccs[0].String(), denom))
+	suite.Require().NoError(err)
+	factoryDenom := res.GetNewTokenDenom()
+
+	// instantiate wasm code
+	tokenFactoryModuleAddr := suite.GetNeutronZoneApp(suite.ChainA).AccountKeeper.GetModuleAddress("tokenfactory")
+	contractKeeper := wasmkeeper.NewDefaultPermissionKeeper(suite.GetNeutronZoneApp(suite.ChainA).WasmKeeper)
+	codeID, _, err := contractKeeper.Create(suite.ChainA.GetContext(), suite.TestAccs[0], wasmCode, nil)
+	suite.Require().NoError(err)
+	initMsg, _ := json.Marshal(
+		map[string]interface{}{
+			"tracked_denom":               factoryDenom,
+			"tokenfactory_module_address": tokenFactoryModuleAddr,
 		},
-	} {
-		suite.Run(fmt.Sprintf("Case %s", tc.name), func() {
-			// setup test
-			suite.Setup()
-
-			senderAddress := suite.ChainA.SenderAccounts[0].SenderAccount.GetAddress()
-			suite.TopUpWallet(suite.ChainA.GetContext(), senderAddress, suite.TestAccs[0])
-
-			// load wasm file
-			wasmCode, err := os.ReadFile(tc.wasmFile)
-			suite.Require().NoError(err)
-
-			// create new denom
-			res, err := suite.msgServer.CreateDenom(suite.ChainA.GetContext(), types.NewMsgCreateDenom(suite.TestAccs[0].String(), "testdenom"))
-			suite.Require().NoError(err, "test: %v", tc.name)
-			factoryDenom := res.GetNewTokenDenom()
-
-			// instantiate wasm code
-			tokenFactoryModuleAddr := suite.GetNeutronZoneApp(suite.ChainA).AccountKeeper.GetModuleAddress("tokenfactory")
-			contractKeeper := wasmkeeper.NewDefaultPermissionKeeper(suite.GetNeutronZoneApp(suite.ChainA).WasmKeeper)
-			codeID, _, err := contractKeeper.Create(suite.ChainA.GetContext(), suite.TestAccs[0], wasmCode, nil)
-			suite.Require().NoError(err, "test: %v", tc.name)
-			initMsg, _ := json.Marshal(
-				map[string]interface{}{
-					"tracked_denom":               factoryDenom,
-					"tokenfactory_module_address": tokenFactoryModuleAddr,
-				},
-			)
-			cosmwasmAddress, _, err := contractKeeper.Instantiate(
-				suite.ChainA.GetContext(), codeID, suite.TestAccs[0], suite.TestAccs[0], initMsg, "", sdk.NewCoins(),
-			)
-			suite.Require().NoError(err, "test: %v", tc.name)
-
-			// set beforesend hook to the new denom
-			_, err = suite.msgServer.SetBeforeSendHook(suite.ChainA.GetContext(), types.NewMsgSetBeforeSendHook(suite.TestAccs[0].String(), factoryDenom, cosmwasmAddress.String()))
-			suite.Require().NoError(err, "test: %v", tc.name)
-
-			tokenToSend := sdk.NewCoin(factoryDenom, math.NewInt(100))
-
-			// mint tokens
-			_, err = suite.msgServer.Mint(suite.ChainA.GetContext(), types.NewMsgMint(suite.TestAccs[0].String(), tokenToSend))
-			suite.Require().NoError(err)
-
-			queryResp, err := suite.GetNeutronZoneApp(suite.ChainA).WasmKeeper.QuerySmart(suite.ChainA.GetContext(), cosmwasmAddress, []byte(`{"total_supply_at":{}}`))
-			suite.Require().NoError(err)
-			suite.Require().Equal("\"100\"", string(queryResp))
-		})
-	}
+	)
+	cosmwasmAddress, _, err := contractKeeper.Instantiate(
+		suite.ChainA.GetContext(), codeID, suite.TestAccs[0], suite.TestAccs[0], initMsg, "", sdk.NewCoins(),
+	)
+	suite.Require().NoError(err)
+
+	return cosmwasmAddress, codeID, factoryDenom
+}
+
+func (suite *KeeperTestSuite) TestTrackBeforeSendWasm() {
+	suite.Setup()
+
+	cosmwasmAddress, codeID, factoryDenom := suite.initBalanceTrackContract("testdenom")
+
+	// Whitelist the hook
+	params := types.DefaultParams()
+	params.WhitelistedHooks = []*types.HookWhitelist{{DenomCreator: suite.TestAccs[0].String(), CodeID: codeID}}
+	err := suite.GetNeutronZoneApp(suite.ChainA).TokenFactoryKeeper.SetParams(suite.ChainA.GetContext(), params)
+	suite.Require().NoError(err)
+
+	// set beforeSendHook for the new denom
+	_, err = suite.msgServer.SetBeforeSendHook(suite.ChainA.GetContext(), types.NewMsgSetBeforeSendHook(suite.TestAccs[0].String(), factoryDenom, cosmwasmAddress.String()))
+	suite.Require().NoError(err)
+
+	tokenToSend := sdk.NewCoin(factoryDenom, math.NewInt(100))
+
+	// mint tokens
+	_, err = suite.msgServer.Mint(suite.ChainA.GetContext(), types.NewMsgMint(suite.TestAccs[0].String(), tokenToSend))
+	suite.Require().NoError(err)
+
+	queryResp, err := suite.GetNeutronZoneApp(suite.ChainA).WasmKeeper.QuerySmart(suite.ChainA.GetContext(), cosmwasmAddress, []byte(`{"total_supply_at":{}}`))
+	suite.Require().NoError(err)
+	// Whitelisted contract is called correctly
+	suite.Require().Equal("\"100\"", string(queryResp))
+}
+
+func (suite *KeeperTestSuite) TestAddNonWhitelistedHooksFails() {
+	suite.Setup()
+
+	cosmwasmAddress, _, factoryDenom := suite.initBalanceTrackContract("testdenom")
+
+	// WHEN to set beforeSendHook
+	_, err := suite.msgServer.SetBeforeSendHook(suite.ChainA.GetContext(), types.NewMsgSetBeforeSendHook(suite.TestAccs[0].String(), factoryDenom, cosmwasmAddress.String()))
+	// THEN fails with BeforeSendHookNotWhitelisted
+	suite.Require().ErrorIs(err, types.ErrBeforeSendHookNotWhitelisted)
+}
+
+func (suite *KeeperTestSuite) TestNonWhitelistedHooksNotCalled() {
+	suite.Setup()
+
+	cosmwasmAddress, codeID, factoryDenom := suite.initBalanceTrackContract("testdenom")
+
+	// Whitelist the hook first so it can be added
+	params := types.DefaultParams()
+	params.WhitelistedHooks = []*types.HookWhitelist{{DenomCreator: suite.TestAccs[0].String(), CodeID: codeID}}
+	err := suite.GetNeutronZoneApp(suite.ChainA).TokenFactoryKeeper.SetParams(suite.ChainA.GetContext(), params)
+	suite.Require().NoError(err)
+
+	// set beforeSendHook for the new denom
+	_, err = suite.msgServer.SetBeforeSendHook(suite.ChainA.GetContext(), types.NewMsgSetBeforeSendHook(suite.TestAccs[0].String(), factoryDenom, cosmwasmAddress.String()))
+	suite.Require().NoError(err)
+
+	// Remove whitelist for the hook
+	params = types.DefaultParams()
+	err = suite.GetNeutronZoneApp(suite.ChainA).TokenFactoryKeeper.SetParams(suite.ChainA.GetContext(), params)
+	suite.Require().NoError(err)
+
+	tokenToSend := sdk.NewCoin(factoryDenom, math.NewInt(100))
+
+	// mint tokens
+	_, err = suite.msgServer.Mint(suite.ChainA.GetContext(), types.NewMsgMint(suite.TestAccs[0].String(), tokenToSend))
+	suite.Require().NoError(err)
+
+	queryResp, err := suite.GetNeutronZoneApp(suite.ChainA).WasmKeeper.QuerySmart(suite.ChainA.GetContext(), cosmwasmAddress, []byte(`{"total_supply_at":{}}`))
+	suite.Require().NoError(err)
+	// Whitelisted contract is not called
+	suite.Require().Equal("\"0\"", string(queryResp))
 }

x/tokenfactory/keeper/keeper_test.go

@@ -57,6 +57,7 @@ func (suite *KeeperTestSuite) Setup() {
 		sdktypes.NewCoins(sdktypes.NewInt64Coin(params.DefaultDenom, TopUpCoinsAmount)),
 		0,
 		FeeCollectorAddress,
+		types.DefaultWhitelistedHooks,
 	))
 	suite.Require().NoError(err)
 

x/tokenfactory/keeper/migrations.go

@@ -0,0 +1,22 @@
+package keeper
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	v2 "github.com/neutron-org/neutron/v4/x/tokenfactory/migrations/v2"
+)
+
+// Migrator is a struct for handling in-place store migrations.
+type Migrator struct {
+	keeper Keeper
+}
+
+// NewMigrator returns a new Migrator.
+func NewMigrator(keeper Keeper) Migrator {
+	return Migrator{keeper: keeper}
+}
+
+// Migrate1to2 migrates from version 1 to 2.
+func (m Migrator) Migrate1to2(ctx sdk.Context) error {
+	return v2.MigrateStore(ctx, m.keeper.cdc, m.keeper.storeKey)
+}

x/tokenfactory/keeper/msg_server.go

@@ -246,6 +246,15 @@ func (server msgServer) SetBeforeSendHook(goCtx context.Context, msg *types.MsgS
 		return nil, types.ErrUnauthorized
 	}
 
+	// If we are not removing a hook make sure it has been already whitelisted
+	if msg.ContractAddr != "" {
+		// msg.ContractAddr has already been validated
+		cwAddr := sdk.MustAccAddressFromBech32(msg.ContractAddr)
+		if err := server.Keeper.AssertIsHookWhitelisted(ctx, msg.Denom, cwAddr); err != nil {
+			return nil, err
+		}
+	}
+
 	err = server.Keeper.setBeforeSendHook(ctx, msg.Denom, msg.ContractAddr)
 	if err != nil {
 		return nil, err

x/tokenfactory/keeper/params.go

@@ -29,3 +29,24 @@ func (k Keeper) SetParams(ctx sdk.Context, params types.Params) error {
 	store.Set(types.ParamsKey, bz)
 	return nil
 }
+
+func (k Keeper) AssertIsHookWhitelisted(ctx sdk.Context, denom string, contractAddress sdk.AccAddress) error {
+	contractInfo := k.contractKeeper.GetContractInfo(ctx, contractAddress)
+	if contractInfo == nil {
+		return types.ErrBeforeSendHookNotWhitelisted.Wrapf("contract with address (%s) does not exist", contractAddress.String())
+	}
+	codeID := contractInfo.CodeID
+	whitelistedHooks := k.GetParams(ctx).WhitelistedHooks
+	denomCreator, _, err := types.DeconstructDenom(denom)
+	if err != nil {
+		return types.ErrBeforeSendHookNotWhitelisted.Wrapf("invalid denom: %s", denom)
+	}
+
+	for _, hook := range whitelistedHooks {
+		if hook.CodeID == codeID && hook.DenomCreator == denomCreator {
+			return nil
+		}
+	}
+
+	return types.ErrBeforeSendHookNotWhitelisted.Wrapf("no whitelist for contract with codeID (%d) and denomCreator (%s) ", codeID, denomCreator)
+}