diff --git a/contracts/AGWRegistryV2.sol b/contracts/AGWRegistryV2.sol new file mode 100644 index 0000000..7da3bbd --- /dev/null +++ b/contracts/AGWRegistryV2.sol @@ -0,0 +1,179 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.17; + +import {Ownable, Ownable2Step} from "@openzeppelin/contracts/access/Ownable2Step.sol"; + +import {Errors} from "./libraries/Errors.sol"; +import {IAGWRegistry} from "./interfaces/IAGWRegistry.sol"; +import {IAGWRegistryV2} from "./interfaces/IAGWRegistryV2.sol"; + +contract AGWRegistryV2 is Ownable2Step, IAGWRegistryV2 { + uint8 public constant VERSION_NONE = 0; + uint8 public constant VERSION_V1 = 1; + + address public immutable v1Registry; + + mapping(address => bool) public isFactory; + + mapping(address => uint8 version) private _agwVersion; + + event AGWRegistered(address indexed account, uint8 indexed version); + event AGWUnregistered( + address indexed account, + uint8 indexed previousVersion + ); + + /** + * @notice Event emitted when a factory contract is set + * @param factory address - Address of the factory contract + */ + event FactorySet(address indexed factory); + + /** + * @notice Event emitted when a factory contract is unset + * @param factory address - Address of the factory contract + */ + event FactoryUnset(address indexed factory); + + error INVALID_VERSION(); + + constructor(address _owner, address _v1Registry) Ownable(_owner) { + v1Registry = _v1Registry; + } + + /** + * @notice Registers an account as an AGW account + * @dev Can only be called by the factory or owner + * @param account address - Address of the account to register + */ + function register( + address account, + uint8 version + ) external override onlyFactoryOrOwner { + _register(account, version); + } + + /** + * @notice Registers multiple accounts as AGW accounts + * @dev Can only be called by the factory or owner + * @param accounts address[] - Array of addresses to register + * @param version uint8 - AGW version to assign to each account + */ + function registerMultiple( + address[] calldata accounts, + uint8 version + ) external onlyFactoryOrOwner { + if (version == VERSION_NONE) revert INVALID_VERSION(); + + for (uint256 i = 0; i < accounts.length; i++) { + _setVersion(accounts[i], version); + } + } + + /** + * @notice Unregisters an account from this registry + * @dev Cannot unregister accounts that are only present in the v1 fallback registry + * @param account address - Address of the account to unregister + */ + function unregister(address account) external onlyFactoryOrOwner { + _unsetVersion(account); + } + + /** + * @notice Unregisters multiple accounts from this registry + * @dev Cannot unregister accounts that are only present in the v1 fallback registry + * @param accounts address[] - Array of addresses to unregister + */ + function unregisterMultiple( + address[] calldata accounts + ) external onlyFactoryOrOwner { + for (uint256 i = 0; i < accounts.length; i++) { + _unsetVersion(accounts[i]); + } + } + + /** + * @notice Returns whether an account is registered as any AGW version + */ + function isAGW(address account) external view override returns (bool) { + return versionOf(account) != VERSION_NONE; + } + + /** + * @notice Returns whether an account is registered as a specific AGW version + */ + function isAGWVersion( + address account, + uint8 version + ) external view override returns (bool) { + return versionOf(account) == version; + } + + /** + * @notice Returns the AGW version for an account + * @dev Local v2 registry state takes precedence over the v1 fallback registry. + */ + function versionOf(address account) public view override returns (uint8) { + uint8 localVersion = _agwVersion[account]; + if (localVersion != VERSION_NONE) { + return localVersion; + } + + if ( + v1Registry != address(0) && IAGWRegistry(v1Registry).isAGW(account) + ) { + return VERSION_V1; + } + + return VERSION_NONE; + } + + /** + * @notice Sets a new factory contract + * @dev Can only be called by the owner + * @param factory_ address - Address of the new factory + */ + function setFactory(address factory_) external onlyOwner { + isFactory[factory_] = true; + + emit FactorySet(factory_); + } + + /** + * @notice Unsets a factory contract + * @dev Can only be called by the owner + * @param factory_ address - Address of the factory + */ + function unsetFactory(address factory_) external onlyOwner { + isFactory[factory_] = false; + + emit FactoryUnset(factory_); + } + + function _register(address account, uint8 version) internal { + if (version == VERSION_NONE) revert INVALID_VERSION(); + + _setVersion(account, version); + } + + function _setVersion(address account, uint8 version) internal { + _agwVersion[account] = version; + + emit AGWRegistered(account, version); + } + + function _unsetVersion(address account) internal { + uint8 previousVersion = _agwVersion[account]; + delete _agwVersion[account]; + + emit AGWUnregistered(account, previousVersion); + } + + modifier onlyFactoryOrOwner() { + if (!isFactory[msg.sender] && msg.sender != owner()) { + revert Errors.NOT_FROM_FACTORY(); + } + + _; + } +} diff --git a/contracts/FeatureRegistry.sol b/contracts/FeatureRegistry.sol new file mode 100644 index 0000000..057f20d --- /dev/null +++ b/contracts/FeatureRegistry.sol @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.17; + +import {Ownable, Ownable2Step} from "@openzeppelin/contracts/access/Ownable2Step.sol"; + +import {IFeatureRegistry} from "./interfaces/IFeatureRegistry.sol"; +import {IFeatureResolver} from "./interfaces/IFeatureResolver.sol"; + +contract FeatureRegistry is Ownable2Step, IFeatureRegistry { + mapping(uint8 id => Feature feature) private _features; + mapping(uint8 id => bool defined) private _isFeatureDefined; + mapping(bytes32 key => uint8 id) private _featureIdByKey; + mapping(bytes32 key => bool exists) private _featureKeyExists; + mapping(uint8 id => uint256 indexPlusOne) private _featureIndexPlusOne; + + uint8[] private _featureIds; + + error EMPTY_FEATURE_KEY(); + error EMPTY_RESOLVER(); + error FEATURE_NOT_DEFINED(); + error FEATURE_KEY_ALREADY_USED(); + + constructor(address _owner) Ownable(_owner) {} + + function setFeature( + uint8 id, + bytes32 key, + string calldata name, + string calldata metadataURI, + address resolver, + bytes calldata resolverData, + bool enabled + ) external onlyOwner { + if (key == bytes32(0)) revert EMPTY_FEATURE_KEY(); + if (resolver == address(0)) revert EMPTY_RESOLVER(); + + bytes32 previousKey = _features[id].key; + if ( + _featureKeyExists[key] && + (!_isFeatureDefined[id] || previousKey != key) + ) { + revert FEATURE_KEY_ALREADY_USED(); + } + + if (!_isFeatureDefined[id]) { + _featureIds.push(id); + _featureIndexPlusOne[id] = _featureIds.length; + _isFeatureDefined[id] = true; + } else if (previousKey != key) { + delete _featureIdByKey[previousKey]; + delete _featureKeyExists[previousKey]; + } + + uint256 mask = _maskFor(id); + _features[id] = Feature({ + id: id, + mask: mask, + key: key, + name: name, + metadataURI: metadataURI, + resolver: resolver, + resolverData: resolverData, + enabled: enabled + }); + _featureIdByKey[key] = id; + _featureKeyExists[key] = true; + + emit FeatureSet( + id, + mask, + key, + name, + metadataURI, + resolver, + resolverData, + enabled + ); + } + + function unsetFeature(uint8 id) external onlyOwner { + if (!_isFeatureDefined[id]) revert FEATURE_NOT_DEFINED(); + + Feature memory feature = _features[id]; + uint256 index = _featureIndexPlusOne[id] - 1; + uint256 lastIndex = _featureIds.length - 1; + + if (index != lastIndex) { + uint8 lastId = _featureIds[lastIndex]; + _featureIds[index] = lastId; + _featureIndexPlusOne[lastId] = index + 1; + } + + _featureIds.pop(); + delete _featureIndexPlusOne[id]; + delete _isFeatureDefined[id]; + delete _featureIdByKey[feature.key]; + delete _featureKeyExists[feature.key]; + delete _features[id]; + + emit FeatureUnset(id, feature.mask, feature.key); + } + + function getFeature(uint8 id) external view returns (Feature memory) { + if (!_isFeatureDefined[id]) revert FEATURE_NOT_DEFINED(); + + return _features[id]; + } + + function getFeatures( + uint8[] calldata ids + ) external view returns (Feature[] memory features) { + features = new Feature[](ids.length); + + for (uint256 i = 0; i < ids.length; i++) { + if (!_isFeatureDefined[ids[i]]) revert FEATURE_NOT_DEFINED(); + + features[i] = _features[ids[i]]; + } + } + + function getFeatureCount() external view returns (uint256) { + return _featureIds.length; + } + + function getFeatureIdAt(uint256 index) external view returns (uint8) { + return _featureIds[index]; + } + + function getFeatureIds() external view returns (uint8[] memory) { + return _featureIds; + } + + function getFeatureMask(uint8 id) external pure returns (uint256) { + return _maskFor(id); + } + + function featureIdOf( + bytes32 key + ) external view returns (uint8 id, bool exists) { + return (_featureIdByKey[key], _featureKeyExists[key]); + } + + function isFeatureDefined(uint8 id) external view returns (bool) { + return _isFeatureDefined[id]; + } + + function supportsFeature(address agw, uint8 id) public view returns (bool) { + if (!_isFeatureDefined[id]) return false; + + Feature storage feature = _features[id]; + if (!feature.enabled) return false; + + try + IFeatureResolver(feature.resolver).supportsFeature( + agw, + id, + feature.resolverData + ) + returns (bool supported) { + return supported; + } catch { + return false; + } + } + + function supportsAllFeatures( + address agw, + uint256 features + ) external view returns (bool) { + if (features == 0) return true; + + for (uint8 id = 0; id < 255; id++) { + uint256 mask = _maskFor(id); + if (features & mask != 0 && !supportsFeature(agw, id)) { + return false; + } + } + + return features & _maskFor(255) == 0 || supportsFeature(agw, 255); + } + + function supportsAnyFeature( + address agw, + uint256 features + ) external view returns (bool) { + if (features == 0) return false; + + for (uint8 id = 0; id < 255; id++) { + uint256 mask = _maskFor(id); + if (features & mask != 0 && supportsFeature(agw, id)) { + return true; + } + } + + return features & _maskFor(255) != 0 && supportsFeature(agw, 255); + } + + function getSupportedFeatures(address agw) external view returns (uint256) { + uint256 supportedFeatures; + + for (uint256 i = 0; i < _featureIds.length; i++) { + uint8 id = _featureIds[i]; + if (supportsFeature(agw, id)) { + supportedFeatures |= _maskFor(id); + } + } + + return supportedFeatures; + } + + function _maskFor(uint8 id) internal pure returns (uint256) { + return uint256(1) << id; + } +} diff --git a/contracts/interfaces/IAGWRegistryV2.sol b/contracts/interfaces/IAGWRegistryV2.sol new file mode 100644 index 0000000..745edfd --- /dev/null +++ b/contracts/interfaces/IAGWRegistryV2.sol @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.17; + +interface IAGWRegistryV2 { + function register(address account, uint8 version) external; + + function isAGW(address account) external view returns (bool); + + function isAGWVersion( + address account, + uint8 version + ) external view returns (bool); + + function versionOf(address account) external view returns (uint8); +} diff --git a/contracts/interfaces/IFeatureRegistry.sol b/contracts/interfaces/IFeatureRegistry.sol new file mode 100644 index 0000000..19374e5 --- /dev/null +++ b/contracts/interfaces/IFeatureRegistry.sol @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.17; + +interface IFeatureRegistry { + struct Feature { + uint8 id; + uint256 mask; + bytes32 key; + string name; + string metadataURI; + address resolver; + bytes resolverData; + bool enabled; + } + + event FeatureSet( + uint8 indexed id, + uint256 indexed mask, + bytes32 indexed key, + string name, + string metadataURI, + address resolver, + bytes resolverData, + bool enabled + ); + event FeatureUnset( + uint8 indexed id, + uint256 indexed mask, + bytes32 indexed key + ); + + function setFeature( + uint8 id, + bytes32 key, + string calldata name, + string calldata metadataURI, + address resolver, + bytes calldata resolverData, + bool enabled + ) external; + + function unsetFeature(uint8 id) external; + + function getFeature(uint8 id) external view returns (Feature memory); + + function getFeatures( + uint8[] calldata ids + ) external view returns (Feature[] memory); + + function getFeatureCount() external view returns (uint256); + + function getFeatureIdAt(uint256 index) external view returns (uint8); + + function getFeatureIds() external view returns (uint8[] memory); + + function getFeatureMask(uint8 id) external pure returns (uint256); + + function featureIdOf( + bytes32 key + ) external view returns (uint8 id, bool exists); + + function isFeatureDefined(uint8 id) external view returns (bool); + + function supportsFeature( + address agw, + uint8 id + ) external view returns (bool); + + function supportsAllFeatures( + address agw, + uint256 features + ) external view returns (bool); + + function supportsAnyFeature( + address agw, + uint256 features + ) external view returns (bool); + + function getSupportedFeatures(address agw) external view returns (uint256); +} diff --git a/contracts/interfaces/IFeatureResolver.sol b/contracts/interfaces/IFeatureResolver.sol new file mode 100644 index 0000000..6d79849 --- /dev/null +++ b/contracts/interfaces/IFeatureResolver.sol @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.17; + +interface IFeatureResolver { + function supportsFeature( + address agw, + uint8 featureId, + bytes calldata resolverData + ) external view returns (bool); +} diff --git a/contracts/test/MockFeatureResolver.sol b/contracts/test/MockFeatureResolver.sol new file mode 100644 index 0000000..078070d --- /dev/null +++ b/contracts/test/MockFeatureResolver.sol @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.17; + +import {IFeatureResolver} from "../interfaces/IFeatureResolver.sol"; + +contract MockFeatureResolver is IFeatureResolver { + mapping(address agw => mapping(uint8 featureId => bool supported)) + public supported; + + bool public shouldRevert; + + function setSupported( + address agw, + uint8 featureId, + bool isSupported + ) external { + supported[agw][featureId] = isSupported; + } + + function setShouldRevert(bool shouldRevert_) external { + shouldRevert = shouldRevert_; + } + + function supportsFeature( + address agw, + uint8 featureId, + bytes calldata + ) external view returns (bool) { + if (shouldRevert) revert(); + + return supported[agw][featureId]; + } +} diff --git a/test/registry/agwregistryv2.test.ts b/test/registry/agwregistryv2.test.ts new file mode 100644 index 0000000..c92a899 --- /dev/null +++ b/test/registry/agwregistryv2.test.ts @@ -0,0 +1,107 @@ +import { expect } from "chai"; +import { getAddress, ZeroAddress } from "ethers"; +import * as hre from "hardhat"; +import type { Contract, Wallet } from "zksync-ethers"; + +import { + deployContract, + getWallet, + LOCAL_RICH_WALLETS, +} from "../../deploy/utils"; + +describe("AGWRegistryV2", () => { + let wallet: Wallet; + + const account = getAddress("0x000000000000000000000000000000000000a001"); + const otherAccount = getAddress("0x000000000000000000000000000000000000a002"); + const VERSION_NONE = 0; + const VERSION_V1 = 1; + const VERSION_V2 = 2; + + before(async () => { + wallet = getWallet(hre, LOCAL_RICH_WALLETS[0].privateKey); + }); + + async function deployRegistryV2( + v1Registry: string = ZeroAddress, + ): Promise { + return (await deployContract( + hre, + "AGWRegistryV2", + [wallet.address, v1Registry], + { + wallet, + silent: true, + }, + )) as unknown as Contract; + } + + it("registers local AGW versions", async () => { + const registry = await deployRegistryV2(); + + const tx = await registry.register(account, VERSION_V2); + await tx.wait(); + + expect(await registry.isAGW(account)).to.eq(true); + expect(await registry.versionOf(account)).to.eq(VERSION_V2); + expect(await registry.isAGWVersion(account, VERSION_V2)).to.eq(true); + expect(await registry.isAGWVersion(account, VERSION_V1)).to.eq(false); + }); + + it("noops fallback reads when the v1 registry is not configured", async () => { + const registry = await deployRegistryV2(); + + expect(await registry.isAGW(otherAccount)).to.eq(false); + expect(await registry.versionOf(otherAccount)).to.eq(VERSION_NONE); + expect(await registry.isAGWVersion(otherAccount, VERSION_V1)).to.eq(false); + }); + + it("falls back to the v1 registry when no local version is set", async () => { + const v1Registry = (await deployContract( + hre, + "AGWRegistry", + [wallet.address], + { + wallet, + silent: true, + }, + )) as unknown as Contract; + const registerV1Tx = await v1Registry.register(account); + await registerV1Tx.wait(); + + const registry = await deployRegistryV2(await v1Registry.getAddress()); + + expect(await registry.isAGW(account)).to.eq(true); + expect(await registry.versionOf(account)).to.eq(VERSION_V1); + expect(await registry.isAGWVersion(account, VERSION_V1)).to.eq(true); + expect(await registry.isAGWVersion(account, VERSION_V2)).to.eq(false); + }); + + it("uses local versions before falling back to the v1 registry", async () => { + const v1Registry = (await deployContract( + hre, + "AGWRegistry", + [wallet.address], + { + wallet, + silent: true, + }, + )) as unknown as Contract; + const registerV1Tx = await v1Registry.register(account); + await registerV1Tx.wait(); + + const registry = await deployRegistryV2(await v1Registry.getAddress()); + const registerV2Tx = await registry.register(account, VERSION_V2); + await registerV2Tx.wait(); + + expect(await registry.versionOf(account)).to.eq(VERSION_V2); + expect(await registry.isAGWVersion(account, VERSION_V1)).to.eq(false); + expect(await registry.isAGWVersion(account, VERSION_V2)).to.eq(true); + + const unregisterTx = await registry.unregister(account); + await unregisterTx.wait(); + + expect(await registry.versionOf(account)).to.eq(VERSION_V1); + expect(await registry.isAGWVersion(account, VERSION_V1)).to.eq(true); + }); +}); diff --git a/test/registry/featureregistry.test.ts b/test/registry/featureregistry.test.ts new file mode 100644 index 0000000..ab954b1 --- /dev/null +++ b/test/registry/featureregistry.test.ts @@ -0,0 +1,133 @@ +import { expect } from "chai"; +import { getAddress, id } from "ethers"; +import * as hre from "hardhat"; +import type { Contract, Wallet } from "zksync-ethers"; + +import { + deployContract, + getWallet, + LOCAL_RICH_WALLETS, +} from "../../deploy/utils"; + +describe("FeatureRegistry", () => { + let wallet: Wallet; + let registry: Contract; + let resolver: Contract; + + const agw = getAddress("0x000000000000000000000000000000000000f001"); + const sessionKeysId = 3; + const recoveryId = 5; + const sessionKeysMask = 1n << BigInt(sessionKeysId); + const recoveryMask = 1n << BigInt(recoveryId); + const sessionKeysKey = id("agw.feature.session-keys"); + const recoveryKey = id("agw.feature.recovery"); + + beforeEach(async () => { + wallet = getWallet(hre, LOCAL_RICH_WALLETS[0].privateKey); + registry = (await deployContract(hre, "FeatureRegistry", [wallet.address], { + wallet, + silent: true, + })) as unknown as Contract; + resolver = (await deployContract(hre, "MockFeatureResolver", [], { + wallet, + silent: true, + })) as unknown as Contract; + }); + + async function setFeature( + featureId: number, + key: string, + name: string, + enabled: boolean, + ): Promise { + const tx = await registry.setFeature( + featureId, + key, + name, + "ipfs://feature-metadata", + await resolver.getAddress(), + "0x1234", + enabled, + ); + await tx.wait(); + } + + it("stores feature catalog entries by id and key", async () => { + await setFeature(sessionKeysId, sessionKeysKey, "Session Keys", true); + + const feature = await registry.getFeature(sessionKeysId); + expect(feature.id).to.eq(sessionKeysId); + expect(feature.mask).to.eq(sessionKeysMask); + expect(feature.key).to.eq(sessionKeysKey); + expect(feature.name).to.eq("Session Keys"); + expect(feature.metadataURI).to.eq("ipfs://feature-metadata"); + expect(feature.resolver).to.eq(await resolver.getAddress()); + expect(feature.resolverData).to.eq("0x1234"); + expect(feature.enabled).to.eq(true); + + const [featureId, exists] = await registry.featureIdOf(sessionKeysKey); + expect(featureId).to.eq(sessionKeysId); + expect(exists).to.eq(true); + expect(await registry.getFeatureMask(sessionKeysId)).to.eq(sessionKeysMask); + expect(await registry.getFeatureCount()).to.eq(1n); + expect(await registry.getFeatureIds()).to.deep.eq([BigInt(sessionKeysId)]); + }); + + it("derives support through feature resolvers", async () => { + await setFeature(sessionKeysId, sessionKeysKey, "Session Keys", true); + await setFeature(recoveryId, recoveryKey, "Recovery", true); + + const setSupportedTx = await resolver.setSupported( + agw, + sessionKeysId, + true, + ); + await setSupportedTx.wait(); + + expect(await registry.supportsFeature(agw, sessionKeysId)).to.eq(true); + expect(await registry.supportsFeature(agw, recoveryId)).to.eq(false); + expect(await registry.supportsAllFeatures(agw, sessionKeysMask)).to.eq( + true, + ); + expect( + await registry.supportsAllFeatures(agw, sessionKeysMask | recoveryMask), + ).to.eq(false); + expect( + await registry.supportsAnyFeature(agw, sessionKeysMask | recoveryMask), + ).to.eq(true); + expect(await registry.getSupportedFeatures(agw)).to.eq(sessionKeysMask); + }); + + it("treats disabled or reverting feature resolvers as unsupported", async () => { + await setFeature(sessionKeysId, sessionKeysKey, "Session Keys", false); + let setSupportedTx = await resolver.setSupported(agw, sessionKeysId, true); + await setSupportedTx.wait(); + + expect(await registry.supportsFeature(agw, sessionKeysId)).to.eq(false); + expect(await registry.getSupportedFeatures(agw)).to.eq(0n); + + await setFeature(sessionKeysId, sessionKeysKey, "Session Keys", true); + const setShouldRevertTx = await resolver.setShouldRevert(true); + await setShouldRevertTx.wait(); + + expect(await registry.supportsFeature(agw, sessionKeysId)).to.eq(false); + expect(await registry.supportsAnyFeature(agw, sessionKeysMask)).to.eq( + false, + ); + }); + + it("removes feature catalog entries", async () => { + await setFeature(sessionKeysId, sessionKeysKey, "Session Keys", true); + + const unsetTx = await registry.unsetFeature(sessionKeysId); + await unsetTx.wait(); + + const [, exists] = await registry.featureIdOf(sessionKeysKey); + expect(exists).to.eq(false); + expect(await registry.isFeatureDefined(sessionKeysId)).to.eq(false); + expect(await registry.getFeatureCount()).to.eq(0n); + await expect( + registry.getFeature(sessionKeysId), + ).to.be.revertedWithCustomError(registry, "FEATURE_NOT_DEFINED"); + }); +});