feat(protect): add onlyInitiated modifier

src/BaseBundler.sol

@@ -21,6 +21,15 @@ abstract contract BaseBundler is IMulticall {
     /// @dev Also prevents interacting with the bundler outside of an initiated execution context.
     address private _initiator = UNSET_INITIATOR;
 
+    /* MODIFIERS */
+
+    /// @dev Prevents a function to be called outside a `multicall` context.
+    modifier onlyInitiated() {
+        require(_initiator != UNSET_INITIATOR, ErrorsLib.UNINITIATED);
+
+        _;
+    }
+
     /* PUBLIC */
 
     /// @notice Returns the address of the initiator of the multicall transaction.
@@ -57,11 +66,6 @@ abstract contract BaseBundler is IMulticall {
         }
     }
 
-    /// @dev Checks that the contract is in an initiated execution context.
-    function _checkInitiated() internal view {
-        require(_initiator != UNSET_INITIATOR, ErrorsLib.UNINITIATED);
-    }
-
     /// @dev Bubbles up the revert reason / custom error encoded in `returnData`.
     /// @dev Assumes `returnData` is the return data of any kind of failing CALL to a contract.
     function _revert(bytes memory returnData) internal pure {

src/ERC4626Bundler.sol

@@ -72,14 +72,17 @@ abstract contract ERC4626Bundler is BaseBundler {
 
     /// @notice Withdraws the given amount of `assets` from the given ERC4626 `vault`, transferring assets to
     /// `receiver`.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Warning: `vault` can re-enter the bundler flow.
     /// @dev Assumes the given `vault` implements EIP-4626.
     /// @param vault The address of the vault.
     /// @param assets The amount of assets to withdraw. Pass `type(uint256).max` to withdraw max.
     /// @param maxShares The maximum amount of shares to redeem in exchange for `assets`.
     /// @param receiver The address that will receive the withdrawn assets.
-    function erc4626Withdraw(address vault, uint256 assets, uint256 maxShares, address receiver) external payable {
+    function erc4626Withdraw(address vault, uint256 assets, uint256 maxShares, address receiver)
+        external
+        payable
+        onlyInitiated
+    {
         require(receiver != address(0), ErrorsLib.ZERO_ADDRESS);
         /// Do not check `receiver != address(this)` to allow the bundler to receive the underlying asset.
 
@@ -94,14 +97,17 @@ abstract contract ERC4626Bundler is BaseBundler {
     }
 
     /// @notice Redeems the given amount of `shares` from the given ERC4626 `vault`, transferring assets to `receiver`.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Warning: `vault` can re-enter the bundler flow.
     /// @dev Assumes the given `vault` implements EIP-4626.
     /// @param vault The address of the vault.
     /// @param shares The amount of shares to burn. Pass `type(uint256).max` to redeem max.
     /// @param minAssets The minimum amount of assets to withdraw in exchange for `shares`.
     /// @param receiver The address that will receive the withdrawn assets.
-    function erc4626Redeem(address vault, uint256 shares, uint256 minAssets, address receiver) external payable {
+    function erc4626Redeem(address vault, uint256 shares, uint256 minAssets, address receiver)
+        external
+        payable
+        onlyInitiated
+    {
         require(receiver != address(0), ErrorsLib.ZERO_ADDRESS);
         /// Do not check `receiver != address(this)` to allow the bundler to receive the underlying asset.
 

src/MorphoBundler.sol

@@ -121,8 +121,7 @@ abstract contract MorphoBundler is BaseBundler, IMorphoBundler {
         MORPHO.supplyCollateral(marketParams, amount, onBehalf, data);
     }
 
-    /// @notice Borrows `amount` of the loan asset on behalf of the sender.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
+    /// @notice Borrows `amount` of `asset` on behalf of the sender.
     /// @dev Initiator must have previously authorized the bundler to act on their behalf on Morpho.
     /// @param marketParams The Morpho market to borrow assets from.
     /// @param amount The amount of assets to borrow.
@@ -131,6 +130,7 @@ abstract contract MorphoBundler is BaseBundler, IMorphoBundler {
     function morphoBorrow(MarketParams calldata marketParams, uint256 amount, uint256 shares, address receiver)
         external
         payable
+        onlyInitiated
     {
         MORPHO.borrow(marketParams, amount, shares, initiator(), receiver);
     }
@@ -162,7 +162,6 @@ abstract contract MorphoBundler is BaseBundler, IMorphoBundler {
     }
 
     /// @notice Withdraws `amount` of the loan asset on behalf of `onBehalf`.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Initiator must have previously authorized the bundler to act on their behalf on Morpho.
     /// @param marketParams The Morpho market to withdraw assets from.
     /// @param amount The amount of assets to withdraw.
@@ -171,19 +170,20 @@ abstract contract MorphoBundler is BaseBundler, IMorphoBundler {
     function morphoWithdraw(MarketParams calldata marketParams, uint256 amount, uint256 shares, address receiver)
         external
         payable
+        onlyInitiated
     {
         MORPHO.withdraw(marketParams, amount, shares, initiator(), receiver);
     }
 
     /// @notice Withdraws `amount` of the collateral asset on behalf of sender.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Initiator must have previously authorized the bundler to act on their behalf on Morpho.
     /// @param marketParams The Morpho market to withdraw collateral from.
     /// @param amount The amount of collateral to withdraw.
     /// @param receiver The address that will receive the collateral assets.
     function morphoWithdrawCollateral(MarketParams calldata marketParams, uint256 amount, address receiver)
         external
         payable
+        onlyInitiated
     {
         MORPHO.withdrawCollateral(marketParams, amount, initiator(), receiver);
     }
@@ -219,9 +219,7 @@ abstract contract MorphoBundler is BaseBundler, IMorphoBundler {
     /* INTERNAL */
 
     /// @dev Triggers `_multicall` logic during a callback.
-    function _callback(bytes calldata data) internal {
-        _checkInitiated();
-
+    function _callback(bytes calldata data) internal onlyInitiated {
         _multicall(abi.decode(data, (bytes[])));
     }
 

src/Permit2Bundler.sol

@@ -18,14 +18,14 @@ abstract contract Permit2Bundler is BaseBundler {
     /* ACTIONS */
 
     /// @notice Permits and performs a transfer from the initiator to the recipient via Permit2.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Warning: `permit.permitted.token` can re-enter the bundler flow.
     /// @dev Pass `permit.permitted.amount = type(uint256).max` to transfer all.
     /// @param permit The `PermitTransferFrom` struct.
     /// @param signature The signature.
     function permit2TransferFrom(ISignatureTransfer.PermitTransferFrom memory permit, bytes memory signature)
         external
         payable
+        onlyInitiated
     {
         address initiator = initiator();
         uint256 amount = Math.min(permit.permitted.amount, ERC20(permit.permitted.token).balanceOf(initiator));

src/PermitBundler.sol

@@ -12,7 +12,6 @@ import {BaseBundler} from "./BaseBundler.sol";
 abstract contract PermitBundler is BaseBundler {
     /// @notice Permits the given `amount` of `asset` from sender to be spent by the bundler via EIP-2612 Permit with
     /// the given `deadline` & EIP-712 signature's `v`, `r` & `s`.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Warning: `asset` can re-enter the bundler flow.
     /// @param asset The address of the token to be permitted.
     /// @param amount The amount of `asset` to be permitted.
@@ -24,6 +23,7 @@ abstract contract PermitBundler is BaseBundler {
     function permit(address asset, uint256 amount, uint256 deadline, uint8 v, bytes32 r, bytes32 s, bool skipRevert)
         external
         payable
+        onlyInitiated
     {
         try IERC20Permit(asset).permit(initiator(), address(this), amount, deadline, v, r, s) {}
         catch (bytes memory returnData) {

src/TransferBundler.sol

@@ -53,12 +53,11 @@ abstract contract TransferBundler is BaseBundler {
     }
 
     /// @notice Transfers the given `amount` of `asset` from sender to this contract via ERC20 transferFrom.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Warning: `asset` can re-enter the bundler flow.
     /// @dev Pass `amount = type(uint256).max` to transfer all.
     /// @param asset The address of the ERC20 token to transfer.
     /// @param amount The amount of `asset` to transfer from the initiator.
-    function erc20TransferFrom(address asset, uint256 amount) external payable {
+    function erc20TransferFrom(address asset, uint256 amount) external payable onlyInitiated {
         address initiator = initiator();
         amount = Math.min(amount, ERC20(asset).balanceOf(initiator));
 

src/WNativeBundler.sol

@@ -40,10 +40,9 @@ abstract contract WNativeBundler is BaseBundler {
     /* ACTIONS */
 
     /// @notice Wraps the given `amount` of the native token to wNative.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Pass `amount = type(uint256).max` to wrap all.
     /// @param amount The amount of native token to wrap.
-    function wrapNative(uint256 amount) external payable {
+    function wrapNative(uint256 amount) external payable onlyInitiated {
         amount = Math.min(amount, address(this).balance);
 
         require(amount != 0, ErrorsLib.ZERO_AMOUNT);
@@ -52,10 +51,9 @@ abstract contract WNativeBundler is BaseBundler {
     }
 
     /// @notice Unwraps the given `amount` of wNative to the native token.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Pass `amount = type(uint256).max` to unwrap all.
     /// @param amount The amount of wrapped native token to unwrap.
-    function unwrapNative(uint256 amount) external payable {
+    function unwrapNative(uint256 amount) external payable onlyInitiated {
         amount = Math.min(amount, ERC20(WRAPPED_NATIVE).balanceOf(address(this)));
 
         require(amount != 0, ErrorsLib.ZERO_AMOUNT);

src/ethereum/EthereumPermitBundler.sol

@@ -14,7 +14,6 @@ import {PermitBundler} from "../PermitBundler.sol";
 abstract contract EthereumPermitBundler is PermitBundler {
     /// @notice Permits DAI from sender to be spent by the bundler with the given `nonce`, `expiry` & EIP-712
     /// signature's `v`, `r` & `s`.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @param nonce The nonce of the signed message.
     /// @param expiry The expiry of the signed message.
     /// @param allowed Whether the initiator gives the bundler infinite Dai approval or not.
@@ -25,6 +24,7 @@ abstract contract EthereumPermitBundler is PermitBundler {
     function permitDai(uint256 nonce, uint256 expiry, bool allowed, uint8 v, bytes32 r, bytes32 s, bool skipRevert)
         external
         payable
+        onlyInitiated
     {
         try IDaiPermit(MainnetLib.DAI).permit(initiator(), address(this), nonce, expiry, allowed, v, r, s) {}
         catch (bytes memory returnData) {

src/migration/AaveV2MigrationBundler.sol

@@ -31,13 +31,12 @@ contract AaveV2MigrationBundler is MigrationBundler {
     /* ACTIONS */
 
     /// @notice Repays `amount` of `asset` on AaveV2, on behalf of the initiator.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Warning: `asset` can re-enter the bundler flow.
     /// @dev Pass `amount = type(uint256).max` to repay all.
     /// @param asset The address of the token to repay.
     /// @param amount The amount of `asset` to repay.
     /// @param interestRateMode The interest rate mode of the position.
-    function aaveV2Repay(address asset, uint256 amount, uint256 interestRateMode) external payable {
+    function aaveV2Repay(address asset, uint256 amount, uint256 interestRateMode) external payable onlyInitiated {
         if (amount != type(uint256).max) amount = Math.min(amount, ERC20(asset).balanceOf(address(this)));
 
         require(amount != 0, ErrorsLib.ZERO_AMOUNT);

src/migration/AaveV3MigrationBundler.sol

@@ -30,13 +30,12 @@ contract AaveV3MigrationBundler is MigrationBundler {
     /* ACTIONS */
 
     /// @notice Repays `amount` of `asset` on AaveV3, on behalf of the initiator.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Warning: `asset` can re-enter the bundler flow.
     /// @dev Pass `amount = type(uint256).max` to repay all.
     /// @param asset The address of the token to repay.
     /// @param amount The amount of `asset` to repay.
     /// @param interestRateMode The interest rate mode of the position.
-    function aaveV3Repay(address asset, uint256 amount, uint256 interestRateMode) external payable {
+    function aaveV3Repay(address asset, uint256 amount, uint256 interestRateMode) external payable onlyInitiated {
         if (amount != type(uint256).max) amount = Math.min(amount, ERC20(asset).balanceOf(address(this)));
 
         require(amount != 0, ErrorsLib.ZERO_AMOUNT);

src/migration/AaveV3OptimizerMigrationBundler.sol

@@ -30,12 +30,10 @@ contract AaveV3OptimizerMigrationBundler is MigrationBundler {
     /* ACTIONS */
 
     /// @notice Repays `amount` of `underlying` on the AaveV3 Optimizer, on behalf of the initiator.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Warning: `underlying` can re-enter the bundler flow.
-    /// @dev Pass `amount = type(uint256).max` to repay all.
     /// @param underlying The address of the underlying asset to repay.
     /// @param amount The amount of `underlying` to repay.
-    function aaveV3OptimizerRepay(address underlying, uint256 amount) external payable {
+    function aaveV3OptimizerRepay(address underlying, uint256 amount) external payable onlyInitiated {
         if (amount != type(uint256).max) amount = Math.min(amount, ERC20(underlying).balanceOf(address(this)));
 
         require(amount != 0, ErrorsLib.ZERO_AMOUNT);
@@ -46,31 +44,32 @@ contract AaveV3OptimizerMigrationBundler is MigrationBundler {
     }
 
     /// @notice Withdraws `amount` of `underlying` on the AaveV3 Optimizer, on behalf of the initiator`.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Initiator must have previously approved the bundler to manage their AaveV3 Optimizer position.
     /// @dev Pass `amount = type(uint256).max` to withdraw all.
     /// @param underlying The address of the underlying asset to withdraw.
     /// @param amount The amount of `underlying` to withdraw.
     /// @param maxIterations The maximum number of iterations allowed during the matching process. If it is less than
     /// `_defaultIterations.withdraw`, the latter will be used. Pass 0 to fallback to the `_defaultIterations.withdraw`.
-    function aaveV3OptimizerWithdraw(address underlying, uint256 amount, uint256 maxIterations) external payable {
+    function aaveV3OptimizerWithdraw(address underlying, uint256 amount, uint256 maxIterations)
+        external
+        payable
+        onlyInitiated
+    {
         AAVE_V3_OPTIMIZER.withdraw(underlying, amount, initiator(), address(this), maxIterations);
     }
 
     /// @notice Withdraws `amount` of `underlying` used as collateral on the AaveV3 Optimizer, on behalf of the
     /// initiator.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Initiator must have previously approved the bundler to manage their AaveV3 Optimizer position.
     /// @dev Pass `amount = type(uint256).max` to withdraw all.
     /// @param underlying The address of the underlying asset to withdraw.
     /// @param amount The amount of `underlying` to withdraw.
-    function aaveV3OptimizerWithdrawCollateral(address underlying, uint256 amount) external payable {
+    function aaveV3OptimizerWithdrawCollateral(address underlying, uint256 amount) external payable onlyInitiated {
         AAVE_V3_OPTIMIZER.withdrawCollateral(underlying, amount, initiator(), address(this));
     }
 
     /// @notice Approves the bundler to act on behalf of the initiator on the AaveV3 Optimizer, given a signed EIP-712
     /// approval message.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @param isApproved Whether the bundler is allowed to manage the initiator's position or not.
     /// @param nonce The nonce of the signed message.
     /// @param deadline The deadline of the signed message.
@@ -82,7 +81,7 @@ contract AaveV3OptimizerMigrationBundler is MigrationBundler {
         uint256 deadline,
         Signature calldata signature,
         bool skipRevert
-    ) external payable {
+    ) external payable onlyInitiated {
         try AAVE_V3_OPTIMIZER.approveManagerWithSig(initiator(), address(this), isApproved, nonce, deadline, signature)
         {} catch (bytes memory returnData) {
             if (!skipRevert) _revert(returnData);

src/migration/CompoundV2MigrationBundler.sol

@@ -40,12 +40,11 @@ contract CompoundV2MigrationBundler is WNativeBundler, MigrationBundler {
     /* ACTIONS */
 
     /// @notice Repays `amount` of `cToken`'s underlying asset, on behalf of the initiator.
-    /// @notice Warning: should only be called via the bundler's `multicall` function.
     /// @dev Warning: `cToken` can re-enter the bundler flow.
     /// @dev Pass `amount = type(uint256).max` to repay all.
     /// @param cToken The address of the cToken contract
     /// @param amount The amount of `cToken` to repay.
-    function compoundV2Repay(address cToken, uint256 amount) external payable {
+    function compoundV2Repay(address cToken, uint256 amount) external payable onlyInitiated {
         if (cToken == C_ETH) {
             amount = Math.min(amount, address(this).balance);