diff --git a/script/deploy/DeployUpgradeTimelock.s.sol b/script/deploy/DeployUpgradeTimelock.s.sol new file mode 100644 index 00000000..134ab40f --- /dev/null +++ b/script/deploy/DeployUpgradeTimelock.s.sol @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.25; + +import {Script} from "forge-std/Script.sol"; +import {console} from "forge-std/console.sol"; +import {TransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol"; +import {UpgradeTimelock} from "../../src/UpgradeTimelock.sol"; + +/** + * @title DeployUpgradeTimelock + * @notice Script to deploy the UpgradeTimelock contract for managing upgrades + */ +contract DeployUpgradeTimelock is Script { + + struct TimelockConfig { + uint256 minDelay; // Minimum delay in seconds (e.g., 86400 for 24 hours) + address[] proposers; // Addresses that can propose operations + address[] executors; // Addresses that can execute operations + address admin; // Optional admin address (can be zero address) + } + + function run() public { + // Load configuration + TimelockConfig memory config = getConfig(); + + // Deploy the timelock + vm.startBroadcast(); + UpgradeTimelock timelock = deployTimelock(config); + vm.stopBroadcast(); + + // Log deployment details + console.log("UpgradeTimelock deployed at:", address(timelock)); + console.log("Minimum delay:", config.minDelay, "seconds"); + console.log("Number of proposers:", config.proposers.length); + console.log("Number of executors:", config.executors.length); + + if (config.admin != address(0)) { + console.log("Admin address:", config.admin); + } + } + + /** + * @notice Deploy the UpgradeTimelock contract + * @param config The timelock configuration + * @return timelock The deployed timelock contract + */ + function deployTimelock(TimelockConfig memory config) public returns (UpgradeTimelock) { + // Deploy implementation + UpgradeTimelock implementation = new UpgradeTimelock(); + console.log("UpgradeTimelock implementation deployed at:", address(implementation)); + + // Prepare initialization data + bytes memory initData = abi.encodeCall( + UpgradeTimelock.initialize, + (config.minDelay, config.proposers, config.executors, config.admin) + ); + + // Deploy proxy + TransparentUpgradeableProxy proxy = new TransparentUpgradeableProxy( + address(implementation), + msg.sender, // Proxy admin (should be changed to timelock itself or multisig later) + initData + ); + + return UpgradeTimelock(payable(address(proxy))); + } + + /** + * @notice Get timelock configuration + * @dev Override this function or set environment variables for different deployments + * @return config The timelock configuration + */ + function getConfig() public view returns (TimelockConfig memory config) { + // Default configuration - can be overridden via environment variables or JSON config + try vm.envUint("TIMELOCK_MIN_DELAY") returns (uint256 delay) { + config.minDelay = delay; + } catch { + config.minDelay = 86400; // 24 hours default + } + + // Get proposers from environment or use deployer as default + try vm.envString("TIMELOCK_PROPOSERS") returns (string memory proposersEnv) { + if (bytes(proposersEnv).length > 0) { + // Parse comma-separated addresses from environment + config.proposers = parseAddresses(proposersEnv); + } else { + // Default: deployer can propose + config.proposers = new address[](1); + config.proposers[0] = msg.sender; + } + } catch { + // Default: deployer can propose + config.proposers = new address[](1); + config.proposers[0] = msg.sender; + } + + // Get executors from environment or use deployer as default + try vm.envString("TIMELOCK_EXECUTORS") returns (string memory executorsEnv) { + if (bytes(executorsEnv).length > 0) { + // Parse comma-separated addresses from environment + config.executors = parseAddresses(executorsEnv); + } else { + // Default: deployer can execute + config.executors = new address[](1); + config.executors[0] = msg.sender; + } + } catch { + // Default: deployer can execute + config.executors = new address[](1); + config.executors[0] = msg.sender; + } + + // Admin address (optional) + try vm.envAddress("TIMELOCK_ADMIN") returns (address adminAddr) { + config.admin = adminAddr; + } catch { + config.admin = address(0); + } + } + + /** + * @notice Parse comma-separated addresses from a string + * @param addressesStr Comma-separated address string + * @return addresses Array of parsed addresses + */ + function parseAddresses(string memory addressesStr) internal pure returns (address[] memory addresses) { + // Simple implementation - in production, use a more robust parser + // For now, assume single address (can be extended) + addresses = new address[](1); + addresses[0] = vm.parseAddress(addressesStr); + } + + /** + * @notice Alternative deployment with JSON config + * @param configPath Path to JSON configuration file + * @return timelock The deployed timelock contract + */ + function deployWithConfig(string memory configPath) public returns (UpgradeTimelock) { + string memory json = vm.readFile(configPath); + + TimelockConfig memory config; + config.minDelay = vm.parseJsonUint(json, ".minDelay"); + + // For simplicity, assume single proposer and executor in JSON + // In production, extend to support arrays + config.proposers = new address[](1); + config.proposers[0] = vm.parseJsonAddress(json, ".proposer"); + + config.executors = new address[](1); + config.executors[0] = vm.parseJsonAddress(json, ".executor"); + + config.admin = vm.parseJsonAddress(json, ".admin"); + + return deployTimelock(config); + } +} \ No newline at end of file diff --git a/script/upgrades/TimelockUpgrade.s.sol b/script/upgrades/TimelockUpgrade.s.sol new file mode 100644 index 00000000..83e775d8 --- /dev/null +++ b/script/upgrades/TimelockUpgrade.s.sol @@ -0,0 +1,270 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.25; + +import {Script} from "forge-std/Script.sol"; +import {console} from "forge-std/console.sol"; +import {UpgradeTimelock} from "../../src/UpgradeTimelock.sol"; +import {ITransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol"; + +/** + * @title TimelockUpgrade + * @notice Script for scheduling and executing upgrades through the timelock controller + */ +contract TimelockUpgrade is Script { + + struct UpgradeParams { + address timelock; // Address of the UpgradeTimelock contract + address proxy; // Address of the proxy to upgrade + address newImplementation; // Address of the new implementation + bytes upgradeData; // Optional data for upgradeToAndCall + } + + /** + * @notice Schedule an upgrade operation + * @dev This function schedules an upgrade that must wait for the timelock delay + */ + function scheduleUpgrade() public { + UpgradeParams memory params = getUpgradeParams(); + + vm.startBroadcast(); + + UpgradeTimelock timelock = UpgradeTimelock(payable(params.timelock)); + + // Schedule the upgrade + // Prepare upgrade calldata + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", params.newImplementation); + bytes32 salt = timelock.getUpgradeSalt(params.proxy, params.newImplementation); + bytes32 predecessor = bytes32(0); + uint256 delay = timelock.getMinDelay(); + + // Schedule the upgrade + timelock.schedule( + params.proxy, + 0, + upgradeCalldata, + predecessor, + salt, + delay + ); + + bytes32 operationId = timelock.getUpgradeOperationId(params.proxy, params.newImplementation); + + vm.stopBroadcast(); + + console.log("Upgrade scheduled!"); + console.log("Operation ID:", vm.toString(operationId)); + console.log("Proxy:", params.proxy); + console.log("New Implementation:", params.newImplementation); + console.log("Timelock:", params.timelock); + + uint256 readyTimestamp = timelock.getUpgradeTimestamp( + params.proxy, + params.newImplementation + ); + + console.log("Ready for execution at timestamp:", readyTimestamp); + console.log("Ready for execution at (human readable):", timestampToString(readyTimestamp)); + } + + /** + * @notice Execute a previously scheduled upgrade + * @dev This function executes an upgrade that has passed the timelock delay + */ + function executeUpgrade() public { + UpgradeParams memory params = getUpgradeParams(); + + UpgradeTimelock timelock = UpgradeTimelock(payable(params.timelock)); + + // Check if upgrade is ready + bool ready = timelock.isUpgradeReady(params.proxy, params.newImplementation); + require(ready, "Upgrade is not ready for execution yet"); + + // Prepare upgrade parameters + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", params.newImplementation); + bytes32 salt = timelock.getUpgradeSalt(params.proxy, params.newImplementation); + bytes32 predecessor = bytes32(0); + + vm.startBroadcast(); + + // Execute the upgrade + timelock.execute( + params.proxy, + 0, + upgradeCalldata, + predecessor, + salt + ); + + vm.stopBroadcast(); + + console.log("Upgrade executed successfully!"); + console.log("Proxy:", params.proxy); + console.log("New Implementation:", params.newImplementation); + } + + /** + * @notice Check the status of an upgrade operation + */ + function checkUpgradeStatus() public view { + UpgradeParams memory params = getUpgradeParams(); + + UpgradeTimelock timelock = UpgradeTimelock(payable(params.timelock)); + + bool ready = timelock.isUpgradeReady(params.proxy, params.newImplementation); + uint256 readyTimestamp = timelock.getUpgradeTimestamp(params.proxy, params.newImplementation); + + console.log("=== Upgrade Status ==="); + console.log("Proxy:", params.proxy); + console.log("New Implementation:", params.newImplementation); + console.log("Timelock:", params.timelock); + console.log("Ready for execution:", ready); + console.log("Ready timestamp:", readyTimestamp); + + if (readyTimestamp > 0) { + console.log("Ready at (human readable):", timestampToString(readyTimestamp)); + + if (block.timestamp < readyTimestamp) { + uint256 timeLeft = readyTimestamp - block.timestamp; + console.log("Time remaining:", timeLeft, "seconds"); + console.log("Time remaining (human readable):", secondsToString(timeLeft)); + } + } else { + console.log("Operation not scheduled"); + } + } + + /** + * @notice Get upgrade parameters from environment variables or override this function + * @return params The upgrade parameters + */ + function getUpgradeParams() public view returns (UpgradeParams memory params) { + params.timelock = vm.envAddress("UPGRADE_TIMELOCK"); + params.proxy = vm.envAddress("UPGRADE_PROXY"); + params.newImplementation = vm.envAddress("UPGRADE_NEW_IMPLEMENTATION"); + + // Optional upgrade data + try vm.envString("UPGRADE_DATA") returns (string memory upgradeDataHex) { + if (bytes(upgradeDataHex).length > 0) { + params.upgradeData = vm.parseBytes(upgradeDataHex); + } else { + params.upgradeData = ""; + } + } catch { + params.upgradeData = ""; + } + } + + /** + * @notice Convert timestamp to human readable string + * @param timestamp The timestamp to convert + * @return Human readable date/time string (simplified) + */ + function timestampToString(uint256 timestamp) internal pure returns (string memory) { + if (timestamp == 0) return "Not scheduled"; + return string(abi.encodePacked("Timestamp: ", vm.toString(timestamp))); + } + + /** + * @notice Convert seconds to human readable duration + * @param secondsAmount The number of seconds + * @return Human readable duration string + */ + function secondsToString(uint256 secondsAmount) internal pure returns (string memory) { + if (secondsAmount < 60) { + return string(abi.encodePacked(vm.toString(secondsAmount), " seconds")); + } else if (secondsAmount < 3600) { + uint256 minutesAmount = secondsAmount / 60; + uint256 remainingSecs = secondsAmount % 60; + return string(abi.encodePacked( + vm.toString(minutesAmount), " minutes, ", + vm.toString(remainingSecs), " seconds" + )); + } else if (secondsAmount < 86400) { + uint256 hoursAmount = secondsAmount / 3600; + uint256 remainingMins = (secondsAmount % 3600) / 60; + return string(abi.encodePacked( + vm.toString(hoursAmount), " hours, ", + vm.toString(remainingMins), " minutes" + )); + } else { + uint256 daysAmount = secondsAmount / 86400; + uint256 remainingHrs = (secondsAmount % 86400) / 3600; + return string(abi.encodePacked( + vm.toString(daysAmount), " days, ", + vm.toString(remainingHrs), " hours" + )); + } + } + + /** + * @notice Helper function for YieldDistributor upgrades + */ + function scheduleYieldDistributorUpgrade( + address timelock, + address yieldDistributorProxy, + address newImplementation + ) public { + vm.startBroadcast(); + + UpgradeTimelock timelockContract = UpgradeTimelock(payable(timelock)); + + // Prepare upgrade parameters + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", newImplementation); + bytes32 salt = timelockContract.getUpgradeSalt(yieldDistributorProxy, newImplementation); + bytes32 predecessor = bytes32(0); + uint256 delay = timelockContract.getMinDelay(); + + // Schedule the upgrade + timelockContract.schedule( + yieldDistributorProxy, + 0, + upgradeCalldata, + predecessor, + salt, + delay + ); + + bytes32 operationId = timelockContract.getUpgradeOperationId(yieldDistributorProxy, newImplementation); + + vm.stopBroadcast(); + + console.log("YieldDistributor upgrade scheduled!"); + console.log("Operation ID:", vm.toString(operationId)); + } + + /** + * @notice Helper function for ButteredBread upgrades + */ + function scheduleButteredBreadUpgrade( + address timelock, + address butteredBreadProxy, + address newImplementation + ) public { + vm.startBroadcast(); + + UpgradeTimelock timelockContract = UpgradeTimelock(payable(timelock)); + + // Prepare upgrade parameters + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", newImplementation); + bytes32 salt = timelockContract.getUpgradeSalt(butteredBreadProxy, newImplementation); + bytes32 predecessor = bytes32(0); + uint256 delay = timelockContract.getMinDelay(); + + // Schedule the upgrade + timelockContract.schedule( + butteredBreadProxy, + 0, + upgradeCalldata, + predecessor, + salt, + delay + ); + + bytes32 operationId = timelockContract.getUpgradeOperationId(butteredBreadProxy, newImplementation); + + vm.stopBroadcast(); + + console.log("ButteredBread upgrade scheduled!"); + console.log("Operation ID:", vm.toString(operationId)); + } +} \ No newline at end of file diff --git a/src/UpgradeTimelock.sol b/src/UpgradeTimelock.sol new file mode 100644 index 00000000..372e6396 --- /dev/null +++ b/src/UpgradeTimelock.sol @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.25; + +import {TimelockControllerUpgradeable} from "@openzeppelin/contracts-upgradeable/governance/TimelockControllerUpgradeable.sol"; + +/** + * @title UpgradeTimelock + * @notice Timelock controller specifically designed for managing upgrades of transparent proxies. + * This contract enforces a time delay between proposal and execution of proxy upgrades. + * + * @dev This contract is a wrapper around OpenZeppelin's TimelockControllerUpgradeable, + * providing convenience functions for proxy upgrades while maintaining full compatibility + * with the standard timelock interface. + */ +contract UpgradeTimelock is TimelockControllerUpgradeable { + + /** + * @notice Initialize the timelock controller + * @param minDelay The minimum delay for operations + * @param proposers Array of addresses that can propose operations + * @param executors Array of addresses that can execute operations + * @param admin Optional account to be granted admin role (can be zero address) + */ + function initialize( + uint256 minDelay, + address[] memory proposers, + address[] memory executors, + address admin + ) external initializer { + __TimelockController_init(minDelay, proposers, executors, admin); + } + + /** + * @notice Get a deterministic salt for upgrade operations + * @param proxy The proxy address + * @param newImplementation The new implementation address + * @return salt The deterministic salt + */ + function getUpgradeSalt(address proxy, address newImplementation) external pure returns (bytes32 salt) { + return keccak256(abi.encodePacked("UPGRADE", proxy, newImplementation)); + } + + /** + * @notice Helper to get the operation ID for a proxy upgrade + * @param proxy The proxy address + * @param newImplementation The new implementation address + * @return operationId The operation identifier + */ + function getUpgradeOperationId(address proxy, address newImplementation) external pure returns (bytes32 operationId) { + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", newImplementation); + bytes32 salt = keccak256(abi.encodePacked("UPGRADE", proxy, newImplementation)); + // Manual hash calculation to avoid type conversion issues + return keccak256(abi.encode(proxy, 0, keccak256(upgradeCalldata), bytes32(0), salt)); + } + + /** + * @notice Check if a proxy upgrade is ready for execution + * @param proxy The proxy address + * @param newImplementation The new implementation address + * @return ready True if the upgrade can be executed + */ + function isUpgradeReady(address proxy, address newImplementation) external view returns (bool ready) { + bytes32 operationId = this.getUpgradeOperationId(proxy, newImplementation); + return isOperationReady(operationId); + } + + /** + * @notice Get the timestamp when an upgrade will be ready + * @param proxy The proxy address + * @param newImplementation The new implementation address + * @return timestamp The ready timestamp + */ + function getUpgradeTimestamp(address proxy, address newImplementation) external view returns (uint256 timestamp) { + bytes32 operationId = this.getUpgradeOperationId(proxy, newImplementation); + return getTimestamp(operationId); + } +} \ No newline at end of file diff --git a/test/UpgradeTimelock.t.sol b/test/UpgradeTimelock.t.sol new file mode 100644 index 00000000..63d57d1c --- /dev/null +++ b/test/UpgradeTimelock.t.sol @@ -0,0 +1,322 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.25; + +import {Test} from "forge-std/Test.sol"; +import {console} from "forge-std/console.sol"; +import {TransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol"; +import {UpgradeTimelock} from "../src/UpgradeTimelock.sol"; +import {YieldDistributor} from "../src/YieldDistributor.sol"; + +/** + * @title UpgradeTimelockTest + * @notice Tests for the UpgradeTimelock contract functionality + */ +contract UpgradeTimelockTest is Test { + UpgradeTimelock public timelock; + TransparentUpgradeableProxy public timelockProxy; + + // Test proxy and implementations + TransparentUpgradeableProxy public testProxy; + YieldDistributor public implementationV1; + YieldDistributor public implementationV2; + + // Test accounts + address public admin = makeAddr("admin"); + address public proposer = makeAddr("proposer"); + address public executor = makeAddr("executor"); + address public unauthorized = makeAddr("unauthorized"); + + // Timelock configuration + uint256 public constant MIN_DELAY = 1 days; + + function setUp() public { + // Deploy timelock implementation + UpgradeTimelock timelockImpl = new UpgradeTimelock(); + + // Setup roles + address[] memory proposers = new address[](1); + proposers[0] = proposer; + + address[] memory executors = new address[](1); + executors[0] = executor; + + // Initialize timelock via proxy + bytes memory initData = abi.encodeCall( + UpgradeTimelock.initialize, + (MIN_DELAY, proposers, executors, admin) + ); + + timelockProxy = new TransparentUpgradeableProxy( + address(timelockImpl), + admin, + initData + ); + + timelock = UpgradeTimelock(payable(address(timelockProxy))); + + // Deploy test implementations + implementationV1 = new YieldDistributor(); + implementationV2 = new YieldDistributor(); + + // Deploy a test proxy that will be managed by the timelock + testProxy = new TransparentUpgradeableProxy( + address(implementationV1), + address(timelock), // Timelock is the proxy admin + "" + ); + } + + function test_Initialize() public view { + assertEq(timelock.getMinDelay(), MIN_DELAY); + assertTrue(timelock.hasRole(timelock.PROPOSER_ROLE(), proposer)); + assertTrue(timelock.hasRole(timelock.EXECUTOR_ROLE(), executor)); + assertTrue(timelock.hasRole(timelock.DEFAULT_ADMIN_ROLE(), admin)); + } + + function test_GetUpgradeSalt() public view { + bytes32 salt = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + assertTrue(salt != bytes32(0)); + + // Salt should be deterministic + bytes32 salt2 = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + assertEq(salt, salt2); + } + + function test_GetUpgradeOperationId() public view { + bytes32 operationId = timelock.getUpgradeOperationId(address(testProxy), address(implementationV2)); + assertTrue(operationId != bytes32(0)); + } + + function test_ScheduleUpgrade() public { + // Prepare upgrade parameters + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", address(implementationV2)); + bytes32 salt = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + bytes32 predecessor = bytes32(0); + uint256 delay = timelock.getMinDelay(); + + // Schedule the upgrade + vm.prank(proposer); + timelock.schedule( + address(testProxy), + 0, + upgradeCalldata, + predecessor, + salt, + delay + ); + + // Verify the operation was scheduled + bytes32 operationId = timelock.getUpgradeOperationId(address(testProxy), address(implementationV2)); + assertTrue(timelock.isOperation(operationId)); + assertFalse(timelock.isUpgradeReady(address(testProxy), address(implementationV2))); + } + + function test_ExecuteUpgrade_AfterDelay() public { + // First schedule the upgrade + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", address(implementationV2)); + bytes32 salt = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + bytes32 predecessor = bytes32(0); + uint256 delay = timelock.getMinDelay(); + + vm.prank(proposer); + timelock.schedule( + address(testProxy), + 0, + upgradeCalldata, + predecessor, + salt, + delay + ); + + // Fast forward time + vm.warp(block.timestamp + MIN_DELAY + 1); + + // Verify upgrade is ready + assertTrue(timelock.isUpgradeReady(address(testProxy), address(implementationV2))); + + // Execute upgrade + vm.prank(executor); + timelock.execute( + address(testProxy), + 0, + upgradeCalldata, + predecessor, + salt + ); + + // Verify operation is marked as done + bytes32 operationId = timelock.getUpgradeOperationId(address(testProxy), address(implementationV2)); + assertTrue(timelock.isOperationDone(operationId)); + } + + function test_ExecuteUpgrade_BeforeDelayFails() public { + // Schedule upgrade + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", address(implementationV2)); + bytes32 salt = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + bytes32 predecessor = bytes32(0); + uint256 delay = timelock.getMinDelay(); + + vm.prank(proposer); + timelock.schedule( + address(testProxy), + 0, + upgradeCalldata, + predecessor, + salt, + delay + ); + + // Try to execute immediately (should fail) + vm.prank(executor); + vm.expectRevert(); + timelock.execute( + address(testProxy), + 0, + upgradeCalldata, + predecessor, + salt + ); + } + + function test_UnauthorizedScheduleFails() public { + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", address(implementationV2)); + bytes32 salt = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + + vm.prank(unauthorized); + vm.expectRevert(); + timelock.schedule( + address(testProxy), + 0, + upgradeCalldata, + bytes32(0), + salt, + timelock.getMinDelay() + ); + } + + function test_UnauthorizedExecuteFails() public { + // Schedule upgrade first + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", address(implementationV2)); + bytes32 salt = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + bytes32 predecessor = bytes32(0); + uint256 delay = timelock.getMinDelay(); + + vm.prank(proposer); + timelock.schedule( + address(testProxy), + 0, + upgradeCalldata, + predecessor, + salt, + delay + ); + + // Fast forward time + vm.warp(block.timestamp + MIN_DELAY + 1); + + // Try to execute with unauthorized account + vm.prank(unauthorized); + vm.expectRevert(); + timelock.execute( + address(testProxy), + 0, + upgradeCalldata, + predecessor, + salt + ); + } + + function test_GetUpgradeTimestamp() public { + uint256 scheduledTime = block.timestamp; + + // Schedule upgrade + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", address(implementationV2)); + bytes32 salt = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + + vm.prank(proposer); + timelock.schedule( + address(testProxy), + 0, + upgradeCalldata, + bytes32(0), + salt, + timelock.getMinDelay() + ); + + uint256 readyTimestamp = timelock.getUpgradeTimestamp( + address(testProxy), + address(implementationV2) + ); + + assertEq(readyTimestamp, scheduledTime + MIN_DELAY); + } + + function test_CancelOperation() public { + // Schedule upgrade + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", address(implementationV2)); + bytes32 salt = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + bytes32 predecessor = bytes32(0); + uint256 delay = timelock.getMinDelay(); + + vm.prank(proposer); + timelock.schedule( + address(testProxy), + 0, + upgradeCalldata, + predecessor, + salt, + delay + ); + + bytes32 operationId = timelock.getUpgradeOperationId(address(testProxy), address(implementationV2)); + + // Admin can cancel the operation + vm.prank(admin); + timelock.cancel(operationId); + + // Fast forward time + vm.warp(block.timestamp + MIN_DELAY + 1); + + // Should not be ready anymore + assertFalse(timelock.isUpgradeReady(address(testProxy), address(implementationV2))); + + // Execution should fail + vm.prank(executor); + vm.expectRevert(); + timelock.execute( + address(testProxy), + 0, + upgradeCalldata, + predecessor, + salt + ); + } + + function test_RoleManagement() public { + address newProposer = makeAddr("newProposer"); + + // Admin can grant roles + vm.prank(admin); + timelock.grantRole(timelock.PROPOSER_ROLE(), newProposer); + + assertTrue(timelock.hasRole(timelock.PROPOSER_ROLE(), newProposer)); + + // New proposer can schedule upgrades + bytes memory upgradeCalldata = abi.encodeWithSignature("upgradeTo(address)", address(implementationV2)); + bytes32 salt = timelock.getUpgradeSalt(address(testProxy), address(implementationV2)); + + vm.prank(newProposer); + timelock.schedule( + address(testProxy), + 0, + upgradeCalldata, + bytes32(0), + salt, + timelock.getMinDelay() + ); + + // Verify operation was scheduled + bytes32 operationId = timelock.getUpgradeOperationId(address(testProxy), address(implementationV2)); + assertTrue(timelock.isOperation(operationId)); + } +} \ No newline at end of file