refactor(sdk): Remove NetworkArg + Make functions that accept Network methods of Network

batcher/aligned-sdk/Cargo.toml

@@ -24,4 +24,4 @@ url = "2.5.0"
 hex = "0.4.3"
 ciborium = "=0.2.2"
 serde_repr = "0.1.19"
-dialoguer = "0.11.0"
+dialoguer = "0.11.0"

batcher/aligned-sdk/src/communication/batch.rs

@@ -41,7 +41,7 @@ pub async fn await_batch_verification(
     network: Network,
 ) -> Result<(), errors::SubmitError> {
     for _ in 0..RETRIES {
-        if is_proof_verified(aligned_verification_data, network, rpc_url)
+        if is_proof_verified(aligned_verification_data, network.clone(), rpc_url)
             .await
             .is_ok_and(|r| r)
         {

batcher/aligned-sdk/src/core/constants.rs

@@ -37,3 +37,20 @@ pub const BUMP_MIN_RETRY_DELAY: u64 = 500; // milliseconds
 pub const BUMP_MAX_RETRIES: usize = 33; // ~ 1 day
 pub const BUMP_BACKOFF_FACTOR: f32 = 2.0;
 pub const BUMP_MAX_RETRY_DELAY: u64 = 3600; // seconds
+
+/// NETWORK ADDRESSES ///
+/// BatcherPaymentService
+pub const BATCHER_PAYMENT_SERVICE_ADDRESS_DEVNET: &str =
+    "0x7bc06c482DEAd17c0e297aFbC32f6e63d3846650";
+pub const BATCHER_PAYMENT_SERVICE_ADDRESS_HOLESKY: &str =
+    "0x815aeCA64a974297942D2Bbf034ABEe22a38A003";
+pub const BATCHER_PAYMENT_SERVICE_ADDRESS_HOLESKY_STAGE: &str =
+    "0x7577Ec4ccC1E6C529162ec8019A49C13F6DAd98b";
+pub const BATCHER_PAYMENT_SERVICE_ADDRESS_MAINNET: &str =
+    "0xb0567184A52cB40956df6333510d6eF35B89C8de";
+/// AlignedServiceManager
+pub const ALIGNED_SERVICE_MANAGER_DEVNET: &str = "0x851356ae760d987E095750cCeb3bC6014560891C";
+pub const ALIGNED_SERVICE_MANAGER_HOLESKY: &str = "0x58F280BeBE9B34c9939C3C39e0890C81f163B623";
+pub const ALIGNED_SERVICE_MANAGER_HOLESKY_STAGE: &str =
+    "0x9C5231FC88059C086Ea95712d105A2026048c39B";
+pub const ALIGNED_SERVICE_MANAGER_MAINNET: &str = "0xeF2A435e5EE44B2041100EF8cbC8ae035166606c";

batcher/aligned-sdk/src/core/types.rs

@@ -11,13 +11,20 @@ use ethers::types::transaction::eip712::Eip712;
 use ethers::types::transaction::eip712::Eip712Error;
 use ethers::types::Address;
 use ethers::types::Signature;
+use ethers::types::H160;
 use ethers::types::U256;
 use lambdaworks_crypto::merkle_tree::{
     merkle::MerkleTree, proof::Proof, traits::IsMerkleTreeBackend,
 };
 use serde::{Deserialize, Serialize};
 use sha3::{Digest, Keccak256};
 
+use super::constants::{
+    ALIGNED_SERVICE_MANAGER_DEVNET, ALIGNED_SERVICE_MANAGER_HOLESKY,
+    ALIGNED_SERVICE_MANAGER_HOLESKY_STAGE, ALIGNED_SERVICE_MANAGER_MAINNET,
+    BATCHER_PAYMENT_SERVICE_ADDRESS_DEVNET, BATCHER_PAYMENT_SERVICE_ADDRESS_HOLESKY,
+    BATCHER_PAYMENT_SERVICE_ADDRESS_HOLESKY_STAGE, BATCHER_PAYMENT_SERVICE_ADDRESS_MAINNET,
+};
 use super::errors::VerifySignatureError;
 
 // VerificationData is a bytes32 instead of a VerificationData struct because in the BatcherPaymentService contract
@@ -396,12 +403,37 @@ pub enum GetNonceResponseMessage {
     InvalidRequest(String),
 }
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone)]
 pub enum Network {
     Devnet,
     Holesky,
     HoleskyStage,
     Mainnet,
+    Custom(String, String),
+}
+
+impl Network {
+    pub fn get_batcher_payment_service_address(&self) -> ethers::types::H160 {
+        match self {
+            Self::Devnet => H160::from_str(BATCHER_PAYMENT_SERVICE_ADDRESS_DEVNET).unwrap(),
+            Self::Holesky => H160::from_str(BATCHER_PAYMENT_SERVICE_ADDRESS_HOLESKY).unwrap(),
+            Self::HoleskyStage => {
+                H160::from_str(BATCHER_PAYMENT_SERVICE_ADDRESS_HOLESKY_STAGE).unwrap()
+            }
+            Self::Mainnet => H160::from_str(BATCHER_PAYMENT_SERVICE_ADDRESS_MAINNET).unwrap(),
+            Self::Custom(s, _) => H160::from_str(s.as_str()).unwrap(),
+        }
+    }
+
+    pub fn get_aligned_service_manager_address(&self) -> ethers::types::H160 {
+        match self {
+            Self::Devnet => H160::from_str(ALIGNED_SERVICE_MANAGER_DEVNET).unwrap(),
+            Self::Holesky => H160::from_str(ALIGNED_SERVICE_MANAGER_HOLESKY).unwrap(),
+            Self::HoleskyStage => H160::from_str(ALIGNED_SERVICE_MANAGER_HOLESKY_STAGE).unwrap(),
+            Self::Mainnet => H160::from_str(ALIGNED_SERVICE_MANAGER_MAINNET).unwrap(),
+            Self::Custom(_, s) => H160::from_str(s.as_str()).unwrap(),
+        }
+    }
 }
 
 impl FromStr for Network {
@@ -413,10 +445,24 @@ impl FromStr for Network {
             "holesky-stage" => Ok(Network::HoleskyStage),
             "devnet" => Ok(Network::Devnet),
             "mainnet" => Ok(Network::Mainnet),
-            _ => Err(
-                "Invalid network, possible values are: \"holesky\", \"holesky-stage\", \"devnet\", \"mainnet\""
-                    .to_string(),
-            ),
+            s => {
+                if !s.contains("custom") {
+                    return Err(
+                        "Invalid network, possible values are: \"holesky\", \"holesky-stage\", \"devnet\", \"mainnet\", \"custom <BATCHER_PAYMENT_SERVICE_ADDRESS> <ALIGNED_SERVICE_MANAGER_ADDRESS>\""
+                            .to_string(),
+                    );
+                }
+                let parts: Vec<&str> = s.split_whitespace().collect();
+
+                if parts.len() == 3 && parts[0].contains("custom") {
+                    Ok(Network::Custom(parts[1].to_string(), parts[2].to_string()))
+                } else {
+                    Err(
+                        "Invalid custom network, \"custom <BATCHER_PAYMENT_SERVICE_ADDRESS, ALIGNED_SERVICE_MANAGER_ADDRESS>\""
+                            .to_string()
+                    )
+                }
+            }
         }
     }
 }

batcher/aligned-sdk/src/sdk.rs

@@ -28,10 +28,10 @@ use ethers::{
     prelude::k256::ecdsa::SigningKey,
     providers::{Http, Middleware, Provider},
     signers::{LocalWallet, Wallet},
-    types::{Address, H160, U256},
+    types::{Address, U256},
 };
 use sha3::{Digest, Keccak256};
-use std::{str::FromStr, sync::Arc};
+use std::sync::Arc;
 use tokio::{net::TcpStream, sync::Mutex};
 use tokio_tungstenite::{connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream};
 
@@ -89,7 +89,7 @@ pub async fn submit_multiple_and_wait_verification(
 ) -> Vec<Result<AlignedVerificationData, errors::SubmitError>> {
     let mut aligned_verification_data = submit_multiple(
         batcher_url,
-        network,
+        network.clone(),
         verification_data,
         max_fee,
         wallet,
@@ -102,7 +102,8 @@ pub async fn submit_multiple_and_wait_verification(
     let mut error_awaiting_batch_verification: Option<errors::SubmitError> = None;
     for aligned_verification_data_item in aligned_verification_data.iter().flatten() {
         if let Err(e) =
-            await_batch_verification(aligned_verification_data_item, eth_rpc_url, network).await
+            await_batch_verification(aligned_verification_data_item, eth_rpc_url, network.clone())
+                .await
         {
             error_awaiting_batch_verification = Some(e);
             break;
@@ -271,28 +272,6 @@ pub async fn submit_multiple(
     .await
 }
 
-pub fn get_payment_service_address(network: Network) -> ethers::types::H160 {
-    match network {
-        Network::Devnet => H160::from_str("0x7bc06c482DEAd17c0e297aFbC32f6e63d3846650").unwrap(),
-        Network::Holesky => H160::from_str("0x815aeCA64a974297942D2Bbf034ABEe22a38A003").unwrap(),
-        Network::HoleskyStage => {
-            H160::from_str("0x7577Ec4ccC1E6C529162ec8019A49C13F6DAd98b").unwrap()
-        }
-        Network::Mainnet => H160::from_str("0xb0567184A52cB40956df6333510d6eF35B89C8de").unwrap(),
-    }
-}
-
-pub fn get_aligned_service_manager_address(network: Network) -> ethers::types::H160 {
-    match network {
-        Network::Devnet => H160::from_str("0x1613beB3B2C4f22Ee086B2b38C1476A3cE7f78E8").unwrap(),
-        Network::Holesky => H160::from_str("0x58F280BeBE9B34c9939C3C39e0890C81f163B623").unwrap(),
-        Network::HoleskyStage => {
-            H160::from_str("0x9C5231FC88059C086Ea95712d105A2026048c39B").unwrap()
-        }
-        Network::Mainnet => H160::from_str("0xeF2A435e5EE44B2041100EF8cbC8ae035166606c").unwrap(),
-    }
-}
-
 // Will submit the proofs to the batcher and wait for their responses
 // Will return once all proofs are responded, or up to a proof that is responded with an error
 async fn _submit_multiple(
@@ -328,7 +307,7 @@ async fn _submit_multiple(
 
     let response_stream = Arc::new(Mutex::new(response_stream));
 
-    let payment_service_addr = get_payment_service_address(network);
+    let payment_service_addr = network.get_batcher_payment_service_address();
 
     let result = async {
         let sent_verification_data_rev = send_messages(
@@ -498,8 +477,8 @@ async fn _is_proof_verified(
     network: Network,
     eth_rpc_provider: Provider<Http>,
 ) -> Result<bool, errors::VerificationError> {
-    let contract_address = get_aligned_service_manager_address(network);
-    let payment_service_addr = get_payment_service_address(network);
+    let contract_address = network.clone().get_aligned_service_manager_address();
+    let payment_service_addr = network.get_batcher_payment_service_address();
 
     // All the elements from the merkle proof have to be concatenated
     let merkle_proof: Vec<u8> = aligned_verification_data
@@ -640,7 +619,7 @@ pub async fn get_nonce_from_ethereum(
     let eth_rpc_provider = Provider::<Http>::try_from(eth_rpc_url)
         .map_err(|e| GetNonceError::EthRpcError(e.to_string()))?;
 
-    let payment_service_address = get_payment_service_address(network);
+    let payment_service_address = network.get_batcher_payment_service_address();
 
     match batcher_payment_service(eth_rpc_provider, payment_service_address).await {
         Ok(contract) => {
@@ -691,7 +670,7 @@ pub async fn deposit_to_aligned(
     signer: SignerMiddleware<Provider<Http>, LocalWallet>,
     network: Network,
 ) -> Result<ethers::types::TransactionReceipt, errors::PaymentError> {
-    let payment_service_address = get_payment_service_address(network);
+    let payment_service_address = network.get_batcher_payment_service_address();
     let from = signer.address();
 
     let tx = TransactionRequest::new()
@@ -729,7 +708,7 @@ pub async fn get_balance_in_aligned(
     let eth_rpc_provider = Provider::<Http>::try_from(eth_rpc_url)
         .map_err(|e| errors::BalanceError::EthereumProviderError(e.to_string()))?;
 
-    let payment_service_address = get_payment_service_address(network);
+    let payment_service_address = network.get_batcher_payment_service_address();
 
     match batcher_payment_service(eth_rpc_provider, payment_service_address).await {
         Ok(batcher_payment_service) => {

batcher/aligned/src/main.rs

@@ -14,6 +14,7 @@ use aligned_sdk::sdk::get_chain_id;
 use aligned_sdk::sdk::get_nonce_from_batcher;
 use aligned_sdk::sdk::{deposit_to_aligned, get_balance_in_aligned};
 use aligned_sdk::sdk::{get_vk_commitment, is_proof_verified, save_response, submit_multiple};
+use clap::value_parser;
 use clap::Parser;
 use clap::Subcommand;
 use clap::ValueEnum;
@@ -118,9 +119,10 @@ pub struct SubmitArgs {
     #[arg(
         name = "The working network's name",
         long = "network",
-        default_value = "devnet"
+        default_value = "devnet",
+        value_parser = value_parser!(Network)
     )]
-    network: NetworkArg,
+    network: Network,
 }
 
 #[derive(Parser, Debug)]
@@ -141,9 +143,10 @@ pub struct DepositToBatcherArgs {
     #[arg(
         name = "The working network's name",
         long = "network",
-        default_value = "devnet"
+        default_value = "devnet",
+        value_parser = value_parser!(Network)
     )]
-    network: NetworkArg,
+    network: Network,
     #[arg(name = "Amount to deposit", long = "amount", required = true)]
     amount: String,
 }
@@ -162,9 +165,10 @@ pub struct VerifyProofOnchainArgs {
     #[arg(
         name = "The working network's name",
         long = "network",
-        default_value = "devnet"
+        default_value = "devnet",
+        value_parser = value_parser!(Network)
     )]
-    network: NetworkArg,
+    network: Network,
 }
 
 #[derive(Parser, Debug)]
@@ -184,9 +188,10 @@ pub struct GetUserBalanceArgs {
     #[arg(
         name = "The working network's name",
         long = "network",
-        default_value = "devnet"
+        default_value = "devnet",
+        value_parser = value_parser!(Network)
     )]
-    network: NetworkArg,
+    network: Network,
     #[arg(
         name = "Ethereum RPC provider address",
         long = "rpc_url",
@@ -218,25 +223,6 @@ pub struct GetUserNonceArgs {
     address: String,
 }
 
-#[derive(Debug, Clone, ValueEnum, Copy)]
-enum NetworkArg {
-    Devnet,
-    Holesky,
-    HoleskyStage,
-    Mainnet,
-}
-
-impl From<NetworkArg> for Network {
-    fn from(env_arg: NetworkArg) -> Self {
-        match env_arg {
-            NetworkArg::Devnet => Network::Devnet,
-            NetworkArg::Holesky => Network::Holesky,
-            NetworkArg::HoleskyStage => Network::HoleskyStage,
-            NetworkArg::Mainnet => Network::Mainnet,
-        }
-    }
-}
-
 #[derive(Debug, Clone, ValueEnum)]
 pub enum ProvingSystemArg {
     #[clap(name = "GnarkPlonkBls12_381")]
@@ -366,7 +352,7 @@ async fn main() -> Result<(), AlignedError> {
 
             let aligned_verification_data_vec = submit_multiple(
                 &connect_addr,
-                submit_args.network.into(),
+                submit_args.network,
                 &verification_data_arr,
                 max_fee_wei,
                 wallet.clone(),
@@ -425,7 +411,7 @@ async fn main() -> Result<(), AlignedError> {
             info!("Verifying response data matches sent proof data...");
             let response = is_proof_verified(
                 &aligned_verification_data,
-                verify_inclusion_args.network.into(),
+                verify_inclusion_args.network,
                 &verify_inclusion_args.eth_rpc_url,
             )
             .await?;
@@ -490,9 +476,7 @@ async fn main() -> Result<(), AlignedError> {
 
             let client = SignerMiddleware::new(eth_rpc_provider.clone(), wallet.clone());
 
-            match deposit_to_aligned(amount_wei, client, deposit_to_batcher_args.network.into())
-                .await
-            {
+            match deposit_to_aligned(amount_wei, client, deposit_to_batcher_args.network).await {
                 Ok(receipt) => {
                     info!(
                         "Payment sent to the batcher successfully. Tx: 0x{:x}",
@@ -509,7 +493,7 @@ async fn main() -> Result<(), AlignedError> {
             match get_balance_in_aligned(
                 user_address,
                 &get_user_balance_args.eth_rpc_url,
-                get_user_balance_args.network.into(),
+                get_user_balance_args.network,
             )
             .await
             {