From 423d8b1ae947976fbd8a409e7af206057c195a5e Mon Sep 17 00:00:00 2001
From: Jon Cinque <jon.cinque@gmail.com>
Date: Fri, 9 Oct 2020 19:23:43 +0200
Subject: [PATCH] Add checks for 0 token output and tests (#595)

---
 token-swap/program/src/curve.rs     | 23 ++++++--
 token-swap/program/src/processor.rs | 86 +++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+), 5 deletions(-)

diff --git a/token-swap/program/src/curve.rs b/token-swap/program/src/curve.rs
index 8d7d7d44a67d89..76e9006dc68c0b 100644
--- a/token-swap/program/src/curve.rs
+++ b/token-swap/program/src/curve.rs
@@ -29,9 +29,12 @@ impl SwapResult {
         let invariant = swap_source_amount.checked_mul(swap_destination_amount)?;
 
         // debit the fee to calculate the amount swapped
-        let fee = source_amount
+        let mut fee = source_amount
             .checked_mul(fee_numerator)?
             .checked_div(fee_denominator)?;
+        if fee == 0 {
+            fee = 1; // minimum fee of one token
+        }
         let new_source_amount_less_fee = swap_source_amount
             .checked_add(source_amount)?
             .checked_sub(fee)?;
@@ -48,6 +51,14 @@ impl SwapResult {
     }
 }
 
+fn map_zero_to_none(x: u64) -> Option<u64> {
+    if x == 0 {
+        None
+    } else {
+        Some(x)
+    }
+}
+
 /// The Uniswap invariant calculator.
 pub struct ConstantProduct {
     /// Token A
@@ -72,7 +83,7 @@ impl ConstantProduct {
         )?;
         self.token_a = result.new_source_amount;
         self.token_b = result.new_destination_amount;
-        Some(result.amount_swapped)
+        map_zero_to_none(result.amount_swapped)
     }
 
     /// Swap token b to a
@@ -86,7 +97,7 @@ impl ConstantProduct {
         )?;
         self.token_b = result.new_source_amount;
         self.token_a = result.new_destination_amount;
-        Some(result.amount_swapped)
+        map_zero_to_none(result.amount_swapped)
     }
 }
 
@@ -123,18 +134,20 @@ impl PoolTokenConverter {
         }
     }
 
-    /// A tokens for pool tokens
+    /// A tokens for pool tokens, returns None if output is less than 0
     pub fn token_a_rate(&self, pool_tokens: u64) -> Option<u64> {
         pool_tokens
             .checked_mul(self.token_a)?
             .checked_div(self.supply)
+            .and_then(map_zero_to_none)
     }
 
-    /// B tokens for pool tokens
+    /// B tokens for pool tokens, returns None is output is less than 0
     pub fn token_b_rate(&self, pool_tokens: u64) -> Option<u64> {
         pool_tokens
             .checked_mul(self.token_b)?
             .checked_div(self.supply)
+            .and_then(map_zero_to_none)
     }
 }
 
diff --git a/token-swap/program/src/processor.rs b/token-swap/program/src/processor.rs
index 08f2cf172116ac..e71f9d5ae1df2b 100644
--- a/token-swap/program/src/processor.rs
+++ b/token-swap/program/src/processor.rs
@@ -1938,6 +1938,33 @@ mod tests {
             accounts.pool_mint_account = old_pool_account;
         }
 
+        // deposit 1 pool token fails beacuse it equates to 0 swap tokens
+        {
+            let (
+                token_a_key,
+                mut token_a_account,
+                token_b_key,
+                mut token_b_account,
+                pool_key,
+                mut pool_account,
+            ) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
+            assert_eq!(
+                Err(SwapError::CalculationFailure.into()),
+                accounts.deposit(
+                    &depositor_key,
+                    &pool_key,
+                    &mut pool_account,
+                    &token_a_key,
+                    &mut token_a_account,
+                    &token_b_key,
+                    &mut token_b_account,
+                    1,
+                    deposit_a,
+                    deposit_b / 10,
+                )
+            );
+        }
+
         // slippage exceeeded
         {
             let (
@@ -2435,6 +2462,39 @@ mod tests {
             accounts.pool_mint_account = old_pool_account;
         }
 
+        // withdrawing 1 pool token fails because it equates to 0 output tokens
+        {
+            let (
+                token_a_key,
+                mut token_a_account,
+                token_b_key,
+                mut token_b_account,
+                pool_key,
+                mut pool_account,
+            ) = accounts.setup_token_accounts(
+                &user_key,
+                &withdrawer_key,
+                initial_a,
+                initial_b,
+                initial_pool,
+            );
+            assert_eq!(
+                Err(SwapError::CalculationFailure.into()),
+                accounts.withdraw(
+                    &withdrawer_key,
+                    &pool_key,
+                    &mut pool_account,
+                    &token_a_key,
+                    &mut token_a_account,
+                    &token_b_key,
+                    &mut token_b_account,
+                    1,
+                    minimum_a_amount * 10,
+                    minimum_b_amount,
+                )
+            );
+        }
+
         // slippage exceeeded
         {
             let (
@@ -2820,6 +2880,32 @@ mod tests {
             );
         }
 
+        // output token value 0
+        {
+            let (
+                token_a_key,
+                mut token_a_account,
+                token_b_key,
+                mut token_b_account,
+                _pool_key,
+                _pool_account,
+            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
+            assert_eq!(
+                Err(SwapError::CalculationFailure.into()),
+                accounts.swap(
+                    &swapper_key,
+                    &token_b_key,
+                    &mut token_b_account,
+                    &swap_token_b_key,
+                    &swap_token_a_key,
+                    &token_a_key,
+                    &mut token_a_account,
+                    1,
+                    1,
+                )
+            );
+        }
+
         // slippage exceeeded: minimum out amount too high
         {
             let (