From b19bfe604bd16e23706ff960bb040f502f7ece0f Mon Sep 17 00:00:00 2001 From: 0xchin Date: Mon, 18 Nov 2024 22:15:37 -0300 Subject: [PATCH] fix: enhance security with reentrancy guard and constants --- src/contracts/Grateful.sol | 70 +++++++++++++++++++++++------------- src/interfaces/IGrateful.sol | 14 ++++++++ 2 files changed, 60 insertions(+), 24 deletions(-) diff --git a/src/contracts/Grateful.sol b/src/contracts/Grateful.sol index 443cd09..d5f2167 100644 --- a/src/contracts/Grateful.sol +++ b/src/contracts/Grateful.sol @@ -6,6 +6,7 @@ import {Ownable2Step} from "@openzeppelin/contracts/access/Ownable2Step.sol"; import {IERC20} from "@openzeppelin/contracts/interfaces/IERC20.sol"; import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; +import {ReentrancyGuard} from "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; import {OneTime} from "contracts/OneTime.sol"; import {AaveV3Vault} from "contracts/vaults/AaveV3Vault.sol"; @@ -20,11 +21,20 @@ import {IPool} from "yield-daddy/aave-v3/AaveV3ERC4626.sol"; * @title Grateful Contract * @notice Allows payments in whitelisted tokens with optional yield via AAVE */ -contract Grateful is IGrateful, Ownable2Step { +contract Grateful is IGrateful, Ownable2Step, ReentrancyGuard { using Bytes32AddressLib for bytes32; using FixedPointMathLib for uint256; using SafeERC20 for IERC20; + /*////////////////////////////////////////////////////////////// + CONSTANTS + //////////////////////////////////////////////////////////////*/ + /// @inheritdoc IGrateful + uint256 public constant MAX_FEE = 10_000; // Max 100% fee (10000 basis points) + + /// @inheritdoc IGrateful + uint256 public constant MAX_PERFORMANCE_FEE = 5000; // Max 50% performance fee (5000 basis points) + /*////////////////////////////////////////////////////////////// STATE VARIABLES //////////////////////////////////////////////////////////////*/ @@ -94,7 +104,7 @@ contract Grateful is IGrateful, Ownable2Step { fee = _initialFee; for (uint256 i = 0; i < _tokens.length; i++) { tokensWhitelisted[_tokens[i]] = true; - IERC20 token = IERC20(_tokens[i]); // Renamed from _token to token + IERC20 token = IERC20(_tokens[i]); token.forceApprove(address(_aavePool), type(uint256).max); } } @@ -105,8 +115,8 @@ contract Grateful is IGrateful, Ownable2Step { /// @inheritdoc IGrateful function calculateAssets(address _merchant, address _token) public view returns (uint256 assets) { - AaveV3Vault vault = vaults[_token]; // Renamed from _vault to vault - uint256 sharesAmount = shares[_merchant][_token]; // Renamed from _shares to sharesAmount + AaveV3Vault vault = vaults[_token]; + uint256 sharesAmount = shares[_merchant][_token]; assets = vault.convertToAssets(sharesAmount); } @@ -137,8 +147,8 @@ contract Grateful is IGrateful, Ownable2Step { /// @inheritdoc IGrateful function calculateProfit(address _user, address _token) public view returns (uint256 profit) { - AaveV3Vault vault = vaults[_token]; // Renamed from _vault to vault - uint256 sharesAmount = shares[_user][_token]; // Renamed from _shares to sharesAmount + AaveV3Vault vault = vaults[_token]; + uint256 sharesAmount = shares[_user][_token]; uint256 assets = vault.previewRedeem(sharesAmount); // Current value of user's shares uint256 initialDeposit = userDeposits[_user][_token]; // User's initial deposit if (assets > initialDeposit) { @@ -164,7 +174,7 @@ contract Grateful is IGrateful, Ownable2Step { address _token ) external onlyOwner { tokensWhitelisted[_token] = true; - IERC20 token = IERC20(_token); // Renamed from _token to token + IERC20 token = IERC20(_token); token.forceApprove(address(aavePool), type(uint256).max); emit TokenAdded(_token); } @@ -177,7 +187,7 @@ contract Grateful is IGrateful, Ownable2Step { revert Grateful_TokenOrVaultNotFound(); } delete tokensWhitelisted[_token]; - IERC20 token = IERC20(_token); // Renamed from _token to token + IERC20 token = IERC20(_token); token.forceApprove(address(aavePool), 0); token.forceApprove(address(vaults[_token]), 0); emit TokenRemoved(_token); @@ -186,7 +196,7 @@ contract Grateful is IGrateful, Ownable2Step { /// @inheritdoc IGrateful function addVault(address _token, address _vault) external onlyOwner onlyWhenTokenWhitelisted(_token) { vaults[_token] = AaveV3Vault(_vault); - IERC20 token = IERC20(_token); // Renamed from _token to token + IERC20 token = IERC20(_token); token.safeIncreaseAllowance(address(_vault), type(uint256).max); emit VaultAdded(_token, _vault); } @@ -199,7 +209,7 @@ contract Grateful is IGrateful, Ownable2Step { if (address(vault) == address(0)) { revert Grateful_TokenOrVaultNotFound(); } - IERC20 token = IERC20(_token); // Renamed from _token to token + IERC20 token = IERC20(_token); token.forceApprove(address(vault), 0); emit VaultRemoved(_token, address(vault)); delete vaults[_token]; @@ -304,6 +314,9 @@ contract Grateful is IGrateful, Ownable2Step { function setFee( uint256 _newFee ) external onlyOwner { + if (_newFee > MAX_FEE) { + revert Grateful_FeeRateTooHigh(); + } fee = _newFee; emit FeeUpdated(_newFee); } @@ -312,7 +325,7 @@ contract Grateful is IGrateful, Ownable2Step { function setPerformanceFeeRate( uint256 _newPerformanceFeeRate ) external onlyOwner { - if (_newPerformanceFeeRate > 5000) { + if (_newPerformanceFeeRate > MAX_PERFORMANCE_FEE) { revert Grateful_FeeRateTooHigh(); } performanceFeeRate = _newPerformanceFeeRate; @@ -321,6 +334,9 @@ contract Grateful is IGrateful, Ownable2Step { /// @inheritdoc IGrateful function setCustomFee(uint256 _newFee, address _merchant) external onlyOwner { + if (_newFee > MAX_FEE) { + revert Grateful_FeeRateTooHigh(); + } customFees[_merchant] = CustomFee({isSet: true, fee: _newFee}); emit CustomFeeUpdated(_merchant, _newFee); } @@ -365,23 +381,23 @@ contract Grateful is IGrateful, Ownable2Step { uint256 _amount, uint256 _paymentId, bool _yieldFunds - ) private { - // Transfer the full amount from the sender to this contract - IERC20 token = IERC20(_token); // Renamed from _token to token - token.safeTransferFrom(_sender, address(this), _amount); - + ) private nonReentrant { // Check payment id if (paymentIds[_paymentId]) { revert Grateful_PaymentIdAlreadyUsed(); } - paymentIds[_paymentId] = true; // Apply the fee uint256 amountWithFee = applyFee(_merchant, _amount); + uint256 feeAmount = _amount - amountWithFee; + + IERC20 token = IERC20(_token); + + // Transfer the full amount from the sender to this contract + token.safeTransferFrom(_sender, address(this), _amount); // Transfer fee to owner - uint256 feeAmount = _amount - amountWithFee; token.safeTransfer(owner(), feeAmount); if (_yieldFunds) { @@ -390,6 +406,7 @@ contract Grateful is IGrateful, Ownable2Step { token.safeTransfer(_merchant, amountWithFee); } else { uint256 sharesAmount = vault.deposit(amountWithFee, address(this)); + // Update state after receiving sharesAmount shares[_merchant][_token] += sharesAmount; userDeposits[_merchant][_token] += amountWithFee; } @@ -402,12 +419,12 @@ contract Grateful is IGrateful, Ownable2Step { } /** - * @dev Internal function to handle withdrawals. + * @notice Handles a withdrawal. * @param _token The address of the token to withdraw. * @param _assets The amount of assets to withdraw (ignored if full withdrawal). * @param _isFullWithdrawal Indicates if it's a full withdrawal. */ - function _withdraw(address _token, uint256 _assets, bool _isFullWithdrawal) internal { + function _withdraw(address _token, uint256 _assets, bool _isFullWithdrawal) private nonReentrant { AaveV3Vault vault = vaults[_token]; if (address(vault) == address(0)) { revert Grateful_VaultNotSet(); @@ -442,12 +459,9 @@ contract Grateful is IGrateful, Ownable2Step { profit = assetsToWithdraw - initialDepositToWithdraw; performanceFeeAmount = calculatePerformanceFee(profit); assetsToWithdraw -= performanceFeeAmount; // Deduct fee from assets - - // Withdraw performance fee to fee recipient (owner) - vault.withdraw(performanceFeeAmount, owner(), address(this)); } - // Update user's shares and deposits + // Update user's shares and deposits before external calls shares[msg.sender][_token] = totalShares - sharesToWithdraw; userDeposits[msg.sender][_token] = initialDeposit - initialDepositToWithdraw; @@ -457,7 +471,15 @@ contract Grateful is IGrateful, Ownable2Step { userDeposits[msg.sender][_token] = 0; } + // Withdraw performance fee to fee recipient (owner) + if (performanceFeeAmount > 0) { + vault.withdraw(performanceFeeAmount, owner(), address(this)); + } + // Withdraw assets to user vault.withdraw(assetsToWithdraw, msg.sender, address(this)); + + // Emit an event for the withdrawal + emit Withdrawal(msg.sender, _token, assetsToWithdraw, performanceFeeAmount); } } diff --git a/src/interfaces/IGrateful.sol b/src/interfaces/IGrateful.sol index 8058991..bf2191d 100644 --- a/src/interfaces/IGrateful.sol +++ b/src/interfaces/IGrateful.sol @@ -110,6 +110,15 @@ interface IGrateful { */ event VaultRemoved(address indexed token, address indexed vault); + /** + * @notice Emitted when a withdrawal is made. + * @param user Address of the user making the withdrawal. + * @param token Address of the token being withdrawn. + * @param amount Amount of the token being withdrawn. + * @param performanceFee Amount of the performance fee deducted. + */ + event Withdrawal(address indexed user, address indexed token, uint256 amount, uint256 performanceFee); + /*/////////////////////////////////////////////////////////////// ERRORS //////////////////////////////////////////////////////////////*/ @@ -140,6 +149,11 @@ interface IGrateful { /*/////////////////////////////////////////////////////////////// VARIABLES //////////////////////////////////////////////////////////////*/ + /// @notice Returns the maximum fee in basis points (10000 = 100%). + function MAX_FEE() external pure returns (uint256); + + /// @notice Returns the maximum performance fee in basis points (5000 = 50%). + function MAX_PERFORMANCE_FEE() external pure returns (uint256); /// @notice Returns the owner of the contract. function owner() external view returns (address);