fix: nonce handling (#669)

Author: taturosati

batcher/aligned-batcher/src/config/mod.rs

@@ -1,14 +1,5 @@
-use std::sync::Arc;
-
-use ethers::{
-    core::k256::ecdsa::SigningKey,
-    signers::{Signer, Wallet},
-    types::{Address, U256},
-};
+use ethers::{core::k256::ecdsa::SigningKey, signers::Wallet, types::Address};
 use serde::Deserialize;
-use tokio::sync::Mutex;
-
-use crate::eth::BatcherPaymentService;
 
 #[derive(Clone, Debug, Deserialize)]
 pub struct ECDSAConfig {
@@ -20,7 +11,6 @@ pub struct ECDSAConfig {
 pub struct NonPayingConfig {
     pub address: Address,
     pub replacement: Wallet<SigningKey>,
-    pub nonce: Arc<Mutex<U256>>,
 }
 
 #[derive(Debug, Deserialize)]
@@ -30,28 +20,17 @@ pub struct NonPayingConfigFromYaml {
 }
 
 impl NonPayingConfig {
-    pub async fn from_yaml_config(
-        config: NonPayingConfigFromYaml,
-        batcher_payments_contract: &BatcherPaymentService,
-    ) -> Self {
+    pub async fn from_yaml_config(config: NonPayingConfigFromYaml) -> Self {
         let replacement = Wallet::from_bytes(
             hex::decode(config.replacement_private_key)
                 .expect("Failed to decode replacement private key")
                 .as_slice(),
         )
         .expect("Failed to create replacement wallet");
 
-        let nonce = batcher_payments_contract
-            .user_nonces(replacement.address())
-            .await
-            .expect("Failed to get nonce");
-
-        let nonce = Arc::new(Mutex::new(nonce));
-
         NonPayingConfig {
             address: config.address,
             replacement,
-            nonce,
         }
     }
 }

batcher/aligned-batcher/src/lib.rs

@@ -4,15 +4,17 @@ use aligned_sdk::eth::batcher_payment_service::SignatureData;
 use config::NonPayingConfig;
 use dotenv::dotenv;
 use ethers::signers::Signer;
+use serde::Serialize;
 
+use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::env;
 use std::net::SocketAddr;
 use std::sync::Arc;
 
 use aligned_sdk::core::types::{
     BatchInclusionData, ClientMessage, NoncedVerificationData, ResponseMessage,
-    VerificationCommitmentBatch, VerificationDataCommitment,
+    ValidityResponseMessage, VerificationCommitmentBatch, VerificationDataCommitment,
 };
 use aws_sdk_s3::client::Client as S3Client;
 use eth::BatcherPaymentService;
@@ -106,7 +108,7 @@ impl Batcher {
         let non_paying_config = if let Some(non_paying_config) = config.batcher.non_paying {
             warn!("Non-paying address configuration detected. Will replace non-paying address {} with configured address.",
                 non_paying_config.address);
-            Some(NonPayingConfig::from_yaml_config(non_paying_config, &payment_service).await)
+            Some(NonPayingConfig::from_yaml_config(non_paying_config).await)
         } else {
             None
         };
@@ -220,9 +222,9 @@ impl Batcher {
                     .await;
             } else {
                 if !self.check_user_balance(&addr).await {
-                    send_error_message(
+                    send_message(
                         ws_conn_sink.clone(),
-                        ResponseMessage::InsufficientBalanceError(addr),
+                        ValidityResponseMessage::InsufficientBalance(addr),
                     )
                     .await;
 
@@ -237,21 +239,15 @@ impl Batcher {
                         && !zk_utils::verify(&nonced_verification_data.verification_data)
                     {
                         error!("Invalid proof detected. Verification failed.");
-                        send_error_message(
-                            ws_conn_sink.clone(),
-                            ResponseMessage::VerificationError(),
-                        )
-                        .await;
+                        send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidProof)
+                            .await;
                         return Ok(()); // Send error message to the client and return
                     }
 
                     // Doing nonce verification after proof verification to avoid unnecessary nonce increment
                     if !self.check_nonce_and_increment(addr, nonce).await {
-                        send_error_message(
-                            ws_conn_sink.clone(),
-                            ResponseMessage::InvalidNonceError,
-                        )
-                        .await;
+                        send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidNonce)
+                            .await;
                         return Ok(()); // Send error message to the client and return
                     }
 
@@ -263,20 +259,21 @@ impl Batcher {
                     .await;
                 } else {
                     error!("Proof is too large");
-                    send_error_message(ws_conn_sink.clone(), ResponseMessage::ProofTooLargeError())
+                    send_message(ws_conn_sink.clone(), ValidityResponseMessage::ProofTooLarge)
                         .await;
                     return Ok(()); // Send error message to the client and return
                 };
 
                 info!("Verification data message handled");
 
+                send_message(ws_conn_sink, ValidityResponseMessage::Valid).await;
                 return Ok(());
             }
         } else {
             error!("Signature verification error");
-            send_error_message(
+            send_message(
                 ws_conn_sink.clone(),
-                ResponseMessage::SignatureVerificationError(),
+                ValidityResponseMessage::InvalidSignature,
             )
             .await;
             Ok(()) // Send error message to the client and return
@@ -482,12 +479,15 @@ impl Batcher {
         {
             for (_, _, ws_sink, _) in finalized_batch.iter() {
                 let merkle_root = hex::encode(batch_merkle_tree.root);
-                send_error_message(
+                send_message(
                     ws_sink.clone(),
                     ResponseMessage::CreateNewTaskError(merkle_root),
                 )
                 .await
             }
+
+            self.flush_queue_and_clear_nonce_cache().await;
+
             return Err(e);
         };
 
@@ -496,6 +496,22 @@ impl Batcher {
         Ok(())
     }
 
+    async fn flush_queue_and_clear_nonce_cache(&self) {
+        warn!("Resetting state... Flushing queue and nonces");
+
+        let mut batch_queue = self.batch_queue.lock().await;
+        let mut user_nonces = self.user_nonces.lock().await;
+        let mut user_proof_count_in_batch = self.user_proof_count_in_batch.lock().await;
+
+        for (_, _, ws_sink, _) in batch_queue.iter() {
+            send_message(ws_sink.clone(), ResponseMessage::BatchReset).await;
+        }
+
+        batch_queue.clear();
+        user_nonces.clear();
+        user_proof_count_in_batch.clear();
+    }
+
     /// Receives new block numbers, checks if conditions are met for submission and
     /// finalizes the batch.
     async fn handle_new_block(&self, block_number: u64) -> Result<(), BatcherError> {
@@ -576,24 +592,8 @@ impl Batcher {
         client_msg: ClientMessage,
     ) -> Result<(), Error> {
         let non_paying_config = self.non_paying_config.as_ref().unwrap();
-
-        // The nonpaying nonce is locked through the entire message processing so that
-        // another incoming connections using the nonpaying address don't desync its nonce
-        let mut nonpaying_nonce = non_paying_config.nonce.lock().await;
         let addr = non_paying_config.replacement.address();
 
-        let mut nonce_bytes = [0u8; 32];
-        nonpaying_nonce.to_big_endian(&mut nonce_bytes);
-        *nonpaying_nonce += U256::one();
-
-        let verifcation_data = NoncedVerificationData::new(
-            client_msg.verification_data.verification_data.clone(),
-            nonce_bytes,
-        );
-
-        let client_msg =
-            ClientMessage::new(verifcation_data, non_paying_config.replacement.clone());
-
         let user_balance = self
             .payment_service
             .user_balances(addr)
@@ -603,32 +603,57 @@ impl Batcher {
 
         if user_balance == U256::from(0) {
             error!("Insufficient funds for address {:?}", addr);
-            send_error_message(
+            send_message(
                 ws_conn_sink.clone(),
-                ResponseMessage::InsufficientBalanceError(addr),
+                ValidityResponseMessage::InsufficientBalance(addr),
             )
             .await;
             return Ok(()); // Send error message to the client and return
         }
 
-        let nonce = U256::from_big_endian(client_msg.verification_data.nonce.as_slice());
-        let nonced_verification_data = client_msg.verification_data;
-        if nonced_verification_data.verification_data.proof.len() <= self.max_proof_size {
+        if client_msg.verification_data.verification_data.proof.len() <= self.max_proof_size {
             // When pre-verification is enabled, batcher will verify proofs for faster feedback with clients
             if self.pre_verification_is_enabled
-                && !zk_utils::verify(&nonced_verification_data.verification_data)
+                && !zk_utils::verify(&client_msg.verification_data.verification_data)
             {
                 error!("Invalid proof detected. Verification failed.");
-                send_error_message(ws_conn_sink.clone(), ResponseMessage::VerificationError())
-                    .await;
+                send_message(ws_conn_sink.clone(), ValidityResponseMessage::InvalidProof).await;
                 return Ok(()); // Send error message to the client and return
             }
 
-            // Doing nonce verification after proof verification to avoid unnecessary nonce increment
-            if !self.check_nonce_and_increment(addr, nonce).await {
-                send_error_message(ws_conn_sink.clone(), ResponseMessage::InvalidNonceError).await;
-                return Ok(()); // Send error message to the client and return
-            }
+            let nonced_verification_data = {
+                let mut user_nonces = self.user_nonces.lock().await;
+
+                let nonpaying_nonce = match user_nonces.entry(addr) {
+                    Entry::Occupied(o) => o.into_mut(),
+                    Entry::Vacant(vacant) => {
+                        let nonce = self
+                            .payment_service
+                            .user_nonces(addr)
+                            .call()
+                            .await
+                            .expect("Failed to get nonce");
+
+                        vacant.insert(nonce)
+                    }
+                };
+
+                debug!("non paying nonce: {:?}", nonpaying_nonce);
+
+                let mut nonce_bytes = [0u8; 32];
+                nonpaying_nonce.to_big_endian(&mut nonce_bytes);
+                *nonpaying_nonce += U256::one();
+
+                NoncedVerificationData::new(
+                    client_msg.verification_data.verification_data.clone(),
+                    nonce_bytes,
+                )
+            };
+
+            let client_msg = ClientMessage::new(
+                nonced_verification_data.clone(),
+                non_paying_config.replacement.clone(),
+            );
 
             self.clone()
                 .add_to_batch(
@@ -639,12 +664,13 @@ impl Batcher {
                 .await;
         } else {
             error!("Proof is too large");
-            send_error_message(ws_conn_sink.clone(), ResponseMessage::ProofTooLargeError()).await;
+            send_message(ws_conn_sink.clone(), ValidityResponseMessage::ProofTooLarge).await;
             return Ok(()); // Send error message to the client and return
         };
 
         info!("Verification data message handled");
 
+        send_message(ws_conn_sink, ValidityResponseMessage::Valid).await;
         Ok(())
     }
 
@@ -687,18 +713,17 @@ async fn send_batch_inclusion_data_responses(
         .await;
 }
 
-async fn send_error_message(
+async fn send_message<T: Serialize>(
     ws_conn_sink: Arc<RwLock<SplitSink<WebSocketStream<TcpStream>, Message>>>,
-    error_message: ResponseMessage,
+    message: T,
 ) {
-    let serialized_response =
-        serde_json::to_vec(&error_message).expect("Could not serialize response");
+    let serialized_response = serde_json::to_vec(&message).expect("Could not serialize response");
 
     // Send error message
     ws_conn_sink
         .write()
         .await
         .send(Message::binary(serialized_response))
         .await
-        .expect("Failed to send error message");
+        .expect("Failed to send message");
 }