From 46d8bed8b9a349d59077fd378c797c838cdcd4c2 Mon Sep 17 00:00:00 2001 From: Pierre Hay Date: Wed, 13 Nov 2024 23:31:30 +0100 Subject: [PATCH] refactor: relocate session key validation from execution to validation hook --- contracts/core/base/BaseOpenfortAccount.sol | 73 ++++-- tasks/deployFactory.ts | 2 +- test/upgradeableOpenfortAccountTest.ts | 262 +++++++++++++++++--- 3 files changed, 272 insertions(+), 65 deletions(-) diff --git a/contracts/core/base/BaseOpenfortAccount.sol b/contracts/core/base/BaseOpenfortAccount.sol index 432c5d9..341348f 100644 --- a/contracts/core/base/BaseOpenfortAccount.sol +++ b/contracts/core/base/BaseOpenfortAccount.sol @@ -28,6 +28,17 @@ import {TokenCallbackHandler} from "./TokenCallbackHandler.sol"; import {OpenfortErrors} from "../../interfaces/OpenfortErrors.sol"; import {OpenfortEvents} from "../../interfaces/OpenfortEvents.sol"; +// Each call data for batches +struct Call { + address target; // Target contract address + uint256 value; // Amount of ETH to send with call + bytes callData; // Calldata to send +} + +interface ITimestampAsserter { + function assertTimestampInRange(uint256 start, uint256 end) external view; +} + /** * @title BaseOpenfortAccount (Non upgradeable by default) @@ -50,7 +61,6 @@ abstract contract BaseOpenfortAccount is { using TransactionHelper for Transaction; using ECDSAUpgradeable for bytes32; - // bytes4(keccak256("isValidSignature(bytes32,bytes)") bytes4 internal constant EIP1271_SUCCESS_RETURN_VALUE = 0x1626ba7e; // bytes4(keccak256("executeTransaction(bytes32,bytes32,Transaction)") @@ -59,6 +69,8 @@ abstract contract BaseOpenfortAccount is bytes32 internal constant OF_MSG_TYPEHASH = 0x57159f03b9efda178eab2037b2ec0b51ce11be0051b8a2a9992c29dc260e4a30; address internal constant BATCH_CALLER_ADDRESS = 0x7219257B57d9546c1BC0649617d557Db09C92D23; + address internal constant TIMESTAMP_ASSERTER_ADDRESS = 0x5ea4C1df68Fd54EA9242bC6C405E7699EBbcb5F1; + /** * Struct to keep track of session keys' data @@ -131,21 +143,27 @@ abstract contract BaseOpenfortAccount is require(totalRequiredBalance <= address(this).balance, "Not enough balance for fee + value"); - if (_validateSignature(txHash, _transaction.signature, _transaction.to) == EIP1271_SUCCESS_RETURN_VALUE) { + if ( + _validateSignature(txHash, _transaction.signature, _transaction.to, _transaction.data) + == EIP1271_SUCCESS_RETURN_VALUE + ) { magic = ACCOUNT_VALIDATION_SUCCESS_MAGIC; } else { magic = ""; } } - function _validateSignature(bytes32 _hash, bytes memory _signature, uint256 _to) internal returns (bytes4 magic) { + function _validateSignature(bytes32 _hash, bytes memory _signature, uint256 _to, bytes calldata _data) + internal + returns (bytes4 magic) + { magic = EIP1271_SUCCESS_RETURN_VALUE; address signerAddress = _recoverECDSAsignature(_hash, _signature); // Note, that we should abstain from using the require here in order to allow for fee estimation to work if (signerAddress != owner() && signerAddress != address(0)) { // if not owner, try session key validation - if (!isValidSessionKey(signerAddress, _to)) { + if (!isValidSessionKey(signerAddress, _to, _data)) { magic = ""; } } @@ -184,27 +202,45 @@ abstract contract BaseOpenfortAccount is /* * @notice Return whether a sessionKey is valid. */ - function isValidSessionKey(address _sessionKey, uint256 _to) internal virtual returns (bool) { + function isValidSessionKey(address _sessionKey, uint256 _to, bytes calldata _data) internal virtual returns (bool) { SessionKeyStruct storage sessionKey = sessionKeys[_sessionKey]; // If not owner and the session key is revoked, return false if (sessionKey.validUntil == 0) return false; // If the sessionKey was not registered by the owner, return false // If the account is transferred or sold, isValidSessionKey() will return false with old session keys; if (sessionKey.registrarAddress != owner()) return false; + // If the session key is out of time range, return false + + ITimestampAsserter timestampAsserter = ITimestampAsserter(TIMESTAMP_ASSERTER_ADDRESS); + timestampAsserter.assertTimestampInRange(sessionKey.validAfter, sessionKey.validUntil); + // master session key bypasses whitelist and limit checks + if (sessionKey.masterSessionKey) return true; address to = address(uint160(_to)); + if (to != BATCH_CALLER_ADDRESS) { + return _validateSessionKeyCall(sessionKey, to); + } + Call[] memory calls = abi.decode(_data[4:], (Call[])); + for (uint256 i; i < calls.length;) { + if (!_validateSessionKeyCall(sessionKey, calls[i].target)) return false; + unchecked { + ++i; + } + } + return true; + } + + function _validateSessionKeyCall(SessionKeyStruct storage _sessionKey, address _to) internal returns (bool) { // Check if reenter, do not allow - if (to == address(this)) return false; - // Check if it is a masterSessionKey - if (sessionKey.masterSessionKey) return true; + if (_to == address(this)) return false; // Limit of transactions per sessionKey reached - if (sessionKey.limit == 0) return false; + if (_sessionKey.limit == 0) return false; // Deduct one use of the limit for the given session key unchecked { - sessionKey.limit = sessionKey.limit - 1; + _sessionKey.limit = _sessionKey.limit - 1; } // If there is no whitelist or there is, but the target is whitelisted, return true - if (!sessionKey.whitelisting || sessionKey.whitelist[to]) { + if (!_sessionKey.whitelisting || _sessionKey.whitelist[_to]) { return true; } // All other cases, deny @@ -261,21 +297,6 @@ abstract contract BaseOpenfortAccount is function _executeTransaction(Transaction calldata _transaction) internal { address to = address(uint160(_transaction.to)); require(to != address(DEPLOYER_SYSTEM_CONTRACT), "Deployment from smart account not supported"); - // ! WARNING ! SESSION KEY VALIDATION - // THIS CODE DOES NOT BELONG IN THE EXECUTION LOGIC - // SIGNATURE HAS ALREADY BEEN RECOVERED IN THE VALIDATION FLOW - - // TODO: MOVE IT TO THE VALIDATION LOGIC WHEN CONTEXTUAL VARIABLES ARE AVAILABLE - bytes32 txHash = _transaction.encodeHash(); - address signer = _recoverECDSAsignature(txHash, _transaction.signature); - - if (signer != owner() && signer != address(0)) { - SessionKeyStruct storage sk = sessionKeys[signer]; - if (block.timestamp < sk.validAfter || block.timestamp > sk.validUntil) { - revert SessionKeyOutOfTimeRange(); - } - } - to == BATCH_CALLER_ADDRESS ? _executeDelegatecall(to, _transaction.data) : _call(to, _transaction.value, _transaction.data); diff --git a/tasks/deployFactory.ts b/tasks/deployFactory.ts index 00643cc..76b83bc 100644 --- a/tasks/deployFactory.ts +++ b/tasks/deployFactory.ts @@ -55,7 +55,7 @@ task("deploy-factory", "Deploy an Openfort Factory") gasPerPubdata: utils.DEFAULT_GAS_PER_PUBDATA_LIMIT, }, }, - [proxyArtifact.bytecode] // Pass the bytecode for additional factory deps + [proxyArtifact.bytecode] ) const FACTORY_ADDRESS = await contract.getAddress() diff --git a/test/upgradeableOpenfortAccountTest.ts b/test/upgradeableOpenfortAccountTest.ts index 591ab30..b224ae9 100644 --- a/test/upgradeableOpenfortAccountTest.ts +++ b/test/upgradeableOpenfortAccountTest.ts @@ -1,5 +1,7 @@ -import { expect } from "chai" -import { encodeFunctionData, encodePacked, parseAbi } from "viem" +import chai from "chai" +import chaiAsPromised from "chai-as-promised" + +import { encodeFunctionData, parseAbi } from "viem" import { generatePrivateKey, privateKeyToAccount } from "viem/accounts" import { eip712WalletActions, toSinglesigSmartAccount } from "viem/zksync" import { createWalletClient, createPublicClient, hashTypedData, http } from "viem" @@ -7,6 +9,9 @@ import { getViemChainFromConfig, writeContract } from "../tasks/utils" import { getGeneralPaymasterInput, serializeTransaction } from "viem/zksync" import hre from "hardhat"; +chai.use(chaiAsPromised); +const expect = chai.expect; + // Global test config const owner = privateKeyToAccount(hre.network.config.accounts[0]) const chain = getViemChainFromConfig() @@ -164,14 +169,12 @@ describe("ERC20 interactions from Openfort Account", function () { address: openfortAccountAddress, abi: parseAbi(["function registerSessionKey(address, uint48, uint48, uint48, address[]) external"]), functionName: "registerSessionKey", - // Session Key is valid for 24 hours args: [sessionKeyAccount.address, blockTimestamp, blockTimestamp + BigInt(24 * 60 * 60), 100, []], }) // sign with the new sessionKey const amount = BigInt(42) - // Make sure we sign with the session key const sessionKeyWalletClient = createWalletClient({ account: sessionKeyAccount, @@ -209,7 +212,7 @@ describe("ERC20 interactions from Openfort Account", function () { value: 0n, callData: encodeFunctionData({ abi: mintAbi, - functionName: 'mint', + functionName: "mint", args: [openfortAccountAddress, 10n] }) }, @@ -218,7 +221,7 @@ describe("ERC20 interactions from Openfort Account", function () { value: 0n, callData: encodeFunctionData({ abi: mintAbi, - functionName: 'mint', + functionName: "mint", args: [openfortAccountAddress, 20n] }) }, @@ -227,7 +230,7 @@ describe("ERC20 interactions from Openfort Account", function () { value: 0n, callData: encodeFunctionData({ abi: mintAbi, - functionName: 'mint', + functionName: "mint", args: [openfortAccountAddress, 30n] }) }, @@ -236,46 +239,45 @@ describe("ERC20 interactions from Openfort Account", function () { value: 0n, callData: encodeFunctionData({ abi: mintAbi, - functionName: 'mint', + functionName: "mint", args: [openfortAccountAddress, 40n] }) } ]; const abi = [ - { - inputs: [ - { - components: [ - { - internalType: "address", - name: "target", - type: "address" - }, - { - internalType: "uint256", - name: "value", - type: "uint256" - }, - { - internalType: "bytes", - name: "callData", - type: "bytes" - } + { + inputs: [ + { + components: [ + { + internalType: "address", + name: "target", + type: "address" + }, + { + internalType: "uint256", + name: "value", + type: "uint256" + }, + { + internalType: "bytes", + name: "callData", + type: "bytes" + } + ], + internalType: "struct Call[]", + name: "calls", + type: "tuple[]" + } ], - internalType: "struct Call[]", - name: "calls", - type: "tuple[]" - } - ], - name: "batchCall", - outputs: [], - stateMutability: "nonpayable", - type: "function" - } + name: "batchCall", + outputs: [], + stateMutability: "nonpayable", + type: "function" + } ]; - const data = encodeFunctionData({ abi: abi, functionName: "batchCall", @@ -337,4 +339,188 @@ describe("ERC20 interactions from Openfort Account", function () { // Assert that the final balance is the initial balance plus the sum of all minted amounts expect(finalBalance - initialBalance).to.equal(10n + 20n + 30n + 40n); }); + + + + it("batch calls to mint with session key should update balance accordingly", async function () { + await deployTokens() + const batchCallerAddress = "0x7219257B57d9546c1BC0649617d557Db09C92D23"; // salt 0x666 + const blockTimestamp = (await publicClient.getBlock()).timestamp + // generate a new private key + // to avoid Account contract reverts with "SessionKey already registered" + const sessionKey = generatePrivateKey() + const sessionKeyAccount = privateKeyToAccount(sessionKey) + + // setup openfort smart account with session key as signer + const accountWithSessionKey = toSinglesigSmartAccount({ + address: openfortAccountAddress, + privateKey: sessionKey + }) + + // register a new random sessionKey + await writeContract(walletClient, { + account: owner, + address: openfortAccountAddress, + abi: parseAbi(["function registerSessionKey(address, uint48, uint48, uint48, address[]) external"]), + functionName: "registerSessionKey", + // Session Key is valid for 24 hours + // limit to 4 calls, just enough for the batch -- the next call SHOULD fail + args: [sessionKeyAccount.address, blockTimestamp, blockTimestamp + BigInt(24 * 60 * 60), 4, []], + }) + + // Each call data for batches + const mintAbi = parseAbi(["function mint(address sender, uint256 amount) external"]) + + const initialBalance = await publicClient.readContract({ + account: accountWithOwner, + address: tokens.mockERC20, + abi: parseAbi(["function balanceOf(address owner) external view returns (uint256)"]), + functionName: "balanceOf", + args: [openfortAccountAddress], + }); + + const calls = [ + { + target: tokens.mockERC20, + value: 0n, + callData: encodeFunctionData({ + abi: mintAbi, + functionName: "mint", + args: [openfortAccountAddress, 10n] + }) + }, + { + target: tokens.mockERC20, + value: 0n, + callData: encodeFunctionData({ + abi: mintAbi, + functionName: "mint", + args: [openfortAccountAddress, 20n] + }) + }, + { + target: tokens.mockERC20, + value: 0n, + callData: encodeFunctionData({ + abi: mintAbi, + functionName: "mint", + args: [openfortAccountAddress, 30n] + }) + }, + { + target: tokens.mockERC20, + value: 0n, + callData: encodeFunctionData({ + abi: mintAbi, + functionName: "mint", + args: [openfortAccountAddress, 40n] + }) + } + ]; + + const abi = [ + { + inputs: [ + { + components: [ + { + internalType: "address", + name: "target", + type: "address" + }, + { + internalType: "uint256", + name: "value", + type: "uint256" + }, + { + internalType: "bytes", + name: "callData", + type: "bytes" + } + ], + internalType: "struct Call[]", + name: "calls", + type: "tuple[]" + } + ], + name: "batchCall", + outputs: [], + stateMutability: "nonpayable", + type: "function" + } + ]; + + const data = encodeFunctionData({ + abi: abi, + functionName: "batchCall", + args: [calls] + }) + + const paymaster = { + paymaster: process.env.SOPHON_TESTNET_PAYMASTER_ADDRESS as `0x${string}`, + paymasterInput: getGeneralPaymasterInput({ innerInput: new Uint8Array() }), + }; + + const transactionRequest = await walletClient.prepareTransactionRequest({ + type: "eip712", + account: accountWithOwner, + from: accountWithOwner.address, + chainId: chain.id, + to: batchCallerAddress, + data: data, + ...(chain.name === "Sophon" ? paymaster : {}), + + }) + + const signableTransaction = { + type: "eip712", + from: accountWithOwner.address, + chainId: chain.id, + // preparedTransactionRequest + + nonce: transactionRequest.nonce, + gas: transactionRequest.gas, + maxFeePerGas: transactionRequest.maxFeePerGas, + maxPriorityFeePerGas: transactionRequest.maxPriorityFeePerGas, + to: batchCallerAddress, + data: data, + ...(chain.name === "Sophon" ? paymaster : {}), + + }; + + const EIP712hash = hashTypedData(chain.custom.getEip712Domain(signableTransaction)) + const signature = await accountWithSessionKey.sign({ hash: EIP712hash }) + const signedTransaction = serializeTransaction({ + ...signableTransaction, + customSignature: signature, + }); + + const thehash = await publicClient.sendRawTransaction({ + serializedTransaction: signedTransaction, + }) + + console.log("Batch Call Transaction Hash", thehash); + + const finalBalance = await publicClient.readContract({ + account: accountWithOwner, + address: tokens.mockERC20, + abi: parseAbi(["function balanceOf(address owner) external view returns (uint256)"]), + functionName: "balanceOf", + args: [openfortAccountAddress], + }); + // Assert that the final balance is the initial balance plus the sum of all minted amounts + expect(finalBalance - initialBalance).to.equal(10n + 20n + 30n + 40n); + // Try to mint one more time with the session key + // SHOULD FAIL: since we set the usage limit to 4 and used them all in the batch. + await expect( + writeContract(walletClient, { + account: accountWithSessionKey, + address: tokens.mockERC20, + abi: parseAbi(["function mint(address sender, uint256 amount) external"]), + functionName: "mint", + args: [openfortAccountAddress, 10n] + }) + ).to.be.rejectedWith("failed to validate the transaction"); + }); }) \ No newline at end of file