From 6213e992d1458919cc42d8e3a767fd29e2600828 Mon Sep 17 00:00:00 2001
From: 0xb337r007 <0xe4e5@proton.me>
Date: Mon, 26 Feb 2024 11:00:31 +0100
Subject: [PATCH] feat(CommunityVault): add erc20 deposit function

---
 contracts/CommunityVault.sol | 32 +++++++++++++++++
 test/CommunityVault.t.sol    | 67 ++++++++++++++++++++++++++----------
 2 files changed, 80 insertions(+), 19 deletions(-)

diff --git a/contracts/CommunityVault.sol b/contracts/CommunityVault.sol
index 4f97978..b52739e 100644
--- a/contracts/CommunityVault.sol
+++ b/contracts/CommunityVault.sol
@@ -17,12 +17,39 @@ import { CommunityOwnable } from "./CommunityOwnable.sol";
 contract CommunityVault is CommunityOwnable {
     using SafeERC20 for IERC20;
 
+    event ERC20Deposited(address indexed depositor, address indexed token, uint256 amount);
+    event ERC721Deposited(address indexed depositor, address indexed token, uint256 tokenId);
+
     error CommunityVault_LengthMismatch();
     error CommunityVault_NoRecipients();
     error CommunityVault_TransferAmountZero();
+    error CommunityVault_ERC20TransferAmountTooBig();
+    error CommunityVault_DepositAmountZero();
+
+    mapping(address => uint256) public erc20TokenBalances;
 
     constructor(address _ownerToken, address _masterToken) CommunityOwnable(_ownerToken, _masterToken) { }
 
+    /**
+     * @dev Allows anyone to deposit ERC20 tokens into the vault.
+     * @param token The address of the ERC20 token to deposit.
+     * @param amount The amount of tokens to deposit.
+     */
+    function depositERC20(address token, uint256 amount) external {
+        if (amount == 0) {
+            revert CommunityVault_DepositAmountZero();
+        }
+
+        // Transfer tokens from the sender to this contract
+        IERC20(token).safeTransferFrom(msg.sender, address(this), amount);
+
+        // Update the total balance of the token in the vault
+        erc20TokenBalances[token] += amount;
+
+        // Emit an event for the deposit (optional, but recommended for tracking)
+        emit ERC20Deposited(msg.sender, token, amount);
+    }
+
     /**
      * @dev Transfers ERC20 tokens to a list of addresses.
      * @param token The ERC20 token address.
@@ -50,6 +77,11 @@ contract CommunityVault is CommunityOwnable {
                 revert CommunityVault_TransferAmountZero();
             }
 
+            if (amounts[i] > erc20TokenBalances[token]) {
+                revert CommunityVault_ERC20TransferAmountTooBig();
+            }
+
+            erc20TokenBalances[token] -= amounts[i];
             IERC20(token).safeTransfer(recipients[i], amounts[i]);
         }
     }
diff --git a/test/CommunityVault.t.sol b/test/CommunityVault.t.sol
index e9826a2..3c0298f 100644
--- a/test/CommunityVault.t.sol
+++ b/test/CommunityVault.t.sol
@@ -52,14 +52,7 @@ contract CommunityVaultTest is Test {
 contract CommunityVaultBaseERC20Test is CommunityVaultTest {
     function setUp() public virtual override {
         CommunityVaultTest.setUp();
-
-        // mint 10 tokens to user
-        address user = accounts[0];
-        erc20Token.mint(user, 10e18);
-
-        // user transfer 10 tokens to the vault
-        vm.prank(user);
-        erc20Token.transfer(address(vault), 10e18);
+        erc20Token.mint(accounts[0], 10e18);
     }
 }
 
@@ -69,8 +62,6 @@ contract TransferERC20ByNonAdminTest is CommunityVaultBaseERC20Test {
     }
 
     function test_revertIfCalledByNonAdmin() public {
-        assertEq(erc20Token.balanceOf(address(vault)), 10e18);
-
         uint256[] memory amounts = new uint256[](2);
         amounts[0] = 1;
         amounts[1] = 1;
@@ -83,13 +74,18 @@ contract TransferERC20ByNonAdminTest is CommunityVaultBaseERC20Test {
 }
 
 contract TransferERC20ByAdminTest is CommunityVaultBaseERC20Test {
+    uint256 depositAmount = 10e18;
+
     function setUp() public virtual override {
         CommunityVaultBaseERC20Test.setUp();
+
+        vm.startPrank(accounts[0]);
+        erc20Token.approve(address(vault), depositAmount);
+        vault.depositERC20(address(erc20Token), depositAmount);
+        vm.stopPrank();
     }
 
     function test_LengthMismatch() public {
-        assertEq(erc20Token.balanceOf(address(vault)), 10e18);
-
         uint256[] memory amounts = new uint256[](1);
         amounts[0] = 5e18;
 
@@ -99,8 +95,6 @@ contract TransferERC20ByAdminTest is CommunityVaultBaseERC20Test {
     }
 
     function test_TransferAmountZero() public {
-        assertEq(erc20Token.balanceOf(address(vault)), 10e18);
-
         uint256[] memory amounts = new uint256[](2);
         amounts[0] = 5e18;
         amounts[1] = 0;
@@ -111,8 +105,6 @@ contract TransferERC20ByAdminTest is CommunityVaultBaseERC20Test {
     }
 
     function test_NoRecipients() public {
-        assertEq(erc20Token.balanceOf(address(vault)), 10e18);
-
         uint256[] memory amounts = new uint256[](0);
         address[] memory tmpAccounts = new address[](0);
 
@@ -123,15 +115,52 @@ contract TransferERC20ByAdminTest is CommunityVaultBaseERC20Test {
 
     function test_AdminCanTransferERC20() public {
         assertEq(erc20Token.balanceOf(address(vault)), 10e18);
+        assertEq(vault.erc20TokenBalances(address(erc20Token)), 10e18);
 
         uint256[] memory amounts = new uint256[](2);
-        amounts[0] = 5e18;
-        amounts[1] = 5e18;
+        amounts[0] = 3e18;
+        amounts[1] = 3e18;
 
         vm.prank(deployer);
         vault.transferERC20(address(erc20Token), accounts, amounts);
 
-        assertEq(erc20Token.balanceOf(address(vault)), 0);
+        assertEq(erc20Token.balanceOf(address(vault)), 4e18);
+        assertEq(vault.erc20TokenBalances(address(erc20Token)), 4e18);
+    }
+
+    function test_TransferERC20AmountTooBig() public {
+        assertEq(erc20Token.balanceOf(address(vault)), 10e18);
+        assertEq(vault.erc20TokenBalances(address(erc20Token)), 10e18);
+
+        uint256[] memory amounts = new uint256[](2);
+        amounts[0] = 10e18;
+        amounts[1] = 10e18;
+
+        vm.prank(deployer);
+        vm.expectRevert(CommunityVault.CommunityVault_ERC20TransferAmountTooBig.selector);
+        vault.transferERC20(address(erc20Token), accounts, amounts);
+    }
+}
+
+contract DepositERC20Test is CommunityVaultBaseERC20Test {
+    function testSuccessfulDepositERC20() public {
+        uint256 depositAmount = 5;
+        uint256 initialVaultBalance = erc20Token.balanceOf(address(vault));
+        uint256 initialTokenBalanceValue = vault.erc20TokenBalances(address(erc20Token));
+
+        vm.startPrank(accounts[0]);
+        erc20Token.approve(address(vault), depositAmount);
+        vault.depositERC20(address(erc20Token), depositAmount);
+        vm.stopPrank();
+
+        assertEq(erc20Token.balanceOf(address(vault)), initialVaultBalance + depositAmount);
+        assertEq(vault.erc20TokenBalances(address(erc20Token)), initialTokenBalanceValue + depositAmount);
+    }
+
+    function testDepositZeroTokens() public {
+        vm.prank(accounts[0]);
+        vm.expectRevert(CommunityVault.CommunityVault_DepositAmountZero.selector);
+        vault.depositERC20(address(erc20Token), 0);
     }
 }