1use alloy::{
2 hex::{FromHex, ToHexExt},
3 network::TransactionBuilder,
4 primitives::{Address, Bytes, U256},
5 providers::Provider,
6};
7use anyhow::{Context, Result, anyhow};
8use espresso_types::v0_1::L1Client;
9use hotshot_contract_adapter::sol_types::{
10 EspToken, EspTokenV2, FeeContract, LightClientV2, LightClientV2Mock, LightClientV3,
11 LightClientV3Mock, OwnableUpgradeable, PlonkVerifierV2, PlonkVerifierV3, StakeTable,
12 StakeTableV2,
13};
14
15use crate::{
16 Contract, Contracts, LIBRARY_PLACEHOLDER_ADDRESS,
17 output::{CalldataInfo, FunctionInfo},
18};
19
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21pub enum MultisigOwnerCheck {
22 RequireContract,
23 Skip,
24}
25
26#[derive(Clone)]
27pub struct TransferOwnershipParams {
28 pub new_owner: Address,
29}
30
31pub fn encode_upgrade_calldata(
33 proxy_addr: Address,
34 new_impl_addr: Address,
35 init_data: Bytes,
36) -> Result<CalldataInfo> {
37 let sig = "upgradeToAndCall(address newImplementation, bytes data)";
38 let args = vec![new_impl_addr.to_string(), init_data.to_string()];
39 let data = crate::encode_function_call(sig, args.clone())
40 .context("Failed to encode upgradeToAndCall calldata")?;
41 Ok(CalldataInfo::with_method(
42 proxy_addr,
43 data,
44 U256::ZERO,
45 FunctionInfo {
46 signature: sig.to_string(),
47 args,
48 },
49 ))
50}
51
52pub fn encode_transfer_ownership_calldata(
54 proxy_addr: Address,
55 new_owner: Address,
56) -> Result<CalldataInfo> {
57 let sig = "transferOwnership(address newOwner)";
58 let args = vec![new_owner.to_string()];
59 let data = crate::encode_function_call(sig, args.clone())
60 .context("Failed to encode transferOwnership calldata")?;
61 Ok(CalldataInfo::with_method(
62 proxy_addr,
63 data,
64 U256::ZERO,
65 FunctionInfo {
66 signature: sig.to_string(),
67 args,
68 },
69 ))
70}
71
72pub fn encode_generic_calldata(
74 target: Address,
75 function_signature: &str,
76 function_args: Vec<String>,
77 value: U256,
78) -> Result<CalldataInfo> {
79 let data = crate::encode_function_call(function_signature, function_args.clone())
80 .context("Failed to encode generic calldata")?;
81 Ok(CalldataInfo::with_method(
82 target,
83 data,
84 value,
85 FunctionInfo {
86 signature: function_signature.to_string(),
87 args: function_args,
88 },
89 ))
90}
91
92pub fn transfer_ownership_from_multisig_to_timelock(
93 contracts: &mut Contracts,
94 contract: Contract,
95 params: TransferOwnershipParams,
96) -> Result<CalldataInfo> {
97 tracing::info!(
98 "Encoding ownership transfer for {} to timelock {}",
99 contract,
100 params.new_owner
101 );
102
103 let proxy_addr = match contract {
104 Contract::LightClientProxy
105 | Contract::FeeContractProxy
106 | Contract::EspTokenProxy
107 | Contract::StakeTableProxy
108 | Contract::RewardClaimProxy => contracts
109 .address(contract)
110 .ok_or_else(|| anyhow!("{contract} (multisig owner) not found, can't upgrade"))?,
111 _ => anyhow::bail!("Not a proxy contract, can't transfer ownership"),
112 };
113 tracing::info!("{} found at {proxy_addr:#x}", contract);
114
115 encode_transfer_ownership_calldata(proxy_addr, params.new_owner)
116}
117
118pub struct LightClientV2UpgradeParams {
120 pub blocks_per_epoch: u64,
121 pub epoch_start_block: u64,
122}
123
124pub async fn upgrade_light_client_v2_multisig_owner(
127 provider: impl Provider,
128 contracts: &mut Contracts,
129 params: LightClientV2UpgradeParams,
130 is_mock: bool,
131 multisig_owner_check: MultisigOwnerCheck,
132) -> Result<CalldataInfo> {
133 let expected_major_version: u8 = 2;
134
135 let proxy_addr = contracts
136 .address(Contract::LightClientProxy)
137 .ok_or_else(|| anyhow!("LightClientProxy (multisig owner) not found, can't upgrade"))?;
138 tracing::info!("LightClientProxy found at {proxy_addr:#x}");
139
140 let owner_addr = OwnableUpgradeable::new(proxy_addr, &provider)
141 .owner()
142 .call()
143 .await?;
144 if multisig_owner_check == MultisigOwnerCheck::RequireContract
145 && !crate::is_contract(&provider, owner_addr).await?
146 {
147 anyhow::bail!(
148 "LightClientProxy owner {owner_addr:#x} is not a contract (expected multisig)"
149 );
150 }
151
152 let pv2_addr = contracts
153 .deploy(
154 Contract::PlonkVerifierV2,
155 PlonkVerifierV2::deploy_builder(&provider),
156 )
157 .await?;
158
159 let target_lcv2_bytecode = if is_mock {
160 LightClientV2Mock::BYTECODE.encode_hex()
161 } else {
162 LightClientV2::BYTECODE.encode_hex()
163 };
164 let lcv2_linked_bytecode = {
165 match target_lcv2_bytecode
166 .matches(LIBRARY_PLACEHOLDER_ADDRESS)
167 .count()
168 {
169 0 => return Err(anyhow!("lib placeholder not found")),
170 1 => Bytes::from_hex(target_lcv2_bytecode.replacen(
171 LIBRARY_PLACEHOLDER_ADDRESS,
172 &pv2_addr.encode_hex(),
173 1,
174 ))?,
175 _ => {
176 return Err(anyhow!(
177 "more than one lib placeholder found, consider using a different value"
178 ));
179 },
180 }
181 };
182 let lcv2_addr = if is_mock {
183 let addr = LightClientV2Mock::deploy_builder(&provider)
184 .map(|req| req.with_deploy_code(lcv2_linked_bytecode))
185 .deploy()
186 .await?;
187 tracing::info!("deployed LightClientV2Mock at {addr:#x}");
188 addr
189 } else {
190 contracts
191 .deploy(
192 Contract::LightClientV2,
193 LightClientV2::deploy_builder(&provider)
194 .map(|req| req.with_deploy_code(lcv2_linked_bytecode)),
195 )
196 .await?
197 };
198
199 let init_data =
200 if crate::already_initialized(&provider, proxy_addr, expected_major_version).await? {
201 tracing::info!(
202 "Proxy was already initialized for version {}",
203 expected_major_version
204 );
205 vec![].into()
206 } else {
207 tracing::info!(
208 "Init Data to be signed.\n Function: initializeV2\n Arguments:\n \
209 blocks_per_epoch: {:?}\n epoch_start_block: {:?}",
210 params.blocks_per_epoch,
211 params.epoch_start_block
212 );
213 LightClientV2::new(lcv2_addr, &provider)
214 .initializeV2(params.blocks_per_epoch, params.epoch_start_block)
215 .calldata()
216 .to_owned()
217 };
218
219 encode_upgrade_calldata(proxy_addr, lcv2_addr, init_data)
220}
221
222pub async fn upgrade_light_client_v3_multisig_owner(
225 provider: impl Provider,
226 contracts: &mut Contracts,
227 is_mock: bool,
228 multisig_owner_check: MultisigOwnerCheck,
229) -> Result<CalldataInfo> {
230 let expected_major_version: u8 = 3;
231
232 let proxy_addr = contracts
233 .address(Contract::LightClientProxy)
234 .ok_or_else(|| anyhow!("LightClientProxy (multisig owner) not found, can't upgrade"))?;
235 tracing::info!("LightClientProxy found at {proxy_addr:#x}");
236
237 let owner_addr = OwnableUpgradeable::new(proxy_addr, &provider)
238 .owner()
239 .call()
240 .await?;
241 if multisig_owner_check == MultisigOwnerCheck::RequireContract
242 && !crate::is_contract(&provider, owner_addr).await?
243 {
244 anyhow::bail!(
245 "LightClientProxy owner {owner_addr:#x} is not a contract (expected multisig)"
246 );
247 }
248
249 let pv3_addr = contracts
250 .deploy(
251 Contract::PlonkVerifierV3,
252 PlonkVerifierV3::deploy_builder(&provider),
253 )
254 .await?;
255
256 let target_lcv3_bytecode = if is_mock {
257 LightClientV3Mock::BYTECODE.encode_hex()
258 } else {
259 LightClientV3::BYTECODE.encode_hex()
260 };
261 let lcv3_linked_bytecode = {
262 match target_lcv3_bytecode
263 .matches(LIBRARY_PLACEHOLDER_ADDRESS)
264 .count()
265 {
266 0 => return Err(anyhow!("lib placeholder not found")),
267 1 => Bytes::from_hex(target_lcv3_bytecode.replacen(
268 LIBRARY_PLACEHOLDER_ADDRESS,
269 &pv3_addr.encode_hex(),
270 1,
271 ))?,
272 _ => {
273 return Err(anyhow!(
274 "more than one lib placeholder found, consider using a different value"
275 ));
276 },
277 }
278 };
279 let lcv3_addr = if is_mock {
280 let addr = LightClientV3Mock::deploy_builder(&provider)
281 .map(|req| req.with_deploy_code(lcv3_linked_bytecode))
282 .deploy()
283 .await?;
284 tracing::info!("deployed LightClientV3Mock at {addr:#x}");
285 addr
286 } else {
287 contracts
288 .deploy(
289 Contract::LightClientV3,
290 LightClientV3::deploy_builder(&provider)
291 .map(|req| req.with_deploy_code(lcv3_linked_bytecode)),
292 )
293 .await?
294 };
295
296 let init_data =
297 if crate::already_initialized(&provider, proxy_addr, expected_major_version).await? {
298 tracing::info!(
299 "Proxy was already initialized for version {}",
300 expected_major_version
301 );
302 vec![].into()
303 } else {
304 tracing::info!(
305 "Init Data to be signed.\n Function: initializeV3\n Arguments: none (V3 inherits \
306 from V2)"
307 );
308 LightClientV3::new(lcv3_addr, &provider)
309 .initializeV3()
310 .calldata()
311 .to_owned()
312 };
313
314 encode_upgrade_calldata(proxy_addr, lcv3_addr, init_data)
315}
316
317pub async fn upgrade_esp_token_v2_multisig_owner(
320 provider: impl Provider,
321 contracts: &mut Contracts,
322 multisig_owner_check: MultisigOwnerCheck,
323) -> Result<CalldataInfo> {
324 let proxy_addr = contracts
325 .address(Contract::EspTokenProxy)
326 .ok_or_else(|| anyhow!("EspTokenProxy (multisig owner) not found, can't upgrade"))?;
327 tracing::info!("EspTokenProxy found at {proxy_addr:#x}");
328 let proxy = EspToken::new(proxy_addr, &provider);
329 let owner_addr = proxy.owner().call().await?;
330 if multisig_owner_check == MultisigOwnerCheck::RequireContract
331 && !crate::is_contract(&provider, owner_addr).await?
332 {
333 anyhow::bail!("EspTokenProxy owner {owner_addr:#x} is not a contract (expected multisig)");
334 }
335
336 let esp_token_v2_addr = contracts
337 .deploy(Contract::EspTokenV2, EspTokenV2::deploy_builder(&provider))
338 .await?;
339
340 let reward_claim_addr = contracts
341 .address(Contract::RewardClaimProxy)
342 .ok_or_else(|| anyhow!("RewardClaimProxy not found"))?;
343 let proxy_as_v2 = EspTokenV2::new(proxy_addr, &provider);
344 let init_data = proxy_as_v2
345 .initializeV2(reward_claim_addr)
346 .calldata()
347 .to_owned();
348
349 tracing::info!(
350 %reward_claim_addr,
351 "Data to be signed: Function: initializeV2 Arguments:"
352 );
353
354 encode_upgrade_calldata(proxy_addr, esp_token_v2_addr, init_data)
355}
356
357#[derive(Clone, Debug)]
358pub struct StakeTableV2UpgradeParams {
359 pub multisig_address: Address,
360 pub pauser: Address,
361}
362
363pub async fn upgrade_stake_table_v2_multisig_owner(
366 provider: impl Provider,
367 l1_client: L1Client,
368 contracts: &mut Contracts,
369 params: StakeTableV2UpgradeParams,
370 multisig_owner_check: MultisigOwnerCheck,
371) -> Result<CalldataInfo> {
372 tracing::info!("Upgrading StakeTableProxy to StakeTableV2 using multisig owner");
373 let Some(proxy_addr) = contracts.address(Contract::StakeTableProxy) else {
374 anyhow::bail!("StakeTableProxy not found, can't upgrade")
375 };
376
377 let proxy = StakeTable::new(proxy_addr, &provider);
378 let owner_addr = proxy.owner().call().await?;
379
380 if owner_addr != params.multisig_address {
381 anyhow::bail!(
382 "Proxy not owned by multisig. expected: {:#x}, got: {owner_addr:#x}",
383 params.multisig_address
384 );
385 }
386 if multisig_owner_check == MultisigOwnerCheck::RequireContract
387 && !crate::is_contract(&provider, owner_addr).await?
388 {
389 anyhow::bail!(
390 "StakeTableProxy owner {owner_addr:#x} is not a contract (expected multisig)"
391 );
392 }
393
394 let (_init_commissions, _init_active_stake, init_data) =
395 crate::prepare_stake_table_v2_upgrade(l1_client, proxy_addr, params.pauser, owner_addr)
396 .await?;
397
398 let stake_table_v2_addr = contracts
399 .deploy(
400 Contract::StakeTableV2,
401 StakeTableV2::deploy_builder(&provider),
402 )
403 .await?;
404
405 encode_upgrade_calldata(
406 proxy_addr,
407 stake_table_v2_addr,
408 init_data.unwrap_or_default(),
409 )
410}
411
412pub async fn upgrade_fee_contract_multisig_owner(
415 provider: impl Provider,
416 contracts: &mut Contracts,
417 multisig_owner_check: MultisigOwnerCheck,
418) -> Result<CalldataInfo> {
419 let proxy_addr = contracts
420 .address(Contract::FeeContractProxy)
421 .ok_or_else(|| anyhow!("FeeContractProxy (multisig owner) not found, can't upgrade"))?;
422 tracing::info!("FeeContractProxy found at {proxy_addr:#x}");
423 let proxy = FeeContract::new(proxy_addr, &provider);
424 let owner_addr = proxy.owner().call().await?;
425 if multisig_owner_check == MultisigOwnerCheck::RequireContract
426 && !crate::is_contract(&provider, owner_addr).await?
427 {
428 anyhow::bail!(
429 "FeeContractProxy owner {owner_addr:#x} is not a contract (expected multisig)"
430 );
431 }
432
433 let curr_version = proxy.getVersion().call().await?;
434 if curr_version.majorVersion != 1 {
435 anyhow::bail!(
436 "Expected FeeContract V1.x for upgrade to V1.0.1, found V{}.{}.{}",
437 curr_version.majorVersion,
438 curr_version.minorVersion,
439 curr_version.patchVersion
440 );
441 }
442
443 let cached_fee_contract_addr = contracts.address(Contract::FeeContract);
444 if let Some(cached_fee_contract_addr) = cached_fee_contract_addr {
445 anyhow::bail!(
446 "FeeContract implementation address is already set in cache ({:#x}). For patch \
447 upgrades, the implementation must be redeployed. Please unset \
448 ESPRESSO_FEE_CONTRACT_ADDRESS or remove it from the cache first.",
449 cached_fee_contract_addr
450 );
451 }
452
453 let fee_contract_addr = contracts
454 .deploy(
455 Contract::FeeContract,
456 FeeContract::deploy_builder(&provider),
457 )
458 .await?;
459
460 encode_upgrade_calldata(proxy_addr, fee_contract_addr, Bytes::new())
461}
462
463#[cfg(test)]
464mod tests {
465 use alloy::primitives::{Address, Bytes, U256};
466
467 use super::*;
468
469 #[test]
470 fn test_encode_upgrade_calldata() {
471 let proxy = Address::random();
472 let impl_addr = Address::random();
473 let info = encode_upgrade_calldata(proxy, impl_addr, Bytes::new()).unwrap();
474 assert_eq!(info.to, proxy);
475 assert!(info.data.len() > 4);
476 assert_eq!(info.value, U256::ZERO);
477 let fi = info.function_info.unwrap();
478 assert_eq!(
479 fi.signature,
480 "upgradeToAndCall(address newImplementation, bytes data)"
481 );
482 assert_eq!(fi.args.len(), 2);
483 }
484
485 #[test]
486 fn test_encode_upgrade_calldata_with_init_data() {
487 let proxy = Address::random();
488 let impl_addr = Address::random();
489 let empty_calldata = encode_upgrade_calldata(proxy, impl_addr, Bytes::new()).unwrap();
490 let with_data =
491 encode_upgrade_calldata(proxy, impl_addr, Bytes::from(vec![1, 2, 3, 4])).unwrap();
492 assert!(with_data.data.len() > empty_calldata.data.len());
493 }
494
495 #[test]
496 fn test_encode_transfer_ownership_calldata() {
497 let proxy = Address::random();
498 let new_owner = Address::random();
499 let info = encode_transfer_ownership_calldata(proxy, new_owner).unwrap();
500 assert_eq!(info.to, proxy);
501 assert!(info.data.len() > 4);
502 assert_eq!(info.value, U256::ZERO);
503 let fi = info.function_info.unwrap();
504 assert_eq!(fi.signature, "transferOwnership(address newOwner)");
505 assert_eq!(fi.args, vec![new_owner.to_string()]);
506 }
507
508 #[test]
509 fn test_encode_generic_calldata() {
510 let target = Address::random();
511 let addr = Address::random();
512 let info = encode_generic_calldata(
513 target,
514 "transfer(address to, uint256 amount)",
515 vec![addr.to_string(), "1000".to_string()],
516 U256::ZERO,
517 )
518 .unwrap();
519 assert_eq!(info.to, target);
520 assert!(info.data.len() > 4);
521 let fi = info.function_info.unwrap();
522 assert_eq!(fi.signature, "transfer(address to, uint256 amount)");
523 assert_eq!(fi.args, vec![addr.to_string(), "1000".to_string()]);
524 }
525
526 #[test]
527 fn test_encode_generic_calldata_arg_mismatch() {
528 let target = Address::random();
529 let result = encode_generic_calldata(
530 target,
531 "transfer(address to, uint256 amount)",
532 vec!["0x000000000000000000000000000000000000dead".to_string()], U256::ZERO,
534 );
535 assert!(result.is_err());
536 }
537}