From a6dcb24ec51554bf16b4726da21dab2e7519009f Mon Sep 17 00:00:00 2001
From: Carlton Hanna <nullpointer0x00@gmail.com>
Date: Mon, 20 Mar 2023 13:37:06 -0600
Subject: [PATCH] 1256 support for tokens restricted with attributes (#563)

* add injection of attribute checking...still POC atm

* expect restricted function to only return err, refactor names

* move send restrictions check to SendCoins

* remove line that was missed during refactor

* add code comments

* move accept logic to msg server

* Revert "add code comments"

This reverts commit d88da386badf7405ff9045f6b32a3f5202426fd0.

* revert ensure restirctions function into send keeper, add restrictions check to mult-send

* add back send enable checks to msg server

* add change log
---
 CHANGELOG.md                |  4 ++++
 x/bank/keeper/keeper.go     |  2 ++
 x/bank/keeper/msg_server.go |  1 -
 x/bank/keeper/send.go       | 44 ++++++++++++++++++++++++++++++++++++-
 4 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6a2e8f9177d9..ee386f99a496 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -39,6 +39,10 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ## Unreleased
 
+### Features
+
+* [PR 563](https://github.com/provenance-io/cosmos-sdk/pull/563) Add support for applying send restrictions before doing a bank send
+
 ### Improvements
 
 * [PR 565](https://github.com/provenance-io/cosmos-sdk/pull/565) Removes locks around state listening. There's no concurrency risk.
diff --git a/x/bank/keeper/keeper.go b/x/bank/keeper/keeper.go
index dc34a1235577..ff56d90c7153 100644
--- a/x/bank/keeper/keeper.go
+++ b/x/bank/keeper/keeper.go
@@ -51,6 +51,8 @@ type Keeper interface {
 
 	GetAuthority() string
 
+	SetSendRestrictionsFunc(sendRestrictionsFunc func(sdk.Context, string, string, string) error)
+
 	types.QueryServer
 }
 
diff --git a/x/bank/keeper/msg_server.go b/x/bank/keeper/msg_server.go
index a98500506420..9011a0aa9a72 100644
--- a/x/bank/keeper/msg_server.go
+++ b/x/bank/keeper/msg_server.go
@@ -115,4 +115,3 @@ func (k msgServer) MultiSend(goCtx context.Context, msg *types.MsgMultiSend) (*t
 
 	return &types.MsgMultiSendResponse{}, nil
 }
-
diff --git a/x/bank/keeper/send.go b/x/bank/keeper/send.go
index b83275d83886..e519874bf447 100644
--- a/x/bank/keeper/send.go
+++ b/x/bank/keeper/send.go
@@ -37,6 +37,7 @@ type SendKeeper interface {
 
 	IsSendEnabledCoin(ctx sdk.Context, coin sdk.Coin) bool
 	IsSendEnabledCoins(ctx sdk.Context, coins ...sdk.Coin) error
+	EnsureSendRestrictions(ctx sdk.Context, from, to string, coins ...sdk.Coin) error
 
 	BlockedAddr(addr sdk.AccAddress) bool
 }
@@ -58,6 +59,8 @@ type BaseSendKeeper struct {
 
 	qk types.QuarantineKeeper
 	sk types.SanctionKeeper
+
+	sendRestrictionsFunc func(sdk.Context, string, string, string) error
 }
 
 func NewBaseSendKeeper(
@@ -109,6 +112,15 @@ func (k BaseSendKeeper) GetParams(ctx sdk.Context) (params types.Params) {
 	return params
 }
 
+// SetSendRestrictionsFunc set a function to be called before sends can occur
+// if not set, it is a no-op
+func (k *BaseSendKeeper) SetSendRestrictionsFunc(sendRestrictionsFunc func(sdk.Context, string, string, string) error) {
+	if k.sendRestrictionsFunc != nil {
+		panic("the send restrictions function has already been set")
+	}
+	k.sendRestrictionsFunc = sendRestrictionsFunc
+}
+
 // SetParams sets the total set of bank parameters.
 func (k BaseSendKeeper) SetParams(ctx sdk.Context, params types.Params) {
 	if len(params.SendEnabled) > 0 {
@@ -128,6 +140,16 @@ func (k BaseSendKeeper) InputOutputCoins(ctx sdk.Context, inputs []types.Input,
 		return err
 	}
 
+	// ensure all inputs and outputs pass any restrictions
+	for _, input := range inputs {
+		for _, output := range outputs {
+			err := k.EnsureSendRestrictions(ctx, input.Address, output.Address, input.Coins...)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
 	allInputAddrs := make([]sdk.AccAddress, len(inputs))
 
 	for i, in := range inputs {
@@ -229,7 +251,12 @@ func (k BaseSendKeeper) SendCoins(ctx sdk.Context, fromAddr sdk.AccAddress, toAd
 // of possible quarantine on the toAddr.
 // An error is returned upon failure.
 func (k BaseSendKeeper) SendCoinsBypassQuarantine(ctx sdk.Context, fromAddr sdk.AccAddress, toAddr sdk.AccAddress, amt sdk.Coins) error {
-	err := k.subUnlockedCoins(ctx, fromAddr, amt)
+	err := k.EnsureSendRestrictions(ctx, fromAddr.String(), toAddr.String(), amt...)
+	if err != nil {
+		return err
+	}
+
+	err = k.subUnlockedCoins(ctx, fromAddr, amt)
 	if err != nil {
 		return err
 	}
@@ -424,6 +451,21 @@ func (k BaseSendKeeper) IsSendEnabledCoins(ctx sdk.Context, coins ...sdk.Coin) e
 	return nil
 }
 
+// EnsureSendRestrictions applies the send restrictions function returns error if send restrictions do not pass
+// no-op if the send restrictions function is not set
+func (k *BaseSendKeeper) EnsureSendRestrictions(ctx sdk.Context, from, to string, coins ...sdk.Coin) error {
+	if k.sendRestrictionsFunc == nil {
+		return nil
+	}
+	for _, coin := range coins {
+		err := k.sendRestrictionsFunc(ctx, from, to, coin.Denom)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // IsSendEnabledCoin returns the current SendEnabled status of the provided coin's denom
 func (k BaseSendKeeper) IsSendEnabledCoin(ctx sdk.Context, coin sdk.Coin) bool {
 	return k.IsSendEnabledDenom(ctx, coin.Denom)