WIP: feat: OIDC account recovery

cspell-config/cspell-misc.txt

@@ -13,6 +13,7 @@ dockerized
 ethereum
 sepolia
 
+
 // examples/bank-demo
 ctap
 Aave
@@ -28,3 +29,10 @@ usdc
 // examples/nft-quest
 Fren
 fren
+
+// src/OidcKeyRegistry
+Oidc
+oidc
+
+// src/validators/OidcRecoveryValidator
+pkop

cspell-config/cspell-sol.txt

@@ -34,3 +34,4 @@ solady
 xbatch
 tload
 tstore
+Groth

cspell.json

@@ -26,7 +26,8 @@
     "**/test-results/**",
     "**/playwright-report/**",
     "**/blob-report/**",
-    "**/playwright/.cache/**"
+    "**/playwright/.cache/**",
+    "src/autogenerated/**"
   ],
   "caseSensitive": true,
   "dictionaries": [

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,12 +6,16 @@ import { Wallet } from "zksync-ethers";
 
 const WEBAUTH_NAME = "WebAuthValidator";
 const SESSIONS_NAME = "SessionKeyValidator";
+const GUARDIAN_RECOVERY_NAME = "GuardianRecoveryValidator";
+const OIDC_RECOVERY_NAME = "OidcRecoveryValidator";
+const OIDC_VERIFIER_NAME = "Groth16Verifier";
 const ACCOUNT_IMPL_NAME = "SsoAccount";
 const FACTORY_NAME = "AAFactory";
 const PAYMASTER_NAME = "ExampleAuthServerPaymaster";
 const BEACON_NAME = "SsoBeacon";
+const OIDC_KEY_REGISTRY_NAME = "OidcKeyRegistry";
 
-async function deploy(name: string, deployer: Wallet, proxy: boolean, args?: any[]): Promise<string> {
+async function deploy(name: string, deployer: Wallet, proxy: boolean, args?: any[], initArgs?: any): Promise<string> {
   // eslint-disable-next-line @typescript-eslint/no-require-imports
   const { deployFactory, create2, ethersStaticSalt } = require("../test/utils");
   console.log("Deploying", name, "contract...");
@@ -26,13 +30,12 @@ async function deploy(name: string, deployer: Wallet, proxy: boolean, args?: any
     console.log(name, "contract deployed at:", implAddress, "\n");
     return implAddress;
   }
-  const proxyContract = await create2("TransparentProxy", deployer, ethersStaticSalt, [implAddress]);
+  const proxyContract = await create2("TransparentProxy", deployer, ethersStaticSalt, [implAddress, initArgs ?? "0x"]);
   const proxyAddress = await proxyContract.getAddress();
   console.log(name, "proxy contract deployed at:", proxyAddress, "\n");
   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")
@@ -76,12 +79,19 @@ 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 oidcKeyRegistryInterface = new ethers.Interface((await hre.artifacts.readArtifact(OIDC_KEY_REGISTRY_NAME)).abi);
+      const oidcKeyRegistry = await deploy(OIDC_KEY_REGISTRY_NAME, deployer, !cmd.noProxy, [], oidcKeyRegistryInterface.encodeFunctionData("initialize", []));
+      const oidcRecoveryInterface = new ethers.Interface((await hre.artifacts.readArtifact(OIDC_RECOVERY_NAME)).abi);
+      const oidcVerifier = await deploy(OIDC_VERIFIER_NAME, deployer, false, []);
+      const oidcRecovery = await deploy(OIDC_RECOVERY_NAME, deployer, !cmd.noProxy, [oidcKeyRegistry, oidcVerifier], oidcRecoveryInterface.encodeFunctionData("initialize", [oidcKeyRegistry, oidcVerifier]));
+      const paymaster = await deploy(PAYMASTER_NAME, deployer, false, [factory, sessions, recovery, oidcRecovery]);
 
       await fundPaymaster(paymaster, cmd.fund);
     } else {
@@ -105,6 +115,9 @@ task("deploy", "Deploys ZKsync SSO contracts")
         }
         args = [cmd.factory, cmd.sessions];
       }
+      if (cmd.only == OIDC_KEY_REGISTRY_NAME) {
+        args = [];
+      }
       const deployedContract = await deploy(cmd.only, deployer, false, args);
 
       if (cmd.only == PAYMASTER_NAME) {

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/OidcKeyRegistry.sol

@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.24;
+
+import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+
+contract OidcKeyRegistry is Initializable, OwnableUpgradeable {
+  uint8 public constant MAX_KEYS = 5;
+
+  struct Key {
+    bytes32 kid; // Key ID
+    bytes n; // RSA modulus
+    bytes e; // RSA exponent
+  }
+
+  // Mapping uses keccak256(iss) as the key
+  mapping(bytes32 => Key[MAX_KEYS]) public OIDCKeys; // Stores up to MAX_KEYS per issuer
+  mapping(bytes32 => uint8) public keyIndexes; // Tracks the latest key index for each issuer
+
+  constructor() {
+    initialize();
+  }
+
+  function initialize() public initializer {
+    __Ownable_init();
+  }
+
+  function hashIssuer(string memory iss) public pure returns (bytes32) {
+    return keccak256(abi.encodePacked(iss));
+  }
+
+  function setKey(bytes32 issHash, Key memory key) public onlyOwner {
+    uint8 index = keyIndexes[issHash];
+    uint8 nextIndex = (index + 1) % MAX_KEYS; // Circular buffer
+    OIDCKeys[issHash][nextIndex] = key;
+    keyIndexes[issHash] = nextIndex;
+  }
+
+  function setKeys(bytes32 issHash, Key[] memory keys) public onlyOwner {
+    for (uint8 i = 0; i < keys.length; i++) {
+      setKey(issHash, keys[i]);
+    }
+  }
+
+  function getKey(bytes32 issHash, bytes32 kid) public view returns (Key memory) {
+    require(kid != 0, "Invalid kid");
+    Key[MAX_KEYS] storage keys = OIDCKeys[issHash];
+    for (uint8 i = 0; i < MAX_KEYS; i++) {
+      if (keys[i].kid == kid) {
+        return keys[i];
+      }
+    }
+    revert("Key not found");
+  }
+}

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);