feat: various sdk improvements (#1313)

Co-authored-by: Marcos Nicolau <[email protected]>
Co-authored-by: Avila Gastón <[email protected]>
Co-authored-by: avilagaston9 <[email protected]>

Author: 3 people

batcher/aligned-batcher/src/connection.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use crate::types::{batch_queue::BatchQueueEntry, errors::BatcherError};
 use aligned_sdk::{
     communication::serialization::cbor_serialize,
-    core::types::{BatchInclusionData, ResponseMessage, VerificationCommitmentBatch},
+    core::types::{BatchInclusionData, SubmitProofResponseMessage, VerificationCommitmentBatch},
 };
 use futures_util::{stream::SplitSink, SinkExt};
 use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
@@ -27,7 +27,7 @@ pub(crate) async fn send_batch_inclusion_data_responses(
             batch_merkle_tree,
             entry.nonced_verification_data.nonce,
         );
-        let response = ResponseMessage::BatchInclusionData(batch_inclusion_data);
+        let response = SubmitProofResponseMessage::BatchInclusionData(batch_inclusion_data);
 
         let serialized_response = cbor_serialize(&response)
             .map_err(|e| BatcherError::SerializationError(e.to_string()))?;

batcher/aligned-batcher/src/lib.rs

@@ -11,14 +11,14 @@ use futures_util::stream::{SplitSink, TryFilter};
 use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream};
 
 use crate::communication::serialization::{cbor_deserialize, cbor_serialize};
-use crate::core::types::BatchInclusionData;
+use crate::core::types::{BatchInclusionData, SubmitProofMessage};
 use crate::{
     communication::batch::process_batcher_response,
     core::{
         errors::SubmitError,
         types::{
-            AlignedVerificationData, ClientMessage, NoncedVerificationData, ResponseMessage,
-            VerificationData, VerificationDataCommitment,
+            AlignedVerificationData, ClientMessage, NoncedVerificationData,
+            SubmitProofResponseMessage, VerificationData, VerificationDataCommitment,
         },
     },
 };
@@ -36,26 +36,28 @@ pub async fn send_messages(
     ws_write: Arc<Mutex<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>,
     payment_service_addr: Address,
     verification_data: &[VerificationData],
-    max_fees: &[U256],
+    max_fee: U256,
     wallet: Wallet<SigningKey>,
     mut nonce: U256,
 ) -> Vec<Result<NoncedVerificationData, SubmitError>> {
     let chain_id = U256::from(wallet.chain_id());
     let mut ws_write = ws_write.lock().await;
     let mut sent_verification_data: Vec<Result<NoncedVerificationData, SubmitError>> = Vec::new();
 
-    for (idx, verification_data) in verification_data.iter().enumerate() {
+    for (idx, verification_data_i) in verification_data.iter().enumerate() {
         // Build each message to send
         let verification_data = NoncedVerificationData::new(
-            verification_data.clone(),
+            verification_data_i.clone(),
             nonce,
-            max_fees[idx],
+            max_fee,
             chain_id,
             payment_service_addr,
         );
 
         nonce += U256::one();
-        let msg = ClientMessage::new(verification_data.clone(), wallet.clone()).await;
+        let data = SubmitProofMessage::new(verification_data.clone(), wallet.clone()).await;
+        let msg = ClientMessage::SubmitProof(Box::new(data));
+
         let msg_bin = match cbor_serialize(&msg) {
             Ok(bin) => bin,
             Err(e) => {
@@ -168,49 +170,52 @@ pub async fn receive(
 async fn handle_batcher_response(msg: Message) -> Result<BatchInclusionData, SubmitError> {
     let data = msg.into_data();
     match cbor_deserialize(data.as_slice()) {
-        Ok(ResponseMessage::BatchInclusionData(batch_inclusion_data)) => {
+        Ok(SubmitProofResponseMessage::BatchInclusionData(batch_inclusion_data)) => {
             //OK case. Proofs was valid and it was included in this batch.
             Ok(batch_inclusion_data)
         }
-        Ok(ResponseMessage::InvalidNonce) => {
+        Ok(SubmitProofResponseMessage::InvalidNonce) => {
             error!("Batcher responded with invalid nonce");
             Err(SubmitError::InvalidNonce)
         }
-        Ok(ResponseMessage::InvalidSignature) => {
+        Ok(SubmitProofResponseMessage::InvalidSignature) => {
             error!("Batcher responded with invalid signature");
             Err(SubmitError::InvalidSignature)
         }
-        Ok(ResponseMessage::ProofTooLarge) => {
+        Ok(SubmitProofResponseMessage::ProofTooLarge) => {
             error!("Batcher responded with proof too large");
             Err(SubmitError::ProofTooLarge)
         }
-        Ok(ResponseMessage::InvalidMaxFee) => {
+        Ok(SubmitProofResponseMessage::InvalidMaxFee) => {
             error!("Batcher responded with invalid max fee");
             Err(SubmitError::InvalidMaxFee)
         }
-        Ok(ResponseMessage::InsufficientBalance(addr)) => {
+        Ok(SubmitProofResponseMessage::InsufficientBalance(addr)) => {
             error!("Batcher responded with insufficient balance");
             Err(SubmitError::InsufficientBalance(addr))
         }
-        Ok(ResponseMessage::InvalidChainId) => {
+        Ok(SubmitProofResponseMessage::InvalidChainId) => {
             error!("Batcher responded with invalid chain id");
             Err(SubmitError::InvalidChainId)
         }
-        Ok(ResponseMessage::InvalidReplacementMessage) => {
+        Ok(SubmitProofResponseMessage::InvalidReplacementMessage) => {
             error!("Batcher responded with invalid replacement message");
             Err(SubmitError::InvalidReplacementMessage)
         }
-        Ok(ResponseMessage::AddToBatchError) => {
+        Ok(SubmitProofResponseMessage::AddToBatchError) => {
             error!("Batcher responded with add to batch error");
             Err(SubmitError::AddToBatchError)
         }
-        Ok(ResponseMessage::EthRpcError) => {
+        Ok(SubmitProofResponseMessage::EthRpcError) => {
             error!("Batcher experienced Eth RPC connection error");
             Err(SubmitError::EthereumProviderError(
                 "Batcher experienced Eth RPC connection error".to_string(),
             ))
         }
-        Ok(ResponseMessage::InvalidPaymentServiceAddress(received_addr, expected_addr)) => {
+        Ok(SubmitProofResponseMessage::InvalidPaymentServiceAddress(
+            received_addr,
+            expected_addr,
+        )) => {
             error!(
                 "Batcher responded with invalid payment service address: {:?}, expected: {:?}",
                 received_addr, expected_addr
@@ -220,11 +225,11 @@ async fn handle_batcher_response(msg: Message) -> Result<BatchInclusionData, Sub
                 expected_addr,
             ))
         }
-        Ok(ResponseMessage::InvalidProof(reason)) => {
+        Ok(SubmitProofResponseMessage::InvalidProof(reason)) => {
             error!("Batcher responded with invalid proof: {}", reason);
             Err(SubmitError::InvalidProof(reason))
         }
-        Ok(ResponseMessage::CreateNewTaskError(merkle_root, error)) => {
+        Ok(SubmitProofResponseMessage::CreateNewTaskError(merkle_root, error)) => {
             error!("Batcher responded with create new task error: {}", error);
             Err(SubmitError::BatchSubmissionFailed(
                 "Could not create task with merkle root ".to_owned()
@@ -233,18 +238,18 @@ async fn handle_batcher_response(msg: Message) -> Result<BatchInclusionData, Sub
                     + &error,
             ))
         }
-        Ok(ResponseMessage::ProtocolVersion(_)) => {
+        Ok(SubmitProofResponseMessage::ProtocolVersion(_)) => {
             error!("Batcher responded with protocol version instead of batch inclusion data");
             Err(SubmitError::UnexpectedBatcherResponse(
                 "Batcher responded with protocol version instead of batch inclusion data"
                     .to_string(),
             ))
         }
-        Ok(ResponseMessage::BatchReset) => {
+        Ok(SubmitProofResponseMessage::BatchReset) => {
             error!("Batcher responded with batch reset");
             Err(SubmitError::ProofQueueFlushed)
         }
-        Ok(ResponseMessage::Error(e)) => {
+        Ok(SubmitProofResponseMessage::Error(e)) => {
             error!("Batcher responded with error: {}", e);
             Err(SubmitError::GenericError(e))
         }

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

@@ -2,7 +2,7 @@ use futures_util::{stream::SplitStream, StreamExt};
 use tokio::net::TcpStream;
 use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
 
-use crate::core::{errors::SubmitError, types::ResponseMessage};
+use crate::core::{errors::SubmitError, types::SubmitProofResponseMessage};
 
 use super::serialization::cbor_deserialize;
 
@@ -13,7 +13,7 @@ pub async fn check_protocol_version(
 ) -> Result<(), SubmitError> {
     if let Some(Ok(msg)) = ws_read.next().await {
         match cbor_deserialize(msg.into_data().as_slice()) {
-            Ok(ResponseMessage::ProtocolVersion(protocol_version)) => {
+            Ok(SubmitProofResponseMessage::ProtocolVersion(protocol_version)) => {
                 if protocol_version > EXPECTED_PROTOCOL_VERSION {
                     return Err(SubmitError::ProtocolVersionMismatch {
                         current: protocol_version,

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

@@ -1,6 +1,6 @@
 use std::io::Read;
 
-use serde::{de::DeserializeOwned, Serialize};
+use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 pub fn cbor_serialize<T: Serialize>(value: &T) -> Result<Vec<u8>, SerializationError> {
     let mut buf = Vec::new();
@@ -12,7 +12,7 @@ pub fn cbor_deserialize<R: Read, T: DeserializeOwned>(buf: R) -> Result<T, Seria
     ciborium::from_reader(buf).map_err(|_| SerializationError)
 }
 
-#[derive(Debug)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct SerializationError;
 
 impl std::fmt::Display for SerializationError {

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

@@ -3,6 +3,7 @@ use ethers::providers::ProviderError;
 use ethers::signers::WalletError;
 use ethers::types::transaction::eip712::Eip712Error;
 use ethers::types::{SignatureError, H160};
+use serde::{Deserialize, Serialize};
 use std::io;
 use std::path::PathBuf;
 use tokio_tungstenite::tungstenite::protocol::CloseFrame;
@@ -15,7 +16,6 @@ use super::types::ProofInvalidReason;
 pub enum AlignedError {
     SubmitError(SubmitError),
     VerificationError(VerificationError),
-    NonceError(NonceError),
     ChainIdError(ChainIdError),
     MaxFeeEstimateError(MaxFeeEstimateError),
     FileError(FileError),
@@ -33,12 +33,6 @@ impl From<VerificationError> for AlignedError {
     }
 }
 
-impl From<NonceError> for AlignedError {
-    fn from(e: NonceError) -> Self {
-        AlignedError::NonceError(e)
-    }
-}
-
 impl From<ChainIdError> for AlignedError {
     fn from(e: ChainIdError) -> Self {
         AlignedError::ChainIdError(e)
@@ -62,7 +56,6 @@ impl fmt::Display for AlignedError {
         match self {
             AlignedError::SubmitError(e) => write!(f, "Submit error: {}", e),
             AlignedError::VerificationError(e) => write!(f, "Verification error: {}", e),
-            AlignedError::NonceError(e) => write!(f, "Nonce error: {}", e),
             AlignedError::ChainIdError(e) => write!(f, "Chain ID error: {}", e),
             AlignedError::MaxFeeEstimateError(e) => write!(f, "Max fee estimate error: {}", e),
             AlignedError::FileError(e) => write!(f, "File error: {}", e),
@@ -102,6 +95,7 @@ pub enum SubmitError {
     BatchSubmissionFailed(String),
     AddToBatchError,
     InvalidProofInclusionData,
+    GetNonceError(String),
     GenericError(String),
 }
 
@@ -216,6 +210,7 @@ impl fmt::Display for SubmitError {
             SubmitError::InvalidProofInclusionData => {
                 write!(f, "Batcher responded with invalid batch inclusion data. Can't verify your proof was correctly included in the batch.")
             }
+            SubmitError::GetNonceError(e) => write!(f, "Error while getting nonce {}", e),
         }
     }
 }
@@ -243,21 +238,14 @@ impl fmt::Display for VerificationError {
     }
 }
 
-#[derive(Debug)]
-pub enum NonceError {
-    EthereumProviderError(String),
-    EthereumCallError(String),
-}
-
-impl fmt::Display for NonceError {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self {
-            NonceError::EthereumProviderError(e) => {
-                write!(f, "Ethereum provider error: {}", e)
-            }
-            NonceError::EthereumCallError(e) => write!(f, "Ethereum call error: {}", e),
-        }
-    }
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub enum GetNonceError {
+    EthRpcError(String),
+    ConnectionFailed(String),
+    SerializationError(String),
+    UnexpectedResponse(String),
+    InvalidRequest(String),
+    ProtocolMismatch { current: u16, expected: u16 },
 }
 
 #[derive(Debug)]

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

@@ -220,11 +220,29 @@ impl BatchInclusionData {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ClientMessage {
+pub struct SubmitProofMessage {
     pub verification_data: NoncedVerificationData,
     pub signature: Signature,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum ClientMessage {
+    GetNonceForAddress(Address),
+    // Needs to be wrapped in box as the message is 3x bigger than the others
+    // see https://rust-lang.github.io/rust-clippy/master/index.html#large_enum_variant
+    SubmitProof(Box<SubmitProofMessage>),
+}
+
+impl Display for ClientMessage {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        // Use pattern matching to print the variant name
+        match self {
+            ClientMessage::GetNonceForAddress(_) => write!(f, "GetNonceForAddress"),
+            ClientMessage::SubmitProof(_) => write!(f, "SubmitProof"),
+        }
+    }
+}
+
 impl Eip712 for NoncedVerificationData {
     type Error = Eip712Error;
     fn domain(&self) -> Result<EIP712Domain, Self::Error> {
@@ -275,7 +293,7 @@ impl Eip712 for NoncedVerificationData {
     }
 }
 
-impl ClientMessage {
+impl SubmitProofMessage {
     /// Client message is a wrap around verification data and its signature.
     /// The signature is obtained by calculating the commitments and then hashing them.
     pub async fn new(
@@ -287,7 +305,7 @@ impl ClientMessage {
             .await
             .expect("Failed to sign the verification data");
 
-        ClientMessage {
+        Self {
             verification_data,
             signature,
         }
@@ -352,7 +370,7 @@ impl Display for ProofInvalidReason {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub enum ResponseMessage {
+pub enum SubmitProofResponseMessage {
     BatchInclusionData(BatchInclusionData),
     ProtocolVersion(u16),
     CreateNewTaskError(String, String), //merkle-root, error
@@ -371,6 +389,13 @@ pub enum ResponseMessage {
     InvalidPaymentServiceAddress(Address, Address),
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum GetNonceResponseMessage {
+    Nonce(U256),
+    EthRpcError(String),
+    InvalidRequest(String),
+}
+
 #[derive(Debug, Clone, Copy)]
 pub enum Network {
     Devnet,