From e87f25180d0d0e65df4b2bc495652e9f50cc2cdd Mon Sep 17 00:00:00 2001
From: bayge <alex@fluidity.money>
Date: Mon, 23 Sep 2024 21:29:14 +0930
Subject: [PATCH] Codearena findings 12, If liquidity is insufficient, users
 will need to pay more tokens

---
 pkg/seawater/src/error.rs |  4 ++++
 pkg/seawater/src/lib.rs   | 11 ++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/pkg/seawater/src/error.rs b/pkg/seawater/src/error.rs
index c475c6d2..c56a01b4 100644
--- a/pkg/seawater/src/error.rs
+++ b/pkg/seawater/src/error.rs
@@ -241,6 +241,10 @@ pub enum Error {
     // 43 (0x2b)
     #[error("Trying to swap2 between the same pool")]
     SamePool,
+
+    // 44 (0x2c)
+    #[error("Swap sending fee taking overflow")]
+    TransferToSenderSub,
 }
 
 impl From<Error> for Vec<u8> {
diff --git a/pkg/seawater/src/lib.rs b/pkg/seawater/src/lib.rs
index 7d9262e8..d55292f6 100644
--- a/pkg/seawater/src/lib.rs
+++ b/pkg/seawater/src/lib.rs
@@ -287,9 +287,18 @@ impl Pools {
         ));
 
         // transfer tokens
-        erc20::take(from, original_amount, permit2)?;
+        erc20::take(from, amount_in, permit2)?;
         erc20::transfer_to_sender(to, amount_out)?;
 
+        if original_amount > amount_in {
+            erc20::transfer_to_sender(
+                to,
+                original_amount
+                    .checked_sub(amount_in)
+                    .ok_or(Error::TransferToSenderSub)?,
+            )?;
+        }
+
         #[cfg(feature = "log-events")]
         evm::log(events::Swap2 {
             user: msg::sender(),