Enforce reserve amounts amounts strictly positive

liquidity_pool/src/lib.rs

@@ -217,13 +217,20 @@ impl LiquidityPoolTrait for LiquidityPool {
         let (reserve_a, reserve_b) = (get_reserve_a(&e), get_reserve_b(&e));
 
         // Calculate deposit amounts
-        let amounts = get_deposit_amounts(desired_a, min_a, desired_b, min_b, reserve_a, reserve_b);
+        let (amount_a, amount_b) =
+            get_deposit_amounts(desired_a, min_a, desired_b, min_b, reserve_a, reserve_b);
+
+        if amount_a <= 0 || amount_b <= 0 {
+            // If one of the amounts can be zero, we can get into a situation
+            // where one of the reserves is 0, which leads to a divide by zero.
+            panic!("both amounts must be strictly positive");
+        }
 
         let token_a_client = token::Client::new(&e, &get_token_a(&e));
         let token_b_client = token::Client::new(&e, &get_token_b(&e));
 
-        token_a_client.transfer(&to, &e.current_contract_address(), &amounts.0);
-        token_b_client.transfer(&to, &e.current_contract_address(), &amounts.1);
+        token_a_client.transfer(&to, &e.current_contract_address(), &amount_a);
+        token_b_client.transfer(&to, &e.current_contract_address(), &amount_b);
 
         // Now calculate how many new pool shares to mint
         let (balance_a, balance_b) = (get_balance_a(&e), get_balance_b(&e));
@@ -305,8 +312,15 @@ impl LiquidityPoolTrait for LiquidityPool {
             transfer_b(&e, to, out_b);
         }
 
-        put_reserve_a(&e, balance_a - out_a);
-        put_reserve_b(&e, balance_b - out_b);
+        let new_reserve_a = balance_a - out_a;
+        let new_reserve_b = balance_b - out_b;
+
+        if new_reserve_a <= 0 || new_reserve_b <= 0 {
+            panic!("new reserves must be strictly positive");
+        }
+
+        put_reserve_a(&e, new_reserve_a);
+        put_reserve_b(&e, new_reserve_b);
     }
 
     fn withdraw(e: Env, to: Address, share_amount: i128, min_a: i128, min_b: i128) -> (i128, i128) {

liquidity_pool/src/test.rs

@@ -161,3 +161,76 @@ fn test() {
     assert_eq!(token2.balance(&liqpool.address), 0);
     assert_eq!(token_share.balance(&liqpool.address), 0);
 }
+
+#[test]
+#[should_panic]
+fn deposit_amount_zero_should_panic() {
+    let e = Env::default();
+    e.mock_all_auths();
+
+    // Create contracts
+    let mut admin1 = Address::generate(&e);
+    let mut admin2 = Address::generate(&e);
+
+    let mut token_a = create_token_contract(&e, &admin1);
+    let mut token_b = create_token_contract(&e, &admin2);
+    if &token_b.address < &token_a.address {
+        std::mem::swap(&mut token_a, &mut token_b);
+        std::mem::swap(&mut admin1, &mut admin2);
+    }
+    let liqpool = create_liqpool_contract(
+        &e,
+        &install_token_wasm(&e),
+        &token_a.address,
+        &token_b.address,
+    );
+
+    // Create a user
+    let user1 = Address::generate(&e);
+
+    token_a.mint(&user1, &1000);
+    assert_eq!(token_a.balance(&user1), 1000);
+
+    token_b.mint(&user1, &1000);
+    assert_eq!(token_b.balance(&user1), 1000);
+
+    liqpool.deposit(&user1, &1, &0, &0, &0);
+}
+
+#[test]
+#[should_panic]
+fn swap_reserve_one_nonzero_other_zero() {
+    let e = Env::default();
+    e.mock_all_auths();
+
+    // Create contracts
+    let mut admin1 = Address::generate(&e);
+    let mut admin2 = Address::generate(&e);
+
+    let mut token_a = create_token_contract(&e, &admin1);
+    let mut token_b = create_token_contract(&e, &admin2);
+    if &token_b.address < &token_a.address {
+        std::mem::swap(&mut token_a, &mut token_b);
+        std::mem::swap(&mut admin1, &mut admin2);
+    }
+    let liqpool = create_liqpool_contract(
+        &e,
+        &install_token_wasm(&e),
+        &token_a.address,
+        &token_b.address,
+    );
+
+    // Create a user
+    let user1 = Address::generate(&e);
+
+    token_a.mint(&user1, &1000);
+    assert_eq!(token_a.balance(&user1), 1000);
+
+    token_b.mint(&user1, &1000);
+    assert_eq!(token_b.balance(&user1), 1000);
+
+    // Try to get to a situation where the reserves are 1 and 0.
+    // It shouldn't be possible.
+    token_b.transfer(&user1, &liqpool.address, &1);
+    liqpool.swap(&user1, &false, &1, &1);
+}