From 3c6a27cd718c8aa444be67cbbacd62bd1844e7ca Mon Sep 17 00:00:00 2001
From: Jakub Nowakowski <jakub.nowakowski@thesis.co>
Date: Fri, 9 Aug 2024 11:18:05 +0200
Subject: [PATCH] Add function to prepare proxy upgrade

Proxy upgrades of contracts are performed in two stages:
1. Deploy new implementation contract.
2. Upgrade version of implementation in the Proxy contract.

For upgradable contracts that were already deployed and ownership transferred
to the governance we cannot simply run `upgradeProxy` as the governance
is a multisig.

Here we introduce a solution that will help us execute the upgrade in
setup used across our projects:
1. Deploy new implementation contract.
2. Prepare transaction for the Governance to execute.

The solution is based on Open Zeppelin's [upgradeProxy](https://github.com/OpenZeppelin/openzeppelin-upgrades/blob/49e7ae93ee9be1d6f586517890a83634dea29ebc/packages/plugin-hardhat/src/upgrade-proxy.ts)
function, with the difference that the upgrade implementation transaction
is prepared but not executed.

Implementation upgrade is executed through ProxyAdmin contract which can
exist in two versions: V4 and V5. Currently our new contracts are deployed
with version V5, but for older deployments we still support V5.
The difference between these two version is in the upgrade function that
is called on the ProxyAdmin.
In V4 there were two separate functions `upgradeAndCall` and `upgrade` which
were called depending if the callback should be executed.
In V5 there is just `upgradeAndCall` function, which allows empty calldata.
This is based on https://github.com/OpenZeppelin/openzeppelin-upgrades/blob/49e7ae93ee9be1d6f586517890a83634dea29ebc/packages/plugin-hardhat/src/upgrade-proxy.ts#L62C45-L103
---
 src/upgrades.ts | 138 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 138 insertions(+)

diff --git a/src/upgrades.ts b/src/upgrades.ts
index af1ee24..22edb53 100644
--- a/src/upgrades.ts
+++ b/src/upgrades.ts
@@ -3,6 +3,7 @@ import "@openzeppelin/hardhat-upgrades"
 import type {
   Contract,
   ContractFactory,
+  ContractTransaction,
   ContractTransactionResponse,
 } from "ethers"
 import type {
@@ -16,6 +17,12 @@ import type {
   UpgradeProxyOptions,
 } from "@openzeppelin/hardhat-upgrades/src/utils/options"
 import { Libraries } from "hardhat-deploy/types"
+import {
+  attachProxyAdminV4,
+  attachProxyAdminV5,
+} from "@openzeppelin/hardhat-upgrades/dist/utils"
+
+import { getUpgradeInterfaceVersion } from "@openzeppelin/upgrades-core"
 
 export interface HardhatUpgradesHelpers {
   deployProxy<T extends Contract>(
@@ -27,6 +34,14 @@ export interface HardhatUpgradesHelpers {
     newContractName: string,
     opts?: UpgradesUpgradeOptions
   ): Promise<[T, Deployment]>
+  prepareProxyUpgrade(
+    proxyDeploymentName: string,
+    newContractName: string,
+    opts?: UpgradesPrepareProxyUpgradeOptions
+  ): Promise<{
+    newImplementationAddress: string
+    preparedTransaction: ContractTransaction
+  }>
 }
 
 type CustomFactoryOptions = FactoryOptions & {
@@ -47,6 +62,12 @@ export interface UpgradesUpgradeOptions {
   proxyOpts?: UpgradeProxyOptions
 }
 
+export interface UpgradesPrepareProxyUpgradeOptions {
+  contractName?: string
+  factoryOpts?: CustomFactoryOptions
+  callData?: string
+}
+
 /**
  * Deploys contract as a TransparentProxy.
  *
@@ -206,6 +227,118 @@ async function upgradeProxy<T extends Contract>(
   return [newContractInstance, deployment]
 }
 
+/**
+ * Prepare upgrade of deployed contract.
+ * It deploys new implementation contract and prepares transaction to upgrade
+ * the proxy contract to the new implementation thorough a Proxy Admin instance.
+ * The transaction has to be executed by the owner of the Proxy Admin.
+ *
+ * @param {HardhatRuntimeEnvironment} hre Hardhat runtime environment.
+ * @param {string} proxyDeploymentName Name of the proxy deployment that will be
+ *        upgraded.
+ * @param {string} newContractName Name of the new implementation contract.
+ * @param {UpgradesPrepareProxyUpgradeOptions} opts
+ */
+async function prepareProxyUpgrade(
+  hre: HardhatRuntimeEnvironment,
+  proxyDeploymentName: string,
+  newContractName: string,
+  opts?: UpgradesPrepareProxyUpgradeOptions
+): Promise<{
+  newImplementationAddress: string
+  preparedTransaction: ContractTransaction
+}> {
+  const { ethers, upgrades, deployments, artifacts } = hre
+  const signer = await ethers.provider.getSigner()
+  const { log } = deployments
+
+  const proxyDeployment: Deployment = await deployments.get(proxyDeploymentName)
+
+  const implementationContractFactory: ContractFactory =
+    await ethers.getContractFactory(
+      opts?.contractName || newContractName,
+      opts?.factoryOpts
+    )
+
+  const newImplementationAddress: string = (await upgrades.prepareUpgrade(
+    proxyDeployment.address,
+    implementationContractFactory,
+    {
+      kind: "transparent",
+      getTxResponse: false,
+    }
+  )) as string
+
+  log(`new implementation contract deployed at: ${newImplementationAddress}`)
+
+  const proxyAdminAddress = await hre.upgrades.erc1967.getAdminAddress(
+    proxyDeployment.address
+  )
+
+  let proxyAdmin: Contract
+  let upgradeTxData: string
+
+  const proxyInterfaceVersion = await getUpgradeInterfaceVersion(
+    hre.network.provider,
+    proxyAdminAddress
+  )
+
+  switch (proxyInterfaceVersion) {
+    case "5.0.0": {
+      proxyAdmin = await attachProxyAdminV5(hre, proxyAdminAddress, signer)
+
+      upgradeTxData = proxyAdmin.interface.encodeFunctionData(
+        "upgradeAndCall",
+        [
+          proxyDeployment.address,
+          newImplementationAddress,
+          opts?.callData ?? "0x",
+        ]
+      )
+      break
+    }
+    default: {
+      proxyAdmin = await attachProxyAdminV4(hre, proxyAdminAddress, signer)
+
+      if (opts?.callData) {
+        upgradeTxData = proxyAdmin.interface.encodeFunctionData(
+          "upgradeAndCall",
+          [proxyDeployment.address, newImplementationAddress, opts?.callData]
+        )
+      } else {
+        upgradeTxData = proxyAdmin.interface.encodeFunctionData("upgrade", [
+          proxyDeployment.address,
+          newImplementationAddress,
+        ])
+      }
+    }
+  }
+
+  const preparedTransaction: ContractTransaction = {
+    from: (await proxyAdmin.owner()) as string,
+    to: proxyAdminAddress,
+    data: upgradeTxData,
+  }
+
+  deployments.log(
+    `to upgrade the proxy implementation execute the following ` +
+      `transaction:\n${JSON.stringify(preparedTransaction, null, 2)}`
+  )
+
+  // Update Deployment Artifact
+  const artifact: Artifact = artifacts.readArtifactSync(
+    opts?.contractName || newContractName
+  )
+
+  await deployments.save(proxyDeploymentName, {
+    ...proxyDeployment,
+    abi: artifact.abi,
+    implementation: newImplementationAddress,
+  })
+
+  return { newImplementationAddress, preparedTransaction }
+}
+
 export default function (
   hre: HardhatRuntimeEnvironment
 ): HardhatUpgradesHelpers {
@@ -217,5 +350,10 @@ export default function (
       newContractName: string,
       opts?: UpgradesUpgradeOptions
     ) => upgradeProxy(hre, currentContractName, newContractName, opts),
+    prepareProxyUpgrade: (
+      proxyDeploymentName: string,
+      newContractName: string,
+      opts?: UpgradesPrepareProxyUpgradeOptions
+    ) => prepareProxyUpgrade(hre, proxyDeploymentName, newContractName, opts),
   }
 }