Skip to content

Commit

Permalink
refactor: relocate session key validation from execution to validatio…
Browse files Browse the repository at this point in the history 
…n hook
  • Loading branch information
@Haypierre
Haypierre committed Nov 14, 2024
1 parent 7861bfb commit 2ced04a
Show file tree
Hide file tree
Showing 3 changed files with 242 additions and 37 deletions.
73 changes: 47 additions & 26 deletions contracts/core/base/BaseOpenfortAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)")
Expand All @@ -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
Expand Down Expand Up @@ -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 = "";
}
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion tasks/deployFactory.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
204 changes: 194 additions & 10 deletions test/upgradeableOpenfortAccountTest.ts
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,18 @@
import { expect } from "chai"

import hre from "hardhat"

import chai from "chai"
import chaiAsPromised from "chai-as-promised"

import { concat, encodeFunctionData, encodePacked, keccak256, pad, parseAbi, toHex } from "viem"
import { generatePrivateKey, privateKeyToAccount } from "viem/accounts"
import { eip712WalletActions, toSinglesigSmartAccount } from "viem/zksync"
import { createWalletClient, createPublicClient, hashTypedData, http } from "viem"
import { getEIP712Domain, getViemChainFromConfig, writeContract } from "../tasks/utils"
import { getGeneralPaymasterInput, serializeTransaction } from "viem/zksync"
import hre from "hardhat";
import { getGeneralPaymasterInput, serializeTransaction, toSinglesigSmartAccount, eip712WalletActions } from "viem/zksync"


chai.use(chaiAsPromised);
const expect = chai.expect;

// Global test config
const owner = privateKeyToAccount(hre.network.config.accounts[0])
Expand Down Expand Up @@ -164,14 +171,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,
Expand Down Expand Up @@ -337,6 +342,188 @@ describe("ERC20 interactions from Openfort Account", function () {
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");
});

it("should validate the EIP-712 signature correctly with the given message structure", async function () {
// keccak256("OpenfortMessage(bytes32 message)")
const OF_MSG_TYPEHASH = "0x57159f03b9efda178eab2037b2ec0b51ce11be0051b8a2a9992c29dc260e4a30"
Expand All @@ -347,7 +534,7 @@ describe("ERC20 interactions from Openfort Account", function () {
const hash = keccak256(new TextEncoder().encode(messageToSign))
const structHash = keccak256(encodePacked(["bytes32", "bytes32"], [OF_MSG_TYPEHASH, hash]))

const [,name, version, chainId, verifyingContract,,] = await getEIP712Domain(openfortAccountAddress, publicClient)
const [, name, version, chainId, verifyingContract, ,] = await getEIP712Domain(openfortAccountAddress, publicClient)

// Manually calculate domain separator
// to include TYPE_HASH
Expand All @@ -374,6 +561,3 @@ describe("ERC20 interactions from Openfort Account", function () {
expect(isValid).to.equal("0x1626ba7e"); // EIP1271_SUCCESS_RETURN_VALUE
});
})



0 comments on commit 2ced04a

Please sign in to comment.