Skip to content

Commit

Permalink
wip: validate paymaster balance and postop force-transfer
Browse files Browse the repository at this point in the history
  • Loading branch information
@xenoliss
xenoliss committed Mar 11, 2024
1 parent 32e8d6a commit cee2631
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 31 deletions.
32 changes: 15 additions & 17 deletions src/MagicSpend.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import {UserOperation} from "account-abstraction/interfaces/UserOperation.sol";
import {IPaymaster} from "account-abstraction/interfaces/IPaymaster.sol";
import {IEntryPoint} from "account-abstraction/interfaces/IEntryPoint.sol";

import "forge-std/console.sol";

/// @title Magic Spend
///
/// @author Coinbase (https://github.com/coinbase/magic-spend)
Expand Down Expand Up @@ -75,9 +77,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 nonGasFunds The requested amount excluding gas.
/// @param balance The current contract balance.
error InsufficientBalance(uint256 nonGasFunds, uint256 balance);
/// @param requestedAmount The requested amount excluding gas.
/// @param balance The current contract balance.
error InsufficientBalance(uint256 requestedAmount, uint256 balance);

/// @notice Thrown when trying to withdraw funds but nothing is available.
error NoExcess();
Expand Down Expand Up @@ -105,9 +107,10 @@ contract MagicSpend is Ownable, IPaymaster {
returns (bytes memory context, uint256 validationData)
{
WithdrawRequest memory withdrawRequest = abi.decode(userOp.paymasterAndData[20:], (WithdrawRequest));
uint256 withdrawAmount = withdrawRequest.amount;

if (withdrawRequest.amount < maxCost) {
revert RequestLessThanGasMaxCost(withdrawRequest.amount, maxCost);
if (withdrawAmount < maxCost) {
revert RequestLessThanGasMaxCost(withdrawAmount, maxCost);
}

if (withdrawRequest.asset != address(0)) {
Expand All @@ -119,15 +122,15 @@ contract MagicSpend is Ownable, IPaymaster {
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 "non-gas" requested funds.
// 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 nonGasFunds = withdrawRequest.amount - maxCost;
if (nonGasFunds > address(this).balance) {
revert InsufficientBalance(nonGasFunds, address(this).balance);
if (address(this).balance < withdrawAmount) {
revert InsufficientBalance(withdrawAmount, address(this).balance);
}

withdrawableFunds[userOp.sender] += nonGasFunds;
// NOTE: Do not include the gas part in withdrawable funds as it will be handled in `postOp()`.
withdrawableFunds[userOp.sender] += withdrawAmount - maxCost;
context = abi.encode(maxCost, userOp.sender);
}

Expand All @@ -141,14 +144,9 @@ contract MagicSpend is Ownable, IPaymaster {

(uint256 maxGasCost, address account) = abi.decode(context, (uint256, address));

// Withdraw the gas difference from the Entrypoint.
uint256 gasDiff = maxGasCost - actualGasCost;
if (gasDiff > 0) {
IEntryPoint(entryPoint()).withdrawTo(payable(address(this)), gasDiff);
}

// Compute the total remaining funds available for the user accout.
uint256 withdrawable = withdrawableFunds[account] + gasDiff;
// NOTE: Take into account the user operation gas that was not consummed.
uint256 withdrawable = withdrawableFunds[account] + (maxGasCost - actualGasCost);

// Send the all remaining funds to the user accout.
delete withdrawableFunds[account];
Expand Down
30 changes: 16 additions & 14 deletions test/PostOp.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,23 +9,25 @@ contract PostOpTest is PaymasterMagicSpendBaseTest {
vm.startPrank(magic.entryPoint());
}

// function test_transfersExcess(uint256 mode, uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
// mode = bound(mode, 0, 1);
function test_transfersExcess(uint256 mode, uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
mode = bound(mode, 0, 1);
amount_ = bound(amount_, 0, address(magic).balance);
maxCost_ = bound(maxCost_, 0, amount_);
actualCost = bound(actualCost, 0, maxCost_);

assertEq(withdrawer.balance, 0);
(bytes memory context,) = magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
uint256 expectedBalance = amount - actualCost;
vm.deal(address(magic), expectedBalance);
magic.postOp(IPaymaster.PostOpMode(mode), context, actualCost);
assertEq(withdrawer.balance, expectedBalance);
}

// function test_DoesNotTransferIfPostOpFailed(uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
// amount_ = bound(amount_, 0, address(magic).balance);
// maxCost_ = bound(maxCost_, 0, amount_);
// actualCost = bound(actualCost, 0, maxCost_);
// amount = amount_;
// assertEq(withdrawer.balance, 0);
// (bytes memory context,) = magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
// uint256 expectedBalance = amount - actualCost;
// vm.deal(address(magic), expectedBalance);
// magic.postOp(IPaymaster.PostOpMode(mode), context, actualCost);
// assertEq(withdrawer.balance, expectedBalance);
// }

// function test_DoesNotTransferIfPostOpFailed(uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
// vm.assume(maxCost_ <= amount_);
// vm.assume(actualCost <= maxCost_);
// amount = amount_;
// assertEq(withdrawer.balance, 0);
// (bytes memory context,) = magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
// uint256 expectedBalance = 0;
Expand Down

0 comments on commit cee2631

Please sign in to comment.