Skip to content

Commit

Permalink
test: hook tests (#260)
Browse files Browse the repository at this point in the history
* test: dummy hook test

* test: add hook tests

* chore: add fixme
ly0va authored Jan 30, 2025
1 parent fc0af34 commit e3ecf28
Showing 8 changed files with 426 additions and 13 deletions.
6 changes: 4 additions & 2 deletions src/interfaces/IHook.sol
Original file line number Diff line number Diff line change
@@ -23,8 +23,10 @@ interface IValidationHook is IModule, IERC165 {
interface IExecutionHook is IModule, IERC165 {
/// @notice Hook that triggers before each transaction during the execution phase.
/// @param transaction Transaction that is being executed.
function preExecutionHook(Transaction calldata transaction) external;
/// @return context Context that is passed to the postExecutionHook.
function preExecutionHook(Transaction calldata transaction) external returns (bytes memory context);

/// @notice Hook that triggers after each transaction during the execution phase.
function postExecutionHook() external;
/// @param context Context that was returned by the preExecutionHook.
function postExecutionHook(bytes calldata context) external;
}
19 changes: 13 additions & 6 deletions src/managers/HookManager.sol
Original file line number Diff line number Diff line change
@@ -77,15 +77,22 @@ abstract contract HookManager is IHookManager, Auth {
modifier runExecutionHooks(Transaction calldata transaction) {
EnumerableSet.AddressSet storage hookList = _executionHooks();
uint256 totalHooks = hookList.length();
bytes[] memory context = new bytes[](totalHooks);

for (uint256 i = 0; i < totalHooks; i++) {
IExecutionHook(hookList.at(i)).preExecutionHook(transaction);
context[i] = IExecutionHook(hookList.at(i)).preExecutionHook(transaction);
}

_;

// If we removed any hooks, we have to update totalHooks.
// If we added any hooks, we don't want them to run yet.
if (totalHooks > hookList.length()) {
totalHooks = hookList.length();
}

for (uint256 i = 0; i < totalHooks; i++) {
IExecutionHook(hookList.at(i)).postExecutionHook();
IExecutionHook(hookList.at(i)).postExecutionHook(context[i]);
}
}

@@ -95,9 +102,9 @@ abstract contract HookManager is IHookManager, Auth {
}

if (isValidation) {
_validationHooks().add(hook);
require(_validationHooks().add(hook), "Hook already installed");
} else {
_executionHooks().add(hook);
require(_executionHooks().add(hook), "Hook already installed");
}

IModule(hook).onInstall(initData);
@@ -107,9 +114,9 @@ abstract contract HookManager is IHookManager, Auth {

function _removeHook(address hook, bool isValidation) internal {
if (isValidation) {
_validationHooks().remove(hook);
require(_validationHooks().remove(hook), "Hook not found");
} else {
_executionHooks().remove(hook);
require(_executionHooks().remove(hook), "Hook not found");
}

emit HookRemoved(hook);
4 changes: 2 additions & 2 deletions src/managers/OwnerManager.sol
Original file line number Diff line number Diff line change
@@ -38,13 +38,13 @@ abstract contract OwnerManager is IOwnerManager, Auth {
}

function _k1AddOwner(address addr) internal {
_k1Owners().add(addr);
require(_k1Owners().add(addr), "K1 owner already exists");

emit K1AddOwner(addr);
}

function _k1RemoveOwner(address addr) internal {
_k1Owners().remove(addr);
require(_k1Owners().remove(addr), "K1 owner not found");

emit K1RemoveOwner(addr);
}
4 changes: 2 additions & 2 deletions src/managers/ValidatorManager.sol
Original file line number Diff line number Diff line change
@@ -57,14 +57,14 @@ abstract contract ValidatorManager is IValidatorManager, Auth {
revert Errors.VALIDATOR_ERC165_FAIL(validator);
}

_moduleValidators().add(validator);
require(_moduleValidators().add(validator), "Validator already exists");
IModule(validator).onInstall(initData);

emit ValidatorAdded(validator);
}

function _removeModuleValidator(address validator) internal {
_moduleValidators().remove(validator);
require(_moduleValidators().remove(validator), "Validator not found");

emit ValidatorRemoved(validator);
}
89 changes: 89 additions & 0 deletions src/test/TestExecutionHook.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.28;

import { Transaction, TransactionHelper } from "@matterlabs/zksync-contracts/l2/system-contracts/libraries/TransactionHelper.sol";
import { IExecutionHook } from "../interfaces/IHook.sol";
import { IModule } from "../interfaces/IModule.sol";
import { IERC165 } from "@openzeppelin/contracts/utils/introspection/IERC165.sol";

contract TestExecutionHook is IExecutionHook {
using TransactionHelper for Transaction;

event ExecutionHookInstalled(address indexed account);
event ExecutionHookUninstalled(address indexed account);

event PreExecution(address indexed account, address indexed target);
event PostExecution(address indexed account, address indexed target);

mapping(address => address) public lastTarget;

function onInstall(bytes calldata data) external {
bool shouldRevert = abi.decode(data, (bool));
if (shouldRevert) {
revert("Install hook failed");
}
emit ExecutionHookInstalled(msg.sender);
}

function onUninstall(bytes calldata data) external {
bool shouldRevert = abi.decode(data, (bool));
if (shouldRevert) {
revert("Uninstall hook failed");
}
emit ExecutionHookUninstalled(msg.sender);
}

function supportsInterface(bytes4 interfaceId) external pure override returns (bool) {
return
interfaceId == type(IExecutionHook).interfaceId ||
interfaceId == type(IModule).interfaceId ||
interfaceId == type(IERC165).interfaceId;
}

function preExecutionHook(Transaction calldata transaction) external returns (bytes memory context) {
// arbitrary revert condition
if (transaction.to == 0) {
revert("PreExecution hook failed");
}

// store some data in transient storage
uint256 slot = uint256(uint160(msg.sender));
uint256 storedTarget = transaction.to;
assembly {
tstore(slot, storedTarget)
}

// store some data in regular storage
lastTarget[msg.sender] = address(uint160(transaction.to));

// emit event
emit PreExecution(msg.sender, address(uint160(transaction.to)));

// pass some data via context
return abi.encode(transaction);
}

function postExecutionHook(bytes calldata context) external {
// decode context data
Transaction memory transaction = abi.decode(context, (Transaction));

// arbitrary revert condition
if (transaction.to == uint256(uint160(msg.sender))) {
revert("PostExecution hook failed");
}

// load data from transient storage
uint256 storedTarget;
uint256 slot = uint256(uint160(msg.sender));
assembly {
storedTarget := tload(slot)
}

require(storedTarget != 0, "No data in transient storage");
require(transaction.to == storedTarget, "Targets do not match (tload)");
require(transaction.to == uint256(uint160(lastTarget[msg.sender])), "Targets do not match (sload)");

// emit event
emit PostExecution(msg.sender, address(uint160(transaction.to)));
}
}
55 changes: 55 additions & 0 deletions src/test/TestValidationHook.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.28;

import { Transaction, TransactionHelper } from "@matterlabs/zksync-contracts/l2/system-contracts/libraries/TransactionHelper.sol";
import { IValidationHook } from "../interfaces/IHook.sol";
import { IModule } from "../interfaces/IModule.sol";
import { IERC165 } from "@openzeppelin/contracts/utils/introspection/IERC165.sol";

contract TestValidationHook is IValidationHook {
using TransactionHelper for Transaction;

event ValidationHookInstalled(address indexed account);
event ValidationHookUninstalled(address indexed account);
event ValidationHookTriggered(address indexed account, address indexed target);

mapping(address => address) public lastTarget;

function onInstall(bytes calldata data) external {
bool shouldRevert = abi.decode(data, (bool));
if (shouldRevert) {
revert("Install hook failed");
}
emit ValidationHookInstalled(msg.sender);
}

function onUninstall(bytes calldata data) external {
bool shouldRevert = abi.decode(data, (bool));
if (shouldRevert) {
revert("Uninstall hook failed");
}
emit ValidationHookUninstalled(msg.sender);
}

function supportsInterface(bytes4 interfaceId) external pure override returns (bool) {
return
interfaceId == type(IValidationHook).interfaceId ||
interfaceId == type(IModule).interfaceId ||
interfaceId == type(IERC165).interfaceId;
}

// FIXME: if a validation hook were to always revert, the account would be bricked
function validationHook(bytes32 signedHash, Transaction calldata transaction) external {
if (transaction.data.length == 0) {
revert("Empty calldata not allowed");
}

address target = address(uint160(transaction.to));

// emit event
emit ValidationHookTriggered(msg.sender, target);

// store some data in storage
lastTarget[msg.sender] = target;
}
}
260 changes: 260 additions & 0 deletions test/HooksTest.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,260 @@
import { assert, expect } from "chai";
import { ethers, parseEther, randomBytes } from "ethers";
import { Wallet, ZeroAddress } from "ethers";
import { it } from "mocha";
import { SmartAccount, utils } from "zksync-ethers";
import { create2 } from "./utils";

import { SsoAccount__factory, TestExecutionHook__factory, TestValidationHook__factory } from "../typechain-types";
import type { TestExecutionHook, TestValidationHook, SsoAccount } from "../typechain-types";
import { ContractFixtures, getProvider } from "./utils";

const ssoAccountAbi = SsoAccount__factory.createInterface();

describe("Hook tests", function () {
const fixtures = new ContractFixtures();
const provider = getProvider();
let proxyAccountAddress: string;
let ssoAccount: SsoAccount;
let executionHook: TestExecutionHook;
let validationHook: TestValidationHook;
const abi = new ethers.AbiCoder();
let smartAccount: SmartAccount;

async function aaTxTemplate() {
return {
type: 113,
from: proxyAccountAddress,
data: "0x",
value: 0,
chainId: (await provider.getNetwork()).chainId,
nonce: await provider.getTransactionCount(proxyAccountAddress),
gasPrice: await provider.getGasPrice(),
customData: { gasPerPubdata: utils.DEFAULT_GAS_PER_PUBDATA_LIMIT },
gasLimit: 500_000n,
};
}

it("should deploy proxy account and hooks", async () => {
const accountImplContract = await fixtures.getAccountImplContract();
assert(accountImplContract != null, "No account impl deployed");
const aaFactoryContract = await fixtures.getAaFactory();

const randomSalt = randomBytes(32);
const deployTx = await aaFactoryContract.deployProxySsoAccount(
randomSalt,
"id" + randomBytes(32).toString(),
[],
[fixtures.wallet.address],
);
const deployTxReceipt = await deployTx.wait();
proxyAccountAddress = deployTxReceipt!.contractAddress!;
ssoAccount = SsoAccount__factory.connect(proxyAccountAddress, fixtures.wallet);
smartAccount = new SmartAccount({
address: proxyAccountAddress,
secret: fixtures.wallet.privateKey,
}, provider);

const validationHookContract = await create2("TestValidationHook", fixtures.wallet, randomSalt);
validationHook = TestValidationHook__factory.connect(await validationHookContract.getAddress(), fixtures.wallet);

const executionHookContract = await create2("TestExecutionHook", fixtures.wallet, randomSalt);
executionHook = TestExecutionHook__factory.connect(await executionHookContract.getAddress(), fixtures.wallet);

const fundTx = await fixtures.wallet.sendTransaction({ value: parseEther("0.2"), to: proxyAccountAddress });
await fundTx.wait();
});

describe("Validation hook tests", function () {
it("should revert on install", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("addHook", [await validationHook.getAddress(), true, abi.encode(["bool"], [true])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;
});

it("should install hook", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("addHook", [await validationHook.getAddress(), true, abi.encode(["bool"], [false])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx))
.to.emit(validationHook, "ValidationHookInstalled")
.and.to.emit(ssoAccount, "HookAdded");
expect(await ssoAccount.isHook(await validationHook.getAddress())).to.be.true;
});

it("should fail installing existing hook", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("addHook", [await validationHook.getAddress(), true, abi.encode(["bool"], [false])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;
});

it("should revert while sending a transaction", async () => {
const aaTx = {
...await aaTxTemplate(),
to: Wallet.createRandom().address,
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;
});

it("should send transaction and emit event", async () => {
const aaTx = {
...await aaTxTemplate(),
to: Wallet.createRandom().address,
data: "0x1234",
};

const signedTx = await smartAccount.signTransaction(aaTx);
const tx = await provider.broadcastTransaction(signedTx);
await expect(tx).to.emit(validationHook, "ValidationHookTriggered");
expect(await validationHook.lastTarget(proxyAccountAddress)).to.equal(tx.to);
});

it("should revert on uninstall", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("removeHook", [await validationHook.getAddress(), true, abi.encode(["bool"], [true])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;
});

it("should uninstall hook", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("removeHook", [await validationHook.getAddress(), true, abi.encode(["bool"], [false])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx))
.to.emit(validationHook, "ValidationHookUninstalled")
.and.to.emit(ssoAccount, "HookRemoved");
expect(await ssoAccount.isHook(await validationHook.getAddress())).to.be.false;
});

it("should fail uninstalling already uninstalled hook", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("removeHook", [await validationHook.getAddress(), true, abi.encode(["bool"], [false])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;
});
});

describe("Execution hook tests", function () {
it("should revert on install", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("addHook", [await executionHook.getAddress(), false, abi.encode(["bool"], [true])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;
});

it("should install hook", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("addHook", [await executionHook.getAddress(), false, abi.encode(["bool"], [false])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx))
.to.emit(executionHook, "ExecutionHookInstalled")
.and.to.emit(ssoAccount, "HookAdded");
expect(await ssoAccount.isHook(await executionHook.getAddress())).to.be.true;
});

it("should revert while sending a transaction", async () => {
// pre execution reverts
let aaTx = {
...await aaTxTemplate(),
to: ZeroAddress
};

let signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;

// post execution reverts
aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress
};

signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;
});

it("should send transaction and emit event", async () => {
const aaTx = {
...await aaTxTemplate(),
to: Wallet.createRandom().address,
};

const signedTx = await smartAccount.signTransaction(aaTx);
const tx = await provider.broadcastTransaction(signedTx);
await expect(tx)
.to.emit(executionHook, "PreExecution")
.and.to.emit(executionHook, "PostExecution");
expect(await executionHook.lastTarget(proxyAccountAddress)).to.equal(tx.to);
});

it("should revert on uninstall", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("removeHook", [await executionHook.getAddress(), false, abi.encode(["bool"], [true])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;
});

it("should unlink hook", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("unlinkHook", [await executionHook.getAddress(), false, abi.encode(["bool"], [false])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx))
.and.to.emit(ssoAccount, "HookRemoved");
expect(await ssoAccount.isHook(await executionHook.getAddress())).to.be.false;
});

it("should fail unlinking already unlinked hook", async () => {
const aaTx = {
...await aaTxTemplate(),
to: proxyAccountAddress,
data: ssoAccountAbi.encodeFunctionData("unlinkHook", [await executionHook.getAddress(), false, abi.encode(["bool"], [false])]),
};

const signedTx = await smartAccount.signTransaction(aaTx);
await expect(provider.broadcastTransaction(signedTx)).to.be.reverted;
});
});
});
2 changes: 1 addition & 1 deletion test/PasskeyModule.ts
Original file line number Diff line number Diff line change
@@ -393,7 +393,7 @@ function encodeKeyFromHex(hexStrings: [Hex, Hex], domain: string) {
{ name: 'publicKeys', type: 'bytes32[2]' },
{ name: 'domain', type: 'string' },
],
[[hexStrings[0], hexStrings[1]], domain]
[hexStrings, domain]
)
}

0 comments on commit e3ecf28

Please sign in to comment.