Merge pull request #9 from oraichain/fix/zapper

fix: zap in with no route

contracts/zapper/src/contract.rs

@@ -69,14 +69,18 @@ pub fn execute(
             routes,
             minimum_liquidity,
         ),
+        ExecuteMsg::ZapInAfterSwapOperation {} => internal::zap_in_liquidity(deps, env, info),
+        ExecuteMsg::RefundAfterZapInLiquidity {} => internal::refund_after_zap_in(deps, env, info),
         ExecuteMsg::ZapOutLiquidity {
             position_index,
             routes,
         } => zap_out_liquidity(deps, env, info, position_index, routes),
+        ExecuteMsg::ZapOutAfterSwapOperation {} => internal::zap_out_liquidity(deps, env, info),
         ExecuteMsg::RegisterProtocolFee {
             percent,
             fee_receiver,
         } => execute_register_protocol_fee(deps, info, percent, fee_receiver),
+        ExecuteMsg::Withdraw { assets, recipient } => withdraw(deps, info, assets, recipient),
     }
 }
 
@@ -94,13 +98,3 @@ pub fn migrate(deps: DepsMut, _env: Env, _msg: MigrateMsg) -> Result<Response, C
         cw2::ensure_from_older_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
     Ok(Response::new().add_attribute("new_version", original_version.to_string()))
 }
-
-#[cfg_attr(not(feature = "library"), entry_point)]
-pub fn reply(deps: DepsMut, env: Env, reply: Reply) -> Result<Response, ContractError> {
-    match reply.id {
-        ZAP_IN_LIQUIDITY_REPLY_ID => reply::zap_in_liquidity(deps, env),
-        ZAP_OUT_LIQUIDITY_REPLY_ID => reply::zap_out_liquidity(deps, env),
-        ADD_LIQUIDITY_REPLY_ID => reply::add_liquidity(deps, env),
-        _ => Err(ContractError::UnrecognizedReplyId { id: reply.id }),
-    }
-}

contracts/zapper/src/entrypoints/execute.rs

@@ -14,12 +14,12 @@ use oraiswap_v3_common::{
 use crate::{
     contract::{ZAP_IN_LIQUIDITY_REPLY_ID, ZAP_OUT_LIQUIDITY_REPLY_ID},
     entrypoints::common::get_pool_v3_asset_info,
-    msg::Route,
+    msg::{ExecuteMsg, Route},
     state::{CONFIG, PENDING_POSITION, PROTOCOL_FEE, RECEIVER, SNAP_BALANCES, ZAP_OUT_ROUTES},
     Config, PairBalance, PendingPosition,
 };
 
-use super::{build_swap_msg, validate_fund};
+use super::{build_swap_msg, internal, validate_fund};
 
 pub fn update_config(
     deps: DepsMut,
@@ -89,7 +89,6 @@ pub fn zap_in_liquidity(
 ) -> Result<Response, ContractError> {
     // init messages and submessages
     let mut msgs: Vec<CosmosMsg> = vec![];
-    let mut sub_msgs: Vec<SubMsg> = vec![];
 
     // transfer the amount or check the fund is sent with request
     validate_fund(
@@ -134,6 +133,7 @@ pub fn zap_in_liquidity(
             owner: env.contract.address.clone(),
         },
     )?;
+
     let position = PendingPosition::new(
         position_length,
         pool_key.clone(),
@@ -176,17 +176,16 @@ pub fn zap_in_liquidity(
             None,
             None,
         )?;
-        if i == routes.len() - 1 {
-            sub_msgs.push(SubMsg::reply_on_success(
-                swap_msg,
-                ZAP_IN_LIQUIDITY_REPLY_ID,
-            ));
-        } else {
-            sub_msgs.push(SubMsg::new(swap_msg));
-        }
+        msgs.push(CosmosMsg::Wasm(swap_msg));
     }
 
-    Ok(Response::new().add_messages(msgs).add_submessages(sub_msgs))
+    msgs.push(CosmosMsg::Wasm(WasmMsg::Execute {
+        contract_addr: env.contract.address.to_string(),
+        msg: to_json_binary(&ExecuteMsg::ZapInAfterSwapOperation {})?,
+        funds: vec![],
+    }));
+
+    Ok(Response::new().add_messages(msgs))
 }
 
 pub fn zap_out_liquidity(
@@ -197,7 +196,6 @@ pub fn zap_out_liquidity(
     routes: Vec<Route>,
 ) -> Result<Response, ContractError> {
     let mut msgs: Vec<CosmosMsg> = vec![];
-    let mut sub_msgs: Vec<SubMsg> = vec![];
     let config = CONFIG.load(deps.storage)?;
     let position: Position = deps.querier.query_wasm_smart(
         config.dex_v3.to_string(),
@@ -243,16 +241,42 @@ pub fn zap_out_liquidity(
     ZAP_OUT_ROUTES.save(deps.storage, &routes)?;
 
     // 2. Create SubMsg to process remove liquidity in dex_v3 contract
-    sub_msgs.push(SubMsg::reply_on_success(
-        WasmMsg::Execute {
-            contract_addr: config.dex_v3.to_string(),
-            msg: to_json_binary(&V3ExecuteMsg::Burn {
-                token_id: position.token_id,
-            })?,
-            funds: vec![],
-        },
-        ZAP_OUT_LIQUIDITY_REPLY_ID,
-    ));
+    msgs.push(CosmosMsg::Wasm(WasmMsg::Execute {
+        contract_addr: config.dex_v3.to_string(),
+        msg: to_json_binary(&V3ExecuteMsg::Burn {
+            token_id: position.token_id,
+        })?,
+        funds: vec![],
+    }));
+
+    msgs.push(CosmosMsg::Wasm(WasmMsg::Execute {
+        contract_addr: env.contract.address.to_string(),
+        msg: to_json_binary(&ExecuteMsg::ZapOutAfterSwapOperation {})?,
+        funds: vec![],
+    }));
+
+    Ok(Response::new().add_messages(msgs))
+}
+
+pub fn withdraw(
+    deps: DepsMut,
+    info: MessageInfo,
+    assets: Vec<Asset>,
+    recipient: Option<Addr>,
+) -> Result<Response, ContractError> {
+    let config = CONFIG.load(deps.storage)?;
+    let receiver = recipient.unwrap_or_else(|| info.sender.clone());
+
+    if info.sender != config.admin {
+        return Err(ContractError::Unauthorized {});
+    }
+
+    let mut msgs: Vec<CosmosMsg> = vec![];
+    for asset in assets {
+        asset
+            .info
+            .transfer(&mut msgs, receiver.to_string(), asset.amount)?;
+    }
 
-    Ok(Response::new().add_messages(msgs).add_submessages(sub_msgs))
+    Ok(Response::new().add_messages(msgs))
 }

contracts/zapper/src/entrypoints/reply.rs → contracts/zapper/src/entrypoints/internal.rs

@@ -1,8 +1,7 @@
 use std::vec;
 
 use cosmwasm_std::{
-    to_json_binary, Coin, CosmosMsg, Decimal, DepsMut, Env, Order, Response, StdResult, SubMsg,
-    WasmMsg,
+    to_json_binary, Coin, CosmosMsg, Decimal, DepsMut, Env, MessageInfo, Order, Response, StdResult, SubMsg, WasmMsg
 };
 
 use oraiswap_v3_common::{
@@ -19,6 +18,7 @@ use oraiswap_v3_common::{
 
 use crate::{
     contract::ADD_LIQUIDITY_REPLY_ID,
+    msg::ExecuteMsg,
     state::{
         CONFIG, PENDING_POSITION, PROTOCOL_FEE, RECEIVER, SNAP_BALANCE, SNAP_BALANCES,
         ZAP_OUT_ROUTES,
@@ -28,9 +28,12 @@ use crate::{
 
 use super::build_swap_msg;
 
-pub fn zap_in_liquidity(deps: DepsMut, env: Env) -> Result<Response, ContractError> {
+pub fn zap_in_liquidity(deps: DepsMut, env: Env, info: MessageInfo) -> Result<Response, ContractError> {
+    if info.sender != env.contract.address {
+        return Err(ContractError::Unauthorized {});
+    }
+
     let mut msgs: Vec<CosmosMsg> = vec![];
-    let mut sub_msgs: Vec<SubMsg> = vec![];
 
     // 5. Recheck the balance of tokenX and tokenY in this contract
     let snap_balance = SNAP_BALANCE.load(deps.storage)?;
@@ -49,7 +52,7 @@ pub fn zap_in_liquidity(deps: DepsMut, env: Env) -> Result<Response, ContractErr
 
     // 7. Process create new position with amountX and amountY
     let config = CONFIG.load(deps.storage)?;
-    let pending_position = PENDING_POSITION.load(deps.storage)?;
+    let pending_position: crate::PendingPosition = PENDING_POSITION.load(deps.storage)?;
     let pool_info: Pool = deps.querier.query_wasm_smart(
         config.dex_v3.to_string(),
         &V3QueryMsg::Pool {
@@ -117,30 +120,37 @@ pub fn zap_in_liquidity(deps: DepsMut, env: Env) -> Result<Response, ContractErr
         .info
         .increase_allowance(&mut coins, &mut msgs, config.dex_v3.to_string(), y_amount)?;
 
-    sub_msgs.push(SubMsg::reply_on_success(
-        WasmMsg::Execute {
-            contract_addr: config.dex_v3.to_string(),
-            msg: to_json_binary(&V3ExecuteMsg::CreatePosition {
-                pool_key: pending_position.pool_key.clone(),
-                lower_tick: pending_position.lower_tick,
-                upper_tick: pending_position.upper_tick,
-                liquidity_delta: res.l,
-                slippage_limit_lower: pending_position.slippage_limit_lower.unwrap_or(
-                    get_min_sqrt_price(pending_position.pool_key.fee_tier.tick_spacing),
-                ),
-                slippage_limit_upper: pending_position.slippage_limit_upper.unwrap_or(
-                    get_max_sqrt_price(pending_position.pool_key.fee_tier.tick_spacing),
-                ),
-            })?,
-            funds: coins,
-        },
-        ADD_LIQUIDITY_REPLY_ID,
-    ));
+    msgs.push(CosmosMsg::Wasm(WasmMsg::Execute {
+        contract_addr: config.dex_v3.to_string(),
+        msg: to_json_binary(&V3ExecuteMsg::CreatePosition {
+            pool_key: pending_position.pool_key.clone(),
+            lower_tick: pending_position.lower_tick,
+            upper_tick: pending_position.upper_tick,
+            liquidity_delta: res.l,
+            slippage_limit_lower: pending_position.slippage_limit_lower.unwrap_or(
+                get_min_sqrt_price(pending_position.pool_key.fee_tier.tick_spacing),
+            ),
+            slippage_limit_upper: pending_position.slippage_limit_upper.unwrap_or(
+                get_max_sqrt_price(pending_position.pool_key.fee_tier.tick_spacing),
+            ),
+        })?,
+        funds: coins,
+    }));
 
-    Ok(Response::new().add_messages(msgs).add_submessages(sub_msgs))
+    msgs.push(CosmosMsg::Wasm(WasmMsg::Execute {
+        contract_addr: env.contract.address.to_string(),
+        msg: to_json_binary(&ExecuteMsg::RefundAfterZapInLiquidity {})?,
+        funds: vec![],
+    }));
+
+    Ok(Response::new().add_messages(msgs))
 }
 
-pub fn add_liquidity(deps: DepsMut, env: Env) -> Result<Response, ContractError> {
+pub fn refund_after_zap_in(deps: DepsMut, env: Env, info: MessageInfo) -> Result<Response, ContractError> {
+    if info.sender != env.contract.address {
+        return Err(ContractError::Unauthorized {});
+    }
+
     let mut msgs: Vec<CosmosMsg> = vec![];
 
     // 8. Refund unused tokenX and tokenY to user
@@ -191,7 +201,11 @@ pub fn add_liquidity(deps: DepsMut, env: Env) -> Result<Response, ContractError>
     Ok(Response::new().add_messages(msgs))
 }
 
-pub fn zap_out_liquidity(deps: DepsMut, env: Env) -> Result<Response, ContractError> {
+pub fn zap_out_liquidity(deps: DepsMut, env: Env, info: MessageInfo) -> Result<Response, ContractError> {
+    if info.sender != env.contract.address {
+        return Err(ContractError::Unauthorized {});
+    }
+
     let mut msgs: Vec<CosmosMsg> = vec![];
     let receiver = RECEIVER.load(deps.storage)?;
     let zap_out_routes = ZAP_OUT_ROUTES.load(deps.storage)?;

contracts/zapper/src/entrypoints/mod.rs

@@ -1,6 +1,6 @@
 pub mod execute;
 pub mod query;
-pub mod reply;
+pub mod internal;
 pub mod common;
 
 pub use execute::*;

contracts/zapper/src/msg.rs

@@ -27,14 +27,21 @@ pub enum ExecuteMsg {
         routes: Vec<Route>,
         minimum_liquidity: Option<Liquidity>,
     },
+    ZapInAfterSwapOperation {},
+    RefundAfterZapInLiquidity {},
     ZapOutLiquidity {
         position_index: u32,
         routes: Vec<Route>,
     },
+    ZapOutAfterSwapOperation {},
     RegisterProtocolFee {
         percent: Decimal,
         fee_receiver: Addr,
     },
+    Withdraw {
+        assets: Vec<Asset>,
+        recipient: Option<Addr>,
+    }
 }
 
 #[cw_serde]

contracts/zapper/src/tests/helper.rs

@@ -50,8 +50,7 @@ impl MockApp {
                     crate::contract::execute,
                     crate::contract::instantiate,
                     crate::contract::query,
-                )
-                .with_reply_empty(crate::contract::reply),
+                ),
             ));
         }
         #[cfg(feature = "test-tube")]
@@ -226,6 +225,24 @@ impl MockApp {
         )
     }
 
+    pub fn withdraw(
+        &mut self,
+        sender: &str,
+        zapper: &str,
+        assets: Vec<Asset>,
+        recipient: Option<&str>,
+    ) -> MockResult<ExecuteResponse> {
+        self.execute(
+            Addr::unchecked(sender),
+            Addr::unchecked(zapper),
+            &msg::ExecuteMsg::Withdraw {
+                assets,
+                recipient: recipient.map(Addr::unchecked),
+            },
+            &[],
+        )
+    }
+
     pub fn get_zapper_config(&mut self, zapper: &str) -> StdResult<Config> {
         self.query(Addr::unchecked(zapper), &msg::QueryMsg::Config {})
     }

contracts/zapper/src/tests/mod.rs

@@ -3,3 +3,4 @@ mod config;
 mod helper;
 mod zapin;
 mod zapout;
+mod withdraw;

contracts/zapper/src/tests/withdraw.rs

@@ -0,0 +1,44 @@
+use cosmwasm_std::{coins, Uint128};
+use decimal::*;
+use oraiswap::mixed_router::SwapOperation;
+use oraiswap_v3_common::asset::{Asset, AssetInfo};
+use oraiswap_v3_common::math::percentage::Percentage;
+
+use oraiswap_v3_common::storage::{FeeTier, PoolKey};
+
+use crate::msg::Route;
+use crate::tests::common::init_basic_v3_pool;
+use crate::tests::helper::MockApp;
+use crate::tests::helper::{macros::*, FEE_DENOM};
+
+#[test]
+fn test_withdraw() {
+    let (mut app, accounts) = MockApp::new(&[
+        ("alice", &coins(100_000_000_000, FEE_DENOM)),
+        ("bob", &coins(100_000_000_000, FEE_DENOM)),
+    ]);
+    let alice = &accounts[0];
+    let bob = &accounts[1];
+    let initial_amount = 10u128.pow(20);
+    let (token_x, _, _) =
+        create_3_tokens!(app, initial_amount, initial_amount, initial_amount, alice);
+
+    let zapper = create_zapper!(app, alice);
+
+    app.mint_token(alice, zapper.as_str(), token_x.as_str(), 1)
+        .unwrap();
+    let assets: Vec<Asset> = vec![Asset {
+        info: AssetInfo::Token {
+            contract_addr: token_x.clone(),
+        },
+        amount: Uint128::new(1),
+    }];
+    let err = app.withdraw(bob, zapper.as_str(), assets.clone(), Some(bob))
+        .unwrap_err();
+    assert!(err.root_cause().to_string().contains("Unauthorized"));
+
+    app.withdraw(alice, zapper.as_str(), assets, Some(bob))
+        .unwrap();
+    let balance = balance_of!(app, token_x, bob);
+    assert_eq!(balance, 1);
+}