Skip to content

Commit ebd83c8

Browse files
isaacdecodedisaacdecoded
authored
refactor(host-contracts): add artifacts checks and new param for reinitialize args in upgrade tasks
* refactor(host-contracts): add artifacts checks and new param for reinitialize args in upgrade tasks * refactor(host-contracts): get reinitialize args from env variables
1 parent f43a210 commit ebd83c8

3 files changed

Lines changed: 313 additions & 262 deletions

File tree

‎host-contracts/hardhat.config.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import './tasks/accounts';
1313
import './tasks/etherscanVerify';
1414
import './tasks/taskDeploy';
1515
import './tasks/taskUtils';
16-
import './tasks/upgradeProxy';
16+
import './tasks/upgradeContracts';
1717

1818
const NUM_ACCOUNTS = 15;
1919

‎host-contracts/tasks/upgradeContracts.ts‎

Lines changed: 312 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,312 @@
1+
import { HardhatUpgrades } from '@openzeppelin/hardhat-upgrades';
2+
import dotenv from 'dotenv';
3+
import { Wallet } from 'ethers';
4+
import fs from 'fs';
5+
import { task, types } from 'hardhat/config';
6+
import { HardhatRuntimeEnvironment, RunTaskFunction, TaskArguments } from 'hardhat/types';
7+
8+
import { getRequiredEnvVar } from './utils/loadVariables';
9+
10+
const REINITIALIZE_FUNCTION_PREFIX = 'reinitializeV'; // Prefix for reinitialize functions
11+
12+
// This file defines generic tasks that can be used to upgrade the implementation of already deployed contracts.
13+
14+
function getImplementationDirectory(input: string): string {
15+
const colonIndex = input.lastIndexOf('/');
16+
if (colonIndex !== -1) {
17+
return input.substring(0, colonIndex);
18+
}
19+
return input;
20+
}
21+
22+
async function upgradeCurrentToNew(
23+
proxyAddress: string,
24+
currentImplementation: string,
25+
newImplementation: string,
26+
verifyContract: boolean,
27+
hre: HardhatRuntimeEnvironment,
28+
reinitializeArgs: unknown[] = [],
29+
) {
30+
const deployerPrivateKey = getRequiredEnvVar('DEPLOYER_PRIVATE_KEY');
31+
const deployer = new Wallet(deployerPrivateKey).connect(hre.ethers.provider);
32+
33+
console.log(`Importing ${currentImplementation} contract implementation at address ${proxyAddress}...`);
34+
const currentImplementationFactory = await hre.ethers.getContractFactory(currentImplementation, deployer);
35+
const currentProxyContract = await hre.upgrades.forceImport(proxyAddress, currentImplementationFactory);
36+
console.log('Proxy contract successfully loaded!');
37+
38+
console.log(
39+
`Upgrading proxy to "${newImplementation}" implementation with reinitialize arguments:`,
40+
reinitializeArgs,
41+
);
42+
43+
// Get reinitialize function from the new implementation artifact
44+
const newImplementationArtifact = await hre.artifacts.readArtifact(newImplementation);
45+
const reinitializeFunction = newImplementationArtifact.abi.find(
46+
(item) => item.type === 'function' && item.name.includes(REINITIALIZE_FUNCTION_PREFIX),
47+
);
48+
49+
// Prepare the new implementation factory and execute the upgrade by calling the reinitialize function
50+
const newImplementationFactory = await hre.ethers.getContractFactory(newImplementation, deployer);
51+
52+
await hre.upgrades.upgradeProxy(currentProxyContract, newImplementationFactory, {
53+
call: {
54+
fn: reinitializeFunction.name,
55+
args: reinitializeArgs,
56+
},
57+
});
58+
console.log('Proxy contract successfully upgraded!');
59+
60+
if (verifyContract) {
61+
console.log('Waiting 2 minutes before contract verification... Please wait...');
62+
await new Promise((resolve) => setTimeout(resolve, 2 * 60 * 1000));
63+
const implementationAddress = await hre.upgrades.erc1967.getImplementationAddress(proxyAddress);
64+
await hre.run('verify:verify', {
65+
address: implementationAddress,
66+
constructorArguments: [],
67+
});
68+
}
69+
}
70+
71+
async function compileImplementations(
72+
currentImplementation: string,
73+
newImplementation: string,
74+
hre: HardhatRuntimeEnvironment,
75+
): Promise<void> {
76+
await hre.run('compile:specific', { contract: getImplementationDirectory(currentImplementation) });
77+
await hre.run('compile:specific', { contract: getImplementationDirectory(newImplementation) });
78+
}
79+
80+
async function checkImplementationArtifacts(
81+
expectedArtifactName: string,
82+
currentImplementation: string,
83+
newImplementation: string,
84+
hre: HardhatRuntimeEnvironment,
85+
): Promise<void> {
86+
const currentImplementationArtifact = await hre.artifacts.readArtifact(currentImplementation);
87+
if (currentImplementationArtifact.contractName !== expectedArtifactName) {
88+
throw new Error(
89+
`The current implementation artifact does not match the expected contract name "${expectedArtifactName}". Found: ${currentImplementationArtifact.contractName}`,
90+
);
91+
}
92+
93+
const newImplementationArtifact = await hre.artifacts.readArtifact(newImplementation);
94+
if (newImplementationArtifact.contractName !== expectedArtifactName) {
95+
throw new Error(
96+
`The new implementation artifact does not match the expected contract name "${expectedArtifactName}". Found: ${newImplementationArtifact.contractName}`,
97+
);
98+
}
99+
100+
const hasReinitializeFunction = newImplementationArtifact.abi.some(
101+
(item) => item.type === 'function' && item.name.includes(REINITIALIZE_FUNCTION_PREFIX),
102+
);
103+
if (!hasReinitializeFunction) {
104+
throw new Error(
105+
`The new implementation artifact does not contain a reinitialize function. Please ensure the contract has a reinitialize function defined.`,
106+
);
107+
}
108+
}
109+
110+
task('task:upgradeACL')
111+
.addParam(
112+
'currentImplementation',
113+
'The currently deployed implementation solidity contract path and name, eg: contracts/ACL.sol:ACL',
114+
)
115+
.addParam(
116+
'newImplementation',
117+
'The new implementation solidity contract path and name, eg: examples/ACLUpgradedExample.sol:ACLUpgradedExample',
118+
)
119+
.addOptionalParam(
120+
'useInternalProxyAddress',
121+
'If proxy address from the /addresses directory should be used',
122+
false,
123+
types.boolean,
124+
)
125+
.addOptionalParam(
126+
'verifyContract',
127+
'Verify new implementation on Etherscan (for eg if deploying on Sepolia or Mainnet)',
128+
true,
129+
types.boolean,
130+
)
131+
.setAction(async function (
132+
{ currentImplementation, newImplementation, useInternalProxyAddress, verifyContract }: TaskArguments,
133+
hre,
134+
) {
135+
await compileImplementations(currentImplementation, newImplementation, hre);
136+
137+
await checkImplementationArtifacts('ACL', currentImplementation, newImplementation, hre);
138+
139+
let proxyAddress: string;
140+
if (useInternalProxyAddress) {
141+
const parsedEnv = dotenv.parse(fs.readFileSync('addresses/.env.acl'));
142+
proxyAddress = parsedEnv.ACL_CONTRACT_ADDRESS;
143+
} else {
144+
proxyAddress = getRequiredEnvVar('ACL_CONTRACT_ADDRESS');
145+
}
146+
147+
const pauserAddress = getRequiredEnvVar('PAUSER_ADDRESS');
148+
149+
await upgradeCurrentToNew(proxyAddress, currentImplementation, newImplementation, verifyContract, hre, [
150+
pauserAddress,
151+
]);
152+
});
153+
154+
task('task:upgradeFHEVMExecutor')
155+
.addParam(
156+
'currentImplementation',
157+
'The currently deployed implementation solidity contract path and name, eg: contracts/FHEVMExecutor.sol:FHEVMExecutor',
158+
)
159+
.addParam(
160+
'newImplementation',
161+
'The new implementation solidity contract path and name, eg: examples/FHEVMExecutorUpgradedExample.sol:FHEVMExecutorUpgradedExample',
162+
)
163+
.addOptionalParam(
164+
'useInternalProxyAddress',
165+
'If proxy address from the /addresses directory should be used',
166+
false,
167+
types.boolean,
168+
)
169+
.addOptionalParam(
170+
'verifyContract',
171+
'Verify new implementation on Etherscan (for eg if deploying on Sepolia or Mainnet)',
172+
true,
173+
types.boolean,
174+
)
175+
.setAction(async function (
176+
{ currentImplementation, newImplementation, useInternalProxyAddress, verifyContract }: TaskArguments,
177+
hre,
178+
) {
179+
await compileImplementations(currentImplementation, newImplementation, hre);
180+
181+
await checkImplementationArtifacts('FHEVMExecutor', currentImplementation, newImplementation, hre);
182+
183+
let proxyAddress: string;
184+
if (useInternalProxyAddress) {
185+
const parsedEnv = dotenv.parse(fs.readFileSync('addresses/.env.exec'));
186+
proxyAddress = parsedEnv.FHEVM_EXECUTOR_CONTRACT_ADDRESS;
187+
} else {
188+
proxyAddress = getRequiredEnvVar('FHEVM_EXECUTOR_CONTRACT_ADDRESS');
189+
}
190+
191+
await upgradeCurrentToNew(proxyAddress, currentImplementation, newImplementation, verifyContract, hre);
192+
});
193+
194+
task('task:upgradeKMSVerifier')
195+
.addParam(
196+
'currentImplementation',
197+
'The currently deployed implementation solidity contract path and name, eg: contracts/KMSVerifier.sol:KMSVerifier',
198+
)
199+
.addParam(
200+
'newImplementation',
201+
'The new implementation solidity contract path and name, eg: examples/KMSVerifierUpgradedExample.sol:KMSVerifierUpgradedExample',
202+
)
203+
.addOptionalParam(
204+
'useInternalProxyAddress',
205+
'If proxy address from the /addresses directory should be used',
206+
false,
207+
types.boolean,
208+
)
209+
.addOptionalParam(
210+
'verifyContract',
211+
'Verify new implementation on Etherscan (for eg if deploying on Sepolia or Mainnet)',
212+
true,
213+
types.boolean,
214+
)
215+
.setAction(async function (
216+
{ currentImplementation, newImplementation, useInternalProxyAddress, verifyContract }: TaskArguments,
217+
hre,
218+
) {
219+
await compileImplementations(currentImplementation, newImplementation, hre);
220+
221+
await checkImplementationArtifacts('KMSVerifier', currentImplementation, newImplementation, hre);
222+
223+
let proxyAddress: string;
224+
if (useInternalProxyAddress) {
225+
const parsedEnv = dotenv.parse(fs.readFileSync('addresses/.env.kmsverifier'));
226+
proxyAddress = parsedEnv.KMS_VERIFIER_CONTRACT_ADDRESS;
227+
} else {
228+
proxyAddress = getRequiredEnvVar('KMS_VERIFIER_CONTRACT_ADDRESS');
229+
}
230+
231+
await upgradeCurrentToNew(proxyAddress, currentImplementation, newImplementation, verifyContract, hre);
232+
});
233+
234+
task('task:upgradeInputVerifier')
235+
.addParam(
236+
'currentImplementation',
237+
'The currently deployed implementation solidity contract path and name, eg: contracts/InputVerifier.sol:InputVerifier',
238+
)
239+
.addParam(
240+
'newImplementation',
241+
'The new implementation solidity contract path and name, eg: contracts/InputVerifier2.sol:InputVerifier',
242+
)
243+
.addOptionalParam(
244+
'useInternalProxyAddress',
245+
'If proxy address from the /addresses directory should be used',
246+
false,
247+
types.boolean,
248+
)
249+
.addOptionalParam(
250+
'verifyContract',
251+
'Verify new implementation on Etherscan (for eg if deploying on Sepolia or Mainnet)',
252+
true,
253+
types.boolean,
254+
)
255+
.setAction(async function (
256+
{ currentImplementation, newImplementation, useInternalProxyAddress, verifyContract }: TaskArguments,
257+
hre,
258+
) {
259+
await compileImplementations(currentImplementation, newImplementation, hre);
260+
261+
await checkImplementationArtifacts('InputVerifier', currentImplementation, newImplementation, hre);
262+
263+
let proxyAddress: string;
264+
if (useInternalProxyAddress) {
265+
const parsedEnv = dotenv.parse(fs.readFileSync('addresses/.env.inputverifier'));
266+
proxyAddress = parsedEnv.INPUT_VERIFIER_CONTRACT_ADDRESS;
267+
} else {
268+
proxyAddress = getRequiredEnvVar('INPUT_VERIFIER_CONTRACT_ADDRESS');
269+
}
270+
271+
await upgradeCurrentToNew(proxyAddress, currentImplementation, newImplementation, verifyContract, hre);
272+
});
273+
274+
task('task:upgradeHCULimit')
275+
.addParam(
276+
'currentImplementation',
277+
'The currently deployed implementation solidity contract path and name, eg: contracts/HCULimit.sol:HCULimit',
278+
)
279+
.addParam(
280+
'newImplementation',
281+
'The new implementation solidity contract path and name, eg: examples/HCULimitUpgradedExample.sol:HCULimitUpgradedExample',
282+
)
283+
.addOptionalParam(
284+
'useInternalProxyAddress',
285+
'If proxy address from the /addresses directory should be used',
286+
false,
287+
types.boolean,
288+
)
289+
.addOptionalParam(
290+
'verifyContract',
291+
'Verify new implementation on Etherscan (for eg if deploying on Sepolia or Mainnet)',
292+
true,
293+
types.boolean,
294+
)
295+
.setAction(async function (
296+
{ currentImplementation, newImplementation, useInternalProxyAddress, verifyContract }: TaskArguments,
297+
hre,
298+
) {
299+
await compileImplementations(currentImplementation, newImplementation, hre);
300+
301+
await checkImplementationArtifacts('HCULimit', currentImplementation, newImplementation, hre);
302+
303+
let proxyAddress: string;
304+
if (useInternalProxyAddress) {
305+
const parsedEnv = dotenv.parse(fs.readFileSync('addresses/.env.hculimit'));
306+
proxyAddress = parsedEnv.HCU_LIMIT_CONTRACT_ADDRESS;
307+
} else {
308+
proxyAddress = getRequiredEnvVar('HCU_LIMIT_CONTRACT_ADDRESS');
309+
}
310+
311+
await upgradeCurrentToNew(proxyAddress, currentImplementation, newImplementation, verifyContract, hre);
312+
});

0 commit comments

Comments
 (0)