From 31f96c57deb2265cfc8a2beab609039af55b9ab6 Mon Sep 17 00:00:00 2001
From: Josh Levine <24902242+jparklev@users.noreply.github.com>
Date: Wed, 4 Sep 2024 15:55:12 -0700
Subject: [PATCH] fix: separate deposit cap tracking from contract balance
 changes (#38)

* Feat/mapping comment update (#18)

* feat: update mapping comments

* feat: clean comment

* Update contracts/PointTokenVault.sol

Co-authored-by: Josh Levine <24902242+jparklev@users.noreply.github.com>

---------

Co-authored-by: Josh Levine <24902242+jparklev@users.noreply.github.com>

* fix: separate deposit cap tracking from contract balance changes

* chore: move interaction below effects for deposit

---------

Co-authored-by: Steven Valeri <steven@sense.finance>
---
 contracts/PointTokenVault.sol        | 12 +++++---
 contracts/test/PointTokenVault.t.sol | 41 ++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/contracts/PointTokenVault.sol b/contracts/PointTokenVault.sol
index 6f65b9a..ffa948e 100644
--- a/contracts/PointTokenVault.sol
+++ b/contracts/PointTokenVault.sol
@@ -43,6 +43,8 @@ contract PointTokenVault is UUPSUpgradeable, AccessControlUpgradeable, Multicall
 
     mapping(address => mapping(address => bool)) public trustedReceivers; // owner => delegate => trustedReceivers
 
+    mapping(address => uint256) public totalDeposited; // token => total deposited amount
+
     // Fees
     uint256 public mintFee;
     uint256 public redemptionFee;
@@ -110,25 +112,27 @@ contract PointTokenVault is UUPSUpgradeable, AccessControlUpgradeable, Multicall
         _setFeeCollector(_feeCollector);
     }
 
-    // Rebasing and fee-on-transfer tokens must be wrapped before depositing.
+    // Rebasing and fee-on-transfer tokens must be wrapped before depositing. ie, they are not supported natively.
     function deposit(ERC20 _token, uint256 _amount, address _receiver) public {
         uint256 cap = caps[address(_token)];
 
         if (cap != type(uint256).max) {
-            if (_amount + _token.balanceOf(address(this)) > cap) {
+            if (totalDeposited[address(_token)] + _amount > cap) {
                 revert DepositExceedsCap();
             }
         }
 
-        _token.safeTransferFrom(msg.sender, address(this), _amount);
-
         balances[_receiver][_token] += _amount;
+        totalDeposited[address(_token)] += _amount;
+
+        _token.safeTransferFrom(msg.sender, address(this), _amount);
 
         emit Deposit(msg.sender, _receiver, address(_token), _amount);
     }
 
     function withdraw(ERC20 _token, uint256 _amount, address _receiver) public {
         balances[msg.sender][_token] -= _amount;
+        totalDeposited[address(_token)] -= _amount;
 
         _token.safeTransfer(_receiver, _amount);
 
diff --git a/contracts/test/PointTokenVault.t.sol b/contracts/test/PointTokenVault.t.sol
index 51191dd..f5c74cd 100644
--- a/contracts/test/PointTokenVault.t.sol
+++ b/contracts/test/PointTokenVault.t.sol
@@ -146,6 +146,47 @@ contract PointTokenVaultTest is Test {
         assertEq(pointTokenVault.balances(vitalik, newMockToken), 2e18); // Total 2 tokens deposited
     }
 
+    function test_DepositCapRewardSameAsDeposit() public {
+        // Set up an 18 decimal token as both deposit and reward token
+        MockERC20 token = new MockERC20("Example Token", "EX", 18);
+
+        vm.startPrank(operator);
+        // Set deposit cap for token to 5000
+        pointTokenVault.setCap(address(token), 5000e18);
+
+        // Set token as reward token with 1:1 ratio
+        pointTokenVault.setRedemption(eigenPointsId, token, 1e18, false);
+        vm.stopPrank();
+
+        // Mint tokens to users
+        token.mint(vitalik, 5000e18);
+        token.mint(toly, 2000e18);
+
+        // Vitalik deposits 4000 tokens
+        vm.startPrank(vitalik);
+        token.approve(address(pointTokenVault), 4000e18);
+        pointTokenVault.deposit(token, 4000e18, vitalik);
+        vm.stopPrank();
+
+        // Toly converts 2000 tokens to pTokens
+        vm.startPrank(toly);
+        token.approve(address(pointTokenVault), 2000e18);
+        pointTokenVault.convertRewardsToPTokens(toly, eigenPointsId, 2000e18);
+        vm.stopPrank();
+
+        // Assert current token balance in vault
+        assertEq(token.balanceOf(address(pointTokenVault)), 6000e18);
+
+        // Try to deposit 1000 tokens, which should succeed
+        vm.startPrank(vitalik);
+        token.approve(address(pointTokenVault), 1000e18);
+        pointTokenVault.deposit(token, 1000e18, vitalik);
+        vm.stopPrank();
+
+        // Assert that 5000 tokens have been deposited
+        assertEq(pointTokenVault.balances(vitalik, token), 5000e18);
+    }
+
     function test_DeployPToken() public {
         // Can't deploy the same token twice
         vm.expectRevert(PointTokenVault.PTokenAlreadyDeployed.selector);