feat(bank/v2): Introduce global send restriction (#21925)

x/bank/proto/cosmos/bank/module/v2/module.proto

@@ -14,4 +14,10 @@ message Module {
 
   // authority defines the custom module authority. If not set, defaults to the governance module.
   string authority = 1;
+
+  // restrictions_order specifies the order of send restrictions and should be
+  // a list of module names which provide a send restriction instance. If no
+  // order is provided, then restrictions will be applied in alphabetical order
+  // of module names.
+  repeated string restrictions_order = 2;
 }

x/bank/v2/depinject.go

@@ -1,6 +1,11 @@
 package bankv2
 
 import (
+	"fmt"
+	"maps"
+	"slices"
+	"sort"
+
 	"cosmossdk.io/core/address"
 	"cosmossdk.io/core/appmodule"
 	"cosmossdk.io/depinject"
@@ -22,6 +27,7 @@ func init() {
 	appconfig.RegisterModule(
 		&moduletypes.Module{},
 		appconfig.Provide(ProvideModule),
+		appconfig.Invoke(InvokeSetSendRestrictions),
 	)
 }
 
@@ -61,3 +67,39 @@ func ProvideModule(in ModuleInputs) ModuleOutputs {
 		Module: m,
 	}
 }
+
+func InvokeSetSendRestrictions(
+	config *moduletypes.Module,
+	keeper keeper.Keeper,
+	restrictions map[string]types.SendRestrictionFn,
+) error {
+	if config == nil {
+		return nil
+	}
+
+	modules := slices.Collect(maps.Keys(restrictions))
+	order := config.RestrictionsOrder
+	if len(order) == 0 {
+		order = modules
+		sort.Strings(order)
+	}
+
+	if len(order) != len(modules) {
+		return fmt.Errorf("len(restrictions order: %v) != len(restriction modules: %v)", order, modules)
+	}
+
+	if len(modules) == 0 {
+		return nil
+	}
+
+	for _, module := range order {
+		restriction, ok := restrictions[module]
+		if !ok {
+			return fmt.Errorf("can't find send restriction for module %s", module)
+		}
+
+		keeper.AppendGlobalSendRestriction(restriction)
+	}
+
+	return nil
+}

x/bank/v2/keeper/keeper.go

@@ -28,18 +28,21 @@ type Keeper struct {
 	params       collections.Item[types.Params]
 	balances     *collections.IndexedMap[collections.Pair[[]byte, string], math.Int, BalancesIndexes]
 	supply       collections.Map[string, math.Int]
+
+	sendRestriction *sendRestriction
 }
 
 func NewKeeper(authority []byte, addressCodec address.Codec, env appmodulev2.Environment, cdc codec.BinaryCodec) *Keeper {
 	sb := collections.NewSchemaBuilder(env.KVStoreService)
 
 	k := &Keeper{
-		Environment:  env,
-		authority:    authority,
-		addressCodec: addressCodec, // TODO(@julienrbrt): Should we add address codec to the environment?
-		params:       collections.NewItem(sb, types.ParamsKey, "params", codec.CollValue[types.Params](cdc)),
-		balances:     collections.NewIndexedMap(sb, types.BalancesPrefix, "balances", collections.PairKeyCodec(collections.BytesKey, collections.StringKey), sdk.IntValue, newBalancesIndexes(sb)),
-		supply:       collections.NewMap(sb, types.SupplyKey, "supply", collections.StringKey, sdk.IntValue),
+		Environment:     env,
+		authority:       authority,
+		addressCodec:    addressCodec, // TODO(@julienrbrt): Should we add address codec to the environment?
+		params:          collections.NewItem(sb, types.ParamsKey, "params", codec.CollValue[types.Params](cdc)),
+		balances:        collections.NewIndexedMap(sb, types.BalancesPrefix, "balances", collections.PairKeyCodec(collections.BytesKey, collections.StringKey), sdk.IntValue, newBalancesIndexes(sb)),
+		supply:          collections.NewMap(sb, types.SupplyKey, "supply", collections.StringKey, sdk.IntValue),
+		sendRestriction: newSendRestriction(),
 	}
 
 	schema, err := sb.Build()
@@ -94,7 +97,10 @@ func (k Keeper) SendCoins(ctx context.Context, from, to []byte, amt sdk.Coins) e
 	}
 
 	var err error
-	// TODO: Send restriction
+	to, err = k.sendRestriction.apply(ctx, from, to, amt)
+	if err != nil {
+		return err
+	}
 
 	err = k.subUnlockedCoins(ctx, from, amt)
 	if err != nil {

x/bank/v2/keeper/keeper_test.go

@@ -1,7 +1,9 @@
 package keeper_test
 
 import (
+	"bytes"
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -184,3 +186,53 @@ func (suite *KeeperTestSuite) TestSendCoins_Module_To_Module() {
 	mintBarBalance := suite.bankKeeper.GetBalance(ctx, mintAcc.GetAddress(), barDenom)
 	require.Equal(mintBarBalance.Amount, math.NewInt(0))
 }
+
+func (suite *KeeperTestSuite) TestSendCoins_WithRestriction() {
+	ctx := suite.ctx
+	require := suite.Require()
+	balances := sdk.NewCoins(newFooCoin(100), newBarCoin(50))
+	sendAmt := sdk.NewCoins(newFooCoin(10), newBarCoin(10))
+
+	require.NoError(banktestutil.FundAccount(ctx, suite.bankKeeper, accAddrs[0], balances))
+
+	// Add first restriction
+	addrRestrictFunc := func(ctx context.Context, from, to []byte, amount sdk.Coins) ([]byte, error) {
+		if bytes.Equal(from, to) {
+			return nil, fmt.Errorf("Can not send to same address")
+		}
+		return to, nil
+	}
+	suite.bankKeeper.AppendGlobalSendRestriction(addrRestrictFunc)
+
+	err := suite.bankKeeper.SendCoins(ctx, accAddrs[0], accAddrs[0], sendAmt)
+	require.Error(err)
+	require.Contains(err.Error(), "Can not send to same address")
+
+	// Add second restriction
+	amtRestrictFunc := func(ctx context.Context, from, to []byte, amount sdk.Coins) ([]byte, error) {
+		if len(amount) > 1 {
+			return nil, fmt.Errorf("Allow only one denom per one send")
+		}
+		return to, nil
+	}
+	suite.bankKeeper.AppendGlobalSendRestriction(amtRestrictFunc)
+
+	// Pass the 1st but failt at the 2nd
+	err = suite.bankKeeper.SendCoins(ctx, accAddrs[0], accAddrs[1], sendAmt)
+	require.Error(err)
+	require.Contains(err.Error(), "Allow only one denom per one send")
+
+	// Pass both 2 restrictions
+	err = suite.bankKeeper.SendCoins(ctx, accAddrs[0], accAddrs[1], sdk.NewCoins(newFooCoin(10)))
+	require.NoError(err)
+
+	// Check balances
+	acc0FooBalance := suite.bankKeeper.GetBalance(ctx, accAddrs[0], fooDenom)
+	require.Equal(acc0FooBalance.Amount, math.NewInt(90))
+	acc0BarBalance := suite.bankKeeper.GetBalance(ctx, accAddrs[0], barDenom)
+	require.Equal(acc0BarBalance.Amount, math.NewInt(50))
+	acc1FooBalance := suite.bankKeeper.GetBalance(ctx, accAddrs[1], fooDenom)
+	require.Equal(acc1FooBalance.Amount, math.NewInt(10))
+	acc1BarBalance := suite.bankKeeper.GetBalance(ctx, accAddrs[1], barDenom)
+	require.Equal(acc1BarBalance.Amount, math.ZeroInt())
+}

x/bank/v2/keeper/restriction.go

@@ -0,0 +1,62 @@
+package keeper
+
+import (
+	"context"
+
+	"cosmossdk.io/x/bank/v2/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+// sendRestriction is a struct that houses a SendRestrictionFn.
+// It exists so that the SendRestrictionFn can be updated in the SendKeeper without needing to have a pointer receiver.
+type sendRestriction struct {
+	fn types.SendRestrictionFn
+}
+
+// newSendRestriction creates a new sendRestriction with nil send restriction.
+func newSendRestriction() *sendRestriction {
+	return &sendRestriction{
+		fn: nil,
+	}
+}
+
+// append adds the provided restriction to this, to be run after the existing function.
+func (r *sendRestriction) append(restriction types.SendRestrictionFn) {
+	r.fn = r.fn.Then(restriction)
+}
+
+// prepend adds the provided restriction to this, to be run before the existing function.
+func (r *sendRestriction) prepend(restriction types.SendRestrictionFn) {
+	r.fn = restriction.Then(r.fn)
+}
+
+// clear removes the send restriction (sets it to nil).
+func (r *sendRestriction) clear() {
+	r.fn = nil
+}
+
+var _ types.SendRestrictionFn = (*sendRestriction)(nil).apply
+
+// apply applies the send restriction if there is one. If not, it's a no-op.
+func (r *sendRestriction) apply(ctx context.Context, fromAddr, toAddr []byte, amt sdk.Coins) ([]byte, error) {
+	if r == nil || r.fn == nil {
+		return toAddr, nil
+	}
+	return r.fn(ctx, fromAddr, toAddr, amt)
+}
+
+// AppendSendRestriction adds the provided SendRestrictionFn to run after previously provided restrictions.
+func (k Keeper) AppendGlobalSendRestriction(restriction types.SendRestrictionFn) {
+	k.sendRestriction.append(restriction)
+}
+
+// PrependSendRestriction adds the provided SendRestrictionFn to run before previously provided restrictions.
+func (k Keeper) PrependGlobalSendRestriction(restriction types.SendRestrictionFn) {
+	k.sendRestriction.prepend(restriction)
+}
+
+// ClearSendRestriction removes the send restriction (if there is one).
+func (k Keeper) ClearGlobalSendRestriction() {
+	k.sendRestriction.clear()
+}