Skip to content

Commit

Permalink
fix: make factory proxy work (#216)
Browse files Browse the repository at this point in the history 
* fix: make factory proxy work

* fix: renamings

* fix: make tests runnable multiple times in a row
  • Loading branch information
@ly0va
ly0va authored Dec 12, 2024
1 parent 4628172 commit e14bbf3
Show file tree
Hide file tree
Showing 5 changed files with 88 additions and 64 deletions.
61 changes: 33 additions & 28 deletions scripts/deploy.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<string> {
// 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);
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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 <value>) address must be provided to deploy factory";
throw "Account implementation (--implementation <value>) address must be provided to deploy beacon";
}
args = [cmd.implementation];
}
if (cmd.only == FACTORY_NAME) {
if (!cmd.implementation) {
throw "Beacon (--beacon <value>) address must be provided to deploy factory";
}
args = [cmd.implementation];
}
if (cmd.only == PAYMASTER_NAME) {
if (!cmd.factory || !cmd.sessions) {
throw "⛔️ Factory (--factory <value>) and SessionModule (--sessions <value>) addresses must be provided to deploy paymaster";
throw "Factory (--factory <value>) and SessionModule (--sessions <value>) 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);
}
}
});
16 changes: 6 additions & 10 deletions src/AAFactory.sol
View file
Original file line number Diff line number Diff line change
@@ -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";

Expand All @@ -11,23 +10,25 @@ import { ISsoAccount } from "./interfaces/ISsoAccount.sol";
/// @author Matter Labs
/// @custom:security-contact [email protected]
/// @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.
event AccountCreated(address indexed accountAddress, string uniqueAccountId);

/// @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.
Expand All @@ -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");
Expand Down
11 changes: 11 additions & 0 deletions src/SsoBeacon.sol
View file
Original file line number Diff line number Diff line change
@@ -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 [email protected]
contract SsoBeacon is UpgradeableBeacon {
constructor(address _implementation) UpgradeableBeacon(_implementation) {}
}
17 changes: 13 additions & 4 deletions test/SessionKeyTest.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand Down Expand Up @@ -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(
Expand All @@ -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()}`);
Expand Down Expand Up @@ -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());
});

Expand All @@ -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
Expand Down
47 changes: 25 additions & 22 deletions test/utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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;
}
Expand All @@ -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() {
Expand All @@ -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;
Expand All @@ -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<ERC20> {
Expand Down Expand Up @@ -141,15 +144,15 @@ export const getProviderL1 = () => {
return provider;
};

export async function deployFactory(wallet: Wallet, implAddress: string): Promise<AAFactory> {
export async function deployFactory(wallet: Wallet, beaconAddress: string): Promise<AAFactory> {
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"));

const deployer = new ContractFactory(factoryArtifact.abi, factoryArtifact.bytecode, wallet);
const bytecodeHash = utils.hashBytecode(proxyAaArtifact.bytecode);
const factory = await deployer.deploy(
bytecodeHash,
implAddress,
beaconAddress,
{ customData: { factoryDeps: [proxyAaArtifact.bytecode] } },
);
const factoryAddress = await factory.getAddress();
Expand All @@ -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,
});
}
Expand All @@ -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();
Expand All @@ -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`;
};

/**
Expand Down

0 comments on commit e14bbf3

Please sign in to comment.