diff --git a/artifacts/checksums.txt b/artifacts/checksums.txt index 32eb809..277788a 100644 --- a/artifacts/checksums.txt +++ b/artifacts/checksums.txt @@ -1,2 +1,2 @@ -021ddace1433405834abf2692662a677e3c2b233ff9f7e8c6a677e4c8538aaa9 dex_aggregator.wasm +40bc067d21fb1a80e05b57d8f95cb10296dc1db322dddb56daf68d6f53ad64c3 dex_aggregator.wasm 46dcce8409830a30b424dd342887f1cc203ef9c203f4b754b172c2203a7c0ad7 mock_swap.wasm diff --git a/artifacts/dex_aggregator.wasm b/artifacts/dex_aggregator.wasm index f3a40b5..b331c03 100644 Binary files a/artifacts/dex_aggregator.wasm and b/artifacts/dex_aggregator.wasm differ diff --git a/contracts/dex_aggregator/src/contract.rs b/contracts/dex_aggregator/src/contract.rs index 19c2e3d..65cf9e1 100644 --- a/contracts/dex_aggregator/src/contract.rs +++ b/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,7 +44,7 @@ pub fn execute( msg: ExecuteMsg, ) -> Result, ContractError> { match msg { - ExecuteMsg::AggregateSwaps { + ExecuteMsg::ExecuteRoute { stages, minimum_receive, } => { @@ -52,8 +52,8 @@ pub fn execute( 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::(&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 { 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) + } } } diff --git a/contracts/dex_aggregator/src/error.rs b/contracts/dex_aggregator/src/error.rs index c761589..ae799fd 100644 --- a/contracts/dex_aggregator/src/error.rs +++ b/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 {}, diff --git a/contracts/dex_aggregator/src/execute.rs b/contracts/dex_aggregator/src/execute.rs index 63dd47d..4f96095 100644 --- a/contracts/dex_aggregator/src/execute.rs +++ b/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, @@ -40,7 +42,7 @@ pub fn execute_aggregate_swaps_internal( _info: MessageInfo, stages: Vec, minimum_receive_str: Option, - offer_asset: external::Asset, + offer_asset: amm::Asset, initiator: Addr, ) -> Result, ContractError> { if offer_asset.amount.is_zero() { @@ -64,26 +66,29 @@ 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, operation: &Operation, - offer_asset_info: &external::AssetInfo, + offer_asset_info: &amm::AssetInfo, amount: Uint128, env: &Env, ) -> Result, ContractError> { @@ -91,8 +96,8 @@ pub fn create_swap_cosmos_msg( 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,7 +108,7 @@ 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 { @@ -111,7 +116,7 @@ pub fn create_swap_cosmos_msg( 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, + env: Env, + info: MessageInfo, + asset_info: amm::AssetInfo, +) -> Result, 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) +} diff --git a/contracts/dex_aggregator/src/msg.rs b/contracts/dex_aggregator/src/msg.rs index a1719f0..b4c7d5f 100644 --- a/contracts/dex_aggregator/src/msg.rs +++ b/contracts/dex_aggregator/src/msg.rs @@ -3,12 +3,10 @@ use crate::state::Config; use cosmwasm_schema::{cw_serde, QueryResponses}; use cosmwasm_std::{Addr, Coin, Decimal, Uint128}; use cw20::Cw20ReceiveMsg; -use injective_math::FPDecimal; pub mod cw20_adapter { - use cosmwasm_std::Binary; - use super::*; + use cosmwasm_std::Binary; #[cw_serde] pub struct InstantiateMsg {} @@ -47,41 +45,15 @@ pub mod cw20_adapter { } } -pub mod external { - +pub mod amm { use super::*; - #[cw_serde] - pub enum SwapOperation { - Choice { - offer_asset_info: AssetInfo, - ask_asset_info: AssetInfo, - }, - DojoSwap { - offer_asset_info: AssetInfo, - ask_asset_info: AssetInfo, - }, - TerraSwap { - offer_asset_info: AssetInfo, - ask_asset_info: AssetInfo, - }, - AstroSwap { - offer_asset_info: AssetInfo, - ask_asset_info: AssetInfo, - }, - } - #[cw_serde] pub enum AssetInfo { Token { contract_addr: String }, NativeToken { denom: String }, } - #[cw_serde] - pub struct SimulateSwapOperationsResponse { - pub amount: Uint128, - } - #[cw_serde] pub struct Asset { pub info: AssetInfo, @@ -100,6 +72,17 @@ pub mod external { pub spread_amount: Uint128, pub commission_amount: Uint128, } + + #[cw_serde] + pub enum AmmPairExecuteMsg { + Swap { + offer_asset: amm::Asset, + belief_price: Option, + max_spread: Option, + to: Option, + deadline: Option, + }, + } } pub mod orderbook { @@ -126,20 +109,28 @@ pub mod orderbook { pub expected_fees: Vec, pub result_quantity: FPDecimal, } + + #[cw_serde] + pub enum OrderbookExecuteMsg { + SwapMinOutput { + target_denom: String, + min_output_quantity: FPDecimal, + }, + } } #[cw_serde] pub struct AmmSwapOp { pub pool_address: String, - pub offer_asset_info: external::AssetInfo, - pub ask_asset_info: external::AssetInfo, + pub offer_asset_info: amm::AssetInfo, + pub ask_asset_info: amm::AssetInfo, } #[cw_serde] pub struct OrderbookSwapOp { pub swap_contract: String, - pub offer_asset_info: external::AssetInfo, - pub ask_asset_info: external::AssetInfo, + pub offer_asset_info: amm::AssetInfo, + pub ask_asset_info: amm::AssetInfo, } #[cw_serde] @@ -167,12 +158,12 @@ pub struct PlannedSwap { pub struct StagePlan { pub swaps_to_execute: Vec, - pub conversions_needed: Vec<(external::Asset, external::AssetInfo)>, + pub conversions_needed: Vec<(amm::Asset, amm::AssetInfo)>, } #[cw_serde] pub enum Cw20HookMsg { - AggregateSwaps { + ExecuteRoute { stages: Vec, minimum_receive: Option, }, @@ -187,11 +178,12 @@ pub struct InstantiateMsg { #[cw_serde] pub enum ExecuteMsg { - AggregateSwaps { + ExecuteRoute { stages: Vec, minimum_receive: Option, }, Receive(Cw20ReceiveMsg), + // Admin-only UpdateAdmin { new_admin: String, }, @@ -199,14 +191,31 @@ pub enum ExecuteMsg { pool_address: String, fee_percent: Decimal, }, - /// Admin-only. Removes the fee configuration for a specific pool. RemoveFee { pool_address: String, }, - /// Admin-only. Updates the address where fees are sent. UpdateFeeCollector { new_fee_collector: String, }, + EmergencyWithdraw { + asset_info: amm::AssetInfo, + }, +} + +#[cw_serde] +pub struct FeeInfo { + pub pool_address: String, + pub fee_percent: Decimal, +} + +#[cw_serde] +pub struct FeeResponse { + pub fee: Option, +} + +#[cw_serde] +pub struct AllFeesResponse { + pub fees: Vec, } #[cw_serde] @@ -216,29 +225,16 @@ pub enum QueryMsg { SimulateRoute { stages: Vec, amount_in: Coin }, #[returns(Config)] Config {}, + #[returns(FeeResponse)] + FeeForPool { pool_address: String }, + #[returns(AllFeesResponse)] + AllFees { + start_after: Option, + limit: Option, + }, } #[cw_serde] pub struct SimulateRouteResponse { pub output_amount: Uint128, } - -#[cosmwasm_schema::cw_serde] -pub enum AmmPairExecuteMsg { - Swap { - offer_asset: external::Asset, - belief_price: Option, - max_spread: Option, - to: Option, - deadline: Option, - }, -} - -/// The ExecuteMsg format for the Orderbook swap contract. -#[cosmwasm_schema::cw_serde] -pub enum OrderbookExecuteMsg { - SwapMinOutput { - target_denom: String, - min_output_quantity: FPDecimal, - }, -} diff --git a/contracts/dex_aggregator/src/query.rs b/contracts/dex_aggregator/src/query.rs index 41d13ed..46a9190 100644 --- a/contracts/dex_aggregator/src/query.rs +++ b/contracts/dex_aggregator/src/query.rs @@ -1,9 +1,12 @@ -use crate::msg::{external, orderbook, Operation, SimulateRouteResponse, Stage}; -use crate::state::Config; +use crate::msg::{ + amm, orderbook, AllFeesResponse, FeeInfo, FeeResponse, Operation, SimulateRouteResponse, Stage, +}; +use crate::state::{Config, FEE_MAP}; use cosmwasm_std::{ - to_json_binary, Binary, Coin, Deps, Env, QuerierWrapper, StdError, StdResult, Uint128, + to_json_binary, Binary, Coin, Deps, Env, Order, QuerierWrapper, StdError, StdResult, Uint128, WasmQuery, }; +use cw_storage_plus::Bound; pub fn query_config(deps: Deps) -> StdResult { let config: Config = crate::state::CONFIG.load(deps.storage)?; @@ -22,18 +25,18 @@ pub fn simulate_route( }); } - let mut current_assets: Vec = vec![external::Asset { - info: external::AssetInfo::NativeToken { + let mut current_assets: Vec = vec![amm::Asset { + info: amm::AssetInfo::NativeToken { denom: amount_in.denom, }, amount: amount_in.amount, }]; for stage in stages { - let mut next_stage_outputs: Vec = vec![]; + let mut next_stage_outputs: Vec = vec![]; // Group the current assets by their type to get the total for each pile. - let mut grouped_inputs: Vec<(external::AssetInfo, Uint128)> = vec![]; + let mut grouped_inputs: Vec<(amm::AssetInfo, Uint128)> = vec![]; for asset in current_assets { if let Some((_, amount)) = grouped_inputs .iter_mut() @@ -45,7 +48,7 @@ pub fn simulate_route( } } - let mut amounts_allocated: Vec<(external::AssetInfo, Uint128)> = vec![]; + let mut amounts_allocated: Vec<(amm::AssetInfo, Uint128)> = vec![]; for (i, split) in stage.splits.iter().enumerate() { let path_input_info = get_path_start_info(&split.path)?; @@ -78,7 +81,7 @@ pub fn simulate_route( amounts_allocated.push((path_input_info.clone(), amount_for_split)); } - let mut current_path_asset = external::Asset { + let mut current_path_asset = amm::Asset { info: path_input_info, amount: amount_for_split, }; @@ -107,16 +110,16 @@ pub fn simulate_route( fn simulate_single_operation( querier: &QuerierWrapper, operation: &Operation, - offer_asset: &external::Asset, -) -> StdResult { + offer_asset: &amm::Asset, +) -> StdResult { match operation { Operation::AmmSwap(op) => { - let pair_query = external::QueryMsg::Simulation { + let pair_query = amm::QueryMsg::Simulation { offer_asset: offer_asset.clone(), }; let contract_addr = op.pool_address.to_string(); - let sim_response: external::SimulationResponse = querier.query( + let sim_response: amm::SimulationResponse = querier.query( &WasmQuery::Smart { contract_addr, msg: to_json_binary(&pair_query)?, @@ -124,14 +127,14 @@ fn simulate_single_operation( .into(), )?; - Ok(external::Asset { + Ok(amm::Asset { info: op.ask_asset_info.clone(), amount: sim_response.return_amount, }) } Operation::OrderbookSwap(op) => { let source_denom = match &offer_asset.info { - external::AssetInfo::NativeToken { denom } => denom.clone(), + amm::AssetInfo::NativeToken { denom } => denom.clone(), _ => { return Err(StdError::generic_err( "Orderbook simulation only supports native token inputs", @@ -139,7 +142,7 @@ fn simulate_single_operation( } }; let target_denom = match &op.ask_asset_info { - external::AssetInfo::NativeToken { denom } => denom.clone(), + amm::AssetInfo::NativeToken { denom } => denom.clone(), _ => { return Err(StdError::generic_err( "Orderbook simulation only supports native token outputs", @@ -162,7 +165,7 @@ fn simulate_single_operation( .into(), )?; - Ok(external::Asset { + Ok(amm::Asset { info: op.ask_asset_info.clone(), amount: sim_response.result_quantity.into(), }) @@ -170,7 +173,7 @@ fn simulate_single_operation( } } -fn get_path_start_info(path: &[Operation]) -> StdResult { +fn get_path_start_info(path: &[Operation]) -> StdResult { let first_op = path .first() .ok_or_else(|| StdError::generic_err("Path cannot be empty"))?; @@ -180,21 +183,68 @@ fn get_path_start_info(path: &[Operation]) -> StdResult { }) } +/// Queries the fee percentage for a specific pool address. +pub fn query_fee_for_pool(deps: Deps, pool_address: String) -> StdResult { + let pool_addr = deps.api.addr_validate(&pool_address)?; + let fee = FEE_MAP.may_load(deps.storage, &pool_addr)?; + + to_json_binary(&FeeResponse { fee }) +} + +// Pagination constants +const DEFAULT_LIMIT: u32 = 10; +const MAX_LIMIT: u32 = 30; + +/// Queries all configured fees with pagination. +pub fn query_all_fees( + deps: Deps, + start_after: Option, + limit: Option, +) -> StdResult { + let limit = limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT) as usize; + + // Validate the start_after address if provided + let start = start_after + .map(|addr| deps.api.addr_validate(&addr)) + .transpose()?; + + let fees: Vec = FEE_MAP + .range( + deps.storage, + start.as_ref().map(Bound::exclusive), // Use exclusive bound for start_after + None, + Order::Ascending, + ) + .take(limit) + .map(|item| { + let (pool_addr, fee_percent) = item?; + Ok(FeeInfo { + pool_address: pool_addr.to_string(), + fee_percent, + }) + }) + .collect::>()?; + + to_json_binary(&AllFeesResponse { fees }) +} + #[cfg(test)] mod tests { use super::*; - use crate::msg::{AmmSwapOp, Split, Stage}; - use cosmwasm_std::testing::{mock_dependencies, mock_env, MockQuerier}; - use cosmwasm_std::{from_json, ContractResult, SystemResult}; - use external::AssetInfo; + use crate::contract::query; + use crate::msg::{AmmSwapOp, QueryMsg, Split, Stage}; + use amm::AssetInfo; + use cosmwasm_std::testing::{mock_dependencies, mock_env, MockApi, MockQuerier}; + use cosmwasm_std::{from_json, ContractResult, Decimal, SystemResult}; + use std::str::FromStr; - const FAKE_POOL_A: &str = "inj1pool_a"; - const FAKE_POOL_B: &str = "inj1pool_b"; + const POOL_A_ADDR: &str = "inj1hkhdaj2ts42k2x53h3w0f26g2xvy3a52e0u4gp"; + const POOL_B_ADDR: &str = "inj12sqy2n5qt52n5q2n5qt52n5q2n5qt52n5q2n5qt"; #[test] fn test_simulate_simple_path() { let mut querier = MockQuerier::new(&[]); - let mock_response = external::SimulationResponse { + let mock_response = amm::SimulationResponse { return_amount: Uint128::new(50000), spread_amount: Uint128::zero(), commission_amount: Uint128::zero(), @@ -208,7 +258,7 @@ mod tests { contract_addr, msg: _, } => { - if contract_addr == FAKE_POOL_A { + if contract_addr == POOL_A_ADDR { SystemResult::Ok(ContractResult::Ok(mock_response_binary.clone())) } else { panic!("Unexpected contract call to {}", contract_addr); @@ -225,7 +275,7 @@ mod tests { splits: vec![Split { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { - pool_address: FAKE_POOL_A.to_string(), + pool_address: POOL_A_ADDR.to_string(), offer_asset_info: AssetInfo::NativeToken { denom: "inj".to_string(), }, @@ -251,12 +301,12 @@ mod tests { fn test_simulate_multi_hop_path() { let mut querier = MockQuerier::new(&[]); - let mock_response_hop1 = external::SimulationResponse { + let mock_response_hop1 = amm::SimulationResponse { return_amount: Uint128::new(20000), // 1000 INJ -> 20000 USDT spread_amount: Uint128::zero(), commission_amount: Uint128::zero(), }; - let mock_response_hop2 = external::SimulationResponse { + let mock_response_hop2 = amm::SimulationResponse { return_amount: Uint128::new(5000), // 20000 USDT -> 5000 AUSD spread_amount: Uint128::zero(), commission_amount: Uint128::zero(), @@ -266,15 +316,15 @@ mod tests { WasmQuery::Smart { contract_addr, msg, .. } => { - let decoded: external::QueryMsg = from_json(msg).unwrap(); - if contract_addr == FAKE_POOL_A { - let external::QueryMsg::Simulation { offer_asset } = decoded; + let decoded: amm::QueryMsg = from_json(msg).unwrap(); + if contract_addr == POOL_A_ADDR { + let amm::QueryMsg::Simulation { offer_asset } = decoded; assert_eq!(offer_asset.amount, Uint128::new(1000)); SystemResult::Ok(ContractResult::Ok( to_json_binary(&mock_response_hop1).unwrap(), )) - } else if contract_addr == FAKE_POOL_B { - let external::QueryMsg::Simulation { offer_asset } = decoded; + } else if contract_addr == POOL_B_ADDR { + let amm::QueryMsg::Simulation { offer_asset } = decoded; assert_eq!(offer_asset.amount, Uint128::new(20000)); SystemResult::Ok(ContractResult::Ok( to_json_binary(&mock_response_hop2).unwrap(), @@ -294,7 +344,7 @@ mod tests { percent: 100, path: vec![ Operation::AmmSwap(AmmSwapOp { - pool_address: FAKE_POOL_A.to_string(), + pool_address: POOL_A_ADDR.to_string(), offer_asset_info: AssetInfo::NativeToken { denom: "inj".to_string(), }, @@ -303,7 +353,7 @@ mod tests { }, }), Operation::AmmSwap(AmmSwapOp { - pool_address: FAKE_POOL_B.to_string(), + pool_address: POOL_B_ADDR.to_string(), offer_asset_info: AssetInfo::NativeToken { denom: "usdt".to_string(), }, @@ -335,22 +385,22 @@ mod tests { WasmQuery::Smart { contract_addr, msg, .. } => { - let decoded: external::QueryMsg = from_json(msg).unwrap(); - let external::QueryMsg::Simulation { offer_asset } = decoded; + let decoded: amm::QueryMsg = from_json(msg).unwrap(); + let amm::QueryMsg::Simulation { offer_asset } = decoded; let response_amount = match (contract_addr.as_str(), offer_asset.amount.u128()) { // Stage 1 - (FAKE_POOL_A, 500) => 10000, // 50% of 1000 INJ -> 10000 USDT - (FAKE_POOL_B, 500) => 20000, // 50% of 1000 INJ -> 20000 AUSD + (POOL_A_ADDR, 500) => 10000, // 50% of 1000 INJ -> 10000 USDT + (POOL_B_ADDR, 500) => 20000, // 50% of 1000 INJ -> 20000 AUSD // Stage 2 (Totals: 10k USDT, 20k AUSD) - (FAKE_POOL_A, 10000) => 5000, // 10000 USDT -> 5000 SHROOM - (FAKE_POOL_B, 20000) => 8000, // 20000 AUSD -> 8000 SHROOM + (POOL_A_ADDR, 10000) => 5000, // 10000 USDT -> 5000 SHROOM + (POOL_B_ADDR, 20000) => 8000, // 20000 AUSD -> 8000 SHROOM _ => panic!( "Unexpected query: {} with amount {}", contract_addr, offer_asset.amount ), }; - let mock_response = external::SimulationResponse { + let mock_response = amm::SimulationResponse { return_amount: Uint128::new(response_amount), ..Default::default() }; @@ -369,7 +419,7 @@ mod tests { Split { percent: 50, path: vec![Operation::AmmSwap(AmmSwapOp { - pool_address: FAKE_POOL_A.to_string(), + pool_address: POOL_A_ADDR.to_string(), offer_asset_info: AssetInfo::NativeToken { denom: "inj".to_string(), }, @@ -381,7 +431,7 @@ mod tests { Split { percent: 50, path: vec![Operation::AmmSwap(AmmSwapOp { - pool_address: FAKE_POOL_B.to_string(), + pool_address: POOL_B_ADDR.to_string(), offer_asset_info: AssetInfo::NativeToken { denom: "inj".to_string(), }, @@ -398,7 +448,7 @@ mod tests { Split { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { - pool_address: FAKE_POOL_A.to_string(), + pool_address: POOL_A_ADDR.to_string(), offer_asset_info: AssetInfo::NativeToken { denom: "usdt".to_string(), }, @@ -410,7 +460,7 @@ mod tests { Split { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { - pool_address: FAKE_POOL_B.to_string(), + pool_address: POOL_B_ADDR.to_string(), offer_asset_info: AssetInfo::NativeToken { denom: "ausd".to_string(), }, @@ -434,4 +484,93 @@ mod tests { // Final output is the sum of the shroom from both paths assert_eq!(result.output_amount, Uint128::new(5000 + 8000)); } + + #[test] + fn test_query_fee_for_pool() { + // --- Setup using the proven litmus test pattern --- + let mut deps = mock_dependencies(); + deps.api = MockApi::default().with_prefix("inj"); + + // Use the API to generate valid addresses for the test + let pool_a_addr = deps.api.addr_make("pool_a"); + let pool_c_addr = deps.api.addr_make("pool_c_no_fee"); + + // Arrange: Set up the state needed for this test + let fee_a = Decimal::from_str("0.003").unwrap(); + FEE_MAP + .save(deps.as_mut().storage, &pool_a_addr, &fee_a) + .unwrap(); + + // --- Act & Assert: Test Case 1 (Fee exists) --- + let msg = QueryMsg::FeeForPool { + pool_address: pool_a_addr.to_string(), + }; + let res_binary = query(deps.as_ref(), mock_env(), msg).unwrap(); + let res: FeeResponse = from_json(&res_binary).unwrap(); + assert_eq!(res.fee, Some(Decimal::from_str("0.003").unwrap())); + + // --- Act & Assert: Test Case 2 (Fee does not exist) --- + let msg = QueryMsg::FeeForPool { + pool_address: pool_c_addr.to_string(), + }; + let res_binary = query(deps.as_ref(), mock_env(), msg).unwrap(); + let res: FeeResponse = from_json(&res_binary).unwrap(); + assert_eq!(res.fee, None); + } + + #[test] + fn test_query_all_fees_with_pagination() { + // --- Setup using the proven litmus test pattern --- + let mut deps = mock_dependencies(); + deps.api = MockApi::default().with_prefix("inj"); + + // Use the API to generate valid addresses for the test. + // We will generate them in a way that we can predict the alphabetical order. + let pool_addr_1 = deps.api.addr_make("pool_alpha"); // Starts with "inj1..." + let pool_addr_2 = deps.api.addr_make("pool_zulu"); // Starts with a different "inj1..." + + // To make the test robust, we must determine the correct order programmatically. + let (first_addr, second_addr) = if pool_addr_1 < pool_addr_2 { + (pool_addr_1.clone(), pool_addr_2.clone()) + } else { + (pool_addr_2.clone(), pool_addr_1.clone()) + }; + + // Arrange: Set up the state + let fee_1 = Decimal::from_str("0.003").unwrap(); + let fee_2 = Decimal::from_str("0.015").unwrap(); + FEE_MAP + .save(deps.as_mut().storage, &pool_addr_1, &fee_1) + .unwrap(); + FEE_MAP + .save(deps.as_mut().storage, &pool_addr_2, &fee_2) + .unwrap(); + + // --- Act & Assert: Query all and check the determined order --- + let msg = QueryMsg::AllFees { + start_after: None, + limit: None, + }; + let res_binary = query(deps.as_ref(), mock_env(), msg).unwrap(); + let res: AllFeesResponse = from_json(&res_binary).unwrap(); + + assert_eq!(res.fees.len(), 2); + assert_eq!(res.fees[0].pool_address, first_addr.to_string()); + assert_eq!(res.fees[1].pool_address, second_addr.to_string()); + } + + #[test] + fn test_query_all_fees_empty() { + // --- Setup using the proven litmus test pattern --- + let mut deps = mock_dependencies(); + deps.api = MockApi::default().with_prefix("inj"); + + let msg = QueryMsg::AllFees { + start_after: None, + limit: None, + }; + let res_binary = query(deps.as_ref(), mock_env(), msg).unwrap(); + let res: AllFeesResponse = from_json(&res_binary).unwrap(); + assert_eq!(res.fees.len(), 0); + } } diff --git a/contracts/dex_aggregator/src/reply.rs b/contracts/dex_aggregator/src/reply.rs index d015e49..4947809 100644 --- a/contracts/dex_aggregator/src/reply.rs +++ b/contracts/dex_aggregator/src/reply.rs @@ -1,7 +1,10 @@ use crate::error::ContractError; use crate::execute::create_swap_cosmos_msg; -use crate::msg::{cw20_adapter, external, Operation, PlannedSwap, Stage, StagePlan}; -use crate::state::{Awaiting, Config, PendingPathOp, ReplyState, CONFIG, FEE_MAP, REPLY_STATES}; +use crate::msg::{amm, cw20_adapter, Operation, PlannedSwap, Stage, StagePlan}; +use crate::state::{ + Awaiting, Config, ExecutionState, PendingPathOp, RoutePlan, CONFIG, EXECUTION_STATES, FEE_MAP, + ROUTE_PLANS, +}; use cosmwasm_std::{ to_json_binary, Addr, Coin, CosmosMsg, DepsMut, Env, Reply, Response, StdError, SubMsg, Uint128, WasmMsg, @@ -14,45 +17,64 @@ pub fn handle_reply( env: Env, msg: Reply, ) -> Result, ContractError> { - let mut state = REPLY_STATES.load(deps.storage, msg.id)?; - match state.awaiting { - Awaiting::Swaps => handle_swap_reply(deps, env, msg, &mut state), - Awaiting::Conversions => handle_conversion_reply(deps, env, msg, &mut state), - Awaiting::FinalConversions => handle_final_conversion_reply(deps, env, msg, &mut state), - Awaiting::PathConversion => handle_path_conversion_reply(deps, env, msg, &mut state), + let reply_id = msg.id; + let mut exec_state = EXECUTION_STATES.load(deps.storage, reply_id)?; + let plan = ROUTE_PLANS.load(deps.storage, reply_id)?; + + match exec_state.awaiting { + Awaiting::Swaps => handle_swap_reply(deps, env, msg, &mut exec_state, &plan), + Awaiting::Conversions => handle_conversion_reply(deps, env, msg, &mut exec_state, &plan), + Awaiting::FinalConversions => { + handle_final_conversion_reply(deps, env, msg, &mut exec_state, &plan) + } + Awaiting::PathConversion => { + handle_path_conversion_reply(deps, env, msg, &mut exec_state, &plan) + } } } pub(crate) fn proceed_to_next_step( deps: &mut DepsMut, env: Env, - state: &mut ReplyState, + exec_state: &mut ExecutionState, + plan: &RoutePlan, master_reply_id: u64, ) -> Result, ContractError> { - if state.current_stage_index as usize >= state.stages.len() { - return handle_final_stage(deps, env, master_reply_id, state); + if exec_state.current_stage_index as usize >= plan.stages.len() { + return handle_final_stage(deps, env, master_reply_id, exec_state, plan); } - let next_stage_to_execute = state + let next_stage_to_execute = plan .stages - .get(state.current_stage_index as usize) + .get(exec_state.current_stage_index as usize) .unwrap(); - let plan = plan_next_stage(&state.accumulated_assets, next_stage_to_execute)?; - state.accumulated_assets.clear(); - if plan.conversions_needed.is_empty() { - execute_planned_swaps(deps, env, state, master_reply_id, plan.swaps_to_execute) + let stage_plan = plan_next_stage(&exec_state.accumulated_assets, next_stage_to_execute)?; + exec_state.accumulated_assets.clear(); + + if stage_plan.conversions_needed.is_empty() { + execute_planned_swaps( + deps, + env, + exec_state, + plan, + master_reply_id, + stage_plan.swaps_to_execute, + ) } else { let config = CONFIG.load(deps.storage)?; let mut conversion_submsgs = vec![]; - for (asset_to_convert, _target_info) in &plan.conversions_needed { + for (asset_to_convert, _target_info) in &stage_plan.conversions_needed { let msg = create_conversion_msg(asset_to_convert, &config, &env)?; conversion_submsgs.push(SubMsg::reply_on_success(msg, master_reply_id)); } - state.awaiting = Awaiting::Conversions; - state.replies_expected = conversion_submsgs.len() as u64; - state.pending_swaps = plan.swaps_to_execute; - REPLY_STATES.save(deps.storage, master_reply_id, state)?; + + exec_state.awaiting = Awaiting::Conversions; + exec_state.replies_expected = conversion_submsgs.len() as u64; + exec_state.pending_swaps = stage_plan.swaps_to_execute; + + EXECUTION_STATES.save(deps.storage, master_reply_id, exec_state)?; + Ok(Response::new() .add_submessages(conversion_submsgs) .add_attribute("action", "performing_minimal_conversions")) @@ -63,7 +85,8 @@ fn handle_swap_reply( mut deps: DepsMut, env: Env, msg: Reply, - state: &mut ReplyState, + exec_state: &mut ExecutionState, + plan: &RoutePlan, ) -> Result, ContractError> { let master_reply_id = msg.id; @@ -84,16 +107,17 @@ fn handle_swap_reply( // If there is no swap event, we assume the output was zero. // In this case, the path cannot continue, so we treat it as "complete" with a zero value. if swap_event_opt.is_none() { - state.replies_expected -= 1; + exec_state.replies_expected -= 1; // Mutate exec_state - let response = if state.replies_expected > 0 { - REPLY_STATES.save(deps.storage, master_reply_id, state)?; + let response = if exec_state.replies_expected > 0 { + EXECUTION_STATES.save(deps.storage, master_reply_id, exec_state)?; // Save exec_state Response::new() .add_attribute("action", "accumulating_path_outputs") .add_attribute("info", "zero_value_path_completed") } else { - state.current_stage_index += 1; - proceed_to_next_step(&mut deps, env, state, master_reply_id)? + exec_state.current_stage_index += 1; // Mutate exec_state + // Call proceeds with exec_state and plan + proceed_to_next_step(&mut deps, env, exec_state, plan, master_reply_id)? }; return Ok(response); } @@ -110,9 +134,9 @@ fn handle_swap_reply( let replying_pool_addr = deps.api.addr_validate(&replying_pool_addr_str)?; - let current_stage = state + let current_stage = plan .stages - .get(state.current_stage_index as usize) + .get(exec_state.current_stage_index as usize) .ok_or(ContractError::EmptyRoute {})?; // Now, find the operation that matches this validated address. @@ -141,20 +165,21 @@ fn handle_swap_reply( if let Some(next_op) = replied_path.get(op_index + 1) { let required_input_info = get_operation_input(next_op)?; - let offer_asset_for_next_op = external::Asset { + let offer_asset_for_next_op = amm::Asset { info: received_asset_info, amount: received_amount, }; if offer_asset_for_next_op.info != required_input_info { - state.awaiting = Awaiting::PathConversion; - state.pending_path_op = Some(PendingPathOp { + exec_state.awaiting = Awaiting::PathConversion; // Mutate exec_state + exec_state.pending_path_op = Some(PendingPathOp { + // Mutate exec_state operation: next_op.clone(), amount: received_amount, }); let config = CONFIG.load(deps.storage)?; let conversion_msg = create_conversion_msg(&offer_asset_for_next_op, &config, &env)?; let sub_msg = SubMsg::reply_on_success(conversion_msg, master_reply_id); - REPLY_STATES.save(deps.storage, master_reply_id, state)?; + EXECUTION_STATES.save(deps.storage, master_reply_id, exec_state)?; // Save exec_state return Ok(Response::new() .add_submessage(sub_msg) .add_attribute("action", "performing_path_conversion")); @@ -167,7 +192,7 @@ fn handle_swap_reply( &env, )?; let sub_msg = SubMsg::reply_on_success(next_msg, master_reply_id); - REPLY_STATES.save(deps.storage, master_reply_id, state)?; + EXECUTION_STATES.save(deps.storage, master_reply_id, exec_state)?; // Save exec_state Ok(Response::new() .add_submessage(sub_msg) .add_attribute("action", "proceeding_to_next_op_in_path") @@ -183,18 +208,19 @@ fn handle_swap_reply( None => Uint128::zero(), }; let amount_after_fee = received_amount.checked_sub(fee).map_err(StdError::from)?; - state.accumulated_assets.push(external::Asset { + exec_state.accumulated_assets.push(amm::Asset { + // Mutate exec_state info: received_asset_info.clone(), amount: amount_after_fee, }); - state.replies_expected -= 1; + exec_state.replies_expected -= 1; let mut response; - if state.replies_expected > 0 { - REPLY_STATES.save(deps.storage, master_reply_id, state)?; + if exec_state.replies_expected > 0 { + EXECUTION_STATES.save(deps.storage, master_reply_id, exec_state)?; // Save exec_state response = Response::new().add_attribute("action", "accumulating_path_outputs"); } else { - state.current_stage_index += 1; - response = proceed_to_next_step(&mut deps, env, state, master_reply_id)?; + exec_state.current_stage_index += 1; // Mutate exec_state + response = proceed_to_next_step(&mut deps, env, exec_state, plan, master_reply_id)?; } if !fee.is_zero() { let config = CONFIG.load(deps.storage)?; @@ -211,20 +237,18 @@ fn handle_swap_reply( // A helper to create the final transfer message. fn create_send_msg( recipient: &Addr, - asset_info: &external::AssetInfo, + asset_info: &amm::AssetInfo, amount: Uint128, ) -> Result, ContractError> { match asset_info { - external::AssetInfo::NativeToken { denom } => { - Ok(CosmosMsg::Bank(cosmwasm_std::BankMsg::Send { - to_address: recipient.to_string(), - amount: vec![Coin { - denom: denom.clone(), - amount, - }], - })) - } - external::AssetInfo::Token { contract_addr } => Ok(CosmosMsg::Wasm(WasmMsg::Execute { + amm::AssetInfo::NativeToken { denom } => Ok(CosmosMsg::Bank(cosmwasm_std::BankMsg::Send { + to_address: recipient.to_string(), + amount: vec![Coin { + denom: denom.clone(), + amount, + }], + })), + amm::AssetInfo::Token { contract_addr } => Ok(CosmosMsg::Wasm(WasmMsg::Execute { contract_addr: contract_addr.clone(), msg: to_json_binary(&Cw20ExecuteMsg::Transfer { recipient: recipient.to_string(), @@ -239,29 +263,30 @@ fn handle_final_stage( deps: &mut DepsMut, env: Env, reply_id: u64, - state: &mut ReplyState, + exec_state: &mut ExecutionState, + plan: &RoutePlan, ) -> Result, ContractError> { - // If there are no assets, we're done. - if state.accumulated_assets.is_empty() { - if !state.minimum_receive.is_zero() { + if exec_state.accumulated_assets.is_empty() { + if !plan.minimum_receive.is_zero() { return Err(ContractError::MinimumReceiveNotMet {}); } - REPLY_STATES.remove(deps.storage, reply_id); + // CLEANUP HERE + EXECUTION_STATES.remove(deps.storage, reply_id); + ROUTE_PLANS.remove(deps.storage, reply_id); return Ok(Response::new().add_attribute("action", "aggregate_swap_complete_empty")); } // The target asset for normalization is the type of the first asset in the final list. - let target_asset_info = state.accumulated_assets[0].info.clone(); + let target_asset_info = exec_state.accumulated_assets[0].info.clone(); let mut conversion_submsgs = vec![]; let mut ready_amount = Uint128::zero(); let config = CONFIG.load(deps.storage)?; - for asset in &state.accumulated_assets { + for asset in &exec_state.accumulated_assets { if asset.info == target_asset_info { ready_amount += asset.amount; } else { - // This asset needs to be converted to the target type. let msg = create_conversion_msg(asset, &config, &env)?; conversion_submsgs.push(SubMsg::reply_on_success(msg, reply_id)); } @@ -270,31 +295,37 @@ fn handle_final_stage( if conversion_submsgs.is_empty() { // SCENARIO A: All assets were already the same type. We are done. let total_final_amount = ready_amount; - if total_final_amount < state.minimum_receive { + // Check against minimum_receive from the immutable plan + if total_final_amount < plan.minimum_receive { return Err(ContractError::MinimumReceiveNotMet {}); } let mut response = Response::new(); - // Only create and add the send message if there is a non-zero amount to send. if !total_final_amount.is_zero() { - let send_msg = create_send_msg(&state.sender, &target_asset_info, total_final_amount)?; + // Use the sender address from the immutable plan + let send_msg = create_send_msg(&plan.sender, &target_asset_info, total_final_amount)?; response = response.add_message(send_msg); } - REPLY_STATES.remove(deps.storage, reply_id); + EXECUTION_STATES.remove(deps.storage, reply_id); + ROUTE_PLANS.remove(deps.storage, reply_id); + + // State cleanup is now handled in the main `handle_reply` function Ok(response .add_attribute("action", "aggregate_swap_complete") .add_attribute("final_received", total_final_amount.to_string())) } else { - // SCENARIO B: Conversions are needed. Set up the state for the final reply. - state.awaiting = Awaiting::FinalConversions; - state.replies_expected = conversion_submsgs.len() as u64; - state.accumulated_assets = vec![external::Asset { + // SCENARIO B: Conversions are needed. Set up the exec_state for the final reply. + exec_state.awaiting = Awaiting::FinalConversions; + exec_state.replies_expected = conversion_submsgs.len() as u64; + exec_state.accumulated_assets = vec![amm::Asset { info: target_asset_info, amount: ready_amount, }]; - REPLY_STATES.save(deps.storage, reply_id, state)?; + // Save the small, mutated exec_state + EXECUTION_STATES.save(deps.storage, reply_id, exec_state)?; + Ok(Response::new() .add_submessages(conversion_submsgs) .add_attribute("action", "final_asset_normalization_started")) @@ -305,21 +336,22 @@ fn handle_final_conversion_reply( deps: DepsMut, env: Env, msg: Reply, - state: &mut ReplyState, + exec_state: &mut ExecutionState, + plan: &RoutePlan, ) -> Result, ContractError> { let reply_id = msg.id; let converted_amount = parse_amount_from_conversion_reply(&msg, &env)?; - let running_total_asset = state.accumulated_assets.get_mut(0).ok_or_else(|| { + let running_total_asset = exec_state.accumulated_assets.get_mut(0).ok_or_else(|| { StdError::generic_err("Final conversion state is invalid: no accumulated asset found") })?; running_total_asset.amount += converted_amount; - state.replies_expected -= 1; + exec_state.replies_expected -= 1; - if state.replies_expected > 0 { - // Still waiting for more conversions to finish. - REPLY_STATES.save(deps.storage, reply_id, state)?; + if exec_state.replies_expected > 0 { + // Still waiting for more conversions to finish. Save the updated exec_state. + EXECUTION_STATES.save(deps.storage, reply_id, exec_state)?; return Ok(Response::new().add_attribute("action", "accumulating_final_conversions")); } @@ -327,18 +359,21 @@ fn handle_final_conversion_reply( let total_final_amount = running_total_asset.amount; let final_asset_info = running_total_asset.info.clone(); - if total_final_amount < state.minimum_receive { + if total_final_amount < plan.minimum_receive { return Err(ContractError::MinimumReceiveNotMet {}); } let mut response = Response::new(); - // Only create and add the send message if there is a non-zero amount to send. if !total_final_amount.is_zero() { - let send_msg = create_send_msg(&state.sender, &final_asset_info, total_final_amount)?; + // Get the sender address from the immutable plan + let send_msg = create_send_msg(&plan.sender, &final_asset_info, total_final_amount)?; response = response.add_message(send_msg); } - REPLY_STATES.remove(deps.storage, reply_id); + EXECUTION_STATES.remove(deps.storage, reply_id); + ROUTE_PLANS.remove(deps.storage, reply_id); + + // State cleanup is now handled in the main `handle_reply` function Ok(response .add_attribute("action", "aggregate_swap_complete") .add_attribute("final_received", total_final_amount.to_string())) @@ -348,28 +383,40 @@ fn handle_conversion_reply( mut deps: DepsMut, env: Env, msg: Reply, - state: &mut ReplyState, + exec_state: &mut ExecutionState, + plan: &RoutePlan, ) -> Result, ContractError> { let master_reply_id = msg.id; - state.replies_expected -= 1; + exec_state.replies_expected -= 1; // Mutate exec_state - if state.replies_expected > 0 { - REPLY_STATES.save(deps.storage, master_reply_id, state)?; + if exec_state.replies_expected > 0 { + // Save the small, mutated exec_state + EXECUTION_STATES.save(deps.storage, master_reply_id, exec_state)?; return Ok(Response::new().add_attribute("action", "accumulating_conversion_outputs")); } - let swaps_to_execute = std::mem::take(&mut state.pending_swaps); - execute_planned_swaps(&mut deps, env, state, master_reply_id, swaps_to_execute) + // Take pending_swaps from the mutated exec_state + let swaps_to_execute = std::mem::take(&mut exec_state.pending_swaps); + + // Call the updated execute_planned_swaps with both state objects + execute_planned_swaps( + &mut deps, + env, + exec_state, + plan, + master_reply_id, + swaps_to_execute, + ) } fn create_conversion_msg( - from: &external::Asset, + from: &amm::Asset, config: &Config, env: &Env, ) -> Result, ContractError> { match &from.info { // Convert CW20 -> Native - external::AssetInfo::Token { contract_addr } => { + amm::AssetInfo::Token { contract_addr } => { // This flow uses Cw20::Send which calls the adapter's `Receive` hook. let send_msg = Cw20ExecuteMsg::Send { contract: config.cw20_adapter_address.to_string(), @@ -385,7 +432,7 @@ fn create_conversion_msg( })) } // Convert Native -> CW20 - external::AssetInfo::NativeToken { denom } => Ok(CosmosMsg::Wasm(WasmMsg::Execute { + amm::AssetInfo::NativeToken { denom } => Ok(CosmosMsg::Wasm(WasmMsg::Execute { contract_addr: config.cw20_adapter_address.to_string(), msg: to_json_binary(&cw20_adapter::ExecuteMsg::RedeemAndTransfer { recipient: Some(env.contract.address.to_string()), @@ -398,7 +445,7 @@ fn create_conversion_msg( } } -fn get_operation_output(op: &Operation) -> Result { +fn get_operation_output(op: &Operation) -> Result { Ok(match op { Operation::AmmSwap(o) => o.ask_asset_info.clone(), Operation::OrderbookSwap(o) => o.ask_asset_info.clone(), @@ -496,22 +543,22 @@ fn parse_amount_from_conversion_reply(msg: &Reply, env: &Env) -> Result Result { - let mut native_info: Option = None; - let mut cw20_info: Option = None; + let mut native_info: Option = None; + let mut cw20_info: Option = None; for split in &next_stage.splits { let first_op = split.path.first().ok_or(ContractError::EmptyRoute {})?; let offer_info = get_operation_input(first_op)?; match offer_info { - external::AssetInfo::NativeToken { .. } => { + amm::AssetInfo::NativeToken { .. } => { if native_info.is_none() { native_info = Some(offer_info); } } - external::AssetInfo::Token { .. } => { + amm::AssetInfo::Token { .. } => { if cw20_info.is_none() { cw20_info = Some(offer_info); } @@ -523,10 +570,10 @@ fn plan_next_stage( let mut cw20_have = Uint128::zero(); for asset in accumulated_assets { match &asset.info { - external::AssetInfo::NativeToken { .. } => { + amm::AssetInfo::NativeToken { .. } => { native_have += asset.amount; } - external::AssetInfo::Token { .. } => { + amm::AssetInfo::Token { .. } => { cw20_have += asset.amount; } } @@ -540,17 +587,17 @@ fn plan_next_stage( let first_op = split.path.first().ok_or(ContractError::EmptyRoute {})?; let offer_info = get_operation_input(first_op)?; match offer_info { - external::AssetInfo::NativeToken { .. } => total_native_needs += amount_for_split, - external::AssetInfo::Token { .. } => total_cw20_needs += amount_for_split, + amm::AssetInfo::NativeToken { .. } => total_native_needs += amount_for_split, + amm::AssetInfo::Token { .. } => total_cw20_needs += amount_for_split, } } - let mut conversions_needed: Vec<(external::Asset, external::AssetInfo)> = vec![]; + let mut conversions_needed: Vec<(amm::Asset, amm::AssetInfo)> = vec![]; if native_have > total_native_needs { if let Some(target_info) = &cw20_info { let native_asset_to_convert_info = accumulated_assets .iter() - .find(|a| matches!(a.info, external::AssetInfo::NativeToken { .. })) + .find(|a| matches!(a.info, amm::AssetInfo::NativeToken { .. })) .map(|a| a.info.clone()) .ok_or_else(|| { StdError::generic_err( @@ -559,7 +606,7 @@ fn plan_next_stage( })?; conversions_needed.push(( - external::Asset { + amm::Asset { info: native_asset_to_convert_info, amount: native_have - total_native_needs, }, @@ -571,7 +618,7 @@ fn plan_next_stage( if let Some(target_info) = &native_info { let cw20_asset_to_convert_info = accumulated_assets .iter() - .find(|a| matches!(a.info, external::AssetInfo::Token { .. })) + .find(|a| matches!(a.info, amm::AssetInfo::Token { .. })) .map(|a| a.info.clone()) .ok_or_else(|| { StdError::generic_err( @@ -580,7 +627,7 @@ fn plan_next_stage( })?; conversions_needed.push(( - external::Asset { + amm::Asset { info: cw20_asset_to_convert_info, amount: cw20_have - total_cw20_needs, }, @@ -604,8 +651,8 @@ fn plan_next_stage( .map_err(StdError::from)? }; match offer_info { - external::AssetInfo::NativeToken { .. } => native_allocated += amount_for_split, - external::AssetInfo::Token { .. } => cw20_allocated += amount_for_split, + amm::AssetInfo::NativeToken { .. } => native_allocated += amount_for_split, + amm::AssetInfo::Token { .. } => cw20_allocated += amount_for_split, } swaps_to_execute.push(PlannedSwap { operation: first_op.clone(), @@ -619,7 +666,7 @@ fn plan_next_stage( }) } -fn get_operation_input(op: &Operation) -> Result { +fn get_operation_input(op: &Operation) -> Result { Ok(match op { Operation::AmmSwap(o) => o.offer_asset_info.clone(), Operation::OrderbookSwap(o) => o.offer_asset_info.clone(), @@ -629,13 +676,13 @@ fn get_operation_input(op: &Operation) -> Result, env: Env, - state: &mut ReplyState, + exec_state: &mut ExecutionState, + plan: &RoutePlan, reply_id: u64, swaps: Vec, ) -> Result, ContractError> { let mut submessages = vec![]; - // Filter out zero-amount swaps before creating messages --- for swap in swaps.into_iter().filter(|s| !s.amount.is_zero()) { let offer_asset_info = get_operation_input(&swap.operation)?; let msg = @@ -643,21 +690,20 @@ fn execute_planned_swaps( submessages.push(SubMsg::reply_on_success(msg, reply_id)); } - // If all swaps were zero-amount, we might not have any messages. - // We need to handle this case by proceeding directly to the next step. if submessages.is_empty() { - state.current_stage_index += 1; - return proceed_to_next_step(deps, env, state, reply_id); + exec_state.current_stage_index += 1; + return proceed_to_next_step(deps, env, exec_state, plan, reply_id); } - state.awaiting = Awaiting::Swaps; - state.replies_expected = submessages.len() as u64; - REPLY_STATES.save(deps.storage, reply_id, state)?; + exec_state.awaiting = Awaiting::Swaps; + exec_state.replies_expected = submessages.len() as u64; + + EXECUTION_STATES.save(deps.storage, reply_id, exec_state)?; Ok(Response::new() .add_submessages(submessages) .add_attribute("action", "executing_planned_swaps") - .add_attribute("stage_index", state.current_stage_index.to_string())) + .add_attribute("stage_index", exec_state.current_stage_index.to_string())) } fn get_operation_address(op: &Operation) -> &String { @@ -671,13 +717,14 @@ fn handle_path_conversion_reply( mut deps: DepsMut, env: Env, msg: Reply, - state: &mut ReplyState, + exec_state: &mut ExecutionState, + _plan: &RoutePlan, ) -> Result, ContractError> { let master_reply_id = msg.id; let converted_amount = parse_amount_from_conversion_reply(&msg, &env)?; - let pending_op_details = state.pending_path_op.take().ok_or_else(|| { + let pending_op_details = exec_state.pending_path_op.take().ok_or_else(|| { StdError::generic_err("Path conversion state is invalid: no pending operation found") })?; @@ -692,9 +739,9 @@ fn handle_path_conversion_reply( )?; let sub_msg = SubMsg::reply_on_success(swap_msg, master_reply_id); - // Reset the state back to Awaiting::Swaps to continue the stage. - state.awaiting = Awaiting::Swaps; - REPLY_STATES.save(deps.storage, master_reply_id, state)?; + exec_state.awaiting = Awaiting::Swaps; + + EXECUTION_STATES.save(deps.storage, master_reply_id, exec_state)?; Ok(Response::new() .add_submessage(sub_msg) diff --git a/contracts/dex_aggregator/src/state.rs b/contracts/dex_aggregator/src/state.rs index 0b41b63..e81556e 100644 --- a/contracts/dex_aggregator/src/state.rs +++ b/contracts/dex_aggregator/src/state.rs @@ -1,4 +1,4 @@ -use crate::msg::{external, Operation, PlannedSwap}; +use crate::msg::{amm, Operation, PlannedSwap, Stage}; use cosmwasm_schema::cw_serde; use cosmwasm_std::{Addr, Decimal, Uint128}; use cw_storage_plus::{Item, Map}; @@ -30,19 +30,23 @@ pub struct PendingPathOp { } #[cw_serde] -pub struct ReplyState { +pub struct RoutePlan { pub sender: Addr, pub minimum_receive: Uint128, - pub stages: Vec, + pub stages: Vec, +} + +#[cw_serde] +pub struct ExecutionState { pub awaiting: Awaiting, pub current_stage_index: u64, pub replies_expected: u64, - pub accumulated_assets: Vec, + pub accumulated_assets: Vec, pub pending_swaps: Vec, - pub conversion_target_asset: Option, pub pending_path_op: Option, } -pub const REPLY_STATES: Map = Map::new("reply_states"); +pub const ROUTE_PLANS: Map = Map::new("route_plans"); +pub const EXECUTION_STATES: Map = Map::new("execution_states"); pub const REPLY_ID_COUNTER: Item = Item::new("reply_id_counter"); diff --git a/tests/integration.rs b/tests/integration.rs index 3bb65e9..5b84d2d 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -6,7 +6,7 @@ use cosmwasm_std::{to_json_binary, Addr, Coin, Decimal, Uint128}; use cw20::{BalanceResponse, Cw20QueryMsg}; use cw20_base::msg::InstantiateMsg as Cw20InstantiateMsg; use dex_aggregator::msg::{ - cw20_adapter, external, AmmSwapOp, Cw20HookMsg, ExecuteMsg, InstantiateMsg, Operation, + amm, cw20_adapter, AmmSwapOp, Cw20HookMsg, ExecuteMsg, InstantiateMsg, Operation, OrderbookSwapOp, QueryMsg, Split, Stage, }; use dex_aggregator::state::Config as AggregatorConfig; @@ -284,17 +284,17 @@ fn test_aggregate_swap_success() { // Split 3 (25%): 25 INJ -> OB @ 30.0 = 750 USDT // Total Output: 330 + 840 + 750 = 1920 USDT - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![Stage { splits: vec![ Split { percent: 33, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_1_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -303,10 +303,10 @@ fn test_aggregate_swap_success() { percent: 42, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_2_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -315,10 +315,10 @@ fn test_aggregate_swap_success() { percent: 25, path: vec![Operation::OrderbookSwap(OrderbookSwapOp { swap_contract: env.mock_ob_inj_usdt_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -393,7 +393,7 @@ fn test_multi_stage_aggregate_swap_success() { // Split 2 (51%): 51,000 INJ -> AMM2 @ 20.0 = 1,020,000 USDT // Total Final Output: 490,000 + 1,020,000 = 1,510,000 USDT - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![ // Stage 1: 100% of USDT to the Orderbook to get INJ. Stage { @@ -401,10 +401,10 @@ fn test_multi_stage_aggregate_swap_success() { percent: 100, path: vec![Operation::OrderbookSwap(OrderbookSwapOp { swap_contract: env.mock_ob_usdt_inj_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, })], @@ -417,10 +417,10 @@ fn test_multi_stage_aggregate_swap_success() { percent: 49, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_1_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -429,10 +429,10 @@ fn test_multi_stage_aggregate_swap_success() { percent: 51, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_2_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -970,7 +970,7 @@ fn test_full_normalization_route() { let native_shroom_denom = format!("factory/{}/{}", setup.adapter_addr, setup.shroom_cw20_addr); - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![ // Stage 1: INJ -> SHROOM (mixed native/cw20 output) Stage { @@ -979,10 +979,10 @@ fn test_full_normalization_route() { percent: 50, path: vec![Operation::OrderbookSwap(OrderbookSwapOp { swap_contract: setup.mock_inj_to_native_shroom_ob.clone(), - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: native_shroom_denom.clone(), }, })], @@ -991,10 +991,10 @@ fn test_full_normalization_route() { percent: 50, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: setup.mock_inj_to_cw20_shroom_amm.clone(), - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, - ask_asset_info: external::AssetInfo::Token { + ask_asset_info: amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }, })], @@ -1007,10 +1007,10 @@ fn test_full_normalization_route() { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: setup.mock_cw20_shroom_to_cw20_sai_amm.clone(), - offer_asset_info: external::AssetInfo::Token { + offer_asset_info: amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }, - ask_asset_info: external::AssetInfo::Token { + ask_asset_info: amm::AssetInfo::Token { contract_addr: setup.sai_cw20_addr.clone(), }, })], @@ -1056,7 +1056,7 @@ fn test_multi_stage_with_final_normalization() { let native_shroom_denom = format!("factory/{}/{}", setup.adapter_addr, setup.shroom_cw20_addr); - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![ // Stage 1: 100% of USDT to the Orderbook to get INJ. Stage { @@ -1064,10 +1064,10 @@ fn test_multi_stage_with_final_normalization() { percent: 100, path: vec![Operation::OrderbookSwap(OrderbookSwapOp { swap_contract: setup.mock_usdt_to_inj_ob.clone(), - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -1080,10 +1080,10 @@ fn test_multi_stage_with_final_normalization() { percent: 10, // 10% to CW20 SHROOM path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: setup.mock_inj_to_cw20_shroom_amm.clone(), - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, - ask_asset_info: external::AssetInfo::Token { + ask_asset_info: amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }, })], @@ -1092,10 +1092,10 @@ fn test_multi_stage_with_final_normalization() { percent: 90, // 90% to Native SHROOM path: vec![Operation::OrderbookSwap(OrderbookSwapOp { swap_contract: setup.mock_inj_to_native_shroom_ob.clone(), - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: native_shroom_denom.clone(), }, })], @@ -1173,16 +1173,16 @@ fn test_cw20_entry_point_swap_success() { // --- Define the Swap --- // The user wants to swap 1,000 SHROOM for SAI. // The mock AMM rate is 0.1, so they expect 100 SAI in return. - let hook_msg = Cw20HookMsg::AggregateSwaps { + let hook_msg = Cw20HookMsg::ExecuteRoute { stages: vec![Stage { splits: vec![Split { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: setup.mock_cw20_shroom_to_cw20_sai_amm.clone(), - offer_asset_info: external::AssetInfo::Token { + offer_asset_info: amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }, - ask_asset_info: external::AssetInfo::Token { + ask_asset_info: amm::AssetInfo::Token { contract_addr: setup.sai_cw20_addr.clone(), }, })], @@ -1248,7 +1248,7 @@ fn test_reverse_normalization_route() { let native_shroom_denom = format!("factory/{}/{}", setup.adapter_addr, setup.shroom_cw20_addr); - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![ // Stage 1: Get CW20 SHROOM Stage { @@ -1256,10 +1256,10 @@ fn test_reverse_normalization_route() { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: setup.mock_inj_to_cw20_shroom_amm.clone(), - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, - ask_asset_info: external::AssetInfo::Token { + ask_asset_info: amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }, })], @@ -1272,10 +1272,10 @@ fn test_reverse_normalization_route() { path: vec![Operation::OrderbookSwap(OrderbookSwapOp { swap_contract: setup.mock_native_shroom_to_usdt_ob.clone(), // This is the key part of the test: the offer asset is NATIVE - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: native_shroom_denom.clone(), }, - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, })], @@ -1340,17 +1340,17 @@ fn test_failure_if_minimum_receive_not_met() { .unwrap(); let initial_inj_amount = Uint128::from_str(&initial_inj_balance.amount).unwrap(); - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![Stage { splits: vec![ Split { percent: 33, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_1_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -1359,10 +1359,10 @@ fn test_failure_if_minimum_receive_not_met() { percent: 42, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_2_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -1371,10 +1371,10 @@ fn test_failure_if_minimum_receive_not_met() { percent: 25, path: vec![Operation::OrderbookSwap(OrderbookSwapOp { swap_contract: env.mock_ob_inj_usdt_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -1433,17 +1433,17 @@ fn test_failure_on_invalid_percentage_sum() { .unwrap(); let initial_inj_amount = Uint128::from_str(&initial_inj_balance.amount).unwrap(); - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![Stage { splits: vec![ Split { percent: 50, // 50% path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_1_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -1452,10 +1452,10 @@ fn test_failure_on_invalid_percentage_sum() { percent: 49, // + 49% = 99% (Invalid!) path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_2_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -1519,13 +1519,13 @@ fn test_mixed_input_unified_output_reconciliation() { // - TOTAL: 460 USDT // Asset definitions - let cw20_shroom_info = external::AssetInfo::Token { + let cw20_shroom_info = amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }; - let native_shroom_info = external::AssetInfo::NativeToken { + let native_shroom_info = amm::AssetInfo::NativeToken { denom: format!("factory/{}/{}", setup.adapter_addr, setup.shroom_cw20_addr), }; - let usdt_info = external::AssetInfo::NativeToken { + let usdt_info = amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }; @@ -1535,7 +1535,7 @@ fn test_mixed_input_unified_output_reconciliation() { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: setup.mock_inj_to_cw20_shroom_amm.clone(), - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, ask_asset_info: cw20_shroom_info.clone(), @@ -1567,7 +1567,7 @@ fn test_mixed_input_unified_output_reconciliation() { ], }; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { minimum_receive: Some("459000000".to_string()), // Min 459 USDT (Target is 460) stages: vec![stage1, stage2], }; @@ -1642,13 +1642,13 @@ fn test_cw20_input_with_initial_reconciliation() { .unwrap(); // Asset definitions for the stage - let cw20_shroom_info = external::AssetInfo::Token { + let cw20_shroom_info = amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }; - let native_shroom_info = external::AssetInfo::NativeToken { + let native_shroom_info = amm::AssetInfo::NativeToken { denom: format!("factory/{}/{}", setup.adapter_addr, setup.shroom_cw20_addr), }; - let usdt_info = external::AssetInfo::NativeToken { + let usdt_info = amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }; @@ -1676,7 +1676,7 @@ fn test_cw20_input_with_initial_reconciliation() { }; // The hook message sent with the CW20 token - let hook_msg = Cw20HookMsg::AggregateSwaps { + let hook_msg = Cw20HookMsg::ExecuteRoute { minimum_receive: Some("469000000".to_string()), // Min 469 USDT (Target is 470) stages: vec![stage1], }; @@ -1742,16 +1742,16 @@ fn test_complex_reconciliation_mixed_to_mixed() { // - TOTAL: 425 USDT // Asset definitions for clarity - let cw20_shroom_info = external::AssetInfo::Token { + let cw20_shroom_info = amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }; - let native_shroom_info = external::AssetInfo::NativeToken { + let native_shroom_info = amm::AssetInfo::NativeToken { denom: format!("factory/{}/{}", setup.adapter_addr, setup.shroom_cw20_addr), }; - let usdt_info = external::AssetInfo::NativeToken { + let usdt_info = amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }; - let inj_info = external::AssetInfo::NativeToken { + let inj_info = amm::AssetInfo::NativeToken { denom: "inj".to_string(), }; @@ -1803,7 +1803,7 @@ fn test_complex_reconciliation_mixed_to_mixed() { ], }; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { minimum_receive: Some("424000000".to_string()), // Min 424 USDT (Target is 425) stages: vec![stage1, stage2], }; @@ -1858,13 +1858,13 @@ fn test_final_output_is_cw20_token() { // Stage 2: 1000 CW20 SHROOM -> 100 CW20 SAI // Asset definitions for clarity - let inj_info = external::AssetInfo::NativeToken { + let inj_info = amm::AssetInfo::NativeToken { denom: "inj".to_string(), }; - let cw20_shroom_info = external::AssetInfo::Token { + let cw20_shroom_info = amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }; - let cw20_sai_info = external::AssetInfo::Token { + let cw20_sai_info = amm::AssetInfo::Token { contract_addr: setup.sai_cw20_addr.clone(), }; @@ -1890,7 +1890,7 @@ fn test_final_output_is_cw20_token() { }], }; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { minimum_receive: Some("99000000".to_string()), // Min 99 SAI (Target is 100) stages: vec![stage1, stage2], }; @@ -1948,10 +1948,10 @@ fn test_native_input_with_initial_cw20_requirement() { // Asset definitions let native_shroom_denom = format!("factory/{}/{}", setup.adapter_addr, setup.shroom_cw20_addr); - let cw20_shroom_info = external::AssetInfo::Token { + let cw20_shroom_info = amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }; - let cw20_sai_info = external::AssetInfo::Token { + let cw20_sai_info = amm::AssetInfo::Token { contract_addr: setup.sai_cw20_addr.clone(), }; @@ -2007,7 +2007,7 @@ fn test_native_input_with_initial_cw20_requirement() { }], }; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { minimum_receive: Some("99000000".to_string()), // Min 99 SAI (Target is 100) stages: vec![stage1], }; @@ -2061,10 +2061,10 @@ fn test_zero_amount_from_split_is_handled_gracefully() { percent: 50, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_1_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -2073,10 +2073,10 @@ fn test_zero_amount_from_split_is_handled_gracefully() { percent: 50, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_2_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -2084,7 +2084,7 @@ fn test_zero_amount_from_split_is_handled_gracefully() { ], }; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![stage1], minimum_receive: None, // We don't care about the output amount, only that it doesn't fail. }; @@ -2140,10 +2140,10 @@ fn test_stage_with_single_hundred_percent_split() { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_1_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -2155,17 +2155,17 @@ fn test_stage_with_single_hundred_percent_split() { percent: 100, path: vec![Operation::OrderbookSwap(OrderbookSwapOp { swap_contract: env.mock_ob_usdt_inj_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, })], }], }; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![stage1, stage2], minimum_receive: Some("99000000000000000000".to_string()), // Min 99 INJ }; @@ -2233,10 +2233,10 @@ fn test_intermediate_swap_failure_reverts_transaction() { percent: 100, path: vec![Operation::OrderbookSwap(OrderbookSwapOp { swap_contract: env.mock_ob_usdt_inj_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, })], @@ -2251,10 +2251,10 @@ fn test_intermediate_swap_failure_reverts_transaction() { percent: 50, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: env.mock_amm_1_addr.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -2264,10 +2264,10 @@ fn test_intermediate_swap_failure_reverts_transaction() { percent: 50, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: "inj1invalidcontractaddressxxxxxxxxxxxxxx".to_string(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -2275,7 +2275,7 @@ fn test_intermediate_swap_failure_reverts_transaction() { ], }; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![stage1, stage2], minimum_receive: None, // Not relevant, as the transaction should fail. }; @@ -2338,16 +2338,16 @@ fn test_fee_collection_on_single_swap() { // Fee: 1,000 USDT * 0.3% = 3 USDT. // Net Output to User: 1000 - 3 = 997 USDT. - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![Stage { splits: vec![Split { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: fee_pool_address.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -2477,17 +2477,17 @@ fn test_fee_collection_on_cw20_output() { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: fee_pool_address, - ask_asset_info: external::AssetInfo::Token { + ask_asset_info: amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], }], }; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![stage1], minimum_receive: Some("984000000".to_string()), // Min 984 SHROOM }; @@ -2616,16 +2616,16 @@ fn test_full_admin_fee_lifecycle() { .unwrap(); // --- 2. User swaps, fee goes to ORIGINAL collector --- - let swap_msg = ExecuteMsg::AggregateSwaps { + let swap_msg = ExecuteMsg::ExecuteRoute { stages: vec![Stage { splits: vec![Split { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: fee_pool_address.clone(), - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -2773,10 +2773,10 @@ fn test_multi_split_with_mixed_fees() { percent: 40, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: taxed_pool, - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -2786,10 +2786,10 @@ fn test_multi_split_with_mixed_fees() { percent: 60, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: untaxed_pool, - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -2797,7 +2797,7 @@ fn test_multi_split_with_mixed_fees() { ], }; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![stage1], minimum_receive: Some("1595000000".to_string()), // Min 1595 USDT }; @@ -2885,16 +2885,16 @@ fn test_fee_truncates_to_zero() { // Fee Calculation: 100,000 * 0.000001 = 0.1, which truncates to 0. let input_amount = Uint128::new(10_000_000_000_000_000u128); // 10^16 - let swap_msg = ExecuteMsg::AggregateSwaps { + let swap_msg = ExecuteMsg::ExecuteRoute { stages: vec![Stage { splits: vec![Split { percent: 100, path: vec![Operation::AmmSwap(AmmSwapOp { pool_address: fee_pool_address, - ask_asset_info: external::AssetInfo::NativeToken { + ask_asset_info: amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }, - offer_asset_info: external::AssetInfo::NativeToken { + offer_asset_info: amm::AssetInfo::NativeToken { denom: "inj".to_string(), }, })], @@ -3050,16 +3050,16 @@ fn test_multi_hop_path_with_mid_path_conversion() { // match the required input of Hop 2 (Native SHROOM), forcing a conversion. // Asset definitions for clarity - let inj_info = external::AssetInfo::NativeToken { + let inj_info = amm::AssetInfo::NativeToken { denom: "inj".to_string(), }; - let cw20_shroom_info = external::AssetInfo::Token { + let cw20_shroom_info = amm::AssetInfo::Token { contract_addr: setup.shroom_cw20_addr.clone(), }; - let native_shroom_info = external::AssetInfo::NativeToken { + let native_shroom_info = amm::AssetInfo::NativeToken { denom: format!("factory/{}/{}", setup.adapter_addr, setup.shroom_cw20_addr), }; - let usdt_info = external::AssetInfo::NativeToken { + let usdt_info = amm::AssetInfo::NativeToken { denom: "usdt".to_string(), }; @@ -3085,7 +3085,7 @@ fn test_multi_hop_path_with_mid_path_conversion() { }), ]; - let msg = ExecuteMsg::AggregateSwaps { + let msg = ExecuteMsg::ExecuteRoute { stages: vec![Stage { splits: vec![Split { percent: 100, @@ -3115,11 +3115,13 @@ fn test_multi_hop_path_with_mid_path_conversion() { &[funds_to_send.clone()], user, ); + assert!( res.is_ok(), "Execution with mid-path conversion failed: {:?}", res.unwrap_err() ); + println!("Gas Used: {}", res.unwrap().gas_info.gas_used); // --- ASSERT FINAL BALANCE --- let final_inj_balance_response = bank @@ -3147,3 +3149,149 @@ fn test_multi_hop_path_with_mid_path_conversion() { "Final balance should be greater than initial after a profitable swap" ); } + +#[test] +fn test_emergency_withdraw() { + // 1. --- SETUP --- + let setup = setup_for_conversion_test(); + let wasm = Wasm::new(&setup.env.app); + let bank = Bank::new(&setup.env.app); + let admin = &setup.env.admin; + let unauthorized_user = &setup.env.user; + let aggregator_addr = &setup.env.aggregator_addr; + let shroom_cw20_addr = &setup.shroom_cw20_addr; + + // 2. --- ARRANGE: Fund the aggregator contract with assets to withdraw --- + let native_inj_to_send = Coin::new(100_000_000_000_000_000_000u128, "inj"); // 100 INJ + let cw20_shroom_to_send = Uint128::new(500_000_000); // 500 SHROOM + + // Admin sends 100 INJ to the aggregator contract + bank.send( + MsgSend { + from_address: admin.address(), + to_address: aggregator_addr.clone(), + amount: vec![ProtoCoin { + denom: native_inj_to_send.denom.clone(), + amount: native_inj_to_send.amount.to_string(), + }], + }, + admin, + ) + .unwrap(); + + // Admin mints and sends 500 SHROOM to the aggregator contract + wasm.execute( + shroom_cw20_addr, + &cw20_base::msg::ExecuteMsg::Mint { + recipient: aggregator_addr.clone(), + amount: cw20_shroom_to_send, + }, + &[], + admin, + ) + .unwrap(); + + // 3. --- ACT & ASSERT --- + + // Test Case 1: Fails if called by an unauthorized user + let msg_unauthorized = ExecuteMsg::EmergencyWithdraw { + asset_info: amm::AssetInfo::NativeToken { + denom: "inj".to_string(), + }, + }; + let res_unauthorized = wasm.execute(aggregator_addr, &msg_unauthorized, &[], unauthorized_user); + assert!(res_unauthorized.is_err()); + assert!(res_unauthorized + .unwrap_err() + .to_string() + .contains("Unauthorized")); + + // Test Case 2: Admin successfully withdraws the native INJ + let admin_inj_balance_before = bank + .query_balance(&QueryBalanceRequest { + address: admin.address(), + denom: "inj".to_string(), + }) + .unwrap() + .balance + .unwrap() + .amount + .parse::() + .unwrap(); + + let msg_inj = ExecuteMsg::EmergencyWithdraw { + asset_info: amm::AssetInfo::NativeToken { + denom: "inj".to_string(), + }, + }; + wasm.execute(aggregator_addr, &msg_inj, &[], admin).unwrap(); + + let admin_inj_balance_after = bank + .query_balance(&QueryBalanceRequest { + address: admin.address(), + denom: "inj".to_string(), + }) + .unwrap() + .balance + .unwrap() + .amount + .parse::() + .unwrap(); + + // Admin's balance should increase by (exactly 100 INJ - gas fees) + // A simple check is to ensure it increased significantly. + assert!(admin_inj_balance_after > admin_inj_balance_before); + + // Contract's INJ balance should now be zero + let contract_inj_balance = bank + .query_balance(&QueryBalanceRequest { + address: aggregator_addr.clone(), + denom: "inj".to_string(), + }) + .unwrap() + .balance; + assert!(contract_inj_balance.is_none() || contract_inj_balance.unwrap().amount == "0"); + + // Test Case 3: Admin successfully withdraws the CW20 SHROOM + let admin_shroom_balance_before: BalanceResponse = wasm + .query( + shroom_cw20_addr, + &Cw20QueryMsg::Balance { + address: admin.address(), + }, + ) + .unwrap(); + + let msg_shroom = ExecuteMsg::EmergencyWithdraw { + asset_info: amm::AssetInfo::Token { + contract_addr: shroom_cw20_addr.clone(), + }, + }; + wasm.execute(aggregator_addr, &msg_shroom, &[], admin) + .unwrap(); + + let admin_shroom_balance_after: BalanceResponse = wasm + .query( + shroom_cw20_addr, + &Cw20QueryMsg::Balance { + address: admin.address(), + }, + ) + .unwrap(); + + assert_eq!( + admin_shroom_balance_after.balance, + admin_shroom_balance_before.balance + cw20_shroom_to_send + ); + + // Contract's SHROOM balance should now be zero + let contract_shroom_balance: BalanceResponse = wasm + .query( + shroom_cw20_addr, + &Cw20QueryMsg::Balance { + address: aggregator_addr.clone(), + }, + ) + .unwrap(); + assert_eq!(contract_shroom_balance.balance, Uint128::zero()); +}