make unified proxy upgrade script

addresses-mainnet.json

@@ -12,7 +12,7 @@
     "LBTC": "0xecAc9C5F704e954931349Da37F60E39f515c11c1",
     "Owner": "0xAa4bc534bc7Be0E28a0686Ab6910A9B21dFdc2B1",
     "Consortium": "0x1c600B19C7A5d397BEaaf03ea3646898f1E2D02e",
-    "TimeLock": "",
+    "TimeLock": "0x055E84e7FE8955E2781010B866f10Ef6E1E77e59",
     "BTCB": "0x7130d2a12b9bcbfae4f2634d864a1ee1ce3ead9c",
     "BTCBPMM": "0xE4ff44a615dF38e37cdF475833c1d57774CC9D4A",
     "ProxyFactory": "0xA24b1EB85810F68153721007E5A57d495C117918"

scripts/cre/index.ts

@@ -0,0 +1,49 @@
+import { HardhatRuntimeEnvironment } from 'hardhat/types';
+import { verify } from '../helpers';
+
+export class CustomRuntimeEnvironment {
+    constructor(private readonly hre: HardhatRuntimeEnvironment) {}
+
+    async getTimelock(address: string) {
+        const timelock = await this.hre.ethers.getContractAt(
+            'ITimelockController',
+            address
+        );
+
+        console.log(`Timelock found at ${await timelock.getAddress()}`);
+
+        return timelock;
+    }
+
+    async getProxyAdmin(proxy: string) {
+        const proxyAdmin = await this.hre.ethers.getContractAt(
+            'IProxyAdmin',
+            await this.hre.upgrades.erc1967.getAdminAddress(proxy)
+        );
+
+        console.log(
+            `ProxyAdmin of ${proxy} found at ${await proxyAdmin.getAddress()}`
+        );
+
+        return proxyAdmin;
+    }
+
+    async deployImplementation(contract: string): Promise<string> {
+        // TODO: verify storage
+
+        console.log(`Deploying ${contract} implementation`);
+
+        const impl = await this.hre.ethers.deployContract(contract);
+        await impl.waitForDeployment();
+
+        const addr = await impl.getAddress();
+
+        console.log(
+            `New implementation of ${contract} deployed at ${addr}, verifying...`
+        );
+
+        await verify(run, addr);
+
+        return addr;
+    }
+}

scripts/deployments/upgrade.ts

@@ -1,13 +1,85 @@
 import { task } from 'hardhat/config';
-import { upgradeProxy } from '../helpers/upgrade';
+import { CustomRuntimeEnvironment } from '../cre';
+import { populateSchedule, schedule } from '../helpers';
 
 task('upgrade-proxy', 'Upgrades proxy contract')
     .addParam('proxy', 'The address of the proxy contract')
-    .addParam('contract', 'The name of the contract')
+    .addPositionalParam('contractName', 'The name of the contract')
+    .addOptionalParam('timelock', 'The address of timelock contract')
+    .addOptionalParam('timelockDelay', 'The timelock delay in seconds')
+    .addParam('calldata', 'The calldata to execute during upgrade', '0x')
+    .addFlag('populate', 'Show tx data to execute later')
     .setAction(async (taskArgs, hre) => {
-        let { proxy, contract } = taskArgs;
+        const {
+            proxy,
+            contractName,
+            calldata,
+            populate,
+            timelock: timelockArg,
+            timelockDelay: timelockDelayArg,
+        } = taskArgs;
 
-        const res = await upgradeProxy(contract, proxy, hre);
+        const cre = new CustomRuntimeEnvironment(hre);
 
-        console.log(`Implementation address is ${res.implementation}`);
+        const proxyAdmin = await cre.getProxyAdmin(proxy);
+
+        const impl = await cre.deployImplementation(contractName);
+
+        if (!populate && !timelockArg) {
+            const upgradeTx = await proxyAdmin.upgradeAndCall(
+                proxy,
+                impl,
+                calldata
+            );
+            await upgradeTx.wait(1);
+
+            console.log(
+                `Contract ${contractName} at ${proxy} successfully upgraded in ${upgradeTx.hash}`
+            );
+
+            return;
+        }
+
+        const upgradeTx = await proxyAdmin.upgradeAndCall.populateTransaction(
+            proxy,
+            impl,
+            calldata
+        );
+
+        if (populate && !timelockArg) {
+            console.log(
+                `Upgrade tx data:\n${JSON.stringify(upgradeTx, null, 2)}`
+            );
+
+            return;
+        }
+
+        const timelock = await cre.getTimelock(timelockArg);
+
+        if (populate) {
+            const timelockTx = await populateSchedule(
+                hre,
+                timelock,
+                upgradeTx,
+                timelockDelayArg
+            );
+            console.log(
+                `Upgrade via timelock tx data:\n${JSON.stringify(timelockTx, null, 2)}`
+            );
+
+            return;
+        }
+
+        const timelockTx = await schedule(
+            hre,
+            timelock,
+            upgradeTx,
+            timelockDelayArg
+        );
+
+        await timelockTx.wait(1);
+
+        console.log(
+            `Scheduled upgrade for ${contractName} contract at ${proxy} in ${timelockTx.hash}`
+        );
     });

scripts/helpers/index.ts

@@ -1,6 +1,9 @@
 import { BigNumberish, ContractTransaction } from 'ethers';
 import { BytesLike } from 'ethers/lib.commonjs/utils/data';
 import { DEFAULT_PROXY_FACTORY } from './constants';
+import { ITimelockController } from '../../typechain-types';
+import Ethers from '@typechain/ethers-v6';
+import { HardhatRuntimeEnvironment } from 'hardhat/types';
 
 type TAddressesWithNetwork = {
     [k: string]: TAddresses;
@@ -40,43 +43,48 @@ export async function verify(run: any, address: string, options: any = {}) {
     } catch (e) {
         console.error(`Verification failed: ${e}`);
     }
+
+    console.log('\n');
 }
 
 export async function schedule(
-    ethers: any,
-    {
-        timelockAddr,
-        transaction,
-        predecessor,
-        salt,
-        delay,
-    }: {
-        timelockAddr: string;
-        transaction: ContractTransaction;
-        predecessor?: BytesLike;
-        salt?: BytesLike;
-        delay?: BigNumberish;
-    }
+    hre: HardhatRuntimeEnvironment,
+    timelock: ITimelockController,
+    transaction: ContractTransaction,
+    predecessor?: BytesLike,
+    salt?: BytesLike,
+    delay?: BigNumberish
 ) {
-    const timelock = await ethers.getContractAt(
-        'ITimelockController',
-        timelockAddr
+    delay = delay || (await timelock.getMinDelay());
+
+    return await timelock.schedule(
+        transaction.to,
+        transaction.value || '0',
+        transaction.data,
+        predecessor || hre.ethers.ZeroHash,
+        salt || hre.ethers.ZeroHash,
+        delay
     );
+}
 
-    if (!delay) {
-        delay = await timelock.getMinDelay();
-    }
+export async function populateSchedule(
+    hre: HardhatRuntimeEnvironment,
+    timelock: ITimelockController,
+    transaction: ContractTransaction,
+    predecessor?: BytesLike,
+    salt?: BytesLike,
+    delay?: BigNumberish
+) {
+    delay = delay || (await timelock.getMinDelay());
 
-    const res = await timelock.schedule(
+    return await timelock.schedule.populateTransaction(
         transaction.to,
         transaction.value || '0',
         transaction.data,
-        predecessor || ethers.ZeroHash,
-        salt || ethers.ZeroHash,
+        predecessor || hre.ethers.ZeroHash,
+        salt || hre.ethers.ZeroHash,
         delay
     );
-    await res.wait();
-    console.log(res.hash);
 }
 
 export async function getProxyFactoryAt(