feat: Upgradeable Proxy (#15)

m0-foundation · Jun 20, 2024 · 0d82736 · 0d82736
1 parent 05230c0
commit 0d82736
Show file tree

Hide file tree

Showing 6 changed files with 168 additions and 6 deletions.
diff --git a/src/Migratable.sol b/src/Migratable.sol
@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: GPL-3.0
+
+pragma solidity 0.8.23;
+
+import { IMigratable } from "./interfaces/IMigratable.sol";
+
+abstract contract Migratable is IMigratable {
+    /* ============ Variables ============ */
+
+    /// @dev Storage slot with the address of the current factory. `keccak256('eip1967.proxy.implementation') - 1`.
+    bytes32 private constant _IMPLEMENTATION_SLOT =
+        bytes32(0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc);
+
+    /* ============ Interactive Functions ============ */
+
+    function migrate() external {
+        address migrator_ = _getMigrator();
+
+        if (migrator_ == address(0)) revert ZeroMigrator();
+
+        address oldImplementation_ = implementation();
+
+        migrator_.delegatecall("");
+
+        emit Migrate(migrator_, oldImplementation_, implementation());
+    }
+
+    /* ============ View/Pure Functions ============ */
+
+    function implementation() public view returns (address implementation_) {
+        bytes32 slot_ = _IMPLEMENTATION_SLOT;
+
+        assembly {
+            implementation_ := sload(slot_)
+        }
+    }
+
+    /* ============ Internal View/Pure Functions ============ */
+
+    function _getMigrator() internal view virtual returns (address migrator_);
+}
diff --git a/src/Proxy.sol b/src/Proxy.sol
@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: GPL-3.0
+
+pragma solidity 0.8.23;
+
+contract Proxy {
+    /// @dev Storage slot with the address of the current factory. `keccak256('eip1967.proxy.implementation') - 1`.
+    bytes32 private constant _IMPLEMENTATION_SLOT =
+        bytes32(0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc);
+
+    constructor(address implementation_) {
+        if (implementation_ == address(0)) revert();
+
+        bytes32 slot_ = _IMPLEMENTATION_SLOT;
+
+        assembly {
+            sstore(slot_, implementation_)
+        }
+    }
+
+    fallback() external payable virtual {
+        bytes32 slot_ = _IMPLEMENTATION_SLOT;
+        bytes32 implementation_;
+
+        assembly {
+            implementation_ := sload(slot_)
+        }
+
+        assembly {
+            calldatacopy(0, 0, calldatasize())
+
+            let result_ := delegatecall(gas(), implementation_, 0, calldatasize(), 0, 0)
+
+            returndatacopy(0, 0, returndatasize())
+
+            switch result_
+            case 0 {
+                revert(0, returndatasize())
+            }
+            default {
+                return(0, returndatasize())
+            }
+        }
+    }
+}
diff --git a/src/WrappedM.sol b/src/WrappedM.sol
@@ -12,7 +12,9 @@ import { IMTokenLike } from "./interfaces/IMTokenLike.sol";
 import { IWrappedM } from "./interfaces/IWrappedM.sol";
 import { IRegistrarLike } from "./interfaces/IRegistrarLike.sol";
 
-contract WrappedM is IWrappedM, ERC20Extended {
+import { Migratable } from "./Migratable.sol";
+
+contract WrappedM is IWrappedM, Migratable, ERC20Extended {
     type BalanceInfo is uint256;
 
     /* ============ Variables ============ */
@@ -21,7 +23,8 @@ contract WrappedM is IWrappedM, ERC20Extended {
 
     bytes32 internal constant _EARNERS_LIST_IGNORED = "earners_list_ignored";
     bytes32 internal constant _EARNERS_LIST = "earners";
-    bytes32 internal constant _CLAIM_DESTINATION_PREFIX = "claim_destination";
+    bytes32 internal constant _CLAIM_DESTINATION_PREFIX = "wm_claim_destination";
+    bytes32 internal constant _MIGRATOR_V1_PREFIX = "wm_migrator_v1";
 
     address public immutable mToken;
     address public immutable registrar;
@@ -320,6 +323,15 @@ contract WrappedM is IWrappedM, ERC20Extended {
             );
     }
 
+    function _getMigrator() internal view override returns (address migrator_) {
+        return
+            address(
+                uint160(
+                    uint256(IRegistrarLike(registrar).get(keccak256(abi.encode(_MIGRATOR_V1_PREFIX, address(this)))))
+                )
+            );
+    }
+
     function _getTotalAccruedYield(uint128 currentIndex_) internal view returns (uint240 yield_) {
         uint240 totalProjectedSupply_ = IndexingMath.getPresentAmountRoundedUp(
             _principalOfTotalEarningSupply,

diff --git a/src/interfaces/IMigratable.sol b/src/interfaces/IMigratable.sol
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: GPL-3.0
+
+pragma solidity 0.8.23;
+
+interface IMigratable {
+    event Migrate(address indexed migrator, address indexed oldImplementation, address indexed newImplementation);
+
+    error ZeroMigrator();
+
+    function migrate() external;
+
+    function implementation() external view returns (address implementation);
+}
diff --git a/src/interfaces/IWrappedM.sol b/src/interfaces/IWrappedM.sol
@@ -4,7 +4,9 @@ pragma solidity 0.8.23;
 
 import { IERC20Extended } from "../../lib/common/src/interfaces/IERC20Extended.sol";
 
-interface IWrappedM is IERC20Extended {
+import { IMigratable } from "./IMigratable.sol";
+
+interface IWrappedM is IMigratable, IERC20Extended {
     /* ============ Events ============ */
 
     event Claim(address indexed account, uint256 yield);

diff --git a/test/Test.t.sol b/test/Test.t.sol
@@ -4,7 +4,10 @@ pragma solidity 0.8.23;
 
 import { Test, console2 } from "../lib/forge-std/src/Test.sol";
 
+import { IWrappedM } from "../src/interfaces/IWrappedM.sol";
+
 import { WrappedM } from "../src/WrappedM.sol";
+import { Proxy } from "../src/Proxy.sol";
 
 contract MockM {
     address public ttgRegistrar;
@@ -54,11 +57,38 @@ contract MockRegistrar {
     }
 }
 
+contract WrappedMV2 {
+    function foo() external pure returns (uint256) {
+        return 1;
+    }
+}
+
+contract WrappedMMigratorV1 {
+    bytes32 private constant _IMPLEMENTATION_SLOT =
+        bytes32(0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc);
+
+    address public immutable implementationV2;
+
+    constructor(address implementationV2_) {
+        implementationV2 = implementationV2_;
+    }
+
+    fallback() external virtual {
+        bytes32 slot_ = _IMPLEMENTATION_SLOT;
+        address implementationV2_ = implementationV2;
+
+        assembly {
+            sstore(slot_, implementationV2_)
+        }
+    }
+}
+
 contract Tests is Test {
     uint56 internal constant _EXP_SCALED_ONE = 1e12;
 
     bytes32 internal constant _EARNERS_LIST = "earners";
-    bytes32 internal constant _CLAIM_DESTINATION_PREFIX = "claim_destination";
+    bytes32 internal constant _CLAIM_DESTINATION_PREFIX = "wm_claim_destination";
+    bytes32 internal constant _MIGRATOR_V1_PREFIX = "wm_migrator_v1";
 
     address internal _alice = makeAddr("alice");
     address internal _bob = makeAddr("bob");
@@ -69,7 +99,8 @@ contract Tests is Test {
 
     MockM internal _mToken;
     MockRegistrar internal _registrar;
-    WrappedM internal _wrappedM;
+    WrappedM internal _implementation;
+    IWrappedM internal _wrappedM;
 
     function setUp() external {
         _registrar = new MockRegistrar();
@@ -79,7 +110,9 @@ contract Tests is Test {
         _mToken.setCurrentIndex(_EXP_SCALED_ONE);
         _mToken.setTtgRegistrar(address(_registrar));
 
-        _wrappedM = new WrappedM(address(_mToken));
+        _implementation = new WrappedM(address(_mToken));
+
+        _wrappedM = IWrappedM(address(new Proxy(address(_implementation))));
     }
 
     function test_story() external {
@@ -409,4 +442,21 @@ contract Tests is Test {
         assertEq(_wrappedM.totalAccruedYield(), 1);
         assertEq(_wrappedM.excess(), 6);
     }
+
+    function test_migration() external {
+        WrappedMV2 implementationV2_ = new WrappedMV2();
+        address migrator_ = address(new WrappedMMigratorV1(address(implementationV2_)));
+
+        _registrar.set(
+            keccak256(abi.encode(_MIGRATOR_V1_PREFIX, address(_wrappedM))),
+            bytes32(uint256(uint160(migrator_)))
+        );
+
+        vm.expectRevert();
+        WrappedMV2(address(_wrappedM)).foo();
+
+        _wrappedM.migrate();
+
+        assertEq(WrappedMV2(address(_wrappedM)).foo(), 1);
+    }
 }