diff --git a/src/cache.rs b/src/cache.rs index f469809da..dcf7cfcaf 100644 --- a/src/cache.rs +++ b/src/cache.rs @@ -5,3 +5,16 @@ //! `get_with`) that builds on every target, including wasm32. pub use crate::portable_cache::PortableCache as Cache; + +/// Selects whether an operation may use an existing snapshot or must refresh it +/// from its authoritative source before returning. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +#[non_exhaustive] +pub enum Freshness { + /// Return a cached snapshot when available and consult the source on a miss. + #[default] + CachePreferred, + /// Consult the source and publish the resulting snapshot without clearing the + /// previous one first. + Refresh, +} diff --git a/src/client/app_state.rs b/src/client/app_state.rs index d7137a32a..5669a26c9 100644 --- a/src/client/app_state.rs +++ b/src/client/app_state.rs @@ -1030,12 +1030,11 @@ impl Client { self.send_message_impl( peer, msg, - Some(self.generate_message_id()), - true, - false, - None, - Vec::new(), - None, + crate::send::SendPipelineOptions { + request_id: Some(self.generate_message_id()), + peer: true, + ..Default::default() + }, ) .await } diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index dedd18e93..0fda96b60 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -132,8 +132,12 @@ impl Client { // doesn't hold there, so it is disabled and every send resolves. if !self.group_devices_memo_enabled { return Ok(Arc::new(wacore::send::ResolvedGroupDevices::new( - self.resolve_group_devices_uncached(group_info, own_sending_jid) - .await?, + self.resolve_group_devices_uncached( + group_info, + own_sending_jid, + crate::cache::Freshness::CachePreferred, + ) + .await?, ))); } // Load the generation BEFORE resolving (do NOT move this after @@ -175,7 +179,11 @@ impl Client { } let devices = self - .resolve_group_devices_uncached(group_info, own_sending_jid) + .resolve_group_devices_uncached( + group_info, + own_sending_jid, + crate::cache::Freshness::CachePreferred, + ) .await?; // Member identifiers in both namespaces, so the scoped-invalidation @@ -221,10 +229,11 @@ impl Client { /// (participants + LID normalization, appending self when the server /// snapshot omitted it — mirroring `ensure_self_in_group`, so keying the /// memo off the pre-ensure Arc stays equivalent) and resolve it. - async fn resolve_group_devices_uncached( + pub(crate) async fn resolve_group_devices_uncached( &self, - group_info: &Arc, + group_info: &wacore::client::context::GroupInfo, own_sending_jid: &Jid, + freshness: crate::cache::Freshness, ) -> Result, anyhow::Error> { let is_lid_mode = group_info.addressing_mode == wacore::types::message::AddressingMode::Lid; let mut jids_to_resolve: Vec = group_info @@ -256,7 +265,12 @@ impl Client { jids_to_resolve.push(own); } - let mut devices = self.get_user_devices(&jids_to_resolve).await?; + let mut devices = match freshness { + crate::cache::Freshness::CachePreferred => { + self.get_user_devices_owned(jids_to_resolve).await? + } + crate::cache::Freshness::Refresh => self.refresh_user_devices(jids_to_resolve).await?, + }; if is_lid_mode { // WA Web expects LID addressing in SKDM nodes for LID groups. devices = devices @@ -394,8 +408,17 @@ impl Client { ) )] pub(crate) async fn update_device_list( + &self, + record: wacore::store::traits::DeviceListRecord, + ) -> Result<()> { + let guard = self.device_topology.lock_registry().await; + self.update_device_list_guarded(record, &guard).await + } + + pub(crate) async fn update_device_list_guarded( &self, mut record: wacore::store::traits::DeviceListRecord, + guard: &crate::client::device_topology::DeviceRegistryMutationGuard<'_>, ) -> Result<()> { use anyhow::Context; @@ -413,6 +436,7 @@ impl Client { // (whose member set only knows the PN side) would re-stamp stale. self.device_registry_cache .insert( + guard, canonical_key.clone(), Arc::new(record_for_cache), lookup @@ -433,14 +457,18 @@ impl Client { // gets cleared. Run the second invalidate unconditionally: even // if delete fails, the cache may have been repopulated with data // that no longer reflects our intent. - self.device_registry_cache.invalidate(&original_user).await; + self.device_registry_cache + .invalidate(guard, &original_user) + .await; if let Err(e) = backend.delete_devices(&original_user).await { warn!( "Failed to delete stale device row under {} after canonical flip: {e}", original_user ); } - self.device_registry_cache.invalidate(&original_user).await; + self.device_registry_cache + .invalidate(guard, &original_user) + .await; debug!( "Device registry: stored under LID {} (resolved from {})", canonical_key, original_user @@ -455,10 +483,20 @@ impl Client { /// collapses into a single transaction. Used by usync after fetching /// device lists for many users at once, where the per-row commit /// dominated wall-clock time on large groups. + #[cfg(test)] #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.session.update_device_lists", level = "debug", skip_all, fields(count = records.len()), err(Debug)))] pub(crate) async fn update_device_lists( &self, records: Vec, + ) -> Result<()> { + let guard = self.device_topology.lock_registry().await; + self.update_device_lists_guarded(records, &guard).await + } + + pub(crate) async fn update_device_lists_guarded( + &self, + records: Vec, + guard: &crate::client::device_topology::DeviceRegistryMutationGuard<'_>, ) -> Result<()> { use anyhow::Context; @@ -479,6 +517,7 @@ impl Client { // Same alias rule as update_device_list: record every lookup key. self.device_registry_cache .insert( + guard, canonical_key.clone(), Arc::new(record_for_cache), lookup @@ -504,14 +543,18 @@ impl Client { // rather than batching deletes. On error we log and continue so a // single bad row doesn't drop the rest of the batch. for original_user in to_delete { - self.device_registry_cache.invalidate(&original_user).await; + self.device_registry_cache + .invalidate(guard, &original_user) + .await; if let Err(e) = backend.delete_devices(&original_user).await { warn!( "Failed to delete stale device row under {} after canonical flip: {e}", original_user ); } - self.device_registry_cache.invalidate(&original_user).await; + self.device_registry_cache + .invalidate(guard, &original_user) + .await; } Ok(()) @@ -553,10 +596,19 @@ impl Client { ) )] pub(crate) async fn invalidate_device_cache(&self, user: &str) { + let guard = self.device_topology.lock_registry().await; + self.invalidate_device_cache_guarded(user, &guard).await; + } + + pub(crate) async fn invalidate_device_cache_guarded( + &self, + user: &str, + guard: &crate::client::device_topology::DeviceRegistryMutationGuard<'_>, + ) { let lookup = self.resolve_lookup_keys(user).await; for key in lookup.all_keys() { - self.device_registry_cache.invalidate(key).await; + self.device_registry_cache.invalidate(guard, key).await; // Also delete from DB so get_devices_from_registry doesn't // fall back to stale persisted data — forces a network re-fetch if let Err(e) = self.persistence_manager.backend().delete_devices(key).await { @@ -566,7 +618,7 @@ impl Client { // the doomed DB row can promote() it back between the first // invalidate and the delete commit (same guard as the canonical // flip path in update_device_list). - self.device_registry_cache.invalidate(key).await; + self.device_registry_cache.invalidate(guard, key).await; } debug!("Invalidated device cache for user: {} ({:?})", user, lookup); @@ -595,6 +647,7 @@ impl Client { device: &wacore::stanza::devices::DeviceElement, key_index_info: Option<&wacore::stanza::devices::KeyIndexInfo>, ) { + let guard = self.device_topology.lock_registry().await; let device_id = device.device_id(); let is_hosted = wacore_binary::JidExt::is_hosted(&device.jid); @@ -623,8 +676,7 @@ impl Client { } else { // Filter stale devices by valid_indexes. A raw_id reset already // removed every companion while preserving primary metadata. - record.devices = - wacore::adv::filter_devices_by_key_index(&record.devices, &decoded); + wacore::adv::retain_devices_by_key_index(&mut record.devices, &decoded); } record.raw_id = Some(decoded.raw_id); @@ -666,7 +718,7 @@ impl Client { // unknown device → `device_has_key()` returns `None` → falls into // `needs_skdm`. No global cache invalidation needed. - if let Err(e) = self.update_device_list(record).await { + if let Err(e) = self.update_device_list_guarded(record, &guard).await { warn!("patch_device_add: failed to persist: {e}"); } } @@ -763,6 +815,7 @@ impl Client { if device_id == 0 { return; } + let guard = self.device_topology.lock_registry().await; if let Some(mut record) = self.load_device_record(user).await { let before = record.devices.len(); record.devices.retain(|d| d.device_id != device_id); @@ -775,7 +828,7 @@ impl Client { "patch_device_remove: device_id {device_id} > u16::MAX — skipping \ session/SKDM cleanup but still persisting registry removal" ); - if let Err(e) = self.update_device_list(record).await { + if let Err(e) = self.update_device_list_guarded(record, &guard).await { warn!("patch_device_remove: failed to persist: {e}"); } return; @@ -800,7 +853,7 @@ impl Client { ); return; } - if let Err(e) = self.update_device_list(record).await { + if let Err(e) = self.update_device_list_guarded(record, &guard).await { warn!("patch_device_remove: failed to persist: {e}"); } } @@ -865,13 +918,14 @@ impl Client { user: &str, device: &wacore::stanza::devices::DeviceElement, ) { + let guard = self.device_topology.lock_registry().await; let device_id = device.device_id(); if let Some(mut record) = self.load_device_record(user).await && let Some(d) = record.devices.iter_mut().find(|d| d.device_id == device_id) { d.key_index = device.key_index; - if let Err(e) = self.update_device_list(record).await { + if let Err(e) = self.update_device_list_guarded(record, &guard).await { warn!("patch_device_update: failed to persist: {e}"); } } @@ -998,6 +1052,7 @@ impl Client { ) )] pub(crate) async fn migrate_device_registry_on_lid_discovery(&self, pn: &str, lid: &str) { + let guard = self.device_topology.lock_registry().await; let backend = self.persistence_manager.backend(); match backend.get_devices(pn).await { @@ -1015,13 +1070,13 @@ impl Client { // The backend row may have changed even on error, so the // change is recorded before the early return; the success // path records once via the fused cache insert below. - self.device_topology.record([pn, lid]); + self.device_topology.record_registry(&guard, [pn, lid]); warn!("Failed to migrate device registry to LID: {}", e); return; } self.device_registry_cache - .insert(lid.to_string(), Arc::new(record), [lid, pn]) + .insert(&guard, lid.to_string(), Arc::new(record), [lid, pn]) .await; // Drop the PN-keyed row in both cache and DB. Invalidate @@ -1029,11 +1084,11 @@ impl Client { // resurrect the cache from the DB row between the two calls. // Always run the second invalidate; even if delete fails, the // cache may carry resurrected data that shouldn't stick. - self.device_registry_cache.invalidate(pn).await; + self.device_registry_cache.invalidate(&guard, pn).await; if let Err(e) = backend.delete_devices(pn).await { warn!("Failed to delete PN-keyed device row during LID migration: {e}"); } - self.device_registry_cache.invalidate(pn).await; + self.device_registry_cache.invalidate(&guard, pn).await; } Ok(None) => {} Err(e) => { diff --git a/src/client/device_topology.rs b/src/client/device_topology.rs index 826e9091e..ad874aafc 100644 --- a/src/client/device_topology.rs +++ b/src/client/device_topology.rs @@ -25,7 +25,7 @@ use wacore_binary::CompactString; const TOPOLOGY_LOG_CAPACITY: usize = 256; struct TopologyLog { - /// (generation that the change produced, canonical user touched). + /// (generation, canonical user touched). entries: VecDeque<(u64, CompactString)>, /// Highest generation evicted from `entries` (0 = nothing evicted). /// A memo older than this cannot be proven clean and must recompute. @@ -36,6 +36,14 @@ struct TopologyLog { pub(crate) struct DeviceTopology { generation: AtomicU64, log: std::sync::Mutex, + registry_mutation: async_lock::Mutex<()>, +} + +/// Proof that a device-registry mutation is serialized with authoritative +/// refresh publication. The cache write API requires this token so new write +/// paths cannot accidentally bypass the ordering invariant. +pub(crate) struct DeviceRegistryMutationGuard<'a> { + _guard: async_lock::MutexGuard<'a, ()>, } impl DeviceTopology { @@ -46,6 +54,7 @@ impl DeviceTopology { entries: VecDeque::with_capacity(TOPOLOGY_LOG_CAPACITY), floor: 0, }), + registry_mutation: async_lock::Mutex::new(()), }) } @@ -53,10 +62,20 @@ impl DeviceTopology { self.generation.load(Ordering::Acquire) } + pub(crate) async fn lock_registry(&self) -> DeviceRegistryMutationGuard<'_> { + DeviceRegistryMutationGuard { + _guard: self.registry_mutation.lock().await, + } + } + /// Record one topology change touching the given users (pass BOTH /// namespaces of an identity when known: a mapping change alters which /// canonical record either key resolves to). pub(crate) fn record<'a>(&self, users: impl IntoIterator) { + self.record_change(users); + } + + fn record_change<'a>(&self, users: impl IntoIterator) { let mut log = self.log.lock().unwrap_or_else(|p| p.into_inner()); let generation = self.generation.load(Ordering::Acquire) + 1; for user in users { @@ -74,6 +93,16 @@ impl DeviceTopology { self.generation.store(generation, Ordering::Release); } + /// Record a registry mutation while holding its serialization guard, so a + /// refresh can compare-and-publish without a check/write race. + pub(crate) fn record_registry<'a>( + &self, + _guard: &DeviceRegistryMutationGuard<'_>, + users: impl IntoIterator, + ) { + self.record_change(users); + } + /// Record a change whose blast radius is unknown (bulk warm-up, cache /// clear): bumps and poisons the scoped fast path so every memo /// recomputes once. @@ -130,17 +159,18 @@ impl DeviceRegistryCache { /// canonical flipped). pub(crate) async fn insert<'a>( &self, + guard: &DeviceRegistryMutationGuard<'_>, key: String, record: Arc, touched: impl IntoIterator, ) { self.cache.insert(key, record).await; - self.topology.record(touched); + self.topology.record_registry(guard, touched); } - pub(crate) async fn invalidate(&self, key: &str) { + pub(crate) async fn invalidate(&self, guard: &DeviceRegistryMutationGuard<'_>, key: &str) { self.cache.invalidate(key).await; - self.topology.record([key]); + self.topology.record_registry(guard, [key]); } /// Cache-fill from the DB row the fallback path would have returned: the @@ -179,4 +209,9 @@ impl DeviceRegistryCache { ) { self.cache.insert(key, record).await; } + + #[cfg(test)] + pub(crate) async fn raw_invalidate_for_tests(&self, key: &str) { + self.cache.invalidate(key).await; + } } diff --git a/src/client/lid_pn.rs b/src/client/lid_pn.rs index b35d208c0..e9ee39025 100644 --- a/src/client/lid_pn.rs +++ b/src/client/lid_pn.rs @@ -292,12 +292,34 @@ impl Client { mappings: Vec<(String, String)>, source: LearningSource, is_offline: bool, + ) { + let outcome = self.record_lid_pn_batch_in_memory(mappings, source).await; + self.finish_lid_pn_batch_learning(outcome, is_offline); + } + + pub(crate) async fn learn_lid_pn_mappings_batch_guarded( + self: &Arc, + mappings: Vec<(String, String)>, + source: LearningSource, + is_offline: bool, + guard: &crate::lid_pn_cache::LidPnMutationGuard<'_>, + ) { + let outcome = self + .record_lid_pn_batch_in_memory_guarded(mappings, source, guard) + .await; + self.finish_lid_pn_batch_learning(outcome, is_offline); + } + + fn finish_lid_pn_batch_learning( + self: &Arc, + outcome: BatchRecordOutcome, + is_offline: bool, ) { let BatchRecordOutcome { entries, migration_flags, usync_phones, - } = self.record_lid_pn_batch_in_memory(mappings, source).await; + } = outcome; // Conflicting observational pairs re-resolve live, independent of the // flush gate (WA Web fires syncContactListJob regardless of @@ -382,6 +404,17 @@ impl Client { &self, mappings: Vec<(String, String)>, source: LearningSource, + ) -> BatchRecordOutcome { + let guard = self.lid_pn_cache.lock_mutation().await; + self.record_lid_pn_batch_in_memory_guarded(mappings, source, &guard) + .await + } + + async fn record_lid_pn_batch_in_memory_guarded( + &self, + mappings: Vec<(String, String)>, + source: LearningSource, + guard: &crate::lid_pn_cache::LidPnMutationGuard<'_>, ) -> BatchRecordOutcome { let cap = mappings.len(); let mut deduped: std::collections::HashMap = @@ -395,7 +428,7 @@ impl Client { let mut usync_phones: Vec = Vec::new(); for (phone_number, lid) in deduped { match self - .record_lid_pn_in_memory(&lid, &phone_number, source) + .record_lid_pn_in_memory_guarded(&lid, &phone_number, source, guard) .await { RecordOutcome::Skipped => {} @@ -424,6 +457,18 @@ impl Client { lid: &str, phone_number: &str, source: LearningSource, + ) -> RecordOutcome { + let guard = self.lid_pn_cache.lock_mutation().await; + self.record_lid_pn_in_memory_guarded(lid, phone_number, source, &guard) + .await + } + + async fn record_lid_pn_in_memory_guarded( + &self, + lid: &str, + phone_number: &str, + source: LearningSource, + guard: &crate::lid_pn_cache::LidPnMutationGuard<'_>, ) -> RecordOutcome { // Fully durable and resolvable both ways: nothing to re-add or persist. if self.lid_pn_cache.can_skip_relearn(phone_number, lid).await { @@ -451,7 +496,7 @@ impl Client { }; return match existing { Some(entry) => { - self.lid_pn_cache.add(&entry).await; + self.lid_pn_cache.add_guarded(&entry, guard).await; RecordOutcome::Written { entry, needs_migration, @@ -472,7 +517,7 @@ impl Client { wacore::time::now_secs() }; let entry = LidPnEntry::with_timestamp(lid, phone_number, created_at, source); - self.lid_pn_cache.add(&entry).await; + self.lid_pn_cache.add_guarded(&entry, guard).await; return RecordOutcome::Written { entry, needs_migration: current_lid.is_none(), @@ -1048,23 +1093,37 @@ impl Client { /// (`Ok(None)`) from "lookup failed" (`Err(_)`). #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.session.get_lid_pn_entry", level = "trace", skip_all, fields(peer = %jid.observe()), err(Debug)))] pub async fn get_lid_pn_entry(&self, jid: &Jid) -> Result> { - let (hit, is_lid) = if jid.is_lid() { - (self.lid_pn_cache.get_entry_by_lid(&jid.user).await, true) + let is_lid = if jid.is_lid() { + true } else if jid.is_pn() { - (self.lid_pn_cache.get_entry_by_phone(&jid.user).await, false) + false } else { return Ok(None); }; + self.get_lid_pn_entry_by_user(&jid.user, is_lid).await + } + + async fn get_lid_pn_entry_by_user( + &self, + user: &str, + is_lid: bool, + ) -> Result> { + let hit = if is_lid { + self.lid_pn_cache.get_entry_by_lid(user).await + } else { + self.lid_pn_cache.get_entry_by_phone(user).await + }; + if let Some(entry) = hit { return Ok(Some(entry)); } let backend = self.persistence_manager.backend(); let mapping = if is_lid { - backend.get_lid_mapping(&jid.user).await? + backend.get_lid_mapping(user).await? } else { - backend.get_pn_mapping(&jid.user).await? + backend.get_pn_mapping(user).await? }; let Some(mapping) = mapping else { @@ -1076,6 +1135,43 @@ impl Client { Ok(Some(entry)) } + /// Whether two user JIDs identify the same account, ignoring device + /// suffixes and resolving PN/LID aliases through the canonical mapping + /// cache-aside path. Hosted namespaces belong to their corresponding PN or + /// LID family; unrelated namespaces only match exactly. + pub(crate) async fn jids_share_user_identity(&self, left: &Jid, right: &Jid) -> Result { + if left.is_same_chat_as(right) { + return Ok(true); + } + + let same_user_and_integrator = + left.user == right.user && left.integrator == right.integrator; + if same_user_and_integrator + && ((left.server.is_pn_family() && right.server.is_pn_family()) + || (left.server.is_lid_family() && right.server.is_lid_family())) + { + return Ok(true); + } + + let (lid, pn) = if left.server.is_lid_family() && right.server.is_pn_family() { + (left, right) + } else if right.server.is_lid_family() && left.server.is_pn_family() { + (right, left) + } else { + return Ok(false); + }; + if lid.integrator != pn.integrator { + return Ok(false); + } + + Ok(self + .get_lid_pn_entry_by_user(&lid.user, true) + .await? + .is_some_and(|mapping| { + &*mapping.lid == lid.user.as_str() && &*mapping.phone_number == pn.user.as_str() + })) + } + /// Resolve any user JID to its bare LID form, or `None` when no LID is /// available. Mirrors WA Web's `WAWebLidMigrationUtils.toUserLid`: LID /// passes through, PN goes through the cache-aside mapping, anything diff --git a/src/client/messaging.rs b/src/client/messaging.rs index 5aeefd891..2b06e928e 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -108,12 +108,10 @@ impl Client { self.send_message_impl( to, &edit_container_message, - None, - false, - false, - Some(crate::types::message::EditAttribute::MessageEdit), - vec![], - None, + crate::send::SendPipelineOptions { + edit: Some(crate::types::message::EditAttribute::MessageEdit), + ..Default::default() + }, ) .await .map_err(crate::send::SendError::from_anyhow)?; @@ -196,12 +194,10 @@ impl Client { self.send_message_impl( to, &envelope, - None, - false, - false, - Some(crate::types::message::EditAttribute::MessageEdit), - vec![], - None, + crate::send::SendPipelineOptions { + edit: Some(crate::types::message::EditAttribute::MessageEdit), + ..Default::default() + }, ) .await .map_err(SendError::from_anyhow)?; diff --git a/src/client/sender_keys.rs b/src/client/sender_keys.rs index 6488e74cf..5f3d5580a 100644 --- a/src/client/sender_keys.rs +++ b/src/client/sender_keys.rs @@ -30,20 +30,20 @@ impl Client { .and_then(|s| s.pn.as_ref()) .map(|j| j.user.as_str()); - let kept: Vec<&Jid> = device_jids + let keep = |jid: &&Jid| { + !exclude_own_devices + || !(own_lid_user.is_some_and(|user| user == jid.user) + || own_pn_user.is_some_and(|user| user == jid.user)) + }; + let device_ids: Vec = device_jids .iter() - .filter(|jid| { - !exclude_own_devices - || !(own_lid_user.is_some_and(|u| u == jid.user) - || own_pn_user.is_some_and(|u| u == jid.user)) - }) + .filter(keep) + .map(ToString::to_string) .collect(); - - if kept.is_empty() { + if device_ids.is_empty() { return Ok(()); } - let device_ids: Vec = kept.iter().map(|jid| jid.to_string()).collect(); let entries: Vec<(&str, bool)> = device_ids.iter().map(|s| (s.as_str(), has_key)).collect(); self.persistence_manager .set_sender_key_status(group_jid, &entries) @@ -61,9 +61,8 @@ impl Client { // the skdm_warm_memo compares, so a warm send re-runs its target // filter and re-sends the now-cold device's SKDM — no separate memo // invalidation, hence no cross-cache ordering window. - let jids: Vec = kept.into_iter().cloned().collect(); self.sender_key_device_cache - .mark_forgotten(group_jid, &jids) + .mark_forgotten(group_jid, device_jids.iter().filter(keep)) .await; } Ok(()) diff --git a/src/features/groups.rs b/src/features/groups.rs index fbb99205b..149e2dec9 100644 --- a/src/features/groups.rs +++ b/src/features/groups.rs @@ -287,6 +287,75 @@ pub struct Groups<'a> { client: &'a Client, } +/// Serializes one group's metadata publication with participant mutations and +/// sender-key distribution, reusing the client's existing per-group lane. +/// Keeping the guard in the type prevents persistence and cache publication +/// from accidentally being split across an unlocked await. +pub(crate) struct GroupMetadataGuard<'a> { + client: &'a Client, + jid: &'a Jid, + _guard: async_lock::MutexGuardArc<()>, +} + +impl GroupMetadataGuard<'_> { + pub(crate) async fn current(&self) -> Option> { + self.client.get_group_cache().await.get(self.jid).await + } + + async fn cache(&self, info: Arc) { + self.client + .get_group_cache() + .await + .insert(self.jid.clone(), info) + .await; + } + + pub(crate) async fn publish(&self, info: Arc) { + let jid = self.jid.to_string(); + match serde_json::to_vec(info.as_ref()) { + Ok(blob) => { + if let Err(error) = self + .client + .persistence_manager + .backend() + .put_group_metadata(&jid, &blob) + .await + { + log::warn!("Failed to persist group metadata for {}: {error}", self.jid); + } + } + Err(error) => { + log::warn!( + "Failed to serialize group metadata for {}: {error}", + self.jid + ); + } + } + + self.cache(info).await; + } + + pub(crate) async fn invalidate(&self) { + if let Err(error) = self + .client + .persistence_manager + .backend() + .delete_group_metadata(&self.jid.to_string()) + .await + { + log::warn!( + "Failed to invalidate persisted group metadata for {}: {error}", + self.jid + ); + } + self.client + .get_group_cache() + .await + .invalidate(self.jid) + .await; + } +} + #[derive(Clone, Copy)] enum ParticipantRemovalScope { Group, @@ -299,109 +368,179 @@ impl<'a> Groups<'a> { } pub async fn query_info(&self, jid: &Jid) -> Result, GroupError> { - if let Some(cached) = self.client.get_group_cache().await.get(jid).await { - return Ok(cached); - } - - // Send the persisted participant phash (WA Web queryGroup phash) so the - // server can answer "not-modified" by omitting for an unchanged - // group, letting us reuse the persisted metadata instead of re-parsing it. - let jid_str = jid.to_string(); - let backend = self.client.persistence_manager.backend(); - let persisted: Option = match backend.get_group_metadata(&jid_str).await { - Ok(Some(blob)) => serde_json::from_slice(&blob).ok(), - _ => None, - }; - let phash = persisted.as_ref().and_then(|info| { - wacore::messages::MessageUtils::participant_list_hash(&info.participants).ok() - }); + self.query_info_with_freshness(jid, crate::cache::Freshness::CachePreferred) + .await + } - let group = match self - .client - .execute(GroupQueryIq::with_phash(jid, phash)) - .await? + /// Query group metadata using the requested cache freshness policy. + /// + /// A refresh leaves the current snapshot readable while the network request + /// is in flight, then atomically replaces it after a successful response. + pub async fn query_info_with_freshness( + &self, + jid: &Jid, + freshness: crate::cache::Freshness, + ) -> Result, GroupError> { + let cache = self.client.get_group_cache().await; + let mut cached = cache.get(jid).await; + if freshness == crate::cache::Freshness::CachePreferred + && let Some(cached) = cached.take() { - GroupInfoOutcome::NotModified => { - let info = Arc::new(persisted.ok_or_else(|| { - GroupError::InvalidRequest( - "server returned not-modified group but nothing was cached".into(), - ) - })?); - self.client - .get_group_cache() - .await - .insert(jid.clone(), info.clone()) - .await; - return Ok(info); - } - GroupInfoOutcome::Full(group) => *group, - }; - - // Single pass: move participants out and build lid_to_pn_map alongside. - let n = group.participants.len(); - let is_lid = group.addressing_mode == AddressingMode::Lid; - let mut participants: Vec = Vec::with_capacity(n); - let mut lid_to_pn_map: HashMap = if is_lid { - HashMap::with_capacity(n) - } else { - HashMap::new() - }; - for p in group.participants { - if is_lid && let Some(pn) = p.phone_number { - lid_to_pn_map.insert(p.jid.user.clone(), pn); - } - participants.push(p.jid); + return Ok(cached); } - // Populate lid_pn_cache so silent-observer participants (no messages - // from them) get their mapping; otherwise `invalidate_device_cache` - // can't resolve the PN alias and leaves zombie registry entries. - // One batched call mirrors WA Web's single `createLidPnMappings` - // invocation from `QueryGroupJob`, so N participants = 1 persist - // task + 1 DB transaction instead of N detached tasks. - if !lid_to_pn_map.is_empty() - && let Some(client_arc) = self.client.self_weak.get().and_then(|w| w.upgrade()) - { - let mut batch: Vec<(String, String)> = Vec::with_capacity(lid_to_pn_map.len()); - for (lid_user, pn_jid) in &lid_to_pn_map { - if pn_jid.is_pn() { - batch.push((lid_user.as_str().to_string(), pn_jid.user.to_string())); + self.query_info_from_source(jid, cached).await + } + + #[expect( + clippy::manual_async_fn, + reason = "the explicit async block keeps the network-bound state machine out of line" + )] + fn query_info_from_source<'b>( + &'b self, + jid: &'b Jid, + mut cached: Option>, + ) -> impl Future, GroupError>> + 'b { + // Keep the large, network-bound state machine shared between refresh + // and cache-miss callers. The cache-hit fast path stays in + // `query_info_with_freshness`, while this boundary prevents LTO from + // cloning the slow path for each statically known freshness policy. + #[inline(never)] + async move { + let jid_str = jid.to_string(); + let backend = self.client.persistence_manager.backend(); + loop { + // Send the persisted participant phash (WA Web queryGroup phash) so + // the server can omit for an unchanged snapshot. On a cold + // L1, keep the lane through the request: without an Arc to compare, + // this is the only way to distinguish "still absent" from a + // notification that invalidated an already-absent snapshot. + let (persisted, mut cold_metadata) = if cached.is_some() { + (None, None) + } else { + let metadata = self.client.lock_group_metadata(jid).await; + if let Some(current) = metadata.current().await { + cached = Some(current); + drop(metadata); + continue; + } + let persisted = match backend.get_group_metadata(&jid_str).await { + Ok(Some(blob)) => serde_json::from_slice(&blob).ok(), + _ => None, + }; + (persisted, Some(metadata)) + }; + let phash = cached.as_deref().or(persisted.as_ref()).and_then(|info| { + wacore::messages::MessageUtils::participant_list_hash(&info.participants).ok() + }); + + let group = match self + .client + .execute(GroupQueryIq::with_phash(jid, phash)) + .await? + { + GroupInfoOutcome::NotModified => { + if let Some(metadata) = cold_metadata.take() { + let info = Arc::new(persisted.ok_or_else(|| { + GroupError::InvalidRequest( + "server returned not-modified group but nothing was cached" + .into(), + ) + })?); + metadata.cache(Arc::clone(&info)).await; + return Ok(info); + } + + // Participant mutations use the same per-group lane, so the + // snapshot and its persisted blob cannot change between this + // check and the decision below. + let metadata = self.client.lock_group_metadata(jid).await; + if let Some(current) = metadata.current().await { + return Ok(current); + } + + // The warm snapshot used for the conditional request was + // invalidated while the IQ was in flight. Retry without it + // instead of resurrecting pre-notification membership. + drop(metadata); + cached = None; + continue; + } + GroupInfoOutcome::Full(group) => *group, + }; + + // Single pass: move participants out and build lid_to_pn_map alongside. + let participant_count = group.participants.len(); + let is_lid = group.addressing_mode == AddressingMode::Lid; + let mut participants = Vec::with_capacity(participant_count); + let mut lid_to_pn_map: HashMap = if is_lid { + HashMap::with_capacity(participant_count) + } else { + HashMap::new() + }; + for participant in group.participants { + if is_lid && let Some(pn) = participant.phone_number { + lid_to_pn_map.insert(participant.jid.user.clone(), pn); + } + participants.push(participant.jid); } - } - client_arc - .learn_lid_pn_mappings_batch( - batch, - crate::lid_pn_cache::LearningSource::Other, - false, - ) - .await; - } - let mut info = GroupInfo::new(participants, group.addressing_mode); - info.is_community_announce = Some(group.is_default_sub_group); - if !lid_to_pn_map.is_empty() { - info.set_lid_to_pn_map(lid_to_pn_map); - } + // Populate lid_pn_cache so silent-observer participants (no messages + // from them) get their mapping; otherwise `invalidate_device_cache` + // can't resolve the PN alias and leaves zombie registry entries. + if !lid_to_pn_map.is_empty() + && let Some(client_arc) = self.client.self_weak.get().and_then(|w| w.upgrade()) + { + let mut batch = Vec::with_capacity(lid_to_pn_map.len()); + for (lid_user, pn_jid) in &lid_to_pn_map { + if pn_jid.is_pn() { + batch.push((lid_user.as_str().to_string(), pn_jid.user.to_string())); + } + } + client_arc + .learn_lid_pn_mappings_batch( + batch, + crate::lid_pn_cache::LearningSource::Other, + false, + ) + .await; + } - // Persist so the next query can send this group's participant phash and - // skip the full re-query when membership is unchanged. - match serde_json::to_vec(&info) { - Ok(blob) => { - if let Err(e) = backend.put_group_metadata(&jid_str, &blob).await { - log::warn!("Failed to persist group metadata for {jid}: {e}"); + let mut info = GroupInfo::new(participants, group.addressing_mode); + info.is_community_announce = Some(group.is_default_sub_group); + if !lid_to_pn_map.is_empty() { + info.set_lid_to_pn_map(lid_to_pn_map); } + let info = Arc::new(info); + + // Compare and publish while holding the same lane as participant + // mutations. Persisting inside the guard prevents the durable blob + // and L1 snapshot from being committed in opposite orders. + let metadata = match cold_metadata { + Some(metadata) => metadata, + None => { + let metadata = self.client.lock_group_metadata(jid).await; + let current = metadata.current().await; + let unchanged = matches!( + (cached.as_ref(), current.as_ref()), + (Some(expected), Some(current)) if Arc::ptr_eq(expected, current) + ); + if !unchanged { + drop(metadata); + if let Some(current) = current { + return Ok(current); + } + cached = None; + continue; + } + metadata + } + }; + + metadata.publish(Arc::clone(&info)).await; + return Ok(info); } - Err(e) => log::warn!("Failed to serialize group metadata for {jid}: {e}"), } - - let info = Arc::new(info); - self.client - .get_group_cache() - .await - .insert(jid.clone(), info.clone()) - .await; - - Ok(info) } /// Backfills each LID participant's `phone_number` from the client's LID-PN @@ -561,18 +700,11 @@ impl<'a> Groups<'a> { pub async fn leave(&self, jid: impl Into) -> Result<(), GroupError> { let jid = &jid.into(); self.client.execute(LeaveGroupIq::new(jid)).await?; - self.client.get_group_cache().await.invalidate(jid).await; - // Drop the persisted blob too: we're no longer in the group, so a stale - // phash from it would only force a needless full re-query if ever read. - if let Err(e) = self - .client - .persistence_manager - .backend() - .delete_group_metadata(&jid.to_string()) + self.client + .lock_group_metadata(jid) .await - { - log::warn!("Failed to delete persisted group metadata for {jid}: {e}"); - } + .invalidate() + .await; Ok(()) } @@ -596,8 +728,8 @@ impl<'a> Groups<'a> { let result = self.client.execute(iq).await?; if result.iter().any(|r| r.is_ok()) { - let group_cache = self.client.get_group_cache().await; - if let Some(info) = group_cache.get(jid).await { + let metadata = self.client.lock_group_metadata(jid).await; + if let Some(info) = metadata.current().await { let mut info = Arc::unwrap_or_clone(info); info.add_participants( result @@ -605,11 +737,10 @@ impl<'a> Groups<'a> { .filter(|r| r.is_ok()) .map(|r| (&r.jid, r.phone_number.as_ref())), ); - self.client.persist_group_metadata(jid, &info).await; - group_cache.insert(jid.clone(), Arc::new(info)).await; + metadata.publish(Arc::new(info)).await; } else { // Cache expired: can't patch in place, so drop the now-stale blob. - self.client.invalidate_persisted_group_metadata(jid).await; + metadata.invalidate().await; } } Ok(result) @@ -642,17 +773,16 @@ impl<'a> Groups<'a> { .map(|r| r.jid.user.as_str()) .collect(); if !accepted.is_empty() { - let group_cache = self.client.get_group_cache().await; match scope { ParticipantRemovalScope::Group => { - if let Some(info) = group_cache.get(jid).await { + let metadata = self.client.lock_group_metadata(jid).await; + if let Some(info) = metadata.current().await { let mut info = Arc::unwrap_or_clone(info); info.remove_participants(&accepted); - self.client.persist_group_metadata(jid, &info).await; - group_cache.insert(jid.clone(), Arc::new(info)).await; + metadata.publish(Arc::new(info)).await; } else { // Cache expired: can't patch in place, so drop the now-stale blob. - self.client.invalidate_persisted_group_metadata(jid).await; + metadata.invalidate().await; } } ParticipantRemovalScope::LinkedGroups => { @@ -662,8 +792,11 @@ impl<'a> Groups<'a> { // carry the affected JIDs, patch their own cache entries, // and rotate their sender-key chains without evicting // unrelated groups. - group_cache.invalidate(jid).await; - self.client.invalidate_persisted_group_metadata(jid).await; + self.client + .lock_group_metadata(jid) + .await + .invalidate() + .await; } } self.client @@ -1147,12 +1280,11 @@ impl<'a> Groups<'a> { .send_message_impl( group_jid.clone(), &msg, - Some(message_id.clone()), - false, - false, - None, - meta.into_iter().collect(), - None, + crate::send::SendPipelineOptions { + request_id: Some(message_id.clone()), + extra_stanza_nodes: meta.into_iter().collect(), + ..Default::default() + }, ) .await?; Ok(message_id) @@ -1266,35 +1398,11 @@ impl Client { Groups::new(self) } - /// Re-serialize and persist a group's metadata after a local membership change - /// so the phash fast-path stays consistent: the in-memory cache expires after - /// ~1h, after which a stale persisted blob would force a needless full re-query - /// (or be compared against the server as an out-of-date phash). Shared by the - /// participant-mutation API and the inbound group-notification handler. - pub(crate) async fn persist_group_metadata(&self, jid: &Jid, info: &GroupInfo) { - let backend = self.persistence_manager.backend(); - match serde_json::to_vec(info) { - Ok(blob) => { - if let Err(e) = backend.put_group_metadata(&jid.to_string(), &blob).await { - log::warn!("Failed to persist group metadata for {jid}: {e}"); - } - } - Err(e) => log::warn!("Failed to serialize group metadata for {jid}: {e}"), - } - } - - /// Drop the persisted group metadata on a membership change we can't patch in - /// place (the in-memory cache had already expired), so the next query re-fetches - /// fresh instead of comparing a now-stale phash. Without this, persisting only on - /// a cache hit would miss the exact post-expiry case this fix targets. - pub(crate) async fn invalidate_persisted_group_metadata(&self, jid: &Jid) { - if let Err(e) = self - .persistence_manager - .backend() - .delete_group_metadata(&jid.to_string()) - .await - { - log::warn!("Failed to invalidate persisted group metadata for {jid}: {e}"); + pub(crate) async fn lock_group_metadata<'a>(&'a self, jid: &'a Jid) -> GroupMetadataGuard<'a> { + GroupMetadataGuard { + client: self, + jid, + _guard: self.group_distribution_lock(jid).await, } } } @@ -1536,6 +1644,33 @@ mod tests { assert_eq!(a.participants.len(), 2); } + #[tokio::test] + async fn refresh_keeps_the_previous_group_snapshot_on_source_failure() { + let client = crate::test_utils::create_test_client().await; + let group: Jid = "120363000000000099@g.us".parse().unwrap(); + let previous = Arc::new(GroupInfo::new( + vec!["12025550101@s.whatsapp.net".parse().unwrap()], + AddressingMode::Pn, + )); + let cache = client.get_group_cache().await; + cache.insert(group.clone(), Arc::clone(&previous)).await; + + let result = client + .groups() + .query_info_with_freshness(&group, crate::cache::Freshness::Refresh) + .await; + assert!( + result.is_err(), + "the offline fixture proves refresh consulted the source" + ); + + let preserved = cache + .get(&group) + .await + .expect("refresh failure must not clear the current snapshot"); + assert!(Arc::ptr_eq(&previous, &preserved)); + } + #[tokio::test] async fn linked_removal_preserves_unrelated_group_cache_entries() { use wacore::protocol::ProtocolNode; @@ -1544,7 +1679,7 @@ mod tests { let client = crate::test_utils::create_test_client().await; let parent: Jid = "120363000000000001@g.us".parse().unwrap(); let unrelated: Jid = "120363000000000002@g.us".parse().unwrap(); - let removed: Jid = "15550000001@s.whatsapp.net".parse().unwrap(); + let removed: Jid = "12025550103@s.whatsapp.net".parse().unwrap(); let cache = client.get_group_cache().await; for jid in [&parent, &unrelated] { cache @@ -1590,7 +1725,11 @@ mod tests { .is_some() ); - client.invalidate_persisted_group_metadata(&group_jid).await; + client + .lock_group_metadata(&group_jid) + .await + .invalidate() + .await; assert!( backend diff --git a/src/features/mod.rs b/src/features/mod.rs index ec34d93c3..a38fec25e 100644 --- a/src/features/mod.rs +++ b/src/features/mod.rs @@ -80,8 +80,8 @@ pub use status::{Status, StatusPrivacySetting, StatusSendOptions}; pub use signal::{Signal, SignalError, SignalSessionInfo, SignalSessionMigration}; pub(crate) use stanza::required_stanza_attr; pub use stanza::{ - NackReason, RetryReason, RetryRequestError, RetryRequestOptions, RetryRequestOutcome, - StanzaRejection, StanzaResponseError, + MessageRetransmission, NackReason, RetryReason, RetryRequestError, RetryRequestOptions, + RetryRequestOutcome, StanzaRejection, StanzaResponseError, }; pub use wacore::message_processing::EncType; diff --git a/src/features/stanza.rs b/src/features/stanza.rs index f75a5d1f2..2a0385eb0 100644 --- a/src/features/stanza.rs +++ b/src/features/stanza.rs @@ -2,7 +2,10 @@ use thiserror::Error; +use crate::cache::Freshness; use crate::client::ClientError; +use wacore_binary::Jid; +use waproto::whatsapp as wa; pub(crate) fn required_stanza_attr<'node, 'data>( node: &'node wacore_binary::NodeRef<'data>, @@ -158,6 +161,84 @@ pub enum RetryRequestError { Internal(#[from] anyhow::Error), } +/// A request to retransmit an already-sent message to one requesting device. +/// +/// The client derives the wire stanza and owns all routing, encryption, session, +/// sender-key, and persistence decisions. This type never accepts a pre-built +/// retry stanza. +#[derive(Debug)] +#[non_exhaustive] +pub struct MessageRetransmission { + pub(crate) chat: Jid, + pub(crate) requester: Jid, + pub(crate) message: wa::Message, + pub(crate) message_id: String, + pub(crate) retry_count: u8, + pub(crate) recipient: Option, + pub(crate) group_metadata_freshness: Freshness, +} + +impl MessageRetransmission { + /// Describe a retransmission to the device that requested it. + pub fn new( + chat: Jid, + requester: Jid, + message: wa::Message, + message_id: String, + retry_count: u8, + ) -> Self { + Self { + chat, + requester, + message, + message_id, + retry_count, + recipient: None, + group_metadata_freshness: Freshness::CachePreferred, + } + } + + /// Preserve the receipt's recipient for self-device and bot retry routes. + pub fn with_recipient(mut self, recipient: Jid) -> Self { + self.recipient = Some(recipient); + self + } + + /// Select how group metadata is obtained for this operation. + pub fn with_group_metadata_freshness(mut self, freshness: Freshness) -> Self { + self.group_metadata_freshness = freshness; + self + } + + pub fn chat(&self) -> &Jid { + &self.chat + } + + pub fn requester(&self) -> &Jid { + &self.requester + } + + pub fn message(&self) -> &wa::Message { + &self.message + } + + pub fn message_id(&self) -> &str { + &self.message_id + } + + pub const fn retry_count(&self) -> u8 { + self.retry_count + } + + pub fn recipient(&self) -> Option<&Jid> { + self.recipient.as_ref() + } + + pub const fn group_metadata_freshness(&self) -> Freshness { + self.group_metadata_freshness + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/features/status.rs b/src/features/status.rs index f1714c9a2..5dea6d304 100644 --- a/src/features/status.rs +++ b/src/features/status.rs @@ -2,9 +2,11 @@ use wacore::WireEnum; use wacore_binary::Jid; use waproto::whatsapp as wa; +use crate::cache::Freshness; use crate::client::Client; use crate::send::{SendError, SendResult}; use crate::upload::UploadResponse; +use wacore_binary::Node; /// Privacy setting sent in the `` node of the status stanza. /// Matches WhatsApp Web's `status_setting` attribute. @@ -28,6 +30,12 @@ pub enum StatusPrivacySetting { pub struct StatusSendOptions { /// Privacy setting for this status. Sent in the `` stanza node. pub privacy: StatusPrivacySetting, + /// Override the generated message ID. + pub message_id: Option, + /// Extra child nodes appended to the status stanza. + pub extra_stanza_nodes: Vec, + /// Freshness policy for the recipient device lists used by this send. + pub device_freshness: Freshness, } /// High-level API for WhatsApp status/story updates. diff --git a/src/handlers/notification/groups.rs b/src/handlers/notification/groups.rs index 731e38d3f..bf36903c5 100644 --- a/src/handlers/notification/groups.rs +++ b/src/handlers/notification/groups.rs @@ -135,21 +135,6 @@ pub(crate) async fn handle_group_notification(client: &Arc, node: Arc Some(cache.await), - None => None, - }; let action_count = actions.len(); for (action_index, action) in actions.into_iter().enumerate() { @@ -158,22 +143,15 @@ pub(crate) async fn handle_group_notification(client: &Arc, node: Arc { - let group_cache = group_cache - .as_ref() - .expect("participant actions initialize the group cache"); - if let Some(info) = group_cache.get(¬ification.group_jid).await { + let metadata = client.lock_group_metadata(¬ification.group_jid).await; + if let Some(info) = metadata.current().await { let mut info = Arc::unwrap_or_clone(info); info.add_participants( participants .iter() .map(|p| (&p.jid, p.phone_number.as_ref())), ); - client - .persist_group_metadata(¬ification.group_jid, &info) - .await; - group_cache - .insert(notification.group_jid.clone(), Arc::new(info)) - .await; + metadata.publish(Arc::new(info)).await; debug!( target: "Client/Group", "Patched group cache for {}: added {} participants", @@ -186,25 +164,16 @@ pub(crate) async fn handle_group_notification(client: &Arc, node: Arc { let users: Vec<&str> = participants.iter().map(|p| p.jid.user.as_str()).collect(); - let group_cache = group_cache - .as_ref() - .expect("participant actions initialize the group cache"); - if let Some(info) = group_cache.get(¬ification.group_jid).await { + let metadata = client.lock_group_metadata(¬ification.group_jid).await; + if let Some(info) = metadata.current().await { let mut info = Arc::unwrap_or_clone(info); info.remove_participants(&users); - client - .persist_group_metadata(¬ification.group_jid, &info) - .await; - group_cache - .insert(notification.group_jid.clone(), Arc::new(info)) - .await; + metadata.publish(Arc::new(info)).await; debug!( target: "Client/Group", "Patched group cache for {}: removed {} participants", @@ -217,10 +186,9 @@ pub(crate) async fn handle_group_notification(client: &Arc, node: Arc, node: Arc here would let the - // next send rebuild the sender key against the old participant - // list (missing the migrated JID, still targeting the old one). - group_cache - .as_ref() - .expect("participant actions initialize the group cache") - .invalidate(¬ification.group_jid) - .await; + let metadata = client.lock_group_metadata(¬ification.group_jid).await; + metadata.invalidate().await; + drop(metadata); client .force_rotate_own_sender_key(¬ification.group_jid) .await; diff --git a/src/lib.rs b/src/lib.rs index 3e733ccb6..0ed27905c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -67,6 +67,7 @@ pub use wacore::chrono; pub use waproto::buffa; pub mod cache; +pub use cache::Freshness; pub mod portable_cache; pub(crate) mod resend_rate_limiter; @@ -150,13 +151,13 @@ pub use features::{ GroupParticipantOptions, GroupProfilePicture, GroupSubject, GroupType, Groups, GrowthLockInfo, InviteInfoError, IsOnWhatsAppResult, JoinGroupResult, Labels, LinkSubgroupsResult, MediaRetryResult, MediaReupload, MediaReuploadError, MediaReuploadRequest, MemberAddMode, - MemberLinkMode, MemberShareHistoryMode, MembershipApprovalMode, MembershipRequest, Mex, - MexError, MexErrorExtensions, MexRequest, MexResponse, NackReason, Newsletter, NewsletterError, - NewsletterMessage, NewsletterMessageType, NewsletterMetadata, NewsletterReactionCount, - NewsletterRole, NewsletterState, NewsletterVerification, ParticipantChangeResponse, - ParticipantType, PictureType, PollError, Presence, PresenceError, PresenceStatus, Profile, - ProfileError, ProfilePicture, ReachoutTimelock, RetryReason, RetryRequestError, - RetryRequestOptions, RetryRequestOutcome, SecretEncKind, SecretEncrypted, + MemberLinkMode, MemberShareHistoryMode, MembershipApprovalMode, MembershipRequest, + MessageRetransmission, Mex, MexError, MexErrorExtensions, MexRequest, MexResponse, NackReason, + Newsletter, NewsletterError, NewsletterMessage, NewsletterMessageType, NewsletterMetadata, + NewsletterReactionCount, NewsletterRole, NewsletterState, NewsletterVerification, + ParticipantChangeResponse, ParticipantType, PictureType, PollError, Presence, PresenceError, + PresenceStatus, Profile, ProfileError, ProfilePicture, ReachoutTimelock, RetryReason, + RetryRequestError, RetryRequestOptions, RetryRequestOutcome, SecretEncKind, SecretEncrypted, SetProfilePictureResponse, Signal, SignalError, SignalSessionInfo, SignalSessionMigration, StanzaRejection, StanzaResponseError, Status, StatusPrivacySetting, StatusSendOptions, SyncActionMessageRange, TcToken, TcTokenError, UnlinkSubgroupsResult, UserInfo, diff --git a/src/lid_pn_cache.rs b/src/lid_pn_cache.rs index cf94f0cd9..6036d3825 100644 --- a/src/lid_pn_cache.rs +++ b/src/lid_pn_cache.rs @@ -46,6 +46,11 @@ pub struct LidPnCache { lid_to_entry: TypedCache, Arc>, /// Phone number -> Entry mapping (stores the most recent LID for that PN) pn_to_entry: TypedCache, Arc>, + /// Serializes mapping writers. Reads remain lock-free; the guard exists so + /// an authoritative device refresh can compare topology and publish its + /// mappings plus registry snapshot without a concurrent writer slipping + /// between those steps. + mutation: async_lock::Mutex<()>, /// Device-topology tracker (attached by Client construction): a mapping /// change alters which canonical record either key resolves to, so adds /// record both identifiers. Recording lives here, at the write @@ -61,6 +66,11 @@ pub struct LidPnCache { persisted: TypedCache, Arc>, } +/// Proof that LID/PN mapping mutations are serialized. +pub(crate) struct LidPnMutationGuard<'a> { + _guard: async_lock::MutexGuard<'a, ()>, +} + impl Default for LidPnCache { fn default() -> Self { Self::new() @@ -87,6 +97,7 @@ impl LidPnCache { Some(s) => Self { lid_to_entry: TypedCache::from_store(s.clone(), NS_LID, config.timeout), pn_to_entry: TypedCache::from_store(s, NS_PN, config.timeout), + mutation: async_lock::Mutex::new(()), // Always in-memory: tracks per-process persist state, never the // mapping itself, so it must not go through the shared store. persisted: TypedCache::from_local(config.build_with_tti()), @@ -95,6 +106,7 @@ impl LidPnCache { None => Self { lid_to_entry: TypedCache::from_local(config.build_with_tti()), pn_to_entry: TypedCache::from_local(config.build_with_tti()), + mutation: async_lock::Mutex::new(()), persisted: TypedCache::from_local(config.build_with_tti()), topology: std::sync::OnceLock::new(), }, @@ -111,6 +123,12 @@ impl LidPnCache { let _ = self.topology.set(topology); } + pub(crate) async fn lock_mutation(&self) -> LidPnMutationGuard<'_> { + LidPnMutationGuard { + _guard: self.mutation.lock().await, + } + } + /// Approximate entry counts plus estimated retained bytes for the LID and /// PN maps. Bytes are `0` when backed by a custom store (entries live /// outside this process). @@ -211,11 +229,14 @@ impl LidPnCache { /// For the PN -> Entry map, this only updates if the new entry has a /// newer or equal `created_at` timestamp (matching WhatsApp Web behavior). /// - /// Note: the get-then-insert on the PN map is not atomic. With external - /// backends (e.g., Redis), concurrent `add()` calls for the same phone - /// number can race. This is acceptable because the cache is best-effort - /// and backed by persistent storage for correctness. + /// The writer guard keeps the get-then-insert sequence ordered within this + /// client, including when the cache uses an external backend. pub async fn add(&self, entry: &LidPnEntry) { + let guard = self.lock_mutation().await; + self.add_guarded(entry, &guard).await; + } + + pub(crate) async fn add_guarded(&self, entry: &LidPnEntry, _guard: &LidPnMutationGuard<'_>) { let should_update_pn = match self.pn_to_entry.get(&*entry.phone_number).await { Some(existing) => existing.created_at <= entry.created_at, None => true, @@ -262,9 +283,10 @@ impl LidPnCache { pub async fn warm_up(&self, entries: impl IntoIterator) { let start = wacore::time::Instant::now(); let mut count = 0; + let guard = self.lock_mutation().await; for entry in entries { - self.add(&entry).await; + self.add_guarded(&entry, &guard).await; // `warm_up` only accepts durable rows. Mark the pair that won the // PN-side timestamp resolution so a live re-learn neither writes // it again nor repeats discovery migrations. @@ -291,6 +313,7 @@ impl LidPnCache { /// Awaits the actual clear operation on custom backends (unlike /// `invalidate_all` which is fire-and-forget). pub async fn clear(&self) { + let _guard = self.lock_mutation().await; self.lid_to_entry.clear().await; self.pn_to_entry.clear().await; self.persisted.clear().await; diff --git a/src/message/special.rs b/src/message/special.rs index 16b567fc7..23912fc89 100644 --- a/src/message/special.rs +++ b/src/message/special.rs @@ -363,12 +363,11 @@ impl Client { self.send_message_impl( requester.clone(), message, - Some(message_id.to_owned()), - true, - false, - None, - Vec::new(), - None, + crate::send::SendPipelineOptions { + request_id: Some(message_id.to_owned()), + peer: true, + ..Default::default() + }, ) .await } diff --git a/src/pdo.rs b/src/pdo.rs index 74ebc5f90..a008899a5 100644 --- a/src/pdo.rs +++ b/src/pdo.rs @@ -272,12 +272,11 @@ impl Client { self.send_message_impl( to, msg, - Some(msg_id.clone()), - true, // is_peer_message - false, // is_retry - None, - vec![], // No extra stanza nodes for peer messages - None, + crate::send::SendPipelineOptions { + request_id: Some(msg_id.clone()), + peer: true, + ..Default::default() + }, ) .await?; diff --git a/src/retry.rs b/src/retry.rs index 8060c5481..3ffbef9cd 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -1,6 +1,7 @@ use crate::client::Client; -use crate::features::RetryRequestError; +use crate::features::{MessageRetransmission, RetryRequestError}; use crate::message::RetryReason; +use crate::send::SendError; use crate::types::events::Receipt; use log::{debug, info, warn}; use wacore::types::message::MessageCategory; @@ -18,6 +19,7 @@ use wacore_binary::NodeContent; use wacore_binary::builder::NodeBuilder; use wacore_binary::{Jid, Node, OwnedNodeRef}; use wacore_binary::{NodeContentRef, NodeRef}; +use waproto::whatsapp as wa; /// Helper to extract bytes content from a Node (used in tests). #[cfg(test)] @@ -40,6 +42,140 @@ fn get_bytes_content_ref<'a>(node: &'a NodeRef<'_>) -> Option<&'a [u8]> { /// whatsmeow's `recreateSessionTimeout` (`retry.go:156`). const RECREATE_SESSION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3600); +#[derive(Clone, Copy)] +enum RetransmissionRoute { + Direct, + Group, + Status, + BroadcastList, +} + +impl RetransmissionRoute { + const fn uses_sender_key(self) -> bool { + matches!(self, Self::Group | Self::Status) + } +} + +#[inline] +fn is_own_account_jid(jid: &Jid, own_pn: Option<&Jid>, own_lid: Option<&Jid>) -> bool { + own_pn.is_some_and(|pn| jid.is_same_user_as(pn)) + || own_lid.is_some_and(|lid| jid.is_same_user_as(lid)) +} + +struct PreparedRetransmission { + route: RetransmissionRoute, + chat: Jid, + wire_requester: Jid, + encryption_jid: Jid, + message: wa::Message, + message_id: String, + retry_count: u8, + recipient: Option, + group_info: Option>, + /// Canonical unpadded protobuf bytes shared with the recent-message cache. + /// Public retransmissions provide them; the automatic path may fall back + /// to its already-decoded message when the cache bytes are unavailable. + pre_encoded: Option>>, +} + +fn validate_retransmission( + chat: &Jid, + requester: &Jid, + message_id: &str, + retry_count: u8, + recipient: Option<&Jid>, +) -> Result { + if chat.is_empty() || requester.is_empty() { + return Err(SendError::InvalidRequest( + "retransmission JIDs must not be empty".into(), + )); + } + if message_id.is_empty() { + return Err(SendError::InvalidRequest( + "retransmission message ID must not be empty".into(), + )); + } + if !(1..MAX_RETRY_COUNT).contains(&retry_count) { + return Err(SendError::InvalidRequest(format!( + "retry count must be in 1..{MAX_RETRY_COUNT}" + ))); + } + + let requester_is_user = matches!( + requester.server, + wacore_binary::Server::Pn + | wacore_binary::Server::Lid + | wacore_binary::Server::Hosted + | wacore_binary::Server::HostedLid + | wacore_binary::Server::Bot + ); + if !requester_is_user { + return Err(SendError::InvalidRequest( + "retransmission requester must be a user device JID".into(), + )); + } + + let route = if chat.is_group() { + RetransmissionRoute::Group + } else if chat.is_status_broadcast() { + if !matches!( + requester.server, + wacore_binary::Server::Pn | wacore_binary::Server::Lid + ) { + return Err(SendError::InvalidRequest( + "status retransmission requester must be a PN or LID device".into(), + )); + } + RetransmissionRoute::Status + } else if chat.is_broadcast_list() { + if !matches!( + requester.server, + wacore_binary::Server::Pn | wacore_binary::Server::Lid + ) { + return Err(SendError::InvalidRequest( + "broadcast retransmission requester must be a PN or LID device".into(), + )); + } + RetransmissionRoute::BroadcastList + } else if matches!( + chat.server, + wacore_binary::Server::Pn + | wacore_binary::Server::Lid + | wacore_binary::Server::Hosted + | wacore_binary::Server::HostedLid + | wacore_binary::Server::Bot + ) { + RetransmissionRoute::Direct + } else { + return Err(SendError::InvalidRequest( + "unsupported retransmission chat class".into(), + )); + }; + + if recipient.is_some() && !matches!(route, RetransmissionRoute::Direct) { + return Err(SendError::InvalidRequest( + "recipient is only valid for direct retransmissions".into(), + )); + } + if recipient.is_some_and(|recipient| { + recipient.is_empty() + || !matches!( + recipient.server, + wacore_binary::Server::Pn + | wacore_binary::Server::Lid + | wacore_binary::Server::Hosted + | wacore_binary::Server::HostedLid + | wacore_binary::Server::Bot + ) + }) { + return Err(SendError::InvalidRequest( + "retransmission recipient must be a user JID".into(), + )); + } + + Ok(route) +} + pub(crate) enum RetryReceiptSendOutcome { Sent { included_keys: bool }, Suppressed, @@ -80,8 +216,8 @@ fn resolve_retry_chat_info( ) -> Option { let from = &receipt.source.chat; - if from.is_group() || from.is_status_broadcast() { - // Groups/status: chat is already the group/broadcast JID. + if from.is_group() || from.is_status_broadcast() || from.is_broadcast_list() { + // Group-like chats: chat is already the group/broadcast JID. // Requester is the participant attr (the actual retrying device). let participant = node.attrs().optional_jid("participant"); let is_fbid_bot_retry = @@ -106,8 +242,7 @@ fn resolve_retry_chat_info( // 2. Peer device + recipient → chat = recipient // 3. Peer device without recipient → WA Web aborts (returns null). // 4. Normal user → chat = asUserWidOrThrow(from) = from.to_non_ad() - let is_peer = own_pn.is_some_and(|pn| from.is_same_user_as(pn)) - || own_lid.is_some_and(|lid| from.is_same_user_as(lid)); + let is_peer = is_own_account_jid(from, own_pn, own_lid); let chat = if is_bot && let Some(r) = recipient.as_ref() { r.to_non_ad() @@ -163,6 +298,126 @@ fn build_retry_processing_key(chat: &Jid, message_id: &str, participant_jid: &Ji } impl Client { + async fn resolve_retransmission_encryption_jid( + &self, + route: RetransmissionRoute, + requester: &Jid, + ) -> Result { + if matches!(route, RetransmissionRoute::Status) && requester.is_pn() { + return match self.get_lid_pn_entry(requester).await? { + Some(mapping) => Ok(Jid { + user: wacore_binary::CompactString::new(&mapping.lid), + server: wacore_binary::Server::Lid, + device: requester.device, + agent: requester.agent, + integrator: requester.integrator, + }), + // WAWebResendStatusMsg explicitly falls back to the PN device + // when no LID mapping is available. + None => Ok(requester.clone()), + }; + } + Ok(self.resolve_encryption_jid(requester).await) + } + + /// Retransmit a message to one requesting device. + /// + /// The client derives the stanza from native protocol data and retains + /// ownership of routing, encryption, sender-key tracking, persistence, and + /// transport. The original message ID and retry count are preserved. + pub async fn retransmit_message( + &self, + request: MessageRetransmission, + ) -> Result<(), SendError> { + let route = validate_retransmission( + &request.chat, + &request.requester, + &request.message_id, + request.retry_count, + request.recipient.as_ref(), + )?; + + if matches!(route, RetransmissionRoute::Direct) { + let snapshot = self.persistence_manager.get_device_snapshot(); + let requester_is_local = is_own_account_jid( + &request.requester, + snapshot.pn.as_ref(), + snapshot.lid.as_ref(), + ); + if request.recipient.is_some() { + if !requester_is_local && !request.requester.is_bot() { + return Err(SendError::InvalidRequest( + "a direct retransmission recipient is only valid for a local device or bot" + .into(), + )); + } + } else if requester_is_local { + return Err(SendError::InvalidRequest( + "a direct retransmission to another local device requires a recipient".into(), + )); + } + + let routing_chat = request.recipient.as_ref().unwrap_or(&request.requester); + if !self + .jids_share_user_identity(&request.chat, routing_chat) + .await + .map_err(SendError::from_anyhow)? + { + return Err(SendError::InvalidRequest( + "direct retransmission chat does not match its routing identity".into(), + )); + } + } + + let group_info = if matches!(route, RetransmissionRoute::Group) { + Some( + self.groups() + .query_info_with_freshness(&request.chat, request.group_metadata_freshness) + .await?, + ) + } else { + None + }; + + let encryption_jid = self + .resolve_retransmission_encryption_jid(route, &request.requester) + .await + .map_err(SendError::from_anyhow)?; + if route.uses_sender_key() { + let chat_key = request.chat.to_string(); + self.mark_forget_sender_key(&chat_key, std::slice::from_ref(&encryption_jid)) + .await + .map_err(SendError::from_anyhow)?; + } + + let MessageRetransmission { + chat, + requester: wire_requester, + message, + message_id, + retry_count, + recipient, + group_metadata_freshness: _, + } = request; + let pre_encoded = Arc::new(waproto::codec::message_to_vec(&message)); + self.add_recent_message(&chat, &message_id, &message, Some(Arc::clone(&pre_encoded))) + .await; + self.retransmit_message_prepared(PreparedRetransmission { + route, + wire_requester, + encryption_jid, + chat, + message, + message_id, + retry_count, + recipient, + group_info, + pre_encoded: Some(pre_encoded), + }) + .await + .map_err(SendError::from_anyhow) + } + /// Handle an inbound ``. /// /// WA Web authorizes these through `isRetryEligible` (`WAWebApiMessageInfoStore`). @@ -224,7 +479,20 @@ impl Client { ) else { return Ok(()); }; - let is_group_or_status = info.chat.is_group() || info.chat.is_status_broadcast(); + let route = match validate_retransmission( + &info.chat, + &info.requester, + &message_id, + retry_count, + info.recipient.as_ref(), + ) { + Ok(route) => route, + Err(error) => { + debug!("Ignoring malformed retry request: {error}"); + return Ok(()); + } + }; + let uses_sender_key = route.uses_sender_key(); // WA Web doesn't dedupe receipts (Message/Queue.js just serializes per-chat); // MAX_RETRY_COUNT covers loop prevention. This lock only guards against @@ -301,7 +569,7 @@ impl Client { // WA Web: `e.from.isBot() ? (p = e.from) : (p = d.isLid() ? toLid(e.from) : toPn(e.from))` // Bots skip namespace normalization (WAWebHandleRetryRequest:311-312). let resolved_jid = if let Some(alt_chat) = alt_chat - && !is_group_or_status + && !uses_sender_key && !info.is_bot { let requester = &info.requester; @@ -314,7 +582,8 @@ impl Client { }; info.requester.clone() } else { - self.resolve_encryption_jid(&info.requester).await + self.resolve_retransmission_encryption_jid(route, &info.requester) + .await? }; let keys_node_present = nr.get_optional_child("keys").is_some(); @@ -330,20 +599,17 @@ impl Client { } // Check if this is a retry from our own device (peer). - let is_peer = device_snapshot - .pn - .as_ref() - .is_some_and(|our_pn| info.requester.is_same_user_as(our_pn)) - || device_snapshot - .lid - .as_ref() - .is_some_and(|our_lid| info.requester.is_same_user_as(our_lid)); + let is_peer = is_own_account_jid( + &info.requester, + device_snapshot.pn.as_ref(), + device_snapshot.lid.as_ref(), + ); // Volume-throttling inbound retries diverges from WA Web (which // processes every receipt), so it is an operator opt-in, gated here // before the expensive repair stages below. Own devices (`is_peer`) and // DMs are never gated: dropping their retries has no safe SKDM fallback. - if is_group_or_status + if uses_sender_key && !is_peer && let Some(policy) = self.retry_admission.get() && !policy.admit(&info.chat, &info.requester, retry_count) @@ -378,7 +644,7 @@ impl Client { // force full sender key rotation by clearing all sender key device tracking. // This is separate from updateLocalSignalSession and specific to group retries. let mut rotated_sender_key = false; - if is_group_or_status && !info.requester.is_lid() && !info.chat.is_status_broadcast() { + if matches!(route, RetransmissionRoute::Group) && !info.requester.is_lid() { let group_jid = info.chat.to_string(); let is_known_participant = cached_group_info .as_ref() @@ -486,17 +752,6 @@ impl Client { } } - // Status broadcasts can't resend (requires explicit recipient list). - // Participant already marked for fresh SKDM above; next status send includes them. - if info.chat.is_status_broadcast() { - info!( - "Status broadcast retry for {} — participant marked for fresh SKDM, \ - will be included in next status send", - message_id - ); - return Ok(()); - } - // Bound the aggregate resend rate per group (the anti-abuse signal): a // PN to LID fan-out has many distinct devices retry the same messages, // which per-device/per-message caps miss. Group-only: the requester was @@ -520,94 +775,211 @@ impl Client { retry_count ); - if info.chat.is_group() { - // Group retry: pairwise encrypt to failing device only (RetryMsgJob.js:71). - // Using sender-key broadcast would resend to ALL participants → duplicates. - // - // WA Web calls ensureE2ESessions for all chat types, not just DMs - // (RetryRequest.js:200). Without this, a reg-ID mismatch or unknown - // device whose session was deleted above would fail `prepare_group_retry_stanza` - // with "session not found", silencing subsequent retries via the duplicate filter. - self.ensure_e2e_sessions_resolved(std::slice::from_ref(&resolved_jid)) - .await?; + let wire_requester = if matches!(route, RetransmissionRoute::Direct) { + info.original_from + } else { + info.requester + }; + self.retransmit_message_prepared(PreparedRetransmission { + route, + chat: info.chat, + wire_requester, + encryption_jid: resolved_jid, + message: original_msg, + message_id, + retry_count, + recipient: info.recipient, + group_info: cached_group_info, + pre_encoded: None, + }) + .await?; - let device_snapshot = self.persistence_manager.get_device_snapshot(); + Ok(()) + } - let addressing_mode = cached_group_info - .as_ref() - .map(|g| g.addressing_mode) - .unwrap_or_default(); + async fn send_retry_stanza(&self, stanza: Node) -> Result<(), anyhow::Error> { + self.persist_signal_state_pre_wire().await?; + self.send_node(stanza).await?; + Ok(()) + } - 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; - let mut store_adapter = self.signal_adapter().await; - - let edit_attr = - wacore::types::message::EditAttribute::infer_from_message(&original_msg); - let stanza = wacore::send::prepare_group_retry_stanza( - &mut store_adapter.session_store, - &mut store_adapter.identity_store, - info.chat, - info.requester, - resolved_jid.clone(), - &original_msg, - message_id, - retry_count, - device_snapshot.account.as_deref(), - addressing_mode, - edit_attr, - ) - .await?; + async fn retransmit_message_prepared( + &self, + request: PreparedRetransmission, + ) -> Result<(), anyhow::Error> { + let PreparedRetransmission { + route, + chat, + wire_requester, + encryption_jid, + message, + message_id, + retry_count, + recipient, + group_info, + pre_encoded, + } = request; + + if matches!(route, RetransmissionRoute::Status) { + return self + .retransmit_status_message( + chat, + encryption_jid, + message, + message_id, + pre_encoded.as_deref().map(Vec::as_slice), + ) + .await; + } - // The pre-wire gate can take the processing permit, whose holder - // may need this session lock. - drop(_session_guard); - self.send_retry_stanza(stanza).await?; - } else { - // DM retry: pairwise resend to the requesting device only. - // Use _resolved variant: resolved_jid is already in the correct - // namespace (including alternate PN/LID normalization). - // WA Web's ensureE2ESessions also uses already-normalized JIDs. - self.ensure_e2e_sessions_resolved(std::slice::from_ref(&resolved_jid)) - .await?; + // Every remaining route is pairwise, including broadcast-list + // participants, and shares the normal session recovery path. + self.ensure_e2e_sessions_resolved(std::slice::from_ref(&encryption_jid)) + .await?; + let signal_address = encryption_jid.to_protocol_address(); + let session_mutex = self.session_lock_for(signal_address.as_str()).await; + let session_guard = session_mutex.lock().await; + let mut store_adapter = self.signal_adapter().await; + let device_snapshot = self.persistence_manager.get_device_snapshot(); + let edit = wacore::types::message::EditAttribute::infer_from_message(&message); - 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; - let mut store_adapter = self.signal_adapter().await; - - let edit_attr = - wacore::types::message::EditAttribute::infer_from_message(&original_msg); - // WA Web forwards the receipt's `recipient` verbatim - // (`f && (k.recipient = f)` in handleRetryRequest); for non-self - // DM receipts the attribute is absent and the resend drops it. - let stanza = wacore::send::prepare_dm_retry_stanza( - &mut store_adapter.session_store, - &mut store_adapter.identity_store, - info.original_from, - info.recipient.clone(), - resolved_jid.clone(), - &original_msg, + let destination = match route { + RetransmissionRoute::Direct => wacore::send::PairwiseRetryDestination::Direct { + to: wire_requester, + recipient, + }, + RetransmissionRoute::Group => { + let addressing_mode = group_info + .as_ref() + .map(|info| info.addressing_mode) + .unwrap_or_default(); + wacore::send::PairwiseRetryDestination::Participant { + to: chat, + participant: wire_requester, + addressing_mode: Some(addressing_mode), + } + } + RetransmissionRoute::BroadcastList => { + wacore::send::PairwiseRetryDestination::Participant { + to: chat, + participant: wire_requester, + addressing_mode: None, + } + } + RetransmissionRoute::Status => unreachable!("status handled above"), + }; + let stanza = wacore::send::prepare_pairwise_retry_stanza( + &mut store_adapter.session_store, + &mut store_adapter.identity_store, + wacore::send::PairwiseRetryRequest { + destination, + encryption_jid, + message: &message, message_id, retry_count, - device_snapshot.account.as_deref(), - edit_attr, - ) - .await?; + account: device_snapshot.account.as_deref(), + edit, + pre_encoded: pre_encoded.as_deref().map(Vec::as_slice), + }, + ) + .await?; + + // Persistence may need the processing permit, whose holder may in turn + // need this session lock. Release it before the durability gate. + drop(session_guard); + self.send_retry_stanza(stanza).await + } - // Same lock-ordering rule as the group branch above. - drop(_session_guard); - self.send_retry_stanza(stanza).await?; + /// Rebuild a status message for exactly the requesting device. The retry + /// count remains an operation-level guard; the captured status wire does not + /// encode it on either the skmsg or SKDM `` node. + async fn retransmit_status_message( + &self, + chat: Jid, + requester: Jid, + message: wa::Message, + message_id: String, + pre_encoded: Option<&[u8]>, + ) -> Result<(), anyhow::Error> { + let snapshot = self.persistence_manager.get_device_snapshot(); + let own_pn = snapshot + .pn + .as_ref() + .ok_or(crate::client::ClientError::NotLoggedIn)?; + let own_lid = snapshot + .lid + .as_ref() + .ok_or_else(|| anyhow::anyhow!("cannot retransmit status without a device LID"))?; + let is_sending_device = (requester.is_same_user_as(own_pn) + && requester.device == own_pn.device) + || (requester.is_same_user_as(own_lid) && requester.device == own_lid.device); + if is_sending_device { + anyhow::bail!("cannot retransmit a status to the sending device itself"); } - Ok(()) - } + let chat_key = chat.to_string(); + let distribution_guard = self.group_distribution_lock(&chat).await; + let group_info = wacore::client::context::GroupInfo::new( + Vec::new(), + wacore::types::message::AddressingMode::Lid, + ); - async fn send_retry_stanza(&self, stanza: Node) -> Result<(), anyhow::Error> { - self.persist_signal_state_pre_wire().await?; - self.send_node(stanza).await?; + let can_reuse_encoding = message.message_context_info.is_unset(); + let encoded_fallback = (pre_encoded.is_none() && can_reuse_encoding) + .then(|| waproto::codec::message_to_vec(&message)); + let encoded = pre_encoded + .filter(|_| can_reuse_encoding) + .or(encoded_fallback.as_deref()); + let device_store = self.persistence_manager.get_device_arc().await; + let mut store_adapter = self.signal_adapter_from(device_store); + let mut stores = store_adapter.as_signal_stores(); + let edit = wacore::types::message::EditAttribute::infer_from_message(&message); + let prepared = match wacore::send::prepare_group_stanza( + &*self.runtime, + &mut stores, + self, + wacore::send::GroupStanzaRequest { + group: &group_info, + own_jid: own_pn, + own_lid, + account: snapshot.account.as_deref(), + to: &chat, + message: &message, + message_id: &message_id, + force_distribution: false, + distribution_targets: Some(vec![requester]), + distribution_policy: wacore::send::SenderKeyDistributionPolicy::Required, + phash_devices: None, + edit: edit.as_ref(), + extra_nodes: &[], + pre_encoded: encoded, + }, + ) + .await + { + Ok(prepared) => prepared, + Err(error) => { + // Do not hold the sender-key distribution lane across registry + // I/O. The typed failure retains the original source chain and + // identifies only users whose pre-key lookup returned 406. + drop(distribution_guard); + if let Some(failure) = + error.downcast_ref::() + { + for user in failure.stale_device_users() { + self.invalidate_device_cache(user).await; + } + } + return Err(error); + } + }; + self.send_retry_stanza(prepared.node).await?; + self.update_sender_key_devices(&chat_key, &prepared.skdm_devices) + .await; + drop(distribution_guard); + for user in &prepared.stale_device_users { + self.invalidate_device_cache(user).await; + } Ok(()) } @@ -641,7 +1013,7 @@ impl Client { if info.chat.is_group() || info.chat.is_status_broadcast() { let group_jid = info.chat.to_string(); match self - .mark_forget_sender_key(&group_jid, std::slice::from_ref(&info.requester)) + .mark_forget_sender_key(&group_jid, std::slice::from_ref(resolved_jid)) .await { Ok(()) => { @@ -655,7 +1027,7 @@ impl Client { // (WA Web logs the same event at its verbose LOG level). debug!( "Marked {} for fresh SKDM in {} {} due to retry receipt", - info.requester.observe(), + resolved_jid.observe(), chat_type, group_jid ); @@ -2293,6 +2665,105 @@ mod tests { ); } + #[tokio::test] + async fn update_local_signal_session_cools_resolved_sender_key_namespace() { + let client = crate::test_utils::create_test_client_with_failing_http( + "retry_sender_key_resolved_namespace", + ) + .await; + let group = "120363000000000006@g.us"; + let requester_pn: Jid = "12025550108:33@s.whatsapp.net".parse().unwrap(); + let resolved_lid: Jid = "100000000000088:33@lid".parse().unwrap(); + client + .persistence_manager + .set_sender_key_status( + group, + &[ + ("12025550108:33@s.whatsapp.net", true), + ("100000000000088:33@lid", true), + ], + ) + .await + .unwrap(); + + let rows = client + .persistence_manager + .get_sender_key_devices(group) + .await + .unwrap(); + let cached = client + .sender_key_device_cache + .get_or_init(group, async { + Arc::new(crate::sender_key_device_cache::SenderKeyDeviceMap::from_db_rows(&rows)) + }) + .await; + + let info = RetryChatInfo { + chat: group.parse().unwrap(), + requester: requester_pn, + original_from: group.parse().unwrap(), + recipient: None, + is_bot: false, + is_fbid_bot_retry: false, + }; + let node = build_retry_receipt_without_keys(); + assert!( + client + .update_local_signal_session( + &info, + &resolved_lid, + "MSG-GRP-NAMESPACE", + 1, + &node.as_node_ref(), + false, + ) + .await + ); + + assert_eq!(cached.device_has_key("100000000000088", 33), Some(false)); + assert_eq!(cached.device_has_key("12025550108", 33), Some(true)); + let persisted = crate::sender_key_device_cache::SenderKeyDeviceMap::from_db_rows( + &client + .persistence_manager + .get_sender_key_devices(group) + .await + .unwrap(), + ); + assert_eq!(persisted.device_has_key("100000000000088", 33), Some(false)); + assert_eq!(persisted.device_has_key("12025550108", 33), Some(true)); + } + + #[tokio::test] + async fn status_retransmission_resolution_is_cache_aside_with_pn_fallback() { + let client = crate::test_utils::create_test_client_with_failing_http( + "retry_status_requester_resolution", + ) + .await; + client + .add_lid_pn_mapping( + "100000000000089", + "12025550109", + crate::lid_pn_cache::LearningSource::Usync, + ) + .await + .unwrap(); + client.lid_pn_cache.clear().await; + + let mapped_pn: Jid = "12025550109:19@s.whatsapp.net".parse().unwrap(); + let mapped = client + .resolve_retransmission_encryption_jid(RetransmissionRoute::Status, &mapped_pn) + .await + .unwrap(); + assert_eq!(mapped, "100000000000089:19@lid".parse::().unwrap()); + + let unmapped_pn: Jid = "12025550110:20@s.whatsapp.net".parse().unwrap(); + let fallback = client + .resolve_retransmission_encryption_jid(RetransmissionRoute::Status, &unmapped_pn) + .await + .unwrap(); + assert_eq!(fallback, unmapped_pn); + } + /// `should_recreate_session` mirrors whatsmeow `shouldRecreateSession`: /// 1) no session → always recreate; /// 2) session exists + retry<2 → never recreate; @@ -2909,6 +3380,204 @@ mod tests { assert!(!(dm.is_group() || dm.is_status_broadcast())); } + #[test] + fn retransmission_route_validation_is_strict_and_typed() { + let direct: Jid = "12025550100@s.whatsapp.net".parse().unwrap(); + let requester: Jid = "12025550100:7@s.whatsapp.net".parse().unwrap(); + let group: Jid = "120363000000000001@g.us".parse().unwrap(); + let status = Jid::status_broadcast(); + let broadcast: Jid = "1234567890@broadcast".parse().unwrap(); + + assert!(matches!( + validate_retransmission(&direct, &requester, "DM1", 1, Some(&direct)), + Ok(RetransmissionRoute::Direct) + )); + assert!(matches!( + validate_retransmission(&group, &requester, "GROUP1", 1, None), + Ok(RetransmissionRoute::Group) + )); + assert!(matches!( + validate_retransmission(&status, &requester, "STATUS1", 1, None), + Ok(RetransmissionRoute::Status) + )); + assert!(matches!( + validate_retransmission(&broadcast, &requester, "BROADCAST1", 1, None), + Ok(RetransmissionRoute::BroadcastList) + )); + + for (id, count) in [("ZERO", 0), ("", 1), ("MAX", MAX_RETRY_COUNT)] { + assert!( + validate_retransmission(&direct, &requester, id, count, None).is_err(), + "invalid id/count pair must fail: {id:?}/{count}" + ); + } + assert!( + validate_retransmission(&group, &requester, "GROUP2", 1, Some(&direct)).is_err(), + "recipient is only meaningful on a direct retry" + ); + assert!( + validate_retransmission(&status, &group, "STATUS2", 1, None).is_err(), + "a group JID cannot be a requesting status device" + ); + } + + #[tokio::test] + async fn public_peer_retransmission_requires_a_recipient() { + let client = crate::test_utils::create_test_client().await; + let own_pn: Jid = "12025550100:13@s.whatsapp.net".parse().unwrap(); + client + .persistence_manager + .process_command(crate::store::commands::DeviceCommand::SetId(Some( + own_pn.clone(), + ))) + .await; + + let chat: Jid = "12025550101@s.whatsapp.net".parse().unwrap(); + let requester = own_pn.with_device(7); + let request = MessageRetransmission::new( + chat, + requester, + wa::Message::default(), + "PEER-RETRY-1".to_string(), + 1, + ); + + let error = client + .retransmit_message(request) + .await + .expect_err("a peer route without its actual chat cannot be sent"); + assert!(matches!(error, SendError::InvalidRequest(_))); + assert!(error.to_string().contains("requires a recipient")); + } + + #[tokio::test] + async fn public_direct_retransmission_binds_chat_to_routing_identity() { + let client = crate::test_utils::create_test_client().await; + let chat = Jid::pn("12025550104"); + let requester = Jid::pn_device("12025550105", 7); + let bot_requester: Jid = "200000000000002@bot".parse().unwrap(); + + for request in [ + MessageRetransmission::new( + chat.clone(), + requester, + wa::Message::default(), + "DIRECT-CHAT-MISMATCH-1".to_string(), + 1, + ), + MessageRetransmission::new( + chat.clone(), + bot_requester, + wa::Message::default(), + "DIRECT-RECIPIENT-MISMATCH-1".to_string(), + 1, + ) + .with_recipient(Jid::pn("12025550106")), + ] { + let error = client + .retransmit_message(request) + .await + .expect_err("an unrelated routing identity must be rejected"); + assert!(matches!(error, SendError::InvalidRequest(_))); + assert!(error.to_string().contains("routing identity")); + } + } + + #[tokio::test] + async fn public_direct_recipient_rejects_an_unrelated_requester() { + let client = crate::test_utils::create_test_client().await; + let chat = Jid::pn("12025550108"); + let request = MessageRetransmission::new( + chat.clone(), + Jid::pn_device("12025550109", 7), + wa::Message::default(), + "DIRECT-RECIPIENT-SOURCE-1".to_string(), + 1, + ) + .with_recipient(chat); + + let error = client + .retransmit_message(request) + .await + .expect_err("a normal remote user cannot declare a recipient route"); + assert!(matches!(error, SendError::InvalidRequest(_))); + assert!(error.to_string().contains("local device or bot")); + } + + #[tokio::test] + async fn direct_retransmission_chat_accepts_known_pn_lid_alias() { + let client = crate::test_utils::create_test_client().await; + let pn = Jid::pn("12025550107"); + let lid = Jid::lid("100000000000107"); + client + .lid_pn_cache + .add(&wacore::types::lid_pn::LidPnEntry { + lid: lid.user.as_str().into(), + phone_number: pn.user.as_str().into(), + created_at: 1, + learning_source: wacore::types::lid_pn::LearningSource::Usync, + }) + .await; + + assert!(client.jids_share_user_identity(&pn, &lid).await.unwrap()); + assert!(client.jids_share_user_identity(&lid, &pn).await.unwrap()); + } + + #[tokio::test] + async fn public_retransmission_recaches_the_supplied_message() { + let mut config = crate::cache_config::CacheConfig::default(); + config.recent_messages.capacity = 16; + let client = crate::test_utils::create_test_client_with_config( + "public_retransmission_cache", + Arc::new(MockHttpClient), + config, + ) + .await; + let chat = Jid::pn("12025550103"); + let requester = chat.with_device(7); + crate::test_utils::seed_peer_session(&client, &requester).await; + let message = wa::Message { + conversation: Some("retry me".into()), + ..Default::default() + }; + let message_id = "PUBLIC-RETRY-CACHE-1"; + + // The fresh test session emits pkmsg and this client intentionally has + // no device identity, so the wire attempt fails after the public API has + // accepted and cached the supplied message. + let result = client + .retransmit_message(MessageRetransmission::new( + chat.clone(), + requester, + message, + message_id.to_string(), + 1, + )) + .await; + assert!(result.is_err()); + + let (cached, alternate) = client + .peek_recent_message(&chat, message_id) + .await + .expect("a later retry count must find the retransmitted message"); + assert!(alternate.is_none()); + assert_eq!(cached.conversation.as_deref(), Some("retry me")); + } + + #[test] + fn resolve_retry_chat_info_broadcast_uses_participant_device() { + let broadcast = "1234567890@broadcast"; + let participant = "12025550101:9@s.whatsapp.net"; + let node = NodeBuilder::new("receipt") + .attr("participant", participant) + .build(); + let receipt = make_test_receipt(broadcast); + let info = resolve_retry_chat_info(&receipt, &node.as_node_ref(), None, None); + + assert!(info.chat.is_broadcast_list()); + assert_eq!(info.requester, participant.parse::().unwrap()); + } + /// The key-bundle policy is driven only by explicit force, stateless routing, /// and the retry threshold. The diagnostic reason must not change the wire /// shape of a first retry. diff --git a/src/send/actions.rs b/src/send/actions.rs index e3248c885..0bc9c7252 100644 --- a/src/send/actions.rs +++ b/src/send/actions.rs @@ -74,12 +74,11 @@ impl Client { self.send_message_impl( to, &revoke_message, - None, - false, - force_skdm, - Some(edit_attr), - vec![], - None, + SendPipelineOptions { + force_key_distribution: force_skdm, + edit: Some(edit_attr), + ..Default::default() + }, ) .await .map_err(SendError::from_anyhow)?; @@ -163,12 +162,10 @@ impl Client { self.send_message_impl( chat, &message, - None, - false, - false, - Some(crate::types::message::EditAttribute::PinInChat), - vec![], - None, + SendPipelineOptions { + edit: Some(crate::types::message::EditAttribute::PinInChat), + ..Default::default() + }, ) .await .map_err(SendError::from_anyhow)?; diff --git a/src/send/mod.rs b/src/send/mod.rs index d523fd684..38024f8c5 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -2,7 +2,6 @@ use crate::client::Client; use crate::types::message::EditAttribute; use anyhow::anyhow; use log::debug; -use wacore::client::context::SendContextResolver; use wacore::libsignal::protocol::SignalProtocolError; use wacore::send::StanzaType; use wacore::types::jid::JidExt; @@ -157,6 +156,41 @@ struct SendBranchOutput { dm_phash: Option, } +struct GroupBranchRequest<'a> { + to: Jid, + message: &'a wa::Message, + request_id: String, + force_key_distribution: bool, + edit: Option, + extra_stanza_nodes: &'a [Node], + group_metadata_freshness: crate::cache::Freshness, + device_freshness: crate::cache::Freshness, +} + +struct DmBranchRequest<'a> { + to: Jid, + message: &'a wa::Message, + request_id: String, + edit: Option, + extra_stanza_nodes: Vec, + is_status_addon: bool, + device_freshness: crate::cache::Freshness, +} + +enum GroupDeviceSnapshot { + Owned(wacore::send::ResolvedGroupDevices), + Shared(std::sync::Arc), +} + +impl AsRef for GroupDeviceSnapshot { + fn as_ref(&self) -> &wacore::send::ResolvedGroupDevices { + match self { + Self::Owned(devices) => devices, + Self::Shared(devices) => devices, + } + } +} + /// Keep each branch future out of the shared send frame. In tracing builds the /// dedicated span lets allocation profilers distinguish this deliberate box /// from work performed while polling the selected branch. @@ -189,6 +223,23 @@ fn skdm_needs_only_own_devices(needs: &[Jid], own_pn: Option<&Jid>, own_lid: Opt }) } +const RESERVED_EXTRA_STANZA_CHILDREN: &[&str] = + &["enc", "participants", "device-identity", "plaintext"]; + +fn validate_extra_stanza_nodes(nodes: &[Node]) -> Result<(), SendError> { + if let Some(node) = nodes.iter().find(|node| { + RESERVED_EXTRA_STANZA_CHILDREN + .iter() + .any(|reserved| node.tag == *reserved) + }) { + return Err(SendError::InvalidRequest(format!( + "extra stanza child <{}> is reserved by the send pipeline", + node.tag + ))); + } + Ok(()) +} + impl SendBranchOutput { fn stanza_only(node: Node) -> Self { Self { @@ -218,6 +269,22 @@ pub struct SendOptions { /// Force the `` attribute instead of deriving it from /// content. Escape hatch for a type the classifier can't infer. pub stanza_type_override: Option, + /// Freshness policy for group metadata used by this send. + pub group_metadata_freshness: crate::cache::Freshness, + /// Freshness policy for recipient device lists used by this send. + pub device_freshness: crate::cache::Freshness, +} + +#[derive(Default)] +pub(crate) struct SendPipelineOptions { + pub(crate) request_id: Option, + pub(crate) peer: bool, + pub(crate) force_key_distribution: bool, + pub(crate) edit: Option, + pub(crate) extra_stanza_nodes: Vec, + pub(crate) stanza_type: Option, + pub(crate) group_metadata_freshness: crate::cache::Freshness, + pub(crate) device_freshness: crate::cache::Freshness, } /// Result of a successfully sent message. @@ -341,6 +408,26 @@ pub(crate) fn infer_stanza_metadata(msg: &wa::Message) -> (Option (edit, has_attr.then(|| meta.build())) } +fn validate_status_message_id( + message: &wa::Message, + outer_id: Option<&str>, +) -> Result<(), SendError> { + let Some(outer_id) = outer_id else { + return Ok(()); + }; + if outer_id.is_empty() { + return Err(SendError::InvalidRequest( + "status message ID must not be empty".into(), + )); + } + if wacore::send::status_revoke_target_id(message) == Some(outer_id) { + return Err(SendError::InvalidRequest( + "status revoke stanza ID must differ from the revoked message ID".into(), + )); + } + Ok(()) +} + /// Offset subtracted from the current unix timestamp to produce the /// `privacy_mode_ts` attr value on a `` stanza. Empirically confirmed /// against live WhatsApp servers. @@ -667,6 +754,13 @@ impl Client { #[cfg(feature = "tracing")] self.record_identity_on_span(&tracing::Span::current()); + validate_extra_stanza_nodes(&options.extra_stanza_nodes)?; + if options.message_id.as_ref().is_some_and(String::is_empty) { + return Err(SendError::InvalidRequest( + "message ID must not be empty".into(), + )); + } + let _t = wacore::telemetry::timer(wacore::telemetry::SEND_DURATION); self.stats.record_message_sent(); wacore::telemetry::send(match to.server { @@ -686,6 +780,8 @@ impl Client { } let stanza_type_override = options.stanza_type_override; + let group_metadata_freshness = options.group_metadata_freshness; + let device_freshness = options.device_freshness; let request_id = match options.message_id { Some(id) => id, None => self.generate_message_id(), @@ -737,12 +833,15 @@ impl Client { self.send_message_impl( to, &message, - Some(request_id), - false, - false, - edit, - extra_nodes, - stanza_type_override, + SendPipelineOptions { + request_id: Some(request_id), + edit, + extra_stanza_nodes: extra_nodes, + stanza_type: stanza_type_override, + group_metadata_freshness, + device_freshness, + ..Default::default() + }, ) .await .map_err(SendError::from_anyhow)?; @@ -763,7 +862,7 @@ impl Client { &self, message: wa::Message, recipients: &[Jid], - options: crate::features::status::StatusSendOptions, + mut options: crate::features::status::StatusSendOptions, ) -> Result { use wacore::client::context::GroupInfo; use wacore_binary::builder::NodeBuilder; @@ -773,6 +872,8 @@ impl Client { "cannot send status with no recipients".into(), )); } + validate_extra_stanza_nodes(&options.extra_stanza_nodes)?; + validate_status_message_id(&message, options.message_id.as_deref())?; // Status posts don't go through send_message_with_options, so count them here. let _t = wacore::telemetry::timer(wacore::telemetry::SEND_DURATION); @@ -780,7 +881,10 @@ impl Client { wacore::telemetry::send("status"); let to = Jid::status_broadcast(); - let request_id = self.generate_message_id(); + let request_id = options + .message_id + .take() + .unwrap_or_else(|| 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(); @@ -878,13 +982,19 @@ impl Client { // Determine which devices need SKDM using the unified per-device map. // Status keeps the prior phash behavior, so we drop the full device set // and only use the SKDM-target subset. - let skdm_target_devices: Option> = if force_skdm { - None - } else { - self.resolve_skdm_targets(&to_str, &group_info, own_lid) - .await - .map(|(_all, needs)| needs) - }; + let skdm_target_devices = + if !force_skdm || options.device_freshness == crate::cache::Freshness::Refresh { + self.resolve_status_skdm_targets( + &to_str, + &group_info, + own_lid, + options.device_freshness, + force_skdm, + ) + .await? + } else { + None + }; // prepare_group_stanza and ensure_status_participants both read the // participant list and expect self present. Done after SKDM resolution @@ -902,35 +1012,35 @@ impl Client { // status. Reactions go through WA Web's addon path and never visit // `WAWebEncryptAndSendStatusMsg`; attaching the meta on a reaction // gets the stanza NACK'd with 479 (SmaxInvalid). Revokes also skip it. - let extra_stanza_nodes = if wacore::send::status_carries_privacy_meta(&message) { - vec![ + let mut extra_stanza_nodes = options.extra_stanza_nodes; + if wacore::send::status_carries_privacy_meta(&message) { + extra_stanza_nodes.push( NodeBuilder::new("meta") .attr("status_setting", options.privacy.as_str()) .build(), - ] - } else { - vec![] - }; + ); + } let prepared = match wacore::send::prepare_group_stanza( &*self.runtime, &mut stores, self, - &group_info, - own_jid, - own_lid, - account_info.as_deref(), - to.clone(), - &message, - request_id.clone(), - force_skdm, - skdm_target_devices, - // Status broadcasts keep the prior phash behavior (no full-set/self - // augmentation) — that path is group-only. - None, - None, - &extra_stanza_nodes, - shared_content.clone(), + wacore::send::GroupStanzaRequest { + group: &group_info, + own_jid, + own_lid, + account: account_info.as_deref(), + to: &to, + message: &message, + message_id: &request_id, + force_distribution: force_skdm, + distribution_targets: skdm_target_devices, + distribution_policy: wacore::send::SenderKeyDistributionPolicy::BestEffort, + phash_devices: None, + edit: None, + extra_nodes: &extra_stanza_nodes, + pre_encoded: shared_content.as_deref().map(Vec::as_slice), + }, ) .await { @@ -951,19 +1061,23 @@ impl Client { &*self.runtime, &mut stores_retry, self, - &group_info, - own_jid, - own_lid, - account_info.as_deref(), - to.clone(), - &message, - request_id.clone(), - true, - None, - None, - None, - &extra_stanza_nodes, - shared_content.clone(), + wacore::send::GroupStanzaRequest { + group: &group_info, + own_jid, + own_lid, + account: account_info.as_deref(), + to: &to, + message: &message, + message_id: &request_id, + force_distribution: true, + distribution_targets: None, + distribution_policy: + wacore::send::SenderKeyDistributionPolicy::BestEffort, + phash_devices: None, + edit: None, + extra_nodes: &extra_stanza_nodes, + pre_encoded: shared_content.as_deref().map(Vec::as_slice), + }, ) .await? } else { @@ -1097,13 +1211,19 @@ impl Client { /// SKDM target resolution for the status path, whose `GroupInfo` is built /// fresh per send (no stable identity to memoize against). #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.resolve_skdm_targets", level = "debug", skip_all, fields(group = %wacore_binary::jid::observe_str(group_jid))))] - async fn resolve_skdm_targets( + async fn resolve_status_skdm_targets( &self, group_jid: &str, group_info: &wacore::client::context::GroupInfo, own_sending_jid: &Jid, - ) -> Option<(std::sync::Arc, Vec)> { - let cached_map = self.skdm_device_map(group_jid).await; + freshness: crate::cache::Freshness, + force_distribution: bool, + ) -> Result>, anyhow::Error> { + let cached_map = if force_distribution { + None + } else { + Some(self.skdm_device_map(group_jid).await) + }; let is_lid_mode = group_info.addressing_mode == wacore::types::message::AddressingMode::Lid; let jids_to_resolve: Vec = group_info @@ -1120,34 +1240,45 @@ impl Client { }) .collect(); - match SendContextResolver::resolve_devices(self, &jids_to_resolve).await { - Ok(all_devices) => { - let all_devices: Vec = if is_lid_mode { - all_devices - .into_iter() - .map(|d| group_info.phone_device_jid_into_lid(d)) - .collect() - } else { - all_devices - }; - let all_devices = - std::sync::Arc::new(wacore::send::ResolvedGroupDevices::new(all_devices)); - let needs_skdm = self.filter_skdm_targets( - group_jid, - all_devices.devices(), - &cached_map, - own_sending_jid, + let resolved = match freshness { + crate::cache::Freshness::CachePreferred => { + self.get_user_devices_owned(jids_to_resolve).await + } + crate::cache::Freshness::Refresh => self.refresh_user_devices(jids_to_resolve).await, + }; + match resolved { + Ok(mut devices) => { + if is_lid_mode { + for device in &mut devices { + *device = group_info.phone_device_jid_into_lid(std::mem::take(device)); + } + } + if force_distribution { + wacore::send::retain_skdm_distribution_targets(&mut devices, own_sending_jid); + } else if let Some(cached_map) = cached_map { + devices.retain(|device| { + !device.is_hosted() + && !(device.user == own_sending_jid.user + && device.device == own_sending_jid.device) + && !cached_map.device_and_primary_warm(&device.user, device.device) + }); + } + log::debug!( + "Resolved {} status devices needing SKDM for {}", + devices.len(), + group_jid ); - Some((all_devices, needs_skdm)) + Ok(Some(devices)) } - Err(e) => { + Err(error) if freshness == crate::cache::Freshness::CachePreferred => { log::warn!( "Failed to resolve devices for SKDM check in {}: {:?}", group_jid, - e + error ); - None + Ok(None) } + Err(error) => Err(error), } } @@ -1256,7 +1387,7 @@ impl Client { /// would be one-directional: the retry-receipt forget path also excludes own /// devices (to stop an inbound retry tearing down our own session), so an own /// companion whose one SKDM encryption failed could never be re-sent one. - async fn update_sender_key_devices(&self, group_jid: &str, devices: &[Jid]) { + pub(crate) async fn update_sender_key_devices(&self, group_jid: &str, devices: &[Jid]) { if devices.is_empty() { return; } @@ -1385,7 +1516,7 @@ impl Client { .await; } if invalidate_group_cache { - self.get_group_cache().await.invalidate(jid).await; + self.lock_group_metadata(jid).await.invalidate().await; } } @@ -1403,18 +1534,26 @@ impl Client { } #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.impl", level = "debug", skip_all, fields(to = %to.observe()), err(Debug)))] - #[allow(clippy::too_many_arguments)] pub(crate) async fn send_message_impl( &self, to: Jid, message: &wa::Message, - request_id_override: Option, - peer: bool, - force_key_distribution: bool, - edit: Option, - extra_stanza_nodes: Vec, - stanza_type_override: Option, + options: SendPipelineOptions, ) -> Result<(), anyhow::Error> { + let SendPipelineOptions { + request_id: request_id_override, + peer, + force_key_distribution, + edit, + extra_stanza_nodes, + stanza_type: stanza_type_override, + group_metadata_freshness, + device_freshness, + } = options; + validate_extra_stanza_nodes(&extra_stanza_nodes)?; + if request_id_override.as_ref().is_some_and(String::is_empty) { + return Err(SendError::InvalidRequest("message ID must not be empty".into()).into()); + } // Newsletters are plaintext channels and never use the E2E path. Text // sends go through the branch in send_message_with_options; // edit/revoke have dedicated plaintext methods (newsletter().edit_message @@ -1479,24 +1618,27 @@ impl Client { } = if peer && !to.is_group() { box_send_branch(self.send_peer_branch(to, message, request_id)).await? } else if to.is_group() { - box_send_branch(self.send_group_branch( + box_send_branch(self.send_group_branch(GroupBranchRequest { to, message, request_id, force_key_distribution, edit, - &extra_stanza_nodes, - )) + extra_stanza_nodes: &extra_stanza_nodes, + group_metadata_freshness, + device_freshness, + })) .await? } else { - box_send_branch(self.send_dm_branch( + box_send_branch(self.send_dm_branch(DmBranchRequest { to, message, request_id, edit, extra_stanza_nodes, is_status_addon, - )) + device_freshness, + })) .await? }; @@ -1635,13 +1777,18 @@ impl Client { /// SKDM distribution and the cold/rotation single-flight. async fn send_group_branch( &self, - to: Jid, - message: &wa::Message, - request_id: String, - force_key_distribution: bool, - edit: Option<crate::types::message::EditAttribute>, - extra_stanza_nodes: &[Node], + request: GroupBranchRequest<'_>, ) -> Result<SendBranchOutput, anyhow::Error> { + let GroupBranchRequest { + to, + message, + request_id, + force_key_distribution, + edit, + extra_stanza_nodes, + group_metadata_freshness, + device_freshness, + } = request; // Every arm of the prepare match below assigns these three. let outbound_msg_secret: Option<[u8; 32]>; let outbound_group_sender_identity: Option<Jid>; @@ -1650,7 +1797,10 @@ impl Client { let node = { // No send-level lock: encrypt_group_message serializes the // sender-key chain advance per (group, sender) at the cipher. - let group_info = self.groups().query_info(&to).await?; + let group_info = self + .groups() + .query_info_with_freshness(&to, group_metadata_freshness) + .await?; // Borrow from the held snapshot: no field clones, the Arc keeps it alive. let device_snapshot = self.persistence_manager.get_device_snapshot(); @@ -1686,6 +1836,18 @@ impl Client { // would make the memo miss on every send to such groups. The memoized // resolver applies the same self-append internally. let group_info_for_memo = std::sync::Arc::clone(&group_info); + let refreshed_devices = if device_freshness == crate::cache::Freshness::Refresh { + Some( + self.resolve_group_devices_uncached( + &group_info_for_memo, + &own_sending_jid, + crate::cache::Freshness::Refresh, + ) + .await?, + ) + } else { + None + }; // resolve_skdm_targets and prepare_group_stanza both read the // participant list and expect self to be present. let group_info = ensure_self_in_group(group_info, &own_sending_jid); @@ -1759,20 +1921,45 @@ impl Client { // still missing the key. On the cold/`force_skdm` path both are // `None` and `prepare_group_stanza` resolves the set itself. let (all_devices_for_phash, skdm_target_devices): ( - Option<std::sync::Arc<wacore::send::ResolvedGroupDevices>>, + Option<GroupDeviceSnapshot>, Option<Vec<Jid>>, ) = if force_skdm { - (None, None) + match refreshed_devices { + Some(mut targets) => { + wacore::send::retain_skdm_distribution_targets( + &mut targets, + &own_sending_jid, + ); + (None, Some(targets)) + } + None => (None, None), + } } else { - match self - .resolve_skdm_targets_memoized( - &to, - &to_str, - &group_info_for_memo, - &own_sending_jid, - ) - .await - { + let initial_targets = match refreshed_devices { + Some(all) => { + let all = GroupDeviceSnapshot::Owned( + wacore::send::ResolvedGroupDevices::new(all), + ); + let cached_map = self.skdm_device_map(&to_str).await; + let needs = self.filter_skdm_targets( + &to_str, + all.as_ref().devices(), + &cached_map, + &own_sending_jid, + ); + Some((all, needs)) + } + None => self + .resolve_skdm_targets_memoized( + &to, + &to_str, + &group_info_for_memo, + &own_sending_jid, + ) + .await + .map(|(all, needs)| (GroupDeviceSnapshot::Shared(all), needs)), + }; + match initial_targets { Some((all, needs)) if needs.is_empty() => (Some(all), Some(needs)), // Own devices are never memoized warm, so they re-receive // their SKDM on every send by design — own-only needs IS @@ -1818,7 +2005,7 @@ impl Client { { distribution_guard = None; } - (Some(all), Some(needs)) + (Some(GroupDeviceSnapshot::Shared(all)), Some(needs)) } // Transient re-resolve failure: keep the first // resolve's targets rather than silently sending @@ -1835,19 +2022,22 @@ impl Client { &*self.runtime, &mut stores, self, - &group_info, - own_jid, - own_lid, - account_info.as_deref(), - to.clone(), - message, - request_id.clone(), - force_skdm, - skdm_target_devices, - all_devices_for_phash, - edit.clone(), - extra_stanza_nodes, - shared_content.clone(), + wacore::send::GroupStanzaRequest { + group: &group_info, + own_jid, + own_lid, + account: account_info.as_deref(), + to: &to, + message, + message_id: &request_id, + force_distribution: force_skdm, + distribution_targets: skdm_target_devices, + distribution_policy: wacore::send::SenderKeyDistributionPolicy::BestEffort, + phash_devices: all_devices_for_phash.as_ref().map(AsRef::as_ref), + edit: edit.as_ref(), + extra_nodes: extra_stanza_nodes, + pre_encoded: shared_content.as_deref().map(Vec::as_slice), + }, ) .await { @@ -1895,7 +2085,9 @@ impl Client { None }; let (retry_force, retry_targets, retry_all) = match warm_targets { - Some((all, needs)) => (false, Some(needs), Some(all)), + Some((all, needs)) => { + (false, Some(needs), Some(GroupDeviceSnapshot::Shared(all))) + } None => { self.reset_sender_key_device_tracking(&to_str).await?; (true, None, None) @@ -1910,19 +2102,23 @@ impl Client { &*self.runtime, &mut stores_retry, self, - &group_info, - own_jid, - own_lid, - account_info.as_deref(), - to, - message, - request_id, - retry_force, - retry_targets, - retry_all, - edit.clone(), - extra_stanza_nodes, - shared_content.clone(), + wacore::send::GroupStanzaRequest { + group: &group_info, + own_jid, + own_lid, + account: account_info.as_deref(), + to: &to, + message, + message_id: &request_id, + force_distribution: retry_force, + distribution_targets: retry_targets, + distribution_policy: + wacore::send::SenderKeyDistributionPolicy::BestEffort, + phash_devices: retry_all.as_ref().map(AsRef::as_ref), + edit: edit.as_ref(), + extra_nodes: extra_stanza_nodes, + pre_encoded: shared_content.as_deref().map(Vec::as_slice), + }, ) .await?; @@ -1955,13 +2151,17 @@ impl Client { /// with device fan-out (also used by status-reaction add-ons). async fn send_dm_branch( &self, - to: Jid, - message: &wa::Message, - request_id: String, - edit: Option<crate::types::message::EditAttribute>, - extra_stanza_nodes: Vec<Node>, - is_status_addon: bool, + request: DmBranchRequest<'_>, ) -> Result<SendBranchOutput, anyhow::Error> { + let DmBranchRequest { + to, + message, + request_id, + edit, + extra_stanza_nodes, + is_status_addon, + device_freshness, + } = request; let mut should_issue_tc_token_after_send = false; let prepared = { // Per-device locking to match decrypt path (message.rs:684), @@ -2037,6 +2237,11 @@ impl Client { let stanza_to = dm_stanza_to(&recipient_bare, &to); + if device_freshness == crate::cache::Freshness::Refresh { + self.refresh_user_devices(vec![recipient_bare.to_non_ad(), own_jid.to_non_ad()]) + .await?; + } + // Local registry first; network warm only on miss to avoid // unnecessary LID-migration side effects from get_user_devices let mut recipient_cached = self.get_devices_from_registry(&recipient_bare).await; @@ -2137,16 +2342,18 @@ impl Client { &*self.runtime, &mut stores, self, - own_jid, - device_snapshot.lid.as_ref(), - device_snapshot.account.as_deref(), - stanza_to, - message, - request_id, - edit, - &extra_stanza_nodes, - all_dm_jids, - shared_content, + wacore::send::DmStanzaRequest { + own_jid, + own_lid: device_snapshot.lid.as_ref(), + account: device_snapshot.account.as_deref(), + to: &stanza_to, + message, + message_id: &request_id, + edit: edit.as_ref(), + extra_nodes: &extra_stanza_nodes, + devices: all_dm_jids, + pre_encoded: shared_content.as_deref().map(Vec::as_slice), + }, ) .await? }; @@ -2284,6 +2491,29 @@ mod tests { use crate::test_utils::wait_for_lock_waiter; use std::str::FromStr; + #[test] + fn status_revoke_requires_a_distinct_outer_stanza_id() { + let target_id = "3EB0REVOKETARGET"; + let revoke = wa::Message { + protocol_message: buffa::MessageField::some(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::Revoke), + key: buffa::MessageField::some(wa::MessageKey { + id: Some(target_id.into()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + + assert!(matches!( + validate_status_message_id(&revoke, Some(target_id)), + Err(SendError::InvalidRequest(_)) + )); + assert!(validate_status_message_id(&revoke, Some("3EB0NEWSTANZAID")).is_ok()); + assert!(validate_status_message_id(&revoke, None).is_ok()); + } + #[test] fn dm_stanza_to_follows_resolved_wire_namespace() { let pn: Jid = "5511987650001@s.whatsapp.net".parse().unwrap(); @@ -3050,18 +3280,23 @@ mod tests { let group_info = GroupInfo::new(participants.clone(), AddressingMode::Lid); - let (all_devices, needs_skdm) = client - .resolve_skdm_targets(group_jid, &group_info, &own_lid) + let needs_skdm = client + .resolve_status_skdm_targets( + group_jid, + &group_info, + &own_lid, + crate::cache::Freshness::CachePreferred, + false, + ) .await - .expect("None means device resolution failed"); + .expect("device resolution must succeed") + .expect("missing targets means device resolution failed"); // Empty cache → every participant needs SKDM, and the full set equals // the target set on this cold path. assert_eq!(needs_skdm.len(), participants.len()); - assert_eq!(all_devices.devices().len(), participants.len()); for user in &participant_users { assert!(needs_skdm.iter().any(|j| j.user == *user)); - assert!(all_devices.devices().iter().any(|j| j.user == *user)); } } @@ -4044,6 +4279,22 @@ mod tests { } } + #[test] + fn structural_extra_children_are_rejected_before_send_work() { + for tag in RESERVED_EXTRA_STANZA_CHILDREN { + let error = validate_extra_stanza_nodes(&[NodeBuilder::new(tag).build()]) + .expect_err("send-owned child must be rejected"); + assert!(error.to_string().contains(tag)); + } + + validate_extra_stanza_nodes(&[ + NodeBuilder::new("meta").build(), + NodeBuilder::new("biz").build(), + NodeBuilder::new("custom-extension").build(), + ]) + .expect("non-structural protocol extensions remain available"); + } + /// Regression tests for #462: send path session lock keys must match decrypt path. mod session_lock_regression { use super::*; @@ -4340,12 +4591,11 @@ mod tests { .send_message_impl( peer, &msg, - Some(request_id.to_string()), - true, - false, - None, - vec![], - None, + SendPipelineOptions { + request_id: Some(request_id.to_string()), + peer: true, + ..Default::default() + }, ) .await; assert!( @@ -4394,12 +4644,12 @@ mod tests { .send_message_impl( peer, &msg, - Some(request_id.to_string()), - true, - false, - None, - vec![], - Some(wacore::send::StanzaType::Poll), + SendPipelineOptions { + request_id: Some(request_id.to_string()), + peer: true, + stanza_type: Some(wacore::send::StanzaType::Poll), + ..Default::default() + }, ) .await; assert!( @@ -4559,12 +4809,10 @@ mod tests { .send_message_impl( peer_pn, &msg, - Some(request_id.to_string()), - false, - false, - None, - vec![], - None, + SendPipelineOptions { + request_id: Some(request_id.to_string()), + ..Default::default() + }, ) .await; assert!( @@ -4637,12 +4885,10 @@ mod tests { .send_message_impl( peer_pn.clone(), &msg, - Some(request_id.to_string()), - false, - false, - None, - vec![], - None, + SendPipelineOptions { + request_id: Some(request_id.to_string()), + ..Default::default() + }, ) .await; assert!( @@ -4705,7 +4951,7 @@ mod tests { ..Default::default() }; let err = client - .send_message_impl(channel, &msg, None, false, false, None, vec![], None) + .send_message_impl(channel, &msg, SendPipelineOptions::default()) .await .expect_err("newsletter JID must be rejected on the E2E send path"); assert!( diff --git a/src/sender_key_device_cache.rs b/src/sender_key_device_cache.rs index adb626aa1..d98df1fab 100644 --- a/src/sender_key_device_cache.rs +++ b/src/sender_key_device_cache.rs @@ -126,7 +126,11 @@ impl SenderKeyDeviceCache { /// cold, so it is skipped. The DB write is the source of truth; this only /// keeps a live cache entry consistent with it. On a cache miss the next /// send rebuilds from the DB, which already carries the write. - pub(crate) async fn mark_forgotten(&self, group_jid: &str, devices: &[Jid]) { + pub(crate) async fn mark_forgotten<'a>( + &self, + group_jid: &str, + devices: impl Iterator<Item = &'a Jid> + Send, + ) { let Some(map) = self.inner.get(group_jid).await else { return; }; @@ -219,7 +223,7 @@ mod tests { let gen_before = map0.generation(); let dev5: Jid = "111:5@lid".parse().unwrap(); - c.mark_forgotten(group, std::slice::from_ref(&dev5)).await; + c.mark_forgotten(group, std::iter::once(&dev5)).await; // Still cached (no whole-group invalidation): reading it must not run // the init closure. @@ -241,7 +245,7 @@ mod tests { // generation must not advance (no spurious memo miss). let gen_after = map.generation(); let absent: Jid = "999:0@lid".parse().unwrap(); - c.mark_forgotten(group, std::slice::from_ref(&absent)).await; + c.mark_forgotten(group, std::iter::once(&absent)).await; assert_eq!( map.generation(), gen_after, @@ -250,7 +254,7 @@ mod tests { // Re-marking an already-cold device it DOES hold is also a no-op: the // flag is already false, so a retry storm must not churn the generation. - c.mark_forgotten(group, std::slice::from_ref(&dev5)).await; + c.mark_forgotten(group, std::iter::once(&dev5)).await; assert_eq!( map.generation(), gen_after, @@ -279,7 +283,8 @@ mod tests { let d0: Jid = "111:0@lid".parse().unwrap(); let d5: Jid = "111:5@lid".parse().unwrap(); let absent: Jid = "333:0@lid".parse().unwrap(); - c.mark_forgotten(group, &[d0, d5, absent]).await; + c.mark_forgotten(group, [&d0, &d5, &absent].into_iter()) + .await; assert_eq!(map.device_has_key("111", 0), Some(false)); assert_eq!(map.device_has_key("111", 5), Some(false)); @@ -340,7 +345,7 @@ mod tests { let c = cache(); let dev: Jid = "111:0@lid".parse().unwrap(); // No entry for this group: must not panic or create one. - c.mark_forgotten("120363000000000009@g.us", std::slice::from_ref(&dev)) + c.mark_forgotten("120363000000000009@g.us", std::iter::once(&dev)) .await; let map = c .get_or_init("120363000000000009@g.us", async { diff --git a/src/usync.rs b/src/usync.rs index 1eb1bf0c9..317e48cc7 100644 --- a/src/usync.rs +++ b/src/usync.rs @@ -5,10 +5,26 @@ use crate::client::Client; use crate::request::IqError; use log::{debug, warn}; -use std::collections::HashSet; use wacore::iq::usync::{DeviceListResponse, DeviceListSpec}; use wacore_binary::Jid; +/// An authoritative refresh retries when a newer registry mutation wins while +/// its IQ is in flight. Bound retries so a continuously changing account never +/// turns one send into an unbounded request loop. +const DEVICE_REFRESH_MAX_ATTEMPTS: usize = 3; + +#[inline] +fn device_response_contains_user(response: &DeviceListResponse, user: &str) -> bool { + response + .device_lists + .iter() + .any(|device_list| device_list.user.user == user) + || response + .lid_mappings + .iter() + .any(|mapping| mapping.phone_number == user || mapping.lid == user) +} + pub use wacore::iq::usync::{ UsyncAddressingMode, UsyncBotCommand, UsyncBotProfessionalType, UsyncBotProfileResult, UsyncBotPrompt, UsyncBusinessResult, UsyncContactResult, UsyncContext, UsyncDeviceListResult, @@ -34,71 +50,138 @@ impl Client { #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.usync.get_user_devices", level = "debug", skip_all, fields(users = jids.len()), err(Debug)))] pub(crate) async fn get_user_devices(&self, jids: &[Jid]) -> Result<Vec<Jid>, anyhow::Error> { - let mut jids_to_fetch: HashSet<Jid> = HashSet::with_capacity(jids.len()); - let mut all_devices = Vec::with_capacity(jids.len() * 2); + let mut owned = Vec::with_capacity(jids.len()); + owned.extend(jids.iter().map(Jid::to_non_ad)); + self.get_user_devices_owned(owned).await + } + + pub(crate) async fn get_user_devices_owned( + &self, + jids: Vec<Jid>, + ) -> Result<Vec<Jid>, anyhow::Error> { + let input_len = jids.len(); + let mut jids_to_fetch: Vec<Jid> = Vec::with_capacity(input_len); + let mut all_devices = Vec::with_capacity(input_len * 2); // Resolve the LOCAL registry scan concurrently (the network usync below is // already one batched IQ) — a cold-cache large group would otherwise // serialize 256+ per-user cache/DB reads. Order is irrelevant (phash sorts, // encrypt fan-out is order-agnostic). A None result means an empty/corrupt // record, which falls through to the network below (WA Web always keeps - // device 0). Materialize to an owned Vec first so the stream doesn't borrow - // `jids` through buffer_unordered (Send bound). + // device 0). The stream owns each JID and is drained incrementally, so + // no second result Vec is materialized. use futures::StreamExt; // Bounded fan-out over the independent per-user registry reads. const DEVICE_LIST_RESOLVE_CONCURRENCY: usize = 16; - let non_ad: Vec<Jid> = jids.iter().map(|j| j.to_non_ad()).collect(); - let resolved: Vec<(Jid, Option<Vec<Jid>>)> = futures::stream::iter(non_ad) + let mut resolved = futures::stream::iter(jids.into_iter().map(Jid::into_non_ad)) .map(|jid| async move { let devices = self.get_devices_from_registry(&jid).await; (jid, devices) }) - .buffer_unordered(DEVICE_LIST_RESOLVE_CONCURRENCY) - .collect() - .await; + .buffer_unordered(DEVICE_LIST_RESOLVE_CONCURRENCY); - for (jid, devices) in resolved { + while let Some((jid, devices)) = resolved.next().await { match devices { Some(devices) => all_devices.extend(devices), None => { - jids_to_fetch.insert(jid); + jids_to_fetch.push(jid); } } } if !jids_to_fetch.is_empty() { + wacore::types::jid::sort_dedup_by_user(&mut jids_to_fetch); debug!( "get_user_devices: Cache miss, fetching from network for {} unique users", jids_to_fetch.len() ); + all_devices.extend(self.fetch_user_devices(jids_to_fetch).await?); + } + Ok(all_devices) + } + + pub(crate) async fn refresh_user_devices( + &self, + mut jids: Vec<Jid>, + ) -> Result<Vec<Jid>, anyhow::Error> { + for jid in &mut jids { + jid.agent = 0; + jid.device = 0; + } + wacore::types::jid::sort_dedup_by_user(&mut jids); + self.fetch_user_devices_with_freshness(jids, crate::cache::Freshness::Refresh) + .await + } + + async fn fetch_user_devices(&self, jids: Vec<Jid>) -> Result<Vec<Jid>, anyhow::Error> { + self.fetch_user_devices_with_freshness(jids, crate::cache::Freshness::CachePreferred) + .await + } + + async fn fetch_user_devices_with_freshness( + &self, + mut jids: Vec<Jid>, + freshness: crate::cache::Freshness, + ) -> Result<Vec<Jid>, anyhow::Error> { + if jids.is_empty() { + return Ok(Vec::new()); + } + if freshness == crate::cache::Freshness::CachePreferred { let sid = self.generate_request_id(); - let jids_vec: Vec<Jid> = jids_to_fetch.into_iter().collect(); - let spec = DeviceListSpec::new(jids_vec, sid); + let response = self.execute(DeviceListSpec::new(jids, sid)).await?; + return self + .process_device_list_response(&response, freshness) + .await; + } - let response = self.execute(spec).await?; + for attempt in 0..DEVICE_REFRESH_MAX_ATTEMPTS { + let topology_generation = self.device_topology.current(); + let sid = self.generate_request_id(); + let response = self + .execute(DeviceListSpec::new(jids, sid).require_complete_response()) + .await?; + + if let Some(devices) = self + .try_process_refreshed_device_list_response( + &response, + freshness, + topology_generation, + ) + .await? + { + return Ok(devices); + } + + if attempt + 1 == DEVICE_REFRESH_MAX_ATTEMPTS { + anyhow::bail!( + "device registry kept changing while an authoritative refresh was in flight" + ); + } - let fetched_devices = self.process_device_list_response(&response).await; - all_devices.extend(fetched_devices); + // The complete response contains every requested identity. Move its + // canonical JIDs into the next attempt only on this rare race path; + // the ordinary refresh performs no duplicate query-vector clone. + jids = response + .device_lists + .into_iter() + .map(|user| user.user) + .collect(); } - Ok(all_devices) + unreachable!("bounded device refresh loop always returns") } - /// Apply a usync device-list response to the registry: persist LID mappings, - /// rebuild each returned user's `DeviceListRecord` (preserving key indices and - /// handling raw_id identity changes), and batch-write them. Returns the - /// resolved device JIDs for the users present in the response. - /// - /// Users the server OMITS — unchanged ones, when we sent a `device_hash` — are - /// simply absent here, so their cached records are left untouched (the - /// merge-safe behavior the `device_hash` optimization depends on). - #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.usync.process_device_list", level = "debug", skip_all, fields(users = response.device_lists.len())))] - async fn process_device_list_response(&self, response: &DeviceListResponse) -> Vec<Jid> { + async fn learn_device_list_mappings_guarded( + &self, + response: &DeviceListResponse, + guard: &crate::lid_pn_cache::LidPnMutationGuard<'_>, + ) -> Result<(), anyhow::Error> { // Learn LID↔PN mappings via the same batched, guarded learner query_info // uses (one detached transaction, skipping already-durable pairs), so - // per-mapping DB writes stay off the send's critical path. Falls back to - // the per-mapping path if the owning Arc<Client> isn't available. + // per-mapping DB writes stay off the send's critical path. Client + // construction always installs `self_weak`; failing the impossible + // upgrade keeps mapping and registry publication atomic. // // Ordering: the old per-mapping path AWAITED migrate_signal_sessions_on_lid_discovery. // Detaching it can let a standalone usync (sync_own_device_list / @@ -107,60 +190,105 @@ impl Client { // runs — but the per-address session_lock_for both take is the real // barrier (they can't interleave). The group-send path is unchanged: // query_info already learns these same pairs detached upstream. - if !response.lid_mappings.is_empty() { - if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { - let mappings: Vec<(String, String)> = response - .lid_mappings - .iter() - .map(|m| (m.lid.to_string(), m.phone_number.to_string())) - .collect(); - client - .learn_lid_pn_mappings_batch( - mappings, - crate::lid_pn_cache::LearningSource::Usync, - false, - ) - .await; - } else { - for mapping in &response.lid_mappings { - if let Err(err) = self - .add_lid_pn_mapping( - &mapping.lid, - &mapping.phone_number, - crate::lid_pn_cache::LearningSource::Usync, - ) - .await - { - warn!( - "Failed to persist LID {} -> {} from usync: {err}", - mapping.lid, mapping.phone_number, - ); - } - } - } + if response.lid_mappings.is_empty() { + return Ok(()); + } + let client = self + .self_weak + .get() + .and_then(|weak| weak.upgrade()) + .ok_or_else(|| anyhow::anyhow!("client ownership unavailable during device sync"))?; + let mappings: Vec<(String, String)> = response + .lid_mappings + .iter() + .map(|mapping| (mapping.lid.to_string(), mapping.phone_number.to_string())) + .collect(); + client + .learn_lid_pn_mappings_batch_guarded( + mappings, + crate::lid_pn_cache::LearningSource::Usync, + false, + guard, + ) + .await; + Ok(()) + } + + /// Apply a usync device-list response to the registry: persist LID mappings, + /// rebuild each returned user's `DeviceListRecord` (preserving key indices and + /// handling raw_id identity changes), and batch-write them. Returns the + /// resolved device JIDs for the users present in the response. + /// + /// Users the server OMITS — unchanged ones, when we sent a `device_hash` — are + /// simply absent here, so their cached records are left untouched (the + /// merge-safe behavior the `device_hash` optimization depends on). + #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.usync.process_device_list", level = "debug", skip_all, fields(users = response.device_lists.len())))] + async fn process_device_list_response( + &self, + response: &DeviceListResponse, + freshness: crate::cache::Freshness, + ) -> Result<Vec<Jid>, anyhow::Error> { + // Keep this order stable: mapping writers never acquire the registry + // guard while holding their lock, so mapping -> registry cannot cycle. + let mapping_guard = self.lid_pn_cache.lock_mutation().await; + let registry_guard = self.device_topology.lock_registry().await; + self.learn_device_list_mappings_guarded(response, &mapping_guard) + .await?; + self.process_device_list_response_guarded(response, freshness, &registry_guard) + .await + } + + async fn try_process_refreshed_device_list_response( + &self, + response: &DeviceListResponse, + freshness: crate::cache::Freshness, + topology_generation: u64, + ) -> Result<Option<Vec<Jid>>, anyhow::Error> { + // Serialize both mutation classes across the CAS and publication. The + // response's own mappings are deliberately learned only after the CAS. + let mapping_guard = self.lid_pn_cache.lock_mutation().await; + let registry_guard = self.device_topology.lock_registry().await; + if !self + .device_topology + .unchanged_for(topology_generation, |user| { + device_response_contains_user(response, user) + }) + { + return Ok(None); } + self.learn_device_list_mappings_guarded(response, &mapping_guard) + .await?; + self.process_device_list_response_guarded(response, freshness, &registry_guard) + .await + .map(Some) + } + async fn process_device_list_response_guarded( + &self, + response: &DeviceListResponse, + freshness: crate::cache::Freshness, + guard: &crate::client::device_topology::DeviceRegistryMutationGuard<'_>, + ) -> Result<Vec<Jid>, anyhow::Error> { let mut fetched_devices = Vec::with_capacity(response.device_lists.len()); let mut device_records: Vec<wacore::store::traits::DeviceListRecord> = Vec::with_capacity(response.device_lists.len()); + struct PendingIdentityReset<'a> { + user: &'a Jid, + previous: wacore::store::traits::DeviceListRecord, + invalidate_registry: bool, + } + // Identity changes are rare, so this stays allocation-free for the + // ordinary response. More importantly, deferring the destructive work + // lets an authoritative refresh validate every user before any prior + // Signal sessions or registry snapshot are discarded. + let mut pending_identity_resets = Vec::new(); for user_list in &response.device_lists { // Update device registry (single source of truth for device lists). // Preserve key_index values from existing records (set via account_sync) // Use alias-aware lookup (resolves LID ↔ PN) to find // existing record regardless of which key it was stored under - let existing_record = self.load_device_record(&user_list.user.user).await; - - let mut existing_key_indices: std::collections::HashMap<u32, Option<u32>> = - existing_record - .as_ref() - .map(|r| { - r.devices - .iter() - .map(|d| (d.device_id, d.key_index)) - .collect() - }) - .unwrap_or_default(); + let mut existing_record = self.load_device_record(&user_list.user.user).await; // Decode key-index-list if present (WA Web: handleKeyIndexResult) let decoded_key_index = user_list @@ -171,32 +299,29 @@ impl Client { // Check raw_id mismatch for identity change detection // TODO: also check advAccountType mismatch (see patch_device_add TODO) let mut raw_id = decoded_key_index.as_ref().map(|d| d.raw_id); - if let Some(ref decoded) = decoded_key_index + let pending_identity_reset = if let Some(ref decoded) = decoded_key_index && let Some(ref existing) = existing_record && let Some(stored_raw_id) = existing.raw_id && stored_raw_id != decoded.raw_id { log::info!( - "raw_id mismatch for user {} in usync: stored={stored_raw_id}, received={}. Clearing record.", + "raw_id mismatch for user {} in usync: stored={stored_raw_id}, received={}. Scheduling record reset.", user_list.user.user, decoded.raw_id ); - self.clear_device_record( - &user_list.user.user, - user_list.user.server.as_str(), - existing, - ) - .await; - // Old key indices are from the previous identity — don't reuse - existing_key_indices.clear(); - } + existing_record.take() + } else { + None + }; // Preserve raw_id from existing when usync didn't provide one - // (no key-index-list) and no mismatch cleared the indices. - // existing_key_indices is empty after a mismatch clear, so this - // correctly skips preservation after identity change. - if raw_id.is_none() && !existing_key_indices.is_empty() { - raw_id = existing_record.as_ref().and_then(|r| r.raw_id); + // (no key-index-list). An identity change takes the old record out + // above, so its raw_id and key indices cannot leak into the new one. + if raw_id.is_none() { + raw_id = existing_record + .as_ref() + .filter(|record| !record.devices.is_empty()) + .and_then(|record| record.raw_id); } let mut devices: Vec<wacore::store::traits::DeviceInfo> = user_list @@ -205,10 +330,16 @@ impl Client { .map(|d| { // Server-returned key_index takes priority over cached let key_index = d.key_index.or_else(|| { - existing_key_indices - .get(&(d.device as u32)) - .copied() - .flatten() + // Accounts ordinarily have only a handful of companion + // devices; a short scan avoids allocating a HashMap for + // every user in a large fanout response. + existing_record.as_ref().and_then(|record| { + record + .devices + .iter() + .find(|cached| cached.device_id == d.device as u32) + .and_then(|cached| cached.key_index) + }) }); wacore::store::traits::DeviceInfo::new(d.device as u32, key_index) .with_hosting(d.is_hosted) @@ -217,13 +348,7 @@ impl Client { // Apply valid_indexes filtering if key-index-list was decoded if let Some(ref decoded) = decoded_key_index { - devices = wacore::adv::filter_devices_by_key_index(&devices, decoded); - } - - // Convert filtered DeviceInfo list back to JIDs for return - let user_jid = &user_list.user; - for d in &devices { - fetched_devices.push(user_jid.with_device_hosting(d.device_id as u16, d.is_hosted)); + wacore::adv::retain_devices_by_key_index(&mut devices, decoded); } // An empty device list is never valid — WA Web always keeps the primary @@ -231,9 +356,39 @@ impl Client { // corrupt. Persisting it would clobber a good cached record, or store an // empty one that get_user_devices then re-fetches on every send. if devices.is_empty() { + if freshness == crate::cache::Freshness::Refresh { + anyhow::bail!( + "device-list refresh left no valid devices for {}", + user_list.user + ); + } + if let Some(previous) = pending_identity_reset { + pending_identity_resets.push(PendingIdentityReset { + user: &user_list.user, + previous, + // No replacement record will be written. Keeping the old + // snapshot would pair a new identity with stale devices + // and suppress the next authoritative network fetch. + invalidate_registry: true, + }); + } continue; } + if let Some(previous) = pending_identity_reset { + pending_identity_resets.push(PendingIdentityReset { + user: &user_list.user, + previous, + invalidate_registry: false, + }); + } + + // Convert filtered DeviceInfo list back to JIDs for return + let user_jid = &user_list.user; + for d in &devices { + fetched_devices.push(user_jid.with_device_hosting(d.device_id as u16, d.is_hosted)); + } + device_records.push(wacore::store::traits::DeviceListRecord { user: user_list.user.user.to_string(), devices, @@ -243,14 +398,33 @@ impl Client { }); } + // All strict validation has completed. Apply identity cleanup before + // publishing replacement snapshots so no send can pair a new registry + // record with sessions established under the previous identity. + for reset in pending_identity_resets { + self.clear_device_record( + &reset.user.user, + reset.user.server.as_str(), + &reset.previous, + ) + .await; + if reset.invalidate_registry { + self.invalidate_device_cache_guarded(&reset.user.user, guard) + .await; + } + } + // One batched backend write for the whole usync response — for // large groups this collapses N spawn_blocking SQLite hops into // a single transaction, which dominated the per-send wall-clock. - if let Err(e) = self.update_device_lists(device_records).await { + if let Err(e) = self + .update_device_lists_guarded(device_records, guard) + .await + { warn!("Failed to update device registry batch: {e}"); } - fetched_devices + Ok(fetched_devices) } /// Re-sync own device list from the server. Mirrors WA Web `syncMyDeviceList`: @@ -293,7 +467,9 @@ impl Client { let response = self.execute(spec).await?; // `process_device_list_response` only touches users the server actually // returned, so unchanged (omitted) own devices keep their cache. - let devices = self.process_device_list_response(&response).await; + let devices = self + .process_device_list_response(&response, crate::cache::Freshness::CachePreferred) + .await?; log::info!( "Re-synced own device list: {} device(s) updated", devices.len() @@ -347,8 +523,44 @@ impl Client { #[cfg(test)] mod tests { use super::*; + use crate::cache::Freshness; use crate::test_utils::create_test_client; + use wacore::libsignal::protocol::{ProtocolAddress, SessionRecord}; use wacore::store::traits::{DeviceInfo, DeviceListRecord}; + use wacore::types::jid::JidExt; + + fn signed_key_index_bytes(valid_indexes: Vec<u32>, current_index: u32) -> Vec<u8> { + let key_index = waproto::whatsapp::ADVKeyIndexList { + raw_id: Some(1), + timestamp: Some(1_700_000_000), + current_index: Some(current_index), + valid_indexes, + ..Default::default() + }; + let signed = waproto::whatsapp::ADVSignedKeyIndexList { + details: Some(waproto::codec::adv_key_index_list_to_vec(&key_index)), + ..Default::default() + }; + waproto::codec::adv_signed_key_index_list_to_vec(&signed) + } + + async fn seed_fresh_session(client: &Client, jid: &Jid) -> ProtocolAddress { + let address = jid.to_protocol_address(); + client + .signal_cache + .put_session(&address, SessionRecord::new_fresh()) + .await; + address + } + + async fn has_session(client: &Client, address: &ProtocolAddress) -> bool { + let snapshot = client.persistence_manager.get_device_snapshot(); + client + .signal_cache + .has_session(address, &*snapshot.backend) + .await + .unwrap() + } #[tokio::test] async fn test_device_registry_hit_resolves_devices() { @@ -374,6 +586,45 @@ mod tests { assert!(devices.iter().all(|d| d.is_pn())); } + #[tokio::test] + async fn refresh_bypasses_a_warm_registry_without_clearing_it_first() { + let client = create_test_client().await; + let user: Jid = "12025550102@s.whatsapp.net".parse().unwrap(); + client + .update_device_list(DeviceListRecord { + user: "12025550102".into(), + devices: vec![DeviceInfo::new(0, None), DeviceInfo::new(8, None)], + timestamp: wacore::time::now_secs(), + phash: None, + raw_id: None, + }) + .await + .unwrap(); + + let cached = client + .get_user_devices(std::slice::from_ref(&user)) + .await + .unwrap(); + assert_eq!( + cached.len(), + 2, + "cache-preferred must use the warm snapshot" + ); + + let refresh = client.refresh_user_devices(vec![user.clone()]).await; + assert!( + refresh.is_err(), + "the offline fixture proves refresh consulted the source" + ); + + let preserved = client + .get_devices_from_registry(&user) + .await + .expect("a failed refresh must leave the previous snapshot readable"); + assert_eq!(preserved.len(), 2); + assert!(preserved.iter().any(|device| device.device == 8)); + } + #[tokio::test] async fn test_device_registry_hit_for_lid_jid() { let client = create_test_client().await; @@ -413,7 +664,10 @@ mod tests { client.update_device_list(record).await.unwrap(); // Evict from registry cache to force DB path - client.device_registry_cache.invalidate("9876543210").await; + client + .device_registry_cache + .raw_invalidate_for_tests("9876543210") + .await; client.device_registry_cache.run_pending_tasks().await; // Should still resolve from DB @@ -501,7 +755,10 @@ mod tests { lid_mappings: vec![], }; - let fetched = client.process_device_list_response(&response).await; + let fetched = client + .process_device_list_response(&response, Freshness::CachePreferred) + .await + .unwrap(); assert!( fetched.iter().any(|j| j.user == "1111111111"), "returned user A must be resolved" @@ -540,7 +797,10 @@ mod tests { lid_mappings: Vec::new(), }; - let fetched = client.process_device_list_response(&response).await; + let fetched = client + .process_device_list_response(&response, Freshness::CachePreferred) + .await + .unwrap(); assert!( fetched .iter() @@ -558,6 +818,216 @@ mod tests { ); } + #[tokio::test] + async fn stale_refresh_does_not_overwrite_a_newer_device_notification() { + use wacore::usync::{UserDeviceList, UsyncDevice}; + + let client = create_test_client().await; + let user = Jid::pn("12025550102"); + client + .update_device_list(DeviceListRecord { + user: user.user.to_string(), + devices: vec![DeviceInfo::new(0, None)], + timestamp: wacore::time::now_secs(), + phash: Some("1:before".to_string()), + raw_id: None, + }) + .await + .unwrap(); + + let refresh_started_at = client.device_topology.current(); + let stale_response = DeviceListResponse { + device_lists: vec![UserDeviceList { + user: user.clone(), + devices: vec![UsyncDevice::new(0, None)], + phash: Some("1:stale".to_string()), + key_index_bytes: None, + }], + lid_mappings: Vec::new(), + }; + + client + .patch_device_add( + &user.user, + &wacore::stanza::devices::DeviceElement { + jid: user.with_device(7), + key_index: None, + lid: None, + }, + None, + ) + .await; + + let published = client + .try_process_refreshed_device_list_response( + &stale_response, + Freshness::Refresh, + refresh_started_at, + ) + .await + .unwrap(); + assert!(published.is_none(), "the stale response must be retried"); + + let retained = client + .get_devices_from_registry(&user) + .await + .expect("notification snapshot must remain available"); + assert!(retained.iter().any(|device| device.device == 7)); + } + + #[tokio::test] + async fn stale_refresh_does_not_overwrite_a_newer_lid_mapping() { + use wacore::usync::UsyncLidMapping; + + let client = create_test_client().await; + let phone = "12025550110"; + let current_lid = "100000000000110"; + let stale_lid = "100000000000111"; + let refresh_started_at = client.device_topology.current(); + + client + .add_lid_pn_mapping( + current_lid, + phone, + crate::lid_pn_cache::LearningSource::PeerPnMessage, + ) + .await + .unwrap(); + + let stale_response = DeviceListResponse { + device_lists: Vec::new(), + lid_mappings: vec![UsyncLidMapping { + phone_number: phone.into(), + lid: stale_lid.into(), + }], + }; + let published = client + .try_process_refreshed_device_list_response( + &stale_response, + Freshness::Refresh, + refresh_started_at, + ) + .await + .unwrap(); + + assert!(published.is_none(), "the stale response must be retried"); + assert_eq!( + client.lid_pn_cache.get_current_lid(phone).await.as_deref(), + Some(current_lid), + "the mapping learned after the refresh started must survive" + ); + } + + #[tokio::test] + async fn refresh_commits_its_own_lid_mapping_after_the_cas() { + use wacore::usync::UsyncLidMapping; + + let client = create_test_client().await; + let phone = "12025550112"; + let lid = "100000000000112"; + let refresh_started_at = client.device_topology.current(); + let response = DeviceListResponse { + device_lists: Vec::new(), + lid_mappings: vec![UsyncLidMapping { + phone_number: phone.into(), + lid: lid.into(), + }], + }; + + let published = client + .try_process_refreshed_device_list_response( + &response, + Freshness::Refresh, + refresh_started_at, + ) + .await + .unwrap(); + + assert!( + published.is_some(), + "the response must not conflict with itself" + ); + assert_eq!( + client.lid_pn_cache.get_current_lid(phone).await.as_deref(), + Some(lid) + ); + } + + #[tokio::test] + async fn unrelated_registry_change_does_not_restart_a_refresh() { + use wacore::usync::{UserDeviceList, UsyncDevice}; + + let client = create_test_client().await; + let refreshed = Jid::pn("12025550104"); + let refresh_started_at = client.device_topology.current(); + client + .update_device_list(DeviceListRecord { + user: "12025550105".to_string(), + devices: vec![DeviceInfo::new(0, None)], + timestamp: wacore::time::now_secs(), + phash: None, + raw_id: None, + }) + .await + .unwrap(); + + let response = DeviceListResponse { + device_lists: vec![UserDeviceList { + user: refreshed, + devices: vec![UsyncDevice::new(0, None)], + phash: None, + key_index_bytes: None, + }], + lid_mappings: Vec::new(), + }; + let published = client + .try_process_refreshed_device_list_response( + &response, + Freshness::Refresh, + refresh_started_at, + ) + .await + .unwrap(); + assert!(published.is_some()); + } + + #[tokio::test] + async fn unrelated_mapping_change_does_not_restart_a_refresh() { + use wacore::usync::{UserDeviceList, UsyncDevice}; + + let client = create_test_client().await; + let refreshed = Jid::pn("12025550113"); + let refresh_started_at = client.device_topology.current(); + client + .add_lid_pn_mapping( + "100000000000114", + "12025550114", + crate::lid_pn_cache::LearningSource::PeerPnMessage, + ) + .await + .unwrap(); + + let response = DeviceListResponse { + device_lists: vec![UserDeviceList { + user: refreshed, + devices: vec![UsyncDevice::new(0, None)], + phash: None, + key_index_bytes: None, + }], + lid_mappings: Vec::new(), + }; + let published = client + .try_process_refreshed_device_list_response( + &response, + Freshness::Refresh, + refresh_started_at, + ) + .await + .unwrap(); + + assert!(published.is_some()); + } + /// A usync that returns an empty device list for a user is transient or /// corrupt (WA Web always keeps device 0). `process_device_list_response` /// must not persist it: a good cached record stays intact instead of being @@ -591,7 +1061,10 @@ mod tests { lid_mappings: vec![], }; - let fetched = client.process_device_list_response(&response).await; + let fetched = client + .process_device_list_response(&response, Freshness::CachePreferred) + .await + .unwrap(); assert!( !fetched.iter().any(|j| j.user == "3333333333"), "an empty returned list contributes no devices" @@ -610,6 +1083,153 @@ mod tests { ); } + #[tokio::test] + async fn refresh_rejects_a_device_list_emptied_by_key_index_filtering() { + use wacore::usync::{UserDeviceList, UsyncDevice}; + + let client = create_test_client().await; + let user = Jid::pn("4444444444"); + client + .update_device_list(DeviceListRecord { + user: user.user.to_string(), + devices: vec![DeviceInfo::new(0, None), DeviceInfo::new(7, Some(3))], + timestamp: wacore::time::now_secs(), + phash: Some("2:previous".to_string()), + raw_id: Some(1), + }) + .await + .unwrap(); + + // The wire response is non-empty, but its only companion is outside the + // signed key-index set. This exercises the post-projection completeness + // check rather than the raw USync response check. + let response = DeviceListResponse { + device_lists: vec![UserDeviceList { + user: user.clone(), + devices: vec![UsyncDevice::new(7, Some(3))], + phash: Some("2:incomplete".to_string()), + key_index_bytes: Some(signed_key_index_bytes(Vec::new(), 10)), + }], + lid_mappings: Vec::new(), + }; + + let error = client + .process_device_list_response(&response, Freshness::Refresh) + .await + .expect_err("an authoritative refresh must not return an empty projection"); + assert!(error.to_string().contains("no valid devices")); + + let preserved = client + .get_devices_from_registry(&user) + .await + .expect("a rejected refresh must preserve the previous snapshot"); + assert_eq!(preserved.len(), 2); + assert!(preserved.iter().any(|device| device.device == 7)); + } + + #[tokio::test] + async fn rejected_refresh_defers_identity_cleanup_for_every_user() { + use wacore::usync::{UserDeviceList, UsyncDevice}; + + let client = create_test_client().await; + let identity_changed = Jid::pn("4444444451"); + let invalid = Jid::pn("4444444452"); + + for (user, raw_id) in [(&identity_changed, 2), (&invalid, 1)] { + client + .update_device_list(DeviceListRecord { + user: user.user.to_string(), + devices: vec![DeviceInfo::new(0, None), DeviceInfo::new(7, Some(3))], + timestamp: wacore::time::now_secs(), + phash: Some("2:previous".to_string()), + raw_id: Some(raw_id), + }) + .await + .unwrap(); + } + + let previous_session = seed_fresh_session(&client, &identity_changed.with_device(7)).await; + let response = DeviceListResponse { + device_lists: vec![ + UserDeviceList { + user: identity_changed.clone(), + // A primary device survives key-index filtering, so this + // first user schedules a valid identity replacement. + devices: vec![UsyncDevice::new(0, None)], + phash: Some("1:changed".to_string()), + key_index_bytes: Some(signed_key_index_bytes(Vec::new(), 10)), + }, + UserDeviceList { + user: invalid, + // The second user makes the authoritative response invalid + // only after key-index projection. + devices: vec![UsyncDevice::new(7, Some(3))], + phash: Some("1:invalid".to_string()), + key_index_bytes: Some(signed_key_index_bytes(Vec::new(), 10)), + }, + ], + lid_mappings: Vec::new(), + }; + + client + .process_device_list_response(&response, Freshness::Refresh) + .await + .expect_err("the second user must reject the whole authoritative refresh"); + + assert!( + has_session(&client, &previous_session).await, + "validation failure must not partially clear an earlier user's sessions" + ); + let preserved = client + .get_devices_from_registry(&identity_changed) + .await + .expect("validation failure must preserve the earlier registry snapshot"); + assert!(preserved.iter().any(|device| device.device == 7)); + } + + #[tokio::test] + async fn filtered_identity_change_invalidates_the_stale_registry() { + use wacore::usync::{UserDeviceList, UsyncDevice}; + + let client = create_test_client().await; + let user = Jid::pn("4444444453"); + client + .update_device_list(DeviceListRecord { + user: user.user.to_string(), + devices: vec![DeviceInfo::new(0, None), DeviceInfo::new(7, Some(3))], + timestamp: wacore::time::now_secs(), + phash: Some("2:previous".to_string()), + raw_id: Some(2), + }) + .await + .unwrap(); + let previous_session = seed_fresh_session(&client, &user.with_device(7)).await; + + let response = DeviceListResponse { + device_lists: vec![UserDeviceList { + user: user.clone(), + devices: vec![UsyncDevice::new(7, Some(3))], + phash: Some("1:changed".to_string()), + key_index_bytes: Some(signed_key_index_bytes(Vec::new(), 10)), + }], + lid_mappings: Vec::new(), + }; + + let fetched = client + .process_device_list_response(&response, Freshness::CachePreferred) + .await + .unwrap(); + assert!(fetched.is_empty()); + assert!( + !has_session(&client, &previous_session).await, + "an accepted identity change must clear sessions from the old identity" + ); + assert!( + client.get_devices_from_registry(&user).await.is_none(), + "without a replacement snapshot, the stale registry must be invalidated" + ); + } + /// The batched LID-PN learn path warms the in-memory cache SYNCHRONOUSLY /// (the persist runs detached), so a mapping from the usync response is /// resolvable the moment `process_device_list_response` returns. Locks the @@ -646,7 +1266,10 @@ mod tests { .is_none() ); - let _ = client.process_device_list_response(&response).await; + client + .process_device_list_response(&response, Freshness::CachePreferred) + .await + .unwrap(); assert_eq!( client diff --git a/wacore/benches/send_receive_benchmark.rs b/wacore/benches/send_receive_benchmark.rs index 1a8e5cb83..ba3974401 100644 --- a/wacore/benches/send_receive_benchmark.rs +++ b/wacore/benches/send_receive_benchmark.rs @@ -15,7 +15,10 @@ use std::hint::black_box; use wacore::client::context::{GroupInfo, SendContextResolver}; use wacore::messages::MessageUtils; use wacore::runtime::{AbortHandle, Runtime}; -use wacore::send::{SignalStores, prepare_group_stanza, prepare_peer_stanza}; +use wacore::send::{ + GroupStanzaRequest, SenderKeyDistributionPolicy, SignalStores, prepare_group_stanza, + prepare_peer_stanza, +}; use wacore::types::jid::{JidExt, make_sender_key_name}; use wacore::types::message::AddressingMode; use wacore_binary::JidExt as _; @@ -733,23 +736,27 @@ fn setup_group_recv() -> GrpRecvData { }; let runtime = BenchRuntime; + let message = text_msg(); let result = futures::executor::block_on(prepare_group_stanza( &runtime, &mut stores, &resolver, - &group_info, - &own_jid, - &own_jid, - None, - group_jid.clone(), - &text_msg(), - "bench-grp-recv".into(), - false, - None, - None, - None, - &[], - None, + GroupStanzaRequest { + group: &group_info, + own_jid: &own_jid, + own_lid: &own_jid, + account: None, + to: &group_jid, + message: &message, + message_id: "bench-grp-recv", + force_distribution: false, + distribution_targets: None, + distribution_policy: SenderKeyDistributionPolicy::BestEffort, + phash_devices: None, + edit: None, + extra_nodes: &[], + pre_encoded: None, + }, )) .unwrap(); @@ -803,7 +810,6 @@ fn run_group_send(d: &mut GrpSendData) { // only emits a phash if it gets the full device set. Mirror the real // warm-send caller by passing it; the cold/force_skdm path resolves the set // itself and keeps None. - let all_devices_for_phash = d.resolved_for_phash.clone(); let mut group_info = GroupInfo::new(std::mem::take(&mut d.participants), AddressingMode::Pn); let own_base = own_jid.to_non_ad(); if !group_info @@ -825,19 +831,22 @@ fn run_group_send(d: &mut GrpSendData) { &d.runtime, &mut stores, &d.resolver, - &group_info, - &own_jid, - &own_jid, - None, - d.group_jid.clone(), - &d.msg, - "b-grp".into(), - d.force_skdm, - None, - all_devices_for_phash, - None, - &[], - None, + GroupStanzaRequest { + group: &group_info, + own_jid: &own_jid, + own_lid: &own_jid, + account: None, + to: &d.group_jid, + message: &d.msg, + message_id: "b-grp", + force_distribution: d.force_skdm, + distribution_targets: None, + distribution_policy: SenderKeyDistributionPolicy::BestEffort, + phash_devices: d.resolved_for_phash.as_deref(), + edit: None, + extra_nodes: &[], + pre_encoded: None, + }, )) .unwrap(); diff --git a/wacore/src/adv.rs b/wacore/src/adv.rs index 13ab70ad2..a4c628906 100644 --- a/wacore/src/adv.rs +++ b/wacore/src/adv.rs @@ -149,25 +149,24 @@ pub fn filter_devices_by_key_index( devices: &[DeviceInfo], decoded: &DecodedKeyIndex, ) -> Vec<DeviceInfo> { - let valid_set: std::collections::HashSet<u32> = decoded.valid_indexes.iter().copied().collect(); - devices .iter() - .filter(|d| { - // Primary device always kept - if d.device_id == 0 { - return true; - } - match d.key_index { - Some(ki) => valid_set.contains(&ki) || ki > decoded.current_index, - // WA Web: h.has(null) → false, null > y → false → device removed - None => false, - } - }) + .filter(|device| should_retain_device(device, decoded)) .cloned() .collect() } +/// Filter a device list in place using the same ADV key-index rules as +/// [`filter_devices_by_key_index`]. This avoids allocating and copying a second +/// list when the caller already owns its device snapshot. +pub fn retain_devices_by_key_index(devices: &mut Vec<DeviceInfo>, decoded: &DecodedKeyIndex) { + devices.retain(|device| should_retain_device(device, decoded)); +} + +fn should_retain_device(device: &DeviceInfo, decoded: &DecodedKeyIndex) -> bool { + device.device_id == 0 || is_key_index_valid(device.key_index, decoded) +} + /// Check if a key_index is accepted by the decoded ADV list. /// Used to validate a newly-notified device before adding it to the registry. /// diff --git a/wacore/src/iq/usync.rs b/wacore/src/iq/usync.rs index 37468646d..8f9abfea3 100644 --- a/wacore/src/iq/usync.rs +++ b/wacore/src/iq/usync.rs @@ -286,6 +286,16 @@ fn project_business( pub(crate) fn project_lid_mapping(user: &UsyncUserResult) -> Option<UsyncLidMapping> { let user_jid = user.id.as_ref()?; + if user_jid.server.is_lid_family() { + let pn_jid = user + .pn_jid + .as_ref() + .filter(|jid| jid.server.is_pn_family())?; + return Some(UsyncLidMapping { + phone_number: pn_jid.user.clone(), + lid: user_jid.user.clone(), + }); + } if !user_jid.server.is_pn_family() { return None; } @@ -300,6 +310,28 @@ pub(crate) fn project_lid_mapping(user: &UsyncUserResult) -> Option<UsyncLidMapp }) } +#[inline] +fn device_list_identity_matches( + returned: &Jid, + expected: &Jid, + mappings: &[UsyncLidMapping], +) -> bool { + if returned.user == expected.user && returned.server == expected.server { + return true; + } + + mappings.iter().any(|mapping| { + (returned.server.is_lid_family() + && expected.server.is_pn_family() + && returned.user == mapping.lid + && expected.user == mapping.phone_number) + || (returned.server.is_pn_family() + && expected.server.is_lid_family() + && returned.user == mapping.phone_number + && expected.user == mapping.lid) + }) +} + fn is_on_whatsapp_query(spec: &IsOnWhatsAppSpec) -> UsyncQuery { let mut protocols = Vec::with_capacity(3); if spec.query_type == IsOnWhatsAppQueryType::Pn { @@ -587,8 +619,24 @@ impl DeviceListSpec { hashes, } } + + /// Reject a response that omits any requested user or returns no devices + /// for one. Use for authoritative refreshes; ordinary fanout remains + /// best-effort when an individual user cannot be resolved. + pub fn require_complete_response(self) -> CompleteDeviceListSpec { + CompleteDeviceListSpec(self) + } } +/// A device-list query whose response must contain a usable entry for every +/// requested user. +/// +/// This wrapper keeps [`DeviceListSpec`]'s public data model unchanged while +/// making authoritative refresh semantics explicit in the type system. It has +/// no additional runtime state. +#[derive(Debug, Clone)] +pub struct CompleteDeviceListSpec(DeviceListSpec); + pub(crate) fn device_list_query( jids: &[Jid], hashes: Option<&HashMap<Jid, (String, i64)>>, @@ -711,6 +759,33 @@ impl IqSpec for DeviceListSpec { } } +impl IqSpec for CompleteDeviceListSpec { + type Response = DeviceListResponse; + + fn build_iq(&self) -> InfoQuery<'static> { + self.0.build_iq() + } + + fn encode_iq_direct(&self, request_id: &str, out: &mut Vec<u8>) -> Result<bool, anyhow::Error> { + self.0.encode_iq_direct(request_id, out) + } + + fn parse_response(&self, response: &NodeRef<'_>) -> Result<Self::Response, anyhow::Error> { + let parsed = self.0.parse_response(response)?; + for expected in &self.0.jids { + let Some(returned) = parsed.device_lists.iter().find(|returned| { + device_list_identity_matches(&returned.user, expected, &parsed.lid_mappings) + }) else { + anyhow::bail!("device-list response omitted user {expected}"); + }; + if returned.devices.is_empty() { + anyhow::bail!("device-list response returned no devices for {expected}"); + } + } + Ok(parsed) + } +} + /// Resolve PN→LID mappings for JIDs without a known LID. /// Matches WA Web's `ensurePhoneNumberToLidMapping` (PhoneNumberMappingJob.js). /// Uses a separate usync with only `<lid/>` in the query to avoid side effects @@ -1340,6 +1415,21 @@ mod tests { } } + #[test] + fn device_list_spec_preserves_its_public_data_shape() { + let jid = Jid::pn("1234567890"); + let spec = DeviceListSpec { + jids: vec![jid.clone()], + sid: "test-sid".to_string(), + hashes: std::collections::HashMap::new(), + }; + + let DeviceListSpec { jids, sid, hashes } = spec; + assert_eq!(jids, [jid]); + assert_eq!(sid, "test-sid"); + assert!(hashes.is_empty()); + } + #[test] fn test_device_list_spec_build_iq_with_device_hash() { let jid: Jid = "1234567890@s.whatsapp.net".parse().unwrap(); @@ -1384,7 +1474,11 @@ mod tests { // keeps the omitted user's cached devices (device_hash merge-safety). let a: Jid = "1111111111@s.whatsapp.net".parse().unwrap(); let b: Jid = "2222222222@s.whatsapp.net".parse().unwrap(); - let spec = DeviceListSpec::new(vec![a, b], "sid-omit"); + let best_effort = DeviceListSpec::new(vec![a.clone(), b.clone()], "sid-best-effort"); + let complete = DeviceListSpec::new(vec![a.clone(), b.clone()], "sid-complete") + .require_complete_response(); + let incremental = + DeviceListSpec::with_hashes(vec![a, b], "sid-omit", std::collections::HashMap::new()); let response = NodeBuilder::new("iq") .attr("type", "result") @@ -1403,9 +1497,15 @@ mod tests { .build()]) .build(); - let result = spec.parse_response(&response.as_node_ref()).unwrap(); - assert_eq!(result.device_lists.len(), 1, "omitted user must not appear"); - assert_eq!(result.device_lists[0].user.user, "1111111111"); + assert!( + complete.parse_response(&response.as_node_ref()).is_err(), + "a complete query must not accept a partial snapshot" + ); + for spec in [best_effort, incremental] { + let result = spec.parse_response(&response.as_node_ref()).unwrap(); + assert_eq!(result.device_lists.len(), 1, "omitted user must not appear"); + assert_eq!(result.device_lists[0].user.user, "1111111111"); + } } #[test] @@ -1459,7 +1559,15 @@ mod tests { fn device_list_devices_error_skips_only_that_user() { let jid1: Jid = "1234567890@s.whatsapp.net".parse().unwrap(); let jid2: Jid = "9876543210@s.whatsapp.net".parse().unwrap(); - let spec = DeviceListSpec::new(vec![jid1, jid2], "test-sid"); + let best_effort = + DeviceListSpec::new(vec![jid1.clone(), jid2.clone()], "test-sid-best-effort"); + let complete = DeviceListSpec::new(vec![jid1.clone(), jid2.clone()], "test-sid-complete") + .require_complete_response(); + let incremental = DeviceListSpec::with_hashes( + vec![jid1, jid2], + "test-sid", + std::collections::HashMap::new(), + ); let response = NodeBuilder::new("iq") .attr("type", "result") @@ -1494,9 +1602,15 @@ mod tests { .build()]) .build(); - let result = spec.parse_response(&response.as_node_ref()).unwrap(); - assert_eq!(result.device_lists.len(), 1); - assert_eq!(result.device_lists[0].user.user, "9876543210"); + assert!( + complete.parse_response(&response.as_node_ref()).is_err(), + "a per-user error makes a complete snapshot unusable" + ); + for spec in [best_effort, incremental] { + let result = spec.parse_response(&response.as_node_ref()).unwrap(); + assert_eq!(result.device_lists.len(), 1); + assert_eq!(result.device_lists[0].user.user, "9876543210"); + } } #[test] @@ -1595,9 +1709,48 @@ mod tests { let mapping = project_lid_mapping(&user).expect("expected LID mapping"); assert_eq!(mapping.phone_number, "13135550100"); assert_eq!(mapping.lid, "100000000000100"); + + let canonicalized = UsyncUserResult { + id: Some(Jid::new("100000000000100", lid_server)), + pn_jid: Some(Jid::new("13135550100", pn_server)), + protocols: Vec::new(), + }; + let mapping = project_lid_mapping(&canonicalized).expect("expected pn_jid mapping"); + assert_eq!(mapping.phone_number, "13135550100"); + assert_eq!(mapping.lid, "100000000000100"); } } + #[test] + fn complete_device_list_accepts_server_canonicalized_lid() { + let requested = Jid::pn("12025550100"); + let spec = + DeviceListSpec::new(vec![requested], "sid-canonicalized").require_complete_response(); + + let response = NodeBuilder::new("iq") + .attr("type", "result") + .children([NodeBuilder::new("usync") + .children([NodeBuilder::new("list") + .children([NodeBuilder::new("user") + .attr("jid", "100000000000100@lid") + .attr("pn_jid", "12025550100@s.whatsapp.net") + .children([NodeBuilder::new("devices") + .children([NodeBuilder::new("device-list") + .children([NodeBuilder::new("device").attr("id", "0").build()]) + .build()]) + .build()]) + .build()]) + .build()]) + .build()]) + .build(); + + let parsed = spec.parse_response(&response.as_node_ref()).unwrap(); + assert_eq!(parsed.device_lists[0].user, Jid::lid("100000000000100")); + assert_eq!(parsed.lid_mappings.len(), 1); + assert_eq!(parsed.lid_mappings[0].phone_number, "12025550100"); + assert_eq!(parsed.lid_mappings[0].lid, "100000000000100"); + } + #[test] fn test_device_list_spec_parse_response_multiple_users() { let jid1: Jid = "1111111111@s.whatsapp.net".parse().unwrap(); diff --git a/wacore/src/request.rs b/wacore/src/request.rs index 50305af7f..6a768d18f 100644 --- a/wacore/src/request.rs +++ b/wacore/src/request.rs @@ -122,7 +122,7 @@ pub struct ServerErrorCode { impl ServerErrorCode { pub fn from_anyhow(err: &anyhow::Error) -> Option<&Self> { - err.downcast_ref::<Self>() + err.chain().find_map(|cause| cause.downcast_ref::<Self>()) } } diff --git a/wacore/src/send/dm.rs b/wacore/src/send/dm.rs index c5eb06725..e59804b8f 100644 --- a/wacore/src/send/dm.rs +++ b/wacore/src/send/dm.rs @@ -1,6 +1,7 @@ //! 1:1 (DM) stanza preparation and DM retry stanzas. use super::*; +use anyhow::Context as _; fn is_exact_dm_sender_device(device_jid: &Jid, own_jid: &Jid, own_lid: Option<&Jid>) -> bool { (device_jid.is_same_user_as(own_jid) && device_jid.device == own_jid.device) @@ -75,25 +76,41 @@ pub struct PreparedDmStanza { pub message_secret: Option<[u8; crate::reporting_token::MESSAGE_SECRET_SIZE]>, } -#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.dm_prepare", level = "debug", skip_all, fields(to = %to_jid.observe()), err(Debug)))] -#[allow(clippy::too_many_arguments)] +pub struct DmStanzaRequest<'a> { + pub own_jid: &'a Jid, + pub own_lid: Option<&'a Jid>, + pub account: Option<&'a wa::ADVSignedDeviceIdentity>, + pub to: &'a Jid, + pub message: &'a wa::Message, + pub message_id: &'a str, + pub edit: Option<&'a crate::types::message::EditAttribute>, + pub extra_nodes: &'a [Node], + pub devices: Vec<Jid>, + pub pre_encoded: Option<&'a [u8]>, +} + +#[cfg_attr( + feature = "tracing", + tracing::instrument(name = "wa.send.dm_prepare", level = "debug", skip_all, err(Debug)) +)] pub async fn prepare_dm_stanza( runtime: &dyn Runtime, stores: &mut SignalStores<'_>, resolver: &dyn SendContextResolver, - own_jid: &Jid, - own_lid: Option<&Jid>, - account: Option<&wa::ADVSignedDeviceIdentity>, - to_jid: Jid, - message: &wa::Message, - request_id: String, - edit: Option<crate::types::message::EditAttribute>, - extra_stanza_nodes: &[Node], - all_devices: Vec<Jid>, - // Avoids a second full encode when the caller already serialized the message; - // ignored on the mci-hoist path (see `shared_content`). - pre_encoded: Option<std::sync::Arc<Vec<u8>>>, + request: DmStanzaRequest<'_>, ) -> Result<PreparedDmStanza> { + let DmStanzaRequest { + own_jid, + own_lid, + account, + to: to_jid, + message, + message_id: request_id, + edit, + extra_nodes: extra_stanza_nodes, + devices: all_devices, + pre_encoded, + } = request; // Encode the message at most once (reusing the caller's `pre_encoded` bytes when // provided) and thread those bytes through both the reporting token // (whitelisted-field extraction) and the wire plaintext below. The rare mci-hoist @@ -101,7 +118,10 @@ pub async fn prepare_dm_stanza( // plaintext folds the reporting secret into the existing mci, diverging from the // bytes the token is computed over, so it re-encodes. let shared_content = message.message_context_info.is_unset().then(|| { - pre_encoded.unwrap_or_else(|| std::sync::Arc::new(waproto::codec::message_to_vec(message))) + pre_encoded.map_or_else( + || std::borrow::Cow::Owned(waproto::codec::message_to_vec(message)), + std::borrow::Cow::Borrowed, + ) }); // sender is the author's own jid, remote is the chat jid (WAWebReportingTokenUtils: @@ -115,12 +135,12 @@ pub async fn prepare_dm_stanza( Some(content) => generate_reporting_token_from_encoded( message, content, - &request_id, + request_id, own_jid, - &to_jid, + to_jid, existing_secret, ), - None => generate_reporting_token(message, &request_id, own_jid, &to_jid, existing_secret), + None => generate_reporting_token(message, request_id, own_jid, to_jid, existing_secret), }; // The reporting token's MessageContextInfo (message_secret + version) is spliced @@ -164,7 +184,7 @@ pub async fn prepare_dm_stanza( let mut participant_nodes = Vec::with_capacity(total_devices); let mut includes_prekey_message = false; - let hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message); + let hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit, message); let mediatype = media_type_from_message(message); @@ -248,7 +268,7 @@ pub async fn prepare_dm_stanza( .attr("type", stanza_type); if let Some(edit_attr) = edit - && edit_attr != crate::types::message::EditAttribute::Empty + && *edit_attr != crate::types::message::EditAttribute::Empty { stanza_builder = stanza_builder.attr("edit", edit_attr.to_string_val()); } @@ -297,41 +317,142 @@ where session .commit() .await - .map_err(|e| anyhow!("restoring checked-out session after pre-flight: {e}"))?; + .context("restoring checked-out session after pairwise retry pre-flight")?; } Ok(needs_pkmsg) } +/// Structural destination for the canonical pairwise retry encoder. +#[derive(Debug)] +pub enum PairwiseRetryDestination { + Direct { + to: Jid, + recipient: Option<Jid>, + }, + Participant { + to: Jid, + participant: Jid, + addressing_mode: Option<crate::types::message::AddressingMode>, + }, +} + +/// Native inputs for one pairwise retransmission. Grouping them prevents +/// positional argument drift without allocating or introducing an intermediate +/// wire representation. +pub struct PairwiseRetryRequest<'a> { + pub destination: PairwiseRetryDestination, + pub encryption_jid: Jid, + pub message: &'a wa::Message, + pub message_id: String, + pub retry_count: u8, + pub account: Option<&'a wa::ADVSignedDeviceIdentity>, + pub edit: Option<crate::types::message::EditAttribute>, + /// Canonical, unpadded protobuf bytes for `message`, when the caller already + /// encoded it for persistence or another stanza. Reusing them avoids a + /// second tree walk and allocation before padding. + pub pre_encoded: Option<&'a [u8]>, +} + +#[inline] +fn is_pairwise_user(jid: &Jid) -> bool { + !jid.is_empty() + && matches!( + jid.server, + wacore_binary::Server::Pn + | wacore_binary::Server::Lid + | wacore_binary::Server::Hosted + | wacore_binary::Server::HostedLid + | wacore_binary::Server::Bot + ) +} + +fn validate_pairwise_retry_route( + destination: &PairwiseRetryDestination, + encryption_jid: &Jid, +) -> Result<()> { + if !is_pairwise_user(encryption_jid) { + bail!("pairwise retry encryption target must be a user device JID"); + } + + match destination { + PairwiseRetryDestination::Direct { to, recipient } => { + if !is_pairwise_user(to) { + bail!("direct retry destination must be a user JID"); + } + if recipient.as_ref().is_some_and(|jid| !is_pairwise_user(jid)) { + bail!("direct retry recipient must be a user JID"); + } + } + PairwiseRetryDestination::Participant { + to, + participant, + addressing_mode, + } => { + if !is_pairwise_user(participant) { + bail!("participant retry target must be a user device JID"); + } + if to.is_group() { + if addressing_mode.is_none() { + bail!("group retry requires an addressing mode"); + } + } else if to.is_broadcast_list() { + if addressing_mode.is_some() { + bail!("broadcast retry must not carry a group addressing mode"); + } + } else { + bail!("participant retry destination must be a group or broadcast list"); + } + } + } + Ok(()) +} + /// Mirrors `WAWebSendMsgCreateDeviceStanza.createUserDeviceMsgStanza`. -/// `<enc>` goes directly under `<message>`; the fanout wrapper -/// (`<participants><to>`) is server-rejected with 479 on retries. -/// `recipient_jid` is propagated verbatim from the retry receipt -/// (`f && (k.recipient = f)` in `WAWebHandleRetryRequest`); pass `None` -/// when the incoming receipt didn't carry it. -#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.dm_retry", level = "debug", skip_all, fields(to = %to_jid.observe()), err(Debug)))] -#[allow(clippy::too_many_arguments)] -pub async fn prepare_dm_retry_stanza<S, I>( +/// `<enc>` goes directly under `<message>`; the fanout wrapper is rejected for +/// retries. Routing stays typed and structural attributes remain core-owned. +#[cfg_attr( + feature = "tracing", + tracing::instrument(name = "wa.send.pairwise_retry", level = "debug", skip_all, err(Debug)) +)] +pub async fn prepare_pairwise_retry_stanza<S, I>( session_store: &mut S, identity_store: &mut I, - to_jid: Jid, - recipient_jid: Option<Jid>, - encryption_jid: Jid, - message: &wa::Message, - message_id: String, - retry_count: u8, - account: Option<&wa::ADVSignedDeviceIdentity>, - edit: Option<crate::types::message::EditAttribute>, + request: PairwiseRetryRequest<'_>, ) -> Result<Node> where S: crate::libsignal::protocol::SessionStore, I: crate::libsignal::protocol::IdentityKeyStore, { - let plaintext = MessageUtils::encode_and_pad(message); + let PairwiseRetryRequest { + destination, + encryption_jid, + message, + message_id, + retry_count, + account, + edit, + pre_encoded, + } = request; + if message_id.is_empty() { + bail!("retry message ID must not be empty"); + } + if !(1..crate::protocol::retry::MAX_RETRY_COUNT).contains(&retry_count) { + bail!( + "retry count {retry_count} must be in 1..{}", + crate::protocol::retry::MAX_RETRY_COUNT + ); + } + validate_pairwise_retry_route(&destination, &encryption_jid)?; + + let plaintext = match pre_encoded { + Some(content) => MessageUtils::pad_with_context_from_encoded(content, None), + None => MessageUtils::encode_and_pad(message), + }; let signal_address = encryption_jid.to_protocol_address(); if account.is_none() && pkmsg_would_be_emitted(session_store, &signal_address).await? { bail!( - "DM retry pkmsg requires <device-identity> (account is None); \ + "pairwise retry pkmsg requires <device-identity> (account is None); \ refusing before message_encrypt to avoid advancing the sender chain" ); } @@ -340,7 +461,7 @@ where message_encrypt(&plaintext, &signal_address, session_store, identity_store).await?; let (enc_type, is_prekey, serialized) = extract_ciphertext(encrypted) - .ok_or_else(|| anyhow!("Unexpected encryption message type for DM retry"))?; + .ok_or_else(|| anyhow!("Unexpected encryption message type for pairwise retry"))?; let hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message); let mut enc_builder = NodeBuilder::new("enc") @@ -366,13 +487,30 @@ where ); } - let mut stanza_builder = NodeBuilder::new("message") - .attr("to", to_jid) + let mut stanza_builder = NodeBuilder::new("message"); + match destination { + PairwiseRetryDestination::Direct { to, recipient } => { + stanza_builder = stanza_builder.attr("to", to); + if let Some(recipient) = recipient { + stanza_builder = stanza_builder.attr("recipient", recipient); + } + } + PairwiseRetryDestination::Participant { + to, + participant, + addressing_mode, + } => { + stanza_builder = stanza_builder + .attr("to", to) + .attr("participant", participant); + if let Some(addressing_mode) = addressing_mode { + stanza_builder = stanza_builder.attr("addressing_mode", addressing_mode.as_str()); + } + } + } + stanza_builder = stanza_builder .attr("id", message_id) .attr("type", stanza_type_from_message(message)); - if let Some(r) = recipient_jid { - stanza_builder = stanza_builder.attr("recipient", r); - } // Without `edit`, the resend looks like a normal message and the client never // applies the revoke/edit. diff --git a/wacore/src/send/encrypt.rs b/wacore/src/send/encrypt.rs index 9c41b7469..bbf61e9d4 100644 --- a/wacore/src/send/encrypt.rs +++ b/wacore/src/send/encrypt.rs @@ -89,6 +89,11 @@ pub struct EncryptResult { pub had_unregistered_device: bool, } +pub(crate) struct EncryptAttempt { + pub result: EncryptResult, + pub first_error: Option<anyhow::Error>, +} + /// One device's encrypted ciphertext, node-agnostic. The DM/peer paths map this /// into a `<to><enc>` node; the voip offer maps it into an `<enc>` per device. pub struct EncryptedDevice { @@ -108,6 +113,11 @@ pub struct EncryptForDevicesRaw { pub had_unregistered_device: bool, } +struct RawEncryptAttempt { + result: EncryptForDevicesRaw, + first_error: Option<anyhow::Error>, +} + /// Resolve the `<device-identity>` blob a stanza must carry. A pkmsg recipient /// validates our identity from it; without it a pkmsg advances the sender chain /// while the peer can't consume the pre-key message (the linked-device deadlock). @@ -246,7 +256,7 @@ async fn encrypt_one_device( session_store: &mut dyn crate::libsignal::protocol::SessionStore, identity_store: &mut dyn crate::libsignal::protocol::IdentityKeyStore, device_jid: Jid, -) -> (Jid, Result<Option<EncryptOneResult>, String>) { +) -> (Jid, Result<Option<EncryptOneResult>>) { match message_encrypt(plaintext, addr, session_store, identity_store).await { Ok(encrypted_payload) => { let Some((enc_type, is_prekey, serialized_bytes)) = @@ -264,7 +274,10 @@ async fn encrypt_one_device( })), ) } - Err(e) => (device_jid, Err(format!("{addr}: {e}"))), + Err(error) => ( + device_jid, + Err(anyhow::Error::new(error).context(format!("failed to encrypt for {addr}"))), + ), } } @@ -272,9 +285,10 @@ async fn encrypt_one_device( /// on success, a logged skip on failure. Node-agnostic so both the message /// `<to>` map and the voip offer share it. fn push_raw_result( - (device_jid, res): (Jid, Result<Option<EncryptOneResult>, String>), + (device_jid, res): (Jid, Result<Option<EncryptOneResult>>), devices: &mut Vec<EncryptedDevice>, includes_prekey_message: &mut bool, + first_error: &mut Option<anyhow::Error>, ) { match res { Ok(Some(one)) => { @@ -287,7 +301,12 @@ fn push_raw_result( }); } Ok(None) => {} - Err(msg) => log::warn!("Failed to encrypt for device: {msg}. Skipping."), + Err(error) => { + log::warn!("Failed to encrypt for device: {error:#}. Skipping."); + if first_error.is_none() { + *first_error = Some(error); + } + } } } @@ -360,6 +379,7 @@ pub async fn encrypt_for_devices( pub struct SessionPlan { encryption_overrides: Vec<Option<Jid>>, pub had_unregistered_device: bool, + first_error: Option<anyhow::Error>, } impl SessionPlan { @@ -371,6 +391,7 @@ impl SessionPlan { Self { encryption_overrides: vec![None; device_count], had_unregistered_device: false, + first_error: None, } } } @@ -395,6 +416,7 @@ pub async fn ensure_sessions_for_devices( // Indices into `devices` for those needing prekey fetch. let mut indices_needing_prekeys: Vec<usize> = Vec::with_capacity(devices.len()); let mut had_406 = false; + let mut first_error = None; let mut reusable_addr = crate::types::jid::make_reusable_protocol_address(); @@ -471,6 +493,10 @@ pub async fn ensure_sessions_for_devices( jids_for_fetch.len() ); had_406 = true; + // Best-effort callers still skip these devices, while a required + // distribution can surface the typed server failure without + // reconstructing it or reducing the source chain to a string. + first_error = Some(e); std::collections::HashMap::new() } Err(e) => return Err(e), @@ -537,11 +563,8 @@ pub async fn ensure_sessions_for_devices( // (resolver has no 'static handle into this spawned task). Ok(IdentityChange::ReplacedExisting) => Ok(Some(encryption_jid)), Ok(IdentityChange::NewOrUnchanged) => Ok(None), - Err(e) => Err(anyhow::anyhow!( - "Failed to process pre-key bundle for {}: {:?}", - addr, - e - )), + Err(error) => Err(anyhow::Error::new(error) + .context(format!("failed to process pre-key bundle for {addr}"))), } }) }; @@ -563,11 +586,17 @@ pub async fn ensure_sessions_for_devices( // fan-out below. Ok(Err(e)) => { log::warn!("Group session setup failed for a device, skipping it: {e}"); + if first_error.is_none() { + first_error = Some(e); + } } - Err(SpawnCanceled) => { + Err(error) => { log::warn!( "Session-establishment task did not deliver a result; skipping device." ); + if first_error.is_none() { + first_error = Some(anyhow::Error::new(error)); + } } } if next_spawn < total { @@ -580,6 +609,7 @@ pub async fn ensure_sessions_for_devices( Ok(SessionPlan { encryption_overrides, had_unregistered_device: had_406, + first_error, }) } @@ -589,7 +619,6 @@ pub async fn ensure_sessions_for_devices( /// under locks that must not span I/O. A device whose session is still /// missing (e.g. its bundle was absent) fails its encrypt and is skipped, /// matching the combined path's behavior. -#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.encrypt_fanout", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))] pub async fn encrypt_for_devices_with_sessions( runtime: &dyn Runtime, stores: &mut SignalStores<'_>, @@ -599,9 +628,40 @@ pub async fn encrypt_for_devices_with_sessions( mediatype: Option<&str>, plan: SessionPlan, ) -> Result<EncryptResult> { - let raw = - encrypt_for_devices_with_sessions_raw(runtime, stores, devices, plaintext_to_encrypt, plan) - .await?; + Ok(encrypt_for_devices_with_sessions_detailed( + runtime, + stores, + devices, + plaintext_to_encrypt, + hide_decrypt_fail, + mediatype, + plan, + ) + .await? + .result) +} + +#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.encrypt_fanout", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))] +pub(crate) async fn encrypt_for_devices_with_sessions_detailed( + runtime: &dyn Runtime, + stores: &mut SignalStores<'_>, + devices: &[Jid], + plaintext_to_encrypt: &[u8], + hide_decrypt_fail: bool, + mediatype: Option<&str>, + plan: SessionPlan, +) -> Result<EncryptAttempt> { + let RawEncryptAttempt { + result: raw, + first_error, + } = encrypt_for_devices_with_sessions_raw_detailed( + runtime, + stores, + devices, + plaintext_to_encrypt, + plan, + ) + .await?; // Map each ciphertext to the message path's `<to><enc>` node, preserving the // raw fan-out order so the wire output is identical to the pre-split path. @@ -617,11 +677,14 @@ pub async fn encrypt_for_devices_with_sessions( )); } - Ok(EncryptResult { - participant_nodes, - includes_prekey_message: raw.includes_prekey_message, - encrypted_devices, - had_unregistered_device: raw.had_unregistered_device, + Ok(EncryptAttempt { + result: EncryptResult { + participant_nodes, + includes_prekey_message: raw.includes_prekey_message, + encrypted_devices, + had_unregistered_device: raw.had_unregistered_device, + }, + first_error, }) } @@ -632,7 +695,6 @@ pub async fn encrypt_for_devices_with_sessions( /// run under locks that must not span I/O. A device whose session is still /// missing (e.g. its bundle was absent) fails its encrypt and is skipped. /// Same parallel fan-out + skip-on-fail contract as the message path. -#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.encrypt_fanout", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))] pub async fn encrypt_for_devices_with_sessions_raw( runtime: &dyn Runtime, stores: &mut SignalStores<'_>, @@ -640,6 +702,25 @@ pub async fn encrypt_for_devices_with_sessions_raw( plaintext_to_encrypt: &[u8], plan: SessionPlan, ) -> Result<EncryptForDevicesRaw> { + Ok(encrypt_for_devices_with_sessions_raw_detailed( + runtime, + stores, + devices, + plaintext_to_encrypt, + plan, + ) + .await? + .result) +} + +#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.encrypt_fanout_raw", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))] +async fn encrypt_for_devices_with_sessions_raw_detailed( + runtime: &dyn Runtime, + stores: &mut SignalStores<'_>, + devices: &[Jid], + plaintext_to_encrypt: &[u8], + plan: SessionPlan, +) -> Result<RawEncryptAttempt> { debug_assert_eq!( plan.encryption_overrides.len(), devices.len(), @@ -648,6 +729,7 @@ pub async fn encrypt_for_devices_with_sessions_raw( let SessionPlan { encryption_overrides, had_unregistered_device, + mut first_error, } = plan; let mut encrypted = Vec::with_capacity(devices.len()); @@ -675,7 +757,12 @@ pub async fn encrypt_for_devices_with_sessions_raw( device_jid, ) .await; - push_raw_result(res, &mut encrypted, &mut includes_prekey_message); + push_raw_result( + res, + &mut encrypted, + &mut includes_prekey_message, + &mut first_error, + ); } else { // One task per chunk, not per device: the per-device fan-out allocated a // task + oneshot + two store clones for every recipient. Same parallelism, @@ -730,24 +817,35 @@ pub async fn encrypt_for_devices_with_sessions_raw( match spawn_result { Ok(results) => { for res in results { - push_raw_result(res, &mut encrypted, &mut includes_prekey_message); + push_raw_result( + res, + &mut encrypted, + &mut includes_prekey_message, + &mut first_error, + ); } } - Err(SpawnCanceled) => { + Err(error) => { // A whole chunk drops (not one device); its members stay // un-warm and are re-targeted next send. log::warn!( "Encrypt chunk did not deliver a result; up to ~{} device(s) skipped this send.", total.div_ceil(num_chunks) ); + if first_error.is_none() { + first_error = Some(anyhow::Error::new(error)); + } } } } } - Ok(EncryptForDevicesRaw { - devices: encrypted, - includes_prekey_message, - had_unregistered_device, + Ok(RawEncryptAttempt { + result: EncryptForDevicesRaw { + devices: encrypted, + includes_prekey_message, + had_unregistered_device, + }, + first_error, }) } diff --git a/wacore/src/send/group.rs b/wacore/src/send/group.rs index 403fa7730..947c23083 100644 --- a/wacore/src/send/group.rs +++ b/wacore/src/send/group.rs @@ -2,85 +2,14 @@ use super::*; -/// Pairwise-encrypted retry stanza for a single group participant. -/// WA Web sends retries to the failing device only (RetryMsgJob.js:71), -/// NOT as a sender-key broadcast to all participants. -#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.group_retry", level = "debug", skip_all, fields(group = %group_jid.observe()), err(Debug)))] -#[allow(clippy::too_many_arguments)] -pub async fn prepare_group_retry_stanza<S, I>( - session_store: &mut S, - identity_store: &mut I, - group_jid: Jid, - participant_jid: Jid, - encryption_jid: Jid, - message: &wa::Message, - message_id: String, - retry_count: u8, - account: Option<&wa::ADVSignedDeviceIdentity>, - addressing_mode: crate::types::message::AddressingMode, - edit: Option<crate::types::message::EditAttribute>, -) -> Result<Node> -where - S: crate::libsignal::protocol::SessionStore, - I: crate::libsignal::protocol::IdentityKeyStore, -{ - let plaintext = MessageUtils::encode_and_pad(message); - let signal_address = encryption_jid.to_protocol_address(); - - if account.is_none() && pkmsg_would_be_emitted(session_store, &signal_address).await? { - bail!( - "group retry pkmsg requires <device-identity> (account is None); \ - refusing before message_encrypt to avoid advancing the sender chain" - ); - } - - let encrypted = - message_encrypt(&plaintext, &signal_address, session_store, identity_store).await?; - - let (enc_type, is_prekey, serialized) = extract_ciphertext(encrypted) - .ok_or_else(|| anyhow!("Unexpected encryption message type for group retry"))?; - - // count="N" distinguishes retries from normal sends (MsgCreateDeviceStanza.js:150-153) - let mut enc_builder = NodeBuilder::new("enc") - .attr("v", stanza::ENC_VERSION) - .attr("type", enc_type) - .attr("count", retry_count); - if let Some(mt) = media_type_from_message(message) { - enc_builder = enc_builder.attr("mediatype", mt); - } - let enc_node = enc_builder.bytes(serialized).build(); - - let mut children = vec![enc_node]; - - // Defense in depth: pre-flight should have caught a no-account pkmsg, but a - // corrupt session that triggers a fresh pkmsg mid-call would slip past. - if let Some(device_identity_bytes) = needs_device_identity(is_prekey, account)? { - children.push( - NodeBuilder::new("device-identity") - .bytes(device_identity_bytes) - .build(), - ); - } - - let stanza_type = stanza_type_from_message(message); - let mut stanza_builder = NodeBuilder::new("message") - .attr("to", group_jid) - .attr("participant", participant_jid) - .attr("id", message_id) - .attr("type", stanza_type); - - // WA Web always sets addressing_mode for groups (MsgCreateDeviceStanza.js:131-135) - stanza_builder = stanza_builder.attr("addressing_mode", addressing_mode.as_str()); - - // Without `edit`, the resend looks like a normal message and the client never - // applies the revoke/edit. - if let Some(e) = edit - && e != crate::types::message::EditAttribute::Empty - { - stanza_builder = stanza_builder.attr("edit", e.to_string_val()); - } - - Ok(stanza_builder.children(children).build()) +/// Retain devices that may receive a sender-key distribution. Reuses the +/// caller's allocation and centralizes the exact-sender/hosted exclusions used +/// by both normal and targeted sends. +pub fn retain_skdm_distribution_targets(devices: &mut Vec<Jid>, own_sending_jid: &Jid) { + devices.retain(|device| { + !(device.user == own_sending_jid.user && device.device == own_sending_jid.device) + && !device.is_hosted() + }); } /// Result of `prepare_group_stanza` — carries the stanza node and the exact @@ -110,34 +39,85 @@ pub struct PreparedGroupStanza { pub sender_identity: Jid, } -#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.group_prepare", level = "debug", skip_all, fields(to = %to_jid.observe()), err(Debug)))] -#[allow(clippy::too_many_arguments)] +/// A required sender-key distribution that could not reach every target. +/// +/// The source chain preserves the concrete crypto or transport failure. When +/// the failure followed a 406 pre-key response, `stale_device_users` identifies +/// registry entries the caller should invalidate before the next resolution. +#[derive(Debug, thiserror::Error)] +#[error("required sender-key distribution failed")] +pub struct RequiredSenderKeyDistributionError { + #[source] + source: anyhow::Error, + stale_device_users: Vec<String>, +} + +impl RequiredSenderKeyDistributionError { + fn new(source: anyhow::Error, stale_device_users: Vec<String>) -> Self { + Self { + source, + stale_device_users, + } + } + + pub fn stale_device_users(&self) -> &[String] { + &self.stale_device_users + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum SenderKeyDistributionPolicy { + /// Preserve normal fanout semantics by skipping devices that cannot receive + /// the distribution in this send. + #[default] + BestEffort, + /// Abort unless every requested device receives its distribution. + Required, +} + +pub struct GroupStanzaRequest<'a> { + pub group: &'a GroupInfo, + pub own_jid: &'a Jid, + pub own_lid: &'a Jid, + pub account: Option<&'a wa::ADVSignedDeviceIdentity>, + pub to: &'a Jid, + pub message: &'a wa::Message, + pub message_id: &'a str, + pub force_distribution: bool, + pub distribution_targets: Option<Vec<Jid>>, + pub distribution_policy: SenderKeyDistributionPolicy, + pub phash_devices: Option<&'a super::ResolvedGroupDevices>, + pub edit: Option<&'a crate::types::message::EditAttribute>, + pub extra_nodes: &'a [Node], + pub pre_encoded: Option<&'a [u8]>, +} + +#[cfg_attr( + feature = "tracing", + tracing::instrument(name = "wa.send.group_prepare", level = "debug", skip_all, err(Debug)) +)] pub async fn prepare_group_stanza( runtime: &dyn Runtime, stores: &mut SignalStores<'_>, resolver: &dyn SendContextResolver, - // Caller guarantees `own_base_jid` is already present in `participants`, so - // this reads the shared (Arc-backed) metadata without cloning it. - group_info: &GroupInfo, - own_jid: &Jid, - own_lid: &Jid, - account: Option<&wa::ADVSignedDeviceIdentity>, - to_jid: Jid, - message: &wa::Message, - request_id: String, - force_skdm_distribution: bool, - skdm_target_devices: Option<Vec<Jid>>, - // Full resolved device set for the phash (groups only). `Some` on warm/partial - // sends so the phash covers every device + self even when no SKDM is sent; - // `None` on the cold `force_skdm` path (the set is resolved here) and for - // status broadcasts (which keep the prior phash behavior). - all_devices_for_phash: Option<std::sync::Arc<super::ResolvedGroupDevices>>, - edit: Option<crate::types::message::EditAttribute>, - extra_stanza_nodes: &[Node], - // Avoids a second full encode when the caller already serialized the message; - // ignored on the mci-hoist path (see `shared_content`). - pre_encoded: Option<std::sync::Arc<Vec<u8>>>, + request: GroupStanzaRequest<'_>, ) -> Result<PreparedGroupStanza> { + let GroupStanzaRequest { + group: group_info, + own_jid, + own_lid, + account, + to: to_jid, + message, + message_id: request_id, + force_distribution: force_skdm_distribution, + distribution_targets: skdm_target_devices, + distribution_policy, + phash_devices: all_devices_for_phash, + edit, + extra_nodes: extra_stanza_nodes, + pre_encoded, + } = request; let (own_sending_jid, _) = match group_info.addressing_mode { crate::types::message::AddressingMode::Lid => (own_lid.clone(), "lid"), crate::types::message::AddressingMode::Pn => (own_jid.clone(), "pn"), @@ -150,7 +130,10 @@ pub async fn prepare_group_stanza( // its plaintext folds the reporting secret into the existing mci, diverging from // the bytes the token is computed over, so it re-encodes. let shared_content = message.message_context_info.is_unset().then(|| { - pre_encoded.unwrap_or_else(|| std::sync::Arc::new(waproto::codec::message_to_vec(message))) + pre_encoded.map_or_else( + || std::borrow::Cow::Owned(waproto::codec::message_to_vec(message)), + std::borrow::Cow::Borrowed, + ) }); // Generate reporting token if the message type supports it. @@ -162,12 +145,12 @@ pub async fn prepare_group_stanza( Some(content) => generate_reporting_token_from_encoded( message, content, - &request_id, - &to_jid, - &to_jid, + request_id, + to_jid, + to_jid, existing_secret, ), - None => generate_reporting_token(message, &request_id, &to_jid, &to_jid, existing_secret), + None => generate_reporting_token(message, request_id, to_jid, to_jid, existing_secret), }; // The reporting token's MessageContextInfo (message_secret + version) is spliced @@ -279,12 +262,7 @@ pub async fn prepare_group_stanza( let own_user = &own_sending_jid.user; let own_device = own_sending_jid.device; let before_filter = resolved_list.len(); - resolved_list.retain(|device_jid| { - let is_exact_sender = device_jid.user == *own_user && device_jid.device == own_device; - let is_hosted = device_jid.is_hosted(); - // Exclude the exact sending device and hosted devices - !is_exact_sender && !is_hosted - }); + retain_skdm_distribution_targets(&mut resolved_list, &own_sending_jid); log::debug!( "Filtered SKDM devices from {} to {} (excluded sender {}:{} and hosted devices)", before_filter, @@ -303,20 +281,25 @@ pub async fn prepare_group_stanza( } else { None }; + if distribution_policy == SenderKeyDistributionPolicy::Required + && distribution_list.as_ref().is_none_or(Vec::is_empty) + { + bail!("required sender-key distribution has no targets"); + } // Phash (groups): cover the FULL participant device set + the sending device // on EVERY send, matching WA Web `phashV2([].concat(A, [B]))`. Verified // against a real WA Web capture: the recipient set plus the sending device // reproduced the on-wire phash exactly, the recipient set alone did not. The - // server validates it silently (it is not echoed on a normal ack). Status - // broadcasts keep the prior behavior (phash over the distribution list only, - // when distributing); WA Web's status path does not augment with self. + // server validates it silently (it is not echoed on a normal ack). The + // captured status sender does not attach a phash, including when it carries + // a targeted sender-key distribution. if to_jid.is_group() { // Warm/partial sends pass the complete set in `all_devices_for_phash`, // whose phash memo serves repeat sends with an inline copy; the cold // `force_skdm` path leaves it None and `distribution_list` already // holds the full resolved set. - if let Some(resolved) = all_devices_for_phash.as_deref() { + if let Some(resolved) = all_devices_for_phash { phash_for_stanza = resolved.phash(&own_sending_jid); } else if let Some(src) = distribution_list.as_deref() { let phash_set = build_group_phash_set(src, &own_sending_jid); @@ -331,16 +314,11 @@ pub async fn prepare_group_stanza( } } } - } else if let Some(ref distribution_list) = distribution_list { - match MessageUtils::participant_list_hash(distribution_list) { - Ok(phash) => phash_for_stanza = Some(CompactString::new(&phash)), - Err(e) => log::warn!("Failed to compute phash for {}: {:?}", to_jid.observe(), e), - } } let mut had_unregistered_devices = false; - let sender_key_name = make_sender_key_name(&to_jid, &own_sending_jid.to_protocol_address()); + let sender_key_name = make_sender_key_name(to_jid, &own_sending_jid.to_protocol_address()); // Hold the per-device session locks the DM path uses across BOTH the X3DH setup // and the SKDM fan-out below, so a concurrent DM or group send sharing a device @@ -370,6 +348,9 @@ pub async fn prepare_group_stanza( let _setup_guard = setup_lock.lock().await; match ensure_sessions_for_devices(runtime, stores, resolver, list).await { Ok(plan) => Some(plan), + Err(error) if distribution_policy == SenderKeyDistributionPolicy::Required => { + return Err(error.context("required sender-key session setup failed")); + } Err(e) => { log::warn!( "SKDM session setup failed for group {}, continuing without distribution: {e}", @@ -431,8 +412,8 @@ pub async fn prepare_group_stanza( // Must match the rule applied to the main skmsg payload below: if SKDM carries // `decrypt-fail="hide"` but the payload does not (e.g. AdminRevoke), recipients // without a sender key never decrypt the skmsg and the revoke is silently dropped. - let skdm_hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message); - match encrypt_for_devices_with_sessions( + let skdm_hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit, message); + match encrypt_for_devices_with_sessions_detailed( runtime, stores, distribution_list, @@ -443,27 +424,66 @@ pub async fn prepare_group_stanza( ) .await { - Ok(result) => { - includes_prekey_message = - includes_prekey_message || result.includes_prekey_message; - if result.had_unregistered_device { + Ok(EncryptAttempt { + result, + first_error, + }) => { + let EncryptResult { + participant_nodes, + includes_prekey_message: result_includes_prekey, + encrypted_devices, + had_unregistered_device, + } = result; + if distribution_policy == SenderKeyDistributionPolicy::Required + && (encrypted_devices.len() != distribution_list.len() + || first_error.is_some()) + { + let error = first_error.unwrap_or_else(|| { + anyhow!( + "sender-key distribution encrypted {} of {} required targets", + encrypted_devices.len(), + distribution_list.len() + ) + }); + let stale_device_users = if had_unregistered_device { + collect_stale_device_users( + Some(distribution_list), + &encrypted_devices, + group_info, + ) + } else { + Vec::new() + }; + return Err(RequiredSenderKeyDistributionError::new( + error, + stale_device_users, + ) + .into()); + } + + includes_prekey_message |= result_includes_prekey; + if had_unregistered_device { had_unregistered_devices = true; } - skdm_encrypted_devices = result.encrypted_devices; + skdm_encrypted_devices = encrypted_devices; - if !result.participant_nodes.is_empty() { + if !participant_nodes.is_empty() { message_children.push( NodeBuilder::new("participants") - .children(result.participant_nodes) + .children(participant_nodes) .build(), ); - // Lenient (matches the DM fan-out): a no-account pkmsg - // omits the node rather than failing the whole send. - if let Some(device_identity_bytes) = - needs_device_identity(includes_prekey_message, account) - .ok() - .flatten() - { + let device_identity = match distribution_policy { + SenderKeyDistributionPolicy::BestEffort => { + needs_device_identity(includes_prekey_message, account) + .ok() + .flatten() + } + SenderKeyDistributionPolicy::Required => { + needs_device_identity(includes_prekey_message, account)? + } + }; + if let Some(device_identity_bytes) = device_identity { message_children.push( NodeBuilder::new("device-identity") .bytes(device_identity_bytes) @@ -472,6 +492,9 @@ pub async fn prepare_group_stanza( } } } + Err(error) if distribution_policy == SenderKeyDistributionPolicy::Required => { + return Err(RequiredSenderKeyDistributionError::new(error, Vec::new()).into()); + } Err(e) => { log::warn!( "SKDM distribution failed for group {}, continuing without it: {e}", @@ -504,7 +527,7 @@ pub async fn prepare_group_stanza( let skmsg_ciphertext = skmsg.into_serialized(); let mediatype = media_type_from_message(message); - let hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message); + let hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit, message); let mut enc_builder = NodeBuilder::new("enc") .attr("v", stanza::ENC_VERSION) @@ -532,7 +555,7 @@ pub async fn prepare_group_stanza( stanza_builder.attr("addressing_mode", group_info.addressing_mode.as_str()); } - if let Some(edit_attr) = &edit + if let Some(edit_attr) = edit && *edit_attr != crate::types::message::EditAttribute::Empty { stanza_builder = stanza_builder.attr("edit", edit_attr.to_string_val()); diff --git a/wacore/src/send/status.rs b/wacore/src/send/status.rs index e74330ae7..58bcb5cc0 100644 --- a/wacore/src/send/status.rs +++ b/wacore/src/send/status.rs @@ -95,6 +95,19 @@ pub fn status_carries_privacy_meta(message: &wa::Message) -> bool { !is_revoke && !is_reaction } +/// Return the message ID targeted by a status revoke, if this is a structurally +/// complete revoke protocol message. +/// +/// The target ID belongs to the embedded protocol key. It is distinct from the +/// outer stanza ID generated for the revoke itself. +pub fn status_revoke_target_id(message: &wa::Message) -> Option<&str> { + let protocol_message = unwrap_message(message).protocol_message.as_option()?; + if protocol_message.r#type != Some(wa::message::protocol_message::Type::Revoke) { + return None; + } + protocol_message.key.as_option()?.id.as_deref() +} + /// Dedup a pre-resolved status recipient list by user, then anchor the sender's /// own LID. Errors when no recipient was resolvable (matches WA Web's /// `WAWebLidMigrationUtils.toUserLid` + `compactMap` dropping unresolvable diff --git a/wacore/src/send/tests.rs b/wacore/src/send/tests.rs index ae58c40c0..5c4750856 100644 --- a/wacore/src/send/tests.rs +++ b/wacore/src/send/tests.rs @@ -307,6 +307,52 @@ mod status_carries_privacy_meta { } } +mod status_revoke_target_id { + use super::*; + + #[test] + fn returns_embedded_target_for_revoke() { + let msg = wa::Message { + protocol_message: buffa::MessageField::some(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::Revoke), + key: buffa::MessageField::some(wa::MessageKey { + id: Some("target-id".into()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + + assert_eq!(status_revoke_target_id(&msg), Some("target-id")); + } + + #[test] + fn ignores_other_or_incomplete_protocol_messages() { + let non_revoke = wa::Message { + protocol_message: buffa::MessageField::some(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::EphemeralSetting), + key: buffa::MessageField::some(wa::MessageKey { + id: Some("not-a-revoke".into()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + let incomplete_revoke = wa::Message { + protocol_message: buffa::MessageField::some(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::Revoke), + ..Default::default() + }), + ..Default::default() + }; + + assert_eq!(status_revoke_target_id(&non_revoke), None); + assert_eq!(status_revoke_target_id(&incomplete_revoke), None); + } +} + #[test] fn build_member_label_message_sets_fields() { let msg = build_member_label_message("VIP".to_string(), 1_766_847_151); @@ -376,6 +422,7 @@ struct MockSendContextResolver { /// JIDs reported via `on_local_identity_change` (send-path detection). identity_changes: std::sync::Mutex<Vec<Jid>>, chain_lock_probe: Option<ChainLockProbe>, + prekey_error_code: Option<u16>, } impl MockSendContextResolver { @@ -386,6 +433,7 @@ impl MockSendContextResolver { phone_to_lid: HashMap::new(), identity_changes: std::sync::Mutex::new(Vec::new()), chain_lock_probe: None, + prekey_error_code: None, } } @@ -417,6 +465,11 @@ impl MockSendContextResolver { self.phone_to_lid.insert(phone.to_string(), lid.to_string()); self } + + fn with_prekey_error(mut self, code: u16) -> Self { + self.prekey_error_code = Some(code); + self + } } #[async_trait::async_trait] @@ -441,6 +494,14 @@ impl SendContextResolver for MockSendContextResolver { &self, jids: &[Jid], ) -> Result<HashMap<Jid, PreKeyBundle>> { + if let Some(code) = self.prekey_error_code { + return Err(anyhow::Error::new(crate::request::ServerErrorCode { + code, + text: "injected pre-key failure".to_string(), + error_type: None, + backoff: None, + })); + } if let Some(probe) = &self.chain_lock_probe { probe .fetch_calls @@ -1355,18 +1416,23 @@ mod group_retry { let group: Jid = "120363098765432100@g.us".parse().unwrap(); let p: Jid = jid.to_string().parse().unwrap(); let account = pkmsg_account_proto(); - let n = prepare_group_retry_stanza( + let n = prepare_pairwise_retry_stanza( &mut ss, &mut is, - group.clone(), - p.clone(), - p.clone(), - &wa::Message::default(), - "3EB0ABC".into(), - 1, - Some(&account), - AddressingMode::Pn, - None, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Participant { + to: group.clone(), + participant: p.clone(), + addressing_mode: Some(AddressingMode::Pn), + }, + encryption_jid: p.clone(), + message: &wa::Message::default(), + message_id: "3EB0ABC".into(), + retry_count: 1, + account: Some(&account), + edit: None, + pre_encoded: None, + }, ) .await .unwrap(); @@ -1419,18 +1485,23 @@ mod group_retry { .serialize() .expect("serialize before"); - let result = prepare_group_retry_stanza( + let result = prepare_pairwise_retry_stanza( &mut ss, &mut is, - group, - p.clone(), - p.clone(), - &wa::Message::default(), - "grp-retry-no-account".into(), - 1, - None, - AddressingMode::Pn, - None, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Participant { + to: group, + participant: p.clone(), + addressing_mode: Some(AddressingMode::Pn), + }, + encryption_jid: p.clone(), + message: &wa::Message::default(), + message_id: "grp-retry-no-account".into(), + retry_count: 1, + account: None, + edit: None, + pre_encoded: None, + }, ) .await; let err = result.expect_err("group retry pkmsg must reject missing account"); @@ -1465,17 +1536,22 @@ mod group_retry { let recipient: Jid = "100000000000456@lid".parse().unwrap(); let requester: Jid = jid.to_string().parse().unwrap(); let account = pkmsg_account_proto(); - let n = prepare_dm_retry_stanza( + let n = prepare_pairwise_retry_stanza( &mut ss, &mut is, - to.clone(), - Some(recipient.clone()), - requester, - &wa::Message::default(), - "dm-retry-format-1".into(), - 1, - Some(&account), - None, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Direct { + to: to.clone(), + recipient: Some(recipient.clone()), + }, + encryption_jid: requester, + message: &wa::Message::default(), + message_id: "dm-retry-format-1".into(), + retry_count: 1, + account: Some(&account), + edit: None, + pre_encoded: None, + }, ) .await .unwrap(); @@ -1511,17 +1587,22 @@ mod group_retry { let encryption = jid.clone(); let account = pkmsg_account_proto(); - let n = prepare_dm_retry_stanza( + let n = prepare_pairwise_retry_stanza( &mut ss, &mut is, - to.clone(), - Some(to.clone()), - encryption, - &wa::Message::default(), - "dm-retry-1".into(), - 1, - Some(&account), - None, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Direct { + to: to.clone(), + recipient: Some(to.clone()), + }, + encryption_jid: encryption, + message: &wa::Message::default(), + message_id: "dm-retry-1".into(), + retry_count: 1, + account: Some(&account), + edit: None, + pre_encoded: None, + }, ) .await .unwrap(); @@ -1568,17 +1649,22 @@ mod group_retry { ..Default::default() }; - let n = prepare_dm_retry_stanza( + let n = prepare_pairwise_retry_stanza( &mut ss, &mut is, - to.clone(), - Some(to), - jid, - &wa::Message::default(), - "dm-retry-2".into(), - 2, - Some(&acc), - None, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Direct { + to: to.clone(), + recipient: Some(to), + }, + encryption_jid: jid, + message: &wa::Message::default(), + message_id: "dm-retry-2".into(), + retry_count: 2, + account: Some(&acc), + edit: None, + pre_encoded: None, + }, ) .await .unwrap(); @@ -1601,18 +1687,23 @@ mod group_retry { details: Some(b"t".to_vec()), ..Default::default() }; - let n = prepare_group_retry_stanza( + let n = prepare_pairwise_retry_stanza( &mut ss, &mut is, - group, - p.clone(), - p, - &wa::Message::default(), - "id2".into(), - 2, - Some(&acc), - AddressingMode::Pn, - None, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Participant { + to: group, + participant: p.clone(), + addressing_mode: Some(AddressingMode::Pn), + }, + encryption_jid: p, + message: &wa::Message::default(), + message_id: "id2".into(), + retry_count: 2, + account: Some(&acc), + edit: None, + pre_encoded: None, + }, ) .await .unwrap(); @@ -1641,18 +1732,23 @@ mod group_retry { let group: Jid = "120363098765432100@g.us".parse().unwrap(); let p: Jid = jid.to_string().parse().unwrap(); // Fresh session → pkmsg (pre-key), with LID addressing - let n = prepare_group_retry_stanza( + let n = prepare_pairwise_retry_stanza( &mut ss, &mut is, - group, - p.clone(), - p, - &wa::Message::default(), - "m2".into(), - 3, - Some(&wa::ADVSignedDeviceIdentity::default()), - AddressingMode::Lid, - None, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Participant { + to: group, + participant: p.clone(), + addressing_mode: Some(AddressingMode::Lid), + }, + encryption_jid: p, + message: &wa::Message::default(), + message_id: "m2".into(), + retry_count: 3, + account: Some(&wa::ADVSignedDeviceIdentity::default()), + edit: None, + pre_encoded: None, + }, ) .await .unwrap(); @@ -1673,18 +1769,23 @@ mod group_retry { let group: Jid = "120363098765432100@g.us".parse().unwrap(); let p: Jid = jid.to_string().parse().unwrap(); let account = pkmsg_account_proto(); - let n = prepare_group_retry_stanza( + let n = prepare_pairwise_retry_stanza( &mut ss, &mut is, - group, - p.clone(), - p, - &wa::Message::default(), - "revoke-1".into(), - 1, - Some(&account), - AddressingMode::Lid, - Some(crate::types::message::EditAttribute::AdminRevoke), + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Participant { + to: group, + participant: p.clone(), + addressing_mode: Some(AddressingMode::Lid), + }, + encryption_jid: p, + message: &wa::Message::default(), + message_id: "revoke-1".into(), + retry_count: 1, + account: Some(&account), + edit: Some(crate::types::message::EditAttribute::AdminRevoke), + pre_encoded: None, + }, ) .await .unwrap(); @@ -1696,21 +1797,230 @@ mod group_retry { let (mut ss, mut is, jid) = setup_session().await; let to: Jid = "559922223333@s.whatsapp.net".parse().unwrap(); let account = pkmsg_account_proto(); - let n = prepare_dm_retry_stanza( + let n = prepare_pairwise_retry_stanza( &mut ss, &mut is, - to.clone(), - Some(to), - jid, - &wa::Message::default(), - "edit-1".into(), - 1, - Some(&account), - Some(crate::types::message::EditAttribute::MessageEdit), + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Direct { + to: to.clone(), + recipient: Some(to), + }, + encryption_jid: jid, + message: &wa::Message::default(), + message_id: "edit-1".into(), + retry_count: 1, + account: Some(&account), + edit: Some(crate::types::message::EditAttribute::MessageEdit), + pre_encoded: None, + }, ) .await .unwrap(); assert_eq!(n.attrs().optional_string("edit").unwrap().as_ref(), "1"); + assert_eq!( + n.get_optional_child("enc") + .unwrap() + .attrs() + .optional_string("decrypt-fail") + .unwrap() + .as_ref(), + "hide" + ); + } + + #[tokio::test] + async fn broadcast_retry_preserves_target_and_omits_group_addressing() { + let (mut ss, mut is, jid) = setup_session().await; + let broadcast: Jid = "1234567890@broadcast".parse().unwrap(); + let participant = jid.clone(); + let account = pkmsg_account_proto(); + let node = prepare_pairwise_retry_stanza( + &mut ss, + &mut is, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Participant { + to: broadcast.clone(), + participant: participant.clone(), + addressing_mode: None, + }, + encryption_jid: jid, + message: &wa::Message::default(), + message_id: "broadcast-retry-1".into(), + retry_count: 2, + account: Some(&account), + edit: None, + pre_encoded: None, + }, + ) + .await + .unwrap(); + + let mut attrs = node.attrs(); + assert_eq!( + attrs.optional_string("to").unwrap().as_ref(), + broadcast.to_string() + ); + assert_eq!( + attrs.optional_string("participant").unwrap().as_ref(), + participant.to_string() + ); + assert!(attrs.optional_string("recipient").is_none()); + assert!(attrs.optional_string("addressing_mode").is_none()); + assert_eq!( + node.get_optional_child("enc") + .unwrap() + .attrs() + .optional_string("count") + .unwrap() + .as_ref(), + "2" + ); + } + + #[tokio::test] + async fn invalid_retry_identity_is_rejected_before_ratchet_advance() { + let cases = [ + ("", 1, "message ID"), + ("retry-count-zero", 0, "retry count"), + ( + "retry-count-max", + crate::protocol::retry::MAX_RETRY_COUNT, + "retry count", + ), + ]; + + for (message_id, retry_count, expected_error) in cases { + let (mut sessions, mut identities, jid) = setup_session().await; + let address = jid.to_protocol_address(); + let before = sessions + .load_session(&address) + .await + .unwrap() + .unwrap() + .serialize() + .unwrap(); + let result = prepare_pairwise_retry_stanza( + &mut sessions, + &mut identities, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Direct { + to: jid.clone(), + recipient: None, + }, + encryption_jid: jid, + message: &wa::Message::default(), + message_id: message_id.into(), + retry_count, + account: Some(&pkmsg_account_proto()), + edit: None, + pre_encoded: None, + }, + ) + .await; + let error = result.expect_err("invalid retry must be rejected"); + assert!( + error.to_string().contains(expected_error), + "unexpected error for {message_id:?}/{retry_count}: {error:#}" + ); + let after = sessions + .load_session(&address) + .await + .unwrap() + .unwrap() + .serialize() + .unwrap(); + assert_eq!( + before, after, + "validation must run before the Signal ratchet for {message_id:?}/{retry_count}" + ); + } + + enum InvalidRoute { + DirectGroup, + GroupWithoutAddressingMode, + BroadcastWithAddressingMode, + ParticipantOnDirectChat, + } + + for (case, expected_error) in [ + (InvalidRoute::DirectGroup, "direct retry destination"), + ( + InvalidRoute::GroupWithoutAddressingMode, + "group retry requires an addressing mode", + ), + ( + InvalidRoute::BroadcastWithAddressingMode, + "broadcast retry must not carry", + ), + ( + InvalidRoute::ParticipantOnDirectChat, + "participant retry destination", + ), + ] { + let (mut sessions, mut identities, encryption_jid) = setup_session().await; + let address = encryption_jid.to_protocol_address(); + let before = sessions + .load_session(&address) + .await + .unwrap() + .unwrap() + .serialize() + .unwrap(); + let group: Jid = "120363098765432100@g.us".parse().unwrap(); + let broadcast: Jid = "1234567890@broadcast".parse().unwrap(); + let destination = match case { + InvalidRoute::DirectGroup => PairwiseRetryDestination::Direct { + to: group, + recipient: None, + }, + InvalidRoute::GroupWithoutAddressingMode => PairwiseRetryDestination::Participant { + to: group, + participant: encryption_jid.clone(), + addressing_mode: None, + }, + InvalidRoute::BroadcastWithAddressingMode => { + PairwiseRetryDestination::Participant { + to: broadcast, + participant: encryption_jid.clone(), + addressing_mode: Some(AddressingMode::Pn), + } + } + InvalidRoute::ParticipantOnDirectChat => PairwiseRetryDestination::Participant { + to: encryption_jid.clone(), + participant: encryption_jid.clone(), + addressing_mode: None, + }, + }; + + let result = prepare_pairwise_retry_stanza( + &mut sessions, + &mut identities, + PairwiseRetryRequest { + destination, + encryption_jid, + message: &wa::Message::default(), + message_id: "invalid-route".into(), + retry_count: 1, + account: Some(&pkmsg_account_proto()), + edit: None, + pre_encoded: None, + }, + ) + .await; + let error = result.expect_err("invalid route must be rejected"); + assert!( + error.to_string().contains(expected_error), + "unexpected invalid-route error: {error:#}" + ); + let after = sessions + .load_session(&address) + .await + .unwrap() + .unwrap() + .serialize() + .unwrap(); + assert_eq!(before, after, "route validation must precede the ratchet"); + } } #[tokio::test] @@ -1719,18 +2029,25 @@ mod group_retry { let group: Jid = "120363098765432100@g.us".parse().unwrap(); let p: Jid = jid.to_string().parse().unwrap(); let account = pkmsg_account_proto(); - let n = prepare_group_retry_stanza( + let message = wa::Message::default(); + let encoded = waproto::codec::message_to_vec(&message); + let n = prepare_pairwise_retry_stanza( &mut ss, &mut is, - group, - p.clone(), - p, - &wa::Message::default(), - "plain-1".into(), - 1, - Some(&account), - AddressingMode::Lid, - None, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Participant { + to: group, + participant: p.clone(), + addressing_mode: Some(AddressingMode::Lid), + }, + encryption_jid: p, + message: &message, + message_id: "plain-1".into(), + retry_count: 1, + account: Some(&account), + edit: None, + pre_encoded: Some(&encoded), + }, ) .await .unwrap(); @@ -1961,17 +2278,22 @@ mod group_retry { .expect("serialize before"); let to: Jid = "559922223333@s.whatsapp.net".parse().unwrap(); - let result = prepare_dm_retry_stanza( + let result = prepare_pairwise_retry_stanza( &mut ss, &mut is, - to.clone(), - Some(to), - jid.clone(), - &wa::Message::default(), - "dm-retry-no-account".into(), - 1, - None, - None, + PairwiseRetryRequest { + destination: PairwiseRetryDestination::Direct { + to: to.clone(), + recipient: Some(to), + }, + encryption_jid: jid.clone(), + message: &wa::Message::default(), + message_id: "dm-retry-no-account".into(), + retry_count: 1, + account: None, + edit: None, + pre_encoded: None, + }, ) .await; let err = result.expect_err("DM retry pkmsg path must reject missing account"); @@ -3014,6 +3336,180 @@ mod mark_full_distribution_list { (ss, is) } + #[tokio::test] + async fn targeted_status_retry_sends_only_the_requesting_device() { + let status = Jid::status_broadcast(); + let own_pn: Jid = "12025550120:7@s.whatsapp.net".parse().unwrap(); + let own_lid: Jid = "100000000000000:7@lid".parse().unwrap(); + let requester: Jid = "100000000000001:11@lid".parse().unwrap(); + let (mut sessions, mut identities) = established_stores(&requester).await; + let mut sender_keys = MemSenderKeyStore::default(); + let mut prekeys = UnusedPreKeyStore; + let signed_prekeys = UnusedSignedPreKeyStore; + let mut stores = SignalStores { + sender_key_store: &mut sender_keys, + session_store: &mut sessions, + identity_store: &mut identities, + prekey_store: &mut prekeys, + signed_prekey_store: &signed_prekeys, + }; + let group = GroupInfo::new(Vec::new(), AddressingMode::Lid); + let message = wa::Message { + conversation: Some("status retry".into()), + ..Default::default() + }; + let account = wa::ADVSignedDeviceIdentity::default(); + let extension = NodeBuilder::new("custom-extension") + .attr("version", "1") + .build(); + + let prepared = prepare_group_stanza( + &TokioTestRuntime, + &mut stores, + &MockSendContextResolver::new(), + GroupStanzaRequest { + group: &group, + own_jid: &own_pn, + own_lid: &own_lid, + account: Some(&account), + to: &status, + message: &message, + message_id: "STATUS-RETRY-1", + force_distribution: false, + distribution_targets: Some(vec![requester.clone()]), + distribution_policy: SenderKeyDistributionPolicy::Required, + phash_devices: None, + edit: None, + extra_nodes: std::slice::from_ref(&extension), + pre_encoded: None, + }, + ) + .await + .unwrap(); + + let mut attrs = prepared.node.attrs(); + assert_eq!( + attrs.optional_string("to").unwrap().as_ref(), + "status@broadcast" + ); + assert_eq!( + attrs.optional_string("id").unwrap().as_ref(), + "STATUS-RETRY-1" + ); + assert!(attrs.optional_string("participant").is_none()); + assert!(attrs.optional_string("recipient").is_none()); + assert!(attrs.optional_string("addressing_mode").is_none()); + assert!(attrs.optional_string("phash").is_none()); + assert_eq!( + prepared + .node + .get_optional_child("custom-extension") + .unwrap() + .attrs() + .optional_string("version") + .unwrap() + .as_ref(), + "1" + ); + + let skmsg = prepared.node.get_optional_child("enc").unwrap(); + let mut skmsg_attrs = skmsg.attrs(); + assert_eq!( + skmsg_attrs.optional_string("type").unwrap().as_ref(), + stanza::ENC_TYPE_SKMSG + ); + assert!(skmsg_attrs.optional_string("count").is_none()); + + let participants = prepared.node.get_optional_child("participants").unwrap(); + let targets = participants.children().unwrap(); + assert_eq!(targets.len(), 1, "status retry must not fan out"); + assert_eq!( + targets[0].attrs().optional_string("jid").unwrap().as_ref(), + requester.to_string() + ); + assert!( + targets[0] + .get_optional_child("enc") + .unwrap() + .attrs() + .optional_string("count") + .is_none(), + "captured status SKDM encryption has no retry count" + ); + assert_eq!(prepared.skdm_devices, [requester]); + } + + #[tokio::test] + async fn required_targeted_distribution_reports_an_unregistered_target() { + let status = Jid::status_broadcast(); + let own_pn: Jid = "12025550121:7@s.whatsapp.net".parse().unwrap(); + let own_lid: Jid = "100000000000002:7@lid".parse().unwrap(); + let requester: Jid = "100000000000003:11@lid".parse().unwrap(); + let mut sessions = MemSessionStore::default(); + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + let mut identities = MemIdentityStore { + pair: IdentityKeyPair::generate(&mut rng), + reg_id: 7, + known: Default::default(), + }; + let mut sender_keys = MemSenderKeyStore::default(); + let mut prekeys = UnusedPreKeyStore; + let signed_prekeys = UnusedSignedPreKeyStore; + let mut stores = SignalStores { + sender_key_store: &mut sender_keys, + session_store: &mut sessions, + identity_store: &mut identities, + prekey_store: &mut prekeys, + signed_prekey_store: &signed_prekeys, + }; + let group = GroupInfo::new(Vec::new(), AddressingMode::Lid); + let message = wa::Message { + conversation: Some("status retry".into()), + ..Default::default() + }; + + let result = prepare_group_stanza( + &TokioTestRuntime, + &mut stores, + &MockSendContextResolver::new().with_prekey_error(406), + GroupStanzaRequest { + group: &group, + own_jid: &own_pn, + own_lid: &own_lid, + account: Some(&wa::ADVSignedDeviceIdentity::default()), + to: &status, + message: &message, + message_id: "STATUS-RETRY-MISSING-SESSION", + force_distribution: false, + distribution_targets: Some(vec![requester.clone()]), + distribution_policy: SenderKeyDistributionPolicy::Required, + phash_devices: None, + edit: None, + extra_nodes: &[], + pre_encoded: None, + }, + ) + .await; + let error = match result { + Err(error) => error, + Ok(_) => panic!("a targeted retry must not send without its SKDM"), + }; + + assert!( + format!("{error:#}").contains("required sender-key distribution failed"), + "unexpected error chain: {error:#}" + ); + let failure = error + .downcast_ref::<RequiredSenderKeyDistributionError>() + .expect("required failures must retain typed stale-target metadata"); + assert_eq!(failure.stale_device_users(), [requester.user.as_str()]); + assert_eq!( + crate::request::ServerErrorCode::from_anyhow(&error).map(|server| server.code), + Some(406), + "the typed failure must preserve the original server error chain" + ); + } + #[tokio::test] async fn failed_device_is_still_marked_has_key() { let group: Jid = "120363000000000001@g.us".parse().unwrap(); @@ -3054,19 +3550,22 @@ mod mark_full_distribution_list { &rt, &mut stores, &resolver, - &group_info, - &own_jid, - &own_lid, - None, - group, - &msg, - "TESTREQID".into(), - false, - Some(vec![a.clone(), b.clone()]), - None, - None, - &[], - None, + GroupStanzaRequest { + group: &group_info, + own_jid: &own_jid, + own_lid: &own_lid, + account: None, + to: &group, + message: &msg, + message_id: "TESTREQID", + force_distribution: false, + distribution_targets: Some(vec![a.clone(), b.clone()]), + distribution_policy: SenderKeyDistributionPolicy::BestEffort, + phash_devices: None, + edit: None, + extra_nodes: &[], + pre_encoded: None, + }, ) .await .expect("prepare_group_stanza should succeed even when a device fails to encrypt"); @@ -3139,19 +3638,22 @@ mod mark_full_distribution_list { &rt, &mut stores, &resolver, - group_info, - own_jid, - own_lid, - None, - group.clone(), - msg, - req.into(), - false, - Some(vec![a.clone()]), - None, - None, - &[], - None, + GroupStanzaRequest { + group: group_info, + own_jid, + own_lid, + account: None, + to: group, + message: msg, + message_id: req, + force_distribution: false, + distribution_targets: Some(vec![a.clone()]), + distribution_policy: SenderKeyDistributionPolicy::BestEffort, + phash_devices: None, + edit: None, + extra_nodes: &[], + pre_encoded: None, + }, ) .await .expect("prepare_group_stanza should succeed"); @@ -3263,19 +3765,22 @@ mod mark_full_distribution_list { &rt, &mut stores, &resolver, - &group_info, - &own_jid, - &own_lid, - None, - group, - &msg, - "TESTREQID2".into(), - false, - Some(vec![b.clone()]), - None, - None, - &[], - None, + GroupStanzaRequest { + group: &group_info, + own_jid: &own_jid, + own_lid: &own_lid, + account: None, + to: &group, + message: &msg, + message_id: "TESTREQID2", + force_distribution: false, + distribution_targets: Some(vec![b.clone()]), + distribution_policy: SenderKeyDistributionPolicy::BestEffort, + phash_devices: None, + edit: None, + extra_nodes: &[], + pre_encoded: None, + }, ) .await .expect("prepare_group_stanza should succeed"); @@ -3358,19 +3863,22 @@ mod mark_full_distribution_list { &rt, &mut stores, &resolver, - &group_info, - &own_jid, - &own_lid, - None, - group, - &msg, - "TESTREQID_ISO".into(), - false, - Some(vec![good.clone(), bad.clone()]), - None, - None, - &[], - None, + GroupStanzaRequest { + group: &group_info, + own_jid: &own_jid, + own_lid: &own_lid, + account: None, + to: &group, + message: &msg, + message_id: "TESTREQID_ISO", + force_distribution: false, + distribution_targets: Some(vec![good.clone(), bad.clone()]), + distribution_policy: SenderKeyDistributionPolicy::BestEffort, + phash_devices: None, + edit: None, + extra_nodes: &[], + pre_encoded: None, + }, ) .await .expect("prepare_group_stanza must succeed despite one device's setup failure"); diff --git a/waproto/src/lib.rs b/waproto/src/lib.rs index ebcf4b32d..dad8d95c3 100644 --- a/waproto/src/lib.rs +++ b/waproto/src/lib.rs @@ -364,6 +364,13 @@ pub mod codec { whatsapp::ADVSignedKeyIndexList::decode_from_slice(bytes) } + #[inline(never)] + pub fn adv_signed_key_index_list_to_vec( + key_index: &whatsapp::ADVSignedKeyIndexList, + ) -> Vec<u8> { + key_index.encode_to_vec() + } + #[inline(never)] pub fn adv_key_index_list_decode( bytes: &[u8], @@ -371,6 +378,11 @@ pub mod codec { whatsapp::ADVKeyIndexList::decode_from_slice(bytes) } + #[inline(never)] + pub fn adv_key_index_list_to_vec(key_index: &whatsapp::ADVKeyIndexList) -> Vec<u8> { + key_index.encode_to_vec() + } + #[inline(never)] pub fn client_pairing_props_decode( bytes: &[u8],