refactor: dry-fy withdraw function

refactor: dry-fy renounce function
sablier-labs · Sep 7, 2023 · 70e44f8 · 70e44f8
1 parent d4ccd67
commit 70e44f8
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 75 deletions.
diff --git a/src/SablierV2LockupDynamic.sol b/src/SablierV2LockupDynamic.sol
@@ -630,26 +630,10 @@ contract SablierV2LockupDynamic is
 
         // Effects: renounce the stream by making it not cancelable.
         _streams[streamId].isCancelable = false;
-
-        // Interactions: if the recipient is a contract, try to invoke the renounce hook on the recipient without
-        // reverting if the hook is not implemented, and also without bubbling up any potential revert.
-        address recipient = _ownerOf(streamId);
-        if (recipient.code.length > 0) {
-            try ISablierV2LockupRecipient(recipient).onStreamRenounced(streamId) { } catch { }
-        }
-
-        // Log the renouncement.
-        emit ISablierV2Lockup.RenounceLockupStream(streamId);
     }
 
     /// @dev See the documentation for the user-facing functions that call this internal function.
     function _withdraw(uint256 streamId, address to, uint128 amount) internal override {
-        // Checks: the withdraw amount is not greater than the withdrawable amount.
-        uint128 withdrawableAmount = _withdrawableAmountOf(streamId);
-        if (amount > withdrawableAmount) {
-            revert Errors.SablierV2Lockup_Overdraw(streamId, amount, withdrawableAmount);
-        }
-
         // Effects: update the withdrawn amount.
         _streams[streamId].amounts.withdrawn = _streams[streamId].amounts.withdrawn + amount;
 
@@ -668,23 +652,5 @@ contract SablierV2LockupDynamic is
 
         // Interactions: perform the ERC-20 transfer.
         _streams[streamId].asset.safeTransfer({ to: to, value: amount });
-
-        // Retrieve the recipient from storage.
-        address recipient = _ownerOf(streamId);
-
-        // Interactions: if `msg.sender` is not the recipient and the recipient is a contract, try to invoke the
-        // withdraw hook on it without reverting if the hook is not implemented, and also without bubbling up
-        // any potential revert.
-        if (msg.sender != recipient && recipient.code.length > 0) {
-            try ISablierV2LockupRecipient(recipient).onStreamWithdrawn({
-                streamId: streamId,
-                caller: msg.sender,
-                to: to,
-                amount: amount
-            }) { } catch { }
-        }
-
-        // Log the withdrawal.
-        emit ISablierV2Lockup.WithdrawFromLockupStream(streamId, to, amount);
     }
 }
diff --git a/src/SablierV2LockupLinear.sol b/src/SablierV2LockupLinear.sol
@@ -536,26 +536,10 @@ contract SablierV2LockupLinear is
 
         // Effects: renounce the stream by making it not cancelable.
         _streams[streamId].isCancelable = false;
-
-        // Interactions: if the recipient is a contract, try to invoke the renounce hook on the recipient without
-        // reverting if the hook is not implemented, and also without bubbling up any potential revert.
-        address recipient = _ownerOf(streamId);
-        if (recipient.code.length > 0) {
-            try ISablierV2LockupRecipient(recipient).onStreamRenounced(streamId) { } catch { }
-        }
-
-        // Log the renouncement.
-        emit ISablierV2Lockup.RenounceLockupStream(streamId);
     }
 
     /// @dev See the documentation for the user-facing functions that call this internal function.
     function _withdraw(uint256 streamId, address to, uint128 amount) internal override {
-        // Checks: the withdraw amount is not greater than the withdrawable amount.
-        uint128 withdrawableAmount = _withdrawableAmountOf(streamId);
-        if (amount > withdrawableAmount) {
-            revert Errors.SablierV2Lockup_Overdraw(streamId, amount, withdrawableAmount);
-        }
-
         // Effects: update the withdrawn amount.
         _streams[streamId].amounts.withdrawn = _streams[streamId].amounts.withdrawn + amount;
 
@@ -574,23 +558,5 @@ contract SablierV2LockupLinear is
 
         // Interactions: perform the ERC-20 transfer.
         _streams[streamId].asset.safeTransfer({ to: to, value: amount });
-
-        // Retrieve the recipient from storage.
-        address recipient = _ownerOf(streamId);
-
-        // Interactions: if `msg.sender` is not the recipient and the recipient is a contract, try to invoke the
-        // withdraw hook on it without reverting if the hook is not implemented, and also without bubbling up
-        // any potential revert.
-        if (msg.sender != recipient && recipient.code.length > 0) {
-            try ISablierV2LockupRecipient(recipient).onStreamWithdrawn({
-                streamId: streamId,
-                caller: msg.sender,
-                to: to,
-                amount: amount
-            }) { } catch { }
-        }
-
-        // Log the withdrawal.
-        emit ISablierV2Lockup.WithdrawFromLockupStream(streamId, to, amount);
     }
 }
diff --git a/src/abstracts/SablierV2Lockup.sol b/src/abstracts/SablierV2Lockup.sol
@@ -8,6 +8,7 @@ import { IERC721Metadata } from "@openzeppelin/contracts/token/ERC721/extensions
 import { ISablierV2Comptroller } from "../interfaces/ISablierV2Comptroller.sol";
 import { ISablierV2Lockup } from "../interfaces/ISablierV2Lockup.sol";
 import { ISablierV2NFTDescriptor } from "../interfaces/ISablierV2NFTDescriptor.sol";
+import { ISablierV2LockupRecipient } from "../interfaces/hooks/ISablierV2LockupRecipient.sol";
 import { Errors } from "../libraries/Errors.sol";
 import { Lockup } from "../types/DataTypes.sol";
 import { SablierV2Base } from "./SablierV2Base.sol";
@@ -198,8 +199,18 @@ abstract contract SablierV2Lockup is
             revert Errors.SablierV2Lockup_Unauthorized(streamId, msg.sender);
         }
 
-        // Effects: renounce the stream.
+        // Checks and Effects: renounce the stream.
         _renounce(streamId);
+
+        // Interactions: if the recipient is a contract, try to invoke the renounce hook on the recipient without
+        // reverting if the hook is not implemented, and also without bubbling up any potential revert.
+        address recipient = _ownerOf(streamId);
+        if (recipient.code.length > 0) {
+            try ISablierV2LockupRecipient(recipient).onStreamRenounced(streamId) { } catch { }
+        }
+
+        // Log the renouncement.
+        emit ISablierV2Lockup.RenounceLockupStream(streamId);
     }
 
     /// @inheritdoc ISablierV2Lockup
@@ -250,8 +261,32 @@ abstract contract SablierV2Lockup is
             revert Errors.SablierV2Lockup_WithdrawAmountZero(streamId);
         }
 
-        // Checks, Effects and Interactions: make the withdrawal.
+        // Checks: the withdraw amount is not greater than the withdrawable amount.
+        uint128 withdrawableAmount = _withdrawableAmountOf(streamId);
+        if (amount > withdrawableAmount) {
+            revert Errors.SablierV2Lockup_Overdraw(streamId, amount, withdrawableAmount);
+        }
+
+        // Effects and Interactions: make the withdrawal.
         _withdraw(streamId, to, amount);
+
+        // Retrieve the recipient from storage.
+        address recipient = _ownerOf(streamId);
+
+        // Interactions: if `msg.sender` is not the recipient and the recipient is a contract, try to invoke the
+        // withdraw hook on it without reverting if the hook is not implemented, and also without bubbling up
+        // any potential revert.
+        if (msg.sender != recipient && recipient.code.length > 0) {
+            try ISablierV2LockupRecipient(recipient).onStreamWithdrawn({
+                streamId: streamId,
+                caller: msg.sender,
+                to: to,
+                amount: amount
+            }) { } catch { }
+        }
+
+        // Log the withdrawal.
+        emit ISablierV2Lockup.WithdrawFromLockupStream(streamId, to, amount);
     }
 
     /// @inheritdoc ISablierV2Lockup
@@ -270,16 +305,16 @@ abstract contract SablierV2Lockup is
         notNull(streamId)
         updateMetadata(streamId)
     {
-        // Checks: the caller is the current recipient. This also checks that the NFT was not burned.
         address currentRecipient = _ownerOf(streamId);
-        if (msg.sender != currentRecipient) {
-            revert Errors.SablierV2Lockup_Unauthorized(streamId, msg.sender);
-        }
 
         // Skip the withdrawal if the withdrawable amount is zero.
         uint128 withdrawableAmount = _withdrawableAmountOf(streamId);
         if (withdrawableAmount > 0) {
-            _withdraw({ streamId: streamId, to: currentRecipient, amount: withdrawableAmount });
+            withdraw({ streamId: streamId, to: currentRecipient, amount: withdrawableAmount });
+        }
+        // If the amount is zero, check the caller is authorized. This also checks that the NFT was not burned.
+        else if (msg.sender != currentRecipient) {
+            revert Errors.SablierV2Lockup_Unauthorized(streamId, msg.sender);
         }
 
         // Checks and Effects: transfer the NFT.