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 <[email protected]>
axelarnetwork · Oct 31, 2023 · 1e4e02b · 1e4e02b
1 parent ed9856a
commit 1e4e02b
Showing 1 changed file with 12 additions and 2 deletions.
diff --git 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;
     }
 
     /**