From 23ef08c178453661e72afe6995b4da64f88fe4b9 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 9 Jun 2026 19:46:16 +0000 Subject: [PATCH 1/3] perf(store): cache the device snapshot as Arc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit get_device_snapshot() cloned the whole Device on every call — Jids, push_name, props_hash, edge_routing_info and nct_salt heap allocations — and it is called on every inbound message plus ~70 other sites. The snapshot is now an Arc rebuilt under the device write guard in modify_device (the single mutation funnel), so reads become a refcount bump with no lock against writers and no clone, and can never observe a stale snapshot relative to a committed mutation. Fallout, in the same spirit of avoiding needless work: - get_device_snapshot() and the get_pn/get_lid/get_push_name/require_pn accessors are now sync (no async lock dance); call sites updated. - Read-only get_device_arc() users (retry, sessions, prekeys, signal feature, send force-skdm checks, receive parse_message_info) switch to the snapshot — each was paying an RwLock read just to reach .backend or clone pn/lid. - Send/receive paths borrow pn/lid/account from the held snapshot instead of cloning them; clones remain only where ownership leaves the scope. - Tests that mutated Device through the raw write lock now go through modify_device — direct mutation would bypass the cached snapshot, so the docs now state that invariant explicitly (AGENTS.md, agent_docs). Breaking (pre-1.0): get_device_snapshot returns Arc and is no longer async; the three accessors are no longer async; get_device_arc is documented as store-adapter-only. --- AGENTS.md | 2 +- agent_docs/feature_implementation.md | 2 +- src/bot.rs | 10 ++--- src/client/accessors.rs | 39 +++++++------------ src/client/app_state.rs | 4 +- src/client/device_registry.rs | 10 ++--- src/client/iq_ops.rs | 9 +---- src/client/lifecycle.rs | 6 +-- src/client/messaging.rs | 5 +-- src/client/node_io.rs | 15 ++++---- src/client/sender_keys.rs | 4 +- src/client/sessions.rs | 34 ++++++++--------- src/features/contacts.rs | 2 +- src/features/events.rs | 3 +- src/features/media_reupload.rs | 7 +++- src/features/polls.rs | 3 +- src/features/presence.rs | 12 ++---- src/features/signal.rs | 11 ++---- src/handlers/call.rs | 4 +- src/handlers/notification/device.rs | 7 ++-- src/handshake.rs | 2 +- src/history_sync.rs | 5 +-- src/message/msg_secret.rs | 20 +++++----- src/message/receive.rs | 13 +++---- src/message/special.rs | 2 +- src/message/tests.rs | 56 +++++++++------------------- src/pair.rs | 6 +-- src/pair_code.rs | 4 +- src/pdo.rs | 5 +-- src/prekeys.rs | 44 +++++++++------------- src/receipt.rs | 2 +- src/request.rs | 2 +- src/retry.rs | 48 ++++++++++-------------- src/send.rs | 41 ++++++++++---------- src/store/persistence_manager.rs | 27 +++++++++++++- src/usync.rs | 2 +- src/version.rs | 2 +- tests/e2e/src/lib.rs | 7 +--- tests/e2e/tests/app_state.rs | 16 +++----- tests/e2e/tests/chat_actions.rs | 13 ------- tests/e2e/tests/community.rs | 4 -- tests/e2e/tests/digest_key.rs | 2 +- tests/e2e/tests/groups.rs | 1 - tests/e2e/tests/lid_sessions.rs | 14 +++---- tests/e2e/tests/media.rs | 9 ----- tests/e2e/tests/memory_soak.rs | 20 +++++----- tests/e2e/tests/prekey_sessions.rs | 1 - tests/e2e/tests/privacy_tokens.rs | 13 +------ tests/e2e/tests/profile.rs | 21 +++++------ tests/e2e/tests/profile_picture.rs | 6 --- tests/e2e/tests/session_reuse.rs | 4 +- tests/handshake_integration.rs | 18 ++++----- 52 files changed, 254 insertions(+), 365 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 869fe20db..aefae4b20 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -24,7 +24,7 @@ cargo test -p e2e-tests # requires mock server running ## Critical Conventions -- **State**: Never modify Device state directly. Use `DeviceCommand` + `PersistenceManager::process_command()`. Read via `get_device_snapshot()`. +- **State**: Never modify Device state directly (not even in tests — a write-lock mutation bypasses the cached snapshot). Use `DeviceCommand` + `PersistenceManager::process_command()` (or `modify_device` internally). Read via `get_device_snapshot()` — it returns a cached `Arc` (sync, refcount-cheap, safe to call per message); borrow fields from the held snapshot instead of cloning them. `get_device_arc()` is only for store adapters that need `&mut Device` trait access. - **Async**: All I/O uses Tokio. Wrap blocking I/O (`ureq`) and heavy CPU work in `tokio::task::spawn_blocking`. - **Concurrency**: `session_locks` serializes per-sender Signal encrypt/decrypt. `message_enqueue_locks` serializes per-chat incoming message processing. Outgoing sends are not per-chat locked (matches WA Web). - **Errors**: `thiserror` for typed errors, `anyhow` for multi-failure functions. No `.unwrap()` outside tests. diff --git a/agent_docs/feature_implementation.md b/agent_docs/feature_implementation.md index 2f92cec37..b6a7ba7ca 100644 --- a/agent_docs/feature_implementation.md +++ b/agent_docs/feature_implementation.md @@ -20,7 +20,7 @@ When adding a new feature, follow this flow that mirrors WhatsApp Web behavior w 4. **Keep state changes behind the PersistenceManager** - Use `DeviceCommand` + `PersistenceManager::process_command()` for mutations - - Use `get_device_snapshot()` for read access + - Use `get_device_snapshot()` for read access — sync, returns a cached `Arc` (refcount bump, no Device clone, no lock); hold it and borrow fields rather than cloning them 5. **Confirm concurrency requirements** - Network I/O stays async diff --git a/src/bot.rs b/src/bot.rs index f42cd18b5..24eb58466 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -974,7 +974,7 @@ mod tests { let client = bot.client(); let persistence_manager = client.persistence_manager(); - let device = persistence_manager.get_device_snapshot().await; + let device = persistence_manager.get_device_snapshot(); // Verify the device props were overridden assert_eq!(device.device_props.os, Some(custom_os)); @@ -1001,7 +1001,7 @@ mod tests { let client = bot.client(); let persistence_manager = client.persistence_manager(); - let device = persistence_manager.get_device_snapshot().await; + let device = persistence_manager.get_device_snapshot(); // Verify only OS was overridden, version should be default assert_eq!(device.device_props.os, Some(custom_os)); @@ -1037,7 +1037,7 @@ mod tests { let client = bot.client(); let persistence_manager = client.persistence_manager(); - let device = persistence_manager.get_device_snapshot().await; + let device = persistence_manager.get_device_snapshot(); // Verify only version was overridden, OS should be default ("rust") assert_eq!(device.device_props.version, Some(custom_version)); @@ -1069,7 +1069,7 @@ mod tests { let client = bot.client(); let persistence_manager = client.persistence_manager(); - let device = persistence_manager.get_device_snapshot().await; + let device = persistence_manager.get_device_snapshot(); // Verify platform type was set to Chrome assert_eq!( @@ -1119,7 +1119,7 @@ mod tests { let client = bot.client(); let persistence_manager = client.persistence_manager(); - let device = persistence_manager.get_device_snapshot().await; + let device = persistence_manager.get_device_snapshot(); // Verify all device props were overridden assert_eq!(device.device_props.os, Some(custom_os)); diff --git a/src/client/accessors.rs b/src/client/accessors.rs index 151ed2252..9ffb85ec1 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -114,38 +114,27 @@ impl Client { self.persistence_manager.clone() } - pub async fn get_push_name(&self) -> String { + // The owned returns below are the only clones left: the snapshot read + // itself is an Arc refcount bump (no lock against writers). Callers that + // only need a borrow can hold `persistence_manager().get_device_snapshot()` + // and read fields directly. + pub fn get_push_name(&self) -> String { self.persistence_manager - .get_device_arc() - .await - .read() - .await + .get_device_snapshot() .push_name .clone() } - pub async fn get_pn(&self) -> Option { - self.persistence_manager - .get_device_arc() - .await - .read() - .await - .pn - .clone() + pub fn get_pn(&self) -> Option { + self.persistence_manager.get_device_snapshot().pn.clone() } - pub async fn get_lid(&self) -> Option { - self.persistence_manager - .get_device_arc() - .await - .read() - .await - .lid - .clone() + pub fn get_lid(&self) -> Option { + self.persistence_manager.get_device_snapshot().lid.clone() } - pub(crate) async fn require_pn(&self) -> Result { - self.get_pn().await.ok_or(ClientError::NotLoggedIn.into()) + pub(crate) fn require_pn(&self) -> Result { + self.get_pn().ok_or(ClientError::NotLoggedIn.into()) } /// Resolve our own JID for a group, respecting its addressing mode. @@ -156,7 +145,7 @@ impl Client { &self, group_jid: &Jid, ) -> Result { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let own_pn = device_snapshot .pn .clone() @@ -178,7 +167,7 @@ impl Client { } pub(crate) async fn update_push_name_and_notify(self: &Arc, new_name: String) { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let old_name = device_snapshot.push_name.clone(); if old_name == new_name { diff --git a/src/client/app_state.rs b/src/client/app_state.rs index 041e4e0c0..43d53b16f 100644 --- a/src/client/app_state.rs +++ b/src/client/app_state.rs @@ -579,7 +579,7 @@ impl Client { if raw_key_ids.is_empty() { return Ok(()); } - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let own_jid = match device_snapshot.pn.clone() { Some(j) => j, None => { @@ -721,7 +721,7 @@ impl Client { let new_name = new_name.clone(); let bus = self.core.event_bus.clone(); - let snapshot = self.persistence_manager.get_device_snapshot().await; + let snapshot = self.persistence_manager.get_device_snapshot(); let old = snapshot.push_name.clone(); if old != new_name { debug!(target: "Client/AppState", "Persisting push name from app state mutation: '{}' (old='{}')", new_name, old); diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index 6cf902b48..49d63ddd0 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -2059,11 +2059,10 @@ mod tests { .rotate_sender_key_on_participant_remove(group, &["271060335329480"]) .await; - let device_arc = client.persistence_manager.get_device_arc().await; - let device = device_arc.read().await; + let device_snapshot = client.persistence_manager.get_device_snapshot(); let key = client .signal_cache - .get_sender_key(&sk_name, &*device.backend) + .get_sender_key(&sk_name, &*device_snapshot.backend) .await .unwrap(); assert!( @@ -2114,11 +2113,10 @@ mod tests { .rotate_sender_key_on_participant_remove(group, &["271060335329480"]) .await; - let device_arc = client.persistence_manager.get_device_arc().await; - let device = device_arc.read().await; + let device_snapshot = client.persistence_manager.get_device_snapshot(); let key = client .signal_cache - .get_sender_key(&sk_name, &*device.backend) + .get_sender_key(&sk_name, &*device_snapshot.backend) .await .unwrap(); assert!( diff --git a/src/client/iq_ops.rs b/src/client/iq_ops.rs index 1e16cf897..e7cf17a3d 100644 --- a/src/client/iq_ops.rs +++ b/src/client/iq_ops.rs @@ -15,7 +15,6 @@ impl Client { let stored_hash = self .persistence_manager .get_device_snapshot() - .await .props_hash .clone(); @@ -192,13 +191,7 @@ impl Client { if override_.is_empty() { return; } - if self - .persistence_manager - .get_device_snapshot() - .await - .pn - .is_some() - { + if self.persistence_manager.get_device_snapshot().pn.is_some() { warn!( target: "Client/DeviceProps", "set_device_props called after pairing — stored but not sent on the wire" diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index c151ce945..cc80246a9 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -102,7 +102,7 @@ impl Client { let mut unique_id_bytes = [0u8; 2]; rand::make_rng::().fill_bytes(&mut unique_id_bytes); - let device_snapshot = persistence_manager.get_device_snapshot().await; + let device_snapshot = persistence_manager.get_device_snapshot(); let core = wacore::client::CoreClient::new(device_snapshot.core.clone()); let (tx, rx) = async_channel::bounded(32); @@ -271,7 +271,7 @@ impl Client { // Tag the session-root span with our own (pseudonymous) account id so // connection-lifecycle traces are attributable per account. #[cfg(feature = "tracing")] - if let Some(lid) = self.get_lid().await { + if let Some(lid) = self.get_lid() { tracing::Span::current().record("account", tracing::field::display(lid.observe())); } while self.is_running.load(Ordering::Relaxed) { @@ -459,7 +459,7 @@ impl Client { self.enable_auto_reconnect.store(false, Ordering::Relaxed); if self.is_connected() - && let Ok(jid) = self.require_pn().await + && let Ok(jid) = self.require_pn() && let Err(e) = self.execute(RemoveCompanionDeviceSpec::new(&jid)).await { warn!("Failed to send logout IQ: {e}"); diff --git a/src/client/messaging.rs b/src/client/messaging.rs index 61a5f31fc..792da84ea 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -75,7 +75,7 @@ impl Client { .to_string(), ) } else { - if self.get_pn().await.is_none() { + if self.get_pn().is_none() { return Err(anyhow::Error::from(ClientError::NotLoggedIn)); } None @@ -143,7 +143,6 @@ impl Client { self.get_own_jid_for_group(&to).await?.to_non_ad() } else { self.get_pn() - .await .ok_or_else(|| anyhow::Error::from(ClientError::NotLoggedIn))? .to_non_ad() }; @@ -236,7 +235,7 @@ impl Client { if id.is_empty() { return; } - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); if let Some(own_jid) = &device_snapshot.pn { // Single source of truth for the wire mapping (ReceiptType::Sent is a derived // incoming-only state and is never sent by us). diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 29d9490be..f7608f4ae 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -456,7 +456,7 @@ impl Client { if !self.is_connected() { return Err(ClientError::NotConnected); } - let own_pn = self.get_pn().await; + let own_pn = self.get_pn(); let buf = match encode_ack_bytes(node, own_pn.as_ref()) { Ok(Some(buf)) => buf, Ok(None) => return Ok(()), @@ -473,7 +473,7 @@ impl Client { /// in a single flushed task. pub(crate) async fn send_transport_ack(&self, info: &crate::types::message::MessageInfo) { let source = message_ack_source_node(info); - let own_pn = self.get_pn().await; + let own_pn = self.get_pn(); match encode_ack_bytes(&source.as_node_ref(), own_pn.as_ref()) { Ok(Some(buf)) => { if let Err(e) = self.send_raw_bytes(buf).await @@ -507,7 +507,7 @@ impl Client { self: &Arc, node: &wacore_binary::NodeRef<'_>, ) { - let own_pn = self.get_pn().await; + let own_pn = self.get_pn(); let buf = match encode_ack_bytes(node, own_pn.as_ref()) { Ok(Some(b)) => b, Ok(None) => return, @@ -580,7 +580,7 @@ impl Client { // on Device snapshot + write lock). if let Some(lid) = lid_from_server { let device_snapshot = - client_clone.persistence_manager.get_device_snapshot().await; + client_clone.persistence_manager.get_device_snapshot(); if device_snapshot.lid.as_ref() != Some(&lid) { debug!("Updating LID from server to '{}'", lid.observe()); client_clone @@ -598,7 +598,6 @@ impl Client { let already_paired = client_clone .persistence_manager .get_device_snapshot() - .await .pn .is_some(); if already_paired { @@ -626,7 +625,7 @@ impl Client { // Check if we need initial app state sync (empty pushname indicates fresh pairing // where pushname will come from app state sync's setting_pushName mutation) - let device_snapshot = client_clone.persistence_manager.get_device_snapshot().await; + let device_snapshot = client_clone.persistence_manager.get_device_snapshot(); let needs_pushname_from_sync = device_snapshot.push_name.is_empty(); if needs_pushname_from_sync { debug!("Push name is empty - will be set from app state sync (setting_pushName)"); @@ -805,7 +804,7 @@ impl Client { } // Matches WhatsApp Web's $16(): check if SettingPushName was synced. // If push_name is still empty after 180s, critical sync failed. - let push_name = timeout_client.get_push_name().await; + let push_name = timeout_client.get_push_name(); if push_name.is_empty() { warn!( target: "Client/AppState", @@ -888,7 +887,7 @@ impl Client { } else { // === Reconnection path === // Pushname is already known, send presence and Connected immediately. - let device_snapshot = client_clone.persistence_manager.get_device_snapshot().await; + let device_snapshot = client_clone.persistence_manager.get_device_snapshot(); if !device_snapshot.push_name.is_empty() { if let Err(e) = client_clone.presence().set_available().await { warn!("Failed to send initial presence: {e:?}"); diff --git a/src/client/sender_keys.rs b/src/client/sender_keys.rs index 12228b072..d978c8e22 100644 --- a/src/client/sender_keys.rs +++ b/src/client/sender_keys.rs @@ -17,7 +17,7 @@ impl Client { exclude_own_devices: bool, ) -> Result<()> { let snapshot = if exclude_own_devices { - Some(self.persistence_manager.get_device_snapshot().await) + Some(self.persistence_manager.get_device_snapshot()) } else { None }; @@ -115,7 +115,7 @@ impl Client { use wacore::libsignal::store::sender_key_name::SenderKeyName; use wacore::types::jid::JidExt; - let snapshot = self.persistence_manager.get_device_snapshot().await; + let snapshot = self.persistence_manager.get_device_snapshot(); for own_jid in snapshot.lid.iter().chain(snapshot.pn.iter()) { let sk_name = SenderKeyName::from_parts(group_jid, own_jid.to_protocol_address().as_str()); diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 8b7c0ea00..28ae467f2 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -173,23 +173,20 @@ impl Client { async fn ensure_sessions_inner(&self, jids: Vec) -> Result<()> { use wacore::types::jid::JidExt; - let device_store = self.persistence_manager.get_device_arc().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let mut jids_needing_sessions = Vec::with_capacity(jids.len()); - { - let device_guard = device_store.read().await; - for jid in jids { - let signal_addr = jid.to_protocol_address(); - // Check cache first (includes unflushed sessions), fall back to backend - match self - .signal_cache - .has_session(&signal_addr, &*device_guard.backend) - .await - { - Ok(true) => {} - Ok(false) => jids_needing_sessions.push(jid), - Err(e) => log::warn!("Failed to check session for {}: {}", jid.observe(), e), - } + for jid in jids { + let signal_addr = jid.to_protocol_address(); + // Check cache first (includes unflushed sessions), fall back to backend + match self + .signal_cache + .has_session(&signal_addr, &*device_snapshot.backend) + .await + { + Ok(true) => {} + Ok(false) => jids_needing_sessions.push(jid), + Err(e) => log::warn!("Failed to check session for {}: {}", jid.observe(), e), } } @@ -300,7 +297,7 @@ impl Client { ) )] pub(crate) async fn establish_primary_phone_session_immediate(&self) -> Result<()> { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let own_pn = device_snapshot .pn @@ -339,11 +336,10 @@ impl Client { /// Check if a session exists for the given JID. async fn check_session_exists(&self, jid: &Jid) -> Result { - let device_store = self.persistence_manager.get_device_arc().await; - let device_guard = device_store.read().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let signal_addr = jid.to_protocol_address(); - device_guard + device_snapshot .contains_session(&signal_addr) .await .map_err(|e| anyhow::anyhow!("Failed to check session for {}: {}", jid.observe(), e)) diff --git a/src/features/contacts.rs b/src/features/contacts.rs index 95d3b32f9..d85486fb6 100644 --- a/src/features/contacts.rs +++ b/src/features/contacts.rs @@ -134,7 +134,7 @@ impl<'a> Contacts<'a> { // Skip own JID: server never responds when tctoken is sent for self let is_own_jid = { - let snap = self.client.persistence_manager.get_device_snapshot().await; + let snap = self.client.persistence_manager.get_device_snapshot(); snap.pn.as_ref().is_some_and(|pn| pn.is_same_user_as(jid)) || snap .lid diff --git a/src/features/events.rs b/src/features/events.rs index 207cfdd65..db3a8cb37 100644 --- a/src/features/events.rs +++ b/src/features/events.rs @@ -80,7 +80,6 @@ impl<'a> Events<'a> { let my_jid = self .client .get_pn() - .await .ok_or_else(|| anyhow!("Not logged in — cannot determine own JID"))?; let my_base = my_jid.to_non_ad(); @@ -135,7 +134,7 @@ impl<'a> Events<'a> { if !event_creator_jid.is_lid() { return own_pn.clone(); } - match self.client.get_lid().await { + match self.client.get_lid() { Some(lid) => lid.to_non_ad(), None => own_pn.clone(), } diff --git a/src/features/media_reupload.rs b/src/features/media_reupload.rs index 07ec4db42..31f45bdf7 100644 --- a/src/features/media_reupload.rs +++ b/src/features/media_reupload.rs @@ -67,8 +67,11 @@ impl<'a> MediaReupload<'a> { let (ciphertext, iv) = encrypt_media_retry_receipt(req.media_key, req.msg_id)?; // Get own JID for the receipt's `to` attribute - let device_snapshot = self.client.persistence_manager.get_device_snapshot().await; - let own_jid = device_snapshot.pn.clone().ok_or(ClientError::NotLoggedIn)?; + let device_snapshot = self.client.persistence_manager.get_device_snapshot(); + let own_jid = device_snapshot + .pn + .as_ref() + .ok_or(ClientError::NotLoggedIn)?; // Register waiter BEFORE sending (to avoid race) let waiter = self.client.wait_for_node( diff --git a/src/features/polls.rs b/src/features/polls.rs index fc0270df8..1467bd39d 100644 --- a/src/features/polls.rs +++ b/src/features/polls.rs @@ -107,7 +107,6 @@ impl<'a> Polls<'a> { let my_jid = self .client .get_pn() - .await .ok_or_else(|| anyhow!("Not logged in — cannot determine own JID"))?; let my_base = my_jid.to_non_ad(); @@ -174,7 +173,7 @@ impl<'a> Polls<'a> { if !poll_creator_jid.is_lid() { return own_pn.clone(); } - match self.client.get_lid().await { + match self.client.get_lid() { Some(lid) => lid.to_non_ad(), None => { log::warn!( diff --git a/src/features/presence.rs b/src/features/presence.rs index 730e30e96..b05805a93 100644 --- a/src/features/presence.rs +++ b/src/features/presence.rs @@ -67,11 +67,7 @@ impl<'a> Presence<'a> { /// Set the presence status. pub async fn set(&self, status: PresenceStatus) -> Result<(), PresenceError> { - let device_snapshot = self - .client - .persistence_manager() - .get_device_snapshot() - .await; + let device_snapshot = self.client.persistence_manager().get_device_snapshot(); debug!( "send_presence called with push_name: '{}'", @@ -303,7 +299,7 @@ mod tests { let client = bot.client(); - let snapshot = client.persistence_manager().get_device_snapshot().await; + let snapshot = client.persistence_manager().get_device_snapshot(); assert!( snapshot.push_name.is_empty(), "Pushname should be empty on fresh device" @@ -343,7 +339,7 @@ mod tests { .process_command(DeviceCommand::SetPushName("Test User".to_string())) .await; - let snapshot = client.persistence_manager().get_device_snapshot().await; + let snapshot = client.persistence_manager().get_device_snapshot(); assert_eq!(snapshot.push_name, "Test User"); // Validation passes; error should be connection-related, not pushname @@ -381,7 +377,7 @@ mod tests { let client = bot.client(); // Fresh device has empty pushname - let snapshot = client.persistence_manager().get_device_snapshot().await; + let snapshot = client.persistence_manager().get_device_snapshot(); assert!(snapshot.push_name.is_empty()); // Presence deferred when pushname empty diff --git a/src/features/signal.rs b/src/features/signal.rs index 49f563d8e..34f436dec 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -155,15 +155,13 @@ impl<'a> Signal<'a> { let _chain_guard = chain_lock.lock().await; // Only create SKDM when no sender key exists (matches WA Web behavior) - let device_store = self.client.persistence_manager.get_device_arc().await; - let device_guard = device_store.read().await; + let device_snapshot = self.client.persistence_manager.get_device_snapshot(); let key_exists = self .client .signal_cache - .get_sender_key(&sender_key_name, &*device_guard.backend) + .get_sender_key(&sender_key_name, &*device_snapshot.backend) .await? .is_some(); - drop(device_guard); let mut adapter = self.client.signal_adapter().await; let mut rng = rand::make_rng::(); @@ -230,11 +228,10 @@ impl<'a> Signal<'a> { pub async fn validate_session(&self, jid: &Jid) -> Result { let resolved = self.client.resolve_encryption_jid(jid).await; let signal_addr = resolved.to_protocol_address(); - let device_store = self.client.persistence_manager.get_device_arc().await; - let device_guard = device_store.read().await; + let device_snapshot = self.client.persistence_manager.get_device_snapshot(); self.client .signal_cache - .has_session(&signal_addr, &*device_guard.backend) + .has_session(&signal_addr, &*device_snapshot.backend) .await .map_err(|e| anyhow!("session check failed: {e}")) } diff --git a/src/handlers/call.rs b/src/handlers/call.rs index 676036737..911470c94 100644 --- a/src/handlers/call.rs +++ b/src/handlers/call.rs @@ -58,8 +58,8 @@ impl StanzaHandler for CallHandler { #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.recv.call_offer_ack", level = "debug", skip_all, fields(peer = %call.from.observe()), err(Debug)))] async fn send_offer_ack_receipt(client: &Client, call: &IncomingCall) -> anyhow::Result<()> { let own_from = match call.from.server { - Server::Lid => client.get_lid().await, - _ => client.get_pn().await, + Server::Lid => client.get_lid(), + _ => client.get_pn(), }; let Some(receipt) = build_offer_ack_receipt(call, own_from.as_ref()) else { diff --git a/src/handlers/notification/device.rs b/src/handlers/notification/device.rs index 8dabb9d6c..3bd3539c1 100644 --- a/src/handlers/notification/device.rs +++ b/src/handlers/notification/device.rs @@ -80,7 +80,6 @@ pub(crate) async fn handle_prekey_low(client: &Arc) { if client_clone .persistence_manager .get_device_snapshot() - .await .server_has_prekeys { debug!("Pre-key upload already completed by another task, skipping"); @@ -152,7 +151,7 @@ pub(crate) async fn handle_identity_change(client: &Arc, node: &NodeRef< } // Self-identity changes use a different flow; clearing our own record would break sessions - let device_snapshot = client.persistence_manager.get_device_snapshot().await; + let device_snapshot = client.persistence_manager.get_device_snapshot(); let is_me = device_snapshot .pn .as_ref() @@ -364,7 +363,7 @@ pub(crate) async fn handle_local_identity_change(client: &Arc, sender: J // Self-identity changes use a separate flow; clearing our own record would // break our sessions. - let device_snapshot = client.persistence_manager.get_device_snapshot().await; + let device_snapshot = client.persistence_manager.get_device_snapshot(); let is_me = device_snapshot .pn .as_ref() @@ -555,7 +554,7 @@ pub(crate) async fn handle_account_sync_devices( ); // Get our own JIDs (PN and LID) to verify this is about our account - let device_snapshot = client.persistence_manager.get_device_snapshot().await; + let device_snapshot = client.persistence_manager.get_device_snapshot(); let own_pn = device_snapshot.pn.as_ref(); let own_lid = device_snapshot.lid.as_ref(); diff --git a/src/handshake.rs b/src/handshake.rs index 9813018f9..4a104b533 100644 --- a/src/handshake.rs +++ b/src/handshake.rs @@ -164,7 +164,7 @@ pub async fn do_handshake( transport: Arc, transport_events: &mut async_channel::Receiver, ) -> Result> { - let device_snapshot = persistence_manager.get_device_snapshot().await; + let device_snapshot = persistence_manager.get_device_snapshot(); let now_secs = wacore::time::now_secs(); let pattern = select_pattern( &device_snapshot, diff --git a/src/history_sync.rs b/src/history_sync.rs index d4273934a..78a99d386 100644 --- a/src/history_sync.rs +++ b/src/history_sync.rs @@ -143,7 +143,7 @@ impl Client { }; let own_user = { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); device_snapshot.pn.as_ref().map(|j| j.to_non_ad().user) }; @@ -273,7 +273,7 @@ impl Client { let retention = &self.cache_config.msg_secret_retention; let now = wacore::time::now_secs(); - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let own_pn = device_snapshot.pn.as_ref().map(|j| j.to_non_ad()); let own_lid = device_snapshot.lid.as_ref().map(|j| j.to_non_ad()); @@ -385,7 +385,6 @@ impl Client { ) -> Result<(), anyhow::Error> { let own_jid = self .get_pn() - .await .ok_or(crate::client::ClientError::NotLoggedIn)? .to_non_ad(); let (ciphertext, iv) = diff --git a/src/message/msg_secret.rs b/src/message/msg_secret.rs index 31d62969a..9d6f1e783 100644 --- a/src/message/msg_secret.rs +++ b/src/message/msg_secret.rs @@ -133,23 +133,23 @@ impl Client { } match info.source.addressing_mode { - Some(AddressingMode::Lid) => match self.get_lid().await { + Some(AddressingMode::Lid) => match self.get_lid() { Some(jid) => Some(jid), - None => self.get_pn().await, + None => self.get_pn(), }, - Some(AddressingMode::Pn) => match self.get_pn().await { + Some(AddressingMode::Pn) => match self.get_pn() { Some(jid) => Some(jid), - None => self.get_lid().await, + None => self.get_lid(), }, None if info.source.sender.is_lid() || info.source.chat.is_lid() => { - match self.get_lid().await { + match self.get_lid() { Some(jid) => Some(jid), - None => self.get_pn().await, + None => self.get_pn(), } } - None => match self.get_pn().await { + None => match self.get_pn() { Some(jid) => Some(jid), - None => self.get_lid().await, + None => self.get_lid(), }, } } @@ -653,9 +653,9 @@ impl Client { return Some(ts.clone()); } if info.source.sender.server == wacore_binary::Server::Bot { - self.get_lid().await + self.get_lid() } else { - self.get_pn().await + self.get_pn() } } diff --git a/src/message/receive.rs b/src/message/receive.rs index 703ccb2a2..d2d365bf5 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -62,7 +62,7 @@ impl Client { let participants = nr.get_optional_child_by_tag(&["participants"]); if let Some(participants_node) = participants { - let own_jid = self.get_pn().await; + let own_jid = self.get_pn(); let to_nodes = participants_node.get_children_by_tag("to"); for to_node in to_nodes { let to_jid = match to_node.attrs().optional_jid("jid") { @@ -1566,14 +1566,11 @@ impl Client { &self, node: &wacore_binary::NodeRef<'_>, ) -> Result { - let (own_pn, own_lid) = { - let arc = self.persistence_manager.get_device_arc().await; - let guard = arc.read().await; - (guard.pn.clone(), guard.lid.clone()) - }; + // Per-message path: borrow pn/lid from the snapshot, no lock, no clones. + let device_snapshot = self.persistence_manager.get_device_snapshot(); let default_jid = Jid::default(); - let own_jid = own_pn.as_ref().unwrap_or(&default_jid); - wacore::messages::parse_message_info(node, own_jid, own_lid.as_ref()) + let own_jid = device_snapshot.pn.as_ref().unwrap_or(&default_jid); + wacore::messages::parse_message_info(node, own_jid, device_snapshot.lid.as_ref()) } } diff --git a/src/message/special.rs b/src/message/special.rs index a64c76570..19cfcb015 100644 --- a/src/message/special.rs +++ b/src/message/special.rs @@ -77,7 +77,7 @@ impl Client { }) } - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let key_store = device_snapshot.backend.clone(); let mut stored_count = 0; diff --git a/src/message/tests.rs b/src/message/tests.rs index 986937b3a..117a1793c 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -737,7 +737,7 @@ impl AlicePeer { /// `create_test_client_with_name` returns an unpaired client by default /// so `device_snapshot.lid` / `.pn` are both `None`. async fn ensure_bob_paired(client: &Arc) { - let snapshot = client.persistence_manager.get_device_snapshot().await; + let snapshot = client.persistence_manager.get_device_snapshot(); if snapshot.lid.is_some() || snapshot.pn.is_some() { return; } @@ -761,7 +761,7 @@ async fn ensure_bob_paired(client: &Arc) { async fn bobs_prekey_bundle(client: &Arc) -> (PreKeyBundle, Jid) { use wacore::libsignal::protocol::GenericSignedPreKey; ensure_bob_paired(client).await; - let snapshot = client.persistence_manager.get_device_snapshot().await; + let snapshot = client.persistence_manager.get_device_snapshot(); let identity_kp = snapshot.core.identity_key.clone(); let reg_id = snapshot.core.registration_id; @@ -1035,15 +1035,9 @@ async fn test_badmac_preserves_session() { &client .persistence_manager .get_device_snapshot() - .await .lid .clone() - .or(client - .persistence_manager - .get_device_snapshot() - .await - .pn - .clone()) + .or(client.persistence_manager.get_device_snapshot().pn.clone()) .expect("own jid") .to_protocol_address(), &bob_bundle, @@ -1054,15 +1048,9 @@ async fn test_badmac_preserves_session() { let bob_addr = client .persistence_manager .get_device_snapshot() - .await .lid .clone() - .or(client - .persistence_manager - .get_device_snapshot() - .await - .pn - .clone()) + .or(client.persistence_manager.get_device_snapshot().pn.clone()) .expect("own jid") .to_protocol_address(); let pkmsg = alice.encrypt_text(&bob_addr, "hello").await; @@ -1213,15 +1201,9 @@ async fn test_prod_scenario_pkmsg_archives_old_session_after_badmac() { let bob_addr = client .persistence_manager .get_device_snapshot() - .await .lid .clone() - .or(client - .persistence_manager - .get_device_snapshot() - .await - .pn - .clone()) + .or(client.persistence_manager.get_device_snapshot().pn.clone()) .expect("own jid") .to_protocol_address(); alice.install_bob_session(&bob_addr, &bundle_v1).await; @@ -1768,9 +1750,7 @@ async fn test_parse_message_info_sender_alt_extraction() { ); // Set up own phone number and LID - { - let device_arc = pm.get_device_arc().await; - let mut device = device_arc.write().await; + pm.modify_device(|device| { device.pn = Some( "15551234567@s.whatsapp.net" .parse() @@ -1781,7 +1761,8 @@ async fn test_parse_message_info_sender_alt_extraction() { .parse() .expect("test JID should be valid"), ); - } + }) + .await; let (client, _sync_rx) = Client::new( Arc::new(crate::runtime_impl::TokioRuntime), @@ -2468,9 +2449,7 @@ async fn test_parse_message_info_self_sent_dm_via_lid() { ); // Set up own phone number and LID - { - let device_arc = pm.get_device_arc().await; - let mut device = device_arc.write().await; + pm.modify_device(|device| { device.pn = Some( "15551234567@s.whatsapp.net" .parse() @@ -2481,7 +2460,8 @@ async fn test_parse_message_info_self_sent_dm_via_lid() { .parse() .expect("test JID should be valid"), ); - } + }) + .await; let (client, _sync_rx) = Client::new( Arc::new(crate::runtime_impl::TokioRuntime), @@ -2567,9 +2547,7 @@ async fn test_parse_message_info_dm_from_other_via_lid() { ); // Set up own phone number and LID - { - let device_arc = pm.get_device_arc().await; - let mut device = device_arc.write().await; + pm.modify_device(|device| { device.pn = Some( "15551234567@s.whatsapp.net" .parse() @@ -2580,7 +2558,8 @@ async fn test_parse_message_info_dm_from_other_via_lid() { .parse() .expect("test JID should be valid"), ); - } + }) + .await; let (client, _sync_rx) = Client::new( Arc::new(crate::runtime_impl::TokioRuntime), @@ -2662,9 +2641,7 @@ async fn test_parse_message_info_dm_to_self() { ); // Set up own phone number and LID - { - let device_arc = pm.get_device_arc().await; - let mut device = device_arc.write().await; + pm.modify_device(|device| { device.pn = Some( "15551234567@s.whatsapp.net" .parse() @@ -2675,7 +2652,8 @@ async fn test_parse_message_info_dm_to_self() { .parse() .expect("test JID should be valid"), ); - } + }) + .await; let (client, _sync_rx) = Client::new( Arc::new(crate::runtime_impl::TokioRuntime), diff --git a/src/pair.rs b/src/pair.rs index 61cf1a0cb..d4cab6a66 100644 --- a/src/pair.rs +++ b/src/pair.rs @@ -60,7 +60,7 @@ pub async fn handle_iq(client: &Arc, node: &NodeRef<'_>) -> bool { let mut codes = Vec::new(); - let device_snapshot = client.persistence_manager.get_device_snapshot().await; + let device_snapshot = client.persistence_manager.get_device_snapshot(); let device_state = DeviceState { identity_key: device_snapshot.identity_key.clone(), noise_key: device_snapshot.noise_key.clone(), @@ -222,7 +222,7 @@ async fn handle_pair_success<'a>( (Jid::default(), Jid::default()) }; - let device_snapshot = client.persistence_manager.get_device_snapshot().await; + let device_snapshot = client.persistence_manager.get_device_snapshot(); let device_state = DeviceState { identity_key: device_snapshot.identity_key.clone(), noise_key: device_snapshot.noise_key.clone(), @@ -388,7 +388,7 @@ pub async fn pair_with_qr_code(client: &Arc, qr_code: &str) -> Result<() let master_ephemeral = KeyPair::generate(&mut rand::make_rng::()); - let device_snapshot = client.persistence_manager.get_device_snapshot().await; + let device_snapshot = client.persistence_manager.get_device_snapshot(); let device_state = DeviceState { identity_key: device_snapshot.identity_key.clone(), noise_key: device_snapshot.noise_key.clone(), diff --git a/src/pair_code.rs b/src/pair_code.rs index 2ce74614f..8e99ec0fc 100644 --- a/src/pair_code.rs +++ b/src/pair_code.rs @@ -156,7 +156,7 @@ impl Client { let ephemeral_keypair = KeyPair::generate(&mut rand::make_rng::()); // Get device state for noise key - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let noise_static_pub: [u8; 32] = device_snapshot .noise_key .public_key @@ -334,7 +334,7 @@ pub(crate) async fn handle_pair_code_notification( }; // Get device keys - let device_snapshot = client.persistence_manager.get_device_snapshot().await; + let device_snapshot = client.persistence_manager.get_device_snapshot(); // Prepare encrypted key bundle (includes rotated adv_secret_key) let (wrapped_bundle, new_adv_secret) = match PairCodeUtils::prepare_key_bundle( diff --git a/src/pdo.rs b/src/pdo.rs index caad0411b..bb39b6072 100644 --- a/src/pdo.rs +++ b/src/pdo.rs @@ -65,7 +65,7 @@ impl Client { self: &Arc, info: &Arc, ) -> Result<(), anyhow::Error> { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let peer_target = self_peer_target(&device_snapshot)?; // Resolve to LID for the MessageKey when LID-migrated, matching WA Web's @@ -188,7 +188,7 @@ impl Client { oldest_msg_timestamp_ms: i64, count: i32, ) -> Result { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let peer_target = self_peer_target(&device_snapshot)?; let pdo_request = wa::message::PeerDataOperationRequestMessage { @@ -416,7 +416,6 @@ impl Client { } else if is_from_me { self.persistence_manager .get_device_snapshot() - .await .pn .clone() .unwrap_or_else(|| remote_jid.clone()) diff --git a/src/prekeys.rs b/src/prekeys.rs index d3c2f2c9d..57303d156 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -117,11 +117,11 @@ impl Client { let account_jid = companion_jid.with_device(0); let addr = account_jid.to_protocol_address(); - let backend = { - let device_store = self.persistence_manager.get_device_arc().await; - let guard = device_store.read().await; - guard.backend.clone() - }; + let backend = self + .persistence_manager + .get_device_snapshot() + .backend + .clone(); match self.signal_cache.get_identity(&addr, &*backend).await { Ok(Some(id)) if id.len() == 32 => { let mut arr = [0u8; 32]; @@ -170,7 +170,6 @@ impl Client { let has_prekeys = self .persistence_manager .get_device_snapshot() - .await .server_has_prekeys; if has_prekeys { @@ -185,7 +184,6 @@ impl Client { if self .persistence_manager .get_device_snapshot() - .await .server_has_prekeys { return Ok(()); @@ -240,13 +238,12 @@ impl Client { let next_pre_key_id = self .persistence_manager .get_device_snapshot() - .await .next_pre_key_id; - let backend = { - let device_store = self.persistence_manager.get_device_arc().await; - let guard = device_store.read().await; - guard.backend.clone() - }; + let backend = self + .persistence_manager + .get_device_snapshot() + .backend + .clone(); let max_id = backend.get_max_prekey_id().await?; let id = start_prekey_id(next_pre_key_id, max_id); let next = (id as u64 % MAX_PREKEY_ID as u64) as u32 + 1; @@ -269,13 +266,8 @@ impl Client { ) )] async fn upload_pre_keys_inner(&self) -> Result<(), anyhow::Error> { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; - let device_store = self.persistence_manager.get_device_arc().await; - - let backend = { - let device_guard = device_store.read().await; - device_guard.backend.clone() - }; + let device_snapshot = self.persistence_manager.get_device_snapshot(); + let backend = device_snapshot.backend.clone(); // Use the persistent counter, falling back to max(store_id)+1 for migration. // The counter is the source of truth after the first upload; start_prekey_id wraps @@ -510,7 +502,7 @@ impl Client { // WA Web's validateLocalKeyBundle validates but catches ALL exceptions without // re-uploading. The catch block in digestKey() sets a=false for any throw from y(), // meaning only 404 triggers re-upload. We match that: log warnings, return Ok(()). - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); if response.reg_id != device_snapshot.registration_id { log::warn!( "digestKey: registration ID mismatch (server={}, local={}), skipping", @@ -526,11 +518,11 @@ impl Client { let skey_pub_bytes = device_snapshot.signed_pre_key.public_key.public_key_bytes(); let skey_sig_bytes = &device_snapshot.signed_pre_key_signature; - let device_store = self.persistence_manager.get_device_arc().await; - let backend = { - let guard = device_store.read().await; - guard.backend.clone() - }; + let backend = self + .persistence_manager + .get_device_snapshot() + .backend + .clone(); // Batch-load all prekeys referenced by the server digest let loaded = match backend.load_prekeys_batch(&response.prekey_ids).await { diff --git a/src/receipt.rs b/src/receipt.rs index 7902fbca4..1c945275f 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -464,7 +464,7 @@ impl Client { if info.id.is_empty() { return; } - let Some(own_pn) = self.get_pn().await else { + let Some(own_pn) = self.get_pn() else { log::debug!( "[msg:{}] Skipping nack ({:?}): own PN not yet set", info.id, diff --git a/src/request.rs b/src/request.rs index d8efe419c..9377a3f91 100644 --- a/src/request.rs +++ b/src/request.rs @@ -89,7 +89,7 @@ impl Client { /// /// A string containing the generated message ID in the format expected by WhatsApp. pub async fn generate_message_id(&self) -> String { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); self.get_request_utils() .generate_message_id(device_snapshot.pn.as_ref()) } diff --git a/src/retry.rs b/src/retry.rs index 21cf2d3fd..c907e5be5 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -227,7 +227,7 @@ impl Client { return Ok(()); } - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let mut info = resolve_retry_chat_info( receipt, nr, @@ -491,7 +491,7 @@ impl Client { self.ensure_e2e_sessions_resolved(std::slice::from_ref(&resolved_jid)) .await?; - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let addressing_mode = cached_group_info .as_ref() @@ -530,7 +530,7 @@ impl Client { self.ensure_e2e_sessions_resolved(std::slice::from_ref(&resolved_jid)) .await?; - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let signal_address = resolved_jid.to_protocol_address(); let session_mutex = self.session_lock_for(signal_address.as_str()).await; let _session_guard = session_mutex.lock().await; @@ -649,15 +649,13 @@ impl Client { if let Some(received_reg_id) = extract_registration_id_from_node_ref(node) { let signal_address = resolved_jid.to_protocol_address(); - let device_store = self.persistence_manager.get_device_arc().await; - let device_guard = device_store.read().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let session = self .signal_cache - .peek_session(&signal_address, &*device_guard.backend) + .peek_session(&signal_address, &*device_snapshot.backend) .await .ok() .flatten(); - drop(device_guard); if let Some(session) = session && let Ok(stored_reg_id) = session.remote_registration_id() @@ -684,11 +682,10 @@ impl Client { // 4-5. Base-key collision logic (WA Web L66-80). Applied to ALL chat // types now — previously only ran in the DM branch. let signal_address = resolved_jid.to_protocol_address(); - let device_store = self.persistence_manager.get_device_arc().await; - let device_guard = device_store.read().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let session = self .signal_cache - .peek_session(&signal_address, &*device_guard.backend) + .peek_session(&signal_address, &*device_snapshot.backend) .await .ok() .flatten(); @@ -703,7 +700,7 @@ impl Client { let addr_str = signal_address.as_str(); if retry_count == MIN_RETRY_FOR_BASE_KEY_CHECK { // retry == 2: save base key, do NOT delete (WA Web L66-67). - match device_guard + match device_snapshot .backend .save_base_key(addr_str, message_id, current_base_key) .await @@ -723,7 +720,7 @@ impl Client { } if retry_count > MIN_RETRY_FOR_BASE_KEY_CHECK { - match device_guard + match device_snapshot .backend .has_same_base_key(addr_str, message_id, current_base_key) .await @@ -735,11 +732,10 @@ impl Client { wacore::types::jid::observe_protocol_address(&signal_address), retry_count ); - let _ = device_guard + let _ = device_snapshot .backend .delete_base_key(addr_str, message_id) .await; - drop(device_guard); let lock = self.session_lock_for(signal_address.as_str()).await; let _guard = lock.lock().await; self.signal_cache.delete_session(&signal_address).await; @@ -756,7 +752,7 @@ impl Client { wacore::types::jid::observe_protocol_address(&signal_address), retry_count ); - let _ = device_guard + let _ = device_snapshot .backend .delete_base_key(addr_str, message_id) .await; @@ -797,14 +793,13 @@ impl Client { now: wacore::time::Instant, ) -> Option<&'static str> { let signal_address = jid.to_protocol_address(); - let device_store = self.persistence_manager.get_device_arc().await; - let device_guard = device_store.read().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); // Whatsmeow returns `false` on `ContainsSession` errors so a transient // backend read failure doesn't masquerade as "no session" and trigger // an unnecessary delete + prekey fetch (`retry.go:161-163`). let has_session = match self .signal_cache - .has_session(&signal_address, &*device_guard.backend) + .has_session(&signal_address, &*device_snapshot.backend) .await { Ok(present) => present, @@ -816,7 +811,6 @@ impl Client { return None; } }; - drop(device_guard); let history = &self.session_recreate_history; @@ -895,15 +889,13 @@ impl Client { // Check if the registration ID changed (indicates device reinstall). // Read session through cache for consistent state. { - let device_store = self.persistence_manager.get_device_arc().await; - let device_guard = device_store.read().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let session = self .signal_cache - .peek_session(&signal_address, &*device_guard.backend) + .peek_session(&signal_address, &*device_snapshot.backend) .await .ok() .flatten(); - drop(device_guard); if let Some(session) = session { let existing_reg_id = session.remote_registration_id()?; @@ -1050,7 +1042,7 @@ impl Client { retry_count: u8, reason: RetryReason, ) -> Result<(), anyhow::Error> { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); // WA Web's sendRetryReceipt aborts only when `!to.isBot() && participant.isBot()`, // with participant null for DMs. A bot DM is chat == sender == bot, so it is NOT @@ -1107,15 +1099,13 @@ impl Client { &new_prekey_keypair, ); // This key is not uploaded to the server pool, so mark as false - let device_store = self.persistence_manager.get_device_arc().await; - let device_guard = device_store.read().await; - if let Err(e) = device_guard + let device_snapshot = self.persistence_manager.get_device_snapshot(); + if let Err(e) = device_snapshot .store_prekey(new_prekey_id, new_prekey_record, false) .await { warn!("Failed to store new prekey for retry receipt: {e:?}"); } - drop(device_guard); drop(prekey_guard); let device_identity_bytes = device_snapshot @@ -1213,7 +1203,7 @@ impl Client { call_creator: &wacore_binary::Jid, retry_count: u8, ) -> Result<(), anyhow::Error> { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let registration_id_bytes = device_snapshot.registration_id.to_be_bytes().to_vec(); diff --git a/src/send.rs b/src/send.rs index d4b92312b..04de1c1f4 100644 --- a/src/send.rs +++ b/src/send.rs @@ -547,16 +547,17 @@ impl Client { let to = Jid::status_broadcast(); let request_id = self.generate_message_id().await; - let mut device_snapshot = self.persistence_manager.get_device_snapshot().await; - let account_info = device_snapshot.account.take(); + // Borrow from the held snapshot: no field clones, the Arc keeps it alive. + let device_snapshot = self.persistence_manager.get_device_snapshot(); + let account_info = &device_snapshot.account; let own_jid = device_snapshot .pn - .take() + .as_ref() .ok_or(crate::client::ClientError::NotLoggedIn)?; // Status is LID-addressed (matches WA Web post-LID-migration). Without // a real device LID we can't sign or fan out correctly; refuse rather // than silently emit `addressing_mode="lid"` with a PN sender. - let own_lid = device_snapshot.lid.take().ok_or_else(|| { + let own_lid = device_snapshot.lid.as_ref().ok_or_else(|| { anyhow!( "Cannot send status: device has no LID yet. Finish pairing / LID \ migration before posting status." @@ -610,10 +611,9 @@ impl Client { let sender_address = own_lid.to_protocol_address(); let sender_key_name = SenderKeyName::from_parts(&to_str, sender_address.as_str()); - let device_guard = device_store_arc.read().await; let key_exists = self .signal_cache - .get_sender_key(&sender_key_name, &*device_guard.backend) + .get_sender_key(&sender_key_name, &*device_snapshot.backend) .await? .is_some(); @@ -976,7 +976,7 @@ impl Client { // (WA Web: syncDeviceListJob([recipient, me])) if !jid.is_group() && !jid.is_status_broadcast() { self.invalidate_device_cache(&jid.user).await; - if let Some(own_pn) = &self.persistence_manager.get_device_snapshot().await.pn { + if let Some(own_pn) = &self.persistence_manager.get_device_snapshot().pn { self.invalidate_device_cache(&own_pn.user).await; } } @@ -998,7 +998,7 @@ impl Client { ); use wacore::libsignal::store::sender_key_name::SenderKeyName; use wacore::types::jid::JidExt; - let snapshot = self.persistence_manager.get_device_snapshot().await; + let snapshot = self.persistence_manager.get_device_snapshot(); for own in snapshot.lid.iter().chain(snapshot.pn.iter()) { let sk = SenderKeyName::from_parts(&jid_str, own.to_protocol_address().as_str()); @@ -1046,7 +1046,7 @@ impl Client { revoke_type: RevokeType, ) -> Result<(), anyhow::Error> { let message_id = message_id.into(); - self.require_pn().await?; + self.require_pn()?; let (from_me, participant, edit_attr) = match &revoke_type { RevokeType::Sender => { @@ -1270,7 +1270,7 @@ impl Client { let mut store_adapter = self.signal_adapter().await; - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); wacore::send::prepare_peer_stanza( &mut store_adapter.session_store, &mut store_adapter.identity_store, @@ -1286,15 +1286,16 @@ impl Client { // sender-key chain advance per (group, sender) at the cipher. let group_info = self.groups().query_info(&to).await?; - let mut device_snapshot = self.persistence_manager.get_device_snapshot().await; - let account_info = device_snapshot.account.take(); + // Borrow from the held snapshot: no field clones, the Arc keeps it alive. + let device_snapshot = self.persistence_manager.get_device_snapshot(); + let account_info = &device_snapshot.account; let own_jid = device_snapshot .pn - .take() + .as_ref() .ok_or(crate::client::ClientError::NotLoggedIn)?; let own_lid = device_snapshot .lid - .take() + .as_ref() .ok_or_else(|| anyhow!("LID not set, cannot send to group"))?; // Store serialized message bytes for retry (lightweight) @@ -1317,10 +1318,9 @@ impl Client { let sender_address = own_sending_jid.to_protocol_address(); let sender_key_name = SenderKeyName::from_parts(&to_str, sender_address.as_str()); - let device_guard = device_store_arc.read().await; let record = self .signal_cache - .get_sender_key(&sender_key_name, &*device_guard.backend) + .get_sender_key(&sender_key_name, &*device_snapshot.backend) .await?; let key_exists = record.is_some(); @@ -1336,7 +1336,6 @@ impl Client { .and_then(|state| state.sender_chain_key()) .map(|ck| ck.iteration()) .is_some_and(|iter| iter >= SENDER_KEY_ROTATION_THRESHOLD); - drop(device_guard); if needs_rotation { log::info!( @@ -1479,7 +1478,7 @@ impl Client { self.add_recent_message(&to, &request_id, message).await; } - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let own_jid = device_snapshot .pn .as_ref() @@ -1805,9 +1804,9 @@ impl Client { /// `PreparedGroupStanza.sender_identity` directly instead of this. pub(crate) async fn dm_sender_identity_for(&self, to: &Jid) -> Option { if to.server == wacore_binary::Server::Bot { - self.get_lid().await + self.get_lid() } else { - self.get_pn().await + self.get_pn() } } @@ -1833,7 +1832,7 @@ impl Client { }; // Skip for own JID — no need to send privacy token to ourselves - let snapshot = self.persistence_manager.get_device_snapshot().await; + let snapshot = self.persistence_manager.get_device_snapshot(); let is_self = snapshot .pn .as_ref() diff --git a/src/store/persistence_manager.rs b/src/store/persistence_manager.rs index 9ec24d926..b1862edb5 100644 --- a/src/store/persistence_manager.rs +++ b/src/store/persistence_manager.rs @@ -12,6 +12,11 @@ use wacore::runtime::{AbortHandle, Runtime, ShutdownSignal, wait_for_shutdown}; pub struct PersistenceManager { device: Arc>, + /// Read-mostly snapshot, rebuilt under the device write guard in + /// `modify_device` so it can never lag a committed mutation. Turns every + /// `get_device_snapshot` into an Arc refcount bump instead of a full + /// Device clone (the snapshot is read on every inbound message). + device_snapshot: std::sync::RwLock>, backend: Arc, dirty: Arc, save_notify: Arc, @@ -50,8 +55,10 @@ impl PersistenceManager { Device::new(backend.clone()) }; + let snapshot = Arc::new(device.clone()); Ok(Self { device: Arc::new(RwLock::new(device)), + device_snapshot: std::sync::RwLock::new(snapshot), backend, dirty: Arc::new(AtomicBool::new(false)), save_notify: Arc::new(Event::new()), @@ -59,12 +66,20 @@ impl PersistenceManager { }) } + /// Handle for store adapters that need `&mut Device` trait access. + /// For plain reads, prefer [`get_device_snapshot`](Self::get_device_snapshot). pub async fn get_device_arc(&self) -> Arc> { self.device.clone() } - pub async fn get_device_snapshot(&self) -> Device { - self.device.read().await.clone() + /// Cheap point-in-time view of the device state: an Arc refcount bump, + /// no locking against writers and no Device clone. Always reflects the + /// last committed `modify_device`/`process_command` mutation. + pub fn get_device_snapshot(&self) -> Arc { + self.device_snapshot + .read() + .unwrap_or_else(|p| p.into_inner()) + .clone() } pub fn backend(&self) -> Arc { @@ -83,6 +98,14 @@ impl PersistenceManager { let mut device_guard = self.device.write().await; let result = modifier(&mut device_guard); + // Rebuild while still holding the write guard so no reader can + // observe post-mutation effects with a pre-mutation snapshot. + *self + .device_snapshot + .write() + .unwrap_or_else(|p| p.into_inner()) = Arc::new(device_guard.clone()); + drop(device_guard); + self.dirty.store(true, Ordering::Relaxed); self.save_notify.notify(1); diff --git a/src/usync.rs b/src/usync.rs index 69bdb0452..d6a454511 100644 --- a/src/usync.rs +++ b/src/usync.rs @@ -229,7 +229,7 @@ impl Client { ) )] pub(crate) async fn sync_own_device_list(&self) -> Result<(), anyhow::Error> { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); let mut jids = Vec::with_capacity(2); let mut hashes: std::collections::HashMap = diff --git a/src/version.rs b/src/version.rs index bb1ede996..6edc5c637 100644 --- a/src/version.rs +++ b/src/version.rs @@ -43,7 +43,7 @@ pub async fn resolve_and_update_version( return Ok(()); } - let device = persistence_manager.get_device_snapshot().await; + let device = persistence_manager.get_device_snapshot(); let last_fetched_ms = device.app_version_last_fetched_ms; let needs_fetch = if last_fetched_ms == 0 { diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index c0561804d..ac576e8b8 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -267,7 +267,6 @@ impl TestClient { pub async fn jid(&self) -> Jid { self.client .get_pn() - .await .expect("Client should have a JID after connect") .to_non_ad() } @@ -277,13 +276,12 @@ impl TestClient { /// Notification handling stores tcTokens under the sender's LID when it is /// available, otherwise it falls back to the phone-number user part. pub async fn tc_token_key(&self) -> anyhow::Result { - if let Some(lid) = self.client.get_lid().await { + if let Some(lid) = self.client.get_lid() { return Ok(lid.user.to_string()); } self.client .get_pn() - .await .map(|jid| jid.user.to_string()) .ok_or_else(|| anyhow::anyhow!("Client should have a JID after connect")) } @@ -329,7 +327,6 @@ impl TestClient { self.client .persistence_manager() .get_device_snapshot() - .await .nct_salt .clone() } @@ -446,7 +443,7 @@ impl TestClient { /// Wait for initial app state sync to complete (keys become available). pub async fn wait_for_app_state_sync(&mut self) -> anyhow::Result<()> { - let push_name = self.client.get_push_name().await; + let push_name = self.client.get_push_name(); if !push_name.is_empty() { return Ok(()); } diff --git a/tests/e2e/tests/app_state.rs b/tests/e2e/tests/app_state.rs index 8d87822b2..735cdc26c 100644 --- a/tests/e2e/tests/app_state.rs +++ b/tests/e2e/tests/app_state.rs @@ -16,7 +16,7 @@ async fn test_initial_sync_delivers_push_name() -> anyhow::Result<()> { let mut client = TestClient::connect_without_push_name("e2e_as_init_sync").await?; client.wait_for_app_state_sync().await?; - let push_name = client.client.get_push_name().await; + let push_name = client.client.get_push_name(); assert!( !push_name.is_empty(), "Push name should be set from initial critical_block sync (got empty — app state keys may be broken)" @@ -39,12 +39,12 @@ async fn test_push_name_survives_reconnect() -> anyhow::Result<()> { let name = "ReconnectTest"; client.client.profile().set_push_name(name).await?; - assert_eq!(client.client.get_push_name().await, name); + assert_eq!(client.client.get_push_name(), name); info!("Push name set to '{name}'"); client.reconnect_and_wait().await?; - let after = client.client.get_push_name().await; + let after = client.client.get_push_name(); assert_eq!(after, name, "Push name should survive reconnect"); info!("Push name after reconnect: '{after}'"); @@ -65,7 +65,6 @@ async fn test_mutation_works_after_reconnect() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -97,7 +96,6 @@ async fn test_undo_mutation_after_reconnect() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -134,7 +132,6 @@ async fn test_cross_collection_mutations() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -195,13 +192,11 @@ async fn test_star_received_message() -> anyhow::Result<()> { let jid_a = client_a .client .get_pn() - .await .expect("A should have JID") .to_non_ad(); let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -259,8 +254,8 @@ async fn test_multi_device_app_state_sync() -> anyhow::Result<()> { let mut client_a2 = TestClient::connect_as("e2e_multidev_a2", push_name).await?; // Verify both devices got the same phone number - let phone_a1 = client_a1.client.get_pn().await.expect("A1 should have JID"); - let phone_a2 = client_a2.client.get_pn().await.expect("A2 should have JID"); + let phone_a1 = client_a1.client.get_pn().expect("A1 should have JID"); + let phone_a2 = client_a2.client.get_pn().expect("A2 should have JID"); assert_eq!( phone_a1.user, phone_a2.user, "Both devices should share the same phone number" @@ -318,7 +313,6 @@ async fn test_rapid_successive_mutations() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); diff --git a/tests/e2e/tests/chat_actions.rs b/tests/e2e/tests/chat_actions.rs index 23af313b2..ec9b40502 100644 --- a/tests/e2e/tests/chat_actions.rs +++ b/tests/e2e/tests/chat_actions.rs @@ -18,7 +18,6 @@ async fn test_archive_chat() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -46,7 +45,6 @@ async fn test_unarchive_chat() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -84,7 +82,6 @@ async fn test_pin_chat() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -108,7 +105,6 @@ async fn test_unpin_chat() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -137,7 +133,6 @@ async fn test_mute_chat_indefinite() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -161,7 +156,6 @@ async fn test_mute_chat_with_expiry() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -192,7 +186,6 @@ async fn test_unmute_chat() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -222,7 +215,6 @@ async fn test_star_message() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -265,7 +257,6 @@ async fn test_unstar_message() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -315,7 +306,6 @@ async fn test_multiple_chat_actions() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -366,7 +356,6 @@ async fn test_mark_chat_as_read() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -401,7 +390,6 @@ async fn test_delete_chat() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -429,7 +417,6 @@ async fn test_delete_message_for_me() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); diff --git a/tests/e2e/tests/community.rs b/tests/e2e/tests/community.rs index abfaea504..03a0c3ee5 100644 --- a/tests/e2e/tests/community.rs +++ b/tests/e2e/tests/community.rs @@ -381,13 +381,11 @@ async fn test_community_join_subgroup() -> anyhow::Result<()> { let jid_b_pn = client_b .client .get_pn() - .await .expect("Client B should have a PN JID") .to_non_ad(); let jid_b_lid = client_b .client .get_lid() - .await .expect("Client B should have a LID JID") .to_non_ad(); @@ -492,13 +490,11 @@ async fn test_community_get_linked_groups_participants() -> anyhow::Result<()> { let own_pn = client .client .get_pn() - .await .expect("should have PN JID") .to_non_ad(); let own_lid = client .client .get_lid() - .await .expect("should have LID JID") .to_non_ad(); diff --git a/tests/e2e/tests/digest_key.rs b/tests/e2e/tests/digest_key.rs index 22ca59a7d..15e11b067 100644 --- a/tests/e2e/tests/digest_key.rs +++ b/tests/e2e/tests/digest_key.rs @@ -46,7 +46,7 @@ async fn test_digest_key_hash_matches_server() -> anyhow::Result<()> { // Load local key material to reproduce the hash let pm = tc.client.persistence_manager(); - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); assert_eq!( response.reg_id, device.registration_id, diff --git a/tests/e2e/tests/groups.rs b/tests/e2e/tests/groups.rs index 73fc84bd1..3d4f9eb73 100644 --- a/tests/e2e/tests/groups.rs +++ b/tests/e2e/tests/groups.rs @@ -735,7 +735,6 @@ async fn test_query_info_populates_lid_pn_cache_for_participants() -> anyhow::Re let jid_b_lid = client_b .client .get_lid() - .await .expect("B must have a LID after pairing") .to_non_ad(); info!("B pn={jid_b_pn} lid={jid_b_lid}"); diff --git a/tests/e2e/tests/lid_sessions.rs b/tests/e2e/tests/lid_sessions.rs index e2239bb32..721778577 100644 --- a/tests/e2e/tests/lid_sessions.rs +++ b/tests/e2e/tests/lid_sessions.rs @@ -98,8 +98,8 @@ async fn test_sessions_stored_under_lid_not_pn() -> anyhow::Result<()> { let jid_a = client_a.jid().await; let jid_b = client_b.jid().await; - let lid_a = client_a.client.get_lid().await.expect("A should have LID"); - let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + let lid_a = client_a.client.get_lid().expect("A should have LID"); + let lid_b = client_b.client.get_lid().expect("B should have LID"); // Roundtrip to establish sessions in both directions send_and_expect_text( @@ -163,7 +163,7 @@ async fn test_multiple_sends_stay_lid_only() -> anyhow::Result<()> { let mut client_b = TestClient::connect("e2e_lid_multi_send_b").await?; let jid_b = client_b.jid().await; - let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + let lid_b = client_b.client.get_lid().expect("B should have LID"); // Send 5 messages sequentially for i in 1..=5 { @@ -210,7 +210,7 @@ async fn test_stale_pn_session_does_not_break_lid_messaging() -> anyhow::Result< let jid_a = client_a.jid().await; let jid_b = client_b.jid().await; - let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + let lid_b = client_b.client.get_lid().expect("B should have LID"); // First, establish a normal LID session via roundtrip send_and_expect_text( @@ -312,7 +312,7 @@ async fn test_lid_session_survives_reconnect() -> anyhow::Result<()> { let jid_a = client_a.jid().await; let jid_b = client_b.jid().await; - let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + let lid_b = client_b.client.get_lid().expect("B should have LID"); // Establish sessions with a roundtrip send_and_expect_text( @@ -388,12 +388,10 @@ async fn test_own_device_0_has_lid_session_after_login() -> anyhow::Result<()> { let own_pn = client .client .get_pn() - .await .expect("Client should have PN after connect"); let own_lid = client .client .get_lid() - .await .expect("Client should have LID after connect"); let backend = client.client.persistence_manager().backend(); @@ -491,7 +489,7 @@ async fn test_pn_only_session_causes_undecryptable_on_lid_lookup() -> anyhow::Re let jid_a = client_a.jid().await; let jid_b = client_b.jid().await; - let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + let lid_b = client_b.client.get_lid().expect("B should have LID"); // Step 1: Establish sessions via roundtrip send_and_expect_text(&client_a.client, &mut client_b, &jid_b, "Setup A->B", 30).await?; diff --git a/tests/e2e/tests/media.rs b/tests/e2e/tests/media.rs index 84cbde3ed..8490c4b7b 100644 --- a/tests/e2e/tests/media.rs +++ b/tests/e2e/tests/media.rs @@ -391,7 +391,6 @@ async fn test_send_image_message() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -458,7 +457,6 @@ async fn test_send_video_message() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -508,7 +506,6 @@ async fn test_send_document_message() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -557,7 +554,6 @@ async fn test_send_audio_message() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -606,7 +602,6 @@ async fn test_send_ptt_voice_message() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -658,13 +653,11 @@ async fn test_send_image_bidirectional() -> anyhow::Result<()> { let jid_a = client_a .client .get_pn() - .await .expect("A should have JID") .to_non_ad(); let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -737,7 +730,6 @@ async fn test_send_multiple_media_types() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); @@ -901,7 +893,6 @@ async fn test_send_image_no_caption() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have JID") .to_non_ad(); diff --git a/tests/e2e/tests/memory_soak.rs b/tests/e2e/tests/memory_soak.rs index 0d5903f2d..410d1d3f2 100644 --- a/tests/e2e/tests/memory_soak.rs +++ b/tests/e2e/tests/memory_soak.rs @@ -312,8 +312,8 @@ async fn test_heavy_dm_soak() -> anyhow::Result<()> { let mut client_a = TestClient::connect("soak2_dm_a").await?; let mut client_b = TestClient::connect("soak2_dm_b").await?; - let jid_a = client_a.client.get_pn().await.expect("A JID").to_non_ad(); - let jid_b = client_b.client.get_pn().await.expect("B JID").to_non_ad(); + let jid_a = client_a.client.get_pn().expect("A JID").to_non_ad(); + let jid_b = client_b.client.get_pn().expect("B JID").to_non_ad(); // Warm-up for i in 0..5 { @@ -374,8 +374,8 @@ async fn test_heavy_group_soak() -> anyhow::Result<()> { let mut client_b = TestClient::connect("soak2_grp_b").await?; let mut client_c = TestClient::connect("soak2_grp_c").await?; - let jid_b = client_b.client.get_pn().await.expect("B JID").to_non_ad(); - let jid_c = client_c.client.get_pn().await.expect("C JID").to_non_ad(); + let jid_b = client_b.client.get_pn().expect("B JID").to_non_ad(); + let jid_c = client_c.client.get_pn().expect("C JID").to_non_ad(); // Create group 1: A + B + C let g1 = client_a @@ -518,9 +518,9 @@ async fn test_heavy_mixed_soak() -> anyhow::Result<()> { let mut client_b = TestClient::connect("soak2_mix_b").await?; let mut client_c = TestClient::connect("soak2_mix_c").await?; - let jid_a = client_a.client.get_pn().await.expect("A JID").to_non_ad(); - let jid_b = client_b.client.get_pn().await.expect("B JID").to_non_ad(); - let jid_c = client_c.client.get_pn().await.expect("C JID").to_non_ad(); + let jid_a = client_a.client.get_pn().expect("A JID").to_non_ad(); + let jid_b = client_b.client.get_pn().expect("B JID").to_non_ad(); + let jid_c = client_c.client.get_pn().expect("C JID").to_non_ad(); // Create group let group_jid = client_a @@ -664,7 +664,7 @@ async fn test_many_peers_soak() -> anyhow::Result<()> { let mut peer_jids: Vec = Vec::new(); for i in 0..num_peers { let peer = TestClient::connect(&format!("soak2_peers_p{i}")).await?; - let jid = peer.client.get_pn().await.expect("peer JID").to_non_ad(); + let jid = peer.client.get_pn().expect("peer JID").to_non_ad(); peer_jids.push(jid); peers.push(peer); } @@ -732,8 +732,8 @@ async fn test_heavy_reconnect_soak() -> anyhow::Result<()> { let mut client_a = TestClient::connect("soak2_recon_a").await?; let mut client_b = TestClient::connect("soak2_recon_b").await?; - let jid_a = client_a.client.get_pn().await.expect("A JID").to_non_ad(); - let jid_b = client_b.client.get_pn().await.expect("B JID").to_non_ad(); + let jid_a = client_a.client.get_pn().expect("A JID").to_non_ad(); + let jid_b = client_b.client.get_pn().expect("B JID").to_non_ad(); let mut snaps: Vec = Vec::new(); snaps.push(snapshot("A-recon", 0, &client_a.client).await); diff --git a/tests/e2e/tests/prekey_sessions.rs b/tests/e2e/tests/prekey_sessions.rs index 3b9a5c093..719ca173d 100644 --- a/tests/e2e/tests/prekey_sessions.rs +++ b/tests/e2e/tests/prekey_sessions.rs @@ -42,7 +42,6 @@ async fn test_prekey_collision_regression() -> anyhow::Result<()> { let recipient_jid = recipient .client .get_pn() - .await .expect("Recipient should have a JID") .to_non_ad(); info!("Recipient JID: {recipient_jid}"); diff --git a/tests/e2e/tests/privacy_tokens.rs b/tests/e2e/tests/privacy_tokens.rs index 77518fbd2..373d19a23 100644 --- a/tests/e2e/tests/privacy_tokens.rs +++ b/tests/e2e/tests/privacy_tokens.rs @@ -124,7 +124,6 @@ async fn test_issue_tokens_api_delivers_notification_and_updates_index() -> anyh let jid_b_lid = client_b .client .get_lid() - .await .expect("B should have LID after connect"); let issued = client_a .client @@ -237,8 +236,8 @@ async fn test_tc_token_notification_reaches_all_connected_devices() -> anyhow::R let mut client_b1 = TestClient::connect_as("e2e_tctok_multi_b1", &shared_b_name).await?; let client_b2 = TestClient::connect_as("e2e_tctok_multi_b2", &shared_b_name).await?; - let phone_b1 = client_b1.client.get_pn().await.expect("B1 should have JID"); - let phone_b2 = client_b2.client.get_pn().await.expect("B2 should have JID"); + let phone_b1 = client_b1.client.get_pn().expect("B1 should have JID"); + let phone_b2 = client_b2.client.get_pn().expect("B2 should have JID"); assert_eq!( phone_b1.user, phone_b2.user, "B devices should share a phone" @@ -397,7 +396,6 @@ async fn test_only_nct_send_ab_without_salt_still_receives_463() -> anyhow::Resu let jid_a_lid = client_a .client .get_lid() - .await .expect("restricted recipient should have a LID"); let msg_id = format!("E2ECSNEG1{}", uuid::Uuid::new_v4().simple()); let sent_msg_id = msg_id.clone(); @@ -457,7 +455,6 @@ async fn test_send_and_syncd_ab_without_delivery_still_receives_463() -> anyhow: let jid_a_lid = client_a .client .get_lid() - .await .expect("restricted recipient should have a LID"); let msg_id = format!("E2ECSNEG2{}", uuid::Uuid::new_v4().simple()); let sent_msg_id = msg_id.clone(); @@ -518,7 +515,6 @@ async fn test_history_sync_nct_salt_enables_cstoken_first_contact() -> anyhow::R let jid_a_lid = client_a .client .get_lid() - .await .expect("restricted recipient should have a LID"); let sent_waiter = client_b.next_sent_message_waiter(); client_b @@ -585,7 +581,6 @@ async fn test_cstoken_only_first_contact_succeeds_when_tctoken_disabled() -> any let jid_a_lid = client_a .client .get_lid() - .await .expect("restricted recipient should have a LID"); let sent_waiter = client_b.next_sent_message_waiter(); client_b @@ -646,7 +641,6 @@ async fn test_syncd_nct_salt_enables_cstoken_first_contact() -> anyhow::Result<( let jid_a_lid = client_a .client .get_lid() - .await .expect("restricted recipient should have a LID"); let sent_waiter = client_b.next_sent_message_waiter(); client_b @@ -704,7 +698,6 @@ async fn test_clearing_nct_salt_locally_makes_first_contact_fail_again() -> anyh let jid_a_lid = client_a .client .get_lid() - .await .expect("restricted recipient should have a LID"); client_b .client @@ -734,7 +727,6 @@ async fn test_clearing_nct_salt_locally_makes_first_contact_fail_again() -> anyh let jid_c_lid = client_c .client .get_lid() - .await .expect("restricted recipient should have a LID"); send_first_message_and_expect_463( &client_b, @@ -844,7 +836,6 @@ async fn test_nct_salt_survives_reconnect_and_still_allows_first_contact() -> an let jid_a_lid = client_a .client .get_lid() - .await .expect("restricted recipient should have a LID"); let sent_waiter = client_b.next_sent_message_waiter(); client_b diff --git a/tests/e2e/tests/profile.rs b/tests/e2e/tests/profile.rs index 24ddc1306..1372ec85c 100644 --- a/tests/e2e/tests/profile.rs +++ b/tests/e2e/tests/profile.rs @@ -38,7 +38,7 @@ async fn test_set_push_name() -> anyhow::Result<()> { // The push name mutation requires encryption keys from the critical_block sync. client.wait_for_app_state_sync().await?; - let old_name = client.client.get_push_name().await; + let old_name = client.client.get_push_name(); info!("Current push name: '{}'", old_name); let new_name = "TestBot 🤖"; @@ -46,7 +46,7 @@ async fn test_set_push_name() -> anyhow::Result<()> { client.client.profile().set_push_name(new_name).await?; // Verify it was updated locally - let updated_name = client.client.get_push_name().await; + let updated_name = client.client.get_push_name(); assert_eq!( updated_name, new_name, "Push name should be updated locally" @@ -58,7 +58,7 @@ async fn test_set_push_name() -> anyhow::Result<()> { info!("Setting push name again to '{}'...", second_name); client.client.profile().set_push_name(second_name).await?; - let final_name = client.client.get_push_name().await; + let final_name = client.client.get_push_name(); assert_eq!( final_name, second_name, "Push name should be updated to second value" @@ -164,7 +164,7 @@ async fn test_set_push_name_special_characters() -> anyhow::Result<()> { let name_emoji = "Bot 🤖🦀"; info!("Setting push name with emoji: '{}'...", name_emoji); client.client.profile().set_push_name(name_emoji).await?; - let result = client.client.get_push_name().await; + let result = client.client.get_push_name(); assert_eq!(result, name_emoji, "Push name should support emoji"); info!("Emoji push name set successfully"); @@ -172,7 +172,7 @@ async fn test_set_push_name_special_characters() -> anyhow::Result<()> { let name_russian = "Тест"; info!("Setting push name with Russian: '{}'...", name_russian); client.client.profile().set_push_name(name_russian).await?; - let result = client.client.get_push_name().await; + let result = client.client.get_push_name(); assert_eq!(result, name_russian, "Push name should support Cyrillic"); info!("Russian push name set successfully"); @@ -180,7 +180,7 @@ async fn test_set_push_name_special_characters() -> anyhow::Result<()> { let name_mixed = "Test™ User©"; info!("Setting push name with special chars: '{}'...", name_mixed); client.client.profile().set_push_name(name_mixed).await?; - let result = client.client.get_push_name().await; + let result = client.client.get_push_name(); assert_eq!( result, name_mixed, "Push name should support special characters" @@ -204,7 +204,7 @@ async fn test_set_push_name_long() -> anyhow::Result<()> { let long_name = "A".repeat(25); info!("Setting push name with {} characters...", long_name.len()); client.client.profile().set_push_name(&long_name).await?; - let result = client.client.get_push_name().await; + let result = client.client.get_push_name(); assert_eq!(result, long_name, "Push name should support 25 characters"); info!("Long push name set successfully"); @@ -231,7 +231,7 @@ async fn test_set_push_name_whitespace_only() -> anyhow::Result<()> { .profile() .set_push_name(whitespace_name) .await?; - let result = client.client.get_push_name().await; + let result = client.client.get_push_name(); assert_eq!( result, whitespace_name, "Whitespace-only push name should be accepted (only empty is rejected)" @@ -263,7 +263,6 @@ async fn test_status_text_notification_received() -> anyhow::Result<()> { let jid_a = client_a .client .get_pn() - .await .expect("Client A should have a JID") .to_non_ad(); @@ -303,7 +302,7 @@ async fn test_set_push_name_persists_across_operations() -> anyhow::Result<()> { let push_name = "PersistBot"; info!("Setting push name to '{}'...", push_name); client.client.profile().set_push_name(push_name).await?; - let result = client.client.get_push_name().await; + let result = client.client.get_push_name(); assert_eq!(result, push_name); info!("Push name set successfully"); @@ -317,7 +316,7 @@ async fn test_set_push_name_persists_across_operations() -> anyhow::Result<()> { info!("Status text set successfully"); // Verify push name is still correct - let after_status = client.client.get_push_name().await; + let after_status = client.client.get_push_name(); assert_eq!( after_status, push_name, "Push name should persist after setting status text" diff --git a/tests/e2e/tests/profile_picture.rs b/tests/e2e/tests/profile_picture.rs index edd7ad40d..55fba402f 100644 --- a/tests/e2e/tests/profile_picture.rs +++ b/tests/e2e/tests/profile_picture.rs @@ -23,7 +23,6 @@ async fn test_set_profile_picture() -> anyhow::Result<()> { let own_jid = client .client .get_pn() - .await .expect("should have PN after pairing"); info!("Fetching profile picture for own JID: {}", own_jid); let pic = client @@ -96,7 +95,6 @@ async fn test_set_profile_picture_then_update() -> anyhow::Result<()> { let own_jid = client .client .get_pn() - .await .expect("should have PN after pairing"); let pic = client .client @@ -139,7 +137,6 @@ async fn test_remove_profile_picture() -> anyhow::Result<()> { let own_jid = client .client .get_pn() - .await .expect("should have PN after pairing"); let expected_jid = own_jid.to_non_ad(); let event = client @@ -182,7 +179,6 @@ async fn test_get_nonexistent_profile_picture() -> anyhow::Result<()> { let own_jid = client .client .get_pn() - .await .expect("should have PN after pairing"); let pic = client .client @@ -217,7 +213,6 @@ async fn test_get_contact_profile_picture() -> anyhow::Result<()> { let jid_b = client_b .client .get_pn() - .await .expect("B should have PN") .to_non_ad(); let pic = client_a @@ -257,7 +252,6 @@ async fn test_get_profile_picture_preview_and_full() -> anyhow::Result<()> { let own_jid = client .client .get_pn() - .await .expect("should have PN after pairing"); // Fetch preview diff --git a/tests/e2e/tests/session_reuse.rs b/tests/e2e/tests/session_reuse.rs index c4ef23726..022882a7e 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -157,7 +157,7 @@ async fn test_session_state_after_roundtrip() -> anyhow::Result<()> { // LID sessions: the active sessions used by encrypt_for_devices let mut lid_sessions = Vec::new(); - if let Some(lid) = client_b.client.get_lid().await { + if let Some(lid) = client_b.client.get_lid() { lid_sessions = scan_sessions(&*backend, &lid.user, "lid").await?; for (addr, pending) in &lid_sessions { info!("LID session {addr}: pending_pre_key={pending}"); @@ -235,7 +235,7 @@ async fn test_session_persistence() -> anyhow::Result<()> { // PN→LID mapping was resolved before encryption. let mut post_send = scan_sessions(&*backend, &jid_b.user, "c.us").await?; if post_send.is_empty() - && let Some(lid_b) = client_b.client.get_lid().await + && let Some(lid_b) = client_b.client.get_lid() { post_send = scan_sessions(&*backend, &lid_b.user, "lid").await?; } diff --git a/tests/handshake_integration.rs b/tests/handshake_integration.rs index c4ea52d25..234665cb8 100644 --- a/tests/handshake_integration.rs +++ b/tests/handshake_integration.rs @@ -324,7 +324,7 @@ async fn cold_start_xx_then_cached_ik_reconnect() { ); // The cert chain must now be cached on the device. - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); let chain = device .server_cert_chain .as_ref() @@ -366,7 +366,7 @@ async fn cold_start_xx_then_cached_ik_reconnect() { // After IK Continue, the on-disk cache stays as-is (the orchestrator // does NOT issue SetServerCertChain on the IK path). - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); assert_eq!( device.server_cert_chain.as_ref().unwrap().leaf.key, server_static_pub_expected @@ -470,7 +470,7 @@ async fn post_xxfallback_failure_does_not_invalidate_ik_cache() { // Cache must be untouched: the failure happened post-pivot, so the // orchestrator's invalidation gate must have skipped both the clear // and the counter increment. - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); let chain = device .server_cert_chain .as_ref() @@ -532,7 +532,7 @@ async fn ik_continue_does_not_overwrite_cached_chain() { result.err() ); - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); let chain = device .server_cert_chain .as_ref() @@ -574,7 +574,7 @@ async fn xx_after_pair_success_persists_cert_chain() { .expect("unpaired XX must succeed"); task1.await.unwrap(); - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); assert!(!device.is_registered(), "still unpaired after first XX"); assert!( device.server_cert_chain.is_none(), @@ -603,7 +603,7 @@ async fn xx_after_pair_success_persists_cert_chain() { .expect("paired XX must succeed"); task2.await.unwrap(); - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); assert!(device.is_registered(), "paired after SetId"); let chain = device .server_cert_chain @@ -650,7 +650,7 @@ async fn unpaired_xx_does_not_persist_cert_chain() { result.err() ); - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); assert!( !device.is_registered(), "precondition: device must still be unpaired" @@ -785,7 +785,7 @@ async fn ik_rejected_recovers_via_xxfallback_and_repopulates_cache() { // the fresh chain (same key here, but the orchestration MUST emit // SetServerCertChain regardless). assert_eq!(counter.load(std::sync::atomic::Ordering::Acquire), 0); - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); let chain = device.server_cert_chain.as_ref().expect("cert chain"); assert_eq!(chain.leaf.key, server_static_pub_expected); } @@ -883,7 +883,7 @@ async fn ik_with_stale_cache_invalidates_and_increments_counter() { 1, "ik_handshake_failures must be 1 after one crypto-fatal failure" ); - let device = pm.get_device_snapshot().await; + let device = pm.get_device_snapshot(); assert!( device.server_cert_chain.is_none(), "stale cert chain must be cleared after crypto-fatal IK failure" From 98b092139304799f147b87843eb5914bfb2a3c44 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 9 Jun 2026 20:07:00 +0000 Subject: [PATCH 2/3] fix(store): clippy needless-borrows; mark dirty before snapshot rebuild MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The borrowed pn/lid bindings made the explicit &x at call sites a double reference (clippy needless_borrow under -D warnings). Review follow-up: set the dirty flag immediately after the modifier, before the snapshot clone — a shutdown flush racing that window checked dirty, saw clean, and could exit without persisting the committed mutation. The flush's device read lock still serializes against the held write guard, so it always saves post-mutation state. --- src/features/media_reupload.rs | 2 +- src/send.rs | 20 ++++++++++---------- src/store/persistence_manager.rs | 6 +++++- 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/features/media_reupload.rs b/src/features/media_reupload.rs index 31f45bdf7..cc8b8e89a 100644 --- a/src/features/media_reupload.rs +++ b/src/features/media_reupload.rs @@ -82,7 +82,7 @@ impl<'a> MediaReupload<'a> { // Build and send the receipt node let receipt_node = build_media_retry_receipt( - &own_jid, + own_jid, req.msg_id, req.chat_jid, req.is_from_me, diff --git a/src/send.rs b/src/send.rs index 04de1c1f4..6de51fb70 100644 --- a/src/send.rs +++ b/src/send.rs @@ -593,7 +593,7 @@ impl Client { } lid_to_pn_map.insert(own_lid.user.clone(), own_jid.to_non_ad()); - let participants = wacore::send::assemble_status_participants(resolved, &own_lid)?; + let participants = wacore::send::assemble_status_participants(resolved, own_lid)?; let mut group_info = GroupInfo::with_lid_to_pn_map(participants, AddressingMode::Lid, lid_to_pn_map); @@ -629,7 +629,7 @@ impl Client { let skdm_target_devices: Option> = if force_skdm { None } else { - self.resolve_skdm_targets(&to_str, &group_info, &own_lid) + self.resolve_skdm_targets(&to_str, &group_info, own_lid) .await .map(|(_all, needs)| needs) }; @@ -665,8 +665,8 @@ impl Client { &mut stores, self, &group_info, - &own_jid, - &own_lid, + own_jid, + own_lid, account_info.as_deref(), to.clone(), &message, @@ -710,8 +710,8 @@ impl Client { &mut stores_retry, self, &group_info, - &own_jid, - &own_lid, + own_jid, + own_lid, account_info.as_deref(), to.clone(), &message, @@ -1387,8 +1387,8 @@ impl Client { &mut stores, self, &group_info, - &own_jid, - &own_lid, + own_jid, + own_lid, account_info.as_deref(), to.clone(), message, @@ -1438,8 +1438,8 @@ impl Client { &mut stores_retry, self, &group_info, - &own_jid, - &own_lid, + own_jid, + own_lid, account_info.as_deref(), to, message, diff --git a/src/store/persistence_manager.rs b/src/store/persistence_manager.rs index b1862edb5..2e9328bd3 100644 --- a/src/store/persistence_manager.rs +++ b/src/store/persistence_manager.rs @@ -98,6 +98,11 @@ impl PersistenceManager { let mut device_guard = self.device.write().await; let result = modifier(&mut device_guard); + // Dirty BEFORE the snapshot rebuild: a shutdown flush racing this + // window must see the store dirty, or it would exit clean and drop + // the committed mutation (the clone below is not free). + self.dirty.store(true, Ordering::Relaxed); + // Rebuild while still holding the write guard so no reader can // observe post-mutation effects with a pre-mutation snapshot. *self @@ -106,7 +111,6 @@ impl PersistenceManager { .unwrap_or_else(|p| p.into_inner()) = Arc::new(device_guard.clone()); drop(device_guard); - self.dirty.store(true, Ordering::Relaxed); self.save_notify.notify(1); result From e79f1c6fa24754c10c5021a25c41adae9462edee Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 9 Jun 2026 20:14:14 +0000 Subject: [PATCH 3/3] refactor(client): make generate_message_id sync; single snapshot per test chain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups: generate_message_id lost its only await point with the sync snapshot, so drop the async (call sites updated). The test JID fallback chains re-read the snapshot mid-chain — besides defeating the point-in-time contract, .or(...) evaluated the second snapshot + clone eagerly; they now take one snapshot and borrow. --- src/client/app_state.rs | 2 +- src/client/messaging.rs | 2 +- src/message/tests.rs | 51 ++++++++++++++++++++----------------- src/pdo.rs | 2 +- src/request.rs | 2 +- src/send.rs | 6 ++--- tests/e2e/tests/receipts.rs | 2 +- 7 files changed, 35 insertions(+), 32 deletions(-) diff --git a/src/client/app_state.rs b/src/client/app_state.rs index 43d53b16f..480f8db7e 100644 --- a/src/client/app_state.rs +++ b/src/client/app_state.rs @@ -605,7 +605,7 @@ impl Client { self.send_message_impl( own_jid, &msg, - Some(self.generate_message_id().await), + Some(self.generate_message_id()), true, false, None, diff --git a/src/client/messaging.rs b/src/client/messaging.rs index 792da84ea..0428cabe2 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -184,7 +184,7 @@ impl Client { server_id: u64, reaction: &str, ) -> Result<(), anyhow::Error> { - let request_id = self.generate_message_id().await; + let request_id = self.generate_message_id(); let stanza = NodeBuilder::new("message") .attr("to", to) diff --git a/src/message/tests.rs b/src/message/tests.rs index 117a1793c..0ca1b04c1 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -1032,27 +1032,29 @@ async fn test_badmac_preserves_session() { let (bob_bundle, _) = bobs_prekey_bundle(&client).await; alice .install_bob_session( - &client - .persistence_manager - .get_device_snapshot() - .lid - .clone() - .or(client.persistence_manager.get_device_snapshot().pn.clone()) - .expect("own jid") - .to_protocol_address(), + &{ + let snapshot = client.persistence_manager.get_device_snapshot(); + snapshot + .lid + .as_ref() + .or(snapshot.pn.as_ref()) + .expect("own jid") + .to_protocol_address() + }, &bob_bundle, ) .await; // First message: pkmsg lands on Bob and installs Bob's reciprocal session. - let bob_addr = client - .persistence_manager - .get_device_snapshot() - .lid - .clone() - .or(client.persistence_manager.get_device_snapshot().pn.clone()) - .expect("own jid") - .to_protocol_address(); + let bob_addr = { + let snapshot = client.persistence_manager.get_device_snapshot(); + snapshot + .lid + .as_ref() + .or(snapshot.pn.as_ref()) + .expect("own jid") + .to_protocol_address() + }; let pkmsg = alice.encrypt_text(&bob_addr, "hello").await; let (s1, _, _, still1) = submit_and_check_session(&client, &alice.jid, &pkmsg).await; assert!(s1, "pkmsg should establish session and decrypt"); @@ -1198,14 +1200,15 @@ async fn test_prod_scenario_pkmsg_archives_old_session_after_badmac() { // X3DH round 1 — Alice initiates with Bob's bundle, sends pkmsg. let (bundle_v1, _) = bobs_prekey_bundle(&client).await; - let bob_addr = client - .persistence_manager - .get_device_snapshot() - .lid - .clone() - .or(client.persistence_manager.get_device_snapshot().pn.clone()) - .expect("own jid") - .to_protocol_address(); + let bob_addr = { + let snapshot = client.persistence_manager.get_device_snapshot(); + snapshot + .lid + .as_ref() + .or(snapshot.pn.as_ref()) + .expect("own jid") + .to_protocol_address() + }; alice.install_bob_session(&bob_addr, &bundle_v1).await; let pkmsg_v1 = alice.encrypt_text(&bob_addr, "v1").await; let (s1, _, _, _) = submit_and_check_session(&client, &alice.jid, &pkmsg_v1).await; diff --git a/src/pdo.rs b/src/pdo.rs index bb39b6072..dbaf9b08a 100644 --- a/src/pdo.rs +++ b/src/pdo.rs @@ -241,7 +241,7 @@ impl Client { to: Jid, msg: &wa::Message, ) -> Result { - let msg_id = self.generate_message_id().await; + let msg_id = self.generate_message_id(); // Send with peer category and high priority self.send_message_impl( diff --git a/src/request.rs b/src/request.rs index 9377a3f91..e28722bc8 100644 --- a/src/request.rs +++ b/src/request.rs @@ -88,7 +88,7 @@ impl Client { /// # Returns /// /// A string containing the generated message ID in the format expected by WhatsApp. - pub async fn generate_message_id(&self) -> String { + pub fn generate_message_id(&self) -> String { let device_snapshot = self.persistence_manager.get_device_snapshot(); self.get_request_utils() .generate_message_id(device_snapshot.pn.as_ref()) diff --git a/src/send.rs b/src/send.rs index 6de51fb70..f9118f938 100644 --- a/src/send.rs +++ b/src/send.rs @@ -464,7 +464,7 @@ impl Client { let stanza_type_override = options.stanza_type_override; let request_id = match options.message_id { Some(id) => id, - None => self.generate_message_id().await, + None => self.generate_message_id(), }; // Both paths below consume `to` and `request_id`, so save copies for the result. let result = SendResult { @@ -545,7 +545,7 @@ impl Client { wacore::telemetry::send("status"); let to = Jid::status_broadcast(); - let request_id = self.generate_message_id().await; + let request_id = self.generate_message_id(); // Borrow from the held snapshot: no field clones, the Arc keeps it alive. let device_snapshot = self.persistence_manager.get_device_snapshot(); @@ -1232,7 +1232,7 @@ impl Client { // Generate request ID early (doesn't need lock) let request_id = match request_id_override { Some(id) => id, - None => self.generate_message_id().await, + None => self.generate_message_id(), }; // `request_id` is moved into the branch-specific stanza builders below; // keep a copy for the post-send messageSecret persistence (the secret diff --git a/tests/e2e/tests/receipts.rs b/tests/e2e/tests/receipts.rs index b816287c5..2db102481 100644 --- a/tests/e2e/tests/receipts.rs +++ b/tests/e2e/tests/receipts.rs @@ -373,7 +373,7 @@ async fn test_delivery_receipts_flushed_on_disconnect() -> anyhow::Result<()> { let mut msg_ids: Vec = Vec::with_capacity(N); let mut receipt_waiters = Vec::with_capacity(N); for i in 0..N { - let id = client_a.client.generate_message_id().await; + let id = client_a.client.generate_message_id(); let receipt_waiter = client_b .client .wait_for_sent_node(NodeFilter::tag("receipt").attr("id", id.clone()));