Skip to content

Commit

Permalink
Merge pull request #12 from coinbase/postop-refactor
Browse files Browse the repository at this point in the history 
Postop refactor
  • Loading branch information
@wilsoncusack
wilsoncusack authored Mar 14, 2024
2 parents f4ce575 + b71735a commit 9e040b8
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 86 deletions.
47 changes: 22 additions & 25 deletions .gas-snapshot
View file
Original file line number Diff line number Diff line change
@@ -1,30 +1,27 @@
PostOpTest:test_DoesNotTransferIfPostOpFailed(uint256,uint256,uint256) (runs: 256, μ: 94728, ~: 94980)
PostOpTest:test_PersistsExcessIfPostOpFailed(uint256,uint256,uint256) (runs: 256, μ: 110321, ~: 110735)
PostOpTest:test_RevertsIfPostOpFailed(uint256,uint256,uint256) (runs: 256, μ: 94981, ~: 97095)
PostOpTest:test_entryPointAddStake(uint112,uint32) (runs: 256, μ: 56233, ~: 56233)
PostOpTest:test_entryPointDeposit(uint112) (runs: 256, μ: 32479, ~: 32962)
PostOpTest:test_entryPointDeposit(uint112) (runs: 256, μ: 32331, ~: 32962)
PostOpTest:test_entryPointUnlockStake() (gas: 54910)
PostOpTest:test_entryPointWithdraw(uint112) (runs: 256, μ: 63350, ~: 65300)
PostOpTest:test_entryPointWithdraw(uint112) (runs: 256, μ: 63500, ~: 65300)
PostOpTest:test_entryPointWithdrawStake() (gas: 71825)
PostOpTest:test_paymasterPaysForOp() (gas: 208461)
PostOpTest:test_transfersExcess(uint256,uint256,uint256,uint256) (runs: 256, μ: 112688, ~: 113692)
PostOpTest:test_paymasterPaysForOp() (gas: 208402)
PostOpTest:test_transfersExcess(uint256,uint256,uint256,uint256) (runs: 256, μ: 110216, ~: 111227)
Simulate:test() (gas: 12114)
ValidatePaymasterUserOpTest:test_doesNotOverwriteExcess(uint256,uint256,uint256) (runs: 256, μ: 155887, ~: 158210)
ValidatePaymasterUserOpTest:test_emitsCorrectly(address,uint256,uint256) (runs: 256, μ: 102634, ~: 103867)
ValidatePaymasterUserOpTest:test_recordsNonceUsed(uint256) (runs: 256, μ: 109968, ~: 111057)
ValidatePaymasterUserOpTest:test_returns1sIfWrongSignature() (gas: 94108)
ValidatePaymasterUserOpTest:test_returnsCorrectly() (gas: 88974)
ValidatePaymasterUserOpTest:test_revertsIfMaxCostMoreThanRequested() (gas: 39164)
ValidatePaymasterUserOpTest:test_revertsIfNonceUsed() (gas: 101328)
ValidatePaymasterUserOpTest:test_revertsIfWithdrawAssetNotZero() (gas: 55990)
ValidatePaymasterUserOpTest:test_setsExcess(uint256,uint256,uint256) (runs: 256, μ: 89980, ~: 91165)
ValidatePaymasterUserOpTest:test_emitsCorrectly(address,uint256,uint256) (runs: 256, μ: 102564, ~: 103835)
ValidatePaymasterUserOpTest:test_recordsNonceUsed(uint256) (runs: 256, μ: 109936, ~: 111025)
ValidatePaymasterUserOpTest:test_returns1sIfWrongSignature() (gas: 94076)
ValidatePaymasterUserOpTest:test_returnsCorrectly() (gas: 88920)
ValidatePaymasterUserOpTest:test_revertsIfMaxCostMoreThanRequested() (gas: 39141)
ValidatePaymasterUserOpTest:test_revertsIfNonceUsed() (gas: 101280)
ValidatePaymasterUserOpTest:test_revertsIfWithdrawAssetNotZero() (gas: 55974)
ValidateTest:test_receive() (gas: 14687)
WithdrawGasExcess:test_RevertsIfNoExcess(uint256) (runs: 256, μ: 68946, ~: 69209)
WithdrawGasExcess:test_transferExcess(uint256,uint256,uint256) (runs: 256, μ: 102106, ~: 102106)
WithdrawTest:test_emitsCorrectly(address,uint256,uint256) (runs: 256, μ: 116015, ~: 119113)
WithdrawTest:test_emitsCorrectlyERC20Withdraw(address,uint256,uint256) (runs: 256, μ: 158226, ~: 160553)
WithdrawTest:test_recordsNonceUsed(uint256) (runs: 256, μ: 109118, ~: 110207)
WithdrawTest:test_revertsIfExpired(uint48,uint256) (runs: 256, μ: 59454, ~: 59575)
WithdrawTest:test_revertsIfNonceUsed() (gas: 93546)
WithdrawTest:test_revertsIfWrongSignature() (gas: 61576)
WithdrawTest:test_transfersERC20Successfully(uint256) (runs: 256, μ: 136425, ~: 138865)
WithdrawTest:test_transfersETHSuccessfully(uint256) (runs: 256, μ: 89438, ~: 91435)
WithdrawGasExcess:test_RevertsIfNoExcess(uint256) (runs: 256, μ: 68914, ~: 69177)
WithdrawGasExcess:test_transferExcess(uint256,uint256,uint256) (runs: 256, μ: 102074, ~: 102074)
WithdrawTest:test_emitsCorrectly(address,uint256,uint256) (runs: 256, μ: 115868, ~: 119073)
WithdrawTest:test_emitsCorrectlyERC20Withdraw(address,uint256,uint256) (runs: 256, μ: 158052, ~: 160513)
WithdrawTest:test_recordsNonceUsed(uint256) (runs: 256, μ: 109078, ~: 110167)
WithdrawTest:test_revertsIfExpired(uint48,uint256) (runs: 256, μ: 59403, ~: 59535)
WithdrawTest:test_revertsIfNonceUsed() (gas: 93506)
WithdrawTest:test_revertsIfWrongSignature() (gas: 61536)
WithdrawTest:test_transfersERC20Successfully(uint256) (runs: 256, μ: 136385, ~: 138825)
WithdrawTest:test_transfersETHSuccessfully(uint256) (runs: 256, μ: 89398, ~: 91395)
75 changes: 45 additions & 30 deletions src/MagicSpend.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,18 +30,18 @@ contract MagicSpend is Ownable, IPaymaster {
uint48 expiry;
}

/// @notice Track the funds available to be withdrawn per user.
mapping(address user => uint256 amount) public withdrawableFunds;
/// @notice Track the ETH available to be withdrawn per user.
mapping(address user => uint256 amount) internal _withdrawableETH;

/// @dev Mappings keeping track of already used nonces per user to prevent replays of withdraw requests.
mapping(uint256 nonce => mapping(address user => bool used)) internal _nonceUsed;

/// @notice Emitted after validating a withdraw request and funds are about to be withdrawn.
///
/// @param account The account address.
/// @param asset The asset withdrawn.
/// @param amount The amount withdrawn.
/// @param nonce The request nonce.
/// @param asset The asset withdrawn.
/// @param amount The amount withdrawn.
/// @param nonce The request nonce.
event MagicSpendWithdrawal(address indexed account, address indexed asset, uint256 amount, uint256 nonce);

/// @notice Thrown when the withdraw request signature is invalid.
Expand All @@ -64,14 +64,21 @@ contract MagicSpend is Ownable, IPaymaster {
/// to sponsor the transaction gas.
///
/// @param requested The withdraw request amount.
/// @param maxCost The max gas cost required by the Entrypoint.
/// @param maxCost The max gas cost required by the Entrypoint.
error RequestLessThanGasMaxCost(uint256 requested, uint256 maxCost);

/// @notice Thrown when the withdraw request asset is not ETH (zero address).
///
/// @param asset The requested asset.
error UnsupportedPaymasterAsset(address asset);

/// @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);

/// @notice Thrown when trying to withdraw funds but nothing is available.
error NoExcess();

Expand Down Expand Up @@ -105,9 +112,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,9 +127,15 @@ contract MagicSpend is Ownable, IPaymaster {
bool sigFailed = !isValidWithdrawSignature(userOp.sender, withdrawRequest);
validationData = (sigFailed ? 1 : 0) | (uint256(withdrawRequest.expiry) << 160);

uint256 excess = withdrawRequest.amount - maxCost;
withdrawableFunds[userOp.sender] += excess;
// 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);
}

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

Expand All @@ -130,18 +144,19 @@ contract MagicSpend is Ownable, IPaymaster {
external
onlyEntryPoint
{
// If `postOp` is called back after a first fail then revert entire `UserOperation`.
if (mode == IPaymaster.PostOpMode.postOpReverted) {
revert UnexpectedPostOpRevertedMode();
}
// `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);

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

// credit user difference between actual and maxCost
// and unwithdrawn excess
uint256 withdrawable = withdrawableFunds[account] + (maxCost - actualGasCost);
delete withdrawableFunds[account];
// Compute the total remaining funds available for the user accout.
// NOTE: Take into account the user operation gas that was not consummed.
uint256 withdrawable = _withdrawableETH[account] + (maxGasCost - actualGasCost);

// Send the all remaining funds to the user accout.
delete _withdrawableETH[account];
if (withdrawable > 0) {
SafeTransferLib.forceSafeTransferETH(account, withdrawable, SafeTransferLib.GAS_STIPEND_NO_STORAGE_WRITES);
}
Expand All @@ -152,11 +167,11 @@ contract MagicSpend is Ownable, IPaymaster {
/// @dev Can be called back during the `UserOperation` execution to sponsor funds for non-gas related
/// use cases (e.g., swap or mint).
function withdrawGasExcess() external {
uint256 amount = withdrawableFunds[msg.sender];
uint256 amount = _withdrawableETH[msg.sender];
// we could allow 0 value transfers, but prefer to be explicit
if (amount == 0) revert NoExcess();

delete withdrawableFunds[msg.sender];
delete _withdrawableETH[msg.sender];
_withdraw(address(0), msg.sender, amount);
}

Expand Down Expand Up @@ -202,7 +217,7 @@ contract MagicSpend is Ownable, IPaymaster {
///
/// @dev Reverts if not called by the owner of the contract.
///
/// @param to The beneficiary address.
/// @param to The beneficiary address.
/// @param amount The amount to withdraw from the Entrypoint.
function entryPointWithdraw(address payable to, uint256 amount) external onlyOwner {
IEntryPoint(entryPoint()).withdrawTo(to, amount);
Expand All @@ -212,8 +227,8 @@ contract MagicSpend is Ownable, IPaymaster {
///
/// @dev Reverts if not called by the owner of the contract.
///
/// @param amount The amount to stake in the Entrypoint.
/// @param unstakeDelaySeconds XXX
/// @param amount The amount to stake in the Entrypoint.
/// @param unstakeDelaySeconds The duration for which the stake cannot be withdrawn.
function entryPointAddStake(uint256 amount, uint32 unstakeDelaySeconds) external payable onlyOwner {
IEntryPoint(entryPoint()).addStake{value: amount}(unstakeDelaySeconds);
}
Expand All @@ -238,7 +253,7 @@ contract MagicSpend is Ownable, IPaymaster {
///
/// @dev Does not validate nonce or expiry.
///
/// @param account The account address.
/// @param account The account address.
/// @param withdrawRequest The withdraw request.
///
/// @return `true` if the signature is valid, else `false`.
Expand All @@ -257,7 +272,7 @@ contract MagicSpend is Ownable, IPaymaster {
/// @dev Returns an EIP-191 compliant Ethereum Signed Message (version 0x45), see
/// https://eips.ethereum.org/EIPS/eip-191.
///
/// @param account The account address.
/// @param account The account address.
/// @param withdrawRequest The withdraw request.
///
/// @return The hash to be signed for the given `account` and `withdrawRequest`.
Expand All @@ -278,7 +293,7 @@ contract MagicSpend is Ownable, IPaymaster {
/// @notice Returns whether the `nonce` has been used by the given `account`.
///
/// @param account The account address.
/// @param nonce The nonce to check.
/// @param nonce The nonce to check.
///
/// @return `true` if the nonce has already been used by the account, else `false`.
function nonceUsed(address account, uint256 nonce) external view returns (bool) {
Expand All @@ -295,7 +310,7 @@ contract MagicSpend is Ownable, IPaymaster {
/// @dev Runs all non-signature validation checks.
/// @dev Reverts if the withdraw request nonce has already been used.
///
/// @param account The account address.
/// @param account The account address.
/// @param withdrawRequest The withdraw request to validate.
function _validateRequest(address account, WithdrawRequest memory withdrawRequest) internal {
if (_nonceUsed[withdrawRequest.nonce][account]) {
Expand All @@ -313,8 +328,8 @@ contract MagicSpend is Ownable, IPaymaster {
/// @dev Callers MUST validate that the withdraw is legitimate before calling this method as
/// no validation is performed here.
///
/// @param asset The asset to withdraw.
/// @param to The beneficiary address.
/// @param asset The asset to withdraw.
/// @param to The beneficiary address.
/// @param amount The amount to withdraw.
function _withdraw(address asset, address to, uint256 amount) internal {
if (asset == address(0)) {
Expand Down
17 changes: 8 additions & 9 deletions test/PostOp.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,27 +14,26 @@ contract PostOpTest is PaymasterMagicSpendBaseTest {
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 = amount - actualCost;
vm.deal(address(magic), expectedBalance);
magic.postOp(IPaymaster.PostOpMode(mode), context, actualCost);
assertEq(withdrawer.balance, expectedBalance);
}

function testRevert_PostOp_WhenCalledWithPostOpRevertedMode(uint256 amount_, uint256 maxCost_, uint256 actualCost)
public
{
vm.assume(maxCost_ <= amount_);
vm.assume(actualCost <= maxCost_);
function test_RevertsIfPostOpFailed(uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
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(MagicSpend.UnexpectedPostOpRevertedMode.selector);
vm.expectRevert();
magic.postOp(IPaymaster.PostOpMode.postOpReverted, context, actualCost);

assertEq(withdrawer.balance, expectedBalance);
}
}
22 changes: 0 additions & 22 deletions test/ValidatePaymasterUserOp.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,28 +40,6 @@ contract ValidatePaymasterUserOpTest is PaymasterMagicSpendBaseTest, ValidateTes
assertEq(uint160(validationData), 1);
}

function test_setsExcess(uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
vm.assume(maxCost_ <= amount_);
vm.assume(actualCost <= maxCost_);
amount = amount_;
magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
uint256 expected = amount - maxCost_;
assertEq(magic.withdrawableFunds(withdrawer), expected);
}

function test_doesNotOverwriteExcess(uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
maxCost_ = bound(maxCost_, 0, amount_);
actualCost = bound(actualCost, 0, maxCost_);
uint256 expected = amount_ - maxCost_;
vm.assume(expected < type(uint256).max / 2);
amount = amount_;
magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
assertEq(magic.withdrawableFunds(withdrawer), expected);
nonce += 1;
magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
assertEq(magic.withdrawableFunds(withdrawer), expected * 2);
}

function test_emitsCorrectly(address, uint256 amount_, uint256 nonce_) public override {
maxCost = amount_;
super.test_emitsCorrectly(magic.entryPoint(), amount_, nonce_);
Expand Down

0 comments on commit 9e040b8

Please sign in to comment.