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.
@@ -130,11 +136,18 @@ contract MagicSpend is Ownable, IPaymaster {
         // 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);
+        uint256 newPendingWithdrawals = _pendingWithdrawals + withdrawAmount;
+
+        if ((newPendingWithdrawals - 1) > address(this).balance) {
+            uint256 availableBalance = address(this).balance - (_pendingWithdrawals - 1);
+            revert InsufficientAvailableBalance(withdrawAmount, availableBalance);
         }
 
-        // NOTE: Do not include the gas part in withdrawable funds as it will be handled in `postOp()`.
+        // Register the new pending withdrawals.
+        _pendingWithdrawals = newPendingWithdrawals;
+
+        // 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.
         _withdrawableETH[userOp.sender] += withdrawAmount - maxCost;
         context = abi.encode(maxCost, userOp.sender);
     }
@@ -147,18 +160,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
+            );
         }
     }
 

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,25 @@ 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 pendingWithdrawals,
+        uint256 withdrawAmount
+    ) public {
+        initialBalance = bound(initialBalance, 0, type(uint128).max - 2);
+        pendingWithdrawals = bound(pendingWithdrawals, 0, initialBalance);
+        uint256 availableBalance = initialBalance - pendingWithdrawals;
+        withdrawAmount = bound(withdrawAmount, availableBalance + 1, type(uint128).max - 1);
+
+        maxCost = withdrawAmount;
+        amount = withdrawAmount;
+
+        vm.deal(address(magic), initialBalance);
+        vm.store(address(magic), bytes32(uint256(0)), bytes32(pendingWithdrawals + 1));
+
+        vm.expectRevert(
+            abi.encodeWithSelector(MagicSpend.InsufficientAvailableBalance.selector, amount, availableBalance)
+        );
         magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost);
     }
 
@@ -47,6 +63,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_);