feat(x/cosmwasmpool): Sending token_in_max_amount to the contract before running contract msg

CHANGELOG.md

@@ -55,6 +55,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * [#5534](https://github.com/osmosis-labs/osmosis/pull/5534) fix: fix the account number of x/tokenfactory module account
 * [#5750](https://github.com/osmosis-labs/osmosis/pull/5750) feat: add cli commmand for converting proto structs to proto marshalled bytes
+* [#5855](https://github.com/osmosis-labs/osmosis/pull/5855) feat(x/cosmwasmpool): Sending token_in_max_amount to the contract before running contract msg 
 
 ## v16.0.0
 Osmosis Labs is excited to announce the release of v16.0.0, a major upgrade that includes a number of new features and improvements like introduction of new modules, updates existing APIs, and dependency updates. This upgrade aims to enhance capital efficiency by introducing SuperCharged Liquidity, introduce custom liquidity pools backed by CosmWasm smart contracts, and improve overall functionality.

x/cosmwasmpool/README.md

@@ -92,18 +92,6 @@ It's important to note that the _**contract itselfs hold tokens that are provide
 
 One of the main reason why CosmWasm pool is implemented as a module + contract rather than a contract only is that it allows us to use the existing pool manager module to handle swap, which means things like swap routing, cross chain swap, and other functionality that depends on existing pool interface works out of the box.
 
-```mermaid
-graph TD;
-  Sender((Sender))
-  Sender -- swap --> x/poolmanager
-  x/poolmanager -- route msg to --> x/cosmwasmpool
-  x/cosmwasmpool -- sudo execute contract --> x/wasm
-  x/wasm -- sudo --> wasm/pool
-
-  x/cosmwasmpool -- send token_in from sender to wasm/pool --> x/bank
-  wasm/pool -- send token_out to sender --> x/bank
-```
-
 Pool contract's sudo endpoint expect the following message variant:
 
 ```rs
@@ -131,12 +119,45 @@ SwapExactAmountOut {
 },
 ```
 
+`SwapExactAmountIn`
+
+
+```mermaid
+graph TD;
+  Sender((Sender))
+  Sender -- 1. swap --> x/poolmanager
+  x/poolmanager -- 2. route msg to --> x/cosmwasmpool
+  x/cosmwasmpool -- 3. sudo execute contract --> x/wasm
+
+  x/cosmwasmpool -- 4. send token_in from sender to wasm/pool --> x/bank
+  x/wasm -- 5. sudo --> wasm/pool
+  wasm/pool -- 6. send token_out to sender --> x/bank
+```
+
+
+`SwapExactAmountOut`
+
+```mermaid
+graph TD;
+  Sender((Sender))
+  Sender -- 1. swap --> x/poolmanager
+  x/poolmanager -- 2. route msg to --> x/cosmwasmpool
+  x/cosmwasmpool -- 3. sudo execute contract --> x/wasm
+
+  x/cosmwasmpool -- 4. send token_in_max_amount from sender to wasm/pool --> x/bank
+  x/wasm -- 5. sudo --> wasm/pool
+  x/cosmwasmpool -- 6. send remaining wasm/pool to sender --> x/bank
+
+  wasm/pool -- 7. send token_out to sender --> x/bank
+```
+
 The reason why this needs to be sudo endpoint, which can only be called by the chain itself, is that the chain can provide correct information about `swap_fee`, which can be deviated from contract defined `swap_fee` in multihop scenario.
 
 `swap_fee` in this context is intended to be fee that is collected by liquidity providers. If the contract provider wants to collect fee for itself, it should implement its own fee collection mechanism.
 
 And because sudo message can't attach funds like execute message, chain-side is required to perform sending token to the contract and ensure that `token_in` and `token_in_max_amount` is exactly the same amount of token that gets sent to the contract.
 
+And the reason why the sequence is a little bit different for `SwapExactAmountIn` and `SwapExactAmountOut` is that, for `SwapExactAmountIn`, it is known beforehand how much `token_in_amount` to be sent to the contract and let it process, but for `SwapExactAmountOut`, it isn't, so we need to sent `token_in_max_amount` to the contract and let it process, then send the remaining token back to the sender.
 
 ## Deactivating
 

x/cosmwasmpool/pool_module.go

@@ -282,18 +282,41 @@ func (k Keeper) SwapExactAmountOut(
 		return sdk.Int{}, err
 	}
 
+	contractAddr := sdk.MustAccAddressFromBech32(cosmwasmPool.GetContractAddress())
+
+	// Send token in max amount from sender to the pool
+	// We do this because sudo message does not support sending coins from the sender
+	// And we need to send the max amount because we do not know how much the contract will use.
+	if err := k.bankKeeper.SendCoins(ctx, sender, contractAddr, sdk.NewCoins(sdk.NewCoin(tokenInDenom, tokenInMaxAmount))); err != nil {
+		return sdk.Int{}, err
+	}
+
+	// Note that the contract sends the token out back to the sender after the swap
+	// As a result, we do not need to worry about sending token out here.
 	request := msg.NewSwapExactAmountOutSudoMsg(sender.String(), tokenInDenom, tokenOut, tokenInMaxAmount, swapFee)
 	response, err := cosmwasm.Sudo[msg.SwapExactAmountOutSudoMsg, msg.SwapExactAmountOutSudoMsgResponse](ctx, k.contractKeeper, cosmwasmPool.GetContractAddress(), request)
 	if err != nil {
 		return sdk.Int{}, err
 	}
 
-	// Send token in from sender to the pool
-	// We do this because sudo message does not support sending coins from the sender
-	// However, note that the contract sends the token back to the sender after the swap
-	// As a result, we do not need to worry about sending it back here.
-	if err := k.bankKeeper.SendCoins(ctx, sender, sdk.MustAccAddressFromBech32(cosmwasmPool.GetContractAddress()), sdk.NewCoins(sdk.NewCoin(tokenInDenom, response.TokenInAmount))); err != nil {
-		return sdk.Int{}, err
+	tokenInExcessiveAmount := tokenInMaxAmount.Sub(response.TokenInAmount)
+
+	// Do not send any coins if excessive amount is zero
+
+	// required amount should be less than or equal to max amount
+	if tokenInExcessiveAmount.IsNegative() {
+		return sdk.Int{}, types.NegativeExcessiveTokenInAmountError{
+			TokenInMaxAmount:       tokenInMaxAmount,
+			TokenInRequiredAmount:  response.TokenInAmount,
+			TokenInExcessiveAmount: tokenInExcessiveAmount,
+		}
+	}
+
+	// Send excessibe token in from pool back to sender
+	if tokenInExcessiveAmount.IsPositive() {
+		if err := k.bankKeeper.SendCoins(ctx, contractAddr, sender, sdk.NewCoins(sdk.NewCoin(tokenInDenom, tokenInExcessiveAmount))); err != nil {
+			return sdk.Int{}, err
+		}
 	}
 
 	return response.TokenInAmount, nil

x/cosmwasmpool/pool_module_test.go

@@ -18,7 +18,7 @@ const (
 	denomA        = apptesting.DefaultTransmuterDenomA
 	denomB        = apptesting.DefaultTransmuterDenomB
 	validCodeId   = uint64(1)
-	invalidCodeId        = validCodeId + 1
+	invalidCodeId = validCodeId + 1
 	defaultPoolId = uint64(1)
 	nonZeroFeeStr = "0.01"
 )
@@ -316,20 +316,23 @@ func (s *PoolModuleSuite) TestCalcInAmtGivenOut_SwapInAmtGivenOut() {
 			initialCoins:         initalDefaultSupply,
 			tokenOut:             sdk.NewCoin(denomA, defaultAmount.Add(sdk.OneInt())),
 			tokenInDenom:         denomB,
+			tokenInMaxAmount:     defaultAmount.Sub(sdk.OneInt()),
 			expectedErrorMessage: fmt.Sprintf("Insufficient pool asset: required: %s, available: %s", sdk.NewCoin(denomA, defaultAmount.Add(sdk.OneInt())), sdk.NewCoin(denomA, defaultAmount)),
 		},
 		"non-zero swap fee": {
 			initialCoins:         initalDefaultSupply,
 			tokenOut:             sdk.NewCoin(denomA, defaultAmount.Sub(sdk.OneInt())),
 			tokenInDenom:         denomB,
+			tokenInMaxAmount:     defaultAmount.Sub(sdk.OneInt()),
 			swapFee:              sdk.MustNewDecFromStr(nonZeroFeeStr),
 			expectedErrorMessage: fmt.Sprintf("Invalid swap fee: expected: %s, actual: %s", sdk.ZeroInt(), nonZeroFeeStr),
 		},
 		"invalid pool given": {
-			initialCoins:  sdk.NewCoins(sdk.NewCoin(denomA, defaultAmount), sdk.NewCoin(denomB, defaultAmount)),
-			tokenOut:      sdk.NewCoin(denomA, defaultAmount.Sub(sdk.OneInt())),
-			tokenInDenom:  denomB,
-			isInvalidPool: true,
+			initialCoins:     sdk.NewCoins(sdk.NewCoin(denomA, defaultAmount), sdk.NewCoin(denomB, defaultAmount)),
+			tokenOut:         sdk.NewCoin(denomA, defaultAmount.Sub(sdk.OneInt())),
+			tokenInDenom:     denomB,
+			tokenInMaxAmount: defaultAmount.Sub(sdk.OneInt()),
+			isInvalidPool:    true,
 
 			expectedErrorMessage: types.InvalidPoolTypeError{
 				ActualPool: &clmodel.Pool{},
@@ -382,17 +385,11 @@ func (s *PoolModuleSuite) TestCalcInAmtGivenOut_SwapInAmtGivenOut() {
 			s.Require().Equal(originalPoolBalances.String(), afterCalcPoolBalances.String())
 
 			swapper := s.TestAccs[1]
+
 			// fund swapper
-			if !tc.expectedTokenIn.IsNil() {
-				// Fund with expected token in
-				s.FundAcc(swapper, sdk.NewCoins(tc.expectedTokenIn))
-			} else {
-				// Fund with pool reserve of token in denom
-				// This case happens in the error case, and we want
-				// to make sure that the error that we get is not
-				// due to insufficient funds.
-				s.FundAcc(swapper, sdk.NewCoins(sdk.NewCoin(tc.tokenInDenom, defaultAmount)))
-			}
+			s.FundAcc(swapper, sdk.NewCoins(sdk.NewCoin(tc.tokenInDenom, tc.tokenInMaxAmount)))
+
+			beforeSwapSwapperBalances := s.App.BankKeeper.GetAllBalances(s.Ctx, swapper)
 
 			// system under test non-mutative.
 			actualSwapTokenIn, err := cosmwasmPoolKeeper.SwapExactAmountOut(s.Ctx, swapper, poolIn, tc.tokenInDenom, tc.tokenInMaxAmount, tc.tokenOut, tc.swapFee)
@@ -411,7 +408,7 @@ func (s *PoolModuleSuite) TestCalcInAmtGivenOut_SwapInAmtGivenOut() {
 			s.Require().Equal(expectedPoolBalances.String(), afterSwapPoolBalances.String())
 
 			// Assert that swapper balance is updated correctly
-			expectedSwapperBalances := sdk.NewCoins(tc.tokenOut)
+			expectedSwapperBalances := beforeSwapSwapperBalances.Sub(sdk.NewCoins(tc.expectedTokenIn)).Add(tc.tokenOut)
 			afterSwapSwapperBalances := s.App.BankKeeper.GetAllBalances(s.Ctx, swapper)
 			s.Require().Equal(expectedSwapperBalances.String(), afterSwapSwapperBalances.String())
 		})

x/cosmwasmpool/types/errors.go

@@ -3,6 +3,8 @@ package types
 import (
 	"errors"
 	"fmt"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
 var (
@@ -34,3 +36,13 @@ type CodeIdNotWhitelistedError struct {
 func (e CodeIdNotWhitelistedError) Error() string {
 	return fmt.Sprintf("cannot create coswasm pool with the given code id (%d). Please whitelist it via governance", e.CodeId)
 }
+
+type NegativeExcessiveTokenInAmountError struct {
+	TokenInMaxAmount       sdk.Int
+	TokenInRequiredAmount  sdk.Int
+	TokenInExcessiveAmount sdk.Int
+}
+
+func (e NegativeExcessiveTokenInAmountError) Error() string {
+	return fmt.Sprintf("excessive token in amount cannot be negative. token in max amount = %d, token in required amount = %d, token in excessive amount = %d", e.TokenInMaxAmount, e.TokenInRequiredAmount, e.TokenInExcessiveAmount)
+}