Tidy

artifacts/checksums.txt

@@ -1,2 +1,2 @@
-021ddace1433405834abf2692662a677e3c2b233ff9f7e8c6a677e4c8538aaa9  dex_aggregator.wasm
+40bc067d21fb1a80e05b57d8f95cb10296dc1db322dddb56daf68d6f53ad64c3  dex_aggregator.wasm
 46dcce8409830a30b424dd342887f1cc203ef9c203f4b754b172c2203a7c0ad7  mock_swap.wasm

contracts/dex_aggregator/src/contract.rs

@@ -5,7 +5,7 @@ use injective_cosmwasm::{InjectiveMsgWrapper, InjectiveQueryWrapper};
 
 use crate::error::ContractError;
 use crate::execute::{self, remove_fee, set_fee, update_fee_collector};
-use crate::msg::{external, Cw20HookMsg, ExecuteMsg, InstantiateMsg, QueryMsg};
+use crate::msg::{amm, Cw20HookMsg, ExecuteMsg, InstantiateMsg, QueryMsg};
 use crate::state::{Config, CONFIG};
 use cw20::Cw20ReceiveMsg;
 
@@ -44,16 +44,16 @@ pub fn execute(
     msg: ExecuteMsg,
 ) -> Result<Response<InjectiveMsgWrapper>, ContractError> {
     match msg {
-        ExecuteMsg::AggregateSwaps {
+        ExecuteMsg::ExecuteRoute {
             stages,
             minimum_receive,
         } => {
             // This is the entry point for NATIVE token swaps
             if info.funds.len() != 1 {
                 return Err(ContractError::InvalidFunds {});
             }
-            let offer_asset = external::Asset {
-                info: external::AssetInfo::NativeToken {
+            let offer_asset = amm::Asset {
+                info: amm::AssetInfo::NativeToken {
                     denom: info.funds[0].denom.clone(),
                 },
                 amount: info.funds[0].amount,
@@ -76,12 +76,12 @@ pub fn execute(
             if let Ok(hook_msg) = cosmwasm_std::from_json::<Cw20HookMsg>(&msg) {
                 // This is a user-initiated swap starting with a CW20 token.
                 match hook_msg {
-                    Cw20HookMsg::AggregateSwaps {
+                    Cw20HookMsg::ExecuteRoute {
                         stages,
                         minimum_receive,
                     } => {
-                        let offer_asset = external::Asset {
-                            info: external::AssetInfo::Token {
+                        let offer_asset = amm::Asset {
+                            info: amm::AssetInfo::Token {
                                 contract_addr: info.sender.to_string(),
                             },
                             amount,
@@ -120,6 +120,9 @@ pub fn execute(
         ExecuteMsg::UpdateFeeCollector { new_fee_collector } => {
             update_fee_collector(deps, info, new_fee_collector)
         }
+        ExecuteMsg::EmergencyWithdraw { asset_info } => {
+            crate::execute::emergency_withdraw(deps, env, info, asset_info)
+        }
     }
 }
 
@@ -130,6 +133,12 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
             crate::query::simulate_route(deps, env, stages, amount_in)
         }
         QueryMsg::Config {} => crate::query::query_config(deps),
+        QueryMsg::FeeForPool { pool_address } => {
+            crate::query::query_fee_for_pool(deps, pool_address)
+        }
+        QueryMsg::AllFees { start_after, limit } => {
+            crate::query::query_all_fees(deps, start_after, limit)
+        }
     }
 }
 

contracts/dex_aggregator/src/error.rs

@@ -15,9 +15,6 @@ pub enum ContractError {
     #[error("Route cannot be empty")]
     EmptyRoute {},
 
-    #[error("Execution state not found for sender")]
-    ExecutionStateNotFound {},
-
     #[error("Input amount must be greater than zero")]
     ZeroAmount {},
 

contracts/dex_aggregator/src/execute.rs

@@ -1,16 +1,18 @@
 use cosmwasm_std::{
-    to_json_binary, Addr, Coin, CosmosMsg, Decimal, DepsMut, Env, MessageInfo, Response, StdError,
-    Uint128, WasmMsg,
+    to_json_binary, Addr, BankMsg, Coin, CosmosMsg, Decimal, DepsMut, Env, MessageInfo, Response,
+    StdError, Uint128, WasmMsg,
 };
-use cw20::Cw20ExecuteMsg;
+use cw20::{BalanceResponse, Cw20ExecuteMsg, Cw20QueryMsg};
 use injective_cosmwasm::{InjectiveMsgWrapper, InjectiveQueryWrapper};
 use injective_math::FPDecimal;
 use std::str::FromStr;
 
 use crate::error::ContractError;
-use crate::msg::{self, external, AmmPairExecuteMsg, Operation, OrderbookExecuteMsg, Stage};
+use crate::msg::{self, amm, orderbook, Operation, Stage};
 use crate::reply::proceed_to_next_step;
-use crate::state::{Awaiting, ReplyState, CONFIG, FEE_MAP, REPLY_ID_COUNTER};
+use crate::state::{
+    Awaiting, ExecutionState, RoutePlan, CONFIG, FEE_MAP, REPLY_ID_COUNTER, ROUTE_PLANS,
+};
 
 pub fn update_admin(
     deps: DepsMut<InjectiveQueryWrapper>,
@@ -40,7 +42,7 @@ pub fn execute_aggregate_swaps_internal(
     _info: MessageInfo,
     stages: Vec<Stage>,
     minimum_receive_str: Option<String>,
-    offer_asset: external::Asset,
+    offer_asset: amm::Asset,
     initiator: Addr,
 ) -> Result<Response<InjectiveMsgWrapper>, ContractError> {
     if offer_asset.amount.is_zero() {
@@ -64,35 +66,38 @@ pub fn execute_aggregate_swaps_internal(
         None => Uint128::zero(),
     };
 
-    let mut initial_state = ReplyState {
+    let plan = RoutePlan {
         sender: initiator.clone(),
         minimum_receive,
         stages,
+    };
+    ROUTE_PLANS.save(deps.storage, reply_id, &plan)?;
+
+    let mut initial_exec_state = ExecutionState {
         awaiting: Awaiting::Swaps,
         current_stage_index: 0,
         replies_expected: 0,
         accumulated_assets: vec![offer_asset],
         pending_swaps: vec![],
-        conversion_target_asset: None,
         pending_path_op: None,
     };
 
-    proceed_to_next_step(&mut deps, env, &mut initial_state, reply_id)
+    proceed_to_next_step(&mut deps, env, &mut initial_exec_state, &plan, reply_id)
 }
 
 pub fn create_swap_cosmos_msg(
     deps: &mut DepsMut<InjectiveQueryWrapper>,
     operation: &Operation,
-    offer_asset_info: &external::AssetInfo,
+    offer_asset_info: &amm::AssetInfo,
     amount: Uint128,
     env: &Env,
 ) -> Result<CosmosMsg<InjectiveMsgWrapper>, ContractError> {
     let recipient = env.contract.address.to_string();
 
     let cosmos_msg = match operation {
         Operation::AmmSwap(amm_op) => {
-            let amm_swap_msg = AmmPairExecuteMsg::Swap {
-                offer_asset: external::Asset {
+            let amm_swap_msg = amm::AmmPairExecuteMsg::Swap {
+                offer_asset: amm::Asset {
                     info: offer_asset_info.clone(),
                     amount,
                 },
@@ -103,15 +108,15 @@ pub fn create_swap_cosmos_msg(
             };
 
             match offer_asset_info {
-                external::AssetInfo::NativeToken { denom } => CosmosMsg::Wasm(WasmMsg::Execute {
+                amm::AssetInfo::NativeToken { denom } => CosmosMsg::Wasm(WasmMsg::Execute {
                     contract_addr: amm_op.pool_address.clone(),
                     msg: to_json_binary(&amm_swap_msg)?,
                     funds: vec![Coin {
                         denom: denom.clone(),
                         amount,
                     }],
                 }),
-                external::AssetInfo::Token { contract_addr } => {
+                amm::AssetInfo::Token { contract_addr } => {
                     let cw20_send_msg = Cw20ExecuteMsg::Send {
                         contract: amm_op.pool_address.clone(),
                         amount,
@@ -129,14 +134,14 @@ pub fn create_swap_cosmos_msg(
         Operation::OrderbookSwap(ob_op) => {
             let offer_denom =
                 match &ob_op.offer_asset_info {
-                    external::AssetInfo::NativeToken { denom } => denom.clone(),
+                    amm::AssetInfo::NativeToken { denom } => denom.clone(),
                     _ => return Err(ContractError::Std(StdError::generic_err(
                         "This OrderbookSwapOp implementation only supports native token inputs.",
                     ))),
                 };
 
             let target_denom = match &ob_op.ask_asset_info {
-                external::AssetInfo::NativeToken { denom } => denom.clone(),
+                amm::AssetInfo::NativeToken { denom } => denom.clone(),
                 _ => {
                     return Err(ContractError::Std(StdError::generic_err(
                         "Orderbook swaps only support native token (bank) outputs.",
@@ -155,14 +160,14 @@ pub fn create_swap_cosmos_msg(
             let expected_output_fp = simulation_response.result_quantity;
             let slippage = FPDecimal::from_str("0.005")?;
             let min_output_fp = expected_output_fp * (FPDecimal::ONE - slippage);
-            let swap_msg = OrderbookExecuteMsg::SwapMinOutput {
+            let swap_msg = orderbook::OrderbookExecuteMsg::SwapMinOutput {
                 target_denom,
                 min_output_quantity: min_output_fp,
             };
 
             let funds = vec![Coin {
                 denom: match &ob_op.offer_asset_info {
-                    external::AssetInfo::NativeToken { denom } => denom.clone(),
+                    amm::AssetInfo::NativeToken { denom } => denom.clone(),
                     _ => unreachable!(),
                 },
                 amount,
@@ -245,3 +250,72 @@ pub fn update_fee_collector(
         .add_attribute("action", "update_fee_collector")
         .add_attribute("new_fee_collector", new_collector_addr))
 }
+
+pub fn emergency_withdraw(
+    deps: DepsMut<InjectiveQueryWrapper>,
+    env: Env,
+    info: MessageInfo,
+    asset_info: amm::AssetInfo,
+) -> Result<Response<InjectiveMsgWrapper>, ContractError> {
+    // 1. Authorization Check
+    let config = CONFIG.load(deps.storage)?;
+    if info.sender != config.admin {
+        return Err(ContractError::Unauthorized {});
+    }
+
+    let (amount_to_withdraw, send_msg) = match asset_info.clone() {
+        amm::AssetInfo::NativeToken { denom } => {
+            // 2a. Query the contract's native token balance
+            let balance = deps.querier.query_balance(&env.contract.address, denom)?;
+
+            if balance.amount.is_zero() {
+                // Return success but do nothing if balance is zero
+                (balance.amount, None)
+            } else {
+                // 3a. Create a BankMsg to send the full balance to the admin
+                let msg = CosmosMsg::Bank(BankMsg::Send {
+                    to_address: info.sender.to_string(),
+                    amount: vec![balance.clone()],
+                });
+                (balance.amount, Some(msg))
+            }
+        }
+        amm::AssetInfo::Token { contract_addr } => {
+            // 2b. Query the contract's CW20 token balance
+            let balance_response: BalanceResponse = deps.querier.query_wasm_smart(
+                contract_addr.clone(),
+                &Cw20QueryMsg::Balance {
+                    address: env.contract.address.to_string(),
+                },
+            )?;
+
+            if balance_response.balance.is_zero() {
+                // Return success but do nothing if balance is zero
+                (balance_response.balance, None)
+            } else {
+                // 3b. Create a WasmMsg to transfer the full balance to the admin
+                let msg = CosmosMsg::Wasm(WasmMsg::Execute {
+                    contract_addr,
+                    msg: to_json_binary(&Cw20ExecuteMsg::Transfer {
+                        recipient: info.sender.to_string(),
+                        amount: balance_response.balance,
+                    })?,
+                    funds: vec![],
+                });
+                (balance_response.balance, Some(msg))
+            }
+        }
+    };
+
+    let mut response = Response::new()
+        .add_attribute("action", "emergency_withdraw")
+        .add_attribute("recipient", info.sender.to_string())
+        .add_attribute("asset", format!("{:?}", asset_info))
+        .add_attribute("withdrawn_amount", amount_to_withdraw.to_string());
+
+    if let Some(msg) = send_msg {
+        response = response.add_message(msg);
+    }
+
+    Ok(response)
+}