Skip to content

Commit

Permalink
Merge pull request #8 from xenoliss/cleaning
Browse files Browse the repository at this point in the history 
Improve NatSpec documentation and clean code
  • Loading branch information
@wilsoncusack
wilsoncusack authored Mar 11, 2024
2 parents b3301ad + 1eadab3 commit 37fc113
Show file tree
Hide file tree
Showing 3 changed files with 140 additions and 35 deletions.
167 changes: 136 additions & 31 deletions src/MagicSpend.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,29 +6,73 @@ import {SignatureCheckerLib} from "solady/src/utils/SignatureCheckerLib.sol";
import {SafeTransferLib} from "solady/src/utils/SafeTransferLib.sol";
import {UserOperation} from "account-abstraction/interfaces/UserOperation.sol";
import {IPaymaster} from "account-abstraction/interfaces/IPaymaster.sol";
import {Strings} from "openzeppelin-contracts/contracts/utils/Strings.sol";
import {IEntryPoint} from "account-abstraction/interfaces/IEntryPoint.sol";

/// @title Magic Spend
///
/// @author Coinbase (https://github.com/coinbase/magic-spend)
///
/// @notice ERC4337 Paymaster implementation compatible with Entrypoint v0.6.
///
/// @dev See https://eips.ethereum.org/EIPS/eip-4337#extension-paymasters.
contract MagicSpend is Ownable, IPaymaster {
/// @notice Signed withdraw request allowing accounts to withdraw funds from this contract.
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.
address asset;
/// @dev The requested amount to withdraw.
uint256 amount;
/// @dev Unique nonce used to prevent replays.
uint256 nonce;
/// @dev The maximum expiry the withdraw request remains valid for.
uint48 expiry;
}

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

/// @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.
event MagicSpendWithdrawal(address indexed account, address indexed asset, uint256 amount, uint256 nonce);

/// @notice Thrown when the withdraw request signature is invalid.
///
/// @dev The withdraw request signature MUST be:
/// - an ECDSA signature following EIP-191 (version 0x45)
/// - performed over the content specified in `getHash()`
/// - signed by the current owner of this contract
error InvalidSignature();

/// @notice Thrown when trying to use a withdraw request after its expiry has been reched.
error Expired();

/// @notice Thrown when trying to replay a withdraw request with the same nonce.
///
/// @param nonce The already used nonce.
error InvalidNonce(uint256 nonce);

/// @notice Thrown during validation in the context of ERC4337, when the withraw reques amount is insufficient
/// to sponsor the transaction gas.
///
/// @param requested The withdraw request amount.
/// @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);
error NoPrevalidatedWithdrawForAccount();

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

/// @dev Requires that the caller is the EntryPoint.
Expand All @@ -37,10 +81,14 @@ 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);
}

/// @notice Receive function allowing ETH to be deposited in this contract.
receive() external payable {}

/// @inheritdoc IPaymaster
Expand All @@ -65,7 +113,7 @@ contract MagicSpend is Ownable, IPaymaster {
validationData = (sigFailed ? 1 : 0) | (uint256(withdrawRequest.expiry) << 160);

uint256 excess = withdrawRequest.amount - maxCost;
gasExcessBalance[userOp.sender] += excess;
withdrawableFunds[userOp.sender] += excess;

context = abi.encode(maxCost, userOp.sender);
}
Expand All @@ -79,33 +127,36 @@ contract MagicSpend is Ownable, IPaymaster {

if (mode == IPaymaster.PostOpMode.postOpReverted) {
// we failed to payout the excess, save it so the user can call withdrawGasExcess later
gasExcessBalance[account] += (withheld - actualGasCost);
withdrawableFunds[account] += (withheld - actualGasCost);
return;
}

// credit user difference between actual and withheld
// and unwithdrawn excess
uint256 excess = gasExcessBalance[account] + (withheld - actualGasCost);
delete gasExcessBalance[account];
uint256 excess = withdrawableFunds[account] + (withheld - actualGasCost);
delete withdrawableFunds[account];

if (excess > 0) {
_withdraw(address(0), account, excess);
}
}

/// @dev allows an account, during execution, to withdraw
/// withdrawRequest.amount minus maxCost held for gas
/// allows account to get funding for gas + other purposes with one withdrawRequest
/// @notice Allows the sender to withdraw any available funds associated with him.
///
/// @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 = gasExcessBalance[msg.sender];
uint256 amount = withdrawableFunds[msg.sender];
// we could allow 0 value transfers, but prefer to be explicit
if (amount == 0) revert NoExcess();

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

/// @notice Allows caller to withdraw funds by calling with a valid `withdrawRequest`
///
/// @param withdrawRequest The withdraw request.
function withdraw(WithdrawRequest memory withdrawRequest) external {
_validateRequest(msg.sender, withdrawRequest);

Expand All @@ -121,38 +172,70 @@ contract MagicSpend is Ownable, IPaymaster {
_withdraw(withdrawRequest.asset, msg.sender, withdrawRequest.amount);
}

/// @notice allows owner to withdraw funds
/// @notice Withdraws funds from this contract.
///
/// @dev Reverts if not called by the owner of the contract.
///
/// @param asset The asset to withdraw.
/// @param to The beneficiary address.
/// @param amount The amount to withdraw.
function ownerWithdraw(address asset, address to, uint256 amount) external onlyOwner {
_withdraw(asset, to, amount);
}

/// @notice Deposits ETH from this contract funds to EntryPoint
/// @notice Deposits ETH from this contract funds into the EntryPoint.
///
/// @dev Reverts if not called by the owner of the contract.
///
/// @param amount The amount to deposit on the the Entrypoint.
function entryPointDeposit(uint256 amount) external payable onlyOwner {
SafeTransferLib.safeTransferETH(entryPoint(), amount);
}

/// @notice Withdraws ETH from EntryPoint to `to`
/// @notice Withdraws ETH from the EntryPoint.
///
/// @dev Reverts if not called by the owner of the contract.
///
/// @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);
}

/// @notice Adds stake to EntryPoint
/// @notice Adds stake to the EntryPoint.
///
/// @dev Reverts if not called by the owner of the contract.
///
/// @param amount The amount to stake in the Entrypoint.
/// @param unstakeDelaySeconds XXX
function entryPointAddStake(uint256 amount, uint32 unstakeDelaySeconds) external payable onlyOwner {
IEntryPoint(entryPoint()).addStake{value: amount}(unstakeDelaySeconds);
}

/// @notice Unlocks stake in EntryPoint
/// @notice Unlocks stake in the EntryPoint.
///
/// @dev Reverts if not called by the owner of the contract.
function entryPointUnlockStake() external onlyOwner {
IEntryPoint(entryPoint()).unlockStake();
}

/// @notice Withdraws stake from EntryPoint to `to`
/// @notice Withdraws stake from the EntryPoint.
///
/// @dev Reverts if not called by the owner of the contract.
///
/// @param to The beneficiary address.
function entryPointWithdrawStake(address payable to) external onlyOwner {
IEntryPoint(entryPoint()).withdrawStake(to);
}

/// @notice Returns whether the withdrawRequest signature is valid for the given account
/// @dev Does not validate nonce or expiry
/// @notice Returns whether the `withdrawRequest` signature is valid for the given `account`.
///
/// @dev Does not validate nonce or expiry.
///
/// @param account The account address.
/// @param withdrawRequest The withdraw request.
///
/// @return `true` if the signature is valid, else `false`.
function isValidWithdrawSignature(address account, WithdrawRequest memory withdrawRequest)
public
view
Expand All @@ -163,9 +246,17 @@ contract MagicSpend is Ownable, IPaymaster {
);
}

/// @notice Returns the hash to be signed for a given withdrawRequest
/// @notice Returns the hash to be signed for a given `account` and `withdrawRequest` pair.
///
/// @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 withdrawRequest The withdraw request.
///
/// @return The hash to be signed for the given `account` and `withdrawRequest`.
function getHash(address account, WithdrawRequest memory withdrawRequest) public view returns (bytes32) {
return _toPrefixedMessageHash(
return SignatureCheckerLib.toEthSignedMessageHash(
abi.encode(
address(this),
account,
Expand All @@ -178,22 +269,28 @@ contract MagicSpend is Ownable, IPaymaster {
);
}

/// @notice Returns whether the nonce has been used for the given account
/// @notice Returns whether the `nonce` has been used by the given `account`.
///
/// @param account The account address.
/// @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) {
return _nonceUsed[nonce][account];
}

/// @dev Returns the canonical ERC4337 EntryPoint contract.
/// @notice Returns the canonical ERC-4337 EntryPoint v0.6 contract.
function entryPoint() public pure returns (address) {
return 0x5FF137D4b0FDCD49DcA30c7CF57E578a026d2789;
}

function _toPrefixedMessageHash(bytes memory message) internal pure returns (bytes32) {
return SignatureCheckerLib.toEthSignedMessageHash(message);
}

/// @dev runs all non-signature validation checks
/// signature validation done separately so we can not revert in validatePaymasterUserOp
/// @notice Validate the `withdrawRequest` against the given `account`.
///
/// @dev Runs all non-signature validation checks.
/// @dev Reverts if the withdraw request nonce has already been used.
///
/// @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]) {
revert InvalidNonce(withdrawRequest.nonce);
Expand All @@ -205,6 +302,14 @@ contract MagicSpend is Ownable, IPaymaster {
emit MagicSpendWithdrawal(account, withdrawRequest.asset, withdrawRequest.amount, withdrawRequest.nonce);
}

/// @notice Withdraws funds from this contract.
///
/// @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 amount The amount to withdraw.
function _withdraw(address asset, address to, uint256 amount) internal {
if (asset == address(0)) {
SafeTransferLib.safeTransferETH(to, amount);
Expand Down
2 changes: 1 addition & 1 deletion test/PostOp.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ contract PostOpTest is PaymasterMagicSpendBaseTest {
magic.postOp(IPaymaster.PostOpMode.postOpReverted, context, actualCost);
uint256 expectedBalance = amount - actualCost;
assertEq(withdrawer.balance, 0);
assertEq(magic.gasExcessBalance(withdrawer), expectedBalance);
assertEq(magic.withdrawableFunds(withdrawer), expectedBalance);
vm.stopPrank();
vm.deal(address(magic), expectedBalance);
if (expectedBalance > 0) {
Expand Down
6 changes: 3 additions & 3 deletions test/ValidatePaymasterUserOp.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ contract ValidatePaymasterUserOpTest is PaymasterMagicSpendBaseTest, ValidateTes
amount = amount_;
magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
uint256 expected = amount - maxCost_;
assertEq(magic.gasExcessBalance(withdrawer), expected);
assertEq(magic.withdrawableFunds(withdrawer), expected);
}

function test_doesNotOverwriteExcess(uint256 amount_, uint256 maxCost_, uint256 actualCost) public {
Expand All @@ -56,10 +56,10 @@ contract ValidatePaymasterUserOpTest is PaymasterMagicSpendBaseTest, ValidateTes
vm.assume(expected < type(uint256).max / 2);
amount = amount_;
magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
assertEq(magic.gasExcessBalance(withdrawer), expected);
assertEq(magic.withdrawableFunds(withdrawer), expected);
nonce += 1;
magic.validatePaymasterUserOp(_getUserOp(), userOpHash, maxCost_);
assertEq(magic.gasExcessBalance(withdrawer), expected * 2);
assertEq(magic.withdrawableFunds(withdrawer), expected * 2);
}

function test_emitsCorrectly(address, uint256 amount_, uint256 nonce_) public override {
Expand Down

0 comments on commit 37fc113

Please sign in to comment.