diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 43abf4cd8e..cad97618d7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -343,6 +343,8 @@ jobs: - name: Install foundry uses: foundry-rs/foundry-toolchain@v1 + with: + version: 'v1.2.3' - name: Show Forge version working-directory: ./tee-worker/omni-executor/contracts/aa diff --git a/tee-worker/omni-executor/contracts/aa/README.md b/tee-worker/omni-executor/contracts/aa/README.md index 6241e528f7..0e295b5abe 100644 --- a/tee-worker/omni-executor/contracts/aa/README.md +++ b/tee-worker/omni-executor/contracts/aa/README.md @@ -7,6 +7,10 @@ licensed under the GNU General Public License v3.0. ## Compiling +Install all required git submodules with + +`git submodule update --init --recursive` + `forge compile` ## Running tests @@ -125,3 +129,9 @@ Each deployment artifact includes: The included bytecode can be used to verify deployed contracts on block explorers or to ensure the deployed code matches the source. Also see [DEPLOYMENT.md](./DEPLOYMENT.md) for more information. + +### Versioning + +All changes resulting in bytecode change should be properly versioned using new files suffixed with V{N} where N is a version number. +Each version directory should contain all files required to build contracts, there should be no cross version imports. +Work on new version begins when first change is introduced after previous version deployment and continues until deployment. diff --git a/tee-worker/omni-executor/contracts/aa/foundry.lock b/tee-worker/omni-executor/contracts/aa/foundry.lock new file mode 100644 index 0000000000..a8c053fc00 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/foundry.lock @@ -0,0 +1,8 @@ +{ + "../lib/forge-libs/forge-std": { + "rev": "60acb7aaadcce2d68e52986a0a66fe79f07d138f" + }, + "../lib/forge-libs/openzeppelin-contracts": { + "rev": "f27019d48eee32551e5c9d31849afcaa99944545" + } +} \ No newline at end of file diff --git a/tee-worker/omni-executor/contracts/aa/script/Deploy.s.sol b/tee-worker/omni-executor/contracts/aa/script/Deploy.s.sol index 0f2f51010d..69149e5e10 100644 --- a/tee-worker/omni-executor/contracts/aa/script/Deploy.s.sol +++ b/tee-worker/omni-executor/contracts/aa/script/Deploy.s.sol @@ -3,10 +3,10 @@ pragma solidity ^0.8.28; import "forge-std/Script.sol"; import "forge-std/console.sol"; -import "../src/core/EntryPointV1.sol"; -import "../src/accounts/OmniAccountFactoryV1.sol"; -import "../src/core/SimplePaymaster.sol"; -import "../src/core/ERC20PaymasterV1.sol"; +import "../src/v1/core/EntryPointV1.sol"; +import "../src/v1/accounts/OmniAccountFactoryV1.sol"; +import "../src/v1/core/SimplePaymaster.sol"; +import "../src/v1/core/ERC20PaymasterV1.sol"; import "./DeploymentHelper.sol"; /** diff --git a/tee-worker/omni-executor/contracts/aa/script/DeployLocal.s.sol b/tee-worker/omni-executor/contracts/aa/script/DeployLocal.s.sol index 8aac6e89eb..8a0e98bc0a 100644 --- a/tee-worker/omni-executor/contracts/aa/script/DeployLocal.s.sol +++ b/tee-worker/omni-executor/contracts/aa/script/DeployLocal.s.sol @@ -2,9 +2,9 @@ pragma solidity ^0.8.28; import "forge-std/Script.sol"; -import "../src/core/EntryPointV1.sol"; -import "../src/accounts/OmniAccountFactoryV1.sol"; -import "../src/core/SimplePaymaster.sol"; +import "../src/v1/core/EntryPointV1.sol"; +import "../src/v1/accounts/OmniAccountFactoryV1.sol"; +import "../src/v1/core/SimplePaymaster.sol"; import "../src/TestToken.sol"; contract DeployLocal is Script { diff --git a/tee-worker/omni-executor/contracts/aa/script/DeployLocalWithPaymaster.s.sol b/tee-worker/omni-executor/contracts/aa/script/DeployLocalWithPaymaster.s.sol index 45414b924e..a179cc9b3b 100644 --- a/tee-worker/omni-executor/contracts/aa/script/DeployLocalWithPaymaster.s.sol +++ b/tee-worker/omni-executor/contracts/aa/script/DeployLocalWithPaymaster.s.sol @@ -3,11 +3,11 @@ pragma solidity ^0.8.28; import "forge-std/Script.sol"; import "forge-std/console.sol"; -import "../src/core/EntryPointV1.sol"; -import "../src/accounts/OmniAccountFactoryV1.sol"; -import "../src/core/SimplePaymaster.sol"; -import "../src/core/DemoPaymaster.sol"; -import "../src/core/ERC20PaymasterV1.sol"; +import "../src/v1/core/EntryPointV1.sol"; +import "../src/v1/accounts/OmniAccountFactoryV1.sol"; +import "../src/v1/core/SimplePaymaster.sol"; +import "../src/v1/core/DemoPaymaster.sol"; +import "../src/v1/core/ERC20PaymasterV1.sol"; import "../src/TestToken.sol"; import "./DeploymentHelper.sol"; diff --git a/tee-worker/omni-executor/contracts/aa/script/DeploymentHelper.sol b/tee-worker/omni-executor/contracts/aa/script/DeploymentHelper.sol index c67646cbce..09ccd6d6df 100644 --- a/tee-worker/omni-executor/contracts/aa/script/DeploymentHelper.sol +++ b/tee-worker/omni-executor/contracts/aa/script/DeploymentHelper.sol @@ -376,7 +376,15 @@ library DeploymentHelper { * @notice Reads the ABI from Foundry artifacts * @return The ABI as a JSON string */ - function getContractAbi(Vm, /* vm */ string memory /* contractName */ ) internal pure returns (string memory) { + function getContractAbi( + Vm, + /* vm */ + string memory /* contractName */ + ) + internal + pure + returns (string memory) + { // Due to Solidity limitations with JSON parsing, especially for arrays, // we'll return a placeholder. In production, use a post-deployment script // to enrich the deployment artifacts with ABIs from the Foundry artifacts. @@ -481,8 +489,9 @@ library DeploymentHelper { returns (uint256) { try vm.projectRoot() returns (string memory root) { - string memory broadcastPath = - string(abi.encodePacked(root, "/broadcast/", scriptName, "/", vm.toString(chainId), "/run-latest.json")); + string memory broadcastPath = string( + abi.encodePacked(root, "/broadcast/", scriptName, "/", vm.toString(chainId), "/run-latest.json") + ); try vm.readFile(broadcastPath) returns (string memory json) { return parseBlockNumberFromBroadcast(vm, json, contractAddress); @@ -517,26 +526,34 @@ library DeploymentHelper { return block.number; } catch { // .receipts should be an array, try to get its length - try vm.parseJson(broadcastJson, ".receipts") returns (bytes memory /* receiptsData */ ) { + try vm.parseJson(broadcastJson, ".receipts") returns ( + bytes memory /* receiptsData */ + ) { // Look for transactions that created contracts // We'll check the first few receipts for contract creation (to field is null or empty) for (uint256 i = 0; i < 20; i++) { // Check up to 20 transactions try vm.parseJsonString( broadcastJson, string(abi.encodePacked(".receipts[", vm.toString(i), "].to")) - ) returns (string memory toAddr) { + ) returns ( + string memory toAddr + ) { // If 'to' is null/empty, this is a contract creation transaction if (bytes(toAddr).length == 0 || keccak256(bytes(toAddr)) == keccak256(bytes("null"))) { try vm.parseJsonString( broadcastJson, string(abi.encodePacked(".receipts[", vm.toString(i), "].contractAddress")) - ) returns (string memory contractAddr) { + ) returns ( + string memory contractAddr + ) { if (keccak256(bytes(toLowerCase(contractAddr))) == keccak256(bytes(addressLower))) { // Found our contract! Get its block number try vm.parseJsonString( broadcastJson, string(abi.encodePacked(".receipts[", vm.toString(i), "].blockNumber")) - ) returns (string memory blockHex) { + ) returns ( + string memory blockHex + ) { return hexStringToUint(blockHex); } catch { return block.number; diff --git a/tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountFactoryV1.sol b/tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountFactoryV1.sol similarity index 96% rename from tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountFactoryV1.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountFactoryV1.sol index 028474a1f5..923ffd3a65 100644 --- a/tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountFactoryV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountFactoryV1.sol @@ -40,12 +40,10 @@ contract OmniAccountFactoryV1 { return OmniAccountV1(payable(addr)); } ret = OmniAccountV1( - payable( - new ERC1967Proxy{salt: oa}( + payable(new ERC1967Proxy{salt: oa}( address(accountImplementation), abi.encodeCall(OmniAccountV1.initialize, (oa, oaType, clientId, root)) - ) - ) + )) ); } diff --git a/tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountV1.sol b/tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountV1.sol similarity index 98% rename from tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountV1.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountV1.sol index 68bdbadbe7..1da2aa959b 100644 --- a/tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountV1.sol @@ -40,7 +40,8 @@ contract OmniAccountV1 is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, In bytes4 private constant ADD_ROOT_SIGNER_SELECTOR = bytes4(keccak256("addRootSigner(address)")); bytes4 private constant REMOVE_ROOT_SIGNER_SELECTOR = bytes4(keccak256("removeRootSigner(address)")); bytes4 private constant ADD_PASSKEY_SIGNER_SELECTOR = bytes4(keccak256("addPasskeySigner((uint256,uint256))")); - bytes4 private constant REMOVE_PASSKEY_SIGNER_SELECTOR = bytes4(keccak256("removePasskeySigner((uint256,uint256))")); + bytes4 private constant REMOVE_PASSKEY_SIGNER_SELECTOR = + bytes4(keccak256("removePasskeySigner((uint256,uint256))")); bytes4 private constant WITHDRAW_DEPOSIT_SELECTOR = bytes4(keccak256("withdrawDepositTo(address,uint256)")); bytes4 private constant UPGRADE_TO_AND_CALL_SELECTOR = bytes4(keccak256("upgradeToAndCall(address,bytes)")); diff --git a/tee-worker/omni-executor/contracts/aa/src/accounts/callback/TokenCallbackHandler.sol b/tee-worker/omni-executor/contracts/aa/src/v1/accounts/callback/TokenCallbackHandler.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/accounts/callback/TokenCallbackHandler.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/accounts/callback/TokenCallbackHandler.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/core/BaseAccount.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/BaseAccount.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/core/BaseAccount.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/BaseAccount.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/core/BasePaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/BasePaymaster.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/core/BasePaymaster.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/BasePaymaster.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/DemoPaymaster.sol similarity index 91% rename from tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/DemoPaymaster.sol index 7f6626adcf..6da5caae6d 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/core/DemoPaymaster.sol @@ -24,7 +24,12 @@ contract DemoPaymaster is BasePaymaster { * Validate a user operation. * Sponsors any operation as long as we have sufficient deposit. */ - function _validatePaymasterUserOp(PackedUserOperation calldata userOp, bytes32, /* userOpHash */ uint256 maxCost) + function _validatePaymasterUserOp( + PackedUserOperation calldata userOp, + bytes32, + /* userOpHash */ + uint256 maxCost + ) internal view override @@ -51,7 +56,10 @@ contract DemoPaymaster is BasePaymaster { bytes calldata context, uint256 actualGasCost, uint256 /* actualUserOpFeePerGas */ - ) internal override { + ) + internal + override + { // Decode sender from context address sender = abi.decode(context, (address)); diff --git a/tee-worker/omni-executor/contracts/aa/src/core/ERC20PaymasterV1.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/ERC20PaymasterV1.sol similarity index 95% rename from tee-worker/omni-executor/contracts/aa/src/core/ERC20PaymasterV1.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/ERC20PaymasterV1.sol index f1f54cad3f..00c14d658e 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/ERC20PaymasterV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/core/ERC20PaymasterV1.sol @@ -28,10 +28,10 @@ contract ERC20PaymasterV1 is BasePaymaster, ReentrancyGuard { struct PaymasterData { address token; // ERC20 token address (must not be address(0)) uint256 exchangeRate; // Exchange rate: how many token units per 1 wei of ETH - // For tokens with different decimals, this should account for the difference - // Example: For 6-decimal USDC at $2000/ETH: rate = 2000 * 10^6 = 2000000000 - // Example: For 18-decimal token at 1500:1 ratio: rate = 1500 * 10^18 - // Set to 0 for full sponsorship (no token charge) + // For tokens with different decimals, this should account for the difference + // Example: For 6-decimal USDC at $2000/ETH: rate = 2000 * 10^6 = 2000000000 + // Example: For 18-decimal token at 1500:1 ratio: rate = 1500 * 10^18 + // Set to 0 for full sponsorship (no token charge) uint256 validUntil; // Timestamp until when this exchange rate is valid uint256 validAfter; // Timestamp after which this exchange rate is valid } @@ -83,7 +83,12 @@ contract ERC20PaymasterV1 is BasePaymaster, ReentrancyGuard { /** * Validate a user operation and handle ERC20 token prefunding */ - function _validatePaymasterUserOp(PackedUserOperation calldata userOp, bytes32, /* userOpHash */ uint256 maxCost) + function _validatePaymasterUserOp( + PackedUserOperation calldata userOp, + bytes32, + /* userOpHash */ + uint256 maxCost + ) internal override returns (bytes memory context, uint256 validationData) @@ -207,11 +212,12 @@ contract ERC20PaymasterV1 is BasePaymaster, ReentrancyGuard { // since no prefunding occurred during validation if (actualTokenCost > 0) { // Use low-level call to prevent revert on failed charge - (bool success,) = postOpContext.token.call( - abi.encodeWithSelector( - IERC20.transferFrom.selector, postOpContext.sender, beneficiary, actualTokenCost - ) - ); + (bool success,) = postOpContext.token + .call( + abi.encodeWithSelector( + IERC20.transferFrom.selector, postOpContext.sender, beneficiary, actualTokenCost + ) + ); if (!success) { // Charge failed, but don't revert the entire operation // The approval succeeded, but gas payment failed @@ -227,9 +233,8 @@ contract ERC20PaymasterV1 is BasePaymaster, ReentrancyGuard { if (beneficiary == address(this)) { // If beneficiary is this contract, we can refund directly // Use low-level call to prevent revert on failed refund - (bool success,) = postOpContext.token.call( - abi.encodeWithSelector(IERC20.transfer.selector, postOpContext.sender, refundAmount) - ); + (bool success,) = postOpContext.token + .call(abi.encodeWithSelector(IERC20.transfer.selector, postOpContext.sender, refundAmount)); if (!success) { // Refund failed, but don't revert the entire operation emit UserOpSponsored( diff --git a/tee-worker/omni-executor/contracts/aa/src/core/Eip7702Support.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/Eip7702Support.sol similarity index 98% rename from tee-worker/omni-executor/contracts/aa/src/core/Eip7702Support.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/Eip7702Support.sol index 641d1981ae..9cd39638e8 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/Eip7702Support.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/core/Eip7702Support.sol @@ -3,7 +3,7 @@ pragma solidity ^0.8.28; // solhint-disable no-inline-assembly import "../interfaces/PackedUserOperation.sol"; -import "../core/UserOperationLib.sol"; +import "./UserOperationLib.sol"; library Eip7702Support { // EIP-7702 code prefix before delegate address. diff --git a/tee-worker/omni-executor/contracts/aa/src/core/EntryPointSimulations.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/EntryPointSimulations.sol similarity index 98% rename from tee-worker/omni-executor/contracts/aa/src/core/EntryPointSimulations.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/EntryPointSimulations.sol index 372f5a4a6b..6a19042462 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/EntryPointSimulations.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/core/EntryPointSimulations.sol @@ -139,8 +139,9 @@ contract EntryPointSimulations is EntryPointV1, IEntryPointSimulations { initSenderCreator(); try this.validateSenderAndPaymaster(userOp.initCode, userOp.sender, userOp.paymasterAndData) { - // solhint-disable-next-line no-empty-blocks - } catch Error(string memory revertReason) { + // solhint-disable-next-line no-empty-blocks + } + catch Error(string memory revertReason) { if (bytes(revertReason).length != 0) { revert FailedOp(0, revertReason); } diff --git a/tee-worker/omni-executor/contracts/aa/src/core/EntryPointV1.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/EntryPointV1.sol similarity index 98% rename from tee-worker/omni-executor/contracts/aa/src/core/EntryPointV1.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/EntryPointV1.sol index ef54216bb7..9bad2f8cd9 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/EntryPointV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/core/EntryPointV1.sol @@ -168,7 +168,7 @@ contract EntryPointV1 is IEntryPoint, StakeManager, NonceManager, ReentrancyGuar function supportsInterface(bytes4 interfaceId) public view virtual override returns (bool) { // note: solidity "type(IEntryPoint).interfaceId" is without inherited methods but we want to check everything return interfaceId - == (type(IEntryPoint).interfaceId ^ type(IStakeManager).interfaceId ^ type(INonceManager).interfaceId) + == (type(IEntryPoint).interfaceId ^ type(IStakeManager).interfaceId ^ type(INonceManager).interfaceId) || interfaceId == type(IEntryPoint).interfaceId || interfaceId == type(IStakeManager).interfaceId || interfaceId == type(INonceManager).interfaceId || super.supportsInterface(interfaceId); } @@ -593,16 +593,15 @@ contract EntryPointV1 is IEntryPoint, StakeManager, NonceManager, ReentrancyGuar uint256 maxContextLength; uint256 len; assembly ("memory-safe") { - success := - call( - paymasterVerificationGasLimit, - paymaster, - 0, - add(validatePaymasterCall, 0x20), - mload(validatePaymasterCall), - 0, - 0 - ) + success := call( + paymasterVerificationGasLimit, + paymaster, + 0, + add(validatePaymasterCall, 0x20), + mload(validatePaymasterCall), + 0, + 0 + ) len := returndatasize() // return data from validatePaymasterUserOp is (bytes context, validationData) // encoded as: @@ -777,8 +776,9 @@ contract EntryPointV1 is IEntryPoint, StakeManager, NonceManager, ReentrancyGuar try IPaymaster(paymaster).postOp{gas: mUserOp.paymasterPostOpGasLimit}( mode, context, actualGasCost, gasPrice ) { - // solhint-disable-next-line no-empty-blocks - } catch { + // solhint-disable-next-line no-empty-blocks + } + catch { bytes memory reason = Exec.getReturnData(REVERT_REASON_MAX_LEN); revert PostOpReverted(reason); } diff --git a/tee-worker/omni-executor/contracts/aa/src/core/Helpers.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/Helpers.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/core/Helpers.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/Helpers.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/core/NonceManager.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/NonceManager.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/core/NonceManager.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/NonceManager.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/core/SenderCreator.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/SenderCreator.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/core/SenderCreator.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/SenderCreator.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/SimplePaymaster.sol similarity index 93% rename from tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/SimplePaymaster.sol index 4ecad4aead..892a081228 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/core/SimplePaymaster.sol @@ -28,7 +28,12 @@ contract SimplePaymaster is BasePaymaster { * Validate a user operation. * Only sponsor operations submitted by authorized bundlers. */ - function _validatePaymasterUserOp(PackedUserOperation calldata userOp, bytes32, /* userOpHash */ uint256 maxCost) + function _validatePaymasterUserOp( + PackedUserOperation calldata userOp, + bytes32, + /* userOpHash */ + uint256 maxCost + ) internal view override @@ -61,7 +66,10 @@ contract SimplePaymaster is BasePaymaster { bytes calldata context, uint256 actualGasCost, uint256 /* actualUserOpFeePerGas */ - ) internal override { + ) + internal + override + { // Decode sender from context address sender = abi.decode(context, (address)); diff --git a/tee-worker/omni-executor/contracts/aa/src/core/StakeManager.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/StakeManager.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/core/StakeManager.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/StakeManager.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/core/UserOperationLib.sol b/tee-worker/omni-executor/contracts/aa/src/v1/core/UserOperationLib.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/core/UserOperationLib.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/core/UserOperationLib.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/IAccount.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IAccount.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/IAccount.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IAccount.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/IAccountExecute.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IAccountExecute.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/IAccountExecute.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IAccountExecute.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/IAggregator.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IAggregator.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/IAggregator.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IAggregator.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/IEntryPoint.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IEntryPoint.sol similarity index 99% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/IEntryPoint.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IEntryPoint.sol index 56725a78f6..6b9d30d6a9 100644 --- a/tee-worker/omni-executor/contracts/aa/src/interfaces/IEntryPoint.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IEntryPoint.sol @@ -144,8 +144,7 @@ interface IEntryPoint is IStakeManager, INonceManager { * @param opsPerAggregator - The operations to execute, grouped by aggregator (or address(0) for no-aggregator accounts). * @param beneficiary - The address to receive the fees. */ - function handleAggregatedOps(UserOpsPerAggregator[] calldata opsPerAggregator, address payable beneficiary) - external; + function handleAggregatedOps(UserOpsPerAggregator[] calldata opsPerAggregator, address payable beneficiary) external; /** * Generate a request Id - unique identifier for this request. diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/IEntryPointSimulations.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IEntryPointSimulations.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/IEntryPointSimulations.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IEntryPointSimulations.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/INonceManager.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/INonceManager.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/INonceManager.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/INonceManager.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/IPaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IPaymaster.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/IPaymaster.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IPaymaster.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/ISenderCreator.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/ISenderCreator.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/ISenderCreator.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/ISenderCreator.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/IStakeManager.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IStakeManager.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/IStakeManager.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/IStakeManager.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/OwnerType.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/OwnerType.sol similarity index 99% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/OwnerType.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/OwnerType.sol index 9376234043..e0821bfe68 100644 --- a/tee-worker/omni-executor/contracts/aa/src/interfaces/OwnerType.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/OwnerType.sol @@ -23,5 +23,4 @@ enum OwnerType { Solana, // 0x08 Google, // 0x09 Passkey // 0x0a - } diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/PackedUserOperation.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/PackedUserOperation.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/PackedUserOperation.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/PackedUserOperation.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/Passkey.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/Passkey.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/Passkey.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/Passkey.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/interfaces/UserOpSigner.sol b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/UserOpSigner.sol similarity index 99% rename from tee-worker/omni-executor/contracts/aa/src/interfaces/UserOpSigner.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/interfaces/UserOpSigner.sol index 6e0f4fe24e..f00918d663 100644 --- a/tee-worker/omni-executor/contracts/aa/src/interfaces/UserOpSigner.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/interfaces/UserOpSigner.sol @@ -22,5 +22,4 @@ enum UserOpSigner { RootKey, // 0x01 SessionKey, // 0x02 Passkey // 0x03 - } diff --git a/tee-worker/omni-executor/contracts/aa/src/utils/Exec.sol b/tee-worker/omni-executor/contracts/aa/src/v1/utils/Exec.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/utils/Exec.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/utils/Exec.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountFactoryV2.sol b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountFactoryV2.sol new file mode 100644 index 0000000000..f76e458749 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountFactoryV2.sol @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +import "@openzeppelin/contracts/utils/Create2.sol"; +import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; + +import "../interfaces/ISenderCreator.sol"; +import "../interfaces/OwnerType.sol"; +import "./OmniAccountV2.sol"; + +contract OmniAccountFactoryV2 { + OmniAccountV2 public immutable accountImplementation; + ISenderCreator public immutable senderCreator; + + event AccountCreated(address indexed account, bytes32 indexed oa, OwnerType oaType, address root); + + constructor(IEntryPoint _entryPoint) { + accountImplementation = new OmniAccountV2(_entryPoint); + senderCreator = _entryPoint.senderCreator(); + } + + function createAccount(bytes32 oa, OwnerType oaType, bytes memory clientId, address root) + public + returns (OmniAccountV2 ret) + { + require(msg.sender == address(senderCreator), "only callable from SenderCreator"); + address addr = getAddress(oa, oaType, clientId, root); + uint256 codeSize = addr.code.length; + if (codeSize > 0) { + return OmniAccountV2(payable(addr)); + } + ret = OmniAccountV2( + payable(new ERC1967Proxy{salt: oa}( + address(accountImplementation), + abi.encodeCall(OmniAccountV2.initialize, (oa, oaType, clientId, root)) + )) + ); + emit AccountCreated(address(ret), oa, oaType, root); + } + + function getAddress(bytes32 oa, OwnerType oaType, bytes memory clientId, address root) + public + view + returns (address) + { + return Create2.computeAddress( + oa, + keccak256( + abi.encodePacked( + type(ERC1967Proxy).creationCode, + abi.encode( + address(accountImplementation), + abi.encodeCall(OmniAccountV2.initialize, (oa, oaType, clientId, root)) + ) + ) + ) + ); + } + + function version() public pure returns (string memory) { + return "2.0.0"; + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountV2.sol b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountV2.sol new file mode 100644 index 0000000000..834e2dc93b --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountV2.sol @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable avoid-low-level-calls */ +/* solhint-disable no-inline-assembly */ +/* solhint-disable reason-string */ + +import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; +import "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol"; +import "@openzeppelin/contracts/utils/Strings.sol"; +import "@openzeppelin/contracts/proxy/utils/Initializable.sol"; +import "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol"; +import "../core/BaseAccount.sol"; +import "../interfaces/OwnerType.sol"; +import "../interfaces/UserOpSigner.sol"; +import "../interfaces/Passkey.sol"; +import "../core/Helpers.sol"; +import "./callback/TokenCallbackHandler.sol"; +import "../utils/Exec.sol"; +import "../core/LibModuleManager.sol"; + +contract OmniAccountV2 is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, Initializable { + using Passkey for Passkey.PublicKey; + using LibModuleManager for LibModuleManager.ModuleStorage; + + bytes32 public owner; + bytes public clientId; + OwnerType public ownerType; + + mapping(address => bool) public rootSigners; + mapping(bytes32 => bool) public passkeySigners; + uint256 public passkeySignerCount; + + IEntryPoint private immutable _entryPoint; + + bytes4 private constant ADD_ROOT_SIGNER_SELECTOR = bytes4(keccak256("addRootSigner(address)")); + bytes4 private constant REMOVE_ROOT_SIGNER_SELECTOR = bytes4(keccak256("removeRootSigner(address)")); + bytes4 private constant ADD_PASSKEY_SIGNER_SELECTOR = bytes4(keccak256("addPasskeySigner((uint256,uint256))")); + bytes4 private constant REMOVE_PASSKEY_SIGNER_SELECTOR = + bytes4(keccak256("removePasskeySigner((uint256,uint256))")); + bytes4 private constant WITHDRAW_DEPOSIT_SELECTOR = bytes4(keccak256("withdrawDepositTo(address,uint256)")); + bytes4 private constant UPGRADE_TO_AND_CALL_SELECTOR = bytes4(keccak256("upgradeToAndCall(address,bytes)")); + bytes4 private constant REGISTER_MODULE_SELECTOR = bytes4(keccak256("registerModule(address)")); + bytes4 private constant UNREGISTER_MODULE_SELECTOR = bytes4(keccak256("unregisterModule(address)")); + + event AccountInitialized( + IEntryPoint indexed entryPoint, bytes32 indexed owner, OwnerType ownerType, bytes clientId, address indexed root + ); + event RootSignerAdded(address root); + event RootSignerRemoved(address root); + event PasskeySignerAdded(Passkey.PublicKey pk); + event PasskeySignerRemoved(Passkey.PublicKey pk); + event ModuleRegistered(address indexed module); + event ModuleUnregistered(address indexed module); + + modifier onlyOwner() { + _onlyOwner(); + _; + } + + function entryPoint() public view virtual override returns (IEntryPoint) { + return _entryPoint; + } + + receive() external payable {} + + constructor(IEntryPoint anEntryPoint) { + _entryPoint = anEntryPoint; + _disableInitializers(); + } + + function _onlyOwner() internal view { + require( + _determineOa(msg.sender) == owner || msg.sender == address(entryPoint()) + || (ownerType != OwnerType.Evm && passkeySignerCount == 0 && isRootSigner(msg.sender)), + "only owner" + ); + } + + function initialize(bytes32 anOwner, OwnerType anOwnerType, bytes memory aClientId, address aRoot) + public + virtual + initializer + { + _initialize(anOwner, anOwnerType, aClientId, aRoot); + } + + function _initialize(bytes32 anOwner, OwnerType anOwnerType, bytes memory aClientId, address aRoot) + internal + virtual + { + owner = anOwner; + rootSigners[aRoot] = true; + clientId = aClientId; + ownerType = anOwnerType; + emit AccountInitialized(_entryPoint, owner, ownerType, clientId, aRoot); + } + + function _determineOa(address sender) internal view returns (bytes32) { + bytes3 oaType = 0x65766d; + return sha256(abi.encodePacked(clientId, oaType, sender)); + } + + function isRootSigner(address sender) public view returns (bool) { + return rootSigners[sender]; + } + + function _validateSignature(PackedUserOperation calldata userOp, bytes32 userOpHash) + internal + virtual + override + returns (uint256 validationData) + { + require(userOp.signature.length >= 1, "signature too short"); + + UserOpSigner signer = UserOpSigner(uint8(userOp.signature[0])); + bytes calldata sig = userOp.signature[1:]; + + if (signer == UserOpSigner.Owner) { + return _validateOwner(userOpHash, sig); + } else if (signer == UserOpSigner.RootKey) { + if (ownerType != OwnerType.Evm && passkeySignerCount == 0) { + return _validateRootKey(userOpHash, sig); + } + if (_isRestrictedCall(userOp.callData)) { + return SIG_VALIDATION_FAILED; + } + return _validateRootKey(userOpHash, sig); + } else if (signer == UserOpSigner.SessionKey) { + if (_isRestrictedCall(userOp.callData)) { + return SIG_VALIDATION_FAILED; + } + return _validateSessionKey(userOpHash, sig); + } else if (signer == UserOpSigner.Passkey) { + if (ownerType != OwnerType.Evm && passkeySignerCount > 0) { + return _validatePasskey(userOpHash, sig); + } + if (_isRestrictedCall(userOp.callData)) { + return SIG_VALIDATION_FAILED; + } + return _validatePasskey(userOpHash, sig); + } else { + revert("unsupported signer type"); + } + } + + function _validateOwner(bytes32 userOpHash, bytes calldata sig) internal view returns (uint256 validationData) { + require(sig.length == 65, "Owner signature length invalid"); + address signer = ECDSA.recover(userOpHash, sig); + return owner == _determineOa(signer) ? SIG_VALIDATION_SUCCESS : SIG_VALIDATION_FAILED; + } + + function _validateRootKey(bytes32 userOpHash, bytes calldata sig) internal view returns (uint256 validationData) { + require(sig.length == 65, "RootKey signature length invalid"); + address signer = ECDSA.recover(userOpHash, sig); + return isRootSigner(signer) ? SIG_VALIDATION_SUCCESS : SIG_VALIDATION_FAILED; + } + + function _validateSessionKey(bytes32 userOpHash, bytes calldata sig) + internal + view + returns (uint256 validationData) + { + require(sig.length == 162, "SessionKey signature length invalid"); + bytes memory sessionSig = sig[:65]; + address sessionKey = ECDSA.recover(userOpHash, sessionSig); + uint256 sessionExpiration = uint256(bytes32(sig[65:97])); + + if (block.timestamp > sessionExpiration) { + return SIG_VALIDATION_FAILED; + } + + bytes memory sessionProof = sig[97:162]; + bytes32 sessionDigest = sha256(abi.encodePacked(sessionKey, sessionExpiration)); + address sessionProofSigner = ECDSA.recover(sessionDigest, sessionProof); + + return isRootSigner(sessionProofSigner) ? SIG_VALIDATION_SUCCESS : SIG_VALIDATION_FAILED; + } + + function _validatePasskey(bytes32 userOpHash, bytes calldata sig) internal view returns (uint256 validationData) { + ( + Passkey.PublicKey memory publicKey, + Passkey.Signature memory passkeySignature, + Passkey.Metadata memory metadata + ) = abi.decode(sig, (Passkey.PublicKey, Passkey.Signature, Passkey.Metadata)); + + if (!passkeySigners[publicKey.toKey()]) { + return SIG_VALIDATION_FAILED; + } + + bool isValid = Passkey.verify(userOpHash, metadata, passkeySignature, publicKey); + + return isValid ? SIG_VALIDATION_SUCCESS : SIG_VALIDATION_FAILED; + } + + function getDeposit() public view returns (uint256) { + return entryPoint().balanceOf(address(this)); + } + + function getOwner() public view returns (bytes32) { + return owner; + } + + function addDeposit() public payable { + entryPoint().depositTo{value: msg.value}(address(this)); + } + + function withdrawDepositTo(address payable withdrawAddress, uint256 amount) public onlyOwner { + entryPoint().withdrawTo(withdrawAddress, amount); + } + + function addRootSigner(address root) public onlyOwner { + rootSigners[root] = true; + emit RootSignerAdded(root); + } + + function removeRootSigner(address root) public onlyOwner { + rootSigners[root] = false; + emit RootSignerRemoved(root); + } + + function addPasskeySigner(Passkey.PublicKey memory pk) public onlyOwner { + bytes32 key = pk.toKey(); + if (!passkeySigners[key]) { + passkeySigners[key] = true; + passkeySignerCount++; + emit PasskeySignerAdded(pk); + } + } + + function removePasskeySigner(Passkey.PublicKey memory pk) public onlyOwner { + bytes32 key = pk.toKey(); + if (passkeySigners[key]) { + passkeySigners[key] = false; + passkeySignerCount--; + emit PasskeySignerRemoved(pk); + } + } + + function _authorizeUpgrade(address newImplementation) internal view override { + (newImplementation); + _onlyOwner(); + } + + function _isRestrictedCall(bytes calldata callData) internal pure returns (bool) { + if (callData.length < 4) return false; + + bytes4 selector = bytes4(callData[0:4]); + + return selector == ADD_ROOT_SIGNER_SELECTOR || selector == REMOVE_ROOT_SIGNER_SELECTOR + || selector == ADD_PASSKEY_SIGNER_SELECTOR || selector == REMOVE_PASSKEY_SIGNER_SELECTOR + || selector == WITHDRAW_DEPOSIT_SELECTOR || selector == UPGRADE_TO_AND_CALL_SELECTOR + || selector == REGISTER_MODULE_SELECTOR || selector == UNREGISTER_MODULE_SELECTOR; + } + + function registerModule(address module) external onlyOwner { + LibModuleManager.registerModule(module); + emit ModuleRegistered(module); + } + + function unregisterModule(address module) external onlyOwner { + LibModuleManager.unregisterModule(module); + emit ModuleUnregistered(module); + } + + function isModuleRegistered(address module) external view returns (bool) { + return LibModuleManager.isModuleRegistered(module); + } + + function executeModuleCall(address module, bytes calldata data) external virtual returns (bytes memory) { + _requireFromEntryPoint(); + require(LibModuleManager.isModuleRegistered(module), "Module not registered"); + + bool ok = Exec.delegateCall(module, data, gasleft()); + if (!ok) { + Exec.revertWithReturnData(); + } + + return Exec.getReturnData(0); + } + + function version() public pure virtual returns (string memory) { + return "2.0.0"; + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/accounts/callback/TokenCallbackHandler.sol b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/callback/TokenCallbackHandler.sol new file mode 100644 index 0000000000..ae0c9cd7f9 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/callback/TokenCallbackHandler.sol @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable no-empty-blocks */ + +import "@openzeppelin/contracts/utils/introspection/IERC165.sol"; +import "@openzeppelin/contracts/token/ERC721/IERC721Receiver.sol"; +import "@openzeppelin/contracts/token/ERC1155/IERC1155Receiver.sol"; + +/** + * Token callback handler. + * Handles supported tokens' callbacks, allowing account receiving these tokens. + */ +abstract contract TokenCallbackHandler is IERC721Receiver, IERC1155Receiver { + function onERC721Received(address, address, uint256, bytes calldata) external pure override returns (bytes4) { + return IERC721Receiver.onERC721Received.selector; + } + + function onERC1155Received(address, address, uint256, uint256, bytes calldata) + external + pure + override + returns (bytes4) + { + return IERC1155Receiver.onERC1155Received.selector; + } + + function onERC1155BatchReceived(address, address, uint256[] calldata, uint256[] calldata, bytes calldata) + external + pure + override + returns (bytes4) + { + return IERC1155Receiver.onERC1155BatchReceived.selector; + } + + function supportsInterface(bytes4 interfaceId) external view virtual override returns (bool) { + return interfaceId == type(IERC721Receiver).interfaceId || interfaceId == type(IERC1155Receiver).interfaceId + || interfaceId == type(IERC165).interfaceId; + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/BaseAccount.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/BaseAccount.sol new file mode 100644 index 0000000000..adf1cacf45 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/BaseAccount.sol @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable avoid-low-level-calls */ +/* solhint-disable no-empty-blocks */ +/* solhint-disable no-inline-assembly */ + +import "../interfaces/IAccount.sol"; +import "../interfaces/IEntryPoint.sol"; +import "../utils/Exec.sol"; +import "./UserOperationLib.sol"; + +/** + * Basic account implementation. + * This contract provides the basic logic for implementing the IAccount interface - validateUserOp + * Specific account implementation should inherit it and provide the account-specific logic. + */ +abstract contract BaseAccount is IAccount { + using UserOperationLib for PackedUserOperation; + + struct Call { + address target; + uint256 value; + bytes data; + } + + error ExecuteError(uint256 index, bytes error); + + /** + * Return the account nonce. + * This method returns the next sequential nonce. + * For a nonce of a specific key, use `entrypoint.getNonce(account, key)` + */ + function getNonce() public view virtual returns (uint256) { + return entryPoint().getNonce(address(this), 0); + } + + /** + * Return the entryPoint used by this account. + * Subclass should return the current entryPoint used by this account. + */ + function entryPoint() public view virtual returns (IEntryPoint); + + /** + * execute a single call from the account. + */ + function execute(address target, uint256 value, bytes calldata data) external virtual { + _requireForExecute(); + + bool ok = Exec.call(target, value, data, gasleft()); + if (!ok) { + Exec.revertWithReturnData(); + } + } + + /** + * execute a batch of calls. + * revert on the first call that fails. + * If the batch reverts, and it contains more than a single call, then wrap the revert with ExecuteError, + * to mark the failing call index. + */ + function executeBatch(Call[] calldata calls) external virtual { + _requireForExecute(); + + uint256 callsLength = calls.length; + for (uint256 i = 0; i < callsLength; i++) { + Call calldata call = calls[i]; + bool ok = Exec.call(call.target, call.value, call.data, gasleft()); + if (!ok) { + if (callsLength == 1) { + Exec.revertWithReturnData(); + } else { + revert ExecuteError(i, Exec.getReturnData(0)); + } + } + } + } + + /// @inheritdoc IAccount + function validateUserOp(PackedUserOperation calldata userOp, bytes32 userOpHash, uint256 missingAccountFunds) + external + virtual + override + returns (uint256 validationData) + { + _requireFromEntryPoint(); + validationData = _validateSignature(userOp, userOpHash); + _validateNonce(userOp.nonce); + _payPrefund(missingAccountFunds); + } + + /** + * Ensure the request comes from the known entrypoint. + */ + function _requireFromEntryPoint() internal view virtual { + require(msg.sender == address(entryPoint()), "account: not from EntryPoint"); + } + + function _requireForExecute() internal view virtual { + _requireFromEntryPoint(); + } + + /** + * Validate the signature is valid for this message. + * @param userOp - Validate the userOp.signature field. + * @param userOpHash - Convenient field: the hash of the request, to check the signature against. + * (also hashes the entrypoint and chain id) + * @return validationData - Signature and time-range of this operation. + * <20-byte> aggregatorOrSigFail - 0 for valid signature, 1 to mark signature failure, + * otherwise, an address of an aggregator contract. + * <6-byte> validUntil - Last timestamp this operation is valid at, or 0 for "indefinitely" + * <6-byte> validAfter - first timestamp this operation is valid + * If the account doesn't use time-range, it is enough to return + * SIG_VALIDATION_FAILED value (1) for signature failure. + * Note that the validation code cannot use block.timestamp (or block.number) directly. + */ + function _validateSignature(PackedUserOperation calldata userOp, bytes32 userOpHash) + internal + virtual + returns (uint256 validationData); + + /** + * Validate the nonce of the UserOperation. + * This method may validate the nonce requirement of this account. + * e.g. + * To limit the nonce to use sequenced UserOps only (no "out of order" UserOps): + * `require(nonce < type(uint64).max)` + * For a hypothetical account that *requires* the nonce to be out-of-order: + * `require(nonce & type(uint64).max == 0)` + * + * The actual nonce uniqueness is managed by the EntryPoint, and thus no other + * action is needed by the account itself. + * + * @param nonce to validate + * + * solhint-disable-next-line no-empty-blocks + */ + function _validateNonce(uint256 nonce) internal view virtual {} + + /** + * Sends to the entrypoint (msg.sender) the missing funds for this transaction. + * SubClass MAY override this method for better funds management + * (e.g. send to the entryPoint more than the minimum required, so that in future transactions + * it will not be required to send again). + * @param missingAccountFunds - The minimum value this method should send the entrypoint. + * This value MAY be zero, in case there is enough deposit, + * or the userOp has a paymaster. + */ + function _payPrefund(uint256 missingAccountFunds) internal virtual { + if (missingAccountFunds != 0) { + (bool success,) = payable(msg.sender).call{value: missingAccountFunds}(""); + (success); + // Ignore failure (its EntryPoint's job to verify, not account.) + } + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/BasePaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/BasePaymaster.sol new file mode 100644 index 0000000000..d583df6781 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/BasePaymaster.sol @@ -0,0 +1,164 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable reason-string */ + +import "@openzeppelin/contracts/access/Ownable2Step.sol"; +import "@openzeppelin/contracts/utils/introspection/IERC165.sol"; +import "@openzeppelin/contracts/utils/Pausable.sol"; +import "../interfaces/IPaymaster.sol"; +import "../interfaces/IEntryPoint.sol"; +import "./UserOperationLib.sol"; +/** + * Helper class for creating a paymaster. + * provides helper methods for staking. + * Validates that the postOp is called only by the entryPoint. + */ + +abstract contract BasePaymaster is IPaymaster, Ownable2Step, Pausable { + IEntryPoint public immutable entryPoint; + + uint256 internal constant PAYMASTER_VALIDATION_GAS_OFFSET = UserOperationLib.PAYMASTER_VALIDATION_GAS_OFFSET; + uint256 internal constant PAYMASTER_POSTOP_GAS_OFFSET = UserOperationLib.PAYMASTER_POSTOP_GAS_OFFSET; + uint256 internal constant PAYMASTER_DATA_OFFSET = UserOperationLib.PAYMASTER_DATA_OFFSET; + + constructor(IEntryPoint _entryPoint) Ownable(msg.sender) { + _validateEntryPointInterface(_entryPoint); + entryPoint = _entryPoint; + } + + // Sanity check: make sure this EntryPoint was compiled against the same + // IEntryPoint of this paymaster + function _validateEntryPointInterface(IEntryPoint _entryPoint) internal virtual { + require( + IERC165(address(_entryPoint)).supportsInterface(type(IEntryPoint).interfaceId), + "IEntryPoint interface mismatch" + ); + } + + /// @inheritdoc IPaymaster + function validatePaymasterUserOp(PackedUserOperation calldata userOp, bytes32 userOpHash, uint256 maxCost) + external + override + whenNotPaused + returns (bytes memory context, uint256 validationData) + { + _requireFromEntryPoint(); + return _validatePaymasterUserOp(userOp, userOpHash, maxCost); + } + + /** + * Validate a user operation. + * @param userOp - The user operation. + * @param userOpHash - The hash of the user operation. + * @param maxCost - The maximum cost of the user operation. + */ + function _validatePaymasterUserOp(PackedUserOperation calldata userOp, bytes32 userOpHash, uint256 maxCost) + internal + virtual + returns (bytes memory context, uint256 validationData); + + /// @inheritdoc IPaymaster + function postOp(PostOpMode mode, bytes calldata context, uint256 actualGasCost, uint256 actualUserOpFeePerGas) + external + override + { + _requireFromEntryPoint(); + _postOp(mode, context, actualGasCost, actualUserOpFeePerGas); + } + + /** + * Post-operation handler. + * (verified to be called only through the entryPoint) + * @dev If subclass returns a non-empty context from validatePaymasterUserOp, + * it must also implement this method. + * @param mode - Enum with the following options: + * opSucceeded - User operation succeeded. + * opReverted - User op reverted. The paymaster still has to pay for gas. + * postOpReverted - never passed in a call to postOp(). + * @param context - The context value returned by validatePaymasterUserOp + * @param actualGasCost - Actual cost of gas used so far (without this postOp call). + * @param actualUserOpFeePerGas - the gas price this UserOp pays. This value is based on the UserOp's maxFeePerGas + * and maxPriorityFee (and basefee) + * It is not the same as tx.gasprice, which is what the bundler pays. + */ + function _postOp(PostOpMode mode, bytes calldata context, uint256 actualGasCost, uint256 actualUserOpFeePerGas) + internal + virtual + { + (mode, context, actualGasCost, actualUserOpFeePerGas); // unused params + // subclass must override this method if validatePaymasterUserOp returns a context + revert("must override"); + } + + /** + * Add a deposit for this paymaster, used for paying for transaction fees. + */ + function deposit() public payable { + entryPoint.depositTo{value: msg.value}(address(this)); + } + + /** + * Withdraw value from the deposit. + * @param withdrawAddress - Target to send to. + * @param amount - Amount to withdraw. + */ + function withdrawTo(address payable withdrawAddress, uint256 amount) public onlyOwner { + require(amount > 0, "Amount must be greater than 0"); + entryPoint.withdrawTo(withdrawAddress, amount); + } + + /** + * Add stake for this paymaster. + * This method can also carry eth value to add to the current stake. + * @param unstakeDelaySec - The unstake delay for this paymaster. Can only be increased. + */ + function addStake(uint32 unstakeDelaySec) external payable onlyOwner { + entryPoint.addStake{value: msg.value}(unstakeDelaySec); + } + + /** + * Return current paymaster's deposit on the entryPoint. + */ + function getDeposit() public view returns (uint256) { + return entryPoint.balanceOf(address(this)); + } + + /** + * Unlock the stake, in order to withdraw it. + * The paymaster can't serve requests once unlocked, until it calls addStake again + */ + function unlockStake() external onlyOwner { + entryPoint.unlockStake(); + } + + /** + * Withdraw the entire paymaster's stake. + * stake must be unlocked first (and then wait for the unstakeDelay to be over) + * @param withdrawAddress - The address to send withdrawn value. + */ + function withdrawStake(address payable withdrawAddress) external onlyOwner { + entryPoint.withdrawStake(withdrawAddress); + } + + /** + * Validate the call is made from a valid entrypoint + */ + function _requireFromEntryPoint() internal virtual { + require(msg.sender == address(entryPoint), "Sender not EntryPoint"); + } + + /** + * Pause the paymaster, preventing new user operation validations + */ + function pause() external onlyOwner { + _pause(); + } + + /** + * Unpause the paymaster, allowing user operation validations + */ + function unpause() external onlyOwner { + _unpause(); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/DemoPaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/DemoPaymaster.sol new file mode 100644 index 0000000000..6da5caae6d --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/DemoPaymaster.sol @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: GPL-3.0 +/* + * Based on DemoPaymaster.sol from https://github.com/eth-infinitism/account-abstraction + * Licensed under GNU General Public License v3.0 + */ +pragma solidity ^0.8.28; + +import "./BasePaymaster.sol"; +import "../interfaces/PackedUserOperation.sol"; + +/** + * Demo paymaster implementation that sponsors gas for any user operation + * without bundler restrictions. For demonstration purposes only. + * + * SECURITY WARNING: This paymaster accepts any operation as long as it has + * sufficient deposit. Do not use in production. + */ +contract DemoPaymaster is BasePaymaster { + event UserOpSponsored(address indexed account, uint256 actualGasCost); + + constructor(IEntryPoint _entryPoint) BasePaymaster(_entryPoint) {} + + /** + * Validate a user operation. + * Sponsors any operation as long as we have sufficient deposit. + */ + function _validatePaymasterUserOp( + PackedUserOperation calldata userOp, + bytes32, + /* userOpHash */ + uint256 maxCost + ) + internal + view + override + returns (bytes memory context, uint256 validationData) + { + // Check if we have enough deposit to cover the cost + uint256 ourDeposit = entryPoint.balanceOf(address(this)); + if (ourDeposit < maxCost) { + // Reject - insufficient funds + return ("", 1); + } + + // Accept the operation - sponsor any account + // Return smart account address in context for postOp logging + return (abi.encode(userOp.sender), 0); + } + + /** + * Post-operation handler. + * Log the sponsored operation. + */ + function _postOp( + PostOpMode, /* mode */ + bytes calldata context, + uint256 actualGasCost, + uint256 /* actualUserOpFeePerGas */ + ) + internal + override + { + // Decode sender from context + address sender = abi.decode(context, (address)); + + // Log the sponsored operation + emit UserOpSponsored(sender, actualGasCost); + } + + /** + * Allow contract to receive ETH and automatically deposit to EntryPoint. + */ + receive() external payable { + entryPoint.depositTo{value: msg.value}(address(this)); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/ERC20PaymasterV1.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/ERC20PaymasterV1.sol new file mode 100644 index 0000000000..00c14d658e --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/ERC20PaymasterV1.sol @@ -0,0 +1,500 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import "./BasePaymaster.sol"; +import "./UserOperationLib.sol"; +import "../interfaces/PackedUserOperation.sol"; +import "../interfaces/IPaymaster.sol"; +import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol"; +import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; +import "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; +import "@openzeppelin/contracts/utils/math/Math.sol"; + +/** + * ERC20PaymasterV1 - A paymaster that allows users to pay gas fees with ERC20 tokens + * while reimbursing bundlers with ETH. Only supports ERC20 token payments. + * Only accepts operations from authorized bundlers for security. + */ +contract ERC20PaymasterV1 is BasePaymaster, ReentrancyGuard { + using SafeERC20 for IERC20; + using UserOperationLib for PackedUserOperation; + using Math for uint256; + + // Minimum required postOp gas limit to prevent paymaster losses + uint256 private constant MIN_POST_OP_GAS_LIMIT = 50000; + + // Struct to decode paymaster data + struct PaymasterData { + address token; // ERC20 token address (must not be address(0)) + uint256 exchangeRate; // Exchange rate: how many token units per 1 wei of ETH + // For tokens with different decimals, this should account for the difference + // Example: For 6-decimal USDC at $2000/ETH: rate = 2000 * 10^6 = 2000000000 + // Example: For 18-decimal token at 1500:1 ratio: rate = 1500 * 10^18 + // Set to 0 for full sponsorship (no token charge) + uint256 validUntil; // Timestamp until when this exchange rate is valid + uint256 validAfter; // Timestamp after which this exchange rate is valid + } + + // Context for postOp + struct PostOpContext { + address sender; // User's account address + address token; // ERC20 token address + uint256 exchangeRate; // Exchange rate used + uint256 maxCost; // Maximum cost in wei + uint256 prefundAmount; // Amount prefunded in tokens + uint256 postOpGasLimit; // PostOp gas limit from paymasterAndData + bool isApprovalOp; // Whether this was an approval operation + } + + // Mapping of authorized bundler addresses + mapping(address => bool) public authorizedBundlers; + + // Beneficiary address for collected ERC20 tokens (default: this contract) + address public beneficiary; + + // Events + event UserOpSponsored(address indexed account, address indexed token, uint256 actualGasCost, uint256 tokenAmount); + event AuthorizedBundlerUpdated(address indexed bundler, bool authorized); + event BeneficiaryUpdated(address indexed oldBeneficiary, address indexed newBeneficiary); + event TokensWithdrawn(address indexed token, address indexed to, uint256 amount); + + // Errors + error UnauthorizedBundler(); + error InsufficientDeposit(); + error PostOpGasLimitTooLow(); + error InvalidToken(); + error InsufficientTokenBalance(); + error InsufficientTokenAllowance(); + error InvalidExchangeRate(); + error TimestampValidationFailed(); + error InvalidPaymasterData(); + error TokenTransferFailed(); + error RefundFailed(); + + constructor(IEntryPoint _entryPoint, address _initialBundler) BasePaymaster(_entryPoint) { + authorizedBundlers[_initialBundler] = true; + beneficiary = address(this); + + emit AuthorizedBundlerUpdated(_initialBundler, true); + emit BeneficiaryUpdated(address(0), beneficiary); + } + + /** + * Validate a user operation and handle ERC20 token prefunding + */ + function _validatePaymasterUserOp( + PackedUserOperation calldata userOp, + bytes32, + /* userOpHash */ + uint256 maxCost + ) + internal + override + returns (bytes memory context, uint256 validationData) + { + // Check if the transaction is being submitted by an authorized bundler + if (!authorizedBundlers[tx.origin]) { + revert UnauthorizedBundler(); + } + + // Check if we have enough ETH deposit to cover the cost + uint256 ourDeposit = entryPoint.balanceOf(address(this)); + if (ourDeposit < maxCost) { + revert InsufficientDeposit(); + } + + // Decode paymaster data first to validate structure + PaymasterData memory data = _decodePaymasterData(userOp.paymasterAndData); + + // Extract and validate postOp gas limit after data validation + uint256 postOpGasLimit = userOp.unpackPostOpGasLimit(); + if (postOpGasLimit < MIN_POST_OP_GAS_LIMIT) { + revert PostOpGasLimitTooLow(); + } + + // Validate timestamps + uint256 currentTime = block.timestamp; + if (currentTime < data.validAfter || currentTime > data.validUntil) { + return ("", 1); // Reject with validation failure + } + + // Validate token address - native tokens (address(0)) are not supported + if (data.token == address(0)) { + revert InvalidToken(); + } + + // Handle full sponsorship case (0 exchange rate means paymaster fully sponsors) + bool isFullSponsorship = data.exchangeRate == 0; + + // Handle ERC20 token payment + IERC20 token = IERC20(data.token); + + // Calculate required token amount with overflow protection + uint256 requiredTokenAmount; + if (isFullSponsorship) { + requiredTokenAmount = 0; // No charge for full sponsorship + } else { + // Use Math.mulDiv to prevent overflow: (maxCost * exchangeRate) / 1e18 + requiredTokenAmount = maxCost.mulDiv(data.exchangeRate, 1e18, Math.Rounding.Ceil); + } + + // Check user's token balance + uint256 userBalance = token.balanceOf(userOp.sender); + if (userBalance < requiredTokenAmount) { + revert InsufficientTokenBalance(); + } + + // Check if the userOp is an approval transaction for this paymaster + bool isApprovalOp = _isApprovalOperation(userOp, data.token); + + // For approval operations, still check balance to prevent "free approval" attacks + // but skip allowance and prefunding checks since the approval is happening in this operation + if (isApprovalOp) { + // Even for approval ops, user must have sufficient balance to cover the gas + // This prevents attackers from spamming free approval operations to drain paymaster + if (userBalance < requiredTokenAmount) { + revert InsufficientTokenBalance(); + } + } else if (!isFullSponsorship) { + // For non-approval operations with token charges, check allowance and prefund + uint256 currentAllowance = token.allowance(userOp.sender, address(this)); + if (currentAllowance < requiredTokenAmount) { + revert InsufficientTokenAllowance(); + } + + // Prefund by transferring tokens from user to beneficiary + token.safeTransferFrom(userOp.sender, beneficiary, requiredTokenAmount); + } + + PostOpContext memory postOpContext = PostOpContext({ + sender: userOp.sender, + token: data.token, + exchangeRate: data.exchangeRate, + maxCost: maxCost, + prefundAmount: requiredTokenAmount, + postOpGasLimit: postOpGasLimit, + isApprovalOp: isApprovalOp + }); + + return (abi.encode(postOpContext), 0); + } + + /** + * Post-operation handler to handle refunds and logging + */ + function _postOp( + IPaymaster.PostOpMode mode, + bytes calldata context, + uint256 actualGasCost, + uint256 actualUserOpFeePerGas + ) internal override nonReentrant { + PostOpContext memory postOpContext = abi.decode(context, (PostOpContext)); + + // Add postOp gas overhead to actual gas cost to prevent paymaster losses + // The actualGasCost parameter doesn't include gas consumed by postOp itself + // Use the postOp gas limit from paymasterAndData for accurate accounting + uint256 totalGasCost = actualGasCost + (postOpContext.postOpGasLimit * actualUserOpFeePerGas); + + // Calculate token costs and refunds for ERC20 payments + uint256 actualTokenCost; + if (postOpContext.exchangeRate == 0) { + actualTokenCost = 0; // Full sponsorship + } else { + // Use Math.mulDiv to prevent overflow, include postOp gas overhead + actualTokenCost = totalGasCost.mulDiv(postOpContext.exchangeRate, 1e18, Math.Rounding.Ceil); + } + + // Handle refunds/charges based on operation type + if (mode == IPaymaster.PostOpMode.opSucceeded) { + if (postOpContext.isApprovalOp) { + // For approval operations, charge the user the actual token cost + // since no prefunding occurred during validation + if (actualTokenCost > 0) { + // Use low-level call to prevent revert on failed charge + (bool success,) = postOpContext.token + .call( + abi.encodeWithSelector( + IERC20.transferFrom.selector, postOpContext.sender, beneficiary, actualTokenCost + ) + ); + if (!success) { + // Charge failed, but don't revert the entire operation + // The approval succeeded, but gas payment failed + emit UserOpSponsored(postOpContext.sender, postOpContext.token, totalGasCost, 0); + return; + } + } + } else if (postOpContext.prefundAmount > actualTokenCost) { + // For non-approval operations, refund excess tokens + uint256 refundAmount = postOpContext.prefundAmount - actualTokenCost; + + // Transfer refund from beneficiary back to user + if (beneficiary == address(this)) { + // If beneficiary is this contract, we can refund directly + // Use low-level call to prevent revert on failed refund + (bool success,) = postOpContext.token + .call(abi.encodeWithSelector(IERC20.transfer.selector, postOpContext.sender, refundAmount)); + if (!success) { + // Refund failed, but don't revert the entire operation + emit UserOpSponsored( + postOpContext.sender, postOpContext.token, totalGasCost, postOpContext.prefundAmount + ); + return; + } + } + // If beneficiary is external, they need to handle their own refunds + } + } + + emit UserOpSponsored(postOpContext.sender, postOpContext.token, totalGasCost, actualTokenCost); + } + + /** + * Decode paymaster data from paymasterAndData field + */ + function _decodePaymasterData(bytes calldata paymasterAndData) internal pure returns (PaymasterData memory data) { + // paymasterAndData format: + // paymaster_address (20) + validation_gas_limit (16) + postop_gas_limit (16) + paymaster_data + // Our data starting at PAYMASTER_DATA_OFFSET (52): token(20) + exchangeRate(32) + validUntil(32) + validAfter(32) + if (paymasterAndData.length < PAYMASTER_DATA_OFFSET + 20 + 32 + 32 + 32) { + revert InvalidPaymasterData(); + } + + bytes calldata paymasterData = paymasterAndData[PAYMASTER_DATA_OFFSET:]; + + // Read 20 bytes token address, then 32 bytes each for other fields + data.token = address(bytes20(paymasterData[0:20])); + data.exchangeRate = uint256(bytes32(paymasterData[20:52])); + data.validUntil = uint256(bytes32(paymasterData[52:84])); + data.validAfter = uint256(bytes32(paymasterData[84:116])); + } + + /** + * Check if the user operation is an ERC20 approval operation for this paymaster + * + * In Account Abstraction, users call functions directly on the account contract. + * For ERC20 approvals, the callData will be either: + * 1. execute(tokenAddress, 0, approve(paymaster, amount)) + * 2. executeBatch([(tokenAddress, 0, approve(paymaster, amount))]) + */ + function _isApprovalOperation(PackedUserOperation calldata userOp, address tokenAddress) + internal + view + returns (bool) + { + if (userOp.callData.length < 4) return false; + + bytes4 selector = bytes4(userOp.callData[0:4]); + + // Check if it's an execute() call that might contain approve + if (selector == bytes4(keccak256("execute(address,uint256,bytes)"))) { + return _checkExecuteApproval(userOp.callData, tokenAddress); + } + + // Check if it's a batch execution containing approve + // Using common batch execution selector from BaseAccount + if (selector == bytes4(keccak256("executeBatch((address,uint256,bytes)[])"))) { + return _checkBatchApproval(userOp.callData, tokenAddress); + } + + return false; + } + + /** + * Check if execute callData contains approve(address(this), amount) to the correct token + * Decodes: execute(tokenAddress, 0, approve(paymaster, amount)) + */ + function _checkExecuteApproval(bytes calldata callData, address tokenAddress) internal view returns (bool) { + // execute(address,uint256,bytes) = 4 + 32 + 32 + 32 (offset) + 32 (length) + data + if (callData.length < 132) return false; // Minimum length for execute with some data + + // Decode execute parameters: target (bytes 4:36), value (bytes 36:68), data offset (bytes 68:100) + address target = address(bytes20(callData[16:36])); // Skip 4-byte selector + 12 bytes padding + uint256 value = uint256(bytes32(callData[36:68])); + + // Check if target is the expected token and value is 0 + if (target != tokenAddress || value != 0) { + return false; + } + + // Get data offset and length + uint256 dataOffset = uint256(bytes32(callData[68:100])); + uint256 absoluteDataOffset = 4 + dataOffset; // Add 4 for function selector + + if (callData.length < absoluteDataOffset + 32) return false; // Need at least length field + + uint256 dataLength = uint256(bytes32(callData[absoluteDataOffset:absoluteDataOffset + 32])); + uint256 dataStart = absoluteDataOffset + 32; + + if (callData.length < dataStart + dataLength || dataLength < 68) return false; // approve needs 4+32+32 bytes + + // Check if the inner data is approve(paymaster, amount) + bytes4 innerSelector = bytes4(callData[dataStart:dataStart + 4]); + if (innerSelector != IERC20.approve.selector) { + return false; + } + + // Check if spender is this paymaster + address spender = address(bytes20(callData[dataStart + 16:dataStart + 36])); // Skip selector + 12 bytes padding + return spender == address(this); + } + + /** + * Check if batch execution contains approve(address(this), amount) for the correct token + * Decodes: executeBatch([(tokenAddress, 0, approve(paymaster, amount))]) + */ + function _checkBatchApproval(bytes calldata callData, address tokenAddress) internal view returns (bool) { + // executeBatch((address,uint256,bytes)[]) = 4 + 32 (array offset) + 32 (array length) + Call structs + if (callData.length < 100) return false; // Minimum for non-empty batch + + // Get array offset and length + uint256 arrayOffset = uint256(bytes32(callData[4:36])); + uint256 absoluteArrayOffset = 4 + arrayOffset; + + if (callData.length < absoluteArrayOffset + 32) return false; + + uint256 arrayLength = uint256(bytes32(callData[absoluteArrayOffset:absoluteArrayOffset + 32])); + if (arrayLength == 0) return false; + + // For simplicity, only check if the first call is an approval to the correct token + // Each Call struct has: target(32) + value(32) + data_offset(32) + data_length(32) + data + uint256 firstCallOffset = absoluteArrayOffset + 32; + + if (callData.length < firstCallOffset + 96) return false; // Need at least target+value+offset + + address target = address(bytes20(callData[firstCallOffset + 12:firstCallOffset + 32])); + uint256 value = uint256(bytes32(callData[firstCallOffset + 32:firstCallOffset + 64])); + + // Check if target is the expected token and value is 0 + if (target != tokenAddress || value != 0) { + return false; + } + + // Get data offset and length for the first call + uint256 dataOffset = uint256(bytes32(callData[firstCallOffset + 64:firstCallOffset + 96])); + uint256 absoluteDataOffset = firstCallOffset + dataOffset; + + if (callData.length < absoluteDataOffset + 32) return false; + + uint256 dataLength = uint256(bytes32(callData[absoluteDataOffset:absoluteDataOffset + 32])); + uint256 dataStart = absoluteDataOffset + 32; + + if (callData.length < dataStart + dataLength || dataLength < 68) return false; + + // Check if the inner data is approve(paymaster, amount) + bytes4 innerSelector = bytes4(callData[dataStart:dataStart + 4]); + if (innerSelector != IERC20.approve.selector) { + return false; + } + + // Check if spender is this paymaster + address spender = address(bytes20(callData[dataStart + 16:dataStart + 36])); + return spender == address(this); + } + + /** + * Get token decimals for a given token (helper function for calculating exchange rates) + * @param token The token address to query decimals for + * @return decimals The number of decimals for the token, defaults to 18 if not available + */ + function getTokenDecimals(address token) external view returns (uint8 decimals) { + if (token == address(0)) { + revert InvalidToken(); // Native tokens are not supported + } + + // Check if the address has code + if (token.code.length == 0) { + return 18; // Not a contract, default to 18 + } + + try IERC20Metadata(token).decimals() returns (uint8 result) { + return result; + } catch { + return 18; // Default to 18 if decimals() call fails + } + } + + /** + * Calculate exchange rate for a token given how many tokens equal 1 ETH + * @param tokenDecimals The number of decimals the token has + * @param tokensPerEth How many tokens you get for 1 ETH (before decimal scaling) + * @return exchangeRate The exchange rate to use in PaymasterData + * + * Example: USDC (6 decimals) where 1 ETH = 2000 USDC + * exchangeRate = tokensPerEth * 10^tokenDecimals + * exchangeRate = 2000 * 10^6 = 2000000000 + */ + function calculateExchangeRate(uint8 tokenDecimals, uint256 tokensPerEth) + external + pure + returns (uint256 exchangeRate) + { + if (tokensPerEth == 0) { + revert InvalidExchangeRate(); + } + + // exchangeRate = tokensPerEth * 10^tokenDecimals + // This gives us how many token units per 1 ETH + return tokensPerEth * (10 ** tokenDecimals); + } + + /** + * Add or remove an authorized bundler + */ + function setAuthorizedBundler(address bundler, bool authorized) external onlyOwner { + authorizedBundlers[bundler] = authorized; + emit AuthorizedBundlerUpdated(bundler, authorized); + } + + /** + * Set the beneficiary address for collected ERC20 tokens + */ + function setBeneficiary(address _beneficiary) external onlyOwner { + require(_beneficiary != address(0), "Invalid beneficiary"); + address oldBeneficiary = beneficiary; + beneficiary = _beneficiary; + emit BeneficiaryUpdated(oldBeneficiary, _beneficiary); + } + + /** + * Withdraw accumulated ERC20 tokens (only if beneficiary is this contract) + */ + function withdrawTokens(address token, address to, uint256 amount) external onlyOwner { + require(beneficiary == address(this), "Beneficiary is not this contract"); + require(to != address(0), "Invalid recipient"); + + // Only ERC20 tokens are supported for withdrawal + if (token == address(0)) { + revert InvalidToken(); + } + + // Withdraw ERC20 + IERC20(token).safeTransfer(to, amount); + + emit TokensWithdrawn(token, to, amount); + } + + /** + * Get the version of this paymaster contract + */ + function version() public pure virtual returns (string memory) { + return "1.0.0"; + } + + /** + * Allow contract to receive ETH and automatically deposit to EntryPoint + */ + receive() external payable { + entryPoint.depositTo{value: msg.value}(address(this)); + } + // Test helper function to expose _isApprovalOperation for testing + + function _isApprovalOperation_exposed(PackedUserOperation calldata userOp, address tokenAddress) + external + view + returns (bool) + { + return _isApprovalOperation(userOp, tokenAddress); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/Eip7702Support.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/Eip7702Support.sol new file mode 100644 index 0000000000..9cd39638e8 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/Eip7702Support.sol @@ -0,0 +1,79 @@ +pragma solidity ^0.8.28; +// SPDX-License-Identifier: MIT +// solhint-disable no-inline-assembly + +import "../interfaces/PackedUserOperation.sol"; +import "./UserOperationLib.sol"; + +library Eip7702Support { + // EIP-7702 code prefix before delegate address. + bytes3 internal constant EIP7702_PREFIX = 0xef0100; + + // EIP-7702 initCode marker, to specify this account is EIP-7702. + bytes2 internal constant INITCODE_EIP7702_MARKER = 0x7702; + + using UserOperationLib for PackedUserOperation; + + /** + * Get the alternative 'InitCodeHash' value for the UserOp hash calculation when using EIP-7702. + * + * @param userOp - the UserOperation to for the 'InitCodeHash' calculation. + * @return the 'InitCodeHash' value. + */ + function _getEip7702InitCodeHashOverride(PackedUserOperation calldata userOp) internal view returns (bytes32) { + bytes calldata initCode = userOp.initCode; + if (!_isEip7702InitCode(initCode)) { + return 0; + } + address delegate = _getEip7702Delegate(userOp.sender); + if (initCode.length <= 20) { + return keccak256(abi.encodePacked(delegate)); + } else { + return keccak256(abi.encodePacked(delegate, initCode[20:])); + } + } + + /** + * Check if this 'initCode' is actually an EIP-7702 authorization. + * This is indicated by 'initCode' that starts with INITCODE_EIP7702_MARKER. + * + * @param initCode - the 'initCode' to check. + * @return true if the 'initCode' is EIP-7702 authorization, false otherwise. + */ + function _isEip7702InitCode(bytes calldata initCode) internal pure returns (bool) { + if (initCode.length < 2) { + return false; + } + bytes20 initCodeStart; + // non-empty calldata bytes are always zero-padded to 32-bytes, so can be safely casted to "bytes20" + assembly ("memory-safe") { + initCodeStart := calldataload(initCode.offset) + } + // make sure first 20 bytes of initCode are "0x7702" (padded with zeros) + return initCodeStart == bytes20(INITCODE_EIP7702_MARKER); + } + + /** + * Get the EIP-7702 delegate from contract code. + * Must only be used if _isEip7702InitCode(initCode) is true. + * + * @param sender - the EIP-7702 'sender' account to get the delegated contract code address. + * @return the address of the EIP-7702 authorized contract. + */ + function _getEip7702Delegate(address sender) internal view returns (address) { + bytes32 senderCode; + + assembly ("memory-safe") { + extcodecopy(sender, 0, 0, 23) + senderCode := mload(0) + } + // To be a valid EIP-7702 delegate, the first 3 bytes are EIP7702_PREFIX + // followed by the delegate address + if (bytes3(senderCode) != EIP7702_PREFIX) { + // instead of just "not an EIP-7702 delegate", if some info. + require(sender.code.length > 0, "sender has no code"); + revert("not an EIP-7702 delegate"); + } + return address(bytes20(senderCode << 24)); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/EntryPointSimulations.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/EntryPointSimulations.sol new file mode 100644 index 0000000000..6a19042462 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/EntryPointSimulations.sol @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: GPL-3.0 +/* + * Based on EntryPointSimulations.sol from https://github.com/eth-infinitism/account-abstraction + * Licensed under GNU General Public License v3.0 + */ +pragma solidity ^0.8.28; + +/* solhint-disable avoid-low-level-calls */ +/* solhint-disable no-inline-assembly */ + +import "./EntryPointV1.sol"; +import "../interfaces/IEntryPointSimulations.sol"; + +/* + * This contract inherits the EntryPoint and extends it with the view-only methods that are executed by + * the bundler in order to check UserOperation validity and estimate its gas consumption. + * This contract should never be deployed on-chain and is only used as a parameter for the "eth_call" request. + */ +contract EntryPointSimulations is EntryPointV1, IEntryPointSimulations { + SenderCreator private _senderCreator; + + bytes32 private __domainSeparatorV4; + + function initSenderCreator() internal virtual { + // This is the address of the first contract created with CREATE by this address. + address createdObj = address(uint160(uint256(keccak256(abi.encodePacked(hex"d694", address(this), hex"01"))))); + _senderCreator = SenderCreator(createdObj); + + _initDomainSeparator(); + } + + function senderCreator() public view virtual override(EntryPointV1, IEntryPoint) returns (ISenderCreator) { + // return the same senderCreator as real EntryPoint. + // this call is slightly (100) more expensive than EntryPoint's access to immutable member + return _senderCreator; + } + + /** + * simulation contract should not be deployed, and specifically, accounts should not trust + * it as entrypoint, since the simulation functions don't check the signatures + */ + constructor() { + require(block.number < 1000, "should not be deployed"); + } + + /// @inheritdoc IEntryPointSimulations + function simulateValidation(PackedUserOperation calldata userOp) external returns (ValidationResult memory) { + UserOpInfo memory outOpInfo; + + _simulationOnlyValidations(userOp); + (uint256 validationData, uint256 paymasterValidationData) = _validatePrepayment(0, userOp, outOpInfo); + StakeInfo memory paymasterInfo = _getStakeInfo(outOpInfo.mUserOp.paymaster); + StakeInfo memory senderInfo = _getStakeInfo(outOpInfo.mUserOp.sender); + StakeInfo memory factoryInfo; + { + bytes calldata initCode = userOp.initCode; + address factory = initCode.length >= 20 ? address(bytes20(initCode[0:20])) : address(0); + factoryInfo = _getStakeInfo(factory); + } + + address aggregator = address(uint160(validationData)); + ReturnInfo memory returnInfo = ReturnInfo( + outOpInfo.preOpGas, + outOpInfo.prefund, + validationData, + paymasterValidationData, + _getMemoryBytesFromOffset(outOpInfo.contextOffset) + ); + + AggregatorStakeInfo memory aggregatorInfo; // = NOT_AGGREGATED; + if (uint160(aggregator) != SIG_VALIDATION_SUCCESS && uint160(aggregator) != SIG_VALIDATION_FAILED) { + aggregatorInfo = AggregatorStakeInfo(aggregator, _getStakeInfo(aggregator)); + } + return ValidationResult(returnInfo, senderInfo, factoryInfo, paymasterInfo, aggregatorInfo); + } + + /// @inheritdoc IEntryPointSimulations + function simulateHandleOp(PackedUserOperation calldata op, address target, bytes calldata targetCallData) + external + nonReentrant + returns (ExecutionResult memory) + { + UserOpInfo memory opInfo; + _simulationOnlyValidations(op); + (uint256 validationData, uint256 paymasterValidationData) = _validatePrepayment(0, op, opInfo); + + uint256 paid = _executeUserOp(0, op, opInfo); + bool targetSuccess; + bytes memory targetResult; + if (target != address(0)) { + (targetSuccess, targetResult) = target.call(targetCallData); + } + return + ExecutionResult(opInfo.preOpGas, paid, validationData, paymasterValidationData, targetSuccess, targetResult); + } + + /// @inheritdoc IEntryPointSimulations + function simulateHandleOps(PackedUserOperation[] calldata ops, address payable beneficiary) + external + nonReentrant + returns (ExecutionResult[] memory results) + { + uint256 opslen = ops.length; + results = new ExecutionResult[](opslen); + UserOpInfo[] memory opInfos = new UserOpInfo[](opslen); + uint256 collected = 0; + + unchecked { + // For simulation, we need to run validation on each operation first + for (uint256 i = 0; i < opslen; i++) { + _simulationOnlyValidations(ops[i]); + (uint256 validationData, uint256 paymasterValidationData) = _validatePrepayment(i, ops[i], opInfos[i]); + + // Execute the user operation + uint256 paid = _executeUserOp(i, ops[i], opInfos[i]); + collected += paid; + + // Create execution result for this operation + results[i] = ExecutionResult( + opInfos[i].preOpGas, + paid, + validationData, + paymasterValidationData, + true, // targetSuccess - no target call in batch simulation + "" // targetResult - no target call in batch simulation + ); + } + } + + // Compensate beneficiary to accurately simulate the actual handleOps behavior + // This ensures gas estimation includes the transfer cost and validates the beneficiary + _compensate(beneficiary, collected); + + return results; + } + + function _simulationOnlyValidations(PackedUserOperation calldata userOp) internal { + // Initialize senderCreator(). we can't rely on constructor + initSenderCreator(); + + try this.validateSenderAndPaymaster(userOp.initCode, userOp.sender, userOp.paymasterAndData) { + // solhint-disable-next-line no-empty-blocks + } + catch Error(string memory revertReason) { + if (bytes(revertReason).length != 0) { + revert FailedOp(0, revertReason); + } + } + } + + /** + * Called only during simulation by the EntryPointSimulation contract itself and is not meant to be called by external contracts. + * This function always reverts to prevent warm/cold storage differentiation in simulation vs execution. + * @param initCode - The smart account constructor code. + * @param sender - The sender address. + * @param paymasterAndData - The paymaster address (followed by other params, ignored by this method) + */ + function validateSenderAndPaymaster(bytes calldata initCode, address sender, bytes calldata paymasterAndData) + external + view + { + if (initCode.length == 0 && sender.code.length == 0) { + // it would revert anyway. but give a meaningful message + revert("AA20 account not deployed"); + } + if (paymasterAndData.length >= 20) { + address paymaster = address(bytes20(paymasterAndData[0:20])); + if (paymaster.code.length == 0) { + // It would revert anyway. but give a meaningful message. + revert("AA30 paymaster not deployed"); + } + } + // always revert + revert(""); + } + + // Make sure depositTo cost is more than normal EntryPoint's cost, + // to mitigate DoS vector on the bundler + // empiric test showed that without this wrapper, simulation depositTo costs less.. + function depositTo(address account) public payable override(IStakeManager, StakeManager) { + unchecked { + // silly code, to waste some gas to make sure depositTo is always little more + // expensive than on-chain call + uint256 x = 1; + while (x < 5) { + x++; + } + StakeManager.depositTo(account); + } + } + + // Copied from EIP712.sol + bytes32 private constant TYPE_HASH = + keccak256("EIP712Domain(string name,string version,uint256 chainId,address verifyingContract)"); + + function __buildDomainSeparator() private view returns (bytes32) { + bytes32 _hashedName = keccak256(bytes(DOMAIN_NAME)); + bytes32 _hashedVersion = keccak256(bytes(DOMAIN_VERSION)); + return keccak256(abi.encode(TYPE_HASH, _hashedName, _hashedVersion, block.chainid, address(this))); + } + + // Can't rely on "immutable" (constructor-initialized) variables" in simulation + function _initDomainSeparator() internal { + __domainSeparatorV4 = __buildDomainSeparator(); + } + + function getDomainSeparatorV4() public view override returns (bytes32) { + return __domainSeparatorV4; + } + + function supportsInterface(bytes4) public view virtual override returns (bool) { + return false; + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/EntryPointV1.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/EntryPointV1.sol new file mode 100644 index 0000000000..9bad2f8cd9 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/EntryPointV1.sol @@ -0,0 +1,880 @@ +// SPDX-License-Identifier: GPL-3.0 +/* + * Based on EntryPoint.sol from https://github.com/eth-infinitism/account-abstraction + * Licensed under GNU General Public License v3.0 + */ +pragma solidity ^0.8.28; +/* solhint-disable avoid-low-level-calls */ +/* solhint-disable no-inline-assembly */ + +import "../interfaces/IAccount.sol"; +import "../interfaces/IAccountExecute.sol"; +import "../interfaces/IEntryPoint.sol"; +import "../interfaces/IPaymaster.sol"; + +import "./UserOperationLib.sol"; +import "./StakeManager.sol"; +import "./NonceManager.sol"; +import "./Helpers.sol"; +import "./SenderCreator.sol"; +import "./Eip7702Support.sol"; +import "../utils/Exec.sol"; + +import "@openzeppelin/contracts/utils/ReentrancyGuardTransient.sol"; +import "@openzeppelin/contracts/utils/introspection/ERC165.sol"; +import "@openzeppelin/contracts/utils/cryptography/EIP712.sol"; + +/** + * Account-Abstraction (EIP-4337) singleton EntryPoint v0.8 implementation. + * Only one instance required on each chain. + * + * @custom:security-contact https://bounty.ethereum.org + */ +contract EntryPointV1 is IEntryPoint, StakeManager, NonceManager, ReentrancyGuardTransient, ERC165, EIP712 { + using UserOperationLib for PackedUserOperation; + + /** + * internal-use constants + */ + + // allow some slack for future gas price changes. + uint256 private constant INNER_GAS_OVERHEAD = 10000; + + // Marker for inner call revert on out of gas + bytes32 private constant INNER_OUT_OF_GAS = hex"deaddead"; + bytes32 private constant INNER_REVERT_LOW_PREFUND = hex"deadaa51"; + + uint256 private constant REVERT_REASON_MAX_LEN = 2048; + // Penalty charged for either unused execution gas or postOp gas + uint256 private constant UNUSED_GAS_PENALTY_PERCENT = 10; + // Threshold below which no penalty would be charged + uint256 private constant PENALTY_GAS_THRESHOLD = 40000; + + SenderCreator private immutable _senderCreator = new SenderCreator(); + + string internal constant DOMAIN_NAME = "ERC4337"; + string internal constant DOMAIN_VERSION = "1"; + + constructor() EIP712(DOMAIN_NAME, DOMAIN_VERSION) {} + + /// @inheritdoc IEntryPoint + function handleOps(PackedUserOperation[] calldata ops, address payable beneficiary) external nonReentrant { + uint256 opslen = ops.length; + UserOpInfo[] memory opInfos = new UserOpInfo[](opslen); + unchecked { + _iterateValidationPhase(ops, opInfos, address(0), 0); + + uint256 collected = 0; + emit BeforeExecution(); + + for (uint256 i = 0; i < opslen; i++) { + collected += _executeUserOp(i, ops[i], opInfos[i]); + } + + _compensate(beneficiary, collected); + } + } + + /// @inheritdoc IEntryPoint + function handleAggregatedOps(UserOpsPerAggregator[] calldata opsPerAggregator, address payable beneficiary) + external + nonReentrant + { + unchecked { + uint256 opasLen = opsPerAggregator.length; + uint256 totalOps = 0; + for (uint256 i = 0; i < opasLen; i++) { + UserOpsPerAggregator calldata opa = opsPerAggregator[i]; + PackedUserOperation[] calldata ops = opa.userOps; + IAggregator aggregator = opa.aggregator; + + // address(1) is special marker of "signature error" + require(address(aggregator) != address(1), SignatureValidationFailed(address(aggregator))); + + if (address(aggregator) != address(0)) { + // solhint-disable-next-line no-empty-blocks + try aggregator.validateSignatures(ops, opa.signature) {} + catch { + revert SignatureValidationFailed(address(aggregator)); + } + } + + totalOps += ops.length; + } + + UserOpInfo[] memory opInfos = new UserOpInfo[](totalOps); + + uint256 opIndex = 0; + for (uint256 a = 0; a < opasLen; a++) { + UserOpsPerAggregator calldata opa = opsPerAggregator[a]; + PackedUserOperation[] calldata ops = opa.userOps; + IAggregator aggregator = opa.aggregator; + + opIndex += _iterateValidationPhase(ops, opInfos, address(aggregator), opIndex); + } + + emit BeforeExecution(); + + uint256 collected = 0; + opIndex = 0; + for (uint256 a = 0; a < opasLen; a++) { + UserOpsPerAggregator calldata opa = opsPerAggregator[a]; + emit SignatureAggregatorChanged(address(opa.aggregator)); + PackedUserOperation[] calldata ops = opa.userOps; + uint256 opslen = ops.length; + + for (uint256 i = 0; i < opslen; i++) { + collected += _executeUserOp(opIndex, ops[i], opInfos[opIndex]); + opIndex++; + } + } + + _compensate(beneficiary, collected); + } + } + + /// @inheritdoc IEntryPoint + function getUserOpHash(PackedUserOperation calldata userOp) public view returns (bytes32) { + bytes32 overrideInitCodeHash = Eip7702Support._getEip7702InitCodeHashOverride(userOp); + return MessageHashUtils.toTypedDataHash(getDomainSeparatorV4(), userOp.hash(overrideInitCodeHash)); + } + + /// @inheritdoc IEntryPoint + function getSenderAddress(bytes calldata initCode) external { + address sender = senderCreator().createSender(initCode); + revert SenderAddressResult(sender); + } + + /// @inheritdoc IEntryPoint + function senderCreator() public view virtual returns (ISenderCreator) { + return _senderCreator; + } + + /// @inheritdoc IEntryPoint + function delegateAndRevert(address target, bytes calldata data) external { + (bool success, bytes memory ret) = target.delegatecall(data); + revert DelegateAndRevert(success, ret); + } + + function getPackedUserOpTypeHash() external pure returns (bytes32) { + return UserOperationLib.PACKED_USEROP_TYPEHASH; + } + + function getDomainSeparatorV4() public view virtual returns (bytes32) { + return _domainSeparatorV4(); + } + + /// @inheritdoc IERC165 + function supportsInterface(bytes4 interfaceId) public view virtual override returns (bool) { + // note: solidity "type(IEntryPoint).interfaceId" is without inherited methods but we want to check everything + return interfaceId + == (type(IEntryPoint).interfaceId ^ type(IStakeManager).interfaceId ^ type(INonceManager).interfaceId) + || interfaceId == type(IEntryPoint).interfaceId || interfaceId == type(IStakeManager).interfaceId + || interfaceId == type(INonceManager).interfaceId || super.supportsInterface(interfaceId); + } + + /** + * Compensate the caller's beneficiary address with the collected fees of all UserOperations. + * @param beneficiary - The address to receive the fees. + * @param amount - Amount to transfer. + */ + function _compensate(address payable beneficiary, uint256 amount) internal virtual { + require(beneficiary != address(0), "AA90 invalid beneficiary"); + (bool success,) = beneficiary.call{value: amount}(""); + require(success, "AA91 failed send to beneficiary"); + } + + /** + * Execute a user operation. + * @param opIndex - Index into the opInfo array. + * @param userOp - The userOp to execute. + * @param opInfo - The opInfo filled by validatePrepayment for this userOp. + * @return collected - The total amount this userOp paid. + */ + function _executeUserOp(uint256 opIndex, PackedUserOperation calldata userOp, UserOpInfo memory opInfo) + internal + virtual + returns (uint256 collected) + { + uint256 preGas = gasleft(); + bytes memory context = _getMemoryBytesFromOffset(opInfo.contextOffset); + bool success; + { + uint256 saveFreePtr = _getFreePtr(); + bytes calldata callData = userOp.callData; + bytes memory innerCall; + bytes4 methodSig; + assembly ("memory-safe") { + let len := callData.length + if gt(len, 3) { methodSig := calldataload(callData.offset) } + } + if (methodSig == IAccountExecute.executeUserOp.selector) { + bytes memory executeUserOp = abi.encodeCall(IAccountExecute.executeUserOp, (userOp, opInfo.userOpHash)); + innerCall = abi.encodeCall(this.innerHandleOp, (executeUserOp, opInfo, context)); + } else { + innerCall = abi.encodeCall(this.innerHandleOp, (callData, opInfo, context)); + } + assembly ("memory-safe") { + success := call(gas(), address(), 0, add(innerCall, 0x20), mload(innerCall), 0, 32) + collected := mload(0) + } + _restoreFreePtr(saveFreePtr); + } + if (!success) { + bytes32 innerRevertCode; + assembly ("memory-safe") { + let len := returndatasize() + if eq(32, len) { + returndatacopy(0, 0, 32) + innerRevertCode := mload(0) + } + } + if (innerRevertCode == INNER_OUT_OF_GAS) { + // handleOps was called with gas limit too low. abort entire bundle. + // can only be caused by bundler (leaving not enough gas for inner call) + revert FailedOp(opIndex, "AA95 out of gas"); + } else if (innerRevertCode == INNER_REVERT_LOW_PREFUND) { + // innerCall reverted on prefund too low. treat entire prefund as "gas cost" + uint256 actualGas = preGas - gasleft() + opInfo.preOpGas; + uint256 actualGasCost = opInfo.prefund; + _emitPrefundTooLow(opInfo); + _emitUserOperationEvent(opInfo, false, actualGasCost, actualGas); + collected = actualGasCost; + } else { + uint256 freePtr = _getFreePtr(); + emit PostOpRevertReason( + opInfo.userOpHash, + opInfo.mUserOp.sender, + opInfo.mUserOp.nonce, + Exec.getReturnData(REVERT_REASON_MAX_LEN) + ); + _restoreFreePtr(freePtr); + + uint256 actualGas = preGas - gasleft() + opInfo.preOpGas; + collected = _postExecution(IPaymaster.PostOpMode.postOpReverted, opInfo, context, actualGas); + } + } + } + + /** + * Emit the UserOperationEvent for the given UserOperation. + * + * @param opInfo - The details of the current UserOperation. + * @param success - Whether the execution of the UserOperation has succeeded or not. + * @param actualGasCost - The actual cost of the consumed gas charged from the sender or the paymaster. + * @param actualGas - The actual amount of gas used. + */ + function _emitUserOperationEvent(UserOpInfo memory opInfo, bool success, uint256 actualGasCost, uint256 actualGas) + internal + virtual + { + emit UserOperationEvent( + opInfo.userOpHash, + opInfo.mUserOp.sender, + opInfo.mUserOp.paymaster, + opInfo.mUserOp.nonce, + success, + actualGasCost, + actualGas + ); + } + + /** + * Emit the UserOperationPrefundTooLow event for the given UserOperation. + * + * @param opInfo - The details of the current UserOperation. + */ + function _emitPrefundTooLow(UserOpInfo memory opInfo) internal virtual { + emit UserOperationPrefundTooLow(opInfo.userOpHash, opInfo.mUserOp.sender, opInfo.mUserOp.nonce); + } + + /** + * Iterate over calldata PackedUserOperation array and perform account and paymaster validation. + * @notice UserOpInfo is a global array of all UserOps while PackedUserOperation is grouped per aggregator. + * + * @param ops - an array of UserOps to be validated + * @param opInfos - an array of UserOp metadata being read and filled in during this function's execution + * @param expectedAggregator - an address of the aggregator specified for a given UserOp if any, or address(0) + * @param opIndexOffset - an offset for the index between 'ops' and 'opInfos' arrays, see the notice. + * @return opsLen - processed UserOps (length of "ops" array) + */ + function _iterateValidationPhase( + PackedUserOperation[] calldata ops, + UserOpInfo[] memory opInfos, + address expectedAggregator, + uint256 opIndexOffset + ) internal returns (uint256 opsLen) { + unchecked { + opsLen = ops.length; + for (uint256 i = 0; i < opsLen; i++) { + UserOpInfo memory opInfo = opInfos[opIndexOffset + i]; + (uint256 validationData, uint256 pmValidationData) = + _validatePrepayment(opIndexOffset + i, ops[i], opInfo); + _validateAccountAndPaymasterValidationData( + opIndexOffset + i, validationData, pmValidationData, expectedAggregator + ); + } + } + } + + /** + * A memory copy of UserOp static fields only. + * Excluding: callData, initCode and signature. Replacing paymasterAndData with paymaster. + */ + struct MemoryUserOp { + address sender; + uint256 nonce; + uint256 verificationGasLimit; + uint256 callGasLimit; + uint256 paymasterVerificationGasLimit; + uint256 paymasterPostOpGasLimit; + uint256 preVerificationGas; + address paymaster; + uint256 maxFeePerGas; + uint256 maxPriorityFeePerGas; + } + + struct UserOpInfo { + MemoryUserOp mUserOp; + bytes32 userOpHash; + uint256 prefund; + uint256 contextOffset; + uint256 preOpGas; + } + + /** + * Inner function to handle a UserOperation. + * Must be declared "external" to open a call context, but it can only be called by handleOps. + * @param callData - The callData to execute. + * @param opInfo - The UserOpInfo struct. + * @param context - The context bytes. + * @return actualGasCost - the actual cost in eth this UserOperation paid for gas + */ + function innerHandleOp(bytes memory callData, UserOpInfo memory opInfo, bytes calldata context) + external + returns (uint256 actualGasCost) + { + uint256 preGas = gasleft(); + require(msg.sender == address(this), "AA92 internal call only"); + MemoryUserOp memory mUserOp = opInfo.mUserOp; + + uint256 callGasLimit = mUserOp.callGasLimit; + unchecked { + // handleOps was called with gas limit too low. abort entire bundle. + if (gasleft() * 63 / 64 < callGasLimit + mUserOp.paymasterPostOpGasLimit + INNER_GAS_OVERHEAD) { + assembly ("memory-safe") { + mstore(0, INNER_OUT_OF_GAS) + revert(0, 32) + } + } + } + + IPaymaster.PostOpMode mode = IPaymaster.PostOpMode.opSucceeded; + if (callData.length > 0) { + bool success = Exec.call(mUserOp.sender, 0, callData, callGasLimit); + if (!success) { + uint256 freePtr = _getFreePtr(); + bytes memory result = Exec.getReturnData(REVERT_REASON_MAX_LEN); + if (result.length > 0) { + emit UserOperationRevertReason(opInfo.userOpHash, mUserOp.sender, mUserOp.nonce, result); + } + _restoreFreePtr(freePtr); + mode = IPaymaster.PostOpMode.opReverted; + } + } + + unchecked { + uint256 actualGas = preGas - gasleft() + opInfo.preOpGas; + return _postExecution(mode, opInfo, context, actualGas); + } + } + + /** + * Copy general fields from userOp into the memory opInfo structure. + * @param userOp - The user operation. + * @param mUserOp - The memory user operation. + */ + function _copyUserOpToMemory(PackedUserOperation calldata userOp, MemoryUserOp memory mUserOp) + internal + pure + virtual + { + mUserOp.sender = userOp.sender; + mUserOp.nonce = userOp.nonce; + (mUserOp.verificationGasLimit, mUserOp.callGasLimit) = UserOperationLib.unpackUints(userOp.accountGasLimits); + mUserOp.preVerificationGas = userOp.preVerificationGas; + (mUserOp.maxPriorityFeePerGas, mUserOp.maxFeePerGas) = UserOperationLib.unpackUints(userOp.gasFees); + bytes calldata paymasterAndData = userOp.paymasterAndData; + if (paymasterAndData.length > 0) { + require(paymasterAndData.length >= UserOperationLib.PAYMASTER_DATA_OFFSET, "AA93 invalid paymasterAndData"); + address paymaster; + (paymaster, mUserOp.paymasterVerificationGasLimit, mUserOp.paymasterPostOpGasLimit) = + UserOperationLib.unpackPaymasterStaticFields(paymasterAndData); + require(paymaster != address(0), "AA98 invalid paymaster"); + mUserOp.paymaster = paymaster; + } + } + + /** + * Get the required prefunded gas fee amount for an operation. + * + * @param mUserOp - The user operation in memory. + * @return requiredPrefund - the required amount. + */ + function _getRequiredPrefund(MemoryUserOp memory mUserOp) internal pure virtual returns (uint256 requiredPrefund) { + unchecked { + uint256 requiredGas = mUserOp.verificationGasLimit + mUserOp.callGasLimit + + mUserOp.paymasterVerificationGasLimit + mUserOp.paymasterPostOpGasLimit + mUserOp.preVerificationGas; + + requiredPrefund = requiredGas * mUserOp.maxFeePerGas; + } + } + + /** + * Create sender smart contract account if init code is provided. + * @param opIndex - The operation index. + * @param opInfo - The operation info. + * @param initCode - The init code for the smart contract account. + */ + function _createSenderIfNeeded(uint256 opIndex, UserOpInfo memory opInfo, bytes calldata initCode) + internal + virtual + { + if (initCode.length != 0) { + address sender = opInfo.mUserOp.sender; + if (Eip7702Support._isEip7702InitCode(initCode)) { + if (initCode.length > 20) { + // Already validated it is an EIP-7702 delegate (and hence, already has code) - see getUserOpHash() + // Note: Can be called multiple times as long as an appropriate initCode is supplied + senderCreator().initEip7702Sender{gas: opInfo.mUserOp.verificationGasLimit}(sender, initCode[20:]); + } + return; + } + if (sender.code.length != 0) { + revert FailedOp(opIndex, "AA10 sender already constructed"); + } + if (initCode.length < 20) { + revert FailedOp(opIndex, "AA99 initCode too small"); + } + address sender1 = senderCreator().createSender{gas: opInfo.mUserOp.verificationGasLimit}(initCode); + if (sender1 == address(0)) { + revert FailedOp(opIndex, "AA13 initCode failed or OOG"); + } + if (sender1 != sender) { + revert FailedOp(opIndex, "AA14 initCode must return sender"); + } + if (sender1.code.length == 0) { + revert FailedOp(opIndex, "AA15 initCode must create sender"); + } + address factory = address(bytes20(initCode[0:20])); + emit AccountDeployed(opInfo.userOpHash, sender, factory, opInfo.mUserOp.paymaster); + } + } + + /** + * Call account.validateUserOp. + * Revert (with FailedOp) in case validateUserOp reverts, or account didn't send required prefund. + * Decrement account's deposit if needed. + * @param opIndex - The operation index. + * @param op - The user operation. + * @param opInfo - The operation info. + * @param requiredPrefund - The required prefund amount. + * @return validationData - The account's validationData. + */ + function _validateAccountPrepayment( + uint256 opIndex, + PackedUserOperation calldata op, + UserOpInfo memory opInfo, + uint256 requiredPrefund + ) internal virtual returns (uint256 validationData) { + unchecked { + MemoryUserOp memory mUserOp = opInfo.mUserOp; + address sender = mUserOp.sender; + _createSenderIfNeeded(opIndex, opInfo, op.initCode); + address paymaster = mUserOp.paymaster; + uint256 missingAccountFunds = 0; + if (paymaster == address(0)) { + uint256 bal = balanceOf(sender); + missingAccountFunds = bal > requiredPrefund ? 0 : requiredPrefund - bal; + } + validationData = _callValidateUserOp(opIndex, op, opInfo, missingAccountFunds); + if (paymaster == address(0)) { + if (!_tryDecrementDeposit(sender, requiredPrefund)) { + revert FailedOp(opIndex, "AA21 didn't pay prefund"); + } + } + } + } + + /** + * Make a call to the sender.validateUserOp() function. + * Handle wrong output size by reverting with a FailedOp error. + * + * @param opIndex - index of the UserOperation in the bundle. + * @param op - the packed UserOperation object. + * @param opInfo - the in-memory UserOperation information. + * @param missingAccountFunds - the amount of deposit the account has to make to cover the UserOperation gas. + */ + function _callValidateUserOp( + uint256 opIndex, + PackedUserOperation calldata op, + UserOpInfo memory opInfo, + uint256 missingAccountFunds + ) internal virtual returns (uint256 validationData) { + uint256 gasLimit = opInfo.mUserOp.verificationGasLimit; + address sender = opInfo.mUserOp.sender; + bool success; + { + uint256 saveFreePtr = _getFreePtr(); + bytes memory callData = + abi.encodeCall(IAccount.validateUserOp, (op, opInfo.userOpHash, missingAccountFunds)); + assembly ("memory-safe") { + success := call(gasLimit, sender, 0, add(callData, 0x20), mload(callData), 0, 32) + validationData := mload(0) + // any return data size other than 32 is considered failure + if iszero(eq(returndatasize(), 32)) { success := 0 } + } + _restoreFreePtr(saveFreePtr); + } + if (!success) { + if (sender.code.length == 0) { + revert FailedOp(opIndex, "AA20 account not deployed"); + } else { + revert FailedOpWithRevert(opIndex, "AA23 reverted", Exec.getReturnData(REVERT_REASON_MAX_LEN)); + } + } + } + + /** + * In case the request has a paymaster: + * - Validate paymaster has enough deposit. + * - Call paymaster.validatePaymasterUserOp. + * - Revert with proper FailedOp in case paymaster reverts. + * - Decrement paymaster's deposit. + * @param opIndex - The operation index. + * @param op - The user operation. + * @param opInfo - The operation info. + * @return context - The Paymaster-provided value to be passed to the 'postOp' function later + * @return validationData - The Paymaster's validationData. + */ + function _validatePaymasterPrepayment(uint256 opIndex, PackedUserOperation calldata op, UserOpInfo memory opInfo) + internal + virtual + returns (bytes memory context, uint256 validationData) + { + unchecked { + uint256 preGas = gasleft(); + MemoryUserOp memory mUserOp = opInfo.mUserOp; + address paymaster = mUserOp.paymaster; + uint256 requiredPreFund = opInfo.prefund; + if (!_tryDecrementDeposit(paymaster, requiredPreFund)) { + revert FailedOp(opIndex, "AA31 paymaster deposit too low"); + } + uint256 pmVerificationGasLimit = mUserOp.paymasterVerificationGasLimit; + (context, validationData) = _callValidatePaymasterUserOp(opIndex, op, opInfo); + if (preGas - gasleft() > pmVerificationGasLimit) { + revert FailedOp(opIndex, "AA36 over paymasterVerificationGasLimit"); + } + } + } + + function _callValidatePaymasterUserOp(uint256 opIndex, PackedUserOperation calldata op, UserOpInfo memory opInfo) + internal + returns (bytes memory context, uint256 validationData) + { + uint256 freePtr = _getFreePtr(); + bytes memory validatePaymasterCall = + abi.encodeCall(IPaymaster.validatePaymasterUserOp, (op, opInfo.userOpHash, opInfo.prefund)); + address paymaster = opInfo.mUserOp.paymaster; + uint256 paymasterVerificationGasLimit = opInfo.mUserOp.paymasterVerificationGasLimit; + bool success; + uint256 contextLength; + uint256 contextOffset; + uint256 maxContextLength; + uint256 len; + assembly ("memory-safe") { + success := call( + paymasterVerificationGasLimit, + paymaster, + 0, + add(validatePaymasterCall, 0x20), + mload(validatePaymasterCall), + 0, + 0 + ) + len := returndatasize() + // return data from validatePaymasterUserOp is (bytes context, validationData) + // encoded as: + // 32 bytes offset of context (always 64) + // 32 bytes of validationData + // 32 bytes of context length + // context data (rounded up, to 32 bytes boundary) + // so entire buffer size is (at least) 96+content.length. + // + // we use freePtr, fetched before calling encodeCall, as return data pointer. + // this way we reuse that memory without unnecessary memory expansion + returndatacopy(freePtr, 0, len) + validationData := mload(add(freePtr, 32)) + contextOffset := mload(freePtr) + maxContextLength := sub(len, 96) + context := add(freePtr, 64) + contextLength := mload(context) + } + + unchecked { + if (!success || contextOffset != 64 || contextLength + 31 < maxContextLength) { + revert FailedOpWithRevert(opIndex, "AA33 reverted", Exec.getReturnData(REVERT_REASON_MAX_LEN)); + } + } + finalizeAllocation(freePtr, len); + } + + /** + * Revert if either account validationData or paymaster validationData is expired. + * @param opIndex - The operation index. + * @param validationData - The account validationData. + * @param paymasterValidationData - The paymaster validationData. + * @param expectedAggregator - The expected aggregator. + */ + function _validateAccountAndPaymasterValidationData( + uint256 opIndex, + uint256 validationData, + uint256 paymasterValidationData, + address expectedAggregator + ) internal view virtual { + (address aggregator, bool outOfTimeRange) = _getValidationData(validationData); + if (expectedAggregator != aggregator) { + revert FailedOp(opIndex, "AA24 signature error"); + } + if (outOfTimeRange) { + revert FailedOp(opIndex, "AA22 expired or not due"); + } + // pmAggregator is not a real signature aggregator: we don't have logic to handle it as address. + // Non-zero address means that the paymaster fails due to some signature check (which is ok only during estimation). + address pmAggregator; + (pmAggregator, outOfTimeRange) = _getValidationData(paymasterValidationData); + if (pmAggregator != address(0)) { + revert FailedOp(opIndex, "AA34 signature error"); + } + if (outOfTimeRange) { + revert FailedOp(opIndex, "AA32 paymaster expired or not due"); + } + } + + /** + * Parse validationData into its components. + * @param validationData - The packed validation data (sigFailed, validAfter, validUntil). + * @return aggregator the aggregator of the validationData + * @return outOfTimeRange true if current time is outside the time range of this validationData. + */ + function _getValidationData(uint256 validationData) + internal + view + virtual + returns (address aggregator, bool outOfTimeRange) + { + if (validationData == 0) { + return (address(0), false); + } + ValidationData memory data = _parseValidationData(validationData); + // solhint-disable-next-line not-rely-on-time + outOfTimeRange = block.timestamp > data.validUntil || block.timestamp <= data.validAfter; + aggregator = data.aggregator; + } + + /** + * Validate account and paymaster (if defined) and + * also make sure total validation doesn't exceed verificationGasLimit. + * This method is called off-chain (simulateValidation()) and on-chain (from handleOps) + * @param opIndex - The index of this userOp into the "opInfos" array. + * @param userOp - The packed calldata UserOperation structure to validate. + * @param outOpInfo - The empty unpacked in-memory UserOperation structure that will be filled in here. + * + * @return validationData - The account's validationData. + * @return paymasterValidationData - The paymaster's validationData. + */ + function _validatePrepayment(uint256 opIndex, PackedUserOperation calldata userOp, UserOpInfo memory outOpInfo) + internal + virtual + returns (uint256 validationData, uint256 paymasterValidationData) + { + uint256 preGas = gasleft(); + MemoryUserOp memory mUserOp = outOpInfo.mUserOp; + _copyUserOpToMemory(userOp, mUserOp); + + // getUserOpHash uses temporary allocations, no required after it returns + uint256 freePtr = _getFreePtr(); + outOpInfo.userOpHash = getUserOpHash(userOp); + _restoreFreePtr(freePtr); + // Validate all numeric values in userOp are well below 128 bit, so they can safely be added + // and multiplied without causing overflow. + uint256 verificationGasLimit = mUserOp.verificationGasLimit; + uint256 maxGasValues = mUserOp.preVerificationGas | verificationGasLimit | mUserOp.callGasLimit + | mUserOp.paymasterVerificationGasLimit | mUserOp.paymasterPostOpGasLimit | mUserOp.maxFeePerGas + | mUserOp.maxPriorityFeePerGas; + require(maxGasValues <= type(uint120).max, FailedOp(opIndex, "AA94 gas values overflow")); + uint256 requiredPreFund = _getRequiredPrefund(mUserOp); + outOpInfo.prefund = requiredPreFund; + validationData = _validateAccountPrepayment(opIndex, userOp, outOpInfo, requiredPreFund); + + require(_validateAndUpdateNonce(mUserOp.sender, mUserOp.nonce), FailedOp(opIndex, "AA25 invalid account nonce")); + unchecked { + if (preGas - gasleft() > verificationGasLimit) { + revert FailedOp(opIndex, "AA26 over verificationGasLimit"); + } + } + + bytes memory context; + if (mUserOp.paymaster != address(0)) { + (context, paymasterValidationData) = _validatePaymasterPrepayment(opIndex, userOp, outOpInfo); + } + unchecked { + outOpInfo.contextOffset = _getOffsetOfMemoryBytes(context); + outOpInfo.preOpGas = preGas - gasleft() + userOp.preVerificationGas; + } + } + + /** + * Process post-operation, called just after the callData is executed. + * If a paymaster is defined and its validation returned a non-empty context, its postOp is called. + * The excess amount is refunded to the account (or paymaster - if it was used in the request). + * @param mode - Whether is called from innerHandleOp, or outside (postOpReverted). + * @param opInfo - UserOp fields and info collected during validation. + * @param context - The context returned in validatePaymasterUserOp. + * @param actualGas - The gas used so far by this user operation. + * + * @return actualGasCost - the actual cost in eth this UserOperation paid for gas + */ + function _postExecution( + IPaymaster.PostOpMode mode, + UserOpInfo memory opInfo, + bytes memory context, + uint256 actualGas + ) internal virtual returns (uint256 actualGasCost) { + uint256 preGas = gasleft(); + unchecked { + address refundAddress; + MemoryUserOp memory mUserOp = opInfo.mUserOp; + uint256 gasPrice = _getUserOpGasPrice(mUserOp); + + address paymaster = mUserOp.paymaster; + // Calculating a penalty for unused execution gas + { + uint256 executionGasUsed = actualGas - opInfo.preOpGas; + // this check is required for the gas used within EntryPoint and not covered by explicit gas limits + actualGas += _getUnusedGasPenalty(executionGasUsed, mUserOp.callGasLimit); + } + uint256 postOpUnusedGasPenalty; + if (paymaster == address(0)) { + refundAddress = mUserOp.sender; + } else { + refundAddress = paymaster; + if (context.length > 0) { + actualGasCost = actualGas * gasPrice; + uint256 postOpPreGas = gasleft(); + if (mode != IPaymaster.PostOpMode.postOpReverted) { + try IPaymaster(paymaster).postOp{gas: mUserOp.paymasterPostOpGasLimit}( + mode, context, actualGasCost, gasPrice + ) { + // solhint-disable-next-line no-empty-blocks + } + catch { + bytes memory reason = Exec.getReturnData(REVERT_REASON_MAX_LEN); + revert PostOpReverted(reason); + } + } + // Calculating a penalty for unused postOp gas + // note that if postOp is reverted, the maximum penalty (10% of postOpGasLimit) is charged. + uint256 postOpGasUsed = postOpPreGas - gasleft(); + postOpUnusedGasPenalty = _getUnusedGasPenalty(postOpGasUsed, mUserOp.paymasterPostOpGasLimit); + } + } + actualGas += preGas - gasleft() + postOpUnusedGasPenalty; + actualGasCost = actualGas * gasPrice; + uint256 prefund = opInfo.prefund; + if (prefund < actualGasCost) { + if (mode == IPaymaster.PostOpMode.postOpReverted) { + actualGasCost = prefund; + _emitPrefundTooLow(opInfo); + _emitUserOperationEvent(opInfo, false, actualGasCost, actualGas); + } else { + assembly ("memory-safe") { + mstore(0, INNER_REVERT_LOW_PREFUND) + revert(0, 32) + } + } + } else { + uint256 refund = prefund - actualGasCost; + _incrementDeposit(refundAddress, refund); + bool success = mode == IPaymaster.PostOpMode.opSucceeded; + _emitUserOperationEvent(opInfo, success, actualGasCost, actualGas); + } + } // unchecked + } + + /** + * The gas price this UserOp agrees to pay. + * Relayer/block builder might submit the TX with higher priorityFee, but the user should not be affected. + * @param mUserOp - The userOp to get the gas price from. + */ + function _getUserOpGasPrice(MemoryUserOp memory mUserOp) internal view returns (uint256) { + unchecked { + uint256 maxFeePerGas = mUserOp.maxFeePerGas; + uint256 maxPriorityFeePerGas = mUserOp.maxPriorityFeePerGas; + return min(maxFeePerGas, maxPriorityFeePerGas + block.basefee); + } + } + + /** + * The offset of the given bytes in memory. + * @param data - The bytes to get the offset of. + */ + function _getOffsetOfMemoryBytes(bytes memory data) internal pure returns (uint256 offset) { + assembly ("memory-safe") { + offset := data + } + } + + /** + * The bytes in memory at the given offset. + * @param offset - The offset to get the bytes from. + */ + function _getMemoryBytesFromOffset(uint256 offset) internal pure returns (bytes memory data) { + assembly ("memory-safe") { + data := offset + } + } + + /** + * save free memory pointer. + * save "free memory" pointer, so that it can be restored later using restoreFreePtr. + * This reduce unneeded memory expansion, and reduce memory expansion cost. + * NOTE: all dynamic allocations between saveFreePtr and restoreFreePtr MUST NOT be used after restoreFreePtr is called. + */ + function _getFreePtr() internal pure returns (uint256 ptr) { + assembly ("memory-safe") { + ptr := mload(0x40) + } + } + + /** + * restore free memory pointer. + * any allocated memory since saveFreePtr is cleared, and MUST NOT be accessed later. + */ + function _restoreFreePtr(uint256 ptr) internal pure { + assembly ("memory-safe") { + mstore(0x40, ptr) + } + } + + function _getUnusedGasPenalty(uint256 gasUsed, uint256 gasLimit) internal pure returns (uint256) { + unchecked { + if (gasLimit <= gasUsed + PENALTY_GAS_THRESHOLD) { + return 0; + } + uint256 unusedGas = gasLimit - gasUsed; + uint256 unusedGasPenalty = (unusedGas * UNUSED_GAS_PENALTY_PERCENT) / 100; + return unusedGasPenalty; + } + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/Helpers.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/Helpers.sol new file mode 100644 index 0000000000..1eb4ce6354 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/Helpers.sol @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable no-inline-assembly */ + +/* + * For simulation purposes, validateUserOp (and validatePaymasterUserOp) + * must return this value in case of signature failure, instead of revert. + */ +uint256 constant SIG_VALIDATION_FAILED = 1; + +/* + * For simulation purposes, validateUserOp (and validatePaymasterUserOp) + * return this value on success. + */ +uint256 constant SIG_VALIDATION_SUCCESS = 0; + +/** + * Returned data from validateUserOp. + * validateUserOp returns a uint256, which is created by `_packedValidationData` and + * parsed by `_parseValidationData`. + * @param aggregator - address(0) - The account validated the signature by itself. + * address(1) - The account failed to validate the signature. + * otherwise - This is an address of a signature aggregator that must + * be used to validate the signature. + * @param validAfter - This UserOp is valid only after this timestamp. + * @param validUntil - Last timestamp this operation is valid at, or 0 for "indefinitely". + */ +struct ValidationData { + address aggregator; + uint48 validAfter; + uint48 validUntil; +} + +/** + * Extract aggregator/sigFailed, validAfter, validUntil. + * Also convert zero validUntil to type(uint48).max. + * @param validationData - The packed validation data. + * @return data - The unpacked in-memory validation data. + */ +function _parseValidationData(uint256 validationData) pure returns (ValidationData memory data) { + address aggregator = address(uint160(validationData)); + uint48 validUntil = uint48(validationData >> 160); + if (validUntil == 0) { + validUntil = type(uint48).max; + } + uint48 validAfter = uint48(validationData >> (48 + 160)); + return ValidationData(aggregator, validAfter, validUntil); +} + +/** + * Helper to pack the return value for validateUserOp. + * @param data - The ValidationData to pack. + * @return the packed validation data. + */ +function _packValidationData(ValidationData memory data) pure returns (uint256) { + return uint160(data.aggregator) | (uint256(data.validUntil) << 160) | (uint256(data.validAfter) << (160 + 48)); +} + +/** + * Helper to pack the return value for validateUserOp, when not using an aggregator. + * @param sigFailed - True for signature failure, false for success. + * @param validUntil - Last timestamp this operation is valid at, or 0 for "indefinitely". + * @param validAfter - First timestamp this UserOperation is valid. + * @return the packed validation data. + */ +function _packValidationData(bool sigFailed, uint48 validUntil, uint48 validAfter) pure returns (uint256) { + return (sigFailed ? SIG_VALIDATION_FAILED : SIG_VALIDATION_SUCCESS) | (uint256(validUntil) << 160) + | (uint256(validAfter) << (160 + 48)); +} + +/** + * keccak function over calldata. + * @dev copy calldata into memory, do keccak and drop allocated memory. Strangely, this is more efficient than letting solidity do it. + * + * @param data - the calldata bytes array to perform keccak on. + * @return ret - the keccak hash of the 'data' array. + */ +function calldataKeccak(bytes calldata data) pure returns (bytes32 ret) { + assembly ("memory-safe") { + let mem := mload(0x40) + let len := data.length + calldatacopy(mem, data.offset, len) + ret := keccak256(mem, len) + } +} + +/** + * The minimum of two numbers. + * @param a - First number. + * @param b - Second number. + * @return - the minimum value. + */ +function min(uint256 a, uint256 b) pure returns (uint256) { + return a < b ? a : b; +} + +/** + * standard solidity memory allocation finalization. + * copied from solidity generated code + * @param memPointer - The current memory pointer + * @param allocationSize - Bytes allocated from memPointer. + */ +function finalizeAllocation(uint256 memPointer, uint256 allocationSize) pure { + assembly ("memory-safe") { + finalize_allocation(memPointer, allocationSize) + + function finalize_allocation(memPtr, size) { + let newFreePtr := add(memPtr, round_up_to_mul_of_32(size)) + mstore(64, newFreePtr) + } + + function round_up_to_mul_of_32(value) -> result { + result := and(add(value, 31), not(31)) + } + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/LibModuleManager.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/LibModuleManager.sol new file mode 100644 index 0000000000..2a2f6374da --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/LibModuleManager.sol @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +library LibModuleManager { + bytes32 constant MODULE_STORAGE_POSITION = keccak256("omni.account.module.storage"); + + struct ModuleStorage { + mapping(address => bool) registeredModules; + } + + function moduleStorage() internal pure returns (ModuleStorage storage ms) { + bytes32 position = MODULE_STORAGE_POSITION; + assembly { + ms.slot := position + } + } + + function isModuleRegistered(address module) internal view returns (bool) { + return moduleStorage().registeredModules[module]; + } + + function registerModule(address module) internal { + require(module != address(0), "Invalid module address"); + require(!isModuleRegistered(module), "Module already registered"); + + uint256 codeSize; + assembly { + codeSize := extcodesize(module) + } + require(codeSize > 0, "Module must be a contract"); + + moduleStorage().registeredModules[module] = true; + } + + function unregisterModule(address module) internal { + require(isModuleRegistered(module), "Module not registered"); + moduleStorage().registeredModules[module] = false; + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/NonceManager.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/NonceManager.sol new file mode 100644 index 0000000000..2919e6212a --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/NonceManager.sol @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: GPL-3.0 +/* + * Based on NonceManager.sol from https://github.com/eth-infinitism/account-abstraction + * Licensed under GNU General Public License v3.0 + */ +pragma solidity ^0.8.28; + +import "../interfaces/INonceManager.sol"; + +/** + * nonce management functionality + */ +abstract contract NonceManager is INonceManager { + /** + * The next valid sequence number for a given nonce key. + */ + mapping(address => mapping(uint192 => uint256)) public nonceSequenceNumber; + + /// @inheritdoc INonceManager + function getNonce(address sender, uint192 key) public view override returns (uint256 nonce) { + return nonceSequenceNumber[sender][key] | (uint256(key) << 64); + } + + /// @inheritdoc INonceManager + function incrementNonce(uint192 key) external override { + nonceSequenceNumber[msg.sender][key]++; + } + + /** + * validate nonce uniqueness for this account. + * called just after validateUserOp() + * @return true if the nonce was incremented successfully. + * false if the current nonce doesn't match the given one. + */ + function _validateAndUpdateNonce(address sender, uint256 nonce) internal returns (bool) { + uint192 key = uint192(nonce >> 64); + uint64 seq = uint64(nonce); + return nonceSequenceNumber[sender][key]++ == seq; + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/SenderCreator.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/SenderCreator.sol new file mode 100644 index 0000000000..40edf5fc2f --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/SenderCreator.sol @@ -0,0 +1,56 @@ +// SPDX-License-Identifier: GPL-3.0 +/* + * Based on SenderCreator.sol from https://github.com/eth-infinitism/account-abstraction + * Licensed under GNU General Public License v3.0 + */ +pragma solidity ^0.8.28; +/* solhint-disable avoid-low-level-calls */ +/* solhint-disable no-inline-assembly */ + +import "../interfaces/ISenderCreator.sol"; +import "../interfaces/IEntryPoint.sol"; +import "../utils/Exec.sol"; + +/** + * Helper contract for EntryPoint, to call userOp.initCode from a "neutral" address, + * which is explicitly not the entryPoint itself. + */ +contract SenderCreator is ISenderCreator { + address public immutable entryPoint; + + constructor() { + entryPoint = msg.sender; + } + + uint256 private constant REVERT_REASON_MAX_LEN = 2048; + + /** + * Call the "initCode" factory to create and return the sender account address. + * @param initCode - The initCode value from a UserOp. contains 20 bytes of factory address, + * followed by calldata. + * @return sender - The returned address of the created account, or zero address on failure. + */ + function createSender(bytes calldata initCode) external returns (address sender) { + require(msg.sender == entryPoint, "AA97 should call from EntryPoint"); + address factory = address(bytes20(initCode[0:20])); + bytes memory initCallData = initCode[20:]; + bool success; + assembly ("memory-safe") { + success := call(gas(), factory, 0, add(initCallData, 0x20), mload(initCallData), 0, 32) + if success { sender := mload(0) } + } + } + + /// @inheritdoc ISenderCreator + function initEip7702Sender(address sender, bytes memory initCallData) external { + require(msg.sender == entryPoint, "AA97 should call from EntryPoint"); + bool success; + assembly ("memory-safe") { + success := call(gas(), sender, 0, add(initCallData, 0x20), mload(initCallData), 0, 0) + } + if (!success) { + bytes memory result = Exec.getReturnData(REVERT_REASON_MAX_LEN); + revert IEntryPoint.FailedOpWithRevert(0, "AA13 EIP7702 sender init failed", result); + } + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/SimplePaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/SimplePaymaster.sol new file mode 100644 index 0000000000..892a081228 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/SimplePaymaster.sol @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: GPL-3.0 +/* + * Based on SimplePaymaster.sol from https://github.com/eth-infinitism/account-abstraction + * Licensed under GNU General Public License v3.0 + */ +pragma solidity ^0.8.28; + +import "./BasePaymaster.sol"; +import "../interfaces/PackedUserOperation.sol"; + +/** + * Simple paymaster implementation that sponsors gas for any user operation + * submitted by authorized off-chain bundlers. + */ +contract SimplePaymaster is BasePaymaster { + // Mapping of authorized bundler addresses + mapping(address => bool) public authorizedBundlers; + + event UserOpSponsored(address indexed account, uint256 actualGasCost); + event AuthorizedBundlerUpdated(address indexed bundler, bool authorized); + + constructor(IEntryPoint _entryPoint, address _initialBundler) BasePaymaster(_entryPoint) { + authorizedBundlers[_initialBundler] = true; + emit AuthorizedBundlerUpdated(_initialBundler, true); + } + + /** + * Validate a user operation. + * Only sponsor operations submitted by authorized bundlers. + */ + function _validatePaymasterUserOp( + PackedUserOperation calldata userOp, + bytes32, + /* userOpHash */ + uint256 maxCost + ) + internal + view + override + returns (bytes memory context, uint256 validationData) + { + // Check if the transaction is being submitted by an authorized bundler + if (!authorizedBundlers[tx.origin]) { + // Reject - not from an authorized bundler + return ("", 1); + } + + // Check if we have enough deposit to cover the cost + uint256 ourDeposit = entryPoint.balanceOf(address(this)); + if (ourDeposit < maxCost) { + // Reject - insufficient funds + return ("", 1); + } + + // Accept the operation - sponsor any account as long as bundler is authorized + // Return smart account address in context for postOp logging + return (abi.encode(userOp.sender), 0); + } + + /** + * Post-operation handler. + * Log the sponsored operation. + */ + function _postOp( + PostOpMode, /* mode */ + bytes calldata context, + uint256 actualGasCost, + uint256 /* actualUserOpFeePerGas */ + ) + internal + override + { + // Decode sender from context + address sender = abi.decode(context, (address)); + + // Log the sponsored operation + emit UserOpSponsored(sender, actualGasCost); + } + + /** + * Add or remove an authorized bundler. + */ + function setAuthorizedBundler(address bundler, bool authorized) external onlyOwner { + authorizedBundlers[bundler] = authorized; + emit AuthorizedBundlerUpdated(bundler, authorized); + } + + /** + * Allow contract to receive ETH and automatically deposit to EntryPoint. + */ + receive() external payable { + entryPoint.depositTo{value: msg.value}(address(this)); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/StakeManager.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/StakeManager.sol new file mode 100644 index 0000000000..1e4f16aa81 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/StakeManager.sol @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: GPL-3.0 +/* + * Based on StakeManager.sol from https://github.com/eth-infinitism/account-abstraction + * Licensed under GNU General Public License v3.0 + */ +pragma solidity ^0.8.28; + +import "../interfaces/IStakeManager.sol"; + +/* solhint-disable avoid-low-level-calls */ +/* solhint-disable not-rely-on-time */ + +/** + * Manage deposits and stakes. + * Deposit is just a balance used to pay for UserOperations (either by a paymaster or an account). + * Stake is value locked for at least "unstakeDelay" by a paymaster. + */ +abstract contract StakeManager is IStakeManager { + /// maps paymaster to their deposits and stakes + mapping(address => DepositInfo) private deposits; + + /// @inheritdoc IStakeManager + function getDepositInfo(address account) external view returns (DepositInfo memory info) { + return deposits[account]; + } + + /** + * Internal method to return just the stake info. + * @param addr - The account to query. + */ + function _getStakeInfo(address addr) internal view returns (StakeInfo memory info) { + DepositInfo storage depositInfo = deposits[addr]; + info.stake = depositInfo.stake; + info.unstakeDelaySec = depositInfo.unstakeDelaySec; + } + + /// @inheritdoc IStakeManager + function balanceOf(address account) public view returns (uint256) { + return deposits[account].deposit; + } + + receive() external payable { + depositTo(msg.sender); + } + + /** + * Increments an account's deposit. + * @param account - The account to increment. + * @param amount - The amount to increment by. + * @return the updated deposit of this account + */ + function _incrementDeposit(address account, uint256 amount) internal returns (uint256) { + unchecked { + DepositInfo storage info = deposits[account]; + uint256 newAmount = info.deposit + amount; + info.deposit = newAmount; + return newAmount; + } + } + + /** + * Try to decrement the account's deposit. + * @param account - The account to decrement. + * @param amount - The amount to decrement by. + * @return true if the decrement succeeded (that is, previous balance was at least that amount) + */ + function _tryDecrementDeposit(address account, uint256 amount) internal returns (bool) { + unchecked { + DepositInfo storage info = deposits[account]; + uint256 currentDeposit = info.deposit; + if (currentDeposit < amount) { + return false; + } + info.deposit = currentDeposit - amount; + return true; + } + } + + /// @inheritdoc IStakeManager + function depositTo(address account) public payable virtual { + uint256 newDeposit = _incrementDeposit(account, msg.value); + emit Deposited(account, newDeposit); + } + + /// @inheritdoc IStakeManager + function addStake(uint32 unstakeDelaySec) external payable { + DepositInfo storage info = deposits[msg.sender]; + require(unstakeDelaySec > 0, "must specify unstake delay"); + require(unstakeDelaySec >= info.unstakeDelaySec, "cannot decrease unstake time"); + uint256 stake = info.stake + msg.value; + require(stake > 0, "no stake specified"); + require(stake <= type(uint112).max, "stake overflow"); + deposits[msg.sender] = DepositInfo(info.deposit, true, uint112(stake), unstakeDelaySec, 0); + emit StakeLocked(msg.sender, stake, unstakeDelaySec); + } + + /// @inheritdoc IStakeManager + function unlockStake() external { + DepositInfo storage info = deposits[msg.sender]; + require(info.unstakeDelaySec != 0, "not staked"); + require(info.staked, "already unstaking"); + uint48 withdrawTime = uint48(block.timestamp) + info.unstakeDelaySec; + info.withdrawTime = withdrawTime; + info.staked = false; + emit StakeUnlocked(msg.sender, withdrawTime); + } + + /// @inheritdoc IStakeManager + function withdrawStake(address payable withdrawAddress) external { + DepositInfo storage info = deposits[msg.sender]; + uint256 stake = info.stake; + require(stake > 0, "No stake to withdraw"); + require(info.withdrawTime > 0, "must call unlockStake() first"); + require(info.withdrawTime <= block.timestamp, "Stake withdrawal is not due"); + info.unstakeDelaySec = 0; + info.withdrawTime = 0; + info.stake = 0; + emit StakeWithdrawn(msg.sender, withdrawAddress, stake); + (bool success,) = withdrawAddress.call{value: stake}(""); + require(success, "failed to withdraw stake"); + } + + /// @inheritdoc IStakeManager + function withdrawTo(address payable withdrawAddress, uint256 withdrawAmount) external { + DepositInfo storage info = deposits[msg.sender]; + uint256 currentDeposit = info.deposit; + require(withdrawAmount <= currentDeposit, "Withdraw amount too large"); + info.deposit = currentDeposit - withdrawAmount; + emit Withdrawn(msg.sender, withdrawAddress, withdrawAmount); + (bool success,) = withdrawAddress.call{value: withdrawAmount}(""); + require(success, "failed to withdraw"); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/core/UserOperationLib.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/UserOperationLib.sol new file mode 100644 index 0000000000..c90f7adf3c --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/core/UserOperationLib.sol @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable no-inline-assembly */ + +import "../interfaces/PackedUserOperation.sol"; +import {calldataKeccak, min} from "./Helpers.sol"; + +/** + * Utility functions helpful when working with UserOperation structs. + */ +library UserOperationLib { + uint256 public constant PAYMASTER_VALIDATION_GAS_OFFSET = 20; + uint256 public constant PAYMASTER_POSTOP_GAS_OFFSET = 36; + uint256 public constant PAYMASTER_DATA_OFFSET = 52; + + /** + * Relayer/block builder might submit the TX with higher priorityFee, + * but the user should not pay above what he signed for. + * @param userOp - The user operation data. + */ + function gasPrice(PackedUserOperation calldata userOp) internal view returns (uint256) { + unchecked { + (uint256 maxPriorityFeePerGas, uint256 maxFeePerGas) = unpackUints(userOp.gasFees); + return min(maxFeePerGas, maxPriorityFeePerGas + block.basefee); + } + } + + bytes32 internal constant PACKED_USEROP_TYPEHASH = keccak256( + "PackedUserOperation(address sender,uint256 nonce,bytes initCode,bytes callData,bytes32 accountGasLimits,uint256 preVerificationGas,bytes32 gasFees,bytes paymasterAndData)" + ); + + /** + * Pack the user operation data into bytes for hashing. + * @param userOp - The user operation data. + * @param overrideInitCodeHash - If set, encode this instead of the initCode field in the userOp. + */ + function encode(PackedUserOperation calldata userOp, bytes32 overrideInitCodeHash) + internal + pure + returns (bytes memory ret) + { + address sender = userOp.sender; + uint256 nonce = userOp.nonce; + bytes32 hashInitCode = overrideInitCodeHash != 0 ? overrideInitCodeHash : calldataKeccak(userOp.initCode); + bytes32 hashCallData = calldataKeccak(userOp.callData); + bytes32 accountGasLimits = userOp.accountGasLimits; + uint256 preVerificationGas = userOp.preVerificationGas; + bytes32 gasFees = userOp.gasFees; + bytes32 hashPaymasterAndData = calldataKeccak(userOp.paymasterAndData); + + return abi.encode( + UserOperationLib.PACKED_USEROP_TYPEHASH, + sender, + nonce, + hashInitCode, + hashCallData, + accountGasLimits, + preVerificationGas, + gasFees, + hashPaymasterAndData + ); + } + + function unpackUints(bytes32 packed) internal pure returns (uint256 high128, uint256 low128) { + return (unpackHigh128(packed), unpackLow128(packed)); + } + + // Unpack just the high 128-bits from a packed value + function unpackHigh128(bytes32 packed) internal pure returns (uint256) { + return uint256(packed) >> 128; + } + + // Unpack just the low 128-bits from a packed value + function unpackLow128(bytes32 packed) internal pure returns (uint256) { + return uint128(uint256(packed)); + } + + function unpackMaxPriorityFeePerGas(PackedUserOperation calldata userOp) internal pure returns (uint256) { + return unpackHigh128(userOp.gasFees); + } + + function unpackMaxFeePerGas(PackedUserOperation calldata userOp) internal pure returns (uint256) { + return unpackLow128(userOp.gasFees); + } + + function unpackVerificationGasLimit(PackedUserOperation calldata userOp) internal pure returns (uint256) { + return unpackHigh128(userOp.accountGasLimits); + } + + function unpackCallGasLimit(PackedUserOperation calldata userOp) internal pure returns (uint256) { + return unpackLow128(userOp.accountGasLimits); + } + + function unpackPaymasterVerificationGasLimit(PackedUserOperation calldata userOp) internal pure returns (uint256) { + return uint128(bytes16(userOp.paymasterAndData[PAYMASTER_VALIDATION_GAS_OFFSET:PAYMASTER_POSTOP_GAS_OFFSET])); + } + + function unpackPostOpGasLimit(PackedUserOperation calldata userOp) internal pure returns (uint256) { + return uint128(bytes16(userOp.paymasterAndData[PAYMASTER_POSTOP_GAS_OFFSET:PAYMASTER_DATA_OFFSET])); + } + + function unpackPaymasterStaticFields(bytes calldata paymasterAndData) + internal + pure + returns (address paymaster, uint256 validationGasLimit, uint256 postOpGasLimit) + { + return ( + address(bytes20(paymasterAndData[:PAYMASTER_VALIDATION_GAS_OFFSET])), + uint128(bytes16(paymasterAndData[PAYMASTER_VALIDATION_GAS_OFFSET:PAYMASTER_POSTOP_GAS_OFFSET])), + uint128(bytes16(paymasterAndData[PAYMASTER_POSTOP_GAS_OFFSET:PAYMASTER_DATA_OFFSET])) + ); + } + + /** + * Hash the user operation data. + * @param userOp - The user operation data. + * @param overrideInitCodeHash - If set, the initCode hash will be replaced with this value just for UserOp hashing. + */ + function hash(PackedUserOperation calldata userOp, bytes32 overrideInitCodeHash) internal pure returns (bytes32) { + return keccak256(encode(userOp, overrideInitCodeHash)); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAccount.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAccount.sol new file mode 100644 index 0000000000..b976f13b8c --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAccount.sol @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +import "./PackedUserOperation.sol"; + +interface IAccount { + /** + * Validate user's signature and nonce + * the entryPoint will make the call to the recipient only if this validation call returns successfully. + * signature failure should be reported by returning SIG_VALIDATION_FAILED (1). + * This allows making a "simulation call" without a valid signature + * Other failures (e.g. nonce mismatch, or invalid signature format) should still revert to signal failure. + * + * @dev Must validate caller is the entryPoint. + * Must validate the signature and nonce + * @param userOp - The operation that is about to be executed. + * @param userOpHash - Hash of the user's request data. can be used as the basis for signature. + * @param missingAccountFunds - Missing funds on the account's deposit in the entrypoint. + * This is the minimum amount to transfer to the sender(entryPoint) to be + * able to make the call. The excess is left as a deposit in the entrypoint + * for future calls. Can be withdrawn anytime using "entryPoint.withdrawTo()". + * In case there is a paymaster in the request (or the current deposit is high + * enough), this value will be zero. + * @return validationData - Packaged ValidationData structure. use `_packValidationData` and + * `_unpackValidationData` to encode and decode. + * <20-byte> aggregatorOrSigFail - 0 for valid signature, 1 to mark signature failure, + * otherwise, an address of an "aggregator" contract. + * <6-byte> validUntil - Last timestamp this operation is valid at, or 0 for "indefinitely" + * <6-byte> validAfter - First timestamp this operation is valid + * If an account doesn't use time-range, it is enough to + * return SIG_VALIDATION_FAILED value (1) for signature failure. + * Note that the validation code cannot use block.timestamp (or block.number) directly. + */ + function validateUserOp(PackedUserOperation calldata userOp, bytes32 userOpHash, uint256 missingAccountFunds) + external + returns (uint256 validationData); +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAccountExecute.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAccountExecute.sol new file mode 100644 index 0000000000..b52c2d897c --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAccountExecute.sol @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +import "./PackedUserOperation.sol"; + +interface IAccountExecute { + /** + * Account may implement this execute method. + * passing this methodSig at the beginning of callData will cause the entryPoint to pass the full UserOp (and hash) + * to the account. + * The account should skip the methodSig, and use the callData (and optionally, other UserOp fields) + * + * @param userOp - The operation that was just validated. + * @param userOpHash - Hash of the user's request data. + */ + function executeUserOp(PackedUserOperation calldata userOp, bytes32 userOpHash) external; +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAggregator.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAggregator.sol new file mode 100644 index 0000000000..a833152288 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAggregator.sol @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +import "./PackedUserOperation.sol"; + +/** + * Aggregated Signatures validator. + */ +interface IAggregator { + /** + * Validate an aggregated signature. + * Reverts if the aggregated signature does not match the given list of operations. + * @param userOps - An array of UserOperations to validate the signature for. + * @param signature - The aggregated signature. + */ + function validateSignatures(PackedUserOperation[] calldata userOps, bytes calldata signature) external; + + /** + * Validate the signature of a single userOp. + * This method should be called by bundler after EntryPointSimulation.simulateValidation() returns + * the aggregator this account uses. + * First it validates the signature over the userOp. Then it returns data to be used when creating the handleOps. + * @param userOp - The userOperation received from the user. + * @return sigForUserOp - The value to put into the signature field of the userOp when calling handleOps. + * (usually empty, unless account and aggregator support some kind of "multisig". + */ + function validateUserOpSignature(PackedUserOperation calldata userOp) + external + view + returns (bytes memory sigForUserOp); + + /** + * Aggregate multiple signatures into a single value. + * This method is called off-chain to calculate the signature to pass with handleOps() + * bundler MAY use optimized custom code to perform this aggregation. + * @param userOps - An array of UserOperations to collect the signatures from. + * @return aggregatedSignature - The aggregated signature. + */ + function aggregateSignatures(PackedUserOperation[] calldata userOps) + external + view + returns (bytes memory aggregatedSignature); +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IEntryPoint.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IEntryPoint.sol new file mode 100644 index 0000000000..6b9d30d6a9 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IEntryPoint.sol @@ -0,0 +1,199 @@ +/** + * Account-Abstraction (EIP-4337) singleton EntryPoint implementation. + * Only one instance required on each chain. + * + */ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable avoid-low-level-calls */ +/* solhint-disable no-inline-assembly */ +/* solhint-disable reason-string */ + +import "./PackedUserOperation.sol"; +import "./IStakeManager.sol"; +import "./IAggregator.sol"; +import "./INonceManager.sol"; +import "./ISenderCreator.sol"; + +interface IEntryPoint is IStakeManager, INonceManager { + /** + * + * An event emitted after each successful request. + * @param userOpHash - Unique identifier for the request (hash its entire content, except signature). + * @param sender - The account that generates this request. + * @param paymaster - If non-null, the paymaster that pays for this request. + * @param nonce - The nonce value from the request. + * @param success - True if the sender transaction succeeded, false if reverted. + * @param actualGasCost - Actual amount paid (by account or paymaster) for this UserOperation. + * @param actualGasUsed - Total gas used by this UserOperation (including preVerification, creation, + * validation and execution). + */ + event UserOperationEvent( + bytes32 indexed userOpHash, + address indexed sender, + address indexed paymaster, + uint256 nonce, + bool success, + uint256 actualGasCost, + uint256 actualGasUsed + ); + + /** + * Account "sender" was deployed. + * @param userOpHash - The userOp that deployed this account. UserOperationEvent will follow. + * @param sender - The account that is deployed + * @param factory - The factory used to deploy this account (in the initCode) + * @param paymaster - The paymaster used by this UserOp + */ + event AccountDeployed(bytes32 indexed userOpHash, address indexed sender, address factory, address paymaster); + + /** + * An event emitted if the UserOperation "callData" reverted with non-zero length. + * @param userOpHash - The request unique identifier. + * @param sender - The sender of this request. + * @param nonce - The nonce used in the request. + * @param revertReason - The return bytes from the reverted "callData" call. + */ + event UserOperationRevertReason( + bytes32 indexed userOpHash, address indexed sender, uint256 nonce, bytes revertReason + ); + + /** + * An event emitted if the UserOperation Paymaster's "postOp" call reverted with non-zero length. + * @param userOpHash - The request unique identifier. + * @param sender - The sender of this request. + * @param nonce - The nonce used in the request. + * @param revertReason - The return bytes from the reverted call to "postOp". + */ + event PostOpRevertReason(bytes32 indexed userOpHash, address indexed sender, uint256 nonce, bytes revertReason); + + /** + * UserOp consumed more than prefund. The UserOperation is reverted, and no refund is made. + * @param userOpHash - The request unique identifier. + * @param sender - The sender of this request. + * @param nonce - The nonce used in the request. + */ + event UserOperationPrefundTooLow(bytes32 indexed userOpHash, address indexed sender, uint256 nonce); + + /** + * An event emitted by handleOps() and handleAggregatedOps(), before starting the execution loop. + * Any event emitted before this event, is part of the validation. + */ + event BeforeExecution(); + + /** + * Signature aggregator used by the following UserOperationEvents within this bundle. + * @param aggregator - The aggregator used for the following UserOperationEvents. + */ + event SignatureAggregatorChanged(address indexed aggregator); + + /** + * A custom revert error of handleOps andhandleAggregatedOps, to identify the offending op. + * Should be caught in off-chain handleOps/handleAggregatedOps simulation and not happen on-chain. + * Useful for mitigating DoS attempts against batchers or for troubleshooting of factory/account/paymaster reverts. + * NOTE: If simulateValidation passes successfully, there should be no reason for handleOps to fail on it. + * @param opIndex - Index into the array of ops to the failed one (in simulateValidation, this is always zero). + * @param reason - Revert reason. The string starts with a unique code "AAmn", + * where "m" is "1" for factory, "2" for account and "3" for paymaster issues, + * so a failure can be attributed to the correct entity. + */ + error FailedOp(uint256 opIndex, string reason); + + /** + * A custom revert error of handleOps and handleAggregatedOps, to report a revert by account or paymaster. + * @param opIndex - Index into the array of ops to the failed one (in simulateValidation, this is always zero). + * @param reason - Revert reason. see FailedOp(uint256,string), above + * @param inner - data from inner cought revert reason + * @dev note that inner is truncated to 2048 bytes + */ + error FailedOpWithRevert(uint256 opIndex, string reason, bytes inner); + + error PostOpReverted(bytes returnData); + + /** + * Error case when a signature aggregator fails to verify the aggregated signature it had created. + * @param aggregator The aggregator that failed to verify the signature + */ + error SignatureValidationFailed(address aggregator); + + // Return value of getSenderAddress. + error SenderAddressResult(address sender); + + // UserOps handled, per aggregator. + struct UserOpsPerAggregator { + PackedUserOperation[] userOps; + // Aggregator address + IAggregator aggregator; + // Aggregated signature + bytes signature; + } + + /** + * Execute a batch of UserOperations. + * No signature aggregator is used. + * If any account requires an aggregator (that is, it returned an aggregator when + * performing simulateValidation), then handleAggregatedOps() must be used instead. + * @param ops - The operations to execute. + * @param beneficiary - The address to receive the fees. + */ + function handleOps(PackedUserOperation[] calldata ops, address payable beneficiary) external; + + /** + * Execute a batch of UserOperation with Aggregators + * @param opsPerAggregator - The operations to execute, grouped by aggregator (or address(0) for no-aggregator accounts). + * @param beneficiary - The address to receive the fees. + */ + function handleAggregatedOps(UserOpsPerAggregator[] calldata opsPerAggregator, address payable beneficiary) external; + + /** + * Generate a request Id - unique identifier for this request. + * The request ID is a hash over the content of the userOp (except the signature), entrypoint address, chainId and (optionally) 7702 delegate address + * @param userOp - The user operation to generate the request ID for. + * @return hash the hash of this UserOperation + */ + function getUserOpHash(PackedUserOperation calldata userOp) external view returns (bytes32); + + /** + * Gas and return values during simulation. + * @param preOpGas - The gas used for validation (including preValidationGas) + * @param prefund - The required prefund for this operation + * @param accountValidationData - returned validationData from account. + * @param paymasterValidationData - return validationData from paymaster. + * @param paymasterContext - Returned by validatePaymasterUserOp (to be passed into postOp) + */ + struct ReturnInfo { + uint256 preOpGas; + uint256 prefund; + uint256 accountValidationData; + uint256 paymasterValidationData; + bytes paymasterContext; + } + + /** + * Get counterfactual sender address. + * Calculate the sender contract address that will be generated by the initCode and salt in the UserOperation. + * This method always revert, and returns the address in SenderAddressResult error. + * @notice this method cannot be used for EIP-7702 derived contracts. + * + * @param initCode - The constructor code to be passed into the UserOperation. + */ + function getSenderAddress(bytes memory initCode) external; + + error DelegateAndRevert(bool success, bytes ret); + + /** + * Helper method for dry-run testing. + * @dev calling this method, the EntryPoint will make a delegatecall to the given data, and report (via revert) the result. + * The method always revert, so is only useful off-chain for dry run calls, in cases where state-override to replace + * actual EntryPoint code is less convenient. + * @param target a target contract to make a delegatecall from entrypoint + * @param data data to pass to target in a delegatecall + */ + function delegateAndRevert(address target, bytes calldata data) external; + + /** + * @notice Retrieves the immutable SenderCreator contract which is responsible for deployment of sender contracts. + */ + function senderCreator() external view returns (ISenderCreator); +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IEntryPointSimulations.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IEntryPointSimulations.sol new file mode 100644 index 0000000000..9f0e414366 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IEntryPointSimulations.sol @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +import "./PackedUserOperation.sol"; +import "./IEntryPoint.sol"; + +interface IEntryPointSimulations is IEntryPoint { + // Return value of simulateHandleOp. + struct ExecutionResult { + uint256 preOpGas; + uint256 paid; + uint256 accountValidationData; + uint256 paymasterValidationData; + bool targetSuccess; + bytes targetResult; + } + + /** + * Returned aggregated signature info: + * The aggregator returned by the account, and its current stake. + */ + struct AggregatorStakeInfo { + address aggregator; + StakeInfo stakeInfo; + } + + /** + * Successful result from simulateValidation. + * If the account returns a signature aggregator the "aggregatorInfo" struct is filled in as well. + * @param returnInfo Gas and time-range returned values + * @param senderInfo Stake information about the sender + * @param factoryInfo Stake information about the factory (if any) + * @param paymasterInfo Stake information about the paymaster (if any) + * @param aggregatorInfo Signature aggregation info (if the account requires signature aggregator) + * Bundler MUST use it to verify the signature, or reject the UserOperation. + */ + struct ValidationResult { + ReturnInfo returnInfo; + StakeInfo senderInfo; + StakeInfo factoryInfo; + StakeInfo paymasterInfo; + AggregatorStakeInfo aggregatorInfo; + } + + /** + * Simulate a call to account.validateUserOp and paymaster.validatePaymasterUserOp. + * @dev The node must also verify it doesn't use banned opcodes, and that it doesn't reference storage + * outside the account's data. + * @param userOp - The user operation to validate. + * @return the validation result structure + */ + function simulateValidation(PackedUserOperation calldata userOp) external returns (ValidationResult memory); + + /** + * Simulate full execution of a UserOperation (including both validation and target execution) + * It performs full validation of the UserOperation, but ignores signature error. + * An optional target address is called after the userop succeeds, + * and its value is returned (before the entire call is reverted). + * Note that in order to collect the the success/failure of the target call, it must be executed + * with trace enabled to track the emitted events. + * @param op The UserOperation to simulate. + * @param target - If nonzero, a target address to call after userop simulation. If called, + * the targetSuccess and targetResult are set to the return from that call. + * @param targetCallData - CallData to pass to target address. + * @return the execution result structure + */ + function simulateHandleOp(PackedUserOperation calldata op, address target, bytes calldata targetCallData) + external + returns (ExecutionResult memory); + + /** + * Simulate a batch of UserOperations (similar to handleOps but for simulation) + * It performs full validation and execution simulation for all operations in the batch. + * This is useful for testing entire batches before submitting them to the actual entry point. + * @param ops The array of UserOperations to simulate. + * @param beneficiary The address that will receive the collected fees. + * @return results Array of execution results, one for each UserOperation in the batch. + */ + function simulateHandleOps(PackedUserOperation[] calldata ops, address payable beneficiary) + external + returns (ExecutionResult[] memory results); +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/INonceManager.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/INonceManager.sol new file mode 100644 index 0000000000..c86b3c82f4 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/INonceManager.sol @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +interface INonceManager { + /** + * Return the next nonce for this sender. + * Within a given key, the nonce values are sequenced (starting with zero, and incremented by one on each userop) + * But UserOp with different keys can come with arbitrary order. + * + * @param sender the account address + * @param key the high 192 bit of the nonce + * @return nonce a full nonce to pass for next UserOp with this sender. + */ + function getNonce(address sender, uint192 key) external view returns (uint256 nonce); + + /** + * Manually increment the nonce of the sender. + * This method is exposed just for completeness.. + * Account does NOT need to call it, neither during validation, nor elsewhere, + * as the EntryPoint will update the nonce regardless. + * Possible use-case is call it with various keys to "initialize" their nonces to one, so that future + * UserOperations will not pay extra for the first transaction with a given key. + * + * @param key - the "nonce key" to increment the "nonce sequence" for. + */ + function incrementNonce(uint192 key) external; +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IPaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IPaymaster.sol new file mode 100644 index 0000000000..2acb6d611b --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IPaymaster.sol @@ -0,0 +1,57 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +import "./PackedUserOperation.sol"; + +/** + * The interface exposed by a paymaster contract, who agrees to pay the gas for user's operations. + * A paymaster must hold a stake to cover the required entrypoint stake and also the gas for the transaction. + */ +interface IPaymaster { + enum PostOpMode { + // User op succeeded. + opSucceeded, + // User op reverted. Still has to pay for gas. + opReverted, + // Only used internally in the EntryPoint (cleanup after postOp reverts). Never calling paymaster with this value + postOpReverted + } + + /** + * Payment validation: check if paymaster agrees to pay. + * Must verify sender is the entryPoint. + * Revert to reject this request. + * Note that bundlers will reject this method if it changes the state, unless the paymaster is trusted (whitelisted). + * The paymaster pre-pays using its deposit, and receive back a refund after the postOp method returns. + * @param userOp - The user operation. + * @param userOpHash - Hash of the user's request data. + * @param maxCost - The maximum cost of this transaction (based on maximum gas and gas price from userOp). + * @return context - Value to send to a postOp. Zero length to signify postOp is not required. + * @return validationData - Signature and time-range of this operation, encoded the same as the return + * value of validateUserOperation. + * <20-byte> aggregatorOrSigFail - 0 for valid signature, 1 to mark signature failure, + * other values are invalid for paymaster. + * <6-byte> validUntil - Last timestamp this operation is valid at, or 0 for "indefinitely" + * <6-byte> validAfter - first timestamp this operation is valid + * Note that the validation code cannot use block.timestamp (or block.number) directly. + */ + function validatePaymasterUserOp(PackedUserOperation calldata userOp, bytes32 userOpHash, uint256 maxCost) + external + returns (bytes memory context, uint256 validationData); + + /** + * Post-operation handler. + * Must verify sender is the entryPoint. + * @param mode - Enum with the following options: + * opSucceeded - User operation succeeded. + * opReverted - User op reverted. The paymaster still has to pay for gas. + * postOpReverted - never passed in a call to postOp(). + * @param context - The context value returned by validatePaymasterUserOp + * @param actualGasCost - Actual cost of gas used so far (without this postOp call). + * @param actualUserOpFeePerGas - the gas price this UserOp pays. This value is based on the UserOp's maxFeePerGas + * and maxPriorityFee (and basefee) + * It is not the same as tx.gasprice, which is what the bundler pays. + */ + function postOp(PostOpMode mode, bytes calldata context, uint256 actualGasCost, uint256 actualUserOpFeePerGas) + external; +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/ISenderCreator.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/ISenderCreator.sol new file mode 100644 index 0000000000..42d637ae9b --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/ISenderCreator.sol @@ -0,0 +1,20 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +interface ISenderCreator { + /** + * @dev Creates a new sender contract. + * @return sender Address of the newly created sender contract. + */ + function createSender(bytes calldata initCode) external returns (address sender); + + /** + * Use initCallData to initialize an EIP-7702 account. + * The caller is the EntryPoint contract and it is already verified to be an EIP-7702 account. + * Note: Can be called multiple times as long as an appropriate initCode is supplied + * + * @param sender - the 'sender' EIP-7702 account to be initialized. + * @param initCallData - the call data to be passed to the sender account call. + */ + function initEip7702Sender(address sender, bytes calldata initCallData) external; +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IStakeManager.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IStakeManager.sol new file mode 100644 index 0000000000..86e54955ba --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IStakeManager.sol @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/** + * Manage deposits and stakes. + * Deposit is just a balance used to pay for UserOperations (either by a paymaster or an account). + * Stake is value locked for at least "unstakeDelay" by the staked entity. + */ +interface IStakeManager { + event Deposited(address indexed account, uint256 totalDeposit); + + event Withdrawn(address indexed account, address withdrawAddress, uint256 amount); + + // Emitted when stake or unstake delay are modified. + event StakeLocked(address indexed account, uint256 totalStaked, uint256 unstakeDelaySec); + + // Emitted once a stake is scheduled for withdrawal. + event StakeUnlocked(address indexed account, uint256 withdrawTime); + + event StakeWithdrawn(address indexed account, address withdrawAddress, uint256 amount); + + /** + * @param deposit - The entity's deposit. + * @param staked - True if this entity is staked. + * @param stake - Actual amount of ether staked for this entity. + * @param unstakeDelaySec - Minimum delay to withdraw the stake. + * @param withdrawTime - First block timestamp where 'withdrawStake' will be callable, or zero if already locked. + * @dev Sizes were chosen so that deposit fits into one cell (used during handleOp) + * and the rest fit into a 2nd cell (used during stake/unstake) + * - 112 bit allows for 10^15 eth + * - 48 bit for full timestamp + * - 32 bit allows 150 years for unstake delay + */ + struct DepositInfo { + uint256 deposit; + bool staked; + uint112 stake; + uint32 unstakeDelaySec; + uint48 withdrawTime; + } + + // API struct used by getStakeInfo and simulateValidation. + struct StakeInfo { + uint256 stake; + uint256 unstakeDelaySec; + } + + /** + * Get deposit info. + * @param account - The account to query. + * @return info - Full deposit information of given account. + */ + function getDepositInfo(address account) external view returns (DepositInfo memory info); + + /** + * Get account balance. + * @param account - The account to query. + * @return - The deposit (for gas payment) of the account. + */ + function balanceOf(address account) external view returns (uint256); + + /** + * Add to the deposit of the given account. + * @param account - The account to add to. + */ + function depositTo(address account) external payable; + + /** + * Add to the account's stake - amount and delay + * any pending unstake is first cancelled. + * @param unstakeDelaySec - The new lock duration before the deposit can be withdrawn. + */ + function addStake(uint32 unstakeDelaySec) external payable; + + /** + * Attempt to unlock the stake. + * The value can be withdrawn (using withdrawStake) after the unstake delay. + */ + function unlockStake() external; + + /** + * Withdraw from the (unlocked) stake. + * Must first call unlockStake and wait for the unstakeDelay to pass. + * @param withdrawAddress - The address to send withdrawn value. + */ + function withdrawStake(address payable withdrawAddress) external; + + /** + * Withdraw from the deposit. + * @param withdrawAddress - The address to send withdrawn value. + * @param withdrawAmount - The amount to withdraw. + */ + function withdrawTo(address payable withdrawAddress, uint256 withdrawAmount) external; +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/OwnerType.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/OwnerType.sol new file mode 100644 index 0000000000..e0821bfe68 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/OwnerType.sol @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable no-inline-assembly */ + +/** + * The owner type + * Required as a single byte in the `OmniAccount.initialize` to allow for more fine-grained access control + * + * See P-1644 for more context + * + * The fields of this enum are 1 to 1 mapped to the rust `UserId` type in tee-worker/omni-executor/executor-primitives/src/auth.rs + */ +enum OwnerType { + Pumpx, // 0x00 + Email, // 0x01 + Twitter, // 0x02 + Discord, // 0x03 + Apple, // 0x04 + Substrate, // 0x05 + Evm, // 0x06 + Bitcoin, // 0x07 + Solana, // 0x08 + Google, // 0x09 + Passkey // 0x0a +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/PackedUserOperation.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/PackedUserOperation.sol new file mode 100644 index 0000000000..eb12fb75f9 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/PackedUserOperation.sol @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/** + * User Operation struct + * @param sender - The sender account of this request. + * @param nonce - Unique value the sender uses to verify it is not a replay. + * @param initCode - If set, the account contract will be created by this constructor + * @param callData - The method call to execute on this account. + * @param accountGasLimits - Packed gas limits for validateUserOp and gas limit passed to the callData method call. + * @param preVerificationGas - Gas not calculated by the handleOps method, but added to the gas paid. + * Covers batch overhead. + * @param gasFees - packed gas fields maxPriorityFeePerGas and maxFeePerGas - Same as EIP-1559 gas parameters. + * @param paymasterAndData - If set, this field holds the paymaster address, verification gas limit, postOp gas limit and paymaster-specific extra data + * The paymaster will pay for the transaction instead of the sender. + * @param signature - Sender-verified signature over the entire request, the EntryPoint address and the chain ID. + */ +struct PackedUserOperation { + address sender; + uint256 nonce; + bytes initCode; + bytes callData; + bytes32 accountGasLimits; + uint256 preVerificationGas; + bytes32 gasFees; + bytes paymasterAndData; + bytes signature; +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/Passkey.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/Passkey.sol new file mode 100644 index 0000000000..ee24dcd0cf --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/Passkey.sol @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable no-inline-assembly */ + +/** + * A library to define Passkey related primitives and signature verification methods + * Adapted from https://github.com/ithacaxyz/exp-0001/blob/main/contracts/src/utils/WebAuthnP256.sol + * + * The P256 verification is from openzeppelin implementatons + */ +import "@openzeppelin/contracts/utils/Base64.sol"; +import "@openzeppelin/contracts/utils/cryptography/P256.sol"; + +library Passkey { + struct Metadata { + bytes authData; + /// We assume clientDataJSON.challenge should be UserOpHash + string clientDataJSON; + /// Having these fields to avoid expensive on-chain JSON parsing + uint16 challengeIndex; + uint16 typeIndex; + bool userVerificationRequired; + } + + struct PublicKey { + uint256 x; + uint256 y; + } + + struct Signature { + uint256 r; + uint256 s; + } + + /// Convert PublicKey to hashable storage key + function toKey(PublicKey memory pk) internal pure returns (bytes32) { + return sha256(abi.encodePacked(pk.x, pk.y)); + } + + /// Checks whether substr occurs in str starting at a given byte offset. + function contains(string memory substr, string memory str, uint256 location) internal pure returns (bool) { + bytes memory substrBytes = bytes(substr); + bytes memory strBytes = bytes(str); + + uint256 substrLen = substrBytes.length; + uint256 strLen = strBytes.length; + + for (uint256 i = 0; i < substrLen; i++) { + if (location + i >= strLen) { + return false; + } + + if (substrBytes[i] != strBytes[location + i]) { + return false; + } + } + + return true; + } + + bytes1 constant AUTH_DATA_FLAGS_UP = 0x01; // Bit 0 + bytes1 constant AUTH_DATA_FLAGS_UV = 0x04; // Bit 2 + bytes1 constant AUTH_DATA_FLAGS_BE = 0x08; // Bit 3 + bytes1 constant AUTH_DATA_FLAGS_BS = 0x10; // Bit 4 + + /// Verifies the authFlags in authData. Numbers in inline comment + /// correspond to the same numbered bullets in + /// https://www.w3.org/TR/webauthn-2/#sctn-verifying-assertion. + function checkAuthFlags(bytes1 flags, bool requireUserVerification) internal pure returns (bool) { + // 17. Verify that the UP bit of the flags in authData is set. + if (flags & AUTH_DATA_FLAGS_UP != AUTH_DATA_FLAGS_UP) { + return false; + } + + // 18. If user verification was determined to be required, verify that + // the UV bit of the flags in authData is set. Otherwise, ignore the + // value of the UV flag. + if (requireUserVerification && (flags & AUTH_DATA_FLAGS_UV) != AUTH_DATA_FLAGS_UV) { + return false; + } + + // 19. If the BE bit of the flags in authData is not set, verify that + // the BS bit is not set. + if (flags & AUTH_DATA_FLAGS_BE != AUTH_DATA_FLAGS_BE) { + if (flags & AUTH_DATA_FLAGS_BS == AUTH_DATA_FLAGS_BS) { + return false; + } + } + + return true; + } + + /** + * Verifies a Webauthn P256 signature (Authentication Assertion) as described + * in https://www.w3.org/TR/webauthn-2/#sctn-verifying-assertion. We do not + * verify all the steps as described in the specification, only ones relevant + * to our context. Please carefully read through this list before usage. + * Specifically, we do verify the following: + * - Verify that authData (which comes from the authenticator, + * such as iCloud Keychain) indicates a well-formed assertion. If + * requireUserVerification is set, checks that the authenticator enforced + * user verification. User verification should be required if, + * and only if, options.userVerification is set to required in the request + * - Verifies that the client JSON is of type "webauthn.get", i.e. the client + * was responding to a request to assert authentication. + * - Verifies that the client JSON contains the requested challenge. + * - Finally, verifies that (r, s) constitute a valid signature over both + * the authenicatorData and client JSON, for public key (x, y). + * + * We make some assumptions about the particular use case of this verifier, + * so we do NOT verify the following: + * - Does NOT verify that the origin in the clientDataJSON matches the + * Relying Party's origin: It is considered the authenticator's + * responsibility to ensure that the user is interacting with the correct + * RP. This is enforced by most high quality authenticators properly, + * particularly the iCloud Keychain and Google Password Manager were + * tested. + * - Does NOT verify That c.topOrigin is well-formed: We assume c.topOrigin + * would never be present, i.e. the credentials are never used in a + * cross-origin/iframe context. The website/app set up should disallow + * cross-origin usage of the credentials. This is the default behaviour for + * created credentials in common settings. + * - Does NOT verify that the rpIdHash in authData is the SHA-256 hash of an + * RP ID expected by the Relying Party: This means that we rely on the + * authenticator to properly enforce credentials to be used only by the + * correct RP. This is generally enforced with features like Apple App Site + * Association and Google Asset Links. To protect from edge cases in which + * a previously-linked RP ID is removed from the authorised RP IDs, + * we recommend that messages signed by the authenticator include some + * expiry mechanism. + * - Does NOT verify the credential backup state: This assumes the credential + * backup state is NOT used as part of Relying Party business logic or + * policy. + * - Does NOT verify the values of the client extension outputs: This assumes + * that the Relying Party does not use client extension outputs. + * - Does NOT verify the signature counter: Signature counters are intended + * to enable risk scoring for the Relying Party. This assumes risk scoring + * is not used as part of Relying Party business logic or policy. + * - Does NOT verify the attestation object: This assumes that + * response.attestationObject is NOT present in the response, i.e. the + * RP does not intend to verify an attestation. + */ + function verify(bytes32 challenge, Metadata memory metadata, Signature memory signature, PublicKey memory pk) + internal + view + returns (bool) + { + // Check that authData has good flags + if (metadata.authData.length < 37 || !checkAuthFlags(metadata.authData[32], metadata.userVerificationRequired)) + { + return false; + } + + // Check that response is for an authentication assertion + string memory responseType = '"type":"webauthn.get"'; + if (!contains(responseType, metadata.clientDataJSON, metadata.typeIndex)) { + return false; + } + + // Check that challenge is in the clientDataJSON + string memory challengeB64url = Base64.encodeURL(abi.encodePacked(challenge)); + string memory challengeProperty = string.concat('"challenge":"', challengeB64url, '"'); + + if (!contains(challengeProperty, metadata.clientDataJSON, metadata.challengeIndex)) { + return false; + } + + // Check that the public key signed sha256(authData || sha256(clientDataJSON)) + bytes32 clientDataJSONHash = sha256(bytes(metadata.clientDataJSON)); + bytes32 messageHash = sha256(abi.encodePacked(metadata.authData, clientDataJSONHash)); + + return P256.verify(messageHash, bytes32(signature.r), bytes32(signature.s), bytes32(pk.x), bytes32(pk.y)); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/UserOpSigner.sol b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/UserOpSigner.sol new file mode 100644 index 0000000000..f00918d663 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/interfaces/UserOpSigner.sol @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +/* solhint-disable no-inline-assembly */ + +/** + * The UserOp signature type + * It's the first byte in PackedUserOperation.signature that stands for the signature type. + * Different validation method should be called accordingly. + * + * @param Owner - Signed by user's own EOA address + * @param RootKey - Signed by an authorised evm signer + * @param SessionKey - Signed by a short-lived session key, additional authorisation proof + * (by RootKey) must be provided along. The concatenated signature length + * should be 162 bytes excluding the leading UserOpSigner byte: + * sessionSig (65) | sessionExpiration (32 = uint256) | sessionProof (65) + * @param Passkey - Signed by an authorised Passkey + * Note it's Passkey not PassKey as it's an integral term + */ +enum UserOpSigner { + Owner, // 0x00 + RootKey, // 0x01 + SessionKey, // 0x02 + Passkey // 0x03 +} diff --git a/tee-worker/omni-executor/contracts/aa/src/v2/utils/Exec.sol b/tee-worker/omni-executor/contracts/aa/src/v2/utils/Exec.sol new file mode 100644 index 0000000000..911a7630fe --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/src/v2/utils/Exec.sol @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.28; + +// solhint-disable no-inline-assembly + +/** + * Utility functions helpful when making different kinds of contract calls in Solidity. + */ +library Exec { + function call(address to, uint256 value, bytes memory data, uint256 txGas) internal returns (bool success) { + assembly ("memory-safe") { + success := call(txGas, to, value, add(data, 0x20), mload(data), 0, 0) + } + } + + function staticcall(address to, bytes memory data, uint256 txGas) internal view returns (bool success) { + assembly ("memory-safe") { + success := staticcall(txGas, to, add(data, 0x20), mload(data), 0, 0) + } + } + + function delegateCall(address to, bytes memory data, uint256 txGas) internal returns (bool success) { + assembly ("memory-safe") { + success := delegatecall(txGas, to, add(data, 0x20), mload(data), 0, 0) + } + } + + // get returned data from last call or delegateCall + // maxLen - maximum length of data to return, or zero, for the full length + function getReturnData(uint256 maxLen) internal pure returns (bytes memory returnData) { + assembly ("memory-safe") { + let len := returndatasize() + if gt(maxLen, 0) { if gt(len, maxLen) { len := maxLen } } + let ptr := mload(0x40) + mstore(0x40, add(ptr, add(len, 0x20))) + mstore(ptr, len) + returndatacopy(add(ptr, 0x20), 0, len) + returnData := ptr + } + } + + // revert with explicit byte array (probably reverted info from call) + function revertWithData(bytes memory returnData) internal pure { + assembly ("memory-safe") { + revert(add(returnData, 32), mload(returnData)) + } + } + + // Propagate revert data from last call + function revertWithReturnData() internal pure { + revertWithData(getReturnData(0)); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/DeploymentHelper.t.sol b/tee-worker/omni-executor/contracts/aa/test/DeploymentHelper.t.sol index b895e095f6..ac594094c9 100644 --- a/tee-worker/omni-executor/contracts/aa/test/DeploymentHelper.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/DeploymentHelper.t.sol @@ -202,7 +202,8 @@ contract DeploymentHelperTest is Test { string memory testDir = getTestDir("fallback"); // Create an invalid JSON file first string memory filename = string(abi.encodePacked(testDir, "/local.json")); - try vm.createDir(testDir, true) {} catch {} // Allow creation to fail if directory exists + // Allow creation to fail if directory exists + try vm.createDir(testDir, true) {} catch {} vm.writeFile(filename, "{ invalid json structure without proper closing"); DeploymentHelper.ContractDeployment[] memory newDeployments = new DeploymentHelper.ContractDeployment[](1); diff --git a/tee-worker/omni-executor/contracts/aa/test/ERC20PaymasterV1.t.sol b/tee-worker/omni-executor/contracts/aa/test/ERC20PaymasterV1.t.sol index bba4f1925e..51fa5c7f1a 100644 --- a/tee-worker/omni-executor/contracts/aa/test/ERC20PaymasterV1.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/ERC20PaymasterV1.t.sol @@ -2,14 +2,14 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {ERC20PaymasterV1} from "../src/core/ERC20PaymasterV1.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; +import {ERC20PaymasterV1} from "../src/v1/core/ERC20PaymasterV1.sol"; +import {EntryPointV1} from "../src/v1/core/EntryPointV1.sol"; import {TestToken} from "../src/TestToken.sol"; -import {IPaymaster} from "../src/interfaces/IPaymaster.sol"; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; +import {IPaymaster} from "../src/v1/interfaces/IPaymaster.sol"; +import {PackedUserOperation} from "../src/v1/interfaces/PackedUserOperation.sol"; import {TestUtils} from "./TestUtils.sol"; import {Math} from "@openzeppelin/contracts/utils/math/Math.sol"; -import {BaseAccount} from "../src/core/BaseAccount.sol"; +import {BaseAccount} from "../src/v1/core/BaseAccount.sol"; import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; contract ERC20PaymasterV1Test is Test { diff --git a/tee-worker/omni-executor/contracts/aa/test/EntryPoint.t.sol b/tee-worker/omni-executor/contracts/aa/test/EntryPoint.t.sol index 0d44a1e3bb..122198c28a 100644 --- a/tee-worker/omni-executor/contracts/aa/test/EntryPoint.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/EntryPoint.t.sol @@ -2,11 +2,11 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../src/interfaces/UserOpSigner.sol"; -import {OmniAccountFactoryV1} from "../src/accounts/OmniAccountFactoryV1.sol"; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; -import {OwnerType} from "../src/interfaces/OwnerType.sol"; +import {EntryPointV1} from "../src/v1/core/EntryPointV1.sol"; +import {UserOpSigner} from "../src/v1/interfaces/UserOpSigner.sol"; +import {OmniAccountFactoryV1} from "../src/v1/accounts/OmniAccountFactoryV1.sol"; +import {PackedUserOperation} from "../src/v1/interfaces/PackedUserOperation.sol"; +import {OwnerType} from "../src/v1/interfaces/OwnerType.sol"; import {TestUtils} from "./TestUtils.sol"; contract EntryPointTest is Test { diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountFactory.t.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccountFactory.t.sol deleted file mode 100644 index 9b989885e9..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountFactory.t.sol +++ /dev/null @@ -1,34 +0,0 @@ -// SPDX-License-Identifier: GPL-3.0 -pragma solidity ^0.8.28; - -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {OmniAccountFactoryV1} from "../src/accounts/OmniAccountFactoryV1.sol"; -import {OmniAccountV1} from "../src/accounts/OmniAccountV1.sol"; -import {OwnerType} from "../src/interfaces/OwnerType.sol"; -import {Test, console} from "forge-std/Test.sol"; -import {TestUtils} from "./TestUtils.sol"; - -contract OmniAccountFactoryTest is Test { - EntryPointV1 public entryPoint; - OmniAccountFactoryV1 public omniAccountFactory; - address ownerAddress = 0x0000000000000000000000000000000000000000; - address rootAddress = 0x0000000000000000000000000000000000000001; - bytes clientId = bytes("test_client"); - bytes32 oa; - - function setUp() public { - entryPoint = new EntryPointV1(); - oa = TestUtils.prepare_evm_oa(ownerAddress, clientId); - omniAccountFactory = new OmniAccountFactoryV1(entryPoint); - } - - function test_CreateAccountReturnsSameAddress() public { - address senderCreator = address(entryPoint.senderCreator()); - vm.prank(senderCreator); - OmniAccountV1 account1 = omniAccountFactory.createAccount(oa, OwnerType.Evm, clientId, rootAddress); - vm.prank(senderCreator); - OmniAccountV1 account2 = omniAccountFactory.createAccount(oa, OwnerType.Evm, clientId, rootAddress); - - assertEq(address(account1), address(account2)); - } -} diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountUpgradability.t.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccountUpgradability.t.sol deleted file mode 100644 index b29cf5a829..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountUpgradability.t.sol +++ /dev/null @@ -1,363 +0,0 @@ -// SPDX-License-Identifier: GPL-3.0 -pragma solidity ^0.8.28; - -import {Test} from "forge-std/Test.sol"; -import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; -import "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol"; -import "@openzeppelin/contracts/interfaces/IERC1271.sol"; -import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; -import "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol"; -import {OmniAccountV1} from "../src/accounts/OmniAccountV1.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {OwnerType} from "../src/interfaces/OwnerType.sol"; -import {Passkey} from "../src/interfaces/Passkey.sol"; -import {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; -import {TestUtils} from "./TestUtils.sol"; - -contract OmniAccountV2 is OmniAccountV1, IERC1271 { - // New storage variable to test that new storage doesn't affect old storage - uint256 public newFeatureCounter; - - // ERC1271 magic value for valid signature - bytes4 private constant ERC1271_MAGIC_VALUE = 0x1626ba7e; - - constructor(EntryPointV1 anEntryPoint) OmniAccountV1(anEntryPoint) {} - - function version() public pure override returns (string memory) { - return "2.0.0"; - } - - // New function only available in V2 - function incrementNewFeature() public onlyOwner { - newFeatureCounter++; - } - - // New function to get feature counter - function getNewFeatureCounter() public view returns (uint256) { - return newFeatureCounter; - } - - /** - * @dev ERC1271 signature validation - * @param hash Hash of the data to be signed - * @param signature Signature byte array - * @return magicValue 0x1626ba7e if valid, 0xffffffff otherwise - */ - function isValidSignature(bytes32 hash, bytes memory signature) public view override returns (bytes4 magicValue) { - // Signature must be 65 bytes (r, s, v) - if (signature.length != 65) { - return 0xffffffff; - } - - // Recover the signer from the signature - address signer = ECDSA.recover(hash, signature); - - // Check if the signer is the owner or root signer - if (_determineOa(signer) == owner || isRootSigner(signer)) { - return ERC1271_MAGIC_VALUE; - } - - return 0xffffffff; - } -} - -contract OmniAccountUpgradeable is Test { - // Test addresses - address public owner = address(0x1234); - address public rootSigner = address(0x5678); - address public unauthorizedUser = address(0x9abc); - - // Test data - // bytes32 public ownerOa; - bytes clientId = bytes("test_client"); - - function setUp() public {} - - // Helper function to verify account version - function assertAccountVersion(address account) internal { - (bool success, bytes memory result) = account.call(abi.encodeWithSignature("version()")); - assertTrue(success, "Should be able to call version()"); - string memory version = abi.decode(result, (string)); - assertEq(version, "2.0.0", "Should be upgraded to version 2.0.0"); - } - - function testOaOwnerCanUpgradeOA() public { - (, EntryPointV1 entryPoint, OmniAccountV1 account) = OmniAccountTestUtils.setUp(owner, clientId, rootSigner); - OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); - - vm.prank(owner); - UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); - assertAccountVersion(address(account)); - } - - function testEntryPointCanUpgradeOA() public { - (, EntryPointV1 entryPoint, OmniAccountV1 account) = OmniAccountTestUtils.setUp(owner, clientId, rootSigner); - OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); - - vm.prank(address(entryPoint)); - UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); - assertAccountVersion(address(account)); - } - - function testUnauthorizedSenderCannotUpgradeOA() public { - (, EntryPointV1 entryPoint, OmniAccountV1 account) = - OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); - OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); - - vm.expectRevert("only owner"); - vm.prank(unauthorizedUser); - UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); - } - - function testRootSignerCannotUpgradeEvmOA() public { - (, EntryPointV1 entryPoint, OmniAccountV1 account) = - OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Evm); - OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); - - vm.expectRevert("only owner"); - vm.prank(rootSigner); - UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); - } - - function testRootSignerCannotUpgradeNonEvmOAWithPassKey() public { - (, EntryPointV1 entryPoint, OmniAccountV1 account) = - OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); - OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); - - Passkey.PublicKey memory pk = Passkey.PublicKey({x: 1, y: 2}); - vm.prank(owner); - account.addPasskeySigner(pk); - assertEq(account.passkeySignerCount(), 1); - - vm.expectRevert("only owner"); - vm.prank(rootSigner); - UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); - } - - function testRootSignerCanUpgradeNonEvmOAWithoutPassKey() public { - (, EntryPointV1 entryPoint, OmniAccountV1 account) = - OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); - OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); - - assertEq(account.passkeySignerCount(), 0); - - vm.prank(rootSigner); - UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); - assertAccountVersion(address(account)); - } - - function testUnauthorizedSenderCannotUpgradeNonEvmOA() public { - (, EntryPointV1 entryPoint, OmniAccountV1 account) = - OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); - OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); - - vm.expectRevert("only owner"); - vm.prank(unauthorizedUser); - UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); - } - - // Comprehensive test that verifies: - // 1. Address is preserved after upgrade - // 2. Account continues to be operable - // 3. New logic (version and new functions) works - // 4. Old logic/storage is not affected - function testUpgradePreservesStateAndOperability() public { - // Setup: Create account with rich state - (, EntryPointV1 entryPoint, OmniAccountV1 account) = - OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); - - // Add additional root signers to test state preservation - address additionalRootSigner = address(0xABCD); - vm.prank(owner); - account.addRootSigner(additionalRootSigner); - - // Add passkey signers to test state preservation - Passkey.PublicKey memory pk1 = Passkey.PublicKey({x: 12345, y: 67890}); - Passkey.PublicKey memory pk2 = Passkey.PublicKey({x: 11111, y: 22222}); - vm.prank(owner); - account.addPasskeySigner(pk1); - vm.prank(owner); - account.addPasskeySigner(pk2); - - // Add ETH balance to the account - vm.deal(address(account), 5 ether); - - // Add deposit to EntryPoint - vm.prank(address(account)); - account.addDeposit{value: 2 ether}(); - - // Record pre-upgrade state - address accountAddress = address(account); - bytes32 ownerBefore = account.owner(); - bytes memory clientIdBefore = account.clientId(); - OwnerType ownerTypeBefore = account.ownerType(); - uint256 passkeySignerCountBefore = account.passkeySignerCount(); - uint256 ethBalanceBefore = address(account).balance; - uint256 depositBefore = account.getDeposit(); - bool isRootSignerBefore = account.isRootSigner(rootSigner); - bool isAdditionalRootSignerBefore = account.isRootSigner(additionalRootSigner); - - // Verify pre-upgrade state - assertEq(ownerBefore, account.getOwner(), "Owner mismatch before upgrade"); - assertEq(passkeySignerCountBefore, 2, "Should have 2 passkey signers"); - assertTrue(isRootSignerBefore, "Root signer should exist"); - assertTrue(isAdditionalRootSignerBefore, "Additional root signer should exist"); - assertEq(ethBalanceBefore, 3 ether, "ETH balance should be 3 ether (5 - 2 deposited)"); - assertEq(depositBefore, 2 ether, "Deposit should be 2 ether"); - - // Perform upgrade - OmniAccountV2 accountV2Impl = new OmniAccountV2(entryPoint); - vm.prank(owner); - UUPSUpgradeable(account).upgradeToAndCall(address(accountV2Impl), ""); - - // Cast to V2 for accessing new functions - OmniAccountV2 accountV2 = OmniAccountV2(payable(address(account))); - - // TEST 1: Verify address is preserved - assertEq(address(accountV2), accountAddress, "Address should remain the same after upgrade"); - - // TEST 2: Verify new logic works (version) - assertAccountVersion(address(accountV2)); - - // TEST 3: Verify old storage is preserved - assertEq(accountV2.owner(), ownerBefore, "Owner should be preserved"); - assertEq(accountV2.clientId(), clientIdBefore, "Client ID should be preserved"); - assertTrue(accountV2.ownerType() == ownerTypeBefore, "Owner type should be preserved"); - assertEq(accountV2.passkeySignerCount(), passkeySignerCountBefore, "Passkey signer count should be preserved"); - assertTrue(accountV2.isRootSigner(rootSigner), "Root signer should be preserved"); - assertTrue(accountV2.isRootSigner(additionalRootSigner), "Additional root signer should be preserved"); - assertEq(address(accountV2).balance, ethBalanceBefore, "ETH balance should be preserved"); - assertEq(accountV2.getDeposit(), depositBefore, "Deposit should be preserved"); - - // TEST 4: Verify account is still operable - can add/remove signers - address newRootSigner = address(0xDEAD); - vm.prank(owner); - accountV2.addRootSigner(newRootSigner); - assertTrue(accountV2.isRootSigner(newRootSigner), "Should be able to add new root signer after upgrade"); - - vm.prank(owner); - accountV2.removeRootSigner(newRootSigner); - assertFalse(accountV2.isRootSigner(newRootSigner), "Should be able to remove root signer after upgrade"); - - // TEST 5: Verify can still add/remove passkey signers - Passkey.PublicKey memory pk3 = Passkey.PublicKey({x: 33333, y: 44444}); - vm.prank(owner); - accountV2.addPasskeySigner(pk3); - assertEq(accountV2.passkeySignerCount(), 3, "Should be able to add passkey signer after upgrade"); - - vm.prank(owner); - accountV2.removePasskeySigner(pk3); - assertEq(accountV2.passkeySignerCount(), 2, "Should be able to remove passkey signer after upgrade"); - - // TEST 6: Verify can withdraw deposit - address payable withdrawAddress = payable(address(0xBEEF)); - uint256 withdrawAmount = 0.5 ether; - vm.prank(owner); - accountV2.withdrawDepositTo(withdrawAddress, withdrawAmount); - assertEq(accountV2.getDeposit(), depositBefore - withdrawAmount, "Should be able to withdraw after upgrade"); - - // TEST 7: Verify new V2 functionality works - assertEq(accountV2.getNewFeatureCounter(), 0, "New feature counter should start at 0"); - vm.prank(owner); - accountV2.incrementNewFeature(); - assertEq(accountV2.getNewFeatureCounter(), 1, "New feature should work after upgrade"); - vm.prank(owner); - accountV2.incrementNewFeature(); - assertEq(accountV2.getNewFeatureCounter(), 2, "New feature should continue to work"); - - // TEST 8: Verify unauthorized users still cannot call owner-only functions - vm.expectRevert("only owner"); - vm.prank(unauthorizedUser); - accountV2.incrementNewFeature(); - } - - // Test ERC1271 support is added after upgrade and works correctly - function testERC1271WorksAfterUpgrade() public { - // Setup: Create private keys for owner and root signer - uint256 ownerPrivateKey = 0x1234; - uint256 rootSignerPrivateKey = 0x5678; - uint256 unauthorizedPrivateKey = 0x9abc; - - address ownerAddr = vm.addr(ownerPrivateKey); - address rootSignerAddr = vm.addr(rootSignerPrivateKey); - - // Create account with owner and root signer - (, EntryPointV1 entryPoint, OmniAccountV1 account) = - OmniAccountTestUtils.setUp(ownerAddr, clientId, rootSignerAddr); - - bytes32 messageHash = keccak256("Hello world!"); - - // Before upgrade: V1 doesn't have isValidSignature, attempting to call it should fail - (bool success,) = address(account) - .call(abi.encodeWithSignature("isValidSignature(bytes32,bytes)", messageHash, new bytes(65))); - assertFalse(success, "V1 should not have isValidSignature"); - - // Perform upgrade to V2 - OmniAccountV2 accountV2Impl = new OmniAccountV2(entryPoint); - vm.prank(ownerAddr); - UUPSUpgradeable(account).upgradeToAndCall(address(accountV2Impl), ""); - - // Cast to V2 for accessing ERC1271 - OmniAccountV2 accountV2 = OmniAccountV2(payable(address(account))); - - // Verify upgrade was successful - assertAccountVersion(address(accountV2)); - - // TEST 1: Valid signature from owner should be accepted - (uint8 v1, bytes32 r1, bytes32 s1) = vm.sign(ownerPrivateKey, messageHash); - bytes memory ownerSignature = abi.encodePacked(r1, s1, v1); - - bytes4 result1 = accountV2.isValidSignature(messageHash, ownerSignature); - assertEq(result1, bytes4(0x1626ba7e), "Owner signature should be valid"); - - // TEST 2: Valid signature from root signer should be accepted - (uint8 v2, bytes32 r2, bytes32 s2) = vm.sign(rootSignerPrivateKey, messageHash); - bytes memory rootSignature = abi.encodePacked(r2, s2, v2); - - bytes4 result2 = accountV2.isValidSignature(messageHash, rootSignature); - assertEq(result2, bytes4(0x1626ba7e), "Root signer signature should be valid"); - - // TEST 3: Invalid signature from unauthorized user should be rejected - (uint8 v3, bytes32 r3, bytes32 s3) = vm.sign(unauthorizedPrivateKey, messageHash); - bytes memory unauthorizedSignature = abi.encodePacked(r3, s3, v3); - - bytes4 result3 = accountV2.isValidSignature(messageHash, unauthorizedSignature); - assertEq(result3, bytes4(0xffffffff), "Unauthorized signature should be invalid"); - - // TEST 4: Invalid signature length should be rejected - bytes memory shortSignature = new bytes(32); - bytes4 result4 = accountV2.isValidSignature(messageHash, shortSignature); - assertEq(result4, bytes4(0xffffffff), "Short signature should be invalid"); - - // TEST 5: Test with different message hashes to ensure proper validation - bytes32 differentMessageHash = keccak256("Another message"); - (uint8 v5, bytes32 r5, bytes32 s5) = vm.sign(ownerPrivateKey, messageHash); - bytes memory signatureForOriginalMessage = abi.encodePacked(r5, s5, v5); - - // Using signature for original message with different hash should fail - bytes4 result5 = accountV2.isValidSignature(differentMessageHash, signatureForOriginalMessage); - assertEq(result5, bytes4(0xffffffff), "Signature should not validate for different message"); - - // TEST 6: Verify that added root signers can also validate signatures - address newRootSigner = vm.addr(0xDEADBEEF); - vm.prank(ownerAddr); - accountV2.addRootSigner(newRootSigner); - - (uint8 v6, bytes32 r6, bytes32 s6) = vm.sign(0xDEADBEEF, messageHash); - bytes memory newRootSignature = abi.encodePacked(r6, s6, v6); - - bytes4 result6 = accountV2.isValidSignature(messageHash, newRootSignature); - assertEq(result6, bytes4(0x1626ba7e), "Newly added root signer signature should be valid"); - - // TEST 7: Verify that removed root signers can no longer validate signatures - vm.prank(ownerAddr); - accountV2.removeRootSigner(rootSignerAddr); - - bytes4 result7 = accountV2.isValidSignature(messageHash, rootSignature); - assertEq(result7, bytes4(0xffffffff), "Removed root signer signature should be invalid"); - - // TEST 8: Verify ERC1271 interface support - // The contract should properly implement IERC1271 - assertTrue(address(accountV2).code.length > 0, "Account should have code (implementing ERC1271)"); - } -} diff --git a/tee-worker/omni-executor/contracts/aa/test/SimplePaymaster.t.sol b/tee-worker/omni-executor/contracts/aa/test/SimplePaymaster.t.sol index 239ad9ed95..a374b9b666 100644 --- a/tee-worker/omni-executor/contracts/aa/test/SimplePaymaster.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/SimplePaymaster.t.sol @@ -2,10 +2,10 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {SimplePaymaster} from "../src/core/SimplePaymaster.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {IPaymaster} from "../src/interfaces/IPaymaster.sol"; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; +import {SimplePaymaster} from "../src/v1/core/SimplePaymaster.sol"; +import {EntryPointV1} from "../src/v1/core/EntryPointV1.sol"; +import {IPaymaster} from "../src/v1/interfaces/IPaymaster.sol"; +import {PackedUserOperation} from "../src/v1/interfaces/PackedUserOperation.sol"; import {TestUtils} from "./TestUtils.sol"; contract SimplePaymasterTest is Test { diff --git a/tee-worker/omni-executor/contracts/aa/test/StakeManager.t.sol b/tee-worker/omni-executor/contracts/aa/test/StakeManager.t.sol index f98b922493..c53b73854d 100644 --- a/tee-worker/omni-executor/contracts/aa/test/StakeManager.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/StakeManager.t.sol @@ -2,8 +2,8 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {IStakeManager} from "../src/interfaces/IStakeManager.sol"; +import {EntryPointV1} from "../src/v1/core/EntryPointV1.sol"; +import {IStakeManager} from "../src/v1/interfaces/IStakeManager.sol"; contract StakeManagerTest is Test { EntryPointV1 public entryPoint; diff --git a/tee-worker/omni-executor/contracts/aa/test/TestUtils.sol b/tee-worker/omni-executor/contracts/aa/test/TestUtils.sol index df58e3e9c4..e6bb2a6559 100644 --- a/tee-worker/omni-executor/contracts/aa/test/TestUtils.sol +++ b/tee-worker/omni-executor/contracts/aa/test/TestUtils.sol @@ -1,7 +1,7 @@ // SPDX-License-Identifier: GPL-3.0 pragma solidity ^0.8.28; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; +import {PackedUserOperation} from "../src/v1/interfaces/PackedUserOperation.sol"; library TestUtils { function prepare_evm_oa(address account, bytes memory clientId) public pure returns (bytes32) { @@ -10,11 +10,7 @@ library TestUtils { return sha256(abi.encodePacked(clientId, oaType, account)); } - function preparePackedOp(address sender, bytes memory initCode) - internal - pure - returns (PackedUserOperation memory) - { + function preparePackedOp(address sender, bytes memory initCode) internal pure returns (PackedUserOperation memory) { uint256 nonce = 0; bytes memory callData = ""; bytes32 accountGasLimits = 0x0000000000000000000000000004e20000000000000000000000000000005b8d; @@ -23,8 +19,7 @@ library TestUtils { bytes memory paymasterAndData = ""; bytes memory signature = ""; - return ( - PackedUserOperation( + return (PackedUserOperation( sender, nonce, initCode, @@ -34,7 +29,6 @@ library TestUtils { gasFees, paymasterAndData, signature - ) - ); + )); } } diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccount.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccount.t.sol similarity index 84% rename from tee-worker/omni-executor/contracts/aa/test/OmniAccount.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccount.t.sol index 01c8791c8e..0fff674f89 100644 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccount.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccount.t.sol @@ -2,15 +2,15 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../src/accounts/OmniAccountV1.sol"; -import {BaseAccount} from "../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../src/interfaces/UserOpSigner.sol"; -import {Counter} from "../src/Counter.sol"; +import {OmniAccountV1} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {BaseAccount} from "../../src/v1/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v1/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v1/interfaces/UserOpSigner.sol"; +import {Counter} from "../../src/Counter.sol"; import {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; -import {TestUtils} from "./TestUtils.sol"; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; -import {SIG_VALIDATION_FAILED} from "../src//core/Helpers.sol"; +import {TestUtils} from "../TestUtils.sol"; +import {PackedUserOperation} from "../../src/v1/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; // add test cases for revert if called by non authorized address diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsEntryPoint.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsEntryPoint.t.sol similarity index 79% rename from tee-worker/omni-executor/contracts/aa/test/OmniAccountAsEntryPoint.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsEntryPoint.t.sol index 745d44c8f5..b7f01b0863 100644 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsEntryPoint.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsEntryPoint.t.sol @@ -2,10 +2,10 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../src/accounts/OmniAccountV1.sol"; -import {BaseAccount} from "../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {Counter} from "../src/Counter.sol"; +import {OmniAccountV1} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {BaseAccount} from "../../src/v1/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v1/core/EntryPointV1.sol"; +import {Counter} from "../../src/Counter.sol"; import {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; contract OmniAccountAsEntryPoint is Test { diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsOwner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsOwner.t.sol similarity index 97% rename from tee-worker/omni-executor/contracts/aa/test/OmniAccountAsOwner.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsOwner.t.sol index 2e33c7adea..6c202a1d93 100644 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsOwner.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsOwner.t.sol @@ -2,16 +2,16 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../src/accounts/OmniAccountV1.sol"; -import {BaseAccount} from "../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../src/interfaces/UserOpSigner.sol"; -import {Counter} from "../src/Counter.sol"; +import {OmniAccountV1} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {BaseAccount} from "../../src/v1/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v1/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v1/interfaces/UserOpSigner.sol"; +import {Counter} from "../../src/Counter.sol"; import {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; -import {TestUtils} from "./TestUtils.sol"; -import {SIG_VALIDATION_SUCCESS} from "../src//core/Helpers.sol"; -import {Passkey} from "../src/interfaces/Passkey.sol"; +import {PackedUserOperation} from "../../src/v1/interfaces/PackedUserOperation.sol"; +import {TestUtils} from "../TestUtils.sol"; +import {SIG_VALIDATION_SUCCESS} from "../../src/v1/core/Helpers.sol"; +import {Passkey} from "../../src/v1/interfaces/Passkey.sol"; contract OmniAccountAsOwner is Test { OmniAccountV1 public account; diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsPasskey.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsPasskey.t.sol similarity index 92% rename from tee-worker/omni-executor/contracts/aa/test/OmniAccountAsPasskey.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsPasskey.t.sol index 5efafefb6b..876be311d2 100644 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsPasskey.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsPasskey.t.sol @@ -2,16 +2,16 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../src/accounts/OmniAccountV1.sol"; -import {BaseAccount} from "../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../src/interfaces/UserOpSigner.sol"; -import {Counter} from "../src/Counter.sol"; +import {OmniAccountV1} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {BaseAccount} from "../../src/v1/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v1/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v1/interfaces/UserOpSigner.sol"; +import {Counter} from "../../src/Counter.sol"; import {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; -import {TestUtils} from "./TestUtils.sol"; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; -import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../src//core/Helpers.sol"; -import {Passkey} from "../src/interfaces/Passkey.sol"; +import {TestUtils} from "../TestUtils.sol"; +import {PackedUserOperation} from "../../src/v1/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; +import {Passkey} from "../../src/v1/interfaces/Passkey.sol"; import "@openzeppelin/contracts/utils/Base64.sol"; contract OmniAccountAsPasskey is Test { diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRoot.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRoot.t.sol similarity index 97% rename from tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRoot.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRoot.t.sol index c919312c12..4e4c29366c 100644 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRoot.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRoot.t.sol @@ -2,16 +2,16 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../src/accounts/OmniAccountV1.sol"; -import {BaseAccount} from "../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../src/interfaces/UserOpSigner.sol"; -import {Counter} from "../src/Counter.sol"; +import {OmniAccountV1} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {BaseAccount} from "../../src/v1/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v1/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v1/interfaces/UserOpSigner.sol"; +import {Counter} from "../../src/Counter.sol"; import {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; -import {TestUtils} from "./TestUtils.sol"; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; -import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../src//core/Helpers.sol"; -import {Passkey} from "../src/interfaces/Passkey.sol"; +import {TestUtils} from "../TestUtils.sol"; +import {PackedUserOperation} from "../../src/v1/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; +import {Passkey} from "../../src/v1/interfaces/Passkey.sol"; contract OmniAccountAsRoot is Test { OmniAccountV1 public account; diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRootNonEvm.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRootNonEvm.t.sol similarity index 95% rename from tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRootNonEvm.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRootNonEvm.t.sol index 76e9f09ed1..107fc93216 100644 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRootNonEvm.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRootNonEvm.t.sol @@ -2,17 +2,17 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1 as OmniAccount} from "../src/accounts/OmniAccountV1.sol"; -import {BaseAccount} from "../src/core/BaseAccount.sol"; -import {EntryPointV1 as EntryPoint} from "../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../src/interfaces/UserOpSigner.sol"; -import {OwnerType} from "../src/interfaces/OwnerType.sol"; -import {Counter} from "../src/Counter.sol"; +import {OmniAccountV1 as OmniAccount} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {BaseAccount} from "../../src/v1/core/BaseAccount.sol"; +import {EntryPointV1 as EntryPoint} from "../../src/v1/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v1/interfaces/UserOpSigner.sol"; +import {OwnerType} from "../../src/v1/interfaces/OwnerType.sol"; +import {Counter} from "../../src/Counter.sol"; import {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; -import {TestUtils} from "./TestUtils.sol"; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; -import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../src/core/Helpers.sol"; -import {Passkey} from "../src/interfaces/Passkey.sol"; +import {TestUtils} from "../TestUtils.sol"; +import {PackedUserOperation} from "../../src/v1/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; +import {Passkey} from "../../src/v1/interfaces/Passkey.sol"; contract OmniAccountAsRootNonEvm is Test { OmniAccount public account; diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol similarity index 94% rename from tee-worker/omni-executor/contracts/aa/test/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol index c2fd0a72d4..ef3adff56d 100644 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol @@ -2,17 +2,17 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1 as OmniAccount} from "../src/accounts/OmniAccountV1.sol"; -import {BaseAccount} from "../src/core/BaseAccount.sol"; -import {EntryPointV1 as EntryPoint} from "../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../src/interfaces/UserOpSigner.sol"; -import {OwnerType} from "../src/interfaces/OwnerType.sol"; -import {Passkey} from "../src/interfaces/Passkey.sol"; -import {Counter} from "../src/Counter.sol"; +import {OmniAccountV1 as OmniAccount} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {BaseAccount} from "../../src/v1/core/BaseAccount.sol"; +import {EntryPointV1 as EntryPoint} from "../../src/v1/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v1/interfaces/UserOpSigner.sol"; +import {OwnerType} from "../../src/v1/interfaces/OwnerType.sol"; +import {Passkey} from "../../src/v1/interfaces/Passkey.sol"; +import {Counter} from "../../src/Counter.sol"; import {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; -import {TestUtils} from "./TestUtils.sol"; -import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; -import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../src/core/Helpers.sol"; +import {TestUtils} from "../TestUtils.sol"; +import {PackedUserOperation} from "../../src/v1/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; contract OmniAccountNonEvmOwnerWithPasskeySigner is Test { OmniAccount public account; diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountTestUtils.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountTestUtils.sol similarity index 81% rename from tee-worker/omni-executor/contracts/aa/test/OmniAccountTestUtils.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountTestUtils.sol index e618d31487..e59ca2329d 100644 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountTestUtils.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountTestUtils.sol @@ -1,14 +1,14 @@ // SPDX-License-Identifier: GPL-3.0 pragma solidity ^0.8.28; -import {Counter} from "../src/Counter.sol"; +import {Counter} from "../../src/Counter.sol"; import {Vm} from "forge-std/Vm.sol"; -import {OmniAccountV1 as OmniAccount} from "../src/accounts/OmniAccountV1.sol"; -import {BaseAccount} from "../src/core/BaseAccount.sol"; -import {EntryPointV1 as EntryPoint} from "../src/core/EntryPointV1.sol"; -import {OwnerType} from "../src/interfaces/OwnerType.sol"; +import {OmniAccountV1 as OmniAccount} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {BaseAccount} from "../../src/v1/core/BaseAccount.sol"; +import {EntryPointV1 as EntryPoint} from "../../src/v1/core/EntryPointV1.sol"; +import {OwnerType} from "../../src/v1/interfaces/OwnerType.sol"; import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; -import {TestUtils} from "./TestUtils.sol"; +import {TestUtils} from "../TestUtils.sol"; library OmniAccountTestUtils { function setUpWithOwnerType(address ownerAddress, bytes memory clientId, address rootAddress, OwnerType ownerType) @@ -20,11 +20,9 @@ library OmniAccountTestUtils { OmniAccount accountImpl = new OmniAccount(entryPoint); bytes32 oa = TestUtils.prepare_evm_oa(ownerAddress, clientId); OmniAccount account = OmniAccount( - payable( - new ERC1967Proxy{salt: oa}( + payable(new ERC1967Proxy{salt: oa}( address(accountImpl), abi.encodeCall(OmniAccount.initialize, (oa, ownerType, clientId, rootAddress)) - ) - ) + )) ); return (counter, entryPoint, account); @@ -39,12 +37,10 @@ library OmniAccountTestUtils { OmniAccount accountImpl = new OmniAccount(entryPoint); bytes32 oa = TestUtils.prepare_evm_oa(ownerAddress, clientId); OmniAccount account = OmniAccount( - payable( - new ERC1967Proxy{salt: oa}( + payable(new ERC1967Proxy{salt: oa}( address(accountImpl), abi.encodeCall(OmniAccount.initialize, (oa, OwnerType.Evm, clientId, rootAddress)) - ) - ) + )) ); return (counter, entryPoint, account); diff --git a/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountUpgradability.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountUpgradability.t.sol new file mode 100644 index 0000000000..fc13820669 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountUpgradability.t.sol @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; +import "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol"; +import {OmniAccountV1} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {EntryPointV1} from "../../src/v1/core/EntryPointV1.sol"; +import {OwnerType} from "../../src/v1/interfaces/OwnerType.sol"; +import {Passkey} from "../../src/v1/interfaces/Passkey.sol"; +import {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; +import {TestUtils} from "../TestUtils.sol"; + +contract OmniAccountV2 is OmniAccountV1 { + constructor(EntryPointV1 anEntryPoint) OmniAccountV1(anEntryPoint) {} + + function version() public pure override returns (string memory) { + return "2.0.0"; + } +} + +contract OmniAccountUpgradeable is Test { + // Test addresses + address public owner = address(0x1234); + address public rootSigner = address(0x5678); + address public unauthorizedUser = address(0x9abc); + + // Test data + // bytes32 public ownerOa; + bytes clientId = bytes("test_client"); + + function setUp() public {} + + // Helper function to verify account version + function assertAccountVersion(address account) internal { + (bool success, bytes memory result) = account.call(abi.encodeWithSignature("version()")); + assertTrue(success, "Should be able to call version()"); + string memory version = abi.decode(result, (string)); + assertEq(version, "2.0.0", "Should be upgraded to version 2.0.0"); + } + + function testOaOwnerCanUpgradeOA() public { + (, EntryPointV1 entryPoint, OmniAccountV1 account) = OmniAccountTestUtils.setUp(owner, clientId, rootSigner); + OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); + + vm.prank(owner); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); + assertAccountVersion(address(account)); + } + + function testEntryPointCanUpgradeOA() public { + (, EntryPointV1 entryPoint, OmniAccountV1 account) = OmniAccountTestUtils.setUp(owner, clientId, rootSigner); + OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); + + vm.prank(address(entryPoint)); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); + assertAccountVersion(address(account)); + } + + function testUnauthorizedSenderCannotUpgradeOA() public { + (, EntryPointV1 entryPoint, OmniAccountV1 account) = + OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); + OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); + + vm.expectRevert("only owner"); + vm.prank(unauthorizedUser); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); + } + + function testRootSignerCannotUpgradeEvmOA() public { + (, EntryPointV1 entryPoint, OmniAccountV1 account) = + OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Evm); + OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); + + vm.expectRevert("only owner"); + vm.prank(rootSigner); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); + } + + function testRootSignerCannotUpgradeNonEvmOAWithPassKey() public { + (, EntryPointV1 entryPoint, OmniAccountV1 account) = + OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); + OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); + + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 1, y: 2}); + vm.prank(owner); + account.addPasskeySigner(pk); + assertEq(account.passkeySignerCount(), 1); + + vm.expectRevert("only owner"); + vm.prank(rootSigner); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); + } + + function testRootSignerCanUpgradeNonEvmOAWithoutPassKey() public { + (, EntryPointV1 entryPoint, OmniAccountV1 account) = + OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); + OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); + + assertEq(account.passkeySignerCount(), 0); + + vm.prank(rootSigner); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); + assertAccountVersion(address(account)); + } + + function testUnauthorizedSenderCannotUpgradeNonEvmOA() public { + (, EntryPointV1 entryPoint, OmniAccountV1 account) = + OmniAccountTestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); + OmniAccountV2 accountV2 = new OmniAccountV2(entryPoint); + + vm.expectRevert("only owner"); + vm.prank(unauthorizedUser); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV2), ""); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/MockModule.sol b/tee-worker/omni-executor/contracts/aa/test/v2/MockModule.sol new file mode 100644 index 0000000000..e8c42e8ba6 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/MockModule.sol @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +/** + * Mock module for testing module registration and execution functionality + */ +contract MockModule { + uint256 public counter; + address public lastCaller; + bytes public lastCallData; + + event ModuleFunctionCalled(address caller, uint256 value); + event ModuleCounterIncremented(uint256 newValue); + + function increment() external { + counter++; + lastCaller = msg.sender; + emit ModuleCounterIncremented(counter); + } + + function incrementByValue(uint256 value) external { + counter += value; + lastCaller = msg.sender; + emit ModuleCounterIncremented(counter); + } + + function setCounter(uint256 newValue) external { + counter = newValue; + lastCaller = msg.sender; + } + + function getCounter() external view returns (uint256) { + return counter; + } + + function functionWithReturn() external pure returns (uint256) { + return 42; + } + + function functionWithMultipleReturns() external pure returns (uint256, address, bool) { + return (123, address(0x1234), true); + } + + function functionThatReverts() external pure { + revert("Module function reverted"); + } + + function recordCall(bytes calldata data) external { + lastCaller = msg.sender; + lastCallData = data; + emit ModuleFunctionCalled(msg.sender, 0); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV1ToV2Upgrade.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV1ToV2Upgrade.t.sol new file mode 100644 index 0000000000..87222193f2 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV1ToV2Upgrade.t.sol @@ -0,0 +1,327 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; +import "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol"; +import {OmniAccountV1} from "../../src/v1/accounts/OmniAccountV1.sol"; +import {OmniAccountV2} from "../../src/v2/accounts/OmniAccountV2.sol"; +// we need to use v1 types here .... +import {EntryPointV1 as EntryPointV1} from "../../src/v1/core/EntryPointV1.sol"; +import {IEntryPoint as IEntryPointV1} from "../../src/v1/interfaces/IEntryPoint.sol"; +import {OwnerType as OwnerTypeV1} from "../../src/v1/interfaces/OwnerType.sol"; +import {Passkey as PasskeyV1} from "../../src/v1/interfaces/Passkey.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; +import {StorageTestModule} from "./StorageTestModule.sol"; +import {EntryPointV1 as EntryPointV2} from "../../src/v2/core/EntryPointV1.sol"; +import {OwnerType} from "../../src/v2/interfaces/OwnerType.sol"; +import {Passkey} from "../../src/v2/interfaces/Passkey.sol"; + +/** + * Comprehensive test for upgrading OmniAccountV1 to V2 with full state verification. + * Tests: + * 1. V1 account setup with all fields populated + * 2. Upgrade to V2 + * 3. Verify all V1 state is preserved + * 4. Register and execute modules + * 5. Verify module storage operations don't corrupt account storage + * 6. Verify V1 functionality still works after upgrade + */ +contract OmniAccountV1ToV2Upgrade is Test { + EntryPointV1 public entryPoint; + OmniAccountV1 public accountV1; + OmniAccountV2 public accountV2Implementation; + StorageTestModule public module; + + // Test accounts + address public owner = address(0x1111); + address public rootSigner1 = address(0x2222); + address public rootSigner2 = address(0x3333); + address public rootSigner3 = address(0x4444); + bytes public clientId = bytes("comprehensive_test_client_v1_to_v2"); + + // Passkey signers v1 + PasskeyV1.PublicKey public passkey1V1 = PasskeyV1.PublicKey({x: 111111, y: 222222}); + PasskeyV1.PublicKey public passkey2V1 = PasskeyV1.PublicKey({x: 333333, y: 444444}); + + Passkey.PublicKey public passkey1 = Passkey.PublicKey({x: 111111, y: 222222}); + Passkey.PublicKey public passkey2 = Passkey.PublicKey({x: 333333, y: 444444}); + + // Storage state to verify + bytes32 public expectedOwner; + uint256 public depositAmount = 5 ether; + + function setUp() public { + entryPoint = new EntryPointV1(); + + // Deploy V1 implementation and create proxy + OmniAccountV1 v1Implementation = new OmniAccountV1(IEntryPointV1(address(entryPoint))); + expectedOwner = TestUtils.prepare_evm_oa(owner, clientId); + + accountV1 = OmniAccountV1( + payable(new ERC1967Proxy{salt: expectedOwner}( + address(v1Implementation), + abi.encodeCall(OmniAccountV1.initialize, (expectedOwner, OwnerTypeV1.Evm, clientId, rootSigner1)) + )) + ); + + // Prepare V2 implementation for upgrade + // EntryPointV1 and EntryPointV2 are identical so we use old deployment with V2 type + accountV2Implementation = new OmniAccountV2(EntryPointV2(payable(address(entryPoint)))); + + // Prepare module + module = new StorageTestModule(); + } + + function test_ComprehensiveV1ToV2Upgrade() public { + // ============ STEP 1: Setup V1 Account with Full State ============ + + // Add multiple root signers + vm.prank(owner); + accountV1.addRootSigner(rootSigner2); + vm.prank(owner); + accountV1.addRootSigner(rootSigner3); + + // Add passkey signers + vm.prank(owner); + accountV1.addPasskeySigner(passkey1V1); + vm.prank(owner); + accountV1.addPasskeySigner(passkey2V1); + + // Add deposit to EntryPoint + vm.deal(address(accountV1), 10 ether); + vm.prank(address(accountV1)); + accountV1.addDeposit{value: depositAmount}(); + + // Verify V1 initial state + assertEq(accountV1.owner(), expectedOwner, "V1 owner mismatch"); + assertEq(accountV1.clientId(), clientId, "V1 clientId mismatch"); + assertEq(uint256(accountV1.ownerType()), uint256(OwnerTypeV1.Evm), "V1 ownerType mismatch"); + assertTrue(accountV1.isRootSigner(rootSigner1), "rootSigner1 not registered"); + assertTrue(accountV1.isRootSigner(rootSigner2), "rootSigner2 not registered"); + assertTrue(accountV1.isRootSigner(rootSigner3), "rootSigner3 not registered"); + assertTrue(accountV1.passkeySigners(PasskeyV1.toKey(passkey1V1)), "passkey1 not registered"); + assertTrue(accountV1.passkeySigners(PasskeyV1.toKey(passkey2V1)), "passkey2 not registered"); + assertEq(accountV1.passkeySignerCount(), 2, "passkey count mismatch"); + assertEq(accountV1.getDeposit(), depositAmount, "deposit mismatch"); + assertEq(accountV1.version(), "1.0.0", "V1 version mismatch"); + + // ============ STEP 2: Upgrade to V2 ============ + + vm.prank(owner); + UUPSUpgradeable(address(accountV1)).upgradeToAndCall(address(accountV2Implementation), ""); + + // Cast to V2 interface + OmniAccountV2 accountV2 = OmniAccountV2(payable(address(accountV1))); + + // ============ STEP 3: Verify All V1 State is Preserved ============ + + assertEq(accountV2.owner(), expectedOwner, "V2 owner mismatch after upgrade"); + assertEq(accountV2.clientId(), clientId, "V2 clientId mismatch after upgrade"); + assertEq(uint256(accountV2.ownerType()), uint256(OwnerType.Evm), "V2 ownerType mismatch after upgrade"); + assertTrue(accountV2.isRootSigner(rootSigner1), "rootSigner1 lost after upgrade"); + assertTrue(accountV2.isRootSigner(rootSigner2), "rootSigner2 lost after upgrade"); + assertTrue(accountV2.isRootSigner(rootSigner3), "rootSigner3 lost after upgrade"); + assertTrue(accountV2.passkeySigners(Passkey.toKey(passkey1)), "passkey1 lost after upgrade"); + assertTrue(accountV2.passkeySigners(Passkey.toKey(passkey2)), "passkey2 lost after upgrade"); + assertEq(accountV2.passkeySignerCount(), 2, "passkey count changed after upgrade"); + assertEq(accountV2.getDeposit(), depositAmount, "deposit changed after upgrade"); + assertEq(accountV2.version(), "2.0.0", "V2 version mismatch after upgrade"); + + // ============ STEP 4: Test New V2 Module Functionality ============ + + // Register module + vm.prank(owner); + accountV2.registerModule(address(module)); + assertTrue(accountV2.isModuleRegistered(address(module)), "Module not registered"); + + // Execute module functions that write to various storage slots + vm.prank(address(entryPoint)); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("setSimpleValue(uint256)", 12345)); + + vm.prank(address(entryPoint)); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("pushToArray(uint256)", 111)); + + vm.prank(address(entryPoint)); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("pushToArray(uint256)", 222)); + + vm.prank(address(entryPoint)); + accountV2.executeModuleCall( + address(module), abi.encodeWithSignature("setMapping(address,uint256)", address(0x5555), 9999) + ); + + vm.prank(address(entryPoint)); + accountV2.executeModuleCall( + address(module), + abi.encodeWithSignature("setNestedMapping(address,uint256,uint256)", address(0x6666), 42, 7777) + ); + + vm.prank(address(entryPoint)); + accountV2.executeModuleCall( + address(module), abi.encodeWithSignature("setStruct(uint256,string,uint256)", 1, "test_struct", 8888) + ); + + vm.prank(address(entryPoint)); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("setString(string)", "hello_from_module")); + + vm.prank(address(entryPoint)); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("setBytes(bytes)", hex"deadbeef")); + + // Complex operation with multiple storage writes + vm.prank(address(entryPoint)); + bytes memory complexResult = accountV2.executeModuleCall( + address(module), + abi.encodeWithSignature("complexOperation(uint256,uint256,address)", 100, 200, address(0x7777)) + ); + uint256 complexReturnValue = abi.decode(complexResult, (uint256)); + assertEq(complexReturnValue, 300, "Complex operation return value mismatch"); + + // ============ STEP 5: Verify Account State is NOT Corrupted ============ + + // All V1 state should remain intact + assertEq(accountV2.owner(), expectedOwner, "Owner corrupted after module execution"); + assertEq(accountV2.clientId(), clientId, "ClientId corrupted after module execution"); + assertEq(uint256(accountV2.ownerType()), uint256(OwnerTypeV1.Evm), "OwnerType corrupted after module execution"); + assertTrue(accountV2.isRootSigner(rootSigner1), "rootSigner1 corrupted"); + assertTrue(accountV2.isRootSigner(rootSigner2), "rootSigner2 corrupted"); + assertTrue(accountV2.isRootSigner(rootSigner3), "rootSigner3 corrupted"); + assertTrue(accountV2.passkeySigners(PasskeyV1.toKey(passkey1V1)), "passkey1 corrupted"); + assertTrue(accountV2.passkeySigners(PasskeyV1.toKey(passkey2V1)), "passkey2 corrupted"); + assertEq(accountV2.passkeySignerCount(), 2, "Passkey count corrupted"); + assertEq(accountV2.getDeposit(), depositAmount, "Deposit corrupted"); + + // Module registration state should be intact + assertTrue(accountV2.isModuleRegistered(address(module)), "Module registration lost"); + + // ============ STEP 6: Verify V1 Functionality Still Works ============ + + // Add new root signer + address newRootSigner = address(0x8888); + vm.prank(owner); + accountV2.addRootSigner(newRootSigner); + assertTrue(accountV2.isRootSigner(newRootSigner), "Cannot add new root signer after upgrade"); + + // Remove a root signer + vm.prank(owner); + accountV2.removeRootSigner(rootSigner3); + assertFalse(accountV2.isRootSigner(rootSigner3), "Cannot remove root signer after upgrade"); + assertTrue(accountV2.isRootSigner(rootSigner1), "Other root signers affected"); + assertTrue(accountV2.isRootSigner(rootSigner2), "Other root signers affected"); + + // Add new passkey signer + Passkey.PublicKey memory newPasskey = Passkey.PublicKey({x: 555555, y: 666666}); + vm.prank(owner); + accountV2.addPasskeySigner(newPasskey); + assertTrue(accountV2.passkeySigners(Passkey.toKey(newPasskey)), "Cannot add new passkey after upgrade"); + assertEq(accountV2.passkeySignerCount(), 3, "Passkey count not updated"); + + // Remove a passkey signer + vm.prank(owner); + Passkey.PublicKey memory passkey1V2 = Passkey.PublicKey({x: passkey1.x, y: passkey1.y}); + accountV2.removePasskeySigner(passkey1V2); + assertFalse(accountV2.passkeySigners(Passkey.toKey(passkey1)), "Cannot remove passkey after upgrade"); + assertEq(accountV2.passkeySignerCount(), 2, "Passkey count not updated after removal"); + + // Test deposit operations + vm.deal(address(accountV2), 10 ether); + vm.prank(address(accountV2)); + accountV2.addDeposit{value: 2 ether}(); + assertEq(accountV2.getDeposit(), depositAmount + 2 ether, "Cannot add deposit after upgrade"); + + // ============ STEP 7: Execute More Module Operations ============ + + // Increment simple value multiple times + for (uint256 i = 0; i < 5; i++) { + vm.prank(address(entryPoint)); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("incrementSimpleValue()")); + } + + // Push more values to array + for (uint256 i = 0; i < 3; i++) { + vm.prank(address(entryPoint)); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("pushToArray(uint256)", 1000 + i)); + } + + // ============ STEP 8: Final State Verification ============ + + // Verify all account state is still intact + assertEq(accountV2.owner(), expectedOwner, "Final: Owner corrupted"); + assertEq(accountV2.clientId(), clientId, "Final: ClientId corrupted"); + assertTrue(accountV2.isRootSigner(rootSigner1), "Final: rootSigner1 corrupted"); + assertTrue(accountV2.isRootSigner(rootSigner2), "Final: rootSigner2 corrupted"); + assertTrue(accountV2.isRootSigner(newRootSigner), "Final: newRootSigner lost"); + assertFalse(accountV2.isRootSigner(rootSigner3), "Final: removed root signer reappeared"); + assertTrue(accountV2.passkeySigners(Passkey.toKey(passkey2)), "Final: passkey2 corrupted"); + assertTrue(accountV2.passkeySigners(Passkey.toKey(newPasskey)), "Final: newPasskey lost"); + assertFalse(accountV2.passkeySigners(Passkey.toKey(passkey1)), "Final: removed passkey reappeared"); + assertEq(accountV2.passkeySignerCount(), 2, "Final: passkey count incorrect"); + assertTrue(accountV2.isModuleRegistered(address(module)), "Final: module registration lost"); + + // ============ STEP 9: Test Module Unregistration ============ + + vm.prank(owner); + accountV2.unregisterModule(address(module)); + assertFalse(accountV2.isModuleRegistered(address(module)), "Module still registered after unregistration"); + + // Verify account state is still intact after unregistration + assertEq(accountV2.owner(), expectedOwner, "Owner corrupted after module unregistration"); + assertEq(accountV2.passkeySignerCount(), 2, "Passkey count changed after module unregistration"); + + // ============ STEP 10: Verify Module Cannot Be Executed After Unregistration ============ + + vm.expectRevert("Module not registered"); + vm.prank(address(entryPoint)); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("setSimpleValue(uint256)", 99999)); + } + + function test_V1ToV2UpgradeNonEvmAccount() public { + // Test upgrade with non-EVM account type + OmniAccountV1 v1Implementation = new OmniAccountV1(IEntryPointV1(address(entryPoint))); + bytes32 substrateOwner = keccak256("substrate_owner"); + + OmniAccountV1 substrateAccount = OmniAccountV1( + payable(new ERC1967Proxy{salt: substrateOwner}( + address(v1Implementation), + abi.encodeCall( + OmniAccountV1.initialize, (substrateOwner, OwnerTypeV1.Substrate, clientId, rootSigner1) + ) + )) + ); + + // Setup state + vm.prank(rootSigner1); + substrateAccount.addRootSigner(rootSigner2); + + // Verify V1 state + assertEq(substrateAccount.owner(), substrateOwner); + assertEq(uint256(substrateAccount.ownerType()), uint256(OwnerTypeV1.Substrate)); + assertTrue(substrateAccount.isRootSigner(rootSigner1)); + assertTrue(substrateAccount.isRootSigner(rootSigner2)); + + // Upgrade + vm.prank(rootSigner1); + UUPSUpgradeable(address(substrateAccount)).upgradeToAndCall(address(accountV2Implementation), ""); + + OmniAccountV2 substrateAccountV2 = OmniAccountV2(payable(address(substrateAccount))); + + // Verify state preserved + assertEq(substrateAccountV2.owner(), substrateOwner, "Substrate owner lost"); + assertEq(uint256(substrateAccountV2.ownerType()), uint256(OwnerType.Substrate), "OwnerType changed"); + assertTrue(substrateAccountV2.isRootSigner(rootSigner1), "rootSigner1 lost"); + assertTrue(substrateAccountV2.isRootSigner(rootSigner2), "rootSigner2 lost"); + assertEq(substrateAccountV2.version(), "2.0.0", "Version not updated"); + + // Test module functionality + vm.prank(rootSigner1); + substrateAccountV2.registerModule(address(module)); + assertTrue(substrateAccountV2.isModuleRegistered(address(module))); + + // Execute module + vm.prank(address(entryPoint)); + substrateAccountV2.executeModuleCall(address(module), abi.encodeWithSignature("setSimpleValue(uint256)", 777)); + + // Verify account state intact + assertEq(substrateAccountV2.owner(), substrateOwner, "Substrate owner corrupted"); + assertTrue(substrateAccountV2.isRootSigner(rootSigner1), "rootSigner1 corrupted"); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol new file mode 100644 index 0000000000..16853dec3e --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test, console} from "forge-std/Test.sol"; +import {OmniAccountV2} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {BaseAccount} from "../../src/v2/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v2/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v2/interfaces/UserOpSigner.sol"; +import {Counter} from "../../src/Counter.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; +import {PackedUserOperation} from "../../src/v2/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; + +// add test cases for revert if called by non authorized address + +contract OmniAccountV2Test is Test { + OmniAccountV2 public account; + EntryPointV1 public entryPoint; + Counter public counter; + + address ownerAddress = 0x0000000000000000000000000000000000000000; + address rootAddress = 0x0000000000000000000000000000000000000001; + bytes clientId = bytes("test_client"); + + function setUp() public { + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, rootAddress); + } + + function test_Owner() public view { + bytes32 expectedOwner = 0xe502d639feeb199ae332376b050307d76d11b32e93b9fd1310127d2af64923fe; + assertEq(account.owner(), expectedOwner); + } + + function test_Root() public view { + address expectedRoot = 0x0000000000000000000000000000000000000001; + assert(account.isRootSigner(expectedRoot)); + } + + function test_EntryPoint() public view { + assertEq(address(account.entryPoint()), address(entryPoint)); + } + + function test_Execute_As_Not_Allowed() public { + vm.expectRevert("account: not from EntryPoint"); + account.execute(address(counter), 0, abi.encodeWithSignature("increment()")); + } + + function test_ExecuteBatch_As_Not_Allowed() public { + vm.expectRevert("account: not from EntryPoint"); + BaseAccount.Call[] memory calls = new BaseAccount.Call[](0); + account.executeBatch(calls); + } + + function test_AddRootSigner_As_Not_Allowed() public { + vm.expectRevert("only owner"); + account.addRootSigner(0x0000000000000000000000000000000000000000); + } + + function test_RemoveRootSigner_As_Not_Allowed() public { + vm.expectRevert("only owner"); + account.removeRootSigner(0x0000000000000000000000000000000000000000); + } + + function test_ValidateOp() public { + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, rootAddress); + (, uint256 bobPk) = makeAddrAndKey("bob"); + + address sender = 0x0eAfeE130Ab1F6261885eE7080f9e8B2513111d4; + + bytes memory initCode = ""; + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + // sign userOp + (uint8 v, bytes32 r, bytes32 s) = vm.sign(bobPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsEntryPoint.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsEntryPoint.t.sol new file mode 100644 index 0000000000..3a33e9271a --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsEntryPoint.t.sol @@ -0,0 +1,129 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test, console} from "forge-std/Test.sol"; +import {OmniAccountV2} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {BaseAccount} from "../../src/v2/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v2/core/EntryPointV1.sol"; +import {Counter} from "../../src/Counter.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {MockModule} from "./MockModule.sol"; + +contract OmniAccountV2AsEntryPoint is Test { + OmniAccountV2 public account; + EntryPointV1 public entryPoint; + Counter public counter; + MockModule public module; + + address ownerAddress = 0x0000000000000000000000000000000000000000; + address rootAddress = 0x0000000000000000000000000000000000000001; + bytes clientId = bytes("test_client"); + + function setUp() public { + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, rootAddress); + module = new MockModule(); + } + + function test_Execute() public { + OmniAccountV2TestUtils.performExecuteTestAs(vm, address(entryPoint), account, counter); + } + + function test_ExecuteBatch() public { + OmniAccountV2TestUtils.performExecuteBatchTestAs(vm, address(entryPoint), account, counter); + } + + // ============ Module Management Tests ============ + + function test_EntryPointCanRegisterModule() public { + vm.prank(address(entryPoint)); + account.registerModule(address(module)); + + assertTrue(account.isModuleRegistered(address(module))); + } + + function test_EntryPointCanUnregisterModule() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + vm.prank(address(entryPoint)); + account.unregisterModule(address(module)); + + assertFalse(account.isModuleRegistered(address(module))); + } + + function test_OnlyEntryPointCanExecuteModule() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + bytes memory callData = abi.encodeWithSignature("increment()"); + + vm.expectRevert("account: not from EntryPoint"); + vm.prank(ownerAddress); + account.executeModuleCall(address(module), callData); + } + + function test_CanExecuteRegisteredModule() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + bytes memory callData = abi.encodeWithSignature("increment()"); + + vm.prank(address(entryPoint)); + account.executeModuleCall(address(module), callData); + + // Module executed successfully (delegatecall modifies account's storage, not module's) + // We verify success by checking that no revert occurred + } + + function test_ModuleExecutionWithReturnValue() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + bytes memory callData = abi.encodeWithSignature("functionWithReturn()"); + + vm.prank(address(entryPoint)); + bytes memory returnData = account.executeModuleCall(address(module), callData); + + uint256 result = abi.decode(returnData, (uint256)); + assertEq(result, 42); + } + + function test_ModuleExecutionWithMultipleReturns() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + bytes memory callData = abi.encodeWithSignature("functionWithMultipleReturns()"); + + vm.prank(address(entryPoint)); + bytes memory returnData = account.executeModuleCall(address(module), callData); + + (uint256 num, address addr, bool flag) = abi.decode(returnData, (uint256, address, bool)); + assertEq(num, 123); + assertEq(addr, address(0x1234)); + assertTrue(flag); + } + + function test_ModuleExecutionRevertsCorrectly() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + bytes memory callData = abi.encodeWithSignature("functionThatReverts()"); + + vm.expectRevert("Module function reverted"); + vm.prank(address(entryPoint)); + account.executeModuleCall(address(module), callData); + } + + function test_ModuleExecutionWithParameters() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + bytes memory callData = abi.encodeWithSignature("incrementByValue(uint256)", 5); + + vm.prank(address(entryPoint)); + account.executeModuleCall(address(module), callData); + + // Module executed successfully with parameters (delegatecall modifies account's storage, not module's) + // We verify success by checking that no revert occurred + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsOwner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsOwner.t.sol new file mode 100644 index 0000000000..0d39a0dca1 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsOwner.t.sol @@ -0,0 +1,561 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {OmniAccountV2} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {BaseAccount} from "../../src/v2/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v2/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v2/interfaces/UserOpSigner.sol"; +import {Counter} from "../../src/Counter.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {PackedUserOperation} from "../../src/v2/interfaces/PackedUserOperation.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; +import {SIG_VALIDATION_SUCCESS} from "../../src/v1/core/Helpers.sol"; +import {Passkey} from "../../src/v2/interfaces/Passkey.sol"; +import {MockModule} from "./MockModule.sol"; + +contract OmniAccountV2AsOwner is Test { + OmniAccountV2 public account; + EntryPointV1 public entryPoint; + Counter public counter; + MockModule public module; + + address ownerAddress = 0x0000000000000000000000000000000000000000; + address rootAddress = 0x0000000000000000000000000000000000000001; + bytes clientId = bytes("test_client"); + + function setUp() public { + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, rootAddress); + module = new MockModule(); + } + + function test_Execute_As_Not_Allowed() public { + vm.expectRevert("account: not from EntryPoint"); + account.execute(address(counter), 0, abi.encodeWithSignature("increment()")); + } + + function test_ExecuteBatch_As_Not_Allowed() public { + vm.expectRevert("account: not from EntryPoint"); + BaseAccount.Call[] memory calls = new BaseAccount.Call[](0); + account.executeBatch(calls); + } + + function test_AddRemoveRootSigner() public { + address root = 0x0000000000000000000000000000000000000002; + vm.prank(ownerAddress); + account.addRootSigner(root); + assert(account.rootSigners(root)); + vm.prank(ownerAddress); + account.removeRootSigner(root); + assert(!account.rootSigners(root)); + } + + function test_ValidateOp() public { + (address alice, uint256 alicePk) = makeAddrAndKey("alice"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(alice, clientId, rootAddress); + + address sender = 0x0eAfeE130Ab1F6261885eE7080f9e8B2513111d4; + + bytes memory initCode = ""; + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + // sign userOp + (uint8 v, bytes32 r, bytes32 s) = vm.sign(alicePk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + } + + function test_OwnerCanCallAddRootSignerViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + address newRoot = 0x0000000000000000000000000000000000000002; + + // Prepare UserOp that calls execute() with addRootSigner as inner call + address sender = address(account); + bytes memory initCode = ""; + bytes memory innerCallData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + bytes memory callData = + abi.encodeWithSignature("execute(address,uint256,bytes)", address(account), 0, innerCallData); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed because owner can call restricted functions + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute the UserOp - should fail because execute() cannot call restricted functions + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertFalse(success); // Changed to expect failure + + // Verify the root signer was NOT added + assertFalse(account.rootSigners(newRoot)); // Changed to expect it wasn't added + } + + function test_OwnerCanCallRemoveRootSignerViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + // First add a root signer directly + vm.prank(owner); + account.addRootSigner(rootAddress); + assertTrue(account.rootSigners(rootAddress)); + + // Prepare UserOp that calls execute() with removeRootSigner as inner call + address sender = address(account); + bytes memory initCode = ""; + bytes memory innerCallData = abi.encodeWithSignature("removeRootSigner(address)", rootAddress); + bytes memory callData = + abi.encodeWithSignature("execute(address,uint256,bytes)", address(account), 0, innerCallData); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute the UserOp - should fail because execute() cannot call restricted functions + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertFalse(success); // Changed to expect failure + + // Verify the root signer was NOT removed + assertTrue(account.rootSigners(rootAddress)); // Changed to expect it wasn't removed + } + + function test_OwnerCanCallWithdrawDepositViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + // Add some deposit first + vm.deal(address(account), 1 ether); + vm.prank(address(account)); + account.addDeposit{value: 0.5 ether}(); + + address payable withdrawTo = payable(0x0000000000000000000000000000000000000003); + uint256 withdrawAmount = 0.1 ether; + + // Prepare UserOp that calls withdrawDepositTo + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = + abi.encodeWithSignature("withdrawDepositTo(address,uint256)", withdrawTo, withdrawAmount); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + } + + function test_OwnerCanExecuteBatchViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + // Prepare UserOp that calls executeBatch + address sender = address(account); + bytes memory initCode = ""; + + BaseAccount.Call[] memory calls = new BaseAccount.Call[](2); + calls[0] = BaseAccount.Call({target: address(counter), value: 0, data: abi.encodeWithSignature("increment()")}); + calls[1] = BaseAccount.Call({target: address(counter), value: 0, data: abi.encodeWithSignature("increment()")}); + bytes memory callData = abi.encodeWithSignature("executeBatch((address,uint256,bytes)[])", calls); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed because owner can use executeBatch + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute the UserOp to verify it works + uint256 initialCount = counter.number(); + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + // Verify the counter was incremented twice + assertEq(counter.number(), initialCount + 2); + } + + function test_OwnerCanExecuteWithRestrictedInnerCallViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + address newRoot = 0x0000000000000000000000000000000000000004; + + // Prepare UserOp that calls execute() with addRootSigner as inner call + address sender = address(account); + bytes memory initCode = ""; + bytes memory innerCallData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + bytes memory callData = + abi.encodeWithSignature("execute(address,uint256,bytes)", address(account), 0, innerCallData); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed because owner signed the UserOp + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute the UserOp - should fail because execute() cannot call restricted functions + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertFalse(success); // Changed to expect failure + + // Verify the root signer was NOT added + assertFalse(account.rootSigners(newRoot)); // Changed to expect it wasn't added + } + + function test_OwnerCanCallAddRootSignerDirectlyViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + address newRoot = 0x0000000000000000000000000000000000000005; + + // Prepare UserOp that calls addRootSigner directly (without execute wrapper) + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed for owner + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Now execution will succeed because EntryPoint is allowed in _onlyOwner + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + // Verify the root signer was actually added + assertTrue(account.rootSigners(newRoot)); + } + + function test_OwnerCanCallWithdrawDepositDirectlyViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + // Add some deposit first + vm.deal(address(account), 1 ether); + vm.prank(address(account)); + account.addDeposit{value: 0.5 ether}(); + + address payable withdrawTo = payable(0x0000000000000000000000000000000000000006); + uint256 withdrawAmount = 0.1 ether; + + // Prepare UserOp that calls withdrawDepositTo directly + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = + abi.encodeWithSignature("withdrawDepositTo(address,uint256)", withdrawTo, withdrawAmount); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed for owner + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Now execution will succeed because EntryPoint is allowed in _onlyOwner + uint256 balanceBefore = withdrawTo.balance; + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + // Verify the withdrawal actually happened + assertEq(withdrawTo.balance, balanceBefore + withdrawAmount); + } + + function test_OwnerCanCallAddPasskeySignerViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + + // Prepare UserOp that calls addPasskeySigner directly + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSelector(account.addPasskeySigner.selector, pk); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed for owner + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute and verify + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + // Verify the passkey signer was actually added + bytes32 key = Passkey.toKey(pk); + assertTrue(account.passkeySigners(key)); + } + + function test_OwnerCanCallRemovePasskeySignerViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + + // First add a passkey signer directly + vm.prank(owner); + account.addPasskeySigner(pk); + bytes32 key = Passkey.toKey(pk); + assertTrue(account.passkeySigners(key)); + + // Prepare UserOp that calls removePasskeySigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSelector(account.removePasskeySigner.selector, pk); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute and verify + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + // Verify the passkey signer was actually removed + assertFalse(account.passkeySigners(key)); + } + + function test_OwnerCanCallUpgradeToAndCallViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + address newImplementation = address(0x1234567890123456789012345678901234567890); + bytes memory data = ""; + + // Prepare UserOp that calls upgradeToAndCall + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("upgradeToAndCall(address,bytes)", newImplementation, data); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with owner key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + // Validation should succeed for owner + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Note: Actual execution would fail with "ERC1967: new implementation is not a contract" + // but validation passes, which is what we're testing + } + + function test_AddRemovePasskeySigner() public { + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + bytes32 key = Passkey.toKey(pk); + + vm.prank(ownerAddress); + account.addPasskeySigner(pk); + assert(account.passkeySigners(key)); + + vm.prank(ownerAddress); + account.removePasskeySigner(pk); + assert(!account.passkeySigners(key)); + } + + // ============ Module Management Tests ============ + + function test_OwnerCanRegisterModule() public { + assertFalse(account.isModuleRegistered(address(module))); + + vm.prank(ownerAddress); + account.registerModule(address(module)); + + assertTrue(account.isModuleRegistered(address(module))); + } + + function test_OwnerCanUnregisterModule() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + assertTrue(account.isModuleRegistered(address(module))); + + vm.prank(ownerAddress); + account.unregisterModule(address(module)); + assertFalse(account.isModuleRegistered(address(module))); + } + + function test_RegisterModuleEmitsEvent() public { + vm.expectEmit(true, false, false, false); + emit OmniAccountV2.ModuleRegistered(address(module)); + + vm.prank(ownerAddress); + account.registerModule(address(module)); + } + + function test_UnregisterModuleEmitsEvent() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + vm.expectEmit(true, false, false, false); + emit OmniAccountV2.ModuleUnregistered(address(module)); + + vm.prank(ownerAddress); + account.unregisterModule(address(module)); + } + + function test_OwnerCanRegisterModuleViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("registerModule(address)", address(module)); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + assertTrue(account.isModuleRegistered(address(module))); + } + + function test_OwnerCanUnregisterModuleViaUserOp() public { + (address owner, uint256 ownerPk) = makeAddrAndKey("owner"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootAddress); + + vm.prank(owner); + account.registerModule(address(module)); + assertTrue(account.isModuleRegistered(address(module))); + + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("unregisterModule(address)", address(module)); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + (uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); + + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + assertFalse(account.isModuleRegistered(address(module))); + } + + function test_NonOwnerCannotRegisterModule() public { + address unauthorized = address(0x9999); + vm.expectRevert("only owner"); + vm.prank(unauthorized); + account.registerModule(address(module)); + } + + function test_NonOwnerCannotUnregisterModule() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + address unauthorized = address(0x9999); + vm.expectRevert("only owner"); + vm.prank(unauthorized); + account.unregisterModule(address(module)); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsPasskey.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsPasskey.t.sol new file mode 100644 index 0000000000..7037cd9a92 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsPasskey.t.sol @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {OmniAccountV2} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {BaseAccount} from "../../src/v2/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v2/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v2/interfaces/UserOpSigner.sol"; +import {Counter} from "../../src/Counter.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; +import {PackedUserOperation} from "../../src/v2/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; +import {Passkey} from "../../src/v2/interfaces/Passkey.sol"; +import "@openzeppelin/contracts/utils/Base64.sol"; + +contract OmniAccountV2AsPasskey is Test { + OmniAccountV2 public account; + EntryPointV1 public entryPoint; + Counter public counter; + + address ownerAddress = 0x0000000000000000000000000000000000000000; + bytes clientId = bytes("test_client"); + + function test_ValidateOpPasskey() public { + (address root,) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + // Generate a real P256 private key and derive the public key + uint256 passkeyPrivateKey = 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef; + (uint256 publicKeyX, uint256 publicKeyY) = vm.publicKeyP256(passkeyPrivateKey); + + Passkey.PublicKey memory passkeyPubKey = Passkey.PublicKey({x: publicKeyX, y: publicKeyY}); + + // Add passkey as a signer (must be called by owner) + vm.prank(ownerAddress); + account.addPasskeySigner(passkeyPubKey); + + // Prepare PackedUserOperation + address sender = 0x0eAfeE130Ab1F6261885eE7080f9e8B2513111d4; + bytes memory initCode = ""; + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Create WebAuthn clientDataJSON with the challenge + string memory challengeB64url = Base64.encodeURL(abi.encodePacked(packedOpHash)); + string memory clientDataJSON = string.concat( + '{"type":"webauthn.get","challenge":"', + challengeB64url, + '","origin":"https://example.com","crossOrigin":false}' + ); + + // Create authentic authenticatorData (37 bytes minimum) + // rpIdHash (32 bytes) + flags (1 byte) + signCount (4 bytes) + bytes32 rpIdHash = sha256("example.com"); + bytes1 flags = 0x05; // UP=1, UV=1 (user present and verified) + bytes4 signCount = bytes4(uint32(1)); + bytes memory authData = abi.encodePacked(rpIdHash, flags, signCount); + + // Calculate the message hash that will be signed + bytes32 clientDataJSONHash = sha256(bytes(clientDataJSON)); + bytes32 messageHash = sha256(abi.encodePacked(authData, clientDataJSONHash)); + + // Sign the message hash with P256 + (bytes32 r, bytes32 s) = vm.signP256(passkeyPrivateKey, messageHash); + + Passkey.Signature memory passkeySignature = Passkey.Signature({r: uint256(r), s: uint256(s)}); + + // Find the correct positions for type and challenge in clientDataJSON + bytes memory clientDataJSONBytes = bytes(clientDataJSON); + uint16 typeIndex = _findSubstring(clientDataJSONBytes, bytes('"type":"webauthn.get"')); + uint16 challengeIndex = + _findSubstring(clientDataJSONBytes, bytes(string.concat('"challenge":"', challengeB64url, '"'))); + + Passkey.Metadata memory metadata = Passkey.Metadata({ + authData: authData, + clientDataJSON: clientDataJSON, + challengeIndex: challengeIndex, + typeIndex: typeIndex, + userVerificationRequired: true // UV flag is set + }); + + // Encode signature according to Passkey format + bytes memory passkeySignatureData = abi.encode(passkeyPubKey, passkeySignature, metadata); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Passkey), passkeySignatureData); + + // Validate the user operation + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + + // Should pass with real P256 signature + assertEq(SIG_VALIDATION_SUCCESS, validationData); + } + + function test_ValidateOpPasskeyUnauthorized() public { + (address root,) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + // Create passkey public key but don't add it as authorized signer + Passkey.PublicKey memory passkeyPubKey = Passkey.PublicKey({ + x: 0x65eda5a12577c2bae829437fe338701a10aaa375e1bb5b5de108de439c08551d, + y: 0x1e52ed75701163f7f9e40ddf9f341b3dc9ba860af7e0ca7ca7e9eecd0084d19c + }); + + // Prepare PackedUserOperation + address sender = 0x0eAfeE130Ab1F6261885eE7080f9e8B2513111d4; + bytes memory initCode = ""; + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Create mock signature components + Passkey.Signature memory passkeySignature = Passkey.Signature({ + r: 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, + s: 0xfedcba0987654321fedcba0987654321fedcba0987654321fedcba0987654321 + }); + + bytes memory mockAuthData = hex"49960de5880e8c687434170f6476605b8fe4aeb9a28632c7995cf3ba831d97634500000001"; + string memory mockClientDataJSON = '{"type":"webauthn.get","challenge":"test","origin":"https://example.com"}'; + + Passkey.Metadata memory metadata = Passkey.Metadata({ + authData: mockAuthData, + clientDataJSON: mockClientDataJSON, + challengeIndex: 32, + typeIndex: 1, + userVerificationRequired: false + }); + + bytes memory passkeySignatureData = abi.encode(passkeyPubKey, passkeySignature, metadata); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Passkey), passkeySignatureData); + + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + + // Should fail because passkey isn't authorized + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function _findSubstring(bytes memory haystack, bytes memory needle) internal pure returns (uint16) { + require(needle.length > 0, "Empty needle"); + require(haystack.length >= needle.length, "Needle longer than haystack"); + + for (uint256 i = 0; i <= haystack.length - needle.length; i++) { + bool found = true; + for (uint256 j = 0; j < needle.length; j++) { + if (haystack[i + j] != needle[j]) { + found = false; + break; + } + } + if (found) { + return uint16(i); + } + } + revert("Substring not found"); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRoot.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRoot.t.sol new file mode 100644 index 0000000000..eb2c76eb5a --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRoot.t.sol @@ -0,0 +1,540 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {OmniAccountV2} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {BaseAccount} from "../../src/v2/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/v2/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v2/interfaces/UserOpSigner.sol"; +import {Counter} from "../../src/Counter.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; +import {PackedUserOperation} from "../../src/v2/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; +import {Passkey} from "../../src/v2/interfaces/Passkey.sol"; +import {MockModule} from "./MockModule.sol"; + +contract OmniAccountV2AsRoot is Test { + OmniAccountV2 public account; + EntryPointV1 public entryPoint; + Counter public counter; + MockModule public module; + + address ownerAddress = 0x0000000000000000000000000000000000000000; + address rootAddress = 0x0000000000000000000000000000000000000001; + bytes clientId = bytes("test_client"); + + function setUp() public { + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, rootAddress); + module = new MockModule(); + } + + function test_Execute() public { + vm.expectRevert("account: not from EntryPoint"); + vm.prank(rootAddress); + account.execute(address(counter), 0, abi.encodeWithSignature("increment()")); + } + + function test_ExecuteBatch() public { + vm.expectRevert("account: not from EntryPoint"); + vm.prank(rootAddress); + BaseAccount.Call[] memory calls = new BaseAccount.Call[](1); + calls[0] = BaseAccount.Call({target: address(counter), value: 0, data: abi.encodeWithSignature("increment()")}); + account.executeBatch(calls); + } + + function test_ValidateOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address sender = 0x0eAfeE130Ab1F6261885eE7080f9e8B2513111d4; + bytes memory initCode = ""; + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + // sign userOp + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + } + + function test_ValidateOpFailsWithWrongSigner() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address sender = 0x0eAfeE130Ab1F6261885eE7080f9e8B2513111d4; + bytes memory initCode = ""; + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + // sign userOp + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v); // The signer is wrong here + + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_ValidateOpSessionKey() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (address session, uint256 sessionPk) = makeAddrAndKey("session"); + uint256 sessionExpiration = 2; + + bytes memory sessionProof = prepareSession(session, sessionExpiration, rootPk); + + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address sender = 0x0eAfeE130Ab1F6261885eE7080f9e8B2513111d4; + bytes memory initCode = ""; + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + // sign userOp + (uint8 v, bytes32 r, bytes32 s) = vm.sign(sessionPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.SessionKey), r, s, v, sessionExpiration, sessionProof); + + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + } + + function test_ValidateOpExpiredSessionKey() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (address session, uint256 sessionPk) = makeAddrAndKey("session"); + uint256 sessionExpiration = 0; + + bytes memory sessionProof = prepareSession(session, sessionExpiration, rootPk); + + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address sender = 0x0eAfeE130Ab1F6261885eE7080f9e8B2513111d4; + bytes memory initCode = ""; + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + // sign userOp + (uint8 v, bytes32 r, bytes32 s) = vm.sign(sessionPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.SessionKey), r, s, v, sessionExpiration, sessionProof); + + vm.prank(address(entryPoint)); + + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_ValidateOpSessionKeyFailsIfProofNotSignedByRoot() public { + (, uint256 alicePk) = makeAddrAndKey("alice"); + (address root,) = makeAddrAndKey("root"); + (address session, uint256 sessionPk) = makeAddrAndKey("session"); + uint256 sessionExpiration = 2; + + bytes memory sessionProof = prepareSession(session, sessionExpiration, alicePk); + + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address sender = 0x0eAfeE130Ab1F6261885eE7080f9e8B2513111d4; + bytes memory initCode = ""; + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + // sign userOp + (uint8 v, bytes32 r, bytes32 s) = vm.sign(sessionPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.SessionKey), r, s, v, sessionExpiration, sessionProof); + + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_AddRootSigner_As_Not_Allowed() public { + vm.expectRevert("only owner"); + account.addRootSigner(0x0000000000000000000000000000000000000000); + } + + function test_RemoveRootSigner_As_Not_Allowed() public { + vm.expectRevert("only owner"); + account.removeRootSigner(0x0000000000000000000000000000000000000000); + } + + function prepareSession(address session, uint256 expiration, uint256 proofSigner) + internal + pure + returns (bytes memory) + { + bytes32 sessionDigest = sha256(abi.encodePacked(session, expiration)); + (uint8 sv, bytes32 sr, bytes32 ss) = vm.sign(proofSigner, sessionDigest); + bytes memory sessionProof = abi.encodePacked(sr, ss, sv); + return sessionProof; + } + + function test_RootCannotCallAddRootSignerViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address newRoot = 0x0000000000000000000000000000000000000002; + + // Prepare UserOp that calls addRootSigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should fail because root is trying to call a restricted function + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_RootCannotCallRemoveRootSignerViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + // Prepare UserOp that calls removeRootSigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("removeRootSigner(address)", root); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should fail + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_SessionKeyCannotCallAddRootSignerViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (address session, uint256 sessionPk) = makeAddrAndKey("session"); + uint256 sessionExpiration = block.timestamp + 1000; + + bytes memory sessionProof = prepareSession(session, sessionExpiration, rootPk); + + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address newRoot = 0x0000000000000000000000000000000000000002; + + // Prepare UserOp that calls addRootSigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with session key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(sessionPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.SessionKey), r, s, v, sessionExpiration, sessionProof); + + // Validation should fail because session key is trying to call a restricted function + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_RootCannotExecuteBatchViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + // Prepare UserOp that calls executeBatch + address sender = address(account); + bytes memory initCode = ""; + + BaseAccount.Call[] memory calls = new BaseAccount.Call[](1); + calls[0] = BaseAccount.Call({target: address(counter), value: 0, data: abi.encodeWithSignature("increment()")}); + bytes memory callData = abi.encodeWithSignature("executeBatch((address,uint256,bytes)[])", calls); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should succeed because root can use executeBatch + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + } + + function test_RootCannotCallAddRootSignerViaExecute() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address newRoot = 0x0000000000000000000000000000000000000002; + + // Prepare UserOp that calls execute() with addRootSigner as inner call + address sender = address(account); + bytes memory initCode = ""; + bytes memory innerCallData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + bytes memory callData = + abi.encodeWithSignature("execute(address,uint256,bytes)", address(account), 0, innerCallData); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should succeed because root can call execute + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + } + + function test_SessionKeyCannotCallAddRootSignerViaExecute() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (address session, uint256 sessionPk) = makeAddrAndKey("session"); + uint256 sessionExpiration = block.timestamp + 1000; + + bytes memory sessionProof = prepareSession(session, sessionExpiration, rootPk); + + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address newRoot = 0x0000000000000000000000000000000000000002; + + // Prepare UserOp that calls execute() with addRootSigner as inner call + address sender = address(account); + bytes memory initCode = ""; + bytes memory innerCallData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + bytes memory callData = + abi.encodeWithSignature("execute(address,uint256,bytes)", address(account), 0, innerCallData); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with session key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(sessionPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.SessionKey), r, s, v, sessionExpiration, sessionProof); + + // Validation should succeed because session key can call execute + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + } + + function test_RootCannotCallAddPasskeySignerViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + + // Prepare UserOp that calls addPasskeySigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSelector(account.addPasskeySigner.selector, pk); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should fail because root is trying to call a restricted function + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_RootCannotCallRemovePasskeySignerViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + + // Prepare UserOp that calls removePasskeySigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSelector(account.removePasskeySigner.selector, pk); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should fail + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_RootCannotCallWithdrawDepositToViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + // Add some deposit first + vm.deal(address(account), 1 ether); + vm.prank(address(account)); + account.addDeposit{value: 0.5 ether}(); + + address payable withdrawTo = payable(0x0000000000000000000000000000000000000003); + uint256 withdrawAmount = 0.1 ether; + + // Prepare UserOp that calls withdrawDepositTo + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = + abi.encodeWithSignature("withdrawDepositTo(address,uint256)", withdrawTo, withdrawAmount); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should fail because root is trying to call a restricted function + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_RootCannotCallUpgradeToAndCallViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address newImplementation = address(0x1234567890123456789012345678901234567890); + bytes memory data = ""; + + // Prepare UserOp that calls upgradeToAndCall + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("upgradeToAndCall(address,bytes)", newImplementation, data); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should fail because root is trying to call a restricted function + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_AddPasskeySigner_As_Not_Allowed() public { + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + vm.expectRevert("only owner"); + account.addPasskeySigner(pk); + } + + function test_RemovePasskeySigner_As_Not_Allowed() public { + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + vm.expectRevert("only owner"); + account.removePasskeySigner(pk); + } + + function test_WithdrawDepositTo_As_Not_Allowed() public { + vm.expectRevert("only owner"); + account.withdrawDepositTo(payable(address(0x1)), 100); + } + + function test_UpgradeToAndCall_As_Not_Allowed() public { + vm.expectRevert("only owner"); + account.upgradeToAndCall(address(0x1), ""); + } + + // ============ Module Management Tests ============ + + function test_RootCannotRegisterModule() public { + vm.expectRevert("only owner"); + vm.prank(rootAddress); + account.registerModule(address(module)); + } + + function test_RootCannotUnregisterModule() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + vm.expectRevert("only owner"); + vm.prank(rootAddress); + account.unregisterModule(address(module)); + } + + function test_RootCannotRegisterModuleViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("registerModule(address)", address(module)); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should fail because root is trying to call a restricted function + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_SessionKeyCannotRegisterModuleViaUserOp() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (address session, uint256 sessionPk) = makeAddrAndKey("session"); + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root); + + uint256 sessionExpiration = block.timestamp + 1000; + bytes32 sessionDigest = sha256(abi.encodePacked(session, sessionExpiration)); + (uint8 sv, bytes32 sr, bytes32 ss) = vm.sign(rootPk, sessionDigest); + bytes memory sessionProof = abi.encodePacked(sr, ss, sv); + + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("registerModule(address)", address(module)); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + (uint8 v, bytes32 r, bytes32 s) = vm.sign(sessionPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.SessionKey), r, s, v, sessionExpiration, sessionProof); + + // Validation should fail because session key is trying to call a restricted function + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRootNonEvm.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRootNonEvm.t.sol new file mode 100644 index 0000000000..d5ce0bc2c7 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRootNonEvm.t.sol @@ -0,0 +1,455 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {OmniAccountV2 as OmniAccount} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {BaseAccount} from "../../src/v2/core/BaseAccount.sol"; +import {EntryPointV1 as EntryPoint} from "../../src/v2/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v2/interfaces/UserOpSigner.sol"; +import {OwnerType} from "../../src/v2/interfaces/OwnerType.sol"; +import {Counter} from "../../src/Counter.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; +import {PackedUserOperation} from "../../src/v2/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; +import {Passkey} from "../../src/v2/interfaces/Passkey.sol"; +import {MockModule} from "./MockModule.sol"; + +contract OmniAccountV2AsRootNonEvm is Test { + OmniAccount public account; + EntryPoint public entryPoint; + Counter public counter; + MockModule public module; + + address ownerAddress = 0x0000000000000000000000000000000000000000; + address rootAddress = 0x0000000000000000000000000000000000000001; + bytes clientId = bytes("test_email@example.com"); + + function setUp() public { + // Set up with Email owner type (non-EVM) + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, rootAddress, OwnerType.Email); + module = new MockModule(); + } + + function test_RootCanCallAddRootSignerDirectlyForNonEvmOwner() public { + // Direct call from root should succeed for non-EVM owner + address newRoot = 0x0000000000000000000000000000000000000002; + vm.prank(rootAddress); + account.addRootSigner(newRoot); + assertTrue(account.rootSigners(newRoot)); + } + + function test_RootCanCallRemoveRootSignerDirectlyForNonEvmOwner() public { + // First add a root signer + address newRoot = 0x0000000000000000000000000000000000000002; + vm.prank(rootAddress); + account.addRootSigner(newRoot); + assertTrue(account.rootSigners(newRoot)); + + // Direct call from root should succeed for non-EVM owner + vm.prank(rootAddress); + account.removeRootSigner(newRoot); + assertFalse(account.rootSigners(newRoot)); + } + + function test_RootCanCallAddRootSignerViaUserOpForNonEvmOwner() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, root, OwnerType.Email); + + address newRoot = 0x0000000000000000000000000000000000000002; + + // Prepare UserOp that calls addRootSigner directly + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should succeed because root can call restricted functions for non-EVM owners + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute the UserOp + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + // Verify the root signer was actually added + assertTrue(account.rootSigners(newRoot)); + } + + function test_RootCanCallWithdrawDepositViaUserOpForNonEvmOwner() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, root, OwnerType.Google); + + // Add some deposit first + vm.deal(address(account), 1 ether); + vm.prank(address(account)); + account.addDeposit{value: 0.5 ether}(); + + address payable withdrawTo = payable(0x0000000000000000000000000000000000000003); + uint256 withdrawAmount = 0.1 ether; + + // Prepare UserOp that calls withdrawDepositTo + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = + abi.encodeWithSignature("withdrawDepositTo(address,uint256)", withdrawTo, withdrawAmount); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should succeed for root with non-EVM owner + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute and verify + uint256 balanceBefore = withdrawTo.balance; + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + assertEq(withdrawTo.balance, balanceBefore + withdrawAmount); + } + + function test_SessionKeyStillCannotCallRestrictedFunctionsForNonEvmOwner() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (address session, uint256 sessionPk) = makeAddrAndKey("session"); + uint256 sessionExpiration = block.timestamp + 1000; + + bytes memory sessionProof = prepareSession(session, sessionExpiration, rootPk); + + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, root, OwnerType.Twitter); + + address newRoot = 0x0000000000000000000000000000000000000002; + + // Prepare UserOp that calls addRootSigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with session key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(sessionPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.SessionKey), r, s, v, sessionExpiration, sessionProof); + + // Validation should still fail for session keys even with non-EVM owner + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function prepareSession(address session, uint256 expiration, uint256 proofSigner) + internal + pure + returns (bytes memory) + { + bytes32 sessionDigest = sha256(abi.encodePacked(session, expiration)); + (uint8 sv, bytes32 sr, bytes32 ss) = vm.sign(proofSigner, sessionDigest); + bytes memory sessionProof = abi.encodePacked(sr, ss, sv); + return sessionProof; + } + + function test_RootCanCallAddPasskeySignerViaUserOpForNonEvmOwner() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, root, OwnerType.Email); + + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + + // Prepare UserOp that calls addPasskeySigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSelector(account.addPasskeySigner.selector, pk); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should succeed because root can call restricted functions for non-EVM owners + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute the UserOp + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + // Verify the passkey signer was actually added + bytes32 key = Passkey.toKey(pk); + assertTrue(account.passkeySigners(key)); + } + + function test_RootCannotRemovePasskeySignerViaUserOpWhenPasskeySignersExist() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, root, OwnerType.Google); + + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + + // First add a passkey signer directly + vm.prank(root); + account.addPasskeySigner(pk); + bytes32 key = Passkey.toKey(pk); + assertTrue(account.passkeySigners(key)); + + // Prepare UserOp that calls removePasskeySigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSelector(account.removePasskeySigner.selector, pk); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should fail for root when passkey signers exist + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_RootCanCallUpgradeToAndCallViaUserOpForNonEvmOwner() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, root, OwnerType.Twitter); + + address newImplementation = address(0x1234567890123456789012345678901234567890); + bytes memory data = ""; + + // Prepare UserOp that calls upgradeToAndCall + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("upgradeToAndCall(address,bytes)", newImplementation, data); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should succeed for root with non-EVM owner + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + } + + function test_RootCanCallAddPasskeySignerDirectlyForNonEvmOwner() public { + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + + // Direct call from root should succeed for non-EVM owner + vm.prank(rootAddress); + account.addPasskeySigner(pk); + bytes32 key = Passkey.toKey(pk); + assertTrue(account.passkeySigners(key)); + } + + function test_RootCannotRemovePasskeySignerDirectlyWhenPasskeySignersExist() public { + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + + // First add a passkey signer + vm.prank(rootAddress); + account.addPasskeySigner(pk); + bytes32 key = Passkey.toKey(pk); + assertTrue(account.passkeySigners(key)); + + // Direct call from root should fail because passkey signers now exist + vm.expectRevert("only owner"); + vm.prank(rootAddress); + account.removePasskeySigner(pk); + } + + function test_RootCanCallUpgradeToAndCallDirectlyForNonEvmOwner() public { + address newImplementation = address(0x1234567890123456789012345678901234567890); + bytes memory data = ""; + + // Direct call from root should succeed for non-EVM owner + // This test will revert because the implementation address is not a valid contract + // but that's expected - we're testing access control, not actual upgrade + vm.expectRevert(); + vm.prank(rootAddress); + account.upgradeToAndCall(newImplementation, data); + } + + // ============ Module Management Tests for Non-EVM Accounts ============ + + function test_RootCanRegisterModuleDirectlyWhenNoPasskeys() public { + // Root should be able to register module directly for non-EVM owner with no passkeys + vm.prank(rootAddress); + account.registerModule(address(module)); + + assertTrue(account.isModuleRegistered(address(module))); + } + + function test_RootCanUnregisterModuleDirectlyWhenNoPasskeys() public { + // First register a module + vm.prank(rootAddress); + account.registerModule(address(module)); + assertTrue(account.isModuleRegistered(address(module))); + + // Root should be able to unregister module directly + vm.prank(rootAddress); + account.unregisterModule(address(module)); + + assertFalse(account.isModuleRegistered(address(module))); + } + + function test_RootCanRegisterModuleViaUserOpWhenNoPasskeys() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, root, OwnerType.Substrate); + + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("registerModule(address)", address(module)); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should succeed because root can call restricted functions for non-EVM owners with no passkeys + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_SUCCESS, validationData); + + // Execute the UserOp + vm.prank(address(entryPoint)); + (bool success,) = address(account).call(callData); + assertTrue(success); + + // Verify the module was actually registered + assertTrue(account.isModuleRegistered(address(module))); + } + + function test_RootCannotRegisterModuleDirectlyWhenPasskeyExists() public { + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + + // First add a passkey signer (root can do this while they're still owner) + vm.prank(rootAddress); + account.addPasskeySigner(pk); + bytes32 key = Passkey.toKey(pk); + assertTrue(account.passkeySigners(key)); + + // Now root should NOT be able to register module directly + vm.expectRevert("only owner"); + vm.prank(rootAddress); + account.registerModule(address(module)); + } + + function test_RootCannotUnregisterModuleDirectlyWhenPasskeyExists() public { + // First register module while root is owner (no passkeys yet) + vm.prank(rootAddress); + account.registerModule(address(module)); + assertTrue(account.isModuleRegistered(address(module))); + + // Add passkey signer + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + vm.prank(rootAddress); + account.addPasskeySigner(pk); + + // Now root should NOT be able to unregister module + vm.expectRevert("only owner"); + vm.prank(rootAddress); + account.unregisterModule(address(module)); + } + + function test_RootCannotRegisterModuleViaUserOpWhenPasskeyExists() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, root, OwnerType.Google); + + // Add passkey signer first + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + vm.prank(root); + account.addPasskeySigner(pk); + + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("registerModule(address)", address(module)); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should FAIL because passkey signer now exists + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_RootCannotUnregisterModuleViaUserOpWhenPasskeyExists() public { + (address root, uint256 rootPk) = makeAddrAndKey("root"); + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, root, OwnerType.Twitter); + + // Register module first (while root is owner) + vm.prank(root); + account.registerModule(address(module)); + assertTrue(account.isModuleRegistered(address(module))); + + // Add passkey signer + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 12345, y: 67890}); + vm.prank(root); + account.addPasskeySigner(pk); + + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("unregisterModule(address)", address(module)); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should FAIL because passkey signer now exists + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Modules.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Modules.t.sol new file mode 100644 index 0000000000..ca42235521 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Modules.t.sol @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {OmniAccountV2} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {EntryPointV1} from "../../src/v2/core/EntryPointV1.sol"; +import {Counter} from "../../src/Counter.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {MockModule} from "./MockModule.sol"; + +/** + * General module tests that don't fit into actor-specific test files. + * Tests for specific actors (Owner, Root, EntryPoint) are in their respective files: + * - OmniAccountV2AsOwner.t.sol + * - OmniAccountV2AsRoot.t.sol + * - OmniAccountV2AsEntryPoint.t.sol + */ +contract OmniAccountV2Modules is Test { + OmniAccountV2 public account; + EntryPointV1 public entryPoint; + Counter public counter; + MockModule public module; + MockModule public module2; + + address ownerAddress = 0x0000000000000000000000000000000000000000; + address rootAddress = 0x0000000000000000000000000000000000000001; + bytes clientId = bytes("test_client"); + + function setUp() public { + (counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, rootAddress); + module = new MockModule(); + module2 = new MockModule(); + } + + // ============ Module Registration Validation Tests ============ + + function test_CannotRegisterZeroAddressAsModule() public { + vm.expectRevert("Invalid module address"); + vm.prank(ownerAddress); + account.registerModule(address(0)); + } + + function test_CannotRegisterNonContractAsModule() public { + address eoa = address(0x1234); + vm.expectRevert("Module must be a contract"); + vm.prank(ownerAddress); + account.registerModule(eoa); + } + + function test_CannotRegisterModuleTwice() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + vm.expectRevert("Module already registered"); + vm.prank(ownerAddress); + account.registerModule(address(module)); + } + + function test_CannotUnregisterNonRegisteredModule() public { + vm.expectRevert("Module not registered"); + vm.prank(ownerAddress); + account.unregisterModule(address(module)); + } + + // ============ Module Execution Tests ============ + + function test_CannotExecuteUnregisteredModule() public { + bytes memory callData = abi.encodeWithSignature("increment()"); + + vm.expectRevert("Module not registered"); + vm.prank(address(entryPoint)); + account.executeModuleCall(address(module), callData); + } + + function test_ModuleExecutionAccessesAccountContext() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + // Set a value in the account's storage at slot 0 (where MockModule.counter is stored) + // This proves delegatecall reads from account's storage, not module's storage + uint256 expectedValue = 123; + vm.store(address(account), bytes32(uint256(0)), bytes32(expectedValue)); + + bytes memory callData = abi.encodeWithSignature("getCounter()"); + + vm.prank(address(entryPoint)); + bytes memory returnData = account.executeModuleCall(address(module), callData); + + // Delegatecall executes module code in account's context + // The module function should access account's storage and return the value we set + uint256 result = abi.decode(returnData, (uint256)); + assertEq(result, expectedValue, "Module should read from account's storage via delegatecall"); + + // Verify module's own storage is still 0 (unaffected) + assertEq(module.counter(), 0, "Module's own storage should remain unchanged"); + } + + // ============ Multiple Module Tests ============ + + function test_MultipleModulesCanBeRegistered() public { + vm.prank(ownerAddress); + account.registerModule(address(module)); + + vm.prank(ownerAddress); + account.registerModule(address(module2)); + + assertTrue(account.isModuleRegistered(address(module))); + assertTrue(account.isModuleRegistered(address(module2))); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol new file mode 100644 index 0000000000..88a8ae2d7a --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol @@ -0,0 +1,279 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {OmniAccountV2 as OmniAccount} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {BaseAccount} from "../../src/v2/core/BaseAccount.sol"; +import {EntryPointV1 as EntryPoint} from "../../src/v2/core/EntryPointV1.sol"; +import {UserOpSigner} from "../../src/v2/interfaces/UserOpSigner.sol"; +import {OwnerType} from "../../src/v2/interfaces/OwnerType.sol"; +import {Passkey} from "../../src/v2/interfaces/Passkey.sol"; +import {Counter} from "../../src/Counter.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; +import {PackedUserOperation} from "../../src/v2/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/v1/core/Helpers.sol"; + +contract OmniAccountV2NonEvmOwnerWithPasskeySigner is Test { + OmniAccount public account; + EntryPoint public entryPoint; + Counter public counter; + + address ownerAddress = 0x0000000000000000000000000000000000000000; + address rootAddress = 0x0000000000000000000000000000000000000001; + bytes clientId = bytes("test_email@example.com"); + + function setUp() public { + // Set up with Email owner type (non-EVM) + (counter, entryPoint, account) = + OmniAccountV2TestUtils.setUpWithOwnerType(ownerAddress, clientId, rootAddress, OwnerType.Email); + } + + function test_RootCannotCallOnlyOwnerWhenPasskeySignersExist() public { + // First add a passkey signer as owner + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 1, y: 2}); + vm.prank(ownerAddress); + account.addPasskeySigner(pk); + assertEq(account.passkeySignerCount(), 1); + + // Now root should not be able to call onlyOwner functions + address newRoot = 0x0000000000000000000000000000000000000002; + vm.prank(rootAddress); + vm.expectRevert("only owner"); + account.addRootSigner(newRoot); + } + + function test_RootCanCallOnlyOwnerAgainAfterLastPasskeyRemoved() public { + // Add a passkey signer + Passkey.PublicKey memory pk1 = Passkey.PublicKey({x: 1, y: 2}); + vm.prank(ownerAddress); + account.addPasskeySigner(pk1); + assertEq(account.passkeySignerCount(), 1); + + // Root cannot call onlyOwner + address newRoot = 0x0000000000000000000000000000000000000002; + vm.prank(rootAddress); + vm.expectRevert("only owner"); + account.addRootSigner(newRoot); + + // Remove the passkey signer + vm.prank(ownerAddress); + account.removePasskeySigner(pk1); + assertEq(account.passkeySignerCount(), 0); + + // Now root should be able to call onlyOwner again + vm.prank(rootAddress); + account.addRootSigner(newRoot); + assertTrue(account.rootSigners(newRoot)); + } + + function test_MultiplePasskeySignersCount() public { + // Add multiple passkey signers + Passkey.PublicKey memory pk1 = Passkey.PublicKey({x: 1, y: 2}); + Passkey.PublicKey memory pk2 = Passkey.PublicKey({x: 3, y: 4}); + Passkey.PublicKey memory pk3 = Passkey.PublicKey({x: 5, y: 6}); + + vm.startPrank(ownerAddress); + account.addPasskeySigner(pk1); + assertEq(account.passkeySignerCount(), 1); + + account.addPasskeySigner(pk2); + assertEq(account.passkeySignerCount(), 2); + + account.addPasskeySigner(pk3); + assertEq(account.passkeySignerCount(), 3); + + // Remove one + account.removePasskeySigner(pk2); + assertEq(account.passkeySignerCount(), 2); + + // Try to add the same key again (should not increase count) + account.addPasskeySigner(pk1); + assertEq(account.passkeySignerCount(), 2); + + // Try to remove non-existent key (should not decrease count) + account.removePasskeySigner(pk2); + assertEq(account.passkeySignerCount(), 2); + vm.stopPrank(); + } + + function test_RootCannotCallRestrictedFunctionsViaUserOpWhenPasskeyExists() public { + // Add a passkey signer first + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 100, y: 200}); + vm.prank(ownerAddress); + account.addPasskeySigner(pk); + assertEq(account.passkeySignerCount(), 1); + + // Setup root signer + (address root, uint256 rootPk) = makeAddrAndKey("root"); + vm.prank(ownerAddress); + account.addRootSigner(root); + + address newRoot = 0x0000000000000000000000000000000000000003; + + // Prepare UserOp that calls addRootSigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Sign with root key + (uint8 v, bytes32 r, bytes32 s) = vm.sign(rootPk, packedOpHash); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.RootKey), r, s, v); + + // Validation should fail because passkey signers exist + vm.prank(address(entryPoint)); + uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0); + assertEq(SIG_VALIDATION_FAILED, validationData); + } + + function test_PasskeyValidationForAddRootSignerWhenPasskeyExists() public { + // Add a passkey signer + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 111, y: 222}); + vm.prank(ownerAddress); + account.addPasskeySigner(pk); + assertEq(account.passkeySignerCount(), 1); + + address newRoot = 0x0000000000000000000000000000000000000004; + + // Prepare UserOp that calls addRootSigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("addRootSigner(address)", newRoot); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Create passkey signature (simplified - actual implementation will be done by colleague) + bytes memory passkeySignature = abi.encodePacked("passkey_signature_placeholder"); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Passkey), passkeySignature); + + vm.prank(address(entryPoint)); + // TODO: use a legit passkey signature + vm.expectRevert(); + account.validateUserOp(packedOp, packedOpHash, 0); + } + + function test_PasskeyValidationForRemoveRootSignerWhenPasskeyExists() public { + // Add a passkey signer + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 111, y: 222}); + vm.prank(ownerAddress); + account.addPasskeySigner(pk); + assertEq(account.passkeySignerCount(), 1); + + // Prepare UserOp that calls removeRootSigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("removeRootSigner(address)", rootAddress); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Create passkey signature + bytes memory passkeySignature = abi.encodePacked("passkey_signature_placeholder"); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Passkey), passkeySignature); + + vm.prank(address(entryPoint)); + // TODO: use a legit passkey signature + vm.expectRevert(); + account.validateUserOp(packedOp, packedOpHash, 0); + } + + function test_PasskeyValidationForAddPasskeySignerWhenPasskeyExists() public { + // Add a passkey signer + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 111, y: 222}); + vm.prank(ownerAddress); + account.addPasskeySigner(pk); + assertEq(account.passkeySignerCount(), 1); + + // Prepare UserOp that calls addPasskeySigner + Passkey.PublicKey memory newPk = Passkey.PublicKey({x: 333, y: 444}); + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("addPasskeySigner((uint256,uint256))", newPk); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Create passkey signature + bytes memory passkeySignature = abi.encodePacked("passkey_signature_placeholder"); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Passkey), passkeySignature); + + vm.prank(address(entryPoint)); + // TODO: use a legit passkey signature + vm.expectRevert(); + account.validateUserOp(packedOp, packedOpHash, 0); + } + + function test_PasskeyValidationForRemovePasskeySignerWhenPasskeyExists() public { + // Add a passkey signer + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 111, y: 222}); + vm.prank(ownerAddress); + account.addPasskeySigner(pk); + assertEq(account.passkeySignerCount(), 1); + + // Prepare UserOp that calls removePasskeySigner + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = abi.encodeWithSignature("removePasskeySigner((uint256,uint256))", pk); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Create passkey signature + bytes memory passkeySignature = abi.encodePacked("passkey_signature_placeholder"); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Passkey), passkeySignature); + + vm.prank(address(entryPoint)); + // TODO: use a legit passkey signature + vm.expectRevert(); + account.validateUserOp(packedOp, packedOpHash, 0); + } + + function test_PasskeyValidationForWithdrawDepositWhenPasskeyExists() public { + // Add a passkey signer + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 111, y: 222}); + vm.prank(ownerAddress); + account.addPasskeySigner(pk); + assertEq(account.passkeySignerCount(), 1); + + // Add some deposit first + vm.deal(address(account), 1 ether); + vm.prank(address(account)); + account.addDeposit{value: 0.5 ether}(); + + // Prepare UserOp that calls withdrawDepositTo + address payable withdrawTo = payable(0x0000000000000000000000000000000000000005); + uint256 withdrawAmount = 0.1 ether; + address sender = address(account); + bytes memory initCode = ""; + bytes memory callData = + abi.encodeWithSignature("withdrawDepositTo(address,uint256)", withdrawTo, withdrawAmount); + + PackedUserOperation memory packedOp = TestUtils.preparePackedOp(sender, initCode); + packedOp.callData = callData; + + bytes32 packedOpHash = entryPoint.getUserOpHash(packedOp); + + // Create passkey signature + bytes memory passkeySignature = abi.encodePacked("passkey_signature_placeholder"); + packedOp.signature = abi.encodePacked(uint8(UserOpSigner.Passkey), passkeySignature); + + vm.prank(address(entryPoint)); + // TODO: use a legit passkey signature + vm.expectRevert(); + account.validateUserOp(packedOp, packedOpHash, 0); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol new file mode 100644 index 0000000000..111e234012 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Counter} from "../../src/Counter.sol"; +import {Vm} from "forge-std/Vm.sol"; +import {OmniAccountV2 as OmniAccount} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {BaseAccount} from "../../src/v2/core/BaseAccount.sol"; +import {EntryPointV1 as EntryPoint} from "../../src/v2/core/EntryPointV1.sol"; +import {IEntryPoint} from "../../src/v2/interfaces/IEntryPoint.sol"; +import {OwnerType} from "../../src/v2/interfaces/OwnerType.sol"; +import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; + +library OmniAccountV2TestUtils { + function setUpWithOwnerType(address ownerAddress, bytes memory clientId, address rootAddress, OwnerType ownerType) + external + returns (Counter, EntryPoint, OmniAccount) + { + Counter counter = new Counter(); + EntryPoint entryPoint = new EntryPoint(); + OmniAccount accountImpl = new OmniAccount(IEntryPoint(address(entryPoint))); + bytes32 oa = TestUtils.prepare_evm_oa(ownerAddress, clientId); + OmniAccount account = OmniAccount( + payable(new ERC1967Proxy{salt: oa}( + address(accountImpl), abi.encodeCall(OmniAccount.initialize, (oa, ownerType, clientId, rootAddress)) + )) + ); + + return (counter, entryPoint, account); + } + + function setUp(address ownerAddress, bytes memory clientId, address rootAddress) + external + returns (Counter, EntryPoint, OmniAccount) + { + Counter counter = new Counter(); + EntryPoint entryPoint = new EntryPoint(); + OmniAccount accountImpl = new OmniAccount(IEntryPoint(address(entryPoint))); + bytes32 oa = TestUtils.prepare_evm_oa(ownerAddress, clientId); + OmniAccount account = OmniAccount( + payable(new ERC1967Proxy{salt: oa}( + address(accountImpl), + abi.encodeCall(OmniAccount.initialize, (oa, OwnerType.Evm, clientId, rootAddress)) + )) + ); + + return (counter, entryPoint, account); + } + + function performExecuteTestAs(Vm vm, address asAccount, OmniAccount account, Counter counter) internal { + uint256 number = counter.number(); + vm.prank(asAccount); + account.execute(address(counter), 0, abi.encodeWithSignature("increment()")); + vm.assertEq(number + 1, counter.number()); + } + + function performExecuteBatchTestAs(Vm vm, address asAccount, OmniAccount account, Counter counter) internal { + uint256 number = counter.number(); + vm.prank(asAccount); + BaseAccount.Call[] memory calls = new BaseAccount.Call[](2); + calls[0] = BaseAccount.Call(address(counter), 0, abi.encodeWithSignature("increment()")); + calls[1] = BaseAccount.Call(address(counter), 0, abi.encodeWithSignature("increment()")); + account.executeBatch(calls); + vm.assertEq(number + 2, counter.number()); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Upgradability.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Upgradability.t.sol new file mode 100644 index 0000000000..7c08afef21 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Upgradability.t.sol @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; +import "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol"; +import {OmniAccountV2} from "../../src/v2/accounts/OmniAccountV2.sol"; +import {EntryPointV1} from "../../src/v2/core/EntryPointV1.sol"; +import {IEntryPoint} from "../../src/v2/interfaces/IEntryPoint.sol"; +import {OwnerType} from "../../src/v2/interfaces/OwnerType.sol"; +import {Passkey} from "../../src/v2/interfaces/Passkey.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; + +contract OmniAccountV3 is OmniAccountV2 { + constructor(EntryPointV1 anEntryPoint) OmniAccountV2(IEntryPoint(address(anEntryPoint))) {} + + function version() public pure override returns (string memory) { + return "3.0.0"; + } +} + +contract OmniAccountV2Upgradability is Test { + // Test addresses + address public owner = address(0x1234); + address public rootSigner = address(0x5678); + address public unauthorizedUser = address(0x9abc); + + // Test data + // bytes32 public ownerOa; + bytes clientId = bytes("test_client"); + + function setUp() public {} + + // Helper function to verify account version + function assertAccountVersion(address account) internal { + (bool success, bytes memory result) = account.call(abi.encodeWithSignature("version()")); + assertTrue(success, "Should be able to call version()"); + string memory version = abi.decode(result, (string)); + assertEq(version, "3.0.0", "Should be upgraded to version 3.0.0"); + } + + function testOaOwnerCanUpgradeOA() public { + (, EntryPointV1 entryPoint, OmniAccountV2 account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootSigner); + OmniAccountV3 accountV3 = new OmniAccountV3(entryPoint); + + vm.prank(owner); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV3), ""); + assertAccountVersion(address(account)); + } + + function testEntryPointCanUpgradeOA() public { + (, EntryPointV1 entryPoint, OmniAccountV2 account) = OmniAccountV2TestUtils.setUp(owner, clientId, rootSigner); + OmniAccountV3 accountV3 = new OmniAccountV3(entryPoint); + + vm.prank(address(entryPoint)); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV3), ""); + assertAccountVersion(address(account)); + } + + function testUnauthorizedSenderCannotUpgradeOA() public { + (, EntryPointV1 entryPoint, OmniAccountV2 account) = + OmniAccountV2TestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); + OmniAccountV3 accountV3 = new OmniAccountV3(entryPoint); + + vm.expectRevert("only owner"); + vm.prank(unauthorizedUser); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV3), ""); + } + + function testRootSignerCannotUpgradeEvmOA() public { + (, EntryPointV1 entryPoint, OmniAccountV2 account) = + OmniAccountV2TestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Evm); + OmniAccountV3 accountV3 = new OmniAccountV3(entryPoint); + + vm.expectRevert("only owner"); + vm.prank(rootSigner); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV3), ""); + } + + function testRootSignerCannotUpgradeNonEvmOAWithPassKey() public { + (, EntryPointV1 entryPoint, OmniAccountV2 account) = + OmniAccountV2TestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); + OmniAccountV3 accountV3 = new OmniAccountV3(entryPoint); + + Passkey.PublicKey memory pk = Passkey.PublicKey({x: 1, y: 2}); + vm.prank(owner); + account.addPasskeySigner(pk); + assertEq(account.passkeySignerCount(), 1); + + vm.expectRevert("only owner"); + vm.prank(rootSigner); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV3), ""); + } + + function testRootSignerCanUpgradeNonEvmOAWithoutPassKey() public { + (, EntryPointV1 entryPoint, OmniAccountV2 account) = + OmniAccountV2TestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); + OmniAccountV3 accountV3 = new OmniAccountV3(entryPoint); + + assertEq(account.passkeySignerCount(), 0); + + vm.prank(rootSigner); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV3), ""); + assertAccountVersion(address(account)); + } + + function testUnauthorizedSenderCannotUpgradeNonEvmOA() public { + (, EntryPointV1 entryPoint, OmniAccountV2 account) = + OmniAccountV2TestUtils.setUpWithOwnerType(owner, clientId, rootSigner, OwnerType.Substrate); + OmniAccountV3 accountV3 = new OmniAccountV3(entryPoint); + + vm.expectRevert("only owner"); + vm.prank(unauthorizedUser); + UUPSUpgradeable(account).upgradeToAndCall(address(accountV3), ""); + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/StorageTestModule.sol b/tee-worker/omni-executor/contracts/aa/test/v2/StorageTestModule.sol new file mode 100644 index 0000000000..68d8d49899 --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/StorageTestModule.sol @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +/** + * A module that writes to various storage slots to test storage integrity after upgrade. + * This module uses a namespaced storage pattern (EIP-7201) to avoid colliding with account storage. + * + * IMPORTANT: Modules MUST use namespaced storage to avoid corrupting account state! + * The account uses slots 0-8 for its core state, so we use a completely different namespace. + */ +contract StorageTestModule { + // Use a unique storage namespace to avoid collisions with account storage + bytes32 private constant MODULE_STORAGE_POSITION = keccak256("test.module.storage.v1"); + + struct TestStruct { + uint256 id; + string name; + uint256 value; + bool active; + } + + struct ModuleStorage { + uint256 simpleValue; + uint256[] dynamicArray; + mapping(address => uint256) addressToValue; + mapping(address => mapping(uint256 => uint256)) nestedMapping; + mapping(uint256 => TestStruct) structStorage; + string stringValue; + bytes bytesValue; + } + + // Events to track module execution + event ValueSet(string key, uint256 value); + event ArrayPushed(uint256 value); + event MappingSet(address key, uint256 value); + event NestedMappingSet(address key1, uint256 key2, uint256 value); + event StructSet(uint256 id, string name, uint256 value); + + function _getStorage() private pure returns (ModuleStorage storage ms) { + bytes32 position = MODULE_STORAGE_POSITION; + assembly { + ms.slot := position + } + } + + function setSimpleValue(uint256 _value) external { + ModuleStorage storage ms = _getStorage(); + ms.simpleValue = _value; + emit ValueSet("simpleValue", _value); + } + + function getSimpleValue() external view returns (uint256) { + return _getStorage().simpleValue; + } + + function pushToArray(uint256 _value) external { + ModuleStorage storage ms = _getStorage(); + ms.dynamicArray.push(_value); + emit ArrayPushed(_value); + } + + function getArrayLength() external view returns (uint256) { + return _getStorage().dynamicArray.length; + } + + function getArrayValue(uint256 index) external view returns (uint256) { + return _getStorage().dynamicArray[index]; + } + + function setMapping(address _key, uint256 _value) external { + ModuleStorage storage ms = _getStorage(); + ms.addressToValue[_key] = _value; + emit MappingSet(_key, _value); + } + + function getMapping(address _key) external view returns (uint256) { + return _getStorage().addressToValue[_key]; + } + + function setNestedMapping(address _key1, uint256 _key2, uint256 _value) external { + ModuleStorage storage ms = _getStorage(); + ms.nestedMapping[_key1][_key2] = _value; + emit NestedMappingSet(_key1, _key2, _value); + } + + function getNestedMapping(address _key1, uint256 _key2) external view returns (uint256) { + return _getStorage().nestedMapping[_key1][_key2]; + } + + function setStruct(uint256 _id, string calldata _name, uint256 _value) external { + ModuleStorage storage ms = _getStorage(); + ms.structStorage[_id] = TestStruct({id: _id, name: _name, value: _value, active: true}); + emit StructSet(_id, _name, _value); + } + + function getStruct(uint256 _id) external view returns (uint256 id, string memory name, uint256 value, bool active) { + TestStruct memory s = _getStorage().structStorage[_id]; + return (s.id, s.name, s.value, s.active); + } + + function setString(string calldata _value) external { + ModuleStorage storage ms = _getStorage(); + ms.stringValue = _value; + } + + function getString() external view returns (string memory) { + return _getStorage().stringValue; + } + + function setBytes(bytes calldata _value) external { + ModuleStorage storage ms = _getStorage(); + ms.bytesValue = _value; + } + + function getBytes() external view returns (bytes memory) { + return _getStorage().bytesValue; + } + + function complexOperation(uint256 _value1, uint256 _value2, address _addr) external returns (uint256) { + ModuleStorage storage ms = _getStorage(); + // Perform multiple storage operations + ms.simpleValue = _value1 + _value2; + ms.dynamicArray.push(_value1); + ms.dynamicArray.push(_value2); + ms.addressToValue[_addr] = _value1 * _value2; + ms.nestedMapping[_addr][_value1] = _value2; + + return ms.simpleValue; + } + + function incrementSimpleValue() external returns (uint256) { + ModuleStorage storage ms = _getStorage(); + ms.simpleValue++; + return ms.simpleValue; + } +} diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/TestUtilsV2.sol b/tee-worker/omni-executor/contracts/aa/test/v2/TestUtilsV2.sol new file mode 100644 index 0000000000..f67efa0f9c --- /dev/null +++ b/tee-worker/omni-executor/contracts/aa/test/v2/TestUtilsV2.sol @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {PackedUserOperation} from "../../src/v2/interfaces/PackedUserOperation.sol"; + +library TestUtilsV2 { + function prepare_evm_oa(address account, bytes memory clientId) public pure returns (bytes32) { + // bytes("evm"); + bytes3 oaType = 0x65766d; + return sha256(abi.encodePacked(clientId, oaType, account)); + } + + function preparePackedOp(address sender, bytes memory initCode) internal pure returns (PackedUserOperation memory) { + uint256 nonce = 0; + bytes memory callData = ""; + bytes32 accountGasLimits = 0x0000000000000000000000000004e20000000000000000000000000000005b8d; + uint256 preVerificationGas = 21000; + bytes32 gasFees = 0x0000000000000000000000003b9aca00000000000000000000000000b2d05e00; + bytes memory paymasterAndData = ""; + bytes memory signature = ""; + + return (PackedUserOperation( + sender, + nonce, + initCode, + callData, + accountGasLimits, + preVerificationGas, + gasFees, + paymasterAndData, + signature + )); + } +}