feat: adopt announcePlannedUpgrade for PDPVerifier

src/PDPVerifier.sol

@@ -147,6 +147,16 @@ contract PDPVerifier is Initializable, UUPSUpgradeable, OwnableUpgradeable {
 
     FeeStatus private feeStatus;
 
+    // Used for announcing upgrades, packed into one slot
+    struct PlannedUpgrade {
+        // Address of the new implementation contract
+        address nextImplementation;
+        // Upgrade will not occur until at least this epoch
+        uint96 afterEpoch;
+    }
+
+    PlannedUpgrade public nextUpgrade;
+
     // Methods
 
     /// @custom:oz-upgrades-unsafe-allow constructor
@@ -165,12 +175,25 @@ contract PDPVerifier is Initializable, UUPSUpgradeable, OwnableUpgradeable {
     string public constant VERSION = "3.1.0";
 
     event ContractUpgraded(string version, address implementation);
+    event UpgradeAnnounced(PlannedUpgrade plannedUpgrade);
 
-    function migrate() external onlyOwner reinitializer(2) {
+    function migrate() external onlyProxy onlyOwner reinitializer(2) {
         emit ContractUpgraded(VERSION, ERC1967Utils.getImplementation());
     }
 
-    function _authorizeUpgrade(address newImplementation) internal override onlyOwner {}
+    function announcePlannedUpgrade(PlannedUpgrade calldata plannedUpgrade) external onlyOwner {
+        require(plannedUpgrade.nextImplementation.code.length > 3000);
+        require(plannedUpgrade.afterEpoch > block.number);
+        nextUpgrade = plannedUpgrade;
+        emit UpgradeAnnounced(plannedUpgrade);
+    }
+
+    function _authorizeUpgrade(address newImplementation) internal override onlyOwner {
+        // zero address already checked by ERC1967Utils._setImplementation
+        require(newImplementation == nextUpgrade.nextImplementation);
+        require(block.number >= nextUpgrade.afterEpoch);
+        delete nextUpgrade;
+    }
 
     function burnFee(uint256 amount) internal {
         require(msg.value >= amount, "Incorrect fee amount");

test/ERC1967Proxy.t.sol

@@ -45,6 +45,15 @@ contract ERC1967ProxyTest is Test {
         // Deploy new implementation
         PDPVerifier newImplementation = new PDPVerifier();
 
+        // Announce upgrade first (required by new upgrade pattern)
+        PDPVerifier.PlannedUpgrade memory plan;
+        plan.nextImplementation = address(newImplementation);
+        plan.afterEpoch = uint96(vm.getBlockNumber()) + 1;
+        proxy.announcePlannedUpgrade(plan);
+
+        // Roll to afterEpoch
+        vm.roll(plan.afterEpoch);
+
         // Upgrade proxy to new implementation
         proxy.upgradeToAndCall(address(newImplementation), "");
 
@@ -57,6 +66,15 @@ contract ERC1967ProxyTest is Test {
     function testUpgradeFromNonOwnerNoGood() public {
         PDPVerifier newImplementation = new PDPVerifier();
 
+        // Announce upgrade first (as owner)
+        PDPVerifier.PlannedUpgrade memory plan;
+        plan.nextImplementation = address(newImplementation);
+        plan.afterEpoch = uint96(vm.getBlockNumber()) + 1;
+        proxy.announcePlannedUpgrade(plan);
+
+        vm.roll(plan.afterEpoch);
+
+        // Try to upgrade as non-owner
         vm.stopPrank();
         vm.startPrank(address(0xdead));
 

test/PDPVerifier.t.sol

@@ -1911,19 +1911,114 @@ contract PDPVerifierMigrateTest is Test {
     PDPVerifier implementation;
     PDPVerifier newImplementation;
     MyERC1967Proxy proxy;
+    PDPVerifier pdpVerifier;
 
     function setUp() public {
         bytes memory initializeData = abi.encodeWithSelector(PDPVerifier.initialize.selector, 2);
         implementation = new PDPVerifier();
         newImplementation = new PDPVerifier();
         proxy = new MyERC1967Proxy(address(implementation), initializeData);
+        pdpVerifier = PDPVerifier(address(proxy));
+    }
+
+    function testAnnouncePlannedUpgrade() public {
+        // Initially, no upgrade is planned
+        (address nextImplementation, uint96 afterEpoch) = pdpVerifier.nextUpgrade();
+        assertEq(nextImplementation, address(0));
+        assertEq(afterEpoch, uint96(0));
+
+        // Deploy new implementation
+        PDPVerifier newImpl = new PDPVerifier();
+
+        // Announce upgrade
+        PDPVerifier.PlannedUpgrade memory plan;
+        plan.nextImplementation = address(newImpl);
+        plan.afterEpoch = uint96(vm.getBlockNumber()) + 2000;
+
+        vm.expectEmit(false, false, false, true);
+        emit PDPVerifier.UpgradeAnnounced(plan);
+        pdpVerifier.announcePlannedUpgrade(plan);
+
+        // Verify upgrade plan is stored
+        (nextImplementation, afterEpoch) = pdpVerifier.nextUpgrade();
+        assertEq(nextImplementation, plan.nextImplementation);
+        assertEq(afterEpoch, plan.afterEpoch);
+
+        // Cannot upgrade before afterEpoch
+        bytes memory migrateData = abi.encodeWithSelector(PDPVerifier.migrate.selector);
+        vm.expectRevert();
+        pdpVerifier.upgradeToAndCall(plan.nextImplementation, migrateData);
+
+        // Still cannot upgrade at afterEpoch - 1
+        vm.roll(plan.afterEpoch - 1);
+        vm.expectRevert();
+        pdpVerifier.upgradeToAndCall(plan.nextImplementation, migrateData);
+
+        // Can upgrade at afterEpoch
+        vm.roll(plan.afterEpoch);
+        vm.expectEmit(false, false, false, true);
+        emit PDPVerifier.ContractUpgraded(newImpl.VERSION(), plan.nextImplementation);
+        pdpVerifier.upgradeToAndCall(plan.nextImplementation, migrateData);
+
+        // After upgrade, nextUpgrade should be cleared
+        (nextImplementation, afterEpoch) = pdpVerifier.nextUpgrade();
+        assertEq(nextImplementation, address(0));
+        assertEq(afterEpoch, uint96(0));
+    }
+
+    function testAnnouncePlannedUpgradeOnlyOwner() public {
+        PDPVerifier newImpl = new PDPVerifier();
+        PDPVerifier.PlannedUpgrade memory plan;
+        plan.nextImplementation = address(newImpl);
+        plan.afterEpoch = uint96(vm.getBlockNumber()) + 2000;
+
+        // Non-owner cannot announce upgrade
+        vm.prank(address(0x1234));
+        vm.expectRevert();
+        pdpVerifier.announcePlannedUpgrade(plan);
+    }
+
+    function testAnnouncePlannedUpgradeInvalidImplementation() public {
+        PDPVerifier.PlannedUpgrade memory plan;
+        plan.nextImplementation = address(0x123); // Invalid address with no code
+        plan.afterEpoch = uint96(vm.getBlockNumber()) + 2000;
+
+        vm.expectRevert();
+        pdpVerifier.announcePlannedUpgrade(plan);
+    }
+
+    function testAnnouncePlannedUpgradeInvalidEpoch() public {
+        PDPVerifier newImpl = new PDPVerifier();
+        PDPVerifier.PlannedUpgrade memory plan;
+        plan.nextImplementation = address(newImpl);
+        plan.afterEpoch = uint96(vm.getBlockNumber()); // Must be in the future
+
+        vm.expectRevert();
+        pdpVerifier.announcePlannedUpgrade(plan);
     }
 
     function testMigrate() public {
+        // Announce upgrade first (required by new upgrade pattern)
+        PDPVerifier.PlannedUpgrade memory plan;
+        plan.nextImplementation = address(newImplementation);
+        plan.afterEpoch = uint96(vm.getBlockNumber()) + 1;
+        pdpVerifier.announcePlannedUpgrade(plan);
+
+        // Roll to afterEpoch
+        vm.roll(plan.afterEpoch);
+
         vm.expectEmit(true, true, true, true);
         emit IPDPEvents.ContractUpgraded(newImplementation.VERSION(), address(newImplementation));
         bytes memory migrationCall = abi.encodeWithSelector(PDPVerifier.migrate.selector);
         UUPSUpgradeable(address(proxy)).upgradeToAndCall(address(newImplementation), migrationCall);
+
+        // Announce a second upgrade to test reinitializer behavior
+        PDPVerifier.PlannedUpgrade memory plan2;
+        plan2.nextImplementation = address(newImplementation);
+        plan2.afterEpoch = uint96(vm.getBlockNumber()) + 1;
+        pdpVerifier.announcePlannedUpgrade(plan2);
+        vm.roll(plan2.afterEpoch);
+
         // Second call should fail because reinitializer(2) can only be called once
         vm.expectRevert("InvalidInitialization()");
         UUPSUpgradeable(address(proxy)).upgradeToAndCall(address(newImplementation), migrationCall);

tools/announce-planned-upgrade.sh

@@ -0,0 +1,76 @@
+#!/bin/bash
+
+# announce-planned-upgrade.sh: Announces a planned upgrade for PDPVerifier
+# Required args: RPC_URL, PDP_VERIFIER_PROXY_ADDRESS, KEYSTORE, PASSWORD, NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS, AFTER_EPOCH
+
+if [ -z "$RPC_URL" ]; then
+  echo "Error: RPC_URL is not set"
+  exit 1
+fi
+
+if [ -z "$KEYSTORE" ]; then
+  echo "Error: KEYSTORE is not set"
+  exit 1
+fi
+
+if [ -z "$PASSWORD" ]; then
+  echo "Error: PASSWORD is not set"
+  exit 1
+fi
+
+if [ -z "$CHAIN" ]; then
+  CHAIN=$(cast chain-id --rpc-url "$RPC_URL")
+  if [ -z "$CHAIN" ]; then
+    echo "Error: Failed to detect chain ID from RPC"
+    exit 1
+  fi
+fi
+
+if [ -z "$NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS" ]; then
+  echo "NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS is not set"
+  exit 1
+fi
+
+if [ -z "$AFTER_EPOCH" ]; then
+  echo "AFTER_EPOCH is not set"
+  exit 1
+fi
+
+CURRENT_EPOCH=$(cast block-number --rpc-url "$RPC_URL" 2>/dev/null)
+
+if [ "$CURRENT_EPOCH" -gt "$AFTER_EPOCH" ]; then
+  echo "Already past AFTER_EPOCH ($CURRENT_EPOCH > $AFTER_EPOCH)"
+  exit 1
+else
+  echo "Announcing planned upgrade after $(($AFTER_EPOCH - $CURRENT_EPOCH)) epochs"
+fi
+
+
+ADDR=$(cast wallet address --keystore "$KEYSTORE" --password "$PASSWORD")
+echo "Sending announcement from owner address: $ADDR"
+
+# Get current nonce
+NONCE=$(cast nonce --rpc-url "$RPC_URL" "$ADDR")
+
+if [ -z "$PDP_VERIFIER_PROXY_ADDRESS" ]; then
+  echo "Error: PDP_VERIFIER_PROXY_ADDRESS is not set"
+  exit 1
+fi
+
+PROXY_OWNER=$(cast call --rpc-url "$RPC_URL" -f 0x0000000000000000000000000000000000000000 "$PDP_VERIFIER_PROXY_ADDRESS" "owner()(address)" 2>/dev/null)
+if [ "$PROXY_OWNER" != "$ADDR" ]; then
+  echo "Supplied KEYSTORE ($ADDR) is not the proxy owner ($PROXY_OWNER)."
+  exit 1
+fi
+
+TX_HASH=$(cast send --rpc-url "$RPC_URL" --keystore "$KEYSTORE" --password "$PASSWORD" "$PDP_VERIFIER_PROXY_ADDRESS" "announcePlannedUpgrade((address,uint96))" "($NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS,$AFTER_EPOCH)" \
+  --nonce "$NONCE" \
+  --json | jq -r '.transactionHash')
+
+if [ -z "$TX_HASH" ]; then
+  echo "Error: Failed to send announcePlannedUpgrade transaction"
+  exit 1
+fi
+
+echo "announcePlannedUpgrade transaction sent: $TX_HASH"
+