diff --git a/src/contracts/OneTime.sol b/src/contracts/OneTime.sol index 6085013..e385767 100644 --- 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 index a60faf5..81d281a 100644 --- 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 index c936538..fa2f352 100644 --- 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";