add setting of balance to 0 after retrieval

yieldnest · Apr 12, 2024 · c78843c · c78843c
1 parent d2c4131
commit c78843c
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 8 deletions.
diff --git a/src/PooledDeposits.sol b/src/PooledDeposits.sol
@@ -1,5 +1,8 @@
 import "./interfaces/IynETH.sol";
 
+error DepositsPeriodNotEnded();
+error DepositMustBeGreaterThanZero();
+
 contract PooledDeposits {
     mapping(address => uint256) public balances;
     uint256 public depositEndTime;
@@ -15,18 +18,22 @@ contract PooledDeposits {
     }
 
     function deposit() public payable {
-        require(block.timestamp <= depositEndTime, "Deposits period has not ended");
-        require(msg.value > 0, "Deposit must be greater than 0");
+        if (block.timestamp > depositEndTime) revert DepositsPeriodNotEnded();
+        if (msg.value == 0) revert DepositMustBeGreaterThanZero();
         balances[msg.sender] += msg.value;
         emit DepositReceived(msg.sender, msg.value);
     }
 
     function finalizeDeposits(address[] calldata depositors) external {
-        require(block.timestamp > depositEndTime, "Deposits period has not ended");
+        if (block.timestamp <= depositEndTime) revert DepositsPeriodNotEnded();
 
         for (uint i = 0; i < depositors.length; i++) {
             address depositor = depositors[i];
             uint256 depositAmountPerDepositor = balances[depositor];
+            if (depositAmountPerDepositor == 0) {
+                continue;
+            }
+            balances[depositor] = 0;
             uint256 shares = ynETH.depositETH{value: depositAmountPerDepositor}(depositor);
             emit DepositsFinalized(depositor, depositAmountPerDepositor, shares);
         }

diff --git a/test/integration/PooledDeposits.t.sol b/test/integration/PooledDeposits.t.sol
@@ -84,7 +84,7 @@ contract PooledDepositsTest is IntegrationBaseTest {
         vm.warp(block.timestamp + 1); // Move time forward to block deposits
 
         // Act & Assert
-        vm.expectRevert("Deposits period has not ended");
+        vm.expectRevert(DepositsPeriodNotEnded.selector);
         pooledDeposits.deposit{value: depositAmount}();
     }
 
@@ -95,7 +95,7 @@ contract PooledDepositsTest is IntegrationBaseTest {
         depositors[0] = address(this);
 
         // Act & Assert
-        vm.expectRevert("Deposits period has not ended");
+        vm.expectRevert(DepositsPeriodNotEnded.selector);
         pooledDeposits.finalizeDeposits(depositors);
     }
 
@@ -104,7 +104,7 @@ contract PooledDepositsTest is IntegrationBaseTest {
         PooledDeposits pooledDeposits = new PooledDeposits(IynETH(address(yneth)), block.timestamp + 1 days);
 
         // Act & Assert
-        vm.expectRevert("Deposit must be greater than 0");
+        vm.expectRevert(DepositMustBeGreaterThanZero.selector);
         pooledDeposits.deposit{value: 0}();
     }
 
@@ -147,16 +147,27 @@ contract PooledDepositsTest is IntegrationBaseTest {
         pooledDeposits.deposit{value: depositAmount}();
 
         // Act
-        vm.warp(block.timestamp + 1 days); // Move time forward to allow finalizing deposits
+        vm.warp(block.timestamp + 1 days + 1); // Move time forward to allow finalizing deposits
         pooledDeposits.finalizeDeposits(depositors);
 
         // Assert first finalize
         uint256 sharesAfterFirstFinalize = IynETH(address(yneth)).balanceOf(address(this));
         assertTrue(sharesAfterFirstFinalize > 0, "Shares should be allocated after first finalize");
 
+        // Check balance after first finalize
+        uint256 balanceAfterFirstFinalize = pooledDeposits.balances(address(this));
+        assertEq(balanceAfterFirstFinalize, 0, "Balance should be 0 after first finalize");
+
         // Attempt to finalize again
-        vm.expectRevert("Deposits already finalized for user");
         pooledDeposits.finalizeDeposits(depositors);
+
+        // Assert second finalize
+        uint256 sharesAfterSecondFinalize = IynETH(address(yneth)).balanceOf(address(this));
+        assertEq(sharesAfterFirstFinalize, sharesAfterSecondFinalize, "Shares should not change after second finalize");
+
+        // Check balance after second finalize
+        uint256 balanceAfterSecondFinalize = pooledDeposits.balances(address(this));
+        assertEq(balanceAfterSecondFinalize, 0, "Balance should remain 0 after second finalize");
     }
 
     receive() external payable {}