rearrange logic for user_balance check + subtract original_msg_max_fee from computed replacement balance

Author: PatStiles

batcher/aligned-batcher/src/lib.rs

@@ -745,17 +745,6 @@ impl Batcher {
             return Ok(());
         };
 
-        if !self.verify_user_has_enough_balance(user_balance, user_accumulated_fee, msg_max_fee) {
-            std::mem::drop(batch_state_lock);
-            send_message(
-                ws_conn_sink.clone(),
-                SubmitProofResponseMessage::InsufficientBalance(addr),
-            )
-            .await;
-            self.metrics.user_error(&["insufficient_balance", ""]);
-            return Ok(());
-        }
-
         let cached_user_nonce = batch_state_lock.get_user_nonce(&addr).await;
         let Some(expected_nonce) = cached_user_nonce else {
             error!("Failed to get cached user nonce: User not found in user states, but it should have been already inserted");
@@ -791,12 +780,26 @@ impl Batcher {
                 ws_conn_sink.clone(),
                 client_msg.signature,
                 addr,
+                user_balance,
+                user_accumulated_fee,
+                msg_max_fee
             )
             .await;
 
             return Ok(());
         }
 
+        if !self.verify_user_has_enough_balance(user_balance, user_accumulated_fee, msg_max_fee) {
+            std::mem::drop(batch_state_lock);
+            send_message(
+                ws_conn_sink.clone(),
+                SubmitProofResponseMessage::InsufficientBalance(addr),
+            )
+            .await;
+            self.metrics.user_error(&["insufficient_balance", ""]);
+            return Ok(());
+        }
+
         if msg_max_fee > user_last_max_fee_limit {
             std::mem::drop(batch_state_lock);
             warn!("Invalid max fee for address {addr}, had fee limit of {user_last_max_fee_limit:?}, sent {msg_max_fee:?}");
@@ -863,6 +866,9 @@ impl Batcher {
         ws_conn_sink: WsMessageSink,
         signature: Signature,
         addr: Address,
+        user_balance: U256,
+        user_accumulated_fee: U256,
+        msg_max_fee: U256,
     ) {
         let replacement_max_fee = nonced_verification_data.max_fee;
         let nonce = nonced_verification_data.nonce;
@@ -892,6 +898,18 @@ impl Batcher {
             return;
         }
 
+        // For a replacement msg we must subtract the original_max_fee of the message from the user's accumulated max fee. 
+        if !self.verify_user_has_enough_balance(user_balance, user_accumulated_fee - original_max_fee, msg_max_fee) {
+            std::mem::drop(batch_state_lock);
+            send_message(
+                ws_conn_sink.clone(),
+                SubmitProofResponseMessage::InsufficientBalance(addr),
+            )
+            .await;
+            self.metrics.user_error(&["insufficient_balance", ""]);
+            return;
+        }
+
         info!("Replacing message for address {addr} with nonce {nonce} and max fee {replacement_max_fee}");
 
         // The replacement entry is built from the old entry and validated for then to be replaced