Tax token support

contracts/dex_aggregator/schema/execute_msg.json

@@ -156,6 +156,50 @@
         }
       },
       "additionalProperties": false
+    },
+    {
+      "description": "Registers a new tax token that requires special handling.",
+      "type": "object",
+      "required": [
+        "register_tax_token"
+      ],
+      "properties": {
+        "register_tax_token": {
+          "type": "object",
+          "required": [
+            "contract_addr"
+          ],
+          "properties": {
+            "contract_addr": {
+              "type": "string"
+            }
+          },
+          "additionalProperties": false
+        }
+      },
+      "additionalProperties": false
+    },
+    {
+      "description": "Removes a tax token from the registry.",
+      "type": "object",
+      "required": [
+        "deregister_tax_token"
+      ],
+      "properties": {
+        "deregister_tax_token": {
+          "type": "object",
+          "required": [
+            "contract_addr"
+          ],
+          "properties": {
+            "contract_addr": {
+              "type": "string"
+            }
+          },
+          "additionalProperties": false
+        }
+      },
+      "additionalProperties": false
     }
   ],
   "definitions": {

contracts/dex_aggregator/src/contract.rs

@@ -123,6 +123,12 @@ pub fn execute(
         ExecuteMsg::EmergencyWithdraw { asset_info } => {
             crate::execute::emergency_withdraw(deps, env, info, asset_info)
         }
+        ExecuteMsg::RegisterTaxToken { contract_addr } => {
+            crate::execute::register_tax_token(deps, info, contract_addr)
+        }
+        ExecuteMsg::DeregisterTaxToken { contract_addr } => {
+            crate::execute::deregister_tax_token(deps, info, contract_addr)
+        }
     }
 }
 

contracts/dex_aggregator/src/execute.rs

@@ -10,7 +10,9 @@ use std::str::FromStr;
 use crate::error::ContractError;
 use crate::msg::{self, amm, orderbook, Operation, Stage};
 use crate::reply::proceed_to_next_step;
-use crate::state::{Awaiting, ExecutionState, RoutePlan, CONFIG, FEE_MAP, REPLY_ID_COUNTER};
+use crate::state::{
+    Awaiting, ExecutionState, RoutePlan, CONFIG, FEE_MAP, REPLY_ID_COUNTER, TAX_TOKEN_REGISTRY,
+};
 
 pub fn update_admin(
     deps: DepsMut<InjectiveQueryWrapper>,
@@ -109,17 +111,33 @@ pub fn create_swap_cosmos_msg(
                     }],
                 }),
                 amm::AssetInfo::Token { contract_addr } => {
-                    let cw20_send_msg = Cw20ExecuteMsg::Send {
-                        contract: amm_op.pool_address.clone(),
-                        amount,
-                        msg: to_json_binary(&amm_swap_msg)?,
-                    };
-
-                    CosmosMsg::Wasm(WasmMsg::Execute {
-                        contract_addr: contract_addr.clone(),
-                        msg: to_json_binary(&cw20_send_msg)?,
-                        funds: vec![],
-                    })
+                    let token_addr = deps.api.addr_validate(contract_addr)?;
+                    if TAX_TOKEN_REGISTRY.has(deps.storage, &token_addr) {
+                        // It's a tax token. We must use its tax-exempt send function.
+                        CosmosMsg::Wasm(WasmMsg::Execute {
+                            contract_addr: contract_addr.clone(),
+                            msg: to_json_binary(
+                                &crate::msg::reflection::ExecuteMsg::TaxExemptSend {
+                                    contract: amm_op.pool_address.clone(),
+                                    amount,
+                                    msg: to_json_binary(&amm_swap_msg)?,
+                                },
+                            )?,
+                            funds: vec![],
+                        })
+                    } else {
+                        // It's a standard token. Use the normal Cw20::Send.
+                        let cw20_send_msg = Cw20ExecuteMsg::Send {
+                            contract: amm_op.pool_address.clone(),
+                            amount,
+                            msg: to_json_binary(&amm_swap_msg)?,
+                        };
+                        CosmosMsg::Wasm(WasmMsg::Execute {
+                            contract_addr: contract_addr.clone(),
+                            msg: to_json_binary(&cw20_send_msg)?,
+                            funds: vec![],
+                        })
+                    }
                 }
             }
         }
@@ -328,3 +346,35 @@ pub fn emergency_withdraw(
 
     Ok(response)
 }
+
+pub fn register_tax_token(
+    deps: DepsMut<InjectiveQueryWrapper>,
+    info: MessageInfo,
+    contract_addr: String,
+) -> Result<Response<InjectiveMsgWrapper>, ContractError> {
+    let config = CONFIG.load(deps.storage)?;
+    if info.sender != config.admin {
+        return Err(ContractError::Unauthorized {});
+    }
+    let addr = deps.api.addr_validate(&contract_addr)?;
+    TAX_TOKEN_REGISTRY.save(deps.storage, &addr, &true)?;
+    Ok(Response::new()
+        .add_attribute("action", "register_tax_token")
+        .add_attribute("token_addr", addr))
+}
+
+pub fn deregister_tax_token(
+    deps: DepsMut<InjectiveQueryWrapper>,
+    info: MessageInfo,
+    contract_addr: String,
+) -> Result<Response<InjectiveMsgWrapper>, ContractError> {
+    let config = CONFIG.load(deps.storage)?;
+    if info.sender != config.admin {
+        return Err(ContractError::Unauthorized {});
+    }
+    let addr = deps.api.addr_validate(&contract_addr)?;
+    TAX_TOKEN_REGISTRY.remove(deps.storage, &addr);
+    Ok(Response::new()
+        .add_attribute("action", "deregister_tax_token")
+        .add_attribute("token_addr", addr))
+}

contracts/dex_aggregator/src/msg.rs

@@ -118,6 +118,24 @@ pub mod orderbook {
     }
 }
 
+pub mod reflection {
+    use super::*;
+    use cosmwasm_std::Binary;
+
+    #[cw_serde]
+    pub enum ExecuteMsg {
+        TaxExemptTransfer {
+            recipient: String,
+            amount: Uint128,
+        },
+        TaxExemptSend {
+            contract: String,
+            amount: Uint128,
+            msg: Binary,
+        },
+    }
+}
+
 #[cw_serde]
 pub struct AmmSwapOp {
     pub pool_address: String,
@@ -202,6 +220,14 @@ pub enum ExecuteMsg {
     EmergencyWithdraw {
         asset_info: amm::AssetInfo,
     },
+    /// Registers a new tax token that requires special handling.
+    RegisterTaxToken {
+        contract_addr: String,
+    },
+    /// Removes a tax token from the registry.
+    DeregisterTaxToken {
+        contract_addr: String,
+    },
 }
 
 #[cw_serde]

contracts/dex_aggregator/src/reply.rs

@@ -10,7 +10,7 @@ use crate::execute::create_swap_cosmos_msg;
 use crate::msg::{amm, cw20_adapter, Operation, PlannedSwap, Stage, StagePlan};
 use crate::state::{
     Awaiting, Config, ExecutionState, PendingPathOp, SubmsgReplyState, ACTIVE_ROUTES, CONFIG,
-    FEE_MAP, REPLY_ID_COUNTER, SUBMSG_REPLY_STATES,
+    FEE_MAP, REPLY_ID_COUNTER, SUBMSG_REPLY_STATES, TAX_TOKEN_REGISTRY,
 };
 
 const DECIMAL_FRACTIONAL: u128 = 1_000_000_000_000_000_000;
@@ -142,7 +142,7 @@ fn handle_swap_reply(
         }
     }
 
-    let received_amount = parse_amount_from_swap_reply(events)?;
+    let received_amount = parse_amount_from_swap_reply(events, &env)?;
     let received_asset_info = get_operation_output(replied_op)?;
 
     let replied_path = &current_stage.splits[split_index].path;
@@ -230,7 +230,8 @@ fn handle_swap_reply(
 
         if !fee.is_zero() {
             let config = CONFIG.load(deps.storage)?;
-            let fee_send_msg = create_send_msg(&config.fee_collector, &received_asset_info, fee)?;
+            let fee_send_msg =
+                create_send_msg(&deps, &config.fee_collector, &received_asset_info, fee)?;
             response = response
                 .add_message(fee_send_msg)
                 .add_attribute("fee_collected", fee.to_string())
@@ -256,6 +257,7 @@ fn apply_fee(
 
 // A helper to create the final transfer message.
 fn create_send_msg(
+    deps: &DepsMut<InjectiveQueryWrapper>,
     recipient: &Addr,
     asset_info: &amm::AssetInfo,
     amount: Uint128,
@@ -268,14 +270,31 @@ fn create_send_msg(
                 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(),
-                amount,
-            })?,
-            funds: vec![],
-        })),
+        amm::AssetInfo::Token { contract_addr } => {
+            let token_addr = deps.api.addr_validate(contract_addr)?;
+            // Check if we are dealing with a registered tax token.
+            if TAX_TOKEN_REGISTRY.has(deps.storage, &token_addr) {
+                // Use the new tax-exempt message.
+                Ok(CosmosMsg::Wasm(WasmMsg::Execute {
+                    contract_addr: contract_addr.clone(),
+                    msg: to_json_binary(&crate::msg::reflection::ExecuteMsg::TaxExemptTransfer {
+                        recipient: recipient.to_string(),
+                        amount,
+                    })?,
+                    funds: vec![],
+                }))
+            } else {
+                // Use a standard CW20 Transfer for all other tokens.
+                Ok(CosmosMsg::Wasm(WasmMsg::Execute {
+                    contract_addr: contract_addr.clone(),
+                    msg: to_json_binary(&Cw20ExecuteMsg::Transfer {
+                        recipient: recipient.to_string(),
+                        amount,
+                    })?,
+                    funds: vec![],
+                }))
+            }
+        }
     }
 }
 
@@ -327,6 +346,7 @@ fn handle_final_stage(
         if !total_final_amount.is_zero() {
             // Use the sender address from the immutable plan
             let send_msg = create_send_msg(
+                deps,
                 &exec_state.plan.sender,
                 &target_asset_info,
                 total_final_amount,
@@ -400,6 +420,7 @@ fn handle_final_conversion_reply(
     let mut response = Response::new();
     if !total_final_amount.is_zero() {
         let send_msg = create_send_msg(
+            &deps,
             &exec_state.plan.sender,
             &final_asset_info,
             total_final_amount,
@@ -489,7 +510,36 @@ fn get_operation_output(op: &Operation) -> Result<amm::AssetInfo, ContractError>
     })
 }
 
-fn parse_amount_from_swap_reply(events: &[cosmwasm_std::Event]) -> Result<Uint128, ContractError> {
+fn parse_amount_from_swap_reply(
+    events: &[cosmwasm_std::Event],
+    env: &Env,
+) -> Result<Uint128, ContractError> {
+    // Check for `post_tax_amount` from a tax token's transfer event.
+    for event in events.iter().rev() {
+        if event.ty != "wasm" {
+            continue;
+        }
+        let is_recipient_self = event
+            .attributes
+            .iter()
+            .any(|attr| attr.key == "to" && attr.value == env.contract.address.to_string());
+
+        if is_recipient_self {
+            if let Some(amount_attr) = event
+                .attributes
+                .iter()
+                .find(|attr| attr.key == "post_tax_amount")
+            {
+                return amount_attr.value.parse::<Uint128>().map_err(|_| {
+                    ContractError::MalformedAmountInReply {
+                        value: amount_attr.value.clone(),
+                    }
+                });
+            }
+        }
+    }
+
+    // 2. Fallback to original logic for standard, non-taxable tokens.
     let amount_str_opt = events.iter().find_map(|event| {
         if !event.ty.starts_with("wasm") {
             return None;
@@ -513,12 +563,11 @@ fn parse_amount_from_swap_reply(events: &[cosmwasm_std::Event]) -> Result<Uint12
             } else {
                 &amount_str
             };
-
             integer_part_str
                 .parse::<Uint128>()
                 .map_err(|_| ContractError::MalformedAmountInReply { value: amount_str })
         }
-        None => Ok(Uint128::zero()),
+        None => Ok(Uint128::zero()), // Return zero if no relevant amount is found.
     }
 }
 

contracts/dex_aggregator/src/state.rs

@@ -57,3 +57,8 @@ pub struct SubmsgReplyState {
 pub const ACTIVE_ROUTES: Map<u64, ExecutionState> = Map::new("execution_states");
 pub const SUBMSG_REPLY_STATES: Map<u64, SubmsgReplyState> = Map::new("submsg_reply_states");
 pub const REPLY_ID_COUNTER: Item<u64> = Item::new("reply_id_counter");
+
+/// A registry of known tax tokens that require special handling.
+/// The key is the token's contract address.
+/// The value is a simple boolean `true` to indicate it's registered.
+pub const TAX_TOKEN_REGISTRY: Map<&Addr, bool> = Map::new("tax_tokens");