diff --git a/zk_toolbox/crates/zk_inception/src/commands/consensus.rs b/zk_toolbox/crates/zk_inception/src/commands/consensus.rs index 06bfd12b851..41b968e7128 100644 --- a/zk_toolbox/crates/zk_inception/src/commands/consensus.rs +++ b/zk_toolbox/crates/zk_inception/src/commands/consensus.rs @@ -11,7 +11,7 @@ use ethers::{ middleware::{Middleware, NonceManagerMiddleware, SignerMiddleware}, providers::{Http, JsonRpcClient, PendingTransaction, Provider, RawCall as _}, signers::{LocalWallet, Signer as _}, - types::{BlockId, H256}, + types::{Address, BlockId, H256}, }; use xshell::Shell; use zksync_consensus_crypto::ByteFmt; @@ -110,24 +110,6 @@ impl TxSet { } } -async fn get_attester_committee( - consensus_registry: &abi::ConsensusRegistry, -) -> anyhow::Result { - let attesters = consensus_registry - .get_attester_committee() - .call() - .await - .context("get_attester_committee()")?; - let attesters: Vec<_> = attesters - .iter() - .map(decode_weighted_attester) - .collect::>() - .context("decode_weighted_attester")?; - let attesters = - attester::Committee::new(attesters.into_iter()).context("attester::Committee::new()")?; - Ok(attesters) -} - fn print_attesters(committee: &attester::Committee) { logger::success( committee @@ -138,19 +120,43 @@ fn print_attesters(committee: &attester::Committee) { ); } -impl Command { - pub(crate) async fn run(self, shell: &Shell) -> anyhow::Result<()> { - let ecosystem_config = EcosystemConfig::from_file(shell)?; - let chain_name = global_config().chain_name.clone(); - let chain_config = ecosystem_config - .load_chain(chain_name) - .context(messages::MSG_CHAIN_NOT_INITIALIZED)?; +struct Setup { + chain: config::ChainConfig, + contracts: config::ContractsConfig, + general: config::GeneralConfig, + genesis: config::GenesisConfig, +} + +impl Setup { + fn provider(&self) -> anyhow::Result> { + let l2_url = &self + .general + .api_config + .as_ref() + .context("api_config missing")? + .web3_json_rpc + .http_url; + Provider::try_from(l2_url).with_context(|| format!("Provider::try_from({l2_url})")) + } + + fn multicall(&self, m: Arc) -> anyhow::Result> { + Ok(Multicall::new_with_chain_id( + m, + Some( + self.chain + .get_contracts_config() + .context("get_contracts_config()")? + .l2 + .multicall3 + .context("multicall3 contract not configured")?, + ), + Some(self.genesis.l2_chain_id.as_u64()), + )?) + } - let chain_id = chain_config - .get_genesis_config() - .context("get_genesis_config()")? - .l2_chain_id; - let governor = chain_config + fn governor(&self) -> anyhow::Result> { + let governor = self + .chain .get_wallets_config() .context("get_secrets_config()")? .governor @@ -158,153 +164,191 @@ impl Command { .context("governor private key not set")?; let governor = LocalWallet::from_bytes(governor.as_bytes()) .context("LocalWallet::from_bytes()")? - .with_chain_id(chain_id.as_u64()); - - let cfg = chain_config - .get_general_config() - .context("get_general_config()")?; - let l2_url = &cfg - .api_config - .as_ref() - .context("api_config missing")? - .web3_json_rpc - .http_url; - let provider: Provider = l2_url - .try_into() - .with_context(|| format!("{l2_url}.try_into::()"))?; - let block_id = BlockId::from( - provider - .get_block_number() - .await - .context("get_block_number")?, - ); - let signer = SignerMiddleware::new(provider.clone(), governor.clone()); + .with_chain_id(self.genesis.l2_chain_id.as_u64()); + let provider = self.provider().context("provider()")?; + let signer = SignerMiddleware::new(provider, governor.clone()); // Allows us to send next transaction without waiting for the previous to complete. let signer = NonceManagerMiddleware::new(signer, governor.address()); - let signer = Arc::new(signer); + Ok(Arc::new(signer)) + } - let contracts_cfg = chain_config + fn new(shell: &Shell) -> anyhow::Result { + let ecosystem_config = EcosystemConfig::from_file(shell)?; + let chain_name = global_config().chain_name.clone(); + let chain = ecosystem_config + .load_chain(chain_name) + .context(messages::MSG_CHAIN_NOT_INITIALIZED)?; + let contracts = chain .get_contracts_config() .context("get_contracts_config()")?; - // We use it to batch read calls. - let mut multicall = Multicall::new_with_chain_id( - signer.clone(), - Some( - contracts_cfg - .l2 - .multicall3 - .context("multicall3 contract not configured")?, - ), - Some(chain_id.as_u64()), - ) - .context("Multicall::new()")?; - let addr = contracts_cfg + let genesis = chain.get_genesis_config().context("get_genesis_config()")?; + let general = chain.get_general_config().context("get_general_config()")?; + Ok(Self { + chain, + contracts, + general, + genesis, + }) + } + + fn consensus_registry( + &self, + m: Arc, + ) -> anyhow::Result> { + let addr = self + .contracts .l2 .consensus_registry .context("consensus_registry address not configured")?; - let consensus_registry = abi::ConsensusRegistry::new(addr, signer.clone()); - match self { - Self::SetAttesterCommittee => { - // Fetch the desired state. - let want = (|| { - Some( - &cfg.consensus_config - .as_ref()? - .genesis_spec - .as_ref()? - .attesters, - ) - })() - .context("consensus.genesis_spec.attesters missing in general.yaml")?; - let want = parse_attester_committee(want).context("parse_attester_committee()")?; + Ok(abi::ConsensusRegistry::new(addr, m)) + } - // Fetch contract state. - let n: usize = consensus_registry - .num_nodes() - .call_raw() - .block(block_id) - .await - .context("num_nodes()")? - .try_into() - .ok() - .context("overflow")?; - multicall.block = Some(block_id); - let node_owners: Vec = multicall - .add_calls( - false, - (0..n).map(|i| consensus_registry.node_owners(i.into())), - ) - .call_array() - .await - .context("node_owners()")?; - multicall.clear_calls(); - let nodes: Vec = multicall - .add_calls( - false, - node_owners - .iter() - .map(|addr| consensus_registry.nodes(*addr)), - ) - .call_array() - .await - .context("nodes()")?; - multicall.clear_calls(); + async fn last_block(&self, m: &(impl 'static + Middleware)) -> anyhow::Result { + Ok(m.get_block_number() + .await + .context("get_block_number")? + .into()) + } - // Update the state. - let mut txs = TxSet::default(); - let mut to_insert: HashMap<_, _> = - want.iter().map(|a| (a.key.clone(), a.weight)).collect(); - for (i, node) in nodes.into_iter().enumerate() { - if node.attester_latest.removed { - continue; - } - let got = attester::WeightedAttester { - key: decode_attester_key(&node.attester_latest.pub_key) - .context("decode_attester_key")?, - weight: node.attester_latest.weight.into(), - }; - if let Some(weight) = to_insert.remove(&got.key) { - if weight != got.weight { - txs.send( - "changed_attester_weight", - consensus_registry.change_attester_weight( - node_owners[i], - weight.try_into().context("overflow")?, - ), - ) - .await?; - } - if !node.attester_latest.active { - txs.send("activate", consensus_registry.activate(node_owners[i])) - .await?; - } - } else { - txs.send("remove", consensus_registry.remove(node_owners[i])) - .await?; - } - } - for (key, weight) in to_insert { - let vk = validator::SecretKey::generate(); + async fn get_attester_committee(&self) -> anyhow::Result { + let provider = Arc::new(self.provider()?); + let consensus_registry = self + .consensus_registry(provider) + .context("consensus_registry()")?; + let attesters = consensus_registry + .get_attester_committee() + .call() + .await + .context("get_attester_committee()")?; + let attesters: Vec<_> = attesters + .iter() + .map(decode_weighted_attester) + .collect::>() + .context("decode_weighted_attester")?; + attester::Committee::new(attesters.into_iter()).context("attester::Committee::new()") + } + + async fn set_attester_committee(&self) -> anyhow::Result { + // Fetch the desired state. + let want = (|| { + Some( + &self + .general + .consensus_config + .as_ref()? + .genesis_spec + .as_ref()? + .attesters, + ) + })() + .context("consensus.genesis_spec.attesters missing in general.yaml")?; + let want = parse_attester_committee(want).context("parse_attester_committee()")?; + + let provider = self.provider().context("provider()")?; + let block_id = self.last_block(&provider).await.context("last_block()")?; + let governor = self.governor().context("governor()")?; + let consensus_registry = self + .consensus_registry(governor.clone()) + .context("consensus_registry")?; + let mut multicall = self.multicall(governor.clone()).context("multicall()")?; + + // Fetch contract state. + let n: usize = consensus_registry + .num_nodes() + .call_raw() + .block(block_id) + .await + .context("num_nodes()")? + .try_into() + .ok() + .context("overflow")?; + + multicall.block = Some(block_id); + let node_owners: Vec
= multicall + .add_calls( + false, + (0..n).map(|i| consensus_registry.node_owners(i.into())), + ) + .call_array() + .await + .context("node_owners()")?; + multicall.clear_calls(); + let nodes: Vec = multicall + .add_calls( + false, + node_owners + .iter() + .map(|addr| consensus_registry.nodes(*addr)), + ) + .call_array() + .await + .context("nodes()")?; + multicall.clear_calls(); + + // Update the state. + let mut txs = TxSet::default(); + let mut to_insert: HashMap<_, _> = want.iter().map(|a| (a.key.clone(), a.weight)).collect(); + for (i, node) in nodes.into_iter().enumerate() { + if node.attester_latest.removed { + continue; + } + let got = attester::WeightedAttester { + key: decode_attester_key(&node.attester_latest.pub_key) + .context("decode_attester_key")?, + weight: node.attester_latest.weight.into(), + }; + if let Some(weight) = to_insert.remove(&got.key) { + if weight != got.weight { txs.send( - "add", - consensus_registry.add( - ethers::types::Address::random(), - /*validator_weight=*/ 1, - encode_validator_key(&vk.public()), - encode_validator_pop(&vk.sign_pop()), + "changed_attester_weight", + consensus_registry.change_attester_weight( + node_owners[i], weight.try_into().context("overflow")?, - encode_attester_key(&key), ), ) .await?; } - txs.send( - "commit_attester_committee", - consensus_registry.commit_attester_committee(), - ) - .await?; - txs.wait(&provider).await.context("wait()")?; - let got = get_attester_committee(&consensus_registry).await?; + if !node.attester_latest.active { + txs.send("activate", consensus_registry.activate(node_owners[i])) + .await?; + } + } else { + txs.send("remove", consensus_registry.remove(node_owners[i])) + .await?; + } + } + for (key, weight) in to_insert { + let vk = validator::SecretKey::generate(); + txs.send( + "add", + consensus_registry.add( + Address::random(), + /*validator_weight=*/ 1, + encode_validator_key(&vk.public()), + encode_validator_pop(&vk.sign_pop()), + weight.try_into().context("overflow")?, + encode_attester_key(&key), + ), + ) + .await?; + } + txs.send( + "commit_attester_committee", + consensus_registry.commit_attester_committee(), + ) + .await?; + txs.wait(&provider).await.context("wait()")?; + Ok(want) + } +} + +impl Command { + pub(crate) async fn run(self, shell: &Shell) -> anyhow::Result<()> { + let setup = Setup::new(shell).context("Setup::new()")?; + match self { + Self::SetAttesterCommittee => { + let want = setup.set_attester_committee().await?; + let got = setup.get_attester_committee().await?; anyhow::ensure!( got == want, "setting attester committee failed: got {got:?}, want {want:?}" @@ -312,8 +356,8 @@ impl Command { print_attesters(&got); } Self::GetAttesterCommittee => { - let attesters = get_attester_committee(&consensus_registry).await?; - print_attesters(&attesters); + let got = setup.get_attester_committee().await?; + print_attesters(&got); } } Ok(())