From 06494e1b14698cb89a7e1fa6e1f1ca50706473c5 Mon Sep 17 00:00:00 2001
From: r4bbit <445106+0x-r4bbit@users.noreply.github.com>
Date: Mon, 4 May 2026 14:47:50 +0200
Subject: [PATCH] chore(amm): add defensive check for lp token solvency

This check is added to fulfill the program invariant that no more tokens
than owned can be burned. This was not a bug before, because the `token`
program will revert on `Transfer::Burn` when one tries to burn more
tokens than available.

So this change is merely for making the invariant explicit.
---
 amm/src/remove.rs |  4 ++++
 amm/src/tests.rs  | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/amm/src/remove.rs b/amm/src/remove.rs
index 67e33b4..078ef04 100644
--- a/amm/src/remove.rs
+++ b/amm/src/remove.rs
@@ -101,6 +101,10 @@ pub fn remove_liquidity(
         pool_def_data.liquidity_pool_supply > MINIMUM_LIQUIDITY,
         "Pool only contains locked liquidity"
     );
+    assert!(
+        remove_liquidity_amount <= user_lp_balance,
+        "Remove amount exceeds user LP balance"
+    );
     let unlocked_liquidity = pool_def_data.liquidity_pool_supply - MINIMUM_LIQUIDITY;
     // The remove instruction never sees the LP lock account directly, so we must still refuse any
     // request that would burn through the permanent floor even if ownership is already corrupted.
diff --git a/amm/src/tests.rs b/amm/src/tests.rs
index c6d12a0..f925047 100644
--- a/amm/src/tests.rs
+++ b/amm/src/tests.rs
@@ -3411,6 +3411,24 @@ fn test_remove_liquidity_rejects_user_holding_b_wrong_program() {
     );
 }
 
+#[should_panic(expected = "Remove amount exceeds user LP balance")]
+#[test]
+fn test_remove_liquidity_rejects_amount_exceeding_user_lp_balance() {
+    let lp_balance = BalanceForTests::remove_amount_lp() - 1;
+    let _ = remove_liquidity(
+        AccountWithMetadataForTests::pool_definition_init(),
+        AccountWithMetadataForTests::vault_a_init(),
+        AccountWithMetadataForTests::vault_b_init(),
+        AccountWithMetadataForTests::pool_lp_init(),
+        AccountWithMetadataForTests::user_holding_a(),
+        AccountWithMetadataForTests::user_holding_b(),
+        AccountWithMetadataForTests::user_holding_lp_with_balance(lp_balance),
+        NonZero::new(BalanceForTests::remove_amount_lp()).unwrap(),
+        BalanceForTests::remove_min_amount_a(),
+        BalanceForTests::remove_min_amount_b_low(),
+    );
+}
+
 #[should_panic(expected = "User Token A holding must be owned by the vault's Token Program")]
 #[test]
 fn test_swap_exact_input_rejects_user_holding_a_wrong_program() {