From d74f44b437fccb0a885ea5c291bd35a39b92bb7b Mon Sep 17 00:00:00 2001 From: 0xb337r007 <0xe4e5@proton.me> Date: Sat, 2 Mar 2024 15:53:52 +0100 Subject: [PATCH] feat(Vault): allow withdrawal of untracked tokens --- .gas-snapshot | 51 +++++++++++++++++++----------------- contracts/CommunityVault.sol | 49 ++++++++++++++++++++++++++++++++++ test/CommunityVault.t.sol | 44 +++++++++++++++++++++++++++++++ 3 files changed, 120 insertions(+), 24 deletions(-) diff --git a/.gas-snapshot b/.gas-snapshot index 72502c7..5fd2441 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -4,12 +4,15 @@ AddEntryTest:test_RevertWhen_InvalidAddress() (gas: 25133) AddEntryTest:test_RevertWhen_SenderIsNotTokenDeployer() (gas: 14827) CommunityERC20Test:test_Deployment() (gas: 35198) CommunityTokenDeployerTest:test_Deployment() (gas: 14805) -CommunityVaultBaseERC20Test:test_Deployment() (gas: 10572) -CommunityVaultBaseERC721Test:test_Deployment() (gas: 10572) -CommunityVaultBaseTransferERC721Test:test_Deployment() (gas: 10572) -CommunityVaultDepositERC721Test:testSuccessfulDepositERC721() (gas: 184700) -CommunityVaultDepositERC721Test:test_Deployment() (gas: 10714) -CommunityVaultTest:test_Deployment() (gas: 10572) +CommunityVaulWithdrawUntrackedERC721Test:testRevertWithdrawalIfTokenIsTracked() (gas: 37990) +CommunityVaulWithdrawUntrackedERC721Test:testSuccessfulDepositERC721() (gas: 73349) +CommunityVaulWithdrawUntrackedERC721Test:test_Deployment() (gas: 10783) +CommunityVaultBaseERC20Test:test_Deployment() (gas: 10641) +CommunityVaultBaseERC721Test:test_Deployment() (gas: 10641) +CommunityVaultBaseTransferERC721Test:test_Deployment() (gas: 10641) +CommunityVaultDepositERC721Test:testSuccessfulDepositERC721() (gas: 184942) +CommunityVaultDepositERC721Test:test_Deployment() (gas: 10783) +CommunityVaultTest:test_Deployment() (gas: 10641) CreateTest:test_Create() (gas: 2374801) CreateTest:test_Create() (gas: 2661968) CreateTest:test_RevertWhen_InvalidOwnerTokenAddress() (gas: 15523) @@ -33,9 +36,9 @@ DeploymentTest:test_Deployment() (gas: 14671) DeploymentTest:test_Deployment() (gas: 14671) DeploymentTest:test_Deployment() (gas: 17295) DeploymentTest:test_Deployment() (gas: 36430) -DepositERC20Test:testDepositZeroTokens() (gas: 15199) -DepositERC20Test:testSuccessfulDepositERC20() (gas: 85584) -DepositERC20Test:test_Deployment() (gas: 10594) +DepositERC20Test:testDepositZeroTokens() (gas: 15211) +DepositERC20Test:testSuccessfulDepositERC20() (gas: 85703) +DepositERC20Test:test_Deployment() (gas: 10663) GetEntryTest:test_ReturnZeroAddressIfEntryDoesNotExist() (gas: 11906) MintToTest:test_Deployment() (gas: 35220) MintToTest:test_Deployment() (gas: 83308) @@ -88,18 +91,18 @@ SetTokenDeployerAddressTest:test_RevertWhen_SenderIsNotOwner() (gas: 12438) SetTokenDeployerAddressTest:test_RevertWhen_SenderIsNotOwner() (gas: 12438) SetTokenDeployerAddressTest:test_SetTokenDeployerAddress() (gas: 22768) SetTokenDeployerAddressTest:test_SetTokenDeployerAddress() (gas: 22768) -TransferERC20ByAdminTest:test_AdminCanTransferERC20() (gas: 97818) -TransferERC20ByAdminTest:test_Deployment() (gas: 10714) -TransferERC20ByAdminTest:test_LengthMismatch() (gas: 26146) -TransferERC20ByAdminTest:test_NoRecipients() (gas: 19516) -TransferERC20ByAdminTest:test_TransferAmountZero() (gas: 66057) -TransferERC20ByAdminTest:test_TransferERC20AmountTooBig() (gas: 59079) -TransferERC20ByNonAdminTest:test_Deployment() (gas: 10594) -TransferERC20ByNonAdminTest:test_revertIfCalledByNonAdmin() (gas: 29912) -TransferERC721ByAdminTest:test_AdminCanTransferERC721() (gas: 141776) -TransferERC721ByAdminTest:test_Deployment() (gas: 10670) -TransferERC721ByAdminTest:test_LengthMismatch() (gas: 26156) -TransferERC721ByAdminTest:test_NoRecipients() (gas: 19506) -TransferERC721ByAdminTest:test_RevertOnTransferERC721IfNotDeposited() (gas: 32736) -TransferERC721ByNonAdminTest:test_Deployment() (gas: 10714) -TransferERC721ByNonAdminTest:test_RevertIfCalledByNonAdmin() (gas: 29953) +TransferERC20ByAdminTest:test_AdminCanTransferERC20() (gas: 98048) +TransferERC20ByAdminTest:test_Deployment() (gas: 10783) +TransferERC20ByAdminTest:test_LengthMismatch() (gas: 26182) +TransferERC20ByAdminTest:test_NoRecipients() (gas: 19552) +TransferERC20ByAdminTest:test_TransferAmountZero() (gas: 66178) +TransferERC20ByAdminTest:test_TransferERC20AmountTooBig() (gas: 59224) +TransferERC20ByNonAdminTest:test_Deployment() (gas: 10663) +TransferERC20ByNonAdminTest:test_revertIfCalledByNonAdmin() (gas: 29972) +TransferERC721ByAdminTest:test_AdminCanTransferERC721() (gas: 141912) +TransferERC721ByAdminTest:test_Deployment() (gas: 10739) +TransferERC721ByAdminTest:test_LengthMismatch() (gas: 26192) +TransferERC721ByAdminTest:test_NoRecipients() (gas: 19542) +TransferERC721ByAdminTest:test_RevertOnTransferERC721IfNotDeposited() (gas: 32784) +TransferERC721ByNonAdminTest:test_Deployment() (gas: 10783) +TransferERC721ByNonAdminTest:test_RevertIfCalledByNonAdmin() (gas: 30013) \ No newline at end of file diff --git a/contracts/CommunityVault.sol b/contracts/CommunityVault.sol index a4a51c6..7e95c9a 100644 --- a/contracts/CommunityVault.sol +++ b/contracts/CommunityVault.sol @@ -31,6 +31,8 @@ contract CommunityVault is CommunityOwnable, IERC721Receiver { error CommunityVault_IndexOutOfBounds(); error CommunityVault_ERC721TokenAlreadyDeposited(); error CommunityVault_ERC721TokenNotDeposited(); + error CommunityVault_AmountExceedsUntrackedBalanceERC20(); + error CommunityVault_CannotWithdrawTrackedERC721(); mapping(address => uint256) public erc20TokenBalances; mapping(address => EnumerableSet.UintSet) private erc721TokenIds; @@ -169,6 +171,53 @@ contract CommunityVault is CommunityOwnable, IERC721Receiver { } } + /// @notice Withdraws a specified amount of an untracked ERC20 token from the community vault. + /// @dev This function allows the community owner or token master to withdraw untracked ERC20 tokens. It checks if + /// the requested amount does not exceed the untracked balance. If it does, the transaction is reverted. + /// @param tokenAddress The address of the ERC20 token to withdraw. + /// @param amount The amount of the ERC20 token to withdraw. + /// @param to The address to which the ERC20 tokens will be transferred. + function withdrawUntrackedERC20( + address tokenAddress, + uint256 amount, + address to + ) + public + onlyCommunityOwnerOrTokenMaster + { + uint256 contractBalance = IERC20(tokenAddress).balanceOf(address(this)); + uint256 untrackedBalance = contractBalance - erc20TokenBalances[tokenAddress]; + + if (amount > untrackedBalance) { + revert CommunityVault_AmountExceedsUntrackedBalanceERC20(); + } + + IERC20(tokenAddress).safeTransfer(to, amount); + } + + /// @notice Withdraws specified ERC721 tokens that are not tracked by the community vault. + /// @dev This function allows the community owner or token master to withdraw untracked ERC721 tokens by token IDs. + /// It checks each token ID against tracked tokens and if any are found, the transaction is reverted. + /// @param tokenAddress The address of the ERC721 token to withdraw. + /// @param tokenIds An array of token IDs of the ERC721 tokens to withdraw. + /// @param to The address to which the ERC721 tokens will be transferred. + function withdrawUntrackedERC721( + address tokenAddress, + uint256[] memory tokenIds, + address to + ) + public + onlyCommunityOwnerOrTokenMaster + { + for (uint256 i = 0; i < tokenIds.length; i++) { + if (erc721TokenIds[tokenAddress].contains(tokenIds[i])) { + revert CommunityVault_CannotWithdrawTrackedERC721(); + } + + IERC721(tokenAddress).safeTransferFrom(address(this), to, tokenIds[i]); + } + } + /** * @dev Handles the receipt of an ERC721 token. * @return bytes4 Returns `bytes4(keccak256("onERC721Received(address,address,uint256,bytes)"))` diff --git a/test/CommunityVault.t.sol b/test/CommunityVault.t.sol index 433ea1d..8720922 100644 --- a/test/CommunityVault.t.sol +++ b/test/CommunityVault.t.sol @@ -303,3 +303,47 @@ contract CommunityVaultDepositERC721Test is CommunityVaultBaseERC721Test { assertEq(vault.erc721TokenBalances(address(erc721Token)), initialTokenBalanceValue + 2); } } + +contract CommunityVaulWithdrawUntrackedERC721Test is CommunityVaultBaseERC721Test { + function setUp() public virtual override { + CommunityVaultBaseERC721Test.setUp(); + vm.startPrank(accounts[0]); + // trasfer to contract ids 0 and 1 + erc721Token.transferFrom(accounts[0], address(vault), 0); + erc721Token.transferFrom(accounts[0], address(vault), 1); + + // deposit id 2 + uint256[] memory ids = new uint256[](1); + ids[0] = 2; + erc721Token.approve(address(vault), 2); + vault.depositERC721(address(erc721Token), ids); + vm.stopPrank(); + } + + function testRevertWithdrawalIfTokenIsTracked() public { + uint256[] memory ids = new uint256[](1); + ids[0] = 2; + + assertEq(erc721Token.ownerOf(2), address(vault)); + assertEq(vault.getERC721DepositedTokenByIndex(address(erc721Token), 0), 2); + + vm.prank(deployer); + vm.expectRevert(CommunityVault.CommunityVault_CannotWithdrawTrackedERC721.selector); + vault.withdrawUntrackedERC721(address(erc721Token), ids, accounts[0]); + } + + function testSuccessfulDepositERC721() public { + uint256[] memory ids = new uint256[](2); + ids[0] = 0; + ids[1] = 1; + + assertEq(erc721Token.ownerOf(0), address(vault)); + assertEq(erc721Token.ownerOf(1), address(vault)); + + vm.prank(deployer); + vault.withdrawUntrackedERC721(address(erc721Token), ids, accounts[0]); + + assertEq(erc721Token.ownerOf(0), accounts[0]); + assertEq(erc721Token.ownerOf(1), accounts[0]); + } +}