From 1e4e02b993bc35c9da7168a1b5298aacbeb76e70 Mon Sep 17 00:00:00 2001
From: Foivos <foivos@umich.edu>
Date: Tue, 31 Oct 2023 14:59:38 +0200
Subject: [PATCH] feat: fixed some potential attack for the fee token manager
 (#133)

* Fixed some potential attack for the fee token manager

* Update contracts/token-manager/implementations/TokenManagerLockUnlockFee.sol

---------

Co-authored-by: Milap Sheth <milap@axelar.network>
---
 .../implementations/TokenManagerLockUnlockFee.sol  | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/contracts/token-manager/implementations/TokenManagerLockUnlockFee.sol b/contracts/token-manager/implementations/TokenManagerLockUnlockFee.sol
index 606b4b4e..cc6d174e 100644
--- a/contracts/token-manager/implementations/TokenManagerLockUnlockFee.sol
+++ b/contracts/token-manager/implementations/TokenManagerLockUnlockFee.sol
@@ -52,7 +52,12 @@ contract TokenManagerLockUnlockFee is TokenManager, NoReEntrancy, ITokenManagerL
 
         token.safeTransferFrom(from, address(this), amount);
 
-        return token.balanceOf(address(this)) - balanceBefore;
+        uint256 diff = token.balanceOf(address(this)) - balanceBefore;
+        if (diff < amount) {
+            amount = diff;
+        }
+
+        return amount;
     }
 
     /**
@@ -67,7 +72,12 @@ contract TokenManagerLockUnlockFee is TokenManager, NoReEntrancy, ITokenManagerL
 
         token.safeTransfer(to, amount);
 
-        return token.balanceOf(to) - balanceBefore;
+        uint256 diff = token.balanceOf(to) - balanceBefore;
+        if (diff < amount) {
+            amount = diff;
+        }
+
+        return amount;
     }
 
     /**