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 <[email protected]> --------- Co-authored-by: Josh Levine <[email protected]> * fix: separate deposit cap tracking from contract balance changes * chore: move interaction below effects for deposit --------- Co-authored-by: Steven Valeri <[email protected]>
sense-finance · Sep 4, 2024 · 31f96c5 · 31f96c5
1 parent e285244
commit 31f96c5
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 4 deletions.
diff --git 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
@@ -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);