From e14bbf36e36b07aaab0700adf3cb24e4d39c3552 Mon Sep 17 00:00:00 2001 From: Lyova Potyomkin Date: Thu, 12 Dec 2024 12:27:22 +0200 Subject: [PATCH] fix: make factory proxy work (#216) * fix: make factory proxy work * fix: renamings * fix: make tests runnable multiple times in a row --- scripts/deploy.ts | 61 +++++++++++++++++++++++------------------- src/AAFactory.sol | 16 +++++------ src/SsoBeacon.sol | 11 ++++++++ test/SessionKeyTest.ts | 17 +++++++++--- test/utils.ts | 47 +++++++++++++++++--------------- 5 files changed, 88 insertions(+), 64 deletions(-) create mode 100644 src/SsoBeacon.sol diff --git a/scripts/deploy.ts b/scripts/deploy.ts index cb8ec9a7..11fc36a0 100644 --- a/scripts/deploy.ts +++ b/scripts/deploy.ts @@ -11,17 +11,20 @@ const ethersStaticSalt = new Uint8Array([ 62, 140, 111, 128, 47, 32, 21, 177, 177, 174, 166, ]); + +const WEBAUTH_NAME = "WebAuthValidator"; const SESSIONS_NAME = "SessionKeyValidator"; const ACCOUNT_IMPL_NAME = "SsoAccount"; -const AA_FACTORY_NAME = "AAFactory"; +const FACTORY_NAME = "AAFactory"; const PAYMASTER_NAME = "ExampleAuthServerPaymaster"; +const BEACON_NAME = "SsoBeacon"; async function deploy(name: string, deployer: Wallet, proxy: boolean, args?: any[]): Promise { // eslint-disable-next-line @typescript-eslint/no-require-imports const { deployFactory, create2 } = require("../test/utils"); console.log("Deploying", name, "contract..."); let implContract; - if (name == AA_FACTORY_NAME) { + if (name == FACTORY_NAME) { implContract = await deployFactory(deployer, args![0]); } else { implContract = await create2(name, deployer, ethersStaticSalt, args); @@ -37,13 +40,15 @@ async function deploy(name: string, deployer: Wallet, proxy: boolean, args?: any return proxyAddress; } + task("deploy", "Deploys ZKsync SSO contracts") .addOptionalParam("privatekey", "private key to the account to deploy the contracts from") .addOptionalParam("only", "name of a specific contract to deploy") .addFlag("noProxy", "do not deploy transparent proxies for factory and modules") - .addOptionalParam("implementation", "address of the account implementation to use in the factory") + .addOptionalParam("implementation", "address of the account implementation to use in the beacon") .addOptionalParam("factory", "address of the factory to use in the paymaster") .addOptionalParam("sessions", "address of the sessions module to use in the paymaster") + .addOptionalParam("beacon", "address of the beacon to use in the factory") .addOptionalParam("fund", "amount of ETH to send to the paymaster", "0") .setAction(async (cmd, hre) => { // eslint-disable-next-line @typescript-eslint/no-require-imports @@ -64,16 +69,9 @@ task("deploy", "Deploys ZKsync SSO contracts") const deployer = new Wallet(privateKey, provider); console.log("Deployer address:", deployer.address); - if (!cmd.only) { - await deploy("WebAuthValidator", deployer, !cmd.noProxy); - const sessions = await deploy(SESSIONS_NAME, deployer, !cmd.noProxy); - const implementation = await deploy(ACCOUNT_IMPL_NAME, deployer, false); - // TODO: enable proxy for factory -- currently it's not working - const factory = await deploy(AA_FACTORY_NAME, deployer, false, [implementation]); - const paymaster = await deploy(PAYMASTER_NAME, deployer, false, [factory, sessions]); - - if (cmd.fund != 0) { - console.log("Funding paymaster with", cmd.fund, "ETH..."); + async function fundPaymaster(paymaster: string, fund?: string | number) { + if (fund && fund != 0) { + console.log("Funding paymaster with", fund, "ETH..."); await ( await deployer.sendTransaction({ to: paymaster, @@ -84,35 +82,42 @@ task("deploy", "Deploys ZKsync SSO contracts") } else { console.log("--fund flag not provided, skipping funding paymaster\n"); } + } + + if (!cmd.only) { + await deploy(WEBAUTH_NAME, deployer, !cmd.noProxy); + const sessions = await deploy(SESSIONS_NAME, deployer, !cmd.noProxy); + const implementation = await deploy(ACCOUNT_IMPL_NAME, deployer, false); + const beacon = await deploy(BEACON_NAME, deployer, false, [implementation]); + const factory = await deploy(FACTORY_NAME, deployer, !cmd.noProxy, [beacon]); + const paymaster = await deploy(PAYMASTER_NAME, deployer, false, [factory, sessions]); + + await fundPaymaster(paymaster, cmd.fund); } else { let args: any[] = []; - if (cmd.only == AA_FACTORY_NAME) { + + if (cmd.only == BEACON_NAME) { if (!cmd.implementation) { - throw "⛔️ Implementation (--implementation ) address must be provided to deploy factory"; + throw "Account implementation (--implementation ) address must be provided to deploy beacon"; + } + args = [cmd.implementation]; + } + if (cmd.only == FACTORY_NAME) { + if (!cmd.implementation) { + throw "Beacon (--beacon ) address must be provided to deploy factory"; } args = [cmd.implementation]; } if (cmd.only == PAYMASTER_NAME) { if (!cmd.factory || !cmd.sessions) { - throw "⛔️ Factory (--factory ) and SessionModule (--sessions ) addresses must be provided to deploy paymaster"; + throw "Factory (--factory ) and SessionModule (--sessions ) addresses must be provided to deploy paymaster"; } args = [cmd.factory, cmd.sessions]; } const deployedContract = await deploy(cmd.only, deployer, false, args); if (cmd.only == PAYMASTER_NAME) { - if (cmd.fund == 0) { - console.log("⚠️ Paymaster is unfunded ⚠️\n"); - } else { - console.log("Funding paymaster with", cmd.fund, "ETH..."); - await ( - await deployer.sendTransaction({ - to: deployedContract, - value: ethers.parseEther(cmd.fund), - }) - ).wait(); - console.log("Paymaster funded\n"); - } + await fundPaymaster(deployedContract, cmd.fund); } } }); diff --git a/src/AAFactory.sol b/src/AAFactory.sol index 1cb9ccda..5bcaf5ca 100644 --- a/src/AAFactory.sol +++ b/src/AAFactory.sol @@ -1,7 +1,6 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.24; -import { UpgradeableBeacon } from "@openzeppelin/contracts/proxy/beacon/UpgradeableBeacon.sol"; import { DEPLOYER_SYSTEM_CONTRACT, IContractDeployer } from "@matterlabs/zksync-contracts/l2/system-contracts/Constants.sol"; import { SystemContractsCaller } from "@matterlabs/zksync-contracts/l2/system-contracts/libraries/SystemContractsCaller.sol"; @@ -11,7 +10,7 @@ import { ISsoAccount } from "./interfaces/ISsoAccount.sol"; /// @author Matter Labs /// @custom:security-contact security@matterlabs.dev /// @dev This contract is used to deploy SSO accounts as beacon proxies. -contract AAFactory is UpgradeableBeacon { +contract AAFactory { /// @notice Emitted when a new account is successfully created. /// @param accountAddress The address of the newly created account. /// @param uniqueAccountId A unique identifier for the account. @@ -19,15 +18,17 @@ contract AAFactory is UpgradeableBeacon { /// @dev The bytecode hash of the beacon proxy, used for deploying proxy accounts. bytes32 private immutable beaconProxyBytecodeHash; + address private immutable beacon; /// @notice A mapping from unique account IDs to their corresponding deployed account addresses. mapping(string => address) public accountMappings; /// @notice Constructor that initializes the factory with a beacon proxy bytecode hash and implementation contract address. /// @param _beaconProxyBytecodeHash The bytecode hash of the beacon proxy. - /// @param _implementation The address of the implementation contract used by the beacon. - constructor(bytes32 _beaconProxyBytecodeHash, address _implementation) UpgradeableBeacon(_implementation) { + /// @param _beacon The address of the UpgradeableBeacon contract used for the SSO accounts' beacon proxies. + constructor(bytes32 _beaconProxyBytecodeHash, address _beacon) { beaconProxyBytecodeHash = _beaconProxyBytecodeHash; + beacon = _beacon; } /// @notice Deploys a new SSO account as a beacon proxy with the specified parameters. @@ -51,12 +52,7 @@ contract AAFactory is UpgradeableBeacon { uint128(0), abi.encodeCall( DEPLOYER_SYSTEM_CONTRACT.create2Account, - ( - _salt, - beaconProxyBytecodeHash, - abi.encode(address(this)), - IContractDeployer.AccountAbstractionVersion.Version1 - ) + (_salt, beaconProxyBytecodeHash, abi.encode(beacon), IContractDeployer.AccountAbstractionVersion.Version1) ) ); require(success, "Deployment failed"); diff --git a/src/SsoBeacon.sol b/src/SsoBeacon.sol new file mode 100644 index 00000000..712304a3 --- /dev/null +++ b/src/SsoBeacon.sol @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import { UpgradeableBeacon } from "@openzeppelin/contracts/proxy/beacon/UpgradeableBeacon.sol"; + +/// @title SsoBeacon +/// @author Matter Labs +/// @custom:security-contact security@matterlabs.dev +contract SsoBeacon is UpgradeableBeacon { + constructor(address _implementation) UpgradeableBeacon(_implementation) {} +} diff --git a/test/SessionKeyTest.ts b/test/SessionKeyTest.ts index 44393eec..281d0f50 100644 --- a/test/SessionKeyTest.ts +++ b/test/SessionKeyTest.ts @@ -6,8 +6,8 @@ import { it } from "mocha"; import { SmartAccount, utils } from "zksync-ethers"; import type { ERC20 } from "../typechain-types"; -import { AAFactory__factory, SsoAccount__factory } from "../typechain-types"; -import type { AAFactory } from "../typechain-types/src/AAFactory"; +import { SsoBeacon__factory, SsoAccount__factory } from "../typechain-types"; +import type { SsoBeacon } from "../typechain-types/src/SsoBeacon"; import type { SessionLib } from "../typechain-types/src/validators/SessionKeyValidator"; import { ContractFixtures, getProvider, logInfo } from "./utils"; @@ -365,6 +365,8 @@ describe("SessionKeyModule tests", function () { assert(sessionModuleContract != null, "No session module deployed"); const ssoContract = await fixtures.getAccountImplContract(); assert(ssoContract != null, "No SSO Account deployed"); + const beaconContract = await fixtures.getBeaconContract(); + assert(beaconContract != null, "No Beacon deployed"); const factoryContract = await fixtures.getAaFactory(); assert(factoryContract != null, "No AA Factory deployed"); const authServerPaymaster = await fixtures.deployExampleAuthServerPaymaster( @@ -373,6 +375,8 @@ describe("SessionKeyModule tests", function () { ); assert(authServerPaymaster != null, "No Auth Server Paymaster deployed"); + logInfo(`SSO Account Address : ${await ssoContract.getAddress()}`); + logInfo(`Beacon Address : ${await beaconContract.getAddress()}`); logInfo(`Session Address : ${await sessionModuleContract.getAddress()}`); logInfo(`Passkey Address : ${await verifierContract.getAddress()}`); logInfo(`Account Factory Address : ${await factoryContract.getAddress()}`); @@ -579,10 +583,10 @@ describe("SessionKeyModule tests", function () { }); describe("Upgrade tests", function () { - let beacon: AAFactory; + let beacon: SsoBeacon; it("should check implementation address", async () => { - beacon = AAFactory__factory.connect(await fixtures.getAaFactoryAddress(), fixtures.wallet); + beacon = SsoBeacon__factory.connect(await fixtures.getBeaconAddress(), fixtures.wallet); expect(await beacon.implementation()).to.equal(await fixtures.getAccountImplAddress()); }); @@ -596,6 +600,11 @@ describe("SessionKeyModule tests", function () { const proxy = new ethers.Contract(proxyAccountAddress, ["function dummy() pure returns (string)"], provider); expect(await proxy.dummy()).to.equal("dummy"); }); + + it("should upgrade implementation back", async () => { + await beacon.upgradeTo(await fixtures.getAccountImplAddress()); + expect(await beacon.implementation()).to.equal(await fixtures.getAccountImplAddress()); + }); }); // TODO: module uninstall tests diff --git a/test/utils.ts b/test/utils.ts index d4fa7ca4..cdde5a25 100644 --- a/test/utils.ts +++ b/test/utils.ts @@ -9,8 +9,8 @@ import * as hre from "hardhat"; import { ContractFactory, Provider, utils, Wallet } from "zksync-ethers"; import { base64UrlToUint8Array, getPublicKeyBytesFromPasskeySignature, unwrapEC2Signature } from "zksync-sso/utils"; -import { AAFactory, ERC20, ExampleAuthServerPaymaster, SessionKeyValidator, SsoAccount, WebAuthValidator } from "../typechain-types"; -import { AAFactory__factory, ERC20__factory, ExampleAuthServerPaymaster__factory, SessionKeyValidator__factory, SsoAccount__factory, WebAuthValidator__factory } from "../typechain-types"; +import { AAFactory, ERC20, ExampleAuthServerPaymaster, SessionKeyValidator, SsoAccount, WebAuthValidator, SsoBeacon } from "../typechain-types"; +import { AAFactory__factory, ERC20__factory, ExampleAuthServerPaymaster__factory, SessionKeyValidator__factory, SsoAccount__factory, WebAuthValidator__factory, SsoBeacon__factory } from "../typechain-types"; export class ContractFixtures { readonly wallet: Wallet = getWallet(LOCAL_RICH_WALLETS[0].privateKey); @@ -24,9 +24,9 @@ export class ContractFixtures { private _aaFactory: AAFactory; async getAaFactory() { - const implAddress = await this.getAccountImplAddress(); + const beaconAddress = await this.getBeaconAddress(); if (!this._aaFactory) { - this._aaFactory = await deployFactory(this.wallet, implAddress); + this._aaFactory = await deployFactory(this.wallet, beaconAddress); } return this._aaFactory; } @@ -48,6 +48,20 @@ export class ContractFixtures { return (await this.getSessionKeyContract()).getAddress(); } + private _beacon: SsoBeacon; + async getBeaconContract() { + if (!this._beacon) { + const implAddress = await this.getAccountImplAddress(); + const contract = await create2("SsoBeacon", this.wallet, this.ethersStaticSalt, [implAddress]); + this._beacon = SsoBeacon__factory.connect(await contract.getAddress(), this.wallet); + } + return this._beacon; + } + + async getBeaconAddress() { + return (await this.getBeaconContract()).getAddress(); + } + private _webauthnValidatorModule: WebAuthValidator; // does passkey validation via modular interface async getWebAuthnVerifierContract() { @@ -58,13 +72,8 @@ export class ContractFixtures { return this._webauthnValidatorModule; } - private _passkeyModuleAddress: string; async getPasskeyModuleAddress() { - if (!this._passkeyModuleAddress) { - const passkeyModule = await this.getWebAuthnVerifierContract(); - this._passkeyModuleAddress = await passkeyModule.getAddress(); - } - return this._passkeyModuleAddress; + return (await this.getWebAuthnVerifierContract()).getAddress(); } private _accountImplContract: SsoAccount; @@ -76,14 +85,8 @@ export class ContractFixtures { return this._accountImplContract; } - private _accountImplAddress: string; - // deploys the base account for future proxy use async getAccountImplAddress() { - if (!this._accountImplAddress) { - const accountImpl = await this.getAccountImplContract(); - this._accountImplAddress = await accountImpl.getAddress(); - } - return this._accountImplAddress; + return (await this.getAccountImplContract()).getAddress(); } async deployERC20(mintTo: string): Promise { @@ -141,7 +144,7 @@ export const getProviderL1 = () => { return provider; }; -export async function deployFactory(wallet: Wallet, implAddress: string): Promise { +export async function deployFactory(wallet: Wallet, beaconAddress: string): Promise { const factoryArtifact = JSON.parse(await promises.readFile("artifacts-zk/src/AAFactory.sol/AAFactory.json", "utf8")); const proxyAaArtifact = JSON.parse(await promises.readFile("artifacts-zk/src/AccountProxy.sol/AccountProxy.json", "utf8")); @@ -149,7 +152,7 @@ export async function deployFactory(wallet: Wallet, implAddress: string): Promis const bytecodeHash = utils.hashBytecode(proxyAaArtifact.bytecode); const factory = await deployer.deploy( bytecodeHash, - implAddress, + beaconAddress, { customData: { factoryDeps: [proxyAaArtifact.bytecode] } }, ); const factoryAddress = await factory.getAddress(); @@ -160,7 +163,7 @@ export async function deployFactory(wallet: Wallet, implAddress: string): Promis await verifyContract({ address: factoryAddress, contract: `src/AAFactory.sol:AAFactory`, - constructorArguments: deployer.interface.encodeDeploy([bytecodeHash, implAddress]), + constructorArguments: deployer.interface.encodeDeploy([bytecodeHash, beaconAddress]), bytecode: factoryArtifact.bytecode, }); } @@ -171,7 +174,7 @@ export async function deployFactory(wallet: Wallet, implAddress: string): Promis export const getWallet = (privateKey?: string) => { if (!privateKey) { // Get wallet private key from .env file - if (!process.env.WALLET_PRIVATE_KEY) throw "⛔️ Wallet private key wasn't found in .env file!"; + if (!process.env.WALLET_PRIVATE_KEY) throw "Wallet private key wasn't found in .env file!"; } const provider = getProvider(); @@ -186,7 +189,7 @@ export const getWallet = (privateKey?: string) => { export const verifyEnoughBalance = async (wallet: Wallet, amount: bigint) => { // Check if the wallet has enough balance const balance = await wallet.getBalance(); - if (balance < amount) throw `⛔️ Wallet balance is too low! Required ${ethers.formatEther(amount)} ETH, but current ${wallet.address} balance is ${ethers.formatEther(balance)} ETH`; + if (balance < amount) throw `Wallet balance is too low! Required ${ethers.formatEther(amount)} ETH, but current ${wallet.address} balance is ${ethers.formatEther(balance)} ETH`; }; /**