fix: revert in withdraw if the balance is zero

sablier-labs · Apr 19, 2024 · f470114 · f470114
1 parent 59ae9f0
commit f470114
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 22 deletions.
diff --git a/src/SablierV2OpenEnded.sol b/src/SablierV2OpenEnded.sol
@@ -593,6 +593,11 @@ contract SablierV2OpenEnded is ISablierV2OpenEnded, NoDelegateCall, SablierV2Ope
             revert Errors.SablierV2OpenEnded_WithdrawalTimeInTheFuture(time, block.timestamp);
         }
 
+        // Check: the stream balance is not zero.
+        if (_streams[streamId].balance == 0) {
+            revert Errors.SablierV2OpenEnded_WithdrawBalanceZero(streamId);
+        }
+
         // Calculate how much to withdraw based on the time reference.
         uint128 withdrawAmount = _withdrawableAmountOf(streamId, time);
 

diff --git a/src/interfaces/ISablierV2OpenEnded.sol b/src/interfaces/ISablierV2OpenEnded.sol
@@ -308,8 +308,8 @@ interface ISablierV2OpenEnded is ISablierV2OpenEndedState {
     /// - `streamId` must not reference a null or canceled stream.
     /// - `to` must not be the zero address.
     /// - `to` must be the recipient if `msg.sender` is not the stream's recipient.
-    /// - `time` must be greater than the stream's `lastTimeUpdate`.
-    /// - `time` must not be greater than the `block.timestamp`.
+    /// - `time` must be greater than the stream's `lastTimeUpdate` and must not be in the future.
+    /// -  The stream balance must be greater than zero.
     ///
     /// @param streamId The id of the stream to withdraw from.
     /// @param to The address receiving the withdrawn assets.

diff --git a/src/libraries/Errors.sol b/src/libraries/Errors.sol
@@ -71,4 +71,7 @@ library Errors {
 
     /// @notice Thrown when trying to withdraw to the zero address.
     error SablierV2OpenEnded_WithdrawToZeroAddress();
+
+    /// @notice Thrown when trying to withdraw but the stream balance is zero.
+    error SablierV2OpenEnded_WithdrawBalanceZero(uint256 streamId);
 }
diff --git a/test/integration/withdraw/withdraw.t.sol b/test/integration/withdraw/withdraw.t.sol
@@ -120,6 +120,24 @@ contract Withdraw_Integration_Test is Integration_Test {
         openEnded.withdraw({ streamId: defaultStreamId, to: users.recipient, time: futureTime });
     }
 
+    function test_RevertWhen_BalanceZero()
+        external
+        whenNotDelegateCalled
+        givenNotNull
+        givenNotCanceled
+        whenToNonZeroAddress
+        whenWithdrawalAddressIsRecipient
+        whenWithdrawalTimeGreaterThanLastUpdate
+        whenWithdrawalTimeNotInTheFuture
+    {
+        vm.warp({ newTimestamp: WARP_ONE_MONTH - ONE_MONTH });
+        uint256 streamId = createDefaultStream();
+        vm.warp({ newTimestamp: WARP_ONE_MONTH });
+
+        vm.expectRevert(abi.encodeWithSelector(Errors.SablierV2OpenEnded_WithdrawBalanceZero.selector, streamId));
+        openEnded.withdraw({ streamId: streamId, to: users.recipient, time: WITHDRAW_TIME });
+    }
+
     function test_Withdraw_CallerSender()
         external
         whenNotDelegateCalled

diff --git a/test/integration/withdraw/withdraw.tree b/test/integration/withdraw/withdraw.tree
@@ -23,21 +23,24 @@ withdraw.t.sol
                   ├── when the withdrawal time is in the future
                   │  └── it should revert
                   └── when the withdrawal time is not in the future
-                     ├── when the caller is not the recipient
-                     │  ├── when the caller is the sender
-                     │  │  └── it should make the withdrawal
-                     │  └── when the caller is unknown
-                     │     └── it should make the withdrawal
-                     └── when the caller is the recipient
-                           ├── given the asset does not have 18 decimals
-                           │  ├── it should make the withdrawal
-                           │  ├── it should update the time
-                           │  ├── it should update the stream balance
-                           │  ├── it should perform the ERC-20 transfer
-                           │  └── it should emit a {Transfer} and {WithdrawFromOpenEndedStream} event
-                           └── given the asset has 18 decimals
-                              ├── it should make the withdrawal   
-                              ├── it should update the time
-                              ├── it should update the stream balance
-                              ├── it should perform the ERC-20 transfer
-                              └── it should emit a {Transfer} and {WithdrawFromOpenEndedStream} event
+                     ├── when the balance is zero
+                     │  └── it should revert
+                     └── when the balance is not zero
+                        ├── when the caller is not the recipient
+                        │  ├── when the caller is the sender
+                        │  │  └── it should make the withdrawal
+                        │  └── when the caller is unknown
+                        │     └── it should make the withdrawal
+                        └── when the caller is the recipient
+                              ├── given the asset does not have 18 decimals
+                              │  ├── it should make the withdrawal
+                              │  ├── it should update the time
+                              │  ├── it should update the stream balance
+                              │  ├── it should perform the ERC-20 transfer
+                              │  └── it should emit a {Transfer} and {WithdrawFromOpenEndedStream} event
+                              └── given the asset has 18 decimals
+                                 ├── it should make the withdrawal   
+                                 ├── it should update the time
+                                 ├── it should update the stream balance
+                                 ├── it should perform the ERC-20 transfer
+                                 └── it should emit a {Transfer} and {WithdrawFromOpenEndedStream} event
diff --git a/test/invariant/handlers/OpenEndedHandler.sol b/test/invariant/handlers/OpenEndedHandler.sol
@@ -101,13 +101,11 @@ contract OpenEndedHandler is BaseHandler {
     }
 
     function deposit(
-        uint256 timeJumpSeed,
         uint256 streamIndexSeed,
         uint128 depositAmount
     )
         external
         instrument("deposit")
-        adjustTimestamp(timeJumpSeed)
         useFuzzedStream(streamIndexSeed)
         useFuzzedStreamSender
     {
@@ -245,6 +243,10 @@ contract OpenEndedHandler is BaseHandler {
             return;
         }
 
+        if (openEnded.getBalance(currentStreamId) == 0) {
+            return;
+        }
+
         // Bound the time so that it is between last time update and current time.
         time = uint40(_bound(time, openEnded.getLastTimeUpdate(currentStreamId) + 1, block.timestamp));