From bf1f76b051ad4254b7a0810b22db80b79403bc27 Mon Sep 17 00:00:00 2001
From: r4bbit <445106+0x-r4bbit@users.noreply.github.com>
Date: Thu, 6 Aug 2026 14:14:15 +0200
Subject: [PATCH] fix(modules/amm): swap plan must use the pool's stored vault
 ids
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The swap-submission plan derived the pool's vaults from the canonical token
order (compute_vault_pda(pool, canonical_token_a/b) via derive_pair), but the
guest asserts the provided vaults against the pool's stored vault_a_id /
vault_b_id, which are in the pool's *creation* order. compute_pool_pda_seed
canonicalizes, so the pool address is order-independent, but NewDefinition
stores def_a/def_b (and their vaults) as created — so for a pool created in
non-canonical order, the plan put vault_b in the vault_a slot and the guest
panicked with "Vault A was not provided", reverting the swap.

Read the pool account in swapExactInput and pass its data to
swap_exact_in_plan, which now uses pool.vault_a_id / pool.vault_b_id verbatim
for the vault slots (pool / current_tick / clock stay from derive_pair, since
those are order-independent). This restores the order-agnostic behavior the
original swap client had before it was rewired onto the canonical derive_pair
in 737b2f6.

Adds a regression test that builds a non-canonically-created pool (stored
def_a = the smaller-valued token) and asserts the plan emits the pool's stored
vaults, guarding that they differ from the canonical derivation.
---
 modules/amm/ffi/src/api/request.rs  |  4 +++
 modules/amm/ffi/src/api/swap.rs     | 17 +++++++--
 modules/amm/ffi/src/api/tests.rs    | 54 ++++++++++++++++++++++++++++-
 modules/amm/src/amm_module_impl.cpp | 19 ++++++++--
 4 files changed, 89 insertions(+), 5 deletions(-)

diff --git a/modules/amm/ffi/src/api/request.rs b/modules/amm/ffi/src/api/request.rs
index fee626e..4fe9046 100644
--- a/modules/amm/ffi/src/api/request.rs
+++ b/modules/amm/ffi/src/api/request.rs
@@ -111,6 +111,10 @@ pub struct SwapExactInPlanRequest {
     pub amount_in: String,
     pub min_out: String,
     pub deadline_ms: String,
+    /// Pool account data (hex Borsh `PoolDefinition`) — its stored `vault_a_id` /
+    /// `vault_b_id` are used verbatim (the guest asserts the vaults in the pool's
+    /// creation order, which needn't match the canonical token order).
+    pub pool_data: String,
 }
 
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq)]
diff --git a/modules/amm/ffi/src/api/swap.rs b/modules/amm/ffi/src/api/swap.rs
index 0ec2d5c..dba7276 100644
--- a/modules/amm/ffi/src/api/swap.rs
+++ b/modules/amm/ffi/src/api/swap.rs
@@ -287,6 +287,18 @@ pub(super) fn swap_exact_in_plan(request: SwapExactInPlanRequest) -> Result<Valu
         return Ok(json!({ "status": "error", "code": "config_unavailable" }));
     };
 
+    // Take the vaults from the pool's STORED ids, in its `def_a`/`def_b` (creation)
+    // order — the guest asserts the provided vaults against `pool.vault_a_id` /
+    // `vault_b_id`, which needn't match `derive_pair`'s canonical order for a
+    // non-canonically created pool. (`pair.pool`/`current_tick`/`clock` are
+    // order-independent, so they still come from `derive_pair`.)
+    let Some(pool) = hex::decode(&request.pool_data)
+        .ok()
+        .and_then(|bytes| borsh::from_slice::<PoolDefinition>(&bytes).ok())
+    else {
+        return Ok(json!({ "status": "error", "code": "no_pool" }));
+    };
+
     let user_input_holding =
         account_id_from_hex(&request.user_input_holding_id, "user input holding id")?;
     let user_output_holding =
@@ -307,8 +319,8 @@ pub(super) fn swap_exact_in_plan(request: SwapExactInPlanRequest) -> Result<Valu
     let account_ids = [
         pair.config,
         pair.pool,
-        pair.vault_a,
-        pair.vault_b,
+        pool.vault_a_id,
+        pool.vault_b_id,
         user_input_holding,
         user_output_holding,
         pair.current_tick,
@@ -434,6 +446,7 @@ mod tests {
             amount_in: String::new(),
             min_out: String::new(),
             deadline_ms: String::new(),
+            pool_data: String::new(),
         })
         .unwrap();
         assert_eq!(plan, expected);
diff --git a/modules/amm/ffi/src/api/tests.rs b/modules/amm/ffi/src/api/tests.rs
index f622483..da462ec 100644
--- a/modules/amm/ffi/src/api/tests.rs
+++ b/modules/amm/ffi/src/api/tests.rs
@@ -22,8 +22,9 @@ use super::{
     plan::plan,
     position::AccountPlanHoldings,
     quote::{div_ceil_u256, minimum_opening_pair, quote, Q64},
+    swap::swap_exact_in_plan,
     ContextRequest, PairIdsRequest, PairSnapshot, PlanRequest, PositionRequest, QuoteRequest,
-    TokenIdsRequest,
+    SwapExactInPlanRequest, TokenIdsRequest,
 };
 use crate::{
     account::{account_id_hex, account_read, decode_account, parse_base58_id, program_id_bytes},
@@ -700,3 +701,54 @@ fn stale_hash_returns_recomputed_quote_without_plan() {
     assert_eq!(value["code"], "quote_changed");
     assert_eq!(value["quote"]["status"], "ok");
 }
+
+#[test]
+fn swap_plan_uses_the_pool_stored_vaults_not_canonical_order() {
+    // A pool created NON-canonically: its stored def_a is the smaller-valued
+    // token, so pool.vault_a_id is the vault for the smaller token — the opposite
+    // of what canonical_pair (larger first) would derive. The plan must emit the
+    // pool's own stored vaults, which is what the guest asserts against.
+    let token_small = AccountId::new([1; 32]);
+    let token_large = AccountId::new([2; 32]);
+    assert!(is_canonical_pair(token_large, token_small)); // large is canonical token_a
+
+    let pool_id = compute_pool_pda(AMM_PROGRAM, token_small, token_large);
+    let pool = PoolDefinition {
+        definition_token_a_id: token_small, // stored non-canonically (small first)
+        definition_token_b_id: token_large,
+        vault_a_id: compute_vault_pda(AMM_PROGRAM, pool_id, token_small),
+        vault_b_id: compute_vault_pda(AMM_PROGRAM, pool_id, token_large),
+        liquidity_pool_id: compute_liquidity_token_pda(AMM_PROGRAM, pool_id),
+        liquidity_pool_supply: 1_000,
+        reserve_a: 1_000,
+        reserve_b: 1_000,
+        fees: 30,
+    };
+
+    let holding = AccountId::new([9; 32]);
+    let plan = swap_exact_in_plan(SwapExactInPlanRequest {
+        amm_program_id: amm_program_id(),
+        token_in_id: account_id_hex(token_small),
+        token_out_id: account_id_hex(token_large),
+        config: account_read(compute_config_pda(AMM_PROGRAM), &config_account()),
+        user_input_holding_id: account_id_hex(holding),
+        user_output_holding_id: account_id_hex(holding),
+        amount_in: String::from("100"),
+        min_out: String::from("0"),
+        deadline_ms: String::from("0"),
+        pool_data: hex::encode(borsh::to_vec(&pool).unwrap()),
+    })
+    .unwrap();
+
+    // Slots 2 and 3 are vault_a / vault_b — in the pool's stored order, not the
+    // canonical order. (A domain error would leave accountIds absent, so these
+    // also assert the plan succeeded.)
+    assert_eq!(plan["accountIds"][2], account_id_hex(pool.vault_a_id));
+    assert_eq!(plan["accountIds"][3], account_id_hex(pool.vault_b_id));
+    // Guard: the stored vault_a genuinely differs from the canonical derivation
+    // (the pre-fix bug would have emitted this one in slot 2).
+    assert_ne!(
+        pool.vault_a_id,
+        compute_vault_pda(AMM_PROGRAM, pool_id, token_large)
+    );
+}
diff --git a/modules/amm/src/amm_module_impl.cpp b/modules/amm/src/amm_module_impl.cpp
index 3afb9e8..2e6e4f9 100644
--- a/modules/amm/src/amm_module_impl.cpp
+++ b/modules/amm/src/amm_module_impl.cpp
@@ -609,13 +609,28 @@ std::string AmmModuleImpl::swapExactInput(const std::string& def_a_hex,
         return {};
     }
 
-    // amm_swap_exact_in_plan resolves the pool, reorders holdings to the pool's
-    // canonical def order, encodes SwapExactInput, and returns a ready-to-submit plan.
+    // Read the pool so the plan can use its stored vault ids (the guest asserts
+    // the vaults in the pool's creation order — see amm_swap_exact_in_plan).
+    const FfiResult poolId = call(amm_pool_id, json{
+        {"ammProgramId", net.amm_program_id},
+        {"tokenInId", def_a_hex},
+        {"tokenOutId", def_b_hex},
+    });
+    if (!poolId.ok) {
+        AMM_TRACE("swapExactInput: FAIL amm_pool_id");
+        return {};
+    }
+    const json pool = readPublicAccount(jStr(poolId.value, "poolId"));
+    const std::string pool_data = jStr(pool.value("account", json::object()), "data");
+
+    // amm_swap_exact_in_plan resolves the pool accounts, encodes SwapExactInput,
+    // and returns a ready-to-submit plan.
     const FfiResult planResult = call(amm_swap_exact_in_plan, json{
         {"ammProgramId", net.amm_program_id},
         {"tokenInId", def_a_hex},
         {"tokenOutId", def_b_hex},
         {"config", config},
+        {"poolData", pool_data},
         {"userInputHoldingId", user_input_holding_hex},
         {"userOutputHoldingId", user_output_holding_hex},
         {"amountIn", amount_in_decimal},