From f8432e7fb9922e81d985c721fc0929903958401d Mon Sep 17 00:00:00 2001 From: deo Date: Wed, 20 Dec 2023 03:04:59 -0500 Subject: [PATCH] feat: 5.1 done --- solidity/contracts/Flashloan/Flashloan.sol | 75 +++++++++++++++++++ .../contracts/Flashloan/GreedyReceiver.sol | 30 ++++++++ solidity/contracts/Flashloan/Receiver.sol | 35 +++++++++ 3 files changed, 140 insertions(+) create mode 100644 solidity/contracts/Flashloan/Flashloan.sol create mode 100644 solidity/contracts/Flashloan/GreedyReceiver.sol create mode 100644 solidity/contracts/Flashloan/Receiver.sol diff --git a/solidity/contracts/Flashloan/Flashloan.sol b/solidity/contracts/Flashloan/Flashloan.sol new file mode 100644 index 00000000..91d8ee81 --- /dev/null +++ b/solidity/contracts/Flashloan/Flashloan.sol @@ -0,0 +1,75 @@ +pragma solidity ^0.8.19; + +import {IERC20} from '@openzeppelin/contracts/token/ERC20/IERC20.sol'; +import {SafeERC20} from '@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol'; +import {IReceiver} from '../../interfaces/Flashloan/IReceiver.sol'; +import {IFlashloan} from '../../interfaces/Flashloan/IFlashloan.sol'; +import {ReentrancyGuard} from '@openzeppelin/contracts/utils/ReentrancyGuard.sol'; + +// Without reentrancy guard, user can make a flashloan attack by using his flashloan to deposit into the pool +contract Flashloan is IFlashloan, ReentrancyGuard { + using SafeERC20 for IERC20; + + mapping(address _user => mapping(address _token => uint256 _amount)) tokenBalances; + mapping(address _user => uint256 _amount) ethBalances; + + constructor() {} + + /// @dev allows users to flashloan any tokens or eth deposited in the pool + /// @dev If address(0) is passed as _token, this is treated as ETH + /// @param _amount flashloan total desired, must not exceed pool balance + /// @param _token token (or eth) to be flashloaned + function flashloan(uint256 _amount, address _token) external nonReentrant { + if (address(_token) == address(0)) { + uint256 poolBalance = address(this).balance; + if (_amount > poolBalance) revert InsufficientPoolBalance(); + + IReceiver(msg.sender).getETH{value: _amount}(); + + uint256 balanceAfter = address(this).balance; + + if (balanceAfter < poolBalance + (poolBalance * 1e13) / 1e18) revert InsufficientRepayment(); + } else { + uint256 poolBalance = IERC20(_token).balanceOf(address(this)); + if (_amount > poolBalance) revert InsufficientPoolBalance(); + + IERC20(_token).safeTransfer(msg.sender, _amount); + IReceiver(msg.sender).getTokens(_token, _amount); + + uint256 balanceAfter = IERC20(_token).balanceOf(address(this)); + + if (balanceAfter < poolBalance + (poolBalance * 1e13) / 1e18) revert InsufficientRepayment(); + } + } + + /// @dev Deposit function for erc20 tokens or eth + /// @dev if _token address is address(0), this is an eth deposit + /// @param _amount Amount of token to deposit + /// @param _token Address of token to deposit + function deposit(uint256 _amount, address _token) external payable nonReentrant { + if (_token == address(0)) { + ethBalances[msg.sender] += msg.value; + } else { + IERC20(_token).safeTransferFrom(msg.sender, address(this), _amount); + tokenBalances[msg.sender][_token] += _amount; + } + } + + /// @dev Withdraw function for erc20 tokens or eth + /// @dev if _token address is address(0), this is an eth withdrawal + /// @param _amount Amount of token to withdraw + /// @param _token Address of token to withdraw + function withdraw(uint256 _amount, address _token) external nonReentrant { + if (_token == address(0)) { + ethBalances[msg.sender] -= _amount; + (bool success,) = payable(msg.sender).call{value: _amount}(''); + require(success, 'Failed to send ETH'); + } else { + tokenBalances[msg.sender][_token] -= _amount; + IERC20(_token).safeTransfer(msg.sender, _amount); + } + } + + /// @dev users can deposit eth by sending it directly to the contract + receive() external payable {} +} diff --git a/solidity/contracts/Flashloan/GreedyReceiver.sol b/solidity/contracts/Flashloan/GreedyReceiver.sol new file mode 100644 index 00000000..7b3adb1f --- /dev/null +++ b/solidity/contracts/Flashloan/GreedyReceiver.sol @@ -0,0 +1,30 @@ +pragma solidity ^0.8.19; + +import {IFlashloan} from '../../interfaces/Flashloan/IFlashloan.sol'; +import {IERC20} from '@openzeppelin/contracts/token/ERC20/IERC20.sol'; +import {SafeERC20} from '@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol'; + +error OnlyPool(); + +contract GreedyReceiver { + IFlashloan public pool; + + using SafeERC20 for IERC20; + + modifier onlyPool() { + if (msg.sender != address(pool)) revert OnlyPool(); + _; + } + + constructor(address _poolAddress) { + pool = IFlashloan(_poolAddress); + } + + function flashLoan(uint256 _amount, address _token) external { + pool.flashloan(_amount, _token); + } + + function getETH() external payable onlyPool {} + + function getTokens(address _token, uint256 _amount) external onlyPool {} +} diff --git a/solidity/contracts/Flashloan/Receiver.sol b/solidity/contracts/Flashloan/Receiver.sol new file mode 100644 index 00000000..d85a89f9 --- /dev/null +++ b/solidity/contracts/Flashloan/Receiver.sol @@ -0,0 +1,35 @@ +pragma solidity ^0.8.19; + +import {IFlashloan} from '../../interfaces/Flashloan/IFlashloan.sol'; +import {IERC20} from '@openzeppelin/contracts/token/ERC20/IERC20.sol'; +import {SafeERC20} from '@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol'; + +error OnlyPool(); + +contract Receiver { + IFlashloan public pool; + + using SafeERC20 for IERC20; + + modifier onlyPool() { + if (msg.sender != address(pool)) revert OnlyPool(); + _; + } + + constructor(address _poolAddress) { + pool = IFlashloan(_poolAddress); + } + + function flashLoan(uint256 _amount, address _token) external { + pool.flashloan(_amount, _token); + } + + function getETH() external payable onlyPool { + (bool success,) = payable(msg.sender).call{value: msg.value + (msg.value * 1e13) / 1e18}(''); + require(success, 'failed to send ETH'); + } + + function getTokens(address _token, uint256 _amount) external onlyPool { + IERC20(_token).safeTransfer(msg.sender, _amount + (_amount * 1e13) / 1e18); + } +}