diff --git a/src/MagicSpend.sol b/src/MagicSpend.sol index c2023e6..0c4ed51 100644 --- a/src/MagicSpend.sol +++ b/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); diff --git a/test/PostOp.t.sol b/test/PostOp.t.sol index f50cb12..def3ba1 100644 --- a/test/PostOp.t.sol +++ b/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); } } diff --git a/test/Validate.t.sol b/test/Validate.t.sol index d5078bb..5337947 100644 --- a/test/Validate.t.sol +++ b/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); diff --git a/test/ValidatePaymasterUserOp.t.sol b/test/ValidatePaymasterUserOp.t.sol index f6bb052..b9fb087 100644 --- a/test/ValidatePaymasterUserOp.t.sol +++ b/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_); } diff --git a/test/WithdrawGasExcess.t.sol b/test/WithdrawGasExcess.t.sol index 9392664..c9ddb60 100644 --- a/test/WithdrawGasExcess.t.sol +++ b/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_);