Merge branch 'feat/oidc-account-recovery' into matias/test-key-registry

package.json

@@ -37,6 +37,7 @@
     "@nomad-xyz/excessively-safe-call": "0.0.1-rc.1",
     "@nomicfoundation/hardhat-chai-matchers": "2.0.8",
     "@nomicfoundation/hardhat-ethers": "3.0.8",
+    "@nomicfoundation/hardhat-network-helpers": "^1.0.12",
     "@nomicfoundation/hardhat-toolbox": "^5.0.0",
     "@nomicfoundation/hardhat-verify": "2.0.11",
     "@openzeppelin/contracts": "4.9.6",

scripts/deploy.ts

@@ -6,6 +6,7 @@ import { Wallet } from "zksync-ethers";
 
 const WEBAUTH_NAME = "WebAuthValidator";
 const SESSIONS_NAME = "SessionKeyValidator";
+const GUARDIAN_RECOVERY_NAME = "GuardianRecoveryValidator";
 const ACCOUNT_IMPL_NAME = "SsoAccount";
 const FACTORY_NAME = "AAFactory";
 const PAYMASTER_NAME = "ExampleAuthServerPaymaster";
@@ -33,7 +34,6 @@ async function deploy(name: string, deployer: Wallet, proxy: boolean, args?: any
   return proxyAddress;
 }
 
-
 task("deploy", "Deploys ZKsync SSO contracts")
   .addOptionalParam("only", "name of a specific contract to deploy")
   .addFlag("noProxy", "do not deploy transparent proxies for factory and modules")
@@ -77,12 +77,14 @@ task("deploy", "Deploys ZKsync SSO contracts")
     }
 
     if (!cmd.only) {
-      await deploy(WEBAUTH_NAME, deployer, !cmd.noProxy);
+      const webauth = 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]);
+      const guardianInterface = new ethers.Interface((await hre.artifacts.readArtifact(GUARDIAN_RECOVERY_NAME)).abi);
+      const recovery = await deploy(GUARDIAN_RECOVERY_NAME, deployer, !cmd.noProxy, [webauth, factory], guardianInterface.encodeFunctionData("initialize", [webauth, factory]));
+      const paymaster = await deploy(PAYMASTER_NAME, deployer, false, [factory, sessions, recovery]);
       const oidcKeyRegistryInterface = new ethers.Interface((await hre.artifacts.readArtifact(OIDC_KEY_REGISTRY_NAME)).abi);
       await deploy(OIDC_KEY_REGISTRY_NAME, deployer, !cmd.noProxy, [], oidcKeyRegistryInterface.encodeFunctionData("initialize", []));
 

src/AAFactory.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.24;
 import { DEPLOYER_SYSTEM_CONTRACT } from "@matterlabs/zksync-contracts/l2/system-contracts/Constants.sol";
 import { IContractDeployer } from "@matterlabs/zksync-contracts/l2/system-contracts/interfaces/IContractDeployer.sol";
 import { SystemContractsCaller } from "@matterlabs/zksync-contracts/l2/system-contracts/libraries/SystemContractsCaller.sol";
+import { Strings } from "@openzeppelin/contracts/utils/Strings.sol";
 
 import { ISsoAccount } from "./interfaces/ISsoAccount.sol";
 
@@ -12,6 +13,8 @@ import { ISsoAccount } from "./interfaces/ISsoAccount.sol";
 /// @custom:security-contact [email protected]
 /// @dev This contract is used to deploy SSO accounts as beacon proxies.
 contract AAFactory {
+  using Strings for string;
+
   /// @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.
@@ -24,6 +27,17 @@ contract AAFactory {
   /// @notice A mapping from unique account IDs to their corresponding deployed account addresses.
   mapping(string => address) public accountMappings;
 
+  /// @notice A mapping from account addresses to their corresponding unique account IDs.
+  mapping(address => string) public accountIds;
+
+  /// @notice A mapping that marks account IDs as being used for recovery.
+  /// @dev This is used to prevent the same account ID from being used for recovery, deployment and future uses.
+  mapping(string => address) public recoveryAccountIds;
+
+  error AccountAlreadyRegistered(string uniqueAccountId, address accountAddress);
+  error AccountNotRegistered(string uniqueAccountId, address accountAddress);
+  error AccountUsedForRecovery(string uniqueAccountId, address accountAddress);
+
   /// @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 _beacon The address of the UpgradeableBeacon contract used for the SSO accounts' beacon proxies.
@@ -49,8 +63,6 @@ contract AAFactory {
     bytes[] calldata _initialValidators,
     address[] calldata _initialK1Owners
   ) external returns (address accountAddress) {
-    require(accountMappings[_uniqueAccountId] == address(0), "Account already exists");
-
     (bool success, bytes memory returnData) = SystemContractsCaller.systemCallWithReturndata(
       uint32(gasleft()),
       address(DEPLOYER_SYSTEM_CONTRACT),
@@ -63,11 +75,113 @@ contract AAFactory {
     require(success, "Deployment failed");
     (accountAddress) = abi.decode(returnData, (address));
 
-    accountMappings[_uniqueAccountId] = accountAddress;
+    // Check if the account is already registered
+    // Note: this check is done at this point, to use `accountAddress` to process the error message.
+    require(
+      accountMappings[_uniqueAccountId] == address(0),
+      AccountAlreadyRegistered(_uniqueAccountId, accountAddress)
+    );
+    require(accountIds[accountAddress].equal(""), AccountAlreadyRegistered(_uniqueAccountId, accountAddress));
+    require(
+      recoveryAccountIds[_uniqueAccountId] == address(0),
+      AccountUsedForRecovery(_uniqueAccountId, accountAddress)
+    );
 
     // Initialize the newly deployed account with validators, hooks and K1 owners.
     ISsoAccount(accountAddress).initialize(_initialValidators, _initialK1Owners);
 
+    _registerAccount(_uniqueAccountId, accountAddress);
+
     emit AccountCreated(accountAddress, _uniqueAccountId);
   }
+
+  /// @notice Registers an account with a given account ID.
+  /// @dev Can only be called by the account's validators.
+  /// @param _uniqueAccountId The unique identifier for the account.
+  /// @param _accountAddress The address of the account to register.
+  function registerAccount(
+    string calldata _uniqueAccountId,
+    address _accountAddress
+  ) external onlyAccountValidator(_accountAddress) {
+    require(
+      accountMappings[_uniqueAccountId] == address(0),
+      AccountAlreadyRegistered(_uniqueAccountId, _accountAddress)
+    );
+    require(accountIds[_accountAddress].equal(""), AccountAlreadyRegistered(_uniqueAccountId, _accountAddress));
+    require(
+      recoveryAccountIds[_uniqueAccountId] == address(0),
+      AccountUsedForRecovery(_uniqueAccountId, _accountAddress)
+    );
+
+    _registerAccount(_uniqueAccountId, _accountAddress);
+  }
+
+  function _registerAccount(string calldata _uniqueAccountId, address _accountAddress) internal {
+    accountMappings[_uniqueAccountId] = _accountAddress;
+    accountIds[_accountAddress] = _uniqueAccountId;
+  }
+
+  /// @notice Unregisters an account from the factory.
+  /// @dev Can only be called by the account's validators.
+  /// @param _uniqueAccountId The unique identifier for the account.
+  /// @param _accountAddress The address of the account to unregister.
+  function unregisterAccount(
+    string memory _uniqueAccountId,
+    address _accountAddress
+  ) external onlyAccountValidator(_accountAddress) {
+    require(
+      accountMappings[_uniqueAccountId] == _accountAddress,
+      AccountNotRegistered(_uniqueAccountId, _accountAddress)
+    );
+    require(
+      accountIds[_accountAddress].equal(_uniqueAccountId),
+      AccountNotRegistered(_uniqueAccountId, _accountAddress)
+    );
+
+    accountMappings[_uniqueAccountId] = address(0);
+    accountIds[_accountAddress] = "";
+  }
+
+  /// @notice Updates the account mapping for a given account ID during recovery.
+  /// @dev Can only be called by the account's validators.
+  /// @param _uniqueAccountId The unique identifier for the account.
+  /// @param _accountAddress The address of the account to update the mapping for.
+  function registerRecoveryBlockedAccount(
+    string calldata _uniqueAccountId,
+    address _accountAddress
+  ) external onlyAccountValidator(_accountAddress) {
+    require(
+      accountMappings[_uniqueAccountId] == address(0),
+      AccountAlreadyRegistered(_uniqueAccountId, _accountAddress)
+    );
+    require(
+      recoveryAccountIds[_uniqueAccountId] == address(0),
+      AccountUsedForRecovery(_uniqueAccountId, _accountAddress)
+    );
+
+    recoveryAccountIds[_uniqueAccountId] = _accountAddress;
+  }
+
+  /// @notice Unregisters a recovery blocked account from the factory.
+  /// @dev Can only be called by the account's validators.
+  /// @param _uniqueAccountId The unique identifier for the account.
+  /// @param _accountAddress The address of the account to unregister.
+  function unregisterRecoveryBlockedAccount(
+    string calldata _uniqueAccountId,
+    address _accountAddress
+  ) external onlyAccountValidator(_accountAddress) {
+    require(
+      recoveryAccountIds[_uniqueAccountId] == _accountAddress,
+      AccountNotRegistered(_uniqueAccountId, _accountAddress)
+    );
+
+    recoveryAccountIds[_uniqueAccountId] = address(0);
+  }
+
+  /// @notice Modifier that checks if the caller is a validator for the given account.
+  /// @param _accountAddress The address of the account to check the validator for.
+  modifier onlyAccountValidator(address _accountAddress) {
+    require(ISsoAccount(_accountAddress).isModuleValidator(msg.sender), "Unauthorized validator");
+    _;
+  }
 }

src/TransparentProxy.sol

@@ -12,7 +12,10 @@ import { TransparentUpgradeableProxy } from "@openzeppelin/contracts/proxy/trans
 /// cheap delegate calls on ZKsync.
 /// @dev This proxy is placed in front of `AAFactory` and all modules (`WebAuthValidator`, `SessionKeyValidator`).
 contract TransparentProxy is TransparentUpgradeableProxy, EfficientProxy {
-  constructor(address implementation) TransparentUpgradeableProxy(implementation, msg.sender, bytes("")) {}
+  constructor(
+    address implementation,
+    bytes memory data
+  ) TransparentUpgradeableProxy(implementation, msg.sender, data) {}
 
   function _delegate(address implementation) internal override(EfficientProxy, Proxy) {
     EfficientProxy._delegate(implementation);

src/interfaces/IGuardianRecoveryValidator.sol

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import { IModuleValidator } from "./IModuleValidator.sol";
+import { Transaction } from "@matterlabs/zksync-contracts/l2/system-contracts/libraries/TransactionHelper.sol";
+
+interface IGuardianRecoveryValidator is IModuleValidator {
+  struct GuardianConfirmation {
+    address ssoAccount;
+  }
+
+  function proposeValidationKey(address externalAccount) external;
+
+  function removeValidationKey(address externalAccount) external;
+
+  function initRecovery(address accountToRecover, bytes memory passkey, string memory accountId) external;
+
+  function addValidationKey(bytes memory key) external returns (bool);
+
+  function validateTransaction(
+    bytes32 signedHash,
+    bytes memory signature,
+    Transaction calldata transaction
+  ) external returns (bool);
+
+  function validateSignature(bytes32 signedHash, bytes memory signature) external view returns (bool);
+}

src/test/ExampleAuthServerPaymaster.sol

@@ -10,26 +10,33 @@ import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
 
 import { AAFactory } from "../AAFactory.sol";
 import { SessionKeyValidator } from "../validators/SessionKeyValidator.sol";
+import { GuardianRecoveryValidator } from "../validators/GuardianRecoveryValidator.sol";
 
 /// @author Matter Labs
 /// @notice This contract does not include any validations other than using the paymaster general flow.
 contract ExampleAuthServerPaymaster is IPaymaster, Ownable {
   address public immutable AA_FACTORY_CONTRACT_ADDRESS;
   address public immutable SESSION_KEY_VALIDATOR_CONTRACT_ADDRESS;
+  address public immutable ACCOUNT_RECOVERY_VALIDATOR_CONTRACT_ADDRESS;
   bytes4 constant DEPLOY_ACCOUNT_SELECTOR = AAFactory.deployProxySsoAccount.selector;
   bytes4 constant SESSION_CREATE_SELECTOR = SessionKeyValidator.createSession.selector;
   bytes4 constant SESSION_REVOKE_KEY_SELECTOR = SessionKeyValidator.revokeKey.selector;
   bytes4 constant SESSION_REVOKE_KEYS_SELECTOR = SessionKeyValidator.revokeKeys.selector;
+  bytes4 constant GUARDIAN_RECOVERY_ADD_KEY_SELECTOR = GuardianRecoveryValidator.addValidationKey.selector;
+  bytes4 constant GUARDIAN_RECOVERY_PROPOSE_KEY_SELECTOR = GuardianRecoveryValidator.proposeValidationKey.selector;
+  bytes4 constant GUARDIAN_RECOVERY_DISCARD_RECOVERY_SELECTOR = GuardianRecoveryValidator.discardRecovery.selector;
+  bytes4 constant GUARDIAN_RECOVERY_REMOVE_KEY_SELECTOR = GuardianRecoveryValidator.removeValidationKey.selector;
 
   modifier onlyBootloader() {
     require(msg.sender == BOOTLOADER_FORMAL_ADDRESS, "Only bootloader can call this method");
     // Continue execution if called from the bootloader.
     _;
   }
 
-  constructor(address aaFactoryAddress, address sessionKeyValidatorAddress) {
+  constructor(address aaFactoryAddress, address sessionKeyValidatorAddress, address accountRecoveryValidatorAddress) {
     AA_FACTORY_CONTRACT_ADDRESS = aaFactoryAddress;
     SESSION_KEY_VALIDATOR_CONTRACT_ADDRESS = sessionKeyValidatorAddress;
+    ACCOUNT_RECOVERY_VALIDATOR_CONTRACT_ADDRESS = accountRecoveryValidatorAddress;
   }
 
   function validateAndPayForPaymasterTransaction(
@@ -44,7 +51,9 @@ contract ExampleAuthServerPaymaster is IPaymaster, Ownable {
     // Ensure the transaction is calling one of our allowed contracts
     address to = address(uint160(_transaction.to));
     require(
-      to == AA_FACTORY_CONTRACT_ADDRESS || to == SESSION_KEY_VALIDATOR_CONTRACT_ADDRESS,
+      to == AA_FACTORY_CONTRACT_ADDRESS ||
+        to == SESSION_KEY_VALIDATOR_CONTRACT_ADDRESS ||
+        to == ACCOUNT_RECOVERY_VALIDATOR_CONTRACT_ADDRESS,
       "Unsupported contract address"
     );
 
@@ -63,6 +72,16 @@ contract ExampleAuthServerPaymaster is IPaymaster, Ownable {
       );
     }
 
+    if (to == ACCOUNT_RECOVERY_VALIDATOR_CONTRACT_ADDRESS) {
+      require(
+        methodSelector == GUARDIAN_RECOVERY_ADD_KEY_SELECTOR ||
+          methodSelector == GUARDIAN_RECOVERY_PROPOSE_KEY_SELECTOR ||
+          methodSelector == GUARDIAN_RECOVERY_DISCARD_RECOVERY_SELECTOR ||
+          methodSelector == GUARDIAN_RECOVERY_REMOVE_KEY_SELECTOR,
+        "Unsupported method"
+      );
+    }
+
     bytes4 paymasterInputSelector = bytes4(_transaction.paymasterInput[0:4]);
     require(paymasterInputSelector == IPaymasterFlow.general.selector, "Unsupported paymaster flow");