return surplus of native token on liquidity stable mint

Cardinal-Cryptography · Sep 24, 2024 · f8a20ed · f8a20ed
1 parent ea7cc44
commit f8a20ed
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/amm/contracts/router_v2/pair.rs b/amm/contracts/router_v2/pair.rs
@@ -63,6 +63,8 @@ impl Pair {
 
     /// Makes a cross-contract call to fetch the Pair reserves.
     /// Returns reserves `(reserve_0, reserve_1)` in order of `token_0` and `token_1`
+    /// NOTE: before calling this method ensure that `token_0` and `token_1` belong to
+    /// this `Pair` pool
     fn get_reserves(&self, token_0: &AccountId, token_1: &AccountId) -> (u128, u128) {
         let (reserve_0, reserve_1, _) = self.contract_ref().get_reserves();
         if token_0 < token_1 {

diff --git a/amm/contracts/router_v2/stable_pool.rs b/amm/contracts/router_v2/stable_pool.rs
@@ -73,23 +73,23 @@ impl StablePool {
             self.tokens.len() == amounts.len(),
             RouterV2Error::StablePoolError(StablePoolError::IncorrectAmountsCount)
         );
-        let wnative_idx = match wnative {
+        let native_received = transferred_value::<Env>();
+        let (wnative_idx, native_surplus) = match wnative {
             Some(wnative) => {
                 let wnative_idx = self.wnative_idx(wnative)?;
-                let native_received = transferred_value::<Env>();
                 let wnative_amount = amounts[wnative_idx];
                 ensure!(
                     native_received >= wnative_amount,
                     RouterV2Error::InsufficientTransferredAmount
                 );
                 wrap(wnative, wnative_amount)?;
-                if native_received > wnative_amount {
-                    transfer_native(caller::<Env>(), native_received - wnative_amount)?;
-                }
-                wnative_idx
+                (wnative_idx, native_received.saturating_sub(wnative_amount))
             }
-            None => self.tokens.len(),
+            None => (self.tokens.len(), native_received),
         };
+        if native_surplus > 0 {
+            transfer_native(caller::<Env>(), native_surplus)?;
+        }
         for i in (0..self.tokens.len()).filter(|&idx| idx != wnative_idx) {
             psp22_transfer_from(
                 self.tokens[i],