fix: address 9 audit findings across correctness, safety, and performance

src/client.rs

@@ -873,25 +873,36 @@ impl Client {
             if let Err(connect_err) = self.connect().await {
                 error!("Failed to connect: {connect_err:#}. Will retry...");
             } else {
-                if self.read_messages_loop().await.is_err() {
+                let unexpected_disconnect = if self.read_messages_loop().await.is_err() {
                     // Check intentional_reconnect AFTER read loop exits — reconnect()
                     // sets this flag while the loop is running, so it must be read here.
                     if self.expected_disconnect.load(Ordering::Relaxed)
                         || self.intentional_reconnect.swap(false, Ordering::Relaxed)
                     {
                         debug!("Message loop exited during expected disconnect.");
+                        false
                     } else {
                         warn!(
                             "Message loop exited with an error. Will attempt to reconnect if enabled."
                         );
+                        true
                     }
                 } else if self.expected_disconnect.load(Ordering::Relaxed) {
                     debug!("Message loop exited gracefully (expected disconnect).");
+                    false
                 } else {
                     info!("Message loop exited gracefully.");
-                }
+                    false
+                };
 
                 self.cleanup_connection_state().await;
+
+                // Dispatch after cleanup so handlers see cleared connection state.
+                if unexpected_disconnect {
+                    self.core
+                        .event_bus
+                        .dispatch(&Event::Disconnected(crate::types::events::Disconnected));
+                }
             }
 
             if !self.enable_auto_reconnect.load(Ordering::Relaxed) {
@@ -1255,7 +1266,7 @@ impl Client {
                             Ok(crate::transport::TransportEvent::DataReceived(data)) => {
                                 // Update dead-socket timer (WA Web: deadSocketTimer reset)
                                 self.last_data_received_ms.store(
-                                    wacore::time::now_millis() as u64,
+                                    wacore::time::now_millis().max(0) as u64,
                                     Ordering::Relaxed,
                                 );
 
@@ -1308,9 +1319,7 @@ impl Client {
                                 }
                             },
                             Ok(crate::transport::TransportEvent::Disconnected) | Err(_) => {
-                                self.cleanup_connection_state().await;
-                                 if !self.expected_disconnect.load(Ordering::Relaxed) {
-                                    self.core.event_bus.dispatch(&Event::Disconnected(crate::types::events::Disconnected));
+                                if !self.expected_disconnect.load(Ordering::Relaxed) {
                                     debug!("Transport disconnected unexpectedly.");
                                     return Err(anyhow::anyhow!("Transport disconnected unexpectedly"));
                                 } else {
@@ -3318,7 +3327,7 @@ impl Client {
 
         // WA Web: callStanza → deadSocketTimer.onOrBefore(deadSocketTime, socketId)
         self.last_data_sent_ms
-            .store(wacore::time::now_millis() as u64, Ordering::Relaxed);
+            .store(wacore::time::now_millis().max(0) as u64, Ordering::Relaxed);
 
         Ok(())
     }

src/handlers/message.rs

@@ -58,9 +58,9 @@ impl StanzaHandler for MessageHandler {
         let tx = client
             .message_queues
             .get_with_by_ref(&chat_id, async {
-                // Create a channel with backpressure
-                // Increased capacity to handle high message rates without blocking
-                let (tx, rx) = async_channel::bounded::<Arc<Node>>(10000);
+                // Bounded capacity provides backpressure to prevent unbounded memory growth.
+                // 500 is enough for burst handling while limiting per-chat memory.
+                let (tx, rx) = async_channel::bounded::<Arc<Node>>(500);
 
                 let client_for_worker = client.clone();
 

src/message.rs

@@ -132,7 +132,10 @@ impl Client {
     /// Increments the retry count for a message and returns the new count.
     /// Returns `None` if max retries have been reached.
     ///
-    /// Uses get + insert for portability across cache backends.
+    /// Note: get-then-insert has a theoretical TOCTOU window since
+    /// `spawn_retry_receipt` detaches. In practice, retries for the same
+    /// message are rare and a double-send is benign (recipients deduplicate
+    /// by message ID).
     async fn increment_retry_count(&self, cache_key: &str) -> Option<u8> {
         let current = self.message_retry_counts.get(&cache_key.to_string()).await;
         match current {
@@ -537,10 +540,10 @@ impl Client {
             session_enc_nodes.len()
         );
 
-        // Skip session processing for group senders (@c.us, @g.us, @broadcast)
-        // Groups don't use 1:1 Signal Protocol sessions
-        let is_group_sender = sender_encryption_jid.server.contains(".us")
-            || sender_encryption_jid.server.contains("broadcast");
+        // Skip session processing for group/broadcast JIDs — they use sender keys, not 1:1 sessions.
+        let is_group_sender = sender_encryption_jid.is_group()
+            || sender_encryption_jid.is_broadcast_list()
+            || sender_encryption_jid.is_status_broadcast();
 
         let (
             session_decrypted_successfully,
@@ -1034,7 +1037,7 @@ impl Client {
         enc_nodes: &[&wacore_binary::node::Node],
         info: &MessageInfo,
         _sender_encryption_jid: &Jid,
-        _decrypt_fail_mode: crate::types::events::DecryptFailMode,
+        decrypt_fail_mode: crate::types::events::DecryptFailMode,
     ) -> Result<(), DecryptionError> {
         if enc_nodes.is_empty() {
             return Ok(());
@@ -1114,6 +1117,7 @@ impl Client {
                         "No sender key state for group message [msg:{}] from {}: {}. Sending retry receipt.",
                         info.id, info.source.sender, msg
                     );
+                    self.dispatch_undecryptable_event(info, decrypt_fail_mode);
                     self.spawn_retry_receipt(info, RetryReason::NoSession);
                 }
                 Err(e) => {

src/send.rs

@@ -871,15 +871,15 @@ impl Client {
             }
 
             let force_skdm = {
-                use wacore::libsignal::protocol::SenderKeyStore;
                 use wacore::libsignal::store::sender_key_name::SenderKeyName;
-                let mut device_guard = device_store_arc.write().await;
                 let sender_address = own_sending_jid.to_protocol_address();
                 let sender_key_name =
                     SenderKeyName::new(to_str.clone(), sender_address.to_string());
 
-                let key_exists = device_guard
-                    .load_sender_key(&sender_key_name)
+                let device_guard = device_store_arc.read().await;
+                let key_exists = self
+                    .signal_cache
+                    .get_sender_key(&sender_key_name, &*device_guard.backend)
                     .await?
                     .is_some();
 

wacore/src/protocol/keepalive.rs

@@ -23,7 +23,7 @@ pub fn ms_since(timestamp_ms: u64) -> Option<u64> {
     if timestamp_ms == 0 {
         return None;
     }
-    let now = crate::time::now_millis() as u64;
+    let now = crate::time::now_millis().max(0) as u64;
     Some(now.saturating_sub(timestamp_ms))
 }
 
@@ -61,14 +61,14 @@ mod tests {
 
     #[test]
     fn ms_since_recent() {
-        let now_ms = crate::time::now_millis() as u64;
+        let now_ms = crate::time::now_millis().max(0) as u64;
         let elapsed = ms_since(now_ms).unwrap();
         assert!(elapsed < 100, "should be near-zero, got {elapsed}ms");
     }
 
     #[test]
     fn ms_since_stale() {
-        let thirty_sec_ago = (crate::time::now_millis() as u64).saturating_sub(30_000);
+        let thirty_sec_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(30_000);
         let elapsed = ms_since(thirty_sec_ago).unwrap();
         assert!(
             (29_000..=31_000).contains(&elapsed),
@@ -85,33 +85,33 @@ mod tests {
 
     #[test]
     fn dead_socket_received_after_send() {
-        let t = crate::time::now_millis() as u64;
+        let t = crate::time::now_millis().max(0) as u64;
         assert!(!is_dead_socket(t, t + 1));
     }
 
     #[test]
     fn dead_socket_sent_recently() {
-        let now = crate::time::now_millis() as u64;
+        let now = crate::time::now_millis().max(0) as u64;
         assert!(!is_dead_socket(now, 0));
     }
 
     #[test]
     fn dead_socket_sent_long_ago_no_reply() {
-        let thirty_ago = (crate::time::now_millis() as u64).saturating_sub(30_000);
+        let thirty_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(30_000);
         assert!(is_dead_socket(thirty_ago, 0));
     }
 
     #[test]
     fn dead_socket_sent_long_ago_old_reply() {
-        let thirty_ago = (crate::time::now_millis() as u64).saturating_sub(30_000);
+        let thirty_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(30_000);
         let thirty_one_ago = thirty_ago.saturating_sub(1_000);
         assert!(is_dead_socket(thirty_ago, thirty_one_ago));
     }
 
     #[test]
     fn dead_socket_sent_long_ago_recent_reply() {
-        let thirty_ago = (crate::time::now_millis() as u64).saturating_sub(30_000);
-        let one_ago = (crate::time::now_millis() as u64).saturating_sub(1_000);
+        let thirty_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(30_000);
+        let one_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(1_000);
         assert!(!is_dead_socket(thirty_ago, one_ago));
     }
 

wacore/src/store/signal_cache.rs

@@ -294,73 +294,95 @@ impl SignalStoreCache {
 
     // === Flush ===
 
-    /// Flush all dirty state to the backend in a single batch.
-    /// Acquires all 3 mutexes to ensure consistency (matches WhatsApp Web's pattern).
+    /// Flush all dirty state to the backend.
     ///
-    /// Sessions are serialized here (not on every store_session call).
-    /// Dirty sets are only cleared after ALL writes succeed.
+    /// Each store (sessions, identities, sender_keys) is flushed independently
+    /// under its own lock. This means:
+    /// - Only ONE store is locked during its I/O — the other two are free for
+    ///   concurrent encrypt/decrypt operations.
+    /// - No race between snapshot and clear — the lock is held throughout, so
+    ///   mutations to the same store are blocked until the flush completes.
+    /// - Dirty sets are cleared only after successful writes.
     pub async fn flush(&self, backend: &dyn SignalStore) -> Result<()> {
-        let mut sessions = self.sessions.lock().await;
-        let mut identities = self.identities.lock().await;
-        let mut sender_keys = self.sender_keys.lock().await;
-
-        // Snapshot dirty/deleted sets WITHOUT draining — preserve on failure
-        let session_dirty: Vec<_> = sessions.dirty.iter().cloned().collect();
-        let session_deleted: Vec<_> = sessions.deleted.iter().cloned().collect();
-        let identity_dirty: Vec<_> = identities.dirty.iter().cloned().collect();
-        let identity_deleted: Vec<_> = identities.deleted.iter().cloned().collect();
-        let sender_key_dirty: Vec<_> = sender_keys.dirty.iter().cloned().collect();
-
-        // Persist dirty sessions — serialize only here, not on every store_session
-        for address in &session_dirty {
-            if let Some(Some(record)) = sessions.cache.get(address.as_ref()) {
-                let bytes = record
-                    .serialize()
-                    .map_err(|e| anyhow::anyhow!("session serialize for {address}: {e}"))?;
-                backend.put_session(address, &bytes).await?;
+        // Flush sessions
+        {
+            let mut state = self.sessions.lock().await;
+            let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect();
+            let deleted_keys: Vec<_> = state.deleted.iter().cloned().collect();
+
+            for address in &dirty_keys {
+                if let Some(Some(record)) = state.cache.get(address.as_ref()) {
+                    let bytes = record
+                        .serialize()
+                        .map_err(|e| anyhow::anyhow!("session serialize for {address}: {e}"))?;
+                    backend.put_session(address, &bytes).await?;
+                }
+            }
+            for address in &deleted_keys {
+                backend.delete_session(address).await?;
             }
-        }
-        for address in &session_deleted {
-            backend.delete_session(address).await?;
-        }
 
-        for address in &identity_dirty {
-            if let Some(Some(data)) = identities.cache.get(address.as_ref()) {
-                let key: [u8; 32] = data.as_ref().try_into().map_err(|_| {
-                    anyhow::anyhow!(
-                        "Corrupted identity key for {address}: expected 32 bytes, got {}",
-                        data.len()
-                    )
-                })?;
-                backend.put_identity(address, key).await?;
+            for key in &dirty_keys {
+                state.dirty.remove(key);
+            }
+            for key in &deleted_keys {
+                state.deleted.remove(key);
             }
-        }
-        for address in &identity_deleted {
-            backend.delete_identity(address).await?;
         }
 
-        for name in &sender_key_dirty {
-            match sender_keys.cache.get(name.as_ref()) {
-                Some(Some(record)) => {
-                    let bytes = record
-                        .serialize()
-                        .map_err(|e| anyhow::anyhow!("sender key serialize for {name}: {e}"))?;
-                    backend.put_sender_key(name, &bytes).await?;
-                }
-                Some(None) => {
-                    // Deleted via delete_sender_key — propagate to backend
-                    backend.delete_sender_key(name).await?;
+        // Flush identities
+        {
+            let mut state = self.identities.lock().await;
+            let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect();
+            let deleted_keys: Vec<_> = state.deleted.iter().cloned().collect();
+
+            for address in &dirty_keys {
+                if let Some(Some(data)) = state.cache.get(address.as_ref()) {
+                    let key: [u8; 32] = data.as_ref().try_into().map_err(|_| {
+                        anyhow::anyhow!(
+                            "Corrupted identity key for {address}: expected 32 bytes, got {}",
+                            data.len()
+                        )
+                    })?;
+                    backend.put_identity(address, key).await?;
                 }
-                None => {}
+            }
+            for address in &deleted_keys {
+                backend.delete_identity(address).await?;
+            }
+
+            for key in &dirty_keys {
+                state.dirty.remove(key);
+            }
+            for key in &deleted_keys {
+                state.deleted.remove(key);
             }
         }
 
-        // All writes succeeded — clear dirty sets (matches WA Web's clearDirty())
-        sessions.dirty.clear();
-        sessions.deleted.clear();
-        identities.dirty.clear();
-        identities.deleted.clear();
-        sender_keys.dirty.clear();
+        // Flush sender keys
+        {
+            let mut state = self.sender_keys.lock().await;
+            let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect();
+
+            for name in &dirty_keys {
+                match state.cache.get(name.as_ref()) {
+                    Some(Some(record)) => {
+                        let bytes = record
+                            .serialize()
+                            .map_err(|e| anyhow::anyhow!("sender key serialize for {name}: {e}"))?;
+                        backend.put_sender_key(name, &bytes).await?;
+                    }
+                    Some(None) => {
+                        backend.delete_sender_key(name).await?;
+                    }
+                    None => {}
+                }
+            }
+
+            for key in &dirty_keys {
+                state.dirty.remove(key);
+            }
+        }
 
         Ok(())
     }