Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ 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 "@openzeppelin/contracts/interfaces/IERC1271.sol";
import "../core/BaseAccount.sol";
import "../interfaces/OwnerType.sol";
import "../interfaces/UserOpSigner.sol";
Expand All @@ -19,7 +20,7 @@ import "./callback/TokenCallbackHandler.sol";
import "../utils/Exec.sol";
import "../core/LibModuleManager.sol";

contract OmniAccountV2 is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, Initializable {
contract OmniAccountV2 is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, Initializable, IERC1271 {
using Passkey for Passkey.PublicKey;
using LibModuleManager for LibModuleManager.ModuleStorage;

Expand All @@ -36,7 +37,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)"));
Expand Down Expand Up @@ -71,22 +73,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;
Expand All @@ -105,10 +107,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");

Expand Down Expand Up @@ -156,9 +158,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];
Expand Down Expand Up @@ -247,9 +249,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 {
Expand Down Expand Up @@ -278,6 +280,63 @@ contract OmniAccountV2 is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, In
return Exec.getReturnData(0);
}

/**
* @dev ERC-1271 signature validation
* @param hash Hash of the data to be signed
* @param signature Signature byte array
* @return magicValue 0x1626ba7e if signature is valid, 0xffffffff otherwise
*/
function isValidSignature(bytes32 hash, bytes memory signature) external view override returns (bytes4 magicValue) {
if (signature.length < 1) {
return 0xffffffff;
}

UserOpSigner signer = UserOpSigner(uint8(signature[0]));

bool isValid = false;

if (signer == UserOpSigner.Owner) {
if (signature.length == 66) {
// 1 byte signer type + 65 bytes signature
bytes memory sig = new bytes(65);
for (uint256 i = 0; i < 65; i++) {
sig[i] = signature[i + 1];
}
address recovered = ECDSA.recover(hash, sig);
isValid = (owner == _determineOa(recovered));
}
} else if (signer == UserOpSigner.RootKey) {
if (signature.length == 66) {
// 1 byte signer type + 65 bytes signature
bytes memory sig = new bytes(65);
for (uint256 i = 0; i < 65; i++) {
sig[i] = signature[i + 1];
}
address recovered = ECDSA.recover(hash, sig);
isValid = isRootSigner(recovered);
}
} else if (signer == UserOpSigner.Passkey) {
if (signature.length > 1) {
bytes memory sig = new bytes(signature.length - 1);
for (uint256 i = 0; i < signature.length - 1; i++) {
sig[i] = signature[i + 1];
}

(
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()]) {
isValid = Passkey.verify(hash, metadata, passkeySignature, publicKey);
}
}
}

return isValid ? bytes4(0x1626ba7e) : bytes4(0xffffffff);
}

function version() public pure virtual returns (string memory) {
return "2.0.0";
}
Expand Down
52 changes: 52 additions & 0 deletions tee-worker/omni-executor/aa-contracts/test/v2/OmniAccountV2.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,9 @@ contract OmniAccountV2Test is Test {
address rootAddress = 0x0000000000000000000000000000000000000001;
bytes clientId = bytes("test_client");

bytes4 constant ERC1271_MAGIC_VALUE = 0x1626ba7e;
bytes4 constant ERC1271_INVALID_SIGNATURE = 0xffffffff;

function setUp() public {
(counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, rootAddress);
}
Expand Down Expand Up @@ -80,4 +83,53 @@ contract OmniAccountV2Test is Test {
uint256 validationData = account.validateUserOp(packedOp, packedOpHash, 0);
assertEq(SIG_VALIDATION_FAILED, validationData);
}

// ============ ERC-1271 Edge Cases and Invalid Input Tests ============

function test_ERC1271_EmptySignature() public view {
bytes32 messageHash = keccak256("Test message");
bytes memory signature = "";

bytes4 result = account.isValidSignature(messageHash, signature);
assertEq(result, ERC1271_INVALID_SIGNATURE);
}

function test_ERC1271_InvalidSignerType() public {
bytes32 messageHash = keccak256("Test message");
(, uint256 pk) = makeAddrAndKey("signer");

// Use invalid signer type (99) - this will cause a panic in enum conversion
(uint8 v, bytes32 r, bytes32 s) = vm.sign(pk, messageHash);
bytes memory signature = abi.encodePacked(uint8(99), r, s, v);

// Expect a panic (0x21 is the panic code for invalid enum conversion)
vm.expectRevert();
account.isValidSignature(messageHash, signature);
}

function test_ERC1271_OnlySignerTypeByte() public view {
bytes32 messageHash = keccak256("Test message");
bytes memory signature = abi.encodePacked(uint8(UserOpSigner.Owner));

bytes4 result = account.isValidSignature(messageHash, signature);
assertEq(result, ERC1271_INVALID_SIGNATURE);
}

function test_ERC1271_MultipleSignatureTypesForSameMessage() public {
(address owner, uint256 ownerPk) = makeAddrAndKey("owner");
(address root, uint256 rootPk) = makeAddrAndKey("root");
(counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(owner, clientId, root);

bytes32 messageHash = keccak256("Shared message");

// Owner signature
(uint8 v1, bytes32 r1, bytes32 s1) = vm.sign(ownerPk, messageHash);
bytes memory ownerSig = abi.encodePacked(uint8(UserOpSigner.Owner), r1, s1, v1);
assertEq(account.isValidSignature(messageHash, ownerSig), ERC1271_MAGIC_VALUE);

// Root key signature
(uint8 v2, bytes32 r2, bytes32 s2) = vm.sign(rootPk, messageHash);
bytes memory rootSig = abi.encodePacked(uint8(UserOpSigner.RootKey), r2, s2, v2);
assertEq(account.isValidSignature(messageHash, rootSig), ERC1271_MAGIC_VALUE);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,16 @@ contract OmniAccountV2AsOwner is Test {
Counter public counter;
MockModule public module;

address ownerAddress = 0x0000000000000000000000000000000000000000;
address ownerAddress;
uint256 ownerPk;
address rootAddress = 0x0000000000000000000000000000000000000001;
bytes clientId = bytes("test_client");

bytes4 constant ERC1271_MAGIC_VALUE = 0x1626ba7e;
bytes4 constant ERC1271_INVALID_SIGNATURE = 0xffffffff;

function setUp() public {
(ownerAddress, ownerPk) = makeAddrAndKey("owner");
(counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, rootAddress);
module = new MockModule();
}
Expand Down Expand Up @@ -558,4 +563,28 @@ contract OmniAccountV2AsOwner is Test {
vm.prank(unauthorized);
account.unregisterModule(address(module));
}

// ============ ERC-1271 Signature Validation Tests ============

function test_ERC1271_OwneValidSignature() public view {
bytes32 messageHash = keccak256("Hello, World!");

(uint8 v, bytes32 r, bytes32 s) = vm.sign(ownerPk, messageHash);
bytes memory signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v);

bytes4 result = account.isValidSignature(messageHash, signature);
assertEq(result, ERC1271_MAGIC_VALUE);
}

function test_ERC1271_OwnerInvalidSignatureRejected() public {
bytes32 messageHash = keccak256("Hello, World!");

// Sign with wrong key
(, uint256 wrongPk) = makeAddrAndKey("wrong");
(uint8 v, bytes32 r, bytes32 s) = vm.sign(wrongPk, messageHash);
bytes memory signature = abi.encodePacked(uint8(UserOpSigner.Owner), r, s, v);

bytes4 result = account.isValidSignature(messageHash, signature);
assertEq(result, ERC1271_INVALID_SIGNATURE);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,9 @@ contract OmniAccountV2AsPasskey is Test {
address ownerAddress = 0x0000000000000000000000000000000000000000;
bytes clientId = bytes("test_client");

bytes4 constant ERC1271_MAGIC_VALUE = 0x1626ba7e;
bytes4 constant ERC1271_INVALID_SIGNATURE = 0xffffffff;

function test_ValidateOpPasskey() public {
(address root,) = makeAddrAndKey("root");
(counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root);
Expand Down Expand Up @@ -153,4 +156,77 @@ contract OmniAccountV2AsPasskey is Test {
}
revert("Substring not found");
}

// ============ ERC-1271 Signature Validation Tests ============

function test_ERC1271_PasskeyValidSignature() public {
(address root,) = makeAddrAndKey("root");
(counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root);

// Add passkey signer
Passkey.PublicKey memory pk = Passkey.PublicKey({
x: 0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef,
y: 0xfedcba0987654321fedcba0987654321fedcba0987654321fedcba0987654321
});

vm.prank(ownerAddress);
account.addPasskeySigner(pk);

bytes32 messageHash = keccak256("Passkey message");

// Create valid passkey signature (note: in real tests you'd need proper P256 signature)
Passkey.Signature memory passkeySignature = Passkey.Signature({
r: 0xaabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccdd,
s: 0x1122334411223344112233441122334411223344112233441122334411223344
});

Passkey.Metadata memory metadata = Passkey.Metadata({
authData: hex"49960de5880e8c687434170f6476605b8fe4aeb9a28632c7995cf3ba831d97630500000000",
clientDataJSON: '{"type":"webauthn.get","challenge":"AAAA","origin":"https://example.com"}',
challengeIndex: 32,
typeIndex: 1,
userVerificationRequired: false
});

bytes memory signature = abi.encode(pk, passkeySignature, metadata);
signature = abi.encodePacked(uint8(UserOpSigner.Passkey), signature);

// Note: This will fail because we don't have a real P256 signature
// In production, you'd use proper P256 signing or mock the Passkey.verify
bytes4 result = account.isValidSignature(messageHash, signature);
// Expect invalid because we don't have real P256 signature
assertEq(result, ERC1271_INVALID_SIGNATURE);
}

function test_ERC1271_UnregisteredPasskeyCannotSign() public {
(address root,) = makeAddrAndKey("root");
(counter, entryPoint, account) = OmniAccountV2TestUtils.setUp(ownerAddress, clientId, root);

// Use unregistered passkey
Passkey.PublicKey memory pk = Passkey.PublicKey({
x: 0x9999999999999999999999999999999999999999999999999999999999999999,
y: 0x8888888888888888888888888888888888888888888888888888888888888888
});

bytes32 messageHash = keccak256("Passkey message");

Passkey.Signature memory passkeySignature = Passkey.Signature({
r: 0xaabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccdd,
s: 0x1122334411223344112233441122334411223344112233441122334411223344
});

Passkey.Metadata memory metadata = Passkey.Metadata({
authData: hex"49960de5880e8c687434170f6476605b8fe4aeb9a28632c7995cf3ba831d97630500000000",
clientDataJSON: '{"type":"webauthn.get","challenge":"AAAA","origin":"https://example.com"}',
challengeIndex: 32,
typeIndex: 1,
userVerificationRequired: false
});

bytes memory signature = abi.encode(pk, passkeySignature, metadata);
signature = abi.encodePacked(uint8(UserOpSigner.Passkey), signature);

bytes4 result = account.isValidSignature(messageHash, signature);
assertEq(result, ERC1271_INVALID_SIGNATURE);
}
}
Loading