feat: use SafeERC20 for tokens

Grateful-protocol · Nov 1, 2024 · c0c6448 · c0c6448
1 parent e628544
commit c0c6448
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/src/contracts/OneTime.sol b/src/contracts/OneTime.sol
@@ -2,11 +2,15 @@
 pragma solidity 0.8.26;
 
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+
+import {IERC20} from "@openzeppelin/contracts/interfaces/IERC20.sol";
+import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import {console} from "forge-std/console.sol";
-import {IERC20} from "forge-std/interfaces/IERC20.sol";
 import {IGrateful} from "interfaces/IGrateful.sol";
 
 contract OneTime {
+  using SafeERC20 for IERC20;
+
   IGrateful immutable grateful;
 
   constructor(
@@ -23,17 +27,17 @@ contract OneTime {
     for (uint256 i = 0; i < _tokens.length; i++) {
       IERC20 token = IERC20(_tokens[i]);
       if (token.balanceOf(address(this)) >= _AMOUNT_USDC) {
-        token.approve(address(_grateful), _AMOUNT_USDC);
+        token.safeIncreaseAllowance(address(_grateful), _AMOUNT_USDC);
         _grateful.receiveOneTimePayment(_merchant, address(token), _paymentId, _AMOUNT_USDC, _recipients, _percentages);
       }
     }
   }
 
-  function rescueFunds(address _token, address _receiver, uint256 _AMOUNT_USDC) external {
+  function rescueFunds(IERC20 _token, address _receiver, uint256 _AMOUNT_USDC) external {
     if (msg.sender != grateful.owner()) {
       revert Ownable.OwnableUnauthorizedAccount(msg.sender);
     }
 
-    IERC20(_token).transfer(_receiver, _AMOUNT_USDC);
+    IERC20(_token).safeTransfer(_receiver, _AMOUNT_USDC);
   }
 }
diff --git a/test/integration/Grateful.t.sol b/test/integration/Grateful.t.sol
@@ -81,7 +81,7 @@ contract IntegrationGreeter is IntegrationBase {
 
     // Rescue funds
     vm.prank(_owner);
-    _oneTime.rescueFunds(address(_usdc), _payer, _AMOUNT_USDC);
+    _oneTime.rescueFunds(_usdc, _payer, _AMOUNT_USDC);
 
     // Client has received his funds
     assertEq(_usdc.balanceOf(address(_payer)), prevWhaleBalance + _AMOUNT_USDC);

diff --git a/test/integration/IntegrationBase.sol b/test/integration/IntegrationBase.sol
@@ -4,7 +4,7 @@ pragma solidity 0.8.26;
 import {Grateful, IGrateful} from "contracts/Grateful.sol";
 import {Test} from "forge-std/Test.sol";
 
-import {IERC20} from "forge-std/interfaces/IERC20.sol";
+import {IERC20} from "@openzeppelin/contracts/interfaces/IERC20.sol";
 
 import {AaveV3Vault} from "contracts/vaults/AaveV3Vault.sol";
 import {ERC20} from "solmate/tokens/ERC20.sol";