feat(zk_toolbox): improved wallet handling

zk_toolbox/crates/common/src/ethereum.rs

@@ -6,18 +6,17 @@ use ethers::{
     middleware::MiddlewareBuilder,
     prelude::{Http, LocalWallet, Provider, Signer, SignerMiddleware},
     providers::Middleware,
-    types::{Address, TransactionRequest, H256},
+    types::{Address, TransactionRequest},
 };
 use types::TokenInfo;
 
 use crate::{logger, wallets::Wallet};
 
 pub fn create_ethers_client(
-    private_key: H256,
+    mut wallet: LocalWallet,
     l1_rpc: String,
     chain_id: Option<u64>,
 ) -> anyhow::Result<SignerMiddleware<Provider<Http>, ethers::prelude::Wallet<SigningKey>>> {
-    let mut wallet = LocalWallet::from_bytes(private_key.as_bytes())?;
     if let Some(chain_id) = chain_id {
         wallet = wallet.with_chain_id(chain_id);
     }

zk_toolbox/crates/common/src/forge.rs

@@ -143,10 +143,12 @@ impl ForgeScript {
     }
 
     // Do not start the script if balance is not enough
-    pub fn private_key(&self) -> Option<H256> {
+    pub fn private_key(&self) -> Option<LocalWallet> {
         self.args.args.iter().find_map(|a| {
             if let ForgeScriptArg::PrivateKey { private_key } = a {
-                Some(H256::from_str(private_key).unwrap())
+                let key = H256::from_str(private_key).unwrap();
+                let key = LocalWallet::from_bytes(key.as_bytes()).unwrap();
+                Some(key)
             } else {
                 None
             }
@@ -164,11 +166,7 @@ impl ForgeScript {
     }
 
     pub fn address(&self) -> Option<Address> {
-        self.private_key().and_then(|a| {
-            LocalWallet::from_bytes(a.as_bytes())
-                .ok()
-                .map(|a| Address::from_slice(a.address().as_bytes()))
-        })
+        self.private_key().map(|k| k.address())
     }
 
     pub async fn get_the_balance(&self) -> anyhow::Result<Option<U256>> {

zk_toolbox/crates/common/src/wallets.rs

@@ -1,31 +1,70 @@
 use ethers::{
-    core::rand::Rng,
+    core::rand::{CryptoRng, Rng},
     signers::{coins_bip39::English, LocalWallet, MnemonicBuilder, Signer},
     types::{Address, H256},
 };
 use serde::{Deserialize, Serialize};
+use types::parse_h256;
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Wallet {
+#[derive(Serialize, Deserialize)]
+struct WalletSerde {
     pub address: Address,
     pub private_key: Option<H256>,
 }
 
-impl Wallet {
-    pub fn random(rng: &mut impl Rng) -> Self {
-        let private_key = H256::random_using(rng);
-        let local_wallet = LocalWallet::from_bytes(private_key.as_bytes()).unwrap();
+#[derive(Debug, Clone)]
+pub struct Wallet {
+    pub address: Address,
+    pub private_key: Option<LocalWallet>,
+}
 
-        Self {
-            address: Address::from_slice(local_wallet.address().as_bytes()),
-            private_key: Some(private_key),
+impl<'de> Deserialize<'de> for Wallet {
+    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
+        let x = WalletSerde::deserialize(d)?;
+        Ok(match x.private_key {
+            None => Self {
+                address: x.address,
+                private_key: None,
+            },
+            Some(k) => {
+                let k = LocalWallet::from_bytes(k.as_bytes()).map_err(serde::de::Error::custom)?;
+                if k.address() != x.address {
+                    return Err(serde::de::Error::custom(format!(
+                        "address does not match private key: got address {:#x}, want {:#x}",
+                        x.address,
+                        k.address(),
+                    )));
+                }
+                Self::new(k)
+            }
+        })
+    }
+}
+
+impl Serialize for Wallet {
+    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
+        WalletSerde {
+            address: self.address,
+            private_key: self.private_key_h256(),
         }
+        .serialize(s)
+    }
+}
+
+impl Wallet {
+    pub fn private_key_h256(&self) -> Option<H256> {
+        self.private_key
+            .as_ref()
+            .map(|k| parse_h256(k.signer().to_bytes().as_slice()).unwrap())
+    }
+
+    pub fn random(rng: &mut (impl Rng + CryptoRng)) -> Self {
+        Self::new(LocalWallet::new(rng))
     }
 
-    pub fn new_with_key(private_key: H256) -> Self {
-        let local_wallet = LocalWallet::from_bytes(private_key.as_bytes()).unwrap();
+    pub fn new(private_key: LocalWallet) -> Self {
         Self {
-            address: Address::from_slice(local_wallet.address().as_bytes()),
+            address: private_key.address(),
             private_key: Some(private_key),
         }
     }
@@ -35,14 +74,13 @@ impl Wallet {
             .phrase(mnemonic)
             .derivation_path(&format!("{}/{}", base_path, index))?
             .build()?;
-        let private_key = H256::from_slice(&wallet.signer().to_bytes());
-        Ok(Self::new_with_key(private_key))
+        Ok(Self::new(wallet))
     }
 
     pub fn empty() -> Self {
         Self {
             address: Address::zero(),
-            private_key: Some(H256::zero()),
+            private_key: None,
         }
     }
 }

zk_toolbox/crates/config/src/chain.rs

@@ -87,8 +87,8 @@ impl ChainConfig {
 
     pub fn get_wallets_config(&self) -> anyhow::Result<WalletsConfig> {
         let path = self.configs.join(WALLETS_FILE);
-        if let Ok(wallets) = WalletsConfig::read(self.get_shell(), &path) {
-            return Ok(wallets);
+        if self.get_shell().path_exists(&path) {
+            return WalletsConfig::read(self.get_shell(), &path);
         }
         if self.wallet_creation == WalletCreation::Localhost {
             let wallets = create_localhost_wallets(self.get_shell(), &self.link_to_code, self.id)?;

zk_toolbox/crates/config/src/ecosystem.rs

@@ -196,8 +196,8 @@ impl EcosystemConfig {
 
     pub fn get_wallets(&self) -> anyhow::Result<WalletsConfig> {
         let path = self.config.join(WALLETS_FILE);
-        if let Ok(wallets) = WalletsConfig::read(self.get_shell(), &path) {
-            return Ok(wallets);
+        if self.get_shell().path_exists(&path) {
+            return WalletsConfig::read(self.get_shell(), &path);
         }
         if self.wallet_creation == WalletCreation::Localhost {
             // Use 0 id for ecosystem  wallets

zk_toolbox/crates/config/src/wallets.rs

@@ -1,6 +1,5 @@
 use common::wallets::Wallet;
-use ethers::types::H256;
-use rand::Rng;
+use rand::{CryptoRng, Rng};
 use serde::{Deserialize, Serialize};
 
 use crate::{
@@ -20,7 +19,7 @@ pub struct WalletsConfig {
 
 impl WalletsConfig {
     /// Generate random wallets
-    pub fn random(rng: &mut impl Rng) -> Self {
+    pub fn random(rng: &mut (impl CryptoRng + Rng)) -> Self {
         Self {
             deployer: Some(Wallet::random(rng)),
             operator: Wallet::random(rng),
@@ -42,13 +41,6 @@ impl WalletsConfig {
             token_multiplier_setter: Some(Wallet::empty()),
         }
     }
-    pub fn deployer_private_key(&self) -> Option<H256> {
-        self.deployer.as_ref().and_then(|wallet| wallet.private_key)
-    }
-
-    pub fn governor_private_key(&self) -> Option<H256> {
-        self.governor.private_key
-    }
 }
 
 impl FileConfigWithDefaultName for WalletsConfig {

zk_toolbox/crates/types/src/lib.rs

@@ -10,5 +10,5 @@ pub use prover_mode::*;
 pub use token_info::*;
 pub use wallet_creation::*;
 pub use zksync_basic_types::{
-    commitment::L1BatchCommitmentMode, protocol_version::ProtocolSemanticVersion,
+    commitment::L1BatchCommitmentMode, parse_h256, protocol_version::ProtocolSemanticVersion,
 };

zk_toolbox/crates/zk_inception/src/accept_ownership.rs

@@ -1,13 +1,10 @@
 use common::{
     forge::{Forge, ForgeScript, ForgeScriptArgs},
     spinner::Spinner,
+    wallets::Wallet,
 };
 use config::{forge_interface::script_params::ACCEPT_GOVERNANCE_SCRIPT_PARAMS, EcosystemConfig};
-use ethers::{
-    abi::parse_abi,
-    contract::BaseContract,
-    types::{Address, H256},
-};
+use ethers::{abi::parse_abi, contract::BaseContract, types::Address};
 use lazy_static::lazy_static;
 use xshell::Shell;
 
@@ -31,7 +28,7 @@ pub async fn accept_admin(
     shell: &Shell,
     ecosystem_config: &EcosystemConfig,
     admin: Address,
-    governor: Option<H256>,
+    governor: &Wallet,
     target_address: Address,
     forge_args: &ForgeScriptArgs,
     l1_rpc_url: String,
@@ -62,7 +59,7 @@ pub async fn accept_owner(
     shell: &Shell,
     ecosystem_config: &EcosystemConfig,
     governor_contract: Address,
-    governor: Option<H256>,
+    governor: &Wallet,
     target_address: Address,
     forge_args: &ForgeScriptArgs,
     l1_rpc_url: String,
@@ -89,10 +86,10 @@ pub async fn accept_owner(
 
 async fn accept_ownership(
     shell: &Shell,
-    governor: Option<H256>,
+    governor: &Wallet,
     mut forge: ForgeScript,
 ) -> anyhow::Result<()> {
-    forge = fill_forge_private_key(forge, governor)?;
+    forge = fill_forge_private_key(forge, Some(governor))?;
     check_the_balance(&forge).await?;
     let spinner = Spinner::new(MSG_ACCEPTING_GOVERNANCE_SPINNER);
     forge.run(shell)?;

zk_toolbox/crates/zk_inception/src/commands/chain/accept_chain_ownership.rs

@@ -30,7 +30,7 @@ pub async fn run(args: ForgeScriptArgs, shell: &Shell) -> anyhow::Result<()> {
         shell,
         &ecosystem_config,
         contracts.l1.chain_admin_addr,
-        chain_config.get_wallets_config()?.governor_private_key(),
+        &chain_config.get_wallets_config()?.governor,
         contracts.l1.diamond_proxy_addr,
         &args,
         l1_rpc_url.clone(),

zk_toolbox/crates/zk_inception/src/commands/chain/deploy_l2_contracts.rs

@@ -277,10 +277,7 @@ async fn call_forge(
         forge = forge.with_signature(signature);
     }
 
-    forge = fill_forge_private_key(
-        forge,
-        ecosystem_config.get_wallets()?.governor_private_key(),
-    )?;
+    forge = fill_forge_private_key(forge, Some(&ecosystem_config.get_wallets()?.governor))?;
 
     check_the_balance(&forge).await?;
     forge.run(shell)?;

zk_toolbox/crates/zk_inception/src/commands/chain/deploy_paymaster.rs

@@ -56,10 +56,7 @@ pub async fn deploy_paymaster(
     if let Some(address) = sender {
         forge = forge.with_sender(address);
     } else {
-        forge = fill_forge_private_key(
-            forge,
-            chain_config.get_wallets_config()?.governor_private_key(),
-        )?;
+        forge = fill_forge_private_key(forge, Some(&chain_config.get_wallets_config()?.governor))?;
     }
 
     if broadcast {

zk_toolbox/crates/zk_inception/src/commands/chain/init/mod.rs

@@ -107,7 +107,7 @@ pub async fn init(
         shell,
         ecosystem_config,
         contracts_config.l1.chain_admin_addr,
-        chain_config.get_wallets_config()?.governor_private_key(),
+        &chain_config.get_wallets_config()?.governor,
         contracts_config.l1.diamond_proxy_addr,
         &init_args.forge_args.clone(),
         init_args.l1_rpc_url.clone(),
@@ -121,7 +121,7 @@ pub async fn init(
         set_token_multiplier_setter(
             shell,
             ecosystem_config,
-            chain_config.get_wallets_config()?.governor_private_key(),
+            &chain_config.get_wallets_config()?.governor,
             contracts_config.l1.chain_admin_addr,
             chain_config
                 .get_wallets_config()

zk_toolbox/crates/zk_inception/src/commands/chain/register_chain.rs

@@ -81,7 +81,7 @@ pub async fn register_chain(
     if let Some(address) = sender {
         forge = forge.with_sender(address);
     } else {
-        forge = fill_forge_private_key(forge, config.get_wallets()?.governor_private_key())?;
+        forge = fill_forge_private_key(forge, Some(&config.get_wallets()?.governor))?;
         check_the_balance(&forge).await?;
     }
 

zk_toolbox/crates/zk_inception/src/commands/chain/set_token_multiplier_setter.rs

@@ -3,12 +3,13 @@ use common::{
     forge::{Forge, ForgeScript, ForgeScriptArgs},
     logger,
     spinner::Spinner,
+    wallets::Wallet,
 };
 use config::{forge_interface::script_params::ACCEPT_GOVERNANCE_SCRIPT_PARAMS, EcosystemConfig};
 use ethers::{abi::parse_abi, contract::BaseContract, utils::hex};
 use lazy_static::lazy_static;
 use xshell::Shell;
-use zksync_basic_types::{Address, H256};
+use zksync_basic_types::Address;
 
 use crate::{
     messages::{
@@ -52,7 +53,7 @@ pub async fn run(args: ForgeScriptArgs, shell: &Shell) -> anyhow::Result<()> {
     set_token_multiplier_setter(
         shell,
         &ecosystem_config,
-        chain_config.get_wallets_config()?.governor_private_key(),
+        &chain_config.get_wallets_config()?.governor,
         contracts_config.l1.chain_admin_addr,
         token_multiplier_setter_address,
         &args.clone(),
@@ -72,7 +73,7 @@ pub async fn run(args: ForgeScriptArgs, shell: &Shell) -> anyhow::Result<()> {
 pub async fn set_token_multiplier_setter(
     shell: &Shell,
     ecosystem_config: &EcosystemConfig,
-    governor: Option<H256>,
+    governor: &Wallet,
     chain_admin_address: Address,
     target_address: Address,
     forge_args: &ForgeScriptArgs,
@@ -105,10 +106,10 @@ pub async fn set_token_multiplier_setter(
 
 async fn update_token_multiplier_setter(
     shell: &Shell,
-    governor: Option<H256>,
+    governor: &Wallet,
     mut forge: ForgeScript,
 ) -> anyhow::Result<()> {
-    forge = fill_forge_private_key(forge, governor)?;
+    forge = fill_forge_private_key(forge, Some(governor))?;
     check_the_balance(&forge).await?;
     forge.run(shell)?;
     Ok(())

zk_toolbox/crates/zk_inception/src/commands/chain/setup_legacy_bridge.rs

@@ -59,10 +59,7 @@ pub async fn setup_legacy_bridge(
         )
         .with_broadcast();
 
-    forge = fill_forge_private_key(
-        forge,
-        ecosystem_config.get_wallets()?.governor_private_key(),
-    )?;
+    forge = fill_forge_private_key(forge, Some(&ecosystem_config.get_wallets()?.governor))?;
 
     let spinner = Spinner::new(MSG_DEPLOYING_PAYMASTER);
     check_the_balance(&forge).await?;

zk_toolbox/crates/zk_inception/src/commands/consensus/mod.rs

@@ -11,7 +11,7 @@ use ethers::{
     contract::{FunctionCall, Multicall},
     middleware::{Middleware, NonceManagerMiddleware, SignerMiddleware},
     providers::{Http, JsonRpcClient, PendingTransaction, Provider, RawCall as _},
-    signers::{LocalWallet, Signer as _},
+    signers::Signer as _,
     types::{Address, BlockId, H256},
 };
 use xshell::Shell;
@@ -182,9 +182,7 @@ impl Setup {
             .governor
             .private_key
             .context(messages::MSG_GOVERNOR_PRIVATE_KEY_NOT_SET)?;
-        let governor = LocalWallet::from_bytes(governor.as_bytes())
-            .context("LocalWallet::from_bytes()")?
-            .with_chain_id(self.genesis.l2_chain_id.as_u64());
+        let governor = governor.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.