diff --git a/src/PDPVerifier.sol b/src/PDPVerifier.sol index 45357d6..160f148 100644 --- a/src/PDPVerifier.sol +++ b/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"); diff --git a/test/ERC1967Proxy.t.sol b/test/ERC1967Proxy.t.sol index be2db7b..b31b641 100644 --- a/test/ERC1967Proxy.t.sol +++ b/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)); diff --git a/test/PDPVerifier.t.sol b/test/PDPVerifier.t.sol index 7400d01..8e132c1 100644 --- a/test/PDPVerifier.t.sol +++ b/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); diff --git a/tools/announce-planned-upgrade.sh b/tools/announce-planned-upgrade.sh new file mode 100755 index 0000000..b016ccc --- /dev/null +++ b/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" + diff --git a/tools/upgrade.sh b/tools/upgrade.sh new file mode 100755 index 0000000..fc0af5c --- /dev/null +++ b/tools/upgrade.sh @@ -0,0 +1,131 @@ +#!/bin/bash + +# upgrade.sh: Completes a pending upgrade for PDPVerifier +# Required args: RPC_URL, PDP_VERIFIER_PROXY_ADDRESS, KEYSTORE, PASSWORD, NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS + +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 + +ADDR=$(cast wallet address --keystore "$KEYSTORE" --password "$PASSWORD") +echo "Using 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 + +# Get the upgrade plan (if any) +# Try to call nextUpgrade() - this will fail if the method doesn't exist (old contracts) +UPGRADE_PLAN_OUTPUT=$(cast call --rpc-url "$RPC_URL" -f 0x0000000000000000000000000000000000000000 "$PDP_VERIFIER_PROXY_ADDRESS" "nextUpgrade()(address,uint96)" 2>&1) +CAST_CALL_EXIT_CODE=$? + +ZERO_ADDRESS="0x0000000000000000000000000000000000000000" + +# Check if cast call succeeded (method exists) +if [ $CAST_CALL_EXIT_CODE -eq 0 ] && [ -n "$UPGRADE_PLAN_OUTPUT" ]; then + # Method exists - parse the result + UPGRADE_PLAN=($UPGRADE_PLAN_OUTPUT) + PLANNED_PDP_VERIFIER_IMPLEMENTATION_ADDRESS=${UPGRADE_PLAN[0]} + AFTER_EPOCH=${UPGRADE_PLAN[1]} + + # Check if there's a planned upgrade (non-zero address) + # Zero address means either no upgrade was announced or the upgrade was already completed + if [ -n "$PLANNED_PDP_VERIFIER_IMPLEMENTATION_ADDRESS" ] && [ "$PLANNED_PDP_VERIFIER_IMPLEMENTATION_ADDRESS" != "$ZERO_ADDRESS" ]; then + # New two-step mechanism: validate planned upgrade + echo "Detected planned upgrade (two-step mechanism)" + + if [ "$PLANNED_PDP_VERIFIER_IMPLEMENTATION_ADDRESS" != "$NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS" ]; then + echo "NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS ($NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS) != planned ($PLANNED_PDP_VERIFIER_IMPLEMENTATION_ADDRESS)" + exit 1 + else + echo "Upgrade plan matches ($NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS)" + fi + + CURRENT_EPOCH=$(cast block-number --rpc-url "$RPC_URL" 2>/dev/null) + + if [ "$CURRENT_EPOCH" -lt "$AFTER_EPOCH" ]; then + echo "Not time yet ($CURRENT_EPOCH < $AFTER_EPOCH)" + exit 1 + else + echo "Upgrade ready ($CURRENT_EPOCH >= $AFTER_EPOCH)" + fi + else + # Method exists but returns zero - no planned upgrade or already completed + # On new contracts, _authorizeUpgrade requires a planned upgrade, so one-step will fail + echo "No planned upgrade detected (nextUpgrade returns zero)" + echo "Error: This contract requires a planned upgrade. Please call announce-planned-upgrade.sh first." + exit 1 + fi +else + # Method doesn't exist (old contract without nextUpgrade) or call failed + echo "nextUpgrade() method not found or call failed, using one-step mechanism (direct upgrade)" + echo "WARNING: This is the legacy upgrade path. For new deployments, use announce-planned-upgrade.sh first." +fi + +MIGRATE_DATA=$(cast calldata "migrate()") + +# Call upgradeToAndCall on the proxy with migrate function +echo "Upgrading proxy and calling migrate..." +TX_HASH=$(cast send --rpc-url "$RPC_URL" --keystore "$KEYSTORE" --password "$PASSWORD" "$PDP_VERIFIER_PROXY_ADDRESS" "upgradeToAndCall(address,bytes)" "$NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS" "$MIGRATE_DATA" \ + --nonce "$NONCE" \ + --json | jq -r '.transactionHash') + +if [ -z "$TX_HASH" ]; then + echo "Error: Failed to send upgrade transaction" + echo "The transaction may have failed due to:" + echo "- Insufficient permissions (not owner)" + echo "- Proxy is paused or locked" + echo "- Implementation address is invalid" + exit 1 +fi + +echo "Upgrade transaction sent: $TX_HASH" +echo "Waiting for confirmation..." + +# Wait for transaction receipt +cast receipt --rpc-url "$RPC_URL" "$TX_HASH" --confirmations 1 > /dev/null + +# Verify the upgrade by checking the implementation address +echo "Verifying upgrade..." +NEW_IMPL=$(cast rpc --rpc-url "$RPC_URL" eth_getStorageAt "$PDP_VERIFIER_PROXY_ADDRESS" 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc latest | sed 's/"//g' | sed 's/0x000000000000000000000000/0x/') + +# Compare to lowercase +export EXPECTED_IMPL=$(echo $NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS | tr '[:upper:]' '[:lower:]') + +if [ "$NEW_IMPL" = "$EXPECTED_IMPL" ]; then + echo "✅ Upgrade successful! Proxy now points to: $NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS" +else + echo "⚠️ Warning: Could not verify upgrade. Please check manually." + echo "Expected: $NEW_PDP_VERIFIER_IMPLEMENTATION_ADDRESS" + echo "Got: $NEW_IMPL" +fi +