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, StakeTableV3,
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
412#[derive(Clone, Debug)]
413pub struct StakeTableV3UpgradeParams {
414 pub multisig_address: Address,
415}
416
417pub async fn upgrade_stake_table_v3_multisig_owner(
420 provider: impl Provider,
421 contracts: &mut Contracts,
422 params: StakeTableV3UpgradeParams,
423 multisig_owner_check: MultisigOwnerCheck,
424) -> Result<CalldataInfo> {
425 let expected_major_version: u8 = 3;
426
427 tracing::info!("Upgrading StakeTableProxy to StakeTableV3 using multisig owner");
428 let Some(proxy_addr) = contracts.address(Contract::StakeTableProxy) else {
429 anyhow::bail!("StakeTableProxy not found, can't upgrade")
430 };
431
432 let proxy = StakeTableV3::new(proxy_addr, &provider);
433 let owner_addr = proxy.owner().call().await?;
434
435 if owner_addr != params.multisig_address {
436 anyhow::bail!(
437 "Proxy not owned by multisig. expected: {:#x}, got: {owner_addr:#x}",
438 params.multisig_address
439 );
440 }
441 if multisig_owner_check == MultisigOwnerCheck::RequireContract
442 && !crate::is_contract(&provider, owner_addr).await?
443 {
444 anyhow::bail!(
445 "StakeTableProxy owner {owner_addr:#x} is not a contract (expected multisig)"
446 );
447 }
448
449 let version = proxy.getVersion().call().await?;
451 if version.majorVersion < 2 {
452 anyhow::bail!(
453 "StakeTableProxy must be at major version >= 2 to upgrade to V3, found {}",
454 version.majorVersion
455 );
456 }
457
458 let v3_addr = contracts
459 .deploy(
460 Contract::StakeTableV3,
461 StakeTableV3::deploy_builder(&provider),
462 )
463 .await?;
464
465 let init_data =
468 if crate::already_initialized(&provider, proxy_addr, expected_major_version).await? {
469 tracing::info!(
470 "StakeTableProxy already initialized at V{expected_major_version}, skipping \
471 initializeV3()"
472 );
473 vec![].into()
474 } else {
475 StakeTableV3::new(Address::ZERO, &provider)
476 .initializeV3()
477 .calldata()
478 .to_owned()
479 };
480
481 encode_upgrade_calldata(proxy_addr, v3_addr, init_data)
482}
483
484pub async fn upgrade_fee_contract_multisig_owner(
487 provider: impl Provider,
488 contracts: &mut Contracts,
489 multisig_owner_check: MultisigOwnerCheck,
490) -> Result<CalldataInfo> {
491 let proxy_addr = contracts
492 .address(Contract::FeeContractProxy)
493 .ok_or_else(|| anyhow!("FeeContractProxy (multisig owner) not found, can't upgrade"))?;
494 tracing::info!("FeeContractProxy found at {proxy_addr:#x}");
495 let proxy = FeeContract::new(proxy_addr, &provider);
496 let owner_addr = proxy.owner().call().await?;
497 if multisig_owner_check == MultisigOwnerCheck::RequireContract
498 && !crate::is_contract(&provider, owner_addr).await?
499 {
500 anyhow::bail!(
501 "FeeContractProxy owner {owner_addr:#x} is not a contract (expected multisig)"
502 );
503 }
504
505 let curr_version = proxy.getVersion().call().await?;
506 if curr_version.majorVersion != 1 {
507 anyhow::bail!(
508 "Expected FeeContract V1.x for upgrade to V1.0.1, found V{}.{}.{}",
509 curr_version.majorVersion,
510 curr_version.minorVersion,
511 curr_version.patchVersion
512 );
513 }
514
515 let cached_fee_contract_addr = contracts.address(Contract::FeeContract);
516 if let Some(cached_fee_contract_addr) = cached_fee_contract_addr {
517 anyhow::bail!(
518 "FeeContract implementation address is already set in cache ({:#x}). For patch \
519 upgrades, the implementation must be redeployed. Please unset \
520 ESPRESSO_FEE_CONTRACT_ADDRESS or remove it from the cache first.",
521 cached_fee_contract_addr
522 );
523 }
524
525 let fee_contract_addr = contracts
526 .deploy(
527 Contract::FeeContract,
528 FeeContract::deploy_builder(&provider),
529 )
530 .await?;
531
532 encode_upgrade_calldata(proxy_addr, fee_contract_addr, Bytes::new())
533}
534
535#[cfg(test)]
536mod tests {
537 use alloy::primitives::{Address, Bytes, U256};
538
539 use super::*;
540
541 #[test]
542 fn test_encode_upgrade_calldata() {
543 let proxy = Address::random();
544 let impl_addr = Address::random();
545 let info = encode_upgrade_calldata(proxy, impl_addr, Bytes::new()).unwrap();
546 assert_eq!(info.to, proxy);
547 assert!(info.data.len() > 4);
548 assert_eq!(info.value, U256::ZERO);
549 let fi = info.function_info.unwrap();
550 assert_eq!(
551 fi.signature,
552 "upgradeToAndCall(address newImplementation, bytes data)"
553 );
554 assert_eq!(fi.args.len(), 2);
555 }
556
557 #[test]
558 fn test_encode_upgrade_calldata_with_init_data() {
559 let proxy = Address::random();
560 let impl_addr = Address::random();
561 let empty_calldata = encode_upgrade_calldata(proxy, impl_addr, Bytes::new()).unwrap();
562 let with_data =
563 encode_upgrade_calldata(proxy, impl_addr, Bytes::from(vec![1, 2, 3, 4])).unwrap();
564 assert!(with_data.data.len() > empty_calldata.data.len());
565 }
566
567 #[test]
568 fn test_encode_transfer_ownership_calldata() {
569 let proxy = Address::random();
570 let new_owner = Address::random();
571 let info = encode_transfer_ownership_calldata(proxy, new_owner).unwrap();
572 assert_eq!(info.to, proxy);
573 assert!(info.data.len() > 4);
574 assert_eq!(info.value, U256::ZERO);
575 let fi = info.function_info.unwrap();
576 assert_eq!(fi.signature, "transferOwnership(address newOwner)");
577 assert_eq!(fi.args, vec![new_owner.to_string()]);
578 }
579
580 #[test]
581 fn test_encode_generic_calldata() {
582 let target = Address::random();
583 let addr = Address::random();
584 let info = encode_generic_calldata(
585 target,
586 "transfer(address to, uint256 amount)",
587 vec![addr.to_string(), "1000".to_string()],
588 U256::ZERO,
589 )
590 .unwrap();
591 assert_eq!(info.to, target);
592 assert!(info.data.len() > 4);
593 let fi = info.function_info.unwrap();
594 assert_eq!(fi.signature, "transfer(address to, uint256 amount)");
595 assert_eq!(fi.args, vec![addr.to_string(), "1000".to_string()]);
596 }
597
598 #[test]
599 fn test_encode_generic_calldata_arg_mismatch() {
600 let target = Address::random();
601 let result = encode_generic_calldata(
602 target,
603 "transfer(address to, uint256 amount)",
604 vec!["0x000000000000000000000000000000000000dead".to_string()], U256::ZERO,
606 );
607 assert!(result.is_err());
608 }
609}