From 3be94b8d8ba0011138d0965955c2396fb8d18bda Mon Sep 17 00:00:00 2001 From: alhonaut Date: Fri, 4 Oct 2024 01:56:13 +0300 Subject: [PATCH] Add withdraw functions --- .gitignore | 1 + .gitmodules | 9 ++++++ lib/aave-v3-core | 1 + lib/forge-std | 2 +- lib/openzeppelin-contracts | 1 + remappings.txt | 4 +-- src/PaymentV1.sol | 56 ++++++++++++++++++++++++++++++++++---- test/PaymentTest.t.sol | 17 ++++++++++-- 8 files changed, 80 insertions(+), 11 deletions(-) create mode 100644 .gitmodules create mode 160000 lib/aave-v3-core create mode 160000 lib/openzeppelin-contracts diff --git a/.gitignore b/.gitignore index 85198aa..1390913 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ out/ !/broadcast /broadcast/*/31337/ /broadcast/**/dry-run/ +lib/ # Docs docs/ diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..d170a65 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,9 @@ +[submodule "lib/openzeppelin-contracts"] + path = lib/openzeppelin-contracts + url = https://github.com/openzeppelin/openzeppelin-contracts +[submodule "lib/aave-v3-core"] + path = lib/aave-v3-core + url = https://github.com/aave/aave-v3-core +[submodule "lib/forge-std"] + path = lib/forge-std + url = https://github.com/foundry-rs/forge-std diff --git a/lib/aave-v3-core b/lib/aave-v3-core new file mode 160000 index 0000000..b74526a --- /dev/null +++ b/lib/aave-v3-core @@ -0,0 +1 @@ +Subproject commit b74526a7bc67a3a117a1963fc871b3eb8cea8435 diff --git a/lib/forge-std b/lib/forge-std index 1714bee..8f24d6b 160000 --- a/lib/forge-std +++ b/lib/forge-std @@ -1 +1 @@ -Subproject commit 1714bee72e286e73f76e320d110e0eaf5c4e649d +Subproject commit 8f24d6b04c92975e0795b5868aa0d783251cdeaa diff --git a/lib/openzeppelin-contracts b/lib/openzeppelin-contracts new file mode 160000 index 0000000..dbb6104 --- /dev/null +++ b/lib/openzeppelin-contracts @@ -0,0 +1 @@ +Subproject commit dbb6104ce834628e473d2173bbc9d47f81a9eec3 diff --git a/remappings.txt b/remappings.txt index d5a40fb..0a647c9 100644 --- a/remappings.txt +++ b/remappings.txt @@ -1,4 +1,4 @@ forge-std/=lib/forge-std/src/ @chainlink/=/node_modules/@chainlink -@openzeppelin/=/node_modules/@openzeppelin -@aave/=/node_modules/@aave +openzeppelin-contracts/=/home/trauslamen/Blockchain_Projects/gig-contracts/lib/openzeppelin-contracts +aave/aave-v3-core/=/home/trauslamen/Blockchain_Projects/gig-contracts/lib/aave-v3-core diff --git a/src/PaymentV1.sol b/src/PaymentV1.sol index c9c67d6..c378d0b 100644 --- a/src/PaymentV1.sol +++ b/src/PaymentV1.sol @@ -1,9 +1,9 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.19; -import {IERC20} from "@chainlink/contracts-ccip/src/v0.8/vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/IERC20.sol"; -import {SafeERC20} from "@chainlink/contracts-ccip/src/v0.8/vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/utils/SafeERC20.sol"; -import {IPool} from "@aave/core-v3/contracts/interfaces/IPool.sol"; +import {IERC20} from "openzeppelin-contracts/contracts/token/ERC20/IERC20.sol"; +import {SafeERC20} from "openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol"; +import {IPool} from "aave/aave-v3-core/contracts/interfaces/IPool.sol"; contract USDCPayment { using SafeERC20 for IERC20; @@ -13,21 +13,40 @@ contract USDCPayment { IPool public immutable i_pool; IERC20 public immutable i_usdc; + address public team_wallet; // safe wallet address + error AmountIsZero(); + error IncorrectAmount(uint256 amount); error TransferFailed(); error IncorrectTokenAmount(); error IncorrectAmountId(); error NotOwner(); error InvalidUsdcToken(); + error IncorrectWalletAddress(); event ProposalOpened(bytes32 id, uint256 amount, address initiator); event ProposalClosed(bytes32 id, address freelancer); + event WalletChanged(address curr_wallet, address new_wallet); + event WithdrawUSDC(address receiver, uint256 amount); + event WithdrawETH(address receiver, uint256 amount); + + modifier onlyTeam() { + if (msg.sender != team_wallet) { + revert IncorrectWalletAddress(); + } + _; + } - constructor(address _addressPool, address _usdcAddress) { + constructor( + address _addressPool, + address _usdcAddress, + address _teamAddress + ) { if (_usdcAddress == address(0)) revert InvalidUsdcToken(); i_pool = IPool(_addressPool); i_usdc = IERC20(_usdcAddress); - i_usdc.safeApprove(_addressPool, type(uint256).max); + i_usdc.safeIncreaseAllowance(_addressPool, type(uint256).max); + team_wallet = _teamAddress; } function openProposal(uint256 _amount) external { @@ -65,6 +84,33 @@ contract USDCPayment { i_pool.withdraw(address(i_usdc), proposalAmount, receiver); } + function withdrawUSDC(uint256 _amount) public onlyTeam { + (uint256 totalCollateralBase, , , , , ) = i_pool.getUserAccountData( + address(this) + ); // get all lended money from AAVE Pool + if (_amount > totalCollateralBase) { + revert IncorrectAmount(_amount); + } + i_usdc.safeTransfer(msg.sender, _amount); + + emit WithdrawUSDC(msg.sender, _amount); + } + + function withdrawETH() public onlyTeam { + uint256 balance = address(this).balance; + (bool success, ) = msg.sender.call{value: balance}(""); + if (!success) { + revert TransferFailed(); + } + + emit WithdrawETH(msg.sender, balance); + } + + function changeTeamWallet(address _newTeamWallet) public onlyTeam { + team_wallet = _newTeamWallet; + emit WalletChanged(team_wallet, _newTeamWallet); + } + function getBalance(address _tokenAddress) external view returns (uint256) { return IERC20(_tokenAddress).balanceOf(address(this)); } diff --git a/test/PaymentTest.t.sol b/test/PaymentTest.t.sol index cc82411..86f6d07 100644 --- a/test/PaymentTest.t.sol +++ b/test/PaymentTest.t.sol @@ -4,8 +4,8 @@ pragma solidity ^0.8.19; import {Test} from "forge-std/Test.sol"; import {USDCPayment} from "../src/PaymentV1.sol"; import {IUSDC} from "../src/interfaces/IUSDC.sol"; -import {IERC20} from "@chainlink/contracts-ccip/src/v0.8/vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/IERC20.sol"; -import {IPool} from "@aave/core-v3/contracts/interfaces/IPool.sol"; +import {IERC20} from "openzeppelin-contracts/contracts/token/ERC20/IERC20.sol"; +import {IPool} from "aave/aave-v3-core/contracts/interfaces/IPool.sol"; contract USDCPaymentTest is Test { USDCPayment private usdcPayment; @@ -13,9 +13,11 @@ contract USDCPaymentTest is Test { address public constant POOL = 0x794a61358D6845594F94dc1DB02A252b5b4814aD; // Arbitrum mainnet Pool uint256 public constant TESTING_AMOUNT = 1e6; + address public testWallet = makeAddr("TEST_WALLET"); + function setUp() public { vm.createSelectFork(vm.envString("ARBITRUM_RPC_URL")); - usdcPayment = new USDCPayment(POOL, USDC); + usdcPayment = new USDCPayment(POOL, USDC, testWallet); } function mintTokens() public { @@ -224,4 +226,13 @@ contract USDCPaymentTest is Test { "Need to delete proposal budget after closing." ); } + + function test_RevertWithdrawIfNotOwner() public { + address tester = makeAddr("test_teamWallet"); + vm.startPrank(tester); + vm.expectRevert(USDCPayment.IncorrectWalletAddress.selector); + usdcPayment.withdrawUSDC(TESTING_AMOUNT); + vm.expectRevert(USDCPayment.IncorrectWalletAddress.selector); + usdcPayment.withdrawETH(); + } }