migratables: transparent -> uups

src/contracts/MigratableEntity.sol

@@ -1,38 +1,57 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.25;
 
-import {IMigratableEntityProxy} from "src/interfaces/IMigratableEntityProxy.sol";
 import {IMigratableEntity} from "src/interfaces/IMigratableEntity.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";
+import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 
-abstract contract MigratableEntity is IMigratableEntity, Initializable, OwnableUpgradeable {
-    using Address for address;
+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 onlyProxyAdmin() {
+        if (msg.sender != proxyAdmin()) {
+            revert NotProxyAdmin();
+        }
+        _;
+    }
 
     constructor() {
         _disableInitializers();
     }
 
+    /**
+     * @inheritdoc IMigratableEntity
+     */
+    function proxyAdmin() public view returns (address) {
+        return _getMigratableEntityStorage()._proxyAdmin;
+    }
+
     /**
      * @inheritdoc IMigratableEntity
      */
     function initialize(bytes memory data) public virtual initializer {
         address owner = abi.decode(data, (address));
         __Ownable_init(owner);
+        __UUPSUpgradeable_init();
+
+        MigratableEntityStorage storage $ = _getMigratableEntityStorage();
+        $._proxyAdmin = msg.sender;
     }
 
     /**
      * @inheritdoc IMigratableEntity
      */
-    function migrate(bytes memory) public virtual reinitializer(_getInitializedVersion() + 1) {
-        address proxyAdmin = abi.decode(
-            address(this).functionStaticCall(abi.encodeWithSelector(IMigratableEntityProxy.proxyAdmin.selector)),
-            (address)
-        );
-        if (msg.sender != proxyAdmin) {
-            revert NotProxyAdmin();
+    function migrate(bytes memory) public virtual onlyProxyAdmin reinitializer(_getInitializedVersion() + 1) {}
+
+    function _authorizeUpgrade(address newImplementation) internal override onlyProxyAdmin {}
+
+    function _getMigratableEntityStorage() private pure returns (MigratableEntityStorage storage $) {
+        assembly {
+            $.slot := MigratableEntityStorageLocation
         }
     }
 }

src/contracts/MigratablesRegistry.sol

@@ -2,20 +2,16 @@
 pragma solidity 0.8.25;
 
 import {IMigratablesRegistry} from "src/interfaces/IMigratablesRegistry.sol";
-import {IMigratableEntityProxy} from "src/interfaces/IMigratableEntityProxy.sol";
 
-import {MigratableEntityProxy} from "./MigratableEntityProxy.sol";
 import {MigratableEntity} from "./MigratableEntity.sol";
 import {Registry} from "./Registry.sol";
 
-import {ITransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
+import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {Address} from "@openzeppelin/contracts/utils/Address.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
-import {ProxyAdmin} from "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol";
 
 contract MigratablesRegistry is Registry, Ownable, IMigratablesRegistry {
-    using Address for address;
     using EnumerableSet for EnumerableSet.AddressSet;
 
     EnumerableSet.AddressSet private _whitelistedImplementations;
@@ -47,12 +43,32 @@ contract MigratablesRegistry is Registry, Ownable, IMigratablesRegistry {
         }
     }
 
+    /**
+     * @inheritdoc IMigratablesRegistry
+     */
+    function create(uint256 version_, bytes memory data) external returns (address entity_) {
+        uint256 maxVersion_ = maxVersion();
+        if (version_ == 0 || version_ > maxVersion_) {
+            revert InvalidVersion();
+        }
+
+        entity_ = address(
+            new ERC1967Proxy(
+                _whitelistedImplementations.at(version_ - 1),
+                abi.encodeWithSelector(MigratableEntity.initialize.selector, data)
+            )
+        );
+
+        _addEntity(entity_);
+        _versions[entity_] = version_;
+    }
+
     /**
      * @inheritdoc IMigratablesRegistry
      */
     function migrate(address entity_, bytes memory data) external checkEntity(entity_) {
         if (msg.sender != MigratableEntity(entity_).owner()) {
-            revert ImproperOwner();
+            revert NotOwner();
         }
 
         uint256 currentVersion = _versions[entity_];
@@ -63,31 +79,9 @@ contract MigratablesRegistry is Registry, Ownable, IMigratablesRegistry {
 
         _versions[entity_] = currentVersion + 1;
 
-        address proxyAdmin = abi.decode(
-            entity_.functionStaticCall(abi.encodeWithSelector(IMigratableEntityProxy.proxyAdmin.selector)), (address)
-        );
-        ProxyAdmin(proxyAdmin).upgradeAndCall(
-            ITransparentUpgradeableProxy(entity_),
+        UUPSUpgradeable(entity_).upgradeToAndCall(
             _whitelistedImplementations.at(currentVersion),
             abi.encodeWithSelector(MigratableEntity.migrate.selector, data)
         );
     }
-
-    function create(uint256 version_, bytes memory data) external returns (address entity_) {
-        uint256 maxVersion_ = maxVersion();
-        if (version_ == 0 || version_ > maxVersion_) {
-            revert InvalidVersion();
-        }
-
-        entity_ = address(
-            new MigratableEntityProxy(
-                _whitelistedImplementations.at(version_ - 1),
-                address(this),
-                abi.encodeWithSelector(MigratableEntity.initialize.selector, data)
-            )
-        );
-
-        _addEntity(entity_);
-        _versions[entity_] = version_;
-    }
 }

src/interfaces/IMigratableEntity.sol

@@ -4,6 +4,12 @@ pragma solidity 0.8.25;
 interface IMigratableEntity {
     error NotProxyAdmin();
 
+    struct MigratableEntityStorage {
+        address _proxyAdmin;
+    }
+
+    function proxyAdmin() external view returns (address);
+
     /**
      * @notice Initialize this entity contract using a given data.
      * @param data some data to use

src/interfaces/IMigratablesRegistry.sol

@@ -6,7 +6,7 @@ import {IRegistry} from "./IRegistry.sol";
 interface IMigratablesRegistry is IRegistry {
     error AlreadyWhitelisted();
     error AlreadyUpToDate();
-    error ImproperOwner();
+    error NotOwner();
     error InvalidVersion();
 
     /**

test/MigratablesRegistry.sol

@@ -109,7 +109,7 @@ contract MigratablesRegistryTest is Test {
         registry.whitelist(implV2);
 
         vm.startPrank(bob);
-        vm.expectRevert(IMigratablesRegistry.ImproperOwner.selector);
+        vm.expectRevert(IMigratablesRegistry.NotOwner.selector);
         registry.migrate(entity, "");
         vm.stopPrank();
     }