From 59038f7977d037bcf4547361f31d31b493eb205b Mon Sep 17 00:00:00 2001
From: Nicolas Lara <nicolaslara@gmail.com>
Date: Tue, 21 Nov 2023 14:41:11 +0100
Subject: [PATCH] better checks for affiliate swap contract in arb filter
 (#6906)

* better checks for affiliate swap contract

* properly unmarshalling
---
 x/txfees/keeper/txfee_filters/arb_tx.go      | 35 ++++++++++++++++++++
 x/txfees/keeper/txfee_filters/arb_tx_test.go | 15 +++++++++
 2 files changed, 50 insertions(+)

diff --git a/x/txfees/keeper/txfee_filters/arb_tx.go b/x/txfees/keeper/txfee_filters/arb_tx.go
index e6ff07acc74..d04507dc634 100644
--- a/x/txfees/keeper/txfee_filters/arb_tx.go
+++ b/x/txfees/keeper/txfee_filters/arb_tx.go
@@ -43,6 +43,9 @@ func (m AffiliateSwapMsg) TokenInDenom() string {
 
 // TokenOutDenom implements types.SwapMsgRoute.
 func (m AffiliateSwapMsg) TokenOutDenom() string {
+	if len(m.Routes) == 0 {
+		return "no-token-out"
+	}
 	lastPoolInRoute := m.Routes[len(m.Routes)-1]
 	return lastPoolInRoute.TokenOutDenom
 }
@@ -102,6 +105,12 @@ func isArbTxLooseAuthz(msg sdk.Msg, swapInDenom string, lpTypesSeen map[gammtype
 
 		// Get the contract message and attempt to unmarshal it into the affiliate swap message
 		contractMessage := msgExecuteContract.GetMsg()
+
+		// Check that the contract message is an affiliate swap message
+		if ok := isAffiliateSwapMsg(contractMessage); !ok {
+			return swapInDenom, false
+		}
+
 		var affiliateSwapMsg AffiliateSwapMsg
 		if err := json.Unmarshal(contractMessage, &affiliateSwapMsg); err != nil {
 			// If we can't unmarshal it, it's not an affiliate swap message
@@ -160,3 +169,29 @@ func isArbTxLooseSwapMsg(swapMsg poolmanagertypes.SwapMsgRoute, swapInDenom stri
 	swapInDenom = swapMsg.TokenInDenom()
 	return swapInDenom, false
 }
+
+// TODO: Make this generic by using isJsonSuperset from https://github.com/osmosis-labs/osmosis/blob/d56de7365428f0282eeab05c1cc75787370ef997/x/authenticator/authenticator/message_filter.go#L173C6-L173C12
+func isAffiliateSwapMsg(msg []byte) bool {
+	// Check that the contract message is a valid JSON object
+	jsonObject := make(map[string]interface{})
+	err := json.Unmarshal(msg, &jsonObject)
+	if err != nil {
+		return false
+	}
+
+	// check the main key is "swap"
+	swap, ok := jsonObject["swap"].(map[string]interface{})
+	if !ok {
+		return false
+	}
+
+	if routes, ok := swap["routes"].([]interface{}); !ok || len(routes) == 0 {
+		return false
+	}
+
+	if tokenOutMinAmount, ok := swap["token_out_min_amount"].(map[string]interface{}); !ok || len(tokenOutMinAmount) == 0 {
+		return false
+	}
+
+	return true
+}
diff --git a/x/txfees/keeper/txfee_filters/arb_tx_test.go b/x/txfees/keeper/txfee_filters/arb_tx_test.go
index 3494d32704b..d1b6e2efa26 100644
--- a/x/txfees/keeper/txfee_filters/arb_tx_test.go
+++ b/x/txfees/keeper/txfee_filters/arb_tx_test.go
@@ -64,3 +64,18 @@ func (suite *KeeperTestSuite) TestIsArbTxLooseAuthz_AffiliateSwapMsg() {
 	_, isArb := txfee_filters.IsArbTxLooseAuthz(executeMsg, executeMsg.Funds[0].Denom, map[types.LiquidityChangeType]bool{})
 	suite.Require().True(isArb)
 }
+
+func (suite *KeeperTestSuite) TestIsArbTxLooseAuthz_OtherMsg() {
+	otherMsg := []byte(`{"update_feed": {}}`)
+
+	// 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:      otherMsg,
+	}
+
+	_, isArb := txfee_filters.IsArbTxLooseAuthz(executeMsg, executeMsg.Funds[0].Denom, map[types.LiquidityChangeType]bool{})
+	suite.Require().False(isArb)
+}