feat: enable arb filter for affiliate swap messages (#6890)

* feat: enable arb filter for affiliate swap messages

* updates

CHANGELOG.md

@@ -61,6 +61,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#6788](https://github.com/osmosis-labs/osmosis/pull/6788) Improve error message when CL LP fails due to slippage bound hit.
 * [#6858](https://github.com/osmosis-labs/osmosis/pull/6858) Merge mempool improvements from v20
 * [#6861](https://github.com/osmosis-labs/osmosis/pull/6861) Protorev address added to reduced taker fee whitelist
+* [#6890](https://github.com/osmosis-labs/osmosis/pull/6890) Enable arb filter for affiliate swap contract
 
 ### API Breaks
 

x/txfees/keeper/txfee_filters/arb_tx.go

@@ -1,14 +1,54 @@
 package txfee_filters
 
 import (
+	"encoding/json"
+
 	authztypes "github.com/cosmos/cosmos-sdk/x/authz"
 
+	wasmtypes "github.com/CosmWasm/wasmd/x/wasm/types"
+
 	gammtypes "github.com/osmosis-labs/osmosis/v20/x/gamm/types"
 	poolmanagertypes "github.com/osmosis-labs/osmosis/v20/x/poolmanager/types"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
+// See this for reference: https://github.com/osmosis-labs/affiliate-swap
+type Swap struct {
+	Routes            []poolmanagertypes.SwapAmountInRoute `json:"routes"`
+	TokenOutMinAmount sdk.Coin                             `json:"token_out_min_amount"`
+	FeePercentage     sdk.Dec                              `json:"fee_percentage"`
+	FeeCollector      string                               `json:"fee_collector"`
+	TokenIn           string                               `json:"token_in,omitempty"`
+}
+
+type AffiliateSwapMsg struct {
+	Swap `json:"swap"`
+}
+
+// TokenDenomsOnPath implements types.SwapMsgRoute.
+func (m AffiliateSwapMsg) TokenDenomsOnPath() []string {
+	denoms := make([]string, 0, len(m.Routes)+1)
+	denoms = append(denoms, m.TokenInDenom())
+	for i := 0; i < len(m.Routes); i++ {
+		denoms = append(denoms, m.Routes[i].TokenOutDenom)
+	}
+	return denoms
+}
+
+// TokenInDenom implements types.SwapMsgRoute.
+func (m AffiliateSwapMsg) TokenInDenom() string {
+	return m.TokenIn
+}
+
+// TokenOutDenom implements types.SwapMsgRoute.
+func (m AffiliateSwapMsg) TokenOutDenom() string {
+	lastPoolInRoute := m.Routes[len(m.Routes)-1]
+	return lastPoolInRoute.TokenOutDenom
+}
+
+var _ poolmanagertypes.SwapMsgRoute = AffiliateSwapMsg{}
+
 // We check if a tx is an arbitrage for the mempool right now by seeing:
 // 1) does start token of a msg = final token of msg (definitionally correct)
 // 2) does it have multiple swap messages, with different tx ins. If so, we assume its an arb.
@@ -49,6 +89,35 @@ func isArbTxLooseAuthz(msg sdk.Msg, swapInDenom string, lpTypesSeen map[gammtype
 		return swapInDenom, false
 	}
 
+	// Detects the affiliate swap message from the CosmWasm contract
+	// See an example here:
+	// // https://celatone.osmosis.zone/osmosis-1/txs/315EB6284778EBB5BAC0F94CC740F5D7E35DDA5BBE4EC9EC79F012548589C6E5
+	if msgExecuteContract, ok := msg.(*wasmtypes.MsgExecuteContract); ok {
+		// Grab token in from the funds sent to the contract
+		tokensIn := msgExecuteContract.GetFunds()
+		if len(tokensIn) != 1 {
+			return swapInDenom, false
+		}
+		tokenIn := tokensIn[0]
+
+		// Get the contract message and attempt to unmarshal it into the affiliate swap message
+		contractMessage := msgExecuteContract.GetMsg()
+		var affiliateSwapMsg AffiliateSwapMsg
+		if err := json.Unmarshal(contractMessage, &affiliateSwapMsg); err != nil {
+			// If we can't unmarshal it, it's not an affiliate swap message
+			return swapInDenom, false
+		}
+
+		// Otherwise, we have an affiliate swap message, so we check if it's an arb
+		affiliateSwapMsg.TokenIn = tokenIn.Denom
+		swapInDenom, isArb := isArbTxLooseSwapMsg(affiliateSwapMsg, swapInDenom)
+		if isArb {
+			return swapInDenom, true
+		}
+
+		return swapInDenom, false
+	}
+
 	// (4) Check that the tx doesn't have both JoinPool & ExitPool msgs
 	lpMsg, isLpMsg := msg.(gammtypes.LiquidityChangeMsg)
 	if isLpMsg {

x/txfees/keeper/txfee_filters/arb_tx_test.go

@@ -0,0 +1,66 @@
+package txfee_filters_test
+
+import (
+	"encoding/json"
+	"testing"
+
+	wasmtypes "github.com/CosmWasm/wasmd/x/wasm/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/stretchr/testify/suite"
+
+	"github.com/osmosis-labs/osmosis/v20/app/apptesting"
+	"github.com/osmosis-labs/osmosis/v20/x/gamm/types"
+	poolmanagertypes "github.com/osmosis-labs/osmosis/v20/x/poolmanager/types"
+	"github.com/osmosis-labs/osmosis/v20/x/txfees/keeper/txfee_filters"
+)
+
+type KeeperTestSuite struct {
+	apptesting.KeeperTestHelper
+}
+
+func TestTxFeeFilters(t *testing.T) {
+	suite.Run(t, new(KeeperTestSuite))
+}
+
+// Tests that the arb filter is enabled on the affiliate swap msg.
+func (suite *KeeperTestSuite) TestIsArbTxLooseAuthz_AffiliateSwapMsg() {
+	affiliateSwapMsg := &txfee_filters.AffiliateSwapMsg{
+		Swap: txfee_filters.Swap{
+			FeeCollector:  "osmo1dldrxz5p8uezxz3qstpv92de7wgfp7hvr72dcm",
+			FeePercentage: sdk.ZeroDec(),
+			Routes: []poolmanagertypes.SwapAmountInRoute{
+				{
+					PoolId:        1221,
+					TokenOutDenom: "uosmo",
+				},
+				{
+					PoolId:        3,
+					TokenOutDenom: "ibc/1480B8FD20AD5FCAE81EA87584D269547DD4D436843C1D20F15E00EB64743EF4",
+				},
+				{
+					PoolId:        4,
+					TokenOutDenom: "ibc/27394FB092D2ECCD56123C74F36E4C1F926001CEADA9CA97EA622B25F41E5EB",
+				},
+				{
+					PoolId:        1251,
+					TokenOutDenom: "ibc/498A0751C798A0D9A389AA3691123DADA57DAA4FE165D5C75894505B876BA6E4",
+				},
+			},
+			TokenOutMinAmount: sdk.NewCoin("ibc/498A0751C798A0D9A389AA3691123DADA57DAA4FE165D5C75894505B876BA6E4", sdk.NewInt(217084399)),
+		},
+	}
+
+	affiliateSwapMsgBz, err := json.Marshal(affiliateSwapMsg)
+	suite.Require().NoError(err)
+
+	// https://celatone.osmosis.zone/osmosis-1/txs/315EB6284778EBB5BAC0F94CC740F5D7E35DDA5BBE4EC9EC79F012548589C6E5
+	executeMsg := &wasmtypes.MsgExecuteContract{
+		Contract: "osmo1etpha3a65tds0hmn3wfjeag6wgxgrkuwg2zh94cf5hapz7mz04dq6c25s5",
+		Sender:   "osmo1dldrxz5p8uezxz3qstpv92de7wgfp7hvr72dcm",
+		Funds:    sdk.NewCoins(sdk.NewCoin("ibc/498A0751C798A0D9A389AA3691123DADA57DAA4FE165D5C75894505B876BA6E4", sdk.NewInt(217084399))),
+		Msg:      affiliateSwapMsgBz,
+	}
+
+	_, isArb := txfee_filters.IsArbTxLooseAuthz(executeMsg, executeMsg.Funds[0].Denom, map[types.LiquidityChangeType]bool{})
+	suite.Require().True(isArb)
+}

x/txfees/keeper/txfee_filters/export_test.go

@@ -0,0 +1,11 @@
+package txfee_filters
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	gammtypes "github.com/osmosis-labs/osmosis/v20/x/gamm/types"
+)
+
+func IsArbTxLooseAuthz(msg sdk.Msg, swapInDenom string, lpTypesSeen map[gammtypes.LiquidityChangeType]bool) (string, bool) {
+	return isArbTxLooseAuthz(msg, swapInDenom, lpTypesSeen)
+}