From 2dafd4a3ed3b52a1895723b3b178461b2d9a289d Mon Sep 17 00:00:00 2001
From: woocash2 <lukasgniecki@gmail.com>
Date: Fri, 12 Jul 2024 10:35:58 +0200
Subject: [PATCH 1/2] try exploit test

---
 .../stable_swap_tests/tests_add_remove_lp.rs  | 87 +++++++++++++++++++
 1 file changed, 87 insertions(+)

diff --git a/amm/drink-tests/src/stable_swap_tests/tests_add_remove_lp.rs b/amm/drink-tests/src/stable_swap_tests/tests_add_remove_lp.rs
index 1ebb9ae..a3b5476 100644
--- a/amm/drink-tests/src/stable_swap_tests/tests_add_remove_lp.rs
+++ b/amm/drink-tests/src/stable_swap_tests/tests_add_remove_lp.rs
@@ -841,3 +841,90 @@ fn test_07(mut session: Session) {
         "Incorrect Users tokens balances"
     );
 }
+
+#[drink::test]
+fn test_lp_withdraw_all_but_no_more() {
+    seed_account(&mut session, CHARLIE);
+    seed_account(&mut session, DAVE);
+    seed_account(&mut session, EVA);
+
+    let charlie_input = vec![1_234_543 * ONE_USDT, 1_112_323 * ONE_USDC];
+    let dave_input = vec![1131 * 105 * ONE_USDT, 1157 * 105 * ONE_USDC];
+    let initial_supply: Vec<u128> = charlie_input.iter().map(|amount| amount * 10).collect();
+
+    let (stable_swap, tokens) = setup_stable_swap_with_tokens(
+        &mut session,
+        vec![6, 6],
+        initial_supply.clone(),
+        10_000,
+        0,
+        0,
+        BOB,
+        vec![],
+    );
+
+    transfer_and_increase_allowance(
+        &mut session,
+        stable_swap,
+        tokens.clone(),
+        CHARLIE,
+        vec![1_434_543 * ONE_USDT, 1_112_323 * ONE_USDC],
+        BOB,
+    );
+    transfer_and_increase_allowance(
+        &mut session,
+        stable_swap,
+        tokens,
+        DAVE,
+        vec![1131 * 105 * ONE_USDT, 1157 * 105 * ONE_USDC],
+        BOB,
+    );
+
+    _ = stable_swap::add_liquidity(
+        &mut session,
+        stable_swap,
+        CHARLIE,
+        1,
+        charlie_input.clone(),
+        charlie(),
+    )
+    .expect("Charlie should successfully add liquidity");
+
+    let (shares, _) = stable_swap::add_liquidity(
+        &mut session,
+        stable_swap,
+        DAVE,
+        1,
+        dave_input.clone(),
+        dave(),
+    )
+    .expect("Dave should successfully add liquidity");
+
+    for _ in 0..105 {
+        let withdraw = vec![1131 * ONE_USDT, 1157 * ONE_USDC]; // 1/105 of Dave's
+        _ = stable_swap::remove_liquidity_by_amounts(
+            &mut session,
+            stable_swap,
+            DAVE,
+            shares, // just an upper bound
+            withdraw,
+            dave(),
+        )
+        .expect("Should successfully remove liquidity");
+    }
+
+    assert_eq!(
+        psp22_utils::balance_of(&mut session, stable_swap, dave()),
+        0
+    );
+
+    _ = stable_swap::remove_liquidity_by_shares(
+        &mut session,
+        stable_swap,
+        DAVE,
+        10,
+        vec![1, 1], // at least withdraw something
+        dave(),
+    )
+    .expect_err("Should not successfully remove liquidity");
+}

From 62b4a243e3ef16ee943f0b77289ff0be6cc456e2 Mon Sep 17 00:00:00 2001
From: woocash2 <lukasgniecki@gmail.com>
Date: Fri, 12 Jul 2024 10:43:49 +0200
Subject: [PATCH 2/2] add some comments

---
 .../src/stable_swap_tests/tests_add_remove_lp.rs          | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/amm/drink-tests/src/stable_swap_tests/tests_add_remove_lp.rs b/amm/drink-tests/src/stable_swap_tests/tests_add_remove_lp.rs
index a3b5476..bf94baf 100644
--- a/amm/drink-tests/src/stable_swap_tests/tests_add_remove_lp.rs
+++ b/amm/drink-tests/src/stable_swap_tests/tests_add_remove_lp.rs
@@ -848,8 +848,8 @@ fn test_lp_withdraw_all_but_no_more() {
     seed_account(&mut session, DAVE);
     seed_account(&mut session, EVA);
 
-    let charlie_input = vec![1_234_543 * ONE_USDT, 1_112_323 * ONE_USDC];
-    let dave_input = vec![1131 * 105 * ONE_USDT, 1157 * 105 * ONE_USDC];
+    let charlie_input = vec![1_434_543 * ONE_USDT, 1_112_323 * ONE_USDC];
+    let dave_input = vec![1131 * 105 * ONE_USDT, 1157 * 105 * ONE_USDC]; // treat as 105 parts of (1131, 1157)
     let initial_supply: Vec<u128> = charlie_input.iter().map(|amount| amount * 10).collect();
 
     let (stable_swap, tokens) = setup_stable_swap_with_tokens(
@@ -900,6 +900,7 @@ fn test_lp_withdraw_all_but_no_more() {
     )
     .expect("Dave should successfully add liquidity");
 
+    // 105 times withdraw (1131, 1157) which should withdraw the whole Dave's share
     for _ in 0..105 {
         let withdraw = vec![1131 * ONE_USDT, 1157 * ONE_USDC]; // 1/105 of Dave's
         _ = stable_swap::remove_liquidity_by_amounts(
@@ -915,7 +916,8 @@ fn test_lp_withdraw_all_but_no_more() {
 
     assert_eq!(
         psp22_utils::balance_of(&mut session, stable_swap, dave()),
-        0
+        0,
+        "Dave shouldn't have any LPT left",
     );
 
     _ = stable_swap::remove_liquidity_by_shares(