From c789b7fd8953d23206e4fc0cc4c6e4191b624c31 Mon Sep 17 00:00:00 2001 From: Kasper Ziemianek Date: Tue, 4 Nov 2025 10:41:27 +0100 Subject: [PATCH 1/7] OmniAccountV2 with module system --- .../src/accounts/OmniAccountFactoryV2.sol | 65 ++ .../src/accounts/OmniAccountV2.sol | 284 +++++++++ .../src/core/LibModuleManager.sol | 39 ++ .../test/{ => v1}/OmniAccount.t.sol | 16 +- .../{ => v1}/OmniAccountAsEntryPoint.t.sol | 8 +- .../test/{ => v1}/OmniAccountAsOwner.t.sol | 18 +- .../test/{ => v1}/OmniAccountAsPasskey.t.sol | 18 +- .../test/{ => v1}/OmniAccountAsRoot.t.sol | 18 +- .../{ => v1}/OmniAccountAsRootNonEvm.t.sol | 20 +- ...iAccountNonEvmOwnerWithPasskeySigner.t.sol | 20 +- .../test/{ => v1}/OmniAccountTestUtils.sol | 12 +- .../{ => v1}/OmniAccountUpgradability.t.sol | 10 +- .../aa-contracts/test/v2/MockModule.sol | 53 ++ .../test/v2/OmniAccountV1ToV2Upgrade.t.sol | 348 +++++++++++ .../aa-contracts/test/v2/OmniAccountV2.t.sol | 83 +++ .../test/v2/OmniAccountV2AsEntryPoint.t.sol | 129 ++++ .../test/v2/OmniAccountV2AsOwner.t.sol | 561 ++++++++++++++++++ .../test/v2/OmniAccountV2AsPasskey.t.sol | 156 +++++ .../test/v2/OmniAccountV2AsRoot.t.sol | 540 +++++++++++++++++ .../test/v2/OmniAccountV2AsRootNonEvm.t.sol | 455 ++++++++++++++ .../test/v2/OmniAccountV2Modules.t.sol | 111 ++++ ...ccountV2NonEvmOwnerWithPasskeySigner.t.sol | 279 +++++++++ .../test/v2/OmniAccountV2TestUtils.sol | 69 +++ .../test/v2/OmniAccountV2Upgradability.t.sol | 116 ++++ .../test/v2/StorageTestModule.sol | 141 +++++ 25 files changed, 3499 insertions(+), 70 deletions(-) create mode 100644 tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountFactoryV2.sol create mode 100644 tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountV2.sol create mode 100644 tee-worker/omni-executor/aa-contracts/src/core/LibModuleManager.sol rename tee-worker/omni-executor/aa-contracts/test/{ => v1}/OmniAccount.t.sol (84%) rename tee-worker/omni-executor/aa-contracts/test/{ => v1}/OmniAccountAsEntryPoint.t.sol (79%) rename tee-worker/omni-executor/aa-contracts/test/{ => v1}/OmniAccountAsOwner.t.sol (97%) rename tee-worker/omni-executor/aa-contracts/test/{ => v1}/OmniAccountAsPasskey.t.sol (92%) rename tee-worker/omni-executor/aa-contracts/test/{ => v1}/OmniAccountAsRoot.t.sol (97%) rename tee-worker/omni-executor/aa-contracts/test/{ => v1}/OmniAccountAsRootNonEvm.t.sol (95%) rename tee-worker/omni-executor/aa-contracts/test/{ => v1}/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol (94%) rename tee-worker/omni-executor/aa-contracts/test/{ => v1}/OmniAccountTestUtils.sol (87%) rename tee-worker/omni-executor/aa-contracts/test/{ => v1}/OmniAccountUpgradability.t.sol (93%) create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/MockModule.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV1ToV2Upgrade.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsOwner.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsPasskey.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRoot.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRootNonEvm.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Modules.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2TestUtils.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Upgradability.t.sol create mode 100644 tee-worker/omni-executor/aa-contracts/test/v2/StorageTestModule.sol diff --git a/tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountFactoryV2.sol b/tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountFactoryV2.sol new file mode 100644 index 0000000000..6da8dacb33 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountFactoryV2.sol @@ -0,0 +1,65 @@ +// 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/aa-contracts/src/accounts/OmniAccountV2.sol b/tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountV2.sol new file mode 100644 index 0000000000..cec8d76695 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountV2.sol @@ -0,0 +1,284 @@ +// 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/aa-contracts/src/core/LibModuleManager.sol b/tee-worker/omni-executor/aa-contracts/src/core/LibModuleManager.sol new file mode 100644 index 0000000000..2a2f6374da --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/src/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/aa-contracts/test/OmniAccount.t.sol b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccount.t.sol similarity index 84% rename from tee-worker/omni-executor/aa-contracts/test/OmniAccount.t.sol rename to tee-worker/omni-executor/aa-contracts/test/v1/OmniAccount.t.sol index 01c8791c8e..fc85fdfdb9 100644 --- a/tee-worker/omni-executor/aa-contracts/test/OmniAccount.t.sol +++ b/tee-worker/omni-executor/aa-contracts/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/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 {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/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_FAILED} from "../../src//core/Helpers.sol"; // add test cases for revert if called by non authorized address diff --git a/tee-worker/omni-executor/aa-contracts/test/OmniAccountAsEntryPoint.t.sol b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsEntryPoint.t.sol similarity index 79% rename from tee-worker/omni-executor/aa-contracts/test/OmniAccountAsEntryPoint.t.sol rename to tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsEntryPoint.t.sol index 745d44c8f5..0f0fb3727d 100644 --- a/tee-worker/omni-executor/aa-contracts/test/OmniAccountAsEntryPoint.t.sol +++ b/tee-worker/omni-executor/aa-contracts/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/accounts/OmniAccountV1.sol"; +import {BaseAccount} from "../../src/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/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/aa-contracts/test/OmniAccountAsOwner.t.sol b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsOwner.t.sol similarity index 97% rename from tee-worker/omni-executor/aa-contracts/test/OmniAccountAsOwner.t.sol rename to tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsOwner.t.sol index 2e33c7adea..8961761e24 100644 --- a/tee-worker/omni-executor/aa-contracts/test/OmniAccountAsOwner.t.sol +++ b/tee-worker/omni-executor/aa-contracts/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/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 {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/interfaces/PackedUserOperation.sol"; +import {TestUtils} from "../TestUtils.sol"; +import {SIG_VALIDATION_SUCCESS} from "../../src//core/Helpers.sol"; +import {Passkey} from "../../src/interfaces/Passkey.sol"; contract OmniAccountAsOwner is Test { OmniAccountV1 public account; diff --git a/tee-worker/omni-executor/aa-contracts/test/OmniAccountAsPasskey.t.sol b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsPasskey.t.sol similarity index 92% rename from tee-worker/omni-executor/aa-contracts/test/OmniAccountAsPasskey.t.sol rename to tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsPasskey.t.sol index 5efafefb6b..db68925dd8 100644 --- a/tee-worker/omni-executor/aa-contracts/test/OmniAccountAsPasskey.t.sol +++ b/tee-worker/omni-executor/aa-contracts/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/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 {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/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src//core/Helpers.sol"; +import {Passkey} from "../../src/interfaces/Passkey.sol"; import "@openzeppelin/contracts/utils/Base64.sol"; contract OmniAccountAsPasskey is Test { diff --git a/tee-worker/omni-executor/aa-contracts/test/OmniAccountAsRoot.t.sol b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRoot.t.sol similarity index 97% rename from tee-worker/omni-executor/aa-contracts/test/OmniAccountAsRoot.t.sol rename to tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRoot.t.sol index c919312c12..bcf6205c9f 100644 --- a/tee-worker/omni-executor/aa-contracts/test/OmniAccountAsRoot.t.sol +++ b/tee-worker/omni-executor/aa-contracts/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/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 {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/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src//core/Helpers.sol"; +import {Passkey} from "../../src/interfaces/Passkey.sol"; contract OmniAccountAsRoot is Test { OmniAccountV1 public account; diff --git a/tee-worker/omni-executor/aa-contracts/test/OmniAccountAsRootNonEvm.t.sol b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRootNonEvm.t.sol similarity index 95% rename from tee-worker/omni-executor/aa-contracts/test/OmniAccountAsRootNonEvm.t.sol rename to tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRootNonEvm.t.sol index 76e9f09ed1..8ed8760154 100644 --- a/tee-worker/omni-executor/aa-contracts/test/OmniAccountAsRootNonEvm.t.sol +++ b/tee-worker/omni-executor/aa-contracts/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/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 {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/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/core/Helpers.sol"; +import {Passkey} from "../../src/interfaces/Passkey.sol"; contract OmniAccountAsRootNonEvm is Test { OmniAccount public account; diff --git a/tee-worker/omni-executor/aa-contracts/test/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol similarity index 94% rename from tee-worker/omni-executor/aa-contracts/test/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol rename to tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol index c2fd0a72d4..ad731619ff 100644 --- a/tee-worker/omni-executor/aa-contracts/test/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol +++ b/tee-worker/omni-executor/aa-contracts/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/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 {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/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/core/Helpers.sol"; contract OmniAccountNonEvmOwnerWithPasskeySigner is Test { OmniAccount public account; diff --git a/tee-worker/omni-executor/aa-contracts/test/OmniAccountTestUtils.sol b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountTestUtils.sol similarity index 87% rename from tee-worker/omni-executor/aa-contracts/test/OmniAccountTestUtils.sol rename to tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountTestUtils.sol index e618d31487..2aa5b17e9d 100644 --- a/tee-worker/omni-executor/aa-contracts/test/OmniAccountTestUtils.sol +++ b/tee-worker/omni-executor/aa-contracts/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/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 "@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) diff --git a/tee-worker/omni-executor/aa-contracts/test/OmniAccountUpgradability.t.sol b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountUpgradability.t.sol similarity index 93% rename from tee-worker/omni-executor/aa-contracts/test/OmniAccountUpgradability.t.sol rename to tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountUpgradability.t.sol index b881317ed9..e2a29c9eaf 100644 --- a/tee-worker/omni-executor/aa-contracts/test/OmniAccountUpgradability.t.sol +++ b/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountUpgradability.t.sol @@ -4,12 +4,12 @@ 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/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 {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"; +import {TestUtils} from "../TestUtils.sol"; contract OmniAccountV2 is OmniAccountV1 { constructor(EntryPointV1 anEntryPoint) OmniAccountV1(anEntryPoint) {} diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/MockModule.sol b/tee-worker/omni-executor/aa-contracts/test/v2/MockModule.sol new file mode 100644 index 0000000000..e8c42e8ba6 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/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/aa-contracts/test/v2/OmniAccountV1ToV2Upgrade.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV1ToV2Upgrade.t.sol new file mode 100644 index 0000000000..e2eb355a29 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV1ToV2Upgrade.t.sol @@ -0,0 +1,348 @@ +// 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/accounts/OmniAccountV1.sol"; +import {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; +import {OwnerType} from "../../src/interfaces/OwnerType.sol"; +import {Passkey} from "../../src/interfaces/Passkey.sol"; +import {TestUtils} from "../TestUtils.sol"; +import {StorageTestModule} from "./StorageTestModule.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 + 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(entryPoint); + expectedOwner = TestUtils.prepare_evm_oa(owner, clientId); + + accountV1 = OmniAccountV1( + payable( + new ERC1967Proxy{salt: expectedOwner}( + address(v1Implementation), + abi.encodeCall(OmniAccountV1.initialize, (expectedOwner, OwnerType.Evm, clientId, rootSigner1)) + ) + ) + ); + + // Prepare V2 implementation for upgrade + accountV2Implementation = new OmniAccountV2(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(passkey1); + vm.prank(owner); + accountV1.addPasskeySigner(passkey2); + + // 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(OwnerType.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(Passkey.toKey(passkey1)), "passkey1 not registered"); + assertTrue(accountV1.passkeySigners(Passkey.toKey(passkey2)), "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(OwnerType.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(Passkey.toKey(passkey1)), "passkey1 corrupted"); + assertTrue(accountV2.passkeySigners(Passkey.toKey(passkey2)), "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); + accountV2.removePasskeySigner(passkey1); + 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(entryPoint); + bytes32 substrateOwner = keccak256("substrate_owner"); + + OmniAccountV1 substrateAccount = OmniAccountV1( + payable( + new ERC1967Proxy{salt: substrateOwner}( + address(v1Implementation), + abi.encodeCall(OmniAccountV1.initialize, (substrateOwner, OwnerType.Substrate, clientId, rootSigner1)) + ) + ) + ); + + // Setup state + vm.prank(rootSigner1); + substrateAccount.addRootSigner(rootSigner2); + + // Verify V1 state + assertEq(substrateAccount.owner(), substrateOwner); + assertEq(uint256(substrateAccount.ownerType()), uint256(OwnerType.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/aa-contracts/test/v2/OmniAccountV2.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2.t.sol new file mode 100644 index 0000000000..a36b9fa8ba --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/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/accounts/OmniAccountV2.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 {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {TestUtils} from "../TestUtils.sol"; +import {PackedUserOperation} from "../../src/interfaces/PackedUserOperation.sol"; +import {SIG_VALIDATION_FAILED} from "../../src//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/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol new file mode 100644 index 0000000000..0dede7b16f --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/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/accounts/OmniAccountV2.sol"; +import {BaseAccount} from "../../src/core/BaseAccount.sol"; +import {EntryPointV1} from "../../src/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)); + bytes memory returnData = 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/aa-contracts/test/v2/OmniAccountV2AsOwner.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsOwner.t.sol new file mode 100644 index 0000000000..1ddc70afff --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/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/accounts/OmniAccountV2.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 {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.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 {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/aa-contracts/test/v2/OmniAccountV2AsPasskey.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsPasskey.t.sol new file mode 100644 index 0000000000..595506080e --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/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/accounts/OmniAccountV2.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 {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.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 "@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/aa-contracts/test/v2/OmniAccountV2AsRoot.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRoot.t.sol new file mode 100644 index 0000000000..bb310db5d8 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/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/accounts/OmniAccountV2.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 {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.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 {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/aa-contracts/test/v2/OmniAccountV2AsRootNonEvm.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRootNonEvm.t.sol new file mode 100644 index 0000000000..a1bda36a68 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/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/accounts/OmniAccountV2.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 {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.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 {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/aa-contracts/test/v2/OmniAccountV2Modules.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Modules.t.sol new file mode 100644 index 0000000000..cee97cf346 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Modules.t.sol @@ -0,0 +1,111 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {EntryPointV1} from "../../src/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/aa-contracts/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol new file mode 100644 index 0000000000..1e12e984f0 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/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/accounts/OmniAccountV2.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 {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.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"; + +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/aa-contracts/test/v2/OmniAccountV2TestUtils.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2TestUtils.sol new file mode 100644 index 0000000000..00e8f74814 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2TestUtils.sol @@ -0,0 +1,69 @@ +// 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/accounts/OmniAccountV2.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 "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; +import {TestUtils} from "../TestUtils.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(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(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/aa-contracts/test/v2/OmniAccountV2Upgradability.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Upgradability.t.sol new file mode 100644 index 0000000000..ccb4364d49 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Upgradability.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 {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; +import {OwnerType} from "../../src/interfaces/OwnerType.sol"; +import {Passkey} from "../../src/interfaces/Passkey.sol"; +import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; +import {TestUtils} from "../TestUtils.sol"; + +contract OmniAccountV3 is OmniAccountV2 { + constructor(EntryPointV1 anEntryPoint) OmniAccountV2(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/aa-contracts/test/v2/StorageTestModule.sol b/tee-worker/omni-executor/aa-contracts/test/v2/StorageTestModule.sol new file mode 100644 index 0000000000..37265a71d8 --- /dev/null +++ b/tee-worker/omni-executor/aa-contracts/test/v2/StorageTestModule.sol @@ -0,0 +1,141 @@ +// 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; + } +} From c687d05d88d7da7783137f3814299fca8d9bf721 Mon Sep 17 00:00:00 2001 From: Kasper Ziemianek Date: Tue, 4 Nov 2025 11:56:18 +0100 Subject: [PATCH 2/7] remove warning --- .../aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol index 0dede7b16f..20f5d19529 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol +++ b/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol @@ -69,7 +69,7 @@ contract OmniAccountV2AsEntryPoint is Test { bytes memory callData = abi.encodeWithSignature("increment()"); vm.prank(address(entryPoint)); - bytes memory returnData = account.executeModuleCall(address(module), callData); + 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 From 586d5136ee273dd276c224ccefdc085d25a9ecb0 Mon Sep 17 00:00:00 2001 From: Kasper Ziemianek Date: Tue, 4 Nov 2025 12:15:41 +0100 Subject: [PATCH 3/7] specify forge version --- .github/workflows/ci.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9d7642aa0e..9dbccab37c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -605,6 +605,8 @@ jobs: - name: Install foundry uses: foundry-rs/foundry-toolchain@v1 + with: + version: 'v1.2.3' - name: Install solana sdk run: | From 60d3fe782ae1428cd5ecc869f673a896b12e4c0c Mon Sep 17 00:00:00 2001 From: Kasper Ziemianek Date: Wed, 17 Dec 2025 10:43:13 +0100 Subject: [PATCH 4/7] adjust to project structure update + foundry.lock --- .../omni-executor/contracts/aa/README.md | 4 + .../omni-executor/contracts/aa/foundry.lock | 8 + .../contracts/aa/script/Deploy.s.sol | 2 +- .../contracts/aa/script/DeployLocal.s.sol | 2 +- .../aa/script/DeployLocalWithPaymaster.s.sol | 2 +- .../contracts/aa/script/DeploymentHelper.sol | 31 +- .../{ => v1}/OmniAccountFactoryV1.sol | 10 +- .../src/accounts/{ => v1}/OmniAccountV1.sol | 15 +- .../src/accounts/v2}/OmniAccountFactoryV2.sol | 20 +- .../aa/src/accounts/v2}/OmniAccountV2.sol | 51 +- .../contracts/aa/src/core/DemoPaymaster.sol | 17 +- .../aa/src/core/ERC20PaymasterV1.sol | 35 +- .../aa/src/core/EntryPointSimulations.sol | 5 +- .../contracts/aa/src/core/EntryPointV1.sol | 26 +- .../aa}/src/core/LibModuleManager.sol | 0 .../contracts/aa/src/core/SimplePaymaster.sol | 17 +- .../aa/src/interfaces/IEntryPoint.sol | 3 +- .../contracts/aa/src/interfaces/OwnerType.sol | 1 - .../aa/src/interfaces/UserOpSigner.sol | 1 - .../contracts/aa/test/DeploymentHelper.t.sol | 2 +- .../contracts/aa/test/EntryPoint.t.sol | 2 +- .../contracts/aa/test/OmniAccount.t.sol | 83 ---- .../aa/test/OmniAccountAsEntryPoint.t.sol | 31 -- .../aa/test/OmniAccountAsOwner.t.sol | 443 ----------------- .../aa/test/OmniAccountAsPasskey.t.sol | 156 ------ .../contracts/aa/test/OmniAccountAsRoot.t.sol | 470 ------------------ .../aa/test/OmniAccountAsRootNonEvm.t.sol | 305 ------------ .../aa/test/OmniAccountFactory.t.sol | 34 -- ...iAccountNonEvmOwnerWithPasskeySigner.t.sol | 279 ----------- .../aa/test/OmniAccountTestUtils.sol | 69 --- .../aa/test/OmniAccountUpgradability.t.sol | 363 -------------- .../contracts/aa/test/TestUtils.sol | 12 +- .../aa}/test/v1/OmniAccount.t.sol | 2 +- .../aa}/test/v1/OmniAccountAsEntryPoint.t.sol | 2 +- .../aa}/test/v1/OmniAccountAsOwner.t.sol | 2 +- .../aa}/test/v1/OmniAccountAsPasskey.t.sol | 2 +- .../aa}/test/v1/OmniAccountAsRoot.t.sol | 2 +- .../aa}/test/v1/OmniAccountAsRootNonEvm.t.sol | 2 +- ...iAccountNonEvmOwnerWithPasskeySigner.t.sol | 2 +- .../aa}/test/v1/OmniAccountTestUtils.sol | 14 +- .../test/v1/OmniAccountUpgradability.t.sol | 2 +- .../aa}/test/v2/MockModule.sol | 0 .../test/v2/OmniAccountV1ToV2Upgrade.t.sol | 71 +-- .../aa}/test/v2/OmniAccountV2.t.sol | 4 +- .../test/v2/OmniAccountV2AsEntryPoint.t.sol | 2 +- .../aa}/test/v2/OmniAccountV2AsOwner.t.sol | 4 +- .../aa}/test/v2/OmniAccountV2AsPasskey.t.sol | 4 +- .../aa}/test/v2/OmniAccountV2AsRoot.t.sol | 4 +- .../test/v2/OmniAccountV2AsRootNonEvm.t.sol | 2 +- .../aa}/test/v2/OmniAccountV2Modules.t.sol | 3 +- ...ccountV2NonEvmOwnerWithPasskeySigner.t.sol | 2 +- .../aa}/test/v2/OmniAccountV2TestUtils.sol | 14 +- .../test/v2/OmniAccountV2Upgradability.t.sol | 2 +- .../aa}/test/v2/StorageTestModule.sol | 7 +- 54 files changed, 198 insertions(+), 2450 deletions(-) create mode 100644 tee-worker/omni-executor/contracts/aa/foundry.lock rename tee-worker/omni-executor/contracts/aa/src/accounts/{ => v1}/OmniAccountFactoryV1.sol (93%) rename tee-worker/omni-executor/contracts/aa/src/accounts/{ => v1}/OmniAccountV1.sol (96%) rename tee-worker/omni-executor/{aa-contracts/src/accounts => contracts/aa/src/accounts/v2}/OmniAccountFactoryV2.sol (87%) rename tee-worker/omni-executor/{aa-contracts/src/accounts => contracts/aa/src/accounts/v2}/OmniAccountV2.sol (89%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/src/core/LibModuleManager.sol (100%) delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccount.t.sol delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccountAsEntryPoint.t.sol delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccountAsOwner.t.sol delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccountAsPasskey.t.sol delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRoot.t.sol delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRootNonEvm.t.sol delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccountFactory.t.sol delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccountTestUtils.sol delete mode 100644 tee-worker/omni-executor/contracts/aa/test/OmniAccountUpgradability.t.sol rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v1/OmniAccount.t.sol (97%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v1/OmniAccountAsEntryPoint.t.sol (93%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v1/OmniAccountAsOwner.t.sol (99%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v1/OmniAccountAsPasskey.t.sol (99%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v1/OmniAccountAsRoot.t.sol (99%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v1/OmniAccountAsRootNonEvm.t.sol (99%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol (99%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v1/OmniAccountTestUtils.sol (90%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v1/OmniAccountUpgradability.t.sol (98%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/MockModule.sol (100%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV1ToV2Upgrade.t.sol (87%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2.t.sol (95%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2AsEntryPoint.t.sol (98%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2AsOwner.t.sol (99%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2AsPasskey.t.sol (98%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2AsRoot.t.sol (99%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2AsRootNonEvm.t.sol (99%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2Modules.t.sol (98%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol (99%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2TestUtils.sol (90%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/OmniAccountV2Upgradability.t.sol (98%) rename tee-worker/omni-executor/{aa-contracts => contracts/aa}/test/v2/StorageTestModule.sol (96%) diff --git a/tee-worker/omni-executor/contracts/aa/README.md b/tee-worker/omni-executor/contracts/aa/README.md index 6241e528f7..fbba98e03e 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 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..5e38da29fd 100644 --- a/tee-worker/omni-executor/contracts/aa/script/Deploy.s.sol +++ b/tee-worker/omni-executor/contracts/aa/script/Deploy.s.sol @@ -4,7 +4,7 @@ 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/accounts/v1/OmniAccountFactoryV1.sol"; import "../src/core/SimplePaymaster.sol"; import "../src/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..f366fcd955 100644 --- a/tee-worker/omni-executor/contracts/aa/script/DeployLocal.s.sol +++ b/tee-worker/omni-executor/contracts/aa/script/DeployLocal.s.sol @@ -3,7 +3,7 @@ pragma solidity ^0.8.28; import "forge-std/Script.sol"; import "../src/core/EntryPointV1.sol"; -import "../src/accounts/OmniAccountFactoryV1.sol"; +import "../src/accounts/v1/OmniAccountFactoryV1.sol"; import "../src/core/SimplePaymaster.sol"; import "../src/TestToken.sol"; 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..e65ae0ef1e 100644 --- a/tee-worker/omni-executor/contracts/aa/script/DeployLocalWithPaymaster.s.sol +++ b/tee-worker/omni-executor/contracts/aa/script/DeployLocalWithPaymaster.s.sol @@ -4,7 +4,7 @@ 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/accounts/v1/OmniAccountFactoryV1.sol"; import "../src/core/SimplePaymaster.sol"; import "../src/core/DemoPaymaster.sol"; import "../src/core/ERC20PaymasterV1.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/accounts/v1/OmniAccountFactoryV1.sol similarity index 93% rename from tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountFactoryV1.sol rename to tee-worker/omni-executor/contracts/aa/src/accounts/v1/OmniAccountFactoryV1.sol index 028474a1f5..acc7fafeb6 100644 --- a/tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountFactoryV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/accounts/v1/OmniAccountFactoryV1.sol @@ -4,8 +4,8 @@ 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 "../../interfaces/ISenderCreator.sol"; +import "../../interfaces/OwnerType.sol"; import "./OmniAccountV1.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/accounts/v1/OmniAccountV1.sol similarity index 96% rename from tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountV1.sol rename to tee-worker/omni-executor/contracts/aa/src/accounts/v1/OmniAccountV1.sol index 68bdbadbe7..99ed5505ca 100644 --- a/tee-worker/omni-executor/contracts/aa/src/accounts/OmniAccountV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/accounts/v1/OmniAccountV1.sol @@ -10,12 +10,12 @@ 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 "../../core/BaseAccount.sol"; +import "../../interfaces/OwnerType.sol"; +import "../../interfaces/UserOpSigner.sol"; +import "../../interfaces/Passkey.sol"; +import "../../core/Helpers.sol"; +import "./../callback/TokenCallbackHandler.sol"; /** * smart wallet account @@ -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/aa-contracts/src/accounts/OmniAccountFactoryV2.sol b/tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountFactoryV2.sol similarity index 87% rename from tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountFactoryV2.sol rename to tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountFactoryV2.sol index 6da8dacb33..7a50b0b527 100644 --- a/tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountFactoryV2.sol +++ b/tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountFactoryV2.sol @@ -4,8 +4,8 @@ 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 "../../interfaces/ISenderCreator.sol"; +import "../../interfaces/OwnerType.sol"; import "./OmniAccountV2.sol"; contract OmniAccountFactoryV2 { @@ -20,8 +20,8 @@ contract OmniAccountFactoryV2 { } function createAccount(bytes32 oa, OwnerType oaType, bytes memory clientId, address root) - public - returns (OmniAccountV2 ret) + public + returns (OmniAccountV2 ret) { require(msg.sender == address(senderCreator), "only callable from SenderCreator"); address addr = getAddress(oa, oaType, clientId, root); @@ -30,20 +30,18 @@ contract OmniAccountFactoryV2 { return OmniAccountV2(payable(addr)); } ret = OmniAccountV2( - payable( - new ERC1967Proxy{salt: oa}( + 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) + public + view + returns (address) { return Create2.computeAddress( oa, diff --git a/tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountV2.sol b/tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountV2.sol similarity index 89% rename from tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountV2.sol rename to tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountV2.sol index cec8d76695..6d1add472b 100644 --- a/tee-worker/omni-executor/aa-contracts/src/accounts/OmniAccountV2.sol +++ b/tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountV2.sol @@ -10,14 +10,14 @@ 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"; +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; @@ -36,7 +36,8 @@ contract OmniAccountV2 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)")); bytes4 private constant REGISTER_MODULE_SELECTOR = bytes4(keccak256("registerModule(address)")); @@ -71,22 +72,22 @@ contract OmniAccountV2 is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, In function _onlyOwner() internal view { require( _determineOa(msg.sender) == owner || msg.sender == address(entryPoint()) - || (ownerType != OwnerType.Evm && passkeySignerCount == 0 && isRootSigner(msg.sender)), + || (ownerType != OwnerType.Evm && passkeySignerCount == 0 && isRootSigner(msg.sender)), "only owner" ); } function initialize(bytes32 anOwner, OwnerType anOwnerType, bytes memory aClientId, address aRoot) - public - virtual - initializer + public + virtual + initializer { _initialize(anOwner, anOwnerType, aClientId, aRoot); } function _initialize(bytes32 anOwner, OwnerType anOwnerType, bytes memory aClientId, address aRoot) - internal - virtual + internal + virtual { owner = anOwner; rootSigners[aRoot] = true; @@ -105,10 +106,10 @@ contract OmniAccountV2 is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, In } function _validateSignature(PackedUserOperation calldata userOp, bytes32 userOpHash) - internal - virtual - override - returns (uint256 validationData) + internal + virtual + override + returns (uint256 validationData) { require(userOp.signature.length >= 1, "signature too short"); @@ -156,9 +157,9 @@ contract OmniAccountV2 is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, In } function _validateSessionKey(bytes32 userOpHash, bytes calldata sig) - internal - view - returns (uint256 validationData) + internal + view + returns (uint256 validationData) { require(sig.length == 162, "SessionKey signature length invalid"); bytes memory sessionSig = sig[:65]; @@ -247,9 +248,9 @@ contract OmniAccountV2 is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, In 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; + || 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 { diff --git a/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol index 7f6626adcf..70533e3a5d 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol +++ b/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol @@ -24,12 +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) - internal - view - override - returns (bytes memory context, uint256 validationData) - { + 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) { @@ -51,7 +51,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/core/ERC20PaymasterV1.sol index f1f54cad3f..a042324e24 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/ERC20PaymasterV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/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,11 +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) - internal - override - returns (bytes memory context, uint256 validationData) - { + 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(); @@ -207,11 +208,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 +229,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/EntryPointSimulations.sol b/tee-worker/omni-executor/contracts/aa/src/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/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/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/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/aa-contracts/src/core/LibModuleManager.sol b/tee-worker/omni-executor/contracts/aa/src/core/LibModuleManager.sol similarity index 100% rename from tee-worker/omni-executor/aa-contracts/src/core/LibModuleManager.sol rename to tee-worker/omni-executor/contracts/aa/src/core/LibModuleManager.sol diff --git a/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol index 4ecad4aead..dd61552b84 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol +++ b/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol @@ -28,12 +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) - internal - view - override - returns (bytes memory context, uint256 validationData) - { + 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 @@ -61,7 +61,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/interfaces/IEntryPoint.sol b/tee-worker/omni-executor/contracts/aa/src/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/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/OwnerType.sol b/tee-worker/omni-executor/contracts/aa/src/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/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/UserOpSigner.sol b/tee-worker/omni-executor/contracts/aa/src/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/interfaces/UserOpSigner.sol @@ -22,5 +22,4 @@ enum UserOpSigner { RootKey, // 0x01 SessionKey, // 0x02 Passkey // 0x03 - } 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..5bab6b4aa3 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,7 @@ 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 + try vm.createDir(testDir, true) {} catch { // Allow creation to fail if directory exists} 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/EntryPoint.t.sol b/tee-worker/omni-executor/contracts/aa/test/EntryPoint.t.sol index 0d44a1e3bb..01af3b0b8c 100644 --- a/tee-worker/omni-executor/contracts/aa/test/EntryPoint.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/EntryPoint.t.sol @@ -4,7 +4,7 @@ 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 {OmniAccountFactoryV1} from "../src/accounts/v1/OmniAccountFactoryV1.sol"; import {PackedUserOperation} from "../src/interfaces/PackedUserOperation.sol"; import {OwnerType} from "../src/interfaces/OwnerType.sol"; import {TestUtils} from "./TestUtils.sol"; diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccount.t.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccount.t.sol deleted file mode 100644 index 01c8791c8e..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccount.t.sol +++ /dev/null @@ -1,83 +0,0 @@ -// SPDX-License-Identifier: GPL-3.0 -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 {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"; - -// add test cases for revert if called by non authorized address - -contract OmniAccountTest is Test { - OmniAccountV1 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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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/OmniAccountAsEntryPoint.t.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsEntryPoint.t.sol deleted file mode 100644 index 745d44c8f5..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsEntryPoint.t.sol +++ /dev/null @@ -1,31 +0,0 @@ -// SPDX-License-Identifier: GPL-3.0 -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 {OmniAccountTestUtils} from "./OmniAccountTestUtils.sol"; - -contract OmniAccountAsEntryPoint is Test { - OmniAccountV1 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) = OmniAccountTestUtils.setUp(ownerAddress, clientId, rootAddress); - } - - function test_Execute() public { - OmniAccountTestUtils.performExecuteTestAs(vm, address(entryPoint), account, counter); - } - - function test_ExecuteBatch() public { - OmniAccountTestUtils.performExecuteBatchTestAs(vm, address(entryPoint), account, counter); - } -} diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsOwner.t.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsOwner.t.sol deleted file mode 100644 index 2e33c7adea..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsOwner.t.sol +++ /dev/null @@ -1,443 +0,0 @@ -// SPDX-License-Identifier: GPL-3.0 -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 {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"; - -contract OmniAccountAsOwner is Test { - OmniAccountV1 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) = OmniAccountTestUtils.setUp(ownerAddress, clientId, rootAddress); - } - - 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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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)); - } -} diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsPasskey.t.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsPasskey.t.sol deleted file mode 100644 index 5efafefb6b..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsPasskey.t.sol +++ /dev/null @@ -1,156 +0,0 @@ -// SPDX-License-Identifier: GPL-3.0 -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 {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 "@openzeppelin/contracts/utils/Base64.sol"; - -contract OmniAccountAsPasskey is Test { - OmniAccountV1 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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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/OmniAccountAsRoot.t.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRoot.t.sol deleted file mode 100644 index c919312c12..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRoot.t.sol +++ /dev/null @@ -1,470 +0,0 @@ -// SPDX-License-Identifier: GPL-3.0 -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 {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"; - -contract OmniAccountAsRoot is Test { - OmniAccountV1 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) = OmniAccountTestUtils.setUp(ownerAddress, clientId, rootAddress); - } - - 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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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) = OmniAccountTestUtils.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), ""); - } -} diff --git a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRootNonEvm.t.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRootNonEvm.t.sol deleted file mode 100644 index 76e9f09ed1..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountAsRootNonEvm.t.sol +++ /dev/null @@ -1,305 +0,0 @@ -// SPDX-License-Identifier: GPL-3.0 -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 {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"; - -contract OmniAccountAsRootNonEvm 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) = - OmniAccountTestUtils.setUpWithOwnerType(ownerAddress, clientId, rootAddress, OwnerType.Email); - } - - 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) = - OmniAccountTestUtils.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) = - OmniAccountTestUtils.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) = - OmniAccountTestUtils.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) = - OmniAccountTestUtils.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) = - OmniAccountTestUtils.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) = - OmniAccountTestUtils.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); - } -} 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/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol deleted file mode 100644 index c2fd0a72d4..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol +++ /dev/null @@ -1,279 +0,0 @@ -// SPDX-License-Identifier: GPL-3.0 -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 {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"; - -contract OmniAccountNonEvmOwnerWithPasskeySigner 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) = - OmniAccountTestUtils.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/OmniAccountTestUtils.sol b/tee-worker/omni-executor/contracts/aa/test/OmniAccountTestUtils.sol deleted file mode 100644 index e618d31487..0000000000 --- a/tee-worker/omni-executor/contracts/aa/test/OmniAccountTestUtils.sol +++ /dev/null @@ -1,69 +0,0 @@ -// 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 {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 "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; -import {TestUtils} from "./TestUtils.sol"; - -library OmniAccountTestUtils { - 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(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(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/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/TestUtils.sol b/tee-worker/omni-executor/contracts/aa/test/TestUtils.sol index df58e3e9c4..ba958a02a2 100644 --- a/tee-worker/omni-executor/contracts/aa/test/TestUtils.sol +++ b/tee-worker/omni-executor/contracts/aa/test/TestUtils.sol @@ -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/aa-contracts/test/v1/OmniAccount.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccount.t.sol similarity index 97% rename from tee-worker/omni-executor/aa-contracts/test/v1/OmniAccount.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccount.t.sol index fc85fdfdb9..1c3ecc7dda 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccount.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccount.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../../src/accounts/OmniAccountV1.sol"; +import {OmniAccountV1} from "../../src/accounts/v1/OmniAccountV1.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsEntryPoint.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsEntryPoint.t.sol similarity index 93% rename from tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsEntryPoint.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsEntryPoint.t.sol index 0f0fb3727d..cdf8393771 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsEntryPoint.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsEntryPoint.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../../src/accounts/OmniAccountV1.sol"; +import {OmniAccountV1} from "../../src/accounts/v1/OmniAccountV1.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {Counter} from "../../src/Counter.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsOwner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsOwner.t.sol similarity index 99% rename from tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsOwner.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsOwner.t.sol index 8961761e24..564f8b49d4 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsOwner.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsOwner.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../../src/accounts/OmniAccountV1.sol"; +import {OmniAccountV1} from "../../src/accounts/v1/OmniAccountV1.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsPasskey.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsPasskey.t.sol similarity index 99% rename from tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsPasskey.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsPasskey.t.sol index db68925dd8..3c56700712 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsPasskey.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsPasskey.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../../src/accounts/OmniAccountV1.sol"; +import {OmniAccountV1} from "../../src/accounts/v1/OmniAccountV1.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRoot.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRoot.t.sol similarity index 99% rename from tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRoot.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRoot.t.sol index bcf6205c9f..c262e2e0d3 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRoot.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRoot.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../../src/accounts/OmniAccountV1.sol"; +import {OmniAccountV1} from "../../src/accounts/v1/OmniAccountV1.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRootNonEvm.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRootNonEvm.t.sol similarity index 99% rename from tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRootNonEvm.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRootNonEvm.t.sol index 8ed8760154..cf73f47a6c 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountAsRootNonEvm.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRootNonEvm.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1 as OmniAccount} from "../../src/accounts/OmniAccountV1.sol"; +import {OmniAccountV1 as OmniAccount} from "../../src/accounts/v1/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"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol similarity index 99% rename from tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol index ad731619ff..6a4c512327 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV1 as OmniAccount} from "../../src/accounts/OmniAccountV1.sol"; +import {OmniAccountV1 as OmniAccount} from "../../src/accounts/v1/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"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountTestUtils.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountTestUtils.sol similarity index 90% rename from tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountTestUtils.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountTestUtils.sol index 2aa5b17e9d..31ae846e16 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountTestUtils.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountTestUtils.sol @@ -3,7 +3,7 @@ pragma solidity ^0.8.28; import {Counter} from "../../src/Counter.sol"; import {Vm} from "forge-std/Vm.sol"; -import {OmniAccountV1 as OmniAccount} from "../../src/accounts/OmniAccountV1.sol"; +import {OmniAccountV1 as OmniAccount} from "../../src/accounts/v1/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"; @@ -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/aa-contracts/test/v1/OmniAccountUpgradability.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountUpgradability.t.sol similarity index 98% rename from tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountUpgradability.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountUpgradability.t.sol index e2a29c9eaf..ef7df39766 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v1/OmniAccountUpgradability.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountUpgradability.t.sol @@ -4,7 +4,7 @@ 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/accounts/OmniAccountV1.sol"; +import {OmniAccountV1} from "../../src/accounts/v1/OmniAccountV1.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {OwnerType} from "../../src/interfaces/OwnerType.sol"; import {Passkey} from "../../src/interfaces/Passkey.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/MockModule.sol b/tee-worker/omni-executor/contracts/aa/test/v2/MockModule.sol similarity index 100% rename from tee-worker/omni-executor/aa-contracts/test/v2/MockModule.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/MockModule.sol diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV1ToV2Upgrade.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV1ToV2Upgrade.t.sol similarity index 87% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV1ToV2Upgrade.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV1ToV2Upgrade.t.sol index e2eb355a29..4dd5ec3dd4 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV1ToV2Upgrade.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV1ToV2Upgrade.t.sol @@ -4,8 +4,8 @@ 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/accounts/OmniAccountV1.sol"; -import {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV1} from "../../src/accounts/v1/OmniAccountV1.sol"; +import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {OwnerType} from "../../src/interfaces/OwnerType.sol"; import {Passkey} from "../../src/interfaces/Passkey.sol"; @@ -51,12 +51,10 @@ contract OmniAccountV1ToV2Upgrade is Test { expectedOwner = TestUtils.prepare_evm_oa(owner, clientId); accountV1 = OmniAccountV1( - payable( - new ERC1967Proxy{salt: expectedOwner}( + payable(new ERC1967Proxy{salt: expectedOwner}( address(v1Implementation), abi.encodeCall(OmniAccountV1.initialize, (expectedOwner, OwnerType.Evm, clientId, rootSigner1)) - ) - ) + )) ); // Prepare V2 implementation for upgrade @@ -130,27 +128,17 @@ contract OmniAccountV1ToV2Upgrade is Test { // Execute module functions that write to various storage slots vm.prank(address(entryPoint)); - accountV2.executeModuleCall( - address(module), - abi.encodeWithSignature("setSimpleValue(uint256)", 12345) - ); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("setSimpleValue(uint256)", 12345)); vm.prank(address(entryPoint)); - accountV2.executeModuleCall( - address(module), - abi.encodeWithSignature("pushToArray(uint256)", 111) - ); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("pushToArray(uint256)", 111)); vm.prank(address(entryPoint)); - accountV2.executeModuleCall( - address(module), - abi.encodeWithSignature("pushToArray(uint256)", 222) - ); + 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) + address(module), abi.encodeWithSignature("setMapping(address,uint256)", address(0x5555), 9999) ); vm.prank(address(entryPoint)); @@ -161,21 +149,14 @@ contract OmniAccountV1ToV2Upgrade is Test { vm.prank(address(entryPoint)); accountV2.executeModuleCall( - address(module), - abi.encodeWithSignature("setStruct(uint256,string,uint256)", 1, "test_struct", 8888) + 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") - ); + 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") - ); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("setBytes(bytes)", hex"deadbeef")); // Complex operation with multiple storage writes vm.prank(address(entryPoint)); @@ -242,19 +223,13 @@ contract OmniAccountV1ToV2Upgrade is Test { // Increment simple value multiple times for (uint256 i = 0; i < 5; i++) { vm.prank(address(entryPoint)); - accountV2.executeModuleCall( - address(module), - abi.encodeWithSignature("incrementSimpleValue()") - ); + 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) - ); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("pushToArray(uint256)", 1000 + i)); } // ============ STEP 8: Final State Verification ============ @@ -286,10 +261,7 @@ contract OmniAccountV1ToV2Upgrade is Test { vm.expectRevert("Module not registered"); vm.prank(address(entryPoint)); - accountV2.executeModuleCall( - address(module), - abi.encodeWithSignature("setSimpleValue(uint256)", 99999) - ); + accountV2.executeModuleCall(address(module), abi.encodeWithSignature("setSimpleValue(uint256)", 99999)); } function test_V1ToV2UpgradeNonEvmAccount() public { @@ -298,12 +270,12 @@ contract OmniAccountV1ToV2Upgrade is Test { bytes32 substrateOwner = keccak256("substrate_owner"); OmniAccountV1 substrateAccount = OmniAccountV1( - payable( - new ERC1967Proxy{salt: substrateOwner}( + payable(new ERC1967Proxy{salt: substrateOwner}( address(v1Implementation), - abi.encodeCall(OmniAccountV1.initialize, (substrateOwner, OwnerType.Substrate, clientId, rootSigner1)) - ) - ) + abi.encodeCall( + OmniAccountV1.initialize, (substrateOwner, OwnerType.Substrate, clientId, rootSigner1) + ) + )) ); // Setup state @@ -336,10 +308,7 @@ contract OmniAccountV1ToV2Upgrade is Test { // Execute module vm.prank(address(entryPoint)); - substrateAccountV2.executeModuleCall( - address(module), - abi.encodeWithSignature("setSimpleValue(uint256)", 777) - ); + substrateAccountV2.executeModuleCall(address(module), abi.encodeWithSignature("setSimpleValue(uint256)", 777)); // Verify account state intact assertEq(substrateAccountV2.owner(), substrateOwner, "Substrate owner corrupted"); diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol similarity index 95% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol index a36b9fa8ba..f76c747d5e 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2OmniAccountV2.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; @@ -10,7 +10,7 @@ import {Counter} from "../../src/Counter.sol"; import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; import {TestUtils} from "../TestUtils.sol"; import {PackedUserOperation} from "../../src/interfaces/PackedUserOperation.sol"; -import {SIG_VALIDATION_FAILED} from "../../src//core/Helpers.sol"; +import {SIG_VALIDATION_FAILED} from "../../src/core/Helpers.sol"; // add test cases for revert if called by non authorized address diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsEntryPoint.t.sol similarity index 98% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsEntryPoint.t.sol index 20f5d19529..96ab3bdf8a 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsEntryPoint.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsEntryPoint.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {Counter} from "../../src/Counter.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsOwner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsOwner.t.sol similarity index 99% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsOwner.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsOwner.t.sol index 1ddc70afff..2579801de6 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsOwner.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsOwner.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; @@ -10,7 +10,7 @@ import {Counter} from "../../src/Counter.sol"; import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; import {PackedUserOperation} from "../../src/interfaces/PackedUserOperation.sol"; import {TestUtils} from "../TestUtils.sol"; -import {SIG_VALIDATION_SUCCESS} from "../../src//core/Helpers.sol"; +import {SIG_VALIDATION_SUCCESS} from "../../src/core/Helpers.sol"; import {Passkey} from "../../src/interfaces/Passkey.sol"; import {MockModule} from "./MockModule.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsPasskey.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsPasskey.t.sol similarity index 98% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsPasskey.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsPasskey.t.sol index 595506080e..bcd34e2297 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsPasskey.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsPasskey.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; @@ -10,7 +10,7 @@ import {Counter} from "../../src/Counter.sol"; import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.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 {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/core/Helpers.sol"; import {Passkey} from "../../src/interfaces/Passkey.sol"; import "@openzeppelin/contracts/utils/Base64.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRoot.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRoot.t.sol similarity index 99% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRoot.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRoot.t.sol index bb310db5d8..9c02321860 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRoot.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRoot.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; @@ -10,7 +10,7 @@ import {Counter} from "../../src/Counter.sol"; import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.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 {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/core/Helpers.sol"; import {Passkey} from "../../src/interfaces/Passkey.sol"; import {MockModule} from "./MockModule.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRootNonEvm.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRootNonEvm.t.sol similarity index 99% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRootNonEvm.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRootNonEvm.t.sol index a1bda36a68..9b6cb11f15 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2AsRootNonEvm.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRootNonEvm.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2 as OmniAccount} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV2 as OmniAccount} from "../../src/accounts/v2/OmniAccountV2.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1 as EntryPoint} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Modules.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Modules.t.sol similarity index 98% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Modules.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Modules.t.sol index cee97cf346..40eb02d442 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Modules.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Modules.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {Counter} from "../../src/Counter.sol"; import {OmniAccountV2TestUtils} from "./OmniAccountV2TestUtils.sol"; @@ -107,5 +107,4 @@ contract OmniAccountV2Modules is Test { assertTrue(account.isModuleRegistered(address(module))); assertTrue(account.isModuleRegistered(address(module2))); } - } diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol similarity index 99% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol index 1e12e984f0..6f7060644b 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2 as OmniAccount} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV2 as OmniAccount} from "../../src/accounts/v2/OmniAccountV2.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1 as EntryPoint} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2TestUtils.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol similarity index 90% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2TestUtils.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol index 00e8f74814..de0a135588 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2TestUtils.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol @@ -3,7 +3,7 @@ pragma solidity ^0.8.28; import {Counter} from "../../src/Counter.sol"; import {Vm} from "forge-std/Vm.sol"; -import {OmniAccountV2 as OmniAccount} from "../../src/accounts/OmniAccountV2.sol"; +import {OmniAccountV2 as OmniAccount} from "../../src/accounts/v2/OmniAccountV2.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1 as EntryPoint} from "../../src/core/EntryPointV1.sol"; import {OwnerType} from "../../src/interfaces/OwnerType.sol"; @@ -20,11 +20,9 @@ library OmniAccountV2TestUtils { 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 OmniAccountV2TestUtils { 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/aa-contracts/test/v2/OmniAccountV2Upgradability.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Upgradability.t.sol similarity index 98% rename from tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Upgradability.t.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Upgradability.t.sol index ccb4364d49..c61d4cf9cc 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2Upgradability.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Upgradability.t.sol @@ -4,7 +4,7 @@ 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/accounts/OmniAccountV2.sol"; +import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {OwnerType} from "../../src/interfaces/OwnerType.sol"; import {Passkey} from "../../src/interfaces/Passkey.sol"; diff --git a/tee-worker/omni-executor/aa-contracts/test/v2/StorageTestModule.sol b/tee-worker/omni-executor/contracts/aa/test/v2/StorageTestModule.sol similarity index 96% rename from tee-worker/omni-executor/aa-contracts/test/v2/StorageTestModule.sol rename to tee-worker/omni-executor/contracts/aa/test/v2/StorageTestModule.sol index 37265a71d8..68d8d49899 100644 --- a/tee-worker/omni-executor/aa-contracts/test/v2/StorageTestModule.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/StorageTestModule.sol @@ -89,12 +89,7 @@ contract StorageTestModule { 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 - }); + ms.structStorage[_id] = TestStruct({id: _id, name: _name, value: _value, active: true}); emit StructSet(_id, _name, _value); } From 6f0de3882ba2b0b7adbcb5d70a5ca488a8e91e5f Mon Sep 17 00:00:00 2001 From: Kasper Ziemianek Date: Wed, 17 Dec 2025 10:51:57 +0100 Subject: [PATCH 5/7] fix build --- .../omni-executor/contracts/aa/src/core/DemoPaymaster.sol | 7 ++++++- .../contracts/aa/src/core/ERC20PaymasterV1.sol | 6 +++++- .../contracts/aa/src/core/SimplePaymaster.sol | 7 ++++++- .../omni-executor/contracts/aa/test/DeploymentHelper.t.sol | 3 ++- .../omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol | 2 +- 5 files changed, 20 insertions(+), 5 deletions(-) diff --git a/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol index 70533e3a5d..6da5caae6d 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol +++ b/tee-worker/omni-executor/contracts/aa/src/core/DemoPaymaster.sol @@ -29,7 +29,12 @@ contract DemoPaymaster is BasePaymaster { bytes32, /* userOpHash */ uint256 maxCost - ) internal view override returns (bytes memory context, uint256 validationData) { + ) + 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) { diff --git a/tee-worker/omni-executor/contracts/aa/src/core/ERC20PaymasterV1.sol b/tee-worker/omni-executor/contracts/aa/src/core/ERC20PaymasterV1.sol index a042324e24..00c14d658e 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/ERC20PaymasterV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/core/ERC20PaymasterV1.sol @@ -88,7 +88,11 @@ contract ERC20PaymasterV1 is BasePaymaster, ReentrancyGuard { bytes32, /* userOpHash */ uint256 maxCost - ) internal override returns (bytes memory context, uint256 validationData) { + ) + 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(); diff --git a/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol b/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol index dd61552b84..892a081228 100644 --- a/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol +++ b/tee-worker/omni-executor/contracts/aa/src/core/SimplePaymaster.sol @@ -33,7 +33,12 @@ contract SimplePaymaster is BasePaymaster { bytes32, /* userOpHash */ uint256 maxCost - ) internal view override returns (bytes memory context, uint256 validationData) { + ) + 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 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 5bab6b4aa3..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/v2/OmniAccountV2.t.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol index f76c747d5e..74bb81f23c 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol @@ -2,7 +2,7 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2OmniAccountV2.sol"; +import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; import {BaseAccount} from "../../src/core/BaseAccount.sol"; import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; import {UserOpSigner} from "../../src/interfaces/UserOpSigner.sol"; From 7572dd592f8308caae73681f36c8ac7f059b42a6 Mon Sep 17 00:00:00 2001 From: Kasper Ziemianek Date: Wed, 17 Dec 2025 11:05:25 +0100 Subject: [PATCH 6/7] add versioning notes --- tee-worker/omni-executor/contracts/aa/README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tee-worker/omni-executor/contracts/aa/README.md b/tee-worker/omni-executor/contracts/aa/README.md index fbba98e03e..12e686355a 100644 --- a/tee-worker/omni-executor/contracts/aa/README.md +++ b/tee-worker/omni-executor/contracts/aa/README.md @@ -129,3 +129,8 @@ 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. +Work on new version begins when first change is introduced after previous version deployment and continues until deployment. From 185bd0019b1ac8aa8139ac604cda5ea1f97f8472 Mon Sep 17 00:00:00 2001 From: Kasper Ziemianek Date: Wed, 17 Dec 2025 12:42:15 +0100 Subject: [PATCH 7/7] aa contracts split v1 and v2 --- .../omni-executor/contracts/aa/README.md | 1 + .../contracts/aa/script/Deploy.s.sol | 8 +- .../contracts/aa/script/DeployLocal.s.sol | 6 +- .../aa/script/DeployLocalWithPaymaster.s.sol | 10 +- .../accounts}/OmniAccountFactoryV1.sol | 4 +- .../v1 => v1/accounts}/OmniAccountV1.sol | 12 +- .../callback/TokenCallbackHandler.sol | 0 .../aa/src/{ => v1}/core/BaseAccount.sol | 0 .../aa/src/{ => v1}/core/BasePaymaster.sol | 0 .../aa/src/{ => v1}/core/DemoPaymaster.sol | 0 .../aa/src/{ => v1}/core/ERC20PaymasterV1.sol | 0 .../aa/src/{ => v1}/core/Eip7702Support.sol | 2 +- .../{ => v1}/core/EntryPointSimulations.sol | 0 .../aa/src/{ => v1}/core/EntryPointV1.sol | 0 .../aa/src/{ => v1}/core/Helpers.sol | 0 .../aa/src/{ => v1}/core/NonceManager.sol | 0 .../aa/src/{ => v1}/core/SenderCreator.sol | 0 .../aa/src/{ => v1}/core/SimplePaymaster.sol | 0 .../aa/src/{ => v1}/core/StakeManager.sol | 0 .../aa/src/{ => v1}/core/UserOperationLib.sol | 0 .../aa/src/{ => v1}/interfaces/IAccount.sol | 0 .../{ => v1}/interfaces/IAccountExecute.sol | 0 .../src/{ => v1}/interfaces/IAggregator.sol | 0 .../src/{ => v1}/interfaces/IEntryPoint.sol | 0 .../interfaces/IEntryPointSimulations.sol | 0 .../src/{ => v1}/interfaces/INonceManager.sol | 0 .../aa/src/{ => v1}/interfaces/IPaymaster.sol | 0 .../{ => v1}/interfaces/ISenderCreator.sol | 0 .../src/{ => v1}/interfaces/IStakeManager.sol | 0 .../aa/src/{ => v1}/interfaces/OwnerType.sol | 0 .../interfaces/PackedUserOperation.sol | 0 .../aa/src/{ => v1}/interfaces/Passkey.sol | 0 .../src/{ => v1}/interfaces/UserOpSigner.sol | 0 .../contracts/aa/src/{ => v1}/utils/Exec.sol | 0 .../accounts}/OmniAccountFactoryV2.sol | 4 +- .../v2 => v2/accounts}/OmniAccountV2.sol | 16 +- .../callback/TokenCallbackHandler.sol | 41 + .../contracts/aa/src/v2/core/BaseAccount.sol | 156 ++++ .../aa/src/v2/core/BasePaymaster.sol | 164 ++++ .../aa/src/v2/core/DemoPaymaster.sol | 76 ++ .../aa/src/v2/core/ERC20PaymasterV1.sol | 500 ++++++++++ .../aa/src/v2/core/Eip7702Support.sol | 79 ++ .../aa/src/v2/core/EntryPointSimulations.sol | 214 +++++ .../contracts/aa/src/v2/core/EntryPointV1.sol | 880 ++++++++++++++++++ .../contracts/aa/src/v2/core/Helpers.sol | 117 +++ .../aa/src/{ => v2}/core/LibModuleManager.sol | 0 .../contracts/aa/src/v2/core/NonceManager.sol | 40 + .../aa/src/v2/core/SenderCreator.sol | 56 ++ .../aa/src/v2/core/SimplePaymaster.sol | 94 ++ .../contracts/aa/src/v2/core/StakeManager.sol | 133 +++ .../aa/src/v2/core/UserOperationLib.sol | 123 +++ .../aa/src/v2/interfaces/IAccount.sol | 37 + .../aa/src/v2/interfaces/IAccountExecute.sol | 17 + .../aa/src/v2/interfaces/IAggregator.sol | 43 + .../aa/src/v2/interfaces/IEntryPoint.sol | 199 ++++ .../v2/interfaces/IEntryPointSimulations.sol | 82 ++ .../aa/src/v2/interfaces/INonceManager.sol | 27 + .../aa/src/v2/interfaces/IPaymaster.sol | 57 ++ .../aa/src/v2/interfaces/ISenderCreator.sol | 20 + .../aa/src/v2/interfaces/IStakeManager.sol | 94 ++ .../aa/src/v2/interfaces/OwnerType.sol | 26 + .../src/v2/interfaces/PackedUserOperation.sol | 28 + .../aa/src/v2/interfaces/Passkey.sol | 175 ++++ .../aa/src/v2/interfaces/UserOpSigner.sol | 25 + .../contracts/aa/src/v2/utils/Exec.sol | 53 ++ .../contracts/aa/test/ERC20PaymasterV1.t.sol | 10 +- .../contracts/aa/test/EntryPoint.t.sol | 10 +- .../contracts/aa/test/SimplePaymaster.t.sol | 8 +- .../contracts/aa/test/StakeManager.t.sol | 4 +- .../contracts/aa/test/TestUtils.sol | 2 +- .../contracts/aa/test/v1/OmniAccount.t.sol | 12 +- .../aa/test/v1/OmniAccountAsEntryPoint.t.sol | 6 +- .../aa/test/v1/OmniAccountAsOwner.t.sol | 14 +- .../aa/test/v1/OmniAccountAsPasskey.t.sol | 14 +- .../aa/test/v1/OmniAccountAsRoot.t.sol | 14 +- .../aa/test/v1/OmniAccountAsRootNonEvm.t.sol | 16 +- ...iAccountNonEvmOwnerWithPasskeySigner.t.sol | 16 +- .../aa/test/v1/OmniAccountTestUtils.sol | 8 +- .../aa/test/v1/OmniAccountUpgradability.t.sol | 8 +- .../aa/test/v2/OmniAccountV1ToV2Upgrade.t.sol | 54 +- .../contracts/aa/test/v2/OmniAccountV2.t.sol | 14 +- .../test/v2/OmniAccountV2AsEntryPoint.t.sol | 6 +- .../aa/test/v2/OmniAccountV2AsOwner.t.sol | 16 +- .../aa/test/v2/OmniAccountV2AsPasskey.t.sol | 16 +- .../aa/test/v2/OmniAccountV2AsRoot.t.sol | 16 +- .../test/v2/OmniAccountV2AsRootNonEvm.t.sol | 18 +- .../aa/test/v2/OmniAccountV2Modules.t.sol | 4 +- ...ccountV2NonEvmOwnerWithPasskeySigner.t.sol | 18 +- .../aa/test/v2/OmniAccountV2TestUtils.sol | 15 +- .../test/v2/OmniAccountV2Upgradability.t.sol | 13 +- .../contracts/aa/test/v2/TestUtilsV2.sol | 34 + 91 files changed, 3794 insertions(+), 191 deletions(-) rename tee-worker/omni-executor/contracts/aa/src/{accounts/v1 => v1/accounts}/OmniAccountFactoryV1.sol (96%) rename tee-worker/omni-executor/contracts/aa/src/{accounts/v1 => v1/accounts}/OmniAccountV1.sol (97%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/accounts/callback/TokenCallbackHandler.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/BaseAccount.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/BasePaymaster.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/DemoPaymaster.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/ERC20PaymasterV1.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/Eip7702Support.sol (98%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/EntryPointSimulations.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/EntryPointV1.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/Helpers.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/NonceManager.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/SenderCreator.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/SimplePaymaster.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/StakeManager.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/core/UserOperationLib.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/IAccount.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/IAccountExecute.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/IAggregator.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/IEntryPoint.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/IEntryPointSimulations.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/INonceManager.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/IPaymaster.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/ISenderCreator.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/IStakeManager.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/OwnerType.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/PackedUserOperation.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/Passkey.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/interfaces/UserOpSigner.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{ => v1}/utils/Exec.sol (100%) rename tee-worker/omni-executor/contracts/aa/src/{accounts/v2 => v2/accounts}/OmniAccountFactoryV2.sol (95%) rename tee-worker/omni-executor/contracts/aa/src/{accounts/v2 => v2/accounts}/OmniAccountV2.sol (97%) create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/accounts/callback/TokenCallbackHandler.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/BaseAccount.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/BasePaymaster.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/DemoPaymaster.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/ERC20PaymasterV1.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/Eip7702Support.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/EntryPointSimulations.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/EntryPointV1.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/Helpers.sol rename tee-worker/omni-executor/contracts/aa/src/{ => v2}/core/LibModuleManager.sol (100%) create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/NonceManager.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/SenderCreator.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/SimplePaymaster.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/StakeManager.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/core/UserOperationLib.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAccount.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAccountExecute.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IAggregator.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IEntryPoint.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IEntryPointSimulations.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/INonceManager.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IPaymaster.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/ISenderCreator.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/IStakeManager.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/OwnerType.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/PackedUserOperation.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/Passkey.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/interfaces/UserOpSigner.sol create mode 100644 tee-worker/omni-executor/contracts/aa/src/v2/utils/Exec.sol create mode 100644 tee-worker/omni-executor/contracts/aa/test/v2/TestUtilsV2.sol diff --git a/tee-worker/omni-executor/contracts/aa/README.md b/tee-worker/omni-executor/contracts/aa/README.md index 12e686355a..0e295b5abe 100644 --- a/tee-worker/omni-executor/contracts/aa/README.md +++ b/tee-worker/omni-executor/contracts/aa/README.md @@ -133,4 +133,5 @@ 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/script/Deploy.s.sol b/tee-worker/omni-executor/contracts/aa/script/Deploy.s.sol index 5e38da29fd..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/v1/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 f366fcd955..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/v1/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 e65ae0ef1e..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/v1/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/src/accounts/v1/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/v1/OmniAccountFactoryV1.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountFactoryV1.sol index acc7fafeb6..923ffd3a65 100644 --- a/tee-worker/omni-executor/contracts/aa/src/accounts/v1/OmniAccountFactoryV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountFactoryV1.sol @@ -4,8 +4,8 @@ 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 "../interfaces/ISenderCreator.sol"; +import "../interfaces/OwnerType.sol"; import "./OmniAccountV1.sol"; /** diff --git a/tee-worker/omni-executor/contracts/aa/src/accounts/v1/OmniAccountV1.sol b/tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountV1.sol similarity index 97% rename from tee-worker/omni-executor/contracts/aa/src/accounts/v1/OmniAccountV1.sol rename to tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountV1.sol index 99ed5505ca..1da2aa959b 100644 --- a/tee-worker/omni-executor/contracts/aa/src/accounts/v1/OmniAccountV1.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v1/accounts/OmniAccountV1.sol @@ -10,12 +10,12 @@ 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 "../core/BaseAccount.sol"; +import "../interfaces/OwnerType.sol"; +import "../interfaces/UserOpSigner.sol"; +import "../interfaces/Passkey.sol"; +import "../core/Helpers.sol"; +import "./callback/TokenCallbackHandler.sol"; /** * smart wallet account 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 100% 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 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 100% 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 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 100% 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 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 100% 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 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 100% 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 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 100% 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 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 100% 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 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 100% 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 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/accounts/v2/OmniAccountFactoryV2.sol b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountFactoryV2.sol similarity index 95% rename from tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountFactoryV2.sol rename to tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountFactoryV2.sol index 7a50b0b527..f76e458749 100644 --- a/tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountFactoryV2.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountFactoryV2.sol @@ -4,8 +4,8 @@ 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 "../interfaces/ISenderCreator.sol"; +import "../interfaces/OwnerType.sol"; import "./OmniAccountV2.sol"; contract OmniAccountFactoryV2 { diff --git a/tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountV2.sol b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountV2.sol similarity index 97% rename from tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountV2.sol rename to tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountV2.sol index 6d1add472b..834e2dc93b 100644 --- a/tee-worker/omni-executor/contracts/aa/src/accounts/v2/OmniAccountV2.sol +++ b/tee-worker/omni-executor/contracts/aa/src/v2/accounts/OmniAccountV2.sol @@ -10,14 +10,14 @@ 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"; +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; 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/core/LibModuleManager.sol b/tee-worker/omni-executor/contracts/aa/src/v2/core/LibModuleManager.sol similarity index 100% rename from tee-worker/omni-executor/contracts/aa/src/core/LibModuleManager.sol rename to tee-worker/omni-executor/contracts/aa/src/v2/core/LibModuleManager.sol 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/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 01af3b0b8c..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/v1/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/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 ba958a02a2..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) { diff --git a/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccount.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccount.t.sol index 1c3ecc7dda..0fff674f89 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v1/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/v1/OmniAccountV1.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../../src/interfaces/UserOpSigner.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 {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/v1/OmniAccountAsEntryPoint.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsEntryPoint.t.sol index cdf8393771..b7f01b0863 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsEntryPoint.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsEntryPoint.t.sol @@ -2,9 +2,9 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV1} from "../../src/accounts/v1/OmniAccountV1.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.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"; diff --git a/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsOwner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsOwner.t.sol index 564f8b49d4..6c202a1d93 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v1/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/v1/OmniAccountV1.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../../src/interfaces/UserOpSigner.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 {PackedUserOperation} from "../../src/v1/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 {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/v1/OmniAccountAsPasskey.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsPasskey.t.sol index 3c56700712..876be311d2 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v1/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/v1/OmniAccountV1.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../../src/interfaces/UserOpSigner.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 {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/v1/OmniAccountAsRoot.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRoot.t.sol index c262e2e0d3..4e4c29366c 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v1/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/v1/OmniAccountV1.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../../src/interfaces/UserOpSigner.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 {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/v1/OmniAccountAsRootNonEvm.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountAsRootNonEvm.t.sol index cf73f47a6c..107fc93216 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v1/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/v1/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 {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 {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/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountNonEvmOwnerWithPasskeySigner.t.sol index 6a4c512327..ef3adff56d 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v1/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/v1/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 {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 {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/v1/OmniAccountTestUtils.sol b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountTestUtils.sol index 31ae846e16..e59ca2329d 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountTestUtils.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountTestUtils.sol @@ -3,10 +3,10 @@ pragma solidity ^0.8.28; import {Counter} from "../../src/Counter.sol"; import {Vm} from "forge-std/Vm.sol"; -import {OmniAccountV1 as OmniAccount} from "../../src/accounts/v1/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"; 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 index ef7df39766..fc13820669 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountUpgradability.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v1/OmniAccountUpgradability.t.sol @@ -4,10 +4,10 @@ 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/accounts/v1/OmniAccountV1.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {OwnerType} from "../../src/interfaces/OwnerType.sol"; -import {Passkey} from "../../src/interfaces/Passkey.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"; 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 index 4dd5ec3dd4..87222193f2 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV1ToV2Upgrade.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV1ToV2Upgrade.t.sol @@ -4,13 +4,18 @@ 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/accounts/v1/OmniAccountV1.sol"; -import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {OwnerType} from "../../src/interfaces/OwnerType.sol"; -import {Passkey} from "../../src/interfaces/Passkey.sol"; -import {TestUtils} from "../TestUtils.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. @@ -35,7 +40,10 @@ contract OmniAccountV1ToV2Upgrade is Test { address public rootSigner3 = address(0x4444); bytes public clientId = bytes("comprehensive_test_client_v1_to_v2"); - // Passkey signers + // 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}); @@ -47,18 +55,19 @@ contract OmniAccountV1ToV2Upgrade is Test { entryPoint = new EntryPointV1(); // Deploy V1 implementation and create proxy - OmniAccountV1 v1Implementation = new OmniAccountV1(entryPoint); + 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, OwnerType.Evm, clientId, rootSigner1)) + abi.encodeCall(OmniAccountV1.initialize, (expectedOwner, OwnerTypeV1.Evm, clientId, rootSigner1)) )) ); // Prepare V2 implementation for upgrade - accountV2Implementation = new OmniAccountV2(entryPoint); + // 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(); @@ -75,9 +84,9 @@ contract OmniAccountV1ToV2Upgrade is Test { // Add passkey signers vm.prank(owner); - accountV1.addPasskeySigner(passkey1); + accountV1.addPasskeySigner(passkey1V1); vm.prank(owner); - accountV1.addPasskeySigner(passkey2); + accountV1.addPasskeySigner(passkey2V1); // Add deposit to EntryPoint vm.deal(address(accountV1), 10 ether); @@ -87,12 +96,12 @@ contract OmniAccountV1ToV2Upgrade is Test { // Verify V1 initial state assertEq(accountV1.owner(), expectedOwner, "V1 owner mismatch"); assertEq(accountV1.clientId(), clientId, "V1 clientId mismatch"); - assertEq(uint256(accountV1.ownerType()), uint256(OwnerType.Evm), "V1 ownerType 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(Passkey.toKey(passkey1)), "passkey1 not registered"); - assertTrue(accountV1.passkeySigners(Passkey.toKey(passkey2)), "passkey2 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"); @@ -172,12 +181,12 @@ contract OmniAccountV1ToV2Upgrade is Test { // 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(OwnerType.Evm), "OwnerType 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(Passkey.toKey(passkey1)), "passkey1 corrupted"); - assertTrue(accountV2.passkeySigners(Passkey.toKey(passkey2)), "passkey2 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"); @@ -208,7 +217,8 @@ contract OmniAccountV1ToV2Upgrade is Test { // Remove a passkey signer vm.prank(owner); - accountV2.removePasskeySigner(passkey1); + 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"); @@ -266,14 +276,14 @@ contract OmniAccountV1ToV2Upgrade is Test { function test_V1ToV2UpgradeNonEvmAccount() public { // Test upgrade with non-EVM account type - OmniAccountV1 v1Implementation = new OmniAccountV1(entryPoint); + 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, OwnerType.Substrate, clientId, rootSigner1) + OmniAccountV1.initialize, (substrateOwner, OwnerTypeV1.Substrate, clientId, rootSigner1) ) )) ); @@ -284,7 +294,7 @@ contract OmniAccountV1ToV2Upgrade is Test { // Verify V1 state assertEq(substrateAccount.owner(), substrateOwner); - assertEq(uint256(substrateAccount.ownerType()), uint256(OwnerType.Substrate)); + assertEq(uint256(substrateAccount.ownerType()), uint256(OwnerTypeV1.Substrate)); assertTrue(substrateAccount.isRootSigner(rootSigner1)); assertTrue(substrateAccount.isRootSigner(rootSigner2)); 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 index 74bb81f23c..16853dec3e 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2.t.sol @@ -2,15 +2,15 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../../src/interfaces/UserOpSigner.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 {TestUtils} from "../TestUtils.sol"; -import {PackedUserOperation} from "../../src/interfaces/PackedUserOperation.sol"; -import {SIG_VALIDATION_FAILED} from "../../src/core/Helpers.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 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 index 96ab3bdf8a..3a33e9271a 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsEntryPoint.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsEntryPoint.t.sol @@ -2,9 +2,9 @@ pragma solidity ^0.8.28; import {Test, console} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.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"; 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 index 2579801de6..0d39a0dca1 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsOwner.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsOwner.t.sol @@ -2,16 +2,16 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../../src/interfaces/UserOpSigner.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/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/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 { 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 index bcd34e2297..7037cd9a92 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsPasskey.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsPasskey.t.sol @@ -2,16 +2,16 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../../src/interfaces/UserOpSigner.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 {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 {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 { 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 index 9c02321860..eb2c76eb5a 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRoot.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRoot.t.sol @@ -2,16 +2,16 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; -import {BaseAccount} from "../../src/core/BaseAccount.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {UserOpSigner} from "../../src/interfaces/UserOpSigner.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 {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 {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 { 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 index 9b6cb11f15..d5ce0bc2c7 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRootNonEvm.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2AsRootNonEvm.t.sol @@ -2,17 +2,17 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2 as OmniAccount} from "../../src/accounts/v2/OmniAccountV2.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 {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 {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 {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 { 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 index 40eb02d442..ca42235521 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Modules.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Modules.t.sol @@ -2,8 +2,8 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2} from "../../src/accounts/v2/OmniAccountV2.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.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"; 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 index 6f7060644b..88a8ae2d7a 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2NonEvmOwnerWithPasskeySigner.t.sol @@ -2,17 +2,17 @@ pragma solidity ^0.8.28; import {Test} from "forge-std/Test.sol"; -import {OmniAccountV2 as OmniAccount} from "../../src/accounts/v2/OmniAccountV2.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 {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 {TestUtils} from "../TestUtils.sol"; -import {PackedUserOperation} from "../../src/interfaces/PackedUserOperation.sol"; -import {SIG_VALIDATION_SUCCESS, SIG_VALIDATION_FAILED} from "../../src/core/Helpers.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; diff --git a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol index de0a135588..111e234012 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2TestUtils.sol @@ -3,12 +3,13 @@ pragma solidity ^0.8.28; import {Counter} from "../../src/Counter.sol"; import {Vm} from "forge-std/Vm.sol"; -import {OmniAccountV2 as OmniAccount} from "../../src/accounts/v2/OmniAccountV2.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 {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 {TestUtils} from "../TestUtils.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; library OmniAccountV2TestUtils { function setUpWithOwnerType(address ownerAddress, bytes memory clientId, address rootAddress, OwnerType ownerType) @@ -17,7 +18,7 @@ library OmniAccountV2TestUtils { { Counter counter = new Counter(); EntryPoint entryPoint = new EntryPoint(); - OmniAccount accountImpl = new OmniAccount(entryPoint); + OmniAccount accountImpl = new OmniAccount(IEntryPoint(address(entryPoint))); bytes32 oa = TestUtils.prepare_evm_oa(ownerAddress, clientId); OmniAccount account = OmniAccount( payable(new ERC1967Proxy{salt: oa}( @@ -34,7 +35,7 @@ library OmniAccountV2TestUtils { { Counter counter = new Counter(); EntryPoint entryPoint = new EntryPoint(); - OmniAccount accountImpl = new OmniAccount(entryPoint); + OmniAccount accountImpl = new OmniAccount(IEntryPoint(address(entryPoint))); bytes32 oa = TestUtils.prepare_evm_oa(ownerAddress, clientId); OmniAccount account = OmniAccount( payable(new ERC1967Proxy{salt: oa}( 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 index c61d4cf9cc..7c08afef21 100644 --- a/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Upgradability.t.sol +++ b/tee-worker/omni-executor/contracts/aa/test/v2/OmniAccountV2Upgradability.t.sol @@ -4,15 +4,16 @@ 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/accounts/v2/OmniAccountV2.sol"; -import {EntryPointV1} from "../../src/core/EntryPointV1.sol"; -import {OwnerType} from "../../src/interfaces/OwnerType.sol"; -import {Passkey} from "../../src/interfaces/Passkey.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 {TestUtils} from "../TestUtils.sol"; +import {TestUtilsV2 as TestUtils} from "./TestUtilsV2.sol"; contract OmniAccountV3 is OmniAccountV2 { - constructor(EntryPointV1 anEntryPoint) OmniAccountV2(anEntryPoint) {} + constructor(EntryPointV1 anEntryPoint) OmniAccountV2(IEntryPoint(address(anEntryPoint))) {} function version() public pure override returns (string memory) { return "3.0.0"; 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 + )); + } +}