From 3db69ce5e8c0841b5cd63be14a8348c0535c1eee Mon Sep 17 00:00:00 2001
From: Ricardo Guilherme Schmidt <3esmit@gmail.com>
Date: Mon, 10 Aug 2026 11:43:52 -0300
Subject: [PATCH] fix(amm): avoid redundant output holding signatures

---
 modules/amm/ffi/src/api/request.rs  |   6 ++
 modules/amm/ffi/src/api/swap.rs     | 109 +++++++++++++++++++++++++++-
 modules/amm/ffi/src/api/tests.rs    |   3 +
 modules/amm/src/amm_module_impl.cpp |   4 +
 programs/amm/src/swap.rs            |   4 +-
 5 files changed, 121 insertions(+), 5 deletions(-)

diff --git a/modules/amm/ffi/src/api/request.rs b/modules/amm/ffi/src/api/request.rs
index 3d4c4d8..3f0f464 100644
--- a/modules/amm/ffi/src/api/request.rs
+++ b/modules/amm/ffi/src/api/request.rs
@@ -108,6 +108,9 @@ pub struct SwapExactInPlanRequest {
     pub config: AccountRead,
     pub user_input_holding_id: String,
     pub user_output_holding_id: String,
+    /// Current read of the output holding. The plan uses its owner to avoid
+    /// requesting a signature for an already initialized Token holding.
+    pub user_output_holding: AccountRead,
     pub amount_in: String,
     pub min_out: String,
     pub deadline_ms: String,
@@ -126,6 +129,9 @@ pub struct SwapExactOutPlanRequest {
     pub config: AccountRead,
     pub user_input_holding_id: String,
     pub user_output_holding_id: String,
+    /// Current read of the output holding. The plan uses its owner to avoid
+    /// requesting a signature for an already initialized Token holding.
+    pub user_output_holding: AccountRead,
     pub amount_out: String,
     pub max_in: String,
     pub deadline_ms: String,
diff --git a/modules/amm/ffi/src/api/swap.rs b/modules/amm/ffi/src/api/swap.rs
index d2bbed8..52159dc 100644
--- a/modules/amm/ffi/src/api/swap.rs
+++ b/modules/amm/ffi/src/api/swap.rs
@@ -7,7 +7,7 @@ use amm_core::{
     compute_pool_pda, mul_div_ceil, mul_div_floor, price_impact_bps, swap_exact_in_amounts,
     swap_exact_out_amounts, PoolDefinition, FEE_BPS_DENOMINATOR,
 };
-use nssa_core::account::AccountId;
+use nssa_core::{account::AccountId, program::ProgramId};
 use risc0_binfmt::ProgramBinary;
 use serde_json::{json, Value};
 
@@ -18,8 +18,20 @@ use super::{
 };
 use crate::account::{
     account_id_from_hex, account_id_hex, decode_account, parse_program_id, program_id_bytes,
+    AccountRead,
 };
 
+fn output_holding_requires_signature(
+    read: &AccountRead,
+    expected_id: AccountId,
+    token_program_id: ProgramId,
+) -> bool {
+    let Ok((account_id, account)) = decode_account(read) else {
+        return true;
+    };
+    account_id != expected_id || account.program_owner != token_program_id
+}
+
 /// Orders `(token_in, token_out)` into the pool's canonical `(token_a, token_b)`
 /// so derived vault PDAs line up with the pool's stored `vault_a`/`vault_b`.
 fn canonical_pair(token_in: AccountId, token_out: AccountId) -> (AccountId, AccountId) {
@@ -330,7 +342,21 @@ pub(super) fn swap_exact_in_plan(request: SwapExactInPlanRequest) -> Result<Valu
         pair.current_tick,
         pair.clock,
     ];
-    let signing_requirements = [false, false, false, false, true, true, false, false];
+    let output_requires_signature = output_holding_requires_signature(
+        &request.user_output_holding,
+        user_output_holding,
+        pair.token_program,
+    );
+    let signing_requirements = [
+        false,
+        false,
+        false,
+        false,
+        true,
+        output_requires_signature,
+        false,
+        false,
+    ];
 
     Ok(json!({
         "programId": request.amm_program_id,
@@ -397,7 +423,21 @@ pub(super) fn swap_exact_out_plan(request: SwapExactOutPlanRequest) -> Result<Va
         pair.current_tick,
         pair.clock,
     ];
-    let signing_requirements = [false, false, false, false, true, true, false, false];
+    let output_requires_signature = output_holding_requires_signature(
+        &request.user_output_holding,
+        user_output_holding,
+        pair.token_program,
+    );
+    let signing_requirements = [
+        false,
+        false,
+        false,
+        false,
+        true,
+        output_requires_signature,
+        false,
+        false,
+    ];
 
     Ok(json!({
         "programId": request.amm_program_id,
@@ -441,6 +481,54 @@ mod tests {
         }
     }
 
+    #[test]
+    fn output_holding_signature_requires_fresh_or_unreadable_account() {
+        let output_id = AccountId::new([0xAB; 32]);
+        let token_program = [22; 8];
+        let initialized = AccountRead {
+            id: account_id_hex(output_id),
+            status: String::from("ok"),
+            account: Some(WalletAccount {
+                program_owner: crate::account::program_id_hex(token_program),
+                balance: "00".repeat(16),
+                nonce: "00".repeat(16),
+                data: String::new(),
+            }),
+        };
+        assert!(!output_holding_requires_signature(
+            &initialized,
+            output_id,
+            token_program,
+        ));
+
+        let fresh = AccountRead {
+            id: account_id_hex(output_id),
+            status: String::from("ok"),
+            account: Some(WalletAccount {
+                program_owner: "00".repeat(32),
+                balance: "00".repeat(16),
+                nonce: "00".repeat(16),
+                data: String::new(),
+            }),
+        };
+        assert!(output_holding_requires_signature(
+            &fresh,
+            output_id,
+            token_program,
+        ));
+
+        let unreadable = AccountRead {
+            id: account_id_hex(output_id),
+            status: String::from("read_failed"),
+            account: None,
+        };
+        assert!(output_holding_requires_signature(
+            &unreadable,
+            output_id,
+            token_program,
+        ));
+    }
+
     #[test]
     fn resolve_pool_reports_canonical_token_ids() {
         let def_a = AccountId::new([0xAA; 32]);
@@ -515,6 +603,11 @@ mod tests {
             config: dummy_config,
             user_input_holding_id: String::new(),
             user_output_holding_id: String::new(),
+            user_output_holding: AccountRead {
+                id: String::new(),
+                status: String::from("read_failed"),
+                account: None,
+            },
             amount_in: String::new(),
             min_out: String::new(),
             deadline_ms: String::new(),
@@ -540,6 +633,11 @@ mod tests {
             },
             user_input_holding_id: String::new(),
             user_output_holding_id: String::new(),
+            user_output_holding: AccountRead {
+                id: String::new(),
+                status: String::from("read_failed"),
+                account: None,
+            },
             amount_in: String::new(),
             min_out: String::new(),
             deadline_ms: String::new(),
@@ -566,6 +664,11 @@ mod tests {
             config: dummy_config,
             user_input_holding_id: String::new(),
             user_output_holding_id: String::new(),
+            user_output_holding: AccountRead {
+                id: String::new(),
+                status: String::from("read_failed"),
+                account: None,
+            },
             amount_out: String::new(),
             max_in: String::new(),
             deadline_ms: String::new(),
diff --git a/modules/amm/ffi/src/api/tests.rs b/modules/amm/ffi/src/api/tests.rs
index 7f30e9a..90c647d 100644
--- a/modules/amm/ffi/src/api/tests.rs
+++ b/modules/amm/ffi/src/api/tests.rs
@@ -859,6 +859,7 @@ fn swap_plan_uses_the_pool_stored_vaults_not_canonical_order() {
         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),
+        user_output_holding: default_read(holding),
         amount_in: String::from("100"),
         min_out: String::from("0"),
         deadline_ms: String::from("0"),
@@ -896,6 +897,7 @@ fn swap_exact_in_plan_missing_pool_fails_closed_with_err() {
         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),
+        user_output_holding: default_read(holding),
         amount_in: String::from("100"),
         min_out: String::from("0"),
         deadline_ms: String::from("0"),
@@ -934,6 +936,7 @@ fn swap_exact_out_plan_uses_the_pool_stored_vaults_not_canonical_order() {
         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),
+        user_output_holding: default_read(holding),
         amount_out: String::from("100"),
         max_in: String::from("1000"),
         deadline_ms: String::from("0"),
diff --git a/modules/amm/src/amm_module_impl.cpp b/modules/amm/src/amm_module_impl.cpp
index d8e723e..3184071 100644
--- a/modules/amm/src/amm_module_impl.cpp
+++ b/modules/amm/src/amm_module_impl.cpp
@@ -622,6 +622,7 @@ std::string AmmModuleImpl::swapExactInput(const std::string& def_a_hex,
     }
     const json pool = readPublicAccount(jStr(poolId.value, "poolId"));
     const std::string pool_data = jStr(pool.value("account", json::object()), "data");
+    const json user_output_holding = readPublicAccount(user_output_holding_hex);
 
     // amm_swap_exact_in_plan resolves the pool accounts, encodes SwapExactInput,
     // and returns a ready-to-submit plan.
@@ -633,6 +634,7 @@ std::string AmmModuleImpl::swapExactInput(const std::string& def_a_hex,
         {"poolData", pool_data},
         {"userInputHoldingId", user_input_holding_hex},
         {"userOutputHoldingId", user_output_holding_hex},
+        {"userOutputHolding", user_output_holding},
         {"amountIn", amount_in_decimal},
         {"minOut", min_out_decimal},
         {"deadlineMs", deadline_decimal},
@@ -705,6 +707,7 @@ std::string AmmModuleImpl::swapExactOutput(const std::string& def_a_hex,
     }
     const json pool = readPublicAccount(jStr(poolId.value, "poolId"));
     const std::string pool_data = jStr(pool.value("account", json::object()), "data");
+    const json user_output_holding = readPublicAccount(user_output_holding_hex);
 
     // amm_swap_exact_out_plan resolves the pool accounts, encodes SwapExactOutput,
     // and returns a ready-to-submit plan.
@@ -716,6 +719,7 @@ std::string AmmModuleImpl::swapExactOutput(const std::string& def_a_hex,
         {"poolData", pool_data},
         {"userInputHoldingId", user_input_holding_hex},
         {"userOutputHoldingId", user_output_holding_hex},
+        {"userOutputHolding", user_output_holding},
         {"amountOut", amount_out_decimal},
         {"maxIn", max_in_decimal},
         {"deadlineMs", deadline_decimal},
diff --git a/programs/amm/src/swap.rs b/programs/amm/src/swap.rs
index fac5771..f2465bc 100644
--- a/programs/amm/src/swap.rs
+++ b/programs/amm/src/swap.rs
@@ -203,12 +203,12 @@ pub fn swap_exact_input(
     assert_user_holding_owner_or_fresh(
         &user_holding_a,
         token_program_id,
-        "User Token A holding must be owned by the configured Token Program",
+        "User Token A holding must be owned by the configured Token Program or be a fresh authorized account",
     );
     assert_user_holding_owner_or_fresh(
         &user_holding_b,
         token_program_id,
-        "User Token B holding must be owned by the configured Token Program",
+        "User Token B holding must be owned by the configured Token Program or be a fresh authorized account",
     );
     // The current tick is refreshed by a chained call to the oracle; validate its PDA and the
     // clock here so the swap is rejected early with an AMM-level error.