From aacd7fd43966725c1413f2433b6df607acd73a67 Mon Sep 17 00:00:00 2001 From: 0xb337r007 <0xe4e5@proton.me> Date: Wed, 13 Mar 2024 17:11:44 +0100 Subject: [PATCH] feat(CommunityVault): add migration functions for ERC20 and ERC721 tokens --- .gas-snapshot | 45 ++++--- contracts/CommunityVault.sol | 67 +++++++++- contracts/interfaces/ICommunityVault.sol | 8 ++ test/CommunityVault.t.sol | 155 +++++++++++++++++++++++ 4 files changed, 256 insertions(+), 19 deletions(-) create mode 100644 contracts/interfaces/ICommunityVault.sol diff --git a/.gas-snapshot b/.gas-snapshot index 3710ab5..9b1badd 100644 --- a/.gas-snapshot +++ b/.gas-snapshot @@ -7,14 +7,23 @@ CommunityTokenDeployerTest:test_Deployment() (gas: 14805) CommunityVaultBaseERC20Test:test_Deployment() (gas: 10641) CommunityVaultBaseERC721Test:test_Deployment() (gas: 10641) CommunityVaultBaseTransferERC721Test:test_Deployment() (gas: 10641) -CommunityVaultDepositERC721Test:testSuccessfulDepositERC721() (gas: 184942) +CommunityVaultDepositERC721Test:testSuccessfulDepositERC721() (gas: 185076) CommunityVaultDepositERC721Test:test_Deployment() (gas: 10783) +CommunityVaultMigrationTest:test_Deployment() (gas: 10762) +CommunityVaultMigrationTest:test_migrateERC20RevertsIfNewImplementationIsNotSet() (gas: 19992) +CommunityVaultMigrationTest:test_migrateERC20RevertsIfNotAuthorized() (gas: 24578) +CommunityVaultMigrationTest:test_migrateERC20RevertsIfTokenBalanceIsZero() (gas: 54025) +CommunityVaultMigrationTest:test_migrateERC20Tokens() (gas: 190966) +CommunityVaultMigrationTest:test_migrateERC721RevertsIfNewImplementationIsNotSet() (gas: 20103) +CommunityVaultMigrationTest:test_migrateERC721RevertsIfNotAuthorized() (gas: 24731) +CommunityVaultMigrationTest:test_migrateERC721Tokens() (gas: 398356) +CommunityVaultMigrationTest:test_migrateERC721TokensRevertsIfTokenNotDeposited() (gas: 54130) CommunityVaultTest:test_Deployment() (gas: 10641) -CommunityVaultWithdrawUntrackedERC20Test:testRevertWithdrawalIfAmountIsMoreThanTheUntracked() (gas: 30810) -CommunityVaultWithdrawUntrackedERC20Test:testSuccessfulWithdrawal() (gas: 64806) +CommunityVaultWithdrawUntrackedERC20Test:testRevertWithdrawalIfAmountIsMoreThanTheUntracked() (gas: 30800) +CommunityVaultWithdrawUntrackedERC20Test:testSuccessfulWithdrawal() (gas: 64796) CommunityVaultWithdrawUntrackedERC20Test:test_Deployment() (gas: 10761) -CommunityVaultWithdrawUntrackedERC721Test:testRevertWithdrawalIfTokenIsTracked() (gas: 37990) -CommunityVaultWithdrawUntrackedERC721Test:testSuccessfulWithdrUntrackedERC721() (gas: 73328) +CommunityVaultWithdrawUntrackedERC721Test:testRevertWithdrawalIfTokenIsTracked() (gas: 38040) +CommunityVaultWithdrawUntrackedERC721Test:testSuccessfulWithdrUntrackedERC721() (gas: 73356) CommunityVaultWithdrawUntrackedERC721Test:test_Deployment() (gas: 10783) CreateTest:test_Create() (gas: 2374801) CreateTest:test_Create() (gas: 2661968) @@ -39,8 +48,8 @@ DeploymentTest:test_Deployment() (gas: 14671) DeploymentTest:test_Deployment() (gas: 14671) DeploymentTest:test_Deployment() (gas: 17295) DeploymentTest:test_Deployment() (gas: 36430) -DepositERC20Test:testDepositZeroTokens() (gas: 15211) -DepositERC20Test:testSuccessfulDepositERC20() (gas: 85703) +DepositERC20Test:testDepositZeroTokens() (gas: 15233) +DepositERC20Test:testSuccessfulDepositERC20() (gas: 85857) DepositERC20Test:test_Deployment() (gas: 10663) GetEntryTest:test_ReturnZeroAddressIfEntryDoesNotExist() (gas: 11906) MintToTest:test_Deployment() (gas: 35220) @@ -94,18 +103,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: 98048) +TransferERC20ByAdminTest:test_AdminCanTransferERC20() (gas: 98208) 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) +TransferERC20ByAdminTest:test_LengthMismatch() (gas: 26210) +TransferERC20ByAdminTest:test_NoRecipients() (gas: 19580) +TransferERC20ByAdminTest:test_TransferAmountZero() (gas: 66206) +TransferERC20ByAdminTest:test_TransferERC20AmountTooBig() (gas: 59318) TransferERC20ByNonAdminTest:test_Deployment() (gas: 10663) -TransferERC20ByNonAdminTest:test_revertIfCalledByNonAdmin() (gas: 29972) -TransferERC721ByAdminTest:test_AdminCanTransferERC721() (gas: 141912) +TransferERC20ByNonAdminTest:test_revertIfCalledByNonAdmin() (gas: 30006) +TransferERC721ByAdminTest:test_AdminCanTransferERC721() (gas: 142074) TransferERC721ByAdminTest:test_Deployment() (gas: 10739) -TransferERC721ByAdminTest:test_LengthMismatch() (gas: 26192) -TransferERC721ByAdminTest:test_NoRecipients() (gas: 19542) -TransferERC721ByAdminTest:test_RevertOnTransferERC721IfNotDeposited() (gas: 32784) +TransferERC721ByAdminTest:test_LengthMismatch() (gas: 26220) +TransferERC721ByAdminTest:test_NoRecipients() (gas: 19570) +TransferERC721ByAdminTest:test_RevertOnTransferERC721IfNotDeposited() (gas: 32812) TransferERC721ByNonAdminTest:test_Deployment() (gas: 10783) -TransferERC721ByNonAdminTest:test_RevertIfCalledByNonAdmin() (gas: 30013) \ No newline at end of file +TransferERC721ByNonAdminTest:test_RevertIfCalledByNonAdmin() (gas: 30047) \ No newline at end of file diff --git a/contracts/CommunityVault.sol b/contracts/CommunityVault.sol index 7e95c9a..6063247 100644 --- a/contracts/CommunityVault.sol +++ b/contracts/CommunityVault.sol @@ -8,6 +8,7 @@ import { IERC721 } from "@openzeppelin/contracts/token/ERC721/IERC721.sol"; import { IERC721Receiver } from "@openzeppelin/contracts/token/ERC721/IERC721Receiver.sol"; import { EnumerableSet } from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol"; import { CommunityOwnable } from "./CommunityOwnable.sol"; +import { ICommunityVault } from "./interfaces/ICommunityVault.sol"; /** * @title CommunityVault @@ -16,7 +17,7 @@ import { CommunityOwnable } from "./CommunityOwnable.sol"; * Only community owners, as defined in the CommunityOwnable contract, have * permissions to transfer these tokens out of the vault. */ -contract CommunityVault is CommunityOwnable, IERC721Receiver { +contract CommunityVault is ICommunityVault, CommunityOwnable, IERC721Receiver { using SafeERC20 for IERC20; using EnumerableSet for EnumerableSet.UintSet; @@ -33,12 +34,30 @@ contract CommunityVault is CommunityOwnable, IERC721Receiver { error CommunityVault_ERC721TokenNotDeposited(); error CommunityVault_AmountExceedsUntrackedBalanceERC20(); error CommunityVault_CannotWithdrawTrackedERC721(); + error CommunityVault_ZeroAddress(); + error CommunityVault_NewImplementationNotSet(); + error CommunityVault_ZeroBalance(); mapping(address => uint256) public erc20TokenBalances; mapping(address => EnumerableSet.UintSet) private erc721TokenIds; + // New implementation address for ERC20 token migration + address public newImplementation; + constructor(address _ownerToken, address _masterToken) CommunityOwnable(_ownerToken, _masterToken) { } + /** + * @dev Sets the new implementation address. Only callable by the community owner or token master. + * @param _newImplementation The address of the new implementation to which tokens will be migrated. + */ + function setNewImplementation(address _newImplementation) external onlyCommunityOwnerOrTokenMaster { + if (_newImplementation == address(0)) { + revert CommunityVault_ZeroAddress(); + } + + newImplementation = _newImplementation; + } + /** * @dev Allows anyone to deposit ERC20 tokens into the vault. * @param token The address of the ERC20 token to deposit. @@ -226,4 +245,50 @@ contract CommunityVault is CommunityOwnable, IERC721Receiver { function onERC721Received(address, address, uint256, bytes calldata) public pure override returns (bytes4) { return this.onERC721Received.selector; } + + /** + * @dev Migrates ERC20 tokens to the new implementation address. + * @param tokens The addresses of the ERC20 tokens to migrate. + */ + function migrateERC20Tokens(address[] calldata tokens) external onlyCommunityOwnerOrTokenMaster { + if (newImplementation == address(0)) { + revert CommunityVault_NewImplementationNotSet(); + } + + for (uint256 i = 0; i < tokens.length; i++) { + address token = tokens[i]; + uint256 balance = erc20TokenBalances[token]; + + if (balance == 0) { + revert CommunityVault_ZeroBalance(); + } + + erc20TokenBalances[token] = 0; + IERC20(token).approve(newImplementation, balance); + ICommunityVault(newImplementation).depositERC20(token, balance); + } + } + + /** + * @dev Migrates ERC721 tokens to the new implementation address. + * @param token The address of the ERC721 token to migrate. + * @param tokenIds The IDs of the ERC721 tokens to migrate. + */ + function migrateERC721Tokens(address token, uint256[] calldata tokenIds) external onlyCommunityOwnerOrTokenMaster { + if (newImplementation == address(0)) { + revert CommunityVault_NewImplementationNotSet(); + } + + for (uint256 i = 0; i < tokenIds.length; i++) { + uint256 tokenId = tokenIds[i]; + if (!erc721TokenIds[token].contains(tokenId)) { + revert CommunityVault_ERC721TokenNotDeposited(); + } + + erc721TokenIds[token].remove(tokenId); + IERC721(token).approve(newImplementation, tokenId); + } + + ICommunityVault(newImplementation).depositERC721(token, tokenIds); + } } diff --git a/contracts/interfaces/ICommunityVault.sol b/contracts/interfaces/ICommunityVault.sol new file mode 100644 index 0000000..de2823b --- /dev/null +++ b/contracts/interfaces/ICommunityVault.sol @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Mozilla Public License 2.0 + +pragma solidity ^0.8.17; + +interface ICommunityVault { + function depositERC20(address token, uint256 amount) external; + function depositERC721(address token, uint256[] calldata tokenIds) external; +} diff --git a/test/CommunityVault.t.sol b/test/CommunityVault.t.sol index 7123cf4..27af11d 100644 --- a/test/CommunityVault.t.sol +++ b/test/CommunityVault.t.sol @@ -381,3 +381,158 @@ contract CommunityVaultWithdrawUntrackedERC721Test is CommunityVaultBaseERC721Te assertEq(erc721Token.ownerOf(1), accounts[0]); } } + +contract CommunityVaultMigrationTest is CommunityVaultTest { + CommunityVault internal newVault; + TestERC20Token internal erc20Token2; + TestERC20Token internal erc20Token3; + + function setUp() public virtual override { + CommunityVaultTest.setUp(); + + newVault = new CommunityVault(address(ownerToken), address(masterToken)); + erc20Token2 = new TestERC20Token(); + erc20Token3 = new TestERC20Token(); + + vm.startPrank(deployer); + // mint erc20 tokens and deposit + erc20Token.mint(deployer, 10e18); + erc20Token.approve(address(vault), 10e18); + vault.depositERC20(address(erc20Token), 10e18); + erc20Token2.mint(deployer, 5e18); + erc20Token2.approve(address(vault), 5e18); + vault.depositERC20(address(erc20Token2), 5e18); + + // mint erc721 tokens and deposit + erc721Token.mint(deployer); + erc721Token.mint(deployer); + erc721Token.mint(deployer); + erc721Token.mint(deployer); + // id 4 is not deposited + erc721Token.mint(deployer); + + uint256[] memory ids = new uint256[](4); + ids[0] = 0; + ids[1] = 1; + ids[2] = 2; + ids[3] = 3; + erc721Token.approve(address(vault), 0); + erc721Token.approve(address(vault), 1); + erc721Token.approve(address(vault), 2); + erc721Token.approve(address(vault), 3); + vault.depositERC721(address(erc721Token), ids); + + vm.stopPrank(); + } + + function test_migrateERC20RevertsIfNotAuthorized() public { + vm.prank(accounts[0]); + vm.expectRevert(CommunityOwnable.CommunityOwnable_NotAuthorized.selector); + + address[] memory tokens = new address[](0); + vault.migrateERC20Tokens(tokens); + } + + function test_migrateERC721RevertsIfNotAuthorized() public { + vm.prank(accounts[0]); + vm.expectRevert(CommunityOwnable.CommunityOwnable_NotAuthorized.selector); + + uint256[] memory ids = new uint256[](0); + vault.migrateERC721Tokens(address(0), ids); + } + + function test_migrateERC20RevertsIfNewImplementationIsNotSet() public { + assertEq(vault.newImplementation(), address(0)); + vm.prank(deployer); + vm.expectRevert(CommunityVault.CommunityVault_NewImplementationNotSet.selector); + + address[] memory tokens = new address[](0); + vault.migrateERC20Tokens(tokens); + } + + function test_migrateERC721RevertsIfNewImplementationIsNotSet() public { + assertEq(vault.newImplementation(), address(0)); + vm.prank(deployer); + vm.expectRevert(CommunityVault.CommunityVault_NewImplementationNotSet.selector); + + uint256[] memory ids = new uint256[](0); + vault.migrateERC721Tokens(address(0), ids); + } + + function test_migrateERC20RevertsIfTokenBalanceIsZero() public { + vm.startPrank(deployer); + vault.setNewImplementation(address(newVault)); + + assertEq(erc20Token3.balanceOf(address(vault)), 0); + vm.expectRevert(CommunityVault.CommunityVault_ZeroBalance.selector); + + address[] memory tokens = new address[](1); + tokens[0] = address(erc20Token3); + + vault.migrateERC20Tokens(tokens); + + vm.stopPrank(); + } + + function test_migrateERC20Tokens() public { + vm.startPrank(deployer); + + vault.setNewImplementation(address(newVault)); + assertEq(erc20Token.balanceOf(address(vault)), 10e18); + assertEq(erc20Token2.balanceOf(address(vault)), 5e18); + assertEq(erc20Token.balanceOf(address(newVault)), 0); + assertEq(erc20Token2.balanceOf(address(newVault)), 0); + + address[] memory tokens = new address[](2); + tokens[0] = address(erc20Token); + tokens[1] = address(erc20Token2); + vault.migrateERC20Tokens(tokens); + + assertEq(erc20Token.balanceOf(address(vault)), 0); + assertEq(erc20Token2.balanceOf(address(vault)), 0); + assertEq(erc20Token.balanceOf(address(newVault)), 10e18); + assertEq(erc20Token2.balanceOf(address(newVault)), 5e18); + + vm.stopPrank(); + } + + function test_migrateERC721TokensRevertsIfTokenNotDeposited() public { + vm.startPrank(deployer); + + vault.setNewImplementation(address(newVault)); + assertEq(erc721Token.ownerOf(4), deployer); + + uint256[] memory ids = new uint256[](1); + ids[0] = 4; + + vm.expectRevert(CommunityVault.CommunityVault_ERC721TokenNotDeposited.selector); + vault.migrateERC721Tokens(address(erc721Token), ids); + + vm.stopPrank(); + } + + function test_migrateERC721Tokens() public { + vm.startPrank(deployer); + + vault.setNewImplementation(address(newVault)); + assertEq(erc721Token.ownerOf(0), address(vault)); + assertEq(erc721Token.ownerOf(1), address(vault)); + assertEq(erc721Token.ownerOf(2), address(vault)); + assertEq(erc721Token.ownerOf(3), address(vault)); + + uint256[] memory ids = new uint256[](4); + ids[0] = 0; + ids[1] = 1; + ids[2] = 2; + ids[3] = 3; + + vault.migrateERC721Tokens(address(erc721Token), ids); + + assertEq(erc721Token.ownerOf(0), address(newVault)); + assertEq(erc721Token.ownerOf(1), address(newVault)); + assertEq(erc721Token.ownerOf(2), address(newVault)); + assertEq(erc721Token.ownerOf(3), address(newVault)); + + vm.stopPrank(); + } +}