@piplabs/story-contracts 0.1.0-alpha.0

package/test/upgrades/PredeployUpgrades.t.sol

@@ -0,0 +1,219 @@
+// SPDX-License-Identifier: GPL-3.0-only
+pragma solidity 0.8.23;
+
+import { ProxyAdmin } from "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol";
+// solhint-disable-next-line max-line-length
+import { ITransparentUpgradeableProxy } from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
+
+import { IPTokenStaking } from "../../src/protocol/IPTokenStaking.sol";
+import { UpgradeEntrypoint } from "../../src/protocol/UpgradeEntrypoint.sol";
+import { UBIPool } from "../../src/protocol/UBIPool.sol";
+
+import { EIP1967Helper } from "../../script/utils/EIP1967Helper.sol";
+import { Predeploys } from "../../src/libraries/Predeploys.sol";
+import { Test } from "../utils/Test.sol";
+
+import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+
+abstract contract MockNewFeatures {
+    function foo() external pure returns (string memory) {
+        return "bar";
+    }
+}
+
+contract IPTokenStakingV2 is IPTokenStaking, MockNewFeatures {
+    constructor(uint256 defaultMinFee, uint256 maxDataLength) IPTokenStaking(defaultMinFee, maxDataLength) {}
+}
+
+contract UpgradeEntrypointV2 is UpgradeEntrypoint, MockNewFeatures {}
+
+contract UBIPoolV2 is UBIPool, MockNewFeatures {
+    constructor(uint32 maxUBIPercentage) UBIPool(maxUBIPercentage) {}
+}
+
+contract InitialImplementation {
+    function foo() external pure returns (string memory) {
+        return "bar";
+    }
+}
+
+/**
+ * @title PredeployUpgrades
+ * @dev A script to test upgrading the precompile contracts
+ */
+contract PredeployUpgrades is Test {
+    function testUpgradeStaking() public {
+        // ---- Staking
+        address newImpl = address(
+            new IPTokenStakingV2(
+                1 ether, // defaultMinFee
+                256 // maxDataLength
+            )
+        );
+        ProxyAdmin proxyAdmin = ProxyAdmin(EIP1967Helper.getAdmin(Predeploys.Staking));
+        assertEq(proxyAdmin.owner(), address(timelock));
+
+        performTimelocked(
+            address(proxyAdmin),
+            abi.encodeWithSelector(
+                ProxyAdmin.upgradeAndCall.selector,
+                ITransparentUpgradeableProxy(Predeploys.Staking),
+                newImpl,
+                ""
+            )
+        );
+
+        assertEq(EIP1967Helper.getImplementation(Predeploys.Staking), newImpl, "Staking not upgraded");
+        assertEq(
+            keccak256(abi.encode(IPTokenStakingV2(Predeploys.Staking).foo())),
+            keccak256(abi.encode("bar")),
+            "Upgraded to wrong iface"
+        );
+    }
+
+    function testUpgradeUpgradeEntrypoint() public {
+        // ---- Upgrades
+        address newImpl = address(new UpgradeEntrypointV2());
+        ProxyAdmin proxyAdmin = ProxyAdmin(EIP1967Helper.getAdmin(Predeploys.Upgrades));
+        assertEq(proxyAdmin.owner(), address(timelock));
+
+        performTimelocked(
+            address(proxyAdmin),
+            abi.encodeWithSelector(
+                ProxyAdmin.upgradeAndCall.selector,
+                ITransparentUpgradeableProxy(Predeploys.Upgrades),
+                newImpl,
+                ""
+            )
+        );
+        assertEq(EIP1967Helper.getImplementation(Predeploys.Upgrades), newImpl, "Upgrades not upgraded");
+        assertEq(
+            keccak256(abi.encode(IPTokenStakingV2(Predeploys.Upgrades).foo())),
+            keccak256(abi.encode("bar")),
+            "Upgraded to wrong iface"
+        );
+    }
+
+    function testUpgradeUBIPool() public {
+        // ---- UBIPool
+        address newImpl = address(new UBIPoolV2(10_00));
+        ProxyAdmin proxyAdmin = ProxyAdmin(EIP1967Helper.getAdmin(Predeploys.UBIPool));
+        assertEq(proxyAdmin.owner(), address(timelock));
+
+        performTimelocked(
+            address(proxyAdmin),
+            abi.encodeWithSelector(
+                ProxyAdmin.upgradeAndCall.selector,
+                ITransparentUpgradeableProxy(Predeploys.UBIPool),
+                newImpl,
+                ""
+            )
+        );
+        assertEq(EIP1967Helper.getImplementation(Predeploys.UBIPool), newImpl, "Upgrades not upgraded");
+        assertEq(
+            keccak256(abi.encode(IPTokenStakingV2(Predeploys.UBIPool).foo())),
+            keccak256(abi.encode("bar")),
+            "Upgraded to wrong iface"
+        );
+    }
+
+    function testUpgradeUnusedProxies() public {
+        address initialImpl = address(new InitialImplementation());
+
+        for (
+            uint160 i = uint160(Predeploys.Upgrades) + uint160(1);
+            i <= uint160(Predeploys.Namespace) + Predeploys.NamespaceSize;
+            i++
+        ) {
+            // Verify predeploy is proxied and not upgraded
+            address predeploy = address(i);
+            assertTrue(Predeploys.proxied(predeploy), "Predeploy not proxied");
+            assertEq(
+                EIP1967Helper.getImplementation(predeploy),
+                Predeploys.getImplAddress(predeploy),
+                "Predeploy upgraded already"
+            );
+            ProxyAdmin proxyAdmin = ProxyAdmin(EIP1967Helper.getAdmin(predeploy));
+            assertEq(proxyAdmin.owner(), address(timelock));
+
+            // Revert if not owner
+            vm.expectRevert();
+            proxyAdmin.upgradeAndCall(ITransparentUpgradeableProxy(predeploy), initialImpl, "");
+
+            // Upgrade predeploy
+            performTimelocked(
+                address(proxyAdmin),
+                abi.encodeWithSelector(
+                    ProxyAdmin.upgradeAndCall.selector,
+                    ITransparentUpgradeableProxy(predeploy),
+                    initialImpl,
+                    ""
+                )
+            );
+
+            // Verify predeploy is upgraded
+            assertEq(EIP1967Helper.getImplementation(predeploy), initialImpl, "Predeploy not upgraded");
+            assertEq(
+                keccak256(abi.encode(InitialImplementation(predeploy).foo())),
+                keccak256(abi.encode("bar")),
+                "Upgraded to wrong iface"
+            );
+        }
+    }
+
+    function testRenounceUpgradeability() public {
+        address newImpl = address(new InitialImplementation());
+        // For each predeploy, renounce upgradeability
+        for (
+            uint160 i = uint160(Predeploys.Namespace) + uint160(1);
+            i <= uint160(Predeploys.Namespace) + Predeploys.NamespaceSize;
+            i++
+        ) {
+            address predeploy = address(i);
+            ProxyAdmin proxyAdmin = ProxyAdmin(EIP1967Helper.getAdmin(predeploy));
+            assertEq(proxyAdmin.owner(), address(timelock));
+            // Renounce ownership
+            performTimelocked(address(proxyAdmin), abi.encodeWithSelector(Ownable.renounceOwnership.selector));
+
+            // Verify ownership was renounced
+            assertEq(proxyAdmin.owner(), address(0), "Ownership not renounced");
+
+            // Try to upgrade
+            bytes memory upgradeAction = abi.encodeWithSelector(
+                ProxyAdmin.upgradeAndCall.selector,
+                ITransparentUpgradeableProxy(predeploy),
+                newImpl,
+                ""
+            );
+            bytes memory expectedReason = abi.encodeWithSelector(
+                Ownable.OwnableUnauthorizedAccount.selector,
+                address(timelock)
+            );
+            expectRevertTimelocked(address(proxyAdmin), upgradeAction, expectedReason);
+        }
+    }
+
+    function testUpgradeLastPredeploy() public {
+        address newImpl = address(new InitialImplementation());
+        address lastPredeploy = address(uint160(Predeploys.Namespace) + uint160(Predeploys.NamespaceSize));
+        ProxyAdmin proxyAdmin = ProxyAdmin(EIP1967Helper.getAdmin(lastPredeploy));
+        assertEq(proxyAdmin.owner(), address(timelock));
+
+        // Upgrade last predeploy to new implementation
+        performTimelocked(
+            address(proxyAdmin),
+            abi.encodeWithSelector(
+                ProxyAdmin.upgradeAndCall.selector,
+                ITransparentUpgradeableProxy(lastPredeploy),
+                newImpl,
+                ""
+            )
+        );
+        assertEq(EIP1967Helper.getImplementation(lastPredeploy), newImpl, "Last predeploy not upgraded");
+        assertEq(
+            keccak256(abi.encode(InitialImplementation(lastPredeploy).foo())),
+            keccak256(abi.encode("bar")),
+            "Upgraded to wrong iface"
+        );
+    }
+}

package/test/utils/Test.sol

@@ -0,0 +1,170 @@
+// SPDX-License-Identifier: GPL-3.0-only
+pragma solidity 0.8.23;
+/* solhint-disable no-console */
+/* solhint-disable max-line-length */
+
+import { Test as ForgeTest } from "forge-std/Test.sol";
+
+import { IPTokenStaking } from "../../src/protocol/IPTokenStaking.sol";
+import { UpgradeEntrypoint } from "../../src/protocol/UpgradeEntrypoint.sol";
+import { UBIPool } from "../../src/protocol/UBIPool.sol";
+import { Predeploys } from "../../src/libraries/Predeploys.sol";
+import { Create3 } from "../../src/deploy/Create3.sol";
+import { GenerateAlloc } from "../../script/GenerateAlloc.s.sol";
+import { TimelockController } from "@openzeppelin/contracts/governance/TimelockController.sol";
+import { ERC6551Registry } from "erc6551/ERC6551Registry.sol";
+import { WIP } from "../../src/token/WIP.sol";
+
+contract Test is ForgeTest {
+    address internal admin = address(0x123);
+    address internal executor = address(0x456);
+    address internal guardian = address(0x789);
+
+    address internal deployer = address(0xDDdDddDdDdddDDddDDddDDDDdDdDDdDDdDDDDDDd);
+
+    IPTokenStaking internal ipTokenStaking;
+    UpgradeEntrypoint internal upgradeEntrypoint;
+    UBIPool internal ubiPool;
+    Create3 internal create3;
+    ERC6551Registry internal erc6551Registry;
+    TimelockController internal timelock;
+    WIP internal wip;
+
+    function setUp() public virtual {
+        GenerateAlloc initializer = new GenerateAlloc();
+        initializer.disableStateDump(); // Faster tests. Don't call to verify JSON output
+        initializer.setAdminAddresses(admin, executor, guardian);
+        initializer.run();
+        ipTokenStaking = IPTokenStaking(Predeploys.Staking);
+        upgradeEntrypoint = UpgradeEntrypoint(Predeploys.Upgrades);
+        ubiPool = UBIPool(Predeploys.UBIPool);
+        create3 = Create3(Predeploys.Create3);
+        wip = WIP(payable(Predeploys.WIP));
+        erc6551Registry = ERC6551Registry(Predeploys.ERC6551Registry);
+        address timelockAddress = create3.getDeployed(deployer, keccak256("STORY_TIMELOCK_CONTROLLER"));
+        timelock = TimelockController(payable(timelockAddress));
+        require(timelockAddress.code.length > 0, "Timelock not deployed");
+    }
+
+    /// @notice schedules, waits for timelock and executes a timelocked call
+    /// @param target The address to call
+    /// @param data The data to call with
+    function performTimelocked(address target, bytes memory data) internal {
+        uint256 minDelay = timelock.getMinDelay();
+        vm.prank(admin);
+        timelock.schedule(
+            target,
+            0, // value
+            data,
+            bytes32(0), // predecessor: Non Zero if order must be respected
+            bytes32(keccak256("SALT")), // salt
+            minDelay
+        );
+        vm.warp(block.timestamp + minDelay + 1);
+        vm.prank(executor);
+        timelock.execute(
+            target,
+            0, // value
+            data,
+            bytes32(0), // predecessor: Non Zero if order must be respected
+            bytes32(keccak256("SALT")) // salt
+        );
+    }
+
+    /// @notice schedules, waits for timelock and executes a timelocked call, with a custom salt
+    /// @dev This is to be used if we want to call the same target with the same data multiple times
+    /// @param target The address to call
+    /// @param data The data to call with
+    /// @param salt The salt to use for the timelock
+    function performTimelocked(address target, bytes memory data, bytes32 salt) internal {
+        uint256 minDelay = timelock.getMinDelay();
+        vm.prank(admin);
+        timelock.schedule(
+            target,
+            0, // value
+            data,
+            bytes32(0), // predecessor: Non Zero if order must be respected
+            bytes32(salt), // salt
+            minDelay
+        );
+        vm.warp(block.timestamp + minDelay + 1);
+        vm.prank(executor);
+        timelock.execute(
+            target,
+            0, // value
+            data,
+            bytes32(0), // predecessor: Non Zero if order must be respected
+            bytes32(salt) // salt
+        );
+    }
+
+    /// @notice schedules a timelocked call
+    /// @param target The address to call
+    /// @param data The data to call with
+    function schedule(address target, bytes memory data) internal {
+        uint256 minDelay = timelock.getMinDelay();
+        vm.prank(admin);
+        timelock.schedule(
+            target,
+            0, // value
+            data,
+            bytes32(0), // predecessor: Non Zero if order must be respected
+            bytes32(keccak256("SALT")), // salt
+            minDelay
+        );
+    }
+
+    /// @notice waits for timelock (minDelay)
+    /// @dev This is to be called after schedule()
+    /// If the scheduled time > minDelay, this wait time won't be enough
+    /// and the test will revert
+    function waitForTimelock() internal {
+        uint256 minDelay = timelock.getMinDelay();
+        vm.warp(block.timestamp + minDelay + 1);
+    }
+
+    /// @notice executes a timelocked call
+    /// @param target The address to call
+    /// @param data The data to call with
+    function executeTimelocked(address target, bytes memory data) internal {
+        vm.prank(executor);
+        timelock.execute(
+            target,
+            0, // value
+            data,
+            bytes32(0), // predecessor: Non Zero if order must be respected
+            bytes32(keccak256("SALT")) // salt
+        );
+    }
+
+    /// @notice schedules, waits for timelock and executes a timelocked call that is expected to revert
+    /// @param target The address to call
+    /// @param data The data to call with
+    /// @param reason The expected revert reason
+    function expectRevertTimelocked(address target, bytes memory data, bytes memory reason) internal {
+        uint256 minDelay = timelock.getMinDelay();
+        vm.prank(admin);
+        timelock.schedule(
+            target,
+            0, // value
+            data,
+            bytes32(0), // predecessor: Non Zero if order must be respected
+            bytes32(keccak256("SALT")), // salt
+            minDelay
+        );
+        vm.warp(block.timestamp + minDelay + 1);
+        vm.prank(executor);
+        vm.expectRevert(bytes(reason));
+        timelock.execute(
+            target,
+            0, // value
+            data,
+            bytes32(0), // predecessor: Non Zero if order must be respected
+            bytes32(keccak256("SALT")) // salt
+        );
+        // Cancel the scheduled call to clean the hash in the timelock
+        bytes32 id = timelock.hashOperation(target, 0, data, bytes32(0), bytes32(keccak256("SALT")));
+        vm.prank(admin);
+        timelock.cancel(id);
+    }
+}