Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Guard against concurrent withdraws causing unexpected reverts #17

Merged
merged 5 commits into from
Apr 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 34 additions & 23 deletions .gas-snapshot
Original file line number Diff line number Diff line change
@@ -1,27 +1,38 @@
PostOpTest:test_RevertsIfPostOpFailed(uint256,uint256,uint256) (runs: 256, μ: 94981, ~: 97095)
GetHashTest:test_returnsValidHash() (gas: 46531)
IsValidWithdrawalSignature:test_returnsFalseWithInvalidSignature() (gas: 53845)
IsValidWithdrawalSignature:test_returnsTrueWithValidSignature() (gas: 51016)
OwnerWithdrawTest:test_revertsIfNotOwner() (gas: 15305)
OwnerWithdrawTest:test_transfersERC20Successfully(uint256) (runs: 256, μ: 97271, ~: 99711)
OwnerWithdrawTest:test_transfersETHSuccessfully(uint256) (runs: 256, μ: 51418, ~: 53415)
PostOpTest:test_RevertsIfPostOpFailed(uint256,uint256,uint256) (runs: 256, μ: 97259, ~: 99371)
PostOpTest:test_entryPointAddStake(uint112,uint32) (runs: 256, μ: 56233, ~: 56233)
PostOpTest:test_entryPointDeposit(uint112) (runs: 256, μ: 32331, ~: 32962)
PostOpTest:test_entryPointUnlockStake() (gas: 54910)
PostOpTest:test_entryPointWithdraw(uint112) (runs: 256, μ: 63500, ~: 65300)
PostOpTest:test_entryPointWithdrawStake() (gas: 71825)
PostOpTest:test_paymasterPaysForOp() (gas: 208402)
PostOpTest:test_transfersExcess(uint256,uint256,uint256,uint256) (runs: 256, μ: 110216, ~: 111227)
PostOpTest:test_entryPointDeposit(uint112) (runs: 256, μ: 32479, ~: 32962)
PostOpTest:test_entryPointUnlockStake() (gas: 54888)
PostOpTest:test_entryPointWithdraw(uint112) (runs: 256, μ: 63050, ~: 65300)
PostOpTest:test_entryPointWithdrawStake() (gas: 71803)
PostOpTest:test_paymasterPaysForOp() (gas: 210663)
PostOpTest:test_transfersExcess(uint256,uint256,uint256,uint256) (runs: 256, μ: 112442, ~: 113477)
SetMaxWithdrawDenominator:test_emitsCorrectly(uint256) (runs: 256, μ: 19851, ~: 20191)
SetMaxWithdrawDenominator:test_reverts_whenNotCalledByOwner() (gas: 12257)
SetMaxWithdrawDenominator:test_setsMaxWithdrawPercent(uint256) (runs: 256, μ: 19209, ~: 19549)
Simulate:test() (gas: 12114)
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_emitsCorrectly(address,uint256,uint256) (runs: 256, μ: 104847, ~: 106080)
ValidatePaymasterUserOpTest:test_recordsNonceUsed(uint256) (runs: 256, μ: 112158, ~: 113247)
ValidatePaymasterUserOpTest:test_returns1sIfWrongSignature() (gas: 96298)
ValidatePaymasterUserOpTest:test_returnsCorrectly() (gas: 91185)
ValidatePaymasterUserOpTest:test_revertsIfMaxCostMoreThanRequested() (gas: 39141)
ValidatePaymasterUserOpTest:test_revertsIfNonceUsed() (gas: 101280)
ValidatePaymasterUserOpTest:test_revertsIfWithdrawAssetNotZero() (gas: 55974)
ValidatePaymasterUserOpTest:test_revertsIfNonceUsed() (gas: 103502)
ValidatePaymasterUserOpTest:test_revertsIfWithdrawAssetNotZero() (gas: 55952)
ValidatePaymasterUserOpTest:test_reverts_whenWithdrawExceedsMaxAllowed(uint256,uint256,uint256) (runs: 256, μ: 100806, ~: 100992)
ValidateTest:test_receive() (gas: 14687)
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)
WithdrawGasExcess:test_RevertsIfNoExcess(uint256) (runs: 256, μ: 71133, ~: 71396)
WithdrawGasExcess:test_transferExcess(uint256,uint256,uint256) (runs: 256, μ: 104290, ~: 104290)
WithdrawTest:test_emitsCorrectly(address,uint256,uint256) (runs: 256, μ: 118270, ~: 121368)
WithdrawTest:test_emitsCorrectlyERC20Withdraw(address,uint256,uint256) (runs: 256, μ: 162427, ~: 164773)
WithdrawTest:test_recordsNonceUsed(uint256) (runs: 256, μ: 111372, ~: 112461)
WithdrawTest:test_revertsIfExpired(uint48,uint256) (runs: 256, μ: 61664, ~: 61785)
WithdrawTest:test_revertsIfNonceUsed() (gas: 95844)
WithdrawTest:test_revertsIfWrongSignature() (gas: 63830)
WithdrawTest:test_reverts_whenWithdrawExceedsMaxAllowed(uint256,uint256,uint256) (runs: 256, μ: 85119, ~: 85096)
WithdrawTest:test_transfersERC20Successfully(uint256) (runs: 256, μ: 140626, ~: 143066)
WithdrawTest:test_transfersETHSuccessfully(uint256) (runs: 256, μ: 91692, ~: 93689)
2 changes: 1 addition & 1 deletion script/DeployMagicSpend.s.sol
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ contract MagicSpendDeployScript is Script {
address signerAddress = 0x3E0cd4Dc43811888efa242Ab17118FcE0035EFF7;
uint256 deployerPrivateKey = vm.envUint("PRIVATE_KEY");
vm.startBroadcast(deployerPrivateKey);
MagicSpend c = new MagicSpend{salt: "0x1"}(vm.addr(deployerPrivateKey));
MagicSpend c = new MagicSpend{salt: "0x1"}(vm.addr(deployerPrivateKey), 20);
console2.log(address(c));
c.entryPointDeposit{value: 0.01 ether}(0.01 ether);
c.entryPointAddStake{value: 0x16345785d8a0000}(0x16345785d8a0000, 0x15180);
Expand Down
81 changes: 55 additions & 26 deletions src/MagicSpend.sol
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ contract MagicSpend is Ownable, IPaymaster {
struct WithdrawRequest {
/// @dev The signature associated with this withdraw request.
bytes signature;
/// @dev The asset to withdraw. NOTE: Only ETH (associated with zero address) is supported for now.
/// @dev The asset to withdraw.
address asset;
/// @dev The requested amount to withdraw.
uint256 amount;
Expand All @@ -30,8 +30,14 @@ contract MagicSpend is Ownable, IPaymaster {
uint48 expiry;
}

/// @notice Track the ETH available to be withdrawn per user.
mapping(address user => uint256 amount) internal _withdrawableETH;
/// @notice address(this).balance divided by maxWithdrawDenominator expresses the max WithdrawRequest.amount
/// allowed for a native asset withdraw.
///
/// @dev Helps prevent withdraws in the same transaction leading to reverts and hurting paymaster reputation.
uint256 public maxWithdrawDenominator;

/// @notice Track the amount of native asset available to be withdrawn per user.
mapping(address user => uint256 amount) internal _withdrawable;

/// @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;
Expand All @@ -44,6 +50,11 @@ contract MagicSpend is Ownable, IPaymaster {
/// @param nonce The request nonce.
event MagicSpendWithdrawal(address indexed account, address indexed asset, uint256 amount, uint256 nonce);

/// @notice Emitted when the `maxWithdrawDenominator` is set.
///
/// @param newDenominator The new maxWithdrawDenominator value.
event MaxWithdrawDenominatorSet(uint256 newDenominator);

/// @notice Thrown when the withdraw request signature is invalid.
///
/// @dev The withdraw request signature MUST be:
Expand Down Expand Up @@ -72,12 +83,11 @@ contract MagicSpend is Ownable, IPaymaster {
/// @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).
/// @notice Thrown if WithdrawRequest.amount exceeds address(this).balance / maxWithdrawDenominator.
///
/// @param requestedAmount The requested amount excluding gas.
/// @param balance The current contract balance.
error InsufficientBalance(uint256 requestedAmount, uint256 balance);
/// @param maxAllowed The current max allowed withdraw.
error WithdrawTooLarge(uint256 requestedAmount, uint256 maxAllowed);

/// @notice Thrown when trying to withdraw funds but nothing is available.
error NoExcess();
Expand All @@ -97,9 +107,11 @@ contract MagicSpend is Ownable, IPaymaster {

/// @notice Deploy the contract and set its initial owner.
///
/// @param _owner The initial owner of this contract.
constructor(address _owner) {
Ownable._initializeOwner(_owner);
/// @param owner_ The initial owner of this contract.
/// @param maxWithdrawDenominator_ The initial maxWithdrawDenominator.
constructor(address owner_, uint256 maxWithdrawDenominator_) {
Ownable._initializeOwner(owner_);
_setMaxWithdrawDenominator(maxWithdrawDenominator_);
}

/// @notice Receive function allowing ETH to be deposited in this contract.
Expand Down Expand Up @@ -127,15 +139,8 @@ 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 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;
_withdrawable[userOp.sender] += withdrawAmount - maxCost;
context = abi.encode(maxCost, userOp.sender);
}

Expand All @@ -144,19 +149,23 @@ contract MagicSpend is Ownable, IPaymaster {
external
onlyEntryPoint
{
// `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);
// `PostOpMode.postOpReverted` should never happen.
// The flow here can only revert if there are > maxWithdrawDenominator
// withdraws in the same transaction, which should be highly unlikely.
// If the ETH transfer fails, the entire bundle will revert due an issue in the EntryPoint
// https://github.com/eth-infinitism/account-abstraction/pull/293
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);
uint256 withdrawable = _withdrawable[account] + (maxGasCost - actualGasCost);

// Send the all remaining funds to the user accout.
delete _withdrawableETH[account];
delete _withdrawable[account];
if (withdrawable > 0) {
SafeTransferLib.forceSafeTransferETH(account, withdrawable, SafeTransferLib.GAS_STIPEND_NO_STORAGE_WRITES);
}
Expand All @@ -167,11 +176,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 = _withdrawableETH[msg.sender];
uint256 amount = _withdrawable[msg.sender];
// we could allow 0 value transfers, but prefer to be explicit
if (amount == 0) revert NoExcess();

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

Expand Down Expand Up @@ -249,6 +258,15 @@ contract MagicSpend is Ownable, IPaymaster {
IEntryPoint(entryPoint()).withdrawStake(to);
}

/// @notice Sets maxWithdrawDenominator.
///
/// @dev Reverts if not called by the owner of the contract.
///
/// @param newDenominator The new value for maxWithdrawDenominator.
function setMaxWithdrawDenominator(uint256 newDenominator) external onlyOwner {
_setMaxWithdrawDenominator(newDenominator);
}

/// @notice Returns whether the `withdrawRequest` signature is valid for the given `account`.
///
/// @dev Does not validate nonce or expiry.
Expand Down Expand Up @@ -305,6 +323,12 @@ contract MagicSpend is Ownable, IPaymaster {
return 0x5FF137D4b0FDCD49DcA30c7CF57E578a026d2789;
}

function _setMaxWithdrawDenominator(uint256 newDenominator) internal {
maxWithdrawDenominator = newDenominator;

emit MaxWithdrawDenominatorSet(newDenominator);
}
stevieraykatz marked this conversation as resolved.
Show resolved Hide resolved

/// @notice Validate the `withdrawRequest` against the given `account`.
///
/// @dev Runs all non-signature validation checks.
Expand All @@ -317,6 +341,11 @@ contract MagicSpend is Ownable, IPaymaster {
revert InvalidNonce(withdrawRequest.nonce);
}

uint256 maxAllowed = address(this).balance / maxWithdrawDenominator;
if (withdrawRequest.asset == address(0) && withdrawRequest.amount > maxAllowed) {
revert WithdrawTooLarge(withdrawRequest.amount, maxAllowed);
}

_nonceUsed[withdrawRequest.nonce][account] = true;

// This is emitted ahead of fund transfer, but allows a consolidated code path
Expand Down
2 changes: 1 addition & 1 deletion test/MagicSpend.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ contract MagicSpendTest is Test {
address withdrawer = address(0xb0b);
uint256 ownerPrivateKey = 0xa11ce;
address owner = vm.addr(ownerPrivateKey);
MagicSpend magic = new MagicSpend(owner);
MagicSpend magic = new MagicSpend(owner, 1);

// signature params
address asset;
Expand Down
2 changes: 1 addition & 1 deletion test/OwnerWithdraw.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
pragma solidity >=0.8.21;

import {Ownable} from "./MagicSpend.t.sol";
import {MagicSpendTest} from "./Validate.t.sol";
import {MagicSpendTest} from "./MagicSpend.t.sol";
import {MockERC20} from "solady/test/utils/mocks/MockERC20.sol";

contract OwnerWithdrawTest is MagicSpendTest {
Expand Down
2 changes: 1 addition & 1 deletion test/PostOp.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ contract PostOpTest is PaymasterMagicSpendBaseTest {
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);
assertEq(withdrawer.balance, expectedBalance);
}
Expand Down
26 changes: 26 additions & 0 deletions test/SetMaxWithdrawDenominator.t.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.23;

import "./MagicSpend.t.sol";
import {Ownable} from "./MagicSpend.t.sol";

contract SetMaxWithdrawDenominator is MagicSpendTest {
function test_reverts_whenNotCalledByOwner() public {
vm.prank(makeAddr("fake"));
vm.expectRevert(Ownable.Unauthorized.selector);
magic.setMaxWithdrawDenominator(20);
}

function test_setsMaxWithdrawPercent(uint256 newDenominator) public {
vm.prank(owner);
magic.setMaxWithdrawDenominator(newDenominator);
assertEq(magic.maxWithdrawDenominator(), newDenominator);
}

function test_emitsCorrectly(uint256 newDenominator) public {
vm.expectEmit(false, false, false, true);
emit MagicSpend.MaxWithdrawDenominatorSet(newDenominator);
vm.prank(owner);
magic.setMaxWithdrawDenominator(newDenominator);
}
}
28 changes: 28 additions & 0 deletions test/Validate.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ import {Test, console2} from "forge-std/Test.sol";
import "./MagicSpend.t.sol";

abstract contract ValidateTest is MagicSpendTest {
address invoker;

function test_recordsNonceUsed(uint256 nonce_) public {
nonce = nonce_;
assertFalse(magic.nonceUsed(withdrawer, nonce));
Expand Down Expand Up @@ -36,5 +38,31 @@ abstract contract ValidateTest is MagicSpendTest {
_validateInvokingCall();
}

function test_reverts_whenWithdrawExceedsMaxAllowed(
uint256 accountBalance,
uint256 withdraw,
uint256 maxWithdrawDenominator
) public virtual {
vm.assume(maxWithdrawDenominator > 0);
accountBalance = bound(accountBalance, 0, type(uint256).max - 1);
withdraw = bound(withdraw, accountBalance / maxWithdrawDenominator + 1, type(uint256).max);
amount = withdraw;

vm.stopPrank();
vm.startPrank(owner);
magic.setMaxWithdrawDenominator(maxWithdrawDenominator);
magic.ownerWithdraw(address(0), address(1), address(magic).balance);
vm.stopPrank();

vm.deal(address(magic), accountBalance);
vm.expectRevert(
abi.encodeWithSelector(
MagicSpend.WithdrawTooLarge.selector, withdraw, accountBalance / maxWithdrawDenominator
)
);
vm.startPrank(invoker);
_validateInvokingCall();
}

function _validateInvokingCall() internal virtual;
}
16 changes: 10 additions & 6 deletions test/ValidatePaymasterUserOp.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ contract ValidatePaymasterUserOpTest is PaymasterMagicSpendBaseTest, ValidateTes
function setUp() public override {
super.setUp();
vm.startPrank(magic.entryPoint());
invoker = magic.entryPoint();
}

function test_revertsIfMaxCostMoreThanRequested() public {
Expand All @@ -16,12 +17,6 @@ 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));
magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost);
}

function test_returnsCorrectly() public {
(bytes memory context, uint256 validationData) =
magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost);
Expand Down Expand Up @@ -51,6 +46,15 @@ contract ValidatePaymasterUserOpTest is PaymasterMagicSpendBaseTest, ValidateTes
super.test_emitsCorrectly(magic.entryPoint(), amount_, nonce_);
}

function test_reverts_whenWithdrawExceedsMaxAllowed(
uint256 accountBalance,
uint256 withdraw,
uint256 MaxWithdrawDenominator
) public override {
maxCost = withdraw;
super.test_reverts_whenWithdrawExceedsMaxAllowed(accountBalance, withdraw, MaxWithdrawDenominator);
}

function _validateInvokingCall() internal override {
magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost);
}
Expand Down
1 change: 1 addition & 0 deletions test/Withdraw.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ contract WithdrawTest is MagicSpendTest, ValidateTest {
function setUp() public override {
super.setUp();
vm.startPrank(withdrawer);
invoker = withdrawer;
}

function test_transfersETHSuccessfully(uint256 amount_) public {
Expand Down
2 changes: 1 addition & 1 deletion test/echidna/FuzzSetup.sol
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ contract FuzzSetup is FuzzBase {
uint256 totalWithdrawn;

constructor() payable {
magic = new MagicSpend(OWNER);
magic = new MagicSpend(OWNER, 1);
address(magic).call{value: PAYMASTER_STARTING_BALANCE}("");
}

Expand Down
Loading