fix: keep track of pending withdrawals

src/MagicSpend.sol

@@ -30,6 +30,9 @@ contract MagicSpend is Ownable, IPaymaster {
         uint48 expiry;
     }
 
+    /// @notice Track the pending ETH withdrwals.
+    uint256 internal _pendingWithdrawals;
+
     /// @notice Track the ETH available to be withdrawn per user.
     mapping(address user => uint256 amount) internal _withdrawableETH;
 
@@ -75,9 +78,9 @@ contract MagicSpend is Ownable, IPaymaster {
     /// @notice Thrown during `UserOperation` validation when the current balance is insufficient to cover the
     ///         requested amount (exluding the `maxGasCost` set by the Entrypoint).
     ///
-    /// @param requestedAmount The requested amount excluding gas.
-    /// @param balance         The current contract balance.
-    error InsufficientBalance(uint256 requestedAmount, uint256 balance);
+    /// @param requestedAmount  The requested amount excluding gas.
+    /// @param availableBalance The current contract balance.
+    error InsufficientAvailableBalance(uint256 requestedAmount, uint256 availableBalance);
 
     /// @notice Thrown when trying to withdraw funds but nothing is available.
     error NoExcess();
@@ -100,6 +103,9 @@ contract MagicSpend is Ownable, IPaymaster {
     /// @param _owner The initial owner of this contract.
     constructor(address _owner) {
         Ownable._initializeOwner(_owner);
+
+        // Initialize it to 1 to save gas and avoid resetting it to 0.
+        _pendingWithdrawals = 1;
     }
 
     /// @notice Receive function allowing ETH to be deposited in this contract.
@@ -122,20 +128,29 @@ contract MagicSpend is Ownable, IPaymaster {
             revert UnsupportedPaymasterAsset(withdrawRequest.asset);
         }
 
+        // Ensure at validation that the contract has enough balance to cover the requested funds.
+        // NOTE: This check is necessary to enforce that the contract will be able to transfer the remaining funds
+        //       when `postOp()` is called back after the `UserOperation` has been executed.
+        uint256 newPendingWithdrawals = _pendingWithdrawals + withdrawAmount;
+        if ((newPendingWithdrawals - 1) > address(this).balance) {
+            uint256 availableBalance = address(this).balance - (_pendingWithdrawals - 1);
+            revert InsufficientAvailableBalance(withdrawAmount, availableBalance);
+        }
+
         _validateRequest(userOp.sender, withdrawRequest);
 
         bool sigFailed = !isValidWithdrawSignature(userOp.sender, withdrawRequest);
         validationData = (sigFailed ? 1 : 0) | (uint256(withdrawRequest.expiry) << 160);
 
-        // Ensure at validation that the contract has enough balance to cover the requested funds.
-        // NOTE: This check is necessary to enforce that the contract will be able to transfer the remaining funds
-        //       when `postOp()` is called back after the `UserOperation` has been executed.
-        if (address(this).balance < withdrawAmount) {
-            revert InsufficientBalance(withdrawAmount, address(this).balance);
-        }
+        // Register the new pending withdrawals.
+        _pendingWithdrawals = newPendingWithdrawals;
 
-        // NOTE: Do not include the gas part in withdrawable funds as it will be handled in `postOp()`.
-        _withdrawableETH[userOp.sender] += withdrawAmount - maxCost;
+        // NOTE: Do not include the gas part in withdrawable funds because it is considered already consumed at this point.
+        //      `postOp()` will refund the amout of gas that was not consumed.
+        uint256 nonGasAmount = withdrawAmount - maxCost;
+        if (nonGasAmount > 0) {
+            _withdrawableETH[userOp.sender] += nonGasAmount;
+        }
         context = abi.encode(maxCost, userOp.sender);
     }
 
@@ -147,18 +162,27 @@ contract MagicSpend is Ownable, IPaymaster {
         // `PostOpMode.postOpReverted` should be impossible.
         // Only possible cause would be if this contract does not own enough ETH to transfer
         // but this is checked at the validation step.
-        assert(mode != PostOpMode.postOpReverted);
+        if (mode == PostOpMode.postOpReverted) {
+            revert UnexpectedPostOpRevertedMode();
+        }
 
         (uint256 maxGasCost, address account) = abi.decode(context, (uint256, address));
 
-        // Compute the total remaining funds available for the user accout.
-        // NOTE: Take into account the user operation gas that was not consumed.
-        uint256 withdrawable = _withdrawableETH[account] + (maxGasCost - actualGasCost);
+        // Compute the total remaining funds the user can withdraw.
+        // NOTE: Take into account the gas consumed cost.
+        uint256 withdrawableAfterGasRefund = _withdrawableETH[account] + (maxGasCost - actualGasCost);
 
-        // Send the all remaining funds to the user accout.
+        // Update storage.
+        if (_pendingWithdrawals > 1) {
+            _pendingWithdrawals = 1;
+        }
         delete _withdrawableETH[account];
-        if (withdrawable > 0) {
-            SafeTransferLib.forceSafeTransferETH(account, withdrawable, SafeTransferLib.GAS_STIPEND_NO_STORAGE_WRITES);
+
+        // Send the all remaining funds to the user accout.
+        if (withdrawableAfterGasRefund > 0) {
+            SafeTransferLib.forceSafeTransferETH(
+                account, withdrawableAfterGasRefund, SafeTransferLib.GAS_STIPEND_NO_STORAGE_WRITES
+            );
         }
     }
 
@@ -168,6 +192,7 @@ contract MagicSpend is Ownable, IPaymaster {
     ///      use cases (e.g., swap or mint).
     function withdrawGasExcess() external {
         uint256 amount = _withdrawableETH[msg.sender];
+
         // we could allow 0 value transfers, but prefer to be explicit
         if (amount == 0) revert NoExcess();
 
@@ -179,15 +204,15 @@ contract MagicSpend is Ownable, IPaymaster {
     ///
     /// @param withdrawRequest The withdraw request.
     function withdraw(WithdrawRequest memory withdrawRequest) external {
-        _validateRequest(msg.sender, withdrawRequest);
+        if (block.timestamp > withdrawRequest.expiry) {
+            revert Expired();
+        }
 
         if (!isValidWithdrawSignature(msg.sender, withdrawRequest)) {
             revert InvalidSignature();
         }
 
-        if (block.timestamp > withdrawRequest.expiry) {
-            revert Expired();
-        }
+        _validateRequest(msg.sender, withdrawRequest);
 
         // reserve funds for gas, will credit user with difference in post op
         _withdraw(withdrawRequest.asset, msg.sender, withdrawRequest.amount);

test/PostOp.t.sol

@@ -11,29 +11,39 @@ contract PostOpTest is PaymasterMagicSpendBaseTest {
 
     function test_transfersExcess(uint256 mode, uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
         mode = bound(mode, 0, 1);
+        amount_ = bound(amount_, 0, type(uint256).max - 1);
         maxCost_ = bound(maxCost_, 0, amount_);
         actualCost = bound(actualCost, 0, maxCost_);
+
+        maxCost = maxCost_;
         amount = amount_;
 
         assertEq(withdrawer.balance, 0);
         vm.deal(address(magic), amount);
         (bytes memory context,) = magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
-        uint256 expectedBalance = amount - actualCost;
+
         magic.postOp(IPaymaster.PostOpMode(mode), context, actualCost);
+
+        uint256 expectedBalance = amount - actualCost;
         assertEq(withdrawer.balance, expectedBalance);
     }
 
     function test_RevertsIfPostOpFailed(uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
+        amount_ = bound(amount_, 0, type(uint256).max - 1);
         maxCost_ = bound(maxCost_, 0, amount_);
         actualCost = bound(actualCost, 0, maxCost_);
+
         amount = amount_;
 
         assertEq(withdrawer.balance, 0);
         vm.deal(address(magic), amount);
+
         (bytes memory context,) = magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
-        uint256 expectedBalance = 0;
-        vm.expectRevert();
+
+        vm.expectRevert(MagicSpend.UnexpectedPostOpRevertedMode.selector);
         magic.postOp(IPaymaster.PostOpMode.postOpReverted, context, actualCost);
+
+        uint256 expectedBalance = 0;
         assertEq(withdrawer.balance, expectedBalance);
     }
 }

test/Validate.t.sol

@@ -18,10 +18,12 @@ abstract contract ValidateTest is MagicSpendTest {
         // avoid precompiles
         vm.assume(withdrawer_ > address(0x10000));
         vm.assume(withdrawer_ != address(vm));
+
         asset = address(0);
         withdrawer = withdrawer_;
         amount = amount_;
         nonce = nonce_;
+
         vm.deal(address(magic), amount);
         vm.expectEmit(true, true, true, true);
         emit MagicSpend.MagicSpendWithdrawal(withdrawer, asset, amount, nonce);

test/ValidatePaymasterUserOp.t.sol

@@ -16,9 +16,34 @@ contract ValidatePaymasterUserOpTest is PaymasterMagicSpendBaseTest, ValidateTes
         magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost);
     }
 
-    function test_revertsIfWithdrawalExceedsBalance() public {
-        vm.deal(address(magic), 0);
-        vm.expectRevert(abi.encodeWithSelector(MagicSpend.InsufficientBalance.selector, amount, 0));
+    function test_revertsInsufficientAvailableBalance(
+        uint256 initialBalance,
+        uint256 firstWithdrawAmount,
+        uint256 excessAmount
+    ) public {
+        initialBalance = bound(initialBalance, 0, type(uint128).max);
+        excessAmount = bound(excessAmount, 1, type(uint128).max);
+
+        firstWithdrawAmount = bound(firstWithdrawAmount, 0, initialBalance);
+        uint256 availableBalance = initialBalance - firstWithdrawAmount;
+        uint256 secondWithdrawAmount = availableBalance + excessAmount;
+
+        vm.deal(address(magic), initialBalance);
+
+        // 1st validation call to increase the `_pendingWithdrawals` traker.
+        amount = firstWithdrawAmount;
+        maxCost = amount;
+        magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost);
+
+        nonce += 1;
+        amount = secondWithdrawAmount;
+        maxCost = amount;
+
+        // 2nd validation call is expected to revert as `secondWithdrawAmount` is above the MagicSpend `availableBalance`
+        // by the fuzzed `excessAmount`.
+        vm.expectRevert(
+            abi.encodeWithSelector(MagicSpend.InsufficientAvailableBalance.selector, amount, availableBalance)
+        );
         magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost);
     }
 
@@ -47,6 +72,7 @@ contract ValidatePaymasterUserOpTest is PaymasterMagicSpendBaseTest, ValidateTes
     }
 
     function test_emitsCorrectly(address, uint256 amount_, uint256 nonce_) public override {
+        amount_ = bound(amount_, 0, type(uint256).max - 1);
         maxCost = amount_;
         super.test_emitsCorrectly(magic.entryPoint(), amount_, nonce_);
     }

test/WithdrawGasExcess.t.sol

@@ -10,9 +10,12 @@ contract WithdrawGasExcess is PaymasterMagicSpendBaseTest {
     }
 
     function test_transferExcess(uint256 amount_, uint256 maxCost_, uint256 actual) public {
-        vm.assume(maxCost_ < amount_);
-        vm.assume(actual <= maxCost_);
+        maxCost_ = bound(maxCost_, 0, type(uint256).max - 2);
+        actual = bound(actual, 0, maxCost_);
+        amount_ = bound(amount_, maxCost_ + 1, type(uint256).max - 1);
+
         amount = amount_;
+
         vm.deal(address(magic), amount);
         magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
 
@@ -24,7 +27,9 @@ contract WithdrawGasExcess is PaymasterMagicSpendBaseTest {
     }
 
     function test_RevertsIfNoExcess(uint256 maxCost_) public {
+        maxCost_ = bound(maxCost_, 0, type(uint256).max - 1);
         amount = maxCost_;
+
         vm.deal(address(magic), amount);
         magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);