Skip to content

Commit

Permalink
Merge pull request #17 from coinbase/wilson/reputation-fix
Browse files Browse the repository at this point in the history
Guard against concurrent withdraws causing unexpected reverts
  • Loading branch information
wilsoncusack authored Apr 4, 2024
2 parents 2662df1 + ab2361b commit 0041c9d
Show file tree
Hide file tree
Showing 11 changed files with 159 additions and 60 deletions.
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);
}

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

0 comments on commit 0041c9d

Please sign in to comment.