From 1fb6bf9c19070db2b42dbfce0bf68585c105cc95 Mon Sep 17 00:00:00 2001 From: Yash Date: Mon, 18 Nov 2024 19:04:29 +0530 Subject: [PATCH 1/3] custom paymaster --- src/CustomPaymaster.sol | 129 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) create mode 100644 src/CustomPaymaster.sol diff --git a/src/CustomPaymaster.sol b/src/CustomPaymaster.sol new file mode 100644 index 0000000..4e149d8 --- /dev/null +++ b/src/CustomPaymaster.sol @@ -0,0 +1,129 @@ +// 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 + view + 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(); + } + + validationData = success ? 0 : 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) { + ERC20(token).transfer(sender, amount); + } + + if (address(this).balance < actualGasCost) { + revert InsufficientPaymasterBalance(); + } + + payable(address(entryPoint)).transfer(actualGasCost); + } + + function deposit() external payable { + if (msg.value == 0) { + revert ZeroAmount(); + } + } + + 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[20:], (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 + ) + ); + } + +} From c52d42ee7f3d4dc53f7b52fc54280363a75e9fb1 Mon Sep 17 00:00:00 2001 From: Yash Date: Mon, 18 Nov 2024 21:22:16 +0530 Subject: [PATCH 2/3] test poc --- test/CustomPaymaster.t.sol | 122 +++++++++++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 test/CustomPaymaster.t.sol diff --git a/test/CustomPaymaster.t.sol b/test/CustomPaymaster.t.sol new file mode 100644 index 0000000..243a148 --- /dev/null +++ b/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), 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); + } + +} From 95f93780a62f4091ed30182620a0c9c5c467d48a Mon Sep 17 00:00:00 2001 From: Yash Date: Tue, 19 Nov 2024 18:20:35 +0530 Subject: [PATCH 3/3] poc --- src/CustomPaymaster.sol | 44 +++--- test/CustomPaymaster.t.sol | 2 +- test/CustomPaymasterPOC.t.sol | 253 ++++++++++++++++++++++++++++++++++ 3 files changed, 274 insertions(+), 25 deletions(-) create mode 100644 test/CustomPaymasterPOC.t.sol diff --git a/src/CustomPaymaster.sol b/src/CustomPaymaster.sol index 4e149d8..4fe7c0b 100644 --- a/src/CustomPaymaster.sol +++ b/src/CustomPaymaster.sol @@ -50,22 +50,26 @@ contract CustomPaymaster is IPaymaster, Ownable { function validatePaymasterUserOp(PackedUserOperation calldata userOp, bytes32 userOpHash, uint256 maxCost) external - view 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 - ); + bool success = + SignatureCheckerLib.isValidSignatureNow(owner(), getHash(userOp.sender, tokenData), tokenData.signature); - if (address(this).balance < maxCost) { - revert InsufficientPaymasterBalance(); - } + // if (address(this).balance < maxCost) { + // revert InsufficientPaymasterBalance(); + // } - validationData = success ? 0 : 1; + 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); } @@ -77,21 +81,19 @@ contract CustomPaymaster is IPaymaster, Ownable { { (address sender, address token, uint256 amount) = abi.decode(context, (address, address, uint256)); - if (mode == PostOpMode.opSucceeded) { - ERC20(token).transfer(sender, amount); - } + if (mode == PostOpMode.opSucceeded) {} else {} - if (address(this).balance < actualGasCost) { - revert InsufficientPaymasterBalance(); - } - - payable(address(entryPoint)).transfer(actualGasCost); + // 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 { @@ -111,18 +113,12 @@ contract CustomPaymaster is IPaymaster, Ownable { revert InvalidPaymasterAndData(); } - tokenData = abi.decode(paymasterAndData[20:], (TokenData)); + 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 - ) + abi.encode(address(this), account, block.chainid, tokenData.token, tokenData.amount) ); } diff --git a/test/CustomPaymaster.t.sol b/test/CustomPaymaster.t.sol index 243a148..83241e2 100644 --- a/test/CustomPaymaster.t.sol +++ b/test/CustomPaymaster.t.sol @@ -103,7 +103,7 @@ contract CustomPaymasterTest is Test { signature: signature }); - bytes memory paymasterAndData = abi.encodePacked(address(paymaster), abi.encode(tokenData)); + bytes memory paymasterAndData = abi.encodePacked(address(paymaster), uint128(0), uint128(0), abi.encode(tokenData)); userOp.paymasterAndData = paymasterAndData; } diff --git a/test/CustomPaymasterPOC.t.sol b/test/CustomPaymasterPOC.t.sol new file mode 100644 index 0000000..0eb84d4 --- /dev/null +++ b/test/CustomPaymasterPOC.t.sol @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +pragma solidity ^0.8.26; + +import {UserOperationLib} from "account-abstraction/core/UserOperationLib.sol"; +import {IEntryPoint} from "account-abstraction/interfaces/IEntryPoint.sol"; +import {PackedUserOperation} from "account-abstraction/interfaces/PackedUserOperation.sol"; +import {Test} from "forge-std/Test.sol"; +import {ECDSA} from "solady/utils/ECDSA.sol"; +import {SignatureCheckerLib} from "solady/utils/SignatureCheckerLib.sol"; + +import {MockTarget} from "test/mock/MockTarget.sol"; +import {MockValidator} from "test/mock/MockValidator.sol"; +import {EntryPointLib} from "test/util/ERC4337Test.sol"; + +import "solady/tokens/ERC20.sol"; +import "solady/tokens/ERC721.sol"; + +import "src/CustomPaymaster.sol"; + +import "lib/forge-std/src/console.sol"; + +import { + CALLTYPE_SINGLE, + CALLTYPE_STATIC, + Execution, + ExecutionLib, + IERC7579Account, + MODULE_TYPE_FALLBACK, + MODULE_TYPE_VALIDATOR, + ModeLib, + ModularAccount +} from "src/ModularAccount.sol"; +import {ModularAccountFactory} from "src/ModularAccountFactory.sol"; + +import {InitializerInstallModule} from "src/interface/IModularAccount.sol"; + +import {DefaultValidator, SessionKey, SessionKeyParams, SessionKeyType} from "src/DefaultValidator.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 MockNFT is ERC721 { + + address public immutable CURRENCY; + uint256 public immutable PRICE_PER_TOKEN; + + constructor(address _currency, uint256 _pricePerToken) { + CURRENCY = _currency; + PRICE_PER_TOKEN = _pricePerToken; + } + + function name() public view override returns (string memory) { + return "Mock NFT"; + } + + function symbol() public view override returns (string memory) { + return "MNFT"; + } + + function tokenURI(uint256 id) public view override returns (string memory) { + return ""; + } + + function claim(address to, uint256 id) external payable { + ERC20(CURRENCY).transferFrom(msg.sender, address(this), PRICE_PER_TOKEN); + + _mint(to, id); + } + +} + +contract CustomPaymasterPOCTest is Test { + + IEntryPoint public entrypoint; + ModularAccountFactory public factory; + CustomPaymaster public paymaster; + MockERC20 public token; + MockNFT public nft; + + bytes32 public accountSalt = bytes32(uint256(0xdeadbeef)); + + uint256 accountOwnerPKey = 0x2; + address public factoryOwner = vm.addr(0x1); + address public accountOwner = vm.addr(accountOwnerPKey); + uint256 paymasterOwnerPrivateKey = 0x3; + address paymasterOwner = vm.addr(paymasterOwnerPrivateKey); + + address public validator; + MockTarget public target; + + address[] public defaultModules; + + uint256 nftPrice = 10 ether; + + bytes32 internal constant _BEACON_SLOT = 0xa3f0ad74e5423aebfd80d3ef4346578335a9a72aeaee59ff6cb3582b35133d50; + bytes32 internal constant _IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc; + + function setUp() public { + vm.prank(factoryOwner); + entrypoint = IEntryPoint(EntryPointLib.deploy()); + + token = new MockERC20(); + nft = new MockNFT(address(token), nftPrice); + paymaster = new CustomPaymaster(entrypoint, paymasterOwner); + + token.transfer(address(paymaster), 10_000 * 10 ** 18); + vm.deal(paymasterOwner, 1000 ether); + vm.startPrank(paymasterOwner); + paymaster.deposit{value: 100 ether}(); + address(paymaster).call{value: 100 ether}(""); + vm.stopPrank(); + + + target = new MockTarget(); + validator = address(new MockValidator()); + + defaultModules.push(validator); + + ModularAccount accountImpl = new ModularAccount(address(entrypoint)); + factory = new ModularAccountFactory(address(entrypoint), factoryOwner, address(accountImpl), defaultModules); + } + + function test_customPaymaster_poc() public { + Execution[] memory executions = new Execution[](2); + executions[0] = Execution({ + target: address(token), + value: 0, + data: abi.encodeCall(ERC20.approve, (address(nft), nftPrice)) + }); + executions[1] = + Execution({target: address(nft), value: 0, data: abi.encodeCall(MockNFT.claim, (address(accountOwner), 0))}); + bytes memory userOpCalldata = + abi.encodeCall(IERC7579Account.execute, (ModeLib.encodeSimpleBatch(), ExecutionLib.encodeBatch(executions))); + + (address acc, bytes memory initCode) = getAccountInitCode(accountOwner, accountSalt); + uint256 nonce = getNonce(acc, address(validator)); + + PackedUserOperation memory userOp = getDefaultUserOp(); + userOp.sender = acc; + userOp.nonce = nonce; + userOp.callData = userOpCalldata; + userOp.initCode = initCode; + + uint256 amount = 1000 * 10 ** 18; + + { + bytes memory message = abi.encode(address(paymaster), acc, block.chainid, address(token), amount); + (uint8 v, bytes32 r, bytes32 s) = + vm.sign(paymasterOwnerPrivateKey, 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(2e6), uint128(2e6), abi.encode(tokenData)); + userOp.paymasterAndData = paymasterAndData; + } + + // Create userOps array + PackedUserOperation[] memory userOps = new PackedUserOperation[](1); + userOps[0] = userOp; + + vm.deal(acc, 1 ether); + entrypoint.handleOps(userOps, payable(acc)); + + assertTrue(nft.ownerOf(0) == accountOwner); + } + + // test utils + + function getAccountInitCode(address owner, bytes32 salt) internal view returns (address, bytes memory) { + InitializerInstallModule[] memory modules = new InitializerInstallModule[](1); + modules[0] = + InitializerInstallModule({moduleTypeId: MODULE_TYPE_VALIDATOR, module: address(validator), initData: ""}); + return getAccountInitCode(owner, salt, modules); + } + + function getAccountInitCode(address owner, bytes32 salt, InitializerInstallModule[] memory modules) + internal + view + returns (address acc, bytes memory initCode) + { + acc = factory.getAddress(owner, salt); + initCode = abi.encodePacked( + address(factory), abi.encodeWithSelector(factory.createAccount.selector, owner, salt, modules) + ); + } + + function getNonce(address account, address _validator) internal returns (uint256 nonce) { + uint192 key = uint192(bytes24(bytes20(address(_validator)))); + nonce = entrypoint.getNonce(address(account), key); + } + + function getDefaultUserOp() internal returns (PackedUserOperation memory userOp) { + uint128 verificationGasLimit = 500_000; + uint128 callGasLimit = 500_000; + bytes32 packedAccountGasLimits = (bytes32(uint256(verificationGasLimit)) << 128) | + bytes32(uint256(callGasLimit)); + bytes32 packedGasLimits = (bytes32(uint256(1e9)) << 128) | bytes32(uint256(1e9)); + + userOp = PackedUserOperation({ + sender: address(0), + nonce: 0, + initCode: "", + callData: "", + accountGasLimits: packedAccountGasLimits, + preVerificationGas: 2e9, + gasFees: packedGasLimits, + 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)); + } + +}