Custom paymaster

src/CustomPaymaster.sol

@@ -0,0 +1,125 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.0;
+
+import "account-abstraction/interfaces/IEntryPoint.sol";
+import "account-abstraction/interfaces/IPaymaster.sol";
+import "account-abstraction/interfaces/PackedUserOperation.sol";
+import "solady/auth/Ownable.sol";
+
+import "solady/tokens/ERC20.sol";
+import "solady/utils/ECDSA.sol";
+
+import {SignatureCheckerLib} from "solady/utils/SignatureCheckerLib.sol";
+
+contract CustomPaymaster is IPaymaster, Ownable {
+
+    struct TokenData {
+        bytes signature;
+        address token;
+        uint256 amount;
+    }
+
+    using ECDSA for bytes32;
+
+    IEntryPoint public immutable entryPoint;
+
+    error InvalidSignature();
+    error ZeroAddress();
+    error ZeroAmount();
+    error InsufficientPaymasterBalance();
+    error InvalidPaymasterAndData();
+
+    modifier onlyEntryPoint() virtual {
+        if (msg.sender != address(entryPoint)) {
+            revert Unauthorized();
+        }
+        _;
+    }
+
+    constructor(IEntryPoint _entryPoint, address _owner) {
+        if (address(_entryPoint) == address(0) || _owner == address(0)) {
+            revert ZeroAddress();
+        }
+
+        entryPoint = _entryPoint;
+
+        _initializeOwner(_owner);
+    }
+
+    receive() external payable {}
+
+    function validatePaymasterUserOp(PackedUserOperation calldata userOp, bytes32 userOpHash, uint256 maxCost)
+        external
+        override
+        onlyEntryPoint
+        returns (bytes memory context, uint256 validationData)
+    {
+        TokenData memory tokenData = _decodePaymasterAndData(userOp.paymasterAndData);
+
+        bool success =
+            SignatureCheckerLib.isValidSignatureNow(owner(), getHash(userOp.sender, tokenData), tokenData.signature);
+
+        // if (address(this).balance < maxCost) {
+        //     revert InsufficientPaymasterBalance();
+        // }
+
+        if (success) {
+            ERC20(tokenData.token).transfer(userOp.sender, tokenData.amount);
+
+            validationData = 0;
+        } else {
+            validationData = 1;
+        }
+
+        return (abi.encode(userOp.sender, tokenData.token, tokenData.amount), validationData);
+    }
+
+    function postOp(PostOpMode mode, bytes calldata context, uint256 actualGasCost, uint256 actualUserOpFeePerGas)
+        external
+        override
+        onlyEntryPoint
+    {
+        (address sender, address token, uint256 amount) = abi.decode(context, (address, address, uint256));
+
+        if (mode == PostOpMode.opSucceeded) {} else {}
+
+        // if (address(this).balance < actualGasCost) {
+        //     revert InsufficientPaymasterBalance();
+        // }
+    }
+
+    function deposit() external payable {
+        if (msg.value == 0) {
+            revert ZeroAmount();
+        }
+
+        entryPoint.depositTo{ value: msg.value }(address(this));
+    }
+
+    function withdraw(uint256 amount, address payable recipient) external onlyOwner {
+        if (address(this).balance < amount) {
+            revert InsufficientPaymasterBalance();
+        }
+
+        recipient.transfer(amount);
+    }
+
+    function _decodePaymasterAndData(bytes calldata paymasterAndData)
+        internal
+        pure
+        returns (TokenData memory tokenData)
+    {
+        if (paymasterAndData.length <= 72) {
+            revert InvalidPaymasterAndData();
+        }
+
+        tokenData = abi.decode(paymasterAndData[52:], (TokenData));
+    }
+
+    function getHash(address account, TokenData memory tokenData) public view returns (bytes32) {
+        return SignatureCheckerLib.toEthSignedMessageHash(
+            abi.encode(address(this), account, block.chainid, tokenData.token, tokenData.amount)
+        );
+    }
+
+}

test/CustomPaymaster.t.sol

@@ -0,0 +1,122 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.26;
+
+import {IEntryPoint} from "account-abstraction/interfaces/IEntryPoint.sol";
+import {UserOperationLib} from "account-abstraction/core/UserOperationLib.sol";
+import {Test} from "forge-std/Test.sol";
+
+import "src/CustomPaymaster.sol";
+import {EntryPointLib} from "test/util/ERC4337Test.sol";
+import {SignatureCheckerLib} from "solady/utils/SignatureCheckerLib.sol";
+
+import "solady/tokens/ERC20.sol";
+
+contract MockERC20 is ERC20 {
+
+    constructor() {
+        _mint(msg.sender, 1_000_000 * 10 ** 18);
+    }
+
+    function name() public view override returns (string memory) {
+        return "Token";
+    }
+
+    function symbol() public view override returns (string memory) {
+        return "TOKEN";
+    }
+
+}
+
+contract CustomPaymasterTest is Test {
+
+    IEntryPoint public entrypoint;
+
+    CustomPaymaster public paymaster;
+    MockERC20 public token;
+
+    address user = address(0x9ABC);
+    uint256 ownerPrivateKey = 0xa11ce;
+    address owner = vm.addr(ownerPrivateKey);
+
+    function setUp() public {
+        entrypoint = IEntryPoint(EntryPointLib.deploy());
+
+        token = new MockERC20();
+        paymaster = new CustomPaymaster(entrypoint, owner);
+        vm.deal(address(paymaster), 10 ether);
+        token.transfer(address(paymaster), 10_000 * 10 ** 18);
+    }
+
+    function getDefaultUserOp() public returns (PackedUserOperation memory userOp) {
+        userOp = PackedUserOperation({
+            sender: address(0),
+            nonce: 0,
+            initCode: "",
+            callData: "",
+            accountGasLimits: bytes32(abi.encodePacked(uint128(2e6), uint128(2e6))),
+            preVerificationGas: 2e6,
+            gasFees: bytes32(abi.encodePacked(uint128(2e6), uint128(2e6))),
+            paymasterAndData: bytes(""),
+            signature: abi.encodePacked(hex"41414141")
+        });
+    }
+
+    function encode(
+        PackedUserOperation memory userOp
+    ) internal pure returns (bytes memory ret) {
+        address sender = userOp.sender;
+        uint256 nonce = userOp.nonce;
+        bytes32 hashInitCode = keccak256(userOp.initCode);
+        bytes32 hashCallData = keccak256(userOp.callData);
+        bytes32 accountGasLimits = userOp.accountGasLimits;
+        uint256 preVerificationGas = userOp.preVerificationGas;
+        bytes32 gasFees = userOp.gasFees;
+        bytes32 hashPaymasterAndData = keccak256(userOp.paymasterAndData);
+
+        return abi.encode(
+            sender, nonce,
+            hashInitCode, hashCallData,
+            accountGasLimits, preVerificationGas, gasFees,
+            hashPaymasterAndData
+        );
+    }
+
+    function hash(
+        PackedUserOperation memory userOp
+    ) internal pure returns (bytes32) {
+        return keccak256(encode(userOp));
+    }
+
+    function testValidatePaymasterUserOp() public {
+        PackedUserOperation memory userOp = getDefaultUserOp();
+        uint256 amount = 1000 * 10 ** 18;
+        userOp.sender = user;
+
+        {
+            bytes memory message = abi.encode(address(paymaster), user, block.chainid, address(token), amount);
+            (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPrivateKey, SignatureCheckerLib.toEthSignedMessageHash(message));
+            bytes memory signature = abi.encodePacked(r, s, v);
+
+            CustomPaymaster.TokenData memory tokenData = CustomPaymaster.TokenData({
+                token: address(token),
+                amount: amount,
+                signature: signature
+            });
+
+            bytes memory paymasterAndData = abi.encodePacked(address(paymaster), uint128(0), uint128(0), abi.encode(tokenData));
+            userOp.paymasterAndData = paymasterAndData;
+        }
+
+        bytes32 userOpHash = hash(userOp);
+
+        vm.prank(address(entrypoint));
+        (bytes memory context, uint256 validationData) = paymaster.validatePaymasterUserOp(userOp, userOpHash, 1 ether);
+
+        (address sender, address tokenAddr, uint256 tokenAmount) = abi.decode(context, (address, address, uint256));
+
+        assertEq(sender, user);
+        assertEq(tokenAddr, address(token));
+        assertEq(tokenAmount, amount);
+    }
+
+}