uups -> light proxy

src/contracts/base/MigratableEntity.sol

@@ -1,35 +1,20 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.25;
 
+import {IMigratableEntityProxy} from "src/interfaces/base/IMigratableEntityProxy.sol";
 import {IMigratableEntity} from "src/interfaces/base/IMigratableEntity.sol";
 
-import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
-import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import {Address} from "@openzeppelin/contracts/utils/Address.sol";
 import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 
-abstract contract MigratableEntity is Initializable, UUPSUpgradeable, OwnableUpgradeable, IMigratableEntity {
-    // keccak256(abi.encode(uint256(keccak256("symbiotic.storage.MigratableEntity")) - 1)) & ~bytes32(uint256(0xff))
-    bytes32 private constant MigratableEntityStorageLocation =
-        0x22b5f4baea4997f81f8aeb6360e0bdae13f074e0e55c27a8a6fab78cbad46200;
-
-    modifier onlyRegistry() {
-        if (msg.sender != registry()) {
-            revert NotRegistry();
-        }
-        _;
-    }
+abstract contract MigratableEntity is Initializable, OwnableUpgradeable, IMigratableEntity {
+    using Address for address;
 
     constructor() {
         _disableInitializers();
     }
 
-    /**
-     * @inheritdoc IMigratableEntity
-     */
-    function registry() public view returns (address) {
-        return _getMigratableEntityStorage()._registry;
-    }
-
     /**
      * @inheritdoc IMigratableEntity
      */
@@ -54,19 +39,15 @@ abstract contract MigratableEntity is Initializable, UUPSUpgradeable, OwnableUpg
 
     function _initialize(address owner) internal {
         __Ownable_init(owner);
-        __UUPSUpgradeable_init();
-
-        MigratableEntityStorage storage $ = _getMigratableEntityStorage();
-        $._registry = msg.sender;
     }
 
-    function _migrate() internal onlyRegistry {}
-
-    function _authorizeUpgrade(address newImplementation) internal override onlyRegistry {}
-
-    function _getMigratableEntityStorage() private pure returns (MigratableEntityStorage storage $) {
-        assembly {
-            $.slot := MigratableEntityStorageLocation
+    function _migrate() internal view {
+        address proxyAdmin = abi.decode(
+            address(this).functionStaticCall(abi.encodeWithSelector(IMigratableEntityProxy.proxyAdmin.selector)),
+            (address)
+        );
+        if (msg.sender != proxyAdmin) {
+            revert NotProxyAdmin();
         }
     }
 }

src/contracts/base/MigratableEntityProxy.sol

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.25;
+
+import {IMigratableEntityProxy} from "src/interfaces/base/IMigratableEntityProxy.sol";
+
+import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol";
+import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+
+contract MigratableEntityProxy is ERC1967Proxy, IMigratableEntityProxy {
+    // An immutable address for the admin to avoid unnecessary SLOADs before each call.
+    address private immutable _admin;
+
+    /**
+     * @dev The proxy caller is the current admin, and can't fallback to the proxy target.
+     */
+    error ProxyDeniedAdminAccess();
+
+    /**
+     * @dev Initializes an upgradeable proxy managed by `msg.sender`,
+     * backed by the implementation at `_logic`, and optionally initialized with `_data` as explained in
+     * {ERC1967Proxy-constructor}.
+     */
+    constructor(address _logic, bytes memory _data) payable ERC1967Proxy(_logic, _data) {
+        _admin = msg.sender;
+        // Set the storage value and emit an event for ERC-1967 compatibility
+        ERC1967Utils.changeAdmin(_proxyAdmin());
+    }
+
+    /**
+     * @inheritdoc IMigratableEntityProxy
+     */
+    function proxyAdmin() external view returns (address) {
+        return _proxyAdmin();
+    }
+
+    /**
+     * @inheritdoc IMigratableEntityProxy
+     */
+    function upgradeToAndCall(address newImplementation, bytes calldata data) external payable {
+        if (msg.sender != _proxyAdmin()) {
+            revert ProxyDeniedAdminAccess();
+        }
+
+        ERC1967Utils.upgradeToAndCall(newImplementation, data);
+    }
+
+    /**
+     * @dev Returns the admin of this proxy.
+     */
+    function _proxyAdmin() internal view returns (address) {
+        return _admin;
+    }
+}

src/contracts/base/MigratablesRegistry.sol

@@ -3,16 +3,18 @@ pragma solidity 0.8.25;
 
 import {IMigratablesRegistry} from "src/interfaces/base/IMigratablesRegistry.sol";
 import {IMigratableEntity} from "src/interfaces/base/IMigratableEntity.sol";
+import {IMigratableEntityProxy} from "src/interfaces/base/IMigratableEntityProxy.sol";
 
 import {Registry} from "./Registry.sol";
+import {MigratableEntityProxy} from "./MigratableEntityProxy.sol";
 
-import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
-import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
+import {Address} from "@openzeppelin/contracts/utils/Address.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 
 contract MigratablesRegistry is Registry, Ownable, IMigratablesRegistry {
     using EnumerableSet for EnumerableSet.AddressSet;
+    using Address for address;
 
     EnumerableSet.AddressSet private _whitelistedImplementations;
 
@@ -51,9 +53,9 @@ contract MigratablesRegistry is Registry, Ownable, IMigratablesRegistry {
     /**
      * @inheritdoc IMigratablesRegistry
      */
-    function create(uint64 version, bytes memory data) external isValidVersion(version) returns (address entity_) {
+    function create(uint64 version, bytes memory data) external returns (address entity_) {
         entity_ = address(
-            new ERC1967Proxy(
+            new MigratableEntityProxy(
                 implementation(version), abi.encodeWithSelector(IMigratableEntity.initialize.selector, version, data)
             )
         );
@@ -69,7 +71,7 @@ contract MigratablesRegistry is Registry, Ownable, IMigratablesRegistry {
             revert NotOwner();
         }
 
-        UUPSUpgradeable(entity_).upgradeToAndCall(
+        IMigratableEntityProxy(entity_).upgradeToAndCall(
             implementation(IMigratableEntity(entity_).version() + 1),
             abi.encodeWithSelector(IMigratableEntity.migrate.selector, data)
         );

src/interfaces/base/IMigratableEntity.sol

@@ -2,19 +2,7 @@
 pragma solidity 0.8.25;
 
 interface IMigratableEntity {
-    error NotRegistry();
-    error InvalidMigrateCall();
-
-    /// @custom:storage-location erc7201:symbiotic.storage.MigratableEntity
-    struct MigratableEntityStorage {
-        address _registry;
-    }
-
-    /**
-     * @notice Get the entity's registry.
-     * @return registry of the entity
-     */
-    function registry() external view returns (address);
+    error NotProxyAdmin();
 
     /**
      * @notice Get a the entity's version.

src/interfaces/base/IMigratableEntityProxy.sol

@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.25;
+
+import {IERC1967} from "@openzeppelin/contracts/interfaces/IERC1967.sol";
+
+interface IMigratableEntityProxy is IERC1967 {
+    /**
+     * @notice Get a proxy admin.
+     * @return address of the proxy admin
+     */
+    function proxyAdmin() external returns (address);
+
+    /**
+     * @notice Upgrade the proxy to a new implementation and call a function on the new implementation.
+     * @param newImplementation address of the new implementation
+     * @param data data to call on the new implementation
+     */
+    function upgradeToAndCall(address newImplementation, bytes calldata data) external payable;
+}

test/Vault.t.sol

@@ -2009,8 +2009,6 @@ contract VaultTest is Test {
         uint48 vetoDuration = 1;
         vault = _getVault(metadataURL, epochDuration, vetoDuration, slashDuration);
 
-        uint256 blockTimestamp = block.timestamp * block.timestamp / block.timestamp * block.timestamp / block.timestamp;
-
         address network = bob;
         _registerNetwork(network, bob);
 
@@ -2042,8 +2040,6 @@ contract VaultTest is Test {
         uint48 vetoDuration = 1;
         vault = _getVault(metadataURL, epochDuration, vetoDuration, slashDuration);
 
-        uint256 blockTimestamp = block.timestamp * block.timestamp / block.timestamp * block.timestamp / block.timestamp;
-
         address network = bob;
         _registerNetwork(network, bob);
 
@@ -2064,8 +2060,6 @@ contract VaultTest is Test {
         uint48 vetoDuration = 1;
         vault = _getVault(metadataURL, epochDuration, vetoDuration, slashDuration);
 
-        uint256 blockTimestamp = block.timestamp * block.timestamp / block.timestamp * block.timestamp / block.timestamp;
-
         address network = bob;
 
         address resolver = alice;

test/base/MigratablesRegistry.sol

@@ -146,7 +146,7 @@ contract MigratablesRegistryTest is Test {
         vm.stopPrank();
     }
 
-    function test_MigrateRevertNotRegistry() public {
+    function test_MigrateRevertNotProxyAdmin() public {
         address impl = address(new SimpleMigratableEntity());
         registry.whitelist(impl);
 
@@ -160,7 +160,7 @@ contract MigratablesRegistryTest is Test {
         vm.stopPrank();
 
         vm.startPrank(alice);
-        vm.expectRevert(IMigratableEntity.NotRegistry.selector);
+        vm.expectRevert(IMigratableEntity.NotProxyAdmin.selector);
         IMigratableEntity(entity).migrate(abi.encode(0));
         vm.stopPrank();
     }