Upgradeable IpAssetRegistry

contracts/LicenseToken.sol

@@ -51,6 +51,9 @@ contract LicenseToken is ILicenseToken, ERC721EnumerableUpgradeable, AccessManag
     /// @dev Initializes the LicenseToken contract
     function initialize(address accessManager, string memory imageUrl) public initializer {
         __ERC721_init("Programmable IP License Token", "PILicenseToken");
+        if (accessManager == address(0)) {
+            revert Errors.LicenseToken__ZeroAccessManager();
+        }
         __AccessManaged_init(accessManager);
         __UUPSUpgradeable_init();
         _getLicenseTokenStorage().imageUrl = imageUrl;

contracts/access/AccessController.sol

@@ -57,6 +57,9 @@ contract AccessController is IAccessController, AccessManagedUpgradeable, UUPSUp
     /// @notice initializer for this implementation contract
     /// @param accessManager The address of the protocol admin roles contract
     function initialize(address accessManager) external initializer {
+        if (accessManager == address(0)) {
+            revert Errors.AccessController__ZeroAccessManager();
+        }
         __AccessManaged_init(accessManager);
     }
 

contracts/lib/Errors.sol

@@ -29,6 +29,9 @@ library Errors {
     //                               IPAssetRegistry                         //
     ////////////////////////////////////////////////////////////////////////////
 
+    /// @notice zero address provided for access manager in initializer.
+    error IPAssetRegistry__ZeroAccessManager();
+
     /// @notice The IP asset has already been registered.
     error IPAssetRegistry__AlreadyRegistered();
 
@@ -110,6 +113,7 @@ library Errors {
     //                            LicenseRegistry                             //
     ////////////////////////////////////////////////////////////////////////////
 
+    error LicenseRegistry__ZeroAccessManager();
     error LicenseRegistry__CallerNotLicensingModule();
     error LicenseRegistry__ZeroLicensingModule();
     error LicensingModule__CallerNotLicenseRegistry();
@@ -140,6 +144,7 @@ library Errors {
     error LicenseToken__CallerNotLicensingModule();
     error LicenseToken__ZeroLicensingModule();
     error LicenseToken__ZeroDisputeModule();
+    error LicenseToken__ZeroAccessManager();
     error LicenseToken__RevokedLicense(uint256 tokenId);
     error LicenseToken__NotTransferable();
     error LicenseToken__LicenseTokenExpired(uint256 tokenId, uint256 expiredAt, uint256 currentTimestamp);
@@ -152,6 +157,7 @@ library Errors {
     //                            LicensingModule                             //
     ////////////////////////////////////////////////////////////////////////////
 
+    error LicensingModule__ZeroAccessManager();
     error LicensingModule__IpAlreadyLinked();
     error LicensingModule__PolicyAlreadySetForIpId();
     error LicensingModule__FrameworkNotFound();
@@ -242,6 +248,10 @@ library Errors {
     error DisputeModule__ZeroController();
     error DisputeModule__ZeroAccessManager();
 
+    ////////////////////////////////////////////////////////////////////////////
+    //                         ArbitrationPolicy SP                          //
+    ////////////////////////////////////////////////////////////////////////////
+
     error ArbitrationPolicySP__ZeroDisputeModule();
     error ArbitrationPolicySP__ZeroPaymentToken();
     error ArbitrationPolicySP__NotDisputeModule();
@@ -294,6 +304,7 @@ library Errors {
     //                             ModuleRegistry                             //
     ////////////////////////////////////////////////////////////////////////////
 
+    error ModuleRegistry__ZeroAccessManager();
     error ModuleRegistry__ModuleAddressZeroAddress();
     error ModuleRegistry__ModuleAddressNotContract();
     error ModuleRegistry__ModuleAlreadyRegistered();
@@ -311,6 +322,7 @@ library Errors {
     //                             AccessController                           //
     ////////////////////////////////////////////////////////////////////////////
 
+    error AccessController__ZeroAccessManager();
     error AccessController__IPAccountIsZeroAddress();
     error AccessController__IPAccountIsNotValid(address ipAccount);
     error AccessController__SignerIsZeroAddress();

contracts/modules/dispute/DisputeModule.sol

@@ -93,8 +93,9 @@ contract DisputeModule is
     /// @notice Initializer for this implementation contract
     /// @param accessManager The address of the protocol admin roles contract
     function initialize(address accessManager) external initializer {
-        if (accessManager == address(0)) revert Errors.DisputeModule__ZeroAccessManager();
-
+        if (accessManager == address(0)) {
+            revert Errors.DisputeModule__ZeroAccessManager();
+        }
         __AccessManaged_init(accessManager);
         __ReentrancyGuard_init();
         __UUPSUpgradeable_init();

contracts/modules/dispute/policies/ArbitrationPolicySP.sol

@@ -52,8 +52,9 @@ contract ArbitrationPolicySP is IArbitrationPolicy, AccessManagedUpgradeable, UU
     /// @notice initializer for this implementation contract
     /// @param accessManager The address of the protocol admin roles contract
     function initialize(address accessManager) public initializer {
-        if (accessManager == address(0)) revert Errors.ArbitrationPolicySP__ZeroAccessManager();
-
+        if (accessManager == address(0)) {
+            revert Errors.ArbitrationPolicySP__ZeroAccessManager();
+        }
         __AccessManaged_init(accessManager);
         __UUPSUpgradeable_init();
     }

contracts/modules/licensing/LicensingModule.sol

@@ -99,6 +99,9 @@ contract LicensingModule is
     function initialize(address accessManager) public initializer {
         __ReentrancyGuard_init();
         __UUPSUpgradeable_init();
+        if (accessManager == address(0)) {
+            revert Errors.LicensingModule__ZeroAccessManager();
+        }
         __AccessManaged_init(accessManager);
     }
 

contracts/modules/royalty/RoyaltyModule.sol

@@ -58,7 +58,9 @@ contract RoyaltyModule is
     /// @notice Initializer for this implementation contract
     /// @param accessManager The address of the protocol admin roles contract
     function initialize(address accessManager) external initializer {
-        if (accessManager == address(0)) revert Errors.RoyaltyModule__ZeroAccessManager();
+        if (accessManager == address(0)) {
+            revert Errors.RoyaltyModule__ZeroAccessManager();
+        }
         __AccessManaged_init(accessManager);
         __ReentrancyGuard_init();
         __UUPSUpgradeable_init();

contracts/registries/IPAccountRegistry.sol

@@ -11,14 +11,18 @@ import { Errors } from "../lib/Errors.sol";
 /// It leverages a public ERC6551 registry to deploy IPAccount contracts.
 abstract contract IPAccountRegistry is IIPAccountRegistry {
     /// @notice Returns the IPAccount implementation address
+    /// @custom:oz-upgrades-unsafe-allow state-variable-immutable
     address public immutable IP_ACCOUNT_IMPL;
 
     /// @notice Returns the IPAccount salt
+    /// @custom:oz-upgrades-unsafe-allow state-variable-immutable
     bytes32 public immutable IP_ACCOUNT_SALT;
 
     /// @notice Returns the public ERC6551 registry address
+    /// @custom:oz-upgrades-unsafe-allow state-variable-immutable
     address public immutable ERC6551_PUBLIC_REGISTRY;
 
+    /// @custom:oz-upgrades-unsafe-allow constructor
     constructor(address erc6551Registry, address ipAccountImpl) {
         if (ipAccountImpl == address(0)) revert Errors.IPAccountRegistry_InvalidIpAccountImpl();
         IP_ACCOUNT_IMPL = ipAccountImpl;

contracts/registries/IPAssetRegistry.sol

@@ -5,6 +5,9 @@ import { IERC721 } from "@openzeppelin/contracts/token/ERC721/IERC721.sol";
 import { IERC721Metadata } from "@openzeppelin/contracts/token/ERC721/extensions/IERC721Metadata.sol";
 import { ERC165Checker } from "@openzeppelin/contracts/utils/introspection/ERC165Checker.sol";
 import { Strings } from "@openzeppelin/contracts/utils/Strings.sol";
+import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+// solhint-disable-next-line max-line-length
+import { AccessManagedUpgradeable } from "@openzeppelin/contracts-upgradeable/access/manager/AccessManagedUpgradeable.sol";
 
 import { IIPAccount } from "../interfaces/IIPAccount.sol";
 import { IIPAssetRegistry } from "../interfaces/registries/IIPAssetRegistry.sol";
@@ -21,15 +24,36 @@ import { IPAccountStorageOps } from "../lib/IPAccountStorageOps.sol";
 ///         attribution and an IP account for protocol authorization.
 ///         IMPORTANT: The IP account address, besides being used for protocol
 ///                    auth, is also the canonical IP identifier for the IP NFT.
-contract IPAssetRegistry is IIPAssetRegistry, IPAccountRegistry {
+contract IPAssetRegistry is IIPAssetRegistry, IPAccountRegistry, AccessManagedUpgradeable, UUPSUpgradeable {
     using ERC165Checker for address;
     using Strings for *;
     using IPAccountStorageOps for IIPAccount;
 
+    /// @dev Storage structure for the IPAssetRegistry
     /// @notice Tracks the total number of IP assets in existence.
-    uint256 public totalSupply = 0;
+    /// @custom:storage-location erc7201:story-protocol.IPAssetRegistry
+    struct IPAssetRegistryStorage {
+        uint256 totalSupply;
+    }
+
+    // keccak256(abi.encode(uint256(keccak256("story-protocol.IPAssetRegistry")) - 1)) & ~bytes32(uint256(0xff));
+    bytes32 private constant IPAssetRegistryStorageLocation =
+        0x987c61809af5a42943abd137c7acff8426aab6f7a1f5c967a03d1d718ba5cf00;
+
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor(address erc6551Registry, address ipAccountImpl) IPAccountRegistry(erc6551Registry, ipAccountImpl) {
+        _disableInitializers();
+    }
 
-    constructor(address erc6551Registry, address ipAccountImpl) IPAccountRegistry(erc6551Registry, ipAccountImpl) {}
+    /// @notice Initializes the IPAssetRegistry contract.
+    /// @param accessManager The address of the access manager.
+    function initialize(address accessManager) public initializer {
+        if (accessManager == address(0)) {
+            revert Errors.IPAssetRegistry__ZeroAccessManager();
+        }
+        __AccessManaged_init(accessManager);
+        __UUPSUpgradeable_init();
+    }
 
     /// @notice Registers an NFT as an IP asset.
     /// @dev The IP required metadata name and URI are derived from the NFT's metadata.
@@ -69,7 +93,7 @@ contract IPAssetRegistry is IIPAssetRegistry, IPAccountRegistry {
         ipAccount.setString("URI", uri);
         ipAccount.setUint256("REGISTRATION_DATE", registrationDate);
 
-        totalSupply++;
+        _getIPAssetRegistryStorage().totalSupply++;
 
         emit IPRegistered(id, block.chainid, tokenContract, tokenId, name, uri, registrationDate);
     }
@@ -95,4 +119,20 @@ contract IPAssetRegistry is IIPAssetRegistry, IPAccountRegistry {
         if (id != ipAccount(chainId, tokenContract, tokenId)) return false;
         return bytes(IIPAccount(payable(id)).getString("NAME")).length != 0;
     }
+
+    /// @notice Gets the total number of IP assets registered in the protocol.
+    function totalSupply() external view returns (uint256) {
+        return _getIPAssetRegistryStorage().totalSupply;
+    }
+
+    /// @dev Hook to authorize the upgrade according to UUPSUpgradeable
+    /// @param newImplementation The address of the new implementation
+    function _authorizeUpgrade(address newImplementation) internal override restricted {}
+
+    /// @dev Returns the storage struct of IPAssetRegistry.
+    function _getIPAssetRegistryStorage() private pure returns (IPAssetRegistryStorage storage $) {
+        assembly {
+            $.slot := IPAssetRegistryStorageLocation
+        }
+    }
 }

contracts/registries/LicenseRegistry.sol

@@ -80,6 +80,9 @@ contract LicenseRegistry is ILicenseRegistry, AccessManagedUpgradeable, UUPSUpgr
     /// @notice initializer for this implementation contract
     /// @param accessManager The address of the protocol admin roles contract
     function initialize(address accessManager) public initializer {
+        if (accessManager == address(0)) {
+            revert Errors.LicenseRegistry__ZeroAccessManager();
+        }
         __AccessManaged_init(accessManager);
         __UUPSUpgradeable_init();
     }

contracts/registries/ModuleRegistry.sol

@@ -41,6 +41,9 @@ contract ModuleRegistry is IModuleRegistry, AccessManagedUpgradeable, UUPSUpgrad
     /// @notice initializer for this implementation contract
     /// @param accessManager The address of the governance.
     function initialize(address accessManager) public initializer {
+        if (accessManager == address(0)) {
+            revert Errors.ModuleRegistry__ZeroAccessManager();
+        }
         __AccessManaged_init(accessManager);
         __UUPSUpgradeable_init();
 

script/foundry/utils/DeployHelper.sol

@@ -192,7 +192,19 @@ contract DeployHelper is Script, BroadcastManager, JsonDeploymentHandler, Storag
 
         contractKey = "IPAssetRegistry";
         _predeploy(contractKey);
-        ipAssetRegistry = new IPAssetRegistry(address(erc6551Registry), address(ipAccountImpl));
+        impl = address(
+            new IPAssetRegistry(
+                address(erc6551Registry),
+                address(ipAccountImpl)
+            )
+        );
+        ipAssetRegistry = IPAssetRegistry(
+            TestProxyHelper.deployUUPSProxy(
+                impl,
+                abi.encodeCall(IPAssetRegistry.initialize, address(protocolAccessManager))
+            )
+        );
+        impl = address(0); // Make sure we don't deploy wrong impl
         _postdeploy(contractKey, address(ipAssetRegistry));
 
         IPAccountRegistry ipAccountRegistry = IPAccountRegistry(address(ipAssetRegistry));
@@ -420,6 +432,7 @@ contract DeployHelper is Script, BroadcastManager, JsonDeploymentHandler, Storag
         protocolAccessManager.setTargetFunctionRole(address(royaltyPolicyLAP), selectors, ProtocolAdmin.UPGRADER_ROLE);
         protocolAccessManager.setTargetFunctionRole(address(licenseRegistry), selectors, ProtocolAdmin.UPGRADER_ROLE);
         protocolAccessManager.setTargetFunctionRole(address(moduleRegistry), selectors, ProtocolAdmin.UPGRADER_ROLE);
+        protocolAccessManager.setTargetFunctionRole(address(ipAssetRegistry), selectors, ProtocolAdmin.UPGRADER_ROLE);
 
         ///////// Role Granting /////////
         protocolAccessManager.grantRole(ProtocolAdmin.UPGRADER_ROLE, multisig, upgraderExecDelay);

script/foundry/utils/upgrades/ERC7201Helper.s.sol

@@ -12,7 +12,7 @@ import { console2 } from "forge-std/console2.sol";
 contract ERC7201HelperScript is Script {
 
     string constant NAMESPACE = "story-protocol";
-    string constant CONTRACT_NAME = "IpRoyaltyVault";
+    string constant CONTRACT_NAME = "IPAssetRegistry";
 
     function run() external {
         bytes memory erc7201Key = abi.encodePacked(NAMESPACE, ".", CONTRACT_NAME);

test/foundry/integration/flows/licensing/LicensingIntegration.t.sol

@@ -92,7 +92,13 @@ contract e2e is Test {
 
         erc6551Registry = new ERC6551Registry();
         ipAccountImpl = new IPAccountImpl(address(accessController));
-        ipAssetRegistry = new IPAssetRegistry(address(erc6551Registry), address(ipAccountImpl));
+        impl = address(new IPAssetRegistry(address(erc6551Registry), address(ipAccountImpl)));
+        ipAssetRegistry = IPAssetRegistry(
+            TestProxyHelper.deployUUPSProxy(
+                impl,
+                abi.encodeCall(IPAssetRegistry.initialize, address(protocolAccessManager))
+            )
+        );
 
         impl = address(new LicenseRegistry());
         licenseRegistry = LicenseRegistry(