diff --git a/src/client.rs b/src/client.rs index 143dbc85a..be69ae9ee 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,6 +1,13 @@ +mod accessors; +mod adapters; +mod app_state; mod context_impl; mod device_registry; +mod iq_ops; mod lid_pn; +mod lifecycle; +mod messaging; +mod node_io; pub(crate) mod offline_resume; mod sender_keys; mod sessions; @@ -555,6563 +562,249 @@ pub struct Client { raw_node_forwarding: AtomicBool, } -impl Client { - pub fn shutdown_signal(&self) -> wacore::runtime::ShutdownSignal { - self.shutdown_notifier.subscribe() - } - - /// Synchronous flag-only equivalent of the first lines of `disconnect()`. - /// Spawned tasks watching `is_shutting_down()` / `shutdown_notifier` exit - /// on their next poll. Does NOT flush, close the transport, or touch - /// persistence — prefer `disconnect()` whenever you can `await`. Exists - /// for `Drop` impls on FFI wrappers (e.g. `WasmWhatsAppClient`) that - /// can't run async cleanup synchronously. - pub fn signal_shutdown_sync(&self) { - self.expected_disconnect.store(true, Ordering::Relaxed); - self.is_running.store(false, Ordering::Relaxed); - self.shutdown_notifier.notify(); - self.notify_connection_shutdown(); - } - - pub(crate) fn connection_shutdown_signal(&self) -> wacore::runtime::ShutdownSignal { - self.connection_shutdown - .lock() - .unwrap_or_else(|p| p.into_inner()) - .subscribe() +/// Builds a pong response node for a server-initiated ping. +/// +/// Matches WhatsApp Web (`WAWebCommsHandleStanza`): only includes `id` +/// when the server ping carried one. +fn build_pong(to: String, id: Option<&str>) -> wacore_binary::Node { + let mut builder = NodeBuilder::new("iq").attr("to", to).attr("type", "result"); + if let Some(id) = id { + builder = builder.attr("id", id); } + builder.build() +} - /// Fire the per-connection shutdown. Per-connection subscribers exit; - /// the terminal shutdown_notifier is untouched so reconnects still work. - pub(crate) fn notify_connection_shutdown(&self) { - self.connection_shutdown - .lock() - .unwrap_or_else(|p| p.into_inner()) - .notify(); - } +/// Build an `` for the given stanza, matching WA Web / whatsmeow behavior: +/// +/// - `class` = original stanza tag +/// - `id`, `to` (flipped from `from`), `participant` copied from original +/// - `from` = own device PN, only for message acks +/// - `type` echoed for non-message stanzas (whatsmeow: `node.Tag != "message"`), +/// except `notification type="encrypt"` with `` child (WA Web drops type there). +/// +/// For receipt acks, WA Web uses `MAYBE_CUSTOM_STRING(ackString)` where +/// `ackString = maybeAttrString("type")` — so `type` is only included when +/// explicitly present on the incoming receipt (delivery receipts normally +/// have no type attribute, meaning the ack also has no type). +/// Encode an ack stanza directly to bytes, bypassing Node + marshal_auto. +/// Acks are the most frequent outbound stanza (~1 per inbound message). +fn encode_ack_bytes( + node: &wacore_binary::NodeRef<'_>, + own_device_pn: Option<&Jid>, +) -> Result>, wacore_binary::error::BinaryError> { + use wacore_binary::encoder::{ByteWriter, EncodeNode, Encoder}; - /// Reset the per-connection notifier. Call at the start of each new - /// connection so subscribers registered afterwards see a fresh signal. - /// The previous notifier's subscribers have already been woken (either - /// by notify on disconnect, or by falling out of scope). - pub(crate) fn reset_connection_shutdown(&self) { - *self - .connection_shutdown - .lock() - .unwrap_or_else(|p| p.into_inner()) = wacore::runtime::ShutdownNotifier::new(); - } + let Some(id_val) = node.get_attr("id") else { + return Ok(None); + }; + let Some(from_val) = node.get_attr("from") else { + return Ok(None); + }; + // WAWebReceiptAck: `participant: r && r !== e ? DEVICE_JID(r) : DROP_ATTR`. + // Drop the attribute when it would duplicate `to` (which is the flipped `from`). + let participant_val = node.get_attr("participant").filter(|p| { + let p_str = p.as_str(); + let from_str = from_val.as_str(); + p_str.as_ref() != from_str.as_ref() + }); + // Server expects `recipient` echoed back so it can route the ack to the + // origin companion/device (hosted-companion, peer, LID-routed stanzas). + // Dropping it makes the server close the stream with ``. + let recipient_val = node.get_attr("recipient"); + let tag = node.tag.as_ref(); - /// Read the current semaphore generation and Arc atomically under the mutex. - pub(crate) fn read_message_semaphore(&self) -> (u64, Arc) { - let guard = match self.message_processing_semaphore.lock() { - Ok(g) => g, - Err(poisoned) => poisoned.into_inner(), - }; - ( - self.message_semaphore_generation.load(Ordering::SeqCst), - guard.clone(), - ) - } + let typ_val = if tag != "message" && !is_encrypt_identity_notification(node) { + node.get_attr("type") + } else { + None + }; - /// Replace the message processing semaphore and bump the generation counter. - /// - /// Both operations happen under the same mutex hold so readers always see - /// a consistent (generation, Arc) pair. Must be called from a non-async - /// context or inside a scoped block (MutexGuard is !Send). - pub(crate) fn swap_message_semaphore(&self, permits: usize) { - let mut guard = match self.message_processing_semaphore.lock() { - Ok(g) => g, - Err(poisoned) => poisoned.into_inner(), - }; - *guard = Arc::new(async_lock::Semaphore::new(permits)); - self.message_semaphore_generation - .fetch_add(1, Ordering::SeqCst); - } + let include_from = tag == "message" && own_device_pn.is_some(); - fn should_downgrade_sync_error(&self, err: &anyhow::Error) -> bool { - if self.is_shutting_down() { - return true; - } + // Count attrs: class + id + to + optional(from, participant, recipient, type) + let attr_count = 3 + + usize::from(include_from) + + usize::from(participant_val.is_some()) + + usize::from(recipient_val.is_some()) + + usize::from(typ_val.is_some()); - matches!( - err.downcast_ref::(), - Some( - crate::request::IqError::NotConnected - | crate::request::IqError::InternalChannelClosed - ) - ) + struct AckNode<'a> { + id: &'a wacore_binary::node::ValueRef<'a>, + from: &'a wacore_binary::node::ValueRef<'a>, + participant: Option<&'a wacore_binary::node::ValueRef<'a>>, + recipient: Option<&'a wacore_binary::node::ValueRef<'a>>, + typ: Option<&'a wacore_binary::node::ValueRef<'a>>, + own_pn: Option<&'a Jid>, + tag_str: &'a str, + attr_count: usize, } - /// Log a sync error, downgrading to debug level during shutdown/disconnect. - fn log_sync_error(&self, context: &str, err: &anyhow::Error) { - if self.should_downgrade_sync_error(err) { - debug!("Skipping {context} during shutdown: {err}"); - } else { - warn!("Failed {context}: {err}"); + impl EncodeNode for AckNode<'_> { + fn tag(&self) -> &str { + "ack" } - } - - /// Returns `true` when the client has completed its full startup: - /// transport connected, server authenticated, and critical app state synced. - /// This is the condition `wait_for_connected` uses to resolve. - fn is_fully_ready(&self) -> bool { - self.is_connected() && self.is_logged_in() && self.is_ready.load(Ordering::Relaxed) - } - - /// Dispatch the Connected event and notify waiters. - fn dispatch_connected(&self) { - self.is_ready.store(true, Ordering::Relaxed); - self.core - .event_bus - .dispatch(Event::Connected(crate::types::events::Connected)); - self.connected_notifier.notify(usize::MAX); - } - - /// Enable or disable skipping of history sync notifications at runtime. - /// - /// When enabled, the client will acknowledge incoming history sync - /// notifications but will not download or process the data. - pub fn set_skip_history_sync(&self, enabled: bool) { - self.skip_history_sync.store(enabled, Ordering::Relaxed); - } - - /// Override `DeviceProps` fields before the initial pairing. Only fields - /// with `Some` are changed. In-memory only — WA Web regenerates - /// `device_props` at each registration, and it has no wire effect after - /// pairing. Call before `connect()` on every process start that still - /// needs to pair. - pub async fn set_device_props(&self, override_: wacore::store::DevicePropsOverride) { - use wacore::store::commands::DeviceCommand; - if override_.is_empty() { - return; + fn attrs_len(&self) -> usize { + self.attr_count } - if self - .persistence_manager - .get_device_snapshot() - .await - .pn - .is_some() - { - warn!( - target: "Client/DeviceProps", - "set_device_props called after pairing — stored but not sent on the wire" - ); + fn has_content(&self) -> bool { + false } - self.persistence_manager - .process_command(DeviceCommand::SetDeviceProps(override_)) - .await; - } - - /// Set the noise-handshake `ClientPayload` profile. In-memory only; - /// call before each `connect()` on a fresh process. - pub async fn set_client_profile(&self, profile: wacore::client_profile::ClientProfile) { - use wacore::store::commands::DeviceCommand; - self.persistence_manager - .process_command(DeviceCommand::SetClientProfile(profile)) - .await; - } - - /// Public entry point for processing [`MajorSyncTask`] from the sync channel. - pub async fn process_sync_task(self: &Arc, task: crate::sync_task::MajorSyncTask) { - match task { - crate::sync_task::MajorSyncTask::HistorySync { - message_id, - notification, - } => { - self.process_history_sync_task(message_id, *notification) - .await; - self.finish_history_sync_task(); + fn encode_attrs<'a, W: ByteWriter>( + &self, + enc: &mut Encoder<'a, W>, + ) -> wacore_binary::Result<()> { + enc.write_string("class")?; + enc.write_string(self.tag_str)?; + enc.write_string("id")?; + self.id.encode_value(enc)?; + enc.write_string("to")?; + self.from.encode_value(enc)?; + if let Some(pn) = self.own_pn { + enc.write_string("from")?; + enc.write_jid_owned(pn)?; + } + if let Some(p) = self.participant { + enc.write_string("participant")?; + p.encode_value(enc)?; + } + if let Some(r) = self.recipient { + enc.write_string("recipient")?; + r.encode_value(enc)?; } - crate::sync_task::MajorSyncTask::AppStateSync { name, full_sync } => { - if let Err(e) = self.process_app_state_sync_task(name, full_sync).await { - log::warn!("App state sync task for {name:?} failed: {e}"); - } + if let Some(t) = self.typ { + enc.write_string("type")?; + t.encode_value(enc)?; } + Ok(()) + } + fn encode_content<'a, W: ByteWriter>( + &self, + _enc: &mut Encoder<'a, W>, + ) -> wacore_binary::Result<()> { + Ok(()) } } - /// Returns `true` if history sync notifications are currently being skipped. - pub fn skip_history_sync_enabled(&self) -> bool { - self.skip_history_sync.load(Ordering::Relaxed) - } + let ack = AckNode { + id: id_val, + from: from_val, + participant: participant_val, + recipient: recipient_val, + typ: typ_val, + own_pn: if include_from { own_device_pn } else { None }, + tag_str: tag, + attr_count, + }; - /// Set how many one-time pre-keys are generated per upload batch. - /// - /// Defaults to WA Web's UPLOAD_KEYS_COUNT (812). Call before connecting; it - /// takes effect on the next pre-key upload. The value is clamped to the - /// protocol-safe range at upload time, so out-of-range values are coerced - /// (and logged) rather than rejected here. - pub fn set_wanted_pre_key_count(&self, count: usize) { - self.wanted_pre_key_count.store(count, Ordering::Relaxed); - } + let mut buf = Vec::with_capacity(64); + let mut encoder = Encoder::new_vec(&mut buf)?; + encoder.write_node(&ack)?; + Ok(Some(buf)) +} - /// Returns the configured pre-key upload batch size (the raw value, before - /// the upload-time clamp). - pub fn wanted_pre_key_count(&self) -> usize { - self.wanted_pre_key_count.load(Ordering::Relaxed) +/// Minimal `` stanza carrying the attrs `encode_ack_bytes` needs, +/// reconstructed after the node tree has been dropped. The original `from` +/// is the group for group/broadcast stanzas and the sender otherwise (sender +/// keeps the device qualifier; `chat` is device-stripped for DMs). Mirrors +/// whatsmeow's `sendAck` (`to`=from, copy recipient/participant). +fn message_ack_source_node(info: &crate::types::message::MessageInfo) -> Node { + let from = if info.source.is_group { + &info.source.chat + } else { + &info.source.sender + }; + let mut builder = NodeBuilder::new("message") + .attr("id", &info.id) + .attr("from", from); + if let Some(recipient) = &info.source.recipient { + builder = builder.attr("recipient", recipient); } - - pub(crate) fn is_shutting_down(&self) -> bool { - self.expected_disconnect.load(Ordering::Relaxed) || !self.is_running.load(Ordering::Relaxed) + if info.source.is_group { + builder = builder.attr("participant", &info.source.sender); } + builder.build() +} - /// Create a new `Client` with default cache configuration. - /// - /// This is the standard constructor. Use [`Client::new_with_cache_config`] - /// if you need to customise cache TTL / capacity. - pub async fn new( - runtime: Arc, - persistence_manager: Arc, - transport_factory: Arc, - http_client: Arc, - override_version: Option<(u32, u32, u32)>, - ) -> (Arc, async_channel::Receiver) { - Self::new_with_cache_config( - runtime, - persistence_manager, - transport_factory, - http_client, - override_version, - CacheConfig::default(), - ) - .await +/// Build an ack Node (used in tests for structure verification). +#[cfg(test)] +fn build_ack_node(node: &wacore_binary::NodeRef<'_>, own_device_pn: Option<&Jid>) -> Option { + let id = node.get_attr("id")?.to_node_value(); + let from_ref = node.get_attr("from")?; + let from = from_ref.to_node_value(); + // Drop participant when it duplicates `to` (the flipped `from`). + let participant = node + .get_attr("participant") + .filter(|p| p.as_str().as_ref() != from_ref.as_str().as_ref()) + .map(|v| v.to_node_value()); + let recipient = node.get_attr("recipient").map(|v| v.to_node_value()); + let tag = node.tag.as_ref(); + let typ = if tag != "message" && !is_encrypt_identity_notification(node) { + node.get_attr("type").map(|v| v.to_node_value()) + } else { + None + }; + let mut attrs = Attrs::with_capacity(7); + attrs.insert("class", NodeValue::from(tag)); + attrs.insert("id", id); + attrs.insert("to", from); + if tag == "message" + && let Some(own_device_pn) = own_device_pn + { + attrs.insert("from", NodeValue::Jid(own_device_pn.clone())); } - - /// Create a new `Client` with a custom [`CacheConfig`]. - pub async fn new_with_cache_config( - runtime: Arc, - persistence_manager: Arc, - transport_factory: Arc, - http_client: Arc, - override_version: Option<(u32, u32, u32)>, - cache_config: CacheConfig, - ) -> (Arc, async_channel::Receiver) { - let mut unique_id_bytes = [0u8; 2]; - rand::make_rng::().fill_bytes(&mut unique_id_bytes); - - let device_snapshot = persistence_manager.get_device_snapshot().await; - let core = wacore::client::CoreClient::new(device_snapshot.core.clone()); - - let (tx, rx) = async_channel::bounded(32); - - let this = Self { - runtime: runtime.clone(), - core, - persistence_manager: persistence_manager.clone(), - media_conn: Arc::new(RwLock::new(None)), - is_logged_in: Arc::new(AtomicBool::new(false)), - is_connecting: Arc::new(AtomicBool::new(false)), - is_running: Arc::new(AtomicBool::new(false)), - is_connected: Arc::new(AtomicBool::new(false)), - send_active_receipts: AtomicU32::new(0), - ik_handshake_failures: Arc::new(AtomicU32::new(0)), - shutdown_notifier: wacore::runtime::ShutdownNotifier::new(), - connection_shutdown: std::sync::Mutex::new(wacore::runtime::ShutdownNotifier::new()), - last_data_received_ms: Arc::new(AtomicU64::new(0)), - last_data_sent_ms: Arc::new(AtomicU64::new(0)), - - transport: Arc::new(Mutex::new(None)), - transport_events: Arc::new(Mutex::new(None)), - transport_factory, - noise_socket: Arc::new(Mutex::new(None)), - - response_waiters: Arc::new(Mutex::new(HashMap::new())), - node_waiters: std::sync::Mutex::new(Vec::new()), - node_waiter_count: AtomicUsize::new(0), - sent_node_waiters: std::sync::Mutex::new(Vec::new()), - sent_node_waiter_count: AtomicUsize::new(0), - unique_id: format!("{}.{}", unique_id_bytes[0], unique_id_bytes[1]), - id_counter: Arc::new(AtomicU64::new(0)), - unified_session: crate::unified_session::UnifiedSessionManager::new(), - - signal_cache: Arc::new(crate::store::signal_cache::SignalStoreCache::new()), - message_processing_semaphore: std::sync::Mutex::new(Arc::new( - async_lock::Semaphore::new(1), - )), - message_semaphore_generation: Arc::new(AtomicU64::new(0)), - // Coordination caches: capacity-only eviction, no TTL/TTI. - // These hold live mutexes and channel senders; time-based eviction - // while tasks hold references would silently break serialisation. - session_locks: Cache::builder() - .max_capacity(cache_config.session_locks_capacity.max(1)) - .build(), - chat_lanes: Cache::builder() - .max_capacity(cache_config.chat_lanes_capacity.max(1)) - .build(), - lid_pn_cache: Arc::new(LidPnCache::with_config( - &cache_config.lid_pn_cache, - cache_config.cache_stores.lid_pn_cache.clone(), - )), - ab_props: Arc::new(wacore::store::ab_props::AbPropsCache::new()), - group_cache: async_lock::Mutex::new(None), - - expected_disconnect: Arc::new(AtomicBool::new(false)), - intentional_reconnect: AtomicBool::new(false), - connection_generation: Arc::new(AtomicU64::new(0)), - - recent_messages: cache_config.recent_messages.build_with_ttl(), - - sender_key_device_cache: crate::sender_key_device_cache::SenderKeyDeviceCache::new( - &cache_config.sender_key_devices_cache, - ), - - pending_device_sync: crate::pending_device_sync::PendingDeviceSync::new(), - - pending_retries: Arc::new(std::sync::Mutex::new(HashSet::new())), - - message_retry_counts: cache_config.message_retry_counts.build_with_ttl(), - - recent_retry_reasons: cache_config.message_retry_counts.build_with_ttl(), - - session_recreate_history: cache_config.session_recreate_history.build_with_ttl(), - - undecryptable_dispatched: cache_config.undecryptable_dispatched.build_with_ttl(), - - offline_sync_metrics: Arc::new(OfflineSyncMetrics { - active: AtomicBool::new(false), - total_messages: AtomicUsize::new(0), - processed_messages: AtomicUsize::new(0), - start_time: std::sync::Mutex::new(None), - }), - offline_batch: Arc::new(crate::client::offline_resume::OfflineBatchCoordinator::new()), - - enable_auto_reconnect: Arc::new(AtomicBool::new(true)), - auto_reconnect_errors: Arc::new(AtomicU32::new(0)), - - needs_initial_full_sync: Arc::new(AtomicBool::new(false)), - - app_state_processor: async_lock::Mutex::new(None), - app_state_key_requests: Arc::new(Mutex::new(HashMap::new())), - app_state_syncing: Arc::new(Mutex::new(HashSet::new())), - initial_keys_synced_notifier: Arc::new(event_listener::Event::new()), - initial_app_state_keys_received: Arc::new(AtomicBool::new(false)), - prekey_upload_lock: Arc::new(async_lock::Mutex::new(())), - offline_sync_notifier: Arc::new(event_listener::Event::new()), - offline_sync_completed: Arc::new(AtomicBool::new(false)), - history_sync_tasks_in_flight: Arc::new(AtomicUsize::new(0)), - history_sync_idle_notifier: Arc::new(event_listener::Event::new()), - outbound_flush: Arc::new(crate::flush_scope::FlushScope::new()), - presence_subscriptions: Arc::new(async_lock::Mutex::new(HashSet::new())), - socket_ready_notifier: Arc::new(event_listener::Event::new()), - is_ready: Arc::new(AtomicBool::new(false)), - connected_notifier: Arc::new(event_listener::Event::new()), - major_sync_task_sender: tx, - pairing_cancellation_tx: Arc::new(Mutex::new(None)), - pair_code_state: Arc::new(Mutex::new(wacore::pair_code::PairCodeState::default())), - custom_enc_handlers: Arc::new(async_lock::RwLock::new(HashMap::new())), - chatstate_handlers: Arc::new(RwLock::new(Vec::new())), - pdo_pending_requests: cache_config.pdo_pending_requests.build_with_ttl(), - device_registry_cache: cache_config.device_registry_cache.build_typed_ttl( - cache_config.cache_stores.device_registry_cache.clone(), - "device_registry", - ), - stanza_router: Self::create_stanza_router(), - synchronous_ack: false, - http_client, - override_version, - skip_history_sync: AtomicBool::new(false), - wanted_pre_key_count: AtomicUsize::new(crate::prekeys::DEFAULT_WANTED_PRE_KEY_COUNT), - cache_config, - self_weak: std::sync::OnceLock::new(), - saver_handle: std::sync::OnceLock::new(), - raw_node_forwarding: AtomicBool::new(false), - }; - - let arc = Arc::new(this); - let _ = arc.self_weak.set(Arc::downgrade(&arc)); - - // Warm up the LID-PN cache from persistent storage - let warm_up_arc = arc.clone(); - arc.runtime - .spawn(Box::pin(async move { - if let Err(e) = warm_up_arc.warm_up_lid_pn_cache().await { - warn!("Failed to warm up LID-PN cache: {e}"); - } - })) - .detach(); - - // Start background task to clean up stale device registry entries - let cleanup_arc = arc.clone(); - arc.runtime - .spawn(Box::pin(async move { - cleanup_arc.device_registry_cleanup_loop().await; - })) - .detach(); - - (arc, rx) + if let Some(p) = participant { + attrs.insert("participant", p); } - - pub(crate) async fn get_group_cache(&self) -> Arc { - let mut guard = self.group_cache.lock().await; - if let Some(cache) = guard.as_ref() { - return cache.clone(); - } - debug!("Initializing Group Cache for the first time."); - let cache = Arc::new( - self.cache_config - .group_cache - .build_typed_ttl(self.cache_config.cache_stores.group_cache.clone(), "group"), - ); - *guard = Some(cache.clone()); - cache + if let Some(r) = recipient { + attrs.insert("recipient", r); } - - pub(crate) async fn get_app_state_processor(&self) -> Arc { - let mut guard = self.app_state_processor.lock().await; - if let Some(proc) = guard.as_ref() { - return proc.clone(); - } - debug!("Initializing AppStateProcessor for the first time."); - let proc = Arc::new(AppStateProcessor::new( - self.persistence_manager.backend(), - self.runtime.clone(), - )); - *guard = Some(proc.clone()); - proc + if let Some(t) = typ { + attrs.insert("type", t); } + Some(Node { + tag: Cow::Borrowed("ack"), + attrs, + content: None, + }) +} - /// Create and configure the stanza router with all the handlers. - fn create_stanza_router() -> crate::handlers::router::StanzaRouter { - use crate::handlers::{ - basic::{AckHandler, FailureHandler, StreamErrorHandler, SuccessHandler}, - chatstate::ChatstateHandler, - ib::IbHandler, - iq::IqHandler, - message::MessageHandler, - notification::NotificationHandler, - receipt::ReceiptHandler, - router::StanzaRouter, - }; - - let mut router = StanzaRouter::new(); - - // Register all handlers - router.register(Arc::new(MessageHandler)); - router.register(Arc::new(ReceiptHandler)); - router.register(Arc::new(IqHandler)); - router.register(Arc::new(SuccessHandler)); - router.register(Arc::new(FailureHandler)); - router.register(Arc::new(StreamErrorHandler)); - router.register(Arc::new(IbHandler)); - router.register(Arc::new(NotificationHandler)); - router.register(Arc::new(AckHandler)); - router.register(Arc::new(ChatstateHandler)); - - router.register(Arc::new(crate::handlers::call::CallHandler)); +/// WA Web omits `type` when ACKing ``. +fn is_encrypt_identity_notification(node: &wacore_binary::NodeRef<'_>) -> bool { + node.tag == "notification" + && node + .get_attr("type") + .is_some_and(|v| v.as_str() == "encrypt") + && node.get_optional_child("identity").is_some() +} - // Register unimplemented handlers - router.register(Arc::new(crate::handlers::presence::PresenceHandler)); +/// Computes a reconnect delay matching WhatsApp Web's Fibonacci backoff: +/// `{ algo: { type: "fibonacci", first: 1000, second: 1000 }, jitter: 0.1, max: 9e5 }` +/// +/// Sequence: 1s, 1s, 2s, 3s, 5s, 8s, 13s, 21s, 34s, 55s, 89s, 144s, ... capped at 900s. +/// Each value gets ±10% random jitter. +fn fibonacci_backoff(attempt: u32) -> Duration { + const MAX_MS: u64 = 900_000; // WA Web: 9e5 - router + let mut a: u64 = 1000; + let mut b: u64 = 1000; + for _ in 0..attempt { + let next = a.saturating_add(b).min(MAX_MS); + a = b; + b = next; } + let base = a.min(MAX_MS); - /// Registers an external event handler to the core event bus. - pub fn register_handler(&self, handler: Arc) { - self.core.event_bus.add_handler(handler); - } - - /// Enable or disable raw node forwarding. - /// When enabled, `Event::RawNode` is emitted for every decoded stanza before - /// the stanza router dispatches it. Only enable when external consumers need - /// raw protocol access (e.g. voice call stanzas). - pub fn set_raw_node_forwarding(&self, enabled: bool) { - self.raw_node_forwarding.store(enabled, Ordering::Relaxed); - } - - /// Build a [`SignalProtocolStoreAdapter`] from the current device state and signal cache. - pub(crate) async fn signal_adapter( - &self, - ) -> crate::store::signal_adapter::SignalProtocolStoreAdapter { - let device_store = self.persistence_manager.get_device_arc().await; - self.signal_adapter_from(device_store) - } - - /// Build a standalone [`SenderKeyAdapter`] from the current device state and - /// signal cache, avoiding the full five-store adapter on the SKDM path. - pub(crate) async fn sender_key_adapter( - &self, - ) -> crate::store::signal_adapter::SenderKeyAdapter { - crate::store::signal_adapter::SenderKeyAdapter::new( - self.persistence_manager.get_device_arc().await, - self.signal_cache.clone(), - ) - } - - /// Build a [`SignalProtocolStoreAdapter`] from a pre-fetched device arc. - pub(crate) fn signal_adapter_from( - &self, - device_store: Arc>, - ) -> crate::store::signal_adapter::SignalProtocolStoreAdapter { - crate::store::signal_adapter::SignalProtocolStoreAdapter::new( - device_store, - self.signal_cache.clone(), - ) - } - - /// Get the per-address session mutex from the lock cache. - pub(crate) async fn session_lock_for( - &self, - signal_addr_str: &str, - ) -> Arc> { - self.session_locks - .get_with_by_ref(signal_addr_str, async { - Arc::new(async_lock::Mutex::new(())) - }) - .await - } - - /// Get the active noise socket, or error if not connected. - pub(crate) async fn get_noise_socket( - &self, - ) -> Result, ClientError> { - self.noise_socket - .lock() - .await - .clone() - .ok_or(ClientError::NotConnected) - } - - /// Send pre-marshaled plaintext bytes through the noise socket. - /// - /// The bytes must be a valid WABinary-marshaled stanza (as produced by - /// `wacore_binary::marshal::marshal_to`). Sending malformed data will - /// cause the server to close the connection. - /// - /// This bypasses node logging and `sent_node_waiter` resolution — use - /// [`send_node`](Client::send_node) for normal stanza sending. - pub async fn send_raw_bytes(&self, plaintext: Vec) -> Result<(), ClientError> { - let noise_socket = self.get_noise_socket().await?; - noise_socket - .encrypt_and_send(bytes::Bytes::from(plaintext)) - .await?; - self.last_data_sent_ms - .store(wacore::time::now_millis().max(0) as u64, Ordering::Relaxed); - Ok(()) - } - - /// Register a chatstate handler which will be invoked when a `` stanza is received. - /// - /// The handler receives a `ChatStateEvent` with the parsed chat state information. - pub async fn register_chatstate_handler( - &self, - handler: Arc, - ) { - self.chatstate_handlers.write().await.push(handler); - } - - /// Dispatch a parsed chatstate stanza to registered handlers. - /// - /// Called by `ChatstateHandler` after parsing the incoming stanza. - pub(crate) async fn dispatch_chatstate_event( - &self, - stanza: wacore::iq::chatstate::ChatstateStanza, - ) { - use wacore::iq::chatstate::{ChatstateSource, ReceivedChatState}; - use wacore::types::events::ChatPresenceUpdate; - use wacore::types::message::MessageSource; - use wacore::types::presence::{ChatPresence, ChatPresenceMedia}; - - // Dispatch via event bus - let (chat, sender, is_group) = match &stanza.source { - ChatstateSource::User { from } => (from.clone(), from.clone(), false), - ChatstateSource::Group { from, participant } => { - (from.clone(), participant.clone(), true) - } - }; - - let (state, media) = match stanza.state { - ReceivedChatState::Typing => (ChatPresence::Composing, ChatPresenceMedia::Text), - ReceivedChatState::RecordingAudio => { - (ChatPresence::Composing, ChatPresenceMedia::Audio) - } - ReceivedChatState::Idle => (ChatPresence::Paused, ChatPresenceMedia::Text), - }; - - self.core - .event_bus - .dispatch(Event::ChatPresence(ChatPresenceUpdate { - source: MessageSource { - chat, - sender, - is_from_me: false, - is_group, - addressing_mode: None, - sender_alt: None, - recipient_alt: None, - broadcast_list_owner: None, - recipient: None, - }, - state, - media, - })); - - // Invoke legacy callback handlers - let event = ChatStateEvent::from_stanza(stanza); - let handlers = self.chatstate_handlers.read().await.clone(); - for handler in handlers { - let event_clone = event.clone(); - self.runtime - .spawn(Box::pin(async move { - (handler)(event_clone); - })) - .detach(); - } - } - - pub async fn run(self: &Arc) { - if self.is_running.swap(true, Ordering::SeqCst) { - warn!("Client `run` method called while already running."); - return; - } - while self.is_running.load(Ordering::Relaxed) { - self.expected_disconnect.store(false, Ordering::Relaxed); - - if let Err(connect_err) = self.connect().await { - let is_transient = connect_err - .downcast_ref::() - .is_some_and(|e| e.is_transient()); - if is_transient { - debug!("Transient connect failure, will retry: {connect_err:#}"); - } else { - error!("Failed to connect: {connect_err:#}. Will retry..."); - } - } else { - let unexpected_disconnect = if self.read_messages_loop().await.is_err() { - // Check intentional_reconnect AFTER read loop exits — reconnect() - // sets this flag while the loop is running, so it must be read here. - if self.expected_disconnect.load(Ordering::Relaxed) - || self.intentional_reconnect.swap(false, Ordering::Relaxed) - { - debug!("Message loop exited during expected disconnect."); - false - } else { - warn!( - "Message loop exited with an error. Will attempt to reconnect if enabled." - ); - true - } - } else if self.expected_disconnect.load(Ordering::Relaxed) { - debug!("Message loop exited gracefully (expected disconnect)."); - false - } else { - info!("Message loop exited gracefully."); - false - }; - - self.cleanup_connection_state().await; - - // Dispatch after cleanup so handlers see cleared connection state. - if unexpected_disconnect { - self.core - .event_bus - .dispatch(Event::Disconnected(crate::types::events::Disconnected)); - } - } - - if !self.enable_auto_reconnect.load(Ordering::Relaxed) { - info!("Auto-reconnect disabled, shutting down."); - self.is_running.store(false, Ordering::Relaxed); - break; - } - - // If this was an expected disconnect (e.g., 515 after pairing), reconnect immediately - if self.expected_disconnect.load(Ordering::Relaxed) { - self.auto_reconnect_errors.store(0, Ordering::Relaxed); - info!("Expected disconnect (e.g., 515), reconnecting immediately..."); - continue; - } - - let error_count = self.auto_reconnect_errors.fetch_add(1, Ordering::SeqCst); - // WA Web: Fibonacci backoff with 10% jitter, max 900s. - // algo: { type: "fibonacci", first: 1000, second: 1000 } - // jitter: 0.1, max: 9e5 - let delay = fibonacci_backoff(error_count); - info!( - "Will attempt to reconnect in {:?} (attempt {})", - delay, - error_count + 1 - ); - self.runtime.sleep(delay).await; - } - info!("Client run loop has shut down."); - } - - pub async fn connect(self: &Arc) -> Result<(), anyhow::Error> { - if self.is_connecting.swap(true, Ordering::SeqCst) { - return Err(ClientError::AlreadyConnected.into()); - } - - let _guard = scopeguard::guard((), |_| { - self.is_connecting.store(false, Ordering::Relaxed); - }); - - if self.is_connected() { - return Err(ClientError::AlreadyConnected.into()); - } - - // Reset login state for new connection attempt. This ensures that - // handle_success will properly process the stanza even if - // a previous connection's post-login task bailed out early. - self.is_logged_in.store(false, Ordering::Relaxed); - self.is_ready.store(false, Ordering::Relaxed); - self.is_connected.store(false, Ordering::Relaxed); - self.offline_sync_completed.store(false, Ordering::Relaxed); - self.offline_batch.reset(); - self.outbound_flush.reopen(); - - // WA Web: both MQTT and DGW transports use a 20s connect timeout. - // Without this, a dead network blocks on the OS TCP SYN timeout (~60-75s). - // Version fetch is also wrapped so a hung HTTP request doesn't block connect(). - let version_future = rt_timeout( - &*self.runtime, - TRANSPORT_CONNECT_TIMEOUT, - crate::version::resolve_and_update_version( - &self.persistence_manager, - &self.http_client, - self.override_version, - ), - ); - let transport_future = rt_timeout( - &*self.runtime, - TRANSPORT_CONNECT_TIMEOUT, - self.transport_factory.create_transport(), - ); - - debug!("Connecting WebSocket and fetching latest client version in parallel..."); - let (version_result, transport_result) = futures::join!(version_future, transport_future); - - version_result - .map_err(|_| anyhow!("Version fetch timed out after {TRANSPORT_CONNECT_TIMEOUT:?}"))? - .map_err(|e| anyhow!("Failed to resolve app version: {}", e))?; - let (transport, mut transport_events) = transport_result.map_err(|_| { - anyhow!("Transport connect timed out after {TRANSPORT_CONNECT_TIMEOUT:?}") - })??; - debug!("Version fetch and transport connection established."); - - let noise_socket = match handshake::do_handshake( - self.runtime.clone(), - &self.persistence_manager, - &self.ik_handshake_failures, - transport.clone(), - &mut transport_events, - ) - .await - { - Ok(socket) => socket, - Err(e) => { - transport.disconnect().await; - return Err(e.into()); - } - }; - - // Fresh per-connection shutdown so subscribers registered during this - // connection see a clean signal; the previous notifier was already - // fired on the prior cleanup_connection_state. - self.reset_connection_shutdown(); - - *self.transport.lock().await = Some(transport); - *self.transport_events.lock().await = Some(transport_events); - *self.noise_socket.lock().await = Some(noise_socket); - self.is_connected.store(true, Ordering::Release); - - // Notify waiters that socket is ready (before login) - self.socket_ready_notifier.notify(usize::MAX); - - let client_clone = self.clone(); - self.runtime - .spawn(Box::pin(async move { client_clone.keepalive_loop().await })) - .detach(); - - Ok(()) - } - - /// Deregister this companion device and disconnect. - /// Does NOT wipe stored keys. Delete the storage backend to fully clear credentials. - pub async fn logout(self: &Arc) -> Result<()> { - use wacore::iq::devices::RemoveCompanionDeviceSpec; - - self.enable_auto_reconnect.store(false, Ordering::Relaxed); - - if self.is_connected() - && let Ok(jid) = self.require_pn().await - && let Err(e) = self.execute(RemoveCompanionDeviceSpec::new(&jid)).await - { - warn!("Failed to send logout IQ: {e}"); - } - - self.disconnect().await; - - self.core - .event_bus - .dispatch(Event::LoggedOut(crate::types::events::LoggedOut { - on_connect: false, - reason: ConnectFailureReason::LoggedOut, - })); - - Ok(()) - } - - pub async fn disconnect(self: &Arc) { - info!("Disconnecting client intentionally."); - self.expected_disconnect.store(true, Ordering::Relaxed); - self.is_running.store(false, Ordering::Relaxed); - self.shutdown_notifier.notify(); - - // Prevent late receipt producers from escaping the drain window. - self.outbound_flush.close(); - self.outbound_flush - .flush(&*self.runtime, std::time::Duration::from_secs(5)) - .await; - self.notify_connection_shutdown(); - - if let Err(e) = self.persistence_manager.flush().await { - log::error!("Failed to flush device state during disconnect: {e}"); - } - - // Close after flush; cleanup may also win this race on the run loop. - if let Some(transport) = self.transport.lock().await.as_ref() { - transport.disconnect().await; - } - self.cleanup_connection_state().await; - } - - /// Backoff step used by [`reconnect()`] to create an offline window. - /// - /// `fibonacci_backoff(RECONNECT_BACKOFF_STEP)` determines the delay before - /// the run loop re-connects. This must be longer than the mock server's - /// chatstate TTL (`CHATSTATE_TTL_SECS=3`) so TTL-expiry tests pass. - /// - /// Sequence: fib(0)=1s, fib(1)=1s, fib(2)=2s, fib(3)=3s, **fib(4)=5s**. - pub const RECONNECT_BACKOFF_STEP: u32 = 4; - - /// Drop the current connection and trigger the auto-reconnect loop. - /// - /// Unlike [`disconnect`], this does **not** stop the run loop. The client - /// will reconnect automatically using the same persisted identity/store, - /// just as it would after a network interruption. Use - /// [`wait_for_connected`] to wait for the new connection to be ready. - /// - /// This is useful for: - /// - Handling network changes (e.g., Wi-Fi → cellular) - /// - Forcing a fresh server session - /// - Testing offline message delivery - pub async fn reconnect(self: &Arc) { - info!("Reconnecting: dropping transport for auto-reconnect."); - self.intentional_reconnect.store(true, Ordering::Relaxed); - self.auto_reconnect_errors - .store(Self::RECONNECT_BACKOFF_STEP, Ordering::Relaxed); - - self.outbound_flush.close(); - self.outbound_flush - .flush(&*self.runtime, std::time::Duration::from_secs(2)) - .await; - self.notify_connection_shutdown(); - - if let Some(transport) = self.transport.lock().await.as_ref() { - transport.disconnect().await; - } - } - - /// Drop the current connection and reconnect immediately with no delay. - /// - /// Unlike [`reconnect`], which introduces a deliberate offline window, - /// this method sets the `expected_disconnect` flag so the run loop - /// skips the backoff delay and reconnects as fast as possible. - pub async fn reconnect_immediately(self: &Arc) { - info!("Reconnecting immediately (expected disconnect)."); - self.expected_disconnect.store(true, Ordering::Relaxed); - - self.outbound_flush.close(); - self.outbound_flush - .flush(&*self.runtime, std::time::Duration::from_secs(2)) - .await; - self.notify_connection_shutdown(); - - if let Some(transport) = self.transport.lock().await.as_ref() { - transport.disconnect().await; - } - } - - async fn cleanup_connection_state(&self) { - // Note: node_waiters are intentionally NOT cleared here — they are - // cross-connection (callers may register a waiter before an action that - // completes on a subsequent connection, e.g. after 515 reconnect). - // sent_node_waiters ARE cleared because they match pre-encryption - // outgoing stanzas, which are transport-scoped. - self.clear_sent_node_waiters(); - self.is_logged_in.store(false, Ordering::Relaxed); - self.is_ready.store(false, Ordering::Relaxed); - // Signal the keepalive loop (and any other per-connection tasks) to - // exit promptly. Without this, a stale keepalive loop can overlap - // with the next one after reconnect. Uses the PER-CONNECTION signal - // so the terminal shutdown_notifier stays clean for reconnects. - self.notify_connection_shutdown(); - // Close the socket as part of cleanup so this path is authoritative - // even when reached via the run loop's graceful-exit flow (not just - // `Client::disconnect()`). Transport impls make `disconnect()` - // idempotent, so the redundant call from `Client::disconnect()` is - // safe. - if let Some(transport) = self.transport.lock().await.take() { - transport.disconnect().await; - } - *self.transport_events.lock().await = None; - *self.noise_socket.lock().await = None; - // Clear is_connected AFTER noise_socket is None, so no task can see - // is_connected==true with a cleared socket. send_node() independently - // checks the socket, but this ordering avoids a confusing state window. - self.is_connected.store(false, Ordering::Release); - // Presence doesn't survive reconnects: demote presence-driven active - // receipts (1 -> 0), leaving a forced value (2) untouched. - let _ = - self.send_active_receipts - .compare_exchange(1, 0, Ordering::AcqRel, Ordering::Acquire); - // Drop per-chat lanes so workers exit via channel close. - self.chat_lanes.invalidate_all(); - // Clear pending retries so stale keys from detached scopeguard - // cleanup don't suppress the first retry after reconnect. - self.pending_retries - .lock() - .unwrap_or_else(|p| p.into_inner()) - .clear(); - // Flush before clear: clear() drops dirty entries, so a disconnect - // racing an in-flight encrypt would lose the just-advanced sender-key - // chain and force a full SKDM re-fanout. A disconnect is not a logout. - // Only clear on a successful flush; on a backend error keep the cache so - // the dirty state isn't dropped and the next operation can persist it. - match self.flush_signal_cache().await { - Ok(()) => self.signal_cache.clear().await, - Err(e) => log::error!( - "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" - ), - } - // Reset semaphore to 1 permit for next offline sync. - self.swap_message_semaphore(1); - // Reset dead-socket timestamps so stale values from the previous - // connection don't trigger an immediate reconnect on the next one. - self.last_data_received_ms.store(0, Ordering::Relaxed); - self.last_data_sent_ms.store(0, Ordering::Relaxed); - self.pending_device_sync.clear().await; - // Reset offline sync state for next connection - self.offline_sync_completed.store(false, Ordering::Relaxed); - self.offline_batch.reset(); - self.offline_sync_metrics - .active - .store(false, Ordering::Release); - self.offline_sync_metrics - .total_messages - .store(0, Ordering::Release); - self.offline_sync_metrics - .processed_messages - .store(0, Ordering::Release); - match self.offline_sync_metrics.start_time.lock() { - Ok(mut guard) => *guard = None, - Err(poison) => *poison.into_inner() = None, - } - self.history_sync_tasks_in_flight - .store(0, Ordering::Relaxed); - self.history_sync_idle_notifier.notify(usize::MAX); - // Drain all pending IQ waiters so they fail fast with InternalChannelClosed - // instead of hanging until the 75s timeout. - let mut waiters_map = self.response_waiters.lock().await; - let waiter_count = waiters_map.len(); - // Replace with new map to release backing storage; old senders drop here, - // causing receivers to get RecvError → IqError::InternalChannelClosed - *waiters_map = HashMap::new(); - drop(waiters_map); - if waiter_count > 0 { - debug!( - "Dropping {} orphaned IQ response waiter(s) on disconnect", - waiter_count - ); - } - - // Clear app state tracking maps to prevent unbounded growth across reconnections. - // Replace with new collections to release backing storage. - *self.app_state_key_requests.lock().await = HashMap::new(); - *self.app_state_syncing.lock().await = HashSet::new(); - - // Drop stale media connection (auth tokens become invalid on reconnect) - *self.media_conn.write().await = None; - - // Clear app state key cache — keys will be re-fetched from DB on demand - if let Some(proc) = self.app_state_processor.lock().await.as_ref() { - proc.clear_key_cache().await; - } - } - - /// Returns a snapshot of all internal collection sizes for memory leak detection. - /// - /// Moka caches report approximate counts (pending evictions may not be reflected). - /// Call `run_pending_tasks()` on individual caches first if you need exact counts. - /// - /// Requires the `debug-diagnostics` feature. - #[cfg(feature = "debug-diagnostics")] - pub async fn memory_diagnostics(&self) -> MemoryDiagnostics { - let (sig_sessions, sig_identities, sig_sender_keys) = - self.signal_cache.entry_counts().await; - let (lid_lid, lid_pn) = self.lid_pn_cache.entry_counts(); - let pending_retries_count = self - .pending_retries - .lock() - .unwrap_or_else(|p| p.into_inner()) - .len(); - - MemoryDiagnostics { - group_cache: self - .group_cache - .lock() - .await - .as_ref() - .map_or(0, |c| c.entry_count()), - device_registry_cache: self.device_registry_cache.entry_count(), - lid_pn_lid_entries: lid_lid, - lid_pn_pn_entries: lid_pn, - recent_messages: self.recent_messages.entry_count(), - sender_key_device_cache: self.sender_key_device_cache.entry_count(), - message_retry_counts: self.message_retry_counts.entry_count(), - undecryptable_dispatched: self.undecryptable_dispatched.entry_count(), - pdo_pending_requests: self.pdo_pending_requests.entry_count(), - session_locks: self.session_locks.entry_count(), - chat_lanes: self.chat_lanes.entry_count(), - response_waiters: self.response_waiters.lock().await.len(), - node_waiters: self.node_waiter_count.load(Ordering::Relaxed), - pending_retries: pending_retries_count, - presence_subscriptions: self.presence_subscriptions.lock().await.len(), - app_state_key_requests: self.app_state_key_requests.lock().await.len(), - app_state_syncing: self.app_state_syncing.lock().await.len(), - signal_cache_sessions: sig_sessions, - signal_cache_identities: sig_identities, - signal_cache_sender_keys: sig_sender_keys, - chatstate_handlers: self.chatstate_handlers.read().await.len(), - custom_enc_handlers: self.custom_enc_handlers.read().await.len(), - } - } - - /// Flush the in-memory signal cache to the database backend. - /// Called after each message is decrypted or after encryption operations. - pub(crate) async fn flush_signal_cache(&self) -> Result<(), anyhow::Error> { - let device = self.persistence_manager.get_device_arc().await; - let device_guard = device.read().await; - self.signal_cache - .flush(&*device_guard.backend) - .await - .map_err(|e| anyhow::anyhow!("Failed to flush signal cache: {e}")) - } - - /// [`flush_signal_cache`](Self::flush_signal_cache) with error logging instead of propagation. - pub(crate) async fn flush_signal_cache_logged(&self, context: &str, id: Option<&str>) { - if let Err(e) = self.flush_signal_cache().await { - if let Some(id) = id { - log::error!("Failed to flush signal cache ({context} {id}): {e:?}"); - } else { - log::error!("Failed to flush signal cache ({context}): {e:?}"); - } - } - } - - async fn read_messages_loop(self: &Arc) -> Result<(), anyhow::Error> { - debug!("Starting message processing loop..."); - - let mut rx_guard = self.transport_events.lock().await; - let transport_events = rx_guard - .take() - .ok_or_else(|| anyhow::anyhow!("Cannot start message loop: not connected"))?; - drop(rx_guard); - - // Frame decoder to parse incoming data - let mut frame_decoder = wacore::framing::FrameDecoder::new(); - let shutdown = self.connection_shutdown_signal(); - - loop { - futures::select_biased! { - _ = wacore::runtime::wait_for_shutdown(&shutdown).fuse() => { - debug!("Shutdown signaled in message loop. Exiting message loop."); - return Ok(()); - }, - event_result = transport_events.recv().fuse() => { - match event_result { - Ok(crate::transport::TransportEvent::DataReceived(data)) => { - // Update dead-socket timer (WA Web: deadSocketTimer reset) - self.last_data_received_ms.store( - wacore::time::now_millis().max(0) as u64, - Ordering::Relaxed, - ); - - // Feed data into the frame decoder - frame_decoder.feed(&data); - - // Process all complete frames. - // Frame decryption must be sequential (noise protocol counter), - // but we spawn node processing concurrently after decryption. - let mut frames_in_batch: u32 = 0; - - while let Some(encrypted_frame) = frame_decoder.decode_frame() { - // Decrypt the frame synchronously (required for noise counter ordering) - if let Some(node) = self.decrypt_frame(encrypted_frame).await { - // Determine processing mode for this node: - // - Critical nodes (success/failure/stream:error): inline, required for state - // - Message nodes: inline, preserves arrival order for per-chat queues - // (MessageHandler just enqueues + ACKs, heavy crypto runs in workers) - // - ib (in-band): inline, ensures offline sync tracking (expected count) - // is set up before offline messages are processed - // - Everything else: spawned concurrently for parallelism - let process_inline = matches!( - node.tag(), - "success" | "failure" | "stream:error" | "message" | "ib" - ); - - if process_inline { - self.process_decrypted_node(node).await; - } else { - let client = self.clone(); - self.runtime.spawn(Box::pin(async move { - client.process_decrypted_node(node).await; - })).detach(); - } - } - - // Check if we should exit after processing (e.g., after 515 stream error) - if self.expected_disconnect.load(Ordering::Relaxed) { - debug!("Expected disconnect signaled during frame processing. Exiting message loop."); - return Ok(()); - } - - // Cooperative yield — frequency and behavior are runtime-defined. - frames_in_batch += 1; - if frames_in_batch.is_multiple_of(self.runtime.yield_frequency()) - && let Some(yield_fut) = self.runtime.yield_now() - { - yield_fut.await; - } - } - - // Refresh timestamp after processing the entire batch so - // the keepalive loop sees the batch completion time, not - // just the arrival time. Prevents stale reads when a - // large batch (e.g. offline sync) takes seconds to drain. - if frames_in_batch > 1 { - self.last_data_received_ms.store( - wacore::time::now_millis().max(0) as u64, - Ordering::Relaxed, - ); - } - }, - Ok(crate::transport::TransportEvent::Disconnected(reason)) => { - if !self.expected_disconnect.load(Ordering::Relaxed) { - debug!("Transport disconnected unexpectedly: {reason}"); - return Err(anyhow::anyhow!("Transport disconnected: {reason}")); - } else { - debug!("Transport disconnected as expected: {reason}"); - return Ok(()); - } - } - // Event channel closed (no DisconnectReason available). - Err(_) => { - if !self.expected_disconnect.load(Ordering::Relaxed) { - debug!("Transport event channel closed unexpectedly."); - return Err(anyhow::anyhow!("Transport event channel closed")); - } else { - return Ok(()); - } - } - Ok(crate::transport::TransportEvent::Connected) => { - // Already handled during handshake, but could be useful for logging - debug!("Transport connected event received"); - } - } - } - } - } - } - - /// Decrypt a frame and return the parsed node as a zero-copy OwnedNodeRef. - /// This must be called sequentially due to noise protocol counter requirements. - pub(crate) async fn decrypt_frame( - self: &Arc, - encrypted_frame: bytes::BytesMut, - ) -> Option { - let noise_socket = match self.get_noise_socket().await { - Ok(s) => s, - Err(_) => { - log::error!("Cannot process frame: not connected (no noise socket)"); - return None; - } - }; - - let decrypted_payload = match noise_socket.decrypt_frame(encrypted_frame) { - Ok(p) => p, - Err(e) => { - log::error!("Failed to decrypt frame: {e}"); - return None; - } - }; - - let buffer = match wacore_binary::util::unpack_bytes(decrypted_payload) { - Ok(data) => data, - Err(e) => { - log::warn!(target: "Client/Recv", "Failed to decompress frame: {e}"); - return None; - } - }; - - match wacore_binary::OwnedNodeRef::new(buffer) { - Ok(owned) => Some(owned), - Err(e) => { - log::warn!(target: "Client/Recv", "Failed to unmarshal node: {e}"); - None - } - } - } - - /// Process an already-decrypted node. - /// This can be spawned concurrently since it doesn't depend on noise protocol state. - /// The node is wrapped in Arc to avoid cloning when passing through handlers. - pub(crate) async fn process_decrypted_node( - self: &Arc, - node: wacore_binary::OwnedNodeRef, - ) { - // Wrap in Arc once - all handlers will share this same allocation - let node_arc = Arc::new(node); - self.process_node(node_arc).await; - } - - /// Process a node wrapped in Arc. Handlers receive the Arc and can share/store it cheaply. - pub(crate) async fn process_node(self: &Arc, node: Arc) { - use wacore::xml::DisplayableNodeRef; - let nr = node.get(); - - // --- Offline Sync Tracking --- - if nr.tag.as_ref() == "ib" { - // Check for offline_preview child to get expected count - if let Some(preview) = nr.get_optional_child("offline_preview") { - let count: usize = preview - .get_attr("count") - .map(|v| v.as_str()) - .and_then(|s| s.parse().ok()) - .unwrap_or(0); - - if count == 0 { - self.offline_sync_metrics - .active - .store(false, Ordering::Release); - debug!(target: "Client/OfflineSync", "Sync COMPLETED: 0 items."); - } else { - // Use stronger memory ordering for state transitions - self.offline_sync_metrics - .total_messages - .store(count, Ordering::Release); - self.offline_sync_metrics - .processed_messages - .store(0, Ordering::Release); - self.offline_sync_metrics - .active - .store(true, Ordering::Release); - match self.offline_sync_metrics.start_time.lock() { - Ok(mut guard) => *guard = Some(wacore::time::Instant::now()), - Err(poison) => *poison.into_inner() = Some(wacore::time::Instant::now()), - } - debug!(target: "Client/OfflineSync", "Sync STARTED: Expecting {} items.", count); - } - } else if self.offline_sync_metrics.active.load(Ordering::Acquire) - && nr.get_optional_child("offline").is_some() - { - // Handle end marker: signals sync completion - // Only with an child is a real end marker. - // Other children (thread_metadata, edge_routing, dirty) are NOT end markers. - let processed = self - .offline_sync_metrics - .processed_messages - .load(Ordering::Acquire); - let elapsed = match self.offline_sync_metrics.start_time.lock() { - Ok(guard) => guard.map(|t| t.elapsed()).unwrap_or_default(), - Err(poison) => poison.into_inner().map(|t| t.elapsed()).unwrap_or_default(), - }; - debug!(target: "Client/OfflineSync", "Sync COMPLETED: End marker received. Processed {} items in {:.2?}.", processed, elapsed); - self.offline_sync_metrics - .active - .store(false, Ordering::Release); - } - } - - // Track progress if active - if self.offline_sync_metrics.active.load(Ordering::Acquire) { - // Check for 'offline' attribute on relevant stanzas - if nr.get_attr("offline").is_some() { - let processed = self - .offline_sync_metrics - .processed_messages - .fetch_add(1, Ordering::Release) - + 1; - let total = self - .offline_sync_metrics - .total_messages - .load(Ordering::Acquire); - - if processed.is_multiple_of(50) || processed == total { - trace!(target: "Client/OfflineSync", "Sync Progress: {}/{}", processed, total); - } - - // Drive WA Web pull-batch loop (non-adaptive `$13`): when - // remaining drops to <=C and no batch request is in flight, - // schedule the next one. - let pending = total.saturating_sub(processed); - crate::client::offline_resume::on_offline_stanza_arrived(self, pending); - - if processed >= total { - let elapsed = match self.offline_sync_metrics.start_time.lock() { - Ok(guard) => guard.map(|t| t.elapsed()).unwrap_or_default(), - Err(poison) => poison.into_inner().map(|t| t.elapsed()).unwrap_or_default(), - }; - debug!(target: "Client/OfflineSync", "Sync COMPLETED: Processed {} items in {:.2?}.", processed, elapsed); - self.offline_sync_metrics - .active - .store(false, Ordering::Release); - } - } - } - // --- End Tracking --- - - if nr.tag.as_ref() == "iq" - && let Some(sync_node) = nr.get_optional_child("sync") - && let Some(collection_node) = sync_node.get_optional_child("collection") - { - let name = collection_node.attrs().optional_string("name"); - let name = name.as_deref().unwrap_or(""); - debug!(target: "Client/Recv", "Received app state sync response for '{name}' (hiding content)."); - } else { - debug!(target: "Client/Recv","{}", DisplayableNodeRef(nr)); - } - - // Prepare deferred ACK cancellation flag (sent after dispatch unless cancelled) - let mut cancelled = false; - - // Emit raw node before any early returns so all decoded stanzas - // (including IQ responses and xmlstreamend) reach external observers - if self.raw_node_forwarding.load(Ordering::Relaxed) { - self.core - .event_bus - .dispatch(Event::RawNode(Arc::clone(&node))); - } - - if nr.tag.as_ref() == "xmlstreamend" { - if self.expected_disconnect.load(Ordering::Relaxed) { - debug!("Received , expected disconnect."); - } else { - warn!("Received , treating as disconnect."); - } - self.notify_connection_shutdown(); - return; - } - - // Check generic node waiters (zero-cost when none registered) - if self.node_waiter_count.load(Ordering::Acquire) > 0 { - self.resolve_node_waiters(&node); - } - - if nr.tag.as_ref() == "iq" - && let Some(id) = nr.get_attr("id").map(|v| v.as_str()) - { - // Single lock acquisition: try to remove the waiter directly. - let waiter = self.response_waiters.lock().await.remove(id.as_ref()); - if let Some(waiter) = waiter { - if waiter.send(Arc::clone(&node)).is_err() { - warn!(target: "Client/IQ", "Failed to send IQ response to waiter. Receiver was likely dropped."); - } - return; - } - } - - // Dispatch to appropriate handler using the router - // Clone Arc (cheap - just reference count) not the Node itself - if !self - .stanza_router - .dispatch(self.clone(), Arc::clone(&node), &mut cancelled) - .await - { - warn!( - "Received unknown top-level node: {}", - DisplayableNodeRef(nr) - ); - } - - // Send the deferred ACK if applicable and not cancelled by handler - if self.should_ack(nr) && !cancelled { - self.maybe_deferred_ack(node).await; - } - } - - /// Per WA Web (`Handle/MsgSendReceipt.js`), only newsletter `` - /// gets `` on the success path; DM/group use - /// ``. Failure paths (retry/backfill/nack) emit `` from - /// their dedicated handlers, not via this gate. - /// - /// status@broadcast is included as a fallback: drop paths in - /// `process_group_enc_batch` (expired status, missing sender key, generic - /// decrypt error) intentionally skip the delivery receipt to avoid - /// inflating the server-side offline counter for messages we'll never - /// process. Without the transport `` from this gate, the server - /// would redeliver indefinitely. WA Web emits `` - /// in the success path on top of this; the duplicate is tolerated. - fn should_ack(&self, node: &wacore_binary::NodeRef<'_>) -> bool { - let tag = node.tag.as_ref(); - if node.get_attr("id").is_none() { - return false; - } - let Some(from) = node.get_attr("from") else { - return false; - }; - match tag { - "receipt" | "notification" | "call" => true, - "message" => from - .to_jid() - .is_some_and(|j| j.is_newsletter() || j.is_status_broadcast()), - _ => false, - } - } - - /// Possibly send a deferred ack: either immediately or via spawned task. - /// Handlers can cancel by setting `cancelled` to true. - /// Uses Arc to avoid cloning when spawning the async task. - async fn maybe_deferred_ack(self: &Arc, node: Arc) { - if self.synchronous_ack { - if let Err(e) = self.send_ack_for(node.get()).await - && !e.is_transport_unavailable() - { - warn!("Failed to send ack: {e:?}"); - } - } else { - let this = self.clone(); - self.runtime - .spawn(Box::pin(async move { - if let Err(e) = this.send_ack_for(node.get()).await - && !e.is_transport_unavailable() - { - warn!("Failed to send ack: {e:?}"); - } - })) - .detach(); - } - } - - /// Build and send an node corresponding to the given stanza. - async fn send_ack_for(&self, node: &wacore_binary::NodeRef<'_>) -> Result<(), ClientError> { - if self.expected_disconnect.load(Ordering::Relaxed) { - return Ok(()); - } - if !self.is_connected() { - return Err(ClientError::NotConnected); - } - let own_pn = self.get_pn().await; - let buf = match encode_ack_bytes(node, own_pn.as_ref()) { - Ok(Some(buf)) => buf, - Ok(None) => return Ok(()), - Err(e) => { - log::warn!("Failed to encode ack: {e}"); - return Ok(()); - } - }; - self.send_raw_bytes(buf).await - } - - /// Send a transport ack so the server stops replaying a stanza from the - /// offline queue. Awaitable so callers can order it after a retry receipt - /// in a single flushed task. - pub(crate) async fn send_transport_ack(&self, info: &crate::types::message::MessageInfo) { - let source = message_ack_source_node(info); - let own_pn = self.get_pn().await; - match encode_ack_bytes(&source.as_node_ref(), own_pn.as_ref()) { - Ok(Some(buf)) => { - if let Err(e) = self.send_raw_bytes(buf).await - && !e.is_transport_unavailable() - { - log::warn!("Failed to send transport ack for undecryptable message: {e:?}"); - } - } - Ok(None) => {} - Err(e) => log::warn!("Failed to encode transport ack: {e}"), - } - } - - /// Spawn [`Self::send_transport_ack`], tracked via `outbound_flush` so - /// `disconnect()` flushes it (issue #571), same as delivery receipts. - pub(crate) fn spawn_message_ack( - self: &Arc, - info: &Arc, - ) { - let client = Arc::clone(self); - let info = Arc::clone(info); - self.outbound_flush.spawn(&*self.runtime, async move { - client.send_transport_ack(&info).await; - }); - } - - /// Tracked ack encoded from the original node. Use when the stanza carries - /// `recipient` (LID-routed/hosted-companion/peer) since `MessageInfo` - /// drops it on non-self branches and the server needs it for routing. - pub(crate) async fn spawn_node_transport_ack( - self: &Arc, - node: &wacore_binary::NodeRef<'_>, - ) { - let own_pn = self.get_pn().await; - let buf = match encode_ack_bytes(node, own_pn.as_ref()) { - Ok(Some(b)) => b, - Ok(None) => return, - Err(e) => { - log::warn!("Failed to encode node transport ack: {e}"); - return; - } - }; - let client = Arc::clone(self); - self.outbound_flush.spawn(&*self.runtime, async move { - if let Err(e) = client.send_raw_bytes(buf).await - && !e.is_transport_unavailable() - { - log::warn!("Failed to send node transport ack: {e:?}"); - } - }); - } - - pub async fn set_passive(&self, passive: bool) -> Result<(), crate::request::IqError> { - use wacore::iq::passive::PassiveModeSpec; - self.execute(PassiveModeSpec::new(passive)).await - } - - pub async fn clean_dirty_bits( - &self, - bit: wacore::iq::dirty::DirtyBit, - ) -> Result<(), crate::request::IqError> { - use wacore::iq::dirty::CleanDirtyBitsSpec; - - let spec = CleanDirtyBitsSpec::single(bit); - self.execute(spec).await - } - - pub async fn fetch_props(&self) -> Result<(), crate::request::IqError> { - use wacore::iq::props::PropsSpec; - use wacore::store::commands::DeviceCommand; - - let stored_hash = self - .persistence_manager - .get_device_snapshot() - .await - .props_hash - .clone(); - - // Deltas only contain changed props, so they're invalid against an empty cache. - let spec = match &stored_hash { - Some(hash) if self.ab_props.is_seeded() => { - debug!("Fetching props with hash for delta update..."); - PropsSpec::with_hash(hash) - } - _ => { - debug!("Fetching props (full)..."); - PropsSpec::new() - } - }; - - let response = self.execute(spec).await?; - - if response.delta_update { - debug!( - "Props delta update received ({} changed props)", - response.experiment_props.len() - ); - } else { - debug!( - "Props full update received ({} props, hash={:?})", - response.experiment_props.len(), - response.hash - ); - } - - self.ab_props - .apply_props(response.delta_update, response.experiment_props.into_iter()) - .await; - - if let Some(new_hash) = response.hash { - self.persistence_manager - .process_command(DeviceCommand::SetPropsHash(Some(new_hash))) - .await; - } - - Ok(()) - } - - pub(crate) fn ab_props(&self) -> &wacore::store::ab_props::AbPropsCache { - &self.ab_props - } - - pub async fn fetch_privacy_settings( - &self, - ) -> Result { - use wacore::iq::privacy::PrivacySettingsSpec; - - debug!("Fetching privacy settings..."); - - self.execute(PrivacySettingsSpec::new()).await - } - - /// Set a privacy setting. - /// - /// Use [`PrivacyCategory::is_valid_value`] to check valid combinations. - /// - /// # Example - /// ```ignore - /// use wacore::iq::privacy::{PrivacyCategory, PrivacyValue}; - /// client.set_privacy_setting(PrivacyCategory::Last, PrivacyValue::Contacts).await?; - /// ``` - pub async fn set_privacy_setting( - &self, - category: wacore::iq::privacy::PrivacyCategory, - value: wacore::iq::privacy::PrivacyValue, - ) -> Result { - use wacore::iq::privacy::SetPrivacySettingSpec; - self.execute(SetPrivacySettingSpec::new(category, value)) - .await - } - - /// Set a privacy setting to `contact_blacklist` with a disallowed list update. - /// - /// Only `Last`, `Profile`, `Status`, `GroupAdd` support disallowed lists. - /// Returns the server's updated dhash for use in subsequent updates. - pub async fn set_privacy_disallowed_list( - &self, - category: wacore::iq::privacy::PrivacyCategory, - update: wacore::iq::privacy::DisallowedListUpdate, - ) -> Result { - use wacore::iq::privacy::SetPrivacySettingSpec; - self.execute(SetPrivacySettingSpec::with_disallowed_list( - category, update, - )) - .await - } - - /// Set the default disappearing messages duration (seconds). Pass 0 to disable. - pub async fn set_default_disappearing_mode( - &self, - duration: u32, - ) -> Result<(), crate::request::IqError> { - use wacore::iq::privacy::SetDefaultDisappearingModeSpec; - self.execute(SetDefaultDisappearingModeSpec::new(duration)) - .await - } - - /// Get business profile for a WhatsApp Business account. - pub async fn get_business_profile( - &self, - jid: &wacore_binary::Jid, - ) -> Result, crate::request::IqError> { - use wacore::iq::business::BusinessProfileSpec; - self.execute(BusinessProfileSpec::new(jid)).await - } - - /// Reject an incoming call. Fire-and-forget — no server response is expected. - pub async fn reject_call( - &self, - call_id: &str, - call_from: &wacore_binary::Jid, - ) -> Result<(), anyhow::Error> { - anyhow::ensure!(!call_id.is_empty(), "call_id cannot be empty"); - let id = self.generate_request_id(); - - let stanza = wacore_binary::builder::NodeBuilder::new("call") - .attr("to", call_from) - .attr("id", id) - .children([wacore_binary::builder::NodeBuilder::new("reject") - .attr("call-id", call_id) - .attr("call-creator", call_from) - .attr("count", "0") - .build()]) - .build(); - - self.send_node(stanza).await?; - Ok(()) - } - - pub async fn send_digest_key_bundle(&self) -> Result<(), crate::request::IqError> { - use wacore::iq::prekeys::DigestKeyBundleSpec; - - debug!("Sending digest key bundle..."); - - self.execute(DigestKeyBundleSpec::new()).await.map(|_| ()) - } - - pub(crate) async fn handle_success(self: &Arc, node: &wacore_binary::NodeRef<'_>) { - // Skip processing if an expected disconnect is pending (e.g., 515 received). - // This prevents race conditions where a spawned success handler runs after - // cleanup_connection_state has already reset is_logged_in. - if self.expected_disconnect.load(Ordering::Relaxed) { - debug!("Ignoring stanza: expected disconnect pending"); - return; - } - - // Guard against multiple stanzas (WhatsApp may send more than one during - // routing/reconnection). Only process the first one per connection. - if self.is_logged_in.swap(true, Ordering::SeqCst) { - debug!("Ignoring duplicate stanza (already logged in)"); - return; - } - - // Increment connection generation to invalidate any stale post-login tasks - // from previous connections (e.g., during 515 reconnect cycles). - let current_generation = self.connection_generation.fetch_add(1, Ordering::SeqCst) + 1; - - info!( - "Successfully authenticated with WhatsApp servers! (gen={})", - current_generation - ); - self.auto_reconnect_errors.store(0, Ordering::Relaxed); - - self.update_server_time_offset(node); - - // Extract LID from the node before spawning (node isn't Send). - let lid_from_server = match node.get_attr("lid") { - Some(lid_value) => match lid_value.to_jid() { - Some(lid) => Some(lid), - None => { - warn!("Failed to parse LID from success stanza: {lid_value}"); - None - } - }, - None => { - warn!("LID not found in stanza. Group messaging may fail."); - None - } - }; - - let client_clone = self.clone(); - let task_generation = current_generation; - self.runtime.spawn(Box::pin(async move { - // Update LID if changed (moved here to avoid blocking the read loop - // on Device snapshot + write lock). - if let Some(lid) = lid_from_server { - let device_snapshot = - client_clone.persistence_manager.get_device_snapshot().await; - if device_snapshot.lid.as_ref() != Some(&lid) { - debug!("Updating LID from server to '{lid}'"); - client_clone - .persistence_manager - .process_command(DeviceCommand::SetLid(Some(lid))) - .await; - } - } - - // WA Web bumps `lc` after each successful auth (Start/Backend.js - // listener on `onOpenSocketStream`). The Comms `onConnect` handler - // gates the trigger on `isRegistered()`, so the bump only happens - // for already-paired logins — never during the pairing XX - // handshake. We mirror that by skipping when `device.pn` is None. - let already_paired = client_clone - .persistence_manager - .get_device_snapshot() - .await - .pn - .is_some(); - if already_paired { - client_clone - .persistence_manager - .process_command(DeviceCommand::IncrementLoginCounter) - .await; - } - - // Macro to check if this task is still valid (connection hasn't been replaced) - macro_rules! check_generation { - () => { - if client_clone.connection_generation.load(Ordering::SeqCst) != task_generation - { - debug!("Post-login task cancelled: connection generation changed"); - return; - } - }; - } - - debug!( - "Starting post-login initialization sequence (gen={})...", - task_generation - ); - - // Check if we need initial app state sync (empty pushname indicates fresh pairing - // where pushname will come from app state sync's setting_pushName mutation) - let device_snapshot = client_clone.persistence_manager.get_device_snapshot().await; - let needs_pushname_from_sync = device_snapshot.push_name.is_empty(); - if needs_pushname_from_sync { - debug!("Push name is empty - will be set from app state sync (setting_pushName)"); - } - - // Check connection before network operations. - // During pairing, a 515 disconnect happens quickly after success, - // so the socket may already be gone. - if !client_clone.is_connected() { - debug!( - "Skipping post-login init: connection closed (likely pairing phase reconnect)" - ); - return; - } - - check_generation!(); - client_clone.send_unified_session().await; - - // === Establish session with primary phone for PDO === - // This must happen BEFORE we exit passive mode (before offline messages arrive). - // PDO needs a session with device 0 to request decrypted content from our phone. - // Matches WhatsApp Web's bootstrapDeviceCapabilities() pattern. - check_generation!(); - if let Err(e) = client_clone - .establish_primary_phone_session_immediate() - .await - { - warn!(target: "Client/PDO", "Failed to establish session with primary phone on login: {:?}", e); - // Don't fail login - PDO will retry via ensure_e2e_sessions fallback - } - - // Sync own device list so DM fan-out includes all companions - check_generation!(); - if let Err(e) = client_clone.sync_own_device_list().await { - client_clone.log_sync_error("sync own device list", &e); - } - - check_generation!(); - if !client_clone.is_connected() { - debug!("Skipping passive tasks: connection closed"); - return; - } - if let Err(e) = client_clone.upload_pre_keys_at_login().await - && !client_clone.is_shutting_down() - { - warn!("Failed to upload pre-keys during startup: {e:?}"); - } - - // === Send active IQ === - // The server sends AFTER we exit passive mode. - // This matches WhatsApp Web's behavior: executePassiveTasks() -> sendPassiveModeProtocol("active") - check_generation!(); - if !client_clone.is_connected() { - debug!("Skipping active IQ: connection closed"); - return; - } - if let Err(e) = client_clone.set_passive(false).await - && !client_clone.is_shutting_down() - { - warn!("Failed to send post-connect active IQ: {e:?}"); - } - - // === Wait for offline sync to complete === - // The server sends after we exit passive mode. - client_clone.wait_for_offline_delivery_end().await; - - // Check if connection was replaced while waiting - check_generation!(); - - // Re-check connection and generation before sending presence - check_generation!(); - if !client_clone.is_connected() { - debug!("Skipping presence: connection closed"); - return; - } - - // Background initialization queries (can run in parallel, non-blocking) - let bg_client = client_clone.clone(); - let bg_generation = task_generation; - client_clone.runtime.spawn(Box::pin(async move { - // Check connection and generation before starting background queries - if bg_client.connection_generation.load(Ordering::SeqCst) != bg_generation { - debug!("Skipping background init queries: connection generation changed"); - return; - } - if !bg_client.is_connected() { - debug!("Skipping background init queries: connection closed"); - return; - } - - debug!( - "Sending background initialization queries (Props, Blocklist, Privacy, Digest)..." - ); - - let props_fut = bg_client.fetch_props(); - let binding = bg_client.blocking(); - let blocklist_fut = binding.get_blocklist(); - let privacy_fut = bg_client.fetch_privacy_settings(); - let digest_fut = bg_client.validate_digest_key(); - - let (r_props, r_block, r_priv, r_digest) = - futures::join!(props_fut, blocklist_fut, privacy_fut, digest_fut); - - // Suppress warnings if connection closed while queries were in-flight - if !bg_client.is_shutting_down() { - if let Err(e) = r_props { - warn!("Background init: Failed to fetch props: {e:?}"); - } - if let Err(e) = r_block { - warn!("Background init: Failed to fetch blocklist: {e:?}"); - } - if let Err(e) = r_priv { - warn!("Background init: Failed to fetch privacy settings: {e:?}"); - } - if let Err(e) = r_digest { - warn!("Background init: Failed to validate digest key: {e:?}"); - } - } - - // Prune expired tcTokens on connect (matches WhatsApp Web's PrivacyTokenJob) - if let Err(e) = bg_client.tc_token().prune_expired().await - && !bg_client.is_shutting_down() - { - warn!("Background init: Failed to prune expired tc_tokens: {e:?}"); - } - })).detach(); - - check_generation!(); - - let flag_set = client_clone.needs_initial_full_sync.load(Ordering::Relaxed); - let needs_initial_sync = flag_set || needs_pushname_from_sync; - - if needs_initial_sync { - // === Fresh pairing path === - // Like WhatsApp Web's syncCriticalData(): await critical collections before - // dispatching Connected, so blocklist/privacy settings are applied first. - debug!( - target: "Client/AppState", - "Starting Initial App State Sync (flag_set={flag_set}, needs_pushname={needs_pushname_from_sync})" - ); - - if !client_clone - .initial_app_state_keys_received - .load(Ordering::Relaxed) - { - debug!( - target: "Client/AppState", - "Waiting up to 5s for app state keys..." - ); - let _ = rt_timeout( - &*client_clone.runtime, - Duration::from_secs(5), - client_clone.initial_keys_synced_notifier.listen(), - ) - .await; - - // Check if connection was replaced while waiting - check_generation!(); - } - - // Start the critical sync timeout timer matching WhatsApp Web's - // WAWebSyncBootstrap.$15 (setSyncDCriticalDataSyncTimeout). - // WhatsApp Web uses 180s and calls socketLogout(SyncdTimeout) if - // the critical data hasn't synced by then. - const CRITICAL_SYNC_TIMEOUT_SECS: u64 = 180; - let timeout_client = client_clone.clone(); - let timeout_generation = task_generation; - let timeout_rt = client_clone.runtime.clone(); - let critical_sync_timeout_handle = timeout_rt.spawn(Box::pin(async move { - timeout_client.runtime.sleep(Duration::from_secs(CRITICAL_SYNC_TIMEOUT_SECS)).await; - // Check generation — if connection was replaced, this timeout is stale - if timeout_client.connection_generation.load(Ordering::SeqCst) - != timeout_generation - { - return; - } - // Matches WhatsApp Web's $16(): check if SettingPushName was synced. - // If push_name is still empty after 180s, critical sync failed. - let push_name = timeout_client.get_push_name().await; - if push_name.is_empty() { - warn!( - target: "Client/AppState", - "Critical app state sync timed out after {CRITICAL_SYNC_TIMEOUT_SECS}s \ - (push_name not synced). Reconnecting to retry." - ); - // WhatsApp Web does socketLogout here which clears device identity. - // We reconnect instead — preserving credentials and keeping the - // run loop active so auto-reconnect can retry the sync. - timeout_client.reconnect_immediately().await; - } else { - debug!( - target: "Client/AppState", - "Critical sync timeout fired but push_name was already synced" - ); - } - })); - - // Await critical collections via batched IQ before dispatching Connected. - check_generation!(); - match client_clone - .sync_collections_batched(vec![ - WAPatchName::CriticalBlock, - WAPatchName::CriticalUnblockLow, - ]) - .await - { - Ok(()) => { - // Critical sync completed — cancel the timeout timer - critical_sync_timeout_handle.abort(); - - check_generation!(); - - client_clone - .resubscribe_presence_subscriptions(task_generation) - .await; - - check_generation!(); - - // Dispatch Connected after critical sync completes. - // Presence is NOT sent here — WhatsApp Web sends presence from the - // setting_pushName mutation handler (WAWebPushNameSync), not from - // criticalSyncDone. Our setting_pushName handler already does this. - client_clone.dispatch_connected(); - } - Err(e) => { - client_clone.log_sync_error("critical app state sync", &e); - // Don't abort the timeout or dispatch Connected — the sync failed, - // so the timeout watchdog should remain active to force a reconnect - // if needed. Return early to avoid emitting a spurious Connected event. - return; - } - } - - // Spawn remaining non-critical collections in background - let sync_client = client_clone.clone(); - let sync_generation = task_generation; - client_clone.runtime.spawn(Box::pin(async move { - if sync_client.connection_generation.load(Ordering::SeqCst) != sync_generation { - debug!("App state sync cancelled: connection generation changed"); - return; - } - - if let Err(e) = sync_client - .sync_collections_batched(vec![ - WAPatchName::RegularLow, - WAPatchName::RegularHigh, - WAPatchName::Regular, - ]) - .await - { - sync_client.log_sync_error("non-critical app state sync", &e); - } - - sync_client - .needs_initial_full_sync - .store(false, Ordering::Relaxed); - debug!(target: "Client/AppState", "Initial App State Sync Completed."); - })).detach(); - } else { - // === Reconnection path === - // Pushname is already known, send presence and Connected immediately. - let device_snapshot = client_clone.persistence_manager.get_device_snapshot().await; - if !device_snapshot.push_name.is_empty() { - if let Err(e) = client_clone.presence().set_available().await { - warn!("Failed to send initial presence: {e:?}"); - } else { - debug!("Initial presence sent successfully."); - } - } - - client_clone - .resubscribe_presence_subscriptions(task_generation) - .await; - - // Re-check generation after awaits to avoid dispatching Connected - // for an outdated connection that was replaced mid-await. - check_generation!(); - - client_clone.dispatch_connected(); - } - })).detach(); - } - - /// Handles incoming `` stanzas by resolving pending response waiters. - /// - /// If an ack with an ID that matches a pending task in `response_waiters`, - /// the task is resolved and the function returns `true`. Otherwise, returns `false`. - pub(crate) async fn handle_ack_response(&self, node: &wacore_binary::NodeRef<'_>) -> bool { - // Surface server nack codes for diagnosability. A nacked send still - // resolves Ok to the caller, so without this the failure is invisible. - if let Some(error_code) = node.get_attr("error") { - let code = error_code.as_str(); - let id = node.get_attr("id").map(|v| v.as_str().into_owned()); - match code.as_ref() { - "463" => { - warn!( - target: "Client/Ack", - "Received 463 (MissingTcToken) nack for msg {:?}. \ - The recipient requires a valid tctoken or cstoken. \ - This may indicate a reachout timelock on the account.", - id - ); - } - "479" => { - warn!( - target: "Client/Ack", - "Received 479 (SmaxInvalid) nack for msg {:?}. \ - A stanza field has an incorrect format (e.g. wrong JID format or content type).", - id - ); - } - other => { - warn!( - target: "Client/Ack", - "Received {other} nack for msg {:?}; the message was likely \ - not delivered (e.g. 400 = malformed stanza, 404 = recipient \ - not found, 503 = service unavailable).", - id - ); - } - } - } - - let id_opt = node.get_attr("id").map(|v| v.as_str().into_owned()); - if let Some(id) = id_opt - && let Some(waiter) = self.response_waiters.lock().await.remove(&id) - { - // ACK responses are infrequent; re-encode into OwnedNodeRef for the channel. - // marshal_ref prepends a leading 0x00 format byte; OwnedNodeRef::new expects raw - // protocol bytes without it, matching what unpack() produces from the network. - match wacore_binary::marshal::marshal_ref(node) - .and_then(|bytes| wacore_binary::OwnedNodeRef::new(bytes[1..].to_vec())) - { - Ok(onr) => { - if waiter.send(Arc::new(onr)).is_err() { - warn!(target: "Client/Ack", "Failed to send ACK response to waiter for ID {id}. Receiver was likely dropped."); - } - } - Err(e) => { - warn!(target: "Client/Ack", "Failed to re-encode ACK node for waiter: {e}"); - } - } - return true; - } - false - } - - pub(crate) async fn fetch_app_state_with_retry(&self, name: WAPatchName) -> anyhow::Result<()> { - // In-flight dedup: skip if this collection is already being synced. - // Matches WA Web's WAWebSyncdCollectionsStateMachine which tracks in-flight syncs - // and queues new requests to a pending set. - { - let mut syncing = self.app_state_syncing.lock().await; - if !syncing.insert(name) { - debug!(target: "Client/AppState", "Skipping sync for {:?}: already in flight", name); - return Ok(()); - } - } - - let result = self.fetch_app_state_with_retry_inner(name).await; - - // Always remove from in-flight set when done - self.app_state_syncing.lock().await.remove(&name); - - result - } - - async fn fetch_app_state_with_retry_inner(&self, name: WAPatchName) -> anyhow::Result<()> { - let mut attempt = 0u32; - loop { - attempt += 1; - // full_sync=false lets process_app_state_sync_task auto-detect: - // version 0 → snapshot (full sync), version > 0 → incremental patches. - // Matches WA Web which only requests snapshot when version is undefined. - let res = self.process_app_state_sync_task(name, false).await; - match res { - Ok(()) => return Ok(()), - Err(e) => { - if e.downcast_ref::() - .is_some_and(|ase| { - matches!(ase, crate::appstate_sync::AppStateSyncError::KeyNotFound(_)) - }) - && attempt == 1 - { - if !self.initial_app_state_keys_received.load(Ordering::Relaxed) { - debug!(target: "Client/AppState", "App state key missing for {:?}; waiting up to 10s for key share then retrying", name); - if rt_timeout( - &*self.runtime, - Duration::from_secs(10), - self.initial_keys_synced_notifier.listen(), - ) - .await - .is_err() - { - warn!(target: "Client/AppState", "Timeout waiting for key share for {:?}; retrying anyway", name); - } - } - continue; - } - let is_db_locked = e - .downcast_ref::() - .is_some_and(|se| se.is_database_busy_or_locked()) - || e.downcast_ref::() - .is_some_and(|ase| match ase { - crate::appstate_sync::AppStateSyncError::Store(se) => { - se.is_database_busy_or_locked() - } - _ => false, - }); - if is_db_locked && attempt < APP_STATE_RETRY_MAX_ATTEMPTS { - let backoff = Duration::from_millis(200 * attempt as u64 + 150); - warn!(target: "Client/AppState", "Attempt {} for {:?} failed due to locked DB; backing off {:?} and retrying", attempt, name, backoff); - self.runtime.sleep(backoff).await; - continue; - } - return Err(e); - } - } - } - } - - /// Sync multiple collections in a single IQ request, re-fetching those with `has_more_patches`. - /// Matches WA Web's `serverSync()` outer loop (`3JJWKHeu5-P.js:54278-54305`). - /// Max 5 iterations (WA Web's `C=5` constant). - pub(crate) async fn sync_collections_batched( - &self, - collections: Vec, - ) -> anyhow::Result<()> { - if collections.is_empty() { - return Ok(()); - } - - // In-flight dedup: filter out collections already being synced - let pending = { - let mut syncing = self.app_state_syncing.lock().await; - let mut filtered = Vec::with_capacity(collections.len()); - for name in collections { - if syncing.insert(name) { - filtered.push(name); - } else { - debug!(target: "Client/AppState", "Skipping {:?} in batch: already in flight", name); - } - } - filtered - }; - - if pending.is_empty() { - return Ok(()); - } - - // Track all collections for cleanup - let all_collections: Vec = pending.clone(); - - let result = self.sync_collections_batched_inner(pending).await; - - // Always clean up in-flight set - { - let mut syncing = self.app_state_syncing.lock().await; - for name in &all_collections { - syncing.remove(name); - } - } - - result - } - - async fn sync_collections_batched_inner( - &self, - mut pending: Vec, - ) -> anyhow::Result<()> { - use wacore::appstate::patch_decode::CollectionSyncError; - const MAX_ITERATIONS: usize = 5; - let mut iteration = 0; - - while !pending.is_empty() && iteration < MAX_ITERATIONS { - iteration += 1; - debug!( - target: "Client/AppState", - "Batched sync iteration {}/{}: {:?}", - iteration, MAX_ITERATIONS, pending - ); - - let backend = self.persistence_manager.backend(); - - // Build multi-collection IQ, tracking which collections need a snapshot - let mut collection_nodes = Vec::with_capacity(pending.len()); - let mut was_snapshot = std::collections::HashSet::new(); - for &name in &pending { - let state = backend.get_version(name.as_str()).await?; - let want_snapshot = state.version == 0; - if want_snapshot { - was_snapshot.insert(name); - } - let mut builder = NodeBuilder::new("collection") - .attr("name", name.as_str()) - .attr( - "return_snapshot", - if want_snapshot { "true" } else { "false" }, - ); - if !want_snapshot { - builder = builder.attr("version", state.version); - } - collection_nodes.push(builder.build()); - } - - let sync_node = NodeBuilder::new("sync").children(collection_nodes).build(); - let iq = crate::request::InfoQuery { - namespace: "w:sync:app:state", - query_type: crate::request::InfoQueryType::Set, - to: server_jid().clone(), - target: None, - id: None, - content: Some(wacore_binary::NodeContent::Nodes(vec![sync_node])), - timeout: Some(Duration::from_secs(30)), - }; - - let resp = self.send_iq(iq).await?; - - // Pre-download all external blobs for all collections in the response - let mut pre_downloaded: std::collections::HashMap> = - std::collections::HashMap::new(); - - // Parse the response once here for pre-download; the same parsed - // lists are handed to the processor below (no second parse). - let patch_lists = wacore::appstate::patch_decode::parse_patch_lists_ref(resp.get())?; - { - for pl in &patch_lists { - // Download external snapshot - if let Some(ext) = &pl.snapshot_ref - && let Some(path) = &ext.direct_path - { - match self.download(ext).await { - Ok(bytes) => { - pre_downloaded.insert(path.clone(), bytes); - } - Err(e) => { - warn!( - "Failed to download external snapshot for {:?}: {e}", - pl.name - ); - } - } - } - - // Download external mutations - for patch in &pl.patches { - if let Some(ext) = &patch.external_mutations - && let Some(path) = &ext.direct_path - { - match self.download(ext).await { - Ok(bytes) => { - pre_downloaded.insert(path.clone(), bytes); - } - Err(e) => { - let v = - patch.version.as_ref().and_then(|v| v.version).unwrap_or(0); - warn!( - "Failed to download external mutations for patch v{}: {e}", - v - ); - } - } - } - } - } - } - - let download = |ext: &wa::ExternalBlobReference| -> anyhow::Result> { - if let Some(path) = &ext.direct_path { - if let Some(bytes) = pre_downloaded.get(path) { - Ok(bytes.clone()) - } else { - Err(anyhow::anyhow!( - "external blob not pre-downloaded: {}", - path - )) - } - } else { - Err(anyhow::anyhow!("external blob has no directPath")) - } - }; - - // Process the already-parsed collections (no re-parse of the response). - let proc = self.get_app_state_processor().await; - let results = proc - .process_patch_lists(patch_lists, &download, true) - .await?; - - let mut needs_refetch = Vec::new(); - - for (mutations, new_state, list) in results { - let name = list.name; - - // Handle per-collection errors - if let Some(ref err) = list.error { - match err { - CollectionSyncError::Conflict { has_more } => { - if *has_more { - // ConflictHasMore: server has more patches, must refetch. - warn!(target: "Client/AppState", "Collection {:?} conflict (has_more=true), will refetch", name); - needs_refetch.push(name); - } else { - // Conflict without has_more: WA Web treats this as success - // when there are no pending mutations to push (which is - // always the case for us since we don't push app state). - debug!(target: "Client/AppState", "Collection {:?} conflict (has_more=false), treating as success (no pending mutations)", name); - } - continue; - } - CollectionSyncError::Fatal { code, text } => { - warn!(target: "Client/AppState", "Collection {:?} fatal error {}: {}", name, code, text); - continue; - } - CollectionSyncError::Retry { code, text } => { - warn!(target: "Client/AppState", "Collection {:?} retryable error {}: {}, will refetch", name, code, text); - needs_refetch.push(name); - continue; - } - } - } - - // Handle missing keys - let missing = match proc.get_missing_key_ids(&list).await { - Ok(v) => v, - Err(e) => { - warn!("Failed to get missing key IDs for {:?}: {}", name, e); - Vec::new() - } - }; - self.request_missing_keys_with_dedup(missing).await; - - // full_sync is true only when this collection had a snapshot - // (version was 0 before sync). This prevents server_sync-triggered - // incremental syncs from being incorrectly marked as full syncs. - let full_sync = was_snapshot.contains(&name); - for m in mutations { - self.dispatch_app_state_mutation(&m, full_sync).await; - } - - // Save version - backend - .set_version(name.as_str(), new_state.clone()) - .await?; - - // Check if this collection needs more patches - if list.has_more_patches { - needs_refetch.push(name); - } - - debug!( - target: "Client/AppState", - "Batched sync: {:?} done (version={}, has_more={})", - name, new_state.version, list.has_more_patches - ); - } - - pending = needs_refetch; - } - - if !pending.is_empty() { - warn!( - target: "Client/AppState", - "Batched sync: max iterations ({}) reached for {:?}", - MAX_ITERATIONS, pending - ); - } - - Ok(()) - } - - pub(crate) async fn process_app_state_sync_task( - &self, - name: WAPatchName, - full_sync: bool, - ) -> anyhow::Result<()> { - if self.is_shutting_down() { - debug!(target: "Client/AppState", "Skipping app state sync task {:?}: client is shutting down", name); - return Ok(()); - } - - let backend = self.persistence_manager.backend(); - let mut full_sync = full_sync; - - let mut state = backend.get_version(name.as_str()).await?; - if state.version == 0 { - full_sync = true; - } - - let mut has_more = true; - let mut want_snapshot = full_sync; - // Safety cap to prevent infinite loops if the server keeps returning - // has_more_patches=true without advancing the version (WA Web uses 500). - const MAX_PAGINATION_ITERATIONS: u32 = 500; - let mut iteration = 0u32; - - while has_more { - if self.is_shutting_down() { - debug!(target: "Client/AppState", "Stopping app state sync task {:?}: shutdown detected", name); - break; - } - iteration += 1; - if iteration > MAX_PAGINATION_ITERATIONS { - warn!(target: "Client/AppState", "App state sync for {:?} exceeded {} iterations, aborting", name, MAX_PAGINATION_ITERATIONS); - break; - } - debug!(target: "Client/AppState", "Fetching app state patch batch: name={:?} want_snapshot={want_snapshot} version={} full_sync={} has_more_previous={}", name, state.version, full_sync, has_more); - - let mut collection_builder = NodeBuilder::new("collection") - .attr("name", name.as_str()) - .attr( - "return_snapshot", - if want_snapshot { "true" } else { "false" }, - ); - if !want_snapshot { - collection_builder = collection_builder.attr("version", state.version); - } - let sync_node = NodeBuilder::new("sync") - .children([collection_builder.build()]) - .build(); - let iq = crate::request::InfoQuery { - namespace: "w:sync:app:state", - query_type: crate::request::InfoQueryType::Set, - to: server_jid().clone(), - target: None, - id: None, - content: Some(wacore_binary::NodeContent::Nodes(vec![sync_node])), - timeout: None, - }; - - let resp = self.send_iq(iq).await?; - if self.is_shutting_down() { - debug!(target: "Client/AppState", "Discarding app state sync response for {:?}: shutdown detected", name); - break; - } - debug!(target: "Client/AppState", "Received IQ response for {:?}; decoding patches", name); - - let _decode_start = wacore::time::Instant::now(); - - // Pre-download all external blobs (snapshot and patch mutations) - // We use directPath as the key to identify each blob - let mut pre_downloaded: std::collections::HashMap> = - std::collections::HashMap::new(); - - // Parse the response once here for pre-download; the same parsed list - // is handed to the processor below (no second parse). - let pl = wacore::appstate::patch_decode::parse_patch_list_ref(resp.get())?; - { - debug!(target: "Client/AppState", "Parsed patch list for {:?}: has_snapshot_ref={} has_more_patches={} patches_count={}", - name, pl.snapshot_ref.is_some(), pl.has_more_patches, pl.patches.len()); - - // Download external snapshot if present - if let Some(ext) = &pl.snapshot_ref - && let Some(path) = &ext.direct_path - { - match self.download(ext).await { - Ok(bytes) => { - debug!(target: "Client/AppState", "Downloaded external snapshot ({} bytes)", bytes.len()); - pre_downloaded.insert(path.clone(), bytes); - } - Err(e) => { - warn!("Failed to download external snapshot: {e}"); - } - } - } - - // Download external mutations for each patch that has them - for patch in &pl.patches { - if let Some(ext) = &patch.external_mutations - && let Some(path) = &ext.direct_path - { - let patch_version = - patch.version.as_ref().and_then(|v| v.version).unwrap_or(0); - match self.download(ext).await { - Ok(bytes) => { - debug!(target: "Client/AppState", "Downloaded external mutations for patch v{} ({} bytes)", patch_version, bytes.len()); - pre_downloaded.insert(path.clone(), bytes); - } - Err(e) => { - warn!( - "Failed to download external mutations for patch v{}: {e}", - patch_version - ); - } - } - } - } - } - - let download = |ext: &wa::ExternalBlobReference| -> anyhow::Result> { - if let Some(path) = &ext.direct_path { - if let Some(bytes) = pre_downloaded.get(path) { - Ok(bytes.clone()) - } else { - Err(anyhow::anyhow!( - "external blob not pre-downloaded: {}", - path - )) - } - } else { - Err(anyhow::anyhow!("external blob has no directPath")) - } - }; - - let proc = self.get_app_state_processor().await; - let (mutations, new_state, list) = - proc.process_parsed_patch_list(pl, &download, true).await?; - let decode_elapsed = _decode_start.elapsed(); - if decode_elapsed.as_millis() > 500 { - debug!(target: "Client/AppState", "Patch decode for {:?} took {:?}", name, decode_elapsed); - } - - let missing = match proc.get_missing_key_ids(&list).await { - Ok(v) => v, - Err(e) => { - warn!("Failed to get missing key IDs for {:?}: {}", name, e); - Vec::new() - } - }; - self.request_missing_keys_with_dedup(missing).await; - - for m in mutations { - debug!(target: "Client/AppState", "Dispatching mutation kind={} index_len={} full_sync={}", m.index.first().map(|s| s.as_str()).unwrap_or(""), m.index.len(), full_sync); - self.dispatch_app_state_mutation(&m, full_sync).await; - } - - state = new_state; - has_more = list.has_more_patches; - // After the first batch, never request a snapshot again — only incremental patches. - want_snapshot = false; - debug!(target: "Client/AppState", "After processing batch name={:?} has_more={has_more} new_version={}", name, state.version); - } - - backend.set_version(name.as_str(), state.clone()).await?; - - debug!(target: "Client/AppState", "Completed and saved app state sync for {:?} (final version={})", name, state.version); - Ok(()) - } - - /// Request missing app-state keys with dedup stamps. - /// On send failure, removes stamps so keys can be retried next sync. - async fn request_missing_keys_with_dedup(&self, missing: Vec>) { - if missing.is_empty() { - return; - } - let mut to_request: Vec> = Vec::with_capacity(missing.len()); - let mut guard = self.app_state_key_requests.lock().await; - let now = wacore::time::Instant::now(); - for key_id in missing { - let hex_id = hex::encode(&key_id); - let should = guard - .get(&hex_id) - .map(|t| t.elapsed() > std::time::Duration::from_secs(24 * 3600)) - .unwrap_or(true); - if should { - guard.insert(hex_id, now); - to_request.push(key_id); - } - } - guard.retain(|_, t| t.elapsed() < std::time::Duration::from_secs(24 * 3600)); - drop(guard); - if !to_request.is_empty() - && let Err(e) = self.request_app_state_keys(&to_request).await - { - warn!("Failed to send app state key request: {e}"); - let mut guard = self.app_state_key_requests.lock().await; - for key_id in &to_request { - guard.remove(&hex::encode(key_id)); - } - } - } - - async fn request_app_state_keys(&self, raw_key_ids: &[Vec]) -> Result<(), anyhow::Error> { - if raw_key_ids.is_empty() { - return Ok(()); - } - let device_snapshot = self.persistence_manager.get_device_snapshot().await; - let own_jid = match device_snapshot.pn.clone() { - Some(j) => j, - None => { - return Err(anyhow::anyhow!( - "no own JID available for app-state key request" - )); - } - }; - let key_ids: Vec = raw_key_ids - .iter() - .map(|k| wa::message::AppStateSyncKeyId { - key_id: Some(k.clone()), - }) - .collect(); - let msg = wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage { - r#type: Some(wa::message::protocol_message::Type::AppStateSyncKeyRequest as i32), - app_state_sync_key_request: Some(wa::message::AppStateSyncKeyRequest { key_ids }), - ..Default::default() - })), - ..Default::default() - }; - self.send_message_impl( - own_jid, - &msg, - Some(self.generate_message_id().await), - true, - false, - None, - vec![], - None, - ) - .await?; - Ok(()) - } - - /// Send an app state patch to the server for a given collection. - /// - /// Builds the IQ stanza and sends it. Returns the updated hash state. - pub(crate) async fn send_app_state_patch( - &self, - collection_name: &str, - mutations: Vec, - ) -> Result<()> { - let proc = self.get_app_state_processor().await; - let (patch_bytes, base_version) = proc.build_patch(collection_name, mutations).await?; - - let collection_node = NodeBuilder::new("collection") - .attr("name", collection_name) - .attr("version", base_version) - .attr("return_snapshot", "false") - .children([NodeBuilder::new("patch").bytes(patch_bytes).build()]) - .build(); - let sync_node = NodeBuilder::new("sync").children([collection_node]).build(); - let iq = crate::request::InfoQuery { - namespace: "w:sync:app:state", - query_type: crate::request::InfoQueryType::Set, - to: server_jid().clone(), - target: None, - id: None, - content: Some(wacore_binary::NodeContent::Nodes(vec![sync_node])), - timeout: None, - }; - - self.send_iq(iq).await?; - - // Re-sync to get the latest state from the server after our patch was accepted. - // This matches whatsmeow's behavior: fetchAppState after successful send. - if let Ok(patch_name) = collection_name.parse::() - && let Err(e) = self.fetch_app_state_with_retry(patch_name).await - { - log::warn!("Failed to re-sync {collection_name} after patch send: {e}"); - } - - Ok(()) - } - - async fn dispatch_app_state_mutation( - &self, - m: &crate::appstate_sync::Mutation, - full_sync: bool, - ) { - use wacore::types::events::Event; - - if m.index.is_empty() { - return; - } - - // NCT salt sync — handles both "set" (store salt) and "remove" (clear salt). - // Source: WAWebNctSaltSync, syncd collection RegularHigh, action "nct_salt_sync". - if m.index[0] == "nct_salt_sync" { - if m.operation == wa::syncd_mutation::SyncdOperation::Remove { - debug!(target: "Client/AppState", "Removing NCT salt via app state sync"); - self.persistence_manager - .process_command(DeviceCommand::SetNctSalt(None)) - .await; - } else if let Some(val) = &m.action_value - && let Some(act) = &val.nct_salt_sync_action - && let Some(salt) = &act.salt - { - if salt.is_empty() { - warn!(target: "Client/AppState", "nct_salt_sync mutation has empty salt, ignoring"); - } else { - debug!(target: "Client/AppState", "Stored NCT salt via app state sync ({} bytes)", salt.len()); - self.persistence_manager - .process_command(DeviceCommand::SetNctSalt(Some(salt.clone()))) - .await; - } - } else { - warn!(target: "Client/AppState", "nct_salt_sync mutation missing salt in action value"); - } - return; - } - - // All remaining mutations only care about Set operations - if m.operation != wa::syncd_mutation::SyncdOperation::Set { - return; - } - - // Delegate chat-related mutations (mute, pin, archive, star, contact, etc.) - if crate::features::chat_actions::dispatch_chat_mutation(&self.core.event_bus, m, full_sync) - { - return; - } - - // Label mutations have their own index shape (labelId, not a chat JID at - // index[1]), so they are dispatched separately from chat actions. - if crate::features::labels::dispatch_label_mutation(&self.core.event_bus, m, full_sync) { - return; - } - - // Handle client-internal mutations that need persistence/presence access - if m.index[0] == "setting_pushName" - && let Some(val) = &m.action_value - && let Some(act) = &val.push_name_setting - && let Some(new_name) = &act.name - { - let new_name = new_name.clone(); - let bus = self.core.event_bus.clone(); - - let snapshot = self.persistence_manager.get_device_snapshot().await; - let old = snapshot.push_name.clone(); - if old != new_name { - debug!(target: "Client/AppState", "Persisting push name from app state mutation: '{}' (old='{}')", new_name, old); - self.persistence_manager - .process_command(DeviceCommand::SetPushName(new_name.clone())) - .await; - bus.dispatch(Event::SelfPushNameUpdated( - crate::types::events::SelfPushNameUpdated { - from_server: true, - old_name: old.clone(), - new_name: new_name.clone(), - }, - )); - - // WhatsApp Web sends presence immediately when receiving pushname - if old.is_empty() && !new_name.is_empty() { - debug!(target: "Client/AppState", "Sending presence after receiving initial pushname from app state sync"); - if let Err(e) = self.presence().set_available().await { - warn!(target: "Client/AppState", "Failed to send presence after pushname sync: {e:?}"); - } - } - } else { - debug!(target: "Client/AppState", "Push name mutation received but name unchanged: '{}'", new_name); - } - } - } - - pub(crate) async fn handle_stream_error(&self, node: &wacore_binary::NodeRef<'_>) { - // is_logged_in handling: opt-in branches (515/516/401/409/conflict) clear it - // in the disconnect block below; 429/503 clear it inline because the server - // explicitly rejected the session and outgoing sends should bail fast; the - // unknown/code-less catch-all keeps it true so is_fully_ready()-gated work - // (notably prekey uploads) survives ack-shaped routing errors. - let mut attrs = node.attrs(); - let code_cow = attrs.optional_string("code"); - let code = code_cow.as_deref().unwrap_or(""); - let conflict_type = node - .get_optional_child("conflict") - .map(|n| { - n.attrs() - .optional_string("type") - .as_deref() - .unwrap_or("") - .to_string() - }) - .unwrap_or_default(); - - // Whether to proactively disconnect the transport after handling. - let mut should_disconnect = false; - - if !conflict_type.is_empty() { - info!( - "Got stream error indicating client was removed or replaced (conflict={}). Logging out.", - conflict_type - ); - self.expected_disconnect.store(true, Ordering::Relaxed); - self.enable_auto_reconnect.store(false, Ordering::Relaxed); - - let event = if conflict_type == "replaced" { - Event::StreamReplaced(crate::types::events::StreamReplaced) - } else { - Event::LoggedOut(crate::types::events::LoggedOut { - on_connect: false, - reason: ConnectFailureReason::LoggedOut, - }) - }; - self.core.event_bus.dispatch(event); - should_disconnect = true; - } else { - match code { - "515" => { - info!( - "Got 515 stream error, server is closing stream (expected after pairing). Will auto-reconnect." - ); - self.expected_disconnect.store(true, Ordering::Relaxed); - should_disconnect = true; - } - "516" => { - info!("Got 516 stream error (device removed). Logging out."); - self.expected_disconnect.store(true, Ordering::Relaxed); - self.enable_auto_reconnect.store(false, Ordering::Relaxed); - self.core.event_bus.dispatch(Event::LoggedOut( - crate::types::events::LoggedOut { - on_connect: false, - reason: ConnectFailureReason::LoggedOut, - }, - )); - should_disconnect = true; - } - "401" => { - info!("Got 401 stream error (unauthorized). Logging out."); - self.expected_disconnect.store(true, Ordering::Relaxed); - self.enable_auto_reconnect.store(false, Ordering::Relaxed); - self.core.event_bus.dispatch(Event::LoggedOut( - crate::types::events::LoggedOut { - on_connect: false, - reason: ConnectFailureReason::LoggedOut, - }, - )); - should_disconnect = true; - } - "409" => { - info!("Got 409 stream error (conflict). Another session replaced this one."); - self.expected_disconnect.store(true, Ordering::Relaxed); - self.enable_auto_reconnect.store(false, Ordering::Relaxed); - self.core - .event_bus - .dispatch(Event::StreamReplaced(crate::types::events::StreamReplaced)); - should_disconnect = true; - } - "429" => { - // Server signalled rate-limit on this session: mark logged-out so - // outgoing sends bail fast instead of being interpreted as abuse - // while we wait for the (likely-imminent) reconnect. - warn!( - "Got 429 stream error (rate limited). Will auto-reconnect with extended backoff." - ); - self.is_logged_in.store(false, Ordering::Relaxed); - self.auto_reconnect_errors.fetch_add(5, Ordering::Relaxed); - } - "503" => { - // Server is going down/restarting: mark logged-out so sends fail - // fast against the soon-to-die socket. Auto-reconnect handles recovery. - info!("Got 503 service unavailable, will auto-reconnect."); - self.is_logged_in.store(false, Ordering::Relaxed); - } - _ => { - // Server wraps per-stanza routing failures in without a - // code (e.g. ): treat as informational so we don't trigger reconnect - // storms. is_logged_in stays true on purpose — whatsmeow clears it eagerly, - // but here is_fully_ready() gates prekey uploads and we want them to keep - // working while the socket is still alive. Severity is warn!, not error!, - // because the connection is intentionally preserved. - // WA Web (StreamError.js) knows (type "ack"); - // name it instead of "Unknown". Root cause is usually an un-acked - // offline stanza; the server's drives the reconnect. - if let Some(ack) = node.get_optional_child("ack") { - let id = ack - .get_attr("id") - .map(|v| v.as_str().to_string()) - .unwrap_or_default(); - let class = ack - .get_attr("class") - .map(|v| v.as_str().to_string()) - .unwrap_or_default(); - warn!( - "Stream error carrying (class={class:?}, id={id}): server-driven stream rotation, not an ack rejection; reconnect follows on stream end" - ); - } else { - warn!("Unknown stream error: {}", DisplayableNodeRef(node)); - } - self.core.event_bus.dispatch(Event::StreamError( - crate::types::events::StreamError { - code: code.to_string(), - raw: Some(node.to_owned()), - }, - )); - } - } - } - - // Single is_logged_in clear + transport disconnect for every opt-in branch - // (515/516/401/409 and conflict). 429/503/unknown fall through so the - // socket layer notices a real teardown without us forcing one. - if should_disconnect { - self.is_logged_in.store(false, Ordering::Relaxed); - let transport_opt = self.transport.lock().await.clone(); - if let Some(transport) = transport_opt { - self.runtime - .spawn(Box::pin(async move { - transport.disconnect().await; - })) - .detach(); - } - info!("Notifying connection shutdown from stream error handler"); - self.notify_connection_shutdown(); - } - } - - pub(crate) async fn handle_connect_failure(&self, node: &wacore_binary::NodeRef<'_>) { - self.expected_disconnect.store(true, Ordering::Relaxed); - self.notify_connection_shutdown(); - - let mut attrs = node.attrs(); - let reason_code = attrs.optional_u64("reason").unwrap_or(0) as i32; - let reason = ConnectFailureReason::from(reason_code); - - if reason.should_reconnect() { - self.expected_disconnect.store(false, Ordering::Relaxed); - } else { - self.enable_auto_reconnect.store(false, Ordering::Relaxed); - } - - if reason.is_logged_out() { - // Log the full so a server-side lock/ban is diagnosable; - // `location` (e.g. "rva") is a routing token, not the cause. - warn!( - "Got {reason:?} connect failure, logging out: {}", - DisplayableNodeRef(node) - ); - self.core - .event_bus - .dispatch(wacore::types::events::Event::LoggedOut( - crate::types::events::LoggedOut { - on_connect: true, - reason, - }, - )); - } else if let ConnectFailureReason::TempBanned = reason { - let ban_code = attrs.optional_u64("code").unwrap_or(0) as i32; - let expire_secs = attrs.optional_u64("expire").unwrap_or(0); - let expire_duration = - chrono::Duration::try_seconds(expire_secs as i64).unwrap_or_default(); - warn!( - "Temporary ban connect failure: {}", - DisplayableNodeRef(node) - ); - self.core - .event_bus - .dispatch(Event::TemporaryBan(crate::types::events::TemporaryBan { - code: crate::types::events::TempBanReason::from(ban_code), - expire: expire_duration, - })); - } else if let ConnectFailureReason::ClientOutdated = reason { - error!("Client is outdated and was rejected by server."); - self.core - .event_bus - .dispatch(Event::ClientOutdated(crate::types::events::ClientOutdated)); - } else { - warn!("Unknown connect failure: {}", DisplayableNodeRef(node)); - self.core.event_bus.dispatch(Event::ConnectFailure( - crate::types::events::ConnectFailure { - reason, - message: attrs - .optional_string("message") - .as_deref() - .unwrap_or("") - .to_string(), - raw: Some(node.to_owned()), - }, - )); - } - } - - pub(crate) async fn handle_iq(self: &Arc, node: &wacore_binary::NodeRef<'_>) -> bool { - if node.get_attr("type").is_some_and(|s| s.as_str() == "get") - && (node.get_optional_child("ping").is_some() - || node - .get_attr("xmlns") - .is_some_and(|s| s.as_str() == "urn:xmpp:ping")) - { - debug!("Received ping, sending pong."); - let mut parser = node.attrs(); - let from_jid = parser.jid("from"); - let id = parser.optional_string("id").map(|s| s.to_string()); - let pong = build_pong(from_jid.to_string(), id.as_deref()); - if let Err(e) = self.send_node(pong).await { - warn!("Failed to send pong: {e:?}"); - } - return true; - } - - if pair::handle_iq(self, node).await { - return true; - } - - false - } - - pub fn is_connected(&self) -> bool { - self.is_connected.load(Ordering::Acquire) - } - - /// Whether delivery receipts should be sent active (rendered as ticks) vs - /// `type="inactive"`. Mirrors whatsmeow's `sendActiveReceipts != 0`. - pub(crate) fn receipts_are_active(&self) -> bool { - self.send_active_receipts.load(Ordering::Acquire) != 0 - } - - /// Force active delivery receipts even when offline (whatsmeow's - /// `SetForceActiveDeliveryReceipts`); off restores the default. - pub fn set_force_active_delivery_receipts(&self, active: bool) { - self.send_active_receipts - .store(if active { 2 } else { 0 }, Ordering::Release); - } - - /// CAS so a forced value (2) is preserved (whatsmeow's `CompareAndSwap`). - pub(crate) fn mark_receipts_active_on_presence(&self) { - let _ = - self.send_active_receipts - .compare_exchange(0, 1, Ordering::AcqRel, Ordering::Acquire); - } - - pub(crate) fn mark_receipts_inactive_on_presence(&self) { - let _ = - self.send_active_receipts - .compare_exchange(1, 0, Ordering::AcqRel, Ordering::Acquire); - } - - pub fn is_logged_in(&self) -> bool { - self.is_logged_in.load(Ordering::Relaxed) - } - - /// Register a waiter for an incoming node matching the given filter. - /// - /// Returns a receiver that resolves when a matching node arrives. - /// The waiter starts buffering immediately, so register it **before** - /// performing the action that triggers the expected node. - /// - /// When multiple waiters match the same node, each matching waiter - /// receives a clone of the node (broadcast within a single resolve pass). - /// - /// # Example - /// ```ignore - /// let waiter = client.wait_for_node( - /// NodeFilter::tag("notification").attr("type", "w:gp2"), - /// ); - /// client.groups().add_participants(&group_jid, &[jid_c]).await?; - /// let node = waiter.await.expect("notification arrived"); - /// ``` - pub fn wait_for_node( - &self, - filter: NodeFilter, - ) -> futures::channel::oneshot::Receiver> { - let (tx, rx) = futures::channel::oneshot::channel(); - self.node_waiter_count.fetch_add(1, Ordering::Release); - let mut waiters = self - .node_waiters - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - waiters.push(NodeWaiter { filter, tx }); - rx - } - - /// Register a waiter for an outgoing node before it is encrypted and sent. - /// - /// This is intended for tests and diagnostics that need to inspect the raw - /// stanza built by the client, such as asserting whether `` or - /// `` was attached. - pub fn wait_for_sent_node( - &self, - filter: NodeFilter, - ) -> futures::channel::oneshot::Receiver> { - let (tx, rx) = futures::channel::oneshot::channel(); - self.sent_node_waiter_count.fetch_add(1, Ordering::Release); - let mut waiters = self - .sent_node_waiters - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - waiters.push(SentNodeWaiter { filter, tx }); - rx - } - - /// Check pending node waiters against an incoming node. - /// Only called when `node_waiter_count > 0`. - fn resolve_node_waiters(&self, node: &Arc) { - resolve_waiters(&self.node_waiters, &self.node_waiter_count, node); - } - - fn resolve_sent_node_waiters(&self, node: &Arc) { - let nr = node.as_node_ref(); - let mut waiters = self - .sent_node_waiters - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let mut i = 0; - while i < waiters.len() { - if waiters[i].tx.is_canceled() { - waiters.swap_remove(i); - self.sent_node_waiter_count.fetch_sub(1, Ordering::Release); - } else if waiters[i].filter.matches(&nr) { - let w = waiters.swap_remove(i); - self.sent_node_waiter_count.fetch_sub(1, Ordering::Release); - let _ = w.tx.send(Arc::clone(node)); - } else { - i += 1; - } - } - } - - fn clear_sent_node_waiters(&self) { - let mut waiters = self - .sent_node_waiters - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let count = waiters.len(); - if count > 0 { - waiters.clear(); - self.sent_node_waiter_count - .fetch_sub(count, Ordering::Release); - } - } - - pub(crate) fn update_server_time_offset(&self, node: &wacore_binary::NodeRef<'_>) { - self.unified_session.update_server_time_offset(node); - } - - pub(crate) async fn send_unified_session(&self) { - if !self.is_connected() { - debug!(target: "Client/UnifiedSession", "Skipping: not connected"); - return; - } - - let Some((node, _sequence)) = self.unified_session.prepare_send().await else { - return; - }; - - if let Err(e) = self.send_node(node).await { - debug!(target: "Client/UnifiedSession", "Send failed: {e}"); - self.unified_session.clear_last_sent().await; - } - } - - /// Waits for the noise socket to be established. - /// - /// Returns `Ok(())` when the socket is ready, or `Err` on timeout. - /// This is useful for code that needs to send messages before login, - /// such as requesting a pair code during initial pairing. - /// - /// If the socket is already connected, returns immediately. - pub async fn wait_for_socket(&self, timeout: std::time::Duration) -> Result<(), anyhow::Error> { - // Fast path: already connected - if self.is_connected() { - return Ok(()); - } - - // Register waiter and re-check to avoid race condition: - // If socket becomes ready between checks, the notified future captures it. - let notified = self.socket_ready_notifier.listen(); - if self.is_connected() { - return Ok(()); - } - - rt_timeout(&*self.runtime, timeout, notified) - .await - .map_err(|_| anyhow::anyhow!("Timeout waiting for socket")) - } - - /// Waits for the client to establish a connection and complete login. - /// - /// Returns `Ok(())` when connected, or `Err` on timeout. - /// This is useful for code that needs to run after connection is established - /// and authentication is complete. - /// - /// If the client is already connected and logged in, returns immediately. - pub async fn wait_for_connected( - &self, - timeout: std::time::Duration, - ) -> Result<(), anyhow::Error> { - // Fast path: fully ready (connected + logged in + critical sync done). - if self.is_fully_ready() { - return Ok(()); - } - - // Register waiter and re-check to avoid TOCTOU race: - // dispatch_connected() could fire between the check above and notified() registration. - let notified = self.connected_notifier.listen(); - if self.is_fully_ready() { - return Ok(()); - } - - rt_timeout(&*self.runtime, timeout, notified) - .await - .map_err(|_| anyhow::anyhow!("Timeout waiting for connection")) - } - - /// Get access to the PersistenceManager for this client. - /// This is useful for multi-account scenarios to get the device ID. - pub fn persistence_manager(&self) -> Arc { - self.persistence_manager.clone() - } - - pub async fn edit_message( - &self, - to: Jid, - original_id: impl Into, - new_content: wa::Message, - ) -> Result { - let original_id = original_id.into(); - - // WhatsApp Web uses getMeUserLidOrJidForChat(chat, EditMessage) which - // returns LID for LID-addressing groups and PN otherwise. - let participant = if to.is_group() { - Some( - self.get_own_jid_for_group(&to) - .await? - .to_non_ad() - .to_string(), - ) - } else { - if self.get_pn().await.is_none() { - return Err(anyhow::Error::from(ClientError::NotLoggedIn)); - } - None - }; - - let edit_container_message = crate::send::build_edit_message( - &to, - original_id.clone(), - participant, - new_content, - wacore::time::now_millis(), - ); - - // Use a new stanza ID instead of reusing the original message ID. - // The original message ID is already embedded in protocolMessage.key.id - // inside the encrypted payload. Reusing it as the outer stanza ID causes - // the server to deduplicate against the original message and silently - // drop the edit. - self.send_message_impl( - to, - &edit_container_message, - None, - false, - false, - Some(crate::types::message::EditAttribute::MessageEdit), - vec![], - None, - ) - .await?; - - Ok(original_id) - } - - /// Send a server-side reaction (used by both newsletter and status reactions). - pub(crate) async fn send_server_reaction( - &self, - to: &Jid, - server_id: u64, - reaction: &str, - ) -> Result<(), anyhow::Error> { - let request_id = self.generate_message_id().await; - - let stanza = NodeBuilder::new("message") - .attr("to", to) - .attr("type", "reaction") - .attr("id", &request_id) - .attr("server_id", server_id) - .children([NodeBuilder::new("reaction").attr("code", reaction).build()]) - .build(); - - self.send_node(stanza).await?; - Ok(()) - } - - pub async fn send_node(&self, node: Node) -> Result<(), ClientError> { - debug!(target: "Client/Send", "{}", DisplayableNode(&node)); - if self.sent_node_waiter_count.load(Ordering::Acquire) > 0 { - self.resolve_sent_node_waiters(&Arc::new(node.clone())); - } - - let plaintext_buf = wacore_binary::marshal::marshal_auto(&node).map_err(|e| { - error!("Failed to marshal node: {e:?}"); - SocketError::Marshal(e) - })?; - - self.send_raw_bytes(plaintext_buf).await - } - - /// Register a oneshot waiter for a server ack by message ID. - /// Returns the receiver — caller sends the node separately and awaits this in background. - pub(crate) async fn register_ack_waiter( - &self, - message_id: &str, - ) -> futures::channel::oneshot::Receiver> { - let (tx, rx) = futures::channel::oneshot::channel(); - self.response_waiters - .lock() - .await - .insert(message_id.to_string(), tx); - rx - } - - pub(crate) async fn update_push_name_and_notify(self: &Arc, new_name: String) { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; - let old_name = device_snapshot.push_name.clone(); - - if old_name == new_name { - return; - } - - log::debug!("Updating push name from '{}' -> '{}'", old_name, new_name); - self.persistence_manager - .process_command(DeviceCommand::SetPushName(new_name.clone())) - .await; - - self.core.event_bus.dispatch(Event::SelfPushNameUpdated( - crate::types::events::SelfPushNameUpdated { - from_server: true, - old_name, - new_name: new_name.clone(), - }, - )); - - let client_clone = self.clone(); - self.runtime - .spawn(Box::pin(async move { - if let Err(e) = client_clone.presence().set_available().await { - log::warn!("Failed to send presence after push name update: {:?}", e); - } else { - log::debug!("Sent presence after push name update."); - } - })) - .detach(); - } - - pub async fn get_push_name(&self) -> String { - self.persistence_manager - .get_device_arc() - .await - .read() - .await - .push_name - .clone() - } - - pub async fn get_pn(&self) -> Option { - self.persistence_manager - .get_device_arc() - .await - .read() - .await - .pn - .clone() - } - - pub async fn get_lid(&self) -> Option { - self.persistence_manager - .get_device_arc() - .await - .read() - .await - .lid - .clone() - } - - pub(crate) async fn require_pn(&self) -> Result { - self.get_pn().await.ok_or(ClientError::NotLoggedIn.into()) - } - - /// Resolve our own JID for a group, respecting its addressing mode. - /// - /// Returns LID for LID-addressing groups, PN otherwise. - /// Matches WhatsApp Web's `getMeUserLidOrJidForChat`. - pub(crate) async fn get_own_jid_for_group( - &self, - group_jid: &Jid, - ) -> Result { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; - let own_pn = device_snapshot - .pn - .clone() - .ok_or_else(|| anyhow::Error::from(ClientError::NotLoggedIn))?; - - let addressing_mode = self - .groups() - .query_info(group_jid) - .await - .map(|info| info.addressing_mode) - .unwrap_or(crate::types::message::AddressingMode::Pn); - - Ok(match addressing_mode { - crate::types::message::AddressingMode::Lid => { - device_snapshot.lid.clone().unwrap_or(own_pn) - } - crate::types::message::AddressingMode::Pn => own_pn, - }) - } - - /// Creates a normalized ChatMessageId by resolving PN to LID JIDs. - pub(crate) async fn make_chat_message_id(&self, chat: &Jid, id: &str) -> ChatMessageId { - // Resolve chat JID to LID if possible - let chat = self.resolve_encryption_jid(chat).await; - - ChatMessageId { - chat, - id: id.to_owned(), - } - } - - // get_phone_number_from_lid is in client/lid_pn.rs - - pub(crate) async fn send_protocol_receipt( - &self, - id: String, - receipt_type: crate::types::presence::ReceiptType, - ) { - if id.is_empty() { - return; - } - let device_snapshot = self.persistence_manager.get_device_snapshot().await; - if let Some(own_jid) = &device_snapshot.pn { - // Single source of truth for the wire mapping (ReceiptType::Sent is a derived - // incoming-only state and is never sent by us). - let type_str = receipt_type.as_wire_str(); - - let node = NodeBuilder::new("receipt") - .attrs([ - ("id", id), - ("type", type_str.to_string()), - ("to", own_jid.to_non_ad_string()), - ]) - .build(); - - if let Err(e) = self.send_node(node).await { - warn!( - "Failed to send protocol receipt of type {:?} for message ID {}: {:?}", - receipt_type, self.unique_id, e - ); - } - } - } -} - -/// Builds a pong response node for a server-initiated ping. -/// -/// Matches WhatsApp Web (`WAWebCommsHandleStanza`): only includes `id` -/// when the server ping carried one. -fn build_pong(to: String, id: Option<&str>) -> wacore_binary::Node { - let mut builder = NodeBuilder::new("iq").attr("to", to).attr("type", "result"); - if let Some(id) = id { - builder = builder.attr("id", id); - } - builder.build() -} - -/// Build an `` for the given stanza, matching WA Web / whatsmeow behavior: -/// -/// - `class` = original stanza tag -/// - `id`, `to` (flipped from `from`), `participant` copied from original -/// - `from` = own device PN, only for message acks -/// - `type` echoed for non-message stanzas (whatsmeow: `node.Tag != "message"`), -/// except `notification type="encrypt"` with `` child (WA Web drops type there). -/// -/// For receipt acks, WA Web uses `MAYBE_CUSTOM_STRING(ackString)` where -/// `ackString = maybeAttrString("type")` — so `type` is only included when -/// explicitly present on the incoming receipt (delivery receipts normally -/// have no type attribute, meaning the ack also has no type). -/// Encode an ack stanza directly to bytes, bypassing Node + marshal_auto. -/// Acks are the most frequent outbound stanza (~1 per inbound message). -fn encode_ack_bytes( - node: &wacore_binary::NodeRef<'_>, - own_device_pn: Option<&Jid>, -) -> Result>, wacore_binary::error::BinaryError> { - use wacore_binary::encoder::{ByteWriter, EncodeNode, Encoder}; - - let Some(id_val) = node.get_attr("id") else { - return Ok(None); - }; - let Some(from_val) = node.get_attr("from") else { - return Ok(None); - }; - // WAWebReceiptAck: `participant: r && r !== e ? DEVICE_JID(r) : DROP_ATTR`. - // Drop the attribute when it would duplicate `to` (which is the flipped `from`). - let participant_val = node.get_attr("participant").filter(|p| { - let p_str = p.as_str(); - let from_str = from_val.as_str(); - p_str.as_ref() != from_str.as_ref() - }); - // Server expects `recipient` echoed back so it can route the ack to the - // origin companion/device (hosted-companion, peer, LID-routed stanzas). - // Dropping it makes the server close the stream with ``. - let recipient_val = node.get_attr("recipient"); - let tag = node.tag.as_ref(); - - let typ_val = if tag != "message" && !is_encrypt_identity_notification(node) { - node.get_attr("type") - } else { - None - }; - - let include_from = tag == "message" && own_device_pn.is_some(); - - // Count attrs: class + id + to + optional(from, participant, recipient, type) - let attr_count = 3 - + usize::from(include_from) - + usize::from(participant_val.is_some()) - + usize::from(recipient_val.is_some()) - + usize::from(typ_val.is_some()); - - struct AckNode<'a> { - id: &'a wacore_binary::node::ValueRef<'a>, - from: &'a wacore_binary::node::ValueRef<'a>, - participant: Option<&'a wacore_binary::node::ValueRef<'a>>, - recipient: Option<&'a wacore_binary::node::ValueRef<'a>>, - typ: Option<&'a wacore_binary::node::ValueRef<'a>>, - own_pn: Option<&'a Jid>, - tag_str: &'a str, - attr_count: usize, - } - - impl EncodeNode for AckNode<'_> { - fn tag(&self) -> &str { - "ack" - } - fn attrs_len(&self) -> usize { - self.attr_count - } - fn has_content(&self) -> bool { - false - } - fn encode_attrs<'a, W: ByteWriter>( - &self, - enc: &mut Encoder<'a, W>, - ) -> wacore_binary::Result<()> { - enc.write_string("class")?; - enc.write_string(self.tag_str)?; - enc.write_string("id")?; - self.id.encode_value(enc)?; - enc.write_string("to")?; - self.from.encode_value(enc)?; - if let Some(pn) = self.own_pn { - enc.write_string("from")?; - enc.write_jid_owned(pn)?; - } - if let Some(p) = self.participant { - enc.write_string("participant")?; - p.encode_value(enc)?; - } - if let Some(r) = self.recipient { - enc.write_string("recipient")?; - r.encode_value(enc)?; - } - if let Some(t) = self.typ { - enc.write_string("type")?; - t.encode_value(enc)?; - } - Ok(()) - } - fn encode_content<'a, W: ByteWriter>( - &self, - _enc: &mut Encoder<'a, W>, - ) -> wacore_binary::Result<()> { - Ok(()) - } - } - - let ack = AckNode { - id: id_val, - from: from_val, - participant: participant_val, - recipient: recipient_val, - typ: typ_val, - own_pn: if include_from { own_device_pn } else { None }, - tag_str: tag, - attr_count, - }; - - let mut buf = Vec::with_capacity(64); - let mut encoder = Encoder::new_vec(&mut buf)?; - encoder.write_node(&ack)?; - Ok(Some(buf)) -} - -/// Minimal `` stanza carrying the attrs `encode_ack_bytes` needs, -/// reconstructed after the node tree has been dropped. The original `from` -/// is the group for group/broadcast stanzas and the sender otherwise (sender -/// keeps the device qualifier; `chat` is device-stripped for DMs). Mirrors -/// whatsmeow's `sendAck` (`to`=from, copy recipient/participant). -fn message_ack_source_node(info: &crate::types::message::MessageInfo) -> Node { - let from = if info.source.is_group { - &info.source.chat - } else { - &info.source.sender - }; - let mut builder = NodeBuilder::new("message") - .attr("id", &info.id) - .attr("from", from); - if let Some(recipient) = &info.source.recipient { - builder = builder.attr("recipient", recipient); - } - if info.source.is_group { - builder = builder.attr("participant", &info.source.sender); - } - builder.build() -} - -/// Build an ack Node (used in tests for structure verification). -#[cfg(test)] -fn build_ack_node(node: &wacore_binary::NodeRef<'_>, own_device_pn: Option<&Jid>) -> Option { - let id = node.get_attr("id")?.to_node_value(); - let from_ref = node.get_attr("from")?; - let from = from_ref.to_node_value(); - // Drop participant when it duplicates `to` (the flipped `from`). - let participant = node - .get_attr("participant") - .filter(|p| p.as_str().as_ref() != from_ref.as_str().as_ref()) - .map(|v| v.to_node_value()); - let recipient = node.get_attr("recipient").map(|v| v.to_node_value()); - let tag = node.tag.as_ref(); - let typ = if tag != "message" && !is_encrypt_identity_notification(node) { - node.get_attr("type").map(|v| v.to_node_value()) - } else { - None - }; - let mut attrs = Attrs::with_capacity(7); - attrs.insert("class", NodeValue::from(tag)); - attrs.insert("id", id); - attrs.insert("to", from); - if tag == "message" - && let Some(own_device_pn) = own_device_pn - { - attrs.insert("from", NodeValue::Jid(own_device_pn.clone())); - } - if let Some(p) = participant { - attrs.insert("participant", p); - } - if let Some(r) = recipient { - attrs.insert("recipient", r); - } - if let Some(t) = typ { - attrs.insert("type", t); - } - Some(Node { - tag: Cow::Borrowed("ack"), - attrs, - content: None, - }) -} - -/// WA Web omits `type` when ACKing ``. -fn is_encrypt_identity_notification(node: &wacore_binary::NodeRef<'_>) -> bool { - node.tag == "notification" - && node - .get_attr("type") - .is_some_and(|v| v.as_str() == "encrypt") - && node.get_optional_child("identity").is_some() -} - -/// Computes a reconnect delay matching WhatsApp Web's Fibonacci backoff: -/// `{ algo: { type: "fibonacci", first: 1000, second: 1000 }, jitter: 0.1, max: 9e5 }` -/// -/// Sequence: 1s, 1s, 2s, 3s, 5s, 8s, 13s, 21s, 34s, 55s, 89s, 144s, ... capped at 900s. -/// Each value gets ±10% random jitter. -fn fibonacci_backoff(attempt: u32) -> Duration { - const MAX_MS: u64 = 900_000; // WA Web: 9e5 - - let mut a: u64 = 1000; - let mut b: u64 = 1000; - for _ in 0..attempt { - let next = a.saturating_add(b).min(MAX_MS); - a = b; - b = next; - } - let base = a.min(MAX_MS); - - // ±10% jitter (WA Web: jitter: 0.1) - let jitter_range = base / 10; - let jitter = if jitter_range > 0 { - rand::make_rng::().random_range(0..=(jitter_range * 2)) as i64 - - jitter_range as i64 - } else { - 0 - }; - let ms = (base as i64 + jitter).max(0) as u64; - Duration::from_millis(ms) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::lid_pn_cache::LearningSource; - use crate::test_utils::MockHttpClient; - use futures::channel::oneshot; - use wacore_binary::SERVER_JID; - - #[tokio::test] - async fn test_ack_behavior_for_incoming_stanzas() { - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // --- Assertions --- - - // Verify that we still ack other critical stanzas (regression check). - use wacore_binary::{Attrs, Node, NodeContent}; - - let mut receipt_attrs = Attrs::new(); - receipt_attrs.insert("from".to_string(), "@s.whatsapp.net".to_string()); - receipt_attrs.insert("id".to_string(), "RCPT-1".to_string()); - let receipt_node = Node::new( - "receipt", - receipt_attrs, - Some(NodeContent::String("test".into())), - ); - - let mut notification_attrs = Attrs::new(); - notification_attrs.insert("from".to_string(), "@s.whatsapp.net".to_string()); - notification_attrs.insert("id".to_string(), "NOTIF-1".to_string()); - let notification_node = Node::new( - "notification", - notification_attrs, - Some(NodeContent::String("test".into())), - ); - - assert!( - client.should_ack(&receipt_node.as_node_ref()), - "should_ack must still return TRUE for stanzas." - ); - assert!( - client.should_ack(¬ification_node.as_node_ref()), - "should_ack must still return TRUE for stanzas." - ); - - // Regular stanzas (DM / group) are acked via the delivery - // , not a bare . WA Web only emits - // for newsletter deliveries. - let mut dm_attrs = Attrs::new(); - dm_attrs.insert( - "from".to_string(), - "5511999999999@s.whatsapp.net".to_string(), - ); - dm_attrs.insert("id".to_string(), "MSG-DM-1".to_string()); - let dm_message = Node::new("message", dm_attrs, None); - assert!( - !client.should_ack(&dm_message.as_node_ref()), - "should_ack must return FALSE for regular DM (delivery receipt covers it)." - ); - - let mut group_attrs = Attrs::new(); - group_attrs.insert("from".to_string(), "120363098765432100@g.us".to_string()); - group_attrs.insert("id".to_string(), "MSG-GROUP-1".to_string()); - let group_message = Node::new("message", group_attrs, None); - assert!( - !client.should_ack(&group_message.as_node_ref()), - "should_ack must return FALSE for group ." - ); - - let mut newsletter_attrs = Attrs::new(); - newsletter_attrs.insert( - "from".to_string(), - "120363298765432100@newsletter".to_string(), - ); - newsletter_attrs.insert("id".to_string(), "MSG-NL-1".to_string()); - let newsletter_message = Node::new("message", newsletter_attrs, None); - assert!( - client.should_ack(&newsletter_message.as_node_ref()), - "should_ack must return TRUE for newsletter ." - ); - - // status@broadcast gets the transport as a fallback so that - // drop paths in process_group_enc_batch (expired status, missing - // sender key, decrypt error) don't leave the server retransmitting. - // The success path also emits ; the - // duplicate is tolerated. - let mut status_attrs = Attrs::new(); - status_attrs.insert("from".to_string(), "status@broadcast".to_string()); - status_attrs.insert("id".to_string(), "MSG-STATUS-1".to_string()); - let status_message = Node::new("message", status_attrs, None); - assert!( - client.should_ack(&status_message.as_node_ref()), - "should_ack must return TRUE for status@broadcast (fallback for drop paths)." - ); - - info!( - "✅ test_ack_behavior_for_incoming_stanzas passed: Client correctly differentiates which stanzas to acknowledge." - ); - } - - #[tokio::test] - async fn test_ack_waiter_resolves() { - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // 1. Insert a waiter for a specific ID - let test_id = "ack-test-123".to_string(); - let (tx, rx) = oneshot::channel(); - client - .response_waiters - .lock() - .await - .insert(test_id.clone(), tx); - assert!( - client.response_waiters.lock().await.contains_key(&test_id), - "Waiter should be inserted before handling ack" - ); - - // 2. Create a mock node with the test ID - let ack_node = NodeBuilder::new("ack") - .attr("id", test_id.clone()) - .attr("from", SERVER_JID) - .build(); - - // 3. Handle the ack - let handled = client.handle_ack_response(&ack_node.as_node_ref()).await; - assert!( - handled, - "handle_ack_response should return true when waiter exists" - ); - - // 4. Await the receiver with a timeout - match tokio::time::timeout(Duration::from_secs(1), rx).await { - Ok(Ok(response_node)) => { - assert!( - response_node - .get() - .get_attr("id") - .is_some_and(|v| v.as_str() == test_id.as_str()), - "Response node should have correct ID" - ); - } - Ok(Err(_)) => panic!("Receiver was dropped without being sent a value"), - Err(_) => panic!("Test timed out waiting for ack response"), - } - - // 5. Verify the waiter was removed - assert!( - !client.response_waiters.lock().await.contains_key(&test_id), - "Waiter should be removed after handling" - ); - - info!( - "✅ test_ack_waiter_resolves passed: ACK response correctly resolves pending waiters" - ); - } - - #[tokio::test] - async fn test_ack_without_matching_waiter() { - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Create an ack without any matching waiter - let ack_node = NodeBuilder::new("ack") - .attr("id", "non-existent-id") - .attr("from", SERVER_JID) - .build(); - - // Should return false since there's no waiter - let handled = client.handle_ack_response(&ack_node.as_node_ref()).await; - assert!( - !handled, - "handle_ack_response should return false when no waiter exists" - ); - - info!( - "✅ test_ack_without_matching_waiter passed: ACK without matching waiter handled gracefully" - ); - } - - /// Test that the lid_pn_cache correctly stores and retrieves LID mappings. - /// - /// This is critical for the LID-PN session mismatch fix. When we receive a message - /// with sender_lid, we cache the phone->LID mapping so that when sending replies, - /// we can reuse the existing LID session instead of creating a new PN session. - #[tokio::test] - async fn test_lid_pn_cache_basic_operations() { - let backend = Arc::new( - crate::store::SqliteStore::new("file:memdb_lid_cache_basic?mode=memory&cache=shared") - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Initially, the cache should be empty for a phone number - let phone = "559980000001"; - let lid = "100000012345678"; - - assert!( - client.lid_pn_cache.get_current_lid(phone).await.is_none(), - "Cache should be empty initially" - ); - - // Insert a phone->LID mapping using add_lid_pn_mapping - client - .add_lid_pn_mapping(lid, phone, LearningSource::Usync) - .await - .expect("Failed to persist LID-PN mapping in tests"); - - // Verify we can retrieve it (phone -> LID lookup) - let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; - assert!(cached_lid.is_some(), "Cache should contain the mapping"); - assert_eq!( - cached_lid.expect("cache should have LID"), - lid, - "Cached LID should match what we inserted" - ); - - // Verify reverse lookup works (LID -> phone) - let cached_phone = client.lid_pn_cache.get_phone_number(lid).await; - assert!(cached_phone.is_some(), "Reverse lookup should work"); - assert_eq!( - cached_phone.expect("reverse lookup should return phone"), - phone, - "Cached phone should match what we inserted" - ); - - // Verify a different phone number returns None - assert!( - client - .lid_pn_cache - .get_current_lid("559980000002") - .await - .is_none(), - "Different phone number should not have a mapping" - ); - - info!("✅ test_lid_pn_cache_basic_operations passed: LID-PN cache works correctly"); - } - - /// Test that the lid_pn_cache respects timestamp-based conflict resolution. - /// - /// When a phone number has multiple LIDs, the most recent one should be returned. - #[tokio::test] - async fn test_lid_pn_cache_timestamp_resolution() { - let backend = Arc::new( - crate::store::SqliteStore::new( - "file:memdb_lid_cache_timestamp?mode=memory&cache=shared", - ) - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - let phone = "559980000001"; - let lid_old = "100000012345678"; - let lid_new = "100000087654321"; - - // Insert initial mapping - client - .add_lid_pn_mapping(lid_old, phone, LearningSource::Usync) - .await - .expect("Failed to persist LID-PN mapping in tests"); - - assert_eq!( - client - .lid_pn_cache - .get_current_lid(phone) - .await - .expect("cache should have LID"), - lid_old, - "Initial LID should be stored" - ); - - // Small delay to ensure different timestamp - tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; - - // Add new mapping with newer timestamp - client - .add_lid_pn_mapping(lid_new, phone, LearningSource::PeerPnMessage) - .await - .expect("Failed to persist LID-PN mapping in tests"); - - assert_eq!( - client - .lid_pn_cache - .get_current_lid(phone) - .await - .expect("cache should have newer LID"), - lid_new, - "Newer LID should be returned for phone lookup" - ); - - // Both LIDs should still resolve to the same phone - assert_eq!( - client - .lid_pn_cache - .get_phone_number(lid_old) - .await - .expect("reverse lookup should return phone"), - phone, - "Old LID should still map to phone" - ); - assert_eq!( - client - .lid_pn_cache - .get_phone_number(lid_new) - .await - .expect("reverse lookup should return phone"), - phone, - "New LID should also map to phone" - ); - - info!( - "✅ test_lid_pn_cache_timestamp_resolution passed: Timestamp-based resolution works correctly" - ); - } - - /// Test that get_lid_for_phone (from SendContextResolver) returns the cached value. - /// - /// This is the method used by wacore::send to look up LID mappings when encrypting. - #[tokio::test] - async fn test_get_lid_for_phone_via_send_context_resolver() { - use wacore::client::context::SendContextResolver; - - let backend = Arc::new( - crate::store::SqliteStore::new("file:memdb_get_lid_for_phone?mode=memory&cache=shared") - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - let phone = "559980000001"; - let lid = "100000012345678"; - - // Before caching, should return None - assert!( - client.get_lid_for_phone(phone).await.is_none(), - "get_lid_for_phone should return None before caching" - ); - - // Cache the mapping using add_lid_pn_mapping - client - .add_lid_pn_mapping(lid, phone, LearningSource::Usync) - .await - .expect("Failed to persist LID-PN mapping in tests"); - - // Now it should return the LID - let result = client.get_lid_for_phone(phone).await; - assert!( - result.is_some(), - "get_lid_for_phone should return Some after caching" - ); - assert_eq!( - result.expect("get_lid_for_phone should return Some"), - lid, - "get_lid_for_phone should return the cached LID" - ); - - info!( - "✅ test_get_lid_for_phone_via_send_context_resolver passed: SendContextResolver correctly returns cached LID" - ); - } - - /// Test that wait_for_offline_delivery_end returns immediately when the flag is already set. - #[tokio::test] - async fn test_wait_for_offline_delivery_end_returns_immediately_when_flag_set() { - let backend = Arc::new( - crate::store::SqliteStore::new( - "file:memdb_offline_sync_flag_set?mode=memory&cache=shared", - ) - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Set the flag to true (simulating offline sync completed) - client - .offline_sync_completed - .store(true, std::sync::atomic::Ordering::Relaxed); - - // This should return immediately (not wait 10 seconds) - let start = wacore::time::Instant::now(); - client.wait_for_offline_delivery_end().await; - let elapsed = start.elapsed(); - - // Should complete in < 100ms (not 10 second timeout) - assert!( - elapsed.as_millis() < 100, - "wait_for_offline_delivery_end should return immediately when flag is set, took {:?}", - elapsed - ); - - info!("✅ test_wait_for_offline_delivery_end_returns_immediately_when_flag_set passed"); - } - - /// Test that wait_for_offline_delivery_end times out when the flag is NOT set. - /// This verifies the 10-second timeout is working. - #[tokio::test] - async fn test_wait_for_offline_delivery_end_times_out_when_flag_not_set() { - let backend = Arc::new( - crate::store::SqliteStore::new( - "file:memdb_offline_sync_timeout?mode=memory&cache=shared", - ) - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Flag is false by default, so use a short timeout and verify the helper - // marks the sync complete on timeout. - let start = wacore::time::Instant::now(); - client - .wait_for_offline_delivery_end_with_timeout(std::time::Duration::from_millis(50)) - .await; - - let elapsed = start.elapsed(); - // Count available permits by trying to acquire non-blockingly - let semaphore = match client.message_processing_semaphore.lock() { - Ok(guard) => guard.clone(), - Err(poisoned) => poisoned.into_inner().clone(), - }; - let mut guards = Vec::new(); - while let Some(guard) = semaphore.try_acquire() { - guards.push(guard); - } - let permits = guards.len(); - drop(guards); - - assert!( - elapsed.as_millis() >= 45, // Allow small timing variance - "Should have waited for the configured timeout duration, took {:?}", - elapsed - ); - assert!( - client - .offline_sync_completed - .load(std::sync::atomic::Ordering::Relaxed), - "wait_for_offline_delivery_end should mark offline sync complete on timeout" - ); - assert_eq!( - permits, 64, - "timeout completion should restore parallel permits" - ); - - info!("✅ test_wait_for_offline_delivery_end_times_out_when_flag_not_set passed"); - } - - /// Test that wait_for_offline_delivery_end returns when notified. - #[tokio::test] - async fn test_wait_for_offline_delivery_end_returns_on_notify() { - let backend = Arc::new( - crate::store::SqliteStore::new("file:memdb_offline_notify?mode=memory&cache=shared") - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - let client_clone = client.clone(); - - // Spawn a task that will notify after 50ms - tokio::spawn(async move { - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - client_clone.offline_sync_notifier.notify(usize::MAX); - }); - - let start = wacore::time::Instant::now(); - client.wait_for_offline_delivery_end().await; - let elapsed = start.elapsed(); - - // Should complete around 50ms (when notified), not 10 seconds - assert!( - elapsed.as_millis() < 200, - "wait_for_offline_delivery_end should return when notified, took {:?}", - elapsed - ); - assert!( - elapsed.as_millis() >= 45, // Should have waited for the notify - "Should have waited for the notify, only took {:?}", - elapsed - ); - - info!("✅ test_wait_for_offline_delivery_end_returns_on_notify passed"); - } - - /// Test that the offline_sync_completed flag starts as false. - #[tokio::test] - async fn test_offline_sync_flag_initially_false() { - let backend = Arc::new( - crate::store::SqliteStore::new( - "file:memdb_offline_flag_initial?mode=memory&cache=shared", - ) - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // The flag should be false initially - assert!( - !client - .offline_sync_completed - .load(std::sync::atomic::Ordering::Relaxed), - "offline_sync_completed should be false when Client is first created" - ); - - info!("✅ test_offline_sync_flag_initially_false passed"); - } - - /// Test the complete offline sync lifecycle: - /// 1. Flag starts false - /// 2. Flag is set true after IB offline stanza - /// 3. Notify is called - #[tokio::test] - async fn test_offline_sync_lifecycle() { - use std::sync::atomic::Ordering; - - let backend = Arc::new( - crate::store::SqliteStore::new("file:memdb_offline_lifecycle?mode=memory&cache=shared") - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // 1. Initially false - assert!(!client.offline_sync_completed.load(Ordering::Relaxed)); - - // 2. Spawn a waiter - let client_waiter = client.clone(); - let waiter_handle = tokio::spawn(async move { - client_waiter.wait_for_offline_delivery_end().await; - true // Return that we completed - }); - - // Give the waiter time to start waiting - tokio::time::sleep(std::time::Duration::from_millis(10)).await; - - // Verify waiter hasn't completed yet - assert!( - !waiter_handle.is_finished(), - "Waiter should still be waiting" - ); - - // 3. Simulate IB handler behavior (set flag and notify) - client.offline_sync_completed.store(true, Ordering::Relaxed); - client.offline_sync_notifier.notify(usize::MAX); - - // 4. Waiter should complete - let result = tokio::time::timeout(std::time::Duration::from_millis(100), waiter_handle) - .await - .expect("Waiter should complete after notify") - .expect("Waiter task should not panic"); - - assert!(result, "Waiter should have completed successfully"); - assert!(client.offline_sync_completed.load(Ordering::Relaxed)); - - info!("✅ test_offline_sync_lifecycle passed"); - } - - /// Test that establish_primary_phone_session_immediate returns error when no PN is set. - /// This verifies the "not logged in" guard works. - #[tokio::test] - async fn test_establish_primary_phone_session_fails_without_pn() { - let backend = Arc::new( - crate::store::SqliteStore::new("file:memdb_no_pn?mode=memory&cache=shared") - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // No PN set, so this should fail - let result = client.establish_primary_phone_session_immediate().await; - - assert!( - result.is_err(), - "establish_primary_phone_session_immediate should fail when no PN is set" - ); - - let err = result.unwrap_err(); - assert!( - err.downcast_ref::() - .is_some_and(|e| matches!(e, ClientError::NotLoggedIn)), - "Error should be ClientError::NotLoggedIn, got: {}", - err - ); - - info!("✅ test_establish_primary_phone_session_fails_without_pn passed"); - } - - /// Test that ensure_e2e_sessions waits for offline sync to complete. - /// This is the CRITICAL difference between ensure_e2e_sessions and - /// establish_primary_phone_session_immediate. - #[tokio::test] - async fn test_ensure_e2e_sessions_waits_for_offline_sync() { - use std::sync::atomic::Ordering; - use wacore_binary::Jid; - - let backend = Arc::new( - crate::store::SqliteStore::new("file:memdb_ensure_e2e_waits?mode=memory&cache=shared") - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Flag is false (offline sync not complete) - assert!(!client.offline_sync_completed.load(Ordering::Relaxed)); - - // Call ensure_e2e_sessions with an empty list (so it returns early after the wait) - // This lets us test the waiting behavior without needing network - let client_clone = client.clone(); - let ensure_handle = tokio::spawn(async move { - // Start with some JIDs - but since we're testing the wait, we use empty - // to avoid needing actual session establishment - client_clone.ensure_e2e_sessions(&[]).await - }); - - // Wait a bit - ensure_e2e_sessions should return immediately for empty list - tokio::time::sleep(std::time::Duration::from_millis(10)).await; - assert!( - ensure_handle.is_finished(), - "ensure_e2e_sessions should return immediately for empty JID list" - ); - - // Now test with actual JIDs - it should wait for offline sync - let client_clone = client.clone(); - let test_jid = Jid::pn("559999999999"); - let ensure_handle = tokio::spawn(async move { - // This will wait for offline sync before proceeding - let start = wacore::time::Instant::now(); - let _ = client_clone.ensure_e2e_sessions(&[test_jid]).await; - start.elapsed() - }); - - // Give it a moment to start - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - - // It should still be waiting (offline sync not complete) - assert!( - !ensure_handle.is_finished(), - "ensure_e2e_sessions should be waiting for offline sync" - ); - - // Now complete offline sync - client.offline_sync_completed.store(true, Ordering::Relaxed); - client.offline_sync_notifier.notify(usize::MAX); - - // Now it should complete (might fail on session establishment, but that's ok) - let result = tokio::time::timeout(std::time::Duration::from_secs(2), ensure_handle).await; - - assert!( - result.is_ok(), - "ensure_e2e_sessions should complete after offline sync" - ); - - info!("✅ test_ensure_e2e_sessions_waits_for_offline_sync passed"); - } - - /// Integration test: Verify that the immediate session establishment does NOT - /// wait for offline sync. This is critical for PDO to work during offline sync. - /// - /// The flow is: - /// 1. Login -> establish_primary_phone_session_immediate() is called - /// 2. This should NOT wait for offline sync (flag is false at this point) - /// 3. After session is established, offline messages arrive - /// 4. When decryption fails, PDO can immediately send to device 0 - #[tokio::test] - async fn test_immediate_session_does_not_wait_for_offline_sync() { - use std::sync::atomic::Ordering; - use wacore_binary::Jid; - - let backend = Arc::new( - crate::store::SqliteStore::new("file:memdb_immediate_no_wait?mode=memory&cache=shared") - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend.clone()) - .await - .expect("persistence manager should initialize"), - ); - - // Set a PN so establish_primary_phone_session_immediate doesn't fail early - pm.modify_device(|device| { - device.pn = Some(Jid::pn("559999999999")); - }) - .await; - - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Flag is false (offline sync not complete - simulating login state) - assert!(!client.offline_sync_completed.load(Ordering::Relaxed)); - - // Call establish_primary_phone_session_immediate - // It should NOT wait for offline sync - it should proceed immediately - let start = wacore::time::Instant::now(); - - // Note: This will fail because we can't actually fetch prekeys in tests, - // but the important thing is that it doesn't WAIT for offline sync - let result = tokio::time::timeout( - std::time::Duration::from_millis(500), - client.establish_primary_phone_session_immediate(), - ) - .await; - - let elapsed = start.elapsed(); - - // The call should complete (or fail) quickly, NOT wait for 10 second timeout - assert!( - result.is_ok(), - "establish_primary_phone_session_immediate should not wait for offline sync, timed out" - ); - - // It should complete in < 500ms (not 10 second wait) - assert!( - elapsed.as_millis() < 500, - "establish_primary_phone_session_immediate should not wait, took {:?}", - elapsed - ); - - // The actual result might be an error (no network), but that's fine - // The important thing is it didn't wait for offline sync - info!( - "establish_primary_phone_session_immediate completed in {:?} (result: {:?})", - elapsed, - result.unwrap().is_ok() - ); - - info!("✅ test_immediate_session_does_not_wait_for_offline_sync passed"); - } - - /// Integration test: Verify that establish_primary_phone_session_immediate - /// skips establishment when a session already exists. - /// - /// This is the CRITICAL fix for MAC verification failures: - /// - BUG (before fix): Called process_prekey_bundle() unconditionally, - /// replacing the existing session with a new one - /// - RESULT: Remote device still uses old session state, causing MAC failures - #[tokio::test] - async fn test_establish_session_skips_when_exists() { - use wacore::libsignal::protocol::SessionRecord; - use wacore::libsignal::store::SessionStore; - use wacore::types::jid::JidExt; - use wacore_binary::Jid; - - let backend = Arc::new( - crate::store::SqliteStore::new("file:memdb_skip_existing?mode=memory&cache=shared") - .await - .expect("Failed to create in-memory backend for test"), - ); - let pm = Arc::new( - PersistenceManager::new(backend.clone()) - .await - .expect("persistence manager should initialize"), - ); - - // Set a PN so the function doesn't fail early - let own_pn = Jid::pn("559999999999"); - pm.modify_device(|device| { - device.pn = Some(own_pn.clone()); - }) - .await; - - // Pre-populate a session for the primary phone JID (device 0) - let primary_phone_jid = own_pn.with_device(0); - let signal_addr = primary_phone_jid.to_protocol_address(); - - // Create a dummy session record - let dummy_session = SessionRecord::new_fresh(); - { - let device_arc = pm.get_device_arc().await; - let device = device_arc.read().await; - device - .store_session(&signal_addr, &dummy_session) - .await - .expect("Failed to store test session"); - - // Verify session exists - let exists = device - .contains_session(&signal_addr) - .await - .expect("Failed to check session"); - assert!(exists, "Session should exist after store"); - } - - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm.clone(), - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Call establish_primary_phone_session_immediate - // It should return Ok(()) immediately without fetching prekeys - let result = client.establish_primary_phone_session_immediate().await; - - assert!( - result.is_ok(), - "establish_primary_phone_session_immediate should succeed when session exists" - ); - - // Verify the session was NOT replaced (still has the same record) - // This is the critical assertion - if session was replaced, it would cause MAC failures - { - let device_arc = pm.get_device_arc().await; - let device = device_arc.read().await; - let exists = device - .contains_session(&signal_addr) - .await - .expect("Failed to check session"); - assert!(exists, "Session should still exist after the call"); - } - - info!("✅ test_establish_session_skips_when_exists passed"); - } - - /// Integration test: Verify that the session check prevents MAC failures - /// by documenting the exact control flow that caused the bug. - #[test] - fn test_mac_failure_prevention_flow_documentation() { - // Simulate the decision logic - fn should_establish_session( - check_result: Result, - ) -> Result { - match check_result { - Ok(true) => Ok(false), // Session exists → DON'T establish - Ok(false) => Ok(true), // No session → establish - Err(e) => Err(format!("Cannot verify session: {}", e)), // Fail-safe - } - } - - // Test Case 1: Session exists → skip (prevents MAC failure) - let result = should_establish_session(Ok(true)); - assert_eq!(result, Ok(false), "Should skip when session exists"); - - // Test Case 2: No session → establish - let result = should_establish_session(Ok(false)); - assert_eq!(result, Ok(true), "Should establish when no session"); - - // Test Case 3: Check fails → error (fail-safe) - let result = should_establish_session(Err("database error")); - assert!(result.is_err(), "Should fail when check fails"); - - info!("✅ test_mac_failure_prevention_flow_documentation passed"); - } - - #[test] - fn test_unified_session_id_calculation() { - // Test the mathematical calculation of the unified session ID. - // Formula: (now_ms + server_offset_ms + 3_days_ms) % 7_days_ms - - const DAY_MS: i64 = 24 * 60 * 60 * 1000; - const WEEK_MS: i64 = 7 * DAY_MS; - const OFFSET_MS: i64 = 3 * DAY_MS; - - // Helper function matching the implementation - fn calculate_session_id(now_ms: i64, server_offset_ms: i64) -> i64 { - let adjusted_now = now_ms + server_offset_ms; - (adjusted_now + OFFSET_MS) % WEEK_MS - } - - // Test 1: Zero offset - let now_ms = 1706000000000_i64; // Some arbitrary timestamp - let id = calculate_session_id(now_ms, 0); - assert!( - (0..WEEK_MS).contains(&id), - "Session ID should be in [0, WEEK_MS)" - ); - - // Test 2: Positive server offset (server is ahead) - let id_with_positive_offset = calculate_session_id(now_ms, 5000); - assert!( - (0..WEEK_MS).contains(&id_with_positive_offset), - "Session ID should be in [0, WEEK_MS)" - ); - // The ID should be different from zero offset (unless wrap-around) - // Not testing exact value as it depends on the offset - - // Test 3: Negative server offset (server is behind) - let id_with_negative_offset = calculate_session_id(now_ms, -5000); - assert!( - (0..WEEK_MS).contains(&id_with_negative_offset), - "Session ID should be in [0, WEEK_MS)" - ); - - // Test 4: Verify modulo wrap-around - // If adjusted_now + OFFSET_MS >= WEEK_MS, it should wrap - let wrap_test_now = WEEK_MS - OFFSET_MS + 1000; // Should produce small result - let wrapped_id = calculate_session_id(wrap_test_now, 0); - assert_eq!(wrapped_id, 1000, "Should wrap around correctly"); - - // Test 5: Edge case - at exact boundary - let boundary_now = WEEK_MS - OFFSET_MS; - let boundary_id = calculate_session_id(boundary_now, 0); - assert_eq!(boundary_id, 0, "At exact boundary should be 0"); - } - - #[tokio::test] - async fn test_server_time_offset_extraction() { - use wacore_binary::builder::NodeBuilder; - - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Initially, offset should be 0 - assert_eq!( - client.unified_session.server_time_offset_ms(), - 0, - "Initial offset should be 0" - ); - - // Create a node with a 't' attribute - let server_time = wacore::time::now_secs() + 10; // Server is 10 seconds ahead - let node = NodeBuilder::new("success").attr("t", server_time).build(); - - // Update the offset - client.update_server_time_offset(&node.as_node_ref()); - - // The offset should be approximately 10 * 1000 = 10000 ms - // Allow some tolerance for timing differences during the test - let offset = client.unified_session.server_time_offset_ms(); - assert!( - (offset - 10000).abs() < 1000, // Allow 1 second tolerance - "Offset should be approximately 10000ms, got {}", - offset - ); - - // Test with no 't' attribute - should not change offset - let node_no_t = NodeBuilder::new("success").build(); - client.update_server_time_offset(&node_no_t.as_node_ref()); - let offset_after = client.unified_session.server_time_offset_ms(); - assert!( - (offset_after - offset).abs() < 100, // Should be same (or very close) - "Offset should not change when 't' is missing" - ); - - // Test with invalid 't' attribute - should not change offset - let node_invalid = NodeBuilder::new("success") - .attr("t", "not_a_number") - .build(); - client.update_server_time_offset(&node_invalid.as_node_ref()); - let offset_after_invalid = client.unified_session.server_time_offset_ms(); - assert!( - (offset_after_invalid - offset).abs() < 100, - "Offset should not change when 't' is invalid" - ); - - // Test with negative/zero 't' - should not change offset - let node_zero = NodeBuilder::new("success").attr("t", "0").build(); - client.update_server_time_offset(&node_zero.as_node_ref()); - let offset_after_zero = client.unified_session.server_time_offset_ms(); - assert!( - (offset_after_zero - offset).abs() < 100, - "Offset should not change when 't' is 0" - ); - - info!("✅ test_server_time_offset_extraction passed"); - } - - #[tokio::test] - async fn test_unified_session_manager_integration() { - // Test the unified session manager through the client - - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Initially, sequence should be 0 - assert_eq!( - client.unified_session.sequence(), - 0, - "Initial sequence should be 0" - ); - - // Duplicate prevention depends on the session ID staying the same between calls. - // Since the session ID is millisecond-based, use a retry loop to handle - // the rare case where we cross a millisecond boundary between calls. - loop { - client.unified_session.reset().await; - - let result = client.unified_session.prepare_send().await; - assert!(result.is_some(), "First send should succeed"); - let (node, seq) = result.unwrap(); - assert_eq!(node.tag, "ib", "Should be an IB stanza"); - assert_eq!(seq, 1, "First sequence should be 1 (pre-increment)"); - assert_eq!(client.unified_session.sequence(), 1); - - let result2 = client.unified_session.prepare_send().await; - if result2.is_none() { - // Duplicate was prevented within the same millisecond - assert_eq!(client.unified_session.sequence(), 1); - break; - } - // Millisecond boundary crossed, retry - tokio::task::yield_now().await; - } - - // Clear last sent and try again - sequence resets on "new" session ID - client.unified_session.clear_last_sent().await; - let result3 = client.unified_session.prepare_send().await; - assert!(result3.is_some(), "Should succeed after clearing"); - let (_, seq3) = result3.unwrap(); - assert_eq!(seq3, 1, "Sequence resets when session ID changes"); - assert_eq!(client.unified_session.sequence(), 1); - - info!("✅ test_unified_session_manager_integration passed"); - } - - #[test] - fn test_unified_session_protocol_node() { - // Test the type-safe protocol node implementation - use wacore::ib::{IbStanza, UnifiedSession}; - use wacore::protocol::ProtocolNode; - - // Create a unified session - let session = UnifiedSession::new("123456789"); - assert_eq!(session.id, "123456789"); - assert_eq!(session.tag(), "unified_session"); - - // Convert to node - let node = session.into_node(); - assert_eq!(node.tag, "unified_session"); - assert!(node.attrs.get("id").is_some_and(|v| v == "123456789")); - - // Create an IB stanza - let stanza = IbStanza::unified_session(UnifiedSession::new("987654321")); - assert_eq!(stanza.tag(), "ib"); - - // Convert to node and verify structure - let ib_node = stanza.into_node(); - assert_eq!(ib_node.tag, "ib"); - let children = ib_node.children().expect("IB stanza should have children"); - assert_eq!(children.len(), 1); - assert_eq!(children[0].tag, "unified_session"); - assert!( - children[0] - .attrs - .get("id") - .is_some_and(|v| v == "987654321") - ); - - info!("✅ test_unified_session_protocol_node passed"); - } - - fn node_to_owned_ref(node: Node) -> Arc { - crate::test_utils::node_to_owned_ref(&node) - } - - /// Helper to create a test client for offline sync tests - async fn create_offline_sync_test_client() -> Arc { - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - client - } - - /// Regression: a transport disconnect must flush dirty Signal state before - /// clearing the cache, or a just-advanced sender-key chain is lost (forcing - /// a full SKDM re-fanout on the next send). - #[tokio::test] - async fn cleanup_connection_state_flushes_dirty_signal_state() { - use wacore::libsignal::protocol::ProtocolAddress; - let client = create_offline_sync_test_client().await; - - // A dirty identity lives only in the write-back cache until flushed. - let addr = ProtocolAddress::new("5550001000@s.whatsapp.net".to_string(), 1u32.into()); - client.signal_cache.put_identity(&addr, &[7u8; 32]).await; - - client.cleanup_connection_state().await; - - // cleanup cleared the cache, so a hit now can only come from the DB, - // proving the flush ran before the clear. - let device = client.persistence_manager.get_device_arc().await; - let guard = device.read().await; - let persisted = client - .signal_cache - .get_identity(&addr, &*guard.backend) - .await - .expect("get_identity must not error"); - assert!( - persisted.is_some(), - "dirty Signal state must survive a transport disconnect (flush-before-clear)" - ); - } - - /// Same guarantee on the sender-key store, which drives SKDM fanout. - #[tokio::test] - async fn cleanup_connection_state_flushes_dirty_sender_key() { - use wacore::libsignal::protocol::SenderKeyRecord; - use wacore::libsignal::store::sender_key_name::SenderKeyName; - let client = create_offline_sync_test_client().await; - - let name = SenderKeyName::from_parts("group@g.us", "5550001000@s.whatsapp.net:1"); - client - .signal_cache - .put_sender_key(&name, SenderKeyRecord::new_empty()) - .await; - - client.cleanup_connection_state().await; - - let device = client.persistence_manager.get_device_arc().await; - let guard = device.read().await; - let persisted = client - .signal_cache - .get_sender_key(&name, &*guard.backend) - .await - .expect("get_sender_key must not error"); - assert!( - persisted.is_some(), - "dirty sender key must survive a transport disconnect (flush-before-clear)" - ); - } - - /// When the flush itself fails, cleanup must NOT clear the cache, or it would - /// drop the very state the flush was meant to persist. - #[tokio::test] - async fn cleanup_connection_state_keeps_state_when_flush_fails() { - use wacore::libsignal::protocol::{ProtocolAddress, SenderKeyRecord}; - use wacore::libsignal::store::sender_key_name::SenderKeyName; - let client = create_offline_sync_test_client().await; - - // A malformed identity (not 32 bytes) makes flush() error out, standing - // in for a transient backend write failure during cleanup. - let bad = ProtocolAddress::new("5550002000@s.whatsapp.net".to_string(), 1u32.into()); - client.signal_cache.put_identity(&bad, &[0u8; 16]).await; - - // A valid dirty sender key that must not be dropped when the flush fails. - let name = SenderKeyName::from_parts("group@g.us", "5550001000@s.whatsapp.net:1"); - client - .signal_cache - .put_sender_key(&name, SenderKeyRecord::new_empty()) - .await; - - client.cleanup_connection_state().await; - - // flush() failed, so clear() was skipped; the unpersisted sender key - // survives in the write-back cache instead of being dropped. - let device = client.persistence_manager.get_device_arc().await; - let guard = device.read().await; - let persisted = client - .signal_cache - .get_sender_key(&name, &*guard.backend) - .await - .expect("get_sender_key must not error"); - assert!( - persisted.is_some(), - "a flush failure must not drop dirty Signal state" - ); - } - - /// A 403 connect failure is WA Web's REASON_LOCKED: it must surface a logout - /// carrying AccountLocked and disable auto-reconnect (a lock is not transient). - #[tokio::test] - async fn connect_failure_403_dispatches_account_locked_logout() { - use wacore::types::events::ChannelEventHandler; - let client = create_offline_sync_test_client().await; - let (handler, events) = ChannelEventHandler::new(); - client.register_handler(handler); - - // location="rva" is a region routing token and must not change the verdict. - let failure = NodeBuilder::new("failure") - .attr("reason", "403") - .attr("location", "rva") - .build(); - client.handle_connect_failure(&failure.as_node_ref()).await; - - let evt = events - .try_recv() - .expect("403 must dispatch a LoggedOut event"); - match &*evt { - Event::LoggedOut(lo) => { - assert!(lo.on_connect, "403 arrives as a failure-on-connect"); - assert_eq!(lo.reason, ConnectFailureReason::AccountLocked); - } - _ => panic!("expected Event::LoggedOut for reason=403"), - } - assert!( - !client.enable_auto_reconnect.load(Ordering::Relaxed), - "a server-side lock must not auto-reconnect" - ); - } - - #[tokio::test] - async fn delivery_receipt_activity_state_machine() { - let client = create_offline_sync_test_client().await; - assert!( - !client.receipts_are_active(), - "default is inactive (background companion)" - ); - client.mark_receipts_active_on_presence(); - assert!(client.receipts_are_active(), "presence available -> active"); - client.mark_receipts_inactive_on_presence(); - assert!( - !client.receipts_are_active(), - "presence unavailable -> inactive" - ); - client.set_force_active_delivery_receipts(true); - assert!(client.receipts_are_active(), "forced active"); - client.mark_receipts_inactive_on_presence(); - assert!( - client.receipts_are_active(), - "forced (2) survives a presence-unavailable CAS(1,0)" - ); - client.set_force_active_delivery_receipts(false); - assert!(!client.receipts_are_active()); - - // Teardown resets presence-driven active (so it doesn't leak across - // reconnects) but preserves a forced value. - client.mark_receipts_active_on_presence(); - client.cleanup_connection_state().await; - assert!( - !client.receipts_are_active(), - "teardown resets presence-driven active" - ); - client.set_force_active_delivery_receipts(true); - client.cleanup_connection_state().await; - assert!( - client.receipts_are_active(), - "teardown preserves forced active" - ); - } - - #[tokio::test] - async fn test_ib_thread_metadata_does_not_end_sync() { - let client = create_offline_sync_test_client().await; - client - .offline_sync_metrics - .active - .store(true, Ordering::Release); - - let node = NodeBuilder::new("ib") - .children([NodeBuilder::new("thread_metadata") - .children([NodeBuilder::new("item").build()]) - .build()]) - .build(); - - client.process_node(node_to_owned_ref(node)).await; - assert!( - client.offline_sync_metrics.active.load(Ordering::Acquire), - " should NOT end offline sync" - ); - } - - #[tokio::test] - async fn test_ib_edge_routing_does_not_end_sync() { - let client = create_offline_sync_test_client().await; - client - .offline_sync_metrics - .active - .store(true, Ordering::Release); - - let node = NodeBuilder::new("ib") - .children([NodeBuilder::new("edge_routing") - .children([NodeBuilder::new("routing_info") - .bytes(vec![1, 2, 3]) - .build()]) - .build()]) - .build(); - - client.process_node(node_to_owned_ref(node)).await; - assert!( - client.offline_sync_metrics.active.load(Ordering::Acquire), - " should NOT end offline sync" - ); - } - - #[tokio::test] - async fn test_ib_dirty_does_not_end_sync() { - let client = create_offline_sync_test_client().await; - client - .offline_sync_metrics - .active - .store(true, Ordering::Release); - - let node = NodeBuilder::new("ib") - .children([NodeBuilder::new("dirty") - .attr("type", "groups") - .attr("timestamp", "1234") - .build()]) - .build(); - - client.process_node(node_to_owned_ref(node)).await; - assert!( - client.offline_sync_metrics.active.load(Ordering::Acquire), - " should NOT end offline sync" - ); - } - - #[tokio::test] - async fn test_ib_offline_child_ends_sync() { - let client = create_offline_sync_test_client().await; - client - .offline_sync_metrics - .active - .store(true, Ordering::Release); - client - .offline_sync_metrics - .total_messages - .store(301, Ordering::Release); - - let node = NodeBuilder::new("ib") - .children([NodeBuilder::new("offline").attr("count", "301").build()]) - .build(); - - client.process_node(node_to_owned_ref(node)).await; - assert!( - !client.offline_sync_metrics.active.load(Ordering::Acquire), - " should end offline sync" - ); - } - - #[tokio::test] - async fn test_ib_offline_preview_starts_sync() { - let client = create_offline_sync_test_client().await; - - let node = NodeBuilder::new("ib") - .children([NodeBuilder::new("offline_preview") - .attr("count", "301") - .attr("message", "168") - .attr("notification", "62") - .attr("receipt", "68") - .attr("appdata", "0") - .build()]) - .build(); - - client.process_node(node_to_owned_ref(node)).await; - assert!( - client.offline_sync_metrics.active.load(Ordering::Acquire), - "offline_preview with count>0 should activate sync" - ); - assert_eq!( - client - .offline_sync_metrics - .total_messages - .load(Ordering::Acquire), - 301 - ); - } - - #[tokio::test] - async fn test_offline_message_increments_processed() { - let client = create_offline_sync_test_client().await; - client - .offline_sync_metrics - .active - .store(true, Ordering::Release); - client - .offline_sync_metrics - .total_messages - .store(100, Ordering::Release); - - let node = NodeBuilder::new("message") - .attr("offline", "1") - .attr("from", "5551234567@s.whatsapp.net") - .attr("id", "TEST123") - .attr("t", "1772884671") - .attr("type", "text") - .build(); - - client.process_node(node_to_owned_ref(node)).await; - assert_eq!( - client - .offline_sync_metrics - .processed_messages - .load(Ordering::Acquire), - 1, - "offline message should increment processed count" - ); - } - - // --------------------------------------------------------------- - // Server-initiated ping detection tests - // - // The WhatsApp server can send pings in two formats: - // - // 1. Child-element format (legacy/whatsmeow style): - // - // - // - // - // 2. xmlns-attribute format (real WhatsApp Web format): - // - // This is a self-closing tag with NO child elements. - // Verified against captured WhatsApp Web JS (WAWebCommsHandleStanza): - // if (t.xmlns === "urn:xmpp:ping") return wap("iq", { type: "result", to: t.from }); - // - // Both must be recognized and answered with a pong, otherwise the - // server considers the client dead and stops responding to keepalive - // pings — causing a timeout cascade and forced reconnect. - // --------------------------------------------------------------- - - #[tokio::test] - async fn test_handle_iq_ping_with_child_element() { - // Format 1: — the legacy format with a child node. - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - let ping_node = NodeBuilder::new("iq") - .attr("type", "get") - .attr("from", SERVER_JID) - .attr("id", "ping-child-1") - .children([NodeBuilder::new("ping").build()]) - .build(); - - let handled = client.handle_iq(&ping_node.as_node_ref()).await; - assert!( - handled, - "handle_iq must recognize ping with child element" - ); - } - - #[tokio::test] - async fn test_handle_iq_ping_with_xmlns_attribute() { - // Format 2: — the real WhatsApp Web format. - // This is a self-closing IQ with NO children, only an xmlns attribute. - // The server sends this format; failing to respond causes keepalive timeout cascade. - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - let ping_node = NodeBuilder::new("iq") - .attr("type", "get") - .attr("from", SERVER_JID) - .attr("id", "ping-xmlns-1") - .attr("xmlns", "urn:xmpp:ping") - .build(); - - let handled = client.handle_iq(&ping_node.as_node_ref()).await; - assert!( - handled, - "handle_iq must recognize ping with xmlns=\"urn:xmpp:ping\" attribute (no children)" - ); - } - - #[tokio::test] - async fn test_handle_iq_ping_with_both_child_and_xmlns() { - // Edge case: node has BOTH a child AND xmlns="urn:xmpp:ping". - // Should still be handled (OR condition). - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - let ping_node = NodeBuilder::new("iq") - .attr("type", "get") - .attr("from", SERVER_JID) - .attr("id", "ping-both-1") - .attr("xmlns", "urn:xmpp:ping") - .children([NodeBuilder::new("ping").build()]) - .build(); - - let handled = client.handle_iq(&ping_node.as_node_ref()).await; - assert!( - handled, - "handle_iq must handle ping with both child and xmlns" - ); - } - - #[tokio::test] - async fn test_handle_iq_non_ping_returns_false() { - // A type="get" IQ without ping child or xmlns should NOT be handled as ping. - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - let non_ping_node = NodeBuilder::new("iq") - .attr("type", "get") - .attr("from", SERVER_JID) - .attr("id", "not-a-ping") - .attr("xmlns", "some:other:namespace") - .build(); - - let handled = client.handle_iq(&non_ping_node.as_node_ref()).await; - assert!( - !handled, - "handle_iq must NOT treat non-ping xmlns as a ping" - ); - } - - #[tokio::test] - async fn test_handle_iq_ping_wrong_type_returns_false() { - // xmlns="urn:xmpp:ping" but type="result" (not "get") — should NOT be handled as ping. - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - let result_node = NodeBuilder::new("iq") - .attr("type", "result") - .attr("from", SERVER_JID) - .attr("id", "ping-result-1") - .attr("xmlns", "urn:xmpp:ping") - .build(); - - let handled = client.handle_iq(&result_node.as_node_ref()).await; - assert!( - !handled, - "handle_iq must NOT respond to type=\"result\" even with ping xmlns" - ); - } - - // ── build_pong tests ────────────────────────────────────────────── - - #[test] - fn test_build_pong_with_id() { - let pong = build_pong("s.whatsapp.net".to_string(), Some("ping-123")); - assert!( - pong.attrs.get("id").is_some_and(|v| v == "ping-123"), - "pong should include id when server ping has one" - ); - assert!(pong.attrs.get("type").is_some_and(|v| v == "result")); - assert!(pong.attrs.get("to").is_some_and(|v| v == "s.whatsapp.net")); - } - - #[test] - fn test_build_pong_without_id() { - let pong = build_pong("s.whatsapp.net".to_string(), None); - assert!( - !pong.attrs.contains_key("id"), - "pong should NOT include id when server ping has none" - ); - assert!(pong.attrs.get("type").is_some_and(|v| v == "result")); - } - - #[test] - fn test_encrypt_identity_notification_omits_type() { - let node = NodeBuilder::new("notification") - .attr("from", "186303081611421@lid") - .attr("id", "4128735301") - .attr("type", "encrypt") - .children([NodeBuilder::new("identity").build()]) - .build(); - - assert!( - is_encrypt_identity_notification(&node.as_node_ref()), - "identity-change notification ACK must omit type to match WA Web" - ); - } - - #[test] - fn test_device_notification_is_not_encrypt_identity() { - let node = NodeBuilder::new("notification") - .attr("from", "186303081611421@lid") - .attr("id", "269488578") - .attr("type", "devices") - .children([NodeBuilder::new("remove").build()]) - .build(); - - assert!( - !is_encrypt_identity_notification(&node.as_node_ref()), - "device notification is not an encrypt+identity notification" - ); - } - - #[test] - fn test_build_ack_node_for_message_omits_type_includes_from() { - // Whatsmeow: message acks do NOT echo type (node.Tag != "message" guard). - // They DO include `from` with own device PN. - let incoming = NodeBuilder::new("message") - .attr("from", "120363161500776365@g.us") - .attr("id", "A5791A5392EF60E3FB0670098DE010D4") - .attr("type", "text") - .attr("participant", "181531758878822@lid") - .build(); - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) - .expect("message ack should be buildable"); - - assert_eq!(ack.tag, "ack"); - // Use PartialEq on NodeValue — works for both String and Jid variants - // without allocation, so tests don't depend on internal representation. - assert!(ack.attrs.get("class").is_some_and(|v| v == "message")); - assert!( - ack.attrs - .get("to") - .is_some_and(|v| v == "120363161500776365@g.us") - ); - assert!( - ack.attrs - .get("from") - .is_some_and(|v| v == "155500012345:48@s.whatsapp.net") - ); - assert!( - ack.attrs - .get("participant") - .is_some_and(|v| v == "181531758878822@lid") - ); - assert!( - !ack.attrs.contains_key("type"), - "message ACK must NOT echo type (matches whatsmeow behavior)" - ); - } - - #[test] - fn test_build_ack_node_for_identity_change_omits_type_and_from() { - let incoming = NodeBuilder::new("notification") - .attr("from", "186303081611421@lid") - .attr("id", "4128735301") - .attr("type", "encrypt") - .children([NodeBuilder::new("identity").build()]) - .build(); - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) - .expect("notification ack should be buildable"); - - assert!(ack.attrs.get("class").is_some_and(|v| v == "notification")); - assert!( - !ack.attrs.contains_key("type"), - "identity-change notification ACK must omit type" - ); - assert!( - !ack.attrs.contains_key("from"), - "notification ACKs should not include our device PN" - ); - } - - #[test] - fn test_build_ack_node_for_receipt_with_type_echoes_type() { - // Receipt acks should echo the type attribute when present (e.g. "read", "played"). - let incoming = NodeBuilder::new("receipt") - .attr("from", "156535032389744@lid") - .attr("id", "RCPT-WITH-TYPE") - .attr("type", "read") - .build(); - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) - .expect("receipt ack should be buildable"); - - assert!(ack.attrs.get("class").is_some_and(|v| v == "receipt")); - assert!( - ack.attrs.get("type").is_some_and(|v| v == "read"), - "receipt ACK must echo the type attribute when present" - ); - assert!( - !ack.attrs.contains_key("from"), - "receipt ACKs should not include our device PN" - ); - } - - #[test] - fn test_build_ack_node_drops_participant_when_equal_to_from() { - // WAWebReceiptAck: `participant: r && r !== e ? DEVICE_JID(r) : DROP_ATTR`. - // When the incoming stanza carries participant == from (redundant), - // the ack must not echo it. - let incoming = NodeBuilder::new("receipt") - .attr("from", "156535032389744@lid") - .attr("participant", "156535032389744@lid") - .attr("id", "RCPT-PARTICIPANT-EQ-FROM") - .build(); - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net".parse().unwrap(); - - let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) - .expect("ack should build"); - assert!( - !ack.attrs.contains_key("participant"), - "ack must drop participant when it duplicates `to` (the flipped from); got {:?}", - ack.attrs.get("participant") - ); - } - - #[test] - fn test_build_ack_node_keeps_participant_when_distinct_from_from() { - // Group receipt: participant = sender (user), from = group jid; must be kept. - let incoming = NodeBuilder::new("receipt") - .attr("from", "120363098765432100@g.us") - .attr("participant", "5511999999999@s.whatsapp.net") - .attr("id", "RCPT-GROUP") - .build(); - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net".parse().unwrap(); - - let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) - .expect("ack should build"); - assert!( - ack.attrs - .get("participant") - .is_some_and(|v| v == "5511999999999@s.whatsapp.net"), - "ack must keep participant when it differs from `to`" - ); - } - - #[test] - fn test_build_ack_node_for_receipt_without_type_omits_type() { - // Delivery receipts have no type attribute — the ack must also omit it. - // Sending type="delivery" in the ack causes stream:error disconnections. - let incoming = NodeBuilder::new("receipt") - .attr("from", "156535032389744@lid") - .attr("id", "RCPT-NO-TYPE") - .build(); - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) - .expect("receipt ack should be buildable"); - - assert!(ack.attrs.get("class").is_some_and(|v| v == "receipt")); - assert!( - !ack.attrs.contains_key("type"), - "receipt ACK must NOT contain type when the incoming receipt has no type attribute" - ); - assert!( - !ack.attrs.contains_key("from"), - "receipt ACKs should not include our device PN" - ); - } - - #[test] - fn test_build_ack_node_for_message_with_recipient_preserves_recipient() { - // Peer / hosted-companion / LID-routed messages carry `recipient`. - // The server uses it to route the ack back to the origin device; - // without it the stream is torn down with . - let incoming = NodeBuilder::new("message") - .attr("from", "166361967902821@lid") - .attr("id", "2A32F960553696093D99") - .attr("type", "text") - .attr("recipient", "146991363395800@lid") - .build(); - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) - .expect("message ack should be buildable"); - - assert!(ack.attrs.get("class").is_some_and(|v| v == "message")); - assert!( - ack.attrs - .get("recipient") - .is_some_and(|v| v == "146991363395800@lid"), - "message ACK must echo the incoming `recipient` attribute" - ); - } - - #[test] - fn test_build_ack_node_for_receipt_with_recipient_preserves_recipient() { - // Receipt acks must also echo `recipient` when the incoming carries it. - let incoming = NodeBuilder::new("receipt") - .attr("from", "120363098765432100@g.us") - .attr("id", "RCPT-WITH-RECIPIENT") - .attr("type", "read") - .attr("recipient", "242395589390497@lid") - .build(); - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) - .expect("receipt ack should be buildable"); - - assert!(ack.attrs.get("class").is_some_and(|v| v == "receipt")); - assert!( - ack.attrs - .get("recipient") - .is_some_and(|v| v == "242395589390497@lid"), - "receipt ACK must echo the incoming `recipient` attribute" - ); - } - - #[test] - fn test_build_ack_node_for_message_without_recipient_omits_recipient() { - // Regression guard: never synthesise a `recipient` field if the - // incoming stanza did not carry one — server would reject the ack. - let incoming = NodeBuilder::new("message") - .attr("from", "120363161500776365@g.us") - .attr("id", "A5791A5392EF60E3FB06") - .attr("type", "text") - .attr("participant", "181531758878822@lid") - .build(); - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) - .expect("message ack should be buildable"); - - assert!( - !ack.attrs.contains_key("recipient"), - "ACK must NOT add `recipient` when the incoming stanza has none" - ); - } - - #[test] - fn test_encode_ack_bytes_roundtrip_recipient() { - // Exercises the real wire encoder (`encode_ack_bytes`), not just the - // `build_ack_node` test mirror: serialize, decode the bytes back, and - // assert the parsed ACK echoes `recipient` when present and omits it - // when absent. Guards against the two builders silently diverging. - let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let with_recipient = NodeBuilder::new("message") - .attr("from", "166361967902821@lid") - .attr("id", "2A32F960553696093D99") - .attr("type", "text") - .attr("recipient", "146991363395800@lid") - .build(); - let buf = encode_ack_bytes(&with_recipient.as_node_ref(), Some(&own_device_pn)) - .expect("encode_ack_bytes should not error") - .expect("encode_ack_bytes should produce bytes"); - // The Encoder prepends a leading format byte (see `marshal`); the - // decoder wants raw protocol bytes — same handling as `node_to_owned_ref`. - let decoded = - wacore_binary::marshal::unmarshal_ref(&buf[1..]).expect("encoded ack should decode"); - assert_eq!(decoded.tag, "ack"); - assert!( - decoded - .get_attr("class") - .is_some_and(|v| v.as_str() == "message"), - "decoded ack must have class=message" - ); - assert!( - decoded - .get_attr("recipient") - .is_some_and(|v| v.as_str() == "146991363395800@lid"), - "encode_ack_bytes must echo `recipient` onto the wire" - ); - - let without_recipient = NodeBuilder::new("message") - .attr("from", "120363161500776365@g.us") - .attr("id", "A5791A5392EF60E3FB06") - .attr("type", "text") - .attr("participant", "181531758878822@lid") - .build(); - let buf = encode_ack_bytes(&without_recipient.as_node_ref(), Some(&own_device_pn)) - .expect("encode_ack_bytes should not error") - .expect("encode_ack_bytes should produce bytes"); - let decoded = - wacore_binary::marshal::unmarshal_ref(&buf[1..]).expect("encoded ack should decode"); - assert!( - decoded.get_attr("recipient").is_none(), - "encode_ack_bytes must not synthesise `recipient` when absent" - ); - } - - /// Own-account fan-out ack must address back to the original `from` (own - /// LID) echoing `recipient`, not to the chat. Guards against regressing to - /// the chat-addressed `build_nack_node` style. - #[test] - fn test_message_ack_source_node_own_device_addressing() { - use crate::types::message::{MessageInfo, MessageSource}; - // Own-account branch: sender == `from` (device-qualified), chat is the - // device-stripped recipient. `to` must come from sender, not chat. - let info = MessageInfo { - id: "AC055553E56A2C12DE592DAD6353C477".to_string(), - source: MessageSource { - sender: "236395184570386@lid".parse().expect("sender"), - chat: "156535032389744@lid".parse().expect("chat"), - recipient: Some("156535032389744@lid".parse().expect("recipient")), - is_group: false, - ..Default::default() - }, - ..Default::default() - }; - let own_device_pn: Jid = "559984726662:95@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let source = message_ack_source_node(&info); - let built = build_ack_node(&source.as_node_ref(), Some(&own_device_pn)) - .expect("message ack should be buildable"); - - assert!(built.attrs.get("class").is_some_and(|v| v == "message")); - assert!( - built - .attrs - .get("to") - .is_some_and(|v| v == "236395184570386@lid"), - "ack `to` must be the original `from` (own LID), not the chat" - ); - assert!( - built - .attrs - .get("recipient") - .is_some_and(|v| v == "156535032389744@lid"), - "ack must echo `recipient` so the server can route/clear it" - ); - assert!( - !built.attrs.contains_key("type"), - "message-class acks never carry a `type`" - ); - } - - /// Common incoming DM from another user: `to` is the device-qualified - /// sender, with no `recipient`/`participant` synthesised. - #[test] - fn test_message_ack_source_node_incoming_dm_addressing() { - use crate::types::message::{MessageInfo, MessageSource}; - let info = MessageInfo { - id: "MSGID".to_string(), - source: MessageSource { - sender: "5511999998888:3@s.whatsapp.net".parse().expect("sender"), - chat: "5511999998888@s.whatsapp.net".parse().expect("chat"), - is_group: false, - ..Default::default() - }, - ..Default::default() - }; - let own_device_pn: Jid = "559984726662:95@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let source = message_ack_source_node(&info); - let built = build_ack_node(&source.as_node_ref(), Some(&own_device_pn)) - .expect("dm ack should be buildable"); - - assert!( - built - .attrs - .get("to") - .is_some_and(|v| v == "5511999998888:3@s.whatsapp.net"), - "ack `to` must be the device-qualified sender (the original `from`)" - ); - assert!(!built.attrs.contains_key("recipient")); - assert!(!built.attrs.contains_key("participant")); - } - - /// status@broadcast (is_group=true in the parser) addresses the ack to the - /// status chat, with the sender as participant, not to the sender. - #[test] - fn test_message_ack_source_node_status_addressing() { - use crate::types::message::{MessageInfo, MessageSource}; - let info = MessageInfo { - id: "STATUSMSG".to_string(), - source: MessageSource { - chat: "status@broadcast".parse().expect("status chat"), - sender: "181531758878822@lid".parse().expect("participant"), - is_group: true, - ..Default::default() - }, - ..Default::default() - }; - let own_device_pn: Jid = "559984726662:95@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let source = message_ack_source_node(&info); - let built = build_ack_node(&source.as_node_ref(), Some(&own_device_pn)) - .expect("status ack should be buildable"); - - assert!( - built - .attrs - .get("to") - .is_some_and(|v| v == "status@broadcast"), - "status ack `to` must be the status chat, not the sender" - ); - assert!( - built - .attrs - .get("participant") - .is_some_and(|v| v == "181531758878822@lid"), - "status ack must preserve the sending participant" - ); - } - - /// Group failure ack: `to` is the group, `participant` is preserved. - #[test] - fn test_message_ack_source_node_group_addressing() { - use crate::types::message::{MessageInfo, MessageSource}; - // Group branch: chat == group `from`, sender == participant. - let info = MessageInfo { - id: "GROUPMSGID".to_string(), - source: MessageSource { - chat: "120363011111111111@g.us".parse().expect("group"), - sender: "181531758878822@lid".parse().expect("participant"), - is_group: true, - ..Default::default() - }, - ..Default::default() - }; - let own_device_pn: Jid = "559984726662:95@s.whatsapp.net" - .parse() - .expect("own device PN JID should parse"); - - let source = message_ack_source_node(&info); - let built = build_ack_node(&source.as_node_ref(), Some(&own_device_pn)) - .expect("group message ack should be buildable"); - - assert!( - built - .attrs - .get("to") - .is_some_and(|v| v == "120363011111111111@g.us"), - "group ack `to` must be the group JID" - ); - assert!( - built - .attrs - .get("participant") - .is_some_and(|v| v == "181531758878822@lid"), - "group ack must preserve the sending `participant`" - ); - } - - /// Smoke test: server ping with xmlns but no id attribute is handled. - #[tokio::test] - async fn test_handle_iq_ping_without_id() { - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Server ping without id — real format observed in production logs - let ping_node = NodeBuilder::new("iq") - .attr("type", "get") - .attr("from", SERVER_JID) - .attr("xmlns", "urn:xmpp:ping") - .build(); - - let handled = client.handle_iq(&ping_node.as_node_ref()).await; - assert!( - handled, - "handle_iq must recognize ping without id attribute" - ); - } - - // ── fibonacci_backoff tests ──────────────────────────────────────── - - #[test] - fn test_fibonacci_backoff_sequence() { - // WA Web: first=1000, second=1000 → 1,1,2,3,5,8,13,21,34,55,89,144...s - // We test base values without jitter by checking the range (±10%). - let expected_base_ms = [1000, 1000, 2000, 3000, 5000, 8000, 13000, 21000]; - for (attempt, &base) in expected_base_ms.iter().enumerate() { - let delay = fibonacci_backoff(attempt as u32); - let ms = delay.as_millis() as u64; - let low = base - base / 10; - let high = base + base / 10; - assert!( - ms >= low && ms <= high, - "attempt {attempt}: expected {low}..={high}ms, got {ms}ms" - ); - } - } - - #[test] - fn test_fibonacci_backoff_max_900s() { - // After many attempts, should cap at 900s (±10%) - let delay = fibonacci_backoff(100); - let ms = delay.as_millis() as u64; - assert!( - ms <= 990_000, - "should never exceed 900s + 10% jitter, got {ms}ms" - ); - assert!( - ms >= 810_000, - "should be at least 900s - 10% jitter, got {ms}ms" - ); - } - - #[test] - fn test_fibonacci_backoff_first_attempt_is_1s() { - let delay = fibonacci_backoff(0); - let ms = delay.as_millis() as u64; - assert!( - (900..=1100).contains(&ms), - "first attempt should be ~1s (±10%), got {ms}ms" - ); - } - - // ── stream error tests ───────────────────────────────────────────── - - #[tokio::test] - async fn test_stream_error_401_disables_reconnect() { - let client = create_offline_sync_test_client().await; - let node = NodeBuilder::new("stream:error").attr("code", "401").build(); - client.handle_stream_error(&node.as_node_ref()).await; - assert!( - !client.enable_auto_reconnect.load(Ordering::Relaxed), - "401 should disable auto-reconnect" - ); - } - - #[tokio::test] - async fn test_stream_error_409_disables_reconnect() { - let client = create_offline_sync_test_client().await; - let node = NodeBuilder::new("stream:error").attr("code", "409").build(); - client.handle_stream_error(&node.as_node_ref()).await; - assert!( - !client.enable_auto_reconnect.load(Ordering::Relaxed), - "409 should disable auto-reconnect" - ); - } - - #[tokio::test] - async fn test_stream_error_429_keeps_reconnect_with_backoff() { - let client = create_offline_sync_test_client().await; - client.is_logged_in.store(true, Ordering::Relaxed); - let before = client.auto_reconnect_errors.load(Ordering::Relaxed); - let node = NodeBuilder::new("stream:error").attr("code", "429").build(); - client.handle_stream_error(&node.as_node_ref()).await; - assert!( - client.enable_auto_reconnect.load(Ordering::Relaxed), - "429 should keep auto-reconnect enabled" - ); - assert!( - !client.is_logged_in.load(Ordering::Relaxed), - "429 must clear is_logged_in so sends bail before the server flags abuse" - ); - assert!( - !client.expected_disconnect.load(Ordering::Relaxed), - "429 must not mark the disconnect as expected (auto-reconnect path)" - ); - let after = client.auto_reconnect_errors.load(Ordering::Relaxed); - assert_eq!( - after, - before + 5, - "429 should increase backoff by exactly 5: before={before}, after={after}" - ); - } - - #[tokio::test] - async fn test_stream_error_503_keeps_reconnect() { - let client = create_offline_sync_test_client().await; - client.is_logged_in.store(true, Ordering::Relaxed); - let node = NodeBuilder::new("stream:error").attr("code", "503").build(); - client.handle_stream_error(&node.as_node_ref()).await; - assert!( - client.enable_auto_reconnect.load(Ordering::Relaxed), - "503 should keep auto-reconnect enabled" - ); - assert!( - !client.is_logged_in.load(Ordering::Relaxed), - "503 must clear is_logged_in so sends bail against the dying socket" - ); - assert!( - !client.expected_disconnect.load(Ordering::Relaxed), - "503 must not mark the disconnect as expected (auto-reconnect path)" - ); - } - - #[tokio::test] - async fn test_stream_error_unknown_keeps_connection_alive() { - // Unknown stream:error (no `code` attribute) must mirror whatsmeow's - // default branch: log + dispatch event, but NOT mark this as an - // expected disconnect. Setting that flag silently swallows the next - // real disconnect and races the read loop into shutdown. - let client = create_offline_sync_test_client().await; - // Simulate an authenticated session before the stream error arrives. - client.is_logged_in.store(true, Ordering::Relaxed); - let node = NodeBuilder::new("stream:error").build(); - client.handle_stream_error(&node.as_node_ref()).await; - assert!( - client.is_logged_in.load(Ordering::Relaxed), - "unknown stream:error must NOT log the client out" - ); - assert!( - !client.expected_disconnect.load(Ordering::Relaxed), - "unknown stream:error must not mark the disconnect as expected" - ); - assert!( - client.enable_auto_reconnect.load(Ordering::Relaxed), - "unknown stream:error must keep auto-reconnect enabled" - ); - } - - #[tokio::test] - async fn test_stream_error_ack_shaped_does_not_force_shutdown() { - // Server wraps per-stanza routing failures in `` - // with no `code` attribute. Treat as informational, not as a fatal - // stream teardown. - let client = create_offline_sync_test_client().await; - client.is_logged_in.store(true, Ordering::Relaxed); - let ack_child = NodeBuilder::new("ack") - .attr("class", "message") - .attr("type", "text") - .attr("id", "2A32F960553696093D99") - .build(); - let node = NodeBuilder::new("stream:error") - .children([ack_child]) - .build(); - client.handle_stream_error(&node.as_node_ref()).await; - assert!( - client.is_logged_in.load(Ordering::Relaxed), - "ack-shaped stream:error must NOT log the client out" - ); - assert!( - !client.expected_disconnect.load(Ordering::Relaxed), - "ack-shaped stream:error must not mark the disconnect as expected" - ); - } - - #[tokio::test] - async fn test_custom_cache_config_is_respected() { - use crate::cache_config::{CacheConfig, CacheEntryConfig}; - use std::time::Duration; - - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - - let custom_config = CacheConfig { - group_cache: CacheEntryConfig::new(Some(Duration::from_secs(60)), 10), - device_registry_cache: CacheEntryConfig::new(Some(Duration::from_secs(60)), 10), - ..CacheConfig::default() - }; - - // Verify that constructing a client with a custom config does not panic - // and the client is usable. - let (client, _rx) = Client::new_with_cache_config( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - custom_config, - ) - .await; - - assert!(!client.is_logged_in()); - } - - /// Proves that `is_connected()` no longer gives false negatives under mutex - /// contention. Before the fix, `try_lock()` would fail when another task held - /// the noise_socket mutex, causing `is_connected()` to return `false` even - /// though the connection was alive — silently dropping receipt acks. - /// - /// This test sets up a real NoiseSocket (same as socket unit tests) so it - /// accurately models the pre-fix scenario: socket is Some + mutex is held - /// by another task = old is_connected() returned false. - #[tokio::test] - async fn test_is_connected_not_affected_by_mutex_contention() { - use crate::socket::NoiseSocket; - use wacore::handshake::NoiseCipher; - - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Initially not connected - assert!(!client.is_connected(), "should start disconnected"); - - // Simulate a real connection: create a NoiseSocket and store it - let transport: Arc = - Arc::new(crate::transport::mock::MockTransport); - let key = [0u8; 32]; - let write_key = NoiseCipher::new(&key).expect("valid key"); - let read_key = NoiseCipher::new(&key).expect("valid key"); - let noise_socket = NoiseSocket::new( - Arc::new(crate::runtime_impl::TokioRuntime), - transport, - write_key, - read_key, - ); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); - client.is_connected.store(true, Ordering::Release); - - assert!(client.is_connected(), "should report connected"); - - // Hold the noise_socket mutex — this used to make is_connected() return - // false via try_lock() even though the socket was Some(...) - let _guard = client.noise_socket.lock().await; - assert!( - client.is_connected(), - "is_connected() must return true even while noise_socket mutex is held" - ); - } - - #[tokio::test] - async fn disconnect_does_not_signal_connection_cleanup_before_outbound_flush() { - use crate::socket::NoiseSocket; - use async_trait::async_trait; - use bytes::Bytes; - use wacore::handshake::NoiseCipher; - - struct BlockingTransport { - send_started: async_channel::Sender<()>, - release_send: async_channel::Receiver<()>, - send_done: Arc, - disconnect_called: Arc, - disconnect_before_send_done: Arc, - } - - #[async_trait] - impl crate::transport::Transport for BlockingTransport { - async fn send(&self, _data: Bytes) -> Result<(), anyhow::Error> { - let _ = self.send_started.try_send(()); - let _ = self.release_send.recv().await; - self.send_done.store(true, Ordering::Release); - Ok(()) - } - - async fn disconnect(&self) { - if !self.send_done.load(Ordering::Acquire) { - self.disconnect_before_send_done - .store(true, Ordering::Release); - } - self.disconnect_called.store(true, Ordering::Release); - } - } - - let client = crate::test_utils::create_test_client().await; - let (send_started_tx, send_started_rx) = async_channel::bounded(1); - let (release_send_tx, release_send_rx) = async_channel::bounded(1); - let send_done = Arc::new(AtomicBool::new(false)); - let disconnect_called = Arc::new(AtomicBool::new(false)); - let disconnect_before_send_done = Arc::new(AtomicBool::new(false)); - - let transport_impl = Arc::new(BlockingTransport { - send_started: send_started_tx, - release_send: release_send_rx, - send_done: Arc::clone(&send_done), - disconnect_called: Arc::clone(&disconnect_called), - disconnect_before_send_done: Arc::clone(&disconnect_before_send_done), - }); - let transport: Arc = transport_impl; - - let key = [0u8; 32]; - let write_key = NoiseCipher::new(&key).expect("valid key"); - let read_key = NoiseCipher::new(&key).expect("valid key"); - let noise_socket = NoiseSocket::new( - client.runtime.clone(), - Arc::clone(&transport), - write_key, - read_key, - ); - - *client.transport.lock().await = Some(transport); - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); - client.is_connected.store(true, Ordering::Release); - - let cleanup_signal = client.connection_shutdown_signal(); - let cleanup_client = Arc::clone(&client); - let cleanup_task = tokio::spawn(async move { - wacore::runtime::wait_for_shutdown(&cleanup_signal).await; - cleanup_client.cleanup_connection_state().await; - }); - - let send_client = Arc::clone(&client); - client.outbound_flush.spawn(&*client.runtime, async move { - let receipt = NodeBuilder::new("receipt") - .attr("id", "TEST-FLUSH-ORDER") - .attr("to", "1234567890@s.whatsapp.net") - .build(); - let _ = send_client.send_node(receipt).await; - }); - - tokio::time::timeout(Duration::from_secs(1), send_started_rx.recv()) - .await - .expect("tracked send should start") - .expect("send_started sender should stay open"); - - let disconnect_client = Arc::clone(&client); - let disconnect_task = tokio::spawn(async move { - disconnect_client.disconnect().await; - }); - - tokio::time::sleep(Duration::from_millis(50)).await; - assert!( - !client.connection_shutdown_signal().is_fired(), - "connection cleanup must not fire while outbound flush is blocked" - ); - assert!( - !disconnect_called.load(Ordering::Acquire), - "transport must stay open while outbound flush is blocked" - ); - - release_send_tx - .send(()) - .await - .expect("blocked send should still be waiting"); - - tokio::time::timeout(Duration::from_secs(1), disconnect_task) - .await - .expect("disconnect should finish") - .expect("disconnect task should not panic"); - tokio::time::timeout(Duration::from_secs(1), cleanup_task) - .await - .expect("cleanup should finish") - .expect("cleanup task should not panic"); - - assert!(send_done.load(Ordering::Acquire)); - assert!(disconnect_called.load(Ordering::Acquire)); - assert!( - !disconnect_before_send_done.load(Ordering::Acquire), - "cleanup closed the transport before the tracked send completed" - ); - } - - /// Verifies that `send_ack_for` returns an error (not silent Ok) when - /// disconnected. This ensures the caller's `warn!` fires so dropped acks - /// are visible in logs. - #[tokio::test] - async fn test_send_ack_for_returns_error_when_disconnected() { - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Not connected — send_ack_for should return Err, not Ok - let receipt = NodeBuilder::new("receipt") - .attr("from", "120363040237990503@g.us") - .attr("id", "TEST-RECEIPT-ID") - .attr("participant", "236395184570386@lid") - .build(); - - let result = client.send_ack_for(&receipt.as_node_ref()).await; - assert!( - matches!(result, Err(ClientError::NotConnected)), - "send_ack_for must return Err(NotConnected) when disconnected, got: {result:?}" - ); - } - - /// Verifies that `send_ack_for` returns Ok when expected_disconnect is set, - /// since this is an intentional shutdown path. - #[tokio::test] - async fn test_send_ack_for_returns_ok_on_expected_disconnect() { - let backend = crate::test_utils::create_test_backend().await; - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(crate::transport::mock::MockTransportFactory::new()), - Arc::new(MockHttpClient), - None, - ) - .await; - - // Set expected disconnect — send_ack_for should gracefully return Ok - client.expected_disconnect.store(true, Ordering::Relaxed); - - let receipt = NodeBuilder::new("receipt") - .attr("from", "120363040237990503@g.us") - .attr("id", "TEST-RECEIPT-ID") - .build(); - - let result = client.send_ack_for(&receipt.as_node_ref()).await; - assert!( - result.is_ok(), - "send_ack_for should return Ok during expected disconnect" - ); - } - - // Per-connection notify must NOT set the terminal sticky flag; if it did, - // every reconnect would instantly abort subscribers registered on the - // terminal signal. Regression guard for the CI breakage observed on PR #560. - #[tokio::test] - async fn per_connection_notify_leaves_terminal_signal_untouched() { - let client = crate::test_utils::create_test_client().await; - - client.notify_connection_shutdown(); - - assert!( - !client.shutdown_signal().is_fired(), - "terminal shutdown must stay clean when only per-connection fires" - ); - } - - // Subscribers registered AFTER a reset must not see the previous - // notifier's fired state. This is the core property that makes reconnect - // work: after cleanup_connection_state notifies the per-connection - // signal, the next connection replaces it with a fresh one. - #[tokio::test] - async fn reset_gives_fresh_per_connection_notifier() { - let client = crate::test_utils::create_test_client().await; - - client.notify_connection_shutdown(); - assert!( - client.connection_shutdown_signal().is_fired(), - "subscriber BEFORE reset sees the notify on the current notifier" - ); - - client.reset_connection_shutdown(); - - assert!( - !client.connection_shutdown_signal().is_fired(), - "subscribers AFTER reset must NOT see the previous notifier's state" - ); - } - - // Capture-once regression guard: a ShutdownSignal captured before a reset - // must keep observing the pre-reset fired state. Without this, a - // reconnect after the old notifier is replaced in the Mutex would - // strand long-lived tasks (e.g. keepalive) on a new notifier they - // never registered for. See keepalive_loop which captures its signal - // once at task startup. - #[tokio::test] - async fn captured_signal_keeps_observing_old_notifier_after_reset() { - let client = crate::test_utils::create_test_client().await; - - let captured = client.connection_shutdown_signal(); - client.notify_connection_shutdown(); - client.reset_connection_shutdown(); - - assert!( - captured.is_fired(), - "captured signal must retain the pre-reset notifier's fired state" - ); - } - - // Terminal disconnect() must also wake per-connection subscribers via - // cleanup_connection_state, so keepalive/request/read loop exit promptly. - #[tokio::test] - async fn terminal_disconnect_propagates_to_per_connection_signal() { - let client = crate::test_utils::create_test_client().await; - let conn_signal = client.connection_shutdown_signal(); - - client.disconnect().await; - - assert!( - conn_signal.is_fired(), - "disconnect must fire per-connection via cleanup_connection_state" - ); - assert!( - client.shutdown_signal().is_fired(), - "disconnect must also fire terminal" - ); - } + // ±10% jitter (WA Web: jitter: 0.1) + let jitter_range = base / 10; + let jitter = if jitter_range > 0 { + rand::make_rng::().random_range(0..=(jitter_range * 2)) as i64 + - jitter_range as i64 + } else { + 0 + }; + let ms = (base as i64 + jitter).max(0) as u64; + Duration::from_millis(ms) } + +#[cfg(test)] +mod tests; diff --git a/src/client/accessors.rs b/src/client/accessors.rs new file mode 100644 index 000000000..1fd5b1742 --- /dev/null +++ b/src/client/accessors.rs @@ -0,0 +1,360 @@ +//! Small accessors, config setters, node waiters and sync-error helpers. + +use super::*; + +impl Client { + pub(crate) async fn get_group_cache(&self) -> Arc { + let mut guard = self.group_cache.lock().await; + if let Some(cache) = guard.as_ref() { + return cache.clone(); + } + debug!("Initializing Group Cache for the first time."); + let cache = Arc::new( + self.cache_config + .group_cache + .build_typed_ttl(self.cache_config.cache_stores.group_cache.clone(), "group"), + ); + *guard = Some(cache.clone()); + cache + } + + /// Registers an external event handler to the core event bus. + pub fn register_handler(&self, handler: Arc) { + self.core.event_bus.add_handler(handler); + } + + /// Enable or disable raw node forwarding. + /// When enabled, `Event::RawNode` is emitted for every decoded stanza before + /// the stanza router dispatches it. Only enable when external consumers need + /// raw protocol access (e.g. voice call stanzas). + pub fn set_raw_node_forwarding(&self, enabled: bool) { + self.raw_node_forwarding.store(enabled, Ordering::Relaxed); + } + + /// Enable or disable skipping of history sync notifications at runtime. + /// + /// When enabled, the client will acknowledge incoming history sync + /// notifications but will not download or process the data. + pub fn set_skip_history_sync(&self, enabled: bool) { + self.skip_history_sync.store(enabled, Ordering::Relaxed); + } + + /// Returns `true` if history sync notifications are currently being skipped. + pub fn skip_history_sync_enabled(&self) -> bool { + self.skip_history_sync.load(Ordering::Relaxed) + } + + /// Set how many one-time pre-keys are generated per upload batch. + /// + /// Defaults to WA Web's UPLOAD_KEYS_COUNT (812). Call before connecting; it + /// takes effect on the next pre-key upload. The value is clamped to the + /// protocol-safe range at upload time, so out-of-range values are coerced + /// (and logged) rather than rejected here. + pub fn set_wanted_pre_key_count(&self, count: usize) { + self.wanted_pre_key_count.store(count, Ordering::Relaxed); + } + + /// Returns the configured pre-key upload batch size (the raw value, before + /// the upload-time clamp). + pub fn wanted_pre_key_count(&self) -> usize { + self.wanted_pre_key_count.load(Ordering::Relaxed) + } + + /// Returns a snapshot of all internal collection sizes for memory leak detection. + /// + /// Moka caches report approximate counts (pending evictions may not be reflected). + /// Call `run_pending_tasks()` on individual caches first if you need exact counts. + /// + /// Requires the `debug-diagnostics` feature. + #[cfg(feature = "debug-diagnostics")] + pub async fn memory_diagnostics(&self) -> MemoryDiagnostics { + let (sig_sessions, sig_identities, sig_sender_keys) = + self.signal_cache.entry_counts().await; + let (lid_lid, lid_pn) = self.lid_pn_cache.entry_counts(); + let pending_retries_count = self + .pending_retries + .lock() + .unwrap_or_else(|p| p.into_inner()) + .len(); + + MemoryDiagnostics { + group_cache: self + .group_cache + .lock() + .await + .as_ref() + .map_or(0, |c| c.entry_count()), + device_registry_cache: self.device_registry_cache.entry_count(), + lid_pn_lid_entries: lid_lid, + lid_pn_pn_entries: lid_pn, + recent_messages: self.recent_messages.entry_count(), + sender_key_device_cache: self.sender_key_device_cache.entry_count(), + message_retry_counts: self.message_retry_counts.entry_count(), + undecryptable_dispatched: self.undecryptable_dispatched.entry_count(), + pdo_pending_requests: self.pdo_pending_requests.entry_count(), + session_locks: self.session_locks.entry_count(), + chat_lanes: self.chat_lanes.entry_count(), + response_waiters: self.response_waiters.lock().await.len(), + node_waiters: self.node_waiter_count.load(Ordering::Relaxed), + pending_retries: pending_retries_count, + presence_subscriptions: self.presence_subscriptions.lock().await.len(), + app_state_key_requests: self.app_state_key_requests.lock().await.len(), + app_state_syncing: self.app_state_syncing.lock().await.len(), + signal_cache_sessions: sig_sessions, + signal_cache_identities: sig_identities, + signal_cache_sender_keys: sig_sender_keys, + chatstate_handlers: self.chatstate_handlers.read().await.len(), + custom_enc_handlers: self.custom_enc_handlers.read().await.len(), + } + } + + /// Get access to the PersistenceManager for this client. + /// This is useful for multi-account scenarios to get the device ID. + pub fn persistence_manager(&self) -> Arc { + self.persistence_manager.clone() + } + + pub async fn get_push_name(&self) -> String { + self.persistence_manager + .get_device_arc() + .await + .read() + .await + .push_name + .clone() + } + + pub async fn get_pn(&self) -> Option { + self.persistence_manager + .get_device_arc() + .await + .read() + .await + .pn + .clone() + } + + pub async fn get_lid(&self) -> Option { + self.persistence_manager + .get_device_arc() + .await + .read() + .await + .lid + .clone() + } + + pub(crate) async fn require_pn(&self) -> Result { + self.get_pn().await.ok_or(ClientError::NotLoggedIn.into()) + } + + /// Resolve our own JID for a group, respecting its addressing mode. + /// + /// Returns LID for LID-addressing groups, PN otherwise. + /// Matches WhatsApp Web's `getMeUserLidOrJidForChat`. + pub(crate) async fn get_own_jid_for_group( + &self, + group_jid: &Jid, + ) -> Result { + let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let own_pn = device_snapshot + .pn + .clone() + .ok_or_else(|| anyhow::Error::from(ClientError::NotLoggedIn))?; + + let addressing_mode = self + .groups() + .query_info(group_jid) + .await + .map(|info| info.addressing_mode) + .unwrap_or(crate::types::message::AddressingMode::Pn); + + Ok(match addressing_mode { + crate::types::message::AddressingMode::Lid => { + device_snapshot.lid.clone().unwrap_or(own_pn) + } + crate::types::message::AddressingMode::Pn => own_pn, + }) + } + + pub(crate) async fn update_push_name_and_notify(self: &Arc, new_name: String) { + let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let old_name = device_snapshot.push_name.clone(); + + if old_name == new_name { + return; + } + + log::debug!("Updating push name from '{}' -> '{}'", old_name, new_name); + self.persistence_manager + .process_command(DeviceCommand::SetPushName(new_name.clone())) + .await; + + self.core.event_bus.dispatch(Event::SelfPushNameUpdated( + crate::types::events::SelfPushNameUpdated { + from_server: true, + old_name, + new_name: new_name.clone(), + }, + )); + + let client_clone = self.clone(); + self.runtime + .spawn(Box::pin(async move { + if let Err(e) = client_clone.presence().set_available().await { + log::warn!("Failed to send presence after push name update: {:?}", e); + } else { + log::debug!("Sent presence after push name update."); + } + })) + .detach(); + } + + /// Register a waiter for an incoming node matching the given filter. + /// + /// Returns a receiver that resolves when a matching node arrives. + /// The waiter starts buffering immediately, so register it **before** + /// performing the action that triggers the expected node. + /// + /// When multiple waiters match the same node, each matching waiter + /// receives a clone of the node (broadcast within a single resolve pass). + /// + /// # Example + /// ```ignore + /// let waiter = client.wait_for_node( + /// NodeFilter::tag("notification").attr("type", "w:gp2"), + /// ); + /// client.groups().add_participants(&group_jid, &[jid_c]).await?; + /// let node = waiter.await.expect("notification arrived"); + /// ``` + pub fn wait_for_node( + &self, + filter: NodeFilter, + ) -> futures::channel::oneshot::Receiver> { + let (tx, rx) = futures::channel::oneshot::channel(); + self.node_waiter_count.fetch_add(1, Ordering::Release); + let mut waiters = self + .node_waiters + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + waiters.push(NodeWaiter { filter, tx }); + rx + } + + /// Register a waiter for an outgoing node before it is encrypted and sent. + /// + /// This is intended for tests and diagnostics that need to inspect the raw + /// stanza built by the client, such as asserting whether `` or + /// `` was attached. + pub fn wait_for_sent_node( + &self, + filter: NodeFilter, + ) -> futures::channel::oneshot::Receiver> { + let (tx, rx) = futures::channel::oneshot::channel(); + self.sent_node_waiter_count.fetch_add(1, Ordering::Release); + let mut waiters = self + .sent_node_waiters + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + waiters.push(SentNodeWaiter { filter, tx }); + rx + } + + /// Check pending node waiters against an incoming node. + /// Only called when `node_waiter_count > 0`. + pub(crate) fn resolve_node_waiters(&self, node: &Arc) { + resolve_waiters(&self.node_waiters, &self.node_waiter_count, node); + } + + pub(crate) fn resolve_sent_node_waiters(&self, node: &Arc) { + let nr = node.as_node_ref(); + let mut waiters = self + .sent_node_waiters + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut i = 0; + while i < waiters.len() { + if waiters[i].tx.is_canceled() { + waiters.swap_remove(i); + self.sent_node_waiter_count.fetch_sub(1, Ordering::Release); + } else if waiters[i].filter.matches(&nr) { + let w = waiters.swap_remove(i); + self.sent_node_waiter_count.fetch_sub(1, Ordering::Release); + let _ = w.tx.send(Arc::clone(node)); + } else { + i += 1; + } + } + } + + pub(crate) fn clear_sent_node_waiters(&self) { + let mut waiters = self + .sent_node_waiters + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let count = waiters.len(); + if count > 0 { + waiters.clear(); + self.sent_node_waiter_count + .fetch_sub(count, Ordering::Release); + } + } + + fn should_downgrade_sync_error(&self, err: &anyhow::Error) -> bool { + if self.is_shutting_down() { + return true; + } + + matches!( + err.downcast_ref::(), + Some( + crate::request::IqError::NotConnected + | crate::request::IqError::InternalChannelClosed + ) + ) + } + + /// Log a sync error, downgrading to debug level during shutdown/disconnect. + pub(crate) fn log_sync_error(&self, context: &str, err: &anyhow::Error) { + if self.should_downgrade_sync_error(err) { + debug!("Skipping {context} during shutdown: {err}"); + } else { + warn!("Failed {context}: {err}"); + } + } + + /// Create and configure the stanza router with all the handlers. + pub(crate) fn create_stanza_router() -> crate::handlers::router::StanzaRouter { + use crate::handlers::{ + basic::{AckHandler, FailureHandler, StreamErrorHandler, SuccessHandler}, + chatstate::ChatstateHandler, + ib::IbHandler, + iq::IqHandler, + message::MessageHandler, + notification::NotificationHandler, + receipt::ReceiptHandler, + router::StanzaRouter, + }; + + let mut router = StanzaRouter::new(); + + // Register all handlers + router.register(Arc::new(MessageHandler)); + router.register(Arc::new(ReceiptHandler)); + router.register(Arc::new(IqHandler)); + router.register(Arc::new(SuccessHandler)); + router.register(Arc::new(FailureHandler)); + router.register(Arc::new(StreamErrorHandler)); + router.register(Arc::new(IbHandler)); + router.register(Arc::new(NotificationHandler)); + router.register(Arc::new(AckHandler)); + router.register(Arc::new(ChatstateHandler)); + + router.register(Arc::new(crate::handlers::call::CallHandler)); + + // Register unimplemented handlers + router.register(Arc::new(crate::handlers::presence::PresenceHandler)); + + router + } +} diff --git a/src/client/adapters.rs b/src/client/adapters.rs new file mode 100644 index 000000000..878fbcfb1 --- /dev/null +++ b/src/client/adapters.rs @@ -0,0 +1,80 @@ +//! Signal/sender-key store adapters, per-session locks and noise socket access. + +use super::*; + +impl Client { + /// Build a [`SignalProtocolStoreAdapter`] from the current device state and signal cache. + pub(crate) async fn signal_adapter( + &self, + ) -> crate::store::signal_adapter::SignalProtocolStoreAdapter { + let device_store = self.persistence_manager.get_device_arc().await; + self.signal_adapter_from(device_store) + } + + /// Build a standalone [`SenderKeyAdapter`] from the current device state and + /// signal cache, avoiding the full five-store adapter on the SKDM path. + pub(crate) async fn sender_key_adapter( + &self, + ) -> crate::store::signal_adapter::SenderKeyAdapter { + crate::store::signal_adapter::SenderKeyAdapter::new( + self.persistence_manager.get_device_arc().await, + self.signal_cache.clone(), + ) + } + + /// Build a [`SignalProtocolStoreAdapter`] from a pre-fetched device arc. + pub(crate) fn signal_adapter_from( + &self, + device_store: Arc>, + ) -> crate::store::signal_adapter::SignalProtocolStoreAdapter { + crate::store::signal_adapter::SignalProtocolStoreAdapter::new( + device_store, + self.signal_cache.clone(), + ) + } + + /// Get the per-address session mutex from the lock cache. + pub(crate) async fn session_lock_for( + &self, + signal_addr_str: &str, + ) -> Arc> { + self.session_locks + .get_with_by_ref(signal_addr_str, async { + Arc::new(async_lock::Mutex::new(())) + }) + .await + } + + /// Get the active noise socket, or error if not connected. + pub(crate) async fn get_noise_socket( + &self, + ) -> Result, ClientError> { + self.noise_socket + .lock() + .await + .clone() + .ok_or(ClientError::NotConnected) + } + + /// Flush the in-memory signal cache to the database backend. + /// Called after each message is decrypted or after encryption operations. + pub(crate) async fn flush_signal_cache(&self) -> Result<(), anyhow::Error> { + let device = self.persistence_manager.get_device_arc().await; + let device_guard = device.read().await; + self.signal_cache + .flush(&*device_guard.backend) + .await + .map_err(|e| anyhow::anyhow!("Failed to flush signal cache: {e}")) + } + + /// [`flush_signal_cache`](Self::flush_signal_cache) with error logging instead of propagation. + pub(crate) async fn flush_signal_cache_logged(&self, context: &str, id: Option<&str>) { + if let Err(e) = self.flush_signal_cache().await { + if let Some(id) = id { + log::error!("Failed to flush signal cache ({context} {id}): {e:?}"); + } else { + log::error!("Failed to flush signal cache ({context}): {e:?}"); + } + } + } +} diff --git a/src/client/app_state.rs b/src/client/app_state.rs new file mode 100644 index 000000000..597f64841 --- /dev/null +++ b/src/client/app_state.rs @@ -0,0 +1,745 @@ +//! App-state collection sync and mutation dispatch. + +use super::*; + +impl Client { + pub(crate) async fn get_app_state_processor(&self) -> Arc { + let mut guard = self.app_state_processor.lock().await; + if let Some(proc) = guard.as_ref() { + return proc.clone(); + } + debug!("Initializing AppStateProcessor for the first time."); + let proc = Arc::new(AppStateProcessor::new( + self.persistence_manager.backend(), + self.runtime.clone(), + )); + *guard = Some(proc.clone()); + proc + } + + /// Public entry point for processing [`MajorSyncTask`] from the sync channel. + pub async fn process_sync_task(self: &Arc, task: crate::sync_task::MajorSyncTask) { + match task { + crate::sync_task::MajorSyncTask::HistorySync { + message_id, + notification, + } => { + self.process_history_sync_task(message_id, *notification) + .await; + self.finish_history_sync_task(); + } + crate::sync_task::MajorSyncTask::AppStateSync { name, full_sync } => { + if let Err(e) = self.process_app_state_sync_task(name, full_sync).await { + log::warn!("App state sync task for {name:?} failed: {e}"); + } + } + } + } + + pub(crate) async fn fetch_app_state_with_retry(&self, name: WAPatchName) -> anyhow::Result<()> { + // In-flight dedup: skip if this collection is already being synced. + // Matches WA Web's WAWebSyncdCollectionsStateMachine which tracks in-flight syncs + // and queues new requests to a pending set. + { + let mut syncing = self.app_state_syncing.lock().await; + if !syncing.insert(name) { + debug!(target: "Client/AppState", "Skipping sync for {:?}: already in flight", name); + return Ok(()); + } + } + + let result = self.fetch_app_state_with_retry_inner(name).await; + + // Always remove from in-flight set when done + self.app_state_syncing.lock().await.remove(&name); + + result + } + + async fn fetch_app_state_with_retry_inner(&self, name: WAPatchName) -> anyhow::Result<()> { + let mut attempt = 0u32; + loop { + attempt += 1; + // full_sync=false lets process_app_state_sync_task auto-detect: + // version 0 → snapshot (full sync), version > 0 → incremental patches. + // Matches WA Web which only requests snapshot when version is undefined. + let res = self.process_app_state_sync_task(name, false).await; + match res { + Ok(()) => return Ok(()), + Err(e) => { + if e.downcast_ref::() + .is_some_and(|ase| { + matches!(ase, crate::appstate_sync::AppStateSyncError::KeyNotFound(_)) + }) + && attempt == 1 + { + if !self.initial_app_state_keys_received.load(Ordering::Relaxed) { + debug!(target: "Client/AppState", "App state key missing for {:?}; waiting up to 10s for key share then retrying", name); + if rt_timeout( + &*self.runtime, + Duration::from_secs(10), + self.initial_keys_synced_notifier.listen(), + ) + .await + .is_err() + { + warn!(target: "Client/AppState", "Timeout waiting for key share for {:?}; retrying anyway", name); + } + } + continue; + } + let is_db_locked = e + .downcast_ref::() + .is_some_and(|se| se.is_database_busy_or_locked()) + || e.downcast_ref::() + .is_some_and(|ase| match ase { + crate::appstate_sync::AppStateSyncError::Store(se) => { + se.is_database_busy_or_locked() + } + _ => false, + }); + if is_db_locked && attempt < APP_STATE_RETRY_MAX_ATTEMPTS { + let backoff = Duration::from_millis(200 * attempt as u64 + 150); + warn!(target: "Client/AppState", "Attempt {} for {:?} failed due to locked DB; backing off {:?} and retrying", attempt, name, backoff); + self.runtime.sleep(backoff).await; + continue; + } + return Err(e); + } + } + } + } + + /// Sync multiple collections in a single IQ request, re-fetching those with `has_more_patches`. + /// Matches WA Web's `serverSync()` outer loop (`3JJWKHeu5-P.js:54278-54305`). + /// Max 5 iterations (WA Web's `C=5` constant). + pub(crate) async fn sync_collections_batched( + &self, + collections: Vec, + ) -> anyhow::Result<()> { + if collections.is_empty() { + return Ok(()); + } + + // In-flight dedup: filter out collections already being synced + let pending = { + let mut syncing = self.app_state_syncing.lock().await; + let mut filtered = Vec::with_capacity(collections.len()); + for name in collections { + if syncing.insert(name) { + filtered.push(name); + } else { + debug!(target: "Client/AppState", "Skipping {:?} in batch: already in flight", name); + } + } + filtered + }; + + if pending.is_empty() { + return Ok(()); + } + + // Track all collections for cleanup + let all_collections: Vec = pending.clone(); + + let result = self.sync_collections_batched_inner(pending).await; + + // Always clean up in-flight set + { + let mut syncing = self.app_state_syncing.lock().await; + for name in &all_collections { + syncing.remove(name); + } + } + + result + } + + async fn sync_collections_batched_inner( + &self, + mut pending: Vec, + ) -> anyhow::Result<()> { + use wacore::appstate::patch_decode::CollectionSyncError; + const MAX_ITERATIONS: usize = 5; + let mut iteration = 0; + + while !pending.is_empty() && iteration < MAX_ITERATIONS { + iteration += 1; + debug!( + target: "Client/AppState", + "Batched sync iteration {}/{}: {:?}", + iteration, MAX_ITERATIONS, pending + ); + + let backend = self.persistence_manager.backend(); + + // Build multi-collection IQ, tracking which collections need a snapshot + let mut collection_nodes = Vec::with_capacity(pending.len()); + let mut was_snapshot = std::collections::HashSet::new(); + for &name in &pending { + let state = backend.get_version(name.as_str()).await?; + let want_snapshot = state.version == 0; + if want_snapshot { + was_snapshot.insert(name); + } + let mut builder = NodeBuilder::new("collection") + .attr("name", name.as_str()) + .attr( + "return_snapshot", + if want_snapshot { "true" } else { "false" }, + ); + if !want_snapshot { + builder = builder.attr("version", state.version); + } + collection_nodes.push(builder.build()); + } + + let sync_node = NodeBuilder::new("sync").children(collection_nodes).build(); + let iq = crate::request::InfoQuery { + namespace: "w:sync:app:state", + query_type: crate::request::InfoQueryType::Set, + to: server_jid().clone(), + target: None, + id: None, + content: Some(wacore_binary::NodeContent::Nodes(vec![sync_node])), + timeout: Some(Duration::from_secs(30)), + }; + + let resp = self.send_iq(iq).await?; + + // Pre-download all external blobs for all collections in the response + let mut pre_downloaded: std::collections::HashMap> = + std::collections::HashMap::new(); + + // Parse the response once here for pre-download; the same parsed + // lists are handed to the processor below (no second parse). + let patch_lists = wacore::appstate::patch_decode::parse_patch_lists_ref(resp.get())?; + { + for pl in &patch_lists { + // Download external snapshot + if let Some(ext) = &pl.snapshot_ref + && let Some(path) = &ext.direct_path + { + match self.download(ext).await { + Ok(bytes) => { + pre_downloaded.insert(path.clone(), bytes); + } + Err(e) => { + warn!( + "Failed to download external snapshot for {:?}: {e}", + pl.name + ); + } + } + } + + // Download external mutations + for patch in &pl.patches { + if let Some(ext) = &patch.external_mutations + && let Some(path) = &ext.direct_path + { + match self.download(ext).await { + Ok(bytes) => { + pre_downloaded.insert(path.clone(), bytes); + } + Err(e) => { + let v = + patch.version.as_ref().and_then(|v| v.version).unwrap_or(0); + warn!( + "Failed to download external mutations for patch v{}: {e}", + v + ); + } + } + } + } + } + } + + let download = |ext: &wa::ExternalBlobReference| -> anyhow::Result> { + if let Some(path) = &ext.direct_path { + if let Some(bytes) = pre_downloaded.get(path) { + Ok(bytes.clone()) + } else { + Err(anyhow::anyhow!( + "external blob not pre-downloaded: {}", + path + )) + } + } else { + Err(anyhow::anyhow!("external blob has no directPath")) + } + }; + + // Process the already-parsed collections (no re-parse of the response). + let proc = self.get_app_state_processor().await; + let results = proc + .process_patch_lists(patch_lists, &download, true) + .await?; + + let mut needs_refetch = Vec::new(); + + for (mutations, new_state, list) in results { + let name = list.name; + + // Handle per-collection errors + if let Some(ref err) = list.error { + match err { + CollectionSyncError::Conflict { has_more } => { + if *has_more { + // ConflictHasMore: server has more patches, must refetch. + warn!(target: "Client/AppState", "Collection {:?} conflict (has_more=true), will refetch", name); + needs_refetch.push(name); + } else { + // Conflict without has_more: WA Web treats this as success + // when there are no pending mutations to push (which is + // always the case for us since we don't push app state). + debug!(target: "Client/AppState", "Collection {:?} conflict (has_more=false), treating as success (no pending mutations)", name); + } + continue; + } + CollectionSyncError::Fatal { code, text } => { + warn!(target: "Client/AppState", "Collection {:?} fatal error {}: {}", name, code, text); + continue; + } + CollectionSyncError::Retry { code, text } => { + warn!(target: "Client/AppState", "Collection {:?} retryable error {}: {}, will refetch", name, code, text); + needs_refetch.push(name); + continue; + } + } + } + + // Handle missing keys + let missing = match proc.get_missing_key_ids(&list).await { + Ok(v) => v, + Err(e) => { + warn!("Failed to get missing key IDs for {:?}: {}", name, e); + Vec::new() + } + }; + self.request_missing_keys_with_dedup(missing).await; + + // full_sync is true only when this collection had a snapshot + // (version was 0 before sync). This prevents server_sync-triggered + // incremental syncs from being incorrectly marked as full syncs. + let full_sync = was_snapshot.contains(&name); + for m in mutations { + self.dispatch_app_state_mutation(&m, full_sync).await; + } + + // Save version + backend + .set_version(name.as_str(), new_state.clone()) + .await?; + + // Check if this collection needs more patches + if list.has_more_patches { + needs_refetch.push(name); + } + + debug!( + target: "Client/AppState", + "Batched sync: {:?} done (version={}, has_more={})", + name, new_state.version, list.has_more_patches + ); + } + + pending = needs_refetch; + } + + if !pending.is_empty() { + warn!( + target: "Client/AppState", + "Batched sync: max iterations ({}) reached for {:?}", + MAX_ITERATIONS, pending + ); + } + + Ok(()) + } + + pub(crate) async fn process_app_state_sync_task( + &self, + name: WAPatchName, + full_sync: bool, + ) -> anyhow::Result<()> { + if self.is_shutting_down() { + debug!(target: "Client/AppState", "Skipping app state sync task {:?}: client is shutting down", name); + return Ok(()); + } + + let backend = self.persistence_manager.backend(); + let mut full_sync = full_sync; + + let mut state = backend.get_version(name.as_str()).await?; + if state.version == 0 { + full_sync = true; + } + + let mut has_more = true; + let mut want_snapshot = full_sync; + // Safety cap to prevent infinite loops if the server keeps returning + // has_more_patches=true without advancing the version (WA Web uses 500). + const MAX_PAGINATION_ITERATIONS: u32 = 500; + let mut iteration = 0u32; + + while has_more { + if self.is_shutting_down() { + debug!(target: "Client/AppState", "Stopping app state sync task {:?}: shutdown detected", name); + break; + } + iteration += 1; + if iteration > MAX_PAGINATION_ITERATIONS { + warn!(target: "Client/AppState", "App state sync for {:?} exceeded {} iterations, aborting", name, MAX_PAGINATION_ITERATIONS); + break; + } + debug!(target: "Client/AppState", "Fetching app state patch batch: name={:?} want_snapshot={want_snapshot} version={} full_sync={} has_more_previous={}", name, state.version, full_sync, has_more); + + let mut collection_builder = NodeBuilder::new("collection") + .attr("name", name.as_str()) + .attr( + "return_snapshot", + if want_snapshot { "true" } else { "false" }, + ); + if !want_snapshot { + collection_builder = collection_builder.attr("version", state.version); + } + let sync_node = NodeBuilder::new("sync") + .children([collection_builder.build()]) + .build(); + let iq = crate::request::InfoQuery { + namespace: "w:sync:app:state", + query_type: crate::request::InfoQueryType::Set, + to: server_jid().clone(), + target: None, + id: None, + content: Some(wacore_binary::NodeContent::Nodes(vec![sync_node])), + timeout: None, + }; + + let resp = self.send_iq(iq).await?; + if self.is_shutting_down() { + debug!(target: "Client/AppState", "Discarding app state sync response for {:?}: shutdown detected", name); + break; + } + debug!(target: "Client/AppState", "Received IQ response for {:?}; decoding patches", name); + + let _decode_start = wacore::time::Instant::now(); + + // Pre-download all external blobs (snapshot and patch mutations) + // We use directPath as the key to identify each blob + let mut pre_downloaded: std::collections::HashMap> = + std::collections::HashMap::new(); + + // Parse the response once here for pre-download; the same parsed list + // is handed to the processor below (no second parse). + let pl = wacore::appstate::patch_decode::parse_patch_list_ref(resp.get())?; + { + debug!(target: "Client/AppState", "Parsed patch list for {:?}: has_snapshot_ref={} has_more_patches={} patches_count={}", + name, pl.snapshot_ref.is_some(), pl.has_more_patches, pl.patches.len()); + + // Download external snapshot if present + if let Some(ext) = &pl.snapshot_ref + && let Some(path) = &ext.direct_path + { + match self.download(ext).await { + Ok(bytes) => { + debug!(target: "Client/AppState", "Downloaded external snapshot ({} bytes)", bytes.len()); + pre_downloaded.insert(path.clone(), bytes); + } + Err(e) => { + warn!("Failed to download external snapshot: {e}"); + } + } + } + + // Download external mutations for each patch that has them + for patch in &pl.patches { + if let Some(ext) = &patch.external_mutations + && let Some(path) = &ext.direct_path + { + let patch_version = + patch.version.as_ref().and_then(|v| v.version).unwrap_or(0); + match self.download(ext).await { + Ok(bytes) => { + debug!(target: "Client/AppState", "Downloaded external mutations for patch v{} ({} bytes)", patch_version, bytes.len()); + pre_downloaded.insert(path.clone(), bytes); + } + Err(e) => { + warn!( + "Failed to download external mutations for patch v{}: {e}", + patch_version + ); + } + } + } + } + } + + let download = |ext: &wa::ExternalBlobReference| -> anyhow::Result> { + if let Some(path) = &ext.direct_path { + if let Some(bytes) = pre_downloaded.get(path) { + Ok(bytes.clone()) + } else { + Err(anyhow::anyhow!( + "external blob not pre-downloaded: {}", + path + )) + } + } else { + Err(anyhow::anyhow!("external blob has no directPath")) + } + }; + + let proc = self.get_app_state_processor().await; + let (mutations, new_state, list) = + proc.process_parsed_patch_list(pl, &download, true).await?; + let decode_elapsed = _decode_start.elapsed(); + if decode_elapsed.as_millis() > 500 { + debug!(target: "Client/AppState", "Patch decode for {:?} took {:?}", name, decode_elapsed); + } + + let missing = match proc.get_missing_key_ids(&list).await { + Ok(v) => v, + Err(e) => { + warn!("Failed to get missing key IDs for {:?}: {}", name, e); + Vec::new() + } + }; + self.request_missing_keys_with_dedup(missing).await; + + for m in mutations { + debug!(target: "Client/AppState", "Dispatching mutation kind={} index_len={} full_sync={}", m.index.first().map(|s| s.as_str()).unwrap_or(""), m.index.len(), full_sync); + self.dispatch_app_state_mutation(&m, full_sync).await; + } + + state = new_state; + has_more = list.has_more_patches; + // After the first batch, never request a snapshot again — only incremental patches. + want_snapshot = false; + debug!(target: "Client/AppState", "After processing batch name={:?} has_more={has_more} new_version={}", name, state.version); + } + + backend.set_version(name.as_str(), state.clone()).await?; + + debug!(target: "Client/AppState", "Completed and saved app state sync for {:?} (final version={})", name, state.version); + Ok(()) + } + + /// Request missing app-state keys with dedup stamps. + /// On send failure, removes stamps so keys can be retried next sync. + async fn request_missing_keys_with_dedup(&self, missing: Vec>) { + if missing.is_empty() { + return; + } + let mut to_request: Vec> = Vec::with_capacity(missing.len()); + let mut guard = self.app_state_key_requests.lock().await; + let now = wacore::time::Instant::now(); + for key_id in missing { + let hex_id = hex::encode(&key_id); + let should = guard + .get(&hex_id) + .map(|t| t.elapsed() > std::time::Duration::from_secs(24 * 3600)) + .unwrap_or(true); + if should { + guard.insert(hex_id, now); + to_request.push(key_id); + } + } + guard.retain(|_, t| t.elapsed() < std::time::Duration::from_secs(24 * 3600)); + drop(guard); + if !to_request.is_empty() + && let Err(e) = self.request_app_state_keys(&to_request).await + { + warn!("Failed to send app state key request: {e}"); + let mut guard = self.app_state_key_requests.lock().await; + for key_id in &to_request { + guard.remove(&hex::encode(key_id)); + } + } + } + + async fn request_app_state_keys(&self, raw_key_ids: &[Vec]) -> Result<(), anyhow::Error> { + if raw_key_ids.is_empty() { + return Ok(()); + } + let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let own_jid = match device_snapshot.pn.clone() { + Some(j) => j, + None => { + return Err(anyhow::anyhow!( + "no own JID available for app-state key request" + )); + } + }; + let key_ids: Vec = raw_key_ids + .iter() + .map(|k| wa::message::AppStateSyncKeyId { + key_id: Some(k.clone()), + }) + .collect(); + let msg = wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::AppStateSyncKeyRequest as i32), + app_state_sync_key_request: Some(wa::message::AppStateSyncKeyRequest { key_ids }), + ..Default::default() + })), + ..Default::default() + }; + self.send_message_impl( + own_jid, + &msg, + Some(self.generate_message_id().await), + true, + false, + None, + vec![], + None, + ) + .await?; + Ok(()) + } + + /// Send an app state patch to the server for a given collection. + /// + /// Builds the IQ stanza and sends it. Returns the updated hash state. + pub(crate) async fn send_app_state_patch( + &self, + collection_name: &str, + mutations: Vec, + ) -> Result<()> { + let proc = self.get_app_state_processor().await; + let (patch_bytes, base_version) = proc.build_patch(collection_name, mutations).await?; + + let collection_node = NodeBuilder::new("collection") + .attr("name", collection_name) + .attr("version", base_version) + .attr("return_snapshot", "false") + .children([NodeBuilder::new("patch").bytes(patch_bytes).build()]) + .build(); + let sync_node = NodeBuilder::new("sync").children([collection_node]).build(); + let iq = crate::request::InfoQuery { + namespace: "w:sync:app:state", + query_type: crate::request::InfoQueryType::Set, + to: server_jid().clone(), + target: None, + id: None, + content: Some(wacore_binary::NodeContent::Nodes(vec![sync_node])), + timeout: None, + }; + + self.send_iq(iq).await?; + + // Re-sync to get the latest state from the server after our patch was accepted. + // This matches whatsmeow's behavior: fetchAppState after successful send. + if let Ok(patch_name) = collection_name.parse::() + && let Err(e) = self.fetch_app_state_with_retry(patch_name).await + { + log::warn!("Failed to re-sync {collection_name} after patch send: {e}"); + } + + Ok(()) + } + + async fn dispatch_app_state_mutation( + &self, + m: &crate::appstate_sync::Mutation, + full_sync: bool, + ) { + use wacore::types::events::Event; + + if m.index.is_empty() { + return; + } + + // NCT salt sync — handles both "set" (store salt) and "remove" (clear salt). + // Source: WAWebNctSaltSync, syncd collection RegularHigh, action "nct_salt_sync". + if m.index[0] == "nct_salt_sync" { + if m.operation == wa::syncd_mutation::SyncdOperation::Remove { + debug!(target: "Client/AppState", "Removing NCT salt via app state sync"); + self.persistence_manager + .process_command(DeviceCommand::SetNctSalt(None)) + .await; + } else if let Some(val) = &m.action_value + && let Some(act) = &val.nct_salt_sync_action + && let Some(salt) = &act.salt + { + if salt.is_empty() { + warn!(target: "Client/AppState", "nct_salt_sync mutation has empty salt, ignoring"); + } else { + debug!(target: "Client/AppState", "Stored NCT salt via app state sync ({} bytes)", salt.len()); + self.persistence_manager + .process_command(DeviceCommand::SetNctSalt(Some(salt.clone()))) + .await; + } + } else { + warn!(target: "Client/AppState", "nct_salt_sync mutation missing salt in action value"); + } + return; + } + + // All remaining mutations only care about Set operations + if m.operation != wa::syncd_mutation::SyncdOperation::Set { + return; + } + + // Delegate chat-related mutations (mute, pin, archive, star, contact, etc.) + if crate::features::chat_actions::dispatch_chat_mutation(&self.core.event_bus, m, full_sync) + { + return; + } + + // Label mutations have their own index shape (labelId, not a chat JID at + // index[1]), so they are dispatched separately from chat actions. + if crate::features::labels::dispatch_label_mutation(&self.core.event_bus, m, full_sync) { + return; + } + + // Handle client-internal mutations that need persistence/presence access + if m.index[0] == "setting_pushName" + && let Some(val) = &m.action_value + && let Some(act) = &val.push_name_setting + && let Some(new_name) = &act.name + { + let new_name = new_name.clone(); + let bus = self.core.event_bus.clone(); + + let snapshot = self.persistence_manager.get_device_snapshot().await; + let old = snapshot.push_name.clone(); + if old != new_name { + debug!(target: "Client/AppState", "Persisting push name from app state mutation: '{}' (old='{}')", new_name, old); + self.persistence_manager + .process_command(DeviceCommand::SetPushName(new_name.clone())) + .await; + bus.dispatch(Event::SelfPushNameUpdated( + crate::types::events::SelfPushNameUpdated { + from_server: true, + old_name: old.clone(), + new_name: new_name.clone(), + }, + )); + + // WhatsApp Web sends presence immediately when receiving pushname + if old.is_empty() && !new_name.is_empty() { + debug!(target: "Client/AppState", "Sending presence after receiving initial pushname from app state sync"); + if let Err(e) = self.presence().set_available().await { + warn!(target: "Client/AppState", "Failed to send presence after pushname sync: {e:?}"); + } + } + } else { + debug!(target: "Client/AppState", "Push name mutation received but name unchanged: '{}'", new_name); + } + } + } + + pub async fn clean_dirty_bits( + &self, + bit: wacore::iq::dirty::DirtyBit, + ) -> Result<(), crate::request::IqError> { + use wacore::iq::dirty::CleanDirtyBitsSpec; + + let spec = CleanDirtyBitsSpec::single(bit); + self.execute(spec).await + } +} diff --git a/src/client/iq_ops.rs b/src/client/iq_ops.rs new file mode 100644 index 000000000..ce49d6f71 --- /dev/null +++ b/src/client/iq_ops.rs @@ -0,0 +1,196 @@ +//! IQ-based operations: props, privacy settings, profiles and assorted requests. + +use super::*; + +impl Client { + pub async fn set_passive(&self, passive: bool) -> Result<(), crate::request::IqError> { + use wacore::iq::passive::PassiveModeSpec; + self.execute(PassiveModeSpec::new(passive)).await + } + + pub async fn fetch_props(&self) -> Result<(), crate::request::IqError> { + use wacore::iq::props::PropsSpec; + use wacore::store::commands::DeviceCommand; + + let stored_hash = self + .persistence_manager + .get_device_snapshot() + .await + .props_hash + .clone(); + + // Deltas only contain changed props, so they're invalid against an empty cache. + let spec = match &stored_hash { + Some(hash) if self.ab_props.is_seeded() => { + debug!("Fetching props with hash for delta update..."); + PropsSpec::with_hash(hash) + } + _ => { + debug!("Fetching props (full)..."); + PropsSpec::new() + } + }; + + let response = self.execute(spec).await?; + + if response.delta_update { + debug!( + "Props delta update received ({} changed props)", + response.experiment_props.len() + ); + } else { + debug!( + "Props full update received ({} props, hash={:?})", + response.experiment_props.len(), + response.hash + ); + } + + self.ab_props + .apply_props(response.delta_update, response.experiment_props.into_iter()) + .await; + + if let Some(new_hash) = response.hash { + self.persistence_manager + .process_command(DeviceCommand::SetPropsHash(Some(new_hash))) + .await; + } + + Ok(()) + } + + pub(crate) fn ab_props(&self) -> &wacore::store::ab_props::AbPropsCache { + &self.ab_props + } + + pub async fn fetch_privacy_settings( + &self, + ) -> Result { + use wacore::iq::privacy::PrivacySettingsSpec; + + debug!("Fetching privacy settings..."); + + self.execute(PrivacySettingsSpec::new()).await + } + + /// Set a privacy setting. + /// + /// Use [`PrivacyCategory::is_valid_value`] to check valid combinations. + /// + /// # Example + /// ```ignore + /// use wacore::iq::privacy::{PrivacyCategory, PrivacyValue}; + /// client.set_privacy_setting(PrivacyCategory::Last, PrivacyValue::Contacts).await?; + /// ``` + pub async fn set_privacy_setting( + &self, + category: wacore::iq::privacy::PrivacyCategory, + value: wacore::iq::privacy::PrivacyValue, + ) -> Result { + use wacore::iq::privacy::SetPrivacySettingSpec; + self.execute(SetPrivacySettingSpec::new(category, value)) + .await + } + + /// Set a privacy setting to `contact_blacklist` with a disallowed list update. + /// + /// Only `Last`, `Profile`, `Status`, `GroupAdd` support disallowed lists. + /// Returns the server's updated dhash for use in subsequent updates. + pub async fn set_privacy_disallowed_list( + &self, + category: wacore::iq::privacy::PrivacyCategory, + update: wacore::iq::privacy::DisallowedListUpdate, + ) -> Result { + use wacore::iq::privacy::SetPrivacySettingSpec; + self.execute(SetPrivacySettingSpec::with_disallowed_list( + category, update, + )) + .await + } + + /// Set the default disappearing messages duration (seconds). Pass 0 to disable. + pub async fn set_default_disappearing_mode( + &self, + duration: u32, + ) -> Result<(), crate::request::IqError> { + use wacore::iq::privacy::SetDefaultDisappearingModeSpec; + self.execute(SetDefaultDisappearingModeSpec::new(duration)) + .await + } + + /// Get business profile for a WhatsApp Business account. + pub async fn get_business_profile( + &self, + jid: &wacore_binary::Jid, + ) -> Result, crate::request::IqError> { + use wacore::iq::business::BusinessProfileSpec; + self.execute(BusinessProfileSpec::new(jid)).await + } + + /// Reject an incoming call. Fire-and-forget — no server response is expected. + pub async fn reject_call( + &self, + call_id: &str, + call_from: &wacore_binary::Jid, + ) -> Result<(), anyhow::Error> { + anyhow::ensure!(!call_id.is_empty(), "call_id cannot be empty"); + let id = self.generate_request_id(); + + let stanza = wacore_binary::builder::NodeBuilder::new("call") + .attr("to", call_from) + .attr("id", id) + .children([wacore_binary::builder::NodeBuilder::new("reject") + .attr("call-id", call_id) + .attr("call-creator", call_from) + .attr("count", "0") + .build()]) + .build(); + + self.send_node(stanza).await?; + Ok(()) + } + + pub async fn send_digest_key_bundle(&self) -> Result<(), crate::request::IqError> { + use wacore::iq::prekeys::DigestKeyBundleSpec; + + debug!("Sending digest key bundle..."); + + self.execute(DigestKeyBundleSpec::new()).await.map(|_| ()) + } + + /// Override `DeviceProps` fields before the initial pairing. Only fields + /// with `Some` are changed. In-memory only — WA Web regenerates + /// `device_props` at each registration, and it has no wire effect after + /// pairing. Call before `connect()` on every process start that still + /// needs to pair. + pub async fn set_device_props(&self, override_: wacore::store::DevicePropsOverride) { + use wacore::store::commands::DeviceCommand; + if override_.is_empty() { + return; + } + if self + .persistence_manager + .get_device_snapshot() + .await + .pn + .is_some() + { + warn!( + target: "Client/DeviceProps", + "set_device_props called after pairing — stored but not sent on the wire" + ); + } + self.persistence_manager + .process_command(DeviceCommand::SetDeviceProps(override_)) + .await; + } + + /// Set the noise-handshake `ClientPayload` profile. In-memory only; + /// call before each `connect()` on a fresh process. + pub async fn set_client_profile(&self, profile: wacore::client_profile::ClientProfile) { + use wacore::store::commands::DeviceCommand; + self.persistence_manager + .process_command(DeviceCommand::SetClientProfile(profile)) + .await; + } +} diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs new file mode 100644 index 000000000..ce23f1922 --- /dev/null +++ b/src/client/lifecycle.rs @@ -0,0 +1,697 @@ +//! Client construction and connection lifecycle: connect, run, reconnect, shutdown. + +use super::*; + +impl Client { + pub fn shutdown_signal(&self) -> wacore::runtime::ShutdownSignal { + self.shutdown_notifier.subscribe() + } + + /// Synchronous flag-only equivalent of the first lines of `disconnect()`. + /// Spawned tasks watching `is_shutting_down()` / `shutdown_notifier` exit + /// on their next poll. Does NOT flush, close the transport, or touch + /// persistence — prefer `disconnect()` whenever you can `await`. Exists + /// for `Drop` impls on FFI wrappers (e.g. `WasmWhatsAppClient`) that + /// can't run async cleanup synchronously. + pub fn signal_shutdown_sync(&self) { + self.expected_disconnect.store(true, Ordering::Relaxed); + self.is_running.store(false, Ordering::Relaxed); + self.shutdown_notifier.notify(); + self.notify_connection_shutdown(); + } + + pub(crate) fn connection_shutdown_signal(&self) -> wacore::runtime::ShutdownSignal { + self.connection_shutdown + .lock() + .unwrap_or_else(|p| p.into_inner()) + .subscribe() + } + + /// Fire the per-connection shutdown. Per-connection subscribers exit; + /// the terminal shutdown_notifier is untouched so reconnects still work. + pub(crate) fn notify_connection_shutdown(&self) { + self.connection_shutdown + .lock() + .unwrap_or_else(|p| p.into_inner()) + .notify(); + } + + /// Reset the per-connection notifier. Call at the start of each new + /// connection so subscribers registered afterwards see a fresh signal. + /// The previous notifier's subscribers have already been woken (either + /// by notify on disconnect, or by falling out of scope). + pub(crate) fn reset_connection_shutdown(&self) { + *self + .connection_shutdown + .lock() + .unwrap_or_else(|p| p.into_inner()) = wacore::runtime::ShutdownNotifier::new(); + } + + pub(crate) fn is_shutting_down(&self) -> bool { + self.expected_disconnect.load(Ordering::Relaxed) || !self.is_running.load(Ordering::Relaxed) + } + + /// Returns `true` when the client has completed its full startup: + /// transport connected, server authenticated, and critical app state synced. + /// This is the condition `wait_for_connected` uses to resolve. + fn is_fully_ready(&self) -> bool { + self.is_connected() && self.is_logged_in() && self.is_ready.load(Ordering::Relaxed) + } + + /// Dispatch the Connected event and notify waiters. + pub(crate) fn dispatch_connected(&self) { + self.is_ready.store(true, Ordering::Relaxed); + self.core + .event_bus + .dispatch(Event::Connected(crate::types::events::Connected)); + self.connected_notifier.notify(usize::MAX); + } + + /// Create a new `Client` with default cache configuration. + /// + /// This is the standard constructor. Use [`Client::new_with_cache_config`] + /// if you need to customise cache TTL / capacity. + pub async fn new( + runtime: Arc, + persistence_manager: Arc, + transport_factory: Arc, + http_client: Arc, + override_version: Option<(u32, u32, u32)>, + ) -> (Arc, async_channel::Receiver) { + Self::new_with_cache_config( + runtime, + persistence_manager, + transport_factory, + http_client, + override_version, + CacheConfig::default(), + ) + .await + } + + /// Create a new `Client` with a custom [`CacheConfig`]. + pub async fn new_with_cache_config( + runtime: Arc, + persistence_manager: Arc, + transport_factory: Arc, + http_client: Arc, + override_version: Option<(u32, u32, u32)>, + cache_config: CacheConfig, + ) -> (Arc, async_channel::Receiver) { + let mut unique_id_bytes = [0u8; 2]; + rand::make_rng::().fill_bytes(&mut unique_id_bytes); + + let device_snapshot = persistence_manager.get_device_snapshot().await; + let core = wacore::client::CoreClient::new(device_snapshot.core.clone()); + + let (tx, rx) = async_channel::bounded(32); + + let this = Self { + runtime: runtime.clone(), + core, + persistence_manager: persistence_manager.clone(), + media_conn: Arc::new(RwLock::new(None)), + is_logged_in: Arc::new(AtomicBool::new(false)), + is_connecting: Arc::new(AtomicBool::new(false)), + is_running: Arc::new(AtomicBool::new(false)), + is_connected: Arc::new(AtomicBool::new(false)), + send_active_receipts: AtomicU32::new(0), + ik_handshake_failures: Arc::new(AtomicU32::new(0)), + shutdown_notifier: wacore::runtime::ShutdownNotifier::new(), + connection_shutdown: std::sync::Mutex::new(wacore::runtime::ShutdownNotifier::new()), + last_data_received_ms: Arc::new(AtomicU64::new(0)), + last_data_sent_ms: Arc::new(AtomicU64::new(0)), + + transport: Arc::new(Mutex::new(None)), + transport_events: Arc::new(Mutex::new(None)), + transport_factory, + noise_socket: Arc::new(Mutex::new(None)), + + response_waiters: Arc::new(Mutex::new(HashMap::new())), + node_waiters: std::sync::Mutex::new(Vec::new()), + node_waiter_count: AtomicUsize::new(0), + sent_node_waiters: std::sync::Mutex::new(Vec::new()), + sent_node_waiter_count: AtomicUsize::new(0), + unique_id: format!("{}.{}", unique_id_bytes[0], unique_id_bytes[1]), + id_counter: Arc::new(AtomicU64::new(0)), + unified_session: crate::unified_session::UnifiedSessionManager::new(), + + signal_cache: Arc::new(crate::store::signal_cache::SignalStoreCache::new()), + message_processing_semaphore: std::sync::Mutex::new(Arc::new( + async_lock::Semaphore::new(1), + )), + message_semaphore_generation: Arc::new(AtomicU64::new(0)), + // Coordination caches: capacity-only eviction, no TTL/TTI. + // These hold live mutexes and channel senders; time-based eviction + // while tasks hold references would silently break serialisation. + session_locks: Cache::builder() + .max_capacity(cache_config.session_locks_capacity.max(1)) + .build(), + chat_lanes: Cache::builder() + .max_capacity(cache_config.chat_lanes_capacity.max(1)) + .build(), + lid_pn_cache: Arc::new(LidPnCache::with_config( + &cache_config.lid_pn_cache, + cache_config.cache_stores.lid_pn_cache.clone(), + )), + ab_props: Arc::new(wacore::store::ab_props::AbPropsCache::new()), + group_cache: async_lock::Mutex::new(None), + + expected_disconnect: Arc::new(AtomicBool::new(false)), + intentional_reconnect: AtomicBool::new(false), + connection_generation: Arc::new(AtomicU64::new(0)), + + recent_messages: cache_config.recent_messages.build_with_ttl(), + + sender_key_device_cache: crate::sender_key_device_cache::SenderKeyDeviceCache::new( + &cache_config.sender_key_devices_cache, + ), + + pending_device_sync: crate::pending_device_sync::PendingDeviceSync::new(), + + pending_retries: Arc::new(std::sync::Mutex::new(HashSet::new())), + + message_retry_counts: cache_config.message_retry_counts.build_with_ttl(), + + recent_retry_reasons: cache_config.message_retry_counts.build_with_ttl(), + + session_recreate_history: cache_config.session_recreate_history.build_with_ttl(), + + undecryptable_dispatched: cache_config.undecryptable_dispatched.build_with_ttl(), + + offline_sync_metrics: Arc::new(OfflineSyncMetrics { + active: AtomicBool::new(false), + total_messages: AtomicUsize::new(0), + processed_messages: AtomicUsize::new(0), + start_time: std::sync::Mutex::new(None), + }), + offline_batch: Arc::new(crate::client::offline_resume::OfflineBatchCoordinator::new()), + + enable_auto_reconnect: Arc::new(AtomicBool::new(true)), + auto_reconnect_errors: Arc::new(AtomicU32::new(0)), + + needs_initial_full_sync: Arc::new(AtomicBool::new(false)), + + app_state_processor: async_lock::Mutex::new(None), + app_state_key_requests: Arc::new(Mutex::new(HashMap::new())), + app_state_syncing: Arc::new(Mutex::new(HashSet::new())), + initial_keys_synced_notifier: Arc::new(event_listener::Event::new()), + initial_app_state_keys_received: Arc::new(AtomicBool::new(false)), + prekey_upload_lock: Arc::new(async_lock::Mutex::new(())), + offline_sync_notifier: Arc::new(event_listener::Event::new()), + offline_sync_completed: Arc::new(AtomicBool::new(false)), + history_sync_tasks_in_flight: Arc::new(AtomicUsize::new(0)), + history_sync_idle_notifier: Arc::new(event_listener::Event::new()), + outbound_flush: Arc::new(crate::flush_scope::FlushScope::new()), + presence_subscriptions: Arc::new(async_lock::Mutex::new(HashSet::new())), + socket_ready_notifier: Arc::new(event_listener::Event::new()), + is_ready: Arc::new(AtomicBool::new(false)), + connected_notifier: Arc::new(event_listener::Event::new()), + major_sync_task_sender: tx, + pairing_cancellation_tx: Arc::new(Mutex::new(None)), + pair_code_state: Arc::new(Mutex::new(wacore::pair_code::PairCodeState::default())), + custom_enc_handlers: Arc::new(async_lock::RwLock::new(HashMap::new())), + chatstate_handlers: Arc::new(RwLock::new(Vec::new())), + pdo_pending_requests: cache_config.pdo_pending_requests.build_with_ttl(), + device_registry_cache: cache_config.device_registry_cache.build_typed_ttl( + cache_config.cache_stores.device_registry_cache.clone(), + "device_registry", + ), + stanza_router: Self::create_stanza_router(), + synchronous_ack: false, + http_client, + override_version, + skip_history_sync: AtomicBool::new(false), + wanted_pre_key_count: AtomicUsize::new(crate::prekeys::DEFAULT_WANTED_PRE_KEY_COUNT), + cache_config, + self_weak: std::sync::OnceLock::new(), + saver_handle: std::sync::OnceLock::new(), + raw_node_forwarding: AtomicBool::new(false), + }; + + let arc = Arc::new(this); + let _ = arc.self_weak.set(Arc::downgrade(&arc)); + + // Warm up the LID-PN cache from persistent storage + let warm_up_arc = arc.clone(); + arc.runtime + .spawn(Box::pin(async move { + if let Err(e) = warm_up_arc.warm_up_lid_pn_cache().await { + warn!("Failed to warm up LID-PN cache: {e}"); + } + })) + .detach(); + + // Start background task to clean up stale device registry entries + let cleanup_arc = arc.clone(); + arc.runtime + .spawn(Box::pin(async move { + cleanup_arc.device_registry_cleanup_loop().await; + })) + .detach(); + + (arc, rx) + } + + pub async fn run(self: &Arc) { + if self.is_running.swap(true, Ordering::SeqCst) { + warn!("Client `run` method called while already running."); + return; + } + while self.is_running.load(Ordering::Relaxed) { + self.expected_disconnect.store(false, Ordering::Relaxed); + + if let Err(connect_err) = self.connect().await { + let is_transient = connect_err + .downcast_ref::() + .is_some_and(|e| e.is_transient()); + if is_transient { + debug!("Transient connect failure, will retry: {connect_err:#}"); + } else { + error!("Failed to connect: {connect_err:#}. Will retry..."); + } + } else { + let unexpected_disconnect = if self.read_messages_loop().await.is_err() { + // Check intentional_reconnect AFTER read loop exits — reconnect() + // sets this flag while the loop is running, so it must be read here. + if self.expected_disconnect.load(Ordering::Relaxed) + || self.intentional_reconnect.swap(false, Ordering::Relaxed) + { + debug!("Message loop exited during expected disconnect."); + false + } else { + warn!( + "Message loop exited with an error. Will attempt to reconnect if enabled." + ); + true + } + } else if self.expected_disconnect.load(Ordering::Relaxed) { + debug!("Message loop exited gracefully (expected disconnect)."); + false + } else { + info!("Message loop exited gracefully."); + false + }; + + self.cleanup_connection_state().await; + + // Dispatch after cleanup so handlers see cleared connection state. + if unexpected_disconnect { + self.core + .event_bus + .dispatch(Event::Disconnected(crate::types::events::Disconnected)); + } + } + + if !self.enable_auto_reconnect.load(Ordering::Relaxed) { + info!("Auto-reconnect disabled, shutting down."); + self.is_running.store(false, Ordering::Relaxed); + break; + } + + // If this was an expected disconnect (e.g., 515 after pairing), reconnect immediately + if self.expected_disconnect.load(Ordering::Relaxed) { + self.auto_reconnect_errors.store(0, Ordering::Relaxed); + info!("Expected disconnect (e.g., 515), reconnecting immediately..."); + continue; + } + + let error_count = self.auto_reconnect_errors.fetch_add(1, Ordering::SeqCst); + // WA Web: Fibonacci backoff with 10% jitter, max 900s. + // algo: { type: "fibonacci", first: 1000, second: 1000 } + // jitter: 0.1, max: 9e5 + let delay = fibonacci_backoff(error_count); + info!( + "Will attempt to reconnect in {:?} (attempt {})", + delay, + error_count + 1 + ); + self.runtime.sleep(delay).await; + } + info!("Client run loop has shut down."); + } + + pub async fn connect(self: &Arc) -> Result<(), anyhow::Error> { + if self.is_connecting.swap(true, Ordering::SeqCst) { + return Err(ClientError::AlreadyConnected.into()); + } + + let _guard = scopeguard::guard((), |_| { + self.is_connecting.store(false, Ordering::Relaxed); + }); + + if self.is_connected() { + return Err(ClientError::AlreadyConnected.into()); + } + + // Reset login state for new connection attempt. This ensures that + // handle_success will properly process the stanza even if + // a previous connection's post-login task bailed out early. + self.is_logged_in.store(false, Ordering::Relaxed); + self.is_ready.store(false, Ordering::Relaxed); + self.is_connected.store(false, Ordering::Relaxed); + self.offline_sync_completed.store(false, Ordering::Relaxed); + self.offline_batch.reset(); + self.outbound_flush.reopen(); + + // WA Web: both MQTT and DGW transports use a 20s connect timeout. + // Without this, a dead network blocks on the OS TCP SYN timeout (~60-75s). + // Version fetch is also wrapped so a hung HTTP request doesn't block connect(). + let version_future = rt_timeout( + &*self.runtime, + TRANSPORT_CONNECT_TIMEOUT, + crate::version::resolve_and_update_version( + &self.persistence_manager, + &self.http_client, + self.override_version, + ), + ); + let transport_future = rt_timeout( + &*self.runtime, + TRANSPORT_CONNECT_TIMEOUT, + self.transport_factory.create_transport(), + ); + + debug!("Connecting WebSocket and fetching latest client version in parallel..."); + let (version_result, transport_result) = futures::join!(version_future, transport_future); + + version_result + .map_err(|_| anyhow!("Version fetch timed out after {TRANSPORT_CONNECT_TIMEOUT:?}"))? + .map_err(|e| anyhow!("Failed to resolve app version: {}", e))?; + let (transport, mut transport_events) = transport_result.map_err(|_| { + anyhow!("Transport connect timed out after {TRANSPORT_CONNECT_TIMEOUT:?}") + })??; + debug!("Version fetch and transport connection established."); + + let noise_socket = match handshake::do_handshake( + self.runtime.clone(), + &self.persistence_manager, + &self.ik_handshake_failures, + transport.clone(), + &mut transport_events, + ) + .await + { + Ok(socket) => socket, + Err(e) => { + transport.disconnect().await; + return Err(e.into()); + } + }; + + // Fresh per-connection shutdown so subscribers registered during this + // connection see a clean signal; the previous notifier was already + // fired on the prior cleanup_connection_state. + self.reset_connection_shutdown(); + + *self.transport.lock().await = Some(transport); + *self.transport_events.lock().await = Some(transport_events); + *self.noise_socket.lock().await = Some(noise_socket); + self.is_connected.store(true, Ordering::Release); + + // Notify waiters that socket is ready (before login) + self.socket_ready_notifier.notify(usize::MAX); + + let client_clone = self.clone(); + self.runtime + .spawn(Box::pin(async move { client_clone.keepalive_loop().await })) + .detach(); + + Ok(()) + } + + /// Deregister this companion device and disconnect. + /// Does NOT wipe stored keys. Delete the storage backend to fully clear credentials. + pub async fn logout(self: &Arc) -> Result<()> { + use wacore::iq::devices::RemoveCompanionDeviceSpec; + + self.enable_auto_reconnect.store(false, Ordering::Relaxed); + + if self.is_connected() + && let Ok(jid) = self.require_pn().await + && let Err(e) = self.execute(RemoveCompanionDeviceSpec::new(&jid)).await + { + warn!("Failed to send logout IQ: {e}"); + } + + self.disconnect().await; + + self.core + .event_bus + .dispatch(Event::LoggedOut(crate::types::events::LoggedOut { + on_connect: false, + reason: ConnectFailureReason::LoggedOut, + })); + + Ok(()) + } + + pub async fn disconnect(self: &Arc) { + info!("Disconnecting client intentionally."); + self.expected_disconnect.store(true, Ordering::Relaxed); + self.is_running.store(false, Ordering::Relaxed); + self.shutdown_notifier.notify(); + + // Prevent late receipt producers from escaping the drain window. + self.outbound_flush.close(); + self.outbound_flush + .flush(&*self.runtime, std::time::Duration::from_secs(5)) + .await; + self.notify_connection_shutdown(); + + if let Err(e) = self.persistence_manager.flush().await { + log::error!("Failed to flush device state during disconnect: {e}"); + } + + // Close after flush; cleanup may also win this race on the run loop. + if let Some(transport) = self.transport.lock().await.as_ref() { + transport.disconnect().await; + } + self.cleanup_connection_state().await; + } + + /// Backoff step used by [`reconnect()`] to create an offline window. + /// + /// `fibonacci_backoff(RECONNECT_BACKOFF_STEP)` determines the delay before + /// the run loop re-connects. This must be longer than the mock server's + /// chatstate TTL (`CHATSTATE_TTL_SECS=3`) so TTL-expiry tests pass. + /// + /// Sequence: fib(0)=1s, fib(1)=1s, fib(2)=2s, fib(3)=3s, **fib(4)=5s**. + pub const RECONNECT_BACKOFF_STEP: u32 = 4; + + /// Drop the current connection and trigger the auto-reconnect loop. + /// + /// Unlike [`disconnect`], this does **not** stop the run loop. The client + /// will reconnect automatically using the same persisted identity/store, + /// just as it would after a network interruption. Use + /// [`wait_for_connected`] to wait for the new connection to be ready. + /// + /// This is useful for: + /// - Handling network changes (e.g., Wi-Fi → cellular) + /// - Forcing a fresh server session + /// - Testing offline message delivery + pub async fn reconnect(self: &Arc) { + info!("Reconnecting: dropping transport for auto-reconnect."); + self.intentional_reconnect.store(true, Ordering::Relaxed); + self.auto_reconnect_errors + .store(Self::RECONNECT_BACKOFF_STEP, Ordering::Relaxed); + + self.outbound_flush.close(); + self.outbound_flush + .flush(&*self.runtime, std::time::Duration::from_secs(2)) + .await; + self.notify_connection_shutdown(); + + if let Some(transport) = self.transport.lock().await.as_ref() { + transport.disconnect().await; + } + } + + /// Drop the current connection and reconnect immediately with no delay. + /// + /// Unlike [`reconnect`], which introduces a deliberate offline window, + /// this method sets the `expected_disconnect` flag so the run loop + /// skips the backoff delay and reconnects as fast as possible. + pub async fn reconnect_immediately(self: &Arc) { + info!("Reconnecting immediately (expected disconnect)."); + self.expected_disconnect.store(true, Ordering::Relaxed); + + self.outbound_flush.close(); + self.outbound_flush + .flush(&*self.runtime, std::time::Duration::from_secs(2)) + .await; + self.notify_connection_shutdown(); + + if let Some(transport) = self.transport.lock().await.as_ref() { + transport.disconnect().await; + } + } + + pub(crate) async fn cleanup_connection_state(&self) { + // Note: node_waiters are intentionally NOT cleared here — they are + // cross-connection (callers may register a waiter before an action that + // completes on a subsequent connection, e.g. after 515 reconnect). + // sent_node_waiters ARE cleared because they match pre-encryption + // outgoing stanzas, which are transport-scoped. + self.clear_sent_node_waiters(); + self.is_logged_in.store(false, Ordering::Relaxed); + self.is_ready.store(false, Ordering::Relaxed); + // Signal the keepalive loop (and any other per-connection tasks) to + // exit promptly. Without this, a stale keepalive loop can overlap + // with the next one after reconnect. Uses the PER-CONNECTION signal + // so the terminal shutdown_notifier stays clean for reconnects. + self.notify_connection_shutdown(); + // Close the socket as part of cleanup so this path is authoritative + // even when reached via the run loop's graceful-exit flow (not just + // `Client::disconnect()`). Transport impls make `disconnect()` + // idempotent, so the redundant call from `Client::disconnect()` is + // safe. + if let Some(transport) = self.transport.lock().await.take() { + transport.disconnect().await; + } + *self.transport_events.lock().await = None; + *self.noise_socket.lock().await = None; + // Clear is_connected AFTER noise_socket is None, so no task can see + // is_connected==true with a cleared socket. send_node() independently + // checks the socket, but this ordering avoids a confusing state window. + self.is_connected.store(false, Ordering::Release); + // Presence doesn't survive reconnects: demote presence-driven active + // receipts (1 -> 0), leaving a forced value (2) untouched. + let _ = + self.send_active_receipts + .compare_exchange(1, 0, Ordering::AcqRel, Ordering::Acquire); + // Drop per-chat lanes so workers exit via channel close. + self.chat_lanes.invalidate_all(); + // Clear pending retries so stale keys from detached scopeguard + // cleanup don't suppress the first retry after reconnect. + self.pending_retries + .lock() + .unwrap_or_else(|p| p.into_inner()) + .clear(); + // Flush before clear: clear() drops dirty entries, so a disconnect + // racing an in-flight encrypt would lose the just-advanced sender-key + // chain and force a full SKDM re-fanout. A disconnect is not a logout. + // Only clear on a successful flush; on a backend error keep the cache so + // the dirty state isn't dropped and the next operation can persist it. + match self.flush_signal_cache().await { + Ok(()) => self.signal_cache.clear().await, + Err(e) => log::error!( + "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" + ), + } + // Reset semaphore to 1 permit for next offline sync. + self.swap_message_semaphore(1); + // Reset dead-socket timestamps so stale values from the previous + // connection don't trigger an immediate reconnect on the next one. + self.last_data_received_ms.store(0, Ordering::Relaxed); + self.last_data_sent_ms.store(0, Ordering::Relaxed); + self.pending_device_sync.clear().await; + // Reset offline sync state for next connection + self.offline_sync_completed.store(false, Ordering::Relaxed); + self.offline_batch.reset(); + self.offline_sync_metrics + .active + .store(false, Ordering::Release); + self.offline_sync_metrics + .total_messages + .store(0, Ordering::Release); + self.offline_sync_metrics + .processed_messages + .store(0, Ordering::Release); + match self.offline_sync_metrics.start_time.lock() { + Ok(mut guard) => *guard = None, + Err(poison) => *poison.into_inner() = None, + } + self.history_sync_tasks_in_flight + .store(0, Ordering::Relaxed); + self.history_sync_idle_notifier.notify(usize::MAX); + // Drain all pending IQ waiters so they fail fast with InternalChannelClosed + // instead of hanging until the 75s timeout. + let mut waiters_map = self.response_waiters.lock().await; + let waiter_count = waiters_map.len(); + // Replace with new map to release backing storage; old senders drop here, + // causing receivers to get RecvError → IqError::InternalChannelClosed + *waiters_map = HashMap::new(); + drop(waiters_map); + if waiter_count > 0 { + debug!( + "Dropping {} orphaned IQ response waiter(s) on disconnect", + waiter_count + ); + } + + // Clear app state tracking maps to prevent unbounded growth across reconnections. + // Replace with new collections to release backing storage. + *self.app_state_key_requests.lock().await = HashMap::new(); + *self.app_state_syncing.lock().await = HashSet::new(); + + // Drop stale media connection (auth tokens become invalid on reconnect) + *self.media_conn.write().await = None; + + // Clear app state key cache — keys will be re-fetched from DB on demand + if let Some(proc) = self.app_state_processor.lock().await.as_ref() { + proc.clear_key_cache().await; + } + } + + /// Waits for the noise socket to be established. + /// + /// Returns `Ok(())` when the socket is ready, or `Err` on timeout. + /// This is useful for code that needs to send messages before login, + /// such as requesting a pair code during initial pairing. + /// + /// If the socket is already connected, returns immediately. + pub async fn wait_for_socket(&self, timeout: std::time::Duration) -> Result<(), anyhow::Error> { + // Fast path: already connected + if self.is_connected() { + return Ok(()); + } + + // Register waiter and re-check to avoid race condition: + // If socket becomes ready between checks, the notified future captures it. + let notified = self.socket_ready_notifier.listen(); + if self.is_connected() { + return Ok(()); + } + + rt_timeout(&*self.runtime, timeout, notified) + .await + .map_err(|_| anyhow::anyhow!("Timeout waiting for socket")) + } + + /// Waits for the client to establish a connection and complete login. + /// + /// Returns `Ok(())` when connected, or `Err` on timeout. + /// This is useful for code that needs to run after connection is established + /// and authentication is complete. + /// + /// If the client is already connected and logged in, returns immediately. + pub async fn wait_for_connected( + &self, + timeout: std::time::Duration, + ) -> Result<(), anyhow::Error> { + // Fast path: fully ready (connected + logged in + critical sync done). + if self.is_fully_ready() { + return Ok(()); + } + + // Register waiter and re-check to avoid TOCTOU race: + // dispatch_connected() could fire between the check above and notified() registration. + let notified = self.connected_notifier.listen(); + if self.is_fully_ready() { + return Ok(()); + } + + rt_timeout(&*self.runtime, timeout, notified) + .await + .map_err(|_| anyhow::anyhow!("Timeout waiting for connection")) + } + + pub fn is_connected(&self) -> bool { + self.is_connected.load(Ordering::Acquire) + } + + pub fn is_logged_in(&self) -> bool { + self.is_logged_in.load(Ordering::Relaxed) + } +} diff --git a/src/client/messaging.rs b/src/client/messaging.rs new file mode 100644 index 000000000..d060c1527 --- /dev/null +++ b/src/client/messaging.rs @@ -0,0 +1,278 @@ +//! Outgoing send primitives, receipts, reactions, edits and chat-state events. + +use super::*; + +impl Client { + /// Send pre-marshaled plaintext bytes through the noise socket. + /// + /// The bytes must be a valid WABinary-marshaled stanza (as produced by + /// `wacore_binary::marshal::marshal_to`). Sending malformed data will + /// cause the server to close the connection. + /// + /// This bypasses node logging and `sent_node_waiter` resolution — use + /// [`send_node`](Client::send_node) for normal stanza sending. + pub async fn send_raw_bytes(&self, plaintext: Vec) -> Result<(), ClientError> { + let noise_socket = self.get_noise_socket().await?; + noise_socket + .encrypt_and_send(bytes::Bytes::from(plaintext)) + .await?; + self.last_data_sent_ms + .store(wacore::time::now_millis().max(0) as u64, Ordering::Relaxed); + Ok(()) + } + + pub async fn send_node(&self, node: Node) -> Result<(), ClientError> { + debug!(target: "Client/Send", "{}", DisplayableNode(&node)); + if self.sent_node_waiter_count.load(Ordering::Acquire) > 0 { + self.resolve_sent_node_waiters(&Arc::new(node.clone())); + } + + let plaintext_buf = wacore_binary::marshal::marshal_auto(&node).map_err(|e| { + error!("Failed to marshal node: {e:?}"); + SocketError::Marshal(e) + })?; + + self.send_raw_bytes(plaintext_buf).await + } + + pub(crate) async fn send_unified_session(&self) { + if !self.is_connected() { + debug!(target: "Client/UnifiedSession", "Skipping: not connected"); + return; + } + + let Some((node, _sequence)) = self.unified_session.prepare_send().await else { + return; + }; + + if let Err(e) = self.send_node(node).await { + debug!(target: "Client/UnifiedSession", "Send failed: {e}"); + self.unified_session.clear_last_sent().await; + } + } + + pub async fn edit_message( + &self, + to: Jid, + original_id: impl Into, + new_content: wa::Message, + ) -> Result { + let original_id = original_id.into(); + + // WhatsApp Web uses getMeUserLidOrJidForChat(chat, EditMessage) which + // returns LID for LID-addressing groups and PN otherwise. + let participant = if to.is_group() { + Some( + self.get_own_jid_for_group(&to) + .await? + .to_non_ad() + .to_string(), + ) + } else { + if self.get_pn().await.is_none() { + return Err(anyhow::Error::from(ClientError::NotLoggedIn)); + } + None + }; + + let edit_container_message = crate::send::build_edit_message( + &to, + original_id.clone(), + participant, + new_content, + wacore::time::now_millis(), + ); + + // Use a new stanza ID instead of reusing the original message ID. + // The original message ID is already embedded in protocolMessage.key.id + // inside the encrypted payload. Reusing it as the outer stanza ID causes + // the server to deduplicate against the original message and silently + // drop the edit. + self.send_message_impl( + to, + &edit_container_message, + None, + false, + false, + Some(crate::types::message::EditAttribute::MessageEdit), + vec![], + None, + ) + .await?; + + Ok(original_id) + } + + /// Send a server-side reaction (used by both newsletter and status reactions). + pub(crate) async fn send_server_reaction( + &self, + to: &Jid, + server_id: u64, + reaction: &str, + ) -> Result<(), anyhow::Error> { + let request_id = self.generate_message_id().await; + + let stanza = NodeBuilder::new("message") + .attr("to", to) + .attr("type", "reaction") + .attr("id", &request_id) + .attr("server_id", server_id) + .children([NodeBuilder::new("reaction").attr("code", reaction).build()]) + .build(); + + self.send_node(stanza).await?; + Ok(()) + } + + /// Register a oneshot waiter for a server ack by message ID. + /// Returns the receiver — caller sends the node separately and awaits this in background. + pub(crate) async fn register_ack_waiter( + &self, + message_id: &str, + ) -> futures::channel::oneshot::Receiver> { + let (tx, rx) = futures::channel::oneshot::channel(); + self.response_waiters + .lock() + .await + .insert(message_id.to_string(), tx); + rx + } + + /// Creates a normalized ChatMessageId by resolving PN to LID JIDs. + pub(crate) async fn make_chat_message_id(&self, chat: &Jid, id: &str) -> ChatMessageId { + // Resolve chat JID to LID if possible + let chat = self.resolve_encryption_jid(chat).await; + + ChatMessageId { + chat, + id: id.to_owned(), + } + } + + pub(crate) async fn send_protocol_receipt( + &self, + id: String, + receipt_type: crate::types::presence::ReceiptType, + ) { + if id.is_empty() { + return; + } + let device_snapshot = self.persistence_manager.get_device_snapshot().await; + if let Some(own_jid) = &device_snapshot.pn { + // Single source of truth for the wire mapping (ReceiptType::Sent is a derived + // incoming-only state and is never sent by us). + let type_str = receipt_type.as_wire_str(); + + // Borrow `id` for the attr so it stays available for the error log + // below (the warn used to log self.unique_id, the client UUID, by + // mistake). Separate .attr calls avoid cloning into a homogeneous map. + let node = NodeBuilder::new("receipt") + .attr("id", id.as_str()) + .attr("type", type_str) + .attr("to", own_jid.to_non_ad_string()) + .build(); + + if let Err(e) = self.send_node(node).await { + warn!( + "Failed to send protocol receipt of type {:?} for message ID {}: {:?}", + receipt_type, id, e + ); + } + } + } + + /// Register a chatstate handler which will be invoked when a `` stanza is received. + /// + /// The handler receives a `ChatStateEvent` with the parsed chat state information. + pub async fn register_chatstate_handler( + &self, + handler: Arc, + ) { + self.chatstate_handlers.write().await.push(handler); + } + + /// Dispatch a parsed chatstate stanza to registered handlers. + /// + /// Called by `ChatstateHandler` after parsing the incoming stanza. + pub(crate) async fn dispatch_chatstate_event( + &self, + stanza: wacore::iq::chatstate::ChatstateStanza, + ) { + use wacore::iq::chatstate::{ChatstateSource, ReceivedChatState}; + use wacore::types::events::ChatPresenceUpdate; + use wacore::types::message::MessageSource; + use wacore::types::presence::{ChatPresence, ChatPresenceMedia}; + + // Dispatch via event bus + let (chat, sender, is_group) = match &stanza.source { + ChatstateSource::User { from } => (from.clone(), from.clone(), false), + ChatstateSource::Group { from, participant } => { + (from.clone(), participant.clone(), true) + } + }; + + let (state, media) = match stanza.state { + ReceivedChatState::Typing => (ChatPresence::Composing, ChatPresenceMedia::Text), + ReceivedChatState::RecordingAudio => { + (ChatPresence::Composing, ChatPresenceMedia::Audio) + } + ReceivedChatState::Idle => (ChatPresence::Paused, ChatPresenceMedia::Text), + }; + + self.core + .event_bus + .dispatch(Event::ChatPresence(ChatPresenceUpdate { + source: MessageSource { + chat, + sender, + is_from_me: false, + is_group, + addressing_mode: None, + sender_alt: None, + recipient_alt: None, + broadcast_list_owner: None, + recipient: None, + }, + state, + media, + })); + + // Invoke legacy callback handlers + let event = ChatStateEvent::from_stanza(stanza); + let handlers = self.chatstate_handlers.read().await.clone(); + for handler in handlers { + let event_clone = event.clone(); + self.runtime + .spawn(Box::pin(async move { + (handler)(event_clone); + })) + .detach(); + } + } + + /// Whether delivery receipts should be sent active (rendered as ticks) vs + /// `type="inactive"`. Mirrors whatsmeow's `sendActiveReceipts != 0`. + pub(crate) fn receipts_are_active(&self) -> bool { + self.send_active_receipts.load(Ordering::Acquire) != 0 + } + + /// Force active delivery receipts even when offline (whatsmeow's + /// `SetForceActiveDeliveryReceipts`); off restores the default. + pub fn set_force_active_delivery_receipts(&self, active: bool) { + self.send_active_receipts + .store(if active { 2 } else { 0 }, Ordering::Release); + } + + /// CAS so a forced value (2) is preserved (whatsmeow's `CompareAndSwap`). + pub(crate) fn mark_receipts_active_on_presence(&self) { + let _ = + self.send_active_receipts + .compare_exchange(0, 1, Ordering::AcqRel, Ordering::Acquire); + } + + pub(crate) fn mark_receipts_inactive_on_presence(&self) { + let _ = + self.send_active_receipts + .compare_exchange(1, 0, Ordering::AcqRel, Ordering::Acquire); + } +} diff --git a/src/client/node_io.rs b/src/client/node_io.rs new file mode 100644 index 000000000..5f0bf534a --- /dev/null +++ b/src/client/node_io.rs @@ -0,0 +1,1187 @@ +//! Inbound node I/O: read loop, frame decryption, node routing, acks and stream errors. + +use super::*; + +impl Client { + /// Read the current semaphore generation and Arc atomically under the mutex. + pub(crate) fn read_message_semaphore(&self) -> (u64, Arc) { + let guard = match self.message_processing_semaphore.lock() { + Ok(g) => g, + Err(poisoned) => poisoned.into_inner(), + }; + ( + self.message_semaphore_generation.load(Ordering::SeqCst), + guard.clone(), + ) + } + + /// Replace the message processing semaphore and bump the generation counter. + /// + /// Both operations happen under the same mutex hold so readers always see + /// a consistent (generation, Arc) pair. Must be called from a non-async + /// context or inside a scoped block (MutexGuard is !Send). + pub(crate) fn swap_message_semaphore(&self, permits: usize) { + let mut guard = match self.message_processing_semaphore.lock() { + Ok(g) => g, + Err(poisoned) => poisoned.into_inner(), + }; + *guard = Arc::new(async_lock::Semaphore::new(permits)); + self.message_semaphore_generation + .fetch_add(1, Ordering::SeqCst); + } + + pub(crate) async fn read_messages_loop(self: &Arc) -> Result<(), anyhow::Error> { + debug!("Starting message processing loop..."); + + let mut rx_guard = self.transport_events.lock().await; + let transport_events = rx_guard + .take() + .ok_or_else(|| anyhow::anyhow!("Cannot start message loop: not connected"))?; + drop(rx_guard); + + // Frame decoder to parse incoming data + let mut frame_decoder = wacore::framing::FrameDecoder::new(); + let shutdown = self.connection_shutdown_signal(); + + loop { + futures::select_biased! { + _ = wacore::runtime::wait_for_shutdown(&shutdown).fuse() => { + debug!("Shutdown signaled in message loop. Exiting message loop."); + return Ok(()); + }, + event_result = transport_events.recv().fuse() => { + match event_result { + Ok(crate::transport::TransportEvent::DataReceived(data)) => { + // Update dead-socket timer (WA Web: deadSocketTimer reset) + self.last_data_received_ms.store( + wacore::time::now_millis().max(0) as u64, + Ordering::Relaxed, + ); + + // Feed data into the frame decoder + frame_decoder.feed(&data); + + // Process all complete frames. + // Frame decryption must be sequential (noise protocol counter), + // but we spawn node processing concurrently after decryption. + let mut frames_in_batch: u32 = 0; + + while let Some(encrypted_frame) = frame_decoder.decode_frame() { + // Decrypt the frame synchronously (required for noise counter ordering) + if let Some(node) = self.decrypt_frame(encrypted_frame).await { + // Determine processing mode for this node: + // - Critical nodes (success/failure/stream:error): inline, required for state + // - Message nodes: inline, preserves arrival order for per-chat queues + // (MessageHandler just enqueues + ACKs, heavy crypto runs in workers) + // - ib (in-band): inline, ensures offline sync tracking (expected count) + // is set up before offline messages are processed + // - Everything else: spawned concurrently for parallelism + let process_inline = matches!( + node.tag(), + "success" | "failure" | "stream:error" | "message" | "ib" + ); + + if process_inline { + self.process_decrypted_node(node).await; + } else { + let client = self.clone(); + self.runtime.spawn(Box::pin(async move { + client.process_decrypted_node(node).await; + })).detach(); + } + } + + // Check if we should exit after processing (e.g., after 515 stream error) + if self.expected_disconnect.load(Ordering::Relaxed) { + debug!("Expected disconnect signaled during frame processing. Exiting message loop."); + return Ok(()); + } + + // Cooperative yield — frequency and behavior are runtime-defined. + frames_in_batch += 1; + if frames_in_batch.is_multiple_of(self.runtime.yield_frequency()) + && let Some(yield_fut) = self.runtime.yield_now() + { + yield_fut.await; + } + } + + // Refresh timestamp after processing the entire batch so + // the keepalive loop sees the batch completion time, not + // just the arrival time. Prevents stale reads when a + // large batch (e.g. offline sync) takes seconds to drain. + if frames_in_batch > 1 { + self.last_data_received_ms.store( + wacore::time::now_millis().max(0) as u64, + Ordering::Relaxed, + ); + } + }, + Ok(crate::transport::TransportEvent::Disconnected(reason)) => { + if !self.expected_disconnect.load(Ordering::Relaxed) { + debug!("Transport disconnected unexpectedly: {reason}"); + return Err(anyhow::anyhow!("Transport disconnected: {reason}")); + } else { + debug!("Transport disconnected as expected: {reason}"); + return Ok(()); + } + } + // Event channel closed (no DisconnectReason available). + Err(_) => { + if !self.expected_disconnect.load(Ordering::Relaxed) { + debug!("Transport event channel closed unexpectedly."); + return Err(anyhow::anyhow!("Transport event channel closed")); + } else { + return Ok(()); + } + } + Ok(crate::transport::TransportEvent::Connected) => { + // Already handled during handshake, but could be useful for logging + debug!("Transport connected event received"); + } + } + } + } + } + } + + /// Decrypt a frame and return the parsed node as a zero-copy OwnedNodeRef. + /// This must be called sequentially due to noise protocol counter requirements. + pub(crate) async fn decrypt_frame( + self: &Arc, + encrypted_frame: bytes::BytesMut, + ) -> Option { + let noise_socket = match self.get_noise_socket().await { + Ok(s) => s, + Err(_) => { + log::error!("Cannot process frame: not connected (no noise socket)"); + return None; + } + }; + + let decrypted_payload = match noise_socket.decrypt_frame(encrypted_frame) { + Ok(p) => p, + Err(e) => { + log::error!("Failed to decrypt frame: {e}"); + return None; + } + }; + + let buffer = match wacore_binary::util::unpack_bytes(decrypted_payload) { + Ok(data) => data, + Err(e) => { + log::warn!(target: "Client/Recv", "Failed to decompress frame: {e}"); + return None; + } + }; + + match wacore_binary::OwnedNodeRef::new(buffer) { + Ok(owned) => Some(owned), + Err(e) => { + log::warn!(target: "Client/Recv", "Failed to unmarshal node: {e}"); + None + } + } + } + + /// Process an already-decrypted node. + /// This can be spawned concurrently since it doesn't depend on noise protocol state. + /// The node is wrapped in Arc to avoid cloning when passing through handlers. + pub(crate) async fn process_decrypted_node( + self: &Arc, + node: wacore_binary::OwnedNodeRef, + ) { + // Wrap in Arc once - all handlers will share this same allocation + let node_arc = Arc::new(node); + self.process_node(node_arc).await; + } + + /// Process a node wrapped in Arc. Handlers receive the Arc and can share/store it cheaply. + pub(crate) async fn process_node(self: &Arc, node: Arc) { + use wacore::xml::DisplayableNodeRef; + let nr = node.get(); + + // --- Offline Sync Tracking --- + if nr.tag.as_ref() == "ib" { + // Check for offline_preview child to get expected count + if let Some(preview) = nr.get_optional_child("offline_preview") { + let count: usize = preview + .get_attr("count") + .map(|v| v.as_str()) + .and_then(|s| s.parse().ok()) + .unwrap_or(0); + + if count == 0 { + self.offline_sync_metrics + .active + .store(false, Ordering::Release); + debug!(target: "Client/OfflineSync", "Sync COMPLETED: 0 items."); + } else { + // Use stronger memory ordering for state transitions + self.offline_sync_metrics + .total_messages + .store(count, Ordering::Release); + self.offline_sync_metrics + .processed_messages + .store(0, Ordering::Release); + self.offline_sync_metrics + .active + .store(true, Ordering::Release); + match self.offline_sync_metrics.start_time.lock() { + Ok(mut guard) => *guard = Some(wacore::time::Instant::now()), + Err(poison) => *poison.into_inner() = Some(wacore::time::Instant::now()), + } + debug!(target: "Client/OfflineSync", "Sync STARTED: Expecting {} items.", count); + } + } else if self.offline_sync_metrics.active.load(Ordering::Acquire) + && nr.get_optional_child("offline").is_some() + { + // Handle end marker: signals sync completion + // Only with an child is a real end marker. + // Other children (thread_metadata, edge_routing, dirty) are NOT end markers. + let processed = self + .offline_sync_metrics + .processed_messages + .load(Ordering::Acquire); + let elapsed = match self.offline_sync_metrics.start_time.lock() { + Ok(guard) => guard.map(|t| t.elapsed()).unwrap_or_default(), + Err(poison) => poison.into_inner().map(|t| t.elapsed()).unwrap_or_default(), + }; + debug!(target: "Client/OfflineSync", "Sync COMPLETED: End marker received. Processed {} items in {:.2?}.", processed, elapsed); + self.offline_sync_metrics + .active + .store(false, Ordering::Release); + } + } + + // Track progress if active + if self.offline_sync_metrics.active.load(Ordering::Acquire) { + // Check for 'offline' attribute on relevant stanzas + if nr.get_attr("offline").is_some() { + let processed = self + .offline_sync_metrics + .processed_messages + .fetch_add(1, Ordering::Release) + + 1; + let total = self + .offline_sync_metrics + .total_messages + .load(Ordering::Acquire); + + if processed.is_multiple_of(50) || processed == total { + trace!(target: "Client/OfflineSync", "Sync Progress: {}/{}", processed, total); + } + + // Drive WA Web pull-batch loop (non-adaptive `$13`): when + // remaining drops to <=C and no batch request is in flight, + // schedule the next one. + let pending = total.saturating_sub(processed); + crate::client::offline_resume::on_offline_stanza_arrived(self, pending); + + if processed >= total { + let elapsed = match self.offline_sync_metrics.start_time.lock() { + Ok(guard) => guard.map(|t| t.elapsed()).unwrap_or_default(), + Err(poison) => poison.into_inner().map(|t| t.elapsed()).unwrap_or_default(), + }; + debug!(target: "Client/OfflineSync", "Sync COMPLETED: Processed {} items in {:.2?}.", processed, elapsed); + self.offline_sync_metrics + .active + .store(false, Ordering::Release); + } + } + } + // --- End Tracking --- + + if nr.tag.as_ref() == "iq" + && let Some(sync_node) = nr.get_optional_child("sync") + && let Some(collection_node) = sync_node.get_optional_child("collection") + { + let name = collection_node.attrs().optional_string("name"); + let name = name.as_deref().unwrap_or(""); + debug!(target: "Client/Recv", "Received app state sync response for '{name}' (hiding content)."); + } else { + debug!(target: "Client/Recv","{}", DisplayableNodeRef(nr)); + } + + // Prepare deferred ACK cancellation flag (sent after dispatch unless cancelled) + let mut cancelled = false; + + // Emit raw node before any early returns so all decoded stanzas + // (including IQ responses and xmlstreamend) reach external observers + if self.raw_node_forwarding.load(Ordering::Relaxed) { + self.core + .event_bus + .dispatch(Event::RawNode(Arc::clone(&node))); + } + + if nr.tag.as_ref() == "xmlstreamend" { + if self.expected_disconnect.load(Ordering::Relaxed) { + debug!("Received , expected disconnect."); + } else { + warn!("Received , treating as disconnect."); + } + self.notify_connection_shutdown(); + return; + } + + // Check generic node waiters (zero-cost when none registered) + if self.node_waiter_count.load(Ordering::Acquire) > 0 { + self.resolve_node_waiters(&node); + } + + if nr.tag.as_ref() == "iq" + && let Some(id) = nr.get_attr("id").map(|v| v.as_str()) + { + // Single lock acquisition: try to remove the waiter directly. + let waiter = self.response_waiters.lock().await.remove(id.as_ref()); + if let Some(waiter) = waiter { + if waiter.send(Arc::clone(&node)).is_err() { + warn!(target: "Client/IQ", "Failed to send IQ response to waiter. Receiver was likely dropped."); + } + return; + } + } + + // Dispatch to appropriate handler using the router + // Clone Arc (cheap - just reference count) not the Node itself + if !self + .stanza_router + .dispatch(self.clone(), Arc::clone(&node), &mut cancelled) + .await + { + warn!( + "Received unknown top-level node: {}", + DisplayableNodeRef(nr) + ); + } + + // Send the deferred ACK if applicable and not cancelled by handler + if self.should_ack(nr) && !cancelled { + self.maybe_deferred_ack(node).await; + } + } + + /// Per WA Web (`Handle/MsgSendReceipt.js`), only newsletter `` + /// gets `` on the success path; DM/group use + /// ``. Failure paths (retry/backfill/nack) emit `` from + /// their dedicated handlers, not via this gate. + /// + /// status@broadcast is included as a fallback: drop paths in + /// `process_group_enc_batch` (expired status, missing sender key, generic + /// decrypt error) intentionally skip the delivery receipt to avoid + /// inflating the server-side offline counter for messages we'll never + /// process. Without the transport `` from this gate, the server + /// would redeliver indefinitely. WA Web emits `` + /// in the success path on top of this; the duplicate is tolerated. + pub(crate) fn should_ack(&self, node: &wacore_binary::NodeRef<'_>) -> bool { + let tag = node.tag.as_ref(); + if node.get_attr("id").is_none() { + return false; + } + let Some(from) = node.get_attr("from") else { + return false; + }; + match tag { + "receipt" | "notification" | "call" => true, + "message" => from + .to_jid() + .is_some_and(|j| j.is_newsletter() || j.is_status_broadcast()), + _ => false, + } + } + + /// Possibly send a deferred ack: either immediately or via spawned task. + /// Handlers can cancel by setting `cancelled` to true. + /// Uses Arc to avoid cloning when spawning the async task. + async fn maybe_deferred_ack(self: &Arc, node: Arc) { + if self.synchronous_ack { + if let Err(e) = self.send_ack_for(node.get()).await + && !e.is_transport_unavailable() + { + warn!("Failed to send ack: {e:?}"); + } + } else { + let this = self.clone(); + self.runtime + .spawn(Box::pin(async move { + if let Err(e) = this.send_ack_for(node.get()).await + && !e.is_transport_unavailable() + { + warn!("Failed to send ack: {e:?}"); + } + })) + .detach(); + } + } + + /// Build and send an node corresponding to the given stanza. + pub(crate) async fn send_ack_for( + &self, + node: &wacore_binary::NodeRef<'_>, + ) -> Result<(), ClientError> { + if self.expected_disconnect.load(Ordering::Relaxed) { + return Ok(()); + } + if !self.is_connected() { + return Err(ClientError::NotConnected); + } + let own_pn = self.get_pn().await; + let buf = match encode_ack_bytes(node, own_pn.as_ref()) { + Ok(Some(buf)) => buf, + Ok(None) => return Ok(()), + Err(e) => { + log::warn!("Failed to encode ack: {e}"); + return Ok(()); + } + }; + self.send_raw_bytes(buf).await + } + + /// Send a transport ack so the server stops replaying a stanza from the + /// offline queue. Awaitable so callers can order it after a retry receipt + /// in a single flushed task. + pub(crate) async fn send_transport_ack(&self, info: &crate::types::message::MessageInfo) { + let source = message_ack_source_node(info); + let own_pn = self.get_pn().await; + match encode_ack_bytes(&source.as_node_ref(), own_pn.as_ref()) { + Ok(Some(buf)) => { + if let Err(e) = self.send_raw_bytes(buf).await + && !e.is_transport_unavailable() + { + log::warn!("Failed to send transport ack for undecryptable message: {e:?}"); + } + } + Ok(None) => {} + Err(e) => log::warn!("Failed to encode transport ack: {e}"), + } + } + + /// Spawn [`Self::send_transport_ack`], tracked via `outbound_flush` so + /// `disconnect()` flushes it (issue #571), same as delivery receipts. + pub(crate) fn spawn_message_ack( + self: &Arc, + info: &Arc, + ) { + let client = Arc::clone(self); + let info = Arc::clone(info); + self.outbound_flush.spawn(&*self.runtime, async move { + client.send_transport_ack(&info).await; + }); + } + + /// Tracked ack encoded from the original node. Use when the stanza carries + /// `recipient` (LID-routed/hosted-companion/peer) since `MessageInfo` + /// drops it on non-self branches and the server needs it for routing. + pub(crate) async fn spawn_node_transport_ack( + self: &Arc, + node: &wacore_binary::NodeRef<'_>, + ) { + let own_pn = self.get_pn().await; + let buf = match encode_ack_bytes(node, own_pn.as_ref()) { + Ok(Some(b)) => b, + Ok(None) => return, + Err(e) => { + log::warn!("Failed to encode node transport ack: {e}"); + return; + } + }; + let client = Arc::clone(self); + self.outbound_flush.spawn(&*self.runtime, async move { + if let Err(e) = client.send_raw_bytes(buf).await + && !e.is_transport_unavailable() + { + log::warn!("Failed to send node transport ack: {e:?}"); + } + }); + } + + pub(crate) async fn handle_success(self: &Arc, node: &wacore_binary::NodeRef<'_>) { + // Skip processing if an expected disconnect is pending (e.g., 515 received). + // This prevents race conditions where a spawned success handler runs after + // cleanup_connection_state has already reset is_logged_in. + if self.expected_disconnect.load(Ordering::Relaxed) { + debug!("Ignoring stanza: expected disconnect pending"); + return; + } + + // Guard against multiple stanzas (WhatsApp may send more than one during + // routing/reconnection). Only process the first one per connection. + if self.is_logged_in.swap(true, Ordering::SeqCst) { + debug!("Ignoring duplicate stanza (already logged in)"); + return; + } + + // Increment connection generation to invalidate any stale post-login tasks + // from previous connections (e.g., during 515 reconnect cycles). + let current_generation = self.connection_generation.fetch_add(1, Ordering::SeqCst) + 1; + + info!( + "Successfully authenticated with WhatsApp servers! (gen={})", + current_generation + ); + self.auto_reconnect_errors.store(0, Ordering::Relaxed); + + self.update_server_time_offset(node); + + // Extract LID from the node before spawning (node isn't Send). + let lid_from_server = match node.get_attr("lid") { + Some(lid_value) => match lid_value.to_jid() { + Some(lid) => Some(lid), + None => { + warn!("Failed to parse LID from success stanza: {lid_value}"); + None + } + }, + None => { + warn!("LID not found in stanza. Group messaging may fail."); + None + } + }; + + let client_clone = self.clone(); + let task_generation = current_generation; + self.runtime.spawn(Box::pin(async move { + // Update LID if changed (moved here to avoid blocking the read loop + // on Device snapshot + write lock). + if let Some(lid) = lid_from_server { + let device_snapshot = + client_clone.persistence_manager.get_device_snapshot().await; + if device_snapshot.lid.as_ref() != Some(&lid) { + debug!("Updating LID from server to '{lid}'"); + client_clone + .persistence_manager + .process_command(DeviceCommand::SetLid(Some(lid))) + .await; + } + } + + // WA Web bumps `lc` after each successful auth (Start/Backend.js + // listener on `onOpenSocketStream`). The Comms `onConnect` handler + // gates the trigger on `isRegistered()`, so the bump only happens + // for already-paired logins — never during the pairing XX + // handshake. We mirror that by skipping when `device.pn` is None. + let already_paired = client_clone + .persistence_manager + .get_device_snapshot() + .await + .pn + .is_some(); + if already_paired { + client_clone + .persistence_manager + .process_command(DeviceCommand::IncrementLoginCounter) + .await; + } + + // Macro to check if this task is still valid (connection hasn't been replaced) + macro_rules! check_generation { + () => { + if client_clone.connection_generation.load(Ordering::SeqCst) != task_generation + { + debug!("Post-login task cancelled: connection generation changed"); + return; + } + }; + } + + debug!( + "Starting post-login initialization sequence (gen={})...", + task_generation + ); + + // Check if we need initial app state sync (empty pushname indicates fresh pairing + // where pushname will come from app state sync's setting_pushName mutation) + let device_snapshot = client_clone.persistence_manager.get_device_snapshot().await; + let needs_pushname_from_sync = device_snapshot.push_name.is_empty(); + if needs_pushname_from_sync { + debug!("Push name is empty - will be set from app state sync (setting_pushName)"); + } + + // Check connection before network operations. + // During pairing, a 515 disconnect happens quickly after success, + // so the socket may already be gone. + if !client_clone.is_connected() { + debug!( + "Skipping post-login init: connection closed (likely pairing phase reconnect)" + ); + return; + } + + check_generation!(); + client_clone.send_unified_session().await; + + // === Establish session with primary phone for PDO === + // This must happen BEFORE we exit passive mode (before offline messages arrive). + // PDO needs a session with device 0 to request decrypted content from our phone. + // Matches WhatsApp Web's bootstrapDeviceCapabilities() pattern. + check_generation!(); + if let Err(e) = client_clone + .establish_primary_phone_session_immediate() + .await + { + warn!(target: "Client/PDO", "Failed to establish session with primary phone on login: {:?}", e); + // Don't fail login - PDO will retry via ensure_e2e_sessions fallback + } + + // Sync own device list so DM fan-out includes all companions + check_generation!(); + if let Err(e) = client_clone.sync_own_device_list().await { + client_clone.log_sync_error("sync own device list", &e); + } + + check_generation!(); + if !client_clone.is_connected() { + debug!("Skipping passive tasks: connection closed"); + return; + } + if let Err(e) = client_clone.upload_pre_keys_at_login().await + && !client_clone.is_shutting_down() + { + warn!("Failed to upload pre-keys during startup: {e:?}"); + } + + // === Send active IQ === + // The server sends AFTER we exit passive mode. + // This matches WhatsApp Web's behavior: executePassiveTasks() -> sendPassiveModeProtocol("active") + check_generation!(); + if !client_clone.is_connected() { + debug!("Skipping active IQ: connection closed"); + return; + } + if let Err(e) = client_clone.set_passive(false).await + && !client_clone.is_shutting_down() + { + warn!("Failed to send post-connect active IQ: {e:?}"); + } + + // === Wait for offline sync to complete === + // The server sends after we exit passive mode. + client_clone.wait_for_offline_delivery_end().await; + + // Check if connection was replaced while waiting + check_generation!(); + + // Re-check connection and generation before sending presence + check_generation!(); + if !client_clone.is_connected() { + debug!("Skipping presence: connection closed"); + return; + } + + // Background initialization queries (can run in parallel, non-blocking) + let bg_client = client_clone.clone(); + let bg_generation = task_generation; + client_clone.runtime.spawn(Box::pin(async move { + // Check connection and generation before starting background queries + if bg_client.connection_generation.load(Ordering::SeqCst) != bg_generation { + debug!("Skipping background init queries: connection generation changed"); + return; + } + if !bg_client.is_connected() { + debug!("Skipping background init queries: connection closed"); + return; + } + + debug!( + "Sending background initialization queries (Props, Blocklist, Privacy, Digest)..." + ); + + let props_fut = bg_client.fetch_props(); + let binding = bg_client.blocking(); + let blocklist_fut = binding.get_blocklist(); + let privacy_fut = bg_client.fetch_privacy_settings(); + let digest_fut = bg_client.validate_digest_key(); + + let (r_props, r_block, r_priv, r_digest) = + futures::join!(props_fut, blocklist_fut, privacy_fut, digest_fut); + + // Suppress warnings if connection closed while queries were in-flight + if !bg_client.is_shutting_down() { + if let Err(e) = r_props { + warn!("Background init: Failed to fetch props: {e:?}"); + } + if let Err(e) = r_block { + warn!("Background init: Failed to fetch blocklist: {e:?}"); + } + if let Err(e) = r_priv { + warn!("Background init: Failed to fetch privacy settings: {e:?}"); + } + if let Err(e) = r_digest { + warn!("Background init: Failed to validate digest key: {e:?}"); + } + } + + // Prune expired tcTokens on connect (matches WhatsApp Web's PrivacyTokenJob) + if let Err(e) = bg_client.tc_token().prune_expired().await + && !bg_client.is_shutting_down() + { + warn!("Background init: Failed to prune expired tc_tokens: {e:?}"); + } + })).detach(); + + check_generation!(); + + let flag_set = client_clone.needs_initial_full_sync.load(Ordering::Relaxed); + let needs_initial_sync = flag_set || needs_pushname_from_sync; + + if needs_initial_sync { + // === Fresh pairing path === + // Like WhatsApp Web's syncCriticalData(): await critical collections before + // dispatching Connected, so blocklist/privacy settings are applied first. + debug!( + target: "Client/AppState", + "Starting Initial App State Sync (flag_set={flag_set}, needs_pushname={needs_pushname_from_sync})" + ); + + if !client_clone + .initial_app_state_keys_received + .load(Ordering::Relaxed) + { + debug!( + target: "Client/AppState", + "Waiting up to 5s for app state keys..." + ); + let _ = rt_timeout( + &*client_clone.runtime, + Duration::from_secs(5), + client_clone.initial_keys_synced_notifier.listen(), + ) + .await; + + // Check if connection was replaced while waiting + check_generation!(); + } + + // Start the critical sync timeout timer matching WhatsApp Web's + // WAWebSyncBootstrap.$15 (setSyncDCriticalDataSyncTimeout). + // WhatsApp Web uses 180s and calls socketLogout(SyncdTimeout) if + // the critical data hasn't synced by then. + const CRITICAL_SYNC_TIMEOUT_SECS: u64 = 180; + let timeout_client = client_clone.clone(); + let timeout_generation = task_generation; + let timeout_rt = client_clone.runtime.clone(); + let critical_sync_timeout_handle = timeout_rt.spawn(Box::pin(async move { + timeout_client.runtime.sleep(Duration::from_secs(CRITICAL_SYNC_TIMEOUT_SECS)).await; + // Check generation — if connection was replaced, this timeout is stale + if timeout_client.connection_generation.load(Ordering::SeqCst) + != timeout_generation + { + return; + } + // Matches WhatsApp Web's $16(): check if SettingPushName was synced. + // If push_name is still empty after 180s, critical sync failed. + let push_name = timeout_client.get_push_name().await; + if push_name.is_empty() { + warn!( + target: "Client/AppState", + "Critical app state sync timed out after {CRITICAL_SYNC_TIMEOUT_SECS}s \ + (push_name not synced). Reconnecting to retry." + ); + // WhatsApp Web does socketLogout here which clears device identity. + // We reconnect instead — preserving credentials and keeping the + // run loop active so auto-reconnect can retry the sync. + timeout_client.reconnect_immediately().await; + } else { + debug!( + target: "Client/AppState", + "Critical sync timeout fired but push_name was already synced" + ); + } + })); + + // Await critical collections via batched IQ before dispatching Connected. + check_generation!(); + match client_clone + .sync_collections_batched(vec![ + WAPatchName::CriticalBlock, + WAPatchName::CriticalUnblockLow, + ]) + .await + { + Ok(()) => { + // Critical sync completed — cancel the timeout timer + critical_sync_timeout_handle.abort(); + + check_generation!(); + + client_clone + .resubscribe_presence_subscriptions(task_generation) + .await; + + check_generation!(); + + // Dispatch Connected after critical sync completes. + // Presence is NOT sent here — WhatsApp Web sends presence from the + // setting_pushName mutation handler (WAWebPushNameSync), not from + // criticalSyncDone. Our setting_pushName handler already does this. + client_clone.dispatch_connected(); + } + Err(e) => { + client_clone.log_sync_error("critical app state sync", &e); + // Don't abort the timeout or dispatch Connected — the sync failed, + // so the timeout watchdog should remain active to force a reconnect + // if needed. Return early to avoid emitting a spurious Connected event. + return; + } + } + + // Spawn remaining non-critical collections in background + let sync_client = client_clone.clone(); + let sync_generation = task_generation; + client_clone.runtime.spawn(Box::pin(async move { + if sync_client.connection_generation.load(Ordering::SeqCst) != sync_generation { + debug!("App state sync cancelled: connection generation changed"); + return; + } + + if let Err(e) = sync_client + .sync_collections_batched(vec![ + WAPatchName::RegularLow, + WAPatchName::RegularHigh, + WAPatchName::Regular, + ]) + .await + { + sync_client.log_sync_error("non-critical app state sync", &e); + } + + sync_client + .needs_initial_full_sync + .store(false, Ordering::Relaxed); + debug!(target: "Client/AppState", "Initial App State Sync Completed."); + })).detach(); + } else { + // === Reconnection path === + // Pushname is already known, send presence and Connected immediately. + let device_snapshot = client_clone.persistence_manager.get_device_snapshot().await; + if !device_snapshot.push_name.is_empty() { + if let Err(e) = client_clone.presence().set_available().await { + warn!("Failed to send initial presence: {e:?}"); + } else { + debug!("Initial presence sent successfully."); + } + } + + client_clone + .resubscribe_presence_subscriptions(task_generation) + .await; + + // Re-check generation after awaits to avoid dispatching Connected + // for an outdated connection that was replaced mid-await. + check_generation!(); + + client_clone.dispatch_connected(); + } + })).detach(); + } + + /// Handles incoming `` stanzas by resolving pending response waiters. + /// + /// If an ack with an ID that matches a pending task in `response_waiters`, + /// the task is resolved and the function returns `true`. Otherwise, returns `false`. + pub(crate) async fn handle_ack_response(&self, node: &wacore_binary::NodeRef<'_>) -> bool { + // Surface server nack codes for diagnosability. A nacked send still + // resolves Ok to the caller, so without this the failure is invisible. + if let Some(error_code) = node.get_attr("error") { + let code = error_code.as_str(); + let id = node.get_attr("id").map(|v| v.as_str().into_owned()); + match code.as_ref() { + "463" => { + warn!( + target: "Client/Ack", + "Received 463 (MissingTcToken) nack for msg {:?}. \ + The recipient requires a valid tctoken or cstoken. \ + This may indicate a reachout timelock on the account.", + id + ); + } + "479" => { + warn!( + target: "Client/Ack", + "Received 479 (SmaxInvalid) nack for msg {:?}. \ + A stanza field has an incorrect format (e.g. wrong JID format or content type).", + id + ); + } + other => { + warn!( + target: "Client/Ack", + "Received {other} nack for msg {:?}; the message was likely \ + not delivered (e.g. 400 = malformed stanza, 404 = recipient \ + not found, 503 = service unavailable).", + id + ); + } + } + } + + let id_opt = node.get_attr("id").map(|v| v.as_str().into_owned()); + if let Some(id) = id_opt + && let Some(waiter) = self.response_waiters.lock().await.remove(&id) + { + // ACK responses are infrequent; re-encode into OwnedNodeRef for the channel. + // marshal_ref prepends a leading 0x00 format byte; OwnedNodeRef::new expects raw + // protocol bytes without it, matching what unpack() produces from the network. + match wacore_binary::marshal::marshal_ref(node) + .and_then(|bytes| wacore_binary::OwnedNodeRef::new(bytes[1..].to_vec())) + { + Ok(onr) => { + if waiter.send(Arc::new(onr)).is_err() { + warn!(target: "Client/Ack", "Failed to send ACK response to waiter for ID {id}. Receiver was likely dropped."); + } + } + Err(e) => { + warn!(target: "Client/Ack", "Failed to re-encode ACK node for waiter: {e}"); + } + } + return true; + } + false + } + + pub(crate) async fn handle_stream_error(&self, node: &wacore_binary::NodeRef<'_>) { + // is_logged_in handling: opt-in branches (515/516/401/409/conflict) clear it + // in the disconnect block below; 429/503 clear it inline because the server + // explicitly rejected the session and outgoing sends should bail fast; the + // unknown/code-less catch-all keeps it true so is_fully_ready()-gated work + // (notably prekey uploads) survives ack-shaped routing errors. + let mut attrs = node.attrs(); + let code_cow = attrs.optional_string("code"); + let code = code_cow.as_deref().unwrap_or(""); + let conflict_type = node + .get_optional_child("conflict") + .map(|n| { + n.attrs() + .optional_string("type") + .as_deref() + .unwrap_or("") + .to_string() + }) + .unwrap_or_default(); + + // Whether to proactively disconnect the transport after handling. + let mut should_disconnect = false; + + if !conflict_type.is_empty() { + info!( + "Got stream error indicating client was removed or replaced (conflict={}). Logging out.", + conflict_type + ); + self.expected_disconnect.store(true, Ordering::Relaxed); + self.enable_auto_reconnect.store(false, Ordering::Relaxed); + + let event = if conflict_type == "replaced" { + Event::StreamReplaced(crate::types::events::StreamReplaced) + } else { + Event::LoggedOut(crate::types::events::LoggedOut { + on_connect: false, + reason: ConnectFailureReason::LoggedOut, + }) + }; + self.core.event_bus.dispatch(event); + should_disconnect = true; + } else { + match code { + "515" => { + info!( + "Got 515 stream error, server is closing stream (expected after pairing). Will auto-reconnect." + ); + self.expected_disconnect.store(true, Ordering::Relaxed); + should_disconnect = true; + } + "516" => { + info!("Got 516 stream error (device removed). Logging out."); + self.expected_disconnect.store(true, Ordering::Relaxed); + self.enable_auto_reconnect.store(false, Ordering::Relaxed); + self.core.event_bus.dispatch(Event::LoggedOut( + crate::types::events::LoggedOut { + on_connect: false, + reason: ConnectFailureReason::LoggedOut, + }, + )); + should_disconnect = true; + } + "401" => { + info!("Got 401 stream error (unauthorized). Logging out."); + self.expected_disconnect.store(true, Ordering::Relaxed); + self.enable_auto_reconnect.store(false, Ordering::Relaxed); + self.core.event_bus.dispatch(Event::LoggedOut( + crate::types::events::LoggedOut { + on_connect: false, + reason: ConnectFailureReason::LoggedOut, + }, + )); + should_disconnect = true; + } + "409" => { + info!("Got 409 stream error (conflict). Another session replaced this one."); + self.expected_disconnect.store(true, Ordering::Relaxed); + self.enable_auto_reconnect.store(false, Ordering::Relaxed); + self.core + .event_bus + .dispatch(Event::StreamReplaced(crate::types::events::StreamReplaced)); + should_disconnect = true; + } + "429" => { + // Server signalled rate-limit on this session: mark logged-out so + // outgoing sends bail fast instead of being interpreted as abuse + // while we wait for the (likely-imminent) reconnect. + warn!( + "Got 429 stream error (rate limited). Will auto-reconnect with extended backoff." + ); + self.is_logged_in.store(false, Ordering::Relaxed); + self.auto_reconnect_errors.fetch_add(5, Ordering::Relaxed); + } + "503" => { + // Server is going down/restarting: mark logged-out so sends fail + // fast against the soon-to-die socket. Auto-reconnect handles recovery. + info!("Got 503 service unavailable, will auto-reconnect."); + self.is_logged_in.store(false, Ordering::Relaxed); + } + _ => { + // Server wraps per-stanza routing failures in without a + // code (e.g. ): treat as informational so we don't trigger reconnect + // storms. is_logged_in stays true on purpose — whatsmeow clears it eagerly, + // but here is_fully_ready() gates prekey uploads and we want them to keep + // working while the socket is still alive. Severity is warn!, not error!, + // because the connection is intentionally preserved. + // WA Web (StreamError.js) knows (type "ack"); + // name it instead of "Unknown". Root cause is usually an un-acked + // offline stanza; the server's drives the reconnect. + if let Some(ack) = node.get_optional_child("ack") { + let id = ack + .get_attr("id") + .map(|v| v.as_str().to_string()) + .unwrap_or_default(); + let class = ack + .get_attr("class") + .map(|v| v.as_str().to_string()) + .unwrap_or_default(); + warn!( + "Stream error carrying (class={class:?}, id={id}): server-driven stream rotation, not an ack rejection; reconnect follows on stream end" + ); + } else { + warn!("Unknown stream error: {}", DisplayableNodeRef(node)); + } + self.core.event_bus.dispatch(Event::StreamError( + crate::types::events::StreamError { + code: code.to_string(), + raw: Some(node.to_owned()), + }, + )); + } + } + } + + // Single is_logged_in clear + transport disconnect for every opt-in branch + // (515/516/401/409 and conflict). 429/503/unknown fall through so the + // socket layer notices a real teardown without us forcing one. + if should_disconnect { + self.is_logged_in.store(false, Ordering::Relaxed); + let transport_opt = self.transport.lock().await.clone(); + if let Some(transport) = transport_opt { + self.runtime + .spawn(Box::pin(async move { + transport.disconnect().await; + })) + .detach(); + } + info!("Notifying connection shutdown from stream error handler"); + self.notify_connection_shutdown(); + } + } + + pub(crate) async fn handle_connect_failure(&self, node: &wacore_binary::NodeRef<'_>) { + self.expected_disconnect.store(true, Ordering::Relaxed); + self.notify_connection_shutdown(); + + let mut attrs = node.attrs(); + let reason_code = attrs.optional_u64("reason").unwrap_or(0) as i32; + let reason = ConnectFailureReason::from(reason_code); + + if reason.should_reconnect() { + self.expected_disconnect.store(false, Ordering::Relaxed); + } else { + self.enable_auto_reconnect.store(false, Ordering::Relaxed); + } + + if reason.is_logged_out() { + // Log the full so a server-side lock/ban is diagnosable; + // `location` (e.g. "rva") is a routing token, not the cause. + warn!( + "Got {reason:?} connect failure, logging out: {}", + DisplayableNodeRef(node) + ); + self.core + .event_bus + .dispatch(wacore::types::events::Event::LoggedOut( + crate::types::events::LoggedOut { + on_connect: true, + reason, + }, + )); + } else if let ConnectFailureReason::TempBanned = reason { + let ban_code = attrs.optional_u64("code").unwrap_or(0) as i32; + let expire_secs = attrs.optional_u64("expire").unwrap_or(0); + let expire_duration = + chrono::Duration::try_seconds(expire_secs as i64).unwrap_or_default(); + warn!( + "Temporary ban connect failure: {}", + DisplayableNodeRef(node) + ); + self.core + .event_bus + .dispatch(Event::TemporaryBan(crate::types::events::TemporaryBan { + code: crate::types::events::TempBanReason::from(ban_code), + expire: expire_duration, + })); + } else if let ConnectFailureReason::ClientOutdated = reason { + error!("Client is outdated and was rejected by server."); + self.core + .event_bus + .dispatch(Event::ClientOutdated(crate::types::events::ClientOutdated)); + } else { + warn!("Unknown connect failure: {}", DisplayableNodeRef(node)); + self.core.event_bus.dispatch(Event::ConnectFailure( + crate::types::events::ConnectFailure { + reason, + message: attrs + .optional_string("message") + .as_deref() + .unwrap_or("") + .to_string(), + raw: Some(node.to_owned()), + }, + )); + } + } + + pub(crate) async fn handle_iq(self: &Arc, node: &wacore_binary::NodeRef<'_>) -> bool { + if node.get_attr("type").is_some_and(|s| s.as_str() == "get") + && (node.get_optional_child("ping").is_some() + || node + .get_attr("xmlns") + .is_some_and(|s| s.as_str() == "urn:xmpp:ping")) + { + debug!("Received ping, sending pong."); + let mut parser = node.attrs(); + let from_jid = parser.jid("from"); + let id = parser.optional_string("id").map(|s| s.to_string()); + let pong = build_pong(from_jid.to_string(), id.as_deref()); + if let Err(e) = self.send_node(pong).await { + warn!("Failed to send pong: {e:?}"); + } + return true; + } + + if pair::handle_iq(self, node).await { + return true; + } + + false + } + + pub(crate) fn update_server_time_offset(&self, node: &wacore_binary::NodeRef<'_>) { + self.unified_session.update_server_time_offset(node); + } +} diff --git a/src/client/tests.rs b/src/client/tests.rs new file mode 100644 index 000000000..7a5a46486 --- /dev/null +++ b/src/client/tests.rs @@ -0,0 +1,2795 @@ +//! Client integration and unit tests. + +use super::*; +use crate::lid_pn_cache::LearningSource; +use crate::test_utils::MockHttpClient; +use futures::channel::oneshot; +use wacore_binary::SERVER_JID; + +#[tokio::test] +async fn test_ack_behavior_for_incoming_stanzas() { + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // --- Assertions --- + + // Verify that we still ack other critical stanzas (regression check). + use wacore_binary::{Attrs, Node, NodeContent}; + + let mut receipt_attrs = Attrs::new(); + receipt_attrs.insert("from".to_string(), "@s.whatsapp.net".to_string()); + receipt_attrs.insert("id".to_string(), "RCPT-1".to_string()); + let receipt_node = Node::new( + "receipt", + receipt_attrs, + Some(NodeContent::String("test".into())), + ); + + let mut notification_attrs = Attrs::new(); + notification_attrs.insert("from".to_string(), "@s.whatsapp.net".to_string()); + notification_attrs.insert("id".to_string(), "NOTIF-1".to_string()); + let notification_node = Node::new( + "notification", + notification_attrs, + Some(NodeContent::String("test".into())), + ); + + assert!( + client.should_ack(&receipt_node.as_node_ref()), + "should_ack must still return TRUE for stanzas." + ); + assert!( + client.should_ack(¬ification_node.as_node_ref()), + "should_ack must still return TRUE for stanzas." + ); + + // Regular stanzas (DM / group) are acked via the delivery + // , not a bare . WA Web only emits + // for newsletter deliveries. + let mut dm_attrs = Attrs::new(); + dm_attrs.insert( + "from".to_string(), + "5511999999999@s.whatsapp.net".to_string(), + ); + dm_attrs.insert("id".to_string(), "MSG-DM-1".to_string()); + let dm_message = Node::new("message", dm_attrs, None); + assert!( + !client.should_ack(&dm_message.as_node_ref()), + "should_ack must return FALSE for regular DM (delivery receipt covers it)." + ); + + let mut group_attrs = Attrs::new(); + group_attrs.insert("from".to_string(), "120363098765432100@g.us".to_string()); + group_attrs.insert("id".to_string(), "MSG-GROUP-1".to_string()); + let group_message = Node::new("message", group_attrs, None); + assert!( + !client.should_ack(&group_message.as_node_ref()), + "should_ack must return FALSE for group ." + ); + + let mut newsletter_attrs = Attrs::new(); + newsletter_attrs.insert( + "from".to_string(), + "120363298765432100@newsletter".to_string(), + ); + newsletter_attrs.insert("id".to_string(), "MSG-NL-1".to_string()); + let newsletter_message = Node::new("message", newsletter_attrs, None); + assert!( + client.should_ack(&newsletter_message.as_node_ref()), + "should_ack must return TRUE for newsletter ." + ); + + // status@broadcast gets the transport as a fallback so that + // drop paths in process_group_enc_batch (expired status, missing + // sender key, decrypt error) don't leave the server retransmitting. + // The success path also emits ; the + // duplicate is tolerated. + let mut status_attrs = Attrs::new(); + status_attrs.insert("from".to_string(), "status@broadcast".to_string()); + status_attrs.insert("id".to_string(), "MSG-STATUS-1".to_string()); + let status_message = Node::new("message", status_attrs, None); + assert!( + client.should_ack(&status_message.as_node_ref()), + "should_ack must return TRUE for status@broadcast (fallback for drop paths)." + ); + + info!( + "✅ test_ack_behavior_for_incoming_stanzas passed: Client correctly differentiates which stanzas to acknowledge." + ); +} + +#[tokio::test] +async fn test_ack_waiter_resolves() { + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // 1. Insert a waiter for a specific ID + let test_id = "ack-test-123".to_string(); + let (tx, rx) = oneshot::channel(); + client + .response_waiters + .lock() + .await + .insert(test_id.clone(), tx); + assert!( + client.response_waiters.lock().await.contains_key(&test_id), + "Waiter should be inserted before handling ack" + ); + + // 2. Create a mock node with the test ID + let ack_node = NodeBuilder::new("ack") + .attr("id", test_id.clone()) + .attr("from", SERVER_JID) + .build(); + + // 3. Handle the ack + let handled = client.handle_ack_response(&ack_node.as_node_ref()).await; + assert!( + handled, + "handle_ack_response should return true when waiter exists" + ); + + // 4. Await the receiver with a timeout + match tokio::time::timeout(Duration::from_secs(1), rx).await { + Ok(Ok(response_node)) => { + assert!( + response_node + .get() + .get_attr("id") + .is_some_and(|v| v.as_str() == test_id.as_str()), + "Response node should have correct ID" + ); + } + Ok(Err(_)) => panic!("Receiver was dropped without being sent a value"), + Err(_) => panic!("Test timed out waiting for ack response"), + } + + // 5. Verify the waiter was removed + assert!( + !client.response_waiters.lock().await.contains_key(&test_id), + "Waiter should be removed after handling" + ); + + info!("✅ test_ack_waiter_resolves passed: ACK response correctly resolves pending waiters"); +} + +#[tokio::test] +async fn test_ack_without_matching_waiter() { + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Create an ack without any matching waiter + let ack_node = NodeBuilder::new("ack") + .attr("id", "non-existent-id") + .attr("from", SERVER_JID) + .build(); + + // Should return false since there's no waiter + let handled = client.handle_ack_response(&ack_node.as_node_ref()).await; + assert!( + !handled, + "handle_ack_response should return false when no waiter exists" + ); + + info!( + "✅ test_ack_without_matching_waiter passed: ACK without matching waiter handled gracefully" + ); +} + +/// Test that the lid_pn_cache correctly stores and retrieves LID mappings. +/// +/// This is critical for the LID-PN session mismatch fix. When we receive a message +/// with sender_lid, we cache the phone->LID mapping so that when sending replies, +/// we can reuse the existing LID session instead of creating a new PN session. +#[tokio::test] +async fn test_lid_pn_cache_basic_operations() { + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_lid_cache_basic?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Initially, the cache should be empty for a phone number + let phone = "559980000001"; + let lid = "100000012345678"; + + assert!( + client.lid_pn_cache.get_current_lid(phone).await.is_none(), + "Cache should be empty initially" + ); + + // Insert a phone->LID mapping using add_lid_pn_mapping + client + .add_lid_pn_mapping(lid, phone, LearningSource::Usync) + .await + .expect("Failed to persist LID-PN mapping in tests"); + + // Verify we can retrieve it (phone -> LID lookup) + let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; + assert!(cached_lid.is_some(), "Cache should contain the mapping"); + assert_eq!( + cached_lid.expect("cache should have LID"), + lid, + "Cached LID should match what we inserted" + ); + + // Verify reverse lookup works (LID -> phone) + let cached_phone = client.lid_pn_cache.get_phone_number(lid).await; + assert!(cached_phone.is_some(), "Reverse lookup should work"); + assert_eq!( + cached_phone.expect("reverse lookup should return phone"), + phone, + "Cached phone should match what we inserted" + ); + + // Verify a different phone number returns None + assert!( + client + .lid_pn_cache + .get_current_lid("559980000002") + .await + .is_none(), + "Different phone number should not have a mapping" + ); + + info!("✅ test_lid_pn_cache_basic_operations passed: LID-PN cache works correctly"); +} + +/// Test that the lid_pn_cache respects timestamp-based conflict resolution. +/// +/// When a phone number has multiple LIDs, the most recent one should be returned. +#[tokio::test] +async fn test_lid_pn_cache_timestamp_resolution() { + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_lid_cache_timestamp?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + let phone = "559980000001"; + let lid_old = "100000012345678"; + let lid_new = "100000087654321"; + + // Insert initial mapping + client + .add_lid_pn_mapping(lid_old, phone, LearningSource::Usync) + .await + .expect("Failed to persist LID-PN mapping in tests"); + + assert_eq!( + client + .lid_pn_cache + .get_current_lid(phone) + .await + .expect("cache should have LID"), + lid_old, + "Initial LID should be stored" + ); + + // Small delay to ensure different timestamp + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + + // Add new mapping with newer timestamp + client + .add_lid_pn_mapping(lid_new, phone, LearningSource::PeerPnMessage) + .await + .expect("Failed to persist LID-PN mapping in tests"); + + assert_eq!( + client + .lid_pn_cache + .get_current_lid(phone) + .await + .expect("cache should have newer LID"), + lid_new, + "Newer LID should be returned for phone lookup" + ); + + // Both LIDs should still resolve to the same phone + assert_eq!( + client + .lid_pn_cache + .get_phone_number(lid_old) + .await + .expect("reverse lookup should return phone"), + phone, + "Old LID should still map to phone" + ); + assert_eq!( + client + .lid_pn_cache + .get_phone_number(lid_new) + .await + .expect("reverse lookup should return phone"), + phone, + "New LID should also map to phone" + ); + + info!( + "✅ test_lid_pn_cache_timestamp_resolution passed: Timestamp-based resolution works correctly" + ); +} + +/// Test that get_lid_for_phone (from SendContextResolver) returns the cached value. +/// +/// This is the method used by wacore::send to look up LID mappings when encrypting. +#[tokio::test] +async fn test_get_lid_for_phone_via_send_context_resolver() { + use wacore::client::context::SendContextResolver; + + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_get_lid_for_phone?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + let phone = "559980000001"; + let lid = "100000012345678"; + + // Before caching, should return None + assert!( + client.get_lid_for_phone(phone).await.is_none(), + "get_lid_for_phone should return None before caching" + ); + + // Cache the mapping using add_lid_pn_mapping + client + .add_lid_pn_mapping(lid, phone, LearningSource::Usync) + .await + .expect("Failed to persist LID-PN mapping in tests"); + + // Now it should return the LID + let result = client.get_lid_for_phone(phone).await; + assert!( + result.is_some(), + "get_lid_for_phone should return Some after caching" + ); + assert_eq!( + result.expect("get_lid_for_phone should return Some"), + lid, + "get_lid_for_phone should return the cached LID" + ); + + info!( + "✅ test_get_lid_for_phone_via_send_context_resolver passed: SendContextResolver correctly returns cached LID" + ); +} + +/// Test that wait_for_offline_delivery_end returns immediately when the flag is already set. +#[tokio::test] +async fn test_wait_for_offline_delivery_end_returns_immediately_when_flag_set() { + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_offline_sync_flag_set?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Set the flag to true (simulating offline sync completed) + client + .offline_sync_completed + .store(true, std::sync::atomic::Ordering::Relaxed); + + // This should return immediately (not wait 10 seconds) + let start = wacore::time::Instant::now(); + client.wait_for_offline_delivery_end().await; + let elapsed = start.elapsed(); + + // Should complete in < 100ms (not 10 second timeout) + assert!( + elapsed.as_millis() < 100, + "wait_for_offline_delivery_end should return immediately when flag is set, took {:?}", + elapsed + ); + + info!("✅ test_wait_for_offline_delivery_end_returns_immediately_when_flag_set passed"); +} + +/// Test that wait_for_offline_delivery_end times out when the flag is NOT set. +/// This verifies the 10-second timeout is working. +#[tokio::test] +async fn test_wait_for_offline_delivery_end_times_out_when_flag_not_set() { + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_offline_sync_timeout?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Flag is false by default, so use a short timeout and verify the helper + // marks the sync complete on timeout. + let start = wacore::time::Instant::now(); + client + .wait_for_offline_delivery_end_with_timeout(std::time::Duration::from_millis(50)) + .await; + + let elapsed = start.elapsed(); + // Count available permits by trying to acquire non-blockingly + let semaphore = match client.message_processing_semaphore.lock() { + Ok(guard) => guard.clone(), + Err(poisoned) => poisoned.into_inner().clone(), + }; + let mut guards = Vec::new(); + while let Some(guard) = semaphore.try_acquire() { + guards.push(guard); + } + let permits = guards.len(); + drop(guards); + + assert!( + elapsed.as_millis() >= 45, // Allow small timing variance + "Should have waited for the configured timeout duration, took {:?}", + elapsed + ); + assert!( + client + .offline_sync_completed + .load(std::sync::atomic::Ordering::Relaxed), + "wait_for_offline_delivery_end should mark offline sync complete on timeout" + ); + assert_eq!( + permits, 64, + "timeout completion should restore parallel permits" + ); + + info!("✅ test_wait_for_offline_delivery_end_times_out_when_flag_not_set passed"); +} + +/// Test that wait_for_offline_delivery_end returns when notified. +#[tokio::test] +async fn test_wait_for_offline_delivery_end_returns_on_notify() { + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_offline_notify?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + let client_clone = client.clone(); + + // Spawn a task that will notify after 50ms + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + client_clone.offline_sync_notifier.notify(usize::MAX); + }); + + let start = wacore::time::Instant::now(); + client.wait_for_offline_delivery_end().await; + let elapsed = start.elapsed(); + + // Should complete around 50ms (when notified), not 10 seconds + assert!( + elapsed.as_millis() < 200, + "wait_for_offline_delivery_end should return when notified, took {:?}", + elapsed + ); + assert!( + elapsed.as_millis() >= 45, // Should have waited for the notify + "Should have waited for the notify, only took {:?}", + elapsed + ); + + info!("✅ test_wait_for_offline_delivery_end_returns_on_notify passed"); +} + +/// Test that the offline_sync_completed flag starts as false. +#[tokio::test] +async fn test_offline_sync_flag_initially_false() { + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_offline_flag_initial?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // The flag should be false initially + assert!( + !client + .offline_sync_completed + .load(std::sync::atomic::Ordering::Relaxed), + "offline_sync_completed should be false when Client is first created" + ); + + info!("✅ test_offline_sync_flag_initially_false passed"); +} + +/// Test the complete offline sync lifecycle: +/// 1. Flag starts false +/// 2. Flag is set true after IB offline stanza +/// 3. Notify is called +#[tokio::test] +async fn test_offline_sync_lifecycle() { + use std::sync::atomic::Ordering; + + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_offline_lifecycle?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // 1. Initially false + assert!(!client.offline_sync_completed.load(Ordering::Relaxed)); + + // 2. Spawn a waiter + let client_waiter = client.clone(); + let waiter_handle = tokio::spawn(async move { + client_waiter.wait_for_offline_delivery_end().await; + true // Return that we completed + }); + + // Give the waiter time to start waiting + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + + // Verify waiter hasn't completed yet + assert!( + !waiter_handle.is_finished(), + "Waiter should still be waiting" + ); + + // 3. Simulate IB handler behavior (set flag and notify) + client.offline_sync_completed.store(true, Ordering::Relaxed); + client.offline_sync_notifier.notify(usize::MAX); + + // 4. Waiter should complete + let result = tokio::time::timeout(std::time::Duration::from_millis(100), waiter_handle) + .await + .expect("Waiter should complete after notify") + .expect("Waiter task should not panic"); + + assert!(result, "Waiter should have completed successfully"); + assert!(client.offline_sync_completed.load(Ordering::Relaxed)); + + info!("✅ test_offline_sync_lifecycle passed"); +} + +/// Test that establish_primary_phone_session_immediate returns error when no PN is set. +/// This verifies the "not logged in" guard works. +#[tokio::test] +async fn test_establish_primary_phone_session_fails_without_pn() { + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_no_pn?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // No PN set, so this should fail + let result = client.establish_primary_phone_session_immediate().await; + + assert!( + result.is_err(), + "establish_primary_phone_session_immediate should fail when no PN is set" + ); + + let err = result.unwrap_err(); + assert!( + err.downcast_ref::() + .is_some_and(|e| matches!(e, ClientError::NotLoggedIn)), + "Error should be ClientError::NotLoggedIn, got: {}", + err + ); + + info!("✅ test_establish_primary_phone_session_fails_without_pn passed"); +} + +/// Test that ensure_e2e_sessions waits for offline sync to complete. +/// This is the CRITICAL difference between ensure_e2e_sessions and +/// establish_primary_phone_session_immediate. +#[tokio::test] +async fn test_ensure_e2e_sessions_waits_for_offline_sync() { + use std::sync::atomic::Ordering; + use wacore_binary::Jid; + + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_ensure_e2e_waits?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Flag is false (offline sync not complete) + assert!(!client.offline_sync_completed.load(Ordering::Relaxed)); + + // Call ensure_e2e_sessions with an empty list (so it returns early after the wait) + // This lets us test the waiting behavior without needing network + let client_clone = client.clone(); + let ensure_handle = tokio::spawn(async move { + // Start with some JIDs - but since we're testing the wait, we use empty + // to avoid needing actual session establishment + client_clone.ensure_e2e_sessions(&[]).await + }); + + // Wait a bit - ensure_e2e_sessions should return immediately for empty list + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + assert!( + ensure_handle.is_finished(), + "ensure_e2e_sessions should return immediately for empty JID list" + ); + + // Now test with actual JIDs - it should wait for offline sync + let client_clone = client.clone(); + let test_jid = Jid::pn("559999999999"); + let ensure_handle = tokio::spawn(async move { + // This will wait for offline sync before proceeding + let start = wacore::time::Instant::now(); + let _ = client_clone.ensure_e2e_sessions(&[test_jid]).await; + start.elapsed() + }); + + // Give it a moment to start + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + + // It should still be waiting (offline sync not complete) + assert!( + !ensure_handle.is_finished(), + "ensure_e2e_sessions should be waiting for offline sync" + ); + + // Now complete offline sync + client.offline_sync_completed.store(true, Ordering::Relaxed); + client.offline_sync_notifier.notify(usize::MAX); + + // Now it should complete (might fail on session establishment, but that's ok) + let result = tokio::time::timeout(std::time::Duration::from_secs(2), ensure_handle).await; + + assert!( + result.is_ok(), + "ensure_e2e_sessions should complete after offline sync" + ); + + info!("✅ test_ensure_e2e_sessions_waits_for_offline_sync passed"); +} + +/// Integration test: Verify that the immediate session establishment does NOT +/// wait for offline sync. This is critical for PDO to work during offline sync. +/// +/// The flow is: +/// 1. Login -> establish_primary_phone_session_immediate() is called +/// 2. This should NOT wait for offline sync (flag is false at this point) +/// 3. After session is established, offline messages arrive +/// 4. When decryption fails, PDO can immediately send to device 0 +#[tokio::test] +async fn test_immediate_session_does_not_wait_for_offline_sync() { + use std::sync::atomic::Ordering; + use wacore_binary::Jid; + + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_immediate_no_wait?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend.clone()) + .await + .expect("persistence manager should initialize"), + ); + + // Set a PN so establish_primary_phone_session_immediate doesn't fail early + pm.modify_device(|device| { + device.pn = Some(Jid::pn("559999999999")); + }) + .await; + + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Flag is false (offline sync not complete - simulating login state) + assert!(!client.offline_sync_completed.load(Ordering::Relaxed)); + + // Call establish_primary_phone_session_immediate + // It should NOT wait for offline sync - it should proceed immediately + let start = wacore::time::Instant::now(); + + // Note: This will fail because we can't actually fetch prekeys in tests, + // but the important thing is that it doesn't WAIT for offline sync + let result = tokio::time::timeout( + std::time::Duration::from_millis(500), + client.establish_primary_phone_session_immediate(), + ) + .await; + + let elapsed = start.elapsed(); + + // The call should complete (or fail) quickly, NOT wait for 10 second timeout + assert!( + result.is_ok(), + "establish_primary_phone_session_immediate should not wait for offline sync, timed out" + ); + + // It should complete in < 500ms (not 10 second wait) + assert!( + elapsed.as_millis() < 500, + "establish_primary_phone_session_immediate should not wait, took {:?}", + elapsed + ); + + // The actual result might be an error (no network), but that's fine + // The important thing is it didn't wait for offline sync + info!( + "establish_primary_phone_session_immediate completed in {:?} (result: {:?})", + elapsed, + result.unwrap().is_ok() + ); + + info!("✅ test_immediate_session_does_not_wait_for_offline_sync passed"); +} + +/// Integration test: Verify that establish_primary_phone_session_immediate +/// skips establishment when a session already exists. +/// +/// This is the CRITICAL fix for MAC verification failures: +/// - BUG (before fix): Called process_prekey_bundle() unconditionally, +/// replacing the existing session with a new one +/// - RESULT: Remote device still uses old session state, causing MAC failures +#[tokio::test] +async fn test_establish_session_skips_when_exists() { + use wacore::libsignal::protocol::SessionRecord; + use wacore::libsignal::store::SessionStore; + use wacore::types::jid::JidExt; + use wacore_binary::Jid; + + let backend = Arc::new( + crate::store::SqliteStore::new("file:memdb_skip_existing?mode=memory&cache=shared") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend.clone()) + .await + .expect("persistence manager should initialize"), + ); + + // Set a PN so the function doesn't fail early + let own_pn = Jid::pn("559999999999"); + pm.modify_device(|device| { + device.pn = Some(own_pn.clone()); + }) + .await; + + // Pre-populate a session for the primary phone JID (device 0) + let primary_phone_jid = own_pn.with_device(0); + let signal_addr = primary_phone_jid.to_protocol_address(); + + // Create a dummy session record + let dummy_session = SessionRecord::new_fresh(); + { + let device_arc = pm.get_device_arc().await; + let device = device_arc.read().await; + device + .store_session(&signal_addr, &dummy_session) + .await + .expect("Failed to store test session"); + + // Verify session exists + let exists = device + .contains_session(&signal_addr) + .await + .expect("Failed to check session"); + assert!(exists, "Session should exist after store"); + } + + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm.clone(), + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Call establish_primary_phone_session_immediate + // It should return Ok(()) immediately without fetching prekeys + let result = client.establish_primary_phone_session_immediate().await; + + assert!( + result.is_ok(), + "establish_primary_phone_session_immediate should succeed when session exists" + ); + + // Verify the session was NOT replaced (still has the same record) + // This is the critical assertion - if session was replaced, it would cause MAC failures + { + let device_arc = pm.get_device_arc().await; + let device = device_arc.read().await; + let exists = device + .contains_session(&signal_addr) + .await + .expect("Failed to check session"); + assert!(exists, "Session should still exist after the call"); + } + + info!("✅ test_establish_session_skips_when_exists passed"); +} + +/// Integration test: Verify that the session check prevents MAC failures +/// by documenting the exact control flow that caused the bug. +#[test] +fn test_mac_failure_prevention_flow_documentation() { + // Simulate the decision logic + fn should_establish_session(check_result: Result) -> Result { + match check_result { + Ok(true) => Ok(false), // Session exists → DON'T establish + Ok(false) => Ok(true), // No session → establish + Err(e) => Err(format!("Cannot verify session: {}", e)), // Fail-safe + } + } + + // Test Case 1: Session exists → skip (prevents MAC failure) + let result = should_establish_session(Ok(true)); + assert_eq!(result, Ok(false), "Should skip when session exists"); + + // Test Case 2: No session → establish + let result = should_establish_session(Ok(false)); + assert_eq!(result, Ok(true), "Should establish when no session"); + + // Test Case 3: Check fails → error (fail-safe) + let result = should_establish_session(Err("database error")); + assert!(result.is_err(), "Should fail when check fails"); + + info!("✅ test_mac_failure_prevention_flow_documentation passed"); +} + +#[test] +fn test_unified_session_id_calculation() { + // Test the mathematical calculation of the unified session ID. + // Formula: (now_ms + server_offset_ms + 3_days_ms) % 7_days_ms + + const DAY_MS: i64 = 24 * 60 * 60 * 1000; + const WEEK_MS: i64 = 7 * DAY_MS; + const OFFSET_MS: i64 = 3 * DAY_MS; + + // Helper function matching the implementation + fn calculate_session_id(now_ms: i64, server_offset_ms: i64) -> i64 { + let adjusted_now = now_ms + server_offset_ms; + (adjusted_now + OFFSET_MS) % WEEK_MS + } + + // Test 1: Zero offset + let now_ms = 1706000000000_i64; // Some arbitrary timestamp + let id = calculate_session_id(now_ms, 0); + assert!( + (0..WEEK_MS).contains(&id), + "Session ID should be in [0, WEEK_MS)" + ); + + // Test 2: Positive server offset (server is ahead) + let id_with_positive_offset = calculate_session_id(now_ms, 5000); + assert!( + (0..WEEK_MS).contains(&id_with_positive_offset), + "Session ID should be in [0, WEEK_MS)" + ); + // The ID should be different from zero offset (unless wrap-around) + // Not testing exact value as it depends on the offset + + // Test 3: Negative server offset (server is behind) + let id_with_negative_offset = calculate_session_id(now_ms, -5000); + assert!( + (0..WEEK_MS).contains(&id_with_negative_offset), + "Session ID should be in [0, WEEK_MS)" + ); + + // Test 4: Verify modulo wrap-around + // If adjusted_now + OFFSET_MS >= WEEK_MS, it should wrap + let wrap_test_now = WEEK_MS - OFFSET_MS + 1000; // Should produce small result + let wrapped_id = calculate_session_id(wrap_test_now, 0); + assert_eq!(wrapped_id, 1000, "Should wrap around correctly"); + + // Test 5: Edge case - at exact boundary + let boundary_now = WEEK_MS - OFFSET_MS; + let boundary_id = calculate_session_id(boundary_now, 0); + assert_eq!(boundary_id, 0, "At exact boundary should be 0"); +} + +#[tokio::test] +async fn test_server_time_offset_extraction() { + use wacore_binary::builder::NodeBuilder; + + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Initially, offset should be 0 + assert_eq!( + client.unified_session.server_time_offset_ms(), + 0, + "Initial offset should be 0" + ); + + // Create a node with a 't' attribute + let server_time = wacore::time::now_secs() + 10; // Server is 10 seconds ahead + let node = NodeBuilder::new("success").attr("t", server_time).build(); + + // Update the offset + client.update_server_time_offset(&node.as_node_ref()); + + // The offset should be approximately 10 * 1000 = 10000 ms + // Allow some tolerance for timing differences during the test + let offset = client.unified_session.server_time_offset_ms(); + assert!( + (offset - 10000).abs() < 1000, // Allow 1 second tolerance + "Offset should be approximately 10000ms, got {}", + offset + ); + + // Test with no 't' attribute - should not change offset + let node_no_t = NodeBuilder::new("success").build(); + client.update_server_time_offset(&node_no_t.as_node_ref()); + let offset_after = client.unified_session.server_time_offset_ms(); + assert!( + (offset_after - offset).abs() < 100, // Should be same (or very close) + "Offset should not change when 't' is missing" + ); + + // Test with invalid 't' attribute - should not change offset + let node_invalid = NodeBuilder::new("success") + .attr("t", "not_a_number") + .build(); + client.update_server_time_offset(&node_invalid.as_node_ref()); + let offset_after_invalid = client.unified_session.server_time_offset_ms(); + assert!( + (offset_after_invalid - offset).abs() < 100, + "Offset should not change when 't' is invalid" + ); + + // Test with negative/zero 't' - should not change offset + let node_zero = NodeBuilder::new("success").attr("t", "0").build(); + client.update_server_time_offset(&node_zero.as_node_ref()); + let offset_after_zero = client.unified_session.server_time_offset_ms(); + assert!( + (offset_after_zero - offset).abs() < 100, + "Offset should not change when 't' is 0" + ); + + info!("✅ test_server_time_offset_extraction passed"); +} + +#[tokio::test] +async fn test_unified_session_manager_integration() { + // Test the unified session manager through the client + + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Initially, sequence should be 0 + assert_eq!( + client.unified_session.sequence(), + 0, + "Initial sequence should be 0" + ); + + // Duplicate prevention depends on the session ID staying the same between calls. + // Since the session ID is millisecond-based, use a retry loop to handle + // the rare case where we cross a millisecond boundary between calls. + loop { + client.unified_session.reset().await; + + let result = client.unified_session.prepare_send().await; + assert!(result.is_some(), "First send should succeed"); + let (node, seq) = result.unwrap(); + assert_eq!(node.tag, "ib", "Should be an IB stanza"); + assert_eq!(seq, 1, "First sequence should be 1 (pre-increment)"); + assert_eq!(client.unified_session.sequence(), 1); + + let result2 = client.unified_session.prepare_send().await; + if result2.is_none() { + // Duplicate was prevented within the same millisecond + assert_eq!(client.unified_session.sequence(), 1); + break; + } + // Millisecond boundary crossed, retry + tokio::task::yield_now().await; + } + + // Clear last sent and try again - sequence resets on "new" session ID + client.unified_session.clear_last_sent().await; + let result3 = client.unified_session.prepare_send().await; + assert!(result3.is_some(), "Should succeed after clearing"); + let (_, seq3) = result3.unwrap(); + assert_eq!(seq3, 1, "Sequence resets when session ID changes"); + assert_eq!(client.unified_session.sequence(), 1); + + info!("✅ test_unified_session_manager_integration passed"); +} + +#[test] +fn test_unified_session_protocol_node() { + // Test the type-safe protocol node implementation + use wacore::ib::{IbStanza, UnifiedSession}; + use wacore::protocol::ProtocolNode; + + // Create a unified session + let session = UnifiedSession::new("123456789"); + assert_eq!(session.id, "123456789"); + assert_eq!(session.tag(), "unified_session"); + + // Convert to node + let node = session.into_node(); + assert_eq!(node.tag, "unified_session"); + assert!(node.attrs.get("id").is_some_and(|v| v == "123456789")); + + // Create an IB stanza + let stanza = IbStanza::unified_session(UnifiedSession::new("987654321")); + assert_eq!(stanza.tag(), "ib"); + + // Convert to node and verify structure + let ib_node = stanza.into_node(); + assert_eq!(ib_node.tag, "ib"); + let children = ib_node.children().expect("IB stanza should have children"); + assert_eq!(children.len(), 1); + assert_eq!(children[0].tag, "unified_session"); + assert!( + children[0] + .attrs + .get("id") + .is_some_and(|v| v == "987654321") + ); + + info!("✅ test_unified_session_protocol_node passed"); +} + +fn node_to_owned_ref(node: Node) -> Arc { + crate::test_utils::node_to_owned_ref(&node) +} + +/// Helper to create a test client for offline sync tests +async fn create_offline_sync_test_client() -> Arc { + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + client +} + +/// Regression: a transport disconnect must flush dirty Signal state before +/// clearing the cache, or a just-advanced sender-key chain is lost (forcing +/// a full SKDM re-fanout on the next send). +#[tokio::test] +async fn cleanup_connection_state_flushes_dirty_signal_state() { + use wacore::libsignal::protocol::ProtocolAddress; + let client = create_offline_sync_test_client().await; + + // A dirty identity lives only in the write-back cache until flushed. + let addr = ProtocolAddress::new("5550001000@s.whatsapp.net".to_string(), 1u32.into()); + client.signal_cache.put_identity(&addr, &[7u8; 32]).await; + + client.cleanup_connection_state().await; + + // cleanup cleared the cache, so a hit now can only come from the DB, + // proving the flush ran before the clear. + let device = client.persistence_manager.get_device_arc().await; + let guard = device.read().await; + let persisted = client + .signal_cache + .get_identity(&addr, &*guard.backend) + .await + .expect("get_identity must not error"); + assert!( + persisted.is_some(), + "dirty Signal state must survive a transport disconnect (flush-before-clear)" + ); +} + +/// Same guarantee on the sender-key store, which drives SKDM fanout. +#[tokio::test] +async fn cleanup_connection_state_flushes_dirty_sender_key() { + use wacore::libsignal::protocol::SenderKeyRecord; + use wacore::libsignal::store::sender_key_name::SenderKeyName; + let client = create_offline_sync_test_client().await; + + let name = SenderKeyName::from_parts("group@g.us", "5550001000@s.whatsapp.net:1"); + client + .signal_cache + .put_sender_key(&name, SenderKeyRecord::new_empty()) + .await; + + client.cleanup_connection_state().await; + + let device = client.persistence_manager.get_device_arc().await; + let guard = device.read().await; + let persisted = client + .signal_cache + .get_sender_key(&name, &*guard.backend) + .await + .expect("get_sender_key must not error"); + assert!( + persisted.is_some(), + "dirty sender key must survive a transport disconnect (flush-before-clear)" + ); +} + +/// When the flush itself fails, cleanup must NOT clear the cache, or it would +/// drop the very state the flush was meant to persist. +#[tokio::test] +async fn cleanup_connection_state_keeps_state_when_flush_fails() { + use wacore::libsignal::protocol::{ProtocolAddress, SenderKeyRecord}; + use wacore::libsignal::store::sender_key_name::SenderKeyName; + let client = create_offline_sync_test_client().await; + + // A malformed identity (not 32 bytes) makes flush() error out, standing + // in for a transient backend write failure during cleanup. + let bad = ProtocolAddress::new("5550002000@s.whatsapp.net".to_string(), 1u32.into()); + client.signal_cache.put_identity(&bad, &[0u8; 16]).await; + + // A valid dirty sender key that must not be dropped when the flush fails. + let name = SenderKeyName::from_parts("group@g.us", "5550001000@s.whatsapp.net:1"); + client + .signal_cache + .put_sender_key(&name, SenderKeyRecord::new_empty()) + .await; + + client.cleanup_connection_state().await; + + // flush() failed, so clear() was skipped; the unpersisted sender key + // survives in the write-back cache instead of being dropped. + let device = client.persistence_manager.get_device_arc().await; + let guard = device.read().await; + let persisted = client + .signal_cache + .get_sender_key(&name, &*guard.backend) + .await + .expect("get_sender_key must not error"); + assert!( + persisted.is_some(), + "a flush failure must not drop dirty Signal state" + ); +} + +/// A 403 connect failure is WA Web's REASON_LOCKED: it must surface a logout +/// carrying AccountLocked and disable auto-reconnect (a lock is not transient). +#[tokio::test] +async fn connect_failure_403_dispatches_account_locked_logout() { + use wacore::types::events::ChannelEventHandler; + let client = create_offline_sync_test_client().await; + let (handler, events) = ChannelEventHandler::new(); + client.register_handler(handler); + + // location="rva" is a region routing token and must not change the verdict. + let failure = NodeBuilder::new("failure") + .attr("reason", "403") + .attr("location", "rva") + .build(); + client.handle_connect_failure(&failure.as_node_ref()).await; + + let evt = events + .try_recv() + .expect("403 must dispatch a LoggedOut event"); + match &*evt { + Event::LoggedOut(lo) => { + assert!(lo.on_connect, "403 arrives as a failure-on-connect"); + assert_eq!(lo.reason, ConnectFailureReason::AccountLocked); + } + _ => panic!("expected Event::LoggedOut for reason=403"), + } + assert!( + !client.enable_auto_reconnect.load(Ordering::Relaxed), + "a server-side lock must not auto-reconnect" + ); +} + +#[tokio::test] +async fn delivery_receipt_activity_state_machine() { + let client = create_offline_sync_test_client().await; + assert!( + !client.receipts_are_active(), + "default is inactive (background companion)" + ); + client.mark_receipts_active_on_presence(); + assert!(client.receipts_are_active(), "presence available -> active"); + client.mark_receipts_inactive_on_presence(); + assert!( + !client.receipts_are_active(), + "presence unavailable -> inactive" + ); + client.set_force_active_delivery_receipts(true); + assert!(client.receipts_are_active(), "forced active"); + client.mark_receipts_inactive_on_presence(); + assert!( + client.receipts_are_active(), + "forced (2) survives a presence-unavailable CAS(1,0)" + ); + client.set_force_active_delivery_receipts(false); + assert!(!client.receipts_are_active()); + + // Teardown resets presence-driven active (so it doesn't leak across + // reconnects) but preserves a forced value. + client.mark_receipts_active_on_presence(); + client.cleanup_connection_state().await; + assert!( + !client.receipts_are_active(), + "teardown resets presence-driven active" + ); + client.set_force_active_delivery_receipts(true); + client.cleanup_connection_state().await; + assert!( + client.receipts_are_active(), + "teardown preserves forced active" + ); +} + +#[tokio::test] +async fn test_ib_thread_metadata_does_not_end_sync() { + let client = create_offline_sync_test_client().await; + client + .offline_sync_metrics + .active + .store(true, Ordering::Release); + + let node = NodeBuilder::new("ib") + .children([NodeBuilder::new("thread_metadata") + .children([NodeBuilder::new("item").build()]) + .build()]) + .build(); + + client.process_node(node_to_owned_ref(node)).await; + assert!( + client.offline_sync_metrics.active.load(Ordering::Acquire), + " should NOT end offline sync" + ); +} + +#[tokio::test] +async fn test_ib_edge_routing_does_not_end_sync() { + let client = create_offline_sync_test_client().await; + client + .offline_sync_metrics + .active + .store(true, Ordering::Release); + + let node = NodeBuilder::new("ib") + .children([NodeBuilder::new("edge_routing") + .children([NodeBuilder::new("routing_info") + .bytes(vec![1, 2, 3]) + .build()]) + .build()]) + .build(); + + client.process_node(node_to_owned_ref(node)).await; + assert!( + client.offline_sync_metrics.active.load(Ordering::Acquire), + " should NOT end offline sync" + ); +} + +#[tokio::test] +async fn test_ib_dirty_does_not_end_sync() { + let client = create_offline_sync_test_client().await; + client + .offline_sync_metrics + .active + .store(true, Ordering::Release); + + let node = NodeBuilder::new("ib") + .children([NodeBuilder::new("dirty") + .attr("type", "groups") + .attr("timestamp", "1234") + .build()]) + .build(); + + client.process_node(node_to_owned_ref(node)).await; + assert!( + client.offline_sync_metrics.active.load(Ordering::Acquire), + " should NOT end offline sync" + ); +} + +#[tokio::test] +async fn test_ib_offline_child_ends_sync() { + let client = create_offline_sync_test_client().await; + client + .offline_sync_metrics + .active + .store(true, Ordering::Release); + client + .offline_sync_metrics + .total_messages + .store(301, Ordering::Release); + + let node = NodeBuilder::new("ib") + .children([NodeBuilder::new("offline").attr("count", "301").build()]) + .build(); + + client.process_node(node_to_owned_ref(node)).await; + assert!( + !client.offline_sync_metrics.active.load(Ordering::Acquire), + " should end offline sync" + ); +} + +#[tokio::test] +async fn test_ib_offline_preview_starts_sync() { + let client = create_offline_sync_test_client().await; + + let node = NodeBuilder::new("ib") + .children([NodeBuilder::new("offline_preview") + .attr("count", "301") + .attr("message", "168") + .attr("notification", "62") + .attr("receipt", "68") + .attr("appdata", "0") + .build()]) + .build(); + + client.process_node(node_to_owned_ref(node)).await; + assert!( + client.offline_sync_metrics.active.load(Ordering::Acquire), + "offline_preview with count>0 should activate sync" + ); + assert_eq!( + client + .offline_sync_metrics + .total_messages + .load(Ordering::Acquire), + 301 + ); +} + +#[tokio::test] +async fn test_offline_message_increments_processed() { + let client = create_offline_sync_test_client().await; + client + .offline_sync_metrics + .active + .store(true, Ordering::Release); + client + .offline_sync_metrics + .total_messages + .store(100, Ordering::Release); + + let node = NodeBuilder::new("message") + .attr("offline", "1") + .attr("from", "5551234567@s.whatsapp.net") + .attr("id", "TEST123") + .attr("t", "1772884671") + .attr("type", "text") + .build(); + + client.process_node(node_to_owned_ref(node)).await; + assert_eq!( + client + .offline_sync_metrics + .processed_messages + .load(Ordering::Acquire), + 1, + "offline message should increment processed count" + ); +} + +// --------------------------------------------------------------- +// Server-initiated ping detection tests +// +// The WhatsApp server can send pings in two formats: +// +// 1. Child-element format (legacy/whatsmeow style): +// +// +// +// +// 2. xmlns-attribute format (real WhatsApp Web format): +// +// This is a self-closing tag with NO child elements. +// Verified against captured WhatsApp Web JS (WAWebCommsHandleStanza): +// if (t.xmlns === "urn:xmpp:ping") return wap("iq", { type: "result", to: t.from }); +// +// Both must be recognized and answered with a pong, otherwise the +// server considers the client dead and stops responding to keepalive +// pings — causing a timeout cascade and forced reconnect. +// --------------------------------------------------------------- + +#[tokio::test] +async fn test_handle_iq_ping_with_child_element() { + // Format 1: — the legacy format with a child node. + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + let ping_node = NodeBuilder::new("iq") + .attr("type", "get") + .attr("from", SERVER_JID) + .attr("id", "ping-child-1") + .children([NodeBuilder::new("ping").build()]) + .build(); + + let handled = client.handle_iq(&ping_node.as_node_ref()).await; + assert!( + handled, + "handle_iq must recognize ping with child element" + ); +} + +#[tokio::test] +async fn test_handle_iq_ping_with_xmlns_attribute() { + // Format 2: — the real WhatsApp Web format. + // This is a self-closing IQ with NO children, only an xmlns attribute. + // The server sends this format; failing to respond causes keepalive timeout cascade. + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + let ping_node = NodeBuilder::new("iq") + .attr("type", "get") + .attr("from", SERVER_JID) + .attr("id", "ping-xmlns-1") + .attr("xmlns", "urn:xmpp:ping") + .build(); + + let handled = client.handle_iq(&ping_node.as_node_ref()).await; + assert!( + handled, + "handle_iq must recognize ping with xmlns=\"urn:xmpp:ping\" attribute (no children)" + ); +} + +#[tokio::test] +async fn test_handle_iq_ping_with_both_child_and_xmlns() { + // Edge case: node has BOTH a child AND xmlns="urn:xmpp:ping". + // Should still be handled (OR condition). + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + let ping_node = NodeBuilder::new("iq") + .attr("type", "get") + .attr("from", SERVER_JID) + .attr("id", "ping-both-1") + .attr("xmlns", "urn:xmpp:ping") + .children([NodeBuilder::new("ping").build()]) + .build(); + + let handled = client.handle_iq(&ping_node.as_node_ref()).await; + assert!( + handled, + "handle_iq must handle ping with both child and xmlns" + ); +} + +#[tokio::test] +async fn test_handle_iq_non_ping_returns_false() { + // A type="get" IQ without ping child or xmlns should NOT be handled as ping. + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + let non_ping_node = NodeBuilder::new("iq") + .attr("type", "get") + .attr("from", SERVER_JID) + .attr("id", "not-a-ping") + .attr("xmlns", "some:other:namespace") + .build(); + + let handled = client.handle_iq(&non_ping_node.as_node_ref()).await; + assert!( + !handled, + "handle_iq must NOT treat non-ping xmlns as a ping" + ); +} + +#[tokio::test] +async fn test_handle_iq_ping_wrong_type_returns_false() { + // xmlns="urn:xmpp:ping" but type="result" (not "get") — should NOT be handled as ping. + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + let result_node = NodeBuilder::new("iq") + .attr("type", "result") + .attr("from", SERVER_JID) + .attr("id", "ping-result-1") + .attr("xmlns", "urn:xmpp:ping") + .build(); + + let handled = client.handle_iq(&result_node.as_node_ref()).await; + assert!( + !handled, + "handle_iq must NOT respond to type=\"result\" even with ping xmlns" + ); +} + +// ── build_pong tests ────────────────────────────────────────────── + +#[test] +fn test_build_pong_with_id() { + let pong = build_pong("s.whatsapp.net".to_string(), Some("ping-123")); + assert!( + pong.attrs.get("id").is_some_and(|v| v == "ping-123"), + "pong should include id when server ping has one" + ); + assert!(pong.attrs.get("type").is_some_and(|v| v == "result")); + assert!(pong.attrs.get("to").is_some_and(|v| v == "s.whatsapp.net")); +} + +#[test] +fn test_build_pong_without_id() { + let pong = build_pong("s.whatsapp.net".to_string(), None); + assert!( + !pong.attrs.contains_key("id"), + "pong should NOT include id when server ping has none" + ); + assert!(pong.attrs.get("type").is_some_and(|v| v == "result")); +} + +#[test] +fn test_encrypt_identity_notification_omits_type() { + let node = NodeBuilder::new("notification") + .attr("from", "186303081611421@lid") + .attr("id", "4128735301") + .attr("type", "encrypt") + .children([NodeBuilder::new("identity").build()]) + .build(); + + assert!( + is_encrypt_identity_notification(&node.as_node_ref()), + "identity-change notification ACK must omit type to match WA Web" + ); +} + +#[test] +fn test_device_notification_is_not_encrypt_identity() { + let node = NodeBuilder::new("notification") + .attr("from", "186303081611421@lid") + .attr("id", "269488578") + .attr("type", "devices") + .children([NodeBuilder::new("remove").build()]) + .build(); + + assert!( + !is_encrypt_identity_notification(&node.as_node_ref()), + "device notification is not an encrypt+identity notification" + ); +} + +#[test] +fn test_build_ack_node_for_message_omits_type_includes_from() { + // Whatsmeow: message acks do NOT echo type (node.Tag != "message" guard). + // They DO include `from` with own device PN. + let incoming = NodeBuilder::new("message") + .attr("from", "120363161500776365@g.us") + .attr("id", "A5791A5392EF60E3FB0670098DE010D4") + .attr("type", "text") + .attr("participant", "181531758878822@lid") + .build(); + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) + .expect("message ack should be buildable"); + + assert_eq!(ack.tag, "ack"); + // Use PartialEq on NodeValue — works for both String and Jid variants + // without allocation, so tests don't depend on internal representation. + assert!(ack.attrs.get("class").is_some_and(|v| v == "message")); + assert!( + ack.attrs + .get("to") + .is_some_and(|v| v == "120363161500776365@g.us") + ); + assert!( + ack.attrs + .get("from") + .is_some_and(|v| v == "155500012345:48@s.whatsapp.net") + ); + assert!( + ack.attrs + .get("participant") + .is_some_and(|v| v == "181531758878822@lid") + ); + assert!( + !ack.attrs.contains_key("type"), + "message ACK must NOT echo type (matches whatsmeow behavior)" + ); +} + +#[test] +fn test_build_ack_node_for_identity_change_omits_type_and_from() { + let incoming = NodeBuilder::new("notification") + .attr("from", "186303081611421@lid") + .attr("id", "4128735301") + .attr("type", "encrypt") + .children([NodeBuilder::new("identity").build()]) + .build(); + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) + .expect("notification ack should be buildable"); + + assert!(ack.attrs.get("class").is_some_and(|v| v == "notification")); + assert!( + !ack.attrs.contains_key("type"), + "identity-change notification ACK must omit type" + ); + assert!( + !ack.attrs.contains_key("from"), + "notification ACKs should not include our device PN" + ); +} + +#[test] +fn test_build_ack_node_for_receipt_with_type_echoes_type() { + // Receipt acks should echo the type attribute when present (e.g. "read", "played"). + let incoming = NodeBuilder::new("receipt") + .attr("from", "156535032389744@lid") + .attr("id", "RCPT-WITH-TYPE") + .attr("type", "read") + .build(); + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) + .expect("receipt ack should be buildable"); + + assert!(ack.attrs.get("class").is_some_and(|v| v == "receipt")); + assert!( + ack.attrs.get("type").is_some_and(|v| v == "read"), + "receipt ACK must echo the type attribute when present" + ); + assert!( + !ack.attrs.contains_key("from"), + "receipt ACKs should not include our device PN" + ); +} + +#[test] +fn test_build_ack_node_drops_participant_when_equal_to_from() { + // WAWebReceiptAck: `participant: r && r !== e ? DEVICE_JID(r) : DROP_ATTR`. + // When the incoming stanza carries participant == from (redundant), + // the ack must not echo it. + let incoming = NodeBuilder::new("receipt") + .attr("from", "156535032389744@lid") + .attr("participant", "156535032389744@lid") + .attr("id", "RCPT-PARTICIPANT-EQ-FROM") + .build(); + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net".parse().unwrap(); + + let ack = + build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)).expect("ack should build"); + assert!( + !ack.attrs.contains_key("participant"), + "ack must drop participant when it duplicates `to` (the flipped from); got {:?}", + ack.attrs.get("participant") + ); +} + +#[test] +fn test_build_ack_node_keeps_participant_when_distinct_from_from() { + // Group receipt: participant = sender (user), from = group jid; must be kept. + let incoming = NodeBuilder::new("receipt") + .attr("from", "120363098765432100@g.us") + .attr("participant", "5511999999999@s.whatsapp.net") + .attr("id", "RCPT-GROUP") + .build(); + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net".parse().unwrap(); + + let ack = + build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)).expect("ack should build"); + assert!( + ack.attrs + .get("participant") + .is_some_and(|v| v == "5511999999999@s.whatsapp.net"), + "ack must keep participant when it differs from `to`" + ); +} + +#[test] +fn test_build_ack_node_for_receipt_without_type_omits_type() { + // Delivery receipts have no type attribute — the ack must also omit it. + // Sending type="delivery" in the ack causes stream:error disconnections. + let incoming = NodeBuilder::new("receipt") + .attr("from", "156535032389744@lid") + .attr("id", "RCPT-NO-TYPE") + .build(); + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) + .expect("receipt ack should be buildable"); + + assert!(ack.attrs.get("class").is_some_and(|v| v == "receipt")); + assert!( + !ack.attrs.contains_key("type"), + "receipt ACK must NOT contain type when the incoming receipt has no type attribute" + ); + assert!( + !ack.attrs.contains_key("from"), + "receipt ACKs should not include our device PN" + ); +} + +#[test] +fn test_build_ack_node_for_message_with_recipient_preserves_recipient() { + // Peer / hosted-companion / LID-routed messages carry `recipient`. + // The server uses it to route the ack back to the origin device; + // without it the stream is torn down with . + let incoming = NodeBuilder::new("message") + .attr("from", "166361967902821@lid") + .attr("id", "2A32F960553696093D99") + .attr("type", "text") + .attr("recipient", "146991363395800@lid") + .build(); + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) + .expect("message ack should be buildable"); + + assert!(ack.attrs.get("class").is_some_and(|v| v == "message")); + assert!( + ack.attrs + .get("recipient") + .is_some_and(|v| v == "146991363395800@lid"), + "message ACK must echo the incoming `recipient` attribute" + ); +} + +#[test] +fn test_build_ack_node_for_receipt_with_recipient_preserves_recipient() { + // Receipt acks must also echo `recipient` when the incoming carries it. + let incoming = NodeBuilder::new("receipt") + .attr("from", "120363098765432100@g.us") + .attr("id", "RCPT-WITH-RECIPIENT") + .attr("type", "read") + .attr("recipient", "242395589390497@lid") + .build(); + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) + .expect("receipt ack should be buildable"); + + assert!(ack.attrs.get("class").is_some_and(|v| v == "receipt")); + assert!( + ack.attrs + .get("recipient") + .is_some_and(|v| v == "242395589390497@lid"), + "receipt ACK must echo the incoming `recipient` attribute" + ); +} + +#[test] +fn test_build_ack_node_for_message_without_recipient_omits_recipient() { + // Regression guard: never synthesise a `recipient` field if the + // incoming stanza did not carry one — server would reject the ack. + let incoming = NodeBuilder::new("message") + .attr("from", "120363161500776365@g.us") + .attr("id", "A5791A5392EF60E3FB06") + .attr("type", "text") + .attr("participant", "181531758878822@lid") + .build(); + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let ack = build_ack_node(&incoming.as_node_ref(), Some(&own_device_pn)) + .expect("message ack should be buildable"); + + assert!( + !ack.attrs.contains_key("recipient"), + "ACK must NOT add `recipient` when the incoming stanza has none" + ); +} + +#[test] +fn test_encode_ack_bytes_roundtrip_recipient() { + // Exercises the real wire encoder (`encode_ack_bytes`), not just the + // `build_ack_node` test mirror: serialize, decode the bytes back, and + // assert the parsed ACK echoes `recipient` when present and omits it + // when absent. Guards against the two builders silently diverging. + let own_device_pn: Jid = "155500012345:48@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let with_recipient = NodeBuilder::new("message") + .attr("from", "166361967902821@lid") + .attr("id", "2A32F960553696093D99") + .attr("type", "text") + .attr("recipient", "146991363395800@lid") + .build(); + let buf = encode_ack_bytes(&with_recipient.as_node_ref(), Some(&own_device_pn)) + .expect("encode_ack_bytes should not error") + .expect("encode_ack_bytes should produce bytes"); + // The Encoder prepends a leading format byte (see `marshal`); the + // decoder wants raw protocol bytes — same handling as `node_to_owned_ref`. + let decoded = + wacore_binary::marshal::unmarshal_ref(&buf[1..]).expect("encoded ack should decode"); + assert_eq!(decoded.tag, "ack"); + assert!( + decoded + .get_attr("class") + .is_some_and(|v| v.as_str() == "message"), + "decoded ack must have class=message" + ); + assert!( + decoded + .get_attr("recipient") + .is_some_and(|v| v.as_str() == "146991363395800@lid"), + "encode_ack_bytes must echo `recipient` onto the wire" + ); + + let without_recipient = NodeBuilder::new("message") + .attr("from", "120363161500776365@g.us") + .attr("id", "A5791A5392EF60E3FB06") + .attr("type", "text") + .attr("participant", "181531758878822@lid") + .build(); + let buf = encode_ack_bytes(&without_recipient.as_node_ref(), Some(&own_device_pn)) + .expect("encode_ack_bytes should not error") + .expect("encode_ack_bytes should produce bytes"); + let decoded = + wacore_binary::marshal::unmarshal_ref(&buf[1..]).expect("encoded ack should decode"); + assert!( + decoded.get_attr("recipient").is_none(), + "encode_ack_bytes must not synthesise `recipient` when absent" + ); +} + +/// Own-account fan-out ack must address back to the original `from` (own +/// LID) echoing `recipient`, not to the chat. Guards against regressing to +/// the chat-addressed `build_nack_node` style. +#[test] +fn test_message_ack_source_node_own_device_addressing() { + use crate::types::message::{MessageInfo, MessageSource}; + // Own-account branch: sender == `from` (device-qualified), chat is the + // device-stripped recipient. `to` must come from sender, not chat. + let info = MessageInfo { + id: "AC055553E56A2C12DE592DAD6353C477".to_string(), + source: MessageSource { + sender: "236395184570386@lid".parse().expect("sender"), + chat: "156535032389744@lid".parse().expect("chat"), + recipient: Some("156535032389744@lid".parse().expect("recipient")), + is_group: false, + ..Default::default() + }, + ..Default::default() + }; + let own_device_pn: Jid = "559984726662:95@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let source = message_ack_source_node(&info); + let built = build_ack_node(&source.as_node_ref(), Some(&own_device_pn)) + .expect("message ack should be buildable"); + + assert!(built.attrs.get("class").is_some_and(|v| v == "message")); + assert!( + built + .attrs + .get("to") + .is_some_and(|v| v == "236395184570386@lid"), + "ack `to` must be the original `from` (own LID), not the chat" + ); + assert!( + built + .attrs + .get("recipient") + .is_some_and(|v| v == "156535032389744@lid"), + "ack must echo `recipient` so the server can route/clear it" + ); + assert!( + !built.attrs.contains_key("type"), + "message-class acks never carry a `type`" + ); +} + +/// Common incoming DM from another user: `to` is the device-qualified +/// sender, with no `recipient`/`participant` synthesised. +#[test] +fn test_message_ack_source_node_incoming_dm_addressing() { + use crate::types::message::{MessageInfo, MessageSource}; + let info = MessageInfo { + id: "MSGID".to_string(), + source: MessageSource { + sender: "5511999998888:3@s.whatsapp.net".parse().expect("sender"), + chat: "5511999998888@s.whatsapp.net".parse().expect("chat"), + is_group: false, + ..Default::default() + }, + ..Default::default() + }; + let own_device_pn: Jid = "559984726662:95@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let source = message_ack_source_node(&info); + let built = build_ack_node(&source.as_node_ref(), Some(&own_device_pn)) + .expect("dm ack should be buildable"); + + assert!( + built + .attrs + .get("to") + .is_some_and(|v| v == "5511999998888:3@s.whatsapp.net"), + "ack `to` must be the device-qualified sender (the original `from`)" + ); + assert!(!built.attrs.contains_key("recipient")); + assert!(!built.attrs.contains_key("participant")); +} + +/// status@broadcast (is_group=true in the parser) addresses the ack to the +/// status chat, with the sender as participant, not to the sender. +#[test] +fn test_message_ack_source_node_status_addressing() { + use crate::types::message::{MessageInfo, MessageSource}; + let info = MessageInfo { + id: "STATUSMSG".to_string(), + source: MessageSource { + chat: "status@broadcast".parse().expect("status chat"), + sender: "181531758878822@lid".parse().expect("participant"), + is_group: true, + ..Default::default() + }, + ..Default::default() + }; + let own_device_pn: Jid = "559984726662:95@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let source = message_ack_source_node(&info); + let built = build_ack_node(&source.as_node_ref(), Some(&own_device_pn)) + .expect("status ack should be buildable"); + + assert!( + built + .attrs + .get("to") + .is_some_and(|v| v == "status@broadcast"), + "status ack `to` must be the status chat, not the sender" + ); + assert!( + built + .attrs + .get("participant") + .is_some_and(|v| v == "181531758878822@lid"), + "status ack must preserve the sending participant" + ); +} + +/// Group failure ack: `to` is the group, `participant` is preserved. +#[test] +fn test_message_ack_source_node_group_addressing() { + use crate::types::message::{MessageInfo, MessageSource}; + // Group branch: chat == group `from`, sender == participant. + let info = MessageInfo { + id: "GROUPMSGID".to_string(), + source: MessageSource { + chat: "120363011111111111@g.us".parse().expect("group"), + sender: "181531758878822@lid".parse().expect("participant"), + is_group: true, + ..Default::default() + }, + ..Default::default() + }; + let own_device_pn: Jid = "559984726662:95@s.whatsapp.net" + .parse() + .expect("own device PN JID should parse"); + + let source = message_ack_source_node(&info); + let built = build_ack_node(&source.as_node_ref(), Some(&own_device_pn)) + .expect("group message ack should be buildable"); + + assert!( + built + .attrs + .get("to") + .is_some_and(|v| v == "120363011111111111@g.us"), + "group ack `to` must be the group JID" + ); + assert!( + built + .attrs + .get("participant") + .is_some_and(|v| v == "181531758878822@lid"), + "group ack must preserve the sending `participant`" + ); +} + +/// Smoke test: server ping with xmlns but no id attribute is handled. +#[tokio::test] +async fn test_handle_iq_ping_without_id() { + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Server ping without id — real format observed in production logs + let ping_node = NodeBuilder::new("iq") + .attr("type", "get") + .attr("from", SERVER_JID) + .attr("xmlns", "urn:xmpp:ping") + .build(); + + let handled = client.handle_iq(&ping_node.as_node_ref()).await; + assert!( + handled, + "handle_iq must recognize ping without id attribute" + ); +} + +// ── fibonacci_backoff tests ──────────────────────────────────────── + +#[test] +fn test_fibonacci_backoff_sequence() { + // WA Web: first=1000, second=1000 → 1,1,2,3,5,8,13,21,34,55,89,144...s + // We test base values without jitter by checking the range (±10%). + let expected_base_ms = [1000, 1000, 2000, 3000, 5000, 8000, 13000, 21000]; + for (attempt, &base) in expected_base_ms.iter().enumerate() { + let delay = fibonacci_backoff(attempt as u32); + let ms = delay.as_millis() as u64; + let low = base - base / 10; + let high = base + base / 10; + assert!( + ms >= low && ms <= high, + "attempt {attempt}: expected {low}..={high}ms, got {ms}ms" + ); + } +} + +#[test] +fn test_fibonacci_backoff_max_900s() { + // After many attempts, should cap at 900s (±10%) + let delay = fibonacci_backoff(100); + let ms = delay.as_millis() as u64; + assert!( + ms <= 990_000, + "should never exceed 900s + 10% jitter, got {ms}ms" + ); + assert!( + ms >= 810_000, + "should be at least 900s - 10% jitter, got {ms}ms" + ); +} + +#[test] +fn test_fibonacci_backoff_first_attempt_is_1s() { + let delay = fibonacci_backoff(0); + let ms = delay.as_millis() as u64; + assert!( + (900..=1100).contains(&ms), + "first attempt should be ~1s (±10%), got {ms}ms" + ); +} + +// ── stream error tests ───────────────────────────────────────────── + +#[tokio::test] +async fn test_stream_error_401_disables_reconnect() { + let client = create_offline_sync_test_client().await; + let node = NodeBuilder::new("stream:error").attr("code", "401").build(); + client.handle_stream_error(&node.as_node_ref()).await; + assert!( + !client.enable_auto_reconnect.load(Ordering::Relaxed), + "401 should disable auto-reconnect" + ); +} + +#[tokio::test] +async fn test_stream_error_409_disables_reconnect() { + let client = create_offline_sync_test_client().await; + let node = NodeBuilder::new("stream:error").attr("code", "409").build(); + client.handle_stream_error(&node.as_node_ref()).await; + assert!( + !client.enable_auto_reconnect.load(Ordering::Relaxed), + "409 should disable auto-reconnect" + ); +} + +#[tokio::test] +async fn test_stream_error_429_keeps_reconnect_with_backoff() { + let client = create_offline_sync_test_client().await; + client.is_logged_in.store(true, Ordering::Relaxed); + let before = client.auto_reconnect_errors.load(Ordering::Relaxed); + let node = NodeBuilder::new("stream:error").attr("code", "429").build(); + client.handle_stream_error(&node.as_node_ref()).await; + assert!( + client.enable_auto_reconnect.load(Ordering::Relaxed), + "429 should keep auto-reconnect enabled" + ); + assert!( + !client.is_logged_in.load(Ordering::Relaxed), + "429 must clear is_logged_in so sends bail before the server flags abuse" + ); + assert!( + !client.expected_disconnect.load(Ordering::Relaxed), + "429 must not mark the disconnect as expected (auto-reconnect path)" + ); + let after = client.auto_reconnect_errors.load(Ordering::Relaxed); + assert_eq!( + after, + before + 5, + "429 should increase backoff by exactly 5: before={before}, after={after}" + ); +} + +#[tokio::test] +async fn test_stream_error_503_keeps_reconnect() { + let client = create_offline_sync_test_client().await; + client.is_logged_in.store(true, Ordering::Relaxed); + let node = NodeBuilder::new("stream:error").attr("code", "503").build(); + client.handle_stream_error(&node.as_node_ref()).await; + assert!( + client.enable_auto_reconnect.load(Ordering::Relaxed), + "503 should keep auto-reconnect enabled" + ); + assert!( + !client.is_logged_in.load(Ordering::Relaxed), + "503 must clear is_logged_in so sends bail against the dying socket" + ); + assert!( + !client.expected_disconnect.load(Ordering::Relaxed), + "503 must not mark the disconnect as expected (auto-reconnect path)" + ); +} + +#[tokio::test] +async fn test_stream_error_unknown_keeps_connection_alive() { + // Unknown stream:error (no `code` attribute) must mirror whatsmeow's + // default branch: log + dispatch event, but NOT mark this as an + // expected disconnect. Setting that flag silently swallows the next + // real disconnect and races the read loop into shutdown. + let client = create_offline_sync_test_client().await; + // Simulate an authenticated session before the stream error arrives. + client.is_logged_in.store(true, Ordering::Relaxed); + let node = NodeBuilder::new("stream:error").build(); + client.handle_stream_error(&node.as_node_ref()).await; + assert!( + client.is_logged_in.load(Ordering::Relaxed), + "unknown stream:error must NOT log the client out" + ); + assert!( + !client.expected_disconnect.load(Ordering::Relaxed), + "unknown stream:error must not mark the disconnect as expected" + ); + assert!( + client.enable_auto_reconnect.load(Ordering::Relaxed), + "unknown stream:error must keep auto-reconnect enabled" + ); +} + +#[tokio::test] +async fn test_stream_error_ack_shaped_does_not_force_shutdown() { + // Server wraps per-stanza routing failures in `` + // with no `code` attribute. Treat as informational, not as a fatal + // stream teardown. + let client = create_offline_sync_test_client().await; + client.is_logged_in.store(true, Ordering::Relaxed); + let ack_child = NodeBuilder::new("ack") + .attr("class", "message") + .attr("type", "text") + .attr("id", "2A32F960553696093D99") + .build(); + let node = NodeBuilder::new("stream:error") + .children([ack_child]) + .build(); + client.handle_stream_error(&node.as_node_ref()).await; + assert!( + client.is_logged_in.load(Ordering::Relaxed), + "ack-shaped stream:error must NOT log the client out" + ); + assert!( + !client.expected_disconnect.load(Ordering::Relaxed), + "ack-shaped stream:error must not mark the disconnect as expected" + ); +} + +#[tokio::test] +async fn test_custom_cache_config_is_respected() { + use crate::cache_config::{CacheConfig, CacheEntryConfig}; + use std::time::Duration; + + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + + let custom_config = CacheConfig { + group_cache: CacheEntryConfig::new(Some(Duration::from_secs(60)), 10), + device_registry_cache: CacheEntryConfig::new(Some(Duration::from_secs(60)), 10), + ..CacheConfig::default() + }; + + // Verify that constructing a client with a custom config does not panic + // and the client is usable. + let (client, _rx) = Client::new_with_cache_config( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + custom_config, + ) + .await; + + assert!(!client.is_logged_in()); +} + +/// Proves that `is_connected()` no longer gives false negatives under mutex +/// contention. Before the fix, `try_lock()` would fail when another task held +/// the noise_socket mutex, causing `is_connected()` to return `false` even +/// though the connection was alive — silently dropping receipt acks. +/// +/// This test sets up a real NoiseSocket (same as socket unit tests) so it +/// accurately models the pre-fix scenario: socket is Some + mutex is held +/// by another task = old is_connected() returned false. +#[tokio::test] +async fn test_is_connected_not_affected_by_mutex_contention() { + use crate::socket::NoiseSocket; + use wacore::handshake::NoiseCipher; + + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Initially not connected + assert!(!client.is_connected(), "should start disconnected"); + + // Simulate a real connection: create a NoiseSocket and store it + let transport: Arc = + Arc::new(crate::transport::mock::MockTransport); + let key = [0u8; 32]; + let write_key = NoiseCipher::new(&key).expect("valid key"); + let read_key = NoiseCipher::new(&key).expect("valid key"); + let noise_socket = NoiseSocket::new( + Arc::new(crate::runtime_impl::TokioRuntime), + transport, + write_key, + read_key, + ); + *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + client.is_connected.store(true, Ordering::Release); + + assert!(client.is_connected(), "should report connected"); + + // Hold the noise_socket mutex — this used to make is_connected() return + // false via try_lock() even though the socket was Some(...) + let _guard = client.noise_socket.lock().await; + assert!( + client.is_connected(), + "is_connected() must return true even while noise_socket mutex is held" + ); +} + +#[tokio::test] +async fn disconnect_does_not_signal_connection_cleanup_before_outbound_flush() { + use crate::socket::NoiseSocket; + use async_trait::async_trait; + use bytes::Bytes; + use wacore::handshake::NoiseCipher; + + struct BlockingTransport { + send_started: async_channel::Sender<()>, + release_send: async_channel::Receiver<()>, + send_done: Arc, + disconnect_called: Arc, + disconnect_before_send_done: Arc, + } + + #[async_trait] + impl crate::transport::Transport for BlockingTransport { + async fn send(&self, _data: Bytes) -> Result<(), anyhow::Error> { + let _ = self.send_started.try_send(()); + let _ = self.release_send.recv().await; + self.send_done.store(true, Ordering::Release); + Ok(()) + } + + async fn disconnect(&self) { + if !self.send_done.load(Ordering::Acquire) { + self.disconnect_before_send_done + .store(true, Ordering::Release); + } + self.disconnect_called.store(true, Ordering::Release); + } + } + + let client = crate::test_utils::create_test_client().await; + let (send_started_tx, send_started_rx) = async_channel::bounded(1); + let (release_send_tx, release_send_rx) = async_channel::bounded(1); + let send_done = Arc::new(AtomicBool::new(false)); + let disconnect_called = Arc::new(AtomicBool::new(false)); + let disconnect_before_send_done = Arc::new(AtomicBool::new(false)); + + let transport_impl = Arc::new(BlockingTransport { + send_started: send_started_tx, + release_send: release_send_rx, + send_done: Arc::clone(&send_done), + disconnect_called: Arc::clone(&disconnect_called), + disconnect_before_send_done: Arc::clone(&disconnect_before_send_done), + }); + let transport: Arc = transport_impl; + + let key = [0u8; 32]; + let write_key = NoiseCipher::new(&key).expect("valid key"); + let read_key = NoiseCipher::new(&key).expect("valid key"); + let noise_socket = NoiseSocket::new( + client.runtime.clone(), + Arc::clone(&transport), + write_key, + read_key, + ); + + *client.transport.lock().await = Some(transport); + *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + client.is_connected.store(true, Ordering::Release); + + let cleanup_signal = client.connection_shutdown_signal(); + let cleanup_client = Arc::clone(&client); + let cleanup_task = tokio::spawn(async move { + wacore::runtime::wait_for_shutdown(&cleanup_signal).await; + cleanup_client.cleanup_connection_state().await; + }); + + let send_client = Arc::clone(&client); + client.outbound_flush.spawn(&*client.runtime, async move { + let receipt = NodeBuilder::new("receipt") + .attr("id", "TEST-FLUSH-ORDER") + .attr("to", "1234567890@s.whatsapp.net") + .build(); + let _ = send_client.send_node(receipt).await; + }); + + tokio::time::timeout(Duration::from_secs(1), send_started_rx.recv()) + .await + .expect("tracked send should start") + .expect("send_started sender should stay open"); + + let disconnect_client = Arc::clone(&client); + let disconnect_task = tokio::spawn(async move { + disconnect_client.disconnect().await; + }); + + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + !client.connection_shutdown_signal().is_fired(), + "connection cleanup must not fire while outbound flush is blocked" + ); + assert!( + !disconnect_called.load(Ordering::Acquire), + "transport must stay open while outbound flush is blocked" + ); + + release_send_tx + .send(()) + .await + .expect("blocked send should still be waiting"); + + tokio::time::timeout(Duration::from_secs(1), disconnect_task) + .await + .expect("disconnect should finish") + .expect("disconnect task should not panic"); + tokio::time::timeout(Duration::from_secs(1), cleanup_task) + .await + .expect("cleanup should finish") + .expect("cleanup task should not panic"); + + assert!(send_done.load(Ordering::Acquire)); + assert!(disconnect_called.load(Ordering::Acquire)); + assert!( + !disconnect_before_send_done.load(Ordering::Acquire), + "cleanup closed the transport before the tracked send completed" + ); +} + +/// Verifies that `send_ack_for` returns an error (not silent Ok) when +/// disconnected. This ensures the caller's `warn!` fires so dropped acks +/// are visible in logs. +#[tokio::test] +async fn test_send_ack_for_returns_error_when_disconnected() { + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Not connected — send_ack_for should return Err, not Ok + let receipt = NodeBuilder::new("receipt") + .attr("from", "120363040237990503@g.us") + .attr("id", "TEST-RECEIPT-ID") + .attr("participant", "236395184570386@lid") + .build(); + + let result = client.send_ack_for(&receipt.as_node_ref()).await; + assert!( + matches!(result, Err(ClientError::NotConnected)), + "send_ack_for must return Err(NotConnected) when disconnected, got: {result:?}" + ); +} + +/// Verifies that `send_ack_for` returns Ok when expected_disconnect is set, +/// since this is an intentional shutdown path. +#[tokio::test] +async fn test_send_ack_for_returns_ok_on_expected_disconnect() { + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Set expected disconnect — send_ack_for should gracefully return Ok + client.expected_disconnect.store(true, Ordering::Relaxed); + + let receipt = NodeBuilder::new("receipt") + .attr("from", "120363040237990503@g.us") + .attr("id", "TEST-RECEIPT-ID") + .build(); + + let result = client.send_ack_for(&receipt.as_node_ref()).await; + assert!( + result.is_ok(), + "send_ack_for should return Ok during expected disconnect" + ); +} + +// Per-connection notify must NOT set the terminal sticky flag; if it did, +// every reconnect would instantly abort subscribers registered on the +// terminal signal. Regression guard for the CI breakage observed on PR #560. +#[tokio::test] +async fn per_connection_notify_leaves_terminal_signal_untouched() { + let client = crate::test_utils::create_test_client().await; + + client.notify_connection_shutdown(); + + assert!( + !client.shutdown_signal().is_fired(), + "terminal shutdown must stay clean when only per-connection fires" + ); +} + +// Subscribers registered AFTER a reset must not see the previous +// notifier's fired state. This is the core property that makes reconnect +// work: after cleanup_connection_state notifies the per-connection +// signal, the next connection replaces it with a fresh one. +#[tokio::test] +async fn reset_gives_fresh_per_connection_notifier() { + let client = crate::test_utils::create_test_client().await; + + client.notify_connection_shutdown(); + assert!( + client.connection_shutdown_signal().is_fired(), + "subscriber BEFORE reset sees the notify on the current notifier" + ); + + client.reset_connection_shutdown(); + + assert!( + !client.connection_shutdown_signal().is_fired(), + "subscribers AFTER reset must NOT see the previous notifier's state" + ); +} + +// Capture-once regression guard: a ShutdownSignal captured before a reset +// must keep observing the pre-reset fired state. Without this, a +// reconnect after the old notifier is replaced in the Mutex would +// strand long-lived tasks (e.g. keepalive) on a new notifier they +// never registered for. See keepalive_loop which captures its signal +// once at task startup. +#[tokio::test] +async fn captured_signal_keeps_observing_old_notifier_after_reset() { + let client = crate::test_utils::create_test_client().await; + + let captured = client.connection_shutdown_signal(); + client.notify_connection_shutdown(); + client.reset_connection_shutdown(); + + assert!( + captured.is_fired(), + "captured signal must retain the pre-reset notifier's fired state" + ); +} + +// Terminal disconnect() must also wake per-connection subscribers via +// cleanup_connection_state, so keepalive/request/read loop exit promptly. +#[tokio::test] +async fn terminal_disconnect_propagates_to_per_connection_signal() { + let client = crate::test_utils::create_test_client().await; + let conn_signal = client.connection_shutdown_signal(); + + client.disconnect().await; + + assert!( + conn_signal.is_fired(), + "disconnect must fire per-connection via cleanup_connection_state" + ); + assert!( + client.shutdown_signal().is_fired(), + "disconnect must also fire terminal" + ); +} diff --git a/src/message.rs b/src/message.rs index 32e4b5d34..be81b7d22 100644 --- a/src/message.rs +++ b/src/message.rs @@ -72,7 +72,7 @@ pub(crate) struct ClassifiedMessage { } #[derive(Clone, Copy, Debug, Default)] -struct SessionBatchOutcome { +pub(crate) struct SessionBatchOutcome { decrypted: bool, duplicate: bool, undecryptable: bool, @@ -92,7 +92,7 @@ struct MigrationDecryptOutcome { } #[derive(Clone, Copy, Debug, Default)] -struct PlaintextHandleOutcome { +pub(crate) struct PlaintextHandleOutcome { dispatched: bool, skdm_only: bool, } @@ -135,2804 +135,11 @@ fn decrypt_fail_log_level(mode: crate::types::events::DecryptFailMode) -> log::L pub(crate) use wacore::protocol::retry::RetryReason; -impl Client { - /// Dispatches a successfully parsed message to the event bus and sends a delivery receipt. - async fn dispatch_parsed_message(self: &Arc, msg: wa::Message, info: &Arc) { - use wacore::proto_helpers::MessageExt; - - let mut info = Arc::clone(info); - if info.ephemeral_expiration.is_none() - && msg.get_base_message().get_ephemeral_expiration().is_some() - { - Arc::make_mut(&mut info).ephemeral_expiration = - msg.get_base_message().get_ephemeral_expiration(); - } - - // Keep this ordered with dispatch; add-on messages can immediately - // reference the secret from the stanza just processed. - self.maybe_capture_inbound_msg_secret(&msg, &info).await; - let dispatch_msg = self - .maybe_decrypt_secret_encrypted_message(&msg, &info) - .await - .unwrap_or(msg); - self.ack_received_message(&info); - - self.core - .event_bus - .dispatch(Event::Message(Arc::new(dispatch_msg), info)); - } - - /// Acknowledge a received message so the server drops it from the offline - /// queue: a delivery receipt when applicable (incl. the `type="sender"` - /// receipt for own-account self-fanouts), else a transport ack. status is - /// acked by the `should_ack` gate, newsletters/empty ids need nothing here. - fn ack_received_message(self: &Arc, info: &Arc) { - if info.id.is_empty() || info.source.chat.is_newsletter() { - return; - } - // WA Web `sendAggregateReceipts`: for a DELIVERY where the chat is NOT - // a bot but the author IS a bot (a bot reply inside a group), it emits - // a bare `` via `sendBotInvokeResponseAcks`, not a - // ``. A 1:1 bot chat keeps the normal receipt (chat.isBot() → - // the branch's `v` is false). Our transport ack is that bare - // `` (group form carries `participant`). - if info.source.is_bot_authored_non_bot_chat() { - self.spawn_message_ack(info); - return; - } - if Self::should_send_delivery_receipt(info) { - self.spawn_delivery_receipt(info); - } else if !info.source.chat.is_status_broadcast() { - self.spawn_message_ack(info); - } - } - - /// Spawn a delivery receipt, tracked so `disconnect()` can flush it (issue #571). - fn spawn_delivery_receipt(self: &Arc, info: &Arc) { - let client = self.clone(); - let info = Arc::clone(info); - self.outbound_flush.spawn(&*self.runtime, async move { - client.send_delivery_receipt(&info).await; - }); - } - - /// Capture embedded `MessageContextInfo.message_secret` for add-on - /// decrypts. Bot DMs keep the legacy LID key as a second entry. - pub(crate) async fn maybe_capture_inbound_msg_secret( - self: &Arc, - msg: &wa::Message, - info: &Arc, - ) { - use wacore::proto_helpers::MessageExt; - - let mci = msg.message_context_info.as_ref(); - let Some(secret_bytes) = mci.and_then(|m| m.message_secret.as_deref()) else { - return; - }; - if msg.is_forwarded() { - return; - } - - let policy = self.cache_config.msg_secret_policy; - if !policy.persists() { - return; - } - let chat_is_bot = info.source.chat.is_bot(); - // BotOnly enforcement lives in build_msg_secret_entry (the chokepoint), - // which keys off the classified bot context including group bot prompts. - let class = wacore::msg_secret::classify(msg, chat_is_bot); - let message_ts = u64::try_from(info.timestamp.timestamp()).ok(); - - // Build both aliases (primary, plus the bot-DM LID key) and write them - // in one batch so a partial write can't leave only one stored. - let mut entries = Vec::with_capacity(2); - if let Some(entry) = self.build_msg_secret_entry( - &info.source.chat, - &info.source.sender, - &info.id, - secret_bytes, - class, - message_ts, - ) { - entries.push(entry); - } - if chat_is_bot - && let Some(sender) = self.dm_sender_identity_for(&info.source.chat).await - && sender.to_non_ad() != info.source.sender.to_non_ad() - && let Some(entry) = self.build_msg_secret_entry( - &info.source.chat, - &sender, - &info.id, - secret_bytes, - class, - message_ts, - ) - { - entries.push(entry); - } - self.persist_msg_secret_entries(entries).await; - } - - /// Build one retention entry, applying the policy gates and computing the - /// per-row deadline. Returns `None` when the policy skips this write (not - /// persisting, or `BotOnly` and the class isn't `Bot`) or the secret isn't - /// 32 bytes. Pure (no I/O) so callers can batch several aliases atomically. - fn build_msg_secret_entry( - &self, - chat: &Jid, - sender: &Jid, - msg_id: &str, - secret_bytes: &[u8], - class: wacore::msg_secret::RetentionClass, - message_ts: Option, - ) -> Option { - const SECRET_LEN: usize = wacore::reporting_token::MESSAGE_SECRET_SIZE; - let secret = <&[u8; SECRET_LEN]>::try_from(secret_bytes).ok()?; - let policy = self.cache_config.msg_secret_policy; - if !policy.persists() { - return None; - } - // Single chokepoint for the BotOnly invariant: only bot-context secrets - // (class == Bot) are persisted, no matter which write path got here. - if policy.bot_only() && class != wacore::msg_secret::RetentionClass::Bot { - return None; - } - let expires_at = wacore::msg_secret::expires_at( - policy, - &self.cache_config.msg_secret_retention, - class, - message_ts, - wacore::time::now_secs(), - ); - Some(wacore::store::traits::MsgSecretEntry { - chat: chat.to_non_ad_string(), - sender: sender.to_non_ad_string(), - msg_id: msg_id.to_string(), - secret: secret.to_vec(), - expires_at, - message_ts: message_ts.and_then(|t| i64::try_from(t).ok()).unwrap_or(0), - }) - } - - /// Write a batch of secret aliases in one atomic upsert, so a multi-alias - /// capture/re-persist never leaves only some aliases stored. - async fn persist_msg_secret_entries( - &self, - entries: Vec, - ) -> bool { - if entries.is_empty() { - return false; - } - match self - .persistence_manager - .backend() - .put_msg_secrets(entries) - .await - { - Ok(_) => true, - Err(e) => { - log::warn!("failed to persist messageSecrets: {e:?}"); - false - } - } - } - - async fn own_jid_for_secret_encrypted(&self, info: &MessageInfo) -> Option { - use wacore::types::message::AddressingMode; - - if info.source.is_from_me { - return Some(info.source.sender.to_non_ad()); - } - - match info.source.addressing_mode { - Some(AddressingMode::Lid) => match self.get_lid().await { - Some(jid) => Some(jid), - None => self.get_pn().await, - }, - Some(AddressingMode::Pn) => match self.get_pn().await { - Some(jid) => Some(jid), - None => self.get_lid().await, - }, - None if info.source.sender.is_lid() || info.source.chat.is_lid() => { - match self.get_lid().await { - Some(jid) => Some(jid), - None => self.get_pn().await, - } - } - None => match self.get_pn().await { - Some(jid) => Some(jid), - None => self.get_lid().await, - }, - } - } - - async fn maybe_decrypt_secret_encrypted_message( - self: &Arc, - msg: &wa::Message, - info: &Arc, - ) -> Option { - use crate::features::message_edit::{self, SecretEncKind}; - - let env = message_edit::extract_secret_encrypted(msg)?; - let target_id = env.target_id()?; - - let my_jid = self.own_jid_for_secret_encrypted(info).await?; - let original_sender = match env.original_sender_for_dispatch( - info.source.is_from_me, - &info.source.sender, - &my_jid, - ) { - Ok(jid) => jid, - Err(_) => return None, - }; - - let backend = self.persistence_manager.backend(); - let chat_for_lookup = info.source.chat.to_non_ad_string(); - let original_sender_str = original_sender.to_non_ad_string(); - let fallback_original_sender = self - .alternate_msg_secret_jid(&backend, &original_sender) - .await - .unwrap_or_default(); - - // Look up the secret AND the parent's event time (for the edit window - // check below): primary sender, then the LID/PN alternate. - let store_secret = match backend - .get_msg_secret_with_ts(&chat_for_lookup, &original_sender_str, target_id) - .await - { - Ok(Some(found)) => Some(found), - Ok(None) => match fallback_original_sender.as_ref() { - Some(alt) => { - let alt_str = alt.to_non_ad_string(); - match backend - .get_msg_secret_with_ts(&chat_for_lookup, &alt_str, target_id) - .await - { - Ok(found) => found, - Err(e) => { - log::warn!( - "[msg:{}] secret_encrypted_message alternate secret lookup failed: {e:?}", - info.id - ); - None - } - } - } - None => None, - }, - Err(e) => { - log::warn!( - "[msg:{}] backend error reading secret_encrypted_message secret: {e:?}", - info.id - ); - None - } - }; - // On a total store miss, ask the app-supplied resolver (if any) for the - // parent secret. This is what lets the Disabled policy still decrypt. The - // resolver carries no parent timestamp, so parent_ts stays 0 (unknown). - let (secret, parent_ts) = match store_secret { - Some((secret, ts)) => (secret, ts), - None => { - let alternate = fallback_original_sender - .as_ref() - .map(|j| j.to_non_ad_string()); - match self - .resolve_msg_secret_via_app( - &chat_for_lookup, - &original_sender_str, - alternate.as_deref(), - target_id, - ) - .await - { - Some(secret) => (secret, 0), - None => return None, - } - } - }; - - let fallback_editor = match info.source.sender_alt.clone() { - Some(jid) => Some(jid), - None => self - .alternate_msg_secret_jid(&backend, &info.source.sender) - .await - .unwrap_or_default(), - }; - - let inner = match message_edit::decrypt_secret_encrypted( - env.enc_payload, - env.enc_iv, - &secret, - env.kind, - target_id, - &original_sender, - &info.source.sender, - ) { - Ok(inner) => inner, - Err(primary_err) => { - let mut last_err = primary_err; - let mut decrypted = None; - - if let Some(fallback_original) = fallback_original_sender.as_ref() { - match message_edit::decrypt_secret_encrypted( - env.enc_payload, - env.enc_iv, - &secret, - env.kind, - target_id, - fallback_original, - &info.source.sender, - ) { - Ok(inner) => decrypted = Some(inner), - Err(e) => last_err = e, - } - } - - if decrypted.is_none() - && let Some(fallback_editor) = fallback_editor.as_ref() - { - match message_edit::decrypt_secret_encrypted( - env.enc_payload, - env.enc_iv, - &secret, - env.kind, - target_id, - &original_sender, - fallback_editor, - ) { - Ok(inner) => decrypted = Some(inner), - Err(e) => last_err = e, - } - } - - if decrypted.is_none() - && let (Some(fallback_original), Some(fallback_editor)) = - (fallback_original_sender.as_ref(), fallback_editor.as_ref()) - { - match message_edit::decrypt_secret_encrypted( - env.enc_payload, - env.enc_iv, - &secret, - env.kind, - target_id, - fallback_original, - fallback_editor, - ) { - Ok(inner) => decrypted = Some(inner), - Err(e) => last_err = e, - } - } - - match decrypted { - Some(inner) => inner, - None => { - log::warn!( - "[msg:{}] secret_encrypted_message {:?} decrypt failed: {last_err:?}", - info.id, - env.kind - ); - return None; - } - } - } - }; - - // Mirror WA Web `ProcessEditProtocolMsgs`: drop a MESSAGE_EDIT authored - // outside the parent's edit-processing window (editTs >= parentTs + 20m). - // The check is on authored time, not "now", so a validly-authored edit - // still applies after an offline delivery gap. Only enforceable when we - // know the parent's event time; resolver-supplied secrets carry none - // (parent_ts == 0), so we stay permissive there. - if env.kind == SecretEncKind::MessageEdit && parent_ts > 0 { - let edit_ts = info.timestamp.timestamp(); - if edit_ts >= parent_ts + wacore::msg_secret::EDIT_PROCESSING_WINDOW_SECS { - log::debug!( - "[msg:{}] secret edit authored outside the {}s window (editTs={edit_ts}, parentTs={parent_ts}); dropping", - info.id, - wacore::msg_secret::EDIT_PROCESSING_WINDOW_SECS - ); - return None; - } - } - - if let Some(secret_bytes) = inner - .message_context_info - .as_ref() - .and_then(|m| m.message_secret.as_deref()) - { - // The re-persisted secret keys the NEXT add-on on the same parent, - // so its retention class follows the parent kind and the parent's own - // event time (when known) rather than this edit's arrival time. - let class = match env.kind { - SecretEncKind::MessageEdit => wacore::msg_secret::RetentionClass::Text, - _ => wacore::msg_secret::RetentionClass::PollEvent, - }; - let message_ts = if parent_ts > 0 { - u64::try_from(parent_ts).ok() - } else { - u64::try_from(info.timestamp.timestamp()).ok() - }; - // Primary + LID/PN alternate in one batch so both survive together. - let mut entries = Vec::with_capacity(2); - if let Some(entry) = self.build_msg_secret_entry( - &info.source.chat, - &original_sender, - target_id, - secret_bytes, - class, - message_ts, - ) { - entries.push(entry); - } - if let Some(alternate_sender) = fallback_original_sender.as_ref() - && let Some(entry) = self.build_msg_secret_entry( - &info.source.chat, - alternate_sender, - target_id, - secret_bytes, - class, - message_ts, - ) - { - entries.push(entry); - } - self.persist_msg_secret_entries(entries).await; - } - - if env.kind != SecretEncKind::MessageEdit { - return Some(inner); - } - - match message_edit::rewrap_as_legacy_edit(inner) { - Some(rewrapped) => Some(rewrapped), - None => { - log::warn!( - "[msg:{}] decrypted MESSAGE_EDIT missing protocol_message.edited_message", - info.id - ); - None - } - } - } - - /// Decrypt and dispatch a `` bot reply. Looks up the - /// outbound `messageSecret` we persisted at send time and runs the - /// dual-HKDF + AES-GCM open from [`wacore::bot_message`]. Failures - /// (missing secret, GCM tag fail, malformed proto) nack with code 495. - pub(crate) async fn handle_msmsg_payload( - self: &Arc, - info: &Arc, - payload: EncPayload, - ) { - use prost::Message as _; - use wa::MessageSecretMessage; - use wacore::bot_message::{BotMessageContext, decrypt_bot_message}; - use wacore::protocol::nack::NackReason; - - let ms_msg = match MessageSecretMessage::decode(&*payload.ciphertext) { - Ok(m) => m, - Err(e) => { - log::warn!( - "[msg:{}] failed to decode MessageSecretMessage: {e:?}", - info.id - ); - self.spawn_nack(info, NackReason::ParsingError, None); - return; - } - }; - let (Some(enc_iv), Some(enc_payload)) = - (ms_msg.enc_iv.as_deref(), ms_msg.enc_payload.as_deref()) - else { - log::warn!( - "[msg:{}] MessageSecretMessage missing enc_iv/enc_payload", - info.id - ); - self.spawn_nack(info, NackReason::ParsingError, None); - return; - }; - - // Target sender (us): meta echoes our LID/PN. Falls back to our LID - // when sender is on the bot server, our PN otherwise (whatsmeow - // `decryptBotMessage`). - let target_sender = match self.resolve_msmsg_target_sender(info).await { - Some(j) => j, - None => { - log::warn!("[msg:{}] msmsg: no target_sender resolvable", info.id); - self.spawn_nack(info, NackReason::MissingMessageSecret, None); - return; - } - }; - - // Chat scope for the secret lookup: prefer ; - // fall back to the stanza's chat (matches WA Web `decryptMsmsgBotMessage`). - let chat_for_lookup = info - .meta_info - .target_chat - .as_ref() - .unwrap_or(&info.source.chat) - .to_non_ad() - .to_string(); - let target_sender_str = target_sender.to_non_ad_string(); - - // The id used for the SECRET LOOKUP is `meta.target_id` (our outbound - // id); the id used as HKDF input is the bot reply id (or - // `bot_info.edit_target_id` when the bot is editing a prior reply). - let target_id = match info.meta_info.target_id.as_deref() { - Some(id) => id, - None => { - log::warn!( - "[msg:{}] msmsg: missing target_id; cannot look up secret", - info.id - ); - self.spawn_nack(info, NackReason::MissingMessageSecret, None); - return; - } - }; - - // Mirror WA Web `C()` in `WAWebBotMessageSecret.js`: primary lookup - // plus an alternate (PN ↔ LID swap via lid_pn_mapping) so a row - // stored under one identity family is still found if `` echoes the other. Covers LID migration windows - // and asymmetric outbound/inbound identities. - let backend = self.persistence_manager.backend(); - // Store lookup: primary, then the LID/PN alternate. A backend error is - // logged and treated as a miss (not a hard nack) so the resolver still - // gets a chance — mirrors the secret-encrypted edit path. - let store_secret = match backend - .get_msg_secret(&chat_for_lookup, &target_sender_str, target_id) - .await - { - Ok(Some(s)) => Some(s), - Ok(None) => match self - .alternate_msg_secret_lookup(&backend, &chat_for_lookup, &target_sender, target_id) - .await - { - Ok(found) => found, - Err(e) => { - log::warn!("[msg:{}] msmsg: alternate lookup failed: {e:?}", info.id); - None - } - }, - Err(e) => { - log::warn!( - "[msg:{}] backend error reading message_secret: {e:?}", - info.id - ); - None - } - }; - let secret = match store_secret { - Some(s) => s, - None => { - let alternate = self - .alternate_msg_secret_jid(&backend, &target_sender) - .await - .ok() - .flatten() - .map(|j| j.to_non_ad_string()); - match self - .resolve_msg_secret_via_app( - &chat_for_lookup, - &target_sender_str, - alternate.as_deref(), - target_id, - ) - .await - { - Some(s) => s, - None => { - // For a group bot invocation initiated by our PRIMARY - // device, the messageSecret lives in the bot-addressed - // copy the primary sent directly to the bot — it is NOT - // mirrored to companions in the group skmsg. So a - // companion legitimately never holds the secret; this - // miss is expected and benign (we nack 495 and the server - // stops replaying). A miss in a 1:1 bot chat is unexpected - // and worth a warn. - log::log!( - if info.source.is_group { - log::Level::Debug - } else { - log::Level::Warn - }, - "[msg:{}] msmsg: no message_secret stored for target_id={target_id} (primary or alternate)", - info.id - ); - self.spawn_nack(info, NackReason::MissingMessageSecret, None); - return; - } - } - } - }; - - let bot_user_jid = info.source.sender.to_non_ad_string(); - // WA Web `decryptMsmsgBotMessage` dispatches on `isFbidBot()`: - // * fbid path pre-resolves to `edit_target_id` for INNER/LAST edits, - // `externalId` (info.id) otherwise. Single AES-GCM attempt. - // * regular path tries `externalId` first, falls back to - // `edit_target_id` on AES-GCM failure. - // We don't have `isFbidBot()` detection; instead, we unify the two as - // try-then-fallback with the fbid-style id as primary. That's a strict - // superset: for INNER/LAST it usually succeeds on the first try (fbid - // outcome); for any other case primary is `info.id` so we mirror the - // regular path's first attempt. The fallback is only attempted if - // `bot_info.edit_target_id` is present. - let info_id = info.id.as_str(); - let primary_msg_id = info - .bot_info - .as_ref() - .filter(|bi| { - matches!( - bi.edit_type, - Some( - crate::types::message::BotEditType::Inner - | crate::types::message::BotEditType::Last - ) - ) - }) - .and_then(|bi| bi.edit_target_id.as_deref()) - .unwrap_or(info_id); - let fallback_msg_id = if primary_msg_id == info_id { - info.bot_info - .as_ref() - .and_then(|bi| bi.edit_target_id.as_deref()) - } else { - Some(info_id) - } - .filter(|fb| *fb != primary_msg_id); - - let attempt = |msg_id: &str| { - let ctx = BotMessageContext { - msg_id, - target_sender_user_jid: &target_sender_str, - bot_user_jid: &bot_user_jid, - }; - decrypt_bot_message(&secret, enc_iv, enc_payload, &ctx) - }; - - let plaintext = match attempt(primary_msg_id) { - Ok(p) => p, - Err(primary_err) => match fallback_msg_id { - Some(fb) => match attempt(fb) { - Ok(p) => p, - Err(fallback_err) => { - log::warn!( - "[msg:{}] msmsg AES-GCM open failed both attempts (primary={primary_err:?}, fallback={fallback_err:?})", - info.id - ); - self.spawn_nack(info, NackReason::MissingMessageSecret, None); - return; - } - }, - None => { - log::warn!( - "[msg:{}] msmsg AES-GCM open failed and no fallback msg_id: {primary_err:?}", - info.id - ); - self.spawn_nack(info, NackReason::MissingMessageSecret, None); - return; - } - }, - }; - - let msg = match wa::Message::decode(plaintext.as_slice()) { - Ok(m) => m, - Err(e) => { - log::warn!( - "[msg:{}] msmsg plaintext is not a Message proto: {e:?}", - info.id - ); - self.spawn_nack(info, NackReason::ParsingError, None); - return; - } - }; - - log::info!( - "[msg:{}] Successfully decrypted msmsg bot reply from {}", - info.id, - info.source.sender - ); - self.dispatch_parsed_message(msg, info).await; - } - - /// Resolve `target_sender` for a msmsg stanza: echo from `` when - /// present, else fall back to our LID (sender on bot server) or PN. - async fn resolve_msmsg_target_sender(&self, info: &Arc) -> Option { - if let Some(ts) = info.meta_info.target_sender.as_ref() { - return Some(ts.clone()); - } - if info.source.sender.server == wacore_binary::Server::Bot { - self.get_lid().await - } else { - self.get_pn().await - } - } - - /// Second-chance lookup with the alternate identity family. Mirrors - /// `WAWebLidMigrationUtils.getAlternateMsgKey`: swap PN ↔ LID via the - /// `lid_pn_mapping` store and retry. Returns `Ok(None)` when no mapping - /// is known or the alternate row is absent — the caller treats that as - /// a terminal miss. - async fn alternate_msg_secret_jid( - &self, - backend: &Arc, - primary_sender: &Jid, - ) -> Result, crate::store::error::StoreError> { - let alternate = match primary_sender.server { - wacore_binary::Server::Lid => backend - .get_lid_mapping(&primary_sender.user) - .await? - .map(|m| Jid::new(m.phone_number, wacore_binary::Server::Pn)), - wacore_binary::Server::Pn => backend - .get_pn_mapping(&primary_sender.user) - .await? - .map(|m| Jid::new(m.lid, wacore_binary::Server::Lid)), - _ => None, - }; - Ok(alternate) - } - - async fn alternate_msg_secret_lookup( - &self, - backend: &Arc, - chat_for_lookup: &str, - primary_sender: &Jid, - target_id: &str, - ) -> Result>, crate::store::error::StoreError> { - let Some(alternate) = self - .alternate_msg_secret_jid(backend, primary_sender) - .await? - else { - return Ok(None); - }; - let alternate_str = alternate.to_non_ad_string(); - backend - .get_msg_secret(chat_for_lookup, &alternate_str, target_id) - .await - } - - /// On a total store miss, consult the app-supplied resolver for the parent - /// secret, trying the primary then the LID/PN alternate sender. Bounded by a - /// timeout because it runs inside the per-chat receive lane, so a slow app - /// callback degrades to a miss instead of stalling the chat. - async fn resolve_msg_secret_via_app( - &self, - chat: &str, - primary_sender: &str, - alternate_sender: Option<&str>, - msg_id: &str, - ) -> Option> { - let resolver = self.cache_config.original_message_resolver.as_ref()?; - let lookup = async { - if let Some(secret) = resolver - .resolve_msg_secret(chat, primary_sender, msg_id) - .await - { - return Some(secret); - } - if let Some(alt) = alternate_sender - && alt != primary_sender - && let Some(secret) = resolver.resolve_msg_secret(chat, alt, msg_id).await - { - return Some(secret); - } - None - }; - match wacore::runtime::timeout( - &*self.runtime, - self.cache_config.msg_secret_resolver_timeout, - lookup, - ) - .await - { - Ok(Some(secret)) => Some(secret.to_vec()), - Ok(None) => None, - Err(_) => { - log::warn!("[msg:{msg_id}] original_message_resolver timed out"); - None - } - } - } - - /// Handles a newsletter plaintext message. - /// Newsletters are not E2E encrypted and use the tag directly. - /// They never carry a `secret_encrypted_message`, so no messageSecret is - /// stored or retained for newsletter chats (no newsletter retention class). - async fn handle_newsletter_message( - self: &Arc<Self>, - node: &NodeRef<'_>, - info: &Arc<MessageInfo>, - ) { - let Some(plaintext_node) = node.get_optional_child_by_tag(&["plaintext"]) else { - log::warn!( - "[msg:{}] Received newsletter message without <plaintext> child: {}", - info.id, - node.tag - ); - return; - }; - - if let Some(bytes) = plaintext_node.content_bytes() { - match wa::Message::decode(bytes) { - Ok(msg) => { - log::info!( - "[msg:{}] Received newsletter plaintext message from {}", - info.id, - info.source.chat - ); - self.dispatch_parsed_message(msg, info).await; - } - Err(e) => { - log::warn!( - "[msg:{}] Failed to decode newsletter plaintext: {e}", - info.id - ); - } - } - } - } - /// Dispatch an `UndecryptableMessage` event at most once per `(chat, id)` - /// via the single-flight `get_with` semantic on `undecryptable_dispatched`. - /// The atomic arm avoids the get-then-insert race where two concurrent - /// callers would both dispatch. Mirrors WA Web's DB-level placeholder - /// uniqueness in `WAWebMessageProcessPlaceholder`. - /// - /// Returns `true` if this call dispatched the event, `false` if a - /// previous call already did. - async fn dispatch_undecryptable_event( - &self, - info: Arc<MessageInfo>, - is_unavailable: bool, - unavailable_type: crate::types::events::UnavailableType, - decrypt_fail_mode: crate::types::events::DecryptFailMode, - ) -> bool { - let dedup_key = - wacore::types::message::ChatMessageId::new(info.source.chat.clone(), info.id.clone()); - // The init future only runs for the winning caller. Others receive - // the cached `()` and leave the flag as false. - let fresh = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); - let fresh_clone = fresh.clone(); - self.undecryptable_dispatched - .get_with(dedup_key, async move { - fresh_clone.store(true, std::sync::atomic::Ordering::Release); - }) - .await; - let was_fresh = fresh.load(std::sync::atomic::Ordering::Acquire); - if was_fresh { - self.core.event_bus.dispatch(Event::UndecryptableMessage( - crate::types::events::UndecryptableMessage { - info, - is_unavailable, - unavailable_type, - decrypt_fail_mode, - }, - )); - } else { - log::debug!( - "[msg:{}] UndecryptableMessage already dispatched for this id; skipping duplicate event", - info.id, - ); - } - was_fresh - } - - /// Dispatch an undecryptable event, then send the retry receipt and the - /// transport ack in one ordered, flushed task. - /// - /// The retry asks the sender to re-encrypt; the ack clears the stanza from - /// the server's offline queue (the retry alone does not). Both run in a - /// single `outbound_flush` task so `disconnect()` flushes them together and - /// the retry always goes out before the ack: if only one makes it, it is the - /// retry, so the message is never cleared without a resend request. status is - /// also acked here (flushed) rather than relying on the detached `should_ack` - /// gate, which can be dropped mid-flush on disconnect; the server dedups the - /// resulting duplicate ack. - /// - /// Returns `true` to be assigned to `dispatched_undecryptable` flag. - async fn handle_decrypt_failure( - self: &Arc<Self>, - info: &Arc<MessageInfo>, - reason: RetryReason, - decrypt_fail_mode: crate::types::events::DecryptFailMode, - ) -> bool { - self.dispatch_undecryptable_event( - Arc::clone(info), - false, - crate::types::events::UnavailableType::Unknown, - decrypt_fail_mode, - ) - .await; - let client = Arc::clone(self); - let info = Arc::clone(info); - self.outbound_flush.spawn(&*self.runtime, async move { - // A self-fanout is our own message; retrying it to ourselves is - // futile and the server's offline queue ignores a bare transport - // ack, so it would replay forever. Clear it with the sender receipt - // instead (same stanza the success/duplicate paths now emit). Mirror - // ack_received_message: a bot-authored message in a non-bot chat - // takes the bot-invoke-response bare ack (the retry path below), not - // the sender receipt. Gate on the same eligibility as the ack path. - if info.source.is_self_fanout() - && !info.source.is_bot_authored_non_bot_chat() - && Self::should_send_delivery_receipt(&info) - { - client.send_delivery_receipt(&info).await; - return; - } - // Only ack once the resend request is actually out; otherwise leave - // the stanza queued so the server redelivers and we retry. - let resend_sent = client.run_retry_receipt(&info, reason).await; - if resend_sent { - client.send_transport_ack(&info).await; - } - }); - true - } - - async fn handle_plaintext_failure( - self: &Arc<Self>, - info: &Arc<MessageInfo>, - decrypt_fail_mode: crate::types::events::DecryptFailMode, - ) -> bool { - let dispatched = self - .dispatch_undecryptable_event( - Arc::clone(info), - false, - crate::types::events::UnavailableType::Unknown, - decrypt_fail_mode, - ) - .await; - self.spawn_nack(info, NackReason::InvalidProtobuf, None); - dispatched - } - - /// Increments the retry count for a message and returns the new count. - /// Returns `None` if max retries have been reached. - /// - /// Note: get-then-insert has a theoretical TOCTOU window since - /// `spawn_retry_receipt` detaches. In practice, retries for the same - /// message are rare and a double-send is benign (recipients deduplicate - /// by message ID). - async fn increment_retry_count(&self, cache_key: &str, reason: RetryReason) -> Option<u8> { - let cache_key = cache_key.to_owned(); - let current = self.message_retry_counts.get(&cache_key).await; - let new_count = match current { - Some(count) if count >= MAX_DECRYPT_RETRIES => return None, - Some(count) => count + 1, - None => 1, - }; - self.message_retry_counts - .insert(cache_key.clone(), new_count) - .await; - self.recent_retry_reasons.insert(cache_key, reason).await; - Some(new_count) - } - - /// Generate consistent cache key for retry logic. - pub(crate) async fn make_retry_cache_key( - &self, - chat: &Jid, - msg_id: &str, - sender: &Jid, - ) -> String { - let chat = self.resolve_encryption_jid(chat).await; - let sender = self.resolve_encryption_jid(sender).await; - // +40 covers @server suffixes, :device, separators for two JIDs - let mut key = - String::with_capacity(chat.user.len() + msg_id.len() + sender.user.len() + 40); - chat.push_to(&mut key); - key.push(':'); - key.push_str(msg_id); - key.push(':'); - sender.push_to(&mut key); - key - } - - /// Spawns a task that sends a retry receipt for a failed decryption. - /// - /// This is used when sessions are not found or invalid to request the sender to resend - /// the message with a PreKeySignalMessage to re-establish the session. - /// - /// # Retry Count Tracking - /// - /// This method tracks retry counts per message (keyed by `{chat}:{msg_id}:{sender}`) - /// and stops sending retry receipts after `MAX_DECRYPT_RETRIES` (5) attempts to prevent - /// infinite retry loops. This matches WhatsApp Web's behavior. - /// - /// # PDO Backup - /// - /// A PDO (Peer Data Operation) request is spawned only on the FIRST retry attempt. - /// This asks our primary phone to share the already-decrypted message content. - /// PDO is NOT spawned on subsequent retries to avoid duplicate requests. - /// - /// When max retries is reached, an immediate PDO request is sent as a last resort. - /// - /// # Arguments - /// * `info` - The message info for the failed message - /// * `reason` - The retry reason code (matches WhatsApp Web's RetryReason enum) - #[cfg(test)] - fn spawn_retry_receipt(self: &Arc<Self>, info: &Arc<MessageInfo>, reason: RetryReason) { - let client = Arc::clone(self); - let info = Arc::clone(info); - self.outbound_flush.spawn(&*self.runtime, async move { - client.run_retry_receipt(&info, reason).await; - }); - } - - /// Increment the retry count and send the retry receipt (or, at the cap, a - /// last-resort PDO). Awaitable so it can be ordered before the transport ack. - /// - /// Returns whether the caller should send the ack: `false` when we intended - /// to retry but the send failed (so the stanza stays queued for another try), - /// `true` when the resend went out or we deliberately gave up at the cap. - async fn run_retry_receipt( - self: &Arc<Self>, - info: &Arc<MessageInfo>, - reason: RetryReason, - ) -> bool { - let cache_key = self - .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) - .await; - - let Some(retry_count) = self.increment_retry_count(&cache_key, reason).await else { - log::info!( - "Max retries ({}) reached for message {} from {} [{:?}]. Sending immediate PDO request.", - MAX_DECRYPT_RETRIES, - info.id, - info.source.sender, - reason - ); - // Capped: give up and clear the backlog regardless of PDO outcome. - self.run_pdo_request(info).await; - return true; - }; - - if retry_count > HIGH_RETRY_COUNT_THRESHOLD { - log::warn!( - "High retry count ({}) for message {} in chat {} from {} [{:?}]", - retry_count, - info.id, - info.source.chat, - info.source.sender, - reason - ); - } - - let retry_sent = match self.send_retry_receipt(info, retry_count, reason).await { - Ok(()) => { - debug!( - "Sent retry receipt #{} for message {} in chat {} from {} [{:?}]", - retry_count, info.id, info.source.chat, info.source.sender, reason - ); - true - } - Err(e) => { - log::error!( - "Failed to send retry receipt #{} for message {} [{:?}]: {:?}", - retry_count, - info.id, - reason, - e - ); - false - } - }; - - // First retry only, to avoid duplicate PDO requests. Awaited so it runs - // before the caller's ack; the retry receipt already landed first. - if retry_count == 1 { - self.run_pdo_request(info).await; - } - retry_sent - } - - pub(crate) async fn handle_incoming_message(self: Arc<Self>, node: Arc<OwnedNodeRef>) { - // Phase 1: classify borrows the node tree, extracts owned payloads, returns quickly. - // Phase 2: process_classified_message holds no node borrows across heavy .await points, - // keeping the async state machine small. - let classified = match self.classify_incoming_message(&node).await { - Some(c) => c, - None => return, - }; - // node is no longer borrowed here -- drop it before the heavy phase - drop(node); - self.process_classified_message(classified).await; - } - - async fn classify_incoming_message( - self: &Arc<Self>, - node: &OwnedNodeRef, - ) -> Option<ClassifiedMessage> { - let nr = node.get(); - let info = match self.parse_message_info(nr).await { - Ok(info) => Arc::new(info), - Err(e) => { - let id = nr.get_attr("id").map(|v| v.as_str()); - let from = nr.get_attr("from").map(|v| v.as_str()); - log::warn!("Failed to parse message info (id={id:?}, from={from:?}): {e:?}"); - return None; - } - }; - - // Newsletters use <plaintext> instead of <enc> because they are not E2E encrypted. - if info.source.chat.is_newsletter() { - self.handle_newsletter_message(nr, &info).await; - return None; - } - - self.cache_lid_pn_from_message( - &info.source.sender, - info.source.sender_alt.as_ref(), - info.is_offline, - ) - .await; - let sender_encryption_jid = self.resolve_encryption_jid(&info.source.sender).await; - - let unavailable_node = nr.get_optional_child("unavailable"); - - let mut all_enc_nodes: Vec<&NodeRef<'_>> = Vec::with_capacity(4); - - let direct_enc_nodes = nr.get_children_by_tag("enc"); - all_enc_nodes.extend(direct_enc_nodes); - - let participants = nr.get_optional_child_by_tag(&["participants"]); - if let Some(participants_node) = participants { - let own_jid = self.get_pn().await; - let to_nodes = participants_node.get_children_by_tag("to"); - for to_node in to_nodes { - let to_jid = match to_node.attrs().optional_jid("jid") { - Some(jid) => jid, - None => continue, - }; - if own_jid.as_ref().is_some_and(|ours| *ours == to_jid) { - let enc_children = to_node.get_children_by_tag("enc"); - all_enc_nodes.extend(enc_children); - } - } - } - - if all_enc_nodes.is_empty() && unavailable_node.is_none() { - log::warn!( - "[msg:{}] Received non-newsletter message without <enc> child: {}", - info.id, - nr.tag - ); - return None; - } - - if let Some(unavailable) = unavailable_node - && all_enc_nodes.is_empty() - { - let unavailable_type = match unavailable.get_attr("type").map(|v| v.as_str()).as_deref() - { - Some("view_once") => crate::types::events::UnavailableType::ViewOnce, - _ => crate::types::events::UnavailableType::Unknown, - }; - log::info!( - "[msg:{}] Message has <unavailable> child (type: {:?}), requesting from phone via PDO", - info.id, - unavailable_type - ); - // PDO is the only recovery here (no retry receipt), so run it before - // the transport ack in one flush task: the ack must not clear the - // offline queue before the PDO request goes out. status is acked by - // the should_ack gate. Mirrors whatsmeow's request-then-ack. - self.dispatch_undecryptable_event( - Arc::clone(&info), - true, - unavailable_type, - crate::types::events::DecryptFailMode::Show, - ) - .await; - let client = Arc::clone(self); - let info2 = Arc::clone(&info); - let skip_ack = info.source.chat.is_status_broadcast(); - self.outbound_flush.spawn(&*self.runtime, async move { - // Only ack once the PDO request is out (or skipped as ancient); - // a transient send failure leaves it queued for redelivery. - let pdo_sent = client.run_pdo_request(&info2).await; - if !skip_ack && pdo_sent { - client.send_transport_ack(&info2).await; - } - }); - return None; - } - - let mut session_payloads = Vec::with_capacity(all_enc_nodes.len()); - let mut group_payloads = Vec::with_capacity(all_enc_nodes.len()); - let mut bot_payloads = Vec::with_capacity(all_enc_nodes.len()); - let mut max_sender_retry_count: u8 = 0; - let mut has_hide_fail = false; - let mut had_unknown_enc = false; - let mut had_custom_handler = false; - - for enc_node in &all_enc_nodes { - // Parse sender retry count (WA Web: e.maybeAttrInt("count") ?? 0) - // Clamp to MAX_DECRYPT_RETRIES to prevent u64→u8 truncation on unexpected values. - let sender_count = enc_node - .attrs() - .optional_u64("count") - .map(|c| c.min(MAX_DECRYPT_RETRIES as u64) as u8) - .unwrap_or(0); - max_sender_retry_count = max_sender_retry_count.max(sender_count); - - // Parse decrypt-fail attribute (WA Web: e.maybeAttrString("decrypt-fail") === "hide") - if enc_node - .get_attr("decrypt-fail") - .map(|v| v.as_str()) - .is_some_and(|s| s == "hide") - { - has_hide_fail = true; - } - - let enc_type = match enc_node.attrs().optional_string("type") { - Some(t) => t, - None => { - log::warn!("Enc node missing 'type' attribute, skipping"); - had_unknown_enc = true; - continue; - } - }; - - if let Some(handler) = self - .custom_enc_handlers - .read() - .await - .get(enc_type.as_ref()) - .cloned() - { - let handler_clone = handler; - let client_clone = self.clone(); - let info_arc = Arc::clone(&info); - // Custom enc handlers take &Node (public API); convert from NodeRef here. - let enc_node_owned = (*enc_node).to_owned(); - let enc_type_owned = enc_type.to_string(); - - self.runtime - .spawn(Box::pin(async move { - if let Err(e) = handler_clone - .handle(client_clone, &enc_node_owned, &info_arc) - .await - { - log::warn!( - "Custom handler for enc type '{}' failed: {e:?}", - enc_type_owned - ); - } - })) - .detach(); - had_custom_handler = true; - continue; - } - - // `had_unknown_enc` means "produced no usable payload": either the - // type is unrecognized or it's known but the body is empty. - // Either way the stanza needs the fallback ack or the server replays. - if EncType::from_wire(enc_type.as_ref()).is_none() { - log::warn!("Enc node has unknown type: {enc_type}"); - had_unknown_enc = true; - continue; - } - - let payload = match EncPayload::from_owned_node(node, enc_node) { - Some(p) => p, - None => { - log::warn!("Enc node {enc_type} has no content"); - had_unknown_enc = true; - continue; - } - }; - - if payload.enc_type.is_bot_secret() { - bot_payloads.push(payload); - } else if payload.enc_type.is_session() { - session_payloads.push(payload); - } else { - group_payloads.push(payload); - } - } - - // WA Web diagnostic: validate skmsg is not first in multi-enc messages. - if !session_payloads.is_empty() - && !group_payloads.is_empty() - && all_enc_nodes.first().is_some_and(|n| { - n.get_attr("type") - .map(|v| v.as_str()) - .is_some_and(|s| s == EncType::SenderKey.as_wire_str()) - }) - { - log::error!( - "[msg:{}] Protocol violation: skmsg is first in multi-enc message from {}. \ - Expected pkmsg/msg first (containing SKDM).", - info.id, - info.source.sender - ); - } - - // Unknown-only stanzas would loop in the offline queue until - // <stream:error>. Custom handlers ack on their own; status is covered - // by should_ack. Ack from `nr` so `recipient` survives. Skip when any - // bucket has usable payloads (including msmsg) so the regular dispatch - // path runs and the valid enc still decrypts. - if session_payloads.is_empty() - && group_payloads.is_empty() - && bot_payloads.is_empty() - && had_unknown_enc - && !had_custom_handler - { - log::info!( - "[msg:{}] All enc payloads unrecognized; transport-acking to drop from offline queue", - info.id - ); - if !info.source.chat.is_status_broadcast() { - self.spawn_node_transport_ack(nr).await; - } - return None; - } - - Some(ClassifiedMessage { - info, - sender_encryption_jid, - session_payloads, - group_payloads, - bot_payloads, - max_sender_retry_count, - decrypt_fail_mode: if has_hide_fail { - crate::types::events::DecryptFailMode::Hide - } else { - crate::types::events::DecryptFailMode::Show - }, - }) - } - - /// Phase 2: acquire permit, decrypt payloads, flush. No node borrows. - async fn process_classified_message(self: Arc<Self>, msg: ClassifiedMessage) { - let ClassifiedMessage { - info, - sender_encryption_jid, - session_payloads, - group_payloads, - bot_payloads, - max_sender_retry_count, - decrypt_fail_mode, - } = msg; - - if max_sender_retry_count > 0 { - let cache_key = self - .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) - .await; - let existing = self.message_retry_counts.get(&cache_key).await.unwrap_or(0); - if max_sender_retry_count > existing { - self.message_retry_counts - .insert(cache_key, max_sender_retry_count) - .await; - } - log::debug!( - "[msg:{}] Sender retry count {} pre-seeded into cache", - info.id, - max_sender_retry_count - ); - } - - // Acquire global processing permit (1 during offline sync, N after). - // Read generation + clone Arc under the same mutex so the pair is consistent. - // - // When the semaphore transitions from 1→N (offline→online), tasks waiting on - // the old 1-permit semaphore must re-acquire from the new N-permit semaphore. - // Without this re-acquire loop, those tasks would be silently dropped, which - // can lose pkmsg messages carrying SKDM (sender key distribution). If the - // SKDM is lost, ALL subsequent skmsg messages from that sender will fail - // with "No sender key state". - let _global_permit = loop { - let (generation, semaphore) = self.read_message_semaphore(); - let permit = semaphore.acquire_arc().await; - if generation - == self - .message_semaphore_generation - .load(std::sync::atomic::Ordering::SeqCst) - { - break permit; - } - // Generation changed while waiting (e.g. offline→online transition). - // Drop the stale permit and retry with the new semaphore, which has - // more permits and will grant access quickly. - log::debug!( - "Semaphore generation changed during acquire, re-acquiring from new semaphore" - ); - drop(permit); - }; - - log::debug!( - "Starting PASS 1: Processing {} session establishment messages (pkmsg/msg)", - session_payloads.len() - ); - - // Skip session processing for group/broadcast JIDs — they use sender keys, not 1:1 sessions. - let is_group_sender = sender_encryption_jid.is_group() - || sender_encryption_jid.is_broadcast_list() - || sender_encryption_jid.is_status_broadcast(); - - let session_outcome = if !is_group_sender && !session_payloads.is_empty() { - self.clone() - .process_session_enc_batch( - &session_payloads, - &info, - &sender_encryption_jid, - decrypt_fail_mode, - ) - .await - } else { - if is_group_sender && !session_payloads.is_empty() { - log::debug!( - "Skipping {} session messages from group sender {}", - session_payloads.len(), - sender_encryption_jid - ); - } - SessionBatchOutcome::default() - }; - let session_decrypted_successfully = session_outcome.decrypted; - let session_had_duplicates = session_outcome.duplicate; - let session_dispatched_undecryptable = session_outcome.undecryptable; - - log::debug!( - "Starting PASS 2: Processing {} group content messages (skmsg)", - group_payloads.len() - ); - - // Only process group content if: - // 1. There were no session messages (session already exists), OR - // 2. Session messages were successfully decrypted, OR - // 3. Session messages were duplicates (already processed, so session exists) - // Skip only if session messages FAILED to decrypt (not duplicates, not absent). - // Matches WA Web's `canDecryptNext` pattern: if pkmsg fails with a retriable error, - // the SKDM it carried is lost, so skmsg will always fail with NoSenderKey — skip it - // to avoid unnecessary retry receipts. The retry for the pkmsg will cause the sender - // to resend the entire message including SKDM. - if !group_payloads.is_empty() { - let should_process_skmsg = - should_process_skmsg_after_session(session_payloads.is_empty(), session_outcome); - - if should_process_skmsg { - match self - .clone() - .process_group_enc_batch( - &group_payloads, - &info, - &sender_encryption_jid, - decrypt_fail_mode, - ) - .await - { - Ok(()) => { - // Processed successfully or handled errors (e.g. sent retry receipt) - } - Err(e) => { - log::warn!( - "[msg:{}] Batch group decrypt from {} in {} failed: {e:?}", - info.id, - info.source.sender, - info.source.chat - ); - } - } - } else { - // Only show warning if session messages actually FAILED (not duplicates) - if !session_had_duplicates { - if info.is_expired_status() { - log::debug!( - "[msg:{}] Silently dropping expired status from {}", - info.id, - info.source.sender - ); - } else { - log::log!( - decrypt_fail_log_level(decrypt_fail_mode), - "Skipping skmsg decryption for message {} from {} because pkmsg failed to decrypt.", - info.id, - info.source.sender - ); - if !session_dispatched_undecryptable { - self.dispatch_undecryptable_event( - Arc::clone(&info), - false, - crate::types::events::UnavailableType::Unknown, - decrypt_fail_mode, - ) - .await; - } - } - - // Do NOT send a delivery receipt for undecryptable messages. - // Per whatsmeow's implementation, delivery receipts are only sent for - // successfully decrypted/handled messages. Sending a receipt here would - // tell the server we processed it, incrementing the offline counter. - // The transport <ack> is sufficient for acknowledgment. - } - // If session_had_duplicates is true, we silently skip (no warning, no event) - // because the message was already processed in a previous session - } - } else if !session_decrypted_successfully - && !session_had_duplicates - && !session_payloads.is_empty() - { - // Edge case: message with only msg/pkmsg that failed to decrypt, no skmsg - log::log!( - decrypt_fail_log_level(decrypt_fail_mode), - "Message {} from {} failed to decrypt and has no group content. Dispatching UndecryptableMessage event.", - info.id, - info.source.sender - ); - // Dispatch UndecryptableMessage event for messages that failed to decrypt - // (This should not cause double-dispatching since process_session_enc_batch - // already returned dispatched_undecryptable=false for this case) - if !session_dispatched_undecryptable { - self.dispatch_undecryptable_event( - Arc::clone(&info), - false, - crate::types::events::UnavailableType::Unknown, - decrypt_fail_mode, - ) - .await; - } - // Do NOT send delivery receipt - transport ack is sufficient - } else if session_had_duplicates - && !session_decrypted_successfully - && !session_dispatched_undecryptable - && !info.source.chat.is_status_broadcast() - { - // Duplicate (already-processed) with no group content: ack it so the - // server drops it from the offline queue (whatsmeow/WA Web treat - // old-counter like success). status is acked by the should_ack gate - // (a status SKDM pkmsg can reach here), so skip it to avoid a - // redundant receipt. - self.ack_received_message(&info); - } else if should_ack_skdm_only_session_fallback(session_outcome, bot_payloads.is_empty()) { - // SKDM-only session decrypts skip dispatch, so this stanza would - // otherwise stay queued. WA Web and whatsmeow ack every decrypted - // message; the ack shape still comes from the message source. - // Status is intentionally not filtered here, so its success receipt - // still follows the normal WA Web path. - self.ack_received_message(&info); - } - - // Bot-secret (msmsg) payloads run inline here so they're serialised - // with the session/group decrypt batches under the same global - // permit + per-chat enqueue lock acquired upstream. - for payload in bot_payloads { - self.handle_msmsg_payload(&info, payload).await; - } - - // Flush cached Signal state to DB (matches WA Web's flushBufferToDiskIfNotMemOnlyMode) - self.flush_signal_cache_logged("message", Some(&info.id)) - .await; - } - - async fn process_session_enc_batch( - self: Arc<Self>, - payloads: &[EncPayload], - info: &Arc<MessageInfo>, - sender_encryption_jid: &Jid, - decrypt_fail_mode: crate::types::events::DecryptFailMode, - ) -> SessionBatchOutcome { - use wacore::libsignal::protocol::CiphertextMessage; - if payloads.is_empty() { - return SessionBatchOutcome::default(); - } - - // Acquire a per-sender session lock to prevent race conditions when - // multiple messages from the same sender are processed concurrently. - // Use the full Signal protocol address string as the lock key so it matches - // the SignalProtocolStoreAdapter's per-session locks (prevents ratchet counter races). - let signal_address = sender_encryption_jid.to_protocol_address(); - - // `session_guard` is held across the entire batch but dropped around - // calls into `try_pn_to_lid_migration_decrypt` because that function's - // migration loop re-enters this same mutex (non-reentrant). - let session_mutex = self.session_lock_for(signal_address.as_str()).await; - let mut session_guard: Option<async_lock::MutexGuardArc<()>> = - Some(session_mutex.lock_arc().await); - - let mut adapter = self.signal_adapter().await; - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - let mut outcome = SessionBatchOutcome::default(); - // Local identity-change detection fires once per batch: the first pkmsg - // saves the new key (ReplacedExisting); the rest are NewOrUnchanged. - let mut local_identity_reacted = false; - - for payload in payloads { - let ciphertext = &payload.ciphertext[..]; - let enc_type = payload.enc_type; - let enc_type_str = enc_type.as_wire_str(); - let padding_version = payload.padding_version; - - // WA Web `MsgSendReceipt.js` nacks PARSE_ERROR; without it the - // server retransmits the malformed stanza forever. Mirrors the - // `handle_decrypt_failure` shape (dispatch event + spawn wire I/O - // so the session lock isn't held across the send). - let parsed_message = if enc_type == EncType::PreKeyMessage { - match PreKeySignalMessage::try_from(ciphertext) { - Ok(m) => CiphertextMessage::PreKeySignalMessage(m), - Err(e) => { - log::error!( - "[msg:{}] Failed to parse PreKeySignalMessage from {}: {e:?}. Sending nack.", - info.id, - info.source.sender - ); - // |= so a later dedup'd return (false) can't clobber - // a true set by a prior iteration in this batch. - outcome.had_failure = true; - outcome.undecryptable |= self - .dispatch_undecryptable_event( - Arc::clone(info), - false, - crate::types::events::UnavailableType::Unknown, - decrypt_fail_mode, - ) - .await; - self.spawn_nack(info, NackReason::ParsingError, None); - continue; - } - } - } else { - match SignalMessage::try_from(ciphertext) { - Ok(m) => CiphertextMessage::SignalMessage(m), - Err(e) => { - log::error!( - "[msg:{}] Failed to parse SignalMessage from {}: {e:?}. Sending nack.", - info.id, - info.source.sender - ); - outcome.had_failure = true; - outcome.undecryptable |= self - .dispatch_undecryptable_event( - Arc::clone(info), - false, - crate::types::events::UnavailableType::Unknown, - decrypt_fail_mode, - ) - .await; - self.spawn_nack(info, NackReason::ParsingError, None); - continue; - } - } - }; - - if enc_type == EncType::PreKeyMessage { - // FLAGGED FOR DEBUGGING: "Bad Mac" Reproducibility - #[cfg(feature = "debug-snapshots")] - { - use base64::prelude::*; - let payload = serde_json::json!({ - "id": info.id, - "sender_jid": sender_encryption_jid.to_string(), - "timestamp": info.timestamp, - "enc_type": enc_type_str, - "payload_base64": BASE64_STANDARD.encode(ciphertext), - }); - - let content_bytes = serde_json::to_vec_pretty(&payload).unwrap_or_default(); - - if let Err(e) = self - .persistence_manager - .create_snapshot(&format!("pre_pkmsg_{}", info.id), Some(&content_bytes)) - .await - { - log::warn!("Failed to create snapshot for pkmsg: {}", e); - } - } - #[cfg(not(feature = "debug-snapshots"))] - { - // No-op if disabled - } - } - - // Shadow with wire string for all downstream usage (logging, handlers) - let enc_type = enc_type_str; - - let decrypt_res = message_decrypt( - &parsed_message, - &signal_address, - &mut adapter.session_store, - &mut adapter.identity_store, - &mut adapter.pre_key_store, - &adapter.signed_pre_key_store, - &mut rng, - UsePQRatchet::No, - ) - .await; - - match decrypt_res { - Ok(decrypted) => { - // Buffer the prekey this pkmsg consumed: message_decrypt promoted - // the session into the (volatile) cache but no longer deletes the - // prekey itself. The post-loop flush deletes it only once that - // session is durable, keeping a crash from orphaning the prekey. - if let Some(prekey_id) = decrypted.consumed_prekey_id { - adapter - .pre_key_store - .buffer_consumed_prekey(prekey_id, &signal_address) - .await; - } - if decrypted.identity_change == IdentityChange::ReplacedExisting - && !local_identity_reacted - { - local_identity_reacted = true; - self.react_to_local_identity_change(sender_encryption_jid); - } - let padded_plaintext = decrypted.plaintext; - match self - .clone() - .handle_decrypted_plaintext( - enc_type, - &padded_plaintext, - padding_version, - info, - ) - .await - { - Ok(plaintext_outcome) => { - outcome.decrypted = true; - outcome.dispatched |= plaintext_outcome.dispatched; - outcome.skdm_only |= plaintext_outcome.skdm_only; - } - Err(e) => { - log::warn!( - "[msg:{}] Failed processing plaintext from {}: {e:?}", - info.id, - info.source.sender - ); - outcome.decrypted = true; - outcome.plaintext_failed = true; - outcome.had_failure = true; - outcome.undecryptable |= - self.handle_plaintext_failure(info, decrypt_fail_mode).await; - } - } - } - Err(e) => { - // Handle DuplicatedMessage: This is expected when messages are redelivered during reconnection - if let SignalProtocolError::DuplicatedMessage(chain, counter) = e { - log::debug!( - "Skipping already-processed message from {} (chain {}, counter {}). This is normal during reconnection.", - info.source.sender, - chain, - counter - ); - // Mark that we saw a duplicate so we can skip skmsg without showing error - outcome.duplicate = true; - continue; - } - // Handle UntrustedIdentity: This happens when a user re-installs WhatsApp or changes devices. - // The Signal Protocol's security policy rejects messages from new identity keys by default. - // We handle this by clearing the old identity (to trust the new one), then retrying decryption. - // IMPORTANT: We do NOT delete the session! When the PreKeySignalMessage is processed, - // libsignal's `promote_state` will archive the old session as a "previous state". - // This allows us to decrypt any in-flight messages that were encrypted with the old session. - if let SignalProtocolError::UntrustedIdentity(ref address) = e { - log::warn!( - "[msg:{}] Received message from untrusted identity: {}. This typically means the sender re-installed WhatsApp or changed their device. Clearing old identity to trust new key (keeping session for in-flight messages).", - info.id, - address - ); - - // Delete the old, untrusted identity through the signal cache. - // NOTE: We intentionally do NOT delete the session here. The session will be - // archived (not deleted) when the new PreKeySignalMessage is processed, - // allowing decryption of any in-flight messages encrypted with the old session. - self.signal_cache.delete_identity(address).await; - // Flush immediately so the backend is updated BEFORE the retry decrypt below. - // Device::is_trusted_identity reads from backend, not cache. - if let Err(e) = self.flush_signal_cache().await { - log::warn!("Failed to flush identity deletion for {}: {e:?}", address); - outcome.had_failure = true; - continue; - } - log::info!( - "Cleared old identity for {} from cache and backend", - address - ); - - // Re-attempt decryption with the new identity - log::info!( - "[msg:{}] Retrying message decryption for {} after clearing untrusted identity", - info.id, - address - ); - - let retry_decrypt_res = message_decrypt( - &parsed_message, - &signal_address, - &mut adapter.session_store, - &mut adapter.identity_store, - &mut adapter.pre_key_store, - &adapter.signed_pre_key_store, - &mut rng, - UsePQRatchet::No, - ) - .await; - - match retry_decrypt_res { - Ok(decrypted) => { - log::debug!( - "[msg:{}] Successfully decrypted message from {} after handling untrusted identity", - info.id, - address - ); - if let Some(prekey_id) = decrypted.consumed_prekey_id { - adapter - .pre_key_store - .buffer_consumed_prekey(prekey_id, &signal_address) - .await; - } - // Normally NewOrUnchanged here (the untrusted - // identity was deleted+flushed before the retry), - // but mirror the main-decode gate so a concurrent - // re-save can't drop the signal. - if decrypted.identity_change == IdentityChange::ReplacedExisting - && !local_identity_reacted - { - local_identity_reacted = true; - self.react_to_local_identity_change(sender_encryption_jid); - } - let padded_plaintext = decrypted.plaintext; - match self - .clone() - .handle_decrypted_plaintext( - enc_type, - &padded_plaintext, - padding_version, - info, - ) - .await - { - Ok(plaintext_outcome) => { - outcome.decrypted = true; - outcome.dispatched |= plaintext_outcome.dispatched; - outcome.skdm_only |= plaintext_outcome.skdm_only; - } - Err(e) => { - log::warn!( - "Failed processing plaintext after identity retry: {e:?}" - ); - outcome.decrypted = true; - outcome.plaintext_failed = true; - outcome.had_failure = true; - outcome.undecryptable |= self - .handle_plaintext_failure(info, decrypt_fail_mode) - .await; - } - } - } - Err(retry_err) => { - // Handle DuplicatedMessage in retry path: This commonly happens during reconnection - // when the same message is redelivered by the server after we already processed it. - // The first attempt triggered UntrustedIdentity, we cleared the session, but meanwhile - // another message from the same sender re-established the session and consumed the counter. - // This is benign - the message was already successfully processed. - if let SignalProtocolError::DuplicatedMessage(chain, counter) = - retry_err - { - log::debug!( - "Message from {} was already processed (chain {}, counter {}) - detected during untrusted identity retry. This is normal during reconnection.", - address, - chain, - counter - ); - outcome.duplicate = true; - } else if matches!(retry_err, SignalProtocolError::InvalidPreKeyId) - { - // Session may exist under PN address after identity change - let migration_outcome = self - .try_pn_to_lid_migration_decrypt( - sender_encryption_jid, - &signal_address, - &parsed_message, - &mut adapter, - &mut rng, - enc_type, - padding_version, - info, - &session_mutex, - &mut session_guard, - ) - .await; - if migration_outcome.decrypted - || migration_outcome.duplicate - || migration_outcome.plaintext_failed - { - outcome.decrypted |= migration_outcome.decrypted; - outcome.duplicate |= migration_outcome.duplicate; - outcome.dispatched |= migration_outcome.dispatched; - outcome.skdm_only |= migration_outcome.skdm_only; - outcome.plaintext_failed |= - migration_outcome.plaintext_failed; - outcome.had_failure |= migration_outcome.plaintext_failed; - if migration_outcome.plaintext_failed { - outcome.undecryptable |= self - .handle_plaintext_failure(info, decrypt_fail_mode) - .await; - } - } else { - log::debug!( - "[msg:{}] InvalidPreKeyId after identity change for {}. \ - Sending retry receipt with fresh keys.", - info.id, - address - ); - outcome.had_failure = true; - outcome.undecryptable = self - .handle_decrypt_failure( - info, - RetryReason::InvalidKeyId, - decrypt_fail_mode, - ) - .await; - } - } else { - log::error!( - "[msg:{}] Decryption failed even after clearing untrusted identity for {}: {:?}", - info.id, - address, - retry_err - ); - // Send retry receipt so the sender resends with a PreKeySignalMessage - // to establish a new session with the new identity - outcome.had_failure = true; - outcome.undecryptable = self - .handle_decrypt_failure( - info, - RetryReason::InvalidKey, - decrypt_fail_mode, - ) - .await; - } - } - } - - // Re-issue tctoken so the contact still has a valid token for us - let sender_jid = info.source.sender.clone(); - if !sender_jid.is_bot() && !sender_jid.is_status_broadcast() { - let client = self.clone(); - self.runtime - .spawn(Box::pin(async move { - client - .reissue_tc_token_after_identity_change(&sender_jid) - .await; - })) - .detach(); - } - - continue; - } - // Try PN→LID session migration before sending retry receipt - if let SignalProtocolError::SessionNotFound(_) = e { - let migration_outcome = self - .try_pn_to_lid_migration_decrypt( - sender_encryption_jid, - &signal_address, - &parsed_message, - &mut adapter, - &mut rng, - enc_type, - padding_version, - info, - &session_mutex, - &mut session_guard, - ) - .await; - if migration_outcome.decrypted - || migration_outcome.duplicate - || migration_outcome.plaintext_failed - { - outcome.decrypted |= migration_outcome.decrypted; - outcome.duplicate |= migration_outcome.duplicate; - outcome.dispatched |= migration_outcome.dispatched; - outcome.skdm_only |= migration_outcome.skdm_only; - outcome.plaintext_failed |= migration_outcome.plaintext_failed; - outcome.had_failure |= migration_outcome.plaintext_failed; - if migration_outcome.plaintext_failed { - outcome.undecryptable |= - self.handle_plaintext_failure(info, decrypt_fail_mode).await; - } - continue; - } - - debug!( - "[msg:{}] No session found for {} message from {}. Sending retry receipt to request session establishment.", - info.id, enc_type, info.source.sender - ); - outcome.had_failure = true; - outcome.undecryptable = self - .handle_decrypt_failure(info, RetryReason::NoSession, decrypt_fail_mode) - .await; - continue; - } else if matches!( - e, - SignalProtocolError::BadMac(_) | SignalProtocolError::InvalidMessage(_, _) - ) { - // whatsmeow migrates PN sessions before decrypt; a fresh - // LID record can otherwise shadow the sender's PN ratchet. - let migration_outcome = self - .try_pn_to_lid_migration_decrypt( - sender_encryption_jid, - &signal_address, - &parsed_message, - &mut adapter, - &mut rng, - enc_type, - padding_version, - info, - &session_mutex, - &mut session_guard, - ) - .await; - if migration_outcome.decrypted - || migration_outcome.duplicate - || migration_outcome.plaintext_failed - { - outcome.decrypted |= migration_outcome.decrypted; - outcome.duplicate |= migration_outcome.duplicate; - outcome.dispatched |= migration_outcome.dispatched; - outcome.skdm_only |= migration_outcome.skdm_only; - outcome.plaintext_failed |= migration_outcome.plaintext_failed; - outcome.had_failure |= migration_outcome.plaintext_failed; - if migration_outcome.plaintext_failed { - outcome.undecryptable |= - self.handle_plaintext_failure(info, decrypt_fail_mode).await; - } - continue; - } - - // WAWebMsgProcessingDecryptionHandler classifies both as - // SignalRetryable -> sendRetryReceipt only, with no delete. - let (reason, label) = if matches!(e, SignalProtocolError::BadMac(_)) { - (RetryReason::BadMac, "BadMac") - } else { - (RetryReason::InvalidMessage, "InvalidMessage") - }; - log::log!( - decrypt_fail_log_level(decrypt_fail_mode), - "[msg:{}] Decryption failed for {} message from {} due to {label}. \ - Sending retry receipt.", - info.id, - enc_type, - info.source.sender - ); - - outcome.had_failure = true; - outcome.undecryptable = self - .handle_decrypt_failure(info, reason, decrypt_fail_mode) - .await; - continue; - } else if matches!(e, SignalProtocolError::InvalidPreKeyId) { - // InvalidPreKeyId on a PreKeyMessage can also mean the - // session exists under a PN address (legacy migration). - // Migrating lets Signal use the existing ratchet state - // instead of looking up the consumed one-time prekey. - let migration_outcome = self - .try_pn_to_lid_migration_decrypt( - sender_encryption_jid, - &signal_address, - &parsed_message, - &mut adapter, - &mut rng, - enc_type, - padding_version, - info, - &session_mutex, - &mut session_guard, - ) - .await; - if migration_outcome.decrypted - || migration_outcome.duplicate - || migration_outcome.plaintext_failed - { - outcome.decrypted |= migration_outcome.decrypted; - outcome.duplicate |= migration_outcome.duplicate; - outcome.dispatched |= migration_outcome.dispatched; - outcome.skdm_only |= migration_outcome.skdm_only; - outcome.plaintext_failed |= migration_outcome.plaintext_failed; - outcome.had_failure |= migration_outcome.plaintext_failed; - if migration_outcome.plaintext_failed { - outcome.undecryptable |= - self.handle_plaintext_failure(info, decrypt_fail_mode).await; - } - continue; - } - - log::debug!( - "[msg:{}] Decryption failed for {} message from {} due to InvalidPreKeyId. \ - Sender is using a prekey we don't have (likely session established while offline). \ - Sending retry receipt with fresh prekeys.", - info.id, - enc_type, - info.source.sender - ); - - // Send retry receipt with fresh prekeys - outcome.had_failure = true; - outcome.undecryptable = self - .handle_decrypt_failure( - info, - RetryReason::InvalidKeyId, - decrypt_fail_mode, - ) - .await; - continue; - } else { - // Catch-all → WA Web's UnhandledError nack (500). - log::error!( - "[msg:{}] Batch session decrypt failed (type: {}) from {}: {:?}. Sending nack.", - info.id, - enc_type, - info.source.sender, - e - ); - outcome.had_failure = true; - outcome.undecryptable |= self - .dispatch_undecryptable_event( - Arc::clone(info), - false, - crate::types::events::UnavailableType::Unknown, - decrypt_fail_mode, - ) - .await; - self.spawn_nack(info, NackReason::UnhandledError, None); - continue; - } - } - } - } - outcome - } - - async fn process_group_enc_batch( - self: Arc<Self>, - payloads: &[EncPayload], - info: &Arc<MessageInfo>, - _sender_encryption_jid: &Jid, - decrypt_fail_mode: crate::types::events::DecryptFailMode, - ) -> Result<(), DecryptionError> { - if payloads.is_empty() { - return Ok(()); - } - let mut adapter = self.signal_adapter().await; - - // Always use bare sender for sender key operations. Real WA delivers - // skmsg with bare participant but pkmsg (SKDM) with device-qualified - // participant — normalizing to bare ensures consistent lookup. - // Hoisted out of the payload loop: all three are loop-invariant. - let sender_for_sk = info.source.sender.to_non_ad(); - let sender_address = sender_for_sk.to_protocol_address(); - let sender_key_name = make_sender_key_name(&info.source.chat, &sender_address); - - for payload in payloads { - let ciphertext = &payload.ciphertext[..]; - let padding_version = payload.padding_version; - - log::debug!( - "Looking up sender key for group {} with sender address {} (from sender JID: {})", - info.source.chat, - sender_address, - info.source.sender - ); - - let decrypt_result = - group_decrypt(ciphertext, &mut adapter.sender_key_store, &sender_key_name).await; - - match decrypt_result { - Ok(padded_plaintext) => { - // Sync device list if sender is unknown, but still process - // the message. Signal decryption success already proves the - // sender holds the session key — discarding would only add - // latency via an unnecessary retry round-trip. - if !self.is_from_known_device(&info.source.sender).await { - debug!( - "[msg:{}] Unknown device {}, triggering device sync", - info.id, info.source.sender - ); - self.handle_unknown_device_sync(info).await; - } - - if let Err(e) = self - .clone() - .handle_decrypted_plaintext( - "skmsg", - &padded_plaintext, - padding_version, - info, - ) - .await - { - log::warn!("Failed processing group plaintext (batch): {e:?}"); - } - } - Err(SignalProtocolError::DuplicatedMessage(iteration, counter)) => { - log::debug!( - "Skipping already-processed sender key message from {} in group {} (iteration {}, counter {}). This is normal during reconnection.", - info.source.sender, - info.source.chat, - iteration, - counter - ); - // Redelivered duplicate: ack it so the server drops it from the - // offline queue. status is already acked by the should_ack gate, - // so skip it to avoid a redundant receipt. - if !info.source.chat.is_status_broadcast() { - self.ack_received_message(info); - } - } - Err(SignalProtocolError::NoSenderKeyState(msg)) => { - if info.is_expired_status() { - log::debug!( - "[msg:{}] Skipping retry for expired status from {}", - info.id, - info.source.sender - ); - continue; - } - - let is_unknown_device = !self.is_from_known_device(&info.source.sender).await; - let retry_reason = if is_unknown_device { - RetryReason::UnknownCompanionNoPrekey - } else { - RetryReason::NoSession - }; - - debug!( - "No sender key state for group message [msg:{}] from {}: {}. Sending retry receipt.", - info.id, info.source.sender, msg - ); - - if is_unknown_device { - self.handle_unknown_device_sync(info).await; - } - - self.handle_decrypt_failure(info, retry_reason, decrypt_fail_mode) - .await; - } - Err(e) => { - if info.is_expired_status() { - log::debug!( - "[msg:{}] Ignoring decrypt error for expired status from {}: {:?}", - info.id, - info.source.sender, - e - ); - continue; - } - - log::log!( - decrypt_fail_log_level(decrypt_fail_mode), - "Group batch decrypt failed [msg:{}] for group {} sender {}: {:?}", - info.id, - sender_key_name.group_id(), - sender_key_name.sender_id(), - e - ); - // Always surface the failure to consumers; nack only non-status - // (status is acked by the should_ack gate) so the server drops - // it from the offline queue. - self.dispatch_undecryptable_event( - Arc::clone(info), - false, - crate::types::events::UnavailableType::Unknown, - decrypt_fail_mode, - ) - .await; - if !info.source.chat.is_status_broadcast() { - self.spawn_nack(info, NackReason::UnhandledError, None); - } - } - } - } - Ok(()) - } - - /// WA Web: online → `syncDeviceListJob`, offline → `OfflinePendingDeviceCache`. - async fn handle_unknown_device_sync(self: &Arc<Self>, info: &MessageInfo) { - let user_jid = info.source.sender.to_non_ad(); - - // Dedup: skip if we already have a sync pending/in-flight for this user - if !self.pending_device_sync.add(user_jid.clone()).await { - return; - } - - if info.is_offline { - log::debug!("Queueing {} for pending device sync (offline)", user_jid); - } else { - log::debug!("Triggering immediate device sync for {}", user_jid); - let client = Arc::clone(self); - self.runtime - .spawn(Box::pin(async move { - client.invalidate_device_cache(&user_jid.user).await; - if let Err(e) = client.get_user_devices(&[user_jid]).await { - log::warn!("Immediate device sync failed: {e:?}"); - } - })) - .detach(); - } - } - - async fn handle_decrypted_plaintext( - self: Arc<Self>, - enc_type: &str, - padded_plaintext: &[u8], - padding_version: u8, - info: &Arc<MessageInfo>, - ) -> Result<PlaintextHandleOutcome, anyhow::Error> { - let original_msg = wacore::messages::decode_plaintext(padded_plaintext, padding_version)?; - log::debug!( - "[msg:{}] Successfully decrypted message from {}: type={} [batch path]", - info.id, - info.source.sender, - enc_type - ); - - // Validate DSM presence against sender identity - // (WAWebHandleMsgError.DeviceSentMessageError) - if original_msg.device_sent_message.is_some() && !info.source.is_from_me { - warn!( - "[msg:{}] DeviceSentMessage present but sender {} is not self", - info.id, info.source.sender, - ); - } - - // WA Web validateBclHash: a self-synced broadcast/status carries a - // phashV2 of the broadcast recipients in deviceSentMessage.phash. - // Recompute over our <participants> view and warn on divergence. We log - // only (no drop) until the participant hash form is confirmed live. - if let Some(dsm) = &original_msg.device_sent_message - && let Some(expected) = dsm.phash.as_deref() - && !info.bcl_participants.is_empty() - && !wacore::messages::MessageUtils::validate_bcl_hash(&info.bcl_participants, expected) - { - warn!( - "[msg:{}] bcl hash mismatch on device-sent broadcast (expected={expected}); \ - keeping message (validate-only)", - info.id, - ); - } - - // Unwrap DeviceSentMessage wrapper (self-sent messages synced from - // the primary device). The actual content (reactions, text, etc.) - // is nested inside device_sent_message.message and must be - // extracted before protocol checks or dispatch. - let mut msg = wacore::messages::unwrap_device_sent(original_msg); - - // Post-decryption logic (SKDM, sync keys, etc.) - if let Some(skdm) = &msg.sender_key_distribution_message - && let Some(axolotl_bytes) = &skdm.axolotl_sender_key_distribution_message - { - self.handle_sender_key_distribution_message( - &info.source.chat, - &info.source.sender, - axolotl_bytes, - ) - .await; - } - - // app_state_sync_key_share is a self-only protocol message (app-state - // sync keys shared between our own devices). A peer could otherwise - // inject keys and forge app-state mutations, so honour it only from - // self. WA Web `WAWebKeyManagementHandleKeyShareApi` gates on - // `isMeAccountNonLid(from)`; whatsmeow on `info.IsFromMe`. - if let Some(protocol_msg) = &msg.protocol_message - && let Some(keys) = &protocol_msg.app_state_sync_key_share - { - if info.source.is_from_me { - self.handle_app_state_sync_key_share(keys).await; - } else { - warn!( - "[msg:{}] Dropping app_state_sync_key_share from non-self sender {}", - info.id, info.source.sender - ); - } - } - - // PDO responses come from our own account (is_from_me) via device 0 (primary phone) - if info.source.is_from_me - && let Some(protocol_msg) = &msg.protocol_message - && let Some(pdo_response) = &protocol_msg.peer_data_operation_request_response_message - { - self.handle_pdo_response(pdo_response, info).await; - } - - // Note: msg might be modified by take() below - let history_sync_taken = msg - .protocol_message - .as_mut() - .and_then(|pm| pm.history_sync_notification.take()); - - // history_sync_notification is self-only (our phone drives history sync). - // A spoofed one from a peer would force a download of attacker-controlled - // history, so honour it only from self. WA Web - // `WAWebHandleHistorySyncNotification` gates on `isMePrimaryNonLid`. - if let Some(history_sync) = history_sync_taken { - if info.source.is_from_me { - self.handle_history_sync(info.id.clone(), history_sync) - .await; - } else { - warn!( - "[msg:{}] Dropping history_sync_notification from non-self sender {}", - info.id, info.source.sender - ); - } - } - - // Skip dispatch for messages that only carry sender key distribution - // (protocol-level key exchange) with no user-visible content. - // These arrive as a separate pkmsg enc node alongside the actual - // group message (skmsg) and would otherwise surface as "unknown". - if wacore::messages::is_sender_key_distribution_only(&mut msg) { - log::debug!( - "[msg:{}] Skipping event dispatch for sender key distribution message", - info.id - ); - Ok(PlaintextHandleOutcome { - skdm_only: true, - ..Default::default() - }) - } else { - self.dispatch_parsed_message(msg, info).await; - Ok(PlaintextHandleOutcome { - dispatched: true, - ..Default::default() - }) - } - } - - /// Attempt PN→LID session migration and retry decryption. - /// Returns whether decryption succeeded after migration and whether it - /// reached user dispatch. - /// - /// Manages the per-address session lock around the migration loop: - /// drops the caller's guard (migration re-enters that mutex and - /// async_lock is non-reentrant), then reacquires it for the retry - /// decrypt and replaces the caller's `session_guard` on the way out - /// so the next payload in the batch stays serialized. - #[allow(clippy::too_many_arguments)] - async fn try_pn_to_lid_migration_decrypt( - self: &Arc<Self>, - sender_jid: &Jid, - signal_address: &wacore::libsignal::protocol::ProtocolAddress, - parsed_message: &wacore::libsignal::protocol::CiphertextMessage, - adapter: &mut crate::store::signal_adapter::SignalProtocolStoreAdapter, - rng: &mut rand::rngs::StdRng, - enc_type: &str, - padding_version: u8, - info: &Arc<MessageInfo>, - session_mutex: &Arc<async_lock::Mutex<()>>, - session_guard: &mut Option<async_lock::MutexGuardArc<()>>, - ) -> MigrationDecryptOutcome { - use wacore::libsignal::protocol::{UsePQRatchet, message_decrypt}; - - if !sender_jid.is_lid() { - return MigrationDecryptOutcome::default(); - } - - let Some(pn) = self.lid_pn_cache.get_phone_number(&sender_jid.user).await else { - return MigrationDecryptOutcome::default(); - }; - - // Release the address lock so the migration loop can acquire it for - // the matching device without re-entering. - *session_guard = None; - self.migrate_signal_sessions_on_lid_discovery(&pn, &sender_jid.user) - .await; - // Re-acquire for the retry decrypt and hand the guard back to the - // caller for subsequent payloads in the batch. - *session_guard = Some(session_mutex.lock_arc().await); - - match message_decrypt( - parsed_message, - signal_address, - &mut adapter.session_store, - &mut adapter.identity_store, - &mut adapter.pre_key_store, - &adapter.signed_pre_key_store, - rng, - UsePQRatchet::No, - ) - .await - { - // PN→LID migration re-addresses an existing peer; the LID address gets - // the identity for the first time (NewOrUnchanged), so no local - // identity-change reaction is warranted here. - Ok(decrypted) => { - log::info!( - "[msg:{}] Decrypted after PN→LID session migration for {}", - info.id, - info.source.sender - ); - if let Some(prekey_id) = decrypted.consumed_prekey_id { - adapter - .pre_key_store - .buffer_consumed_prekey(prekey_id, signal_address) - .await; - } - let padded_plaintext = decrypted.plaintext; - match self - .clone() - .handle_decrypted_plaintext(enc_type, &padded_plaintext, padding_version, info) - .await - { - Ok(plaintext_outcome) => MigrationDecryptOutcome { - decrypted: true, - dispatched: plaintext_outcome.dispatched, - skdm_only: plaintext_outcome.skdm_only, - ..Default::default() - }, - Err(e) => { - log::warn!( - "[msg:{}] Failed processing plaintext after migration: {e:?}", - info.id - ); - MigrationDecryptOutcome { - decrypted: true, - plaintext_failed: true, - ..Default::default() - } - } - } - } - Err(SignalProtocolError::DuplicatedMessage(chain, counter)) => { - log::debug!( - "[msg:{}] Already processed (chain {chain}, counter {counter}) after migration", - info.id - ); - MigrationDecryptOutcome { - duplicate: true, - ..Default::default() - } - } - Err(retry_err) => { - log::warn!( - "[msg:{}] Decryption still failed after PN→LID migration: {retry_err:?}", - info.id - ); - MigrationDecryptOutcome::default() - } - } - } - - async fn cache_lid_pn_from_message( - self: &Arc<Self>, - sender: &Jid, - alt: Option<&Jid>, - is_offline: bool, - ) { - let (lid_user, pn_user, source) = if sender.server.is_lid_family() { - if let Some(alt_jid) = alt - && alt_jid.server.is_pn_family() - { - ( - &sender.user, - &alt_jid.user, - crate::lid_pn_cache::LearningSource::PeerLidMessage, - ) - } else { - return; - } - } else if sender.server.is_pn_family() { - if let Some(alt_jid) = alt - && alt_jid.server.is_lid_family() - { - ( - &alt_jid.user, - &sender.user, - crate::lid_pn_cache::LearningSource::PeerPnMessage, - ) - } else { - return; - } - } else { - return; - }; - - self.learn_lid_pn_mapping_fast(lid_user, pn_user, source, is_offline) - .await; - } - - pub(crate) async fn parse_message_info( - &self, - node: &wacore_binary::NodeRef<'_>, - ) -> Result<MessageInfo, anyhow::Error> { - let (own_pn, own_lid) = { - let arc = self.persistence_manager.get_device_arc().await; - let guard = arc.read().await; - (guard.pn.clone(), guard.lid.clone()) - }; - let default_jid = Jid::default(); - let own_jid = own_pn.as_ref().unwrap_or(&default_jid); - wacore::messages::parse_message_info(node, own_jid, own_lid.as_ref()) - } - - pub(crate) async fn handle_app_state_sync_key_share( - &self, - keys: &wa::message::AppStateSyncKeyShare, - ) { - struct KeyComponents<'a> { - key_id: &'a [u8], - data: &'a [u8], - fingerprint_bytes: Vec<u8>, - timestamp: i64, - } - - /// Extract components from an AppStateSyncKey for storage. - fn extract_key_components(key: &wa::message::AppStateSyncKey) -> Option<KeyComponents<'_>> { - let key_id = key.key_id.as_ref()?.key_id.as_ref()?; - let key_data = key.key_data.as_ref()?; - let fingerprint = key_data.fingerprint.as_ref()?; - let data = key_data.key_data.as_ref()?; - Some(KeyComponents { - key_id, - data, - fingerprint_bytes: fingerprint.encode_to_vec(), - timestamp: key_data.timestamp(), - }) - } - - let device_snapshot = self.persistence_manager.get_device_snapshot().await; - let key_store = device_snapshot.backend.clone(); - - let mut stored_count = 0; - let mut failed_count = 0; - - for key in &keys.keys { - if let Some(components) = extract_key_components(key) { - let new_key = crate::store::traits::AppStateSyncKey { - key_data: components.data.to_vec(), - fingerprint: components.fingerprint_bytes, - timestamp: components.timestamp, - }; - - if let Err(e) = key_store.set_sync_key(components.key_id, new_key).await { - log::error!( - "Failed to store app state sync key {:?}: {:?}", - hex::encode(components.key_id), - e - ); - failed_count += 1; - } else { - stored_count += 1; - } - } - } - - if stored_count > 0 || failed_count > 0 { - log::info!( - target: "Client/AppState", - "Processed app state key share: {} stored, {} failed.", - stored_count, - failed_count - ); - } - - // Notify any waiters (initial full sync) that at least one key share was processed. - if stored_count > 0 - && !self - .initial_app_state_keys_received - .swap(true, std::sync::atomic::Ordering::Relaxed) - { - // First time setting; notify any waiters - self.initial_keys_synced_notifier.notify(usize::MAX); - } - } - - async fn handle_sender_key_distribution_message( - self: &Arc<Self>, - group_jid: &Jid, - sender_jid: &Jid, - axolotl_bytes: &[u8], - ) { - let skdm = match SenderKeyDistributionMessage::try_from(axolotl_bytes) { - Ok(msg) => msg, - Err(e1) => match wa::SenderKeyDistributionMessage::decode(axolotl_bytes) { - Ok(go_msg) => { - let (Some(signing_key), Some(id), Some(iteration), Some(chain_key)) = ( - go_msg.signing_key.as_ref(), - go_msg.id, - go_msg.iteration, - go_msg.chain_key.as_ref(), - ) else { - log::warn!( - "Go SKDM from {} missing required fields (signing_key={}, id={}, iteration={}, chain_key={})", - sender_jid, - go_msg.signing_key.is_some(), - go_msg.id.is_some(), - go_msg.iteration.is_some(), - go_msg.chain_key.is_some() - ); - return; - }; - let chain_key_arr: [u8; 32] = match chain_key.as_slice().try_into() { - Ok(arr) => arr, - Err(_) => { - log::error!( - "Invalid chain_key length {} from Go SKDM from {}", - chain_key.len(), - sender_jid - ); - return; - } - }; - match SignalPublicKey::from_djb_public_key_bytes(signing_key) { - Ok(pub_key) => { - match SenderKeyDistributionMessage::new( - SENDERKEY_MESSAGE_CURRENT_VERSION, - id, - iteration, - chain_key_arr, - pub_key, - ) { - Ok(skdm) => skdm, - Err(e) => { - log::error!( - "Failed to construct SKDM from Go format from {}: {:?} (original parse error: {:?})", - sender_jid, - e, - e1 - ); - return; - } - } - } - Err(e) => { - log::error!( - "Failed to parse public key from Go SKDM for {}: {:?} (original parse error: {:?})", - sender_jid, - e, - e1 - ); - return; - } - } - } - Err(e2) => { - log::error!( - "Failed to parse SenderKeyDistributionMessage (standard and Go fallback) from {}: primary: {:?}, fallback: {:?}", - sender_jid, - e1, - e2 - ); - return; - } - }, - }; - - // Normalize to bare sender for consistent sender key addressing. - let sender_bare = sender_jid.to_non_ad(); - let sender_address = sender_bare.to_protocol_address(); - - let sender_key_name = make_sender_key_name(group_jid, &sender_address); - - // Route through the signal cache adapter so the sender key is immediately visible - // in the cache for subsequent group_decrypt calls within the same message batch. - // Only the sender-key store is needed here, so build it standalone instead of - // the full five-store adapter. - let mut sender_key_store = self.sender_key_adapter().await; - - if let Err(e) = - process_sender_key_distribution_message(&sender_key_name, &skdm, &mut sender_key_store) - .await - { - log::error!( - "Failed to process SenderKeyDistributionMessage from {}: {:?}", - sender_jid, - e - ); - } else { - log::debug!( - "Successfully processed sender key distribution for group {} from {}", - group_jid, - sender_jid - ); - } - } -} +mod dispatch; +mod msg_secret; +mod receive; +mod retry; +mod special; /// Unwraps a `DeviceSentMessage` wrapper, returning the inner message with /// merged `message_context_info`. @@ -2957,9511 +164,4 @@ fn is_sender_key_distribution_only(msg: &mut wa::Message) -> bool { } #[cfg(test)] -mod tests { - use super::*; - use crate::store::SqliteStore; - use crate::store::persistence_manager::PersistenceManager; - use crate::test_utils::MockHttpClient; - use crate::types::message::EditAttribute; - use std::sync::Arc; - use wacore_binary::builder::NodeBuilder; - - fn node_to_arc(node: wacore_binary::Node) -> Arc<OwnedNodeRef> { - crate::test_utils::node_to_owned_ref(&node) - } - use wacore_binary::{Jid, SERVER_JID}; - - fn mock_transport() -> Arc<dyn crate::transport::TransportFactory> { - Arc::new(crate::transport::mock::MockTransportFactory::new()) - } - - fn mock_http_client() -> Arc<dyn crate::http::HttpClient> { - Arc::new(MockHttpClient) - } - - #[tokio::test] - async fn test_parse_message_info_for_status_broadcast() { - let backend = Arc::new( - SqliteStore::new("file:memdb_status_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let participant_jid_str = "556899336555:42@s.whatsapp.net"; - let status_broadcast_jid_str = "status@broadcast"; - - let node = NodeBuilder::new("message") - .attr("from", status_broadcast_jid_str) - .attr("id", "8A8CCCC7E6E466D9EE8CA11A967E485A") - .attr("participant", participant_jid_str) - .attr("t", "1759295366") - .attr("type", "media") - .build(); - - let info = client - .parse_message_info(&node.as_node_ref()) - .await - .expect("parse_message_info should not fail"); - - let expected_sender: Jid = participant_jid_str - .parse() - .expect("test JID should be valid"); - let expected_chat: Jid = status_broadcast_jid_str - .parse() - .expect("test JID should be valid"); - - assert_eq!( - info.source.sender, expected_sender, - "The sender should be the 'participant' JID, not 'status@broadcast'" - ); - assert_eq!( - info.source.chat, expected_chat, - "The chat should be 'status@broadcast'" - ); - assert!( - info.source.is_group, - "Broadcast messages should be treated as group-like" - ); - } - - #[tokio::test] - async fn test_status_broadcast_cold_cache_resolves_to_lid() { - use wacore::types::jid::JidExt as _; - use wacore_binary::Server; - - let backend = Arc::new( - SqliteStore::new("file:memdb_status_cold_cache?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let pn_user = "559980000001"; - let lid_user = "100000012345678"; - - assert_eq!( - client.lid_pn_cache.get_current_lid(pn_user).await, - None, - "precondition: empty cache for {pn_user}" - ); - - let node = NodeBuilder::new("message") - .attr("from", "status@broadcast") - .attr("id", "TEST_COLD_CACHE_ID") - .attr("participant", format!("{pn_user}@s.whatsapp.net").as_str()) - .attr("participant_lid", format!("{lid_user}@lid").as_str()) - .attr("t", "1777415965") - .attr("type", "media") - .build(); - - let info = client - .parse_message_info(&node.as_node_ref()) - .await - .expect("parse_message_info must succeed"); - - // Fix #1: parser surfaces participant_lid via sender_alt. - let alt = info - .source - .sender_alt - .as_ref() - .expect("sender_alt must be populated from participant_lid"); - assert_eq!(alt.user.as_str(), lid_user); - assert_eq!(alt.server, Server::Lid); - assert_eq!(info.source.sender.user.as_str(), pn_user); - assert_eq!(info.source.sender.server, Server::Pn); - - client - .cache_lid_pn_from_message( - &info.source.sender, - info.source.sender_alt.as_ref(), - info.is_offline, - ) - .await; - - // Cache learned the mapping in both directions. - assert_eq!( - client - .lid_pn_cache - .get_current_lid(pn_user) - .await - .as_deref(), - Some(lid_user), - "PN→LID lookup must hit" - ); - assert_eq!( - client.lid_pn_cache.get_phone_number(lid_user).await, - Some(pn_user.to_string()), - "LID→PN lookup must hit" - ); - - // Resolution upgrades to LID and Signal address is the LID form. - let resolved = client.resolve_encryption_jid(&info.source.sender).await; - assert_eq!(resolved.user.as_str(), lid_user); - assert_eq!(resolved.server, Server::Lid); - assert_eq!(resolved.device, info.source.sender.device); - assert_eq!( - resolved.to_protocol_address().to_string(), - format!("{lid_user}@lid.0"), - "Signal address must be @lid form, not @c.us" - ); - } - - /// Pins the hosted-family branch + the realistic non-zero device shape. - /// Production stanzas almost always have device != 0, and hosted variants - /// (`@hosted` / `@hosted.lid`) must flow through cache_lid_pn_from_message. - #[tokio::test] - async fn test_status_broadcast_hosted_family_with_device_id_resolves_to_hosted_lid() { - use wacore::types::jid::JidExt as _; - use wacore_binary::Server; - - let backend = Arc::new( - SqliteStore::new("file:memdb_status_hosted_device?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let pn_user = "559980000001"; - let lid_user = "100000012345678"; - let device_id: u16 = 99; - - let node = NodeBuilder::new("message") - .attr("from", "status@broadcast") - .attr("id", "HOSTED_TEST_ID") - .attr( - "participant", - format!("{pn_user}:{device_id}@hosted").as_str(), - ) - .attr( - "participant_lid", - format!("{lid_user}:{device_id}@hosted.lid").as_str(), - ) - .attr("t", "1777415965") - .attr("type", "media") - .build(); - - let info = client - .parse_message_info(&node.as_node_ref()) - .await - .expect("parse_message_info must succeed"); - - assert_eq!(info.source.sender.server, Server::Hosted); - assert_eq!(info.source.sender.device, device_id); - let alt = info - .source - .sender_alt - .as_ref() - .expect("sender_alt must be populated for hosted participant"); - assert_eq!(alt.server, Server::HostedLid); - assert_eq!(alt.user.as_str(), lid_user); - assert_eq!(alt.device, device_id); - - client - .cache_lid_pn_from_message( - &info.source.sender, - info.source.sender_alt.as_ref(), - info.is_offline, - ) - .await; - - // Hosted variant must reach the cache; without it, learn_lid_pn_mapping - // is skipped and the hosted-device fix is incomplete. - assert_eq!( - client - .lid_pn_cache - .get_current_lid(pn_user) - .await - .as_deref(), - Some(lid_user), - "PN→LID lookup must work for hosted family" - ); - assert_eq!( - client.lid_pn_cache.get_phone_number(lid_user).await, - Some(pn_user.to_string()), - ); - - let resolved = client.resolve_encryption_jid(&info.source.sender).await; - assert_eq!(resolved.user.as_str(), lid_user); - assert_eq!(resolved.server, Server::HostedLid); - assert_eq!( - resolved.device, device_id, - "device id must be preserved through resolution" - ); - assert_eq!( - resolved.to_protocol_address().to_string(), - format!("{lid_user}:{device_id}@hosted.lid.0"), - "Signal address must be the @hosted.lid form with device suffix" - ); - } - - #[tokio::test] - async fn test_process_session_enc_batch_handles_session_not_found_gracefully() { - use wacore::libsignal::protocol::{IdentityKeyPair, KeyPair, SignalMessage}; - - let backend = Arc::new( - SqliteStore::new("file:memdb_graceful_fail?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let sender_jid: Jid = "1234567890@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - let info = Arc::new(MessageInfo { - source: crate::types::message::MessageSource { - sender: sender_jid.clone(), - chat: sender_jid.clone(), - ..Default::default() - }, - ..Default::default() - }); - - // Create a valid but undecryptable SignalMessage - let dummy_key = [0u8; 32]; - let sender_ratchet = - KeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()).public_key; - let sender_identity_pair = - IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); - let receiver_identity_pair = - IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); - let signal_message = SignalMessage::new( - 4, - &dummy_key, - sender_ratchet, - 0, - 0, - b"test", - sender_identity_pair.identity_key(), - receiver_identity_pair.identity_key(), - ) - .expect("SignalMessage::new should succeed with valid inputs"); - - let enc_node = NodeBuilder::new("enc") - .attr("type", "msg") - .bytes(signal_message.serialized().to_vec()) - .build(); - let enc_node_ref = enc_node.as_node_ref(); - let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; - - let outcome = client - .process_session_enc_batch( - &payloads, - &info, - &sender_jid, - crate::types::events::DecryptFailMode::Show, - ) - .await; - - assert!( - !outcome.decrypted && !outcome.duplicate && outcome.undecryptable, - "process_session_enc_batch should mark SessionNotFound as undecryptable without success or duplicate" - ); - } - - /// P1: An empty session record (exists but no current/previous state) should be - /// treated the same as SessionNotFound — the retry receipt gets error code 1 (NoSession) - /// and includes keys early, instead of producing an unhelpful InvalidMessage error. - #[tokio::test] - async fn test_empty_session_record_treated_as_session_not_found() { - use wacore::libsignal::protocol::{IdentityKeyPair, KeyPair, SessionRecord, SignalMessage}; - - let backend = Arc::new( - SqliteStore::new("file:memdb_empty_session?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let sender_jid: Jid = "0000000000000@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - let info = Arc::new(MessageInfo { - source: crate::types::message::MessageSource { - sender: sender_jid.clone(), - chat: sender_jid.clone(), - ..Default::default() - }, - ..Default::default() - }); - - // Pre-store an empty (degenerate) session record in the signal cache. - // This simulates the bug scenario: record exists but has no usable ratchet state. - let signal_address = sender_jid.to_protocol_address(); - client - .signal_cache - .put_session(&signal_address, SessionRecord::new_fresh()) - .await; - - // Craft a SignalMessage to trigger decryption - let dummy_key = [0u8; 32]; - let sender_ratchet = - KeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()).public_key; - let sender_identity = - IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); - let receiver_identity = - IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); - let signal_message = SignalMessage::new( - 4, - &dummy_key, - sender_ratchet, - 0, - 0, - b"test", - sender_identity.identity_key(), - receiver_identity.identity_key(), - ) - .expect("SignalMessage::new should succeed"); - - let enc_node = NodeBuilder::new("enc") - .attr("type", "msg") - .bytes(signal_message.serialized().to_vec()) - .build(); - let enc_node_ref = enc_node.as_node_ref(); - let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; - - let outcome = client - .clone() - .process_session_enc_batch( - &payloads, - &info, - &sender_jid, - crate::types::events::DecryptFailMode::Show, - ) - .await; - - // Should behave identically to SessionNotFound: failure, no dupe, event dispatched. - assert!( - !outcome.decrypted && !outcome.duplicate && outcome.undecryptable, - "Empty session record should be treated as SessionNotFound: \ - expected undecryptable without success or duplicate, got {outcome:?}" - ); - - // After the WA Web compliance fix (no delete on BadMac/InvalidMessage either), - // every inbound-decrypt failure preserves the session. This still pins - // that the empty-record path does not regress to a delete. - let backend = client.persistence_manager.backend(); - let session_still_exists = client - .signal_cache - .has_session(&signal_address, &*backend) - .await - .expect("has_session should not fail"); - assert!(session_still_exists); - - // Discriminate from the BadMac / InvalidMessage arms (which also - // preserve the session post-fix): the empty-record path must end up - // in the SessionNotFound branch, which fires a retry receipt with - // `RetryReason::NoSession`. Anything else means the libsignal-side - // empty-record short-circuit regressed. - await_retry_receipt(&client, &info, 1, RetryReason::NoSession).await; - } - - // ─── Fixtures for session-preservation tests ───────────────────────────── - // - // Mirrors the WAWebSignalProtocolStore tests in spirit: a synthetic peer - // holds its own Signal stores in memory so the test can drive X3DH end to - // end against the Client. Inlined (not exported from a helper crate) - // because these are message.rs-specific scenarios. - - use async_trait::async_trait; - use std::collections::HashMap; - use wacore::libsignal::protocol::{ - CiphertextMessage, Direction, IdentityChange, IdentityKey, IdentityKeyPair, KeyPair, - PreKeyBundle, PreKeyRecord, PreKeyStore as SigPreKeyStore, ProtocolAddress, SenderKeyName, - SenderKeyRecord, SenderKeyStore as SigSenderKeyStore, SessionRecord, - SessionStore as SigSessionStore, SignedPreKeyStore as SigSignedPreKeyStore, UsePQRatchet, - create_sender_key_distribution_message, group_encrypt, message_encrypt, - process_prekey_bundle, - }; - use wacore::libsignal::protocol::{ - IdentityKeyStore as SigIdentityKeyStore, SignalProtocolError, - }; - - #[derive(Default, Clone)] - struct MemSessionStore(HashMap<ProtocolAddress, SessionRecord>); - - #[async_trait] - impl SigSessionStore for MemSessionStore { - async fn load_session( - &self, - a: &ProtocolAddress, - ) -> Result<Option<SessionRecord>, SignalProtocolError> { - Ok(self.0.get(a).cloned()) - } - async fn has_session(&self, a: &ProtocolAddress) -> Result<bool, SignalProtocolError> { - Ok(self.0.contains_key(a)) - } - async fn store_session( - &mut self, - a: &ProtocolAddress, - r: SessionRecord, - ) -> Result<(), SignalProtocolError> { - self.0.insert(a.clone(), r); - Ok(()) - } - } - - #[derive(Clone)] - struct MemIdentityStore { - kp: IdentityKeyPair, - reg_id: u32, - known: HashMap<ProtocolAddress, IdentityKey>, - } - - #[async_trait] - impl SigIdentityKeyStore for MemIdentityStore { - async fn get_identity_key_pair(&self) -> Result<IdentityKeyPair, SignalProtocolError> { - Ok(self.kp.clone()) - } - async fn get_local_registration_id(&self) -> Result<u32, SignalProtocolError> { - Ok(self.reg_id) - } - async fn save_identity( - &mut self, - a: &ProtocolAddress, - id: &IdentityKey, - ) -> Result<IdentityChange, SignalProtocolError> { - let prev = self.known.insert(a.clone(), *id); - Ok(match prev { - None => IdentityChange::NewOrUnchanged, - Some(p) if &p == id => IdentityChange::NewOrUnchanged, - _ => IdentityChange::ReplacedExisting, - }) - } - async fn is_trusted_identity( - &self, - _: &ProtocolAddress, - _: &IdentityKey, - _: Direction, - ) -> Result<bool, SignalProtocolError> { - Ok(true) - } - async fn get_identity( - &self, - a: &ProtocolAddress, - ) -> Result<Option<IdentityKey>, SignalProtocolError> { - Ok(self.known.get(a).copied()) - } - } - - #[derive(Default, Clone)] - struct MemSenderKeyStore(HashMap<SenderKeyName, SenderKeyRecord>); - - #[async_trait] - impl SigSenderKeyStore for MemSenderKeyStore { - async fn store_sender_key( - &mut self, - name: &SenderKeyName, - record: SenderKeyRecord, - ) -> Result<(), SignalProtocolError> { - self.0.insert(name.clone(), record); - Ok(()) - } - - async fn load_sender_key( - &self, - name: &SenderKeyName, - ) -> Result<Option<SenderKeyRecord>, SignalProtocolError> { - Ok(self.0.get(name).cloned()) - } - } - - #[derive(Clone)] - struct AlicePeer { - jid: Jid, - address: ProtocolAddress, - identity: MemIdentityStore, - sessions: MemSessionStore, - sender_keys: MemSenderKeyStore, - } - - impl AlicePeer { - async fn new(jid_str: &str) -> Self { - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - let kp = IdentityKeyPair::generate(&mut rng); - let jid: Jid = jid_str.parse().expect("valid jid"); - let address = jid.to_protocol_address(); - Self { - jid, - address, - identity: MemIdentityStore { - kp, - reg_id: 12345, - known: HashMap::new(), - }, - sessions: MemSessionStore::default(), - sender_keys: MemSenderKeyStore::default(), - } - } - - async fn install_bob_session(&mut self, bob_addr: &ProtocolAddress, bundle: &PreKeyBundle) { - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - process_prekey_bundle( - bob_addr, - &mut self.sessions, - &mut self.identity, - bundle, - &mut rng, - UsePQRatchet::No, - ) - .await - .expect("process bob bundle"); - } - - async fn encrypt( - &mut self, - bob_addr: &ProtocolAddress, - plaintext: &[u8], - ) -> CiphertextMessage { - message_encrypt(plaintext, bob_addr, &mut self.sessions, &mut self.identity) - .await - .expect("encrypt") - } - - async fn encrypt_text( - &mut self, - bob_addr: &ProtocolAddress, - text: &str, - ) -> CiphertextMessage { - use wacore::messages::MessageUtils; - - let plaintext = MessageUtils::encode_and_pad(&wa::Message { - conversation: Some(text.to_string()), - ..Default::default() - }); - self.encrypt(bob_addr, &plaintext).await - } - - async fn create_group_skdm( - &mut self, - group_jid: &Jid, - ) -> wa::message::SenderKeyDistributionMessage { - let sender = self.jid.to_non_ad(); - let sender_key_name = make_sender_key_name(group_jid, &sender.to_protocol_address()); - let skdm = create_sender_key_distribution_message( - &sender_key_name, - &mut self.sender_keys, - &mut rand::make_rng::<rand::rngs::StdRng>(), - ) - .await - .expect("create sender key distribution"); - wa::message::SenderKeyDistributionMessage { - group_id: Some(group_jid.to_string()), - axolotl_sender_key_distribution_message: Some(skdm.serialized().to_vec()), - } - } - - async fn encrypt_group_message(&mut self, group_jid: &Jid, plaintext: &[u8]) -> Vec<u8> { - let sender = self.jid.to_non_ad(); - let sender_key_name = make_sender_key_name(group_jid, &sender.to_protocol_address()); - let sender_key_message = group_encrypt( - &mut self.sender_keys, - &sender_key_name, - plaintext, - &mut rand::make_rng::<rand::rngs::StdRng>(), - ) - .await - .expect("encrypt sender key message"); - sender_key_message.serialized().to_vec() - } - } - - /// Ensure the test `Client` has an identity (`pn`/`lid`) provisioned — - /// `create_test_client_with_name` returns an unpaired client by default - /// so `device_snapshot.lid` / `.pn` are both `None`. - async fn ensure_bob_paired(client: &Arc<Client>) { - let snapshot = client.persistence_manager.get_device_snapshot().await; - if snapshot.lid.is_some() || snapshot.pn.is_some() { - return; - } - let pn: Jid = "9000000000000:1@s.whatsapp.net".parse().expect("pn"); - let lid: Jid = "999999999999999:1@lid".parse().expect("lid"); - client - .persistence_manager - .process_command(crate::store::commands::DeviceCommand::SetId(Some(pn))) - .await; - client - .persistence_manager - .process_command(crate::store::commands::DeviceCommand::SetLid(Some(lid))) - .await; - } - - /// Read Bob's currently provisioned identity / signed prekey from the test - /// client and build a `PreKeyBundle` that Alice can use to initialize - /// her side of the session. Mirrors how the real `RetryReceiptJob` ships - /// keys back to the sender — assembled through the same - /// `SignalProtocolStoreAdapter` traits production uses. - async fn bobs_prekey_bundle(client: &Arc<Client>) -> (PreKeyBundle, Jid) { - use wacore::libsignal::protocol::GenericSignedPreKey; - ensure_bob_paired(client).await; - let snapshot = client.persistence_manager.get_device_snapshot().await; - let identity_kp = snapshot.core.identity_key.clone(); - let reg_id = snapshot.core.registration_id; - - // Read/write prekeys through the same trait surface production uses - // (see signal_adapter.rs). Avoids reaching past `PersistenceManager` - // to mutate device storage directly. - let mut adapter = client.signal_adapter().await; - let spk_record = adapter - .signed_pre_key_store - .get_signed_pre_key(1.into()) - .await - .expect("spk present"); - let spk_pub = spk_record.public_key().expect("spk pub"); - let spk_sig_vec = spk_record.signature().expect("spk sig"); - - // Provision a fresh one-time prekey for this test through the - // adapter's `PreKeyStore` impl. - let pk_id_u32: u32 = 9001; - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - let pk_pair = KeyPair::generate(&mut rng); - let pk_record = PreKeyRecord::new(pk_id_u32.into(), &pk_pair); - adapter - .pre_key_store - .save_pre_key(pk_id_u32.into(), &pk_record) - .await - .expect("save pk"); - - let own_device_jid: Jid = snapshot - .lid - .clone() - .or_else(|| snapshot.pn.clone()) - .expect("own jid"); - let bob_jid = own_device_jid.to_non_ad(); - let bundle = PreKeyBundle::new( - reg_id, - u32::from(own_device_jid.device).into(), - Some((pk_id_u32.into(), pk_pair.public_key)), - 1.into(), - spk_pub, - spk_sig_vec, - IdentityKey::new(identity_kp.public_key), - ) - .expect("bundle"); - (bundle, bob_jid) - } - - /// Build an EncPayload-style stanza node and run `process_session_enc_batch`. - /// Returns whether the session for `peer_jid` still exists in the cache afterwards. - async fn submit_and_check_session( - client: &Arc<Client>, - peer_jid: &Jid, - ct: &CiphertextMessage, - ) -> (bool, bool, bool, bool) { - let (enc_type, bytes) = match ct { - CiphertextMessage::SignalMessage(m) => ("msg", m.serialized().to_vec()), - CiphertextMessage::PreKeySignalMessage(m) => ("pkmsg", m.serialized().to_vec()), - _ => panic!("unexpected ciphertext type"), - }; - let enc_node = NodeBuilder::new("enc") - .attr("type", enc_type) - .bytes(bytes) - .build(); - let enc_ref = enc_node.as_node_ref(); - let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_ref).unwrap()]; - let info = Arc::new(MessageInfo { - source: crate::types::message::MessageSource { - sender: peer_jid.clone(), - chat: peer_jid.clone(), - ..Default::default() - }, - ..Default::default() - }); - let outcome = client - .clone() - .process_session_enc_batch( - &payloads, - &info, - peer_jid, - crate::types::events::DecryptFailMode::Show, - ) - .await; - let backend = client.persistence_manager.backend(); - let still = client - .signal_cache - .has_session(&peer_jid.to_protocol_address(), &*backend) - .await - .expect("has_session"); - ( - outcome.decrypted, - outcome.duplicate, - outcome.undecryptable, - still, - ) - } - - #[tokio::test] - async fn test_badmac_migrates_pn_session_when_lid_shadow_exists() { - use crate::lid_pn_cache::{LearningSource, LidPnEntry}; - - let client = crate::test_utils::create_test_client_with_name("badmac_lid_shadow").await; - let alice_pn: Jid = "15550001001@s.whatsapp.net".parse().expect("alice pn"); - let alice_lid: Jid = "100000000000002@lid".parse().expect("alice lid"); - let entry = LidPnEntry::new( - alice_lid.user.to_string(), - alice_pn.user.to_string(), - LearningSource::PeerLidMessage, - ); - client.lid_pn_cache.add(&entry).await; - - let (bundle_v1, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let alice_pn_str = alice_pn.to_string(); - let mut alice_old = AlicePeer::new(&alice_pn_str).await; - alice_old.install_bob_session(&bob_addr, &bundle_v1).await; - let pkmsg_v1 = alice_old.encrypt_text(&bob_addr, "pn establish").await; - let (pn_success, _, _, pn_still) = - submit_and_check_session(&client, &alice_pn, &pkmsg_v1).await; - assert!(pn_success, "PN-keyed session should establish"); - assert!( - pn_still, - "PN-keyed session should be present before migration" - ); - - if let Some(record) = alice_old.sessions.0.get_mut(&bob_addr) - && let Some(state) = record.session_state_mut() - { - state.clear_unacknowledged_pre_key_message(); - } - - let mut alice_fresh = alice_old.clone(); - alice_fresh.jid = alice_lid.clone(); - alice_fresh.address = alice_lid.to_protocol_address(); - alice_fresh.sessions = MemSessionStore::default(); - - let (bundle_v2, _) = bobs_prekey_bundle(&client).await; - alice_fresh.install_bob_session(&bob_addr, &bundle_v2).await; - let pkmsg_v2 = alice_fresh.encrypt_text(&bob_addr, "lid shadow").await; - let (lid_success, _, _, lid_still) = - submit_and_check_session(&client, &alice_lid, &pkmsg_v2).await; - assert!(lid_success, "LID-keyed shadow session should establish"); - assert!(lid_still, "LID-keyed shadow session should exist"); - - let old_pn_msg = alice_old.encrypt_text(&bob_addr, "old pn ratchet").await; - assert!(matches!(old_pn_msg, CiphertextMessage::SignalMessage(_))); - let (success, duplicates, dispatched, lid_after) = - submit_and_check_session(&client, &alice_lid, &old_pn_msg).await; - assert!(success, "BadMac path should recover by migrating PN to LID"); - assert!(!duplicates, "message should decrypt, not dedupe"); - assert!( - !dispatched, - "migration recovery must not emit retry failure" - ); - assert!(lid_after, "migrated LID session should remain"); - - let backend = client.persistence_manager.backend(); - let pn_after = client - .signal_cache - .has_session(&alice_pn.to_protocol_address(), &*backend) - .await - .expect("has_session"); - assert!(!pn_after, "PN session should be consumed by migration"); - } - - #[tokio::test] - async fn migration_plaintext_failure_nacks_without_signal_retry() { - use crate::lid_pn_cache::{LearningSource, LidPnEntry}; - - let (client, transport) = capturing_client("migration_plaintext_nack").await; - let alice_pn: Jid = "15550001002@s.whatsapp.net".parse().expect("alice pn"); - let alice_lid: Jid = "100000000000004@lid".parse().expect("alice lid"); - let entry = LidPnEntry::new( - alice_lid.user.to_string(), - alice_pn.user.to_string(), - LearningSource::PeerLidMessage, - ); - client.lid_pn_cache.add(&entry).await; - - let (bundle_v1, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let alice_pn_str = alice_pn.to_string(); - let mut alice_old = AlicePeer::new(&alice_pn_str).await; - alice_old.install_bob_session(&bob_addr, &bundle_v1).await; - let pkmsg_v1 = alice_old.encrypt_text(&bob_addr, "pn establish").await; - let (pn_success, _, _, _) = submit_and_check_session(&client, &alice_pn, &pkmsg_v1).await; - assert!(pn_success, "PN-keyed session should establish"); - - if let Some(record) = alice_old.sessions.0.get_mut(&bob_addr) - && let Some(state) = record.session_state_mut() - { - state.clear_unacknowledged_pre_key_message(); - } - - let mut alice_fresh = alice_old.clone(); - alice_fresh.jid = alice_lid.clone(); - alice_fresh.address = alice_lid.to_protocol_address(); - alice_fresh.sessions = MemSessionStore::default(); - - let (bundle_v2, _) = bobs_prekey_bundle(&client).await; - alice_fresh.install_bob_session(&bob_addr, &bundle_v2).await; - let pkmsg_v2 = alice_fresh.encrypt_text(&bob_addr, "lid shadow").await; - let (lid_success, _, _, _) = submit_and_check_session(&client, &alice_lid, &pkmsg_v2).await; - assert!(lid_success, "LID-keyed shadow session should establish"); - - let bad_old_pn_msg = alice_old.encrypt(&bob_addr, &[0xff, 0x01]).await; - let payloads = vec![enc_payload_from_ciphertext(&bad_old_pn_msg)]; - let info = Arc::new(MessageInfo { - id: "MIGRATION_BAD_PLAINTEXT".to_string(), - source: crate::types::message::MessageSource { - sender: alice_lid.clone(), - chat: alice_lid.clone(), - ..Default::default() - }, - ..Default::default() - }); - - let outcome = client - .clone() - .process_session_enc_batch( - &payloads, - &info, - &alice_lid, - crate::types::events::DecryptFailMode::Show, - ) - .await; - - assert!( - outcome.decrypted, - "Signal decrypt succeeded after migration" - ); - assert!(outcome.plaintext_failed); - assert!(outcome.undecryptable); - assert!(outcome.had_failure); - assert!(!outcome.dispatched); - assert!(!outcome.skdm_only); - - let mut nack_code = None; - for _ in 0..80 { - nack_code = find_message_nack_error(&transport.sent(), &info.id); - if nack_code.is_some() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!(nack_code, Some(491)); - - let cache_key = client - .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) - .await; - assert_eq!( - client.message_retry_counts.get(&cache_key).await, - None, - "local protobuf failure after migration must not request Signal retry" - ); - } - - /// Smoking-gun regression: a `BadMac` on the inbound path must NOT delete - /// the session. Pre-fix, `src/message.rs:1100` called - /// `signal_cache.delete_session(...)` here — this test would fail with - /// `still=false`. WA Web's `RetryReceiptJob` keeps the session untouched - /// (see `docs/captured-js/WAWeb/Send/RetryReceiptJob.js`). - #[tokio::test] - async fn test_badmac_preserves_session() { - let client = crate::test_utils::create_test_client_with_name("badmac_preserves").await; - let mut alice = AlicePeer::new("1111111111111@s.whatsapp.net").await; - let alice_addr = alice.address.clone(); - - // X3DH: Alice consumes Bob's bundle to set up her outgoing session. - let (bob_bundle, _) = bobs_prekey_bundle(&client).await; - alice - .install_bob_session( - &client - .persistence_manager - .get_device_snapshot() - .await - .lid - .clone() - .or(client - .persistence_manager - .get_device_snapshot() - .await - .pn - .clone()) - .expect("own jid") - .to_protocol_address(), - &bob_bundle, - ) - .await; - - // First message: pkmsg lands on Bob and installs Bob's reciprocal session. - let bob_addr = client - .persistence_manager - .get_device_snapshot() - .await - .lid - .clone() - .or(client - .persistence_manager - .get_device_snapshot() - .await - .pn - .clone()) - .expect("own jid") - .to_protocol_address(); - let pkmsg = alice.encrypt_text(&bob_addr, "hello").await; - let (s1, _, _, still1) = submit_and_check_session(&client, &alice.jid, &pkmsg).await; - assert!(s1, "pkmsg should establish session and decrypt"); - assert!(still1, "session must exist after first message"); - - // Force Alice's next encrypt to be a plain SignalMessage rather than a - // pkmsg by clearing her unacknowledged-pkmsg flag. Tampering the trailing - // bytes of a pkmsg breaks the outer protobuf parse (because reg_id / - // signed_pre_key_id varints are encoded *after* the embedded message - // field), which would short-circuit into the parse-error nack path - // before ever reaching the BadMac arm we want to exercise. - { - let record = alice - .sessions - .0 - .get_mut(&bob_addr) - .expect("alice has a session for bob"); - if let Some(state) = record.session_state_mut() { - state.clear_unacknowledged_pre_key_message(); - } - } - - // Second message: tamper the trailing MAC byte of a real SignalMessage. - // The format is `[version][protobuf body][8-byte MAC]`, so the last byte - // is squarely inside the MAC region — parse succeeds, MAC verification - // fails -> libsignal returns BadMac. - let msg2 = alice.encrypt_text(&bob_addr, "world").await; - let mut bytes = match &msg2 { - CiphertextMessage::SignalMessage(m) => m.serialized().to_vec(), - other => panic!( - "expected SignalMessage, got {:?}", - std::mem::discriminant(other) - ), - }; - let last = bytes.len() - 1; - bytes[last] ^= 0xFF; - let enc_node = NodeBuilder::new("enc") - .attr("type", "msg") - .bytes(bytes) - .build(); - let enc_ref = enc_node.as_node_ref(); - let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_ref).unwrap()]; - let info = Arc::new(MessageInfo { - id: "BADMAC_TAMPER_MSG".to_string(), - source: crate::types::message::MessageSource { - sender: alice.jid.clone(), - chat: alice.jid.clone(), - ..Default::default() - }, - ..Default::default() - }); - - let outcome = client - .clone() - .process_session_enc_batch( - &payloads, - &info, - &alice.jid, - crate::types::events::DecryptFailMode::Show, - ) - .await; - assert!(!outcome.decrypted, "tampered MAC must not decrypt"); - assert!( - outcome.undecryptable, - "undecryptable event must be dispatched" - ); - - // The fix asserts the session lives on so the eventual sender pkmsg - // can archive it into previous_sessions[0]. - let backend = client.persistence_manager.backend(); - let still = client - .signal_cache - .has_session(&alice_addr, &*backend) - .await - .expect("has_session"); - assert!(still, "BadMac must NOT delete the session (WA Web parity)"); - - // Discriminate from the parse-error path (which also preserves the - // session): the BadMac/InvalidMessage branch routes through - // `handle_decrypt_failure` -> `spawn_retry_receipt`, which bumps - // both caches with `RetryReason::BadMac`. Parse errors take the - // nack path instead and never touch either cache. - await_retry_receipt(&client, &info, 1, RetryReason::BadMac).await; - } - - /// Poll for `message_retry_counts == expected_count` AND - /// `recent_retry_reasons == expected_reason` (or fail after a short - /// timeout). `spawn_retry_receipt` detaches the increment onto the - /// runtime, so both caches may lag the `process_session_enc_batch` return. - /// Reading both is what tells the BadMac arm apart from a parse-error - /// regression (which never bumps these caches). - async fn await_retry_receipt( - client: &Arc<Client>, - info: &MessageInfo, - expected_count: u8, - expected_reason: RetryReason, - ) { - let cache_key = client - .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) - .await; - for _ in 0..200 { - if let (Some(c), Some(r)) = ( - client.message_retry_counts.get(&cache_key).await, - client.recent_retry_reasons.get(&cache_key).await, - ) && c == expected_count - && r == expected_reason - { - return; - } - tokio::time::sleep(std::time::Duration::from_millis(5)).await; - } - let count = client.message_retry_counts.get(&cache_key).await; - let reason = client.recent_retry_reasons.get(&cache_key).await; - panic!( - "expected retry ({expected_count}, {expected_reason:?}) for {cache_key}, \ - got ({count:?}, {reason:?})" - ); - } - - // NOTE: the `InvalidMessage` arm of the `matches!()` block in - // `process_session_enc_batch` is exercised by `test_badmac_preserves_session` - // too — libsignal returns `BadMac` whenever *any* candidate state derives a - // message key (which is what a random-ratchet `SignalMessage::new(...)` - // ends up doing as well), so a separate "InvalidMessage" regression test - // would be indistinguishable from the BadMac one. Reaching the - // `InvalidMessage` constructor specifically would require crafting a - // SignalMessage that *parses* but where no state derives any message - // key — empirically impractical without major libsignal-side scaffolding. - - /// Integration test: reproduces the production loop observed in - /// `k8awqjsgww2lnkt89urp3de1-191402150615-...`. After a BadMac the bot - /// used to delete the session; when the sender then sent a fresh pkmsg - /// (post-retry-receipt), `process_prekey_bundle` ran on an empty record - /// and `previous_sessions[0]` stayed empty — any in-flight messages on - /// the OLD ratchet failed permanently. With the fix the old session - /// survives the BadMac, the pkmsg's `promote_state` archives it, and - /// the archived state lives in `previous_sessions[0]` exactly as WA Web - /// expects (see `libsignal/src/protocol/state/session.rs:751-768`). - #[tokio::test] - async fn test_prod_scenario_pkmsg_archives_old_session_after_badmac() { - let client = crate::test_utils::create_test_client_with_name("prod_archive").await; - let mut alice = AlicePeer::new("3333333333333@s.whatsapp.net").await; - - // X3DH round 1 — Alice initiates with Bob's bundle, sends pkmsg. - let (bundle_v1, _) = bobs_prekey_bundle(&client).await; - let bob_addr = client - .persistence_manager - .get_device_snapshot() - .await - .lid - .clone() - .or(client - .persistence_manager - .get_device_snapshot() - .await - .pn - .clone()) - .expect("own jid") - .to_protocol_address(); - alice.install_bob_session(&bob_addr, &bundle_v1).await; - let pkmsg_v1 = alice.encrypt_text(&bob_addr, "v1").await; - let (s1, _, _, _) = submit_and_check_session(&client, &alice.jid, &pkmsg_v1).await; - assert!(s1); - - // Snapshot Bob's session_v1 base key for later comparison. Use - // peek (non-destructive): `get_session` marks the cache entry as - // CheckedOut, which would prevent libsignal from re-loading the - // session in the BadMac path that follows. - let alice_addr = alice.address.clone(); - let backend = client.persistence_manager.backend(); - let v1_record = client - .signal_cache - .peek_session(&alice_addr, &*backend) - .await - .expect("peek_session") - .expect("v1 session present"); - let v1_base_key = v1_record - .session_state() - .expect("v1 current state") - .sender_ratchet_key_for_logging() - .expect("v1 base key"); - - // Force Alice's next encrypt to be a plain SignalMessage so tampering - // the last byte lands inside the MAC region (see comment in - // `test_badmac_preserves_session` for why pkmsg cannot be tampered - // at the tail without breaking the outer protobuf parse). - { - let record = alice - .sessions - .0 - .get_mut(&bob_addr) - .expect("alice has a session for bob"); - if let Some(state) = record.session_state_mut() { - state.clear_unacknowledged_pre_key_message(); - } - } - - // Tampered SignalMessage → BadMac branch (with the fix this no longer - // deletes Bob's session). - let msg = alice.encrypt_text(&bob_addr, "stale").await; - let mut bytes = match &msg { - CiphertextMessage::SignalMessage(m) => m.serialized().to_vec(), - other => panic!( - "expected SignalMessage, got {:?}", - std::mem::discriminant(other) - ), - }; - let last = bytes.len() - 1; - bytes[last] ^= 0xFF; - let enc_node = NodeBuilder::new("enc") - .attr("type", "msg") - .bytes(bytes) - .build(); - let enc_ref = enc_node.as_node_ref(); - let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_ref).unwrap()]; - let info = Arc::new(MessageInfo { - id: "PROD_LOOP_REPRO_STALE".to_string(), - source: crate::types::message::MessageSource { - sender: alice.jid.clone(), - chat: alice.jid.clone(), - ..Default::default() - }, - ..Default::default() - }); - let _outcome = client - .clone() - .process_session_enc_batch( - &payloads, - &info, - &alice.jid, - crate::types::events::DecryptFailMode::Show, - ) - .await; - // Confirm the BadMac branch executed (parse-error path would skip - // both retry caches; another arm would record a different reason). - await_retry_receipt(&client, &info, 1, RetryReason::BadMac).await; - // Pre-fix: this assertion would have failed (session deleted). - let preserved = client - .signal_cache - .has_session(&alice_addr, &*backend) - .await - .expect("has_session"); - assert!(preserved, "BadMac must preserve session"); - - // X3DH round 2 — Alice rebuilds her side from a fresh Bob bundle - // (simulates the bot re-issuing prekeys via a retry receipt) and - // sends another pkmsg. Bob's `process_prekey_bundle` must archive - // session_v1 into previous_sessions[0]. - let (bundle_v2, _) = bobs_prekey_bundle(&client).await; - alice.sessions = MemSessionStore::default(); // forget Alice's v1 to force a fresh X3DH - alice.install_bob_session(&bob_addr, &bundle_v2).await; - let pkmsg_v2 = alice.encrypt_text(&bob_addr, "v2").await; - let (s2, _, _, still2) = submit_and_check_session(&client, &alice.jid, &pkmsg_v2).await; - assert!(s2, "pkmsg_v2 should decrypt"); - assert!(still2); - - let v2_record = client - .signal_cache - .peek_session(&alice_addr, &*backend) - .await - .expect("peek_session") - .expect("v2 session present"); - let v2_base_key = v2_record - .session_state() - .expect("v2 current state") - .sender_ratchet_key_for_logging() - .expect("v2 base key"); - assert_ne!( - v1_base_key, v2_base_key, - "current session must be the new v2" - ); - assert_eq!( - v2_record.previous_session_count(), - 1, - "session_v1 must be archived as previous_sessions[0]" - ); - let archived_state = v2_record - .previous_session_states() - .next() - .expect("archived state") - .expect("archived state decodes"); - let archived_base_key = archived_state - .sender_ratchet_key_for_logging() - .expect("archived base key"); - assert_eq!( - archived_base_key, v1_base_key, - "archived previous_sessions[0] must be the original v1" - ); - } - - #[tokio::test] - async fn test_handle_incoming_message_skips_skmsg_after_msg_failure() { - use wacore::libsignal::protocol::{IdentityKeyPair, KeyPair, SignalMessage}; - - let backend = Arc::new( - SqliteStore::new("file:memdb_skip_skmsg_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let sender_jid: Jid = "1234567890@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - let group_jid: Jid = "120363021033254949@g.us" - .parse() - .expect("test JID should be valid"); - - // Create msg + skmsg node; msg will fail (no session), so skmsg should be skipped - let dummy_key = [0u8; 32]; - let sender_ratchet = - KeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()).public_key; - let sender_identity_pair = - IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); - let receiver_identity_pair = - IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); - let signal_message = SignalMessage::new( - 4, - &dummy_key, - sender_ratchet, - 0, - 0, - b"test", - sender_identity_pair.identity_key(), - receiver_identity_pair.identity_key(), - ) - .expect("SignalMessage::new should succeed with valid inputs"); - - let msg_node = NodeBuilder::new("enc") - .attr("type", "msg") - .bytes(signal_message.serialized().to_vec()) - .build(); - - let skmsg_node = NodeBuilder::new("enc") - .attr("type", "skmsg") - .bytes(vec![4, 5, 6]) - .build(); - - let message_node = node_to_arc( - NodeBuilder::new("message") - .attr("from", group_jid) - .attr("participant", sender_jid) - .attr("id", "test-id-123") - .attr("t", "12345") - .children(vec![msg_node, skmsg_node]) - .build(), - ); - - // Should not panic or retry loop - skmsg is skipped after msg failure - client.clone().handle_incoming_message(message_node).await; - } - - /// Test case for reproducing sender key JID mismatch in LID group messages - /// - /// Problem: - /// - When we process sender key distribution from a self-sent LID message, we store it under the LID JID - /// - But when we try to decrypt the group content (skmsg), we look it up using the phone number JID - /// - This causes "No sender key state" errors even though we just processed the sender key! - /// - /// This test verifies the fix by: - /// 1. Creating a sender key and storing it under the LID address (mimicking SKDM processing) - /// 2. Attempting retrieval with phone number address (the bug) - should fail - /// 3. Attempting retrieval with LID address (the fix) - should succeed - #[tokio::test] - async fn test_self_sent_lid_group_message_sender_key_mismatch() { - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore::libsignal::protocol::{ - SenderKeyStore, create_sender_key_distribution_message, - process_sender_key_distribution_message, - }; - - let backend = Arc::new( - SqliteStore::new("file:memdb_sender_key_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (_client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm.clone(), - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let own_lid: Jid = "100000000000001.1:75@lid" - .parse() - .expect("test JID should be valid"); - let own_phone: Jid = "15551234567:75@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - let group_jid: Jid = "120363021033254949@g.us" - .parse() - .expect("test JID should be valid"); - - // Create SKDM using LID address (mimics handle_sender_key_distribution_message) - let lid_protocol_address = own_lid.to_protocol_address(); - let lid_sender_key_name = make_sender_key_name(&group_jid, &lid_protocol_address); - - // Pin serialized form so from_jid stays compatible with persisted records - assert_eq!(lid_sender_key_name.group_id(), group_jid.to_string()); - assert_eq!( - lid_sender_key_name.sender_id(), - lid_protocol_address.to_string() - ); - - let device_arc = pm.get_device_arc().await; - let skdm = { - let mut device_guard = device_arc.write().await; - create_sender_key_distribution_message( - &lid_sender_key_name, - &mut *device_guard, - &mut rand::make_rng::<rand::rngs::StdRng>(), - ) - .await - .expect("Failed to create SKDM") - }; - - { - let mut device_guard = device_arc.write().await; - process_sender_key_distribution_message( - &lid_sender_key_name, - &skdm, - &mut *device_guard, - ) - .await - .expect("Failed to process SKDM with LID address"); - } - - // Try to retrieve using PHONE NUMBER address (THE BUG) - let phone_protocol_address = own_phone.to_protocol_address(); - let phone_sender_key_name = make_sender_key_name(&group_jid, &phone_protocol_address); - - let phone_lookup_result = { - let device_guard = device_arc.read().await; - device_guard.load_sender_key(&phone_sender_key_name).await - }; - - assert!( - phone_lookup_result - .expect("lookup should not error") - .is_none(), - "Sender key should NOT be found when looking up with phone number address (demonstrates the bug)" - ); - - // Try to retrieve using LID address (THE FIX) - let lid_lookup_result = { - let device_guard = device_arc.read().await; - device_guard.load_sender_key(&lid_sender_key_name).await - }; - - assert!( - lid_lookup_result - .expect("lookup should not error") - .is_some(), - "Sender key SHOULD be found when looking up with LID address (same as storage)" - ); - } - - /// Test that sender key consistency is maintained for multiple LID participants - /// - /// Edge case: Group with multiple LID participants, each should have their own - /// sender key stored under their LID address, not mixed up with phone numbers. - #[tokio::test] - async fn test_multiple_lid_participants_sender_key_isolation() { - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore::libsignal::protocol::{ - SenderKeyStore, create_sender_key_distribution_message, - process_sender_key_distribution_message, - }; - - let backend = Arc::new( - SqliteStore::new("file:memdb_multi_lid_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let transport_factory = Arc::new(crate::transport::mock::MockTransportFactory::new()); - let (_client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm.clone(), - transport_factory, - mock_http_client(), - None, - ) - .await; - - let group_jid: Jid = "120363021033254949@g.us" - .parse() - .expect("test JID should be valid"); - - // Simulate three LID participants - let participants = vec![ - ("100000000000001.1:75@lid", "15551234567:75@s.whatsapp.net"), - ("987654321000000.2:42@lid", "551234567890:42@s.whatsapp.net"), - ("111222333444555.3:10@lid", "559876543210:10@s.whatsapp.net"), - ]; - - let device_arc = pm.get_device_arc().await; - - // Create and store sender keys for each participant under their LID address - for (lid_str, _phone_str) in &participants { - let lid_jid: Jid = lid_str.parse().expect("test JID should be valid"); - let lid_protocol_address = lid_jid.to_protocol_address(); - let lid_sender_key_name = make_sender_key_name(&group_jid, &lid_protocol_address); - - let skdm = { - let mut device_guard = device_arc.write().await; - create_sender_key_distribution_message( - &lid_sender_key_name, - &mut *device_guard, - &mut rand::make_rng::<rand::rngs::StdRng>(), - ) - .await - .expect("Failed to create SKDM") - }; - - let mut device_guard = device_arc.write().await; - process_sender_key_distribution_message( - &lid_sender_key_name, - &skdm, - &mut *device_guard, - ) - .await - .expect("Failed to process SKDM"); - } - - // Verify each participant's sender key can be retrieved using their LID address - for (lid_str, phone_str) in &participants { - let lid_jid: Jid = lid_str.parse().expect("test JID should be valid"); - let phone_jid: Jid = phone_str.parse().expect("test JID should be valid"); - - let lid_protocol_address = lid_jid.to_protocol_address(); - let phone_protocol_address = phone_jid.to_protocol_address(); - - let lid_sender_key_name = make_sender_key_name(&group_jid, &lid_protocol_address); - let phone_sender_key_name = make_sender_key_name(&group_jid, &phone_protocol_address); - - // Should find with LID address - let lid_lookup = { - let device_guard = device_arc.read().await; - device_guard.load_sender_key(&lid_sender_key_name).await - }; - assert!( - lid_lookup.expect("lookup should not error").is_some(), - "Sender key for {} should be found with LID address", - lid_str - ); - - // Should NOT find with phone number address (the bug) - let phone_lookup = { - let device_guard = device_arc.read().await; - device_guard.load_sender_key(&phone_sender_key_name).await - }; - assert!( - phone_lookup.expect("lookup should not error").is_none(), - "Sender key for {} should NOT be found with phone number address", - lid_str - ); - } - } - - /// Test that LID JID parsing handles various edge cases correctly - /// - /// Edge cases: - /// - LID with multiple dots in user portion - /// - LID with device numbers - /// - LID without device numbers - #[test] - fn test_lid_jid_parsing_edge_cases() { - use wacore_binary::Jid; - - // Single dot in user portion - let lid1: Jid = "100000000000001.1:75@lid" - .parse() - .expect("test JID should be valid"); - assert_eq!(lid1.user, "100000000000001.1"); - assert_eq!(lid1.device, 75); - assert_eq!(lid1.agent, 0); - - // Multiple dots in user portion (extreme edge case) - let lid2: Jid = "123.456.789.0:50@lid" - .parse() - .expect("test JID should be valid"); - assert_eq!(lid2.user, "123.456.789.0"); - assert_eq!(lid2.device, 50); - assert_eq!(lid2.agent, 0); - - // No device number (device 0) - let lid3: Jid = "987654321000000.5@lid" - .parse() - .expect("test JID should be valid"); - assert_eq!(lid3.user, "987654321000000.5"); - assert_eq!(lid3.device, 0); - assert_eq!(lid3.agent, 0); - - // Very long user portion with dot - let lid4: Jid = "111222333444555666777.999:1@lid" - .parse() - .expect("test JID should be valid"); - assert_eq!(lid4.user, "111222333444555666777.999"); - assert_eq!(lid4.device, 1); - assert_eq!(lid4.agent, 0); - } - - /// Test that protocol address generation from LID JIDs matches WhatsApp Web format - /// - /// WhatsApp Web uses: {user}[:device]@{server}.0 - /// - The device is encoded in the name - /// - device_id is always 0 - #[test] - fn test_lid_protocol_address_consistency() { - use wacore::types::jid::JidExt as CoreJidExt; - use wacore_binary::Jid; - - // Format: (jid_str, expected_name, expected_device_id, expected_to_string) - let test_cases = vec![ - ( - "100000000000001.1:75@lid", - "100000000000001.1:75@lid", - 0, - "100000000000001.1:75@lid.0", - ), - ( - "987654321000000.2:42@lid", - "987654321000000.2:42@lid", - 0, - "987654321000000.2:42@lid.0", - ), - ( - "111.222.333:10@lid", - "111.222.333:10@lid", - 0, - "111.222.333:10@lid.0", - ), - // No device - should not include :0 - ("123456789@lid", "123456789@lid", 0, "123456789@lid.0"), - ]; - - for (jid_str, expected_name, expected_device_id, expected_to_string) in test_cases { - let lid_jid: Jid = jid_str.parse().expect("test JID should be valid"); - let protocol_addr = lid_jid.to_protocol_address(); - - assert_eq!( - protocol_addr.name(), - expected_name, - "Protocol address name should match WhatsApp Web's SignalAddress format for {}", - jid_str - ); - assert_eq!( - u32::from(protocol_addr.device_id()), - expected_device_id, - "Protocol address device_id should always be 0 for {}", - jid_str - ); - assert_eq!( - protocol_addr.to_string(), - expected_to_string, - "Protocol address to_string() should match createSignalLikeAddress format for {}", - jid_str - ); - } - } - - /// Test sender_alt extraction from message attributes in LID groups - /// - /// Edge cases: - /// - LID group with participant_pn attribute - /// - PN group with participant_lid attribute - /// - Mixed addressing modes - #[tokio::test] - async fn test_parse_message_info_sender_alt_extraction() { - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore::types::message::AddressingMode; - use wacore_binary::builder::NodeBuilder; - - let backend = Arc::new( - SqliteStore::new("file:memdb_sender_alt_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - - // Set up own phone number and LID - { - let device_arc = pm.get_device_arc().await; - let mut device = device_arc.write().await; - device.pn = Some( - "15551234567@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), - ); - device.lid = Some( - "100000000000001.1@lid" - .parse() - .expect("test JID should be valid"), - ); - } - - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - // Test case 1: LID group message with participant_pn - let lid_group_node = NodeBuilder::new("message") - .attr("from", "120363021033254949@g.us") - .attr("participant", "987654321000000.2:42@lid") - .attr("participant_pn", "551234567890:42@s.whatsapp.net") - .attr("addressing_mode", AddressingMode::Lid.as_str()) - .attr("id", "test1") - .attr("t", "12345") - .build(); - - let info1 = client - .parse_message_info(&lid_group_node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - assert_eq!(info1.source.sender.user, "987654321000000.2"); - assert!(info1.source.sender_alt.is_some()); - assert_eq!( - info1 - .source - .sender_alt - .as_ref() - .expect("sender_alt should be present") - .user, - "551234567890" - ); - - // Test case 2: Self-sent LID group message - let self_lid_node = NodeBuilder::new("message") - .attr("from", "120363021033254949@g.us") - .attr("participant", "100000000000001.1:75@lid") - .attr("participant_pn", "15551234567:75@s.whatsapp.net") - .attr("addressing_mode", AddressingMode::Lid.as_str()) - .attr("id", "test2") - .attr("t", "12346") - .build(); - - let info2 = client - .parse_message_info(&self_lid_node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - assert!( - info2.source.is_from_me, - "Should detect self-sent LID message" - ); - assert_eq!(info2.source.sender.user, "100000000000001.1"); - assert!(info2.source.sender_alt.is_some()); - assert_eq!( - info2 - .source - .sender_alt - .as_ref() - .expect("sender_alt should be present") - .user, - "15551234567" - ); - } - - /// Test that device query logic uses phone numbers for LID participants - /// - /// This is a unit test for the logic in wacore/src/send.rs that converts - /// LID JIDs to phone number JIDs for device queries. - #[test] - fn test_lid_to_phone_mapping_for_device_queries() { - use std::collections::HashMap; - use wacore::client::context::GroupInfo; - use wacore::types::message::AddressingMode; - use wacore_binary::Jid; - - // Simulate a LID group with phone number mappings - let mut lid_to_pn_map = HashMap::new(); - lid_to_pn_map.insert( - wacore_binary::CompactString::from("100000000000001.1"), - "15551234567@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), - ); - lid_to_pn_map.insert( - wacore_binary::CompactString::from("987654321000000.2"), - "551234567890@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), - ); - - let mut group_info = GroupInfo::new( - vec![ - "100000000000001.1:75@lid" - .parse() - .expect("test JID should be valid"), - "987654321000000.2:42@lid" - .parse() - .expect("test JID should be valid"), - ], - AddressingMode::Lid, - ); - group_info.set_lid_to_pn_map(lid_to_pn_map.clone()); - - // Simulate the device query logic - let jids_to_query: Vec<Jid> = group_info - .participants - .iter() - .map(|jid| { - let base_jid = jid.to_non_ad(); - if base_jid.is_lid() - && let Some(phone_jid) = group_info.phone_jid_for_lid_user(&base_jid.user) - { - return phone_jid.to_non_ad(); - } - base_jid - }) - .collect(); - - // Verify all queries use phone numbers, not LID JIDs - for jid in &jids_to_query { - assert_eq!( - jid.server, SERVER_JID, - "Device query should use phone number, got: {}", - jid - ); - } - - assert_eq!(jids_to_query.len(), 2); - assert!(jids_to_query.iter().any(|j| j.user == "15551234567")); - assert!(jids_to_query.iter().any(|j| j.user == "551234567890")); - } - - /// Test edge case: Group with mixed LID and phone number participants - /// - /// Some participants may still use phone numbers even in a LID group. - /// The code should handle both correctly. - #[test] - fn test_mixed_lid_and_phone_participants() { - use std::collections::HashMap; - use wacore::client::context::GroupInfo; - use wacore::types::message::AddressingMode; - use wacore_binary::Jid; - - let mut lid_to_pn_map = HashMap::new(); - lid_to_pn_map.insert( - wacore_binary::CompactString::from("100000000000001.1"), - "15551234567@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), - ); - - let mut group_info = GroupInfo::new( - vec![ - "100000000000001.1:75@lid" - .parse() - .expect("test JID should be valid"), // LID participant - "551234567890:42@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), // Phone number participant - ], - AddressingMode::Lid, - ); - group_info.set_lid_to_pn_map(lid_to_pn_map.clone()); - - let jids_to_query: Vec<Jid> = group_info - .participants - .iter() - .map(|jid| { - let base_jid = jid.to_non_ad(); - if base_jid.is_lid() - && let Some(phone_jid) = group_info.phone_jid_for_lid_user(&base_jid.user) - { - return phone_jid.to_non_ad(); - } - base_jid - }) - .collect(); - - // Both should end up as phone numbers - assert_eq!(jids_to_query.len(), 2); - for jid in &jids_to_query { - assert_eq!(jid.server, SERVER_JID); - } - } - - /// Test edge case: Own JID check in LID mode - /// - /// When checking if own JID is in the participant list, we must use - /// the phone number equivalent if in LID mode, not the LID itself. - #[test] - fn test_own_jid_check_in_lid_mode() { - use std::collections::HashMap; - use wacore_binary::Jid; - - let own_lid: Jid = "100000000000001.1@lid" - .parse() - .expect("test JID should be valid"); - let own_phone: Jid = "15551234567@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - - let mut lid_to_pn_map = HashMap::new(); - lid_to_pn_map.insert("100000000000001.1".to_string(), own_phone.clone()); - - // Simulate the own JID check logic from wacore/src/send.rs - let own_base_jid = own_lid.to_non_ad(); - let own_jid_to_check = if own_base_jid.is_lid() { - lid_to_pn_map - .get(own_base_jid.user.as_str()) - .map(|pn| pn.to_non_ad()) - .unwrap_or_else(|| own_base_jid.clone()) - } else { - own_base_jid.clone() - }; - - // Verify we're checking using the phone number - assert_eq!(own_jid_to_check.user, "15551234567"); - assert_eq!(own_jid_to_check.server, SERVER_JID); - } - - /// Test that sender key operations always use the display JID (LID) - /// regardless of what JID is used for E2E session decryption - #[tokio::test] - async fn test_sender_key_always_uses_display_jid() { - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore::libsignal::protocol::{SenderKeyStore, create_sender_key_distribution_message}; - - let backend = Arc::new( - SqliteStore::new("file:memdb_display_jid_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (_client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm.clone(), - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let group_jid: Jid = "120363021033254949@g.us" - .parse() - .expect("test JID should be valid"); - let display_jid: Jid = "100000000000001.1:75@lid" - .parse() - .expect("test JID should be valid"); - let encryption_jid: Jid = "15551234567:75@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - - // Store sender key using display JID (LID) - let display_protocol_address = display_jid.to_protocol_address(); - let display_sender_key_name = make_sender_key_name(&group_jid, &display_protocol_address); - - let device_arc = pm.get_device_arc().await; - { - let mut device_guard = device_arc.write().await; - create_sender_key_distribution_message( - &display_sender_key_name, - &mut *device_guard, - &mut rand::make_rng::<rand::rngs::StdRng>(), - ) - .await - .expect("Failed to create SKDM"); - } - - // Verify it's stored under display JID - let lookup_with_display = { - let device_guard = device_arc.read().await; - device_guard.load_sender_key(&display_sender_key_name).await - }; - assert!( - lookup_with_display - .expect("lookup should not error") - .is_some(), - "Sender key should be found with display JID (LID)" - ); - - // Verify it's NOT accessible via encryption JID (phone number) - let encryption_protocol_address = encryption_jid.to_protocol_address(); - let encryption_sender_key_name = - make_sender_key_name(&group_jid, &encryption_protocol_address); - - let lookup_with_encryption = { - let device_guard = device_arc.read().await; - device_guard - .load_sender_key(&encryption_sender_key_name) - .await - }; - assert!( - lookup_with_encryption - .expect("lookup should not error") - .is_none(), - "Sender key should NOT be found with encryption JID (phone number)" - ); - } - - /// Test edge case: Second message with only skmsg (no pkmsg/msg) - /// - /// After the first message establishes a session and sender key, - /// subsequent messages may contain only skmsg. These should still - /// be decrypted successfully, not skipped. - /// - /// Bug: The code was treating "no session messages" as "session failed", - /// causing it to skip skmsg decryption for all messages after the first. - #[tokio::test] - async fn test_second_message_with_only_skmsg_decrypts() { - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore::libsignal::protocol::{ - create_sender_key_distribution_message, process_sender_key_distribution_message, - }; - - use wacore::types::message::AddressingMode; - use wacore_binary::builder::NodeBuilder; - - let backend = Arc::new( - SqliteStore::new("file:memdb_second_msg_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm.clone(), - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let sender_jid: Jid = "100000000000001.1:75@lid" - .parse() - .expect("test JID should be valid"); - let group_jid: Jid = "120363021033254949@g.us" - .parse() - .expect("test JID should be valid"); - - // Step 1: Create and store a sender key (simulating first message processing) - let sender_protocol_address = sender_jid.to_protocol_address(); - let sender_key_name = make_sender_key_name(&group_jid, &sender_protocol_address); - - let device_arc = pm.get_device_arc().await; - { - let mut device_guard = device_arc.write().await; - let skdm = create_sender_key_distribution_message( - &sender_key_name, - &mut *device_guard, - &mut rand::make_rng::<rand::rngs::StdRng>(), - ) - .await - .expect("Failed to create SKDM"); - - process_sender_key_distribution_message(&sender_key_name, &skdm, &mut *device_guard) - .await - .expect("Failed to process SKDM"); - } - - // Create message with ONLY skmsg (simulating second message after session established) - let skmsg_ciphertext = { - let mut device_guard = device_arc.write().await; - let sender_key_msg = wacore::libsignal::protocol::group_encrypt( - &mut *device_guard, - &sender_key_name, - b"ping", - &mut rand::make_rng::<rand::rngs::StdRng>(), - ) - .await - .expect("Failed to encrypt with sender key"); - sender_key_msg.serialized().to_vec() - }; - - let skmsg_node = NodeBuilder::new("enc") - .attr("type", "skmsg") - .attr("v", "2") - .bytes(skmsg_ciphertext) - .build(); - - let message_node = node_to_arc( - NodeBuilder::new("message") - .attr("from", group_jid) - .attr("participant", sender_jid) - .attr("id", "SECOND_MSG_TEST") - .attr("t", "1759306493") - .attr("type", "text") - .attr("addressing_mode", AddressingMode::Lid.as_str()) - .children(vec![skmsg_node]) - .build(), - ); - - // Should NOT skip skmsg - before the fix this would incorrectly skip - client.clone().handle_incoming_message(message_node).await; - } - - /// Test case for UntrustedIdentity error handling and recovery - /// - /// Scenario: - /// - User re-installs WhatsApp or switches devices - /// - Their device generates a new identity key - /// - The bot still has the old identity key stored - /// - When a message arrives, Signal Protocol rejects it as "UntrustedIdentity" - /// - The bot should catch this error, clear the old identity using the FULL protocol address (with device ID), and retry - /// - /// This test verifies that: - /// 1. process_session_enc_batch handles UntrustedIdentity gracefully - /// 2. The deletion uses the correct full address (name.device_id) not just the name - /// 3. No panic occurs when UntrustedIdentity is encountered - /// 4. The error is logged appropriately - /// 5. The bot continues processing instead of propagating the error - #[tokio::test] - async fn test_untrusted_identity_error_is_caught_and_handled() { - use crate::store::SqliteStore; - use std::sync::Arc; - - // Setup - let backend = Arc::new( - SqliteStore::new("file:memdb_untrusted_identity_caught?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm.clone(), - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let sender_jid: Jid = "559981212574@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - - let info = Arc::new(MessageInfo { - source: crate::types::message::MessageSource { - sender: sender_jid.clone(), - chat: sender_jid.clone(), - ..Default::default() - }, - ..Default::default() - }); - - log::info!("Test: UntrustedIdentity scenario for {}", sender_jid); - - // Create a malformed/invalid encrypted node to trigger error handling path - // This won't create UntrustedIdentity specifically, but tests the error handling code path - // The important fix is that when UntrustedIdentity IS raised, the code uses - // address.to_string() (which gives "559981212574.0") instead of address.name() - // (which only gives "559981212574") for the deletion key. - let enc_node = NodeBuilder::new("enc") - .attr("type", "msg") - .attr("v", "2") - .bytes(vec![0xFF; 100]) // Invalid encrypted payload - .build(); - - let enc_node_ref = enc_node.as_node_ref(); - let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; - - // Call process_session_enc_batch - // This should handle any errors gracefully without panicking - let outcome = client - .process_session_enc_batch( - &payloads, - &info, - &sender_jid, - crate::types::events::DecryptFailMode::Show, - ) - .await; - - log::info!( - "Test: process_session_enc_batch completed - success: {}", - outcome.decrypted - ); - - // The key is that this didn't panic - deletion uses full protocol address - } - - /// Test case: Error handling during batch processing - /// - /// When multiple messages are being processed in a batch, if one triggers - /// an error (like UntrustedIdentity), it should be handled without affecting - /// other messages in the batch. - #[tokio::test] - async fn test_untrusted_identity_does_not_break_batch_processing() { - use crate::store::SqliteStore; - use std::sync::Arc; - - let backend = Arc::new( - SqliteStore::new("file:memdb_untrusted_batch?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm.clone(), - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let sender_jid: Jid = "559981212574@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - - let info = Arc::new(MessageInfo { - source: crate::types::message::MessageSource { - sender: sender_jid.clone(), - chat: sender_jid.clone(), - ..Default::default() - }, - ..Default::default() - }); - - log::info!("Test: Batch processing with multiple error messages"); - - // Create multiple invalid encrypted nodes to test batch error handling - let mut enc_nodes = Vec::new(); - - // First message: Invalid encrypted payload - let enc_node_1 = NodeBuilder::new("enc") - .attr("type", "msg") - .attr("v", "2") - .bytes(vec![0xFF; 50]) - .build(); - enc_nodes.push(enc_node_1); - - // Second message: Another invalid encrypted payload - let enc_node_2 = NodeBuilder::new("enc") - .attr("type", "msg") - .attr("v", "2") - .bytes(vec![0xAA; 50]) - .build(); - enc_nodes.push(enc_node_2); - - log::info!("Test: Created batch of 2 messages with invalid data"); - - let payloads: Vec<EncPayload> = enc_nodes - .iter() - .filter_map(|n| EncPayload::from_node_ref(&n.as_node_ref())) - .collect(); - - // Process the batch - // Should handle all errors gracefully without stopping at first error - let outcome = client - .process_session_enc_batch( - &payloads, - &info, - &sender_jid, - crate::types::events::DecryptFailMode::Show, - ) - .await; - - log::info!( - "Test: Batch processing completed - success: {}", - outcome.decrypted - ); - } - - /// Test case: Error handling in group chat context - /// - /// When processing messages from group members, if identity errors occur, - /// they should be handled per-sender without affecting other group members. - #[tokio::test] - async fn test_untrusted_identity_in_group_context() { - use crate::store::SqliteStore; - use std::sync::Arc; - - let backend = Arc::new( - SqliteStore::new("file:memdb_untrusted_group?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm.clone(), - mock_transport(), - mock_http_client(), - None, - ) - .await; - - // Simulate a group chat scenario - let group_jid: Jid = "120363021033254949@g.us" - .parse() - .expect("test JID should be valid"); - let sender_phone: Jid = "559981212574@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - - let info = Arc::new(MessageInfo { - source: crate::types::message::MessageSource { - sender: sender_phone.clone(), - chat: group_jid.clone(), - is_group: true, - ..Default::default() - }, - ..Default::default() - }); - - log::info!("Test: Group context - error handling for {}", sender_phone); - - // Create an invalid encrypted message - let enc_node = NodeBuilder::new("enc") - .attr("type", "msg") - .attr("v", "2") - .bytes(vec![0xFF; 100]) - .build(); - - let enc_node_ref = enc_node.as_node_ref(); - let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; - - // Process the message - // Should handle errors gracefully in group context - let outcome = client - .process_session_enc_batch( - &payloads, - &info, - &sender_phone, - crate::types::events::DecryptFailMode::Show, - ) - .await; - - log::info!( - "Test: Group message processed - success: {}", - outcome.decrypted - ); - } - - /// Test case: DM message parsing for self-sent messages via LID - /// - /// Scenario: - /// - You send a DM to another user from your phone - /// - Your bot receives the echo with from=your_LID, recipient=their_LID - /// - peer_recipient_pn contains the RECIPIENT's phone number (not sender's) - /// - /// The fix ensures: - /// 1. is_from_me is correctly detected for LID senders - /// 2. sender_alt is NOT populated with peer_recipient_pn (that's the recipient's PN) - /// 3. Decryption uses own PN via the is_from_me fallback path - #[tokio::test] - async fn test_parse_message_info_self_sent_dm_via_lid() { - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore_binary::builder::NodeBuilder; - - let backend = Arc::new( - SqliteStore::new("file:memdb_self_dm_lid_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - - // Set up own phone number and LID - { - let device_arc = pm.get_device_arc().await; - let mut device = device_arc.write().await; - device.pn = Some( - "15551234567@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), - ); - device.lid = Some( - "100000000000001@lid" - .parse() - .expect("test JID should be valid"), - ); - } - - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - // Simulate self-sent DM to another user (from your phone to your bot echo) - // Real log example: - // from="100000000000001@lid" recipient="39492358562039@lid" peer_recipient_pn="559985213786@s.whatsapp.net" - let self_dm_node = NodeBuilder::new("message") - .attr("from", "100000000000001@lid") // Your LID - .attr("recipient", "39492358562039@lid") // Recipient's LID - .attr("peer_recipient_pn", "559985213786@s.whatsapp.net") // Recipient's PN (NOT sender's!) - .attr("notify", "jl") - .attr("id", "AC756E00B560721DBC4C0680131827EA") - .attr("t", "1764845025") - .attr("type", "text") - .build(); - - let info = client - .parse_message_info(&self_dm_node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - // Assertions: - // 1. is_from_me should be true (LID matches own_lid) - assert!( - info.source.is_from_me, - "Should detect self-sent DM from own LID" - ); - - // 2. sender_alt should be own PN (derived from own_jid, not message attrs) - assert!( - info.source.sender_alt.is_some(), - "sender_alt should be own PN for self-sent LID messages" - ); - assert_eq!( - info.source.sender_alt.as_ref().unwrap().user, - "15551234567", - "sender_alt should be the own PN user" - ); - - assert_eq!( - info.source.chat.user, "39492358562039", - "Chat should be the recipient's LID" - ); - - assert_eq!( - info.source.sender.user, "100000000000001", - "Sender should be own LID" - ); - } - - /// Test case: DM message parsing for messages from others via LID - /// - /// Scenario: - /// - Another user sends you a DM - /// - Message arrives with from=their_LID, sender_pn=their_phone_number - /// - /// The fix ensures: - /// 1. is_from_me is false - /// 2. sender_alt is populated from sender_pn attribute (if present) - /// 3. Decryption uses sender_alt for session lookup - #[tokio::test] - async fn test_parse_message_info_dm_from_other_via_lid() { - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore_binary::builder::NodeBuilder; - - let backend = Arc::new( - SqliteStore::new("file:memdb_other_dm_lid_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - - // Set up own phone number and LID - { - let device_arc = pm.get_device_arc().await; - let mut device = device_arc.write().await; - device.pn = Some( - "15551234567@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), - ); - device.lid = Some( - "100000000000001@lid" - .parse() - .expect("test JID should be valid"), - ); - } - - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - // Simulate DM from another user via their LID - // The sender_pn attribute should contain their phone number for session lookup - let other_dm_node = NodeBuilder::new("message") - .attr("from", "39492358562039@lid") // Sender's LID (not ours) - .attr("sender_pn", "559985213786@s.whatsapp.net") // Sender's phone number - .attr("notify", "Other User") - .attr("id", "AABBCCDD1234567890") - .attr("t", "1764845100") - .attr("type", "text") - .build(); - - let info = client - .parse_message_info(&other_dm_node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - assert!( - !info.source.is_from_me, - "Should NOT be detected as self-sent" - ); - - assert!( - info.source.sender_alt.is_some(), - "sender_alt should be set from sender_pn attribute" - ); - assert_eq!( - info.source - .sender_alt - .as_ref() - .expect("sender_alt should be present") - .user, - "559985213786", - "sender_alt should contain sender's phone number" - ); - - assert_eq!( - info.source.chat.user, "39492358562039", - "Chat should be the sender's LID (non-AD)" - ); - - assert_eq!( - info.source.sender.user, "39492358562039", - "Sender should be other user's LID" - ); - } - - /// Test case: DM message to self (own chat, like "Notes to Myself") - /// - /// Scenario: - /// - You send a message to yourself (your own chat) - /// - from=your_LID, recipient=your_LID, peer_recipient_pn=your_PN - /// - /// This is the original bug case that was fixed earlier. - #[tokio::test] - async fn test_parse_message_info_dm_to_self() { - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore_binary::builder::NodeBuilder; - - let backend = Arc::new( - SqliteStore::new("file:memdb_dm_to_self_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - - // Set up own phone number and LID - { - let device_arc = pm.get_device_arc().await; - let mut device = device_arc.write().await; - device.pn = Some( - "15551234567@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), - ); - device.lid = Some( - "100000000000001@lid" - .parse() - .expect("test JID should be valid"), - ); - } - - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - // Simulate DM to self (like "Notes to Myself" or pinging yourself) - // from=your_LID, recipient=your_LID, peer_recipient_pn=your_PN - let self_chat_node = NodeBuilder::new("message") - .attr("from", "100000000000001@lid") // Your LID - .attr("recipient", "100000000000001@lid") // Also your LID (self-chat) - .attr("peer_recipient_pn", "15551234567@s.whatsapp.net") // Your PN - .attr("notify", "jl") - .attr("id", "AC391DD54A28E1CE1F3B106DF9951FAD") - .attr("t", "1764822437") - .attr("type", "text") - .build(); - - let info = client - .parse_message_info(&self_chat_node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - assert!( - info.source.is_from_me, - "Should detect self-sent message to self-chat" - ); - - assert!( - info.source.sender_alt.is_some(), - "sender_alt should be own PN for self-sent LID messages" - ); - assert_eq!( - info.source.sender_alt.as_ref().unwrap().user, - "15551234567", - "sender_alt should match own PN" - ); - - assert_eq!( - info.source.chat.user, "100000000000001", - "Chat should be self (recipient)" - ); - - assert_eq!( - info.source.sender.user, "100000000000001", - "Sender should be own LID" - ); - } - - /// Test that receiving a DM with sender_lid populates the lid_pn_cache. - /// - /// This is the key behavior for the LID-PN session mismatch fix: - /// When we receive a message from a phone number with sender_lid attribute, - /// we cache the phone->LID mapping so that when sending replies, we can - /// reuse the existing LID session instead of creating a new PN session. - /// - /// Flow being tested: - /// 1. Receive message from 559980000001@s.whatsapp.net with sender_lid=100000012345678@lid - /// 2. Cache should be populated with: 559980000001 -> 100000012345678 - /// 3. When sending reply to 559980000001, we can look up the LID and use existing session - #[tokio::test] - async fn test_lid_pn_cache_populated_on_message_with_sender_lid() { - // Setup client - let backend = Arc::new( - SqliteStore::new("file:memdb_lid_cache_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let phone = "559980000001"; - let lid = "100000012345678"; - - // Verify cache is empty initially - assert!( - client.lid_pn_cache.get_current_lid(phone).await.is_none(), - "Cache should be empty before receiving message" - ); - - // Create a DM message node with sender_lid attribute - // This simulates receiving a message from WhatsApp Web - let dm_node = NodeBuilder::new("message") - .attr("from", Jid::pn(phone).to_string()) - .attr("sender_lid", Jid::lid(lid).to_string()) - .attr("id", "TEST123456789") - .attr("t", "1765482972") - .attr("type", "text") - .children([NodeBuilder::new("enc") - .attr("type", "pkmsg") - .attr("v", "2") - .bytes(vec![0u8; 100]) // Dummy encrypted content - .build()]) - .build(); - - // Call handle_incoming_message - this will fail to decrypt (no real session) - // but it should still populate the cache before attempting decryption - client - .clone() - .handle_incoming_message(node_to_arc(dm_node)) - .await; - - // Verify the cache was populated - let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; - assert!( - cached_lid.is_some(), - "Cache should be populated after receiving message with sender_lid" - ); - assert_eq!( - cached_lid.expect("cache should have LID"), - lid, - "Cached LID should match the sender_lid from the message" - ); - } - - /// Test that messages without sender_lid do NOT populate the cache. - /// - /// This ensures we don't accidentally cache incorrect mappings. - #[tokio::test] - async fn test_lid_pn_cache_not_populated_without_sender_lid() { - // Setup client - let backend = Arc::new( - SqliteStore::new("file:memdb_no_lid_cache_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let phone = "559980000001"; - - // Create a DM message node WITHOUT sender_lid attribute - let dm_node = NodeBuilder::new("message") - .attr("from", Jid::pn(phone).to_string()) - // Note: NO sender_lid attribute - .attr("id", "TEST123456789") - .attr("t", "1765482972") - .attr("type", "text") - .children([NodeBuilder::new("enc") - .attr("type", "pkmsg") - .attr("v", "2") - .bytes(vec![0u8; 100]) - .build()]) - .build(); - - // Call handle_incoming_message - client - .clone() - .handle_incoming_message(node_to_arc(dm_node)) - .await; - - assert!( - client.lid_pn_cache.get_current_lid(phone).await.is_none(), - "Cache should NOT be populated for messages without sender_lid" - ); - } - - /// Test that messages from LID senders with participant_pn DO populate the cache. - /// - /// When the sender is a LID (e.g., in LID-mode groups), and participant_pn - /// contains their phone number, we SHOULD cache this mapping because: - /// 1. The cache is bidirectional - we need both LID->PN and PN->LID - /// 2. This enables sending to users we've only seen as LID senders - #[tokio::test] - async fn test_lid_pn_cache_populated_for_lid_sender_with_participant_pn() { - use wacore::types::message::AddressingMode; - - // Setup client - let backend = Arc::new( - SqliteStore::new("file:memdb_lid_sender_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let lid = "100000012345678"; - let phone = "559980000001"; - - // Create a message from a LID sender with participant_pn attribute - // This happens in LID-mode groups (addressing_mode="lid") - let group_node = NodeBuilder::new("message") - .attr("from", "120363123456789012@g.us") // Group chat - .attr("participant", Jid::lid(lid).to_string()) // Sender is LID - .attr("participant_pn", Jid::pn(phone).to_string()) // Their phone number - .attr("addressing_mode", AddressingMode::Lid.as_str()) // Required for participant_pn to be parsed - .attr("id", "TEST123456789") - .attr("t", "1765482972") - .attr("type", "text") - .children([NodeBuilder::new("enc") - .attr("type", "skmsg") - .attr("v", "2") - .bytes(vec![0u8; 100]) - .build()]) - .build(); - - // Call handle_incoming_message - client - .clone() - .handle_incoming_message(node_to_arc(group_node)) - .await; - - // Verify the cache WAS populated (bidirectional cache) - let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; - assert!( - cached_lid.is_some(), - "Cache should be populated for LID senders with participant_pn" - ); - assert_eq!( - cached_lid.expect("cache should have LID"), - lid, - "Cached LID should match the sender's LID" - ); - - // Also verify we can look up the phone number from the LID - let cached_pn = client.lid_pn_cache.get_phone_number(lid).await; - assert!(cached_pn.is_some(), "Reverse lookup (LID->PN) should work"); - assert_eq!( - cached_pn.expect("reverse lookup should return phone"), - phone, - "Cached phone number should match" - ); - } - - /// Test that multiple messages from the same sender update the cache correctly. - /// - /// This ensures the cache handles repeated messages gracefully. - #[tokio::test] - async fn test_lid_pn_cache_handles_repeated_messages() { - // Setup client - let backend = Arc::new( - SqliteStore::new("file:memdb_repeated_msg_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let phone = "559980000001"; - let lid = "100000012345678"; - - // Send multiple messages from the same sender - for i in 0..3 { - let dm_node = NodeBuilder::new("message") - .attr("from", Jid::pn(phone).to_string()) - .attr("sender_lid", Jid::lid(lid).to_string()) - .attr("id", format!("TEST{}", i)) - .attr("t", "1765482972") - .attr("type", "text") - .children([NodeBuilder::new("enc") - .attr("type", "pkmsg") - .attr("v", "2") - .bytes(vec![0u8; 100]) - .build()]) - .build(); - - client - .clone() - .handle_incoming_message(node_to_arc(dm_node)) - .await; - } - - // Verify the cache still has the correct mapping - let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; - assert!(cached_lid.is_some(), "Cache should contain the mapping"); - assert_eq!( - cached_lid.expect("cache should have LID"), - lid, - "Cached LID should be correct after multiple messages" - ); - } - - /// Test that PN-addressed messages use LID for session lookup when LID mapping is known. - /// - /// This test verifies the fix for the MAC verification failure bug: - /// WhatsApp Web's SignalAddress.toString() ALWAYS converts PN addresses to LID - /// when a LID mapping is known. The Rust client must do the same to ensure - /// session keys match between clients. - /// - /// Bug scenario: - /// 1. WhatsApp Web Client A sends a group message to our Rust client - /// 2. Rust client creates session under PN address (559980000001@c.us.0) - /// 3. Rust client sends group response, creates session under LID (100000012345678@lid.0) - /// 4. Client A sends DM to Rust client from PN address - /// 5. Rust client tries to decrypt using PN address but session is under LID - /// 6. MAC verification fails because wrong session is used - /// - /// Fix: When receiving a PN-addressed message, if we have a LID mapping, - /// use the LID address for session lookup (matching WhatsApp Web behavior). - #[tokio::test] - async fn test_pn_message_uses_lid_for_session_lookup_when_mapping_known() { - use crate::lid_pn_cache::LidPnEntry; - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore::types::jid::JidExt; - - let backend = Arc::new( - SqliteStore::new("file:memdb_pn_to_lid_session_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let lid = "100000012345678"; - let phone = "559980000001"; - - // Pre-populate the LID-PN cache (simulating a previous group message) - let entry = LidPnEntry::new( - lid.to_string(), - phone.to_string(), - crate::lid_pn_cache::LearningSource::PeerLidMessage, - ); - client.lid_pn_cache.add(&entry).await; - - // Verify the cache has the mapping - let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; - assert_eq!( - cached_lid.as_deref(), - Some(lid), - "Cache should have the LID-PN mapping" - ); - - // Test scenario: Parse a PN-addressed DM message (with sender_lid attribute) - let dm_node_with_sender_lid = wacore_binary::builder::NodeBuilder::new("message") - .attr("from", Jid::pn(phone).to_string()) - .attr("sender_lid", Jid::lid(lid).to_string()) - .attr("id", "test_dm_with_lid") - .attr("t", "1765494882") - .attr("type", "text") - .build(); - - let info = client - .parse_message_info(&dm_node_with_sender_lid.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - // Verify sender is PN but sender_alt is LID - assert_eq!(info.source.sender.user, phone); - assert_eq!(info.source.sender.server, wacore_binary::Server::Pn); - assert!(info.source.sender_alt.is_some()); - assert_eq!( - info.source - .sender_alt - .as_ref() - .expect("sender_alt should be present") - .user, - lid - ); - assert_eq!( - info.source - .sender_alt - .as_ref() - .expect("sender_alt should be present") - .server, - wacore_binary::Server::Lid - ); - - // Now simulate what handle_incoming_message does: determine encryption JID - // We can't easily call handle_incoming_message, so we'll test the logic directly - let sender = &info.source.sender; - let alt = info.source.sender_alt.as_ref(); - // Apply the same logic as in handle_incoming_message - let sender_encryption_jid = if sender.is_lid() { - sender.clone() - } else if sender.is_pn() { - if let Some(alt_jid) = alt - && alt_jid.is_lid() - { - // Use the LID from the message attribute - Jid { - user: alt_jid.user.clone(), - server: wacore_binary::Server::Lid, - device: sender.device, - agent: sender.agent, - integrator: sender.integrator, - } - } else if let Some(lid_user) = client.lid_pn_cache.get_current_lid(&sender.user).await { - // Use the cached LID - Jid { - user: lid_user, - server: wacore_binary::Server::Lid, - device: sender.device, - agent: sender.agent, - integrator: sender.integrator, - } - } else { - sender.clone() - } - } else { - sender.clone() - }; - - // Verify the encryption JID uses the LID, not the PN - assert_eq!( - sender_encryption_jid.user, lid, - "Encryption JID should use LID user" - ); - assert_eq!( - sender_encryption_jid.server, - wacore_binary::Server::Lid, - "Encryption JID should use LID server" - ); - - // Verify the protocol address format - let protocol_address = sender_encryption_jid.to_protocol_address(); - assert_eq!( - protocol_address.to_string(), - format!("{}@lid.0", lid), - "Protocol address should be in LID format" - ); - } - - /// Test that PN-addressed messages use cached LID even without sender_lid attribute. - /// - /// This tests the fallback path where the message doesn't have a sender_lid - /// attribute but we have a previously cached LID mapping. - #[tokio::test] - async fn test_pn_message_uses_cached_lid_without_sender_lid_attribute() { - use crate::lid_pn_cache::LidPnEntry; - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore::types::jid::JidExt; - - let backend = Arc::new( - SqliteStore::new("file:memdb_cached_lid_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let lid = "100000012345678"; - let phone = "559980000001"; - - // Pre-populate the LID-PN cache - let entry = LidPnEntry::new( - lid.to_string(), - phone.to_string(), - crate::lid_pn_cache::LearningSource::PeerLidMessage, - ); - client.lid_pn_cache.add(&entry).await; - - // Parse a PN-addressed DM message WITHOUT sender_lid attribute - let dm_node_without_sender_lid = wacore_binary::builder::NodeBuilder::new("message") - .attr("from", Jid::pn(phone).to_string()) - // Note: No sender_lid attribute! - .attr("id", "test_dm_no_lid") - .attr("t", "1765494882") - .attr("type", "text") - .build(); - - let info = client - .parse_message_info(&dm_node_without_sender_lid.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - // Verify sender is PN and NO sender_alt (since there's no sender_lid attribute) - assert_eq!(info.source.sender.user, phone); - assert_eq!(info.source.sender.server, wacore_binary::Server::Pn); - assert!( - info.source.sender_alt.is_none(), - "Should have no sender_alt without sender_lid attribute" - ); - - // Apply the encryption JID logic (fallback to cached LID) - let sender = &info.source.sender; - let alt = info.source.sender_alt.as_ref(); - let sender_encryption_jid = if sender.is_lid() { - sender.clone() - } else if sender.is_pn() { - if let Some(alt_jid) = alt - && alt_jid.is_lid() - { - Jid { - user: alt_jid.user.clone(), - server: wacore_binary::Server::Lid, - device: sender.device, - agent: sender.agent, - integrator: sender.integrator, - } - } else if let Some(lid_user) = client.lid_pn_cache.get_current_lid(&sender.user).await { - // This is the path we're testing - fallback to cached LID - Jid { - user: lid_user, - server: wacore_binary::Server::Lid, - device: sender.device, - agent: sender.agent, - integrator: sender.integrator, - } - } else { - sender.clone() - } - } else { - sender.clone() - }; - - // Verify the encryption JID uses the cached LID - assert_eq!( - sender_encryption_jid.user, lid, - "Encryption JID should use cached LID user" - ); - assert_eq!( - sender_encryption_jid.server, - wacore_binary::Server::Lid, - "Encryption JID should use LID server" - ); - - let protocol_address = sender_encryption_jid.to_protocol_address(); - assert_eq!( - protocol_address.to_string(), - format!("{}@lid.0", lid), - "Protocol address should be in LID format from cached mapping" - ); - } - - /// Test that PN-addressed messages use PN when no LID mapping is known. - /// - /// When there's no LID mapping available, we should fall back to using - /// the PN address for session lookup. - #[tokio::test] - async fn test_pn_message_uses_pn_when_no_lid_mapping() { - use crate::store::SqliteStore; - use std::sync::Arc; - use wacore::types::jid::JidExt; - - let backend = Arc::new( - SqliteStore::new("file:memdb_no_lid_mapping_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let phone = "559980000001"; - - // Don't populate the cache - simulate first-time contact - - // Parse a PN-addressed DM message without sender_lid - let dm_node = wacore_binary::builder::NodeBuilder::new("message") - .attr("from", Jid::pn(phone).to_string()) - .attr("id", "test_dm_no_mapping") - .attr("t", "1765494882") - .attr("type", "text") - .build(); - - let info = client - .parse_message_info(&dm_node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - // Verify no cached LID - let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; - assert!(cached_lid.is_none(), "Should have no cached LID mapping"); - - // Apply the encryption JID logic - let sender = &info.source.sender; - let alt = info.source.sender_alt.as_ref(); - - let sender_encryption_jid = if sender.is_lid() { - sender.clone() - } else if sender.is_pn() { - if let Some(alt_jid) = alt - && alt_jid.is_lid() - { - Jid { - user: alt_jid.user.clone(), - server: wacore_binary::Server::Lid, - device: sender.device, - agent: sender.agent, - integrator: sender.integrator, - } - } else if let Some(lid_user) = client.lid_pn_cache.get_current_lid(&sender.user).await { - Jid { - user: lid_user, - server: wacore_binary::Server::Lid, - device: sender.device, - agent: sender.agent, - integrator: sender.integrator, - } - } else { - // This is the path we're testing - no LID mapping, use PN - sender.clone() - } - } else { - sender.clone() - }; - - // Verify the encryption JID uses the PN (no LID available) - assert_eq!( - sender_encryption_jid.user, phone, - "Encryption JID should use PN user when no LID mapping" - ); - assert_eq!( - sender_encryption_jid.server, - wacore_binary::Server::Pn, - "Encryption JID should use PN server when no LID mapping" - ); - - let protocol_address = sender_encryption_jid.to_protocol_address(); - assert_eq!( - protocol_address.to_string(), - format!("{}@c.us.0", phone), - "Protocol address should be in PN format when no LID mapping" - ); - } - - // and PDO fallback behavior to ensure robust message recovery. - - /// Helper to create a test MessageInfo with customizable fields - fn create_test_message_info(chat: &str, msg_id: &str, sender: &str) -> MessageInfo { - use wacore::types::message::{EditAttribute, MessageCategory, MessageSource, MsgMetaInfo}; - - let chat_jid: Jid = chat.parse().expect("valid chat JID"); - let sender_jid: Jid = sender.parse().expect("valid sender JID"); - - MessageInfo { - id: msg_id.to_string(), - server_id: 0, - r#type: "text".to_string(), - source: MessageSource { - chat: chat_jid.clone(), - sender: sender_jid, - sender_alt: None, - recipient_alt: None, - is_from_me: false, - is_group: chat_jid.is_group(), - addressing_mode: None, - broadcast_list_owner: None, - recipient: None, - }, - timestamp: wacore::time::now_utc(), - push_name: "Test User".to_string(), - category: MessageCategory::default(), - multicast: false, - media_type: "".to_string(), - edit: EditAttribute::default(), - bot_info: None, - meta_info: MsgMetaInfo::default(), - verified_name: None, - device_sent_meta: None, - ephemeral_expiration: None, - is_offline: false, - unavailable_request_id: None, - server_timestamp_us: None, - verified_level: None, - verified_name_serial: None, - peer_recipient_pn: None, - bcl_participants: Vec::new(), - } - } - - /// Helper to create a test client for retry tests with a unique database - async fn create_test_client_for_retry_with_id(test_id: &str) -> Arc<Client> { - use portable_atomic::AtomicU64; - use std::sync::atomic::Ordering; - static COUNTER: AtomicU64 = AtomicU64::new(0); - - let unique_id = COUNTER.fetch_add(1, Ordering::SeqCst); - let db_name = format!( - "file:memdb_retry_{}_{}_{}?mode=memory&cache=shared", - test_id, - unique_id, - std::process::id() - ); - - let backend = Arc::new( - SqliteStore::new(&db_name) - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - client - } - - #[tokio::test] - async fn test_increment_retry_count_starts_at_one() { - let client = create_test_client_for_retry_with_id("starts_at_one").await; - - let cache_key = "test_chat:msg123:sender456"; - - // First increment should return 1 - let count = client - .increment_retry_count(cache_key, RetryReason::NoSession) - .await; - assert_eq!(count, Some(1), "First retry should be count 1"); - - // Verify it's stored in cache - let stored = client.message_retry_counts.get(cache_key).await; - assert_eq!(stored, Some(1), "Cache should store count 1"); - } - - #[tokio::test] - async fn test_increment_retry_count_increments_correctly() { - let client = create_test_client_for_retry_with_id("increments").await; - - let cache_key = "test_chat:msg456:sender789"; - - // Simulate multiple retries - let count1 = client - .increment_retry_count(cache_key, RetryReason::NoSession) - .await; - let count2 = client - .increment_retry_count(cache_key, RetryReason::NoSession) - .await; - let count3 = client - .increment_retry_count(cache_key, RetryReason::NoSession) - .await; - - assert_eq!(count1, Some(1), "First retry should be 1"); - assert_eq!(count2, Some(2), "Second retry should be 2"); - assert_eq!(count3, Some(3), "Third retry should be 3"); - } - - #[tokio::test] - async fn test_increment_retry_count_respects_max_retries() { - let client = create_test_client_for_retry_with_id("max_retries").await; - - let cache_key = "test_chat:msg_max:sender_max"; - - // Exhaust all retries (MAX_DECRYPT_RETRIES = 5) - for i in 1..=5 { - let count = client - .increment_retry_count(cache_key, RetryReason::NoSession) - .await; - assert_eq!(count, Some(i), "Retry {} should return {}", i, i); - } - - // 6th attempt should return None (max reached) - let count_after_max = client - .increment_retry_count(cache_key, RetryReason::NoSession) - .await; - assert_eq!( - count_after_max, None, - "After max retries, should return None" - ); - - // Verify cache still has max value - let stored = client.message_retry_counts.get(cache_key).await; - assert_eq!(stored, Some(5), "Cache should retain max count"); - } - - #[tokio::test] - async fn test_retry_count_different_messages_are_independent() { - let client = create_test_client_for_retry_with_id("independent").await; - - let key1 = "chat1:msg1:sender1"; - let key2 = "chat1:msg2:sender1"; // Same chat and sender, different message - let key3 = "chat2:msg1:sender2"; // Different chat and sender - - // Increment each independently - let _ = client - .increment_retry_count(key1, RetryReason::NoSession) - .await; - let _ = client - .increment_retry_count(key1, RetryReason::NoSession) - .await; - let _ = client - .increment_retry_count(key1, RetryReason::NoSession) - .await; // key1 = 3 - - let _ = client - .increment_retry_count(key2, RetryReason::NoSession) - .await; // key2 = 1 - - let _ = client - .increment_retry_count(key3, RetryReason::NoSession) - .await; - let _ = client - .increment_retry_count(key3, RetryReason::NoSession) - .await; // key3 = 2 - - // Verify each has independent counts - assert_eq!(client.message_retry_counts.get(key1).await, Some(3)); - assert_eq!(client.message_retry_counts.get(key2).await, Some(1)); - assert_eq!(client.message_retry_counts.get(key3).await, Some(2)); - } - - #[tokio::test] - async fn test_retry_cache_key_format() { - // Verify the cache key format is consistent - let info = create_test_message_info( - "120363021033254949@g.us", - "3EB0ABCD1234", - "5511999998888@s.whatsapp.net", - ); - - let expected_key = format!("{}:{}:{}", info.source.chat, info.id, info.source.sender); - assert_eq!( - expected_key, - "120363021033254949@g.us:3EB0ABCD1234:5511999998888@s.whatsapp.net" - ); - - // Verify key uniqueness for different senders in same group - let info2 = create_test_message_info( - "120363021033254949@g.us", - "3EB0ABCD1234", // Same message ID - "5511888887777@s.whatsapp.net", // Different sender - ); - - let key2 = format!("{}:{}:{}", info2.source.chat, info2.id, info2.source.sender); - assert_ne!( - expected_key, key2, - "Different senders should have different keys" - ); - } - - /// Test concurrent retry increments are properly serialized. - /// - /// The increment operation uses get+insert which is not fully atomic, - /// but is sufficient since message retry processing is serialized per key - /// by the per-chat lock. At most 5 increments should succeed. - #[tokio::test] - async fn test_concurrent_retry_increments() { - use tokio::task::JoinSet; - - let client = create_test_client_for_retry_with_id("concurrent").await; - let cache_key = "concurrent_test:msg:sender"; - - // Spawn 10 concurrent increment tasks - let mut tasks = JoinSet::new(); - for _ in 0..10 { - let client_clone = client.clone(); - let key = cache_key.to_string(); - tasks.spawn(async move { - client_clone - .increment_retry_count(&key, RetryReason::NoSession) - .await - }); - } - - // Collect all results - let mut results = Vec::new(); - while let Some(result) = tasks.join_next().await { - if let Ok(count) = result { - results.push(count); - } - } - - // With atomic operations, exactly 5 should succeed and 5 should fail - let valid_counts: Vec<_> = results.iter().filter(|r| r.is_some()).collect(); - let none_counts: Vec<_> = results.iter().filter(|r| r.is_none()).collect(); - - assert_eq!( - valid_counts.len(), - 5, - "Exactly 5 increments should succeed with atomic operations" - ); - assert_eq!( - none_counts.len(), - 5, - "Exactly 5 should return None (after max is reached)" - ); - - // Verify the successful increments returned values 1-5 - let mut values: Vec<u8> = valid_counts.iter().filter_map(|r| **r).collect(); - values.sort(); - assert_eq!( - values, - vec![1, 2, 3, 4, 5], - "Successful increments should return 1, 2, 3, 4, 5" - ); - - // Final count should be 5 (max) - let final_count = client.message_retry_counts.get(cache_key).await; - assert_eq!(final_count, Some(5), "Final count should be capped at 5"); - } - - #[tokio::test] - async fn test_high_retry_count_threshold() { - // Verify HIGH_RETRY_COUNT_THRESHOLD is set correctly - assert_eq!( - HIGH_RETRY_COUNT_THRESHOLD, 3, - "High retry threshold should be 3" - ); - assert_eq!(MAX_DECRYPT_RETRIES, 5, "Max retries should be 5"); - // Compile-time assertion that threshold < max (avoids clippy warning) - const _: () = assert!(HIGH_RETRY_COUNT_THRESHOLD < MAX_DECRYPT_RETRIES); - } - - #[tokio::test] - async fn test_message_info_creation_for_groups() { - let info = create_test_message_info( - "120363021033254949@g.us", - "MSG123", - "5511999998888@s.whatsapp.net", - ); - - assert!( - info.source.is_group, - "Group JID should be detected as group" - ); - assert!( - !info.source.is_from_me, - "Test messages default to not from me" - ); - assert_eq!(info.id, "MSG123"); - } - - #[tokio::test] - async fn test_message_info_creation_for_dm() { - let info = create_test_message_info( - "5511999998888@s.whatsapp.net", - "DM456", - "5511999998888@s.whatsapp.net", - ); - - assert!( - !info.source.is_group, - "DM JID should not be detected as group" - ); - assert_eq!(info.id, "DM456"); - } - - #[tokio::test] - async fn test_retry_count_cache_expiration() { - // Note: This test verifies cache configuration, not actual TTL (which would be slow) - let client = create_test_client_for_retry_with_id("expiration").await; - - // The cache should have a TTL of 5 minutes (300 seconds) as configured in client.rs - // We can verify entries are being stored and the cache is functional - let cache_key = "expiry_test:msg:sender"; - - let count = client - .increment_retry_count(cache_key, RetryReason::NoSession) - .await; - assert_eq!(count, Some(1)); - - // Entry should still exist immediately after - let stored = client.message_retry_counts.get(cache_key).await; - assert!( - stored.is_some(), - "Entry should exist immediately after insert" - ); - } - - #[tokio::test] - async fn test_spawn_retry_receipt_basic_flow() { - // This is an integration test that verifies spawn_retry_receipt - // doesn't panic and updates the retry count correctly - - let client = create_test_client_for_retry_with_id("spawn_basic").await; - let info = create_test_message_info( - "120363021033254949@g.us", - "SPAWN_TEST_MSG", - "5511999998888@s.whatsapp.net", - ); - - let cache_key = format!("{}:{}:{}", info.source.chat, info.id, info.source.sender); - - // Verify count starts at 0 - assert!( - client.message_retry_counts.get(&cache_key).await.is_none(), - "Cache should be empty initially" - ); - - // Call spawn_retry_receipt (this spawns a task, so we need to wait) - let info = Arc::new(info); - client.spawn_retry_receipt(&info, RetryReason::UnknownError); - - // Give the spawned task time to execute - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - - // Verify count was incremented (the actual send will fail due to no connection, but count should update) - let stored = client.message_retry_counts.get(&cache_key).await; - assert_eq!(stored, Some(1), "Retry count should be 1 after spawn"); - } - - #[tokio::test] - async fn test_spawn_retry_receipt_respects_max_retries() { - let client = create_test_client_for_retry_with_id("spawn_max").await; - let info = create_test_message_info( - "120363021033254949@g.us", - "MAX_RETRY_TEST", - "5511999998888@s.whatsapp.net", - ); - - let cache_key = format!("{}:{}:{}", info.source.chat, info.id, info.source.sender); - - // Pre-fill cache to max retries - client - .message_retry_counts - .insert(cache_key.clone(), MAX_DECRYPT_RETRIES) - .await; - - // Verify count is at max - assert_eq!( - client.message_retry_counts.get(&cache_key).await, - Some(MAX_DECRYPT_RETRIES) - ); - - // Call spawn_retry_receipt - should NOT increment (already at max) - let info = Arc::new(info); - client.spawn_retry_receipt(&info, RetryReason::UnknownError); - - // Give the spawned task time to execute - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - - // Count should still be at max (not incremented) - let stored = client.message_retry_counts.get(&cache_key).await; - assert_eq!( - stored, - Some(MAX_DECRYPT_RETRIES), - "Count should remain at max" - ); - } - - #[tokio::test] - async fn test_pdo_cache_key_format_matches() { - // PDO uses "{chat}:{msg_id}" format - // Retry uses "{chat}:{msg_id}:{sender}" format - // They are intentionally different to track independently - - let info = create_test_message_info( - "120363021033254949@g.us", - "PDO_KEY_TEST", - "5511999998888@s.whatsapp.net", - ); - - let retry_key = format!("{}:{}:{}", info.source.chat, info.id, info.source.sender); - let pdo_key = format!("{}:{}", info.source.chat, info.id); - - assert_ne!(retry_key, pdo_key, "PDO and retry keys should be different"); - assert!( - retry_key.starts_with(&pdo_key), - "Retry key should start with PDO key pattern" - ); - } - - #[tokio::test] - async fn test_multiple_senders_same_message_id_tracked_separately() { - // In a group, multiple senders could theoretically have the same message ID - // (unlikely but the system should handle it) - - let client = create_test_client_for_retry_with_id("multi_sender").await; - - let group = "120363021033254949@g.us"; - let msg_id = "SAME_MSG_ID"; - let sender1 = "5511111111111@s.whatsapp.net"; - let sender2 = "5522222222222@s.whatsapp.net"; - - let key1 = format!("{}:{}:{}", group, msg_id, sender1); - let key2 = format!("{}:{}:{}", group, msg_id, sender2); - - // Increment for sender1 multiple times - client - .increment_retry_count(&key1, RetryReason::NoSession) - .await; - client - .increment_retry_count(&key1, RetryReason::NoSession) - .await; - client - .increment_retry_count(&key1, RetryReason::NoSession) - .await; - - // Increment for sender2 once - client - .increment_retry_count(&key2, RetryReason::NoSession) - .await; - - // Verify independent tracking - assert_eq!( - client.message_retry_counts.get(&key1).await, - Some(3), - "Sender1 should have 3 retries" - ); - assert_eq!( - client.message_retry_counts.get(&key2).await, - Some(1), - "Sender2 should have 1 retry" - ); - } - - /// Test: Verify JID type detection for status broadcasts, broadcast lists, groups, and users. - #[test] - fn test_status_broadcast_jid_detection() { - use wacore_binary::{Jid, JidExt}; - - let status_jid: Jid = "status@broadcast".parse().expect("status JID should parse"); - assert!(status_jid.is_status_broadcast()); - - let broadcast_list: Jid = "123456789@broadcast" - .parse() - .expect("broadcast JID should parse"); - assert!(!broadcast_list.is_status_broadcast()); - assert!(broadcast_list.is_broadcast_list()); - - let group_jid: Jid = "120363021033254949@g.us" - .parse() - .expect("group JID should parse"); - assert!(!group_jid.is_status_broadcast()); - - let user_jid: Jid = "15551234567@s.whatsapp.net" - .parse() - .expect("user JID should parse"); - assert!(!user_jid.is_status_broadcast()); - } - - /// Test: Verify should_process_skmsg logic matches WA Web's canDecryptNext pattern. - /// - /// WA Web applies canDecryptNext uniformly: if pkmsg fails with a retriable error, - /// skmsg is skipped regardless of chat type (group, status, 1:1). No exception for - /// status broadcasts — the retry receipt for the pkmsg will cause the sender to - /// resend the entire message including SKDM. - #[test] - fn test_should_process_skmsg_logic_matches_wa_web() { - // Test cases: (chat_jid, session_empty, session_success, session_dupe, session_failed, expected) - let test_cases = [ - // Status broadcast: same rules as all other chats (WA Web: canDecryptNext is uniform) - ("status@broadcast", false, false, false, false, false), // Fail: session failed → skip skmsg - ("status@broadcast", false, false, true, false, true), // OK: duplicate - ("status@broadcast", false, true, false, false, true), // OK: success - ("status@broadcast", false, true, false, true, false), // Fail: mixed success + failure - ("status@broadcast", true, false, false, false, true), // OK: no session msgs - // Regular group - ("120363021033254949@g.us", false, false, false, false, false), - ("120363021033254949@g.us", false, false, true, false, true), - ("120363021033254949@g.us", false, true, false, false, true), - ("120363021033254949@g.us", false, true, false, true, false), - ("120363021033254949@g.us", true, false, false, false, true), - // 1:1 chat - ( - "15551234567@s.whatsapp.net", - false, - false, - false, - false, - false, - ), - ( - "15551234567@s.whatsapp.net", - true, - false, - false, - false, - true, - ), - ]; - - for (jid_str, session_empty, session_success, session_dupe, session_failed, expected) in - test_cases - { - let should_process_skmsg = should_process_skmsg_after_session( - session_empty, - SessionBatchOutcome { - decrypted: session_success, - duplicate: session_dupe, - had_failure: session_failed, - ..Default::default() - }, - ); - - assert_eq!( - should_process_skmsg, - expected, - "For chat {} with session_empty={}, session_success={}, session_dupe={}, session_failed={}: \ - expected should_process_skmsg={}, got {}", - jid_str, - session_empty, - session_success, - session_dupe, - session_failed, - expected, - should_process_skmsg - ); - } - } - - #[test] - fn skdm_only_fallback_ack_decision_requires_clean_session_batch() { - let clean_skdm = SessionBatchOutcome { - decrypted: true, - skdm_only: true, - ..Default::default() - }; - assert!( - should_ack_skdm_only_session_fallback(clean_skdm, true), - "a clean SKDM-only session batch needs the fallback ack" - ); - - let cases = [ - ( - SessionBatchOutcome { - dispatched: true, - ..clean_skdm - }, - true, - "content dispatch already acked", - ), - ( - SessionBatchOutcome { - had_failure: true, - ..clean_skdm - }, - true, - "local session failure must block positive ack", - ), - ( - SessionBatchOutcome { - plaintext_failed: true, - had_failure: true, - ..clean_skdm - }, - true, - "plaintext handler failure is not SKDM-only success", - ), - ( - SessionBatchOutcome { - undecryptable: true, - had_failure: true, - ..clean_skdm - }, - true, - "failure event must not be paired with positive ack", - ), - ( - SessionBatchOutcome { - decrypted: false, - ..clean_skdm - }, - true, - "fallback only applies after Signal decrypt success", - ), - ( - SessionBatchOutcome { - skdm_only: false, - ..clean_skdm - }, - true, - "regular content must ack via dispatch", - ), - ( - SessionBatchOutcome { - duplicate: true, - decrypted: false, - skdm_only: false, - ..Default::default() - }, - true, - "duplicates use the duplicate branch", - ), - (clean_skdm, false, "msmsg work must own its response"), - ]; - - for (outcome, bot_payloads_empty, reason) in cases { - assert!( - !should_ack_skdm_only_session_fallback(outcome, bot_payloads_empty), - "{reason}: {outcome:?}" - ); - } - } - - /// Test: parse_message_info returns error when message "id" attribute is missing - /// - /// Missing message IDs would cause silent collisions in caches/keys, so this - /// must be a hard error rather than defaulting to an empty string. - #[tokio::test] - async fn test_parse_message_info_missing_id_returns_error() { - let backend = Arc::new( - SqliteStore::new("file:memdb_missing_id_test?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("test backend should initialize"), - ); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let node = NodeBuilder::new("message") - .attr("from", "15551234567@s.whatsapp.net") - .attr("t", "1759295366") - .attr("type", "text") - .build(); - - let result = client.parse_message_info(&node.as_node_ref()).await; - - assert!( - result.is_err(), - "parse_message_info should fail when 'id' is missing" - ); - let err_msg = result.unwrap_err().to_string(); - assert!( - err_msg.contains("id"), - "Error message should mention missing 'id' attribute: {}", - err_msg - ); - } - #[tokio::test] - async fn test_no_sender_key_sends_immediate_retry() { - // Verify that when skmsg decryption fails with NoSenderKeyState, - // a retry receipt is sent immediately (no delay, no re-queue). - // This matches WA Web behavior where NoSenderKey → SignalRetryable → RETRY. - let _ = env_logger::builder().is_test(true).try_init(); - - use crate::store::SqliteStore; - use crate::store::persistence_manager::PersistenceManager; - use wacore_binary::NodeContent; - use wacore_binary::builder::NodeBuilder; - - let backend = Arc::new( - SqliteStore::new("file:memdb_retry_immediate?mode=memory&cache=shared") - .await - .expect("Failed to create test backend"), - ); - let pm = Arc::new( - PersistenceManager::new(backend.clone()) - .await - .expect("test backend should initialize"), - ); - let (client, _rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm.clone(), - mock_transport(), - mock_http_client(), - None, - ) - .await; - - let group_jid: Jid = "120363021033254949@g.us".parse().unwrap(); - let sender_jid: Jid = "1234567890:1@s.whatsapp.net".parse().unwrap(); - let msg_id = "TEST_IMMEDIATE_RETRY"; - - // Pseudo-valid SenderKeyMessage: Version 3 + Protobuf + Fake Sig (64 bytes) - let mut content = vec![0x33, 0x08, 0x01, 0x10, 0x01, 0x1A, 0x00]; - content.extend(vec![0u8; 64]); - - let node = NodeBuilder::new("message") - .attr("id", msg_id) - .attr("from", group_jid.clone()) - .attr("participant", sender_jid.clone()) - .attr("type", "text") - .children(vec![{ - let mut n = NodeBuilder::new("enc") - .attr("type", "skmsg") - .attr("v", "2") - .build(); - n.content = Some(NodeContent::Bytes(content)); - n - }]) - .build(); - - client - .clone() - .handle_incoming_message(node_to_arc(node)) - .await; - - // spawn_retry_receipt runs in a spawned task, wait for it - let retry_key = client - .make_retry_cache_key(&group_jid, msg_id, &sender_jid) - .await; - for _ in 0..20 { - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - if client.message_retry_counts.get(&retry_key).await.is_some() { - break; - } - } - assert_eq!( - client.message_retry_counts.get(&retry_key).await, - Some(1), - "NoSenderKeyState should immediately trigger retry receipt (count=1)" - ); - } - - #[test] - fn test_is_sender_key_distribution_only() { - let skdm = wa::message::SenderKeyDistributionMessage { - group_id: Some("group".into()), - axolotl_sender_key_distribution_message: Some(vec![1, 2, 3]), - }; - - // Empty message → false (no SKDM) - assert!(!is_sender_key_distribution_only(&mut wa::Message::default())); - - // SKDM only → true - assert!(is_sender_key_distribution_only(&mut wa::Message { - sender_key_distribution_message: Some(skdm.clone()), - ..Default::default() - })); - - // SKDM + message_context_info → still true (context_info is metadata) - assert!(is_sender_key_distribution_only(&mut wa::Message { - sender_key_distribution_message: Some(skdm.clone()), - message_context_info: Some(wa::MessageContextInfo::default()), - ..Default::default() - })); - - // SKDM + sticker → false (has user content) - assert!(!is_sender_key_distribution_only(&mut wa::Message { - sender_key_distribution_message: Some(skdm.clone()), - sticker_message: Some(Box::new(wa::message::StickerMessage::default())), - ..Default::default() - })); - - // SKDM + text → false (has user content) - assert!(!is_sender_key_distribution_only(&mut wa::Message { - sender_key_distribution_message: Some(skdm.clone()), - conversation: Some("hello".into()), - ..Default::default() - })); - - // protocol_message only (no SKDM) → false - assert!(!is_sender_key_distribution_only(&mut wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage::default())), - ..Default::default() - })); - } - - #[test] - fn skdm_only_detection_restores_carrier_fields() { - // The slow path takes the carrier fields out to compare the rest against - // default; it must restore them so callers still see the original message. - let mut msg = wa::Message { - sender_key_distribution_message: Some(wa::message::SenderKeyDistributionMessage { - group_id: Some("group".into()), - axolotl_sender_key_distribution_message: Some(vec![1, 2, 3]), - }), - fast_ratchet_key_sender_key_distribution_message: Some( - wa::message::SenderKeyDistributionMessage { - group_id: Some("group".into()), - axolotl_sender_key_distribution_message: Some(vec![4, 5, 6]), - }, - ), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![9, 8, 7]), - ..Default::default() - }), - ..Default::default() - }; - - assert!(is_sender_key_distribution_only(&mut msg)); - - // Pin the exact payloads of all three taken/restored carrier fields, not - // just presence: a buggy restore that put back a fresh default (losing the - // original contents) must fail here. - assert_eq!( - msg.sender_key_distribution_message - .as_ref() - .and_then(|s| s.axolotl_sender_key_distribution_message.as_deref()), - Some([1, 2, 3].as_slice()), - "sender_key_distribution_message payload must be restored unchanged" - ); - assert_eq!( - msg.fast_ratchet_key_sender_key_distribution_message - .as_ref() - .and_then(|s| s.axolotl_sender_key_distribution_message.as_deref()), - Some([4, 5, 6].as_slice()), - "fast_ratchet carrier payload must be restored unchanged" - ); - assert_eq!( - msg.message_context_info - .as_ref() - .and_then(|c| c.message_secret.as_deref()), - Some([9, 8, 7].as_slice()), - "message_context_info payload must be restored unchanged" - ); - } - - /// Test: unwrap_device_sent extracts a reaction from a DeviceSentMessage wrapper. - #[test] - fn test_unwrap_device_sent_extracts_reaction() { - let wrapped = wa::Message { - device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { - destination_jid: Some("5511999999999@s.whatsapp.net".to_string()), - message: Some(Box::new(wa::Message { - reaction_message: Some(wa::message::ReactionMessage { - text: Some("\u{2764}".to_string()), - ..Default::default() - }), - ..Default::default() - })), - phash: None, - })), - ..Default::default() - }; - - let mut unwrapped = unwrap_device_sent(wrapped); - assert!( - unwrapped.device_sent_message.is_none(), - "DSM wrapper should be removed" - ); - assert_eq!( - unwrapped - .reaction_message - .as_ref() - .and_then(|r| r.text.as_deref()), - Some("\u{2764}"), - "reaction should be accessible after unwrapping" - ); - assert!( - !is_sender_key_distribution_only(&mut unwrapped), - "unwrapped reaction should not be filtered as SKDM-only" - ); - } - - /// Test: unwrap_device_sent preserves the wrapper when inner message is None. - #[test] - fn test_unwrap_device_sent_preserves_empty_wrapper() { - let wrapped = wa::Message { - device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { - destination_jid: Some("5511999999999@s.whatsapp.net".to_string()), - message: None, - phash: None, - })), - ..Default::default() - }; - - let result = unwrap_device_sent(wrapped); - assert!( - result.device_sent_message.is_some(), - "empty DSM wrapper should be preserved" - ); - } - - /// Test: unwrap_device_sent passes through a plain message unchanged. - #[test] - fn test_unwrap_device_sent_passthrough() { - let msg = wa::Message { - conversation: Some("hello".to_string()), - ..Default::default() - }; - - let result = unwrap_device_sent(msg); - assert_eq!(result.conversation.as_deref(), Some("hello")); - } - - /// Test: unwrap_device_sent merges messageContextInfo from outer and inner, - /// matching WAWebDeviceSentMessageProtoUtils.unwrapDeviceSentMessage. - #[test] - fn test_unwrap_device_sent_merges_context_info() { - let wrapped = wa::Message { - // Outer message_context_info (from the DSM envelope) - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![10, 20, 30]), - limit_sharing_v2: Some(wa::LimitSharing::default()), - ..Default::default() - }), - device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { - destination_jid: Some("5511999999999@s.whatsapp.net".to_string()), - message: Some(Box::new(wa::Message { - conversation: Some("hello".to_string()), - // Inner has its own message_secret but no limit_sharing_v2 - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![1, 2, 3]), - ..Default::default() - }), - ..Default::default() - })), - phash: None, - })), - ..Default::default() - }; - - let result = unwrap_device_sent(wrapped); - let ctx = result.message_context_info.as_ref().unwrap(); - - assert_eq!( - ctx.message_secret, - Some(vec![1, 2, 3]), - "inner message_secret should be preferred" - ); - assert!( - ctx.limit_sharing_v2.is_some(), - "limit_sharing_v2 should come from outer (always)" - ); - } - - /// Test: unwrap_device_sent falls back to outer message_secret when inner has none. - #[test] - fn test_unwrap_device_sent_secret_fallback() { - let wrapped = wa::Message { - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![10, 20, 30]), - ..Default::default() - }), - device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { - destination_jid: Some("5511999999999@s.whatsapp.net".to_string()), - message: Some(Box::new(wa::Message { - conversation: Some("hello".to_string()), - // Inner has no message_context_info at all - ..Default::default() - })), - phash: None, - })), - ..Default::default() - }; - - let result = unwrap_device_sent(wrapped); - let ctx = result.message_context_info.as_ref().unwrap(); - assert_eq!( - ctx.message_secret, - Some(vec![10, 20, 30]), - "should fall back to outer message_secret" - ); - } - - #[tokio::test] - async fn test_parse_edit_attribute_sender_revoke() { - let client = create_test_client_for_retry_with_id("edit_sender_revoke").await; - - let node = NodeBuilder::new("message") - .attr("from", "status@broadcast") - .attr("id", "TEST123") - .attr("participant", "5551234567@lid") - .attr("t", "1772895198") - .attr("type", "text") - .attr("edit", "7") - .build(); - - let info = client - .parse_message_info(&node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - assert_eq!( - info.edit, - EditAttribute::SenderRevoke, - "edit='7' should parse as SenderRevoke" - ); - } - - #[tokio::test] - async fn test_parse_edit_attribute_admin_revoke() { - let client = create_test_client_for_retry_with_id("edit_admin_revoke").await; - - let node = NodeBuilder::new("message") - .attr("from", "120363999999999999@g.us") - .attr("id", "TEST456") - .attr("participant", "5551234567@lid") - .attr("t", "1772895198") - .attr("type", "text") - .attr("edit", "8") - .build(); - - let info = client - .parse_message_info(&node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - assert_eq!( - info.edit, - EditAttribute::AdminRevoke, - "edit='8' should parse as AdminRevoke" - ); - } - - #[tokio::test] - async fn test_parse_edit_attribute_message_edit() { - let client = create_test_client_for_retry_with_id("edit_message_edit").await; - - let node = NodeBuilder::new("message") - .attr("from", "5551234567@s.whatsapp.net") - .attr("id", "TEST789") - .attr("t", "1772895198") - .attr("type", "text") - .attr("edit", "1") - .build(); - - let info = client - .parse_message_info(&node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - assert_eq!( - info.edit, - EditAttribute::MessageEdit, - "edit='1' should parse as MessageEdit" - ); - } - - #[tokio::test] - async fn test_parse_edit_attribute_missing() { - let client = create_test_client_for_retry_with_id("edit_missing").await; - - let node = NodeBuilder::new("message") - .attr("from", "5551234567@s.whatsapp.net") - .attr("id", "TESTABC") - .attr("t", "1772895198") - .attr("type", "text") - .build(); - - let info = client - .parse_message_info(&node.as_node_ref()) - .await - .expect("parse_message_info should succeed"); - - assert_eq!( - info.edit, - EditAttribute::Empty, - "missing edit attr should default to Empty" - ); - } - - #[tokio::test] - async fn test_revoked_message_still_retries() { - let client = create_test_client_for_retry_with_id("revoke_retry").await; - - let mut info = create_test_message_info( - "status@broadcast", - "REVOKE_MSG1", - "5551234567@s.whatsapp.net", - ); - info.edit = EditAttribute::SenderRevoke; - - // WA Web retries revoked messages the same as any other — the revoke - // protocol message contains the target ID needed to process the deletion - let info = Arc::new(info); - client.spawn_retry_receipt(&info, RetryReason::NoSession); - - // Wait for the spawned task to execute - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - - let cache_key = client - .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) - .await; - assert_eq!( - client.message_retry_counts.get(&cache_key).await, - Some(1), - "revoked message should still have retry count 1 (WA Web retries all messages)" - ); - } - - #[tokio::test] - async fn test_enc_count_preseeds_retry_cache() { - let client = create_test_client_for_retry_with_id("enc_preseed").await; - - let chat_jid: Jid = "5551234567@s.whatsapp.net".parse().unwrap(); - let msg_id = "ENC_COUNT_MSG1"; - - // Pre-seed via the same logic used in handle_incoming_message - let max_sender_retry_count: u8 = 3; - let cache_key = client - .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) - .await; - // Insert only if absent (portable alternative to moka's entry_by_ref().or_insert()) - if client.message_retry_counts.get(&cache_key).await.is_none() { - client - .message_retry_counts - .insert(cache_key.clone(), max_sender_retry_count) - .await; - } - - assert_eq!( - client.message_retry_counts.get(&cache_key).await, - Some(3), - "cache should be pre-seeded with sender retry count" - ); - } - - #[tokio::test] - async fn test_enc_no_count_cache_empty() { - let client = create_test_client_for_retry_with_id("enc_no_count").await; - - let chat_jid: Jid = "5551234567@s.whatsapp.net".parse().unwrap(); - let msg_id = "ENC_NO_COUNT_MSG1"; - - // When max_sender_retry_count is 0, no pre-seeding occurs - let max_sender_retry_count: u8 = 0; - if max_sender_retry_count > 0 { - let cache_key = client - .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) - .await; - if client.message_retry_counts.get(&cache_key).await.is_none() { - client - .message_retry_counts - .insert(cache_key, max_sender_retry_count) - .await; - } - } - - let cache_key = client - .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) - .await; - assert!( - client.message_retry_counts.get(&cache_key).await.is_none(), - "cache should be empty when no count attribute" - ); - } - - #[tokio::test] - async fn test_enc_count_does_not_overwrite_higher() { - let client = create_test_client_for_retry_with_id("enc_no_overwrite").await; - - let chat_jid: Jid = "5551234567@s.whatsapp.net".parse().unwrap(); - let msg_id = "ENC_NOOVERWRITE_MSG1"; - - let cache_key = client - .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) - .await; - - // Pre-insert a higher value - client - .message_retry_counts - .insert(cache_key.clone(), 4) - .await; - - // max(existing, incoming) should NOT overwrite with a lower value - let max_sender_retry_count: u8 = 2; - let existing = client - .message_retry_counts - .get(&cache_key) - .await - .unwrap_or(0); - if max_sender_retry_count > existing { - client - .message_retry_counts - .insert(cache_key.clone(), max_sender_retry_count) - .await; - } - - assert_eq!( - client.message_retry_counts.get(&cache_key).await, - Some(4), - "should not overwrite existing higher value" - ); - } - - #[tokio::test] - async fn test_enc_count_updates_when_sender_higher() { - let client = create_test_client_for_retry_with_id("enc_update_higher").await; - - let chat_jid: Jid = "5551234567@s.whatsapp.net".parse().unwrap(); - let msg_id = "ENC_UPDATE_MSG1"; - - let cache_key = client - .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) - .await; - - // Pre-insert a lower value - client - .message_retry_counts - .insert(cache_key.clone(), 1) - .await; - - // max(existing, incoming) SHOULD update with a higher value - let max_sender_retry_count: u8 = 3; - let existing = client - .message_retry_counts - .get(&cache_key) - .await - .unwrap_or(0); - if max_sender_retry_count > existing { - client - .message_retry_counts - .insert(cache_key.clone(), max_sender_retry_count) - .await; - } - - assert_eq!( - client.message_retry_counts.get(&cache_key).await, - Some(3), - "should update to higher sender count" - ); - } - - /// Shared helper: the OLD semaphore acquire logic that silently dropped tasks - /// on generation mismatch. Used by the bug-demonstration test. - async fn acquire_permit_old_behavior( - semaphore: &std::sync::Mutex<Arc<async_lock::Semaphore>>, - generation: &portable_atomic::AtomicU64, - ) -> bool { - use std::sync::atomic::Ordering; - let (snap_gen, snap_sem) = { - let guard = semaphore.lock().unwrap(); - (generation.load(Ordering::SeqCst), guard.clone()) - }; - let _permit = snap_sem.acquire_arc().await; - // OLD: if generation changed, silently return false (message lost) - snap_gen == generation.load(Ordering::SeqCst) - } - - /// Shared helper: the FIXED semaphore acquire logic that re-acquires from the - /// new semaphore on generation mismatch. Mirrors the production code in - /// handle_incoming_message. - async fn acquire_permit_with_reacquire( - semaphore: &std::sync::Mutex<Arc<async_lock::Semaphore>>, - generation: &portable_atomic::AtomicU64, - ) { - use std::sync::atomic::Ordering; - loop { - let (snap_gen, snap_sem) = { - let guard = semaphore.lock().unwrap(); - (generation.load(Ordering::SeqCst), guard.clone()) - }; - let permit = snap_sem.acquire_arc().await; - if snap_gen == generation.load(Ordering::SeqCst) { - drop(permit); - break; - } - drop(permit); - } - } - - /// Demonstrates the bug: the OLD code silently dropped tasks when generation changed. - #[tokio::test] - async fn test_old_behavior_drops_tasks_on_generation_swap() { - use portable_atomic::AtomicU64; - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - - let semaphore = Arc::new(std::sync::Mutex::new(Arc::new(async_lock::Semaphore::new( - 1, - )))); - let generation = Arc::new(AtomicU64::new(0)); - let completed = Arc::new(AtomicUsize::new(0)); - let ready = Arc::new(AtomicUsize::new(0)); - - let blocker_sem = semaphore.lock().unwrap().clone(); - let blocker_permit = blocker_sem.acquire_arc().await; - - let num_waiters: usize = 8; - let mut handles = Vec::new(); - - for _ in 0..num_waiters { - let sem = semaphore.clone(); - let gen_counter = generation.clone(); - let done = completed.clone(); - let ready_counter = ready.clone(); - - handles.push(tokio::spawn(async move { - // Signal readiness before blocking on semaphore - ready_counter.fetch_add(1, Ordering::SeqCst); - if acquire_permit_old_behavior(&sem, &gen_counter).await { - done.fetch_add(1, Ordering::SeqCst); - } - })); - } - - // Wait until all waiters have signaled readiness (about to block on semaphore) - while ready.load(Ordering::SeqCst) < num_waiters { - tokio::task::yield_now().await; - } - - // Swap semaphore — triggers the bug - { - let mut guard = semaphore.lock().unwrap(); - *guard = Arc::new(async_lock::Semaphore::new(64)); - generation.fetch_add(1, Ordering::SeqCst); - } - - drop(blocker_permit); - - for handle in handles { - let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await; - assert!(result.is_ok(), "Waiter task timed out"); - result.unwrap().unwrap(); - } - - let done = completed.load(Ordering::SeqCst); - assert!( - done < num_waiters, - "Bug demonstration: expected tasks to be dropped, but all {} completed", - num_waiters - ); - } - - /// Verifies the fix: re-acquire loop ensures NO tasks are dropped on generation swap. - #[tokio::test] - async fn test_semaphore_generation_swap_does_not_drop_tasks() { - use portable_atomic::AtomicU64; - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - - let semaphore = Arc::new(std::sync::Mutex::new(Arc::new(async_lock::Semaphore::new( - 1, - )))); - let generation = Arc::new(AtomicU64::new(0)); - let completed = Arc::new(AtomicUsize::new(0)); - let ready = Arc::new(AtomicUsize::new(0)); - - let blocker_sem = semaphore.lock().unwrap().clone(); - let blocker_permit = blocker_sem.acquire_arc().await; - - let num_waiters: usize = 8; - let mut handles = Vec::new(); - - for _ in 0..num_waiters { - let sem = semaphore.clone(); - let gen_counter = generation.clone(); - let done = completed.clone(); - let ready_counter = ready.clone(); - - handles.push(tokio::spawn(async move { - ready_counter.fetch_add(1, Ordering::SeqCst); - acquire_permit_with_reacquire(&sem, &gen_counter).await; - done.fetch_add(1, Ordering::SeqCst); - })); - } - - // Wait until all waiters have signaled readiness - while ready.load(Ordering::SeqCst) < num_waiters { - tokio::task::yield_now().await; - } - - // Swap semaphore (simulates offline sync completion) - { - let mut guard = semaphore.lock().unwrap(); - *guard = Arc::new(async_lock::Semaphore::new(64)); - generation.fetch_add(1, Ordering::SeqCst); - } - - drop(blocker_permit); - - for handle in handles { - let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await; - assert!( - result.is_ok(), - "Waiter task timed out — likely silently dropped by generation check" - ); - result.unwrap().unwrap(); - } - - assert_eq!( - completed.load(Ordering::SeqCst), - num_waiters, - "All {} waiter tasks should complete, but only {} did. \ - Tasks were silently dropped during semaphore generation swap.", - num_waiters, - completed.load(Ordering::SeqCst) - ); - } - - // Dispatch ordering, per-id dedup, and PDO eligibility for - // UndecryptableMessage. Regressing any of these re-opens data loss bugs - // observed in production. - - use crate::types::events::DecryptFailMode; - use wacore::types::events::{Event, EventHandler}; - - #[derive(Default)] - struct EventRecorder { - events: std::sync::Mutex<Vec<Arc<Event>>>, - } - - impl EventHandler for EventRecorder { - fn handle_event(&self, event: Arc<Event>) { - self.events.lock().unwrap().push(event); - } - } - - impl EventRecorder { - fn undecryptable(&self) -> Vec<Arc<Event>> { - self.events - .lock() - .unwrap() - .iter() - .filter(|e| matches!(&***e, Event::UndecryptableMessage(_))) - .cloned() - .collect() - } - - /// Count of `UndecryptableMessage` events marked as the "stub" - /// variant (`is_unavailable=true`, `UnavailableType::ViewOnce`) — - /// i.e. the branch that routes to PDO instead of falling through to - /// decrypt. - fn view_once_unavailable_count(&self) -> usize { - use crate::types::events::UnavailableType; - self.events - .lock() - .unwrap() - .iter() - .filter(|e| { - matches!( - &***e, - Event::UndecryptableMessage(u) - if u.is_unavailable - && matches!(u.unavailable_type, UnavailableType::ViewOnce) - ) - }) - .count() - } - } - - fn build_unavailable_stanza(sender: &str, msg_id: &str, with_enc: bool) -> Arc<OwnedNodeRef> { - let t = wacore::time::now_secs().to_string(); - let unavailable = NodeBuilder::new("unavailable") - .attr("type", "view_once") - .build(); - let children = if with_enc { - vec![ - unavailable, - NodeBuilder::new("enc") - .attr("type", "msg") - .attr("v", "2") - .bytes(vec![0xDE, 0xAD, 0xBE, 0xEF]) - .build(), - ] - } else { - vec![unavailable] - }; - node_to_arc( - NodeBuilder::new("message") - .attr("from", sender) - .attr("id", msg_id) - .attr("t", &t) - .attr("type", "media") - .children(children) - .build(), - ) - } - - /// Locks the dispatch ordering: consumers must see the event before any - /// retry/PDO side effects, otherwise a late subscriber misses the failure. - #[tokio::test] - async fn test_undecryptable_fires_before_retry_task() { - let client = create_test_client_for_retry_with_id("undec_sync").await; - let recorder = Arc::new(EventRecorder::default()); - client.register_handler(recorder.clone()); - - let info = Arc::new(create_test_message_info( - "5511999998888@s.whatsapp.net", - "MSG_SYNC_1", - "5511777776666@s.whatsapp.net", - )); - - let cache_key = client - .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) - .await; - - assert!(recorder.undecryptable().is_empty()); - assert!(client.message_retry_counts.get(&cache_key).await.is_none()); - - let _ = client - .handle_decrypt_failure(&info, RetryReason::InvalidKeyId, DecryptFailMode::Show) - .await; - - assert_eq!( - recorder.undecryptable().len(), - 1, - "UndecryptableMessage dispatched inside handle_decrypt_failure", - ); - assert!( - client.message_retry_counts.get(&cache_key).await.is_none(), - "retry task has not progressed yet", - ); - - tokio::time::sleep(tokio::time::Duration::from_millis(150)).await; - assert_eq!( - client.message_retry_counts.get(&cache_key).await, - Some(1), - "retry task runs after the dispatch", - ); - } - - /// Atomic dedup under concurrency: 32 parallel callers for the same id - /// must produce exactly one event. Catches regressions where the dedup - /// would slip back to a non-atomic get-then-insert pair. - #[tokio::test] - async fn test_undecryptable_dedup_is_atomic() { - let client = create_test_client_for_retry_with_id("undec_atomic").await; - let recorder = Arc::new(EventRecorder::default()); - client.register_handler(recorder.clone()); - - let info = Arc::new(create_test_message_info( - "5511999998888@s.whatsapp.net", - "ATOMIC_MSG_1", - "5511777776666@s.whatsapp.net", - )); - - let mut handles = Vec::with_capacity(32); - for _ in 0..32 { - let c = Arc::clone(&client); - let i = Arc::clone(&info); - handles.push(tokio::spawn(async move { - c.handle_decrypt_failure(&i, RetryReason::InvalidKeyId, DecryptFailMode::Show) - .await; - })); - } - for h in handles { - h.await.unwrap(); - } - - assert_eq!( - recorder.undecryptable().len(), - 1, - "32 concurrent callers must collapse to one UndecryptableMessage", - ); - } - - /// Server resends of the same id must not surface a duplicate event — - /// would otherwise show the user the same failure twice. - #[tokio::test] - async fn test_undecryptable_deduped_across_resends() { - let client = create_test_client_for_retry_with_id("undec_double").await; - let recorder = Arc::new(EventRecorder::default()); - client.register_handler(recorder.clone()); - - let info = Arc::new(create_test_message_info( - "5511999998888@s.whatsapp.net", - "3AD01881AA95F7D81070", - "85010891714716@lid", - )); - - let _ = client - .handle_decrypt_failure(&info, RetryReason::InvalidKeyId, DecryptFailMode::Show) - .await; - let _ = client - .handle_decrypt_failure(&info, RetryReason::InvalidKeyId, DecryptFailMode::Show) - .await; - - let events = recorder.undecryptable(); - assert_eq!( - events.len(), - 1, - "same message id fires UndecryptableMessage only once", - ); - if let Event::UndecryptableMessage(event) = &*events[0] { - assert_eq!(event.info.id, info.id); - } else { - panic!("event was not UndecryptableMessage"); - } - } - - /// Status posts must flow through PDO — excluding them drops any - /// InvalidPreKeyId status permanently (WA Web recovers them). - #[tokio::test] - async fn test_pdo_armed_for_status_broadcast() { - let client = create_test_client_for_retry_with_id("pdo_status").await; - - let info = Arc::new(create_test_message_info( - "status@broadcast", - "STATUS_MSG_1", - "5511777776666@s.whatsapp.net", - )); - - assert_eq!(info.source.chat.server, wacore_binary::Server::Broadcast); - - client.run_pdo_request(&info).await; - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - } - - /// Broadcast lists share the same code path; locks the guard for both. - #[tokio::test] - async fn test_pdo_armed_for_any_broadcast_chat() { - let client = create_test_client_for_retry_with_id("pdo_bcast_list").await; - - let info = Arc::new(create_test_message_info( - "12345@broadcast", - "BCAST_LIST_MSG_1", - "5511777776666@s.whatsapp.net", - )); - - assert_eq!(info.source.chat.server, wacore_binary::Server::Broadcast); - - client.run_pdo_request(&info).await; - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - } - - #[tokio::test] - async fn test_pdo_armed_for_one_on_one() { - let client = create_test_client_for_retry_with_id("pdo_dm").await; - - let info = Arc::new(create_test_message_info( - "85010891714716@lid", - "DM_MSG_1", - "85010891714716@lid", - )); - - assert_ne!(info.source.chat.server, wacore_binary::Server::Broadcast); - - client.run_pdo_request(&info).await; - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - } - - /// fromMe messages fanned out to a linked device can still fail decrypt - /// on the receiver side; PDO is the only recovery path for them. - #[tokio::test] - async fn test_pdo_armed_for_from_me() { - let client = create_test_client_for_retry_with_id("pdo_from_me").await; - - // When fromMe is true the sender is the user's own JID, not a peer. - let own_jid = "5511999998888@s.whatsapp.net"; - let mut info = create_test_message_info("85010891714716@lid", "FROM_ME_MSG_1", own_jid); - info.source.is_from_me = true; - let info = Arc::new(info); - - client.run_pdo_request(&info).await; - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - } - - /// Stops offline-sync / reconnect tails from flooding the phone with - /// resend requests for old messages the user likely no longer cares about. - #[tokio::test] - async fn test_pdo_skipped_for_ancient_messages() { - use wacore::types::message::ChatMessageId; - - let client = create_test_client_for_retry_with_id("pdo_age").await; - - let mut info = - create_test_message_info("85010891714716@lid", "ANCIENT_MSG_1", "85010891714716@lid"); - info.timestamp = wacore::time::now_utc() - chrono::Duration::days(30); - let info = Arc::new(info); - - let cache_key = ChatMessageId::new(info.source.chat.clone(), info.id.clone()); - - client.run_pdo_request(&info).await; - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - - assert!( - client.pdo_pending_requests.get(&cache_key).await.is_none(), - "messages older than 14 days must not register a PDO entry", - ); - } - - /// Boundary check: age of 14d plus a minute must reject (WA Web uses - /// seconds, not days, so 14d1m is already over the limit). Catches a - /// `num_days()` truncation that would otherwise accept this message. - #[tokio::test] - async fn test_pdo_rejects_just_past_14d_boundary() { - use wacore::types::message::ChatMessageId; - - let client = create_test_client_for_retry_with_id("pdo_boundary").await; - - let mut info = - create_test_message_info("85010891714716@lid", "BOUNDARY_MSG_1", "85010891714716@lid"); - info.timestamp = - wacore::time::now_utc() - chrono::Duration::days(14) - chrono::Duration::minutes(1); - let info = Arc::new(info); - - let cache_key = ChatMessageId::new(info.source.chat.clone(), info.id.clone()); - - client.run_pdo_request(&info).await; - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - - assert!( - client.pdo_pending_requests.get(&cache_key).await.is_none(), - "14d+1m must be over the limit, matching WA Web's seconds-based check", - ); - } - - /// Server-trusted companions (Android-class `DeviceProps.PlatformType`) - /// receive `<unavailable>` as a marker alongside `<enc>`. The cipher - /// must still be decrypted — skipping would discard content the server - /// specifically released for this companion. Decrypt eventually fails - /// on the garbage payload, but via the normal decrypt-failure path, - /// not the `ViewOnce` short-circuit. - #[tokio::test] - async fn test_unavailable_with_enc_skips_unavailable_shortcut() { - let client = create_test_client_for_retry_with_id("unavailable_with_enc").await; - let recorder = Arc::new(EventRecorder::default()); - client.register_handler(recorder.clone()); - - let node = - build_unavailable_stanza("5511777776666@s.whatsapp.net", "UNAV_WITH_ENC_1", true); - client.clone().handle_incoming_message(node).await; - - assert_eq!( - recorder.view_once_unavailable_count(), - 0, - "<unavailable> alongside <enc> must fall through to decrypt, \ - not emit a ViewOnce UndecryptableMessage", - ); - } - - /// Untrusted companions (web-class `PlatformType`) get the bare stub — - /// `<unavailable>` without `<enc>`. That path must still emit a - /// `ViewOnce` `UndecryptableMessage` so consumers surface the failure - /// while the phone relays via PDO. - #[tokio::test] - async fn test_unavailable_without_enc_dispatches_view_once_event() { - let client = create_test_client_for_retry_with_id("unavailable_stub").await; - let recorder = Arc::new(EventRecorder::default()); - client.register_handler(recorder.clone()); - - let node = build_unavailable_stanza("5511777776666@s.whatsapp.net", "UNAV_STUB_1", false); - client.clone().handle_incoming_message(node).await; - - assert_eq!( - recorder.view_once_unavailable_count(), - 1, - "bare <unavailable> stub must dispatch exactly one ViewOnce UndecryptableMessage", - ); - } - - /// The event struct has no "recovery pending" flag, so consumers cannot - /// wait for a PDO outcome before surfacing failure — adding a field - /// here forces a conscious UX decision. - #[test] - fn test_undecryptable_event_has_no_pending_pdo_hint() { - use crate::types::events::{UnavailableType, UndecryptableMessage}; - - let info = Arc::new(create_test_message_info( - "5511999998888@s.whatsapp.net", - "SHAPE_MSG", - "5511777776666@s.whatsapp.net", - )); - let event = UndecryptableMessage { - info, - is_unavailable: false, - unavailable_type: UnavailableType::Unknown, - decrypt_fail_mode: DecryptFailMode::Show, - }; - - let _ = ( - &event.info, - &event.is_unavailable, - &event.unavailable_type, - &event.decrypt_fail_mode, - ); - } - - /// Seed `device.pn` so `send_nack` clears its `get_pn()` guard. - async fn seed_test_pn(client: &Arc<Client>) { - use crate::store::commands::DeviceCommand; - client - .persistence_manager - .process_command(DeviceCommand::SetId(Some( - "5511000000001:0@s.whatsapp.net" - .parse() - .expect("test PN should parse"), - ))) - .await; - } - - /// Build a Client wired to a CapturingMockTransport + a noise socket so - /// `send_node` reaches the wire. Returns the transport so the caller can - /// inspect captured frames. - async fn capturing_client( - test_id: &str, - ) -> ( - Arc<Client>, - Arc<crate::transport::mock::CapturingMockTransport>, - ) { - use crate::socket::NoiseSocket; - use crate::store::SqliteStore; - use crate::store::persistence_manager::PersistenceManager; - use crate::transport::mock::CapturingMockTransportFactory; - use portable_atomic::AtomicU64; - use std::sync::atomic::Ordering; - use wacore::handshake::NoiseCipher; - - static COUNTER: AtomicU64 = AtomicU64::new(0); - let unique_id = COUNTER.fetch_add(1, Ordering::SeqCst); - let db_name = format!( - "file:memdb_capt_{}_{}_{}?mode=memory&cache=shared", - test_id, - unique_id, - std::process::id() - ); - - let backend = Arc::new( - SqliteStore::new(&db_name) - .await - .expect("test backend should initialize"), - ); - let pm = Arc::new( - PersistenceManager::new(backend) - .await - .expect("persistence manager should initialize"), - ); - let factory = CapturingMockTransportFactory::new(); - let transport = factory.transport(); - let (client, _sync_rx) = Client::new( - Arc::new(crate::runtime_impl::TokioRuntime), - pm, - Arc::new(factory), - Arc::new(MockHttpClient), - None, - ) - .await; - - let key = [0u8; 32]; - let write_key = NoiseCipher::new(&key).expect("32-byte key"); - let read_key = NoiseCipher::new(&key).expect("32-byte key"); - let noise_socket = NoiseSocket::new( - Arc::new(crate::runtime_impl::TokioRuntime), - transport.clone() as Arc<dyn crate::transport::Transport>, - write_key, - read_key, - ); - // send_node only needs noise_socket Some; is_connected is read by - // other layers but not on this path. - *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); - seed_test_pn(&client).await; - (client, transport) - } - - /// Regression: a malformed pkmsg used to fall through silently. Now - /// it dispatches the consumer event AND emits a nack on the wire so - /// the server stops retransmitting. - #[tokio::test] - async fn pkmsg_parse_error_dispatches_parsing_error_nack() { - use crate::types::events::DecryptFailMode; - use wacore::message_processing::EncType; - - let (client, transport) = capturing_client("pkmsg_parse_nack").await; - let info = Arc::new(create_test_message_info( - "5511999998888@s.whatsapp.net", - "REGRESSION_PKMSG_PARSE", - "5511777776666@s.whatsapp.net", - )); - let sender_jid: Jid = info.source.sender.clone(); - - // 1-byte ciphertext is a guaranteed parse failure. - let bad_payload = EncPayload { - ciphertext: bytes::Bytes::from_static(&[0xFF]), - enc_type: EncType::PreKeyMessage, - padding_version: 2, - }; - - let outcome = client - .process_session_enc_batch(&[bad_payload], &info, &sender_jid, DecryptFailMode::Show) - .await; - - assert!(!outcome.decrypted); - assert!(!outcome.duplicate); - assert!(outcome.undecryptable); - assert!(outcome.had_failure); - - // spawn_nack is detached; give it a tick to flush through the - // noise_socket sender_task to our CapturingMockTransport. - for _ in 0..40 { - if !transport.sent().is_empty() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let sent = transport.sent(); - assert!( - !sent.is_empty(), - "spawn_nack must produce at least one outbound frame on the wire" - ); - } - - #[tokio::test] - async fn signal_message_parse_error_dispatches_parsing_error_nack() { - use crate::types::events::DecryptFailMode; - use wacore::message_processing::EncType; - - let (client, transport) = capturing_client("sig_parse_nack").await; - let info = Arc::new(create_test_message_info( - "5511999998888@s.whatsapp.net", - "REGRESSION_SIG_PARSE", - "5511777776666@s.whatsapp.net", - )); - let sender_jid: Jid = info.source.sender.clone(); - - let bad_payload = EncPayload { - ciphertext: bytes::Bytes::from_static(&[0xFF]), - enc_type: EncType::Message, - padding_version: 2, - }; - - let outcome = client - .process_session_enc_batch(&[bad_payload], &info, &sender_jid, DecryptFailMode::Show) - .await; - - assert!(outcome.undecryptable); - assert!(outcome.had_failure); - - for _ in 0..40 { - if !transport.sent().is_empty() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert!( - !transport.sent().is_empty(), - "spawn_nack must produce at least one outbound frame on the wire" - ); - } - - #[test] - fn test_decrypt_fail_log_level_gated_on_hide() { - use crate::types::events::DecryptFailMode; - assert_eq!( - decrypt_fail_log_level(DecryptFailMode::Hide), - log::Level::Debug - ); - assert_eq!( - decrypt_fail_log_level(DecryptFailMode::Show), - log::Level::Warn - ); - } - - /// Decrypt one captured noise frame (zero-key, counter-based, empty AAD) to - /// its marshalled node bytes; strips the 3-byte frame header. - fn decode_frame(index: usize, frame: &[u8]) -> Option<Vec<u8>> { - use wacore::handshake::NoiseCipher; - if frame.len() <= 3 { - return None; - } - let cipher = NoiseCipher::new(&[0u8; 32]).expect("32-byte key"); - let mut buf = frame[3..].to_vec(); - cipher - .decrypt_in_place_with_counter(index as u32, &mut buf) - .ok()?; - (!buf.is_empty()).then_some(buf) - } - - /// First `<ack class="message">` on the wire as `(to, recipient)`. - fn find_message_ack(frames: &[bytes::Bytes]) -> Option<(String, Option<String>)> { - for (i, frame) in frames.iter().enumerate() { - let Some(buf) = decode_frame(i, frame) else { - continue; - }; - let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { - continue; - }; - if node.tag.as_ref() == "ack" - && node - .get_attr("class") - .is_some_and(|v| v.as_str() == "message") - && node.get_attr("error").is_none() - && let Some(to) = node.get_attr("to") - { - let recipient = node.get_attr("recipient").map(|v| v.as_str().to_string()); - return Some((to.as_str().to_string(), recipient)); - } - } - None - } - - /// First `<receipt>` on the wire for `id` as `(to, type, recipient)`. - fn find_receipt( - frames: &[bytes::Bytes], - id: &str, - ) -> Option<(String, Option<String>, Option<String>)> { - for (i, frame) in frames.iter().enumerate() { - let Some(buf) = decode_frame(i, frame) else { - continue; - }; - let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { - continue; - }; - if node.tag.as_ref() == "receipt" - && node.get_attr("id").is_some_and(|v| v.as_str() == id) - && let Some(to) = node.get_attr("to") - { - let typ = node.get_attr("type").map(|v| v.as_str().to_string()); - let recipient = node.get_attr("recipient").map(|v| v.as_str().to_string()); - return Some((to.as_str().to_string(), typ, recipient)); - } - } - None - } - - #[derive(Debug)] - struct SentReceipt { - to: String, - typ: Option<String>, - recipient: Option<String>, - participant: Option<String>, - context: Option<String>, - } - - fn find_receipt_details(frames: &[bytes::Bytes], id: &str) -> Option<SentReceipt> { - for (i, frame) in frames.iter().enumerate() { - let Some(buf) = decode_frame(i, frame) else { - continue; - }; - let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { - continue; - }; - if node.tag.as_ref() == "receipt" - && node.get_attr("id").is_some_and(|v| v.as_str() == id) - && let Some(to) = node.get_attr("to") - { - return Some(SentReceipt { - to: to.as_str().to_string(), - typ: node.get_attr("type").map(|v| v.as_str().to_string()), - recipient: node.get_attr("recipient").map(|v| v.as_str().to_string()), - participant: node.get_attr("participant").map(|v| v.as_str().to_string()), - context: node.get_attr("context").map(|v| v.as_str().to_string()), - }); - } - } - None - } - - #[derive(Debug)] - struct SentMessageAck { - to: String, - participant: Option<String>, - recipient: Option<String>, - } - - fn find_message_ack_for(frames: &[bytes::Bytes], id: &str) -> Option<SentMessageAck> { - for (i, frame) in frames.iter().enumerate() { - let Some(buf) = decode_frame(i, frame) else { - continue; - }; - let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { - continue; - }; - if node.tag.as_ref() == "ack" - && node - .get_attr("class") - .is_some_and(|v| v.as_str() == "message") - && node.get_attr("error").is_none() - && node.get_attr("id").is_some_and(|v| v.as_str() == id) - && let Some(to) = node.get_attr("to") - { - return Some(SentMessageAck { - to: to.as_str().to_string(), - participant: node.get_attr("participant").map(|v| v.as_str().to_string()), - recipient: node.get_attr("recipient").map(|v| v.as_str().to_string()), - }); - } - } - None - } - - /// Count delivery `<receipt>` (anything but type="retry") on the wire for `id`. - fn delivery_receipts_for(frames: &[bytes::Bytes], id: &str) -> usize { - let mut count = 0; - for (i, frame) in frames.iter().enumerate() { - let Some(buf) = decode_frame(i, frame) else { - continue; - }; - let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { - continue; - }; - if node.tag.as_ref() == "receipt" - && node.get_attr("id").is_some_and(|v| v.as_str() == id) - && node - .get_attr("type") - .as_ref() - .map(|v| v.as_str()) - .as_deref() - != Some("retry") - { - count += 1; - } - } - count - } - - fn message_acks_for(frames: &[bytes::Bytes], id: &str) -> usize { - let mut count = 0; - for (i, frame) in frames.iter().enumerate() { - let Some(buf) = decode_frame(i, frame) else { - continue; - }; - let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { - continue; - }; - if node.tag.as_ref() == "ack" - && node - .get_attr("class") - .is_some_and(|v| v.as_str() == "message") - && node.get_attr("error").is_none() - && node.get_attr("id").is_some_and(|v| v.as_str() == id) - { - count += 1; - } - } - count - } - - fn sender_receipts_for(frames: &[bytes::Bytes], id: &str) -> usize { - let mut count = 0; - for (i, frame) in frames.iter().enumerate() { - let Some(buf) = decode_frame(i, frame) else { - continue; - }; - let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { - continue; - }; - if node.tag.as_ref() == "receipt" - && node.get_attr("id").is_some_and(|v| v.as_str() == id) - && node - .get_attr("type") - .is_some_and(|v| v.as_str() == "sender") - { - count += 1; - } - } - count - } - - fn confirmations_for(frames: &[bytes::Bytes], id: &str) -> usize { - delivery_receipts_for(frames, id) + message_acks_for(frames, id) - } - - async fn wait_for_confirmations( - transport: &crate::transport::mock::CapturingMockTransport, - id: &str, - expected: usize, - ) -> usize { - let mut count = 0; - for _ in 0..80 { - count = confirmations_for(&transport.sent(), id); - if count >= expected { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - count - } - - async fn assert_exactly_one_confirmation( - transport: &crate::transport::mock::CapturingMockTransport, - id: &str, - ) { - let count = wait_for_confirmations(transport, id, 1).await; - assert_eq!(count, 1, "message {id} must be confirmed exactly once"); - for _ in 0..5 { - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - assert_eq!( - confirmations_for(&transport.sent(), id), - 1, - "message {id} must not get a late second confirmation" - ); - } - } - - /// A stanza that fails to decrypt must emit a transport `<ack class="message">` - /// (else the server replays it on every reconnect forever), addressed to the - /// original `from` echoing `recipient`. Uses `Hide` (the production reactions - /// carried `decrypt-fail="hide"`) to also guard that hide does not suppress - /// the ack. BadMac so the retry carries no keys (no device account needed). - #[tokio::test] - async fn decrypt_failure_emits_transport_ack() { - let (client, transport) = capturing_client("decrypt_fail_ack").await; - - let sender: Jid = "236395184570386@lid".parse().expect("sender JID"); - let recipient: Jid = "156535032389744@lid".parse().expect("recipient JID"); - let info = Arc::new(MessageInfo { - id: "AC055553E56A2C12DE592DAD6353C477".to_string(), - source: crate::types::message::MessageSource { - sender: sender.clone(), - chat: recipient.clone(), - recipient: Some(recipient.clone()), - ..Default::default() - }, - ..Default::default() - }); - - client - .handle_decrypt_failure( - &info, - RetryReason::BadMac, - crate::types::events::DecryptFailMode::Hide, - ) - .await; - - // retry + ack are detached spawns; poll the wire until the ack appears. - let mut found = None; - for _ in 0..80 { - if let Some(a) = find_message_ack(&transport.sent()) { - found = Some(a); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let (to, recipient_attr) = found.expect( - "decrypt failure must emit a transport <ack class=message> \ - (else the server redelivers the stanza forever)", - ); - assert_eq!( - to, "236395184570386@lid", - "ack `to` must be the original `from` (own LID), not the chat" - ); - assert_eq!( - recipient_attr.as_deref(), - Some("156535032389744@lid"), - "ack must echo `recipient` for own-account fan-out" - ); - } - - /// Regression for the bot self-fanout loop on the DECRYPT-FAILURE path - /// (BadMac/NoSession): a self-fanout we cannot decrypt must be cleared with - /// a `<receipt type="sender">`, NOT a bare transport `<ack>` (ignored by the - /// server) nor a retry-to-self (futile). Once stuck in the loop the local - /// counter advances past the duplicate state, so this BadMac path is what - /// actually fires for an already-affected account. - #[tokio::test] - async fn self_fanout_decrypt_failure_acked_via_sender_receipt() { - let (client, transport) = capturing_client("self_fanout_badmac").await; - let info = Arc::new(MessageInfo { - id: "AC00000000000000000000000000BEEF".to_string(), - source: crate::types::message::MessageSource { - sender: "100000000000001@lid".parse().expect("sender"), - chat: "200000000000002@bot".parse().expect("chat"), - recipient: Some("200000000000002@bot".parse().expect("recipient")), - is_from_me: true, - ..Default::default() - }, - ..Default::default() - }); - - client - .handle_decrypt_failure( - &info, - RetryReason::BadMac, - crate::types::events::DecryptFailMode::Hide, - ) - .await; - - let mut found = None; - for _ in 0..80 { - if let Some(r) = find_receipt(&transport.sent(), "AC00000000000000000000000000BEEF") { - found = Some(r); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let (to, typ, recipient) = found - .expect("self-fanout decrypt failure must emit a sender <receipt> to drain the queue"); - assert_eq!(to, "100000000000001@lid"); - assert_eq!(typ.as_deref(), Some("sender")); - assert_eq!(recipient.as_deref(), Some("200000000000002@bot")); - - let sent = transport.sent(); - assert!( - find_message_ack(&sent).is_none(), - "must not emit the bare <ack> the server ignores" - ); - let mut saw_retry = false; - for (i, frame) in sent.iter().enumerate() { - if let Some(buf) = decode_frame(i, frame) - && let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) - && node.tag.as_ref() == "receipt" - && node.get_attr("type").is_some_and(|v| { - v.as_str() == crate::types::presence::ReceiptType::Retry.as_wire_str() - }) - { - saw_retry = true; - } - } - assert!( - !saw_retry, - "must not retry our own undecryptable fanout to ourselves" - ); - } - - /// Consistency with the success/duplicate path: a bot-authored own DM in a - /// non-bot chat (sender on `@bot`, user chat) must NOT take the sender - /// receipt on the decrypt-failure path either; it stays on the - /// bot-invoke-response bare-ack path (WA Web `!chat.isBot() && - /// author.isBot()`), matching ack_received_message and the locked - /// own_bot_author_dm_acks_not_sender_receipt test. - #[tokio::test] - async fn bot_author_self_fanout_decrypt_failure_not_sender_receipt() { - let (client, transport) = capturing_client("bot_author_badmac").await; - let info = Arc::new(MessageInfo { - id: "OWNBOTFAIL1".to_string(), - source: crate::types::message::MessageSource { - sender: "100000000000002@bot".parse().expect("sender"), - chat: "300000000000003@lid".parse().expect("chat"), - recipient: Some("300000000000003@lid".parse().expect("recipient")), - is_from_me: true, - ..Default::default() - }, - ..Default::default() - }); - - client - .handle_decrypt_failure( - &info, - RetryReason::BadMac, - crate::types::events::DecryptFailMode::Hide, - ) - .await; - - // Positive: the message IS cleared, via the bot-invoke-response bare - // <ack class="message"> (the retry-to-self is bot-skipped, so the - // transport ack follows), proving we took the ack path, not a no-op. - let mut found_ack = false; - for _ in 0..80 { - if find_message_ack(&transport.sent()).is_some() { - found_ack = true; - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert!( - found_ack, - "bot-authored own DM must still be transport-acked with a bare <ack class=message>" - ); - - // Negative settle: it must NEVER produce a sender receipt on the failure - // path (that would diverge from WA Web's bot-invoke ack and contradict - // the success-path ordering). - for _ in 0..5 { - assert!( - find_receipt(&transport.sent(), "OWNBOTFAIL1").is_none(), - "bot-authored own DM must not be cleared with a sender <receipt> on decrypt failure" - ); - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - } - - /// If the resend request fails to send, the stanza must NOT be acked, so the - /// server keeps it queued for another try. Here NoSession needs keys, which - /// need a device account this harness lacks, so send_retry_receipt errors. - #[tokio::test] - async fn decrypt_failure_does_not_ack_when_retry_send_fails() { - let (client, transport) = capturing_client("retry_fail_no_ack").await; - let sender: Jid = "5511777776666@s.whatsapp.net".parse().expect("sender"); - let info = Arc::new(MessageInfo { - id: "NOACK1".to_string(), - source: crate::types::message::MessageSource { - sender: sender.clone(), - chat: sender.clone(), - ..Default::default() - }, - ..Default::default() - }); - client - .handle_decrypt_failure( - &info, - RetryReason::NoSession, - crate::types::events::DecryptFailMode::Show, - ) - .await; - tokio::time::sleep(std::time::Duration::from_millis(150)).await; - assert!( - find_message_ack(&transport.sent()).is_none(), - "must not ack when the resend request failed to send" - ); - } - - /// The retry receipt must be sent before the transport ack (one ordered - /// flushed task), so a disconnect mid-flush can never clear the stanza from - /// the offline queue without the sender having received a resend request. - #[tokio::test] - async fn decrypt_failure_sends_retry_before_ack() { - let (client, transport) = capturing_client("retry_before_ack").await; - let sender: Jid = "5511777776666@s.whatsapp.net".parse().expect("sender"); - let info = Arc::new(MessageInfo { - id: "RBA1".to_string(), - source: crate::types::message::MessageSource { - sender: sender.clone(), - chat: sender.clone(), - ..Default::default() - }, - ..Default::default() - }); - // BadMac (not NoSession) so the retry receipt carries no keys and needs - // no device account in this harness. - client - .handle_decrypt_failure( - &info, - RetryReason::BadMac, - crate::types::events::DecryptFailMode::Show, - ) - .await; - - let find = |tag: &str, retry: bool| -> Option<usize> { - let frames = transport.sent(); - for (i, frame) in frames.iter().enumerate() { - let Some(buf) = decode_frame(i, frame) else { - continue; - }; - let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { - continue; - }; - let is_retry = node.get_attr("type").is_some_and(|v| v.as_str() == "retry"); - if node.tag.as_ref() == tag && is_retry == retry { - return Some(i); - } - } - None - }; - - let mut retry_idx = None; - let mut ack_idx = None; - for _ in 0..80 { - retry_idx = find("receipt", true); - ack_idx = find("ack", false); - if retry_idx.is_some() && ack_idx.is_some() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let retry_idx = retry_idx.expect("retry receipt must be sent"); - let ack_idx = ack_idx.expect("transport ack must be sent"); - assert!( - retry_idx < ack_idx, - "retry receipt (frame {retry_idx}) must be sent before the ack (frame {ack_idx})" - ); - } - - /// status@broadcast is already acked by the `should_ack` gate post-dispatch, - /// so the decrypt-failure path must NOT emit a second transport ack - /// (whatsmeow/WA Web send exactly one per message). The retry receipt still - /// goes out. - #[tokio::test] - async fn status_broadcast_decrypt_failure_acks_to_chat() { - let (client, transport) = capturing_client("status_fail_ack").await; - let info = Arc::new(MessageInfo { - id: "STATUSMSGID".to_string(), - source: crate::types::message::MessageSource { - sender: "236395184570386@lid".parse().expect("sender"), - chat: "status@broadcast".parse().expect("status chat"), - is_group: true, - ..Default::default() - }, - ..Default::default() - }); - - client - .handle_decrypt_failure( - &info, - RetryReason::BadMac, - crate::types::events::DecryptFailMode::Show, - ) - .await; - - // status failures are acked from the flushed task (not just the detached - // should_ack gate), so the ack survives a disconnect mid-flush. - let mut found = None; - for _ in 0..80 { - if let Some(a) = find_message_ack(&transport.sent()) { - found = Some(a); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let (to, _) = found.expect("status failure must emit a flushed transport ack"); - assert_eq!( - to, "status@broadcast", - "status ack `to` must be the status chat" - ); - } - - /// Run a single session ciphertext through the full classify->process path. - async fn process_session_ct( - client: &Arc<Client>, - sender: &Jid, - id: &str, - ct: &wacore::libsignal::protocol::CiphertextMessage, - ) { - use wacore::libsignal::protocol::CiphertextMessage; - let (enc_type, bytes) = match ct { - CiphertextMessage::PreKeySignalMessage(m) => ("pkmsg", m.serialized().to_vec()), - CiphertextMessage::SignalMessage(m) => ("msg", m.serialized().to_vec()), - _ => panic!("unexpected ciphertext type"), - }; - let enc = NodeBuilder::new("enc") - .attr("type", enc_type) - .bytes(bytes) - .build(); - let enc_ref = enc.as_node_ref(); - let payload = EncPayload::from_node_ref(&enc_ref).unwrap(); - let info = Arc::new(MessageInfo { - id: id.to_string(), - source: crate::types::message::MessageSource { - sender: sender.clone(), - chat: sender.clone(), - ..Default::default() - }, - ..Default::default() - }); - client - .clone() - .process_classified_message(ClassifiedMessage { - info, - sender_encryption_jid: sender.clone(), - session_payloads: vec![payload], - group_payloads: vec![], - bot_payloads: vec![], - max_sender_retry_count: 0, - decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, - }) - .await; - } - - fn enc_payload_from_ciphertext(ct: &CiphertextMessage) -> EncPayload { - let (enc_type, bytes) = match ct { - CiphertextMessage::PreKeySignalMessage(m) => ("pkmsg", m.serialized().to_vec()), - CiphertextMessage::SignalMessage(m) => ("msg", m.serialized().to_vec()), - _ => panic!("unexpected ciphertext type"), - }; - let enc = NodeBuilder::new("enc") - .attr("type", enc_type) - .bytes(bytes) - .build(); - EncPayload::from_node_ref(&enc.as_node_ref()).expect("ciphertext payload") - } - - fn skmsg_payload_from_bytes(bytes: Vec<u8>) -> EncPayload { - let enc = NodeBuilder::new("enc") - .attr("type", "skmsg") - .bytes(bytes) - .build(); - EncPayload::from_node_ref(&enc.as_node_ref()).expect("skmsg payload") - } - - fn msmsg_payload_from_bytes(bytes: Vec<u8>) -> EncPayload { - let enc = NodeBuilder::new("enc") - .attr("type", "msmsg") - .bytes(bytes) - .build(); - EncPayload::from_node_ref(&enc.as_node_ref()).expect("msmsg payload") - } - - fn group_message_info( - id: &str, - group: &Jid, - sender: &Jid, - is_from_me: bool, - ) -> Arc<MessageInfo> { - Arc::new(MessageInfo { - id: id.to_string(), - source: crate::types::message::MessageSource { - sender: sender.clone(), - chat: group.clone(), - is_from_me, - is_group: true, - ..Default::default() - }, - ..Default::default() - }) - } - - async fn process_group_classified( - client: &Arc<Client>, - info: Arc<MessageInfo>, - sender: &Jid, - session_payload: EncPayload, - group_payloads: Vec<EncPayload>, - ) { - process_group_classified_with_sessions( - client, - info, - sender, - vec![session_payload], - group_payloads, - ) - .await; - } - - async fn process_group_classified_with_sessions( - client: &Arc<Client>, - info: Arc<MessageInfo>, - sender: &Jid, - session_payloads: Vec<EncPayload>, - group_payloads: Vec<EncPayload>, - ) { - process_group_classified_with_payloads( - client, - info, - sender, - session_payloads, - group_payloads, - vec![], - ) - .await; - } - - async fn process_group_classified_with_payloads( - client: &Arc<Client>, - info: Arc<MessageInfo>, - sender: &Jid, - session_payloads: Vec<EncPayload>, - group_payloads: Vec<EncPayload>, - bot_payloads: Vec<EncPayload>, - ) { - client - .clone() - .process_classified_message(ClassifiedMessage { - info, - sender_encryption_jid: sender.clone(), - session_payloads, - group_payloads, - bot_payloads, - max_sender_retry_count: 0, - decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, - }) - .await; - } - - fn message_events_for_id(rx: &async_channel::Receiver<Arc<Event>>, id: &str) -> (usize, usize) { - let mut count = 0; - let mut visible_content = 0; - while let Ok(event) = rx.try_recv() { - if let Event::Message(msg, info) = event.as_ref() - && info.id == id - { - count += 1; - if msg.conversation.is_some() { - visible_content += 1; - } - } - } - (count, visible_content) - } - - fn message_texts_for_id(rx: &async_channel::Receiver<Arc<Event>>, id: &str) -> Vec<String> { - let mut texts = Vec::new(); - while let Ok(event) = rx.try_recv() { - if let Event::Message(msg, info) = event.as_ref() - && info.id == id - && let Some(text) = &msg.conversation - { - texts.push(text.clone()); - } - } - texts - } - - #[tokio::test] - async fn skdm_only_group_session_acknowledged_once_without_message_event() { - use wacore::messages::MessageUtils; - use wacore::types::events::ChannelEventHandler; - - let (client, transport) = capturing_client("skdm_only_group_ack").await; - let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); - - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("146824178450525@lid").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - let group: Jid = "120363408782575443@g.us".parse().expect("group"); - let skdm = alice.create_group_skdm(&group).await; - let plaintext = MessageUtils::encode_and_pad(&wa::Message { - sender_key_distribution_message: Some(skdm), - ..Default::default() - }); - let session_ct = alice.encrypt(&bob_addr, &plaintext).await; - let id = "SKDM_ONLY_SESSION"; - let info = group_message_info(id, &group, &alice.jid, false); - - process_group_classified( - &client, - info, - &alice.jid, - enc_payload_from_ciphertext(&session_ct), - vec![], - ) - .await; - - assert_exactly_one_confirmation(&transport, id).await; - assert_eq!( - delivery_receipts_for(&transport.sent(), id), - 1, - "incoming group SKDM-only session message should drain via delivery receipt" - ); - let sent = transport.sent(); - let receipt = find_receipt_details(&sent, id).expect("delivery receipt"); - let sender_str = alice.jid.to_string(); - assert_eq!(receipt.to, group.to_string()); - assert_eq!(receipt.participant.as_deref(), Some(sender_str.as_str())); - assert_eq!(receipt.recipient, None); - assert_ne!( - receipt.typ.as_deref(), - Some("sender"), - "incoming group SKDM-only must not be cleared as a sender receipt" - ); - assert_eq!( - message_acks_for(&sent, id), - 0, - "incoming group SKDM-only must not also emit a transport ack" - ); - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - assert_eq!( - message_events_for_id(&rx, id), - (0, 0), - "SKDM-only messages must not surface Event::Message" - ); - } - - #[tokio::test] - async fn session_plaintext_decode_error_is_not_acked_as_skdm_only() { - use wacore::types::events::ChannelEventHandler; - - let (client, transport) = capturing_client("bad_plaintext_no_skdm_ack").await; - let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); - let recorder = Arc::new(EventRecorder::default()); - client.register_handler(recorder.clone()); - - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("146824178450527@lid").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - let group: Jid = "120363408782575446@g.us".parse().expect("group"); - let invalid_padded_plaintext = vec![0xff, 0x01]; - let session_ct = alice.encrypt(&bob_addr, &invalid_padded_plaintext).await; - let id = "BAD_SESSION_PLAINTEXT"; - let info = group_message_info(id, &group, &alice.jid, false); - - process_group_classified( - &client, - info, - &alice.jid, - enc_payload_from_ciphertext(&session_ct), - vec![], - ) - .await; - - for _ in 0..5 { - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - assert_eq!( - confirmations_for(&transport.sent(), id), - 0, - "invalid plaintext must not be counted as a successful SKDM-only ack" - ); - } - assert_eq!( - recorder.undecryptable().len(), - 1, - "plaintext handler failures must stay on the undecryptable path" - ); - assert_eq!( - message_events_for_id(&rx, id), - (0, 0), - "invalid plaintext must not surface Event::Message" - ); - let mut nack_code = None; - for _ in 0..80 { - nack_code = find_message_nack_error(&transport.sent(), id); - if nack_code.is_some() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!( - nack_code, - Some(491), - "invalid decrypted protobuf must be drained with InvalidProtobuf nack" - ); - } - - #[tokio::test] - async fn mixed_skdm_and_bad_plaintext_session_is_nacked_not_positive_acked() { - use wacore::messages::MessageUtils; - use wacore::types::events::ChannelEventHandler; - - let (client, transport) = capturing_client("mixed_skdm_bad_plaintext").await; - let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); - let recorder = Arc::new(EventRecorder::default()); - client.register_handler(recorder.clone()); - - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("146824178450531@lid").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - let group: Jid = "120363408782575449@g.us".parse().expect("group"); - let skdm = alice.create_group_skdm(&group).await; - let skdm_plaintext = MessageUtils::encode_and_pad(&wa::Message { - sender_key_distribution_message: Some(skdm), - ..Default::default() - }); - let skdm_ct = alice.encrypt(&bob_addr, &skdm_plaintext).await; - let bad_ct = alice.encrypt(&bob_addr, &[0xff, 0x01]).await; - let id = "SKDM_WITH_BAD_SESSION"; - let info = group_message_info(id, &group, &alice.jid, false); - - process_group_classified_with_sessions( - &client, - info, - &alice.jid, - vec![ - enc_payload_from_ciphertext(&skdm_ct), - enc_payload_from_ciphertext(&bad_ct), - ], - vec![], - ) - .await; - - let mut nack_code = None; - for _ in 0..80 { - nack_code = find_message_nack_error(&transport.sent(), id); - if nack_code.is_some() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!(nack_code, Some(491)); - assert_eq!( - confirmations_for(&transport.sent(), id), - 0, - "SKDM-only fallback must not positive-ack a mixed malformed session batch" - ); - assert_eq!( - recorder.undecryptable().len(), - 1, - "the malformed sibling must still surface as undecryptable" - ); - assert_eq!( - message_events_for_id(&rx, id), - (0, 0), - "mixed SKDM and bad plaintext must not dispatch user content" - ); - } - - #[tokio::test] - async fn bad_session_plaintext_skips_skmsg_sibling_after_nack() { - use wacore::messages::MessageUtils; - use wacore::types::events::ChannelEventHandler; - - let (client, transport) = capturing_client("bad_session_skips_skmsg").await; - let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); - let recorder = Arc::new(EventRecorder::default()); - client.register_handler(recorder.clone()); - - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("146824178450532@lid").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - let group: Jid = "120363408782575450@g.us".parse().expect("group"); - let skdm = alice.create_group_skdm(&group).await; - let skdm_plaintext = MessageUtils::encode_and_pad(&wa::Message { - sender_key_distribution_message: Some(skdm), - ..Default::default() - }); - let skdm_ct = alice.encrypt(&bob_addr, &skdm_plaintext).await; - let bad_ct = alice.encrypt(&bob_addr, &[0xff, 0x01]).await; - let content_plaintext = MessageUtils::encode_and_pad(&wa::Message { - conversation: Some("must not dispatch".to_string()), - ..Default::default() - }); - let skmsg = alice - .encrypt_group_message(&group, &content_plaintext) - .await; - let id = "BAD_SESSION_WITH_SKMSG"; - let info = group_message_info(id, &group, &alice.jid, false); - - process_group_classified_with_sessions( - &client, - info, - &alice.jid, - vec![ - enc_payload_from_ciphertext(&skdm_ct), - enc_payload_from_ciphertext(&bad_ct), - ], - vec![skmsg_payload_from_bytes(skmsg)], - ) - .await; - - let mut nack_code = None; - for _ in 0..80 { - nack_code = find_message_nack_error(&transport.sent(), id); - if nack_code.is_some() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!(nack_code, Some(491)); - assert_eq!( - confirmations_for(&transport.sent(), id), - 0, - "skmsg must not ack after a session InvalidProtobuf nack" - ); - assert_eq!( - recorder.undecryptable().len(), - 1, - "session plaintext failure should own the only user-visible failure" - ); - assert_eq!( - message_texts_for_id(&rx, id), - Vec::<String>::new(), - "skmsg content must be skipped after session plaintext failure" - ); - } - - #[tokio::test] - async fn skdm_only_session_with_msmsg_waits_for_bot_payload_response() { - use wacore::messages::MessageUtils; - - let (client, transport) = capturing_client("skdm_msmsg_no_fallback_ack").await; - - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("146824178450533@lid").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - let group: Jid = "120363408782575451@g.us".parse().expect("group"); - let skdm = alice.create_group_skdm(&group).await; - let plaintext = MessageUtils::encode_and_pad(&wa::Message { - sender_key_distribution_message: Some(skdm), - ..Default::default() - }); - let session_ct = alice.encrypt(&bob_addr, &plaintext).await; - let id = "SKDM_WITH_MSMSG"; - let info = group_message_info(id, &group, &alice.jid, false); - - process_group_classified_with_payloads( - &client, - info, - &alice.jid, - vec![enc_payload_from_ciphertext(&session_ct)], - vec![], - vec![msmsg_payload_from_bytes(vec![0xff])], - ) - .await; - - let mut nack_code = None; - for _ in 0..80 { - nack_code = find_message_nack_error(&transport.sent(), id); - if nack_code.is_some() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!(nack_code, Some(487)); - assert_eq!( - confirmations_for(&transport.sent(), id), - 0, - "SKDM-only fallback must not pre-ack a stanza with msmsg work" - ); - } - - #[tokio::test] - async fn session_content_group_message_acknowledged_once_without_fallback() { - use wacore::messages::MessageUtils; - use wacore::types::events::ChannelEventHandler; - - let (client, transport) = capturing_client("session_content_group_ack").await; - let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); - - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("146824178450528@lid").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - let group: Jid = "120363408782575447@g.us".parse().expect("group"); - let plaintext = MessageUtils::encode_and_pad(&wa::Message { - conversation: Some("session content".to_string()), - ..Default::default() - }); - let session_ct = alice.encrypt(&bob_addr, &plaintext).await; - let id = "SESSION_CONTENT_GROUP"; - let info = group_message_info(id, &group, &alice.jid, false); - - process_group_classified( - &client, - info, - &alice.jid, - enc_payload_from_ciphertext(&session_ct), - vec![], - ) - .await; - - assert_exactly_one_confirmation(&transport, id).await; - let sent = transport.sent(); - assert_eq!( - delivery_receipts_for(&sent, id), - 1, - "session content dispatch should own the only delivery receipt" - ); - assert_eq!( - message_acks_for(&sent, id), - 0, - "normal session content must not also use the SKDM-only transport ack" - ); - assert_eq!( - message_texts_for_id(&rx, id), - vec!["session content".to_string()], - "normal session content must dispatch exactly once" - ); - } - - #[tokio::test] - async fn status_skdm_only_session_uses_one_status_receipt() { - use wacore::messages::MessageUtils; - use wacore::types::events::ChannelEventHandler; - - let (client, transport) = capturing_client("status_skdm_only_ack").await; - let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); - - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("146824178450529@lid").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - let status: Jid = "status@broadcast".parse().expect("status"); - let skdm = alice.create_group_skdm(&status).await; - let plaintext = MessageUtils::encode_and_pad(&wa::Message { - sender_key_distribution_message: Some(skdm), - ..Default::default() - }); - let session_ct = alice.encrypt(&bob_addr, &plaintext).await; - let id = "STATUS_SKDM_ONLY"; - let info = group_message_info(id, &status, &alice.jid, false); - - process_group_classified( - &client, - info, - &alice.jid, - enc_payload_from_ciphertext(&session_ct), - vec![], - ) - .await; - - assert_exactly_one_confirmation(&transport, id).await; - let sent = transport.sent(); - assert_eq!( - delivery_receipts_for(&sent, id), - 1, - "status SKDM-only success must still send the WA Web status receipt" - ); - assert_eq!( - message_acks_for(&sent, id), - 0, - "status success path should not use a transport ack" - ); - let receipt = find_receipt_details(&sent, id).expect("status delivery receipt"); - let sender_str = alice.jid.to_string(); - assert_eq!(receipt.to, status.to_string()); - assert_eq!(receipt.participant.as_deref(), Some(sender_str.as_str())); - assert_eq!(receipt.context.as_deref(), Some("status")); - assert_eq!( - message_events_for_id(&rx, id), - (0, 0), - "status SKDM-only messages must not surface Event::Message" - ); - } - - #[tokio::test] - async fn error_message_ack_is_not_counted_as_positive_confirmation() { - let (client, transport) = capturing_client("error_ack_not_positive").await; - let id = "ERROR_ACK_NOT_POSITIVE"; - let info = Arc::new(MessageInfo { - id: id.to_string(), - source: crate::types::message::MessageSource { - sender: "146824178450530@lid".parse().expect("sender"), - chat: "120363408782575448@g.us".parse().expect("group"), - is_group: true, - ..Default::default() - }, - ..Default::default() - }); - - client.spawn_nack( - &info, - wacore::protocol::nack::NackReason::ParsingError, - None, - ); - - let mut nack_code = None; - for _ in 0..80 { - nack_code = find_message_nack_error(&transport.sent(), id); - if nack_code.is_some() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!(nack_code, Some(487)); - - let sent = transport.sent(); - assert_eq!( - message_acks_for(&sent, id), - 0, - "nacks carry class=message but must not count as positive transport acks" - ); - assert_eq!( - confirmations_for(&sent, id), - 0, - "nacks must not satisfy exactly-one positive confirmation assertions" - ); - assert!( - find_message_ack_for(&sent, id).is_none(), - "error acks must be excluded from positive ack lookup" - ); - } - - #[tokio::test] - async fn skdm_session_with_skmsg_sibling_acknowledged_once() { - use wacore::messages::MessageUtils; - use wacore::types::events::ChannelEventHandler; - - let (client, transport) = capturing_client("skdm_plus_skmsg_ack").await; - let (handler, rx) = ChannelEventHandler::new(); - client.core.event_bus.add_handler(handler); - - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("146824178450526@lid").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - let group: Jid = "120363408782575444@g.us".parse().expect("group"); - let skdm = alice.create_group_skdm(&group).await; - let skdm_plaintext = MessageUtils::encode_and_pad(&wa::Message { - sender_key_distribution_message: Some(skdm), - ..Default::default() - }); - let session_ct = alice.encrypt(&bob_addr, &skdm_plaintext).await; - - let content_plaintext = MessageUtils::encode_and_pad(&wa::Message { - conversation: Some("group content".to_string()), - ..Default::default() - }); - let skmsg = alice - .encrypt_group_message(&group, &content_plaintext) - .await; - let id = "SKDM_WITH_SKMSG"; - let info = group_message_info(id, &group, &alice.jid, false); - - process_group_classified( - &client, - info, - &alice.jid, - enc_payload_from_ciphertext(&session_ct), - vec![skmsg_payload_from_bytes(skmsg)], - ) - .await; - - assert_exactly_one_confirmation(&transport, id).await; - assert_eq!( - delivery_receipts_for(&transport.sent(), id), - 1, - "the skmsg content dispatch should own the only receipt" - ); - let sent = transport.sent(); - let receipt = find_receipt_details(&sent, id).expect("delivery receipt"); - let sender_str = alice.jid.to_string(); - assert_eq!(receipt.to, group.to_string()); - assert_eq!(receipt.participant.as_deref(), Some(sender_str.as_str())); - assert_eq!( - message_acks_for(&sent, id), - 0, - "SKDM+skmsg sibling must not get an extra transport ack" - ); - assert_eq!( - message_texts_for_id(&rx, id), - vec!["group content".to_string()], - "only the skmsg content should dispatch a user message" - ); - } - - #[tokio::test] - async fn own_group_skdm_only_session_uses_transport_ack_once() { - use wacore::messages::MessageUtils; - - let (client, transport) = capturing_client("own_group_skdm_ack").await; - - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("999999999999999@lid").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - let group: Jid = "120363408782575445@g.us".parse().expect("group"); - let skdm = alice.create_group_skdm(&group).await; - let plaintext = MessageUtils::encode_and_pad(&wa::Message { - sender_key_distribution_message: Some(skdm), - ..Default::default() - }); - let session_ct = alice.encrypt(&bob_addr, &plaintext).await; - let id = "OWN_GROUP_SKDM_ONLY"; - let info = group_message_info(id, &group, &alice.jid, true); - - process_group_classified( - &client, - info, - &alice.jid, - enc_payload_from_ciphertext(&session_ct), - vec![], - ) - .await; - - assert_exactly_one_confirmation(&transport, id).await; - let sent = transport.sent(); - assert_eq!( - message_acks_for(&sent, id), - 1, - "own group SKDM-only session message should use transport ack" - ); - let ack = find_message_ack_for(&sent, id).expect("transport ack"); - let sender_str = alice.jid.to_string(); - assert_eq!(ack.to, group.to_string()); - assert_eq!(ack.participant.as_deref(), Some(sender_str.as_str())); - assert_eq!(ack.recipient, None); - assert_eq!( - delivery_receipts_for(&sent, id), - 0, - "own group SKDM-only session message must not use a delivery receipt" - ); - assert_eq!( - sender_receipts_for(&sent, id), - 0, - "group self-fanout must not use type=sender receipt" - ); - } - - /// Regression for the offline-backlog disconnect: an already-processed - /// (duplicate) message must get its own delivery receipt, else the server - /// replays it every reconnect until it force-closes the stream. Pre-fix only - /// the first (success) delivery was acked; the duplicate was skipped silently. - #[tokio::test] - async fn duplicate_message_is_acked_with_delivery_receipt() { - let (client, transport) = capturing_client("dup_receipt").await; - let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; - let bob_addr = bob_jid.to_protocol_address(); - let mut alice = AlicePeer::new("5511888887777@s.whatsapp.net").await; - alice.install_bob_session(&bob_addr, &bundle).await; - - // Establish the session, then mark Alice's prekey acked so her next message - // is a plain SignalMessage. Re-submitting it is a clean duplicate. - let establish = alice.encrypt_text(&bob_addr, "establish").await; - process_session_ct(&client, &alice.jid, "EST", &establish).await; - if let Some(record) = alice.sessions.0.get_mut(&bob_addr) - && let Some(state) = record.session_state_mut() - { - state.clear_unacknowledged_pre_key_message(); - } - - // A real (padded) Message so the success path also emits its receipt. - let plaintext = wacore::messages::MessageUtils::encode_and_pad(&wa::Message { - conversation: Some("hi".to_string()), - ..Default::default() - }); - let msg = alice.encrypt(&bob_addr, &plaintext).await; - process_session_ct(&client, &alice.jid, "DUP", &msg).await; // success - process_session_ct(&client, &alice.jid, "DUP", &msg).await; // duplicate - - let mut count = 0; - for _ in 0..80 { - count = delivery_receipts_for(&transport.sent(), "DUP"); - if count >= 2 { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!( - count, 2, - "duplicate must get its own delivery receipt (pre-fix: only the first send was acked)" - ); - } - - /// Own-account self-fanout (is_from_me, non-peer, carries a `recipient`): - /// our own outgoing message echoed back to this device. WA Web - /// (`isMeAccount(author) => SENDER`) and whatsmeow (`IsFromMe => "sender"`) - /// clear it with a `<receipt type="sender" recipient=...>`, NOT a bare - /// transport `<ack>`. The server's offline queue ignores the bare ack and - /// replays the stanza forever (the ~50min disconnect loop). - #[tokio::test] - async fn own_self_fanout_acked_via_sender_receipt() { - let (client, transport) = capturing_client("own_ack").await; - let own = Arc::new(MessageInfo { - id: "OWN1".to_string(), - source: crate::types::message::MessageSource { - sender: "100000000000001@lid".parse().expect("sender"), - chat: "300000000000003@lid".parse().expect("chat"), - recipient: Some("300000000000003@lid".parse().expect("recipient")), - is_from_me: true, - ..Default::default() - }, - ..Default::default() - }); - client.ack_received_message(&own); - - let mut found = None; - for _ in 0..80 { - if let Some(r) = find_receipt(&transport.sent(), "OWN1") { - found = Some(r); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let (to, typ, recipient) = found.expect("own self-fanout must get a sender <receipt>"); - assert_eq!( - to, "100000000000001@lid", - "receipt `to` must echo the own LID (the fanout sender)" - ); - assert_eq!( - typ.as_deref(), - Some("sender"), - "own self-fanout receipt must be type=sender" - ); - assert_eq!( - recipient.as_deref(), - Some("300000000000003@lid"), - "receipt must echo the fanout recipient" - ); - assert!( - find_message_ack(&transport.sent()).is_none(), - "self-fanout must NOT also emit a bare transport <ack> (the server rejects it)" - ); - } - - /// Regression for the bot self-fanout disconnect loop: our own message to a - /// `@bot` recipient, echoed back as a duplicate/undecryptable stanza, must - /// be cleared with a `<receipt type="sender" recipient=@bot>`. Pre-fix it - /// got a bare `<ack class="message">` which the server ignored, replaying - /// the stanza every reconnect until a ~50min `<stream:error><ack/>` GC - /// force-closed the connection (the exact production symptom). - #[tokio::test] - async fn bot_self_fanout_acked_via_sender_receipt() { - let (client, transport) = capturing_client("bot_self_fanout").await; - let own = Arc::new(MessageInfo { - id: "AC00000000000000000000000000BEEF".to_string(), - source: crate::types::message::MessageSource { - // from = our own LID with its device (the server fans our - // outgoing bot prompt back to this device); chat = the bot - // (recipient.to_non_ad). The device on the sender must survive - // into the receipt `to`, or the LID server rejects it (#649). - sender: "100000000000001:11@lid".parse().expect("sender"), - chat: "200000000000002@bot".parse().expect("chat"), - recipient: Some("200000000000002@bot".parse().expect("recipient")), - is_from_me: true, - ..Default::default() - }, - ..Default::default() - }); - client.ack_received_message(&own); - - let mut found = None; - for _ in 0..80 { - if let Some(r) = find_receipt(&transport.sent(), "AC00000000000000000000000000BEEF") { - found = Some(r); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let (to, typ, recipient) = - found.expect("bot self-fanout must get a sender <receipt> to drain the offline queue"); - assert_eq!( - to, "100000000000001:11@lid", - "receipt `to` must preserve the own LID device" - ); - assert_eq!(typ.as_deref(), Some("sender")); - assert_eq!( - recipient.as_deref(), - Some("200000000000002@bot"), - "receipt must route to the bot recipient" - ); - assert!( - find_message_ack(&transport.sent()).is_none(), - "the bare <ack> that triggered <stream:error><ack/> must no longer be emitted" - ); - } - - /// When WE are the bot author (own DM, sender on the `@bot` server, to a - /// user), WA Web's `MsgSendReceipt` takes the `!chat.isBot() && - /// author.isBot()` branch and emits a bot-invoke-response `<ack>`, NOT a - /// sender `<receipt>`. So the bot-author branch in ack_received_message must - /// keep running before the self-fanout receipt: this locks that ordering - /// against a regression that would wrongly route it to a sender receipt. - #[tokio::test] - async fn own_bot_author_dm_acks_not_sender_receipt() { - let (client, transport) = capturing_client("own_bot_author").await; - let own = Arc::new(MessageInfo { - id: "OWNBOT1".to_string(), - source: crate::types::message::MessageSource { - sender: "100000000000002@bot".parse().expect("sender"), - chat: "300000000000003@lid".parse().expect("chat"), - recipient: Some("300000000000003@lid".parse().expect("recipient")), - is_from_me: true, - ..Default::default() - }, - ..Default::default() - }); - client.ack_received_message(&own); - - let mut found = None; - for _ in 0..80 { - if let Some(a) = find_message_ack(&transport.sent()) { - found = Some(a); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert!( - found.is_some(), - "own bot-author DM must emit a bare <ack class=message> (WA Web bot-invoke-response ack), not a receipt" - ); - // No current race (ack_received_message is synchronous and the - // bot-author branch returns before the receipt branch), but settle - // briefly so a future regression that spawned a receipt on a later tick - // can't slip past this negative assertion. - for _ in 0..5 { - assert!( - find_receipt(&transport.sent(), "OWNBOT1").is_none(), - "must NOT route to a sender <receipt> (would diverge from WA Web's bot-invoke-response ack path)" - ); - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - } - - /// An `<unavailable>` message (no `<enc>`) must be transport-acked so the - /// server stops replaying it (DM/group aren't covered by the should_ack gate). - #[tokio::test] - async fn unavailable_message_is_transport_acked() { - let (client, transport) = capturing_client("unavail_ack").await; - let node = NodeBuilder::new("message") - .attr("from", "5511777776666@s.whatsapp.net") - .attr("id", "UNAVAIL1") - .attr("type", "text") - .children([NodeBuilder::new("unavailable") - .attr("type", "view_once") - .build()]) - .build(); - let owned = node_to_arc(node); - let classified = client.classify_incoming_message(&owned).await; - assert!(classified.is_none(), "unavailable path returns None"); - - let mut found = None; - for _ in 0..80 { - if let Some(a) = find_message_ack(&transport.sent()) { - found = Some(a); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let (to, _) = found.expect("unavailable message must get a transport ack"); - assert_eq!(to, "5511777776666@s.whatsapp.net"); - } - - /// Unknown-only stanzas (e.g. msmsg) must be acked or they loop the queue. - #[tokio::test] - async fn unknown_only_enc_is_transport_acked() { - let (client, transport) = capturing_client("msmsg_ack").await; - let node = NodeBuilder::new("message") - .attr("from", "5511777776666@s.whatsapp.net") - .attr("id", "MSMSG1") - .attr("type", "text") - .children([NodeBuilder::new("enc") - .attr("type", "frskmsg") - .bytes(vec![0u8; 8]) - .build()]) - .build(); - let owned = node_to_arc(node); - let classified = client.classify_incoming_message(&owned).await; - assert!( - classified.is_none(), - "unknown-only enc must short-circuit before the process phase" - ); - - let mut found = None; - for _ in 0..80 { - if let Some(a) = find_message_ack(&transport.sent()) { - found = Some(a); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let (to, _) = found.expect("unknown-only enc must emit a transport ack"); - assert_eq!(to, "5511777776666@s.whatsapp.net"); - } - - /// `recipient` must be echoed verbatim or the server replies <stream:error>. - #[tokio::test] - async fn unknown_only_enc_ack_preserves_recipient() { - let (client, transport) = capturing_client("msmsg_recipient").await; - let node = NodeBuilder::new("message") - .attr("from", "236395184570386@lid") - .attr("recipient", "156535032389744@lid") - .attr("id", "MSMSG_LID") - .attr("type", "text") - .children([NodeBuilder::new("enc") - .attr("type", "frskmsg") - .bytes(vec![0u8; 8]) - .build()]) - .build(); - let owned = node_to_arc(node); - let classified = client.classify_incoming_message(&owned).await; - assert!(classified.is_none()); - - let mut found = None; - for _ in 0..80 { - if let Some(a) = find_message_ack(&transport.sent()) { - found = Some(a); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let (to, recipient) = found.expect("unknown-only enc must emit a transport ack"); - assert_eq!(to, "236395184570386@lid"); - assert_eq!( - recipient.as_deref(), - Some("156535032389744@lid"), - "ack must echo the incoming `recipient` attr or the server replies with <stream:error><ack/>" - ); - } - - /// Known type with empty content still has no usable payload; ack it. - #[tokio::test] - async fn known_enc_type_with_empty_content_is_transport_acked() { - let (client, transport) = capturing_client("known_empty").await; - let node = NodeBuilder::new("message") - .attr("from", "5511777776666@s.whatsapp.net") - .attr("id", "EMPTY1") - .attr("type", "text") - .children([NodeBuilder::new("enc").attr("type", "pkmsg").build()]) - .build(); - let owned = node_to_arc(node); - let classified = client.classify_incoming_message(&owned).await; - assert!(classified.is_none()); - - let mut found = None; - for _ in 0..80 { - if let Some(a) = find_message_ack(&transport.sent()) { - found = Some(a); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let (to, _) = found.expect("known-but-empty enc must emit a transport ack"); - assert_eq!(to, "5511777776666@s.whatsapp.net"); - } - - /// status is covered by should_ack; the fallback must not double-ack it. - #[tokio::test] - async fn unknown_only_enc_on_status_skips_fallback_ack() { - let (client, transport) = capturing_client("msmsg_status_skip").await; - let node = NodeBuilder::new("message") - .attr("from", "status@broadcast") - .attr("id", "MSMSG_STATUS") - .attr("type", "text") - .attr("participant", "5511777776666@s.whatsapp.net") - .children([NodeBuilder::new("enc") - .attr("type", "frskmsg") - .bytes(vec![0u8; 8]) - .build()]) - .build(); - let owned = node_to_arc(node); - let classified = client.classify_incoming_message(&owned).await; - assert!(classified.is_none()); - - // Give any rogue spawned task time to land on the wire. - for _ in 0..16 { - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - if find_message_ack(&transport.sent()).is_some() { - break; - } - } - assert!( - find_message_ack(&transport.sent()).is_none(), - "status@broadcast must not get a fallback transport ack from classify" - ); - } - - /// One recognized enc + one unknown must still go through the normal path. - #[tokio::test] - async fn mixed_recognized_and_unknown_enc_still_classifies() { - let (client, _transport) = capturing_client("msmsg_mixed").await; - let node = NodeBuilder::new("message") - .attr("from", "5511777776666@s.whatsapp.net") - .attr("id", "MIXED1") - .attr("type", "text") - .children([ - NodeBuilder::new("enc") - .attr("type", "pkmsg") - .bytes(vec![0u8; 8]) - .build(), - NodeBuilder::new("enc") - .attr("type", "frskmsg") - .bytes(vec![0u8; 8]) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - let classified = client - .classify_incoming_message(&owned) - .await - .expect("mixed enc must produce a ClassifiedMessage"); - assert_eq!(classified.session_payloads.len(), 1); - assert!(classified.group_payloads.is_empty()); - } - - /// A custom handler owns its ack; the fallback must not double-ack. - #[tokio::test] - async fn custom_handler_only_skips_fallback_ack() { - use crate::types::enc_handler::EncHandler; - use async_lock::Mutex as AsyncMutex; - - #[derive(Default)] - struct NoopHandler { - calls: Arc<AsyncMutex<usize>>, - } - #[async_trait::async_trait] - impl EncHandler for NoopHandler { - async fn handle( - &self, - _client: Arc<Client>, - _enc_node: &wacore_binary::Node, - _info: &crate::types::message::MessageInfo, - ) -> anyhow::Result<()> { - *self.calls.lock().await += 1; - Ok(()) - } - } - - let (client, transport) = capturing_client("msmsg_custom").await; - let calls = Arc::new(AsyncMutex::new(0usize)); - let handler = Arc::new(NoopHandler { - calls: Arc::clone(&calls), - }); - client - .custom_enc_handlers - .write() - .await - .insert("frskmsg".to_string(), handler as Arc<dyn EncHandler>); - - let node = NodeBuilder::new("message") - .attr("from", "5511777776666@s.whatsapp.net") - .attr("id", "CUSTOM1") - .attr("type", "text") - .children([NodeBuilder::new("enc") - .attr("type", "frskmsg") - .bytes(vec![0u8; 8]) - .build()]) - .build(); - let owned = node_to_arc(node); - let classified = client.classify_incoming_message(&owned).await; - assert!( - classified.is_some(), - "custom-handled enc must not be short-circuited by the fallback guard" - ); - - // Let the detached handler + any rogue spawned task run. - for _ in 0..16 { - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert!( - find_message_ack(&transport.sent()).is_none(), - "custom-handled enc must not get a fallback transport ack from classify" - ); - assert_eq!(*calls.lock().await, 1, "custom handler must be invoked"); - } - - /// Security regression: a self-only `app_state_sync_key_share` protocol - /// message must be honoured only when it originates from our own account. - /// A spoofed one from a peer must be dropped (otherwise a peer could inject - /// app-state sync keys). Mirrors WA Web `WAWebKeyManagementHandleKeyShareApi` - /// and whatsmeow's `handleProtocolMessage` self gate. - #[tokio::test] - async fn app_state_sync_key_share_honored_only_from_self() { - use wacore::messages::MessageUtils; - - let client = crate::test_utils::create_test_client().await; - ensure_bob_paired(&client).await; - - let key_id = vec![1u8, 2, 3, 4, 5, 6]; - let share = wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage { - app_state_sync_key_share: Some(wa::message::AppStateSyncKeyShare { - keys: vec![wa::message::AppStateSyncKey { - key_id: Some(wa::message::AppStateSyncKeyId { - key_id: Some(key_id.clone()), - }), - key_data: Some(wa::message::AppStateSyncKeyData { - key_data: Some(vec![7u8; 32]), - fingerprint: Some(wa::message::AppStateSyncKeyFingerprint { - raw_id: Some(1), - current_index: Some(0), - device_indexes: vec![0], - }), - timestamp: Some(123), - }), - }], - }), - ..Default::default() - })), - ..Default::default() - }; - let padded = MessageUtils::encode_and_pad(&share); - let backend = client.persistence_manager.backend(); - - // Non-self sender: the key share must be dropped. - let mut info = - create_test_message_info("5510000@s.whatsapp.net", "AKS1", "5510000@s.whatsapp.net"); - info.source.is_from_me = false; - client - .clone() - .handle_decrypted_plaintext("msg", &padded, 2, &Arc::new(info)) - .await - .unwrap(); - assert!( - backend.get_sync_key(&key_id).await.unwrap().is_none(), - "app-state sync key from a non-self sender must not be stored" - ); - - // Self sender: the key share is honoured and stored. - let mut info = create_test_message_info( - "9000000000000@s.whatsapp.net", - "AKS2", - "9000000000000@s.whatsapp.net", - ); - info.source.is_from_me = true; - client - .clone() - .handle_decrypted_plaintext("msg", &padded, 2, &Arc::new(info)) - .await - .unwrap(); - assert!( - backend.get_sync_key(&key_id).await.unwrap().is_some(), - "app-state sync key from self must be stored" - ); - } - - // ---- msmsg inbound dispatch ----------------------------------------- - - fn find_message_nack_error(frames: &[bytes::Bytes], id: &str) -> Option<u32> { - for (i, frame) in frames.iter().enumerate() { - let Some(buf) = decode_frame(i, frame) else { - continue; - }; - let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { - continue; - }; - if node.tag.as_ref() == "ack" - && node - .get_attr("class") - .is_some_and(|v| v.as_str() == "message") - && node.get_attr("id").is_some_and(|v| v.as_str() == id) - && let Some(err) = node.get_attr("error") - && let Ok(code) = err.as_str().parse::<u32>() - { - return Some(code); - } - } - None - } - - fn encode_message_secret_message(iv: &[u8], payload: &[u8]) -> Vec<u8> { - use prost::Message as _; - let ms = wa::MessageSecretMessage { - version: Some(1), - enc_iv: Some(iv.to_vec()), - enc_payload: Some(payload.to_vec()), - }; - let mut out = Vec::with_capacity(ms.encoded_len()); - ms.encode(&mut out).expect("encode MessageSecretMessage"); - out - } - - async fn collect_event<F>( - client: &Arc<Client>, - collector: Arc<crate::test_utils::TestEventCollector>, - pred: F, - timeout_ms: u64, - ) -> Option<Arc<wacore::types::events::Event>> - where - F: Fn(&wacore::types::events::Event) -> bool, - { - let _ = client; - let mut waited = 0u64; - while waited <= timeout_ms { - for ev in collector.events() { - if pred(&ev) { - return Some(ev); - } - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - waited += 25; - } - None - } - - fn legacy_edit_text(msg: &wa::Message) -> Option<&str> { - msg.protocol_message - .as_ref() - .and_then(|pm| pm.edited_message.as_ref()) - .and_then(|edited| edited.conversation.as_deref()) - } - - fn inner_message_edit(text: &str, next_secret: Option<Vec<u8>>) -> wa::Message { - wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage { - key: Some(wa::MessageKey { - remote_jid: Some("5511777776666@s.whatsapp.net".to_string()), - from_me: Some(false), - id: Some("PARENT_EDIT".to_string()), - participant: None, - }), - r#type: Some(wa::message::protocol_message::Type::MessageEdit as i32), - edited_message: Some(Box::new(wa::Message { - conversation: Some(text.to_string()), - ..Default::default() - })), - timestamp_ms: Some(1_770_000_000_000), - ..Default::default() - })), - message_context_info: next_secret.map(|secret| wa::MessageContextInfo { - message_secret: Some(secret), - ..Default::default() - }), - ..Default::default() - } - } - - fn encrypted_message_edit( - target_key: wa::MessageKey, - original_sender: &str, - editor: &str, - parent_id: &str, - secret: &[u8], - text: &str, - next_secret: Option<Vec<u8>>, - ) -> wa::Message { - let ctx = wacore::message_edit::MessageEditContext { - original_msg_id: parent_id, - original_sender_jid: original_sender, - editor_jid: editor, - }; - let (enc_payload, enc_iv) = wacore::message_edit::encrypt_message_edit( - &inner_message_edit(text, next_secret), - secret, - &ctx, - ) - .expect("test edit encryption"); - - wa::Message { - secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { - target_message_key: Some(target_key), - enc_payload: Some(enc_payload), - enc_iv: Some(enc_iv.to_vec()), - secret_enc_type: Some( - wa::message::secret_encrypted_message::SecretEncType::MessageEdit as i32, - ), - remote_key_id: None, - }), - ..Default::default() - } - } - - #[tokio::test] - async fn secret_encrypted_message_edit_dispatches_legacy_edit() { - let (client, _transport) = capturing_client("secret_edit_dispatch").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "5511777776666@s.whatsapp.net"; - let parent_id = "PARENT_EDIT"; - let edit_id = "EDIT_1"; - let secret = [0x42u8; 32]; - client - .persistence_manager - .backend() - .put_msg_secret(chat, chat, parent_id, &secret) - .await - .unwrap(); - - let info = Arc::new(MessageInfo { - id: edit_id.into(), - source: crate::types::message::MessageSource { - chat: chat.parse().unwrap(), - sender: chat.parse().unwrap(), - ..Default::default() - }, - ..Default::default() - }); - let target_key = wa::MessageKey { - remote_jid: Some(chat.to_string()), - from_me: Some(false), - id: Some(parent_id.to_string()), - participant: None, - }; - let msg = - encrypted_message_edit(target_key, chat, chat, parent_id, &secret, "edited", None); - - client.dispatch_parsed_message(msg, &info).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_id - && legacy_edit_text(msg.as_ref()) == Some("edited") - && msg.secret_encrypted_message.is_none()) - }, - 500, - ) - .await; - assert!(got.is_some(), "encrypted edit must dispatch as legacy edit"); - } - - /// Regression for #667: an incoming peer edit writes `target_message_key` - /// in the editor's frame (`from_me = true`, no `participant`, even in a - /// group), so the target-key resolver maps the parent author to *us* and - /// misses the secret stored under the real author. The dispatch path must - /// take the author from the envelope sender instead. Fails on the pre-fix - /// code (envelope stays encrypted), passes after it. - #[tokio::test] - async fn secret_encrypted_peer_edit_resolves_sender_from_envelope() { - let (client, _transport) = capturing_client("secret_peer_edit_dispatch").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let group = "123456789012345678@g.us"; - let peer = "5511777776666@s.whatsapp.net"; - let parent_id = "PEER_PARENT"; - let edit_id = "PEER_EDIT"; - let secret = [0x42u8; 32]; - // The parent (peer's own message) was stored under the real author. - client - .persistence_manager - .backend() - .put_msg_secret(group, peer, parent_id, &secret) - .await - .unwrap(); - - let info = Arc::new(MessageInfo { - id: edit_id.into(), - source: crate::types::message::MessageSource { - chat: group.parse().unwrap(), - sender: peer.parse().unwrap(), - ..Default::default() - }, - ..Default::default() - }); - // Editor's frame: from_me = true, no participant, even in a group. - let target_key = wa::MessageKey { - remote_jid: Some(group.to_string()), - from_me: Some(true), - id: Some(parent_id.to_string()), - participant: None, - }; - // HKDF binds the real author (peer) as both original sender and editor. - let msg = - encrypted_message_edit(target_key, peer, peer, parent_id, &secret, "edited", None); - - client.dispatch_parsed_message(msg, &info).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_id - && legacy_edit_text(msg.as_ref()) == Some("edited") - && msg.secret_encrypted_message.is_none()) - }, - 500, - ) - .await; - assert!( - got.is_some(), - "incoming peer edit must resolve the author from the envelope and dispatch as legacy edit" - ); - } - - /// Store the parent secret with a known event time, then dispatch a - /// secret-encrypted edit authored `edit_offset` seconds after the parent. - /// Returns whether the decrypted legacy edit was dispatched. - async fn run_secret_edit_with_window(test_id: &str, parent_ts: i64, edit_offset: i64) -> bool { - let (client, _transport) = capturing_client(test_id).await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "5511777776666@s.whatsapp.net"; - let parent_id = "WINDOW_PARENT"; - let edit_id = "WINDOW_EDIT"; - let secret = [0x42u8; 32]; - client - .persistence_manager - .backend() - .put_msg_secrets(vec![wacore::store::traits::MsgSecretEntry { - chat: chat.to_string(), - sender: chat.to_string(), - msg_id: parent_id.to_string(), - secret: secret.to_vec(), - expires_at: 0, - message_ts: parent_ts, - }]) - .await - .unwrap(); - - let info = Arc::new(MessageInfo { - id: edit_id.into(), - timestamp: chrono::DateTime::<chrono::Utc>::from_timestamp(parent_ts + edit_offset, 0) - .unwrap(), - source: crate::types::message::MessageSource { - chat: chat.parse().unwrap(), - sender: chat.parse().unwrap(), - ..Default::default() - }, - ..Default::default() - }); - let target_key = wa::MessageKey { - remote_jid: Some(chat.to_string()), - from_me: Some(false), - id: Some(parent_id.to_string()), - participant: None, - }; - let msg = - encrypted_message_edit(target_key, chat, chat, parent_id, &secret, "edited", None); - client.dispatch_parsed_message(msg, &info).await; - - collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_id - && legacy_edit_text(msg.as_ref()) == Some("edited") - && msg.secret_encrypted_message.is_none()) - }, - 500, - ) - .await - .is_some() - } - - #[tokio::test] - async fn secret_edit_within_window_is_applied() { - // Authored 10 min after the parent — inside the 20 min (1200s) window. - assert!( - run_secret_edit_with_window("secret_edit_in_window", 1_700_000_000, 600).await, - "an in-window edit must dispatch as a legacy edit" - ); - } - - #[tokio::test] - async fn secret_edit_outside_window_is_dropped() { - // Authored 30 min after the parent — past the 1200s window, like WA Web's - // ProcessEditProtocolMsgs, so we drop it (raw envelope surfaces instead). - assert!( - !run_secret_edit_with_window("secret_edit_out_window", 1_700_000_000, 1800).await, - "an out-of-window edit must not dispatch a legacy edit" - ); - } - - #[tokio::test] - async fn secret_edit_unknown_parent_ts_is_permissive() { - // parent_ts = 0 (unknown, e.g. resolver-supplied): no window check, so a - // late edit still applies rather than being silently dropped. - assert!( - run_secret_edit_with_window("secret_edit_unknown_ts", 0, 5_000_000).await, - "with an unknown parent timestamp the edit must still apply" - ); - } - - #[tokio::test] - async fn secret_encrypted_edit_decrypts_via_resolver_when_store_empty() { - use crate::cache_config::{CacheConfig, MsgSecretPolicy}; - - struct StaticResolver { - chat: String, - sender: String, - msg_id: String, - secret: [u8; 32], - } - #[async_trait::async_trait] - impl wacore::msg_secret::OriginalMessageResolver for StaticResolver { - async fn resolve_msg_secret( - &self, - chat: &str, - sender: &str, - msg_id: &str, - ) -> Option<[u8; 32]> { - (chat == self.chat && sender == self.sender && msg_id == self.msg_id) - .then_some(self.secret) - } - } - - let chat = "5511777776666@s.whatsapp.net"; - let parent_id = "RESOLVER_PARENT"; - let edit_id = "RESOLVER_EDIT"; - let secret = [0x7Au8; 32]; - - let resolver = Arc::new(StaticResolver { - chat: chat.to_string(), - sender: chat.to_string(), - msg_id: parent_id.to_string(), - secret, - }); - // Disabled persists nothing, so the only path to the secret is the resolver. - let cfg = CacheConfig { - msg_secret_policy: MsgSecretPolicy::Disabled, - original_message_resolver: Some(resolver), - ..Default::default() - }; - let client = crate::test_utils::create_test_client_with_config( - "resolver_edit", - Arc::new(crate::test_utils::MockHttpClient), - cfg, - ) - .await; - seed_test_pn(&client).await; - - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - assert!( - client - .persistence_manager - .backend() - .get_msg_secret(chat, chat, parent_id) - .await - .unwrap() - .is_none(), - "store must be empty under Disabled" - ); - - let info = Arc::new(MessageInfo { - id: edit_id.into(), - source: crate::types::message::MessageSource { - chat: chat.parse().unwrap(), - sender: chat.parse().unwrap(), - ..Default::default() - }, - ..Default::default() - }); - let target_key = wa::MessageKey { - remote_jid: Some(chat.to_string()), - from_me: Some(false), - id: Some(parent_id.to_string()), - participant: None, - }; - let msg = encrypted_message_edit( - target_key, - chat, - chat, - parent_id, - &secret, - "edited via resolver", - None, - ); - - client.dispatch_parsed_message(msg, &info).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_id - && legacy_edit_text(msg.as_ref()) == Some("edited via resolver") - && msg.secret_encrypted_message.is_none()) - }, - 500, - ) - .await; - assert!( - got.is_some(), - "edit must decrypt via the resolver when the store is empty" - ); - } - - #[tokio::test] - async fn decrypted_message_edit_recaptures_secret_for_next_edit() { - let (client, _transport) = capturing_client("secret_edit_chain").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "5511777776666@s.whatsapp.net"; - let parent_id = "PARENT_EDIT"; - let first_secret = [0x11u8; 32]; - let second_secret = [0x22u8; 32]; - client - .persistence_manager - .backend() - .put_msg_secret(chat, chat, parent_id, &first_secret) - .await - .unwrap(); - - let target_key = wa::MessageKey { - remote_jid: Some(chat.to_string()), - from_me: Some(false), - id: Some(parent_id.to_string()), - participant: None, - }; - let first_info = Arc::new(MessageInfo { - id: "EDIT_CHAIN_1".into(), - source: crate::types::message::MessageSource { - chat: chat.parse().unwrap(), - sender: chat.parse().unwrap(), - ..Default::default() - }, - ..Default::default() - }); - let first_msg = encrypted_message_edit( - target_key.clone(), - chat, - chat, - parent_id, - &first_secret, - "first", - Some(second_secret.to_vec()), - ); - client.dispatch_parsed_message(first_msg, &first_info).await; - - let stored = client - .persistence_manager - .backend() - .get_msg_secret(chat, chat, parent_id) - .await - .unwrap(); - assert_eq!(stored.as_deref(), Some(&second_secret[..])); - - let second_info = Arc::new(MessageInfo { - id: "EDIT_CHAIN_2".into(), - source: crate::types::message::MessageSource { - chat: chat.parse().unwrap(), - sender: chat.parse().unwrap(), - ..Default::default() - }, - ..Default::default() - }); - let second_msg = encrypted_message_edit( - target_key, - chat, - chat, - parent_id, - &second_secret, - "second", - None, - ); - client - .dispatch_parsed_message(second_msg, &second_info) - .await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == "EDIT_CHAIN_2" - && legacy_edit_text(msg.as_ref()) == Some("second")) - }, - 500, - ) - .await; - assert!(got.is_some(), "second edit must use the re-captured secret"); - } - - #[tokio::test] - async fn secret_encrypted_message_edit_uses_lid_pn_fallback_in_group() { - use wacore::store::traits::LidPnMappingEntry; - - let (client, _transport) = capturing_client("secret_edit_alt_group").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "120363021033254949@g.us"; - let parent_id = "GROUP_PARENT_EDIT"; - let sender_lid = "236395184570386@lid"; - let sender_pn = "5511777776666@s.whatsapp.net"; - let secret = [0x77u8; 32]; - - client - .persistence_manager - .backend() - .put_lid_mapping(&LidPnMappingEntry { - lid: "236395184570386".into(), - phone_number: "5511777776666".into(), - created_at: 0, - updated_at: 0, - learning_source: "test".into(), - }) - .await - .unwrap(); - client - .persistence_manager - .backend() - .put_msg_secret(chat, sender_pn, parent_id, &secret) - .await - .unwrap(); - - let info = Arc::new(MessageInfo { - id: "GROUP_EDIT_1".into(), - source: crate::types::message::MessageSource { - chat: chat.parse().unwrap(), - sender: sender_lid.parse().unwrap(), - is_group: true, - addressing_mode: Some(wacore::types::message::AddressingMode::Lid), - ..Default::default() - }, - ..Default::default() - }); - let target_key = wa::MessageKey { - remote_jid: Some(chat.to_string()), - from_me: Some(false), - id: Some(parent_id.to_string()), - participant: Some(sender_lid.to_string()), - }; - let msg = encrypted_message_edit( - target_key, - sender_pn, - sender_lid, - parent_id, - &secret, - "group edited", - None, - ); - - client.dispatch_parsed_message(msg, &info).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == "GROUP_EDIT_1" - && legacy_edit_text(msg.as_ref()) == Some("group edited")) - }, - 500, - ) - .await; - assert!( - got.is_some(), - "group edit must decrypt when the stored secret is under PN" - ); - } - - #[tokio::test] - async fn decrypted_message_edit_refreshes_alternate_secret_alias() { - use wacore::store::traits::LidPnMappingEntry; - - let (client, _transport) = capturing_client("secret_edit_alt_refresh").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "120363021033254949@g.us"; - let parent_id = "GROUP_PARENT_EDIT"; - let sender_lid = "236395184570386@lid"; - let sender_pn = "5511777776666@s.whatsapp.net"; - let first_secret = [0x31u8; 32]; - let second_secret = [0x32u8; 32]; - - client - .persistence_manager - .backend() - .put_lid_mapping(&LidPnMappingEntry { - lid: "236395184570386".into(), - phone_number: "5511777776666".into(), - created_at: 0, - updated_at: 0, - learning_source: "test".into(), - }) - .await - .unwrap(); - for sender in [sender_lid, sender_pn] { - client - .persistence_manager - .backend() - .put_msg_secret(chat, sender, parent_id, &first_secret) - .await - .unwrap(); - } - - let first_info = Arc::new(MessageInfo { - id: "GROUP_EDIT_REFRESH_1".into(), - source: crate::types::message::MessageSource { - chat: chat.parse().unwrap(), - sender: sender_lid.parse().unwrap(), - is_group: true, - addressing_mode: Some(wacore::types::message::AddressingMode::Lid), - ..Default::default() - }, - ..Default::default() - }); - let first_target_key = wa::MessageKey { - remote_jid: Some(chat.to_string()), - from_me: Some(false), - id: Some(parent_id.to_string()), - participant: Some(sender_lid.to_string()), - }; - let first_msg = encrypted_message_edit( - first_target_key, - sender_lid, - sender_lid, - parent_id, - &first_secret, - "first", - Some(second_secret.to_vec()), - ); - client.dispatch_parsed_message(first_msg, &first_info).await; - - for sender in [sender_lid, sender_pn] { - let stored = client - .persistence_manager - .backend() - .get_msg_secret(chat, sender, parent_id) - .await - .unwrap(); - assert_eq!(stored.as_deref(), Some(&second_secret[..])); - } - - let second_info = Arc::new(MessageInfo { - id: "GROUP_EDIT_REFRESH_2".into(), - source: crate::types::message::MessageSource { - chat: chat.parse().unwrap(), - sender: sender_pn.parse().unwrap(), - is_group: true, - addressing_mode: Some(wacore::types::message::AddressingMode::Pn), - ..Default::default() - }, - ..Default::default() - }); - let second_target_key = wa::MessageKey { - remote_jid: Some(chat.to_string()), - from_me: Some(false), - id: Some(parent_id.to_string()), - participant: Some(sender_pn.to_string()), - }; - let second_msg = encrypted_message_edit( - second_target_key, - sender_pn, - sender_pn, - parent_id, - &second_secret, - "second", - None, - ); - client - .dispatch_parsed_message(second_msg, &second_info) - .await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == "GROUP_EDIT_REFRESH_2" - && legacy_edit_text(msg.as_ref()) == Some("second")) - }, - 500, - ) - .await; - assert!( - got.is_some(), - "chained edit must use the refreshed alternate alias" - ); - } - - /// Round-trip: store an outbound messageSecret, build a fake bot reply - /// whose payload we encrypt with the symmetric helper, route it through - /// classify, and assert the decrypted `wa::Message` lands on the bus. - #[tokio::test] - async fn msmsg_decrypts_when_secret_is_stored() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - let (client, _transport) = capturing_client("msmsg_ok").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "867051314767696@bot"; - let our_pn = "5511000000001@s.whatsapp.net"; - let bot_jid = "867051314767696@bot"; - let outbound_id = "OUTBOUND_1"; - let bot_reply_id = "BOT_REPLY_1"; - let secret = [0x42u8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - let plaintext_msg = wa::Message { - conversation: Some("hi from bot".to_string()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - let ctx = BotMessageContext { - msg_id: bot_reply_id, - target_sender_user_jid: our_pn, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - let node = NodeBuilder::new("message") - .attr("from", bot_jid) - .attr("id", bot_reply_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("hi from bot")) - }, - 1500, - ) - .await; - assert!( - got.is_some(), - "msmsg decryption + dispatch must surface Event::Message" - ); - } - - /// No secret stored for `target_id` → nack `error=495`, no Message event. - #[tokio::test] - async fn msmsg_without_stored_secret_nacks_495() { - let (client, transport) = capturing_client("msmsg_nosecret").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let bot_reply_id = "BOT_REPLY_NS"; - let outbound_id = "OUTBOUND_NS"; - let our_pn = "5511000000001@s.whatsapp.net"; - - let ms_msg_proto = encode_message_secret_message(&[0u8; 12], &[0u8; 32]); - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", bot_reply_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let mut code = None; - for _ in 0..80 { - if let Some(c) = find_message_nack_error(&transport.sent(), bot_reply_id) { - code = Some(c); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!( - code, - Some(495), - "missing messageSecret must nack with code 495" - ); - assert!( - collector - .events() - .iter() - .all(|e| !matches!(e.as_ref(), wacore::types::events::Event::Message(_, info) if info.id == bot_reply_id)), - "no Message event must be dispatched when decryption failed" - ); - } - - /// Tampered ciphertext → GCM tag fails → nack 495. - #[tokio::test] - async fn msmsg_with_bad_tag_nacks_495() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - let (client, transport) = capturing_client("msmsg_bad_tag").await; - - let chat = "867051314767696@bot"; - let our_pn = "5511000000001@s.whatsapp.net"; - let outbound_id = "OUTBOUND_BAD"; - let bot_reply_id = "BOT_REPLY_BAD"; - let secret = [0x77u8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - let ctx = BotMessageContext { - msg_id: bot_reply_id, - target_sender_user_jid: our_pn, - bot_user_jid: chat, - }; - let (mut cipher, iv) = encrypt_bot_message(b"hello", &secret, &ctx).unwrap(); - let last = cipher.len() - 1; - cipher[last] ^= 0x01; - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", bot_reply_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let mut code = None; - for _ in 0..80 { - if let Some(c) = find_message_nack_error(&transport.sent(), bot_reply_id) { - code = Some(c); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!(code, Some(495)); - } - - /// Bot edit chain: when `<bot edit="inner">` is set, the HKDF msg_id used - /// for the per-message key swaps to `edit_target_id` so the edited reply - /// decrypts under the same key as the original (whatsmeow / WA Web - /// `decryptMsmsgFbidBotMessage`). - #[tokio::test] - async fn msmsg_bot_edit_uses_edit_target_id_for_hkdf() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - let (client, _transport) = capturing_client("msmsg_edit").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "867051314767696@bot"; - let our_pn = "5511000000001@s.whatsapp.net"; - let outbound_id = "OUTBOUND_EDIT"; - let original_reply_id = "BOT_REPLY_FIRST"; - let edit_reply_id = "BOT_REPLY_EDIT"; - let secret = [0xAAu8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - // Encrypt as if it's the ORIGINAL reply (msg_id = original_reply_id). - let plaintext_msg = wa::Message { - conversation: Some("edited content".to_string()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - let ctx = BotMessageContext { - msg_id: original_reply_id, - target_sender_user_jid: our_pn, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - // Inbound stanza has id=edit_reply_id but <bot edit="inner" edit_target_id=original> - // so the HKDF must derive against original_reply_id. - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", edit_reply_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("bot") - .attr("edit", "inner") - .attr("edit_target_id", original_reply_id) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_reply_id - && msg.conversation.as_deref() == Some("edited content")) - }, - 1500, - ) - .await; - assert!( - got.is_some(), - "bot edit must use edit_target_id for HKDF msg_id" - ); - } - - /// Same setup as the edit test but WITHOUT `<bot edit>`: the HKDF must - /// fall back to `info.id`, and ciphertext encrypted with the edit-target - /// id must fail to decrypt. - #[tokio::test] - async fn msmsg_without_bot_edit_does_not_swap_msg_id() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - let (client, transport) = capturing_client("msmsg_noedit").await; - let chat = "867051314767696@bot"; - let our_pn = "5511000000001@s.whatsapp.net"; - let outbound_id = "OUTBOUND_NOEDIT"; - let stanza_id = "BOT_REPLY_NOEDIT"; - let other_id = "OTHER_ID"; - let secret = [0xBBu8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - // Encrypt with `other_id` to simulate the wrong key derivation if the - // edit branch were taken without `<bot edit>`. - let ctx = BotMessageContext { - msg_id: other_id, - target_sender_user_jid: our_pn, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(b"x", &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", stanza_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let mut code = None; - for _ in 0..80 { - if let Some(c) = find_message_nack_error(&transport.sent(), stanza_id) { - code = Some(c); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!( - code, - Some(495), - "without <bot edit>, HKDF must use stanza id (not OTHER_ID) → tag fails" - ); - } - - /// `<bot edit="first">` is NOT one of {INNER, LAST}, so the HKDF msg_id - /// must remain `info.id`. - #[tokio::test] - async fn msmsg_bot_edit_first_keeps_info_id() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - let (client, _transport) = capturing_client("msmsg_first").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "867051314767696@bot"; - let our_pn = "5511000000001@s.whatsapp.net"; - let outbound_id = "OUTBOUND_FIRST"; - let stanza_id = "BOT_REPLY_FIRST"; - let secret = [0xCCu8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - // Encrypt with stanza_id; "first" edit must NOT swap. - let plaintext_msg = wa::Message { - conversation: Some("first reply".to_string()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - let ctx = BotMessageContext { - msg_id: stanza_id, - target_sender_user_jid: our_pn, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", stanza_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("bot").attr("edit", "first").build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| matches!(e, wacore::types::events::Event::Message(_, info) if info.id == stanza_id), - 1500, - ) - .await; - assert!(got.is_some(), "edit=first must keep info.id as HKDF msg_id"); - } - - /// Regular bot path (`f()` in WA Web `BotMessageSecret.js`): when the - /// fbid pre-resolve picks the WRONG id (e.g. edit_target_id) but the - /// real ciphertext was minted under `info.id`, the fallback attempt - /// must succeed. Validates the try-then-fallback unification. - #[tokio::test] - async fn msmsg_falls_back_to_info_id_when_primary_uses_edit_target() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - let (client, _transport) = capturing_client("msmsg_fb_to_info").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "867051314767696@bot"; - let our_pn = "5511000000001@s.whatsapp.net"; - let outbound_id = "OUT_FB1"; - let stanza_id = "REPLY_FB1"; - let edit_target_id = "WRONG_EDIT_TARGET"; - let secret = [0xDDu8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - // Encrypt under `stanza_id` even though the stanza will declare - // edit=inner with edit_target_id (forces a primary-attempt mismatch). - let plaintext_msg = wa::Message { - conversation: Some("fallback ok".to_string()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - let ctx = BotMessageContext { - msg_id: stanza_id, - target_sender_user_jid: our_pn, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", stanza_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("bot") - .attr("edit", "inner") - .attr("edit_target_id", edit_target_id) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == stanza_id - && msg.conversation.as_deref() == Some("fallback ok")) - }, - 1500, - ) - .await; - assert!( - got.is_some(), - "primary attempt with edit_target_id must fall back to info.id" - ); - } - - /// Inverse of the `falls_back_to_info_id` test: primary is `info.id` - /// (edit_type isn't INNER/LAST so the fbid pre-resolve picks the stanza - /// id), but the bot encrypted under `edit_target_id`. The fallback must - /// rescue. - #[tokio::test] - async fn msmsg_falls_back_to_edit_target_when_primary_uses_info_id() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - let (client, _transport) = capturing_client("msmsg_fb_to_edit").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "867051314767696@bot"; - let our_pn = "5511000000001@s.whatsapp.net"; - let outbound_id = "OUT_INV"; - let stanza_id = "REPLY_INV"; - let edit_target_id = "EDIT_INV"; - let secret = [0xBEu8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - let plaintext_msg = wa::Message { - conversation: Some("inverse fallback".to_string()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - // Encrypt under edit_target_id even though edit=first → primary - // will pick info.id (stanza id), fail, and the fallback should try - // edit_target_id and succeed (WA Web regular bot path `f()`). - let ctx = BotMessageContext { - msg_id: edit_target_id, - target_sender_user_jid: our_pn, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", stanza_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("bot") - .attr("edit", "first") - .attr("edit_target_id", edit_target_id) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == stanza_id - && msg.conversation.as_deref() == Some("inverse fallback")) - }, - 1500, - ) - .await; - assert!( - got.is_some(), - "primary attempt with info.id must fall back to edit_target_id (WA Web f())" - ); - } - - /// Mirror scenario: no `<bot edit>`, so the parser doesn't populate - /// `edit_target_id`. Primary uses `info.id`; with no fallback id - /// available, a deliberately-wrong-key payload must nack 495 (no second - /// attempt to silently mask the failure). - #[tokio::test] - async fn msmsg_no_fallback_when_no_edit_target_present() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - let (client, transport) = capturing_client("msmsg_nofb").await; - let chat = "867051314767696@bot"; - let our_pn = "5511000000001@s.whatsapp.net"; - let outbound_id = "OUT_NOFB"; - let stanza_id = "REPLY_NOFB"; - let secret = [0xCCu8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - // Encrypt under a DIFFERENT id; no <bot> node so parser leaves - // edit_target_id = None and there's nothing to fall back to. - let ctx = BotMessageContext { - msg_id: "MISMATCHED_ID", - target_sender_user_jid: our_pn, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(b"x", &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", stanza_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let mut code = None; - for _ in 0..80 { - if let Some(c) = find_message_nack_error(&transport.sent(), stanza_id) { - code = Some(c); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!( - code, - Some(495), - "no fallback id available → single AES-GCM failure must nack 495" - ); - } - - /// WA Web `processRenderableMessages` captures the embedded - /// `messageSecret` from any bot-targeted msg (fanout from us OR reply - /// from the bot). Verify the helper persists it under - /// (bot_chat, our_lid, info.id). - #[tokio::test] - async fn maybe_capture_inbound_msg_secret_persists_for_bot_chats() { - use crate::store::commands::DeviceCommand; - let (client, _transport) = capturing_client("capture_bot").await; - client - .persistence_manager - .process_command(DeviceCommand::SetLid(Some( - "999888777666555:0@lid".parse().unwrap(), - ))) - .await; - - let info = Arc::new(MessageInfo { - id: "FANOUT_1".into(), - source: crate::types::message::MessageSource { - chat: "867051314767696@bot".parse().unwrap(), - sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), - is_from_me: true, - ..Default::default() - }, - ..Default::default() - }); - let msg = wa::Message { - conversation: Some("hi bot".into()), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![0xAB; 32]), - ..Default::default() - }), - ..Default::default() - }; - client.maybe_capture_inbound_msg_secret(&msg, &info).await; - - let mut got = None; - for _ in 0..40 { - got = client - .persistence_manager - .backend() - .get_msg_secret("867051314767696@bot", "999888777666555@lid", "FANOUT_1") - .await - .unwrap(); - if got.is_some() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!(got.as_deref(), Some(&[0xABu8; 32][..])); - } - - #[tokio::test] - async fn maybe_capture_inbound_msg_secret_persists_for_non_bot_chats() { - let (client, _transport) = capturing_client("capture_regular_dm").await; - let info = Arc::new(MessageInfo { - id: "DM_1".into(), - source: crate::types::message::MessageSource { - chat: "5511777776666@s.whatsapp.net".parse().unwrap(), - sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), - is_from_me: true, - ..Default::default() - }, - ..Default::default() - }); - let msg = wa::Message { - conversation: Some("hi".into()), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![0xCD; 32]), - ..Default::default() - }), - ..Default::default() - }; - client.maybe_capture_inbound_msg_secret(&msg, &info).await; - - let got = client - .persistence_manager - .backend() - .get_msg_secret( - "5511777776666@s.whatsapp.net", - "5511000000001@s.whatsapp.net", - "DM_1", - ) - .await - .unwrap(); - assert_eq!(got.as_deref(), Some(&[0xCDu8; 32][..])); - } - - /// Group invocation: user mentions @MetaAI in a group → chat is the - /// GROUP (not bot), but mentioned_jid contains the bot. WA Web's - /// `processRenderableMessages` keys off `N` (invokedBotWid derived from - /// `mentionedJidList.find(isBot)`); we must persist too. - #[tokio::test] - async fn maybe_capture_inbound_msg_secret_persists_for_group_with_bot_mention() { - use crate::store::commands::DeviceCommand; - let (client, _transport) = capturing_client("capture_group_mention").await; - client - .persistence_manager - .process_command(DeviceCommand::SetLid(Some( - "999888777666555:0@lid".parse().unwrap(), - ))) - .await; - - let info = Arc::new(MessageInfo { - id: "GRP_MENTION".into(), - source: crate::types::message::MessageSource { - chat: "120363021033254949@g.us".parse().unwrap(), - sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), - is_from_me: true, - is_group: true, - ..Default::default() - }, - ..Default::default() - }); - let msg = wa::Message { - extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { - text: Some("hey @MetaAI tell me a joke".into()), - context_info: Some(Box::new(wa::ContextInfo { - mentioned_jid: vec!["867051314767696@bot".into()], - ..Default::default() - })), - ..Default::default() - })), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![0xEE; 32]), - ..Default::default() - }), - ..Default::default() - }; - client.maybe_capture_inbound_msg_secret(&msg, &info).await; - - let mut got = None; - for _ in 0..40 { - got = client - .persistence_manager - .backend() - .get_msg_secret( - "120363021033254949@g.us", - "5511000000001@s.whatsapp.net", - "GRP_MENTION", - ) - .await - .unwrap(); - if got.is_some() { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!( - got.as_deref(), - Some(&[0xEEu8; 32][..]), - "group invocation via @bot mention must still cache the secret" - ); - } - - /// Forwarded message with a secret must NOT be cached — matches WA Web's - /// `x.isForwarded !== true` guard. A planted forward shouldn't poison - /// the cache. - #[tokio::test] - async fn maybe_capture_inbound_msg_secret_skips_forwarded() { - let (client, _transport) = capturing_client("capture_skip_forwarded").await; - let info = Arc::new(MessageInfo { - id: "FWD_1".into(), - source: crate::types::message::MessageSource { - chat: "867051314767696@bot".parse().unwrap(), - sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), - is_from_me: false, - ..Default::default() - }, - ..Default::default() - }); - let msg = wa::Message { - extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { - text: Some("forwarded".into()), - context_info: Some(Box::new(wa::ContextInfo { - is_forwarded: Some(true), - ..Default::default() - })), - ..Default::default() - })), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![0xFF; 32]), - ..Default::default() - }), - ..Default::default() - }; - client.maybe_capture_inbound_msg_secret(&msg, &info).await; - - for _ in 0..16 { - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let got = client - .persistence_manager - .backend() - .get_msg_secret( - "867051314767696@bot", - "5511000000001@s.whatsapp.net", - "FWD_1", - ) - .await - .unwrap(); - assert!(got.is_none(), "forwarded messages must not seed the cache"); - } - - /// Our own group bot prompt carries the secret but NO mentioned_jid - /// (observed in prod: `mentions_bot=false mentioned_jids=[]`). The bot - /// invocation is signalled by `message_context_info.bot_metadata`, which - /// must let the capture fire (WA Web's `w`/`A` group-participant gate). - #[tokio::test] - async fn maybe_capture_inbound_msg_secret_via_bot_metadata_without_mention() { - let (client, _transport) = capturing_client("capture_bot_meta").await; - client - .persistence_manager - .process_command(crate::store::commands::DeviceCommand::SetLid(Some( - "999888777666555:0@lid".parse().unwrap(), - ))) - .await; - let info = Arc::new(MessageInfo { - id: "GRP_OWN_BOT".into(), - source: crate::types::message::MessageSource { - chat: "120363021033254949@g.us".parse().unwrap(), - sender: "236395184570386:0@lid".parse().unwrap(), - is_from_me: true, - is_group: true, - ..Default::default() - }, - ..Default::default() - }); - let msg = wa::Message { - extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { - text: Some("continue".into()), - // No mention at all — just bot_metadata signals the invocation. - ..Default::default() - })), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![0x7B; 32]), - bot_metadata: Some(wa::BotMetadata { - persona_id: Some("867051314767696".into()), - ..Default::default() - }), - ..Default::default() - }), - ..Default::default() - }; - client.maybe_capture_inbound_msg_secret(&msg, &info).await; - - // Group (non-bot chat) → keyed under info.source.sender (our LID in a - // LID group), which is what the bot reply's target_sender_jid echoes. - let got = client - .persistence_manager - .backend() - .get_msg_secret( - "120363021033254949@g.us", - "236395184570386@lid", - "GRP_OWN_BOT", - ) - .await - .unwrap(); - assert_eq!( - got.as_deref(), - Some(&[0x7Bu8; 32][..]), - "bot_metadata presence must let our own group prompt cache without a mention" - ); - } - - #[tokio::test] - async fn bot_only_captures_group_bot_prompt_skips_plain() { - use crate::cache_config::{CacheConfig, MsgSecretPolicy}; - let cfg = CacheConfig { - msg_secret_policy: MsgSecretPolicy::BotOnly, - ..Default::default() - }; - let client = crate::test_utils::create_test_client_with_config( - "botonly_capture", - Arc::new(crate::test_utils::MockHttpClient), - cfg, - ) - .await; - - let group = "120363021033254949@g.us"; - let sender = "5511888887777@s.whatsapp.net"; - - // A plain group message is not a bot context → skipped under BotOnly. - let plain_info = Arc::new(MessageInfo { - id: "PLAIN".into(), - source: crate::types::message::MessageSource { - chat: group.parse().unwrap(), - sender: sender.parse().unwrap(), - is_group: true, - ..Default::default() - }, - ..Default::default() - }); - let plain_msg = wa::Message { - conversation: Some("hi".into()), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![0x01; 32]), - ..Default::default() - }), - ..Default::default() - }; - client - .maybe_capture_inbound_msg_secret(&plain_msg, &plain_info) - .await; - assert!( - client - .persistence_manager - .backend() - .get_msg_secret(group, sender, "PLAIN") - .await - .unwrap() - .is_none(), - "BotOnly must skip a plain (non-bot) group message" - ); - - // A group message that invokes a bot (bot_metadata) classifies as Bot, - // so its secret is kept and the later bot reply can decrypt. - let bot_info = Arc::new(MessageInfo { - id: "BOTP".into(), - source: crate::types::message::MessageSource { - chat: group.parse().unwrap(), - sender: sender.parse().unwrap(), - is_group: true, - ..Default::default() - }, - ..Default::default() - }); - let bot_msg = wa::Message { - extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { - text: Some("continue".into()), - ..Default::default() - })), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![0x02; 32]), - bot_metadata: Some(wa::BotMetadata { - persona_id: Some("867051314767696".into()), - ..Default::default() - }), - ..Default::default() - }), - ..Default::default() - }; - client - .maybe_capture_inbound_msg_secret(&bot_msg, &bot_info) - .await; - assert_eq!( - client - .persistence_manager - .backend() - .get_msg_secret(group, sender, "BOTP") - .await - .unwrap(), - Some(vec![0x02; 32]), - "BotOnly must capture a group bot invocation" - ); - } - - /// Group flow: another participant invokes the bot, their decrypted prompt - /// carries the secret. We must key it under THE PARTICIPANT (the future - /// reply's `<meta target_sender_jid>`), not our own identity. - #[tokio::test] - async fn maybe_capture_inbound_msg_secret_keys_under_other_participant() { - let (client, _transport) = capturing_client("capture_participant").await; - let participant = "5599111112222:7@s.whatsapp.net"; - let info = Arc::new(MessageInfo { - id: "GRP_OTHER".into(), - source: crate::types::message::MessageSource { - chat: "120363021033254949@g.us".parse().unwrap(), - sender: participant.parse().unwrap(), - is_from_me: false, - is_group: true, - ..Default::default() - }, - ..Default::default() - }); - let msg = wa::Message { - extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { - text: Some("@MetaAI question".into()), - context_info: Some(Box::new(wa::ContextInfo { - mentioned_jid: vec!["867051314767696@bot".into()], - ..Default::default() - })), - ..Default::default() - })), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(vec![0x5A; 32]), - ..Default::default() - }), - ..Default::default() - }; - client.maybe_capture_inbound_msg_secret(&msg, &info).await; - - // Keyed under the participant (non-AD), NOT under our own PN/LID. - let under_participant = client - .persistence_manager - .backend() - .get_msg_secret( - "120363021033254949@g.us", - "5599111112222@s.whatsapp.net", - "GRP_OTHER", - ) - .await - .unwrap(); - assert_eq!( - under_participant.as_deref(), - Some(&[0x5Au8; 32][..]), - "another participant's prompt must key under their sender JID" - ); - } - - /// WA Web `sendAggregateReceipts`: a bot reply in a GROUP (chat not bot, - /// author is bot) must ack with a bare `<ack class="message">` - /// (sendBotInvokeResponseAcks), NOT a `<receipt>`. - #[tokio::test] - async fn bot_reply_in_group_acks_with_bare_ack_not_receipt() { - let (client, transport) = capturing_client("bot_group_ack").await; - let info = Arc::new(MessageInfo { - id: "BOT_GRP_ACK".into(), - source: crate::types::message::MessageSource { - chat: "120363021033254949@g.us".parse().unwrap(), - sender: "867051314767696@bot".parse().unwrap(), - is_from_me: false, - is_group: true, - ..Default::default() - }, - ..Default::default() - }); - client.ack_received_message(&info); - - let mut found = None; - for _ in 0..80 { - if let Some(a) = find_message_ack(&transport.sent()) { - found = Some(a); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert!( - found.is_some(), - "group bot reply must emit a bare <ack class=\"message\">" - ); - assert_eq!( - delivery_receipts_for(&transport.sent(), "BOT_GRP_ACK"), - 0, - "group bot reply must NOT emit a <receipt>" - ); - } - - /// Regression: a 1:1 bot chat (chat IS the bot) keeps the normal delivery - /// `<receipt>` — WA Web's `v` gate is false when chat.isBot(). - #[tokio::test] - async fn bot_dm_reply_keeps_delivery_receipt() { - let (client, transport) = capturing_client("bot_dm_receipt").await; - let info = Arc::new(MessageInfo { - id: "BOT_DM_RCPT".into(), - source: crate::types::message::MessageSource { - chat: "867051314767696@bot".parse().unwrap(), - sender: "867051314767696@bot".parse().unwrap(), - is_from_me: false, - ..Default::default() - }, - ..Default::default() - }); - client.ack_received_message(&info); - - let mut count = 0; - for _ in 0..80 { - count = delivery_receipts_for(&transport.sent(), "BOT_DM_RCPT"); - if count > 0 { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!( - count, 1, - "1:1 bot chat must keep the normal delivery receipt" - ); - } - - #[tokio::test] - async fn maybe_capture_inbound_msg_secret_skips_when_secret_absent() { - let (client, _transport) = capturing_client("capture_no_secret").await; - let info = Arc::new(MessageInfo { - id: "NO_SECRET".into(), - source: crate::types::message::MessageSource { - chat: "867051314767696@bot".parse().unwrap(), - sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), - is_from_me: true, - ..Default::default() - }, - ..Default::default() - }); - let msg = wa::Message { - conversation: Some("hi".into()), - ..Default::default() - }; - client.maybe_capture_inbound_msg_secret(&msg, &info).await; - - for _ in 0..16 { - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - let got = client - .persistence_manager - .backend() - .get_msg_secret( - "867051314767696@bot", - "5511000000001@s.whatsapp.net", - "NO_SECRET", - ) - .await - .unwrap(); - assert!(got.is_none()); - } - - /// A stanza carrying BOTH a valid msmsg AND an unknown sibling enc must - /// still dispatch the msmsg — the unknown-only fallback ack must not - /// short-circuit when `bot_payloads` is non-empty. - #[tokio::test] - async fn mixed_msmsg_and_unknown_enc_still_decrypts_msmsg() { - use crate::store::commands::DeviceCommand; - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - - let (client, _transport) = capturing_client("msmsg_mixed_unknown").await; - client - .persistence_manager - .process_command(DeviceCommand::SetLid(Some( - "999888777666555:0@lid".parse().unwrap(), - ))) - .await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "867051314767696@bot"; - let our_lid = "999888777666555@lid"; - let outbound_id = "OUT_MIX"; - let bot_reply_id = "REPLY_MIX"; - let secret = [0x55u8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_lid, outbound_id, &secret) - .await - .unwrap(); - - let plaintext_msg = wa::Message { - conversation: Some("mixed ok".into()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - let ctx = BotMessageContext { - msg_id: bot_reply_id, - target_sender_user_jid: our_lid, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - // Stanza has a valid msmsg PLUS an unrecognised "frskmsg" sibling. - // The fallback transport-ack must NOT fire (would drop the msmsg). - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", bot_reply_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_lid) - .build(), - NodeBuilder::new("enc") - .attr("type", "frskmsg") - .bytes(vec![0u8; 8]) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("mixed ok")) - }, - 1500, - ) - .await; - assert!( - got.is_some(), - "msmsg sibling of an unknown enc must still decrypt and dispatch" - ); - } - - /// LID↔PN migration window: the secret was stored under our PN, but the - /// bot reply's `<meta target_sender_jid>` echoes our LID. The primary - /// lookup misses; `alternate_msg_secret_lookup` resolves PN via - /// `lid_pn_mapping` and hits. Mirrors WA Web `C()`'s `getAlternateMsgKey`. - #[tokio::test] - async fn msmsg_alternate_lookup_resolves_lid_to_stored_pn() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - use wacore::store::traits::LidPnMappingEntry; - - let (client, _transport) = capturing_client("msmsg_alt_lookup").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "867051314767696@bot"; - let our_lid_user = "999888777666555"; - let our_pn_user = "5511000000001"; - let our_lid = "999888777666555@lid"; - let our_pn = "5511000000001@s.whatsapp.net"; - let outbound_id = "OUT_ALT"; - let bot_reply_id = "REPLY_ALT"; - let secret = [0x3Cu8; 32]; - - // Seed the LID→PN mapping so the alternate lookup can swap. - client - .persistence_manager - .backend() - .put_lid_mapping(&LidPnMappingEntry { - lid: our_lid_user.into(), - phone_number: our_pn_user.into(), - created_at: 0, - updated_at: 0, - learning_source: "test".into(), - }) - .await - .unwrap(); - // Secret stored under PN (as if the outbound went out PN-addressed). - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - // Bot reply encrypts with target = our LID (what <meta> declares). - let plaintext_msg = wa::Message { - conversation: Some("alt ok".into()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - let ctx = BotMessageContext { - msg_id: bot_reply_id, - target_sender_user_jid: our_lid, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", bot_reply_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_lid) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("alt ok")) - }, - 1500, - ) - .await; - assert!( - got.is_some(), - "LID-declared reply must resolve the PN-stored secret via lid_pn_mapping" - ); - } - - /// End-to-end: phone fanout dispatches a wa::Message carrying the - /// outbound `messageSecret`; later the Meta AI bot replies via msmsg - /// referencing the same id. The captured secret must let the reply - /// decrypt and surface `Event::Message`. - #[tokio::test] - async fn fanout_capture_lets_subsequent_msmsg_decrypt() { - use crate::store::commands::DeviceCommand; - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - - let (client, _transport) = capturing_client("fanout_to_msmsg").await; - client - .persistence_manager - .process_command(DeviceCommand::SetLid(Some( - "999888777666555:0@lid".parse().unwrap(), - ))) - .await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let bot_chat: Jid = "867051314767696@bot".parse().unwrap(); - let our_lid_str = "999888777666555@lid"; - let outbound_id = "FANOUT_OUT"; - let bot_reply_id = "BOT_REPLY_PHONE"; - let secret = [0x99u8; 32]; - - // Step 1: simulate the fanout dispatch (what dispatch_parsed_message - // would call when the phone's outbound stanza is mirrored to us). - let fanout_info = Arc::new(MessageInfo { - id: outbound_id.into(), - source: crate::types::message::MessageSource { - chat: bot_chat.clone(), - sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), - is_from_me: true, - ..Default::default() - }, - ..Default::default() - }); - let fanout_msg = wa::Message { - conversation: Some("hi bot".into()), - message_context_info: Some(wa::MessageContextInfo { - message_secret: Some(secret.to_vec()), - ..Default::default() - }), - ..Default::default() - }; - client - .maybe_capture_inbound_msg_secret(&fanout_msg, &fanout_info) - .await; - // Write is awaited inline now, so the secret is already durable here. - for _ in 0..40 { - if client - .persistence_manager - .backend() - .get_msg_secret("867051314767696@bot", our_lid_str, outbound_id) - .await - .unwrap() - .is_some() - { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - - // Step 2: the bot reply arrives as <enc type="msmsg"> referencing - // outbound_id via <meta target_id>. With the secret captured above, - // it must decrypt cleanly. - let plaintext_msg = wa::Message { - conversation: Some("bot reply".into()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - let ctx = BotMessageContext { - msg_id: bot_reply_id, - target_sender_user_jid: our_lid_str, - bot_user_jid: "867051314767696@bot", - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", bot_reply_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_lid_str) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("bot reply")) - }, - 1500, - ) - .await; - assert!( - got.is_some(), - "secret captured from fanout must enable msmsg reply decryption" - ); - } - - /// Coherence: the identity `persist_outbound_msg_secret` writes under - /// (LID for bot chats) must match what `handle_msmsg_payload` reads via - /// `<meta target_sender_jid>`. End-to-end without bypassing the helper. - #[tokio::test] - async fn msmsg_outbound_put_and_inbound_get_match_for_lid_bot() { - use crate::store::commands::DeviceCommand; - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - - let (client, _transport) = capturing_client("msmsg_lid_match").await; - // Seed both PN (already seeded by capturing_client) and LID. - client - .persistence_manager - .process_command(DeviceCommand::SetLid(Some( - "999888777666555:0@lid".parse().unwrap(), - ))) - .await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let bot_chat: Jid = "867051314767696@bot".parse().unwrap(); - let outbound_id = "OUT_LID"; - let bot_reply_id = "REPLY_LID"; - let our_lid = "999888777666555@lid"; - let secret = [0x71u8; 32]; - - // Real outbound path: caller resolves the bot identity to our LID. - let sender_identity = client - .dm_sender_identity_for(&bot_chat) - .await - .expect("LID seeded"); - client - .persist_outbound_msg_secret( - &bot_chat, - &sender_identity, - outbound_id, - &secret, - wacore::msg_secret::RetentionClass::Bot, - ) - .await; - - // Inbound msmsg payload encrypted under the same (msg_id, target, bot) - // tuple the meta will declare on the wire. - let plaintext_msg = wa::Message { - conversation: Some("lid coherent".to_string()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - let ctx = BotMessageContext { - msg_id: bot_reply_id, - target_sender_user_jid: our_lid, - bot_user_jid: "867051314767696@bot", - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", bot_reply_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_lid) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("lid coherent")) - }, - 1500, - ) - .await; - assert!( - got.is_some(), - "outbound PUT and inbound GET must converge on LID for bot chats" - ); - } - - /// Regression for the AD_JID encoder bug: a `from="USER:0@bot"` stanza must - /// survive the marshal/unmarshal round-trip with `server=Bot`, so the - /// secret lookup keys hit and the reply decrypts. - #[tokio::test] - async fn msmsg_with_bot_device_suffix_round_trips() { - use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; - let (client, _transport) = capturing_client("msmsg_bot_device").await; - let collector = Arc::new(crate::test_utils::TestEventCollector::default()); - client.register_handler(collector.clone()); - - let chat = "867051314767696@bot"; - let our_pn = "5511000000001@s.whatsapp.net"; - let outbound_id = "OUTBOUND_DEV"; - let bot_reply_id = "BOT_REPLY_DEV"; - let secret = [0x33u8; 32]; - - client - .persistence_manager - .backend() - .put_msg_secret(chat, our_pn, outbound_id, &secret) - .await - .unwrap(); - - let plaintext_msg = wa::Message { - conversation: Some("with device".to_string()), - ..Default::default() - }; - let pt_bytes = { - use prost::Message as _; - let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); - plaintext_msg.encode(&mut v).unwrap(); - v - }; - let ctx = BotMessageContext { - msg_id: bot_reply_id, - target_sender_user_jid: our_pn, - bot_user_jid: chat, - }; - let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); - let ms_msg_proto = encode_message_secret_message(&iv, &cipher); - - let node = NodeBuilder::new("message") - .attr("from", "867051314767696:0@bot") - .attr("id", bot_reply_id) - .attr("type", "text") - .children([ - NodeBuilder::new("meta") - .attr("target_id", outbound_id) - .attr("target_sender_jid", our_pn) - .build(), - NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build(), - ]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let got = collect_event( - &client, - collector, - |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("with device")) - }, - 1500, - ) - .await; - assert!( - got.is_some(), - "msmsg with `:0@bot` from must round-trip (encoder must not strip the bot server)" - ); - } - - /// `<meta>` without `target_id` → cannot identify the parent message, - /// nack 495 and no dispatch. - #[tokio::test] - async fn msmsg_without_meta_target_id_nacks_495() { - let (client, transport) = capturing_client("msmsg_no_target").await; - let bot_reply_id = "BOT_REPLY_NT"; - let ms_msg_proto = encode_message_secret_message(&[0u8; 12], &[0u8; 32]); - let node = NodeBuilder::new("message") - .attr("from", "867051314767696@bot") - .attr("id", bot_reply_id) - .attr("type", "text") - .children([NodeBuilder::new("enc") - .attr("type", "msmsg") - .attr("v", "2") - .bytes(ms_msg_proto) - .build()]) - .build(); - let owned = node_to_arc(node); - client.clone().handle_incoming_message(owned).await; - - let mut code = None; - for _ in 0..80 { - if let Some(c) = find_message_nack_error(&transport.sent(), bot_reply_id) { - code = Some(c); - break; - } - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } - assert_eq!(code, Some(495)); - } -} +mod tests; diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs new file mode 100644 index 000000000..b8f57461b --- /dev/null +++ b/src/message/dispatch.rs @@ -0,0 +1,69 @@ +//! Post-decrypt dispatch: event emission, acks and delivery receipts. + +use super::*; + +impl Client { + /// Dispatches a successfully parsed message to the event bus and sends a delivery receipt. + pub(crate) async fn dispatch_parsed_message( + self: &Arc<Self>, + msg: wa::Message, + info: &Arc<MessageInfo>, + ) { + use wacore::proto_helpers::MessageExt; + + let mut info = Arc::clone(info); + if info.ephemeral_expiration.is_none() + && msg.get_base_message().get_ephemeral_expiration().is_some() + { + Arc::make_mut(&mut info).ephemeral_expiration = + msg.get_base_message().get_ephemeral_expiration(); + } + + // Keep this ordered with dispatch; add-on messages can immediately + // reference the secret from the stanza just processed. + self.maybe_capture_inbound_msg_secret(&msg, &info).await; + let dispatch_msg = self + .maybe_decrypt_secret_encrypted_message(&msg, &info) + .await + .unwrap_or(msg); + self.ack_received_message(&info); + + self.core + .event_bus + .dispatch(Event::Message(Arc::new(dispatch_msg), info)); + } + + /// Acknowledge a received message so the server drops it from the offline + /// queue: a delivery receipt when applicable (incl. the `type="sender"` + /// receipt for own-account self-fanouts), else a transport ack. status is + /// acked by the `should_ack` gate, newsletters/empty ids need nothing here. + pub(crate) fn ack_received_message(self: &Arc<Self>, info: &Arc<MessageInfo>) { + if info.id.is_empty() || info.source.chat.is_newsletter() { + return; + } + // WA Web `sendAggregateReceipts`: for a DELIVERY where the chat is NOT + // a bot but the author IS a bot (a bot reply inside a group), it emits + // a bare `<ack class="message">` via `sendBotInvokeResponseAcks`, not a + // `<receipt>`. A 1:1 bot chat keeps the normal receipt (chat.isBot() → + // the branch's `v` is false). Our transport ack is that bare + // `<ack class="message">` (group form carries `participant`). + if info.source.is_bot_authored_non_bot_chat() { + self.spawn_message_ack(info); + return; + } + if Self::should_send_delivery_receipt(info) { + self.spawn_delivery_receipt(info); + } else if !info.source.chat.is_status_broadcast() { + self.spawn_message_ack(info); + } + } + + /// Spawn a delivery receipt, tracked so `disconnect()` can flush it (issue #571). + fn spawn_delivery_receipt(self: &Arc<Self>, info: &Arc<MessageInfo>) { + let client = self.clone(); + let info = Arc::clone(info); + self.outbound_flush.spawn(&*self.runtime, async move { + client.send_delivery_receipt(&info).await; + }); + } +} diff --git a/src/message/msg_secret.rs b/src/message/msg_secret.rs new file mode 100644 index 000000000..1637bbedd --- /dev/null +++ b/src/message/msg_secret.rs @@ -0,0 +1,744 @@ +//! Inbound msg-secret capture and secret-encrypted message decryption. + +use super::*; + +impl Client { + /// Capture embedded `MessageContextInfo.message_secret` for add-on + /// decrypts. Bot DMs keep the legacy LID key as a second entry. + pub(crate) async fn maybe_capture_inbound_msg_secret( + self: &Arc<Self>, + msg: &wa::Message, + info: &Arc<MessageInfo>, + ) { + use wacore::proto_helpers::MessageExt; + + let mci = msg.message_context_info.as_ref(); + let Some(secret_bytes) = mci.and_then(|m| m.message_secret.as_deref()) else { + return; + }; + if msg.is_forwarded() { + return; + } + + let policy = self.cache_config.msg_secret_policy; + if !policy.persists() { + return; + } + let chat_is_bot = info.source.chat.is_bot(); + // BotOnly enforcement lives in build_msg_secret_entry (the chokepoint), + // which keys off the classified bot context including group bot prompts. + let class = wacore::msg_secret::classify(msg, chat_is_bot); + let message_ts = u64::try_from(info.timestamp.timestamp()).ok(); + + // Build both aliases (primary, plus the bot-DM LID key) and write them + // in one batch so a partial write can't leave only one stored. + let mut entries = Vec::with_capacity(2); + if let Some(entry) = self.build_msg_secret_entry( + &info.source.chat, + &info.source.sender, + &info.id, + secret_bytes, + class, + message_ts, + ) { + entries.push(entry); + } + if chat_is_bot + && let Some(sender) = self.dm_sender_identity_for(&info.source.chat).await + && sender.to_non_ad() != info.source.sender.to_non_ad() + && let Some(entry) = self.build_msg_secret_entry( + &info.source.chat, + &sender, + &info.id, + secret_bytes, + class, + message_ts, + ) + { + entries.push(entry); + } + self.persist_msg_secret_entries(entries).await; + } + + /// Build one retention entry, applying the policy gates and computing the + /// per-row deadline. Returns `None` when the policy skips this write (not + /// persisting, or `BotOnly` and the class isn't `Bot`) or the secret isn't + /// 32 bytes. Pure (no I/O) so callers can batch several aliases atomically. + fn build_msg_secret_entry( + &self, + chat: &Jid, + sender: &Jid, + msg_id: &str, + secret_bytes: &[u8], + class: wacore::msg_secret::RetentionClass, + message_ts: Option<u64>, + ) -> Option<wacore::store::traits::MsgSecretEntry> { + const SECRET_LEN: usize = wacore::reporting_token::MESSAGE_SECRET_SIZE; + let secret = <&[u8; SECRET_LEN]>::try_from(secret_bytes).ok()?; + let policy = self.cache_config.msg_secret_policy; + if !policy.persists() { + return None; + } + // Single chokepoint for the BotOnly invariant: only bot-context secrets + // (class == Bot) are persisted, no matter which write path got here. + if policy.bot_only() && class != wacore::msg_secret::RetentionClass::Bot { + return None; + } + let expires_at = wacore::msg_secret::expires_at( + policy, + &self.cache_config.msg_secret_retention, + class, + message_ts, + wacore::time::now_secs(), + ); + Some(wacore::store::traits::MsgSecretEntry { + chat: chat.to_non_ad_string(), + sender: sender.to_non_ad_string(), + msg_id: msg_id.to_string(), + secret: secret.to_vec(), + expires_at, + message_ts: message_ts.and_then(|t| i64::try_from(t).ok()).unwrap_or(0), + }) + } + + /// Write a batch of secret aliases in one atomic upsert, so a multi-alias + /// capture/re-persist never leaves only some aliases stored. + async fn persist_msg_secret_entries( + &self, + entries: Vec<wacore::store::traits::MsgSecretEntry>, + ) -> bool { + if entries.is_empty() { + return false; + } + match self + .persistence_manager + .backend() + .put_msg_secrets(entries) + .await + { + Ok(_) => true, + Err(e) => { + log::warn!("failed to persist messageSecrets: {e:?}"); + false + } + } + } + + async fn own_jid_for_secret_encrypted(&self, info: &MessageInfo) -> Option<Jid> { + use wacore::types::message::AddressingMode; + + if info.source.is_from_me { + return Some(info.source.sender.to_non_ad()); + } + + match info.source.addressing_mode { + Some(AddressingMode::Lid) => match self.get_lid().await { + Some(jid) => Some(jid), + None => self.get_pn().await, + }, + Some(AddressingMode::Pn) => match self.get_pn().await { + Some(jid) => Some(jid), + None => self.get_lid().await, + }, + None if info.source.sender.is_lid() || info.source.chat.is_lid() => { + match self.get_lid().await { + Some(jid) => Some(jid), + None => self.get_pn().await, + } + } + None => match self.get_pn().await { + Some(jid) => Some(jid), + None => self.get_lid().await, + }, + } + } + + pub(crate) async fn maybe_decrypt_secret_encrypted_message( + self: &Arc<Self>, + msg: &wa::Message, + info: &Arc<MessageInfo>, + ) -> Option<wa::Message> { + use crate::features::message_edit::{self, SecretEncKind}; + + let env = message_edit::extract_secret_encrypted(msg)?; + let target_id = env.target_id()?; + + let my_jid = self.own_jid_for_secret_encrypted(info).await?; + let original_sender = match env.original_sender_for_dispatch( + info.source.is_from_me, + &info.source.sender, + &my_jid, + ) { + Ok(jid) => jid, + Err(_) => return None, + }; + + let backend = self.persistence_manager.backend(); + let chat_for_lookup = info.source.chat.to_non_ad_string(); + let original_sender_str = original_sender.to_non_ad_string(); + let fallback_original_sender = self + .alternate_msg_secret_jid(&backend, &original_sender) + .await + .unwrap_or_default(); + + // Look up the secret AND the parent's event time (for the edit window + // check below): primary sender, then the LID/PN alternate. + let store_secret = match backend + .get_msg_secret_with_ts(&chat_for_lookup, &original_sender_str, target_id) + .await + { + Ok(Some(found)) => Some(found), + Ok(None) => match fallback_original_sender.as_ref() { + Some(alt) => { + let alt_str = alt.to_non_ad_string(); + match backend + .get_msg_secret_with_ts(&chat_for_lookup, &alt_str, target_id) + .await + { + Ok(found) => found, + Err(e) => { + log::warn!( + "[msg:{}] secret_encrypted_message alternate secret lookup failed: {e:?}", + info.id + ); + None + } + } + } + None => None, + }, + Err(e) => { + log::warn!( + "[msg:{}] backend error reading secret_encrypted_message secret: {e:?}", + info.id + ); + None + } + }; + // On a total store miss, ask the app-supplied resolver (if any) for the + // parent secret. This is what lets the Disabled policy still decrypt. The + // resolver carries no parent timestamp, so parent_ts stays 0 (unknown). + let (secret, parent_ts) = match store_secret { + Some((secret, ts)) => (secret, ts), + None => { + let alternate = fallback_original_sender + .as_ref() + .map(|j| j.to_non_ad_string()); + match self + .resolve_msg_secret_via_app( + &chat_for_lookup, + &original_sender_str, + alternate.as_deref(), + target_id, + ) + .await + { + Some(secret) => (secret, 0), + None => return None, + } + } + }; + + let fallback_editor = match info.source.sender_alt.clone() { + Some(jid) => Some(jid), + None => self + .alternate_msg_secret_jid(&backend, &info.source.sender) + .await + .unwrap_or_default(), + }; + + let inner = match message_edit::decrypt_secret_encrypted( + env.enc_payload, + env.enc_iv, + &secret, + env.kind, + target_id, + &original_sender, + &info.source.sender, + ) { + Ok(inner) => inner, + Err(primary_err) => { + let mut last_err = primary_err; + let mut decrypted = None; + + if let Some(fallback_original) = fallback_original_sender.as_ref() { + match message_edit::decrypt_secret_encrypted( + env.enc_payload, + env.enc_iv, + &secret, + env.kind, + target_id, + fallback_original, + &info.source.sender, + ) { + Ok(inner) => decrypted = Some(inner), + Err(e) => last_err = e, + } + } + + if decrypted.is_none() + && let Some(fallback_editor) = fallback_editor.as_ref() + { + match message_edit::decrypt_secret_encrypted( + env.enc_payload, + env.enc_iv, + &secret, + env.kind, + target_id, + &original_sender, + fallback_editor, + ) { + Ok(inner) => decrypted = Some(inner), + Err(e) => last_err = e, + } + } + + if decrypted.is_none() + && let (Some(fallback_original), Some(fallback_editor)) = + (fallback_original_sender.as_ref(), fallback_editor.as_ref()) + { + match message_edit::decrypt_secret_encrypted( + env.enc_payload, + env.enc_iv, + &secret, + env.kind, + target_id, + fallback_original, + fallback_editor, + ) { + Ok(inner) => decrypted = Some(inner), + Err(e) => last_err = e, + } + } + + match decrypted { + Some(inner) => inner, + None => { + log::warn!( + "[msg:{}] secret_encrypted_message {:?} decrypt failed: {last_err:?}", + info.id, + env.kind + ); + return None; + } + } + } + }; + + // Mirror WA Web `ProcessEditProtocolMsgs`: drop a MESSAGE_EDIT authored + // outside the parent's edit-processing window (editTs >= parentTs + 20m). + // The check is on authored time, not "now", so a validly-authored edit + // still applies after an offline delivery gap. Only enforceable when we + // know the parent's event time; resolver-supplied secrets carry none + // (parent_ts == 0), so we stay permissive there. + if env.kind == SecretEncKind::MessageEdit && parent_ts > 0 { + let edit_ts = info.timestamp.timestamp(); + if edit_ts >= parent_ts + wacore::msg_secret::EDIT_PROCESSING_WINDOW_SECS { + log::debug!( + "[msg:{}] secret edit authored outside the {}s window (editTs={edit_ts}, parentTs={parent_ts}); dropping", + info.id, + wacore::msg_secret::EDIT_PROCESSING_WINDOW_SECS + ); + return None; + } + } + + if let Some(secret_bytes) = inner + .message_context_info + .as_ref() + .and_then(|m| m.message_secret.as_deref()) + { + // The re-persisted secret keys the NEXT add-on on the same parent, + // so its retention class follows the parent kind and the parent's own + // event time (when known) rather than this edit's arrival time. + let class = match env.kind { + SecretEncKind::MessageEdit => wacore::msg_secret::RetentionClass::Text, + _ => wacore::msg_secret::RetentionClass::PollEvent, + }; + let message_ts = if parent_ts > 0 { + u64::try_from(parent_ts).ok() + } else { + u64::try_from(info.timestamp.timestamp()).ok() + }; + // Primary + LID/PN alternate in one batch so both survive together. + let mut entries = Vec::with_capacity(2); + if let Some(entry) = self.build_msg_secret_entry( + &info.source.chat, + &original_sender, + target_id, + secret_bytes, + class, + message_ts, + ) { + entries.push(entry); + } + if let Some(alternate_sender) = fallback_original_sender.as_ref() + && let Some(entry) = self.build_msg_secret_entry( + &info.source.chat, + alternate_sender, + target_id, + secret_bytes, + class, + message_ts, + ) + { + entries.push(entry); + } + self.persist_msg_secret_entries(entries).await; + } + + if env.kind != SecretEncKind::MessageEdit { + return Some(inner); + } + + match message_edit::rewrap_as_legacy_edit(inner) { + Some(rewrapped) => Some(rewrapped), + None => { + log::warn!( + "[msg:{}] decrypted MESSAGE_EDIT missing protocol_message.edited_message", + info.id + ); + None + } + } + } + + /// Decrypt and dispatch a `<enc type="msmsg">` bot reply. Looks up the + /// outbound `messageSecret` we persisted at send time and runs the + /// dual-HKDF + AES-GCM open from [`wacore::bot_message`]. Failures + /// (missing secret, GCM tag fail, malformed proto) nack with code 495. + pub(crate) async fn handle_msmsg_payload( + self: &Arc<Self>, + info: &Arc<MessageInfo>, + payload: EncPayload, + ) { + use prost::Message as _; + use wa::MessageSecretMessage; + use wacore::bot_message::{BotMessageContext, decrypt_bot_message}; + use wacore::protocol::nack::NackReason; + + let ms_msg = match MessageSecretMessage::decode(&*payload.ciphertext) { + Ok(m) => m, + Err(e) => { + log::warn!( + "[msg:{}] failed to decode MessageSecretMessage: {e:?}", + info.id + ); + self.spawn_nack(info, NackReason::ParsingError, None); + return; + } + }; + let (Some(enc_iv), Some(enc_payload)) = + (ms_msg.enc_iv.as_deref(), ms_msg.enc_payload.as_deref()) + else { + log::warn!( + "[msg:{}] MessageSecretMessage missing enc_iv/enc_payload", + info.id + ); + self.spawn_nack(info, NackReason::ParsingError, None); + return; + }; + + // Target sender (us): meta echoes our LID/PN. Falls back to our LID + // when sender is on the bot server, our PN otherwise (whatsmeow + // `decryptBotMessage`). + let target_sender = match self.resolve_msmsg_target_sender(info).await { + Some(j) => j, + None => { + log::warn!("[msg:{}] msmsg: no target_sender resolvable", info.id); + self.spawn_nack(info, NackReason::MissingMessageSecret, None); + return; + } + }; + + // Chat scope for the secret lookup: prefer <meta target_chat_jid>; + // fall back to the stanza's chat (matches WA Web `decryptMsmsgBotMessage`). + let chat_for_lookup = info + .meta_info + .target_chat + .as_ref() + .unwrap_or(&info.source.chat) + .to_non_ad() + .to_string(); + let target_sender_str = target_sender.to_non_ad_string(); + + // The id used for the SECRET LOOKUP is `meta.target_id` (our outbound + // id); the id used as HKDF input is the bot reply id (or + // `bot_info.edit_target_id` when the bot is editing a prior reply). + let target_id = match info.meta_info.target_id.as_deref() { + Some(id) => id, + None => { + log::warn!( + "[msg:{}] msmsg: <meta> missing target_id; cannot look up secret", + info.id + ); + self.spawn_nack(info, NackReason::MissingMessageSecret, None); + return; + } + }; + + // Mirror WA Web `C()` in `WAWebBotMessageSecret.js`: primary lookup + // plus an alternate (PN ↔ LID swap via lid_pn_mapping) so a row + // stored under one identity family is still found if `<meta + // target_sender_jid>` echoes the other. Covers LID migration windows + // and asymmetric outbound/inbound identities. + let backend = self.persistence_manager.backend(); + // Store lookup: primary, then the LID/PN alternate. A backend error is + // logged and treated as a miss (not a hard nack) so the resolver still + // gets a chance — mirrors the secret-encrypted edit path. + let store_secret = match backend + .get_msg_secret(&chat_for_lookup, &target_sender_str, target_id) + .await + { + Ok(Some(s)) => Some(s), + Ok(None) => match self + .alternate_msg_secret_lookup(&backend, &chat_for_lookup, &target_sender, target_id) + .await + { + Ok(found) => found, + Err(e) => { + log::warn!("[msg:{}] msmsg: alternate lookup failed: {e:?}", info.id); + None + } + }, + Err(e) => { + log::warn!( + "[msg:{}] backend error reading message_secret: {e:?}", + info.id + ); + None + } + }; + let secret = match store_secret { + Some(s) => s, + None => { + let alternate = self + .alternate_msg_secret_jid(&backend, &target_sender) + .await + .ok() + .flatten() + .map(|j| j.to_non_ad_string()); + match self + .resolve_msg_secret_via_app( + &chat_for_lookup, + &target_sender_str, + alternate.as_deref(), + target_id, + ) + .await + { + Some(s) => s, + None => { + // For a group bot invocation initiated by our PRIMARY + // device, the messageSecret lives in the bot-addressed + // copy the primary sent directly to the bot — it is NOT + // mirrored to companions in the group skmsg. So a + // companion legitimately never holds the secret; this + // miss is expected and benign (we nack 495 and the server + // stops replaying). A miss in a 1:1 bot chat is unexpected + // and worth a warn. + log::log!( + if info.source.is_group { + log::Level::Debug + } else { + log::Level::Warn + }, + "[msg:{}] msmsg: no message_secret stored for target_id={target_id} (primary or alternate)", + info.id + ); + self.spawn_nack(info, NackReason::MissingMessageSecret, None); + return; + } + } + } + }; + + let bot_user_jid = info.source.sender.to_non_ad_string(); + // WA Web `decryptMsmsgBotMessage` dispatches on `isFbidBot()`: + // * fbid path pre-resolves to `edit_target_id` for INNER/LAST edits, + // `externalId` (info.id) otherwise. Single AES-GCM attempt. + // * regular path tries `externalId` first, falls back to + // `edit_target_id` on AES-GCM failure. + // We don't have `isFbidBot()` detection; instead, we unify the two as + // try-then-fallback with the fbid-style id as primary. That's a strict + // superset: for INNER/LAST it usually succeeds on the first try (fbid + // outcome); for any other case primary is `info.id` so we mirror the + // regular path's first attempt. The fallback is only attempted if + // `bot_info.edit_target_id` is present. + let info_id = info.id.as_str(); + let primary_msg_id = info + .bot_info + .as_ref() + .filter(|bi| { + matches!( + bi.edit_type, + Some( + crate::types::message::BotEditType::Inner + | crate::types::message::BotEditType::Last + ) + ) + }) + .and_then(|bi| bi.edit_target_id.as_deref()) + .unwrap_or(info_id); + let fallback_msg_id = if primary_msg_id == info_id { + info.bot_info + .as_ref() + .and_then(|bi| bi.edit_target_id.as_deref()) + } else { + Some(info_id) + } + .filter(|fb| *fb != primary_msg_id); + + let attempt = |msg_id: &str| { + let ctx = BotMessageContext { + msg_id, + target_sender_user_jid: &target_sender_str, + bot_user_jid: &bot_user_jid, + }; + decrypt_bot_message(&secret, enc_iv, enc_payload, &ctx) + }; + + let plaintext = match attempt(primary_msg_id) { + Ok(p) => p, + Err(primary_err) => match fallback_msg_id { + Some(fb) => match attempt(fb) { + Ok(p) => p, + Err(fallback_err) => { + log::warn!( + "[msg:{}] msmsg AES-GCM open failed both attempts (primary={primary_err:?}, fallback={fallback_err:?})", + info.id + ); + self.spawn_nack(info, NackReason::MissingMessageSecret, None); + return; + } + }, + None => { + log::warn!( + "[msg:{}] msmsg AES-GCM open failed and no fallback msg_id: {primary_err:?}", + info.id + ); + self.spawn_nack(info, NackReason::MissingMessageSecret, None); + return; + } + }, + }; + + let msg = match wa::Message::decode(plaintext.as_slice()) { + Ok(m) => m, + Err(e) => { + log::warn!( + "[msg:{}] msmsg plaintext is not a Message proto: {e:?}", + info.id + ); + self.spawn_nack(info, NackReason::ParsingError, None); + return; + } + }; + + log::info!( + "[msg:{}] Successfully decrypted msmsg bot reply from {}", + info.id, + info.source.sender + ); + self.dispatch_parsed_message(msg, info).await; + } + + /// Resolve `target_sender` for a msmsg stanza: echo from `<meta>` when + /// present, else fall back to our LID (sender on bot server) or PN. + async fn resolve_msmsg_target_sender(&self, info: &Arc<MessageInfo>) -> Option<Jid> { + if let Some(ts) = info.meta_info.target_sender.as_ref() { + return Some(ts.clone()); + } + if info.source.sender.server == wacore_binary::Server::Bot { + self.get_lid().await + } else { + self.get_pn().await + } + } + + /// Second-chance lookup with the alternate identity family. Mirrors + /// `WAWebLidMigrationUtils.getAlternateMsgKey`: swap PN ↔ LID via the + /// `lid_pn_mapping` store and retry. Returns `Ok(None)` when no mapping + /// is known or the alternate row is absent — the caller treats that as + /// a terminal miss. + async fn alternate_msg_secret_jid( + &self, + backend: &Arc<dyn crate::store::traits::Backend>, + primary_sender: &Jid, + ) -> Result<Option<Jid>, crate::store::error::StoreError> { + let alternate = match primary_sender.server { + wacore_binary::Server::Lid => backend + .get_lid_mapping(&primary_sender.user) + .await? + .map(|m| Jid::new(m.phone_number, wacore_binary::Server::Pn)), + wacore_binary::Server::Pn => backend + .get_pn_mapping(&primary_sender.user) + .await? + .map(|m| Jid::new(m.lid, wacore_binary::Server::Lid)), + _ => None, + }; + Ok(alternate) + } + + async fn alternate_msg_secret_lookup( + &self, + backend: &Arc<dyn crate::store::traits::Backend>, + chat_for_lookup: &str, + primary_sender: &Jid, + target_id: &str, + ) -> Result<Option<Vec<u8>>, crate::store::error::StoreError> { + let Some(alternate) = self + .alternate_msg_secret_jid(backend, primary_sender) + .await? + else { + return Ok(None); + }; + let alternate_str = alternate.to_non_ad_string(); + backend + .get_msg_secret(chat_for_lookup, &alternate_str, target_id) + .await + } + + /// On a total store miss, consult the app-supplied resolver for the parent + /// secret, trying the primary then the LID/PN alternate sender. Bounded by a + /// timeout because it runs inside the per-chat receive lane, so a slow app + /// callback degrades to a miss instead of stalling the chat. + async fn resolve_msg_secret_via_app( + &self, + chat: &str, + primary_sender: &str, + alternate_sender: Option<&str>, + msg_id: &str, + ) -> Option<Vec<u8>> { + let resolver = self.cache_config.original_message_resolver.as_ref()?; + let lookup = async { + if let Some(secret) = resolver + .resolve_msg_secret(chat, primary_sender, msg_id) + .await + { + return Some(secret); + } + if let Some(alt) = alternate_sender + && alt != primary_sender + && let Some(secret) = resolver.resolve_msg_secret(chat, alt, msg_id).await + { + return Some(secret); + } + None + }; + match wacore::runtime::timeout( + &*self.runtime, + self.cache_config.msg_secret_resolver_timeout, + lookup, + ) + .await + { + Ok(Some(secret)) => Some(secret.to_vec()), + Ok(None) => None, + Err(_) => { + log::warn!("[msg:{msg_id}] original_message_resolver timed out"); + None + } + } + } +} diff --git a/src/message/receive.rs b/src/message/receive.rs new file mode 100644 index 000000000..98c6a6531 --- /dev/null +++ b/src/message/receive.rs @@ -0,0 +1,1528 @@ +//! Core incoming-message pipeline: classify, decrypt and process. + +use super::*; + +impl Client { + pub(crate) async fn handle_incoming_message(self: Arc<Self>, node: Arc<OwnedNodeRef>) { + // Phase 1: classify borrows the node tree, extracts owned payloads, returns quickly. + // Phase 2: process_classified_message holds no node borrows across heavy .await points, + // keeping the async state machine small. + let classified = match self.classify_incoming_message(&node).await { + Some(c) => c, + None => return, + }; + // node is no longer borrowed here -- drop it before the heavy phase + drop(node); + self.process_classified_message(classified).await; + } + + pub(crate) async fn classify_incoming_message( + self: &Arc<Self>, + node: &OwnedNodeRef, + ) -> Option<ClassifiedMessage> { + let nr = node.get(); + let info = match self.parse_message_info(nr).await { + Ok(info) => Arc::new(info), + Err(e) => { + let id = nr.get_attr("id").map(|v| v.as_str()); + let from = nr.get_attr("from").map(|v| v.as_str()); + log::warn!("Failed to parse message info (id={id:?}, from={from:?}): {e:?}"); + return None; + } + }; + + // Newsletters use <plaintext> instead of <enc> because they are not E2E encrypted. + if info.source.chat.is_newsletter() { + self.handle_newsletter_message(nr, &info).await; + return None; + } + + self.cache_lid_pn_from_message( + &info.source.sender, + info.source.sender_alt.as_ref(), + info.is_offline, + ) + .await; + let sender_encryption_jid = self.resolve_encryption_jid(&info.source.sender).await; + + let unavailable_node = nr.get_optional_child("unavailable"); + + let mut all_enc_nodes: Vec<&NodeRef<'_>> = Vec::with_capacity(4); + + let direct_enc_nodes = nr.get_children_by_tag("enc"); + all_enc_nodes.extend(direct_enc_nodes); + + let participants = nr.get_optional_child_by_tag(&["participants"]); + if let Some(participants_node) = participants { + let own_jid = self.get_pn().await; + let to_nodes = participants_node.get_children_by_tag("to"); + for to_node in to_nodes { + let to_jid = match to_node.attrs().optional_jid("jid") { + Some(jid) => jid, + None => continue, + }; + if own_jid.as_ref().is_some_and(|ours| *ours == to_jid) { + let enc_children = to_node.get_children_by_tag("enc"); + all_enc_nodes.extend(enc_children); + } + } + } + + if all_enc_nodes.is_empty() && unavailable_node.is_none() { + log::warn!( + "[msg:{}] Received non-newsletter message without <enc> child: {}", + info.id, + nr.tag + ); + return None; + } + + if let Some(unavailable) = unavailable_node + && all_enc_nodes.is_empty() + { + let unavailable_type = match unavailable.get_attr("type").map(|v| v.as_str()).as_deref() + { + Some("view_once") => crate::types::events::UnavailableType::ViewOnce, + _ => crate::types::events::UnavailableType::Unknown, + }; + log::info!( + "[msg:{}] Message has <unavailable> child (type: {:?}), requesting from phone via PDO", + info.id, + unavailable_type + ); + // PDO is the only recovery here (no retry receipt), so run it before + // the transport ack in one flush task: the ack must not clear the + // offline queue before the PDO request goes out. status is acked by + // the should_ack gate. Mirrors whatsmeow's request-then-ack. + self.dispatch_undecryptable_event( + Arc::clone(&info), + true, + unavailable_type, + crate::types::events::DecryptFailMode::Show, + ) + .await; + let client = Arc::clone(self); + let info2 = Arc::clone(&info); + let skip_ack = info.source.chat.is_status_broadcast(); + self.outbound_flush.spawn(&*self.runtime, async move { + // Only ack once the PDO request is out (or skipped as ancient); + // a transient send failure leaves it queued for redelivery. + let pdo_sent = client.run_pdo_request(&info2).await; + if !skip_ack && pdo_sent { + client.send_transport_ack(&info2).await; + } + }); + return None; + } + + let mut session_payloads = Vec::with_capacity(all_enc_nodes.len()); + let mut group_payloads = Vec::with_capacity(all_enc_nodes.len()); + let mut bot_payloads = Vec::with_capacity(all_enc_nodes.len()); + let mut max_sender_retry_count: u8 = 0; + let mut has_hide_fail = false; + let mut had_unknown_enc = false; + let mut had_custom_handler = false; + + for enc_node in &all_enc_nodes { + // Parse sender retry count (WA Web: e.maybeAttrInt("count") ?? 0) + // Clamp to MAX_DECRYPT_RETRIES to prevent u64→u8 truncation on unexpected values. + let sender_count = enc_node + .attrs() + .optional_u64("count") + .map(|c| c.min(MAX_DECRYPT_RETRIES as u64) as u8) + .unwrap_or(0); + max_sender_retry_count = max_sender_retry_count.max(sender_count); + + // Parse decrypt-fail attribute (WA Web: e.maybeAttrString("decrypt-fail") === "hide") + if enc_node + .get_attr("decrypt-fail") + .map(|v| v.as_str()) + .is_some_and(|s| s == "hide") + { + has_hide_fail = true; + } + + let enc_type = match enc_node.attrs().optional_string("type") { + Some(t) => t, + None => { + log::warn!("Enc node missing 'type' attribute, skipping"); + had_unknown_enc = true; + continue; + } + }; + + if let Some(handler) = self + .custom_enc_handlers + .read() + .await + .get(enc_type.as_ref()) + .cloned() + { + let handler_clone = handler; + let client_clone = self.clone(); + let info_arc = Arc::clone(&info); + // Custom enc handlers take &Node (public API); convert from NodeRef here. + let enc_node_owned = (*enc_node).to_owned(); + let enc_type_owned = enc_type.to_string(); + + self.runtime + .spawn(Box::pin(async move { + if let Err(e) = handler_clone + .handle(client_clone, &enc_node_owned, &info_arc) + .await + { + log::warn!( + "Custom handler for enc type '{}' failed: {e:?}", + enc_type_owned + ); + } + })) + .detach(); + had_custom_handler = true; + continue; + } + + // `had_unknown_enc` means "produced no usable payload": either the + // type is unrecognized or it's known but the body is empty. + // Either way the stanza needs the fallback ack or the server replays. + if EncType::from_wire(enc_type.as_ref()).is_none() { + log::warn!("Enc node has unknown type: {enc_type}"); + had_unknown_enc = true; + continue; + } + + let payload = match EncPayload::from_owned_node(node, enc_node) { + Some(p) => p, + None => { + log::warn!("Enc node {enc_type} has no content"); + had_unknown_enc = true; + continue; + } + }; + + if payload.enc_type.is_bot_secret() { + bot_payloads.push(payload); + } else if payload.enc_type.is_session() { + session_payloads.push(payload); + } else { + group_payloads.push(payload); + } + } + + // WA Web diagnostic: validate skmsg is not first in multi-enc messages. + if !session_payloads.is_empty() + && !group_payloads.is_empty() + && all_enc_nodes.first().is_some_and(|n| { + n.get_attr("type") + .map(|v| v.as_str()) + .is_some_and(|s| s == EncType::SenderKey.as_wire_str()) + }) + { + log::error!( + "[msg:{}] Protocol violation: skmsg is first in multi-enc message from {}. \ + Expected pkmsg/msg first (containing SKDM).", + info.id, + info.source.sender + ); + } + + // Unknown-only stanzas would loop in the offline queue until + // <stream:error>. Custom handlers ack on their own; status is covered + // by should_ack. Ack from `nr` so `recipient` survives. Skip when any + // bucket has usable payloads (including msmsg) so the regular dispatch + // path runs and the valid enc still decrypts. + if session_payloads.is_empty() + && group_payloads.is_empty() + && bot_payloads.is_empty() + && had_unknown_enc + && !had_custom_handler + { + log::info!( + "[msg:{}] All enc payloads unrecognized; transport-acking to drop from offline queue", + info.id + ); + if !info.source.chat.is_status_broadcast() { + self.spawn_node_transport_ack(nr).await; + } + return None; + } + + Some(ClassifiedMessage { + info, + sender_encryption_jid, + session_payloads, + group_payloads, + bot_payloads, + max_sender_retry_count, + decrypt_fail_mode: if has_hide_fail { + crate::types::events::DecryptFailMode::Hide + } else { + crate::types::events::DecryptFailMode::Show + }, + }) + } + + /// Phase 2: acquire permit, decrypt payloads, flush. No node borrows. + pub(crate) async fn process_classified_message(self: Arc<Self>, msg: ClassifiedMessage) { + let ClassifiedMessage { + info, + sender_encryption_jid, + session_payloads, + group_payloads, + bot_payloads, + max_sender_retry_count, + decrypt_fail_mode, + } = msg; + + if max_sender_retry_count > 0 { + let cache_key = self + .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) + .await; + let existing = self.message_retry_counts.get(&cache_key).await.unwrap_or(0); + if max_sender_retry_count > existing { + self.message_retry_counts + .insert(cache_key, max_sender_retry_count) + .await; + } + log::debug!( + "[msg:{}] Sender retry count {} pre-seeded into cache", + info.id, + max_sender_retry_count + ); + } + + // Acquire global processing permit (1 during offline sync, N after). + // Read generation + clone Arc under the same mutex so the pair is consistent. + // + // When the semaphore transitions from 1→N (offline→online), tasks waiting on + // the old 1-permit semaphore must re-acquire from the new N-permit semaphore. + // Without this re-acquire loop, those tasks would be silently dropped, which + // can lose pkmsg messages carrying SKDM (sender key distribution). If the + // SKDM is lost, ALL subsequent skmsg messages from that sender will fail + // with "No sender key state". + let _global_permit = loop { + let (generation, semaphore) = self.read_message_semaphore(); + let permit = semaphore.acquire_arc().await; + if generation + == self + .message_semaphore_generation + .load(std::sync::atomic::Ordering::SeqCst) + { + break permit; + } + // Generation changed while waiting (e.g. offline→online transition). + // Drop the stale permit and retry with the new semaphore, which has + // more permits and will grant access quickly. + log::debug!( + "Semaphore generation changed during acquire, re-acquiring from new semaphore" + ); + drop(permit); + }; + + log::debug!( + "Starting PASS 1: Processing {} session establishment messages (pkmsg/msg)", + session_payloads.len() + ); + + // Skip session processing for group/broadcast JIDs — they use sender keys, not 1:1 sessions. + let is_group_sender = sender_encryption_jid.is_group() + || sender_encryption_jid.is_broadcast_list() + || sender_encryption_jid.is_status_broadcast(); + + let session_outcome = if !is_group_sender && !session_payloads.is_empty() { + self.clone() + .process_session_enc_batch( + &session_payloads, + &info, + &sender_encryption_jid, + decrypt_fail_mode, + ) + .await + } else { + if is_group_sender && !session_payloads.is_empty() { + log::debug!( + "Skipping {} session messages from group sender {}", + session_payloads.len(), + sender_encryption_jid + ); + } + SessionBatchOutcome::default() + }; + let session_decrypted_successfully = session_outcome.decrypted; + let session_had_duplicates = session_outcome.duplicate; + let session_dispatched_undecryptable = session_outcome.undecryptable; + + log::debug!( + "Starting PASS 2: Processing {} group content messages (skmsg)", + group_payloads.len() + ); + + // Only process group content if: + // 1. There were no session messages (session already exists), OR + // 2. Session messages were successfully decrypted, OR + // 3. Session messages were duplicates (already processed, so session exists) + // Skip only if session messages FAILED to decrypt (not duplicates, not absent). + // Matches WA Web's `canDecryptNext` pattern: if pkmsg fails with a retriable error, + // the SKDM it carried is lost, so skmsg will always fail with NoSenderKey — skip it + // to avoid unnecessary retry receipts. The retry for the pkmsg will cause the sender + // to resend the entire message including SKDM. + if !group_payloads.is_empty() { + let should_process_skmsg = + should_process_skmsg_after_session(session_payloads.is_empty(), session_outcome); + + if should_process_skmsg { + match self + .clone() + .process_group_enc_batch( + &group_payloads, + &info, + &sender_encryption_jid, + decrypt_fail_mode, + ) + .await + { + Ok(()) => { + // Processed successfully or handled errors (e.g. sent retry receipt) + } + Err(e) => { + log::warn!( + "[msg:{}] Batch group decrypt from {} in {} failed: {e:?}", + info.id, + info.source.sender, + info.source.chat + ); + } + } + } else { + // Only show warning if session messages actually FAILED (not duplicates) + if !session_had_duplicates { + if info.is_expired_status() { + log::debug!( + "[msg:{}] Silently dropping expired status from {}", + info.id, + info.source.sender + ); + } else { + log::log!( + decrypt_fail_log_level(decrypt_fail_mode), + "Skipping skmsg decryption for message {} from {} because pkmsg failed to decrypt.", + info.id, + info.source.sender + ); + if !session_dispatched_undecryptable { + self.dispatch_undecryptable_event( + Arc::clone(&info), + false, + crate::types::events::UnavailableType::Unknown, + decrypt_fail_mode, + ) + .await; + } + } + + // Do NOT send a delivery receipt for undecryptable messages. + // Per whatsmeow's implementation, delivery receipts are only sent for + // successfully decrypted/handled messages. Sending a receipt here would + // tell the server we processed it, incrementing the offline counter. + // The transport <ack> is sufficient for acknowledgment. + } + // If session_had_duplicates is true, we silently skip (no warning, no event) + // because the message was already processed in a previous session + } + } else if !session_decrypted_successfully + && !session_had_duplicates + && !session_payloads.is_empty() + { + // Edge case: message with only msg/pkmsg that failed to decrypt, no skmsg + log::log!( + decrypt_fail_log_level(decrypt_fail_mode), + "Message {} from {} failed to decrypt and has no group content. Dispatching UndecryptableMessage event.", + info.id, + info.source.sender + ); + // Dispatch UndecryptableMessage event for messages that failed to decrypt + // (This should not cause double-dispatching since process_session_enc_batch + // already returned dispatched_undecryptable=false for this case) + if !session_dispatched_undecryptable { + self.dispatch_undecryptable_event( + Arc::clone(&info), + false, + crate::types::events::UnavailableType::Unknown, + decrypt_fail_mode, + ) + .await; + } + // Do NOT send delivery receipt - transport ack is sufficient + } else if session_had_duplicates + && !session_decrypted_successfully + && !session_dispatched_undecryptable + && !info.source.chat.is_status_broadcast() + { + // Duplicate (already-processed) with no group content: ack it so the + // server drops it from the offline queue (whatsmeow/WA Web treat + // old-counter like success). status is acked by the should_ack gate + // (a status SKDM pkmsg can reach here), so skip it to avoid a + // redundant receipt. + self.ack_received_message(&info); + } else if should_ack_skdm_only_session_fallback(session_outcome, bot_payloads.is_empty()) { + // SKDM-only session decrypts skip dispatch, so this stanza would + // otherwise stay queued. WA Web and whatsmeow ack every decrypted + // message; the ack shape still comes from the message source. + // Status is intentionally not filtered here, so its success receipt + // still follows the normal WA Web path. + self.ack_received_message(&info); + } + + // Bot-secret (msmsg) payloads run inline here so they're serialised + // with the session/group decrypt batches under the same global + // permit + per-chat enqueue lock acquired upstream. + for payload in bot_payloads { + self.handle_msmsg_payload(&info, payload).await; + } + + // Flush cached Signal state to DB (matches WA Web's flushBufferToDiskIfNotMemOnlyMode) + self.flush_signal_cache_logged("message", Some(&info.id)) + .await; + } + + pub(crate) async fn process_session_enc_batch( + self: Arc<Self>, + payloads: &[EncPayload], + info: &Arc<MessageInfo>, + sender_encryption_jid: &Jid, + decrypt_fail_mode: crate::types::events::DecryptFailMode, + ) -> SessionBatchOutcome { + use wacore::libsignal::protocol::CiphertextMessage; + if payloads.is_empty() { + return SessionBatchOutcome::default(); + } + + // Acquire a per-sender session lock to prevent race conditions when + // multiple messages from the same sender are processed concurrently. + // Use the full Signal protocol address string as the lock key so it matches + // the SignalProtocolStoreAdapter's per-session locks (prevents ratchet counter races). + let signal_address = sender_encryption_jid.to_protocol_address(); + + // `session_guard` is held across the entire batch but dropped around + // calls into `try_pn_to_lid_migration_decrypt` because that function's + // migration loop re-enters this same mutex (non-reentrant). + let session_mutex = self.session_lock_for(signal_address.as_str()).await; + let mut session_guard: Option<async_lock::MutexGuardArc<()>> = + Some(session_mutex.lock_arc().await); + + let mut adapter = self.signal_adapter().await; + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + let mut outcome = SessionBatchOutcome::default(); + // Local identity-change detection fires once per batch: the first pkmsg + // saves the new key (ReplacedExisting); the rest are NewOrUnchanged. + let mut local_identity_reacted = false; + + for payload in payloads { + let ciphertext = &payload.ciphertext[..]; + let enc_type = payload.enc_type; + let enc_type_str = enc_type.as_wire_str(); + let padding_version = payload.padding_version; + + // WA Web `MsgSendReceipt.js` nacks PARSE_ERROR; without it the + // server retransmits the malformed stanza forever. Mirrors the + // `handle_decrypt_failure` shape (dispatch event + spawn wire I/O + // so the session lock isn't held across the send). + let parsed_message = if enc_type == EncType::PreKeyMessage { + match PreKeySignalMessage::try_from(ciphertext) { + Ok(m) => CiphertextMessage::PreKeySignalMessage(m), + Err(e) => { + log::error!( + "[msg:{}] Failed to parse PreKeySignalMessage from {}: {e:?}. Sending nack.", + info.id, + info.source.sender + ); + // |= so a later dedup'd return (false) can't clobber + // a true set by a prior iteration in this batch. + outcome.had_failure = true; + outcome.undecryptable |= self + .dispatch_undecryptable_event( + Arc::clone(info), + false, + crate::types::events::UnavailableType::Unknown, + decrypt_fail_mode, + ) + .await; + self.spawn_nack(info, NackReason::ParsingError, None); + continue; + } + } + } else { + match SignalMessage::try_from(ciphertext) { + Ok(m) => CiphertextMessage::SignalMessage(m), + Err(e) => { + log::error!( + "[msg:{}] Failed to parse SignalMessage from {}: {e:?}. Sending nack.", + info.id, + info.source.sender + ); + outcome.had_failure = true; + outcome.undecryptable |= self + .dispatch_undecryptable_event( + Arc::clone(info), + false, + crate::types::events::UnavailableType::Unknown, + decrypt_fail_mode, + ) + .await; + self.spawn_nack(info, NackReason::ParsingError, None); + continue; + } + } + }; + + if enc_type == EncType::PreKeyMessage { + // FLAGGED FOR DEBUGGING: "Bad Mac" Reproducibility + #[cfg(feature = "debug-snapshots")] + { + use base64::prelude::*; + let payload = serde_json::json!({ + "id": info.id, + "sender_jid": sender_encryption_jid.to_string(), + "timestamp": info.timestamp, + "enc_type": enc_type_str, + "payload_base64": BASE64_STANDARD.encode(ciphertext), + }); + + let content_bytes = serde_json::to_vec_pretty(&payload).unwrap_or_default(); + + if let Err(e) = self + .persistence_manager + .create_snapshot(&format!("pre_pkmsg_{}", info.id), Some(&content_bytes)) + .await + { + log::warn!("Failed to create snapshot for pkmsg: {}", e); + } + } + #[cfg(not(feature = "debug-snapshots"))] + { + // No-op if disabled + } + } + + // Shadow with wire string for all downstream usage (logging, handlers) + let enc_type = enc_type_str; + + let decrypt_res = message_decrypt( + &parsed_message, + &signal_address, + &mut adapter.session_store, + &mut adapter.identity_store, + &mut adapter.pre_key_store, + &adapter.signed_pre_key_store, + &mut rng, + UsePQRatchet::No, + ) + .await; + + match decrypt_res { + Ok(decrypted) => { + // Buffer the prekey this pkmsg consumed: message_decrypt promoted + // the session into the (volatile) cache but no longer deletes the + // prekey itself. The post-loop flush deletes it only once that + // session is durable, keeping a crash from orphaning the prekey. + if let Some(prekey_id) = decrypted.consumed_prekey_id { + adapter + .pre_key_store + .buffer_consumed_prekey(prekey_id, &signal_address) + .await; + } + if decrypted.identity_change == IdentityChange::ReplacedExisting + && !local_identity_reacted + { + local_identity_reacted = true; + self.react_to_local_identity_change(sender_encryption_jid); + } + let padded_plaintext = decrypted.plaintext; + match self + .clone() + .handle_decrypted_plaintext( + enc_type, + &padded_plaintext, + padding_version, + info, + ) + .await + { + Ok(plaintext_outcome) => { + outcome.decrypted = true; + outcome.dispatched |= plaintext_outcome.dispatched; + outcome.skdm_only |= plaintext_outcome.skdm_only; + } + Err(e) => { + log::warn!( + "[msg:{}] Failed processing plaintext from {}: {e:?}", + info.id, + info.source.sender + ); + outcome.decrypted = true; + outcome.plaintext_failed = true; + outcome.had_failure = true; + outcome.undecryptable |= + self.handle_plaintext_failure(info, decrypt_fail_mode).await; + } + } + } + Err(e) => { + // Handle DuplicatedMessage: This is expected when messages are redelivered during reconnection + if let SignalProtocolError::DuplicatedMessage(chain, counter) = e { + log::debug!( + "Skipping already-processed message from {} (chain {}, counter {}). This is normal during reconnection.", + info.source.sender, + chain, + counter + ); + // Mark that we saw a duplicate so we can skip skmsg without showing error + outcome.duplicate = true; + continue; + } + // Handle UntrustedIdentity: This happens when a user re-installs WhatsApp or changes devices. + // The Signal Protocol's security policy rejects messages from new identity keys by default. + // We handle this by clearing the old identity (to trust the new one), then retrying decryption. + // IMPORTANT: We do NOT delete the session! When the PreKeySignalMessage is processed, + // libsignal's `promote_state` will archive the old session as a "previous state". + // This allows us to decrypt any in-flight messages that were encrypted with the old session. + if let SignalProtocolError::UntrustedIdentity(ref address) = e { + log::warn!( + "[msg:{}] Received message from untrusted identity: {}. This typically means the sender re-installed WhatsApp or changed their device. Clearing old identity to trust new key (keeping session for in-flight messages).", + info.id, + address + ); + + // Delete the old, untrusted identity through the signal cache. + // NOTE: We intentionally do NOT delete the session here. The session will be + // archived (not deleted) when the new PreKeySignalMessage is processed, + // allowing decryption of any in-flight messages encrypted with the old session. + self.signal_cache.delete_identity(address).await; + // Flush immediately so the backend is updated BEFORE the retry decrypt below. + // Device::is_trusted_identity reads from backend, not cache. + if let Err(e) = self.flush_signal_cache().await { + log::warn!("Failed to flush identity deletion for {}: {e:?}", address); + outcome.had_failure = true; + continue; + } + log::info!( + "Cleared old identity for {} from cache and backend", + address + ); + + // Re-attempt decryption with the new identity + log::info!( + "[msg:{}] Retrying message decryption for {} after clearing untrusted identity", + info.id, + address + ); + + let retry_decrypt_res = message_decrypt( + &parsed_message, + &signal_address, + &mut adapter.session_store, + &mut adapter.identity_store, + &mut adapter.pre_key_store, + &adapter.signed_pre_key_store, + &mut rng, + UsePQRatchet::No, + ) + .await; + + match retry_decrypt_res { + Ok(decrypted) => { + log::debug!( + "[msg:{}] Successfully decrypted message from {} after handling untrusted identity", + info.id, + address + ); + if let Some(prekey_id) = decrypted.consumed_prekey_id { + adapter + .pre_key_store + .buffer_consumed_prekey(prekey_id, &signal_address) + .await; + } + // Normally NewOrUnchanged here (the untrusted + // identity was deleted+flushed before the retry), + // but mirror the main-decode gate so a concurrent + // re-save can't drop the signal. + if decrypted.identity_change == IdentityChange::ReplacedExisting + && !local_identity_reacted + { + local_identity_reacted = true; + self.react_to_local_identity_change(sender_encryption_jid); + } + let padded_plaintext = decrypted.plaintext; + match self + .clone() + .handle_decrypted_plaintext( + enc_type, + &padded_plaintext, + padding_version, + info, + ) + .await + { + Ok(plaintext_outcome) => { + outcome.decrypted = true; + outcome.dispatched |= plaintext_outcome.dispatched; + outcome.skdm_only |= plaintext_outcome.skdm_only; + } + Err(e) => { + log::warn!( + "Failed processing plaintext after identity retry: {e:?}" + ); + outcome.decrypted = true; + outcome.plaintext_failed = true; + outcome.had_failure = true; + outcome.undecryptable |= self + .handle_plaintext_failure(info, decrypt_fail_mode) + .await; + } + } + } + Err(retry_err) => { + // Handle DuplicatedMessage in retry path: This commonly happens during reconnection + // when the same message is redelivered by the server after we already processed it. + // The first attempt triggered UntrustedIdentity, we cleared the session, but meanwhile + // another message from the same sender re-established the session and consumed the counter. + // This is benign - the message was already successfully processed. + if let SignalProtocolError::DuplicatedMessage(chain, counter) = + retry_err + { + log::debug!( + "Message from {} was already processed (chain {}, counter {}) - detected during untrusted identity retry. This is normal during reconnection.", + address, + chain, + counter + ); + outcome.duplicate = true; + } else if matches!(retry_err, SignalProtocolError::InvalidPreKeyId) + { + // Session may exist under PN address after identity change + let migration_outcome = self + .try_pn_to_lid_migration_decrypt( + sender_encryption_jid, + &signal_address, + &parsed_message, + &mut adapter, + &mut rng, + enc_type, + padding_version, + info, + &session_mutex, + &mut session_guard, + ) + .await; + if migration_outcome.decrypted + || migration_outcome.duplicate + || migration_outcome.plaintext_failed + { + outcome.decrypted |= migration_outcome.decrypted; + outcome.duplicate |= migration_outcome.duplicate; + outcome.dispatched |= migration_outcome.dispatched; + outcome.skdm_only |= migration_outcome.skdm_only; + outcome.plaintext_failed |= + migration_outcome.plaintext_failed; + outcome.had_failure |= migration_outcome.plaintext_failed; + if migration_outcome.plaintext_failed { + outcome.undecryptable |= self + .handle_plaintext_failure(info, decrypt_fail_mode) + .await; + } + } else { + log::debug!( + "[msg:{}] InvalidPreKeyId after identity change for {}. \ + Sending retry receipt with fresh keys.", + info.id, + address + ); + outcome.had_failure = true; + outcome.undecryptable |= self + .handle_decrypt_failure( + info, + RetryReason::InvalidKeyId, + decrypt_fail_mode, + ) + .await; + } + } else { + log::error!( + "[msg:{}] Decryption failed even after clearing untrusted identity for {}: {:?}", + info.id, + address, + retry_err + ); + // Send retry receipt so the sender resends with a PreKeySignalMessage + // to establish a new session with the new identity + outcome.had_failure = true; + outcome.undecryptable |= self + .handle_decrypt_failure( + info, + RetryReason::InvalidKey, + decrypt_fail_mode, + ) + .await; + } + } + } + + // Re-issue tctoken so the contact still has a valid token for us + let sender_jid = info.source.sender.clone(); + if !sender_jid.is_bot() && !sender_jid.is_status_broadcast() { + let client = self.clone(); + self.runtime + .spawn(Box::pin(async move { + client + .reissue_tc_token_after_identity_change(&sender_jid) + .await; + })) + .detach(); + } + + continue; + } + // Try PN→LID session migration before sending retry receipt + if let SignalProtocolError::SessionNotFound(_) = e { + let migration_outcome = self + .try_pn_to_lid_migration_decrypt( + sender_encryption_jid, + &signal_address, + &parsed_message, + &mut adapter, + &mut rng, + enc_type, + padding_version, + info, + &session_mutex, + &mut session_guard, + ) + .await; + if migration_outcome.decrypted + || migration_outcome.duplicate + || migration_outcome.plaintext_failed + { + outcome.decrypted |= migration_outcome.decrypted; + outcome.duplicate |= migration_outcome.duplicate; + outcome.dispatched |= migration_outcome.dispatched; + outcome.skdm_only |= migration_outcome.skdm_only; + outcome.plaintext_failed |= migration_outcome.plaintext_failed; + outcome.had_failure |= migration_outcome.plaintext_failed; + if migration_outcome.plaintext_failed { + outcome.undecryptable |= + self.handle_plaintext_failure(info, decrypt_fail_mode).await; + } + continue; + } + + debug!( + "[msg:{}] No session found for {} message from {}. Sending retry receipt to request session establishment.", + info.id, enc_type, info.source.sender + ); + outcome.had_failure = true; + outcome.undecryptable |= self + .handle_decrypt_failure(info, RetryReason::NoSession, decrypt_fail_mode) + .await; + continue; + } else if matches!( + e, + SignalProtocolError::BadMac(_) | SignalProtocolError::InvalidMessage(_, _) + ) { + // whatsmeow migrates PN sessions before decrypt; a fresh + // LID record can otherwise shadow the sender's PN ratchet. + let migration_outcome = self + .try_pn_to_lid_migration_decrypt( + sender_encryption_jid, + &signal_address, + &parsed_message, + &mut adapter, + &mut rng, + enc_type, + padding_version, + info, + &session_mutex, + &mut session_guard, + ) + .await; + if migration_outcome.decrypted + || migration_outcome.duplicate + || migration_outcome.plaintext_failed + { + outcome.decrypted |= migration_outcome.decrypted; + outcome.duplicate |= migration_outcome.duplicate; + outcome.dispatched |= migration_outcome.dispatched; + outcome.skdm_only |= migration_outcome.skdm_only; + outcome.plaintext_failed |= migration_outcome.plaintext_failed; + outcome.had_failure |= migration_outcome.plaintext_failed; + if migration_outcome.plaintext_failed { + outcome.undecryptable |= + self.handle_plaintext_failure(info, decrypt_fail_mode).await; + } + continue; + } + + // WAWebMsgProcessingDecryptionHandler classifies both as + // SignalRetryable -> sendRetryReceipt only, with no delete. + let (reason, label) = if matches!(e, SignalProtocolError::BadMac(_)) { + (RetryReason::BadMac, "BadMac") + } else { + (RetryReason::InvalidMessage, "InvalidMessage") + }; + log::log!( + decrypt_fail_log_level(decrypt_fail_mode), + "[msg:{}] Decryption failed for {} message from {} due to {label}. \ + Sending retry receipt.", + info.id, + enc_type, + info.source.sender + ); + + outcome.had_failure = true; + outcome.undecryptable |= self + .handle_decrypt_failure(info, reason, decrypt_fail_mode) + .await; + continue; + } else if matches!(e, SignalProtocolError::InvalidPreKeyId) { + // InvalidPreKeyId on a PreKeyMessage can also mean the + // session exists under a PN address (legacy migration). + // Migrating lets Signal use the existing ratchet state + // instead of looking up the consumed one-time prekey. + let migration_outcome = self + .try_pn_to_lid_migration_decrypt( + sender_encryption_jid, + &signal_address, + &parsed_message, + &mut adapter, + &mut rng, + enc_type, + padding_version, + info, + &session_mutex, + &mut session_guard, + ) + .await; + if migration_outcome.decrypted + || migration_outcome.duplicate + || migration_outcome.plaintext_failed + { + outcome.decrypted |= migration_outcome.decrypted; + outcome.duplicate |= migration_outcome.duplicate; + outcome.dispatched |= migration_outcome.dispatched; + outcome.skdm_only |= migration_outcome.skdm_only; + outcome.plaintext_failed |= migration_outcome.plaintext_failed; + outcome.had_failure |= migration_outcome.plaintext_failed; + if migration_outcome.plaintext_failed { + outcome.undecryptable |= + self.handle_plaintext_failure(info, decrypt_fail_mode).await; + } + continue; + } + + log::debug!( + "[msg:{}] Decryption failed for {} message from {} due to InvalidPreKeyId. \ + Sender is using a prekey we don't have (likely session established while offline). \ + Sending retry receipt with fresh prekeys.", + info.id, + enc_type, + info.source.sender + ); + + // Send retry receipt with fresh prekeys + outcome.had_failure = true; + outcome.undecryptable |= self + .handle_decrypt_failure( + info, + RetryReason::InvalidKeyId, + decrypt_fail_mode, + ) + .await; + continue; + } else { + // Catch-all → WA Web's UnhandledError nack (500). + log::error!( + "[msg:{}] Batch session decrypt failed (type: {}) from {}: {:?}. Sending nack.", + info.id, + enc_type, + info.source.sender, + e + ); + outcome.had_failure = true; + outcome.undecryptable |= self + .dispatch_undecryptable_event( + Arc::clone(info), + false, + crate::types::events::UnavailableType::Unknown, + decrypt_fail_mode, + ) + .await; + self.spawn_nack(info, NackReason::UnhandledError, None); + continue; + } + } + } + } + outcome + } + + async fn process_group_enc_batch( + self: Arc<Self>, + payloads: &[EncPayload], + info: &Arc<MessageInfo>, + _sender_encryption_jid: &Jid, + decrypt_fail_mode: crate::types::events::DecryptFailMode, + ) -> Result<(), DecryptionError> { + if payloads.is_empty() { + return Ok(()); + } + let mut adapter = self.signal_adapter().await; + + // Always use bare sender for sender key operations. Real WA delivers + // skmsg with bare participant but pkmsg (SKDM) with device-qualified + // participant — normalizing to bare ensures consistent lookup. + // Hoisted out of the payload loop: all three are loop-invariant. + let sender_for_sk = info.source.sender.to_non_ad(); + let sender_address = sender_for_sk.to_protocol_address(); + let sender_key_name = make_sender_key_name(&info.source.chat, &sender_address); + + for payload in payloads { + let ciphertext = &payload.ciphertext[..]; + let padding_version = payload.padding_version; + + log::debug!( + "Looking up sender key for group {} with sender address {} (from sender JID: {})", + info.source.chat, + sender_address, + info.source.sender + ); + + let decrypt_result = + group_decrypt(ciphertext, &mut adapter.sender_key_store, &sender_key_name).await; + + match decrypt_result { + Ok(padded_plaintext) => { + // Sync device list if sender is unknown, but still process + // the message. Signal decryption success already proves the + // sender holds the session key — discarding would only add + // latency via an unnecessary retry round-trip. + if !self.is_from_known_device(&info.source.sender).await { + debug!( + "[msg:{}] Unknown device {}, triggering device sync", + info.id, info.source.sender + ); + self.handle_unknown_device_sync(info).await; + } + + if let Err(e) = self + .clone() + .handle_decrypted_plaintext( + "skmsg", + &padded_plaintext, + padding_version, + info, + ) + .await + { + log::warn!("Failed processing group plaintext (batch): {e:?}"); + } + } + Err(SignalProtocolError::DuplicatedMessage(iteration, counter)) => { + log::debug!( + "Skipping already-processed sender key message from {} in group {} (iteration {}, counter {}). This is normal during reconnection.", + info.source.sender, + info.source.chat, + iteration, + counter + ); + // Redelivered duplicate: ack it so the server drops it from the + // offline queue. status is already acked by the should_ack gate, + // so skip it to avoid a redundant receipt. + if !info.source.chat.is_status_broadcast() { + self.ack_received_message(info); + } + } + Err(SignalProtocolError::NoSenderKeyState(msg)) => { + if info.is_expired_status() { + log::debug!( + "[msg:{}] Skipping retry for expired status from {}", + info.id, + info.source.sender + ); + continue; + } + + let is_unknown_device = !self.is_from_known_device(&info.source.sender).await; + let retry_reason = if is_unknown_device { + RetryReason::UnknownCompanionNoPrekey + } else { + RetryReason::NoSession + }; + + debug!( + "No sender key state for group message [msg:{}] from {}: {}. Sending retry receipt.", + info.id, info.source.sender, msg + ); + + if is_unknown_device { + self.handle_unknown_device_sync(info).await; + } + + self.handle_decrypt_failure(info, retry_reason, decrypt_fail_mode) + .await; + } + Err(e) => { + if info.is_expired_status() { + log::debug!( + "[msg:{}] Ignoring decrypt error for expired status from {}: {:?}", + info.id, + info.source.sender, + e + ); + continue; + } + + log::log!( + decrypt_fail_log_level(decrypt_fail_mode), + "Group batch decrypt failed [msg:{}] for group {} sender {}: {:?}", + info.id, + sender_key_name.group_id(), + sender_key_name.sender_id(), + e + ); + // Always surface the failure to consumers; nack only non-status + // (status is acked by the should_ack gate) so the server drops + // it from the offline queue. + self.dispatch_undecryptable_event( + Arc::clone(info), + false, + crate::types::events::UnavailableType::Unknown, + decrypt_fail_mode, + ) + .await; + if !info.source.chat.is_status_broadcast() { + self.spawn_nack(info, NackReason::UnhandledError, None); + } + } + } + } + Ok(()) + } + + /// WA Web: online → `syncDeviceListJob`, offline → `OfflinePendingDeviceCache`. + async fn handle_unknown_device_sync(self: &Arc<Self>, info: &MessageInfo) { + let user_jid = info.source.sender.to_non_ad(); + + // Dedup: skip if we already have a sync pending/in-flight for this user + if !self.pending_device_sync.add(user_jid.clone()).await { + return; + } + + if info.is_offline { + log::debug!("Queueing {} for pending device sync (offline)", user_jid); + } else { + log::debug!("Triggering immediate device sync for {}", user_jid); + let client = Arc::clone(self); + self.runtime + .spawn(Box::pin(async move { + client.invalidate_device_cache(&user_jid.user).await; + if let Err(e) = client.get_user_devices(&[user_jid]).await { + log::warn!("Immediate device sync failed: {e:?}"); + } + })) + .detach(); + } + } + + pub(crate) async fn handle_decrypted_plaintext( + self: Arc<Self>, + enc_type: &str, + padded_plaintext: &[u8], + padding_version: u8, + info: &Arc<MessageInfo>, + ) -> Result<PlaintextHandleOutcome, anyhow::Error> { + let original_msg = wacore::messages::decode_plaintext(padded_plaintext, padding_version)?; + log::debug!( + "[msg:{}] Successfully decrypted message from {}: type={} [batch path]", + info.id, + info.source.sender, + enc_type + ); + + // Validate DSM presence against sender identity + // (WAWebHandleMsgError.DeviceSentMessageError) + if original_msg.device_sent_message.is_some() && !info.source.is_from_me { + warn!( + "[msg:{}] DeviceSentMessage present but sender {} is not self", + info.id, info.source.sender, + ); + } + + // WA Web validateBclHash: a self-synced broadcast/status carries a + // phashV2 of the broadcast recipients in deviceSentMessage.phash. + // Recompute over our <participants> view and warn on divergence. We log + // only (no drop) until the participant hash form is confirmed live. + if let Some(dsm) = &original_msg.device_sent_message + && let Some(expected) = dsm.phash.as_deref() + && !info.bcl_participants.is_empty() + && !wacore::messages::MessageUtils::validate_bcl_hash(&info.bcl_participants, expected) + { + warn!( + "[msg:{}] bcl hash mismatch on device-sent broadcast (expected={expected}); \ + keeping message (validate-only)", + info.id, + ); + } + + // Unwrap DeviceSentMessage wrapper (self-sent messages synced from + // the primary device). The actual content (reactions, text, etc.) + // is nested inside device_sent_message.message and must be + // extracted before protocol checks or dispatch. + let mut msg = wacore::messages::unwrap_device_sent(original_msg); + + // Post-decryption logic (SKDM, sync keys, etc.) + if let Some(skdm) = &msg.sender_key_distribution_message + && let Some(axolotl_bytes) = &skdm.axolotl_sender_key_distribution_message + { + self.handle_sender_key_distribution_message( + &info.source.chat, + &info.source.sender, + axolotl_bytes, + ) + .await; + } + + // app_state_sync_key_share is a self-only protocol message (app-state + // sync keys shared between our own devices). A peer could otherwise + // inject keys and forge app-state mutations, so honour it only from + // self. WA Web `WAWebKeyManagementHandleKeyShareApi` gates on + // `isMeAccountNonLid(from)`; whatsmeow on `info.IsFromMe`. + if let Some(protocol_msg) = &msg.protocol_message + && let Some(keys) = &protocol_msg.app_state_sync_key_share + { + if info.source.is_from_me { + self.handle_app_state_sync_key_share(keys).await; + } else { + warn!( + "[msg:{}] Dropping app_state_sync_key_share from non-self sender {}", + info.id, info.source.sender + ); + } + } + + // PDO responses come from our own account (is_from_me) via device 0 (primary phone) + if info.source.is_from_me + && let Some(protocol_msg) = &msg.protocol_message + && let Some(pdo_response) = &protocol_msg.peer_data_operation_request_response_message + { + self.handle_pdo_response(pdo_response, info).await; + } + + // Note: msg might be modified by take() below + let history_sync_taken = msg + .protocol_message + .as_mut() + .and_then(|pm| pm.history_sync_notification.take()); + + // history_sync_notification is self-only (our phone drives history sync). + // A spoofed one from a peer would force a download of attacker-controlled + // history, so honour it only from self. WA Web + // `WAWebHandleHistorySyncNotification` gates on `isMePrimaryNonLid`. + if let Some(history_sync) = history_sync_taken { + if info.source.is_from_me { + self.handle_history_sync(info.id.clone(), history_sync) + .await; + } else { + warn!( + "[msg:{}] Dropping history_sync_notification from non-self sender {}", + info.id, info.source.sender + ); + } + } + + // Skip dispatch for messages that only carry sender key distribution + // (protocol-level key exchange) with no user-visible content. + // These arrive as a separate pkmsg enc node alongside the actual + // group message (skmsg) and would otherwise surface as "unknown". + if wacore::messages::is_sender_key_distribution_only(&mut msg) { + log::debug!( + "[msg:{}] Skipping event dispatch for sender key distribution message", + info.id + ); + Ok(PlaintextHandleOutcome { + skdm_only: true, + ..Default::default() + }) + } else { + self.dispatch_parsed_message(msg, info).await; + Ok(PlaintextHandleOutcome { + dispatched: true, + ..Default::default() + }) + } + } + + /// Attempt PN→LID session migration and retry decryption. + /// Returns whether decryption succeeded after migration and whether it + /// reached user dispatch. + /// + /// Manages the per-address session lock around the migration loop: + /// drops the caller's guard (migration re-enters that mutex and + /// async_lock is non-reentrant), then reacquires it for the retry + /// decrypt and replaces the caller's `session_guard` on the way out + /// so the next payload in the batch stays serialized. + #[allow(clippy::too_many_arguments)] + async fn try_pn_to_lid_migration_decrypt( + self: &Arc<Self>, + sender_jid: &Jid, + signal_address: &wacore::libsignal::protocol::ProtocolAddress, + parsed_message: &wacore::libsignal::protocol::CiphertextMessage, + adapter: &mut crate::store::signal_adapter::SignalProtocolStoreAdapter, + rng: &mut rand::rngs::StdRng, + enc_type: &str, + padding_version: u8, + info: &Arc<MessageInfo>, + session_mutex: &Arc<async_lock::Mutex<()>>, + session_guard: &mut Option<async_lock::MutexGuardArc<()>>, + ) -> MigrationDecryptOutcome { + use wacore::libsignal::protocol::{UsePQRatchet, message_decrypt}; + + if !sender_jid.is_lid() { + return MigrationDecryptOutcome::default(); + } + + let Some(pn) = self.lid_pn_cache.get_phone_number(&sender_jid.user).await else { + return MigrationDecryptOutcome::default(); + }; + + // Release the address lock so the migration loop can acquire it for + // the matching device without re-entering. + *session_guard = None; + self.migrate_signal_sessions_on_lid_discovery(&pn, &sender_jid.user) + .await; + // Re-acquire for the retry decrypt and hand the guard back to the + // caller for subsequent payloads in the batch. + *session_guard = Some(session_mutex.lock_arc().await); + + match message_decrypt( + parsed_message, + signal_address, + &mut adapter.session_store, + &mut adapter.identity_store, + &mut adapter.pre_key_store, + &adapter.signed_pre_key_store, + rng, + UsePQRatchet::No, + ) + .await + { + // PN→LID migration re-addresses an existing peer; the LID address gets + // the identity for the first time (NewOrUnchanged), so no local + // identity-change reaction is warranted here. + Ok(decrypted) => { + log::info!( + "[msg:{}] Decrypted after PN→LID session migration for {}", + info.id, + info.source.sender + ); + if let Some(prekey_id) = decrypted.consumed_prekey_id { + adapter + .pre_key_store + .buffer_consumed_prekey(prekey_id, signal_address) + .await; + } + let padded_plaintext = decrypted.plaintext; + match self + .clone() + .handle_decrypted_plaintext(enc_type, &padded_plaintext, padding_version, info) + .await + { + Ok(plaintext_outcome) => MigrationDecryptOutcome { + decrypted: true, + dispatched: plaintext_outcome.dispatched, + skdm_only: plaintext_outcome.skdm_only, + ..Default::default() + }, + Err(e) => { + log::warn!( + "[msg:{}] Failed processing plaintext after migration: {e:?}", + info.id + ); + MigrationDecryptOutcome { + decrypted: true, + plaintext_failed: true, + ..Default::default() + } + } + } + } + Err(SignalProtocolError::DuplicatedMessage(chain, counter)) => { + log::debug!( + "[msg:{}] Already processed (chain {chain}, counter {counter}) after migration", + info.id + ); + MigrationDecryptOutcome { + duplicate: true, + ..Default::default() + } + } + Err(retry_err) => { + log::warn!( + "[msg:{}] Decryption still failed after PN→LID migration: {retry_err:?}", + info.id + ); + MigrationDecryptOutcome::default() + } + } + } + + pub(crate) async fn cache_lid_pn_from_message( + self: &Arc<Self>, + sender: &Jid, + alt: Option<&Jid>, + is_offline: bool, + ) { + let (lid_user, pn_user, source) = if sender.server.is_lid_family() { + if let Some(alt_jid) = alt + && alt_jid.server.is_pn_family() + { + ( + &sender.user, + &alt_jid.user, + crate::lid_pn_cache::LearningSource::PeerLidMessage, + ) + } else { + return; + } + } else if sender.server.is_pn_family() { + if let Some(alt_jid) = alt + && alt_jid.server.is_lid_family() + { + ( + &alt_jid.user, + &sender.user, + crate::lid_pn_cache::LearningSource::PeerPnMessage, + ) + } else { + return; + } + } else { + return; + }; + + self.learn_lid_pn_mapping_fast(lid_user, pn_user, source, is_offline) + .await; + } + + pub(crate) async fn parse_message_info( + &self, + node: &wacore_binary::NodeRef<'_>, + ) -> Result<MessageInfo, anyhow::Error> { + let (own_pn, own_lid) = { + let arc = self.persistence_manager.get_device_arc().await; + let guard = arc.read().await; + (guard.pn.clone(), guard.lid.clone()) + }; + let default_jid = Jid::default(); + let own_jid = own_pn.as_ref().unwrap_or(&default_jid); + wacore::messages::parse_message_info(node, own_jid, own_lid.as_ref()) + } +} diff --git a/src/message/retry.rs b/src/message/retry.rs new file mode 100644 index 000000000..2111616e4 --- /dev/null +++ b/src/message/retry.rs @@ -0,0 +1,268 @@ +//! Decrypt-failure handling, retry receipts and undecryptable events. + +use super::*; + +impl Client { + /// Dispatch an `UndecryptableMessage` event at most once per `(chat, id)` + /// via the single-flight `get_with` semantic on `undecryptable_dispatched`. + /// The atomic arm avoids the get-then-insert race where two concurrent + /// callers would both dispatch. Mirrors WA Web's DB-level placeholder + /// uniqueness in `WAWebMessageProcessPlaceholder`. + /// + /// Returns `true` if this call dispatched the event, `false` if a + /// previous call already did. + pub(crate) async fn dispatch_undecryptable_event( + &self, + info: Arc<MessageInfo>, + is_unavailable: bool, + unavailable_type: crate::types::events::UnavailableType, + decrypt_fail_mode: crate::types::events::DecryptFailMode, + ) -> bool { + let dedup_key = + wacore::types::message::ChatMessageId::new(info.source.chat.clone(), info.id.clone()); + // The init future only runs for the winning caller. Others receive + // the cached `()` and leave the flag as false. + let fresh = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let fresh_clone = fresh.clone(); + self.undecryptable_dispatched + .get_with(dedup_key, async move { + fresh_clone.store(true, std::sync::atomic::Ordering::Release); + }) + .await; + let was_fresh = fresh.load(std::sync::atomic::Ordering::Acquire); + if was_fresh { + self.core.event_bus.dispatch(Event::UndecryptableMessage( + crate::types::events::UndecryptableMessage { + info, + is_unavailable, + unavailable_type, + decrypt_fail_mode, + }, + )); + } else { + log::debug!( + "[msg:{}] UndecryptableMessage already dispatched for this id; skipping duplicate event", + info.id, + ); + } + was_fresh + } + + /// Dispatch an undecryptable event, then send the retry receipt and the + /// transport ack in one ordered, flushed task. + /// + /// The retry asks the sender to re-encrypt; the ack clears the stanza from + /// the server's offline queue (the retry alone does not). Both run in a + /// single `outbound_flush` task so `disconnect()` flushes them together and + /// the retry always goes out before the ack: if only one makes it, it is the + /// retry, so the message is never cleared without a resend request. status is + /// also acked here (flushed) rather than relying on the detached `should_ack` + /// gate, which can be dropped mid-flush on disconnect; the server dedups the + /// resulting duplicate ack. + /// + /// Returns `true` to be assigned to `dispatched_undecryptable` flag. + pub(crate) async fn handle_decrypt_failure( + self: &Arc<Self>, + info: &Arc<MessageInfo>, + reason: RetryReason, + decrypt_fail_mode: crate::types::events::DecryptFailMode, + ) -> bool { + self.dispatch_undecryptable_event( + Arc::clone(info), + false, + crate::types::events::UnavailableType::Unknown, + decrypt_fail_mode, + ) + .await; + let client = Arc::clone(self); + let info = Arc::clone(info); + self.outbound_flush.spawn(&*self.runtime, async move { + // A self-fanout is our own message; retrying it to ourselves is + // futile and the server's offline queue ignores a bare transport + // ack, so it would replay forever. Clear it with the sender receipt + // instead (same stanza the success/duplicate paths now emit). Mirror + // ack_received_message: a bot-authored message in a non-bot chat + // takes the bot-invoke-response bare ack (the retry path below), not + // the sender receipt. Gate on the same eligibility as the ack path. + if info.source.is_self_fanout() + && !info.source.is_bot_authored_non_bot_chat() + && Self::should_send_delivery_receipt(&info) + { + client.send_delivery_receipt(&info).await; + return; + } + // Only ack once the resend request is actually out; otherwise leave + // the stanza queued so the server redelivers and we retry. + let resend_sent = client.run_retry_receipt(&info, reason).await; + if resend_sent { + client.send_transport_ack(&info).await; + } + }); + true + } + + pub(crate) async fn handle_plaintext_failure( + self: &Arc<Self>, + info: &Arc<MessageInfo>, + decrypt_fail_mode: crate::types::events::DecryptFailMode, + ) -> bool { + let dispatched = self + .dispatch_undecryptable_event( + Arc::clone(info), + false, + crate::types::events::UnavailableType::Unknown, + decrypt_fail_mode, + ) + .await; + self.spawn_nack(info, NackReason::InvalidProtobuf, None); + dispatched + } + + /// Increments the retry count for a message and returns the new count. + /// Returns `None` if max retries have been reached. + /// + /// Note: get-then-insert has a theoretical TOCTOU window since + /// `spawn_retry_receipt` detaches. In practice, retries for the same + /// message are rare and a double-send is benign (recipients deduplicate + /// by message ID). + pub(crate) async fn increment_retry_count( + &self, + cache_key: &str, + reason: RetryReason, + ) -> Option<u8> { + let cache_key = cache_key.to_owned(); + let current = self.message_retry_counts.get(&cache_key).await; + let new_count = match current { + Some(count) if count >= MAX_DECRYPT_RETRIES => return None, + Some(count) => count + 1, + None => 1, + }; + self.message_retry_counts + .insert(cache_key.clone(), new_count) + .await; + self.recent_retry_reasons.insert(cache_key, reason).await; + Some(new_count) + } + + /// Generate consistent cache key for retry logic. + pub(crate) async fn make_retry_cache_key( + &self, + chat: &Jid, + msg_id: &str, + sender: &Jid, + ) -> String { + let chat = self.resolve_encryption_jid(chat).await; + let sender = self.resolve_encryption_jid(sender).await; + // +40 covers @server suffixes, :device, separators for two JIDs + let mut key = + String::with_capacity(chat.user.len() + msg_id.len() + sender.user.len() + 40); + chat.push_to(&mut key); + key.push(':'); + key.push_str(msg_id); + key.push(':'); + sender.push_to(&mut key); + key + } + + /// Spawns a task that sends a retry receipt for a failed decryption. + /// + /// This is used when sessions are not found or invalid to request the sender to resend + /// the message with a PreKeySignalMessage to re-establish the session. + /// + /// # Retry Count Tracking + /// + /// This method tracks retry counts per message (keyed by `{chat}:{msg_id}:{sender}`) + /// and stops sending retry receipts after `MAX_DECRYPT_RETRIES` (5) attempts to prevent + /// infinite retry loops. This matches WhatsApp Web's behavior. + /// + /// # PDO Backup + /// + /// A PDO (Peer Data Operation) request is spawned only on the FIRST retry attempt. + /// This asks our primary phone to share the already-decrypted message content. + /// PDO is NOT spawned on subsequent retries to avoid duplicate requests. + /// + /// When max retries is reached, an immediate PDO request is sent as a last resort. + /// + /// # Arguments + /// * `info` - The message info for the failed message + /// * `reason` - The retry reason code (matches WhatsApp Web's RetryReason enum) + #[cfg(test)] + pub(crate) fn spawn_retry_receipt( + self: &Arc<Self>, + info: &Arc<MessageInfo>, + reason: RetryReason, + ) { + let client = Arc::clone(self); + let info = Arc::clone(info); + self.outbound_flush.spawn(&*self.runtime, async move { + client.run_retry_receipt(&info, reason).await; + }); + } + + /// Increment the retry count and send the retry receipt (or, at the cap, a + /// last-resort PDO). Awaitable so it can be ordered before the transport ack. + /// + /// Returns whether the caller should send the ack: `false` when we intended + /// to retry but the send failed (so the stanza stays queued for another try), + /// `true` when the resend went out or we deliberately gave up at the cap. + async fn run_retry_receipt( + self: &Arc<Self>, + info: &Arc<MessageInfo>, + reason: RetryReason, + ) -> bool { + let cache_key = self + .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) + .await; + + let Some(retry_count) = self.increment_retry_count(&cache_key, reason).await else { + log::info!( + "Max retries ({}) reached for message {} from {} [{:?}]. Sending immediate PDO request.", + MAX_DECRYPT_RETRIES, + info.id, + info.source.sender, + reason + ); + // Capped: give up and clear the backlog regardless of PDO outcome. + self.run_pdo_request(info).await; + return true; + }; + + if retry_count > HIGH_RETRY_COUNT_THRESHOLD { + log::warn!( + "High retry count ({}) for message {} in chat {} from {} [{:?}]", + retry_count, + info.id, + info.source.chat, + info.source.sender, + reason + ); + } + + let retry_sent = match self.send_retry_receipt(info, retry_count, reason).await { + Ok(()) => { + debug!( + "Sent retry receipt #{} for message {} in chat {} from {} [{:?}]", + retry_count, info.id, info.source.chat, info.source.sender, reason + ); + true + } + Err(e) => { + log::error!( + "Failed to send retry receipt #{} for message {} [{:?}]: {:?}", + retry_count, + info.id, + reason, + e + ); + false + } + }; + + // First retry only, to avoid duplicate PDO requests. Awaited so it runs + // before the caller's ack; the retry receipt already landed first. + if retry_count == 1 { + self.run_pdo_request(info).await; + } + retry_sent + } +} diff --git a/src/message/special.rs b/src/message/special.rs new file mode 100644 index 000000000..e942cb83b --- /dev/null +++ b/src/message/special.rs @@ -0,0 +1,232 @@ +//! Special message types: newsletter, app-state key share, sender-key distribution. + +use super::*; + +impl Client { + /// Handles a newsletter plaintext message. + /// Newsletters are not E2E encrypted and use the <plaintext> tag directly. + /// They never carry a `secret_encrypted_message`, so no messageSecret is + /// stored or retained for newsletter chats (no newsletter retention class). + pub(crate) async fn handle_newsletter_message( + self: &Arc<Self>, + node: &NodeRef<'_>, + info: &Arc<MessageInfo>, + ) { + let Some(plaintext_node) = node.get_optional_child_by_tag(&["plaintext"]) else { + log::warn!( + "[msg:{}] Received newsletter message without <plaintext> child: {}", + info.id, + node.tag + ); + return; + }; + + if let Some(bytes) = plaintext_node.content_bytes() { + match wa::Message::decode(bytes) { + Ok(msg) => { + log::info!( + "[msg:{}] Received newsletter plaintext message from {}", + info.id, + info.source.chat + ); + self.dispatch_parsed_message(msg, info).await; + } + Err(e) => { + log::warn!( + "[msg:{}] Failed to decode newsletter plaintext: {e}", + info.id + ); + } + } + } else { + log::debug!( + "[msg:{}] Newsletter <plaintext> node from {} had no content bytes; skipping decode", + info.id, + info.source.chat + ); + } + } + + pub(crate) async fn handle_app_state_sync_key_share( + &self, + keys: &wa::message::AppStateSyncKeyShare, + ) { + struct KeyComponents<'a> { + key_id: &'a [u8], + data: &'a [u8], + fingerprint_bytes: Vec<u8>, + timestamp: i64, + } + + /// Extract components from an AppStateSyncKey for storage. + fn extract_key_components(key: &wa::message::AppStateSyncKey) -> Option<KeyComponents<'_>> { + let key_id = key.key_id.as_ref()?.key_id.as_ref()?; + let key_data = key.key_data.as_ref()?; + let fingerprint = key_data.fingerprint.as_ref()?; + let data = key_data.key_data.as_ref()?; + Some(KeyComponents { + key_id, + data, + fingerprint_bytes: fingerprint.encode_to_vec(), + timestamp: key_data.timestamp(), + }) + } + + let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let key_store = device_snapshot.backend.clone(); + + let mut stored_count = 0; + let mut failed_count = 0; + + for key in &keys.keys { + if let Some(components) = extract_key_components(key) { + let new_key = crate::store::traits::AppStateSyncKey { + key_data: components.data.to_vec(), + fingerprint: components.fingerprint_bytes, + timestamp: components.timestamp, + }; + + if let Err(e) = key_store.set_sync_key(components.key_id, new_key).await { + log::error!( + "Failed to store app state sync key {:?}: {:?}", + hex::encode(components.key_id), + e + ); + failed_count += 1; + } else { + stored_count += 1; + } + } + } + + if stored_count > 0 || failed_count > 0 { + log::info!( + target: "Client/AppState", + "Processed app state key share: {} stored, {} failed.", + stored_count, + failed_count + ); + } + + // Notify any waiters (initial full sync) that at least one key share was processed. + if stored_count > 0 + && !self + .initial_app_state_keys_received + .swap(true, std::sync::atomic::Ordering::Relaxed) + { + // First time setting; notify any waiters + self.initial_keys_synced_notifier.notify(usize::MAX); + } + } + + pub(crate) async fn handle_sender_key_distribution_message( + self: &Arc<Self>, + group_jid: &Jid, + sender_jid: &Jid, + axolotl_bytes: &[u8], + ) { + let skdm = match SenderKeyDistributionMessage::try_from(axolotl_bytes) { + Ok(msg) => msg, + Err(e1) => match wa::SenderKeyDistributionMessage::decode(axolotl_bytes) { + Ok(go_msg) => { + let (Some(signing_key), Some(id), Some(iteration), Some(chain_key)) = ( + go_msg.signing_key.as_ref(), + go_msg.id, + go_msg.iteration, + go_msg.chain_key.as_ref(), + ) else { + log::warn!( + "Go SKDM from {} missing required fields (signing_key={}, id={}, iteration={}, chain_key={})", + sender_jid, + go_msg.signing_key.is_some(), + go_msg.id.is_some(), + go_msg.iteration.is_some(), + go_msg.chain_key.is_some() + ); + return; + }; + let chain_key_arr: [u8; 32] = match chain_key.as_slice().try_into() { + Ok(arr) => arr, + Err(_) => { + log::error!( + "Invalid chain_key length {} from Go SKDM from {}", + chain_key.len(), + sender_jid + ); + return; + } + }; + match SignalPublicKey::from_djb_public_key_bytes(signing_key) { + Ok(pub_key) => { + match SenderKeyDistributionMessage::new( + SENDERKEY_MESSAGE_CURRENT_VERSION, + id, + iteration, + chain_key_arr, + pub_key, + ) { + Ok(skdm) => skdm, + Err(e) => { + log::error!( + "Failed to construct SKDM from Go format from {}: {:?} (original parse error: {:?})", + sender_jid, + e, + e1 + ); + return; + } + } + } + Err(e) => { + log::error!( + "Failed to parse public key from Go SKDM for {}: {:?} (original parse error: {:?})", + sender_jid, + e, + e1 + ); + return; + } + } + } + Err(e2) => { + log::error!( + "Failed to parse SenderKeyDistributionMessage (standard and Go fallback) from {}: primary: {:?}, fallback: {:?}", + sender_jid, + e1, + e2 + ); + return; + } + }, + }; + + // Normalize to bare sender for consistent sender key addressing. + let sender_bare = sender_jid.to_non_ad(); + let sender_address = sender_bare.to_protocol_address(); + + let sender_key_name = make_sender_key_name(group_jid, &sender_address); + + // Route through the signal cache adapter so the sender key is immediately visible + // in the cache for subsequent group_decrypt calls within the same message batch. + // Only the sender-key store is needed here, so build it standalone instead of + // the full five-store adapter. + let mut sender_key_store = self.sender_key_adapter().await; + + if let Err(e) = + process_sender_key_distribution_message(&sender_key_name, &skdm, &mut sender_key_store) + .await + { + log::error!( + "Failed to process SenderKeyDistributionMessage from {}: {:?}", + sender_jid, + e + ); + } else { + log::debug!( + "Successfully processed sender key distribution for group {} from {}", + group_jid, + sender_jid + ); + } + } +} diff --git a/src/message/tests.rs b/src/message/tests.rs new file mode 100644 index 000000000..91769b1a9 --- /dev/null +++ b/src/message/tests.rs @@ -0,0 +1,9553 @@ +//! Tests for the message receive/decrypt pipeline. + +use super::*; +use crate::store::SqliteStore; +use crate::store::persistence_manager::PersistenceManager; +use crate::test_utils::MockHttpClient; +use crate::types::message::EditAttribute; +use std::sync::Arc; +use wacore_binary::builder::NodeBuilder; + +fn node_to_arc(node: wacore_binary::Node) -> Arc<OwnedNodeRef> { + crate::test_utils::node_to_owned_ref(&node) +} +use wacore_binary::{Jid, SERVER_JID}; + +fn mock_transport() -> Arc<dyn crate::transport::TransportFactory> { + Arc::new(crate::transport::mock::MockTransportFactory::new()) +} + +fn mock_http_client() -> Arc<dyn crate::http::HttpClient> { + Arc::new(MockHttpClient) +} + +#[tokio::test] +async fn test_parse_message_info_for_status_broadcast() { + let backend = Arc::new( + SqliteStore::new("file:memdb_status_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let participant_jid_str = "556899336555:42@s.whatsapp.net"; + let status_broadcast_jid_str = "status@broadcast"; + + let node = NodeBuilder::new("message") + .attr("from", status_broadcast_jid_str) + .attr("id", "8A8CCCC7E6E466D9EE8CA11A967E485A") + .attr("participant", participant_jid_str) + .attr("t", "1759295366") + .attr("type", "media") + .build(); + + let info = client + .parse_message_info(&node.as_node_ref()) + .await + .expect("parse_message_info should not fail"); + + let expected_sender: Jid = participant_jid_str + .parse() + .expect("test JID should be valid"); + let expected_chat: Jid = status_broadcast_jid_str + .parse() + .expect("test JID should be valid"); + + assert_eq!( + info.source.sender, expected_sender, + "The sender should be the 'participant' JID, not 'status@broadcast'" + ); + assert_eq!( + info.source.chat, expected_chat, + "The chat should be 'status@broadcast'" + ); + assert!( + info.source.is_group, + "Broadcast messages should be treated as group-like" + ); +} + +#[tokio::test] +async fn test_status_broadcast_cold_cache_resolves_to_lid() { + use wacore::types::jid::JidExt as _; + use wacore_binary::Server; + + let backend = Arc::new( + SqliteStore::new("file:memdb_status_cold_cache?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let pn_user = "559980000001"; + let lid_user = "100000012345678"; + + assert_eq!( + client.lid_pn_cache.get_current_lid(pn_user).await, + None, + "precondition: empty cache for {pn_user}" + ); + + let node = NodeBuilder::new("message") + .attr("from", "status@broadcast") + .attr("id", "TEST_COLD_CACHE_ID") + .attr("participant", format!("{pn_user}@s.whatsapp.net").as_str()) + .attr("participant_lid", format!("{lid_user}@lid").as_str()) + .attr("t", "1777415965") + .attr("type", "media") + .build(); + + let info = client + .parse_message_info(&node.as_node_ref()) + .await + .expect("parse_message_info must succeed"); + + // Fix #1: parser surfaces participant_lid via sender_alt. + let alt = info + .source + .sender_alt + .as_ref() + .expect("sender_alt must be populated from participant_lid"); + assert_eq!(alt.user.as_str(), lid_user); + assert_eq!(alt.server, Server::Lid); + assert_eq!(info.source.sender.user.as_str(), pn_user); + assert_eq!(info.source.sender.server, Server::Pn); + + client + .cache_lid_pn_from_message( + &info.source.sender, + info.source.sender_alt.as_ref(), + info.is_offline, + ) + .await; + + // Cache learned the mapping in both directions. + assert_eq!( + client + .lid_pn_cache + .get_current_lid(pn_user) + .await + .as_deref(), + Some(lid_user), + "PN→LID lookup must hit" + ); + assert_eq!( + client.lid_pn_cache.get_phone_number(lid_user).await, + Some(pn_user.to_string()), + "LID→PN lookup must hit" + ); + + // Resolution upgrades to LID and Signal address is the LID form. + let resolved = client.resolve_encryption_jid(&info.source.sender).await; + assert_eq!(resolved.user.as_str(), lid_user); + assert_eq!(resolved.server, Server::Lid); + assert_eq!(resolved.device, info.source.sender.device); + assert_eq!( + resolved.to_protocol_address().to_string(), + format!("{lid_user}@lid.0"), + "Signal address must be @lid form, not @c.us" + ); +} + +/// Pins the hosted-family branch + the realistic non-zero device shape. +/// Production stanzas almost always have device != 0, and hosted variants +/// (`@hosted` / `@hosted.lid`) must flow through cache_lid_pn_from_message. +#[tokio::test] +async fn test_status_broadcast_hosted_family_with_device_id_resolves_to_hosted_lid() { + use wacore::types::jid::JidExt as _; + use wacore_binary::Server; + + let backend = Arc::new( + SqliteStore::new("file:memdb_status_hosted_device?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let pn_user = "559980000001"; + let lid_user = "100000012345678"; + let device_id: u16 = 99; + + let node = NodeBuilder::new("message") + .attr("from", "status@broadcast") + .attr("id", "HOSTED_TEST_ID") + .attr( + "participant", + format!("{pn_user}:{device_id}@hosted").as_str(), + ) + .attr( + "participant_lid", + format!("{lid_user}:{device_id}@hosted.lid").as_str(), + ) + .attr("t", "1777415965") + .attr("type", "media") + .build(); + + let info = client + .parse_message_info(&node.as_node_ref()) + .await + .expect("parse_message_info must succeed"); + + assert_eq!(info.source.sender.server, Server::Hosted); + assert_eq!(info.source.sender.device, device_id); + let alt = info + .source + .sender_alt + .as_ref() + .expect("sender_alt must be populated for hosted participant"); + assert_eq!(alt.server, Server::HostedLid); + assert_eq!(alt.user.as_str(), lid_user); + assert_eq!(alt.device, device_id); + + client + .cache_lid_pn_from_message( + &info.source.sender, + info.source.sender_alt.as_ref(), + info.is_offline, + ) + .await; + + // Hosted variant must reach the cache; without it, learn_lid_pn_mapping + // is skipped and the hosted-device fix is incomplete. + assert_eq!( + client + .lid_pn_cache + .get_current_lid(pn_user) + .await + .as_deref(), + Some(lid_user), + "PN→LID lookup must work for hosted family" + ); + assert_eq!( + client.lid_pn_cache.get_phone_number(lid_user).await, + Some(pn_user.to_string()), + ); + + let resolved = client.resolve_encryption_jid(&info.source.sender).await; + assert_eq!(resolved.user.as_str(), lid_user); + assert_eq!(resolved.server, Server::HostedLid); + assert_eq!( + resolved.device, device_id, + "device id must be preserved through resolution" + ); + assert_eq!( + resolved.to_protocol_address().to_string(), + format!("{lid_user}:{device_id}@hosted.lid.0"), + "Signal address must be the @hosted.lid form with device suffix" + ); +} + +#[tokio::test] +async fn test_process_session_enc_batch_handles_session_not_found_gracefully() { + use wacore::libsignal::protocol::{IdentityKeyPair, KeyPair, SignalMessage}; + + let backend = Arc::new( + SqliteStore::new("file:memdb_graceful_fail?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let sender_jid: Jid = "1234567890@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let info = Arc::new(MessageInfo { + source: crate::types::message::MessageSource { + sender: sender_jid.clone(), + chat: sender_jid.clone(), + ..Default::default() + }, + ..Default::default() + }); + + // Create a valid but undecryptable SignalMessage + let dummy_key = [0u8; 32]; + let sender_ratchet = KeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()).public_key; + let sender_identity_pair = + IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); + let receiver_identity_pair = + IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); + let signal_message = SignalMessage::new( + 4, + &dummy_key, + sender_ratchet, + 0, + 0, + b"test", + sender_identity_pair.identity_key(), + receiver_identity_pair.identity_key(), + ) + .expect("SignalMessage::new should succeed with valid inputs"); + + let enc_node = NodeBuilder::new("enc") + .attr("type", "msg") + .bytes(signal_message.serialized().to_vec()) + .build(); + let enc_node_ref = enc_node.as_node_ref(); + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; + + let outcome = client + .process_session_enc_batch( + &payloads, + &info, + &sender_jid, + crate::types::events::DecryptFailMode::Show, + ) + .await; + + assert!( + !outcome.decrypted && !outcome.duplicate && outcome.undecryptable, + "process_session_enc_batch should mark SessionNotFound as undecryptable without success or duplicate" + ); +} + +/// Two undecryptable payloads sharing one `(chat, id)` must accumulate +/// `undecryptable` across the batch (monotonic OR) and dispatch exactly one +/// `UndecryptableMessage` event — the single-flight dedup that the accumulator +/// exists to protect, exercised through the batch path with multiple payloads. +#[tokio::test] +async fn batch_accumulates_undecryptable_and_dispatches_once() { + let backend = Arc::new( + SqliteStore::new("file:memdb_batch_undec_once?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let recorder = Arc::new(EventRecorder::default()); + client.register_handler(recorder.clone()); + + let sender_jid: Jid = "1234567890@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let info = Arc::new(MessageInfo { + id: "BATCH_UNDEC_ONCE".to_string(), + source: crate::types::message::MessageSource { + sender: sender_jid.clone(), + chat: sender_jid.clone(), + ..Default::default() + }, + ..Default::default() + }); + + // Two enc payloads whose ciphertext does not parse as a Signal message, so + // each payload independently lands on an undecryptable path. Distinct bytes + // to avoid any incidental dedup on payload content. + let enc1 = NodeBuilder::new("enc") + .attr("type", "msg") + .bytes(vec![0xFF, 0x00, 0x01]) + .build(); + let enc2 = NodeBuilder::new("enc") + .attr("type", "msg") + .bytes(vec![0xFF, 0x00, 0x02]) + .build(); + let enc1_ref = enc1.as_node_ref(); + let enc2_ref = enc2.as_node_ref(); + let payloads: Vec<EncPayload> = vec![ + EncPayload::from_node_ref(&enc1_ref).unwrap(), + EncPayload::from_node_ref(&enc2_ref).unwrap(), + ]; + + let outcome = client + .process_session_enc_batch( + &payloads, + &info, + &sender_jid, + crate::types::events::DecryptFailMode::Show, + ) + .await; + + assert!( + outcome.undecryptable, + "batch must stay undecryptable across both failed payloads" + ); + assert_eq!( + recorder.undecryptable().len(), + 1, + "same (chat, id) dispatches UndecryptableMessage exactly once across the batch" + ); +} + +/// P1: An empty session record (exists but no current/previous state) should be +/// treated the same as SessionNotFound — the retry receipt gets error code 1 (NoSession) +/// and includes keys early, instead of producing an unhelpful InvalidMessage error. +#[tokio::test] +async fn test_empty_session_record_treated_as_session_not_found() { + use wacore::libsignal::protocol::{IdentityKeyPair, KeyPair, SessionRecord, SignalMessage}; + + let backend = Arc::new( + SqliteStore::new("file:memdb_empty_session?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let sender_jid: Jid = "0000000000000@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let info = Arc::new(MessageInfo { + source: crate::types::message::MessageSource { + sender: sender_jid.clone(), + chat: sender_jid.clone(), + ..Default::default() + }, + ..Default::default() + }); + + // Pre-store an empty (degenerate) session record in the signal cache. + // This simulates the bug scenario: record exists but has no usable ratchet state. + let signal_address = sender_jid.to_protocol_address(); + client + .signal_cache + .put_session(&signal_address, SessionRecord::new_fresh()) + .await; + + // Craft a SignalMessage to trigger decryption + let dummy_key = [0u8; 32]; + let sender_ratchet = KeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()).public_key; + let sender_identity = IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); + let receiver_identity = IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); + let signal_message = SignalMessage::new( + 4, + &dummy_key, + sender_ratchet, + 0, + 0, + b"test", + sender_identity.identity_key(), + receiver_identity.identity_key(), + ) + .expect("SignalMessage::new should succeed"); + + let enc_node = NodeBuilder::new("enc") + .attr("type", "msg") + .bytes(signal_message.serialized().to_vec()) + .build(); + let enc_node_ref = enc_node.as_node_ref(); + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; + + let outcome = client + .clone() + .process_session_enc_batch( + &payloads, + &info, + &sender_jid, + crate::types::events::DecryptFailMode::Show, + ) + .await; + + // Should behave identically to SessionNotFound: failure, no dupe, event dispatched. + assert!( + !outcome.decrypted && !outcome.duplicate && outcome.undecryptable, + "Empty session record should be treated as SessionNotFound: \ + expected undecryptable without success or duplicate, got {outcome:?}" + ); + + // After the WA Web compliance fix (no delete on BadMac/InvalidMessage either), + // every inbound-decrypt failure preserves the session. This still pins + // that the empty-record path does not regress to a delete. + let backend = client.persistence_manager.backend(); + let session_still_exists = client + .signal_cache + .has_session(&signal_address, &*backend) + .await + .expect("has_session should not fail"); + assert!(session_still_exists); + + // Discriminate from the BadMac / InvalidMessage arms (which also + // preserve the session post-fix): the empty-record path must end up + // in the SessionNotFound branch, which fires a retry receipt with + // `RetryReason::NoSession`. Anything else means the libsignal-side + // empty-record short-circuit regressed. + await_retry_receipt(&client, &info, 1, RetryReason::NoSession).await; +} + +// ─── Fixtures for session-preservation tests ───────────────────────────── +// +// Mirrors the WAWebSignalProtocolStore tests in spirit: a synthetic peer +// holds its own Signal stores in memory so the test can drive X3DH end to +// end against the Client. Inlined (not exported from a helper crate) +// because these are message.rs-specific scenarios. + +use async_trait::async_trait; +use std::collections::HashMap; +use wacore::libsignal::protocol::{ + CiphertextMessage, Direction, IdentityChange, IdentityKey, IdentityKeyPair, KeyPair, + PreKeyBundle, PreKeyRecord, PreKeyStore as SigPreKeyStore, ProtocolAddress, SenderKeyName, + SenderKeyRecord, SenderKeyStore as SigSenderKeyStore, SessionRecord, + SessionStore as SigSessionStore, SignedPreKeyStore as SigSignedPreKeyStore, UsePQRatchet, + create_sender_key_distribution_message, group_encrypt, message_encrypt, process_prekey_bundle, +}; +use wacore::libsignal::protocol::{IdentityKeyStore as SigIdentityKeyStore, SignalProtocolError}; + +#[derive(Default, Clone)] +struct MemSessionStore(HashMap<ProtocolAddress, SessionRecord>); + +#[async_trait] +impl SigSessionStore for MemSessionStore { + async fn load_session( + &self, + a: &ProtocolAddress, + ) -> Result<Option<SessionRecord>, SignalProtocolError> { + Ok(self.0.get(a).cloned()) + } + async fn has_session(&self, a: &ProtocolAddress) -> Result<bool, SignalProtocolError> { + Ok(self.0.contains_key(a)) + } + async fn store_session( + &mut self, + a: &ProtocolAddress, + r: SessionRecord, + ) -> Result<(), SignalProtocolError> { + self.0.insert(a.clone(), r); + Ok(()) + } +} + +#[derive(Clone)] +struct MemIdentityStore { + kp: IdentityKeyPair, + reg_id: u32, + known: HashMap<ProtocolAddress, IdentityKey>, +} + +#[async_trait] +impl SigIdentityKeyStore for MemIdentityStore { + async fn get_identity_key_pair(&self) -> Result<IdentityKeyPair, SignalProtocolError> { + Ok(self.kp.clone()) + } + async fn get_local_registration_id(&self) -> Result<u32, SignalProtocolError> { + Ok(self.reg_id) + } + async fn save_identity( + &mut self, + a: &ProtocolAddress, + id: &IdentityKey, + ) -> Result<IdentityChange, SignalProtocolError> { + let prev = self.known.insert(a.clone(), *id); + Ok(match prev { + None => IdentityChange::NewOrUnchanged, + Some(p) if &p == id => IdentityChange::NewOrUnchanged, + _ => IdentityChange::ReplacedExisting, + }) + } + async fn is_trusted_identity( + &self, + _: &ProtocolAddress, + _: &IdentityKey, + _: Direction, + ) -> Result<bool, SignalProtocolError> { + Ok(true) + } + async fn get_identity( + &self, + a: &ProtocolAddress, + ) -> Result<Option<IdentityKey>, SignalProtocolError> { + Ok(self.known.get(a).copied()) + } +} + +#[derive(Default, Clone)] +struct MemSenderKeyStore(HashMap<SenderKeyName, SenderKeyRecord>); + +#[async_trait] +impl SigSenderKeyStore for MemSenderKeyStore { + async fn store_sender_key( + &mut self, + name: &SenderKeyName, + record: SenderKeyRecord, + ) -> Result<(), SignalProtocolError> { + self.0.insert(name.clone(), record); + Ok(()) + } + + async fn load_sender_key( + &self, + name: &SenderKeyName, + ) -> Result<Option<SenderKeyRecord>, SignalProtocolError> { + Ok(self.0.get(name).cloned()) + } +} + +#[derive(Clone)] +struct AlicePeer { + jid: Jid, + address: ProtocolAddress, + identity: MemIdentityStore, + sessions: MemSessionStore, + sender_keys: MemSenderKeyStore, +} + +impl AlicePeer { + async fn new(jid_str: &str) -> Self { + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + let kp = IdentityKeyPair::generate(&mut rng); + let jid: Jid = jid_str.parse().expect("valid jid"); + let address = jid.to_protocol_address(); + Self { + jid, + address, + identity: MemIdentityStore { + kp, + reg_id: 12345, + known: HashMap::new(), + }, + sessions: MemSessionStore::default(), + sender_keys: MemSenderKeyStore::default(), + } + } + + async fn install_bob_session(&mut self, bob_addr: &ProtocolAddress, bundle: &PreKeyBundle) { + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + process_prekey_bundle( + bob_addr, + &mut self.sessions, + &mut self.identity, + bundle, + &mut rng, + UsePQRatchet::No, + ) + .await + .expect("process bob bundle"); + } + + async fn encrypt(&mut self, bob_addr: &ProtocolAddress, plaintext: &[u8]) -> CiphertextMessage { + message_encrypt(plaintext, bob_addr, &mut self.sessions, &mut self.identity) + .await + .expect("encrypt") + } + + async fn encrypt_text(&mut self, bob_addr: &ProtocolAddress, text: &str) -> CiphertextMessage { + use wacore::messages::MessageUtils; + + let plaintext = MessageUtils::encode_and_pad(&wa::Message { + conversation: Some(text.to_string()), + ..Default::default() + }); + self.encrypt(bob_addr, &plaintext).await + } + + async fn create_group_skdm( + &mut self, + group_jid: &Jid, + ) -> wa::message::SenderKeyDistributionMessage { + let sender = self.jid.to_non_ad(); + let sender_key_name = make_sender_key_name(group_jid, &sender.to_protocol_address()); + let skdm = create_sender_key_distribution_message( + &sender_key_name, + &mut self.sender_keys, + &mut rand::make_rng::<rand::rngs::StdRng>(), + ) + .await + .expect("create sender key distribution"); + wa::message::SenderKeyDistributionMessage { + group_id: Some(group_jid.to_string()), + axolotl_sender_key_distribution_message: Some(skdm.serialized().to_vec()), + } + } + + async fn encrypt_group_message(&mut self, group_jid: &Jid, plaintext: &[u8]) -> Vec<u8> { + let sender = self.jid.to_non_ad(); + let sender_key_name = make_sender_key_name(group_jid, &sender.to_protocol_address()); + let sender_key_message = group_encrypt( + &mut self.sender_keys, + &sender_key_name, + plaintext, + &mut rand::make_rng::<rand::rngs::StdRng>(), + ) + .await + .expect("encrypt sender key message"); + sender_key_message.serialized().to_vec() + } +} + +/// Ensure the test `Client` has an identity (`pn`/`lid`) provisioned — +/// `create_test_client_with_name` returns an unpaired client by default +/// so `device_snapshot.lid` / `.pn` are both `None`. +async fn ensure_bob_paired(client: &Arc<Client>) { + let snapshot = client.persistence_manager.get_device_snapshot().await; + if snapshot.lid.is_some() || snapshot.pn.is_some() { + return; + } + let pn: Jid = "9000000000000:1@s.whatsapp.net".parse().expect("pn"); + let lid: Jid = "999999999999999:1@lid".parse().expect("lid"); + client + .persistence_manager + .process_command(crate::store::commands::DeviceCommand::SetId(Some(pn))) + .await; + client + .persistence_manager + .process_command(crate::store::commands::DeviceCommand::SetLid(Some(lid))) + .await; +} + +/// Read Bob's currently provisioned identity / signed prekey from the test +/// client and build a `PreKeyBundle` that Alice can use to initialize +/// her side of the session. Mirrors how the real `RetryReceiptJob` ships +/// keys back to the sender — assembled through the same +/// `SignalProtocolStoreAdapter` traits production uses. +async fn bobs_prekey_bundle(client: &Arc<Client>) -> (PreKeyBundle, Jid) { + use wacore::libsignal::protocol::GenericSignedPreKey; + ensure_bob_paired(client).await; + let snapshot = client.persistence_manager.get_device_snapshot().await; + let identity_kp = snapshot.core.identity_key.clone(); + let reg_id = snapshot.core.registration_id; + + // Read/write prekeys through the same trait surface production uses + // (see signal_adapter.rs). Avoids reaching past `PersistenceManager` + // to mutate device storage directly. + let mut adapter = client.signal_adapter().await; + let spk_record = adapter + .signed_pre_key_store + .get_signed_pre_key(1.into()) + .await + .expect("spk present"); + let spk_pub = spk_record.public_key().expect("spk pub"); + let spk_sig_vec = spk_record.signature().expect("spk sig"); + + // Provision a fresh one-time prekey for this test through the + // adapter's `PreKeyStore` impl. + let pk_id_u32: u32 = 9001; + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + let pk_pair = KeyPair::generate(&mut rng); + let pk_record = PreKeyRecord::new(pk_id_u32.into(), &pk_pair); + adapter + .pre_key_store + .save_pre_key(pk_id_u32.into(), &pk_record) + .await + .expect("save pk"); + + let own_device_jid: Jid = snapshot + .lid + .clone() + .or_else(|| snapshot.pn.clone()) + .expect("own jid"); + let bob_jid = own_device_jid.to_non_ad(); + let bundle = PreKeyBundle::new( + reg_id, + u32::from(own_device_jid.device).into(), + Some((pk_id_u32.into(), pk_pair.public_key)), + 1.into(), + spk_pub, + spk_sig_vec, + IdentityKey::new(identity_kp.public_key), + ) + .expect("bundle"); + (bundle, bob_jid) +} + +/// Build an EncPayload-style stanza node and run `process_session_enc_batch`. +/// Returns whether the session for `peer_jid` still exists in the cache afterwards. +async fn submit_and_check_session( + client: &Arc<Client>, + peer_jid: &Jid, + ct: &CiphertextMessage, +) -> (bool, bool, bool, bool) { + let (enc_type, bytes) = match ct { + CiphertextMessage::SignalMessage(m) => ("msg", m.serialized().to_vec()), + CiphertextMessage::PreKeySignalMessage(m) => ("pkmsg", m.serialized().to_vec()), + _ => panic!("unexpected ciphertext type"), + }; + let enc_node = NodeBuilder::new("enc") + .attr("type", enc_type) + .bytes(bytes) + .build(); + let enc_ref = enc_node.as_node_ref(); + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_ref).unwrap()]; + let info = Arc::new(MessageInfo { + source: crate::types::message::MessageSource { + sender: peer_jid.clone(), + chat: peer_jid.clone(), + ..Default::default() + }, + ..Default::default() + }); + let outcome = client + .clone() + .process_session_enc_batch( + &payloads, + &info, + peer_jid, + crate::types::events::DecryptFailMode::Show, + ) + .await; + let backend = client.persistence_manager.backend(); + let still = client + .signal_cache + .has_session(&peer_jid.to_protocol_address(), &*backend) + .await + .expect("has_session"); + ( + outcome.decrypted, + outcome.duplicate, + outcome.undecryptable, + still, + ) +} + +#[tokio::test] +async fn test_badmac_migrates_pn_session_when_lid_shadow_exists() { + use crate::lid_pn_cache::{LearningSource, LidPnEntry}; + + let client = crate::test_utils::create_test_client_with_name("badmac_lid_shadow").await; + let alice_pn: Jid = "15550001001@s.whatsapp.net".parse().expect("alice pn"); + let alice_lid: Jid = "100000000000002@lid".parse().expect("alice lid"); + let entry = LidPnEntry::new( + alice_lid.user.to_string(), + alice_pn.user.to_string(), + LearningSource::PeerLidMessage, + ); + client.lid_pn_cache.add(&entry).await; + + let (bundle_v1, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let alice_pn_str = alice_pn.to_string(); + let mut alice_old = AlicePeer::new(&alice_pn_str).await; + alice_old.install_bob_session(&bob_addr, &bundle_v1).await; + let pkmsg_v1 = alice_old.encrypt_text(&bob_addr, "pn establish").await; + let (pn_success, _, _, pn_still) = + submit_and_check_session(&client, &alice_pn, &pkmsg_v1).await; + assert!(pn_success, "PN-keyed session should establish"); + assert!( + pn_still, + "PN-keyed session should be present before migration" + ); + + if let Some(record) = alice_old.sessions.0.get_mut(&bob_addr) + && let Some(state) = record.session_state_mut() + { + state.clear_unacknowledged_pre_key_message(); + } + + let mut alice_fresh = alice_old.clone(); + alice_fresh.jid = alice_lid.clone(); + alice_fresh.address = alice_lid.to_protocol_address(); + alice_fresh.sessions = MemSessionStore::default(); + + let (bundle_v2, _) = bobs_prekey_bundle(&client).await; + alice_fresh.install_bob_session(&bob_addr, &bundle_v2).await; + let pkmsg_v2 = alice_fresh.encrypt_text(&bob_addr, "lid shadow").await; + let (lid_success, _, _, lid_still) = + submit_and_check_session(&client, &alice_lid, &pkmsg_v2).await; + assert!(lid_success, "LID-keyed shadow session should establish"); + assert!(lid_still, "LID-keyed shadow session should exist"); + + let old_pn_msg = alice_old.encrypt_text(&bob_addr, "old pn ratchet").await; + assert!(matches!(old_pn_msg, CiphertextMessage::SignalMessage(_))); + let (success, duplicates, dispatched, lid_after) = + submit_and_check_session(&client, &alice_lid, &old_pn_msg).await; + assert!(success, "BadMac path should recover by migrating PN to LID"); + assert!(!duplicates, "message should decrypt, not dedupe"); + assert!( + !dispatched, + "migration recovery must not emit retry failure" + ); + assert!(lid_after, "migrated LID session should remain"); + + let backend = client.persistence_manager.backend(); + let pn_after = client + .signal_cache + .has_session(&alice_pn.to_protocol_address(), &*backend) + .await + .expect("has_session"); + assert!(!pn_after, "PN session should be consumed by migration"); +} + +#[tokio::test] +async fn migration_plaintext_failure_nacks_without_signal_retry() { + use crate::lid_pn_cache::{LearningSource, LidPnEntry}; + + let (client, transport) = capturing_client("migration_plaintext_nack").await; + let alice_pn: Jid = "15550001002@s.whatsapp.net".parse().expect("alice pn"); + let alice_lid: Jid = "100000000000004@lid".parse().expect("alice lid"); + let entry = LidPnEntry::new( + alice_lid.user.to_string(), + alice_pn.user.to_string(), + LearningSource::PeerLidMessage, + ); + client.lid_pn_cache.add(&entry).await; + + let (bundle_v1, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let alice_pn_str = alice_pn.to_string(); + let mut alice_old = AlicePeer::new(&alice_pn_str).await; + alice_old.install_bob_session(&bob_addr, &bundle_v1).await; + let pkmsg_v1 = alice_old.encrypt_text(&bob_addr, "pn establish").await; + let (pn_success, _, _, _) = submit_and_check_session(&client, &alice_pn, &pkmsg_v1).await; + assert!(pn_success, "PN-keyed session should establish"); + + if let Some(record) = alice_old.sessions.0.get_mut(&bob_addr) + && let Some(state) = record.session_state_mut() + { + state.clear_unacknowledged_pre_key_message(); + } + + let mut alice_fresh = alice_old.clone(); + alice_fresh.jid = alice_lid.clone(); + alice_fresh.address = alice_lid.to_protocol_address(); + alice_fresh.sessions = MemSessionStore::default(); + + let (bundle_v2, _) = bobs_prekey_bundle(&client).await; + alice_fresh.install_bob_session(&bob_addr, &bundle_v2).await; + let pkmsg_v2 = alice_fresh.encrypt_text(&bob_addr, "lid shadow").await; + let (lid_success, _, _, _) = submit_and_check_session(&client, &alice_lid, &pkmsg_v2).await; + assert!(lid_success, "LID-keyed shadow session should establish"); + + let bad_old_pn_msg = alice_old.encrypt(&bob_addr, &[0xff, 0x01]).await; + let payloads = vec![enc_payload_from_ciphertext(&bad_old_pn_msg)]; + let info = Arc::new(MessageInfo { + id: "MIGRATION_BAD_PLAINTEXT".to_string(), + source: crate::types::message::MessageSource { + sender: alice_lid.clone(), + chat: alice_lid.clone(), + ..Default::default() + }, + ..Default::default() + }); + + let outcome = client + .clone() + .process_session_enc_batch( + &payloads, + &info, + &alice_lid, + crate::types::events::DecryptFailMode::Show, + ) + .await; + + assert!( + outcome.decrypted, + "Signal decrypt succeeded after migration" + ); + assert!(outcome.plaintext_failed); + assert!(outcome.undecryptable); + assert!(outcome.had_failure); + assert!(!outcome.dispatched); + assert!(!outcome.skdm_only); + + let mut nack_code = None; + for _ in 0..80 { + nack_code = find_message_nack_error(&transport.sent(), &info.id); + if nack_code.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!(nack_code, Some(491)); + + let cache_key = client + .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) + .await; + assert_eq!( + client.message_retry_counts.get(&cache_key).await, + None, + "local protobuf failure after migration must not request Signal retry" + ); +} + +/// Smoking-gun regression: a `BadMac` on the inbound path must NOT delete +/// the session. Pre-fix, `src/message.rs:1100` called +/// `signal_cache.delete_session(...)` here — this test would fail with +/// `still=false`. WA Web's `RetryReceiptJob` keeps the session untouched +/// (see `docs/captured-js/WAWeb/Send/RetryReceiptJob.js`). +#[tokio::test] +async fn test_badmac_preserves_session() { + let client = crate::test_utils::create_test_client_with_name("badmac_preserves").await; + let mut alice = AlicePeer::new("1111111111111@s.whatsapp.net").await; + let alice_addr = alice.address.clone(); + + // X3DH: Alice consumes Bob's bundle to set up her outgoing session. + let (bob_bundle, _) = bobs_prekey_bundle(&client).await; + alice + .install_bob_session( + &client + .persistence_manager + .get_device_snapshot() + .await + .lid + .clone() + .or(client + .persistence_manager + .get_device_snapshot() + .await + .pn + .clone()) + .expect("own jid") + .to_protocol_address(), + &bob_bundle, + ) + .await; + + // First message: pkmsg lands on Bob and installs Bob's reciprocal session. + let bob_addr = client + .persistence_manager + .get_device_snapshot() + .await + .lid + .clone() + .or(client + .persistence_manager + .get_device_snapshot() + .await + .pn + .clone()) + .expect("own jid") + .to_protocol_address(); + let pkmsg = alice.encrypt_text(&bob_addr, "hello").await; + let (s1, _, _, still1) = submit_and_check_session(&client, &alice.jid, &pkmsg).await; + assert!(s1, "pkmsg should establish session and decrypt"); + assert!(still1, "session must exist after first message"); + + // Force Alice's next encrypt to be a plain SignalMessage rather than a + // pkmsg by clearing her unacknowledged-pkmsg flag. Tampering the trailing + // bytes of a pkmsg breaks the outer protobuf parse (because reg_id / + // signed_pre_key_id varints are encoded *after* the embedded message + // field), which would short-circuit into the parse-error nack path + // before ever reaching the BadMac arm we want to exercise. + { + let record = alice + .sessions + .0 + .get_mut(&bob_addr) + .expect("alice has a session for bob"); + if let Some(state) = record.session_state_mut() { + state.clear_unacknowledged_pre_key_message(); + } + } + + // Second message: tamper the trailing MAC byte of a real SignalMessage. + // The format is `[version][protobuf body][8-byte MAC]`, so the last byte + // is squarely inside the MAC region — parse succeeds, MAC verification + // fails -> libsignal returns BadMac. + let msg2 = alice.encrypt_text(&bob_addr, "world").await; + let mut bytes = match &msg2 { + CiphertextMessage::SignalMessage(m) => m.serialized().to_vec(), + other => panic!( + "expected SignalMessage, got {:?}", + std::mem::discriminant(other) + ), + }; + let last = bytes.len() - 1; + bytes[last] ^= 0xFF; + let enc_node = NodeBuilder::new("enc") + .attr("type", "msg") + .bytes(bytes) + .build(); + let enc_ref = enc_node.as_node_ref(); + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_ref).unwrap()]; + let info = Arc::new(MessageInfo { + id: "BADMAC_TAMPER_MSG".to_string(), + source: crate::types::message::MessageSource { + sender: alice.jid.clone(), + chat: alice.jid.clone(), + ..Default::default() + }, + ..Default::default() + }); + + let outcome = client + .clone() + .process_session_enc_batch( + &payloads, + &info, + &alice.jid, + crate::types::events::DecryptFailMode::Show, + ) + .await; + assert!(!outcome.decrypted, "tampered MAC must not decrypt"); + assert!( + outcome.undecryptable, + "undecryptable event must be dispatched" + ); + + // The fix asserts the session lives on so the eventual sender pkmsg + // can archive it into previous_sessions[0]. + let backend = client.persistence_manager.backend(); + let still = client + .signal_cache + .has_session(&alice_addr, &*backend) + .await + .expect("has_session"); + assert!(still, "BadMac must NOT delete the session (WA Web parity)"); + + // Discriminate from the parse-error path (which also preserves the + // session): the BadMac/InvalidMessage branch routes through + // `handle_decrypt_failure` -> `spawn_retry_receipt`, which bumps + // both caches with `RetryReason::BadMac`. Parse errors take the + // nack path instead and never touch either cache. + await_retry_receipt(&client, &info, 1, RetryReason::BadMac).await; +} + +/// Poll for `message_retry_counts == expected_count` AND +/// `recent_retry_reasons == expected_reason` (or fail after a short +/// timeout). `spawn_retry_receipt` detaches the increment onto the +/// runtime, so both caches may lag the `process_session_enc_batch` return. +/// Reading both is what tells the BadMac arm apart from a parse-error +/// regression (which never bumps these caches). +async fn await_retry_receipt( + client: &Arc<Client>, + info: &MessageInfo, + expected_count: u8, + expected_reason: RetryReason, +) { + let cache_key = client + .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) + .await; + for _ in 0..200 { + if let (Some(c), Some(r)) = ( + client.message_retry_counts.get(&cache_key).await, + client.recent_retry_reasons.get(&cache_key).await, + ) && c == expected_count + && r == expected_reason + { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + let count = client.message_retry_counts.get(&cache_key).await; + let reason = client.recent_retry_reasons.get(&cache_key).await; + panic!( + "expected retry ({expected_count}, {expected_reason:?}) for {cache_key}, \ + got ({count:?}, {reason:?})" + ); +} + +// NOTE: the `InvalidMessage` arm of the `matches!()` block in +// `process_session_enc_batch` is exercised by `test_badmac_preserves_session` +// too — libsignal returns `BadMac` whenever *any* candidate state derives a +// message key (which is what a random-ratchet `SignalMessage::new(...)` +// ends up doing as well), so a separate "InvalidMessage" regression test +// would be indistinguishable from the BadMac one. Reaching the +// `InvalidMessage` constructor specifically would require crafting a +// SignalMessage that *parses* but where no state derives any message +// key — empirically impractical without major libsignal-side scaffolding. + +/// Integration test: reproduces the production loop observed in +/// `k8awqjsgww2lnkt89urp3de1-191402150615-...`. After a BadMac the bot +/// used to delete the session; when the sender then sent a fresh pkmsg +/// (post-retry-receipt), `process_prekey_bundle` ran on an empty record +/// and `previous_sessions[0]` stayed empty — any in-flight messages on +/// the OLD ratchet failed permanently. With the fix the old session +/// survives the BadMac, the pkmsg's `promote_state` archives it, and +/// the archived state lives in `previous_sessions[0]` exactly as WA Web +/// expects (see `libsignal/src/protocol/state/session.rs:751-768`). +#[tokio::test] +async fn test_prod_scenario_pkmsg_archives_old_session_after_badmac() { + let client = crate::test_utils::create_test_client_with_name("prod_archive").await; + let mut alice = AlicePeer::new("3333333333333@s.whatsapp.net").await; + + // X3DH round 1 — Alice initiates with Bob's bundle, sends pkmsg. + let (bundle_v1, _) = bobs_prekey_bundle(&client).await; + let bob_addr = client + .persistence_manager + .get_device_snapshot() + .await + .lid + .clone() + .or(client + .persistence_manager + .get_device_snapshot() + .await + .pn + .clone()) + .expect("own jid") + .to_protocol_address(); + alice.install_bob_session(&bob_addr, &bundle_v1).await; + let pkmsg_v1 = alice.encrypt_text(&bob_addr, "v1").await; + let (s1, _, _, _) = submit_and_check_session(&client, &alice.jid, &pkmsg_v1).await; + assert!(s1); + + // Snapshot Bob's session_v1 base key for later comparison. Use + // peek (non-destructive): `get_session` marks the cache entry as + // CheckedOut, which would prevent libsignal from re-loading the + // session in the BadMac path that follows. + let alice_addr = alice.address.clone(); + let backend = client.persistence_manager.backend(); + let v1_record = client + .signal_cache + .peek_session(&alice_addr, &*backend) + .await + .expect("peek_session") + .expect("v1 session present"); + let v1_base_key = v1_record + .session_state() + .expect("v1 current state") + .sender_ratchet_key_for_logging() + .expect("v1 base key"); + + // Force Alice's next encrypt to be a plain SignalMessage so tampering + // the last byte lands inside the MAC region (see comment in + // `test_badmac_preserves_session` for why pkmsg cannot be tampered + // at the tail without breaking the outer protobuf parse). + { + let record = alice + .sessions + .0 + .get_mut(&bob_addr) + .expect("alice has a session for bob"); + if let Some(state) = record.session_state_mut() { + state.clear_unacknowledged_pre_key_message(); + } + } + + // Tampered SignalMessage → BadMac branch (with the fix this no longer + // deletes Bob's session). + let msg = alice.encrypt_text(&bob_addr, "stale").await; + let mut bytes = match &msg { + CiphertextMessage::SignalMessage(m) => m.serialized().to_vec(), + other => panic!( + "expected SignalMessage, got {:?}", + std::mem::discriminant(other) + ), + }; + let last = bytes.len() - 1; + bytes[last] ^= 0xFF; + let enc_node = NodeBuilder::new("enc") + .attr("type", "msg") + .bytes(bytes) + .build(); + let enc_ref = enc_node.as_node_ref(); + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_ref).unwrap()]; + let info = Arc::new(MessageInfo { + id: "PROD_LOOP_REPRO_STALE".to_string(), + source: crate::types::message::MessageSource { + sender: alice.jid.clone(), + chat: alice.jid.clone(), + ..Default::default() + }, + ..Default::default() + }); + let _outcome = client + .clone() + .process_session_enc_batch( + &payloads, + &info, + &alice.jid, + crate::types::events::DecryptFailMode::Show, + ) + .await; + // Confirm the BadMac branch executed (parse-error path would skip + // both retry caches; another arm would record a different reason). + await_retry_receipt(&client, &info, 1, RetryReason::BadMac).await; + // Pre-fix: this assertion would have failed (session deleted). + let preserved = client + .signal_cache + .has_session(&alice_addr, &*backend) + .await + .expect("has_session"); + assert!(preserved, "BadMac must preserve session"); + + // X3DH round 2 — Alice rebuilds her side from a fresh Bob bundle + // (simulates the bot re-issuing prekeys via a retry receipt) and + // sends another pkmsg. Bob's `process_prekey_bundle` must archive + // session_v1 into previous_sessions[0]. + let (bundle_v2, _) = bobs_prekey_bundle(&client).await; + alice.sessions = MemSessionStore::default(); // forget Alice's v1 to force a fresh X3DH + alice.install_bob_session(&bob_addr, &bundle_v2).await; + let pkmsg_v2 = alice.encrypt_text(&bob_addr, "v2").await; + let (s2, _, _, still2) = submit_and_check_session(&client, &alice.jid, &pkmsg_v2).await; + assert!(s2, "pkmsg_v2 should decrypt"); + assert!(still2); + + let v2_record = client + .signal_cache + .peek_session(&alice_addr, &*backend) + .await + .expect("peek_session") + .expect("v2 session present"); + let v2_base_key = v2_record + .session_state() + .expect("v2 current state") + .sender_ratchet_key_for_logging() + .expect("v2 base key"); + assert_ne!( + v1_base_key, v2_base_key, + "current session must be the new v2" + ); + assert_eq!( + v2_record.previous_session_count(), + 1, + "session_v1 must be archived as previous_sessions[0]" + ); + let archived_state = v2_record + .previous_session_states() + .next() + .expect("archived state") + .expect("archived state decodes"); + let archived_base_key = archived_state + .sender_ratchet_key_for_logging() + .expect("archived base key"); + assert_eq!( + archived_base_key, v1_base_key, + "archived previous_sessions[0] must be the original v1" + ); +} + +#[tokio::test] +async fn test_handle_incoming_message_skips_skmsg_after_msg_failure() { + use wacore::libsignal::protocol::{IdentityKeyPair, KeyPair, SignalMessage}; + + let backend = Arc::new( + SqliteStore::new("file:memdb_skip_skmsg_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let sender_jid: Jid = "1234567890@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let group_jid: Jid = "120363021033254949@g.us" + .parse() + .expect("test JID should be valid"); + + // Create msg + skmsg node; msg will fail (no session), so skmsg should be skipped + let dummy_key = [0u8; 32]; + let sender_ratchet = KeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()).public_key; + let sender_identity_pair = + IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); + let receiver_identity_pair = + IdentityKeyPair::generate(&mut rand::make_rng::<rand::rngs::StdRng>()); + let signal_message = SignalMessage::new( + 4, + &dummy_key, + sender_ratchet, + 0, + 0, + b"test", + sender_identity_pair.identity_key(), + receiver_identity_pair.identity_key(), + ) + .expect("SignalMessage::new should succeed with valid inputs"); + + let msg_node = NodeBuilder::new("enc") + .attr("type", "msg") + .bytes(signal_message.serialized().to_vec()) + .build(); + + let skmsg_node = NodeBuilder::new("enc") + .attr("type", "skmsg") + .bytes(vec![4, 5, 6]) + .build(); + + let message_node = node_to_arc( + NodeBuilder::new("message") + .attr("from", group_jid) + .attr("participant", sender_jid) + .attr("id", "test-id-123") + .attr("t", "12345") + .children(vec![msg_node, skmsg_node]) + .build(), + ); + + // Should not panic or retry loop - skmsg is skipped after msg failure + client.clone().handle_incoming_message(message_node).await; +} + +/// Test case for reproducing sender key JID mismatch in LID group messages +/// +/// Problem: +/// - When we process sender key distribution from a self-sent LID message, we store it under the LID JID +/// - But when we try to decrypt the group content (skmsg), we look it up using the phone number JID +/// - This causes "No sender key state" errors even though we just processed the sender key! +/// +/// This test verifies the fix by: +/// 1. Creating a sender key and storing it under the LID address (mimicking SKDM processing) +/// 2. Attempting retrieval with phone number address (the bug) - should fail +/// 3. Attempting retrieval with LID address (the fix) - should succeed +#[tokio::test] +async fn test_self_sent_lid_group_message_sender_key_mismatch() { + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore::libsignal::protocol::{ + SenderKeyStore, create_sender_key_distribution_message, + process_sender_key_distribution_message, + }; + + let backend = Arc::new( + SqliteStore::new("file:memdb_sender_key_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (_client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm.clone(), + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let own_lid: Jid = "100000000000001.1:75@lid" + .parse() + .expect("test JID should be valid"); + let own_phone: Jid = "15551234567:75@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let group_jid: Jid = "120363021033254949@g.us" + .parse() + .expect("test JID should be valid"); + + // Create SKDM using LID address (mimics handle_sender_key_distribution_message) + let lid_protocol_address = own_lid.to_protocol_address(); + let lid_sender_key_name = make_sender_key_name(&group_jid, &lid_protocol_address); + + // Pin serialized form so from_jid stays compatible with persisted records + assert_eq!(lid_sender_key_name.group_id(), group_jid.to_string()); + assert_eq!( + lid_sender_key_name.sender_id(), + lid_protocol_address.to_string() + ); + + let device_arc = pm.get_device_arc().await; + let skdm = { + let mut device_guard = device_arc.write().await; + create_sender_key_distribution_message( + &lid_sender_key_name, + &mut *device_guard, + &mut rand::make_rng::<rand::rngs::StdRng>(), + ) + .await + .expect("Failed to create SKDM") + }; + + { + let mut device_guard = device_arc.write().await; + process_sender_key_distribution_message(&lid_sender_key_name, &skdm, &mut *device_guard) + .await + .expect("Failed to process SKDM with LID address"); + } + + // Try to retrieve using PHONE NUMBER address (THE BUG) + let phone_protocol_address = own_phone.to_protocol_address(); + let phone_sender_key_name = make_sender_key_name(&group_jid, &phone_protocol_address); + + let phone_lookup_result = { + let device_guard = device_arc.read().await; + device_guard.load_sender_key(&phone_sender_key_name).await + }; + + assert!( + phone_lookup_result + .expect("lookup should not error") + .is_none(), + "Sender key should NOT be found when looking up with phone number address (demonstrates the bug)" + ); + + // Try to retrieve using LID address (THE FIX) + let lid_lookup_result = { + let device_guard = device_arc.read().await; + device_guard.load_sender_key(&lid_sender_key_name).await + }; + + assert!( + lid_lookup_result + .expect("lookup should not error") + .is_some(), + "Sender key SHOULD be found when looking up with LID address (same as storage)" + ); +} + +/// Test that sender key consistency is maintained for multiple LID participants +/// +/// Edge case: Group with multiple LID participants, each should have their own +/// sender key stored under their LID address, not mixed up with phone numbers. +#[tokio::test] +async fn test_multiple_lid_participants_sender_key_isolation() { + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore::libsignal::protocol::{ + SenderKeyStore, create_sender_key_distribution_message, + process_sender_key_distribution_message, + }; + + let backend = Arc::new( + SqliteStore::new("file:memdb_multi_lid_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let transport_factory = Arc::new(crate::transport::mock::MockTransportFactory::new()); + let (_client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm.clone(), + transport_factory, + mock_http_client(), + None, + ) + .await; + + let group_jid: Jid = "120363021033254949@g.us" + .parse() + .expect("test JID should be valid"); + + // Simulate three LID participants + let participants = vec![ + ("100000000000001.1:75@lid", "15551234567:75@s.whatsapp.net"), + ("987654321000000.2:42@lid", "551234567890:42@s.whatsapp.net"), + ("111222333444555.3:10@lid", "559876543210:10@s.whatsapp.net"), + ]; + + let device_arc = pm.get_device_arc().await; + + // Create and store sender keys for each participant under their LID address + for (lid_str, _phone_str) in &participants { + let lid_jid: Jid = lid_str.parse().expect("test JID should be valid"); + let lid_protocol_address = lid_jid.to_protocol_address(); + let lid_sender_key_name = make_sender_key_name(&group_jid, &lid_protocol_address); + + let skdm = { + let mut device_guard = device_arc.write().await; + create_sender_key_distribution_message( + &lid_sender_key_name, + &mut *device_guard, + &mut rand::make_rng::<rand::rngs::StdRng>(), + ) + .await + .expect("Failed to create SKDM") + }; + + let mut device_guard = device_arc.write().await; + process_sender_key_distribution_message(&lid_sender_key_name, &skdm, &mut *device_guard) + .await + .expect("Failed to process SKDM"); + } + + // Verify each participant's sender key can be retrieved using their LID address + for (lid_str, phone_str) in &participants { + let lid_jid: Jid = lid_str.parse().expect("test JID should be valid"); + let phone_jid: Jid = phone_str.parse().expect("test JID should be valid"); + + let lid_protocol_address = lid_jid.to_protocol_address(); + let phone_protocol_address = phone_jid.to_protocol_address(); + + let lid_sender_key_name = make_sender_key_name(&group_jid, &lid_protocol_address); + let phone_sender_key_name = make_sender_key_name(&group_jid, &phone_protocol_address); + + // Should find with LID address + let lid_lookup = { + let device_guard = device_arc.read().await; + device_guard.load_sender_key(&lid_sender_key_name).await + }; + assert!( + lid_lookup.expect("lookup should not error").is_some(), + "Sender key for {} should be found with LID address", + lid_str + ); + + // Should NOT find with phone number address (the bug) + let phone_lookup = { + let device_guard = device_arc.read().await; + device_guard.load_sender_key(&phone_sender_key_name).await + }; + assert!( + phone_lookup.expect("lookup should not error").is_none(), + "Sender key for {} should NOT be found with phone number address", + lid_str + ); + } +} + +/// Test that LID JID parsing handles various edge cases correctly +/// +/// Edge cases: +/// - LID with multiple dots in user portion +/// - LID with device numbers +/// - LID without device numbers +#[test] +fn test_lid_jid_parsing_edge_cases() { + use wacore_binary::Jid; + + // Single dot in user portion + let lid1: Jid = "100000000000001.1:75@lid" + .parse() + .expect("test JID should be valid"); + assert_eq!(lid1.user, "100000000000001.1"); + assert_eq!(lid1.device, 75); + assert_eq!(lid1.agent, 0); + + // Multiple dots in user portion (extreme edge case) + let lid2: Jid = "123.456.789.0:50@lid" + .parse() + .expect("test JID should be valid"); + assert_eq!(lid2.user, "123.456.789.0"); + assert_eq!(lid2.device, 50); + assert_eq!(lid2.agent, 0); + + // No device number (device 0) + let lid3: Jid = "987654321000000.5@lid" + .parse() + .expect("test JID should be valid"); + assert_eq!(lid3.user, "987654321000000.5"); + assert_eq!(lid3.device, 0); + assert_eq!(lid3.agent, 0); + + // Very long user portion with dot + let lid4: Jid = "111222333444555666777.999:1@lid" + .parse() + .expect("test JID should be valid"); + assert_eq!(lid4.user, "111222333444555666777.999"); + assert_eq!(lid4.device, 1); + assert_eq!(lid4.agent, 0); +} + +/// Test that protocol address generation from LID JIDs matches WhatsApp Web format +/// +/// WhatsApp Web uses: {user}[:device]@{server}.0 +/// - The device is encoded in the name +/// - device_id is always 0 +#[test] +fn test_lid_protocol_address_consistency() { + use wacore::types::jid::JidExt as CoreJidExt; + use wacore_binary::Jid; + + // Format: (jid_str, expected_name, expected_device_id, expected_to_string) + let test_cases = vec![ + ( + "100000000000001.1:75@lid", + "100000000000001.1:75@lid", + 0, + "100000000000001.1:75@lid.0", + ), + ( + "987654321000000.2:42@lid", + "987654321000000.2:42@lid", + 0, + "987654321000000.2:42@lid.0", + ), + ( + "111.222.333:10@lid", + "111.222.333:10@lid", + 0, + "111.222.333:10@lid.0", + ), + // No device - should not include :0 + ("123456789@lid", "123456789@lid", 0, "123456789@lid.0"), + ]; + + for (jid_str, expected_name, expected_device_id, expected_to_string) in test_cases { + let lid_jid: Jid = jid_str.parse().expect("test JID should be valid"); + let protocol_addr = lid_jid.to_protocol_address(); + + assert_eq!( + protocol_addr.name(), + expected_name, + "Protocol address name should match WhatsApp Web's SignalAddress format for {}", + jid_str + ); + assert_eq!( + u32::from(protocol_addr.device_id()), + expected_device_id, + "Protocol address device_id should always be 0 for {}", + jid_str + ); + assert_eq!( + protocol_addr.to_string(), + expected_to_string, + "Protocol address to_string() should match createSignalLikeAddress format for {}", + jid_str + ); + } +} + +/// Test sender_alt extraction from message attributes in LID groups +/// +/// Edge cases: +/// - LID group with participant_pn attribute +/// - PN group with participant_lid attribute +/// - Mixed addressing modes +#[tokio::test] +async fn test_parse_message_info_sender_alt_extraction() { + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore::types::message::AddressingMode; + use wacore_binary::builder::NodeBuilder; + + let backend = Arc::new( + SqliteStore::new("file:memdb_sender_alt_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + + // Set up own phone number and LID + { + let device_arc = pm.get_device_arc().await; + let mut device = device_arc.write().await; + device.pn = Some( + "15551234567@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), + ); + device.lid = Some( + "100000000000001.1@lid" + .parse() + .expect("test JID should be valid"), + ); + } + + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + // Test case 1: LID group message with participant_pn + let lid_group_node = NodeBuilder::new("message") + .attr("from", "120363021033254949@g.us") + .attr("participant", "987654321000000.2:42@lid") + .attr("participant_pn", "551234567890:42@s.whatsapp.net") + .attr("addressing_mode", AddressingMode::Lid.as_str()) + .attr("id", "test1") + .attr("t", "12345") + .build(); + + let info1 = client + .parse_message_info(&lid_group_node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + assert_eq!(info1.source.sender.user, "987654321000000.2"); + assert!(info1.source.sender_alt.is_some()); + assert_eq!( + info1 + .source + .sender_alt + .as_ref() + .expect("sender_alt should be present") + .user, + "551234567890" + ); + + // Test case 2: Self-sent LID group message + let self_lid_node = NodeBuilder::new("message") + .attr("from", "120363021033254949@g.us") + .attr("participant", "100000000000001.1:75@lid") + .attr("participant_pn", "15551234567:75@s.whatsapp.net") + .attr("addressing_mode", AddressingMode::Lid.as_str()) + .attr("id", "test2") + .attr("t", "12346") + .build(); + + let info2 = client + .parse_message_info(&self_lid_node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + assert!( + info2.source.is_from_me, + "Should detect self-sent LID message" + ); + assert_eq!(info2.source.sender.user, "100000000000001.1"); + assert!(info2.source.sender_alt.is_some()); + assert_eq!( + info2 + .source + .sender_alt + .as_ref() + .expect("sender_alt should be present") + .user, + "15551234567" + ); +} + +/// Test that device query logic uses phone numbers for LID participants +/// +/// This is a unit test for the logic in wacore/src/send.rs that converts +/// LID JIDs to phone number JIDs for device queries. +#[test] +fn test_lid_to_phone_mapping_for_device_queries() { + use std::collections::HashMap; + use wacore::client::context::GroupInfo; + use wacore::types::message::AddressingMode; + use wacore_binary::Jid; + + // Simulate a LID group with phone number mappings + let mut lid_to_pn_map = HashMap::new(); + lid_to_pn_map.insert( + wacore_binary::CompactString::from("100000000000001.1"), + "15551234567@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), + ); + lid_to_pn_map.insert( + wacore_binary::CompactString::from("987654321000000.2"), + "551234567890@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), + ); + + let mut group_info = GroupInfo::new( + vec![ + "100000000000001.1:75@lid" + .parse() + .expect("test JID should be valid"), + "987654321000000.2:42@lid" + .parse() + .expect("test JID should be valid"), + ], + AddressingMode::Lid, + ); + group_info.set_lid_to_pn_map(lid_to_pn_map.clone()); + + // Simulate the device query logic + let jids_to_query: Vec<Jid> = group_info + .participants + .iter() + .map(|jid| { + let base_jid = jid.to_non_ad(); + if base_jid.is_lid() + && let Some(phone_jid) = group_info.phone_jid_for_lid_user(&base_jid.user) + { + return phone_jid.to_non_ad(); + } + base_jid + }) + .collect(); + + // Verify all queries use phone numbers, not LID JIDs + for jid in &jids_to_query { + assert_eq!( + jid.server, SERVER_JID, + "Device query should use phone number, got: {}", + jid + ); + } + + assert_eq!(jids_to_query.len(), 2); + assert!(jids_to_query.iter().any(|j| j.user == "15551234567")); + assert!(jids_to_query.iter().any(|j| j.user == "551234567890")); +} + +/// Test edge case: Group with mixed LID and phone number participants +/// +/// Some participants may still use phone numbers even in a LID group. +/// The code should handle both correctly. +#[test] +fn test_mixed_lid_and_phone_participants() { + use std::collections::HashMap; + use wacore::client::context::GroupInfo; + use wacore::types::message::AddressingMode; + use wacore_binary::Jid; + + let mut lid_to_pn_map = HashMap::new(); + lid_to_pn_map.insert( + wacore_binary::CompactString::from("100000000000001.1"), + "15551234567@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), + ); + + let mut group_info = GroupInfo::new( + vec![ + "100000000000001.1:75@lid" + .parse() + .expect("test JID should be valid"), // LID participant + "551234567890:42@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), // Phone number participant + ], + AddressingMode::Lid, + ); + group_info.set_lid_to_pn_map(lid_to_pn_map.clone()); + + let jids_to_query: Vec<Jid> = group_info + .participants + .iter() + .map(|jid| { + let base_jid = jid.to_non_ad(); + if base_jid.is_lid() + && let Some(phone_jid) = group_info.phone_jid_for_lid_user(&base_jid.user) + { + return phone_jid.to_non_ad(); + } + base_jid + }) + .collect(); + + // Both should end up as phone numbers + assert_eq!(jids_to_query.len(), 2); + for jid in &jids_to_query { + assert_eq!(jid.server, SERVER_JID); + } +} + +/// Test edge case: Own JID check in LID mode +/// +/// When checking if own JID is in the participant list, we must use +/// the phone number equivalent if in LID mode, not the LID itself. +#[test] +fn test_own_jid_check_in_lid_mode() { + use std::collections::HashMap; + use wacore_binary::Jid; + + let own_lid: Jid = "100000000000001.1@lid" + .parse() + .expect("test JID should be valid"); + let own_phone: Jid = "15551234567@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + + let mut lid_to_pn_map = HashMap::new(); + lid_to_pn_map.insert("100000000000001.1".to_string(), own_phone.clone()); + + // Simulate the own JID check logic from wacore/src/send.rs + let own_base_jid = own_lid.to_non_ad(); + let own_jid_to_check = if own_base_jid.is_lid() { + lid_to_pn_map + .get(own_base_jid.user.as_str()) + .map(|pn| pn.to_non_ad()) + .unwrap_or_else(|| own_base_jid.clone()) + } else { + own_base_jid.clone() + }; + + // Verify we're checking using the phone number + assert_eq!(own_jid_to_check.user, "15551234567"); + assert_eq!(own_jid_to_check.server, SERVER_JID); +} + +/// Test that sender key operations always use the display JID (LID) +/// regardless of what JID is used for E2E session decryption +#[tokio::test] +async fn test_sender_key_always_uses_display_jid() { + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore::libsignal::protocol::{SenderKeyStore, create_sender_key_distribution_message}; + + let backend = Arc::new( + SqliteStore::new("file:memdb_display_jid_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (_client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm.clone(), + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let group_jid: Jid = "120363021033254949@g.us" + .parse() + .expect("test JID should be valid"); + let display_jid: Jid = "100000000000001.1:75@lid" + .parse() + .expect("test JID should be valid"); + let encryption_jid: Jid = "15551234567:75@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + + // Store sender key using display JID (LID) + let display_protocol_address = display_jid.to_protocol_address(); + let display_sender_key_name = make_sender_key_name(&group_jid, &display_protocol_address); + + let device_arc = pm.get_device_arc().await; + { + let mut device_guard = device_arc.write().await; + create_sender_key_distribution_message( + &display_sender_key_name, + &mut *device_guard, + &mut rand::make_rng::<rand::rngs::StdRng>(), + ) + .await + .expect("Failed to create SKDM"); + } + + // Verify it's stored under display JID + let lookup_with_display = { + let device_guard = device_arc.read().await; + device_guard.load_sender_key(&display_sender_key_name).await + }; + assert!( + lookup_with_display + .expect("lookup should not error") + .is_some(), + "Sender key should be found with display JID (LID)" + ); + + // Verify it's NOT accessible via encryption JID (phone number) + let encryption_protocol_address = encryption_jid.to_protocol_address(); + let encryption_sender_key_name = make_sender_key_name(&group_jid, &encryption_protocol_address); + + let lookup_with_encryption = { + let device_guard = device_arc.read().await; + device_guard + .load_sender_key(&encryption_sender_key_name) + .await + }; + assert!( + lookup_with_encryption + .expect("lookup should not error") + .is_none(), + "Sender key should NOT be found with encryption JID (phone number)" + ); +} + +/// Test edge case: Second message with only skmsg (no pkmsg/msg) +/// +/// After the first message establishes a session and sender key, +/// subsequent messages may contain only skmsg. These should still +/// be decrypted successfully, not skipped. +/// +/// Bug: The code was treating "no session messages" as "session failed", +/// causing it to skip skmsg decryption for all messages after the first. +#[tokio::test] +async fn test_second_message_with_only_skmsg_decrypts() { + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore::libsignal::protocol::{ + create_sender_key_distribution_message, process_sender_key_distribution_message, + }; + + use wacore::types::message::AddressingMode; + use wacore_binary::builder::NodeBuilder; + + let backend = Arc::new( + SqliteStore::new("file:memdb_second_msg_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm.clone(), + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let sender_jid: Jid = "100000000000001.1:75@lid" + .parse() + .expect("test JID should be valid"); + let group_jid: Jid = "120363021033254949@g.us" + .parse() + .expect("test JID should be valid"); + + // Step 1: Create and store a sender key (simulating first message processing) + let sender_protocol_address = sender_jid.to_protocol_address(); + let sender_key_name = make_sender_key_name(&group_jid, &sender_protocol_address); + + let device_arc = pm.get_device_arc().await; + { + let mut device_guard = device_arc.write().await; + let skdm = create_sender_key_distribution_message( + &sender_key_name, + &mut *device_guard, + &mut rand::make_rng::<rand::rngs::StdRng>(), + ) + .await + .expect("Failed to create SKDM"); + + process_sender_key_distribution_message(&sender_key_name, &skdm, &mut *device_guard) + .await + .expect("Failed to process SKDM"); + } + + // Create message with ONLY skmsg (simulating second message after session established) + let skmsg_ciphertext = { + let mut device_guard = device_arc.write().await; + let sender_key_msg = wacore::libsignal::protocol::group_encrypt( + &mut *device_guard, + &sender_key_name, + b"ping", + &mut rand::make_rng::<rand::rngs::StdRng>(), + ) + .await + .expect("Failed to encrypt with sender key"); + sender_key_msg.serialized().to_vec() + }; + + let skmsg_node = NodeBuilder::new("enc") + .attr("type", "skmsg") + .attr("v", "2") + .bytes(skmsg_ciphertext) + .build(); + + let message_node = node_to_arc( + NodeBuilder::new("message") + .attr("from", group_jid) + .attr("participant", sender_jid) + .attr("id", "SECOND_MSG_TEST") + .attr("t", "1759306493") + .attr("type", "text") + .attr("addressing_mode", AddressingMode::Lid.as_str()) + .children(vec![skmsg_node]) + .build(), + ); + + // Should NOT skip skmsg - before the fix this would incorrectly skip + client.clone().handle_incoming_message(message_node).await; +} + +/// Test case for UntrustedIdentity error handling and recovery +/// +/// Scenario: +/// - User re-installs WhatsApp or switches devices +/// - Their device generates a new identity key +/// - The bot still has the old identity key stored +/// - When a message arrives, Signal Protocol rejects it as "UntrustedIdentity" +/// - The bot should catch this error, clear the old identity using the FULL protocol address (with device ID), and retry +/// +/// This test verifies that: +/// 1. process_session_enc_batch handles UntrustedIdentity gracefully +/// 2. The deletion uses the correct full address (name.device_id) not just the name +/// 3. No panic occurs when UntrustedIdentity is encountered +/// 4. The error is logged appropriately +/// 5. The bot continues processing instead of propagating the error +#[tokio::test] +async fn test_untrusted_identity_error_is_caught_and_handled() { + use crate::store::SqliteStore; + use std::sync::Arc; + + // Setup + let backend = Arc::new( + SqliteStore::new("file:memdb_untrusted_identity_caught?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm.clone(), + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let sender_jid: Jid = "559981212574@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + + let info = Arc::new(MessageInfo { + source: crate::types::message::MessageSource { + sender: sender_jid.clone(), + chat: sender_jid.clone(), + ..Default::default() + }, + ..Default::default() + }); + + log::info!("Test: UntrustedIdentity scenario for {}", sender_jid); + + // Create a malformed/invalid encrypted node to trigger error handling path + // This won't create UntrustedIdentity specifically, but tests the error handling code path + // The important fix is that when UntrustedIdentity IS raised, the code uses + // address.to_string() (which gives "559981212574.0") instead of address.name() + // (which only gives "559981212574") for the deletion key. + let enc_node = NodeBuilder::new("enc") + .attr("type", "msg") + .attr("v", "2") + .bytes(vec![0xFF; 100]) // Invalid encrypted payload + .build(); + + let enc_node_ref = enc_node.as_node_ref(); + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; + + // Call process_session_enc_batch + // This should handle any errors gracefully without panicking + let outcome = client + .process_session_enc_batch( + &payloads, + &info, + &sender_jid, + crate::types::events::DecryptFailMode::Show, + ) + .await; + + log::info!( + "Test: process_session_enc_batch completed - success: {}", + outcome.decrypted + ); + + // The key is that this didn't panic - deletion uses full protocol address +} + +/// Test case: Error handling during batch processing +/// +/// When multiple messages are being processed in a batch, if one triggers +/// an error (like UntrustedIdentity), it should be handled without affecting +/// other messages in the batch. +#[tokio::test] +async fn test_untrusted_identity_does_not_break_batch_processing() { + use crate::store::SqliteStore; + use std::sync::Arc; + + let backend = Arc::new( + SqliteStore::new("file:memdb_untrusted_batch?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm.clone(), + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let sender_jid: Jid = "559981212574@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + + let info = Arc::new(MessageInfo { + source: crate::types::message::MessageSource { + sender: sender_jid.clone(), + chat: sender_jid.clone(), + ..Default::default() + }, + ..Default::default() + }); + + log::info!("Test: Batch processing with multiple error messages"); + + // Create multiple invalid encrypted nodes to test batch error handling + let mut enc_nodes = Vec::new(); + + // First message: Invalid encrypted payload + let enc_node_1 = NodeBuilder::new("enc") + .attr("type", "msg") + .attr("v", "2") + .bytes(vec![0xFF; 50]) + .build(); + enc_nodes.push(enc_node_1); + + // Second message: Another invalid encrypted payload + let enc_node_2 = NodeBuilder::new("enc") + .attr("type", "msg") + .attr("v", "2") + .bytes(vec![0xAA; 50]) + .build(); + enc_nodes.push(enc_node_2); + + log::info!("Test: Created batch of 2 messages with invalid data"); + + let payloads: Vec<EncPayload> = enc_nodes + .iter() + .filter_map(|n| EncPayload::from_node_ref(&n.as_node_ref())) + .collect(); + + // Process the batch + // Should handle all errors gracefully without stopping at first error + let outcome = client + .process_session_enc_batch( + &payloads, + &info, + &sender_jid, + crate::types::events::DecryptFailMode::Show, + ) + .await; + + log::info!( + "Test: Batch processing completed - success: {}", + outcome.decrypted + ); +} + +/// Test case: Error handling in group chat context +/// +/// When processing messages from group members, if identity errors occur, +/// they should be handled per-sender without affecting other group members. +#[tokio::test] +async fn test_untrusted_identity_in_group_context() { + use crate::store::SqliteStore; + use std::sync::Arc; + + let backend = Arc::new( + SqliteStore::new("file:memdb_untrusted_group?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm.clone(), + mock_transport(), + mock_http_client(), + None, + ) + .await; + + // Simulate a group chat scenario + let group_jid: Jid = "120363021033254949@g.us" + .parse() + .expect("test JID should be valid"); + let sender_phone: Jid = "559981212574@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + + let info = Arc::new(MessageInfo { + source: crate::types::message::MessageSource { + sender: sender_phone.clone(), + chat: group_jid.clone(), + is_group: true, + ..Default::default() + }, + ..Default::default() + }); + + log::info!("Test: Group context - error handling for {}", sender_phone); + + // Create an invalid encrypted message + let enc_node = NodeBuilder::new("enc") + .attr("type", "msg") + .attr("v", "2") + .bytes(vec![0xFF; 100]) + .build(); + + let enc_node_ref = enc_node.as_node_ref(); + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; + + // Process the message + // Should handle errors gracefully in group context + let outcome = client + .process_session_enc_batch( + &payloads, + &info, + &sender_phone, + crate::types::events::DecryptFailMode::Show, + ) + .await; + + log::info!( + "Test: Group message processed - success: {}", + outcome.decrypted + ); +} + +/// Test case: DM message parsing for self-sent messages via LID +/// +/// Scenario: +/// - You send a DM to another user from your phone +/// - Your bot receives the echo with from=your_LID, recipient=their_LID +/// - peer_recipient_pn contains the RECIPIENT's phone number (not sender's) +/// +/// The fix ensures: +/// 1. is_from_me is correctly detected for LID senders +/// 2. sender_alt is NOT populated with peer_recipient_pn (that's the recipient's PN) +/// 3. Decryption uses own PN via the is_from_me fallback path +#[tokio::test] +async fn test_parse_message_info_self_sent_dm_via_lid() { + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore_binary::builder::NodeBuilder; + + let backend = Arc::new( + SqliteStore::new("file:memdb_self_dm_lid_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + + // Set up own phone number and LID + { + let device_arc = pm.get_device_arc().await; + let mut device = device_arc.write().await; + device.pn = Some( + "15551234567@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), + ); + device.lid = Some( + "100000000000001@lid" + .parse() + .expect("test JID should be valid"), + ); + } + + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + // Simulate self-sent DM to another user (from your phone to your bot echo) + // Real log example: + // from="100000000000001@lid" recipient="39492358562039@lid" peer_recipient_pn="559985213786@s.whatsapp.net" + let self_dm_node = NodeBuilder::new("message") + .attr("from", "100000000000001@lid") // Your LID + .attr("recipient", "39492358562039@lid") // Recipient's LID + .attr("peer_recipient_pn", "559985213786@s.whatsapp.net") // Recipient's PN (NOT sender's!) + .attr("notify", "jl") + .attr("id", "AC756E00B560721DBC4C0680131827EA") + .attr("t", "1764845025") + .attr("type", "text") + .build(); + + let info = client + .parse_message_info(&self_dm_node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + // Assertions: + // 1. is_from_me should be true (LID matches own_lid) + assert!( + info.source.is_from_me, + "Should detect self-sent DM from own LID" + ); + + // 2. sender_alt should be own PN (derived from own_jid, not message attrs) + assert!( + info.source.sender_alt.is_some(), + "sender_alt should be own PN for self-sent LID messages" + ); + assert_eq!( + info.source.sender_alt.as_ref().unwrap().user, + "15551234567", + "sender_alt should be the own PN user" + ); + + assert_eq!( + info.source.chat.user, "39492358562039", + "Chat should be the recipient's LID" + ); + + assert_eq!( + info.source.sender.user, "100000000000001", + "Sender should be own LID" + ); +} + +/// Test case: DM message parsing for messages from others via LID +/// +/// Scenario: +/// - Another user sends you a DM +/// - Message arrives with from=their_LID, sender_pn=their_phone_number +/// +/// The fix ensures: +/// 1. is_from_me is false +/// 2. sender_alt is populated from sender_pn attribute (if present) +/// 3. Decryption uses sender_alt for session lookup +#[tokio::test] +async fn test_parse_message_info_dm_from_other_via_lid() { + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore_binary::builder::NodeBuilder; + + let backend = Arc::new( + SqliteStore::new("file:memdb_other_dm_lid_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + + // Set up own phone number and LID + { + let device_arc = pm.get_device_arc().await; + let mut device = device_arc.write().await; + device.pn = Some( + "15551234567@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), + ); + device.lid = Some( + "100000000000001@lid" + .parse() + .expect("test JID should be valid"), + ); + } + + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + // Simulate DM from another user via their LID + // The sender_pn attribute should contain their phone number for session lookup + let other_dm_node = NodeBuilder::new("message") + .attr("from", "39492358562039@lid") // Sender's LID (not ours) + .attr("sender_pn", "559985213786@s.whatsapp.net") // Sender's phone number + .attr("notify", "Other User") + .attr("id", "AABBCCDD1234567890") + .attr("t", "1764845100") + .attr("type", "text") + .build(); + + let info = client + .parse_message_info(&other_dm_node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + assert!( + !info.source.is_from_me, + "Should NOT be detected as self-sent" + ); + + assert!( + info.source.sender_alt.is_some(), + "sender_alt should be set from sender_pn attribute" + ); + assert_eq!( + info.source + .sender_alt + .as_ref() + .expect("sender_alt should be present") + .user, + "559985213786", + "sender_alt should contain sender's phone number" + ); + + assert_eq!( + info.source.chat.user, "39492358562039", + "Chat should be the sender's LID (non-AD)" + ); + + assert_eq!( + info.source.sender.user, "39492358562039", + "Sender should be other user's LID" + ); +} + +/// Test case: DM message to self (own chat, like "Notes to Myself") +/// +/// Scenario: +/// - You send a message to yourself (your own chat) +/// - from=your_LID, recipient=your_LID, peer_recipient_pn=your_PN +/// +/// This is the original bug case that was fixed earlier. +#[tokio::test] +async fn test_parse_message_info_dm_to_self() { + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore_binary::builder::NodeBuilder; + + let backend = Arc::new( + SqliteStore::new("file:memdb_dm_to_self_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + + // Set up own phone number and LID + { + let device_arc = pm.get_device_arc().await; + let mut device = device_arc.write().await; + device.pn = Some( + "15551234567@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), + ); + device.lid = Some( + "100000000000001@lid" + .parse() + .expect("test JID should be valid"), + ); + } + + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + // Simulate DM to self (like "Notes to Myself" or pinging yourself) + // from=your_LID, recipient=your_LID, peer_recipient_pn=your_PN + let self_chat_node = NodeBuilder::new("message") + .attr("from", "100000000000001@lid") // Your LID + .attr("recipient", "100000000000001@lid") // Also your LID (self-chat) + .attr("peer_recipient_pn", "15551234567@s.whatsapp.net") // Your PN + .attr("notify", "jl") + .attr("id", "AC391DD54A28E1CE1F3B106DF9951FAD") + .attr("t", "1764822437") + .attr("type", "text") + .build(); + + let info = client + .parse_message_info(&self_chat_node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + assert!( + info.source.is_from_me, + "Should detect self-sent message to self-chat" + ); + + assert!( + info.source.sender_alt.is_some(), + "sender_alt should be own PN for self-sent LID messages" + ); + assert_eq!( + info.source.sender_alt.as_ref().unwrap().user, + "15551234567", + "sender_alt should match own PN" + ); + + assert_eq!( + info.source.chat.user, "100000000000001", + "Chat should be self (recipient)" + ); + + assert_eq!( + info.source.sender.user, "100000000000001", + "Sender should be own LID" + ); +} + +/// Test that receiving a DM with sender_lid populates the lid_pn_cache. +/// +/// This is the key behavior for the LID-PN session mismatch fix: +/// When we receive a message from a phone number with sender_lid attribute, +/// we cache the phone->LID mapping so that when sending replies, we can +/// reuse the existing LID session instead of creating a new PN session. +/// +/// Flow being tested: +/// 1. Receive message from 559980000001@s.whatsapp.net with sender_lid=100000012345678@lid +/// 2. Cache should be populated with: 559980000001 -> 100000012345678 +/// 3. When sending reply to 559980000001, we can look up the LID and use existing session +#[tokio::test] +async fn test_lid_pn_cache_populated_on_message_with_sender_lid() { + // Setup client + let backend = Arc::new( + SqliteStore::new("file:memdb_lid_cache_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let phone = "559980000001"; + let lid = "100000012345678"; + + // Verify cache is empty initially + assert!( + client.lid_pn_cache.get_current_lid(phone).await.is_none(), + "Cache should be empty before receiving message" + ); + + // Create a DM message node with sender_lid attribute + // This simulates receiving a message from WhatsApp Web + let dm_node = NodeBuilder::new("message") + .attr("from", Jid::pn(phone).to_string()) + .attr("sender_lid", Jid::lid(lid).to_string()) + .attr("id", "TEST123456789") + .attr("t", "1765482972") + .attr("type", "text") + .children([NodeBuilder::new("enc") + .attr("type", "pkmsg") + .attr("v", "2") + .bytes(vec![0u8; 100]) // Dummy encrypted content + .build()]) + .build(); + + // Call handle_incoming_message - this will fail to decrypt (no real session) + // but it should still populate the cache before attempting decryption + client + .clone() + .handle_incoming_message(node_to_arc(dm_node)) + .await; + + // Verify the cache was populated + let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; + assert!( + cached_lid.is_some(), + "Cache should be populated after receiving message with sender_lid" + ); + assert_eq!( + cached_lid.expect("cache should have LID"), + lid, + "Cached LID should match the sender_lid from the message" + ); +} + +/// Test that messages without sender_lid do NOT populate the cache. +/// +/// This ensures we don't accidentally cache incorrect mappings. +#[tokio::test] +async fn test_lid_pn_cache_not_populated_without_sender_lid() { + // Setup client + let backend = Arc::new( + SqliteStore::new("file:memdb_no_lid_cache_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let phone = "559980000001"; + + // Create a DM message node WITHOUT sender_lid attribute + let dm_node = NodeBuilder::new("message") + .attr("from", Jid::pn(phone).to_string()) + // Note: NO sender_lid attribute + .attr("id", "TEST123456789") + .attr("t", "1765482972") + .attr("type", "text") + .children([NodeBuilder::new("enc") + .attr("type", "pkmsg") + .attr("v", "2") + .bytes(vec![0u8; 100]) + .build()]) + .build(); + + // Call handle_incoming_message + client + .clone() + .handle_incoming_message(node_to_arc(dm_node)) + .await; + + assert!( + client.lid_pn_cache.get_current_lid(phone).await.is_none(), + "Cache should NOT be populated for messages without sender_lid" + ); +} + +/// Test that messages from LID senders with participant_pn DO populate the cache. +/// +/// When the sender is a LID (e.g., in LID-mode groups), and participant_pn +/// contains their phone number, we SHOULD cache this mapping because: +/// 1. The cache is bidirectional - we need both LID->PN and PN->LID +/// 2. This enables sending to users we've only seen as LID senders +#[tokio::test] +async fn test_lid_pn_cache_populated_for_lid_sender_with_participant_pn() { + use wacore::types::message::AddressingMode; + + // Setup client + let backend = Arc::new( + SqliteStore::new("file:memdb_lid_sender_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let lid = "100000012345678"; + let phone = "559980000001"; + + // Create a message from a LID sender with participant_pn attribute + // This happens in LID-mode groups (addressing_mode="lid") + let group_node = NodeBuilder::new("message") + .attr("from", "120363123456789012@g.us") // Group chat + .attr("participant", Jid::lid(lid).to_string()) // Sender is LID + .attr("participant_pn", Jid::pn(phone).to_string()) // Their phone number + .attr("addressing_mode", AddressingMode::Lid.as_str()) // Required for participant_pn to be parsed + .attr("id", "TEST123456789") + .attr("t", "1765482972") + .attr("type", "text") + .children([NodeBuilder::new("enc") + .attr("type", "skmsg") + .attr("v", "2") + .bytes(vec![0u8; 100]) + .build()]) + .build(); + + // Call handle_incoming_message + client + .clone() + .handle_incoming_message(node_to_arc(group_node)) + .await; + + // Verify the cache WAS populated (bidirectional cache) + let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; + assert!( + cached_lid.is_some(), + "Cache should be populated for LID senders with participant_pn" + ); + assert_eq!( + cached_lid.expect("cache should have LID"), + lid, + "Cached LID should match the sender's LID" + ); + + // Also verify we can look up the phone number from the LID + let cached_pn = client.lid_pn_cache.get_phone_number(lid).await; + assert!(cached_pn.is_some(), "Reverse lookup (LID->PN) should work"); + assert_eq!( + cached_pn.expect("reverse lookup should return phone"), + phone, + "Cached phone number should match" + ); +} + +/// Test that multiple messages from the same sender update the cache correctly. +/// +/// This ensures the cache handles repeated messages gracefully. +#[tokio::test] +async fn test_lid_pn_cache_handles_repeated_messages() { + // Setup client + let backend = Arc::new( + SqliteStore::new("file:memdb_repeated_msg_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let phone = "559980000001"; + let lid = "100000012345678"; + + // Send multiple messages from the same sender + for i in 0..3 { + let dm_node = NodeBuilder::new("message") + .attr("from", Jid::pn(phone).to_string()) + .attr("sender_lid", Jid::lid(lid).to_string()) + .attr("id", format!("TEST{}", i)) + .attr("t", "1765482972") + .attr("type", "text") + .children([NodeBuilder::new("enc") + .attr("type", "pkmsg") + .attr("v", "2") + .bytes(vec![0u8; 100]) + .build()]) + .build(); + + client + .clone() + .handle_incoming_message(node_to_arc(dm_node)) + .await; + } + + // Verify the cache still has the correct mapping + let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; + assert!(cached_lid.is_some(), "Cache should contain the mapping"); + assert_eq!( + cached_lid.expect("cache should have LID"), + lid, + "Cached LID should be correct after multiple messages" + ); +} + +/// Test that PN-addressed messages use LID for session lookup when LID mapping is known. +/// +/// This test verifies the fix for the MAC verification failure bug: +/// WhatsApp Web's SignalAddress.toString() ALWAYS converts PN addresses to LID +/// when a LID mapping is known. The Rust client must do the same to ensure +/// session keys match between clients. +/// +/// Bug scenario: +/// 1. WhatsApp Web Client A sends a group message to our Rust client +/// 2. Rust client creates session under PN address (559980000001@c.us.0) +/// 3. Rust client sends group response, creates session under LID (100000012345678@lid.0) +/// 4. Client A sends DM to Rust client from PN address +/// 5. Rust client tries to decrypt using PN address but session is under LID +/// 6. MAC verification fails because wrong session is used +/// +/// Fix: When receiving a PN-addressed message, if we have a LID mapping, +/// use the LID address for session lookup (matching WhatsApp Web behavior). +#[tokio::test] +async fn test_pn_message_uses_lid_for_session_lookup_when_mapping_known() { + use crate::lid_pn_cache::LidPnEntry; + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore::types::jid::JidExt; + + let backend = Arc::new( + SqliteStore::new("file:memdb_pn_to_lid_session_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let lid = "100000012345678"; + let phone = "559980000001"; + + // Pre-populate the LID-PN cache (simulating a previous group message) + let entry = LidPnEntry::new( + lid.to_string(), + phone.to_string(), + crate::lid_pn_cache::LearningSource::PeerLidMessage, + ); + client.lid_pn_cache.add(&entry).await; + + // Verify the cache has the mapping + let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; + assert_eq!( + cached_lid.as_deref(), + Some(lid), + "Cache should have the LID-PN mapping" + ); + + // Test scenario: Parse a PN-addressed DM message (with sender_lid attribute) + let dm_node_with_sender_lid = wacore_binary::builder::NodeBuilder::new("message") + .attr("from", Jid::pn(phone).to_string()) + .attr("sender_lid", Jid::lid(lid).to_string()) + .attr("id", "test_dm_with_lid") + .attr("t", "1765494882") + .attr("type", "text") + .build(); + + let info = client + .parse_message_info(&dm_node_with_sender_lid.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + // Verify sender is PN but sender_alt is LID + assert_eq!(info.source.sender.user, phone); + assert_eq!(info.source.sender.server, wacore_binary::Server::Pn); + assert!(info.source.sender_alt.is_some()); + assert_eq!( + info.source + .sender_alt + .as_ref() + .expect("sender_alt should be present") + .user, + lid + ); + assert_eq!( + info.source + .sender_alt + .as_ref() + .expect("sender_alt should be present") + .server, + wacore_binary::Server::Lid + ); + + // Now simulate what handle_incoming_message does: determine encryption JID + // We can't easily call handle_incoming_message, so we'll test the logic directly + let sender = &info.source.sender; + let alt = info.source.sender_alt.as_ref(); + // Apply the same logic as in handle_incoming_message + let sender_encryption_jid = if sender.is_lid() { + sender.clone() + } else if sender.is_pn() { + if let Some(alt_jid) = alt + && alt_jid.is_lid() + { + // Use the LID from the message attribute + Jid { + user: alt_jid.user.clone(), + server: wacore_binary::Server::Lid, + device: sender.device, + agent: sender.agent, + integrator: sender.integrator, + } + } else if let Some(lid_user) = client.lid_pn_cache.get_current_lid(&sender.user).await { + // Use the cached LID + Jid { + user: lid_user, + server: wacore_binary::Server::Lid, + device: sender.device, + agent: sender.agent, + integrator: sender.integrator, + } + } else { + sender.clone() + } + } else { + sender.clone() + }; + + // Verify the encryption JID uses the LID, not the PN + assert_eq!( + sender_encryption_jid.user, lid, + "Encryption JID should use LID user" + ); + assert_eq!( + sender_encryption_jid.server, + wacore_binary::Server::Lid, + "Encryption JID should use LID server" + ); + + // Verify the protocol address format + let protocol_address = sender_encryption_jid.to_protocol_address(); + assert_eq!( + protocol_address.to_string(), + format!("{}@lid.0", lid), + "Protocol address should be in LID format" + ); +} + +/// Test that PN-addressed messages use cached LID even without sender_lid attribute. +/// +/// This tests the fallback path where the message doesn't have a sender_lid +/// attribute but we have a previously cached LID mapping. +#[tokio::test] +async fn test_pn_message_uses_cached_lid_without_sender_lid_attribute() { + use crate::lid_pn_cache::LidPnEntry; + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore::types::jid::JidExt; + + let backend = Arc::new( + SqliteStore::new("file:memdb_cached_lid_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let lid = "100000012345678"; + let phone = "559980000001"; + + // Pre-populate the LID-PN cache + let entry = LidPnEntry::new( + lid.to_string(), + phone.to_string(), + crate::lid_pn_cache::LearningSource::PeerLidMessage, + ); + client.lid_pn_cache.add(&entry).await; + + // Parse a PN-addressed DM message WITHOUT sender_lid attribute + let dm_node_without_sender_lid = wacore_binary::builder::NodeBuilder::new("message") + .attr("from", Jid::pn(phone).to_string()) + // Note: No sender_lid attribute! + .attr("id", "test_dm_no_lid") + .attr("t", "1765494882") + .attr("type", "text") + .build(); + + let info = client + .parse_message_info(&dm_node_without_sender_lid.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + // Verify sender is PN and NO sender_alt (since there's no sender_lid attribute) + assert_eq!(info.source.sender.user, phone); + assert_eq!(info.source.sender.server, wacore_binary::Server::Pn); + assert!( + info.source.sender_alt.is_none(), + "Should have no sender_alt without sender_lid attribute" + ); + + // Apply the encryption JID logic (fallback to cached LID) + let sender = &info.source.sender; + let alt = info.source.sender_alt.as_ref(); + let sender_encryption_jid = if sender.is_lid() { + sender.clone() + } else if sender.is_pn() { + if let Some(alt_jid) = alt + && alt_jid.is_lid() + { + Jid { + user: alt_jid.user.clone(), + server: wacore_binary::Server::Lid, + device: sender.device, + agent: sender.agent, + integrator: sender.integrator, + } + } else if let Some(lid_user) = client.lid_pn_cache.get_current_lid(&sender.user).await { + // This is the path we're testing - fallback to cached LID + Jid { + user: lid_user, + server: wacore_binary::Server::Lid, + device: sender.device, + agent: sender.agent, + integrator: sender.integrator, + } + } else { + sender.clone() + } + } else { + sender.clone() + }; + + // Verify the encryption JID uses the cached LID + assert_eq!( + sender_encryption_jid.user, lid, + "Encryption JID should use cached LID user" + ); + assert_eq!( + sender_encryption_jid.server, + wacore_binary::Server::Lid, + "Encryption JID should use LID server" + ); + + let protocol_address = sender_encryption_jid.to_protocol_address(); + assert_eq!( + protocol_address.to_string(), + format!("{}@lid.0", lid), + "Protocol address should be in LID format from cached mapping" + ); +} + +/// Test that PN-addressed messages use PN when no LID mapping is known. +/// +/// When there's no LID mapping available, we should fall back to using +/// the PN address for session lookup. +#[tokio::test] +async fn test_pn_message_uses_pn_when_no_lid_mapping() { + use crate::store::SqliteStore; + use std::sync::Arc; + use wacore::types::jid::JidExt; + + let backend = Arc::new( + SqliteStore::new("file:memdb_no_lid_mapping_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let phone = "559980000001"; + + // Don't populate the cache - simulate first-time contact + + // Parse a PN-addressed DM message without sender_lid + let dm_node = wacore_binary::builder::NodeBuilder::new("message") + .attr("from", Jid::pn(phone).to_string()) + .attr("id", "test_dm_no_mapping") + .attr("t", "1765494882") + .attr("type", "text") + .build(); + + let info = client + .parse_message_info(&dm_node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + // Verify no cached LID + let cached_lid = client.lid_pn_cache.get_current_lid(phone).await; + assert!(cached_lid.is_none(), "Should have no cached LID mapping"); + + // Apply the encryption JID logic + let sender = &info.source.sender; + let alt = info.source.sender_alt.as_ref(); + + let sender_encryption_jid = if sender.is_lid() { + sender.clone() + } else if sender.is_pn() { + if let Some(alt_jid) = alt + && alt_jid.is_lid() + { + Jid { + user: alt_jid.user.clone(), + server: wacore_binary::Server::Lid, + device: sender.device, + agent: sender.agent, + integrator: sender.integrator, + } + } else if let Some(lid_user) = client.lid_pn_cache.get_current_lid(&sender.user).await { + Jid { + user: lid_user, + server: wacore_binary::Server::Lid, + device: sender.device, + agent: sender.agent, + integrator: sender.integrator, + } + } else { + // This is the path we're testing - no LID mapping, use PN + sender.clone() + } + } else { + sender.clone() + }; + + // Verify the encryption JID uses the PN (no LID available) + assert_eq!( + sender_encryption_jid.user, phone, + "Encryption JID should use PN user when no LID mapping" + ); + assert_eq!( + sender_encryption_jid.server, + wacore_binary::Server::Pn, + "Encryption JID should use PN server when no LID mapping" + ); + + let protocol_address = sender_encryption_jid.to_protocol_address(); + assert_eq!( + protocol_address.to_string(), + format!("{}@c.us.0", phone), + "Protocol address should be in PN format when no LID mapping" + ); +} + +// and PDO fallback behavior to ensure robust message recovery. + +/// Helper to create a test MessageInfo with customizable fields +fn create_test_message_info(chat: &str, msg_id: &str, sender: &str) -> MessageInfo { + use wacore::types::message::{EditAttribute, MessageCategory, MessageSource, MsgMetaInfo}; + + let chat_jid: Jid = chat.parse().expect("valid chat JID"); + let sender_jid: Jid = sender.parse().expect("valid sender JID"); + + MessageInfo { + id: msg_id.to_string(), + server_id: 0, + r#type: "text".to_string(), + source: MessageSource { + chat: chat_jid.clone(), + sender: sender_jid, + sender_alt: None, + recipient_alt: None, + is_from_me: false, + is_group: chat_jid.is_group(), + addressing_mode: None, + broadcast_list_owner: None, + recipient: None, + }, + timestamp: wacore::time::now_utc(), + push_name: "Test User".to_string(), + category: MessageCategory::default(), + multicast: false, + media_type: "".to_string(), + edit: EditAttribute::default(), + bot_info: None, + meta_info: MsgMetaInfo::default(), + verified_name: None, + device_sent_meta: None, + ephemeral_expiration: None, + is_offline: false, + unavailable_request_id: None, + server_timestamp_us: None, + verified_level: None, + verified_name_serial: None, + peer_recipient_pn: None, + bcl_participants: Vec::new(), + } +} + +/// Helper to create a test client for retry tests with a unique database +async fn create_test_client_for_retry_with_id(test_id: &str) -> Arc<Client> { + use portable_atomic::AtomicU64; + use std::sync::atomic::Ordering; + static COUNTER: AtomicU64 = AtomicU64::new(0); + + let unique_id = COUNTER.fetch_add(1, Ordering::SeqCst); + let db_name = format!( + "file:memdb_retry_{}_{}_{}?mode=memory&cache=shared", + test_id, + unique_id, + std::process::id() + ); + + let backend = Arc::new( + SqliteStore::new(&db_name) + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + client +} + +#[tokio::test] +async fn test_increment_retry_count_starts_at_one() { + let client = create_test_client_for_retry_with_id("starts_at_one").await; + + let cache_key = "test_chat:msg123:sender456"; + + // First increment should return 1 + let count = client + .increment_retry_count(cache_key, RetryReason::NoSession) + .await; + assert_eq!(count, Some(1), "First retry should be count 1"); + + // Verify it's stored in cache + let stored = client.message_retry_counts.get(cache_key).await; + assert_eq!(stored, Some(1), "Cache should store count 1"); +} + +#[tokio::test] +async fn test_increment_retry_count_increments_correctly() { + let client = create_test_client_for_retry_with_id("increments").await; + + let cache_key = "test_chat:msg456:sender789"; + + // Simulate multiple retries + let count1 = client + .increment_retry_count(cache_key, RetryReason::NoSession) + .await; + let count2 = client + .increment_retry_count(cache_key, RetryReason::NoSession) + .await; + let count3 = client + .increment_retry_count(cache_key, RetryReason::NoSession) + .await; + + assert_eq!(count1, Some(1), "First retry should be 1"); + assert_eq!(count2, Some(2), "Second retry should be 2"); + assert_eq!(count3, Some(3), "Third retry should be 3"); +} + +#[tokio::test] +async fn test_increment_retry_count_respects_max_retries() { + let client = create_test_client_for_retry_with_id("max_retries").await; + + let cache_key = "test_chat:msg_max:sender_max"; + + // Exhaust all retries (MAX_DECRYPT_RETRIES = 5) + for i in 1..=5 { + let count = client + .increment_retry_count(cache_key, RetryReason::NoSession) + .await; + assert_eq!(count, Some(i), "Retry {} should return {}", i, i); + } + + // 6th attempt should return None (max reached) + let count_after_max = client + .increment_retry_count(cache_key, RetryReason::NoSession) + .await; + assert_eq!( + count_after_max, None, + "After max retries, should return None" + ); + + // Verify cache still has max value + let stored = client.message_retry_counts.get(cache_key).await; + assert_eq!(stored, Some(5), "Cache should retain max count"); +} + +#[tokio::test] +async fn test_retry_count_different_messages_are_independent() { + let client = create_test_client_for_retry_with_id("independent").await; + + let key1 = "chat1:msg1:sender1"; + let key2 = "chat1:msg2:sender1"; // Same chat and sender, different message + let key3 = "chat2:msg1:sender2"; // Different chat and sender + + // Increment each independently + let _ = client + .increment_retry_count(key1, RetryReason::NoSession) + .await; + let _ = client + .increment_retry_count(key1, RetryReason::NoSession) + .await; + let _ = client + .increment_retry_count(key1, RetryReason::NoSession) + .await; // key1 = 3 + + let _ = client + .increment_retry_count(key2, RetryReason::NoSession) + .await; // key2 = 1 + + let _ = client + .increment_retry_count(key3, RetryReason::NoSession) + .await; + let _ = client + .increment_retry_count(key3, RetryReason::NoSession) + .await; // key3 = 2 + + // Verify each has independent counts + assert_eq!(client.message_retry_counts.get(key1).await, Some(3)); + assert_eq!(client.message_retry_counts.get(key2).await, Some(1)); + assert_eq!(client.message_retry_counts.get(key3).await, Some(2)); +} + +#[tokio::test] +async fn test_retry_cache_key_format() { + // Verify the cache key format is consistent + let info = create_test_message_info( + "120363021033254949@g.us", + "3EB0ABCD1234", + "5511999998888@s.whatsapp.net", + ); + + let expected_key = format!("{}:{}:{}", info.source.chat, info.id, info.source.sender); + assert_eq!( + expected_key, + "120363021033254949@g.us:3EB0ABCD1234:5511999998888@s.whatsapp.net" + ); + + // Verify key uniqueness for different senders in same group + let info2 = create_test_message_info( + "120363021033254949@g.us", + "3EB0ABCD1234", // Same message ID + "5511888887777@s.whatsapp.net", // Different sender + ); + + let key2 = format!("{}:{}:{}", info2.source.chat, info2.id, info2.source.sender); + assert_ne!( + expected_key, key2, + "Different senders should have different keys" + ); +} + +/// Test concurrent retry increments are properly serialized. +/// +/// The increment operation uses get+insert which is not fully atomic, +/// but is sufficient since message retry processing is serialized per key +/// by the per-chat lock. At most 5 increments should succeed. +#[tokio::test] +async fn test_concurrent_retry_increments() { + use tokio::task::JoinSet; + + let client = create_test_client_for_retry_with_id("concurrent").await; + let cache_key = "concurrent_test:msg:sender"; + + // Spawn 10 concurrent increment tasks + let mut tasks = JoinSet::new(); + for _ in 0..10 { + let client_clone = client.clone(); + let key = cache_key.to_string(); + tasks.spawn(async move { + client_clone + .increment_retry_count(&key, RetryReason::NoSession) + .await + }); + } + + // Collect all results + let mut results = Vec::new(); + while let Some(result) = tasks.join_next().await { + if let Ok(count) = result { + results.push(count); + } + } + + // With atomic operations, exactly 5 should succeed and 5 should fail + let valid_counts: Vec<_> = results.iter().filter(|r| r.is_some()).collect(); + let none_counts: Vec<_> = results.iter().filter(|r| r.is_none()).collect(); + + assert_eq!( + valid_counts.len(), + 5, + "Exactly 5 increments should succeed with atomic operations" + ); + assert_eq!( + none_counts.len(), + 5, + "Exactly 5 should return None (after max is reached)" + ); + + // Verify the successful increments returned values 1-5 + let mut values: Vec<u8> = valid_counts.iter().filter_map(|r| **r).collect(); + values.sort(); + assert_eq!( + values, + vec![1, 2, 3, 4, 5], + "Successful increments should return 1, 2, 3, 4, 5" + ); + + // Final count should be 5 (max) + let final_count = client.message_retry_counts.get(cache_key).await; + assert_eq!(final_count, Some(5), "Final count should be capped at 5"); +} + +#[tokio::test] +async fn test_high_retry_count_threshold() { + // Verify HIGH_RETRY_COUNT_THRESHOLD is set correctly + assert_eq!( + HIGH_RETRY_COUNT_THRESHOLD, 3, + "High retry threshold should be 3" + ); + assert_eq!(MAX_DECRYPT_RETRIES, 5, "Max retries should be 5"); + // Compile-time assertion that threshold < max (avoids clippy warning) + const _: () = assert!(HIGH_RETRY_COUNT_THRESHOLD < MAX_DECRYPT_RETRIES); +} + +#[tokio::test] +async fn test_message_info_creation_for_groups() { + let info = create_test_message_info( + "120363021033254949@g.us", + "MSG123", + "5511999998888@s.whatsapp.net", + ); + + assert!( + info.source.is_group, + "Group JID should be detected as group" + ); + assert!( + !info.source.is_from_me, + "Test messages default to not from me" + ); + assert_eq!(info.id, "MSG123"); +} + +#[tokio::test] +async fn test_message_info_creation_for_dm() { + let info = create_test_message_info( + "5511999998888@s.whatsapp.net", + "DM456", + "5511999998888@s.whatsapp.net", + ); + + assert!( + !info.source.is_group, + "DM JID should not be detected as group" + ); + assert_eq!(info.id, "DM456"); +} + +#[tokio::test] +async fn test_retry_count_cache_expiration() { + // Note: This test verifies cache configuration, not actual TTL (which would be slow) + let client = create_test_client_for_retry_with_id("expiration").await; + + // The cache should have a TTL of 5 minutes (300 seconds) as configured in client.rs + // We can verify entries are being stored and the cache is functional + let cache_key = "expiry_test:msg:sender"; + + let count = client + .increment_retry_count(cache_key, RetryReason::NoSession) + .await; + assert_eq!(count, Some(1)); + + // Entry should still exist immediately after + let stored = client.message_retry_counts.get(cache_key).await; + assert!( + stored.is_some(), + "Entry should exist immediately after insert" + ); +} + +#[tokio::test] +async fn test_spawn_retry_receipt_basic_flow() { + // This is an integration test that verifies spawn_retry_receipt + // doesn't panic and updates the retry count correctly + + let client = create_test_client_for_retry_with_id("spawn_basic").await; + let info = create_test_message_info( + "120363021033254949@g.us", + "SPAWN_TEST_MSG", + "5511999998888@s.whatsapp.net", + ); + + let cache_key = format!("{}:{}:{}", info.source.chat, info.id, info.source.sender); + + // Verify count starts at 0 + assert!( + client.message_retry_counts.get(&cache_key).await.is_none(), + "Cache should be empty initially" + ); + + // Call spawn_retry_receipt (this spawns a task, so we need to wait) + let info = Arc::new(info); + client.spawn_retry_receipt(&info, RetryReason::UnknownError); + + // Give the spawned task time to execute + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + // Verify count was incremented (the actual send will fail due to no connection, but count should update) + let stored = client.message_retry_counts.get(&cache_key).await; + assert_eq!(stored, Some(1), "Retry count should be 1 after spawn"); +} + +#[tokio::test] +async fn test_spawn_retry_receipt_respects_max_retries() { + let client = create_test_client_for_retry_with_id("spawn_max").await; + let info = create_test_message_info( + "120363021033254949@g.us", + "MAX_RETRY_TEST", + "5511999998888@s.whatsapp.net", + ); + + let cache_key = format!("{}:{}:{}", info.source.chat, info.id, info.source.sender); + + // Pre-fill cache to max retries + client + .message_retry_counts + .insert(cache_key.clone(), MAX_DECRYPT_RETRIES) + .await; + + // Verify count is at max + assert_eq!( + client.message_retry_counts.get(&cache_key).await, + Some(MAX_DECRYPT_RETRIES) + ); + + // Call spawn_retry_receipt - should NOT increment (already at max) + let info = Arc::new(info); + client.spawn_retry_receipt(&info, RetryReason::UnknownError); + + // Give the spawned task time to execute + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + // Count should still be at max (not incremented) + let stored = client.message_retry_counts.get(&cache_key).await; + assert_eq!( + stored, + Some(MAX_DECRYPT_RETRIES), + "Count should remain at max" + ); +} + +#[tokio::test] +async fn test_pdo_cache_key_format_matches() { + // PDO uses "{chat}:{msg_id}" format + // Retry uses "{chat}:{msg_id}:{sender}" format + // They are intentionally different to track independently + + let info = create_test_message_info( + "120363021033254949@g.us", + "PDO_KEY_TEST", + "5511999998888@s.whatsapp.net", + ); + + let retry_key = format!("{}:{}:{}", info.source.chat, info.id, info.source.sender); + let pdo_key = format!("{}:{}", info.source.chat, info.id); + + assert_ne!(retry_key, pdo_key, "PDO and retry keys should be different"); + assert!( + retry_key.starts_with(&pdo_key), + "Retry key should start with PDO key pattern" + ); +} + +#[tokio::test] +async fn test_multiple_senders_same_message_id_tracked_separately() { + // In a group, multiple senders could theoretically have the same message ID + // (unlikely but the system should handle it) + + let client = create_test_client_for_retry_with_id("multi_sender").await; + + let group = "120363021033254949@g.us"; + let msg_id = "SAME_MSG_ID"; + let sender1 = "5511111111111@s.whatsapp.net"; + let sender2 = "5522222222222@s.whatsapp.net"; + + let key1 = format!("{}:{}:{}", group, msg_id, sender1); + let key2 = format!("{}:{}:{}", group, msg_id, sender2); + + // Increment for sender1 multiple times + client + .increment_retry_count(&key1, RetryReason::NoSession) + .await; + client + .increment_retry_count(&key1, RetryReason::NoSession) + .await; + client + .increment_retry_count(&key1, RetryReason::NoSession) + .await; + + // Increment for sender2 once + client + .increment_retry_count(&key2, RetryReason::NoSession) + .await; + + // Verify independent tracking + assert_eq!( + client.message_retry_counts.get(&key1).await, + Some(3), + "Sender1 should have 3 retries" + ); + assert_eq!( + client.message_retry_counts.get(&key2).await, + Some(1), + "Sender2 should have 1 retry" + ); +} + +/// Test: Verify JID type detection for status broadcasts, broadcast lists, groups, and users. +#[test] +fn test_status_broadcast_jid_detection() { + use wacore_binary::{Jid, JidExt}; + + let status_jid: Jid = "status@broadcast".parse().expect("status JID should parse"); + assert!(status_jid.is_status_broadcast()); + + let broadcast_list: Jid = "123456789@broadcast" + .parse() + .expect("broadcast JID should parse"); + assert!(!broadcast_list.is_status_broadcast()); + assert!(broadcast_list.is_broadcast_list()); + + let group_jid: Jid = "120363021033254949@g.us" + .parse() + .expect("group JID should parse"); + assert!(!group_jid.is_status_broadcast()); + + let user_jid: Jid = "15551234567@s.whatsapp.net" + .parse() + .expect("user JID should parse"); + assert!(!user_jid.is_status_broadcast()); +} + +/// Test: Verify should_process_skmsg logic matches WA Web's canDecryptNext pattern. +/// +/// WA Web applies canDecryptNext uniformly: if pkmsg fails with a retriable error, +/// skmsg is skipped regardless of chat type (group, status, 1:1). No exception for +/// status broadcasts — the retry receipt for the pkmsg will cause the sender to +/// resend the entire message including SKDM. +#[test] +fn test_should_process_skmsg_logic_matches_wa_web() { + // Test cases: (chat_jid, session_empty, session_success, session_dupe, session_failed, expected) + let test_cases = [ + // Status broadcast: same rules as all other chats (WA Web: canDecryptNext is uniform) + ("status@broadcast", false, false, false, false, false), // Fail: session failed → skip skmsg + ("status@broadcast", false, false, true, false, true), // OK: duplicate + ("status@broadcast", false, true, false, false, true), // OK: success + ("status@broadcast", false, true, false, true, false), // Fail: mixed success + failure + ("status@broadcast", true, false, false, false, true), // OK: no session msgs + // Regular group + ("120363021033254949@g.us", false, false, false, false, false), + ("120363021033254949@g.us", false, false, true, false, true), + ("120363021033254949@g.us", false, true, false, false, true), + ("120363021033254949@g.us", false, true, false, true, false), + ("120363021033254949@g.us", true, false, false, false, true), + // 1:1 chat + ( + "15551234567@s.whatsapp.net", + false, + false, + false, + false, + false, + ), + ( + "15551234567@s.whatsapp.net", + true, + false, + false, + false, + true, + ), + ]; + + for (jid_str, session_empty, session_success, session_dupe, session_failed, expected) in + test_cases + { + let should_process_skmsg = should_process_skmsg_after_session( + session_empty, + SessionBatchOutcome { + decrypted: session_success, + duplicate: session_dupe, + had_failure: session_failed, + ..Default::default() + }, + ); + + assert_eq!( + should_process_skmsg, + expected, + "For chat {} with session_empty={}, session_success={}, session_dupe={}, session_failed={}: \ + expected should_process_skmsg={}, got {}", + jid_str, + session_empty, + session_success, + session_dupe, + session_failed, + expected, + should_process_skmsg + ); + } +} + +#[test] +fn skdm_only_fallback_ack_decision_requires_clean_session_batch() { + let clean_skdm = SessionBatchOutcome { + decrypted: true, + skdm_only: true, + ..Default::default() + }; + assert!( + should_ack_skdm_only_session_fallback(clean_skdm, true), + "a clean SKDM-only session batch needs the fallback ack" + ); + + let cases = [ + ( + SessionBatchOutcome { + dispatched: true, + ..clean_skdm + }, + true, + "content dispatch already acked", + ), + ( + SessionBatchOutcome { + had_failure: true, + ..clean_skdm + }, + true, + "local session failure must block positive ack", + ), + ( + SessionBatchOutcome { + plaintext_failed: true, + had_failure: true, + ..clean_skdm + }, + true, + "plaintext handler failure is not SKDM-only success", + ), + ( + SessionBatchOutcome { + undecryptable: true, + had_failure: true, + ..clean_skdm + }, + true, + "failure event must not be paired with positive ack", + ), + ( + SessionBatchOutcome { + decrypted: false, + ..clean_skdm + }, + true, + "fallback only applies after Signal decrypt success", + ), + ( + SessionBatchOutcome { + skdm_only: false, + ..clean_skdm + }, + true, + "regular content must ack via dispatch", + ), + ( + SessionBatchOutcome { + duplicate: true, + decrypted: false, + skdm_only: false, + ..Default::default() + }, + true, + "duplicates use the duplicate branch", + ), + (clean_skdm, false, "msmsg work must own its response"), + ]; + + for (outcome, bot_payloads_empty, reason) in cases { + assert!( + !should_ack_skdm_only_session_fallback(outcome, bot_payloads_empty), + "{reason}: {outcome:?}" + ); + } +} + +/// Test: parse_message_info returns error when message "id" attribute is missing +/// +/// Missing message IDs would cause silent collisions in caches/keys, so this +/// must be a hard error rather than defaulting to an empty string. +#[tokio::test] +async fn test_parse_message_info_missing_id_returns_error() { + let backend = Arc::new( + SqliteStore::new("file:memdb_missing_id_test?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("test backend should initialize"), + ); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let node = NodeBuilder::new("message") + .attr("from", "15551234567@s.whatsapp.net") + .attr("t", "1759295366") + .attr("type", "text") + .build(); + + let result = client.parse_message_info(&node.as_node_ref()).await; + + assert!( + result.is_err(), + "parse_message_info should fail when 'id' is missing" + ); + let err_msg = result.unwrap_err().to_string(); + assert!( + err_msg.contains("id"), + "Error message should mention missing 'id' attribute: {}", + err_msg + ); +} +#[tokio::test] +async fn test_no_sender_key_sends_immediate_retry() { + // Verify that when skmsg decryption fails with NoSenderKeyState, + // a retry receipt is sent immediately (no delay, no re-queue). + // This matches WA Web behavior where NoSenderKey → SignalRetryable → RETRY. + let _ = env_logger::builder().is_test(true).try_init(); + + use crate::store::SqliteStore; + use crate::store::persistence_manager::PersistenceManager; + use wacore_binary::NodeContent; + use wacore_binary::builder::NodeBuilder; + + let backend = Arc::new( + SqliteStore::new("file:memdb_retry_immediate?mode=memory&cache=shared") + .await + .expect("Failed to create test backend"), + ); + let pm = Arc::new( + PersistenceManager::new(backend.clone()) + .await + .expect("test backend should initialize"), + ); + let (client, _rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm.clone(), + mock_transport(), + mock_http_client(), + None, + ) + .await; + + let group_jid: Jid = "120363021033254949@g.us".parse().unwrap(); + let sender_jid: Jid = "1234567890:1@s.whatsapp.net".parse().unwrap(); + let msg_id = "TEST_IMMEDIATE_RETRY"; + + // Pseudo-valid SenderKeyMessage: Version 3 + Protobuf + Fake Sig (64 bytes) + let mut content = vec![0x33, 0x08, 0x01, 0x10, 0x01, 0x1A, 0x00]; + content.extend(vec![0u8; 64]); + + let node = NodeBuilder::new("message") + .attr("id", msg_id) + .attr("from", group_jid.clone()) + .attr("participant", sender_jid.clone()) + .attr("type", "text") + .children(vec![{ + let mut n = NodeBuilder::new("enc") + .attr("type", "skmsg") + .attr("v", "2") + .build(); + n.content = Some(NodeContent::Bytes(content)); + n + }]) + .build(); + + client + .clone() + .handle_incoming_message(node_to_arc(node)) + .await; + + // spawn_retry_receipt runs in a spawned task, wait for it + let retry_key = client + .make_retry_cache_key(&group_jid, msg_id, &sender_jid) + .await; + for _ in 0..20 { + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + if client.message_retry_counts.get(&retry_key).await.is_some() { + break; + } + } + assert_eq!( + client.message_retry_counts.get(&retry_key).await, + Some(1), + "NoSenderKeyState should immediately trigger retry receipt (count=1)" + ); +} + +#[test] +fn test_is_sender_key_distribution_only() { + let skdm = wa::message::SenderKeyDistributionMessage { + group_id: Some("group".into()), + axolotl_sender_key_distribution_message: Some(vec![1, 2, 3]), + }; + + // Empty message → false (no SKDM) + assert!(!is_sender_key_distribution_only(&mut wa::Message::default())); + + // SKDM only → true + assert!(is_sender_key_distribution_only(&mut wa::Message { + sender_key_distribution_message: Some(skdm.clone()), + ..Default::default() + })); + + // SKDM + message_context_info → still true (context_info is metadata) + assert!(is_sender_key_distribution_only(&mut wa::Message { + sender_key_distribution_message: Some(skdm.clone()), + message_context_info: Some(wa::MessageContextInfo::default()), + ..Default::default() + })); + + // SKDM + sticker → false (has user content) + assert!(!is_sender_key_distribution_only(&mut wa::Message { + sender_key_distribution_message: Some(skdm.clone()), + sticker_message: Some(Box::new(wa::message::StickerMessage::default())), + ..Default::default() + })); + + // SKDM + text → false (has user content) + assert!(!is_sender_key_distribution_only(&mut wa::Message { + sender_key_distribution_message: Some(skdm.clone()), + conversation: Some("hello".into()), + ..Default::default() + })); + + // protocol_message only (no SKDM) → false + assert!(!is_sender_key_distribution_only(&mut wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage::default())), + ..Default::default() + })); +} + +#[test] +fn skdm_only_detection_restores_carrier_fields() { + // The slow path takes the carrier fields out to compare the rest against + // default; it must restore them so callers still see the original message. + let mut msg = wa::Message { + sender_key_distribution_message: Some(wa::message::SenderKeyDistributionMessage { + group_id: Some("group".into()), + axolotl_sender_key_distribution_message: Some(vec![1, 2, 3]), + }), + fast_ratchet_key_sender_key_distribution_message: Some( + wa::message::SenderKeyDistributionMessage { + group_id: Some("group".into()), + axolotl_sender_key_distribution_message: Some(vec![4, 5, 6]), + }, + ), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![9, 8, 7]), + ..Default::default() + }), + ..Default::default() + }; + + assert!(is_sender_key_distribution_only(&mut msg)); + + // Pin the exact payloads of all three taken/restored carrier fields, not + // just presence: a buggy restore that put back a fresh default (losing the + // original contents) must fail here. + assert_eq!( + msg.sender_key_distribution_message + .as_ref() + .and_then(|s| s.axolotl_sender_key_distribution_message.as_deref()), + Some([1, 2, 3].as_slice()), + "sender_key_distribution_message payload must be restored unchanged" + ); + assert_eq!( + msg.fast_ratchet_key_sender_key_distribution_message + .as_ref() + .and_then(|s| s.axolotl_sender_key_distribution_message.as_deref()), + Some([4, 5, 6].as_slice()), + "fast_ratchet carrier payload must be restored unchanged" + ); + assert_eq!( + msg.message_context_info + .as_ref() + .and_then(|c| c.message_secret.as_deref()), + Some([9, 8, 7].as_slice()), + "message_context_info payload must be restored unchanged" + ); +} + +/// Test: unwrap_device_sent extracts a reaction from a DeviceSentMessage wrapper. +#[test] +fn test_unwrap_device_sent_extracts_reaction() { + let wrapped = wa::Message { + device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { + destination_jid: Some("5511999999999@s.whatsapp.net".to_string()), + message: Some(Box::new(wa::Message { + reaction_message: Some(wa::message::ReactionMessage { + text: Some("\u{2764}".to_string()), + ..Default::default() + }), + ..Default::default() + })), + phash: None, + })), + ..Default::default() + }; + + let mut unwrapped = unwrap_device_sent(wrapped); + assert!( + unwrapped.device_sent_message.is_none(), + "DSM wrapper should be removed" + ); + assert_eq!( + unwrapped + .reaction_message + .as_ref() + .and_then(|r| r.text.as_deref()), + Some("\u{2764}"), + "reaction should be accessible after unwrapping" + ); + assert!( + !is_sender_key_distribution_only(&mut unwrapped), + "unwrapped reaction should not be filtered as SKDM-only" + ); +} + +/// Test: unwrap_device_sent preserves the wrapper when inner message is None. +#[test] +fn test_unwrap_device_sent_preserves_empty_wrapper() { + let wrapped = wa::Message { + device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { + destination_jid: Some("5511999999999@s.whatsapp.net".to_string()), + message: None, + phash: None, + })), + ..Default::default() + }; + + let result = unwrap_device_sent(wrapped); + assert!( + result.device_sent_message.is_some(), + "empty DSM wrapper should be preserved" + ); +} + +/// Test: unwrap_device_sent passes through a plain message unchanged. +#[test] +fn test_unwrap_device_sent_passthrough() { + let msg = wa::Message { + conversation: Some("hello".to_string()), + ..Default::default() + }; + + let result = unwrap_device_sent(msg); + assert_eq!(result.conversation.as_deref(), Some("hello")); +} + +/// Test: unwrap_device_sent merges messageContextInfo from outer and inner, +/// matching WAWebDeviceSentMessageProtoUtils.unwrapDeviceSentMessage. +#[test] +fn test_unwrap_device_sent_merges_context_info() { + let wrapped = wa::Message { + // Outer message_context_info (from the DSM envelope) + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![10, 20, 30]), + limit_sharing_v2: Some(wa::LimitSharing::default()), + ..Default::default() + }), + device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { + destination_jid: Some("5511999999999@s.whatsapp.net".to_string()), + message: Some(Box::new(wa::Message { + conversation: Some("hello".to_string()), + // Inner has its own message_secret but no limit_sharing_v2 + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![1, 2, 3]), + ..Default::default() + }), + ..Default::default() + })), + phash: None, + })), + ..Default::default() + }; + + let result = unwrap_device_sent(wrapped); + let ctx = result.message_context_info.as_ref().unwrap(); + + assert_eq!( + ctx.message_secret, + Some(vec![1, 2, 3]), + "inner message_secret should be preferred" + ); + assert!( + ctx.limit_sharing_v2.is_some(), + "limit_sharing_v2 should come from outer (always)" + ); +} + +/// Test: unwrap_device_sent falls back to outer message_secret when inner has none. +#[test] +fn test_unwrap_device_sent_secret_fallback() { + let wrapped = wa::Message { + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![10, 20, 30]), + ..Default::default() + }), + device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { + destination_jid: Some("5511999999999@s.whatsapp.net".to_string()), + message: Some(Box::new(wa::Message { + conversation: Some("hello".to_string()), + // Inner has no message_context_info at all + ..Default::default() + })), + phash: None, + })), + ..Default::default() + }; + + let result = unwrap_device_sent(wrapped); + let ctx = result.message_context_info.as_ref().unwrap(); + assert_eq!( + ctx.message_secret, + Some(vec![10, 20, 30]), + "should fall back to outer message_secret" + ); +} + +#[tokio::test] +async fn test_parse_edit_attribute_sender_revoke() { + let client = create_test_client_for_retry_with_id("edit_sender_revoke").await; + + let node = NodeBuilder::new("message") + .attr("from", "status@broadcast") + .attr("id", "TEST123") + .attr("participant", "5551234567@lid") + .attr("t", "1772895198") + .attr("type", "text") + .attr("edit", "7") + .build(); + + let info = client + .parse_message_info(&node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + assert_eq!( + info.edit, + EditAttribute::SenderRevoke, + "edit='7' should parse as SenderRevoke" + ); +} + +#[tokio::test] +async fn test_parse_edit_attribute_admin_revoke() { + let client = create_test_client_for_retry_with_id("edit_admin_revoke").await; + + let node = NodeBuilder::new("message") + .attr("from", "120363999999999999@g.us") + .attr("id", "TEST456") + .attr("participant", "5551234567@lid") + .attr("t", "1772895198") + .attr("type", "text") + .attr("edit", "8") + .build(); + + let info = client + .parse_message_info(&node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + assert_eq!( + info.edit, + EditAttribute::AdminRevoke, + "edit='8' should parse as AdminRevoke" + ); +} + +#[tokio::test] +async fn test_parse_edit_attribute_message_edit() { + let client = create_test_client_for_retry_with_id("edit_message_edit").await; + + let node = NodeBuilder::new("message") + .attr("from", "5551234567@s.whatsapp.net") + .attr("id", "TEST789") + .attr("t", "1772895198") + .attr("type", "text") + .attr("edit", "1") + .build(); + + let info = client + .parse_message_info(&node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + assert_eq!( + info.edit, + EditAttribute::MessageEdit, + "edit='1' should parse as MessageEdit" + ); +} + +#[tokio::test] +async fn test_parse_edit_attribute_missing() { + let client = create_test_client_for_retry_with_id("edit_missing").await; + + let node = NodeBuilder::new("message") + .attr("from", "5551234567@s.whatsapp.net") + .attr("id", "TESTABC") + .attr("t", "1772895198") + .attr("type", "text") + .build(); + + let info = client + .parse_message_info(&node.as_node_ref()) + .await + .expect("parse_message_info should succeed"); + + assert_eq!( + info.edit, + EditAttribute::Empty, + "missing edit attr should default to Empty" + ); +} + +#[tokio::test] +async fn test_revoked_message_still_retries() { + let client = create_test_client_for_retry_with_id("revoke_retry").await; + + let mut info = create_test_message_info( + "status@broadcast", + "REVOKE_MSG1", + "5551234567@s.whatsapp.net", + ); + info.edit = EditAttribute::SenderRevoke; + + // WA Web retries revoked messages the same as any other — the revoke + // protocol message contains the target ID needed to process the deletion + let info = Arc::new(info); + client.spawn_retry_receipt(&info, RetryReason::NoSession); + + // Wait for the spawned task to execute + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + let cache_key = client + .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) + .await; + assert_eq!( + client.message_retry_counts.get(&cache_key).await, + Some(1), + "revoked message should still have retry count 1 (WA Web retries all messages)" + ); +} + +#[tokio::test] +async fn test_enc_count_preseeds_retry_cache() { + let client = create_test_client_for_retry_with_id("enc_preseed").await; + + let chat_jid: Jid = "5551234567@s.whatsapp.net".parse().unwrap(); + let msg_id = "ENC_COUNT_MSG1"; + + // Pre-seed via the same logic used in handle_incoming_message + let max_sender_retry_count: u8 = 3; + let cache_key = client + .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) + .await; + // Insert only if absent (portable alternative to moka's entry_by_ref().or_insert()) + if client.message_retry_counts.get(&cache_key).await.is_none() { + client + .message_retry_counts + .insert(cache_key.clone(), max_sender_retry_count) + .await; + } + + assert_eq!( + client.message_retry_counts.get(&cache_key).await, + Some(3), + "cache should be pre-seeded with sender retry count" + ); +} + +#[tokio::test] +async fn test_enc_no_count_cache_empty() { + let client = create_test_client_for_retry_with_id("enc_no_count").await; + + let chat_jid: Jid = "5551234567@s.whatsapp.net".parse().unwrap(); + let msg_id = "ENC_NO_COUNT_MSG1"; + + // When max_sender_retry_count is 0, no pre-seeding occurs + let max_sender_retry_count: u8 = 0; + if max_sender_retry_count > 0 { + let cache_key = client + .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) + .await; + if client.message_retry_counts.get(&cache_key).await.is_none() { + client + .message_retry_counts + .insert(cache_key, max_sender_retry_count) + .await; + } + } + + let cache_key = client + .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) + .await; + assert!( + client.message_retry_counts.get(&cache_key).await.is_none(), + "cache should be empty when no count attribute" + ); +} + +#[tokio::test] +async fn test_enc_count_does_not_overwrite_higher() { + let client = create_test_client_for_retry_with_id("enc_no_overwrite").await; + + let chat_jid: Jid = "5551234567@s.whatsapp.net".parse().unwrap(); + let msg_id = "ENC_NOOVERWRITE_MSG1"; + + let cache_key = client + .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) + .await; + + // Pre-insert a higher value + client + .message_retry_counts + .insert(cache_key.clone(), 4) + .await; + + // max(existing, incoming) should NOT overwrite with a lower value + let max_sender_retry_count: u8 = 2; + let existing = client + .message_retry_counts + .get(&cache_key) + .await + .unwrap_or(0); + if max_sender_retry_count > existing { + client + .message_retry_counts + .insert(cache_key.clone(), max_sender_retry_count) + .await; + } + + assert_eq!( + client.message_retry_counts.get(&cache_key).await, + Some(4), + "should not overwrite existing higher value" + ); +} + +#[tokio::test] +async fn test_enc_count_updates_when_sender_higher() { + let client = create_test_client_for_retry_with_id("enc_update_higher").await; + + let chat_jid: Jid = "5551234567@s.whatsapp.net".parse().unwrap(); + let msg_id = "ENC_UPDATE_MSG1"; + + let cache_key = client + .make_retry_cache_key(&chat_jid, msg_id, &chat_jid) + .await; + + // Pre-insert a lower value + client + .message_retry_counts + .insert(cache_key.clone(), 1) + .await; + + // max(existing, incoming) SHOULD update with a higher value + let max_sender_retry_count: u8 = 3; + let existing = client + .message_retry_counts + .get(&cache_key) + .await + .unwrap_or(0); + if max_sender_retry_count > existing { + client + .message_retry_counts + .insert(cache_key.clone(), max_sender_retry_count) + .await; + } + + assert_eq!( + client.message_retry_counts.get(&cache_key).await, + Some(3), + "should update to higher sender count" + ); +} + +/// Shared helper: the OLD semaphore acquire logic that silently dropped tasks +/// on generation mismatch. Used by the bug-demonstration test. +async fn acquire_permit_old_behavior( + semaphore: &std::sync::Mutex<Arc<async_lock::Semaphore>>, + generation: &portable_atomic::AtomicU64, +) -> bool { + use std::sync::atomic::Ordering; + let (snap_gen, snap_sem) = { + let guard = semaphore.lock().unwrap(); + (generation.load(Ordering::SeqCst), guard.clone()) + }; + let _permit = snap_sem.acquire_arc().await; + // OLD: if generation changed, silently return false (message lost) + snap_gen == generation.load(Ordering::SeqCst) +} + +/// Shared helper: the FIXED semaphore acquire logic that re-acquires from the +/// new semaphore on generation mismatch. Mirrors the production code in +/// handle_incoming_message. +async fn acquire_permit_with_reacquire( + semaphore: &std::sync::Mutex<Arc<async_lock::Semaphore>>, + generation: &portable_atomic::AtomicU64, +) { + use std::sync::atomic::Ordering; + loop { + let (snap_gen, snap_sem) = { + let guard = semaphore.lock().unwrap(); + (generation.load(Ordering::SeqCst), guard.clone()) + }; + let permit = snap_sem.acquire_arc().await; + if snap_gen == generation.load(Ordering::SeqCst) { + drop(permit); + break; + } + drop(permit); + } +} + +/// Demonstrates the bug: the OLD code silently dropped tasks when generation changed. +#[tokio::test] +async fn test_old_behavior_drops_tasks_on_generation_swap() { + use portable_atomic::AtomicU64; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let semaphore = Arc::new(std::sync::Mutex::new(Arc::new(async_lock::Semaphore::new( + 1, + )))); + let generation = Arc::new(AtomicU64::new(0)); + let completed = Arc::new(AtomicUsize::new(0)); + let ready = Arc::new(AtomicUsize::new(0)); + + let blocker_sem = semaphore.lock().unwrap().clone(); + let blocker_permit = blocker_sem.acquire_arc().await; + + let num_waiters: usize = 8; + let mut handles = Vec::new(); + + for _ in 0..num_waiters { + let sem = semaphore.clone(); + let gen_counter = generation.clone(); + let done = completed.clone(); + let ready_counter = ready.clone(); + + handles.push(tokio::spawn(async move { + // Signal readiness before blocking on semaphore + ready_counter.fetch_add(1, Ordering::SeqCst); + if acquire_permit_old_behavior(&sem, &gen_counter).await { + done.fetch_add(1, Ordering::SeqCst); + } + })); + } + + // Wait until all waiters have signaled readiness (about to block on semaphore) + while ready.load(Ordering::SeqCst) < num_waiters { + tokio::task::yield_now().await; + } + + // Swap semaphore — triggers the bug + { + let mut guard = semaphore.lock().unwrap(); + *guard = Arc::new(async_lock::Semaphore::new(64)); + generation.fetch_add(1, Ordering::SeqCst); + } + + drop(blocker_permit); + + for handle in handles { + let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await; + assert!(result.is_ok(), "Waiter task timed out"); + result.unwrap().unwrap(); + } + + let done = completed.load(Ordering::SeqCst); + assert!( + done < num_waiters, + "Bug demonstration: expected tasks to be dropped, but all {} completed", + num_waiters + ); +} + +/// Verifies the fix: re-acquire loop ensures NO tasks are dropped on generation swap. +#[tokio::test] +async fn test_semaphore_generation_swap_does_not_drop_tasks() { + use portable_atomic::AtomicU64; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let semaphore = Arc::new(std::sync::Mutex::new(Arc::new(async_lock::Semaphore::new( + 1, + )))); + let generation = Arc::new(AtomicU64::new(0)); + let completed = Arc::new(AtomicUsize::new(0)); + let ready = Arc::new(AtomicUsize::new(0)); + + let blocker_sem = semaphore.lock().unwrap().clone(); + let blocker_permit = blocker_sem.acquire_arc().await; + + let num_waiters: usize = 8; + let mut handles = Vec::new(); + + for _ in 0..num_waiters { + let sem = semaphore.clone(); + let gen_counter = generation.clone(); + let done = completed.clone(); + let ready_counter = ready.clone(); + + handles.push(tokio::spawn(async move { + ready_counter.fetch_add(1, Ordering::SeqCst); + acquire_permit_with_reacquire(&sem, &gen_counter).await; + done.fetch_add(1, Ordering::SeqCst); + })); + } + + // Wait until all waiters have signaled readiness + while ready.load(Ordering::SeqCst) < num_waiters { + tokio::task::yield_now().await; + } + + // Swap semaphore (simulates offline sync completion) + { + let mut guard = semaphore.lock().unwrap(); + *guard = Arc::new(async_lock::Semaphore::new(64)); + generation.fetch_add(1, Ordering::SeqCst); + } + + drop(blocker_permit); + + for handle in handles { + let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await; + assert!( + result.is_ok(), + "Waiter task timed out — likely silently dropped by generation check" + ); + result.unwrap().unwrap(); + } + + assert_eq!( + completed.load(Ordering::SeqCst), + num_waiters, + "All {} waiter tasks should complete, but only {} did. \ + Tasks were silently dropped during semaphore generation swap.", + num_waiters, + completed.load(Ordering::SeqCst) + ); +} + +// Dispatch ordering, per-id dedup, and PDO eligibility for +// UndecryptableMessage. Regressing any of these re-opens data loss bugs +// observed in production. + +use crate::types::events::DecryptFailMode; +use wacore::types::events::{Event, EventHandler}; + +#[derive(Default)] +struct EventRecorder { + events: std::sync::Mutex<Vec<Arc<Event>>>, +} + +impl EventHandler for EventRecorder { + fn handle_event(&self, event: Arc<Event>) { + self.events.lock().unwrap().push(event); + } +} + +impl EventRecorder { + fn undecryptable(&self) -> Vec<Arc<Event>> { + self.events + .lock() + .unwrap() + .iter() + .filter(|e| matches!(&***e, Event::UndecryptableMessage(_))) + .cloned() + .collect() + } + + /// Count of `UndecryptableMessage` events marked as the "stub" + /// variant (`is_unavailable=true`, `UnavailableType::ViewOnce`) — + /// i.e. the branch that routes to PDO instead of falling through to + /// decrypt. + fn view_once_unavailable_count(&self) -> usize { + use crate::types::events::UnavailableType; + self.events + .lock() + .unwrap() + .iter() + .filter(|e| { + matches!( + &***e, + Event::UndecryptableMessage(u) + if u.is_unavailable + && matches!(u.unavailable_type, UnavailableType::ViewOnce) + ) + }) + .count() + } +} + +fn build_unavailable_stanza(sender: &str, msg_id: &str, with_enc: bool) -> Arc<OwnedNodeRef> { + let t = wacore::time::now_secs().to_string(); + let unavailable = NodeBuilder::new("unavailable") + .attr("type", "view_once") + .build(); + let children = if with_enc { + vec![ + unavailable, + NodeBuilder::new("enc") + .attr("type", "msg") + .attr("v", "2") + .bytes(vec![0xDE, 0xAD, 0xBE, 0xEF]) + .build(), + ] + } else { + vec![unavailable] + }; + node_to_arc( + NodeBuilder::new("message") + .attr("from", sender) + .attr("id", msg_id) + .attr("t", &t) + .attr("type", "media") + .children(children) + .build(), + ) +} + +/// Locks the dispatch ordering: consumers must see the event before any +/// retry/PDO side effects, otherwise a late subscriber misses the failure. +#[tokio::test] +async fn test_undecryptable_fires_before_retry_task() { + let client = create_test_client_for_retry_with_id("undec_sync").await; + let recorder = Arc::new(EventRecorder::default()); + client.register_handler(recorder.clone()); + + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "MSG_SYNC_1", + "5511777776666@s.whatsapp.net", + )); + + let cache_key = client + .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) + .await; + + assert!(recorder.undecryptable().is_empty()); + assert!(client.message_retry_counts.get(&cache_key).await.is_none()); + + let _ = client + .handle_decrypt_failure(&info, RetryReason::InvalidKeyId, DecryptFailMode::Show) + .await; + + assert_eq!( + recorder.undecryptable().len(), + 1, + "UndecryptableMessage dispatched inside handle_decrypt_failure", + ); + assert!( + client.message_retry_counts.get(&cache_key).await.is_none(), + "retry task has not progressed yet", + ); + + tokio::time::sleep(tokio::time::Duration::from_millis(150)).await; + assert_eq!( + client.message_retry_counts.get(&cache_key).await, + Some(1), + "retry task runs after the dispatch", + ); +} + +/// Atomic dedup under concurrency: 32 parallel callers for the same id +/// must produce exactly one event. Catches regressions where the dedup +/// would slip back to a non-atomic get-then-insert pair. +#[tokio::test] +async fn test_undecryptable_dedup_is_atomic() { + let client = create_test_client_for_retry_with_id("undec_atomic").await; + let recorder = Arc::new(EventRecorder::default()); + client.register_handler(recorder.clone()); + + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "ATOMIC_MSG_1", + "5511777776666@s.whatsapp.net", + )); + + let mut handles = Vec::with_capacity(32); + for _ in 0..32 { + let c = Arc::clone(&client); + let i = Arc::clone(&info); + handles.push(tokio::spawn(async move { + c.handle_decrypt_failure(&i, RetryReason::InvalidKeyId, DecryptFailMode::Show) + .await; + })); + } + for h in handles { + h.await.unwrap(); + } + + assert_eq!( + recorder.undecryptable().len(), + 1, + "32 concurrent callers must collapse to one UndecryptableMessage", + ); +} + +/// Server resends of the same id must not surface a duplicate event — +/// would otherwise show the user the same failure twice. +#[tokio::test] +async fn test_undecryptable_deduped_across_resends() { + let client = create_test_client_for_retry_with_id("undec_double").await; + let recorder = Arc::new(EventRecorder::default()); + client.register_handler(recorder.clone()); + + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "3AD01881AA95F7D81070", + "85010891714716@lid", + )); + + let _ = client + .handle_decrypt_failure(&info, RetryReason::InvalidKeyId, DecryptFailMode::Show) + .await; + let _ = client + .handle_decrypt_failure(&info, RetryReason::InvalidKeyId, DecryptFailMode::Show) + .await; + + let events = recorder.undecryptable(); + assert_eq!( + events.len(), + 1, + "same message id fires UndecryptableMessage only once", + ); + if let Event::UndecryptableMessage(event) = &*events[0] { + assert_eq!(event.info.id, info.id); + } else { + panic!("event was not UndecryptableMessage"); + } +} + +/// Status posts must flow through PDO — excluding them drops any +/// InvalidPreKeyId status permanently (WA Web recovers them). +#[tokio::test] +async fn test_pdo_armed_for_status_broadcast() { + let client = create_test_client_for_retry_with_id("pdo_status").await; + + let info = Arc::new(create_test_message_info( + "status@broadcast", + "STATUS_MSG_1", + "5511777776666@s.whatsapp.net", + )); + + assert_eq!(info.source.chat.server, wacore_binary::Server::Broadcast); + + client.run_pdo_request(&info).await; + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; +} + +/// Broadcast lists share the same code path; locks the guard for both. +#[tokio::test] +async fn test_pdo_armed_for_any_broadcast_chat() { + let client = create_test_client_for_retry_with_id("pdo_bcast_list").await; + + let info = Arc::new(create_test_message_info( + "12345@broadcast", + "BCAST_LIST_MSG_1", + "5511777776666@s.whatsapp.net", + )); + + assert_eq!(info.source.chat.server, wacore_binary::Server::Broadcast); + + client.run_pdo_request(&info).await; + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; +} + +#[tokio::test] +async fn test_pdo_armed_for_one_on_one() { + let client = create_test_client_for_retry_with_id("pdo_dm").await; + + let info = Arc::new(create_test_message_info( + "85010891714716@lid", + "DM_MSG_1", + "85010891714716@lid", + )); + + assert_ne!(info.source.chat.server, wacore_binary::Server::Broadcast); + + client.run_pdo_request(&info).await; + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; +} + +/// fromMe messages fanned out to a linked device can still fail decrypt +/// on the receiver side; PDO is the only recovery path for them. +#[tokio::test] +async fn test_pdo_armed_for_from_me() { + let client = create_test_client_for_retry_with_id("pdo_from_me").await; + + // When fromMe is true the sender is the user's own JID, not a peer. + let own_jid = "5511999998888@s.whatsapp.net"; + let mut info = create_test_message_info("85010891714716@lid", "FROM_ME_MSG_1", own_jid); + info.source.is_from_me = true; + let info = Arc::new(info); + + client.run_pdo_request(&info).await; + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; +} + +/// Stops offline-sync / reconnect tails from flooding the phone with +/// resend requests for old messages the user likely no longer cares about. +#[tokio::test] +async fn test_pdo_skipped_for_ancient_messages() { + use wacore::types::message::ChatMessageId; + + let client = create_test_client_for_retry_with_id("pdo_age").await; + + let mut info = + create_test_message_info("85010891714716@lid", "ANCIENT_MSG_1", "85010891714716@lid"); + info.timestamp = wacore::time::now_utc() - chrono::Duration::days(30); + let info = Arc::new(info); + + let cache_key = ChatMessageId::new(info.source.chat.clone(), info.id.clone()); + + client.run_pdo_request(&info).await; + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + assert!( + client.pdo_pending_requests.get(&cache_key).await.is_none(), + "messages older than 14 days must not register a PDO entry", + ); +} + +/// Boundary check: age of 14d plus a minute must reject (WA Web uses +/// seconds, not days, so 14d1m is already over the limit). Catches a +/// `num_days()` truncation that would otherwise accept this message. +#[tokio::test] +async fn test_pdo_rejects_just_past_14d_boundary() { + use wacore::types::message::ChatMessageId; + + let client = create_test_client_for_retry_with_id("pdo_boundary").await; + + let mut info = + create_test_message_info("85010891714716@lid", "BOUNDARY_MSG_1", "85010891714716@lid"); + info.timestamp = + wacore::time::now_utc() - chrono::Duration::days(14) - chrono::Duration::minutes(1); + let info = Arc::new(info); + + let cache_key = ChatMessageId::new(info.source.chat.clone(), info.id.clone()); + + client.run_pdo_request(&info).await; + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + assert!( + client.pdo_pending_requests.get(&cache_key).await.is_none(), + "14d+1m must be over the limit, matching WA Web's seconds-based check", + ); +} + +/// Server-trusted companions (Android-class `DeviceProps.PlatformType`) +/// receive `<unavailable>` as a marker alongside `<enc>`. The cipher +/// must still be decrypted — skipping would discard content the server +/// specifically released for this companion. Decrypt eventually fails +/// on the garbage payload, but via the normal decrypt-failure path, +/// not the `ViewOnce` short-circuit. +#[tokio::test] +async fn test_unavailable_with_enc_skips_unavailable_shortcut() { + let client = create_test_client_for_retry_with_id("unavailable_with_enc").await; + let recorder = Arc::new(EventRecorder::default()); + client.register_handler(recorder.clone()); + + let node = build_unavailable_stanza("5511777776666@s.whatsapp.net", "UNAV_WITH_ENC_1", true); + client.clone().handle_incoming_message(node).await; + + assert_eq!( + recorder.view_once_unavailable_count(), + 0, + "<unavailable> alongside <enc> must fall through to decrypt, \ + not emit a ViewOnce UndecryptableMessage", + ); +} + +/// Untrusted companions (web-class `PlatformType`) get the bare stub — +/// `<unavailable>` without `<enc>`. That path must still emit a +/// `ViewOnce` `UndecryptableMessage` so consumers surface the failure +/// while the phone relays via PDO. +#[tokio::test] +async fn test_unavailable_without_enc_dispatches_view_once_event() { + let client = create_test_client_for_retry_with_id("unavailable_stub").await; + let recorder = Arc::new(EventRecorder::default()); + client.register_handler(recorder.clone()); + + let node = build_unavailable_stanza("5511777776666@s.whatsapp.net", "UNAV_STUB_1", false); + client.clone().handle_incoming_message(node).await; + + assert_eq!( + recorder.view_once_unavailable_count(), + 1, + "bare <unavailable> stub must dispatch exactly one ViewOnce UndecryptableMessage", + ); +} + +/// The event struct has no "recovery pending" flag, so consumers cannot +/// wait for a PDO outcome before surfacing failure — adding a field +/// here forces a conscious UX decision. +#[test] +fn test_undecryptable_event_has_no_pending_pdo_hint() { + use crate::types::events::{UnavailableType, UndecryptableMessage}; + + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "SHAPE_MSG", + "5511777776666@s.whatsapp.net", + )); + let event = UndecryptableMessage { + info, + is_unavailable: false, + unavailable_type: UnavailableType::Unknown, + decrypt_fail_mode: DecryptFailMode::Show, + }; + + let _ = ( + &event.info, + &event.is_unavailable, + &event.unavailable_type, + &event.decrypt_fail_mode, + ); +} + +/// Seed `device.pn` so `send_nack` clears its `get_pn()` guard. +async fn seed_test_pn(client: &Arc<Client>) { + use crate::store::commands::DeviceCommand; + client + .persistence_manager + .process_command(DeviceCommand::SetId(Some( + "5511000000001:0@s.whatsapp.net" + .parse() + .expect("test PN should parse"), + ))) + .await; +} + +/// Build a Client wired to a CapturingMockTransport + a noise socket so +/// `send_node` reaches the wire. Returns the transport so the caller can +/// inspect captured frames. +async fn capturing_client( + test_id: &str, +) -> ( + Arc<Client>, + Arc<crate::transport::mock::CapturingMockTransport>, +) { + use crate::socket::NoiseSocket; + use crate::store::SqliteStore; + use crate::store::persistence_manager::PersistenceManager; + use crate::transport::mock::CapturingMockTransportFactory; + use portable_atomic::AtomicU64; + use std::sync::atomic::Ordering; + use wacore::handshake::NoiseCipher; + + static COUNTER: AtomicU64 = AtomicU64::new(0); + let unique_id = COUNTER.fetch_add(1, Ordering::SeqCst); + let db_name = format!( + "file:memdb_capt_{}_{}_{}?mode=memory&cache=shared", + test_id, + unique_id, + std::process::id() + ); + + let backend = Arc::new( + SqliteStore::new(&db_name) + .await + .expect("test backend should initialize"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let factory = CapturingMockTransportFactory::new(); + let transport = factory.transport(); + let (client, _sync_rx) = Client::new( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(factory), + Arc::new(MockHttpClient), + None, + ) + .await; + + let key = [0u8; 32]; + let write_key = NoiseCipher::new(&key).expect("32-byte key"); + let read_key = NoiseCipher::new(&key).expect("32-byte key"); + let noise_socket = NoiseSocket::new( + Arc::new(crate::runtime_impl::TokioRuntime), + transport.clone() as Arc<dyn crate::transport::Transport>, + write_key, + read_key, + ); + // send_node only needs noise_socket Some; is_connected is read by + // other layers but not on this path. + *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); + seed_test_pn(&client).await; + (client, transport) +} + +/// Regression: a malformed pkmsg used to fall through silently. Now +/// it dispatches the consumer event AND emits a nack on the wire so +/// the server stops retransmitting. +#[tokio::test] +async fn pkmsg_parse_error_dispatches_parsing_error_nack() { + use crate::types::events::DecryptFailMode; + use wacore::message_processing::EncType; + + let (client, transport) = capturing_client("pkmsg_parse_nack").await; + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "REGRESSION_PKMSG_PARSE", + "5511777776666@s.whatsapp.net", + )); + let sender_jid: Jid = info.source.sender.clone(); + + // 1-byte ciphertext is a guaranteed parse failure. + let bad_payload = EncPayload { + ciphertext: bytes::Bytes::from_static(&[0xFF]), + enc_type: EncType::PreKeyMessage, + padding_version: 2, + }; + + let outcome = client + .process_session_enc_batch(&[bad_payload], &info, &sender_jid, DecryptFailMode::Show) + .await; + + assert!(!outcome.decrypted); + assert!(!outcome.duplicate); + assert!(outcome.undecryptable); + assert!(outcome.had_failure); + + // spawn_nack is detached; give it a tick to flush through the + // noise_socket sender_task to our CapturingMockTransport. + for _ in 0..40 { + if !transport.sent().is_empty() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let sent = transport.sent(); + assert!( + !sent.is_empty(), + "spawn_nack must produce at least one outbound frame on the wire" + ); +} + +#[tokio::test] +async fn signal_message_parse_error_dispatches_parsing_error_nack() { + use crate::types::events::DecryptFailMode; + use wacore::message_processing::EncType; + + let (client, transport) = capturing_client("sig_parse_nack").await; + let info = Arc::new(create_test_message_info( + "5511999998888@s.whatsapp.net", + "REGRESSION_SIG_PARSE", + "5511777776666@s.whatsapp.net", + )); + let sender_jid: Jid = info.source.sender.clone(); + + let bad_payload = EncPayload { + ciphertext: bytes::Bytes::from_static(&[0xFF]), + enc_type: EncType::Message, + padding_version: 2, + }; + + let outcome = client + .process_session_enc_batch(&[bad_payload], &info, &sender_jid, DecryptFailMode::Show) + .await; + + assert!(outcome.undecryptable); + assert!(outcome.had_failure); + + for _ in 0..40 { + if !transport.sent().is_empty() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert!( + !transport.sent().is_empty(), + "spawn_nack must produce at least one outbound frame on the wire" + ); +} + +#[test] +fn test_decrypt_fail_log_level_gated_on_hide() { + use crate::types::events::DecryptFailMode; + assert_eq!( + decrypt_fail_log_level(DecryptFailMode::Hide), + log::Level::Debug + ); + assert_eq!( + decrypt_fail_log_level(DecryptFailMode::Show), + log::Level::Warn + ); +} + +/// Decrypt one captured noise frame (zero-key, counter-based, empty AAD) to +/// its marshalled node bytes; strips the 3-byte frame header. +fn decode_frame(index: usize, frame: &[u8]) -> Option<Vec<u8>> { + use wacore::handshake::NoiseCipher; + if frame.len() <= 3 { + return None; + } + let cipher = NoiseCipher::new(&[0u8; 32]).expect("32-byte key"); + let mut buf = frame[3..].to_vec(); + cipher + .decrypt_in_place_with_counter(index as u32, &mut buf) + .ok()?; + (!buf.is_empty()).then_some(buf) +} + +/// First `<ack class="message">` on the wire as `(to, recipient)`. +fn find_message_ack(frames: &[bytes::Bytes]) -> Option<(String, Option<String>)> { + for (i, frame) in frames.iter().enumerate() { + let Some(buf) = decode_frame(i, frame) else { + continue; + }; + let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { + continue; + }; + if node.tag.as_ref() == "ack" + && node + .get_attr("class") + .is_some_and(|v| v.as_str() == "message") + && node.get_attr("error").is_none() + && let Some(to) = node.get_attr("to") + { + let recipient = node.get_attr("recipient").map(|v| v.as_str().to_string()); + return Some((to.as_str().to_string(), recipient)); + } + } + None +} + +/// First `<receipt>` on the wire for `id` as `(to, type, recipient)`. +fn find_receipt( + frames: &[bytes::Bytes], + id: &str, +) -> Option<(String, Option<String>, Option<String>)> { + for (i, frame) in frames.iter().enumerate() { + let Some(buf) = decode_frame(i, frame) else { + continue; + }; + let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { + continue; + }; + if node.tag.as_ref() == "receipt" + && node.get_attr("id").is_some_and(|v| v.as_str() == id) + && let Some(to) = node.get_attr("to") + { + let typ = node.get_attr("type").map(|v| v.as_str().to_string()); + let recipient = node.get_attr("recipient").map(|v| v.as_str().to_string()); + return Some((to.as_str().to_string(), typ, recipient)); + } + } + None +} + +#[derive(Debug)] +struct SentReceipt { + to: String, + typ: Option<String>, + recipient: Option<String>, + participant: Option<String>, + context: Option<String>, +} + +fn find_receipt_details(frames: &[bytes::Bytes], id: &str) -> Option<SentReceipt> { + for (i, frame) in frames.iter().enumerate() { + let Some(buf) = decode_frame(i, frame) else { + continue; + }; + let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { + continue; + }; + if node.tag.as_ref() == "receipt" + && node.get_attr("id").is_some_and(|v| v.as_str() == id) + && let Some(to) = node.get_attr("to") + { + return Some(SentReceipt { + to: to.as_str().to_string(), + typ: node.get_attr("type").map(|v| v.as_str().to_string()), + recipient: node.get_attr("recipient").map(|v| v.as_str().to_string()), + participant: node.get_attr("participant").map(|v| v.as_str().to_string()), + context: node.get_attr("context").map(|v| v.as_str().to_string()), + }); + } + } + None +} + +#[derive(Debug)] +struct SentMessageAck { + to: String, + participant: Option<String>, + recipient: Option<String>, +} + +fn find_message_ack_for(frames: &[bytes::Bytes], id: &str) -> Option<SentMessageAck> { + for (i, frame) in frames.iter().enumerate() { + let Some(buf) = decode_frame(i, frame) else { + continue; + }; + let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { + continue; + }; + if node.tag.as_ref() == "ack" + && node + .get_attr("class") + .is_some_and(|v| v.as_str() == "message") + && node.get_attr("error").is_none() + && node.get_attr("id").is_some_and(|v| v.as_str() == id) + && let Some(to) = node.get_attr("to") + { + return Some(SentMessageAck { + to: to.as_str().to_string(), + participant: node.get_attr("participant").map(|v| v.as_str().to_string()), + recipient: node.get_attr("recipient").map(|v| v.as_str().to_string()), + }); + } + } + None +} + +/// Count delivery `<receipt>` (anything but type="retry") on the wire for `id`. +fn delivery_receipts_for(frames: &[bytes::Bytes], id: &str) -> usize { + let mut count = 0; + for (i, frame) in frames.iter().enumerate() { + let Some(buf) = decode_frame(i, frame) else { + continue; + }; + let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { + continue; + }; + if node.tag.as_ref() == "receipt" + && node.get_attr("id").is_some_and(|v| v.as_str() == id) + && node + .get_attr("type") + .as_ref() + .map(|v| v.as_str()) + .as_deref() + != Some("retry") + { + count += 1; + } + } + count +} + +fn message_acks_for(frames: &[bytes::Bytes], id: &str) -> usize { + let mut count = 0; + for (i, frame) in frames.iter().enumerate() { + let Some(buf) = decode_frame(i, frame) else { + continue; + }; + let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { + continue; + }; + if node.tag.as_ref() == "ack" + && node + .get_attr("class") + .is_some_and(|v| v.as_str() == "message") + && node.get_attr("error").is_none() + && node.get_attr("id").is_some_and(|v| v.as_str() == id) + { + count += 1; + } + } + count +} + +fn sender_receipts_for(frames: &[bytes::Bytes], id: &str) -> usize { + let mut count = 0; + for (i, frame) in frames.iter().enumerate() { + let Some(buf) = decode_frame(i, frame) else { + continue; + }; + let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { + continue; + }; + if node.tag.as_ref() == "receipt" + && node.get_attr("id").is_some_and(|v| v.as_str() == id) + && node + .get_attr("type") + .is_some_and(|v| v.as_str() == "sender") + { + count += 1; + } + } + count +} + +fn confirmations_for(frames: &[bytes::Bytes], id: &str) -> usize { + delivery_receipts_for(frames, id) + message_acks_for(frames, id) +} + +async fn wait_for_confirmations( + transport: &crate::transport::mock::CapturingMockTransport, + id: &str, + expected: usize, +) -> usize { + let mut count = 0; + for _ in 0..80 { + count = confirmations_for(&transport.sent(), id); + if count >= expected { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + count +} + +async fn assert_exactly_one_confirmation( + transport: &crate::transport::mock::CapturingMockTransport, + id: &str, +) { + let count = wait_for_confirmations(transport, id, 1).await; + assert_eq!(count, 1, "message {id} must be confirmed exactly once"); + for _ in 0..5 { + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + assert_eq!( + confirmations_for(&transport.sent(), id), + 1, + "message {id} must not get a late second confirmation" + ); + } +} + +/// A stanza that fails to decrypt must emit a transport `<ack class="message">` +/// (else the server replays it on every reconnect forever), addressed to the +/// original `from` echoing `recipient`. Uses `Hide` (the production reactions +/// carried `decrypt-fail="hide"`) to also guard that hide does not suppress +/// the ack. BadMac so the retry carries no keys (no device account needed). +#[tokio::test] +async fn decrypt_failure_emits_transport_ack() { + let (client, transport) = capturing_client("decrypt_fail_ack").await; + + let sender: Jid = "236395184570386@lid".parse().expect("sender JID"); + let recipient: Jid = "156535032389744@lid".parse().expect("recipient JID"); + let info = Arc::new(MessageInfo { + id: "AC055553E56A2C12DE592DAD6353C477".to_string(), + source: crate::types::message::MessageSource { + sender: sender.clone(), + chat: recipient.clone(), + recipient: Some(recipient.clone()), + ..Default::default() + }, + ..Default::default() + }); + + client + .handle_decrypt_failure( + &info, + RetryReason::BadMac, + crate::types::events::DecryptFailMode::Hide, + ) + .await; + + // retry + ack are detached spawns; poll the wire until the ack appears. + let mut found = None; + for _ in 0..80 { + if let Some(a) = find_message_ack(&transport.sent()) { + found = Some(a); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let (to, recipient_attr) = found.expect( + "decrypt failure must emit a transport <ack class=message> \ + (else the server redelivers the stanza forever)", + ); + assert_eq!( + to, "236395184570386@lid", + "ack `to` must be the original `from` (own LID), not the chat" + ); + assert_eq!( + recipient_attr.as_deref(), + Some("156535032389744@lid"), + "ack must echo `recipient` for own-account fan-out" + ); +} + +/// Regression for the bot self-fanout loop on the DECRYPT-FAILURE path +/// (BadMac/NoSession): a self-fanout we cannot decrypt must be cleared with +/// a `<receipt type="sender">`, NOT a bare transport `<ack>` (ignored by the +/// server) nor a retry-to-self (futile). Once stuck in the loop the local +/// counter advances past the duplicate state, so this BadMac path is what +/// actually fires for an already-affected account. +#[tokio::test] +async fn self_fanout_decrypt_failure_acked_via_sender_receipt() { + let (client, transport) = capturing_client("self_fanout_badmac").await; + let info = Arc::new(MessageInfo { + id: "AC00000000000000000000000000BEEF".to_string(), + source: crate::types::message::MessageSource { + sender: "100000000000001@lid".parse().expect("sender"), + chat: "200000000000002@bot".parse().expect("chat"), + recipient: Some("200000000000002@bot".parse().expect("recipient")), + is_from_me: true, + ..Default::default() + }, + ..Default::default() + }); + + client + .handle_decrypt_failure( + &info, + RetryReason::BadMac, + crate::types::events::DecryptFailMode::Hide, + ) + .await; + + let mut found = None; + for _ in 0..80 { + if let Some(r) = find_receipt(&transport.sent(), "AC00000000000000000000000000BEEF") { + found = Some(r); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let (to, typ, recipient) = + found.expect("self-fanout decrypt failure must emit a sender <receipt> to drain the queue"); + assert_eq!(to, "100000000000001@lid"); + assert_eq!(typ.as_deref(), Some("sender")); + assert_eq!(recipient.as_deref(), Some("200000000000002@bot")); + + let sent = transport.sent(); + assert!( + find_message_ack(&sent).is_none(), + "must not emit the bare <ack> the server ignores" + ); + let mut saw_retry = false; + for (i, frame) in sent.iter().enumerate() { + if let Some(buf) = decode_frame(i, frame) + && let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) + && node.tag.as_ref() == "receipt" + && node.get_attr("type").is_some_and(|v| { + v.as_str() == crate::types::presence::ReceiptType::Retry.as_wire_str() + }) + { + saw_retry = true; + } + } + assert!( + !saw_retry, + "must not retry our own undecryptable fanout to ourselves" + ); +} + +/// Consistency with the success/duplicate path: a bot-authored own DM in a +/// non-bot chat (sender on `@bot`, user chat) must NOT take the sender +/// receipt on the decrypt-failure path either; it stays on the +/// bot-invoke-response bare-ack path (WA Web `!chat.isBot() && +/// author.isBot()`), matching ack_received_message and the locked +/// own_bot_author_dm_acks_not_sender_receipt test. +#[tokio::test] +async fn bot_author_self_fanout_decrypt_failure_not_sender_receipt() { + let (client, transport) = capturing_client("bot_author_badmac").await; + let info = Arc::new(MessageInfo { + id: "OWNBOTFAIL1".to_string(), + source: crate::types::message::MessageSource { + sender: "100000000000002@bot".parse().expect("sender"), + chat: "300000000000003@lid".parse().expect("chat"), + recipient: Some("300000000000003@lid".parse().expect("recipient")), + is_from_me: true, + ..Default::default() + }, + ..Default::default() + }); + + client + .handle_decrypt_failure( + &info, + RetryReason::BadMac, + crate::types::events::DecryptFailMode::Hide, + ) + .await; + + // Positive: the message IS cleared, via the bot-invoke-response bare + // <ack class="message"> (the retry-to-self is bot-skipped, so the + // transport ack follows), proving we took the ack path, not a no-op. + let mut found_ack = false; + for _ in 0..80 { + if find_message_ack(&transport.sent()).is_some() { + found_ack = true; + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert!( + found_ack, + "bot-authored own DM must still be transport-acked with a bare <ack class=message>" + ); + + // Negative settle: it must NEVER produce a sender receipt on the failure + // path (that would diverge from WA Web's bot-invoke ack and contradict + // the success-path ordering). + for _ in 0..5 { + assert!( + find_receipt(&transport.sent(), "OWNBOTFAIL1").is_none(), + "bot-authored own DM must not be cleared with a sender <receipt> on decrypt failure" + ); + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } +} + +/// If the resend request fails to send, the stanza must NOT be acked, so the +/// server keeps it queued for another try. Here NoSession needs keys, which +/// need a device account this harness lacks, so send_retry_receipt errors. +#[tokio::test] +async fn decrypt_failure_does_not_ack_when_retry_send_fails() { + let (client, transport) = capturing_client("retry_fail_no_ack").await; + let sender: Jid = "5511777776666@s.whatsapp.net".parse().expect("sender"); + let info = Arc::new(MessageInfo { + id: "NOACK1".to_string(), + source: crate::types::message::MessageSource { + sender: sender.clone(), + chat: sender.clone(), + ..Default::default() + }, + ..Default::default() + }); + client + .handle_decrypt_failure( + &info, + RetryReason::NoSession, + crate::types::events::DecryptFailMode::Show, + ) + .await; + tokio::time::sleep(std::time::Duration::from_millis(150)).await; + assert!( + find_message_ack(&transport.sent()).is_none(), + "must not ack when the resend request failed to send" + ); +} + +/// The retry receipt must be sent before the transport ack (one ordered +/// flushed task), so a disconnect mid-flush can never clear the stanza from +/// the offline queue without the sender having received a resend request. +#[tokio::test] +async fn decrypt_failure_sends_retry_before_ack() { + let (client, transport) = capturing_client("retry_before_ack").await; + let sender: Jid = "5511777776666@s.whatsapp.net".parse().expect("sender"); + let info = Arc::new(MessageInfo { + id: "RBA1".to_string(), + source: crate::types::message::MessageSource { + sender: sender.clone(), + chat: sender.clone(), + ..Default::default() + }, + ..Default::default() + }); + // BadMac (not NoSession) so the retry receipt carries no keys and needs + // no device account in this harness. + client + .handle_decrypt_failure( + &info, + RetryReason::BadMac, + crate::types::events::DecryptFailMode::Show, + ) + .await; + + let find = |tag: &str, retry: bool| -> Option<usize> { + let frames = transport.sent(); + for (i, frame) in frames.iter().enumerate() { + let Some(buf) = decode_frame(i, frame) else { + continue; + }; + let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { + continue; + }; + let is_retry = node.get_attr("type").is_some_and(|v| v.as_str() == "retry"); + if node.tag.as_ref() == tag && is_retry == retry { + return Some(i); + } + } + None + }; + + let mut retry_idx = None; + let mut ack_idx = None; + for _ in 0..80 { + retry_idx = find("receipt", true); + ack_idx = find("ack", false); + if retry_idx.is_some() && ack_idx.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let retry_idx = retry_idx.expect("retry receipt must be sent"); + let ack_idx = ack_idx.expect("transport ack must be sent"); + assert!( + retry_idx < ack_idx, + "retry receipt (frame {retry_idx}) must be sent before the ack (frame {ack_idx})" + ); +} + +/// status@broadcast is already acked by the `should_ack` gate post-dispatch, +/// so the decrypt-failure path must NOT emit a second transport ack +/// (whatsmeow/WA Web send exactly one per message). The retry receipt still +/// goes out. +#[tokio::test] +async fn status_broadcast_decrypt_failure_acks_to_chat() { + let (client, transport) = capturing_client("status_fail_ack").await; + let info = Arc::new(MessageInfo { + id: "STATUSMSGID".to_string(), + source: crate::types::message::MessageSource { + sender: "236395184570386@lid".parse().expect("sender"), + chat: "status@broadcast".parse().expect("status chat"), + is_group: true, + ..Default::default() + }, + ..Default::default() + }); + + client + .handle_decrypt_failure( + &info, + RetryReason::BadMac, + crate::types::events::DecryptFailMode::Show, + ) + .await; + + // status failures are acked from the flushed task (not just the detached + // should_ack gate), so the ack survives a disconnect mid-flush. + let mut found = None; + for _ in 0..80 { + if let Some(a) = find_message_ack(&transport.sent()) { + found = Some(a); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let (to, _) = found.expect("status failure must emit a flushed transport ack"); + assert_eq!( + to, "status@broadcast", + "status ack `to` must be the status chat" + ); +} + +/// Run a single session ciphertext through the full classify->process path. +async fn process_session_ct( + client: &Arc<Client>, + sender: &Jid, + id: &str, + ct: &wacore::libsignal::protocol::CiphertextMessage, +) { + use wacore::libsignal::protocol::CiphertextMessage; + let (enc_type, bytes) = match ct { + CiphertextMessage::PreKeySignalMessage(m) => ("pkmsg", m.serialized().to_vec()), + CiphertextMessage::SignalMessage(m) => ("msg", m.serialized().to_vec()), + _ => panic!("unexpected ciphertext type"), + }; + let enc = NodeBuilder::new("enc") + .attr("type", enc_type) + .bytes(bytes) + .build(); + let enc_ref = enc.as_node_ref(); + let payload = EncPayload::from_node_ref(&enc_ref).unwrap(); + let info = Arc::new(MessageInfo { + id: id.to_string(), + source: crate::types::message::MessageSource { + sender: sender.clone(), + chat: sender.clone(), + ..Default::default() + }, + ..Default::default() + }); + client + .clone() + .process_classified_message(ClassifiedMessage { + info, + sender_encryption_jid: sender.clone(), + session_payloads: vec![payload], + group_payloads: vec![], + bot_payloads: vec![], + max_sender_retry_count: 0, + decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, + }) + .await; +} + +fn enc_payload_from_ciphertext(ct: &CiphertextMessage) -> EncPayload { + let (enc_type, bytes) = match ct { + CiphertextMessage::PreKeySignalMessage(m) => ("pkmsg", m.serialized().to_vec()), + CiphertextMessage::SignalMessage(m) => ("msg", m.serialized().to_vec()), + _ => panic!("unexpected ciphertext type"), + }; + let enc = NodeBuilder::new("enc") + .attr("type", enc_type) + .bytes(bytes) + .build(); + EncPayload::from_node_ref(&enc.as_node_ref()).expect("ciphertext payload") +} + +fn skmsg_payload_from_bytes(bytes: Vec<u8>) -> EncPayload { + let enc = NodeBuilder::new("enc") + .attr("type", "skmsg") + .bytes(bytes) + .build(); + EncPayload::from_node_ref(&enc.as_node_ref()).expect("skmsg payload") +} + +fn msmsg_payload_from_bytes(bytes: Vec<u8>) -> EncPayload { + let enc = NodeBuilder::new("enc") + .attr("type", "msmsg") + .bytes(bytes) + .build(); + EncPayload::from_node_ref(&enc.as_node_ref()).expect("msmsg payload") +} + +fn group_message_info(id: &str, group: &Jid, sender: &Jid, is_from_me: bool) -> Arc<MessageInfo> { + Arc::new(MessageInfo { + id: id.to_string(), + source: crate::types::message::MessageSource { + sender: sender.clone(), + chat: group.clone(), + is_from_me, + is_group: true, + ..Default::default() + }, + ..Default::default() + }) +} + +async fn process_group_classified( + client: &Arc<Client>, + info: Arc<MessageInfo>, + sender: &Jid, + session_payload: EncPayload, + group_payloads: Vec<EncPayload>, +) { + process_group_classified_with_sessions( + client, + info, + sender, + vec![session_payload], + group_payloads, + ) + .await; +} + +async fn process_group_classified_with_sessions( + client: &Arc<Client>, + info: Arc<MessageInfo>, + sender: &Jid, + session_payloads: Vec<EncPayload>, + group_payloads: Vec<EncPayload>, +) { + process_group_classified_with_payloads( + client, + info, + sender, + session_payloads, + group_payloads, + vec![], + ) + .await; +} + +async fn process_group_classified_with_payloads( + client: &Arc<Client>, + info: Arc<MessageInfo>, + sender: &Jid, + session_payloads: Vec<EncPayload>, + group_payloads: Vec<EncPayload>, + bot_payloads: Vec<EncPayload>, +) { + client + .clone() + .process_classified_message(ClassifiedMessage { + info, + sender_encryption_jid: sender.clone(), + session_payloads, + group_payloads, + bot_payloads, + max_sender_retry_count: 0, + decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, + }) + .await; +} + +fn message_events_for_id(rx: &async_channel::Receiver<Arc<Event>>, id: &str) -> (usize, usize) { + let mut count = 0; + let mut visible_content = 0; + while let Ok(event) = rx.try_recv() { + if let Event::Message(msg, info) = event.as_ref() + && info.id == id + { + count += 1; + if msg.conversation.is_some() { + visible_content += 1; + } + } + } + (count, visible_content) +} + +fn message_texts_for_id(rx: &async_channel::Receiver<Arc<Event>>, id: &str) -> Vec<String> { + let mut texts = Vec::new(); + while let Ok(event) = rx.try_recv() { + if let Event::Message(msg, info) = event.as_ref() + && info.id == id + && let Some(text) = &msg.conversation + { + texts.push(text.clone()); + } + } + texts +} + +#[tokio::test] +async fn skdm_only_group_session_acknowledged_once_without_message_event() { + use wacore::messages::MessageUtils; + use wacore::types::events::ChannelEventHandler; + + let (client, transport) = capturing_client("skdm_only_group_ack").await; + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("146824178450525@lid").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + let group: Jid = "120363408782575443@g.us".parse().expect("group"); + let skdm = alice.create_group_skdm(&group).await; + let plaintext = MessageUtils::encode_and_pad(&wa::Message { + sender_key_distribution_message: Some(skdm), + ..Default::default() + }); + let session_ct = alice.encrypt(&bob_addr, &plaintext).await; + let id = "SKDM_ONLY_SESSION"; + let info = group_message_info(id, &group, &alice.jid, false); + + process_group_classified( + &client, + info, + &alice.jid, + enc_payload_from_ciphertext(&session_ct), + vec![], + ) + .await; + + assert_exactly_one_confirmation(&transport, id).await; + assert_eq!( + delivery_receipts_for(&transport.sent(), id), + 1, + "incoming group SKDM-only session message should drain via delivery receipt" + ); + let sent = transport.sent(); + let receipt = find_receipt_details(&sent, id).expect("delivery receipt"); + let sender_str = alice.jid.to_string(); + assert_eq!(receipt.to, group.to_string()); + assert_eq!(receipt.participant.as_deref(), Some(sender_str.as_str())); + assert_eq!(receipt.recipient, None); + assert_ne!( + receipt.typ.as_deref(), + Some("sender"), + "incoming group SKDM-only must not be cleared as a sender receipt" + ); + assert_eq!( + message_acks_for(&sent, id), + 0, + "incoming group SKDM-only must not also emit a transport ack" + ); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + assert_eq!( + message_events_for_id(&rx, id), + (0, 0), + "SKDM-only messages must not surface Event::Message" + ); +} + +#[tokio::test] +async fn session_plaintext_decode_error_is_not_acked_as_skdm_only() { + use wacore::types::events::ChannelEventHandler; + + let (client, transport) = capturing_client("bad_plaintext_no_skdm_ack").await; + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + let recorder = Arc::new(EventRecorder::default()); + client.register_handler(recorder.clone()); + + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("146824178450527@lid").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + let group: Jid = "120363408782575446@g.us".parse().expect("group"); + let invalid_padded_plaintext = vec![0xff, 0x01]; + let session_ct = alice.encrypt(&bob_addr, &invalid_padded_plaintext).await; + let id = "BAD_SESSION_PLAINTEXT"; + let info = group_message_info(id, &group, &alice.jid, false); + + process_group_classified( + &client, + info, + &alice.jid, + enc_payload_from_ciphertext(&session_ct), + vec![], + ) + .await; + + for _ in 0..5 { + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + assert_eq!( + confirmations_for(&transport.sent(), id), + 0, + "invalid plaintext must not be counted as a successful SKDM-only ack" + ); + } + assert_eq!( + recorder.undecryptable().len(), + 1, + "plaintext handler failures must stay on the undecryptable path" + ); + assert_eq!( + message_events_for_id(&rx, id), + (0, 0), + "invalid plaintext must not surface Event::Message" + ); + let mut nack_code = None; + for _ in 0..80 { + nack_code = find_message_nack_error(&transport.sent(), id); + if nack_code.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!( + nack_code, + Some(491), + "invalid decrypted protobuf must be drained with InvalidProtobuf nack" + ); +} + +#[tokio::test] +async fn mixed_skdm_and_bad_plaintext_session_is_nacked_not_positive_acked() { + use wacore::messages::MessageUtils; + use wacore::types::events::ChannelEventHandler; + + let (client, transport) = capturing_client("mixed_skdm_bad_plaintext").await; + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + let recorder = Arc::new(EventRecorder::default()); + client.register_handler(recorder.clone()); + + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("146824178450531@lid").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + let group: Jid = "120363408782575449@g.us".parse().expect("group"); + let skdm = alice.create_group_skdm(&group).await; + let skdm_plaintext = MessageUtils::encode_and_pad(&wa::Message { + sender_key_distribution_message: Some(skdm), + ..Default::default() + }); + let skdm_ct = alice.encrypt(&bob_addr, &skdm_plaintext).await; + let bad_ct = alice.encrypt(&bob_addr, &[0xff, 0x01]).await; + let id = "SKDM_WITH_BAD_SESSION"; + let info = group_message_info(id, &group, &alice.jid, false); + + process_group_classified_with_sessions( + &client, + info, + &alice.jid, + vec![ + enc_payload_from_ciphertext(&skdm_ct), + enc_payload_from_ciphertext(&bad_ct), + ], + vec![], + ) + .await; + + let mut nack_code = None; + for _ in 0..80 { + nack_code = find_message_nack_error(&transport.sent(), id); + if nack_code.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!(nack_code, Some(491)); + assert_eq!( + confirmations_for(&transport.sent(), id), + 0, + "SKDM-only fallback must not positive-ack a mixed malformed session batch" + ); + assert_eq!( + recorder.undecryptable().len(), + 1, + "the malformed sibling must still surface as undecryptable" + ); + assert_eq!( + message_events_for_id(&rx, id), + (0, 0), + "mixed SKDM and bad plaintext must not dispatch user content" + ); +} + +#[tokio::test] +async fn bad_session_plaintext_skips_skmsg_sibling_after_nack() { + use wacore::messages::MessageUtils; + use wacore::types::events::ChannelEventHandler; + + let (client, transport) = capturing_client("bad_session_skips_skmsg").await; + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + let recorder = Arc::new(EventRecorder::default()); + client.register_handler(recorder.clone()); + + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("146824178450532@lid").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + let group: Jid = "120363408782575450@g.us".parse().expect("group"); + let skdm = alice.create_group_skdm(&group).await; + let skdm_plaintext = MessageUtils::encode_and_pad(&wa::Message { + sender_key_distribution_message: Some(skdm), + ..Default::default() + }); + let skdm_ct = alice.encrypt(&bob_addr, &skdm_plaintext).await; + let bad_ct = alice.encrypt(&bob_addr, &[0xff, 0x01]).await; + let content_plaintext = MessageUtils::encode_and_pad(&wa::Message { + conversation: Some("must not dispatch".to_string()), + ..Default::default() + }); + let skmsg = alice + .encrypt_group_message(&group, &content_plaintext) + .await; + let id = "BAD_SESSION_WITH_SKMSG"; + let info = group_message_info(id, &group, &alice.jid, false); + + process_group_classified_with_sessions( + &client, + info, + &alice.jid, + vec![ + enc_payload_from_ciphertext(&skdm_ct), + enc_payload_from_ciphertext(&bad_ct), + ], + vec![skmsg_payload_from_bytes(skmsg)], + ) + .await; + + let mut nack_code = None; + for _ in 0..80 { + nack_code = find_message_nack_error(&transport.sent(), id); + if nack_code.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!(nack_code, Some(491)); + assert_eq!( + confirmations_for(&transport.sent(), id), + 0, + "skmsg must not ack after a session InvalidProtobuf nack" + ); + assert_eq!( + recorder.undecryptable().len(), + 1, + "session plaintext failure should own the only user-visible failure" + ); + assert_eq!( + message_texts_for_id(&rx, id), + Vec::<String>::new(), + "skmsg content must be skipped after session plaintext failure" + ); +} + +#[tokio::test] +async fn skdm_only_session_with_msmsg_waits_for_bot_payload_response() { + use wacore::messages::MessageUtils; + + let (client, transport) = capturing_client("skdm_msmsg_no_fallback_ack").await; + + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("146824178450533@lid").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + let group: Jid = "120363408782575451@g.us".parse().expect("group"); + let skdm = alice.create_group_skdm(&group).await; + let plaintext = MessageUtils::encode_and_pad(&wa::Message { + sender_key_distribution_message: Some(skdm), + ..Default::default() + }); + let session_ct = alice.encrypt(&bob_addr, &plaintext).await; + let id = "SKDM_WITH_MSMSG"; + let info = group_message_info(id, &group, &alice.jid, false); + + process_group_classified_with_payloads( + &client, + info, + &alice.jid, + vec![enc_payload_from_ciphertext(&session_ct)], + vec![], + vec![msmsg_payload_from_bytes(vec![0xff])], + ) + .await; + + let mut nack_code = None; + for _ in 0..80 { + nack_code = find_message_nack_error(&transport.sent(), id); + if nack_code.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!(nack_code, Some(487)); + assert_eq!( + confirmations_for(&transport.sent(), id), + 0, + "SKDM-only fallback must not pre-ack a stanza with msmsg work" + ); +} + +#[tokio::test] +async fn session_content_group_message_acknowledged_once_without_fallback() { + use wacore::messages::MessageUtils; + use wacore::types::events::ChannelEventHandler; + + let (client, transport) = capturing_client("session_content_group_ack").await; + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("146824178450528@lid").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + let group: Jid = "120363408782575447@g.us".parse().expect("group"); + let plaintext = MessageUtils::encode_and_pad(&wa::Message { + conversation: Some("session content".to_string()), + ..Default::default() + }); + let session_ct = alice.encrypt(&bob_addr, &plaintext).await; + let id = "SESSION_CONTENT_GROUP"; + let info = group_message_info(id, &group, &alice.jid, false); + + process_group_classified( + &client, + info, + &alice.jid, + enc_payload_from_ciphertext(&session_ct), + vec![], + ) + .await; + + assert_exactly_one_confirmation(&transport, id).await; + let sent = transport.sent(); + assert_eq!( + delivery_receipts_for(&sent, id), + 1, + "session content dispatch should own the only delivery receipt" + ); + assert_eq!( + message_acks_for(&sent, id), + 0, + "normal session content must not also use the SKDM-only transport ack" + ); + assert_eq!( + message_texts_for_id(&rx, id), + vec!["session content".to_string()], + "normal session content must dispatch exactly once" + ); +} + +#[tokio::test] +async fn status_skdm_only_session_uses_one_status_receipt() { + use wacore::messages::MessageUtils; + use wacore::types::events::ChannelEventHandler; + + let (client, transport) = capturing_client("status_skdm_only_ack").await; + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("146824178450529@lid").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + let status: Jid = "status@broadcast".parse().expect("status"); + let skdm = alice.create_group_skdm(&status).await; + let plaintext = MessageUtils::encode_and_pad(&wa::Message { + sender_key_distribution_message: Some(skdm), + ..Default::default() + }); + let session_ct = alice.encrypt(&bob_addr, &plaintext).await; + let id = "STATUS_SKDM_ONLY"; + let info = group_message_info(id, &status, &alice.jid, false); + + process_group_classified( + &client, + info, + &alice.jid, + enc_payload_from_ciphertext(&session_ct), + vec![], + ) + .await; + + assert_exactly_one_confirmation(&transport, id).await; + let sent = transport.sent(); + assert_eq!( + delivery_receipts_for(&sent, id), + 1, + "status SKDM-only success must still send the WA Web status receipt" + ); + assert_eq!( + message_acks_for(&sent, id), + 0, + "status success path should not use a transport ack" + ); + let receipt = find_receipt_details(&sent, id).expect("status delivery receipt"); + let sender_str = alice.jid.to_string(); + assert_eq!(receipt.to, status.to_string()); + assert_eq!(receipt.participant.as_deref(), Some(sender_str.as_str())); + assert_eq!(receipt.context.as_deref(), Some("status")); + assert_eq!( + message_events_for_id(&rx, id), + (0, 0), + "status SKDM-only messages must not surface Event::Message" + ); +} + +#[tokio::test] +async fn error_message_ack_is_not_counted_as_positive_confirmation() { + let (client, transport) = capturing_client("error_ack_not_positive").await; + let id = "ERROR_ACK_NOT_POSITIVE"; + let info = Arc::new(MessageInfo { + id: id.to_string(), + source: crate::types::message::MessageSource { + sender: "146824178450530@lid".parse().expect("sender"), + chat: "120363408782575448@g.us".parse().expect("group"), + is_group: true, + ..Default::default() + }, + ..Default::default() + }); + + client.spawn_nack( + &info, + wacore::protocol::nack::NackReason::ParsingError, + None, + ); + + let mut nack_code = None; + for _ in 0..80 { + nack_code = find_message_nack_error(&transport.sent(), id); + if nack_code.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!(nack_code, Some(487)); + + let sent = transport.sent(); + assert_eq!( + message_acks_for(&sent, id), + 0, + "nacks carry class=message but must not count as positive transport acks" + ); + assert_eq!( + confirmations_for(&sent, id), + 0, + "nacks must not satisfy exactly-one positive confirmation assertions" + ); + assert!( + find_message_ack_for(&sent, id).is_none(), + "error acks must be excluded from positive ack lookup" + ); +} + +#[tokio::test] +async fn skdm_session_with_skmsg_sibling_acknowledged_once() { + use wacore::messages::MessageUtils; + use wacore::types::events::ChannelEventHandler; + + let (client, transport) = capturing_client("skdm_plus_skmsg_ack").await; + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("146824178450526@lid").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + let group: Jid = "120363408782575444@g.us".parse().expect("group"); + let skdm = alice.create_group_skdm(&group).await; + let skdm_plaintext = MessageUtils::encode_and_pad(&wa::Message { + sender_key_distribution_message: Some(skdm), + ..Default::default() + }); + let session_ct = alice.encrypt(&bob_addr, &skdm_plaintext).await; + + let content_plaintext = MessageUtils::encode_and_pad(&wa::Message { + conversation: Some("group content".to_string()), + ..Default::default() + }); + let skmsg = alice + .encrypt_group_message(&group, &content_plaintext) + .await; + let id = "SKDM_WITH_SKMSG"; + let info = group_message_info(id, &group, &alice.jid, false); + + process_group_classified( + &client, + info, + &alice.jid, + enc_payload_from_ciphertext(&session_ct), + vec![skmsg_payload_from_bytes(skmsg)], + ) + .await; + + assert_exactly_one_confirmation(&transport, id).await; + assert_eq!( + delivery_receipts_for(&transport.sent(), id), + 1, + "the skmsg content dispatch should own the only receipt" + ); + let sent = transport.sent(); + let receipt = find_receipt_details(&sent, id).expect("delivery receipt"); + let sender_str = alice.jid.to_string(); + assert_eq!(receipt.to, group.to_string()); + assert_eq!(receipt.participant.as_deref(), Some(sender_str.as_str())); + assert_eq!( + message_acks_for(&sent, id), + 0, + "SKDM+skmsg sibling must not get an extra transport ack" + ); + assert_eq!( + message_texts_for_id(&rx, id), + vec!["group content".to_string()], + "only the skmsg content should dispatch a user message" + ); +} + +#[tokio::test] +async fn own_group_skdm_only_session_uses_transport_ack_once() { + use wacore::messages::MessageUtils; + + let (client, transport) = capturing_client("own_group_skdm_ack").await; + + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("999999999999999@lid").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + let group: Jid = "120363408782575445@g.us".parse().expect("group"); + let skdm = alice.create_group_skdm(&group).await; + let plaintext = MessageUtils::encode_and_pad(&wa::Message { + sender_key_distribution_message: Some(skdm), + ..Default::default() + }); + let session_ct = alice.encrypt(&bob_addr, &plaintext).await; + let id = "OWN_GROUP_SKDM_ONLY"; + let info = group_message_info(id, &group, &alice.jid, true); + + process_group_classified( + &client, + info, + &alice.jid, + enc_payload_from_ciphertext(&session_ct), + vec![], + ) + .await; + + assert_exactly_one_confirmation(&transport, id).await; + let sent = transport.sent(); + assert_eq!( + message_acks_for(&sent, id), + 1, + "own group SKDM-only session message should use transport ack" + ); + let ack = find_message_ack_for(&sent, id).expect("transport ack"); + let sender_str = alice.jid.to_string(); + assert_eq!(ack.to, group.to_string()); + assert_eq!(ack.participant.as_deref(), Some(sender_str.as_str())); + assert_eq!(ack.recipient, None); + assert_eq!( + delivery_receipts_for(&sent, id), + 0, + "own group SKDM-only session message must not use a delivery receipt" + ); + assert_eq!( + sender_receipts_for(&sent, id), + 0, + "group self-fanout must not use type=sender receipt" + ); +} + +/// Regression for the offline-backlog disconnect: an already-processed +/// (duplicate) message must get its own delivery receipt, else the server +/// replays it every reconnect until it force-closes the stream. Pre-fix only +/// the first (success) delivery was acked; the duplicate was skipped silently. +#[tokio::test] +async fn duplicate_message_is_acked_with_delivery_receipt() { + let (client, transport) = capturing_client("dup_receipt").await; + let (bundle, bob_jid) = bobs_prekey_bundle(&client).await; + let bob_addr = bob_jid.to_protocol_address(); + let mut alice = AlicePeer::new("5511888887777@s.whatsapp.net").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + // Establish the session, then mark Alice's prekey acked so her next message + // is a plain SignalMessage. Re-submitting it is a clean duplicate. + let establish = alice.encrypt_text(&bob_addr, "establish").await; + process_session_ct(&client, &alice.jid, "EST", &establish).await; + if let Some(record) = alice.sessions.0.get_mut(&bob_addr) + && let Some(state) = record.session_state_mut() + { + state.clear_unacknowledged_pre_key_message(); + } + + // A real (padded) Message so the success path also emits its receipt. + let plaintext = wacore::messages::MessageUtils::encode_and_pad(&wa::Message { + conversation: Some("hi".to_string()), + ..Default::default() + }); + let msg = alice.encrypt(&bob_addr, &plaintext).await; + process_session_ct(&client, &alice.jid, "DUP", &msg).await; // success + process_session_ct(&client, &alice.jid, "DUP", &msg).await; // duplicate + + let mut count = 0; + for _ in 0..80 { + count = delivery_receipts_for(&transport.sent(), "DUP"); + if count >= 2 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!( + count, 2, + "duplicate must get its own delivery receipt (pre-fix: only the first send was acked)" + ); +} + +/// Own-account self-fanout (is_from_me, non-peer, carries a `recipient`): +/// our own outgoing message echoed back to this device. WA Web +/// (`isMeAccount(author) => SENDER`) and whatsmeow (`IsFromMe => "sender"`) +/// clear it with a `<receipt type="sender" recipient=...>`, NOT a bare +/// transport `<ack>`. The server's offline queue ignores the bare ack and +/// replays the stanza forever (the ~50min disconnect loop). +#[tokio::test] +async fn own_self_fanout_acked_via_sender_receipt() { + let (client, transport) = capturing_client("own_ack").await; + let own = Arc::new(MessageInfo { + id: "OWN1".to_string(), + source: crate::types::message::MessageSource { + sender: "100000000000001@lid".parse().expect("sender"), + chat: "300000000000003@lid".parse().expect("chat"), + recipient: Some("300000000000003@lid".parse().expect("recipient")), + is_from_me: true, + ..Default::default() + }, + ..Default::default() + }); + client.ack_received_message(&own); + + let mut found = None; + for _ in 0..80 { + if let Some(r) = find_receipt(&transport.sent(), "OWN1") { + found = Some(r); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let (to, typ, recipient) = found.expect("own self-fanout must get a sender <receipt>"); + assert_eq!( + to, "100000000000001@lid", + "receipt `to` must echo the own LID (the fanout sender)" + ); + assert_eq!( + typ.as_deref(), + Some("sender"), + "own self-fanout receipt must be type=sender" + ); + assert_eq!( + recipient.as_deref(), + Some("300000000000003@lid"), + "receipt must echo the fanout recipient" + ); + assert!( + find_message_ack(&transport.sent()).is_none(), + "self-fanout must NOT also emit a bare transport <ack> (the server rejects it)" + ); +} + +/// Regression for the bot self-fanout disconnect loop: our own message to a +/// `@bot` recipient, echoed back as a duplicate/undecryptable stanza, must +/// be cleared with a `<receipt type="sender" recipient=@bot>`. Pre-fix it +/// got a bare `<ack class="message">` which the server ignored, replaying +/// the stanza every reconnect until a ~50min `<stream:error><ack/>` GC +/// force-closed the connection (the exact production symptom). +#[tokio::test] +async fn bot_self_fanout_acked_via_sender_receipt() { + let (client, transport) = capturing_client("bot_self_fanout").await; + let own = Arc::new(MessageInfo { + id: "AC00000000000000000000000000BEEF".to_string(), + source: crate::types::message::MessageSource { + // from = our own LID with its device (the server fans our + // outgoing bot prompt back to this device); chat = the bot + // (recipient.to_non_ad). The device on the sender must survive + // into the receipt `to`, or the LID server rejects it (#649). + sender: "100000000000001:11@lid".parse().expect("sender"), + chat: "200000000000002@bot".parse().expect("chat"), + recipient: Some("200000000000002@bot".parse().expect("recipient")), + is_from_me: true, + ..Default::default() + }, + ..Default::default() + }); + client.ack_received_message(&own); + + let mut found = None; + for _ in 0..80 { + if let Some(r) = find_receipt(&transport.sent(), "AC00000000000000000000000000BEEF") { + found = Some(r); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let (to, typ, recipient) = + found.expect("bot self-fanout must get a sender <receipt> to drain the offline queue"); + assert_eq!( + to, "100000000000001:11@lid", + "receipt `to` must preserve the own LID device" + ); + assert_eq!(typ.as_deref(), Some("sender")); + assert_eq!( + recipient.as_deref(), + Some("200000000000002@bot"), + "receipt must route to the bot recipient" + ); + assert!( + find_message_ack(&transport.sent()).is_none(), + "the bare <ack> that triggered <stream:error><ack/> must no longer be emitted" + ); +} + +/// When WE are the bot author (own DM, sender on the `@bot` server, to a +/// user), WA Web's `MsgSendReceipt` takes the `!chat.isBot() && +/// author.isBot()` branch and emits a bot-invoke-response `<ack>`, NOT a +/// sender `<receipt>`. So the bot-author branch in ack_received_message must +/// keep running before the self-fanout receipt: this locks that ordering +/// against a regression that would wrongly route it to a sender receipt. +#[tokio::test] +async fn own_bot_author_dm_acks_not_sender_receipt() { + let (client, transport) = capturing_client("own_bot_author").await; + let own = Arc::new(MessageInfo { + id: "OWNBOT1".to_string(), + source: crate::types::message::MessageSource { + sender: "100000000000002@bot".parse().expect("sender"), + chat: "300000000000003@lid".parse().expect("chat"), + recipient: Some("300000000000003@lid".parse().expect("recipient")), + is_from_me: true, + ..Default::default() + }, + ..Default::default() + }); + client.ack_received_message(&own); + + let mut found = None; + for _ in 0..80 { + if let Some(a) = find_message_ack(&transport.sent()) { + found = Some(a); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert!( + found.is_some(), + "own bot-author DM must emit a bare <ack class=message> (WA Web bot-invoke-response ack), not a receipt" + ); + // No current race (ack_received_message is synchronous and the + // bot-author branch returns before the receipt branch), but settle + // briefly so a future regression that spawned a receipt on a later tick + // can't slip past this negative assertion. + for _ in 0..5 { + assert!( + find_receipt(&transport.sent(), "OWNBOT1").is_none(), + "must NOT route to a sender <receipt> (would diverge from WA Web's bot-invoke-response ack path)" + ); + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } +} + +/// An `<unavailable>` message (no `<enc>`) must be transport-acked so the +/// server stops replaying it (DM/group aren't covered by the should_ack gate). +#[tokio::test] +async fn unavailable_message_is_transport_acked() { + let (client, transport) = capturing_client("unavail_ack").await; + let node = NodeBuilder::new("message") + .attr("from", "5511777776666@s.whatsapp.net") + .attr("id", "UNAVAIL1") + .attr("type", "text") + .children([NodeBuilder::new("unavailable") + .attr("type", "view_once") + .build()]) + .build(); + let owned = node_to_arc(node); + let classified = client.classify_incoming_message(&owned).await; + assert!(classified.is_none(), "unavailable path returns None"); + + let mut found = None; + for _ in 0..80 { + if let Some(a) = find_message_ack(&transport.sent()) { + found = Some(a); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let (to, _) = found.expect("unavailable message must get a transport ack"); + assert_eq!(to, "5511777776666@s.whatsapp.net"); +} + +/// Unknown-only stanzas (e.g. msmsg) must be acked or they loop the queue. +#[tokio::test] +async fn unknown_only_enc_is_transport_acked() { + let (client, transport) = capturing_client("msmsg_ack").await; + let node = NodeBuilder::new("message") + .attr("from", "5511777776666@s.whatsapp.net") + .attr("id", "MSMSG1") + .attr("type", "text") + .children([NodeBuilder::new("enc") + .attr("type", "frskmsg") + .bytes(vec![0u8; 8]) + .build()]) + .build(); + let owned = node_to_arc(node); + let classified = client.classify_incoming_message(&owned).await; + assert!( + classified.is_none(), + "unknown-only enc must short-circuit before the process phase" + ); + + let mut found = None; + for _ in 0..80 { + if let Some(a) = find_message_ack(&transport.sent()) { + found = Some(a); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let (to, _) = found.expect("unknown-only enc must emit a transport ack"); + assert_eq!(to, "5511777776666@s.whatsapp.net"); +} + +/// `recipient` must be echoed verbatim or the server replies <stream:error>. +#[tokio::test] +async fn unknown_only_enc_ack_preserves_recipient() { + let (client, transport) = capturing_client("msmsg_recipient").await; + let node = NodeBuilder::new("message") + .attr("from", "236395184570386@lid") + .attr("recipient", "156535032389744@lid") + .attr("id", "MSMSG_LID") + .attr("type", "text") + .children([NodeBuilder::new("enc") + .attr("type", "frskmsg") + .bytes(vec![0u8; 8]) + .build()]) + .build(); + let owned = node_to_arc(node); + let classified = client.classify_incoming_message(&owned).await; + assert!(classified.is_none()); + + let mut found = None; + for _ in 0..80 { + if let Some(a) = find_message_ack(&transport.sent()) { + found = Some(a); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let (to, recipient) = found.expect("unknown-only enc must emit a transport ack"); + assert_eq!(to, "236395184570386@lid"); + assert_eq!( + recipient.as_deref(), + Some("156535032389744@lid"), + "ack must echo the incoming `recipient` attr or the server replies with <stream:error><ack/>" + ); +} + +/// Known type with empty content still has no usable payload; ack it. +#[tokio::test] +async fn known_enc_type_with_empty_content_is_transport_acked() { + let (client, transport) = capturing_client("known_empty").await; + let node = NodeBuilder::new("message") + .attr("from", "5511777776666@s.whatsapp.net") + .attr("id", "EMPTY1") + .attr("type", "text") + .children([NodeBuilder::new("enc").attr("type", "pkmsg").build()]) + .build(); + let owned = node_to_arc(node); + let classified = client.classify_incoming_message(&owned).await; + assert!(classified.is_none()); + + let mut found = None; + for _ in 0..80 { + if let Some(a) = find_message_ack(&transport.sent()) { + found = Some(a); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let (to, _) = found.expect("known-but-empty enc must emit a transport ack"); + assert_eq!(to, "5511777776666@s.whatsapp.net"); +} + +/// status is covered by should_ack; the fallback must not double-ack it. +#[tokio::test] +async fn unknown_only_enc_on_status_skips_fallback_ack() { + let (client, transport) = capturing_client("msmsg_status_skip").await; + let node = NodeBuilder::new("message") + .attr("from", "status@broadcast") + .attr("id", "MSMSG_STATUS") + .attr("type", "text") + .attr("participant", "5511777776666@s.whatsapp.net") + .children([NodeBuilder::new("enc") + .attr("type", "frskmsg") + .bytes(vec![0u8; 8]) + .build()]) + .build(); + let owned = node_to_arc(node); + let classified = client.classify_incoming_message(&owned).await; + assert!(classified.is_none()); + + // Give any rogue spawned task time to land on the wire. + for _ in 0..16 { + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + if find_message_ack(&transport.sent()).is_some() { + break; + } + } + assert!( + find_message_ack(&transport.sent()).is_none(), + "status@broadcast must not get a fallback transport ack from classify" + ); +} + +/// One recognized enc + one unknown must still go through the normal path. +#[tokio::test] +async fn mixed_recognized_and_unknown_enc_still_classifies() { + let (client, _transport) = capturing_client("msmsg_mixed").await; + let node = NodeBuilder::new("message") + .attr("from", "5511777776666@s.whatsapp.net") + .attr("id", "MIXED1") + .attr("type", "text") + .children([ + NodeBuilder::new("enc") + .attr("type", "pkmsg") + .bytes(vec![0u8; 8]) + .build(), + NodeBuilder::new("enc") + .attr("type", "frskmsg") + .bytes(vec![0u8; 8]) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + let classified = client + .classify_incoming_message(&owned) + .await + .expect("mixed enc must produce a ClassifiedMessage"); + assert_eq!(classified.session_payloads.len(), 1); + assert!(classified.group_payloads.is_empty()); +} + +/// A custom handler owns its ack; the fallback must not double-ack. +#[tokio::test] +async fn custom_handler_only_skips_fallback_ack() { + use crate::types::enc_handler::EncHandler; + use async_lock::Mutex as AsyncMutex; + + #[derive(Default)] + struct NoopHandler { + calls: Arc<AsyncMutex<usize>>, + } + #[async_trait::async_trait] + impl EncHandler for NoopHandler { + async fn handle( + &self, + _client: Arc<Client>, + _enc_node: &wacore_binary::Node, + _info: &crate::types::message::MessageInfo, + ) -> anyhow::Result<()> { + *self.calls.lock().await += 1; + Ok(()) + } + } + + let (client, transport) = capturing_client("msmsg_custom").await; + let calls = Arc::new(AsyncMutex::new(0usize)); + let handler = Arc::new(NoopHandler { + calls: Arc::clone(&calls), + }); + client + .custom_enc_handlers + .write() + .await + .insert("frskmsg".to_string(), handler as Arc<dyn EncHandler>); + + let node = NodeBuilder::new("message") + .attr("from", "5511777776666@s.whatsapp.net") + .attr("id", "CUSTOM1") + .attr("type", "text") + .children([NodeBuilder::new("enc") + .attr("type", "frskmsg") + .bytes(vec![0u8; 8]) + .build()]) + .build(); + let owned = node_to_arc(node); + let classified = client.classify_incoming_message(&owned).await; + assert!( + classified.is_some(), + "custom-handled enc must not be short-circuited by the fallback guard" + ); + + // Let the detached handler + any rogue spawned task run. + for _ in 0..16 { + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert!( + find_message_ack(&transport.sent()).is_none(), + "custom-handled enc must not get a fallback transport ack from classify" + ); + assert_eq!(*calls.lock().await, 1, "custom handler must be invoked"); +} + +/// Security regression: a self-only `app_state_sync_key_share` protocol +/// message must be honoured only when it originates from our own account. +/// A spoofed one from a peer must be dropped (otherwise a peer could inject +/// app-state sync keys). Mirrors WA Web `WAWebKeyManagementHandleKeyShareApi` +/// and whatsmeow's `handleProtocolMessage` self gate. +#[tokio::test] +async fn app_state_sync_key_share_honored_only_from_self() { + use wacore::messages::MessageUtils; + + let client = crate::test_utils::create_test_client().await; + ensure_bob_paired(&client).await; + + let key_id = vec![1u8, 2, 3, 4, 5, 6]; + let share = wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + app_state_sync_key_share: Some(wa::message::AppStateSyncKeyShare { + keys: vec![wa::message::AppStateSyncKey { + key_id: Some(wa::message::AppStateSyncKeyId { + key_id: Some(key_id.clone()), + }), + key_data: Some(wa::message::AppStateSyncKeyData { + key_data: Some(vec![7u8; 32]), + fingerprint: Some(wa::message::AppStateSyncKeyFingerprint { + raw_id: Some(1), + current_index: Some(0), + device_indexes: vec![0], + }), + timestamp: Some(123), + }), + }], + }), + ..Default::default() + })), + ..Default::default() + }; + let padded = MessageUtils::encode_and_pad(&share); + let backend = client.persistence_manager.backend(); + + // Non-self sender: the key share must be dropped. + let mut info = + create_test_message_info("5510000@s.whatsapp.net", "AKS1", "5510000@s.whatsapp.net"); + info.source.is_from_me = false; + client + .clone() + .handle_decrypted_plaintext("msg", &padded, 2, &Arc::new(info)) + .await + .unwrap(); + assert!( + backend.get_sync_key(&key_id).await.unwrap().is_none(), + "app-state sync key from a non-self sender must not be stored" + ); + + // Self sender: the key share is honoured and stored. + let mut info = create_test_message_info( + "9000000000000@s.whatsapp.net", + "AKS2", + "9000000000000@s.whatsapp.net", + ); + info.source.is_from_me = true; + client + .clone() + .handle_decrypted_plaintext("msg", &padded, 2, &Arc::new(info)) + .await + .unwrap(); + assert!( + backend.get_sync_key(&key_id).await.unwrap().is_some(), + "app-state sync key from self must be stored" + ); +} + +// ---- msmsg inbound dispatch ----------------------------------------- + +fn find_message_nack_error(frames: &[bytes::Bytes], id: &str) -> Option<u32> { + for (i, frame) in frames.iter().enumerate() { + let Some(buf) = decode_frame(i, frame) else { + continue; + }; + let Ok(node) = wacore_binary::marshal::unmarshal_ref(&buf[1..]) else { + continue; + }; + if node.tag.as_ref() == "ack" + && node + .get_attr("class") + .is_some_and(|v| v.as_str() == "message") + && node.get_attr("id").is_some_and(|v| v.as_str() == id) + && let Some(err) = node.get_attr("error") + && let Ok(code) = err.as_str().parse::<u32>() + { + return Some(code); + } + } + None +} + +fn encode_message_secret_message(iv: &[u8], payload: &[u8]) -> Vec<u8> { + use prost::Message as _; + let ms = wa::MessageSecretMessage { + version: Some(1), + enc_iv: Some(iv.to_vec()), + enc_payload: Some(payload.to_vec()), + }; + let mut out = Vec::with_capacity(ms.encoded_len()); + ms.encode(&mut out).expect("encode MessageSecretMessage"); + out +} + +async fn collect_event<F>( + client: &Arc<Client>, + collector: Arc<crate::test_utils::TestEventCollector>, + pred: F, + timeout_ms: u64, +) -> Option<Arc<wacore::types::events::Event>> +where + F: Fn(&wacore::types::events::Event) -> bool, +{ + let _ = client; + let mut waited = 0u64; + while waited <= timeout_ms { + for ev in collector.events() { + if pred(&ev) { + return Some(ev); + } + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + waited += 25; + } + None +} + +fn legacy_edit_text(msg: &wa::Message) -> Option<&str> { + msg.protocol_message + .as_ref() + .and_then(|pm| pm.edited_message.as_ref()) + .and_then(|edited| edited.conversation.as_deref()) +} + +fn inner_message_edit(text: &str, next_secret: Option<Vec<u8>>) -> wa::Message { + wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + key: Some(wa::MessageKey { + remote_jid: Some("5511777776666@s.whatsapp.net".to_string()), + from_me: Some(false), + id: Some("PARENT_EDIT".to_string()), + participant: None, + }), + r#type: Some(wa::message::protocol_message::Type::MessageEdit as i32), + edited_message: Some(Box::new(wa::Message { + conversation: Some(text.to_string()), + ..Default::default() + })), + timestamp_ms: Some(1_770_000_000_000), + ..Default::default() + })), + message_context_info: next_secret.map(|secret| wa::MessageContextInfo { + message_secret: Some(secret), + ..Default::default() + }), + ..Default::default() + } +} + +fn encrypted_message_edit( + target_key: wa::MessageKey, + original_sender: &str, + editor: &str, + parent_id: &str, + secret: &[u8], + text: &str, + next_secret: Option<Vec<u8>>, +) -> wa::Message { + let ctx = wacore::message_edit::MessageEditContext { + original_msg_id: parent_id, + original_sender_jid: original_sender, + editor_jid: editor, + }; + let (enc_payload, enc_iv) = wacore::message_edit::encrypt_message_edit( + &inner_message_edit(text, next_secret), + secret, + &ctx, + ) + .expect("test edit encryption"); + + wa::Message { + secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { + target_message_key: Some(target_key), + enc_payload: Some(enc_payload), + enc_iv: Some(enc_iv.to_vec()), + secret_enc_type: Some( + wa::message::secret_encrypted_message::SecretEncType::MessageEdit as i32, + ), + remote_key_id: None, + }), + ..Default::default() + } +} + +#[tokio::test] +async fn secret_encrypted_message_edit_dispatches_legacy_edit() { + let (client, _transport) = capturing_client("secret_edit_dispatch").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "5511777776666@s.whatsapp.net"; + let parent_id = "PARENT_EDIT"; + let edit_id = "EDIT_1"; + let secret = [0x42u8; 32]; + client + .persistence_manager + .backend() + .put_msg_secret(chat, chat, parent_id, &secret) + .await + .unwrap(); + + let info = Arc::new(MessageInfo { + id: edit_id.into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: chat.parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }); + let target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: None, + }; + let msg = encrypted_message_edit(target_key, chat, chat, parent_id, &secret, "edited", None); + + client.dispatch_parsed_message(msg, &info).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == edit_id + && legacy_edit_text(msg.as_ref()) == Some("edited") + && msg.secret_encrypted_message.is_none()) + }, + 500, + ) + .await; + assert!(got.is_some(), "encrypted edit must dispatch as legacy edit"); +} + +/// Regression for #667: an incoming peer edit writes `target_message_key` +/// in the editor's frame (`from_me = true`, no `participant`, even in a +/// group), so the target-key resolver maps the parent author to *us* and +/// misses the secret stored under the real author. The dispatch path must +/// take the author from the envelope sender instead. Fails on the pre-fix +/// code (envelope stays encrypted), passes after it. +#[tokio::test] +async fn secret_encrypted_peer_edit_resolves_sender_from_envelope() { + let (client, _transport) = capturing_client("secret_peer_edit_dispatch").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let group = "123456789012345678@g.us"; + let peer = "5511777776666@s.whatsapp.net"; + let parent_id = "PEER_PARENT"; + let edit_id = "PEER_EDIT"; + let secret = [0x42u8; 32]; + // The parent (peer's own message) was stored under the real author. + client + .persistence_manager + .backend() + .put_msg_secret(group, peer, parent_id, &secret) + .await + .unwrap(); + + let info = Arc::new(MessageInfo { + id: edit_id.into(), + source: crate::types::message::MessageSource { + chat: group.parse().unwrap(), + sender: peer.parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }); + // Editor's frame: from_me = true, no participant, even in a group. + let target_key = wa::MessageKey { + remote_jid: Some(group.to_string()), + from_me: Some(true), + id: Some(parent_id.to_string()), + participant: None, + }; + // HKDF binds the real author (peer) as both original sender and editor. + let msg = encrypted_message_edit(target_key, peer, peer, parent_id, &secret, "edited", None); + + client.dispatch_parsed_message(msg, &info).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == edit_id + && legacy_edit_text(msg.as_ref()) == Some("edited") + && msg.secret_encrypted_message.is_none()) + }, + 500, + ) + .await; + assert!( + got.is_some(), + "incoming peer edit must resolve the author from the envelope and dispatch as legacy edit" + ); +} + +/// Store the parent secret with a known event time, then dispatch a +/// secret-encrypted edit authored `edit_offset` seconds after the parent. +/// Returns whether the decrypted legacy edit was dispatched. +async fn run_secret_edit_with_window(test_id: &str, parent_ts: i64, edit_offset: i64) -> bool { + let (client, _transport) = capturing_client(test_id).await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "5511777776666@s.whatsapp.net"; + let parent_id = "WINDOW_PARENT"; + let edit_id = "WINDOW_EDIT"; + let secret = [0x42u8; 32]; + client + .persistence_manager + .backend() + .put_msg_secrets(vec![wacore::store::traits::MsgSecretEntry { + chat: chat.to_string(), + sender: chat.to_string(), + msg_id: parent_id.to_string(), + secret: secret.to_vec(), + expires_at: 0, + message_ts: parent_ts, + }]) + .await + .unwrap(); + + let info = Arc::new(MessageInfo { + id: edit_id.into(), + timestamp: chrono::DateTime::<chrono::Utc>::from_timestamp(parent_ts + edit_offset, 0) + .unwrap(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: chat.parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }); + let target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: None, + }; + let msg = encrypted_message_edit(target_key, chat, chat, parent_id, &secret, "edited", None); + client.dispatch_parsed_message(msg, &info).await; + + collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == edit_id + && legacy_edit_text(msg.as_ref()) == Some("edited") + && msg.secret_encrypted_message.is_none()) + }, + 500, + ) + .await + .is_some() +} + +#[tokio::test] +async fn secret_edit_within_window_is_applied() { + // Authored 10 min after the parent — inside the 20 min (1200s) window. + assert!( + run_secret_edit_with_window("secret_edit_in_window", 1_700_000_000, 600).await, + "an in-window edit must dispatch as a legacy edit" + ); +} + +#[tokio::test] +async fn secret_edit_outside_window_is_dropped() { + // Authored 30 min after the parent — past the 1200s window, like WA Web's + // ProcessEditProtocolMsgs, so we drop it (raw envelope surfaces instead). + assert!( + !run_secret_edit_with_window("secret_edit_out_window", 1_700_000_000, 1800).await, + "an out-of-window edit must not dispatch a legacy edit" + ); +} + +#[tokio::test] +async fn secret_edit_unknown_parent_ts_is_permissive() { + // parent_ts = 0 (unknown, e.g. resolver-supplied): no window check, so a + // late edit still applies rather than being silently dropped. + assert!( + run_secret_edit_with_window("secret_edit_unknown_ts", 0, 5_000_000).await, + "with an unknown parent timestamp the edit must still apply" + ); +} + +#[tokio::test] +async fn secret_encrypted_edit_decrypts_via_resolver_when_store_empty() { + use crate::cache_config::{CacheConfig, MsgSecretPolicy}; + + struct StaticResolver { + chat: String, + sender: String, + msg_id: String, + secret: [u8; 32], + } + #[async_trait::async_trait] + impl wacore::msg_secret::OriginalMessageResolver for StaticResolver { + async fn resolve_msg_secret( + &self, + chat: &str, + sender: &str, + msg_id: &str, + ) -> Option<[u8; 32]> { + (chat == self.chat && sender == self.sender && msg_id == self.msg_id) + .then_some(self.secret) + } + } + + let chat = "5511777776666@s.whatsapp.net"; + let parent_id = "RESOLVER_PARENT"; + let edit_id = "RESOLVER_EDIT"; + let secret = [0x7Au8; 32]; + + let resolver = Arc::new(StaticResolver { + chat: chat.to_string(), + sender: chat.to_string(), + msg_id: parent_id.to_string(), + secret, + }); + // Disabled persists nothing, so the only path to the secret is the resolver. + let cfg = CacheConfig { + msg_secret_policy: MsgSecretPolicy::Disabled, + original_message_resolver: Some(resolver), + ..Default::default() + }; + let client = crate::test_utils::create_test_client_with_config( + "resolver_edit", + Arc::new(crate::test_utils::MockHttpClient), + cfg, + ) + .await; + seed_test_pn(&client).await; + + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + assert!( + client + .persistence_manager + .backend() + .get_msg_secret(chat, chat, parent_id) + .await + .unwrap() + .is_none(), + "store must be empty under Disabled" + ); + + let info = Arc::new(MessageInfo { + id: edit_id.into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: chat.parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }); + let target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: None, + }; + let msg = encrypted_message_edit( + target_key, + chat, + chat, + parent_id, + &secret, + "edited via resolver", + None, + ); + + client.dispatch_parsed_message(msg, &info).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == edit_id + && legacy_edit_text(msg.as_ref()) == Some("edited via resolver") + && msg.secret_encrypted_message.is_none()) + }, + 500, + ) + .await; + assert!( + got.is_some(), + "edit must decrypt via the resolver when the store is empty" + ); +} + +#[tokio::test] +async fn decrypted_message_edit_recaptures_secret_for_next_edit() { + let (client, _transport) = capturing_client("secret_edit_chain").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "5511777776666@s.whatsapp.net"; + let parent_id = "PARENT_EDIT"; + let first_secret = [0x11u8; 32]; + let second_secret = [0x22u8; 32]; + client + .persistence_manager + .backend() + .put_msg_secret(chat, chat, parent_id, &first_secret) + .await + .unwrap(); + + let target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: None, + }; + let first_info = Arc::new(MessageInfo { + id: "EDIT_CHAIN_1".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: chat.parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }); + let first_msg = encrypted_message_edit( + target_key.clone(), + chat, + chat, + parent_id, + &first_secret, + "first", + Some(second_secret.to_vec()), + ); + client.dispatch_parsed_message(first_msg, &first_info).await; + + let stored = client + .persistence_manager + .backend() + .get_msg_secret(chat, chat, parent_id) + .await + .unwrap(); + assert_eq!(stored.as_deref(), Some(&second_secret[..])); + + let second_info = Arc::new(MessageInfo { + id: "EDIT_CHAIN_2".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: chat.parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }); + let second_msg = encrypted_message_edit( + target_key, + chat, + chat, + parent_id, + &second_secret, + "second", + None, + ); + client + .dispatch_parsed_message(second_msg, &second_info) + .await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == "EDIT_CHAIN_2" + && legacy_edit_text(msg.as_ref()) == Some("second")) + }, + 500, + ) + .await; + assert!(got.is_some(), "second edit must use the re-captured secret"); +} + +#[tokio::test] +async fn secret_encrypted_message_edit_uses_lid_pn_fallback_in_group() { + use wacore::store::traits::LidPnMappingEntry; + + let (client, _transport) = capturing_client("secret_edit_alt_group").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "120363021033254949@g.us"; + let parent_id = "GROUP_PARENT_EDIT"; + let sender_lid = "236395184570386@lid"; + let sender_pn = "5511777776666@s.whatsapp.net"; + let secret = [0x77u8; 32]; + + client + .persistence_manager + .backend() + .put_lid_mapping(&LidPnMappingEntry { + lid: "236395184570386".into(), + phone_number: "5511777776666".into(), + created_at: 0, + updated_at: 0, + learning_source: "test".into(), + }) + .await + .unwrap(); + client + .persistence_manager + .backend() + .put_msg_secret(chat, sender_pn, parent_id, &secret) + .await + .unwrap(); + + let info = Arc::new(MessageInfo { + id: "GROUP_EDIT_1".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: sender_lid.parse().unwrap(), + is_group: true, + addressing_mode: Some(wacore::types::message::AddressingMode::Lid), + ..Default::default() + }, + ..Default::default() + }); + let target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: Some(sender_lid.to_string()), + }; + let msg = encrypted_message_edit( + target_key, + sender_pn, + sender_lid, + parent_id, + &secret, + "group edited", + None, + ); + + client.dispatch_parsed_message(msg, &info).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == "GROUP_EDIT_1" + && legacy_edit_text(msg.as_ref()) == Some("group edited")) + }, + 500, + ) + .await; + assert!( + got.is_some(), + "group edit must decrypt when the stored secret is under PN" + ); +} + +#[tokio::test] +async fn decrypted_message_edit_refreshes_alternate_secret_alias() { + use wacore::store::traits::LidPnMappingEntry; + + let (client, _transport) = capturing_client("secret_edit_alt_refresh").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "120363021033254949@g.us"; + let parent_id = "GROUP_PARENT_EDIT"; + let sender_lid = "236395184570386@lid"; + let sender_pn = "5511777776666@s.whatsapp.net"; + let first_secret = [0x31u8; 32]; + let second_secret = [0x32u8; 32]; + + client + .persistence_manager + .backend() + .put_lid_mapping(&LidPnMappingEntry { + lid: "236395184570386".into(), + phone_number: "5511777776666".into(), + created_at: 0, + updated_at: 0, + learning_source: "test".into(), + }) + .await + .unwrap(); + for sender in [sender_lid, sender_pn] { + client + .persistence_manager + .backend() + .put_msg_secret(chat, sender, parent_id, &first_secret) + .await + .unwrap(); + } + + let first_info = Arc::new(MessageInfo { + id: "GROUP_EDIT_REFRESH_1".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: sender_lid.parse().unwrap(), + is_group: true, + addressing_mode: Some(wacore::types::message::AddressingMode::Lid), + ..Default::default() + }, + ..Default::default() + }); + let first_target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: Some(sender_lid.to_string()), + }; + let first_msg = encrypted_message_edit( + first_target_key, + sender_lid, + sender_lid, + parent_id, + &first_secret, + "first", + Some(second_secret.to_vec()), + ); + client.dispatch_parsed_message(first_msg, &first_info).await; + + for sender in [sender_lid, sender_pn] { + let stored = client + .persistence_manager + .backend() + .get_msg_secret(chat, sender, parent_id) + .await + .unwrap(); + assert_eq!(stored.as_deref(), Some(&second_secret[..])); + } + + let second_info = Arc::new(MessageInfo { + id: "GROUP_EDIT_REFRESH_2".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: sender_pn.parse().unwrap(), + is_group: true, + addressing_mode: Some(wacore::types::message::AddressingMode::Pn), + ..Default::default() + }, + ..Default::default() + }); + let second_target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: Some(sender_pn.to_string()), + }; + let second_msg = encrypted_message_edit( + second_target_key, + sender_pn, + sender_pn, + parent_id, + &second_secret, + "second", + None, + ); + client + .dispatch_parsed_message(second_msg, &second_info) + .await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == "GROUP_EDIT_REFRESH_2" + && legacy_edit_text(msg.as_ref()) == Some("second")) + }, + 500, + ) + .await; + assert!( + got.is_some(), + "chained edit must use the refreshed alternate alias" + ); +} + +/// Round-trip: store an outbound messageSecret, build a fake bot reply +/// whose payload we encrypt with the symmetric helper, route it through +/// classify, and assert the decrypted `wa::Message` lands on the bus. +#[tokio::test] +async fn msmsg_decrypts_when_secret_is_stored() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + let (client, _transport) = capturing_client("msmsg_ok").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "867051314767696@bot"; + let our_pn = "5511000000001@s.whatsapp.net"; + let bot_jid = "867051314767696@bot"; + let outbound_id = "OUTBOUND_1"; + let bot_reply_id = "BOT_REPLY_1"; + let secret = [0x42u8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + let plaintext_msg = wa::Message { + conversation: Some("hi from bot".to_string()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + let ctx = BotMessageContext { + msg_id: bot_reply_id, + target_sender_user_jid: our_pn, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + let node = NodeBuilder::new("message") + .attr("from", bot_jid) + .attr("id", bot_reply_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == bot_reply_id + && msg.conversation.as_deref() == Some("hi from bot")) + }, + 1500, + ) + .await; + assert!( + got.is_some(), + "msmsg decryption + dispatch must surface Event::Message" + ); +} + +/// No secret stored for `target_id` → nack `error=495`, no Message event. +#[tokio::test] +async fn msmsg_without_stored_secret_nacks_495() { + let (client, transport) = capturing_client("msmsg_nosecret").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let bot_reply_id = "BOT_REPLY_NS"; + let outbound_id = "OUTBOUND_NS"; + let our_pn = "5511000000001@s.whatsapp.net"; + + let ms_msg_proto = encode_message_secret_message(&[0u8; 12], &[0u8; 32]); + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", bot_reply_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let mut code = None; + for _ in 0..80 { + if let Some(c) = find_message_nack_error(&transport.sent(), bot_reply_id) { + code = Some(c); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!( + code, + Some(495), + "missing messageSecret must nack with code 495" + ); + assert!( + collector + .events() + .iter() + .all(|e| !matches!(e.as_ref(), wacore::types::events::Event::Message(_, info) if info.id == bot_reply_id)), + "no Message event must be dispatched when decryption failed" + ); +} + +/// Tampered ciphertext → GCM tag fails → nack 495. +#[tokio::test] +async fn msmsg_with_bad_tag_nacks_495() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + let (client, transport) = capturing_client("msmsg_bad_tag").await; + + let chat = "867051314767696@bot"; + let our_pn = "5511000000001@s.whatsapp.net"; + let outbound_id = "OUTBOUND_BAD"; + let bot_reply_id = "BOT_REPLY_BAD"; + let secret = [0x77u8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + let ctx = BotMessageContext { + msg_id: bot_reply_id, + target_sender_user_jid: our_pn, + bot_user_jid: chat, + }; + let (mut cipher, iv) = encrypt_bot_message(b"hello", &secret, &ctx).unwrap(); + let last = cipher.len() - 1; + cipher[last] ^= 0x01; + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", bot_reply_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let mut code = None; + for _ in 0..80 { + if let Some(c) = find_message_nack_error(&transport.sent(), bot_reply_id) { + code = Some(c); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!(code, Some(495)); +} + +/// Bot edit chain: when `<bot edit="inner">` is set, the HKDF msg_id used +/// for the per-message key swaps to `edit_target_id` so the edited reply +/// decrypts under the same key as the original (whatsmeow / WA Web +/// `decryptMsmsgFbidBotMessage`). +#[tokio::test] +async fn msmsg_bot_edit_uses_edit_target_id_for_hkdf() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + let (client, _transport) = capturing_client("msmsg_edit").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "867051314767696@bot"; + let our_pn = "5511000000001@s.whatsapp.net"; + let outbound_id = "OUTBOUND_EDIT"; + let original_reply_id = "BOT_REPLY_FIRST"; + let edit_reply_id = "BOT_REPLY_EDIT"; + let secret = [0xAAu8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + // Encrypt as if it's the ORIGINAL reply (msg_id = original_reply_id). + let plaintext_msg = wa::Message { + conversation: Some("edited content".to_string()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + let ctx = BotMessageContext { + msg_id: original_reply_id, + target_sender_user_jid: our_pn, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + // Inbound stanza has id=edit_reply_id but <bot edit="inner" edit_target_id=original> + // so the HKDF must derive against original_reply_id. + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", edit_reply_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("bot") + .attr("edit", "inner") + .attr("edit_target_id", original_reply_id) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == edit_reply_id + && msg.conversation.as_deref() == Some("edited content")) + }, + 1500, + ) + .await; + assert!( + got.is_some(), + "bot edit must use edit_target_id for HKDF msg_id" + ); +} + +/// Same setup as the edit test but WITHOUT `<bot edit>`: the HKDF must +/// fall back to `info.id`, and ciphertext encrypted with the edit-target +/// id must fail to decrypt. +#[tokio::test] +async fn msmsg_without_bot_edit_does_not_swap_msg_id() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + let (client, transport) = capturing_client("msmsg_noedit").await; + let chat = "867051314767696@bot"; + let our_pn = "5511000000001@s.whatsapp.net"; + let outbound_id = "OUTBOUND_NOEDIT"; + let stanza_id = "BOT_REPLY_NOEDIT"; + let other_id = "OTHER_ID"; + let secret = [0xBBu8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + // Encrypt with `other_id` to simulate the wrong key derivation if the + // edit branch were taken without `<bot edit>`. + let ctx = BotMessageContext { + msg_id: other_id, + target_sender_user_jid: our_pn, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(b"x", &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", stanza_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let mut code = None; + for _ in 0..80 { + if let Some(c) = find_message_nack_error(&transport.sent(), stanza_id) { + code = Some(c); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!( + code, + Some(495), + "without <bot edit>, HKDF must use stanza id (not OTHER_ID) → tag fails" + ); +} + +/// `<bot edit="first">` is NOT one of {INNER, LAST}, so the HKDF msg_id +/// must remain `info.id`. +#[tokio::test] +async fn msmsg_bot_edit_first_keeps_info_id() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + let (client, _transport) = capturing_client("msmsg_first").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "867051314767696@bot"; + let our_pn = "5511000000001@s.whatsapp.net"; + let outbound_id = "OUTBOUND_FIRST"; + let stanza_id = "BOT_REPLY_FIRST"; + let secret = [0xCCu8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + // Encrypt with stanza_id; "first" edit must NOT swap. + let plaintext_msg = wa::Message { + conversation: Some("first reply".to_string()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + let ctx = BotMessageContext { + msg_id: stanza_id, + target_sender_user_jid: our_pn, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", stanza_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("bot").attr("edit", "first").build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| matches!(e, wacore::types::events::Event::Message(_, info) if info.id == stanza_id), + 1500, + ) + .await; + assert!(got.is_some(), "edit=first must keep info.id as HKDF msg_id"); +} + +/// Regular bot path (`f()` in WA Web `BotMessageSecret.js`): when the +/// fbid pre-resolve picks the WRONG id (e.g. edit_target_id) but the +/// real ciphertext was minted under `info.id`, the fallback attempt +/// must succeed. Validates the try-then-fallback unification. +#[tokio::test] +async fn msmsg_falls_back_to_info_id_when_primary_uses_edit_target() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + let (client, _transport) = capturing_client("msmsg_fb_to_info").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "867051314767696@bot"; + let our_pn = "5511000000001@s.whatsapp.net"; + let outbound_id = "OUT_FB1"; + let stanza_id = "REPLY_FB1"; + let edit_target_id = "WRONG_EDIT_TARGET"; + let secret = [0xDDu8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + // Encrypt under `stanza_id` even though the stanza will declare + // edit=inner with edit_target_id (forces a primary-attempt mismatch). + let plaintext_msg = wa::Message { + conversation: Some("fallback ok".to_string()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + let ctx = BotMessageContext { + msg_id: stanza_id, + target_sender_user_jid: our_pn, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", stanza_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("bot") + .attr("edit", "inner") + .attr("edit_target_id", edit_target_id) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == stanza_id + && msg.conversation.as_deref() == Some("fallback ok")) + }, + 1500, + ) + .await; + assert!( + got.is_some(), + "primary attempt with edit_target_id must fall back to info.id" + ); +} + +/// Inverse of the `falls_back_to_info_id` test: primary is `info.id` +/// (edit_type isn't INNER/LAST so the fbid pre-resolve picks the stanza +/// id), but the bot encrypted under `edit_target_id`. The fallback must +/// rescue. +#[tokio::test] +async fn msmsg_falls_back_to_edit_target_when_primary_uses_info_id() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + let (client, _transport) = capturing_client("msmsg_fb_to_edit").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "867051314767696@bot"; + let our_pn = "5511000000001@s.whatsapp.net"; + let outbound_id = "OUT_INV"; + let stanza_id = "REPLY_INV"; + let edit_target_id = "EDIT_INV"; + let secret = [0xBEu8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + let plaintext_msg = wa::Message { + conversation: Some("inverse fallback".to_string()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + // Encrypt under edit_target_id even though edit=first → primary + // will pick info.id (stanza id), fail, and the fallback should try + // edit_target_id and succeed (WA Web regular bot path `f()`). + let ctx = BotMessageContext { + msg_id: edit_target_id, + target_sender_user_jid: our_pn, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", stanza_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("bot") + .attr("edit", "first") + .attr("edit_target_id", edit_target_id) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == stanza_id + && msg.conversation.as_deref() == Some("inverse fallback")) + }, + 1500, + ) + .await; + assert!( + got.is_some(), + "primary attempt with info.id must fall back to edit_target_id (WA Web f())" + ); +} + +/// Mirror scenario: no `<bot edit>`, so the parser doesn't populate +/// `edit_target_id`. Primary uses `info.id`; with no fallback id +/// available, a deliberately-wrong-key payload must nack 495 (no second +/// attempt to silently mask the failure). +#[tokio::test] +async fn msmsg_no_fallback_when_no_edit_target_present() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + let (client, transport) = capturing_client("msmsg_nofb").await; + let chat = "867051314767696@bot"; + let our_pn = "5511000000001@s.whatsapp.net"; + let outbound_id = "OUT_NOFB"; + let stanza_id = "REPLY_NOFB"; + let secret = [0xCCu8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + // Encrypt under a DIFFERENT id; no <bot> node so parser leaves + // edit_target_id = None and there's nothing to fall back to. + let ctx = BotMessageContext { + msg_id: "MISMATCHED_ID", + target_sender_user_jid: our_pn, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(b"x", &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", stanza_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let mut code = None; + for _ in 0..80 { + if let Some(c) = find_message_nack_error(&transport.sent(), stanza_id) { + code = Some(c); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!( + code, + Some(495), + "no fallback id available → single AES-GCM failure must nack 495" + ); +} + +/// WA Web `processRenderableMessages` captures the embedded +/// `messageSecret` from any bot-targeted msg (fanout from us OR reply +/// from the bot). Verify the helper persists it under +/// (bot_chat, our_lid, info.id). +#[tokio::test] +async fn maybe_capture_inbound_msg_secret_persists_for_bot_chats() { + use crate::store::commands::DeviceCommand; + let (client, _transport) = capturing_client("capture_bot").await; + client + .persistence_manager + .process_command(DeviceCommand::SetLid(Some( + "999888777666555:0@lid".parse().unwrap(), + ))) + .await; + + let info = Arc::new(MessageInfo { + id: "FANOUT_1".into(), + source: crate::types::message::MessageSource { + chat: "867051314767696@bot".parse().unwrap(), + sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), + is_from_me: true, + ..Default::default() + }, + ..Default::default() + }); + let msg = wa::Message { + conversation: Some("hi bot".into()), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0xAB; 32]), + ..Default::default() + }), + ..Default::default() + }; + client.maybe_capture_inbound_msg_secret(&msg, &info).await; + + let mut got = None; + for _ in 0..40 { + got = client + .persistence_manager + .backend() + .get_msg_secret("867051314767696@bot", "999888777666555@lid", "FANOUT_1") + .await + .unwrap(); + if got.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!(got.as_deref(), Some(&[0xABu8; 32][..])); +} + +#[tokio::test] +async fn maybe_capture_inbound_msg_secret_persists_for_non_bot_chats() { + let (client, _transport) = capturing_client("capture_regular_dm").await; + let info = Arc::new(MessageInfo { + id: "DM_1".into(), + source: crate::types::message::MessageSource { + chat: "5511777776666@s.whatsapp.net".parse().unwrap(), + sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), + is_from_me: true, + ..Default::default() + }, + ..Default::default() + }); + let msg = wa::Message { + conversation: Some("hi".into()), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0xCD; 32]), + ..Default::default() + }), + ..Default::default() + }; + client.maybe_capture_inbound_msg_secret(&msg, &info).await; + + let got = client + .persistence_manager + .backend() + .get_msg_secret( + "5511777776666@s.whatsapp.net", + "5511000000001@s.whatsapp.net", + "DM_1", + ) + .await + .unwrap(); + assert_eq!(got.as_deref(), Some(&[0xCDu8; 32][..])); +} + +/// Group invocation: user mentions @MetaAI in a group → chat is the +/// GROUP (not bot), but mentioned_jid contains the bot. WA Web's +/// `processRenderableMessages` keys off `N` (invokedBotWid derived from +/// `mentionedJidList.find(isBot)`); we must persist too. +#[tokio::test] +async fn maybe_capture_inbound_msg_secret_persists_for_group_with_bot_mention() { + use crate::store::commands::DeviceCommand; + let (client, _transport) = capturing_client("capture_group_mention").await; + client + .persistence_manager + .process_command(DeviceCommand::SetLid(Some( + "999888777666555:0@lid".parse().unwrap(), + ))) + .await; + + let info = Arc::new(MessageInfo { + id: "GRP_MENTION".into(), + source: crate::types::message::MessageSource { + chat: "120363021033254949@g.us".parse().unwrap(), + sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), + is_from_me: true, + is_group: true, + ..Default::default() + }, + ..Default::default() + }); + let msg = wa::Message { + extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { + text: Some("hey @MetaAI tell me a joke".into()), + context_info: Some(Box::new(wa::ContextInfo { + mentioned_jid: vec!["867051314767696@bot".into()], + ..Default::default() + })), + ..Default::default() + })), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0xEE; 32]), + ..Default::default() + }), + ..Default::default() + }; + client.maybe_capture_inbound_msg_secret(&msg, &info).await; + + let mut got = None; + for _ in 0..40 { + got = client + .persistence_manager + .backend() + .get_msg_secret( + "120363021033254949@g.us", + "5511000000001@s.whatsapp.net", + "GRP_MENTION", + ) + .await + .unwrap(); + if got.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!( + got.as_deref(), + Some(&[0xEEu8; 32][..]), + "group invocation via @bot mention must still cache the secret" + ); +} + +/// Forwarded message with a secret must NOT be cached — matches WA Web's +/// `x.isForwarded !== true` guard. A planted forward shouldn't poison +/// the cache. +#[tokio::test] +async fn maybe_capture_inbound_msg_secret_skips_forwarded() { + let (client, _transport) = capturing_client("capture_skip_forwarded").await; + let info = Arc::new(MessageInfo { + id: "FWD_1".into(), + source: crate::types::message::MessageSource { + chat: "867051314767696@bot".parse().unwrap(), + sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), + is_from_me: false, + ..Default::default() + }, + ..Default::default() + }); + let msg = wa::Message { + extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { + text: Some("forwarded".into()), + context_info: Some(Box::new(wa::ContextInfo { + is_forwarded: Some(true), + ..Default::default() + })), + ..Default::default() + })), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0xFF; 32]), + ..Default::default() + }), + ..Default::default() + }; + client.maybe_capture_inbound_msg_secret(&msg, &info).await; + + for _ in 0..16 { + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let got = client + .persistence_manager + .backend() + .get_msg_secret( + "867051314767696@bot", + "5511000000001@s.whatsapp.net", + "FWD_1", + ) + .await + .unwrap(); + assert!(got.is_none(), "forwarded messages must not seed the cache"); +} + +/// Our own group bot prompt carries the secret but NO mentioned_jid +/// (observed in prod: `mentions_bot=false mentioned_jids=[]`). The bot +/// invocation is signalled by `message_context_info.bot_metadata`, which +/// must let the capture fire (WA Web's `w`/`A` group-participant gate). +#[tokio::test] +async fn maybe_capture_inbound_msg_secret_via_bot_metadata_without_mention() { + let (client, _transport) = capturing_client("capture_bot_meta").await; + client + .persistence_manager + .process_command(crate::store::commands::DeviceCommand::SetLid(Some( + "999888777666555:0@lid".parse().unwrap(), + ))) + .await; + let info = Arc::new(MessageInfo { + id: "GRP_OWN_BOT".into(), + source: crate::types::message::MessageSource { + chat: "120363021033254949@g.us".parse().unwrap(), + sender: "236395184570386:0@lid".parse().unwrap(), + is_from_me: true, + is_group: true, + ..Default::default() + }, + ..Default::default() + }); + let msg = wa::Message { + extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { + text: Some("continue".into()), + // No mention at all — just bot_metadata signals the invocation. + ..Default::default() + })), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0x7B; 32]), + bot_metadata: Some(wa::BotMetadata { + persona_id: Some("867051314767696".into()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + client.maybe_capture_inbound_msg_secret(&msg, &info).await; + + // Group (non-bot chat) → keyed under info.source.sender (our LID in a + // LID group), which is what the bot reply's target_sender_jid echoes. + let got = client + .persistence_manager + .backend() + .get_msg_secret( + "120363021033254949@g.us", + "236395184570386@lid", + "GRP_OWN_BOT", + ) + .await + .unwrap(); + assert_eq!( + got.as_deref(), + Some(&[0x7Bu8; 32][..]), + "bot_metadata presence must let our own group prompt cache without a mention" + ); +} + +#[tokio::test] +async fn bot_only_captures_group_bot_prompt_skips_plain() { + use crate::cache_config::{CacheConfig, MsgSecretPolicy}; + let cfg = CacheConfig { + msg_secret_policy: MsgSecretPolicy::BotOnly, + ..Default::default() + }; + let client = crate::test_utils::create_test_client_with_config( + "botonly_capture", + Arc::new(crate::test_utils::MockHttpClient), + cfg, + ) + .await; + + let group = "120363021033254949@g.us"; + let sender = "5511888887777@s.whatsapp.net"; + + // A plain group message is not a bot context → skipped under BotOnly. + let plain_info = Arc::new(MessageInfo { + id: "PLAIN".into(), + source: crate::types::message::MessageSource { + chat: group.parse().unwrap(), + sender: sender.parse().unwrap(), + is_group: true, + ..Default::default() + }, + ..Default::default() + }); + let plain_msg = wa::Message { + conversation: Some("hi".into()), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0x01; 32]), + ..Default::default() + }), + ..Default::default() + }; + client + .maybe_capture_inbound_msg_secret(&plain_msg, &plain_info) + .await; + assert!( + client + .persistence_manager + .backend() + .get_msg_secret(group, sender, "PLAIN") + .await + .unwrap() + .is_none(), + "BotOnly must skip a plain (non-bot) group message" + ); + + // A group message that invokes a bot (bot_metadata) classifies as Bot, + // so its secret is kept and the later bot reply can decrypt. + let bot_info = Arc::new(MessageInfo { + id: "BOTP".into(), + source: crate::types::message::MessageSource { + chat: group.parse().unwrap(), + sender: sender.parse().unwrap(), + is_group: true, + ..Default::default() + }, + ..Default::default() + }); + let bot_msg = wa::Message { + extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { + text: Some("continue".into()), + ..Default::default() + })), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0x02; 32]), + bot_metadata: Some(wa::BotMetadata { + persona_id: Some("867051314767696".into()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + client + .maybe_capture_inbound_msg_secret(&bot_msg, &bot_info) + .await; + assert_eq!( + client + .persistence_manager + .backend() + .get_msg_secret(group, sender, "BOTP") + .await + .unwrap(), + Some(vec![0x02; 32]), + "BotOnly must capture a group bot invocation" + ); +} + +/// Group flow: another participant invokes the bot, their decrypted prompt +/// carries the secret. We must key it under THE PARTICIPANT (the future +/// reply's `<meta target_sender_jid>`), not our own identity. +#[tokio::test] +async fn maybe_capture_inbound_msg_secret_keys_under_other_participant() { + let (client, _transport) = capturing_client("capture_participant").await; + let participant = "5599111112222:7@s.whatsapp.net"; + let info = Arc::new(MessageInfo { + id: "GRP_OTHER".into(), + source: crate::types::message::MessageSource { + chat: "120363021033254949@g.us".parse().unwrap(), + sender: participant.parse().unwrap(), + is_from_me: false, + is_group: true, + ..Default::default() + }, + ..Default::default() + }); + let msg = wa::Message { + extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { + text: Some("@MetaAI question".into()), + context_info: Some(Box::new(wa::ContextInfo { + mentioned_jid: vec!["867051314767696@bot".into()], + ..Default::default() + })), + ..Default::default() + })), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0x5A; 32]), + ..Default::default() + }), + ..Default::default() + }; + client.maybe_capture_inbound_msg_secret(&msg, &info).await; + + // Keyed under the participant (non-AD), NOT under our own PN/LID. + let under_participant = client + .persistence_manager + .backend() + .get_msg_secret( + "120363021033254949@g.us", + "5599111112222@s.whatsapp.net", + "GRP_OTHER", + ) + .await + .unwrap(); + assert_eq!( + under_participant.as_deref(), + Some(&[0x5Au8; 32][..]), + "another participant's prompt must key under their sender JID" + ); +} + +/// WA Web `sendAggregateReceipts`: a bot reply in a GROUP (chat not bot, +/// author is bot) must ack with a bare `<ack class="message">` +/// (sendBotInvokeResponseAcks), NOT a `<receipt>`. +#[tokio::test] +async fn bot_reply_in_group_acks_with_bare_ack_not_receipt() { + let (client, transport) = capturing_client("bot_group_ack").await; + let info = Arc::new(MessageInfo { + id: "BOT_GRP_ACK".into(), + source: crate::types::message::MessageSource { + chat: "120363021033254949@g.us".parse().unwrap(), + sender: "867051314767696@bot".parse().unwrap(), + is_from_me: false, + is_group: true, + ..Default::default() + }, + ..Default::default() + }); + client.ack_received_message(&info); + + let mut found = None; + for _ in 0..80 { + if let Some(a) = find_message_ack(&transport.sent()) { + found = Some(a); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert!( + found.is_some(), + "group bot reply must emit a bare <ack class=\"message\">" + ); + assert_eq!( + delivery_receipts_for(&transport.sent(), "BOT_GRP_ACK"), + 0, + "group bot reply must NOT emit a <receipt>" + ); +} + +/// Regression: a 1:1 bot chat (chat IS the bot) keeps the normal delivery +/// `<receipt>` — WA Web's `v` gate is false when chat.isBot(). +#[tokio::test] +async fn bot_dm_reply_keeps_delivery_receipt() { + let (client, transport) = capturing_client("bot_dm_receipt").await; + let info = Arc::new(MessageInfo { + id: "BOT_DM_RCPT".into(), + source: crate::types::message::MessageSource { + chat: "867051314767696@bot".parse().unwrap(), + sender: "867051314767696@bot".parse().unwrap(), + is_from_me: false, + ..Default::default() + }, + ..Default::default() + }); + client.ack_received_message(&info); + + let mut count = 0; + for _ in 0..80 { + count = delivery_receipts_for(&transport.sent(), "BOT_DM_RCPT"); + if count > 0 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!( + count, 1, + "1:1 bot chat must keep the normal delivery receipt" + ); +} + +#[tokio::test] +async fn maybe_capture_inbound_msg_secret_skips_when_secret_absent() { + let (client, _transport) = capturing_client("capture_no_secret").await; + let info = Arc::new(MessageInfo { + id: "NO_SECRET".into(), + source: crate::types::message::MessageSource { + chat: "867051314767696@bot".parse().unwrap(), + sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), + is_from_me: true, + ..Default::default() + }, + ..Default::default() + }); + let msg = wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }; + client.maybe_capture_inbound_msg_secret(&msg, &info).await; + + for _ in 0..16 { + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + let got = client + .persistence_manager + .backend() + .get_msg_secret( + "867051314767696@bot", + "5511000000001@s.whatsapp.net", + "NO_SECRET", + ) + .await + .unwrap(); + assert!(got.is_none()); +} + +/// A stanza carrying BOTH a valid msmsg AND an unknown sibling enc must +/// still dispatch the msmsg — the unknown-only fallback ack must not +/// short-circuit when `bot_payloads` is non-empty. +#[tokio::test] +async fn mixed_msmsg_and_unknown_enc_still_decrypts_msmsg() { + use crate::store::commands::DeviceCommand; + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + + let (client, _transport) = capturing_client("msmsg_mixed_unknown").await; + client + .persistence_manager + .process_command(DeviceCommand::SetLid(Some( + "999888777666555:0@lid".parse().unwrap(), + ))) + .await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "867051314767696@bot"; + let our_lid = "999888777666555@lid"; + let outbound_id = "OUT_MIX"; + let bot_reply_id = "REPLY_MIX"; + let secret = [0x55u8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_lid, outbound_id, &secret) + .await + .unwrap(); + + let plaintext_msg = wa::Message { + conversation: Some("mixed ok".into()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + let ctx = BotMessageContext { + msg_id: bot_reply_id, + target_sender_user_jid: our_lid, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + // Stanza has a valid msmsg PLUS an unrecognised "frskmsg" sibling. + // The fallback transport-ack must NOT fire (would drop the msmsg). + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", bot_reply_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_lid) + .build(), + NodeBuilder::new("enc") + .attr("type", "frskmsg") + .bytes(vec![0u8; 8]) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == bot_reply_id + && msg.conversation.as_deref() == Some("mixed ok")) + }, + 1500, + ) + .await; + assert!( + got.is_some(), + "msmsg sibling of an unknown enc must still decrypt and dispatch" + ); +} + +/// LID↔PN migration window: the secret was stored under our PN, but the +/// bot reply's `<meta target_sender_jid>` echoes our LID. The primary +/// lookup misses; `alternate_msg_secret_lookup` resolves PN via +/// `lid_pn_mapping` and hits. Mirrors WA Web `C()`'s `getAlternateMsgKey`. +#[tokio::test] +async fn msmsg_alternate_lookup_resolves_lid_to_stored_pn() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + use wacore::store::traits::LidPnMappingEntry; + + let (client, _transport) = capturing_client("msmsg_alt_lookup").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "867051314767696@bot"; + let our_lid_user = "999888777666555"; + let our_pn_user = "5511000000001"; + let our_lid = "999888777666555@lid"; + let our_pn = "5511000000001@s.whatsapp.net"; + let outbound_id = "OUT_ALT"; + let bot_reply_id = "REPLY_ALT"; + let secret = [0x3Cu8; 32]; + + // Seed the LID→PN mapping so the alternate lookup can swap. + client + .persistence_manager + .backend() + .put_lid_mapping(&LidPnMappingEntry { + lid: our_lid_user.into(), + phone_number: our_pn_user.into(), + created_at: 0, + updated_at: 0, + learning_source: "test".into(), + }) + .await + .unwrap(); + // Secret stored under PN (as if the outbound went out PN-addressed). + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + // Bot reply encrypts with target = our LID (what <meta> declares). + let plaintext_msg = wa::Message { + conversation: Some("alt ok".into()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + let ctx = BotMessageContext { + msg_id: bot_reply_id, + target_sender_user_jid: our_lid, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", bot_reply_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_lid) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == bot_reply_id + && msg.conversation.as_deref() == Some("alt ok")) + }, + 1500, + ) + .await; + assert!( + got.is_some(), + "LID-declared reply must resolve the PN-stored secret via lid_pn_mapping" + ); +} + +/// End-to-end: phone fanout dispatches a wa::Message carrying the +/// outbound `messageSecret`; later the Meta AI bot replies via msmsg +/// referencing the same id. The captured secret must let the reply +/// decrypt and surface `Event::Message`. +#[tokio::test] +async fn fanout_capture_lets_subsequent_msmsg_decrypt() { + use crate::store::commands::DeviceCommand; + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + + let (client, _transport) = capturing_client("fanout_to_msmsg").await; + client + .persistence_manager + .process_command(DeviceCommand::SetLid(Some( + "999888777666555:0@lid".parse().unwrap(), + ))) + .await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let bot_chat: Jid = "867051314767696@bot".parse().unwrap(); + let our_lid_str = "999888777666555@lid"; + let outbound_id = "FANOUT_OUT"; + let bot_reply_id = "BOT_REPLY_PHONE"; + let secret = [0x99u8; 32]; + + // Step 1: simulate the fanout dispatch (what dispatch_parsed_message + // would call when the phone's outbound stanza is mirrored to us). + let fanout_info = Arc::new(MessageInfo { + id: outbound_id.into(), + source: crate::types::message::MessageSource { + chat: bot_chat.clone(), + sender: "5511000000001:0@s.whatsapp.net".parse().unwrap(), + is_from_me: true, + ..Default::default() + }, + ..Default::default() + }); + let fanout_msg = wa::Message { + conversation: Some("hi bot".into()), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(secret.to_vec()), + ..Default::default() + }), + ..Default::default() + }; + client + .maybe_capture_inbound_msg_secret(&fanout_msg, &fanout_info) + .await; + // Write is awaited inline now, so the secret is already durable here. + for _ in 0..40 { + if client + .persistence_manager + .backend() + .get_msg_secret("867051314767696@bot", our_lid_str, outbound_id) + .await + .unwrap() + .is_some() + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + + // Step 2: the bot reply arrives as <enc type="msmsg"> referencing + // outbound_id via <meta target_id>. With the secret captured above, + // it must decrypt cleanly. + let plaintext_msg = wa::Message { + conversation: Some("bot reply".into()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + let ctx = BotMessageContext { + msg_id: bot_reply_id, + target_sender_user_jid: our_lid_str, + bot_user_jid: "867051314767696@bot", + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", bot_reply_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_lid_str) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == bot_reply_id + && msg.conversation.as_deref() == Some("bot reply")) + }, + 1500, + ) + .await; + assert!( + got.is_some(), + "secret captured from fanout must enable msmsg reply decryption" + ); +} + +/// Coherence: the identity `persist_outbound_msg_secret` writes under +/// (LID for bot chats) must match what `handle_msmsg_payload` reads via +/// `<meta target_sender_jid>`. End-to-end without bypassing the helper. +#[tokio::test] +async fn msmsg_outbound_put_and_inbound_get_match_for_lid_bot() { + use crate::store::commands::DeviceCommand; + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + + let (client, _transport) = capturing_client("msmsg_lid_match").await; + // Seed both PN (already seeded by capturing_client) and LID. + client + .persistence_manager + .process_command(DeviceCommand::SetLid(Some( + "999888777666555:0@lid".parse().unwrap(), + ))) + .await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let bot_chat: Jid = "867051314767696@bot".parse().unwrap(); + let outbound_id = "OUT_LID"; + let bot_reply_id = "REPLY_LID"; + let our_lid = "999888777666555@lid"; + let secret = [0x71u8; 32]; + + // Real outbound path: caller resolves the bot identity to our LID. + let sender_identity = client + .dm_sender_identity_for(&bot_chat) + .await + .expect("LID seeded"); + client + .persist_outbound_msg_secret( + &bot_chat, + &sender_identity, + outbound_id, + &secret, + wacore::msg_secret::RetentionClass::Bot, + ) + .await; + + // Inbound msmsg payload encrypted under the same (msg_id, target, bot) + // tuple the meta will declare on the wire. + let plaintext_msg = wa::Message { + conversation: Some("lid coherent".to_string()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + let ctx = BotMessageContext { + msg_id: bot_reply_id, + target_sender_user_jid: our_lid, + bot_user_jid: "867051314767696@bot", + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", bot_reply_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_lid) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == bot_reply_id + && msg.conversation.as_deref() == Some("lid coherent")) + }, + 1500, + ) + .await; + assert!( + got.is_some(), + "outbound PUT and inbound GET must converge on LID for bot chats" + ); +} + +/// Regression for the AD_JID encoder bug: a `from="USER:0@bot"` stanza must +/// survive the marshal/unmarshal round-trip with `server=Bot`, so the +/// secret lookup keys hit and the reply decrypts. +#[tokio::test] +async fn msmsg_with_bot_device_suffix_round_trips() { + use wacore::bot_message::{BotMessageContext, encrypt_bot_message}; + let (client, _transport) = capturing_client("msmsg_bot_device").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "867051314767696@bot"; + let our_pn = "5511000000001@s.whatsapp.net"; + let outbound_id = "OUTBOUND_DEV"; + let bot_reply_id = "BOT_REPLY_DEV"; + let secret = [0x33u8; 32]; + + client + .persistence_manager + .backend() + .put_msg_secret(chat, our_pn, outbound_id, &secret) + .await + .unwrap(); + + let plaintext_msg = wa::Message { + conversation: Some("with device".to_string()), + ..Default::default() + }; + let pt_bytes = { + use prost::Message as _; + let mut v = Vec::with_capacity(plaintext_msg.encoded_len()); + plaintext_msg.encode(&mut v).unwrap(); + v + }; + let ctx = BotMessageContext { + msg_id: bot_reply_id, + target_sender_user_jid: our_pn, + bot_user_jid: chat, + }; + let (cipher, iv) = encrypt_bot_message(&pt_bytes, &secret, &ctx).unwrap(); + let ms_msg_proto = encode_message_secret_message(&iv, &cipher); + + let node = NodeBuilder::new("message") + .attr("from", "867051314767696:0@bot") + .attr("id", bot_reply_id) + .attr("type", "text") + .children([ + NodeBuilder::new("meta") + .attr("target_id", outbound_id) + .attr("target_sender_jid", our_pn) + .build(), + NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build(), + ]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == bot_reply_id + && msg.conversation.as_deref() == Some("with device")) + }, + 1500, + ) + .await; + assert!( + got.is_some(), + "msmsg with `:0@bot` from must round-trip (encoder must not strip the bot server)" + ); +} + +/// `<meta>` without `target_id` → cannot identify the parent message, +/// nack 495 and no dispatch. +#[tokio::test] +async fn msmsg_without_meta_target_id_nacks_495() { + let (client, transport) = capturing_client("msmsg_no_target").await; + let bot_reply_id = "BOT_REPLY_NT"; + let ms_msg_proto = encode_message_secret_message(&[0u8; 12], &[0u8; 32]); + let node = NodeBuilder::new("message") + .attr("from", "867051314767696@bot") + .attr("id", bot_reply_id) + .attr("type", "text") + .children([NodeBuilder::new("enc") + .attr("type", "msmsg") + .attr("v", "2") + .bytes(ms_msg_proto) + .build()]) + .build(); + let owned = node_to_arc(node); + client.clone().handle_incoming_message(owned).await; + + let mut code = None; + for _ in 0..80 { + if let Some(c) = find_message_nack_error(&transport.sent(), bot_reply_id) { + code = Some(c); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + } + assert_eq!(code, Some(495)); +} diff --git a/wacore/src/send.rs b/wacore/src/send.rs index bc7c3ec0b..6928b8277 100644 --- a/wacore/src/send.rs +++ b/wacore/src/send.rs @@ -63,5260 +63,22 @@ impl StanzaType { } } -/// Extract (enc_type, is_prekey, serialized) from a CiphertextMessage. -pub fn extract_ciphertext(msg: CiphertextMessage) -> Option<(&'static str, bool, Box<[u8]>)> { - match msg { - CiphertextMessage::SignalMessage(m) => { - Some((stanza::ENC_TYPE_MSG, false, m.into_serialized())) - } - CiphertextMessage::PreKeySignalMessage(m) => { - Some((stanza::ENC_TYPE_PKMSG, true, m.into_serialized())) - } - _ => None, - } -} - -/// Unwrap wrapper message types to reach the inner message. -/// Matches WA Web's getUnwrappedProtobufMessage. Does not unwrap -/// `edited_message`; that field is itself a signal callers may need. -pub(crate) fn unwrap_message(msg: &wa::Message) -> &wa::Message { - macro_rules! try_unwrap { - ($($field:ident),+ $(,)?) => { - $( - if let Some(ref w) = msg.$field { - if let Some(ref inner) = w.message { - return unwrap_message(inner); - } - } - )+ - }; - } - try_unwrap!( - ephemeral_message, - view_once_message, - view_once_message_v2, - view_once_message_v2_extension, - document_with_caption_message, - group_mentioned_message, - bot_invoke_message, - associated_child_message, - poll_creation_option_image_message, - // Remaining FutureProofMessage wrappers from WA Web's - // getUnwrappedProtobufMessage list; classify by the inner message. - event_cover_image, - group_status_message, - group_status_message_v2, - group_status_mention_message, - status_add_yours, - status_mention_message, - question_message, - question_reply_message, - spoiler_message, - lottie_sticker_message, - limit_sharing_message, - newsletter_admin_profile_message, - newsletter_admin_profile_message_v2, - poll_creation_message_v4, - ); - if let Some(ref dsm) = msg.device_sent_message - && let Some(ref inner) = dsm.message - { - return unwrap_message(inner); - } - msg -} - -/// Matches WAWebE2EProtoUtils.typeAttributeFromProtobuf. -pub fn stanza_type_from_message(msg: &wa::Message) -> &'static str { - let msg = unwrap_message(msg); - - if msg.reaction_message.is_some() || msg.enc_reaction_message.is_some() { - return stanza::MSG_TYPE_REACTION; - } - if msg.event_message.is_some() || msg.enc_event_response_message.is_some() { - return stanza::MSG_TYPE_EVENT; - } - if let Some(ref sec) = msg.secret_encrypted_message { - use wa::message::secret_encrypted_message::SecretEncType; - match SecretEncType::try_from(sec.secret_enc_type.unwrap_or(0)) { - Ok(SecretEncType::EventEdit) => return stanza::MSG_TYPE_EVENT, - Ok(SecretEncType::MessageEdit) => return stanza::MSG_TYPE_TEXT, - Ok(SecretEncType::PollEdit | SecretEncType::PollAddOption) => { - return stanza::MSG_TYPE_POLL; - } - _ => {} - } - } - if msg.poll_creation_message.is_some() - || msg.poll_creation_message_v2.is_some() - || msg.poll_creation_message_v3.is_some() - || msg.poll_creation_message_v5.is_some() - || msg.poll_update_message.is_some() - { - return stanza::MSG_TYPE_POLL; - } - if msg.conversation.is_some() - || msg.protocol_message.is_some() - || msg.keep_in_chat_message.is_some() - || msg.edited_message.is_some() - || msg.pin_in_chat_message.is_some() - || msg.interactive_message.is_some() - || msg.template_button_reply_message.is_some() - || msg.request_phone_number_message.is_some() - || msg.enc_comment_message.is_some() - || msg.newsletter_admin_invite_message.is_some() - || msg.newsletter_follower_invite_message_v2.is_some() - || msg.message_history_notice.is_some() - || msg.album_message.is_some() - // Payment family. WA Web's typeAttributeFromProtobuf leaves these at the media - // default, but media-without-mediatype is dropped by the server (so is a bare - // "pay" stanza); text is what delivers and renders on Android. - || msg.request_payment_message.is_some() - || msg.send_payment_message.is_some() - || msg.payment_invite_message.is_some() - || msg.decline_payment_request_message.is_some() - || msg.cancel_payment_request_message.is_some() - { - return stanza::MSG_TYPE_TEXT; - } - // pollResultSnapshotMessage maps to "text" by default in WA Web - // (gated behind isPollResultSnapshotPollTypeEnvelopeEnabled for "poll") - if msg.poll_result_snapshot_message.is_some() || msg.poll_result_snapshot_message_v3.is_some() { - return stanza::MSG_TYPE_TEXT; - } - if let Some(ref ext) = msg.extended_text_message { - if ext - .matched_text - .as_ref() - .is_some_and(|t| !t.trim().is_empty()) - { - return stanza::MSG_TYPE_MEDIA; - } - return stanza::MSG_TYPE_TEXT; - } - stanza::MSG_TYPE_MEDIA -} - -pub fn peer_message_options_from_message(msg: &wa::Message) -> PeerMessageOptions { - use wa::message::PeerDataOperationRequestType as PdoType; - - // WAWebSendNonMessageDataRequest's A/F helpers gate rollout flags we do - // not model; use the default-on wire shape for supported peer PDO flows. - let request_type = unwrap_message(msg) - .protocol_message - .as_deref() - .and_then(|pm| pm.peer_data_operation_request_message.as_ref()) - .and_then(|pdo| pdo.peer_data_operation_request_type) - .and_then(|raw| PdoType::try_from(raw).ok()); - - match request_type { - Some(PdoType::HistorySyncOnDemand) => PeerMessageOptions::high_force_on_demand(), - Some( - PdoType::GenerateLinkPreview - | PdoType::PlaceholderMessageResend - | PdoType::CompanionCanonicalUserNonceFetch, - ) => PeerMessageOptions::high_force(), - _ => PeerMessageOptions::default(), - } -} - -/// Matches WAWebBackendJobsCommon.mediaTypeFromProtobuf + encodeMaybeMediaType. -/// Returns `None` when the attribute should be omitted. -pub fn media_type_from_message(msg: &wa::Message) -> Option<&'static str> { - // WA Web's mediaTypeFromProtobuf treats a top-level lottieStickerMessage as a - // terminal "sticker" and does NOT recurse into it (unlike typeAttributeFromProtobuf, - // which unwraps it via getUnwrappedProtobufMessage). Check before the shared unwrap. - if msg.lottie_sticker_message.is_some() { - return Some("sticker"); - } - - let msg = unwrap_message(msg); - - if msg.image_message.is_some() { - return Some("image"); - } - if let Some(ref vid) = msg.video_message { - return if vid.gif_playback == Some(true) { - Some("gif") - } else { - Some("video") - }; - } - if msg.ptv_message.is_some() { - return Some("ptv"); - } - if let Some(ref audio) = msg.audio_message { - return if audio.ptt == Some(true) { - Some("ptt") - } else { - Some("audio") - }; - } - if msg.document_message.is_some() { - return Some("document"); - } - if msg.sticker_message.is_some() { - return Some("sticker"); - } - if msg.sticker_pack_message.is_some() { - return Some("sticker_pack"); - } - if let Some(ref loc) = msg.location_message { - return if loc.is_live == Some(true) { - Some("livelocation") - } else { - Some("location") - }; - } - if msg.live_location_message.is_some() { - return Some("livelocation"); - } - if msg.contact_message.is_some() { - return Some("vcard"); - } - if msg.contacts_array_message.is_some() { - return Some("contact_array"); - } - if let Some(ref ext) = msg.extended_text_message - && ext - .matched_text - .as_ref() - .is_some_and(|t| !t.trim().is_empty()) - { - return Some("url"); - } - if msg.group_invite_message.is_some() { - return Some("url"); - } - // Interactive / business message families. WA Web's mediaTypeFromProtobuf maps - // each to a concrete mediatype; without it the server drops the type="media" - // stanza. buttonsMessage is intentionally absent: WA Web maps it to - // EncMediaType.Button, which its string mapper drops (no attribute). - if msg.list_message.is_some() { - return Some("list"); - } - if msg.list_response_message.is_some() { - return Some("list_response"); - } - if msg.buttons_response_message.is_some() { - return Some("buttons_response"); - } - if msg.order_message.is_some() { - return Some("order"); - } - if msg.product_message.is_some() { - return Some("product"); - } - if msg.interactive_response_message.is_some() { - return Some("native_flow_response"); - } - if msg.message_history_bundle.is_some() { - return Some("group_history"); - } - None -} - -/// Canonical rule for `decrypt-fail="hide"` on outgoing `<enc>` nodes. -/// Shared by DM fanout, group SKDM and group SKMSG so the three paths can't drift. -/// Both revoke kinds are excluded: WA Web never hides REVOKE, and the server -/// drops revoke stanzas carrying the hide attribute. -pub fn should_hide_decrypt_fail_for_send( - edit: Option<&crate::types::message::EditAttribute>, - msg: &wa::Message, -) -> bool { - use crate::types::message::EditAttribute; - edit.is_some_and(|e| { - *e != EditAttribute::Empty - && *e != EditAttribute::AdminRevoke - && *e != EditAttribute::SenderRevoke - }) || should_hide_decrypt_fail(msg) -} - -/// Infrastructure messages get decrypt-fail="hide" so recipients don't see -/// "waiting for this message" placeholders for things like reactions or pin changes. -pub fn should_hide_decrypt_fail(msg: &wa::Message) -> bool { - let msg = unwrap_message(msg); - - use wa::message::protocol_message::Type as ProtocolType; - use wa::message::secret_encrypted_message::SecretEncType; - - msg.reaction_message.is_some() - || msg.enc_reaction_message.is_some() - || msg.pin_in_chat_message.is_some() - || msg.edited_message.is_some() - || msg.keep_in_chat_message.is_some() - || msg.enc_event_response_message.is_some() - || msg - .poll_update_message - .as_ref() - .is_some_and(|p| p.vote.is_some()) - || msg.message_history_notice.is_some() - || msg.conditional_reveal_message.is_some() - || msg.secret_encrypted_message.as_ref().is_some_and(|s| { - matches!( - SecretEncType::try_from(s.secret_enc_type.unwrap_or(0)), - Ok(SecretEncType::EventEdit - | SecretEncType::PollEdit - | SecretEncType::PollAddOption) - ) - }) - || msg - .bot_invoke_message - .as_ref() - .and_then(|b| b.message.as_ref()) - .and_then(|m| m.protocol_message.as_ref()) - .is_some_and(|p| p.r#type == Some(ProtocolType::RequestWelcomeMessage as i32)) - || msg.protocol_message.as_ref().is_some_and(|p| { - matches!( - p.r#type, - Some(t) if t == ProtocolType::EphemeralSyncResponse as i32 - || t == ProtocolType::RequestWelcomeMessage as i32 - || t == ProtocolType::GroupMemberLabelChange as i32 - ) || p.edited_message.is_some() - }) -} - -/// Caller must hold `SenderKeyStore::sender_key_lock` for `sender_key_name` -/// across the surrounding SKDM creation + this encrypt, so a concurrent send -/// can't split the key between the SKDM and the skmsg. -pub async fn encrypt_group_message<S, R>( - sender_key_store: &mut S, - sender_key_name: &SenderKeyName, - plaintext: &[u8], - csprng: &mut R, -) -> Result<SenderKeyMessage> -where - S: SenderKeyStore + ?Sized, - R: Rng + CryptoRng, -{ - log::debug!( - "Attempting to load sender key for group {} sender {}", - sender_key_name.group_id(), - sender_key_name.sender_id() - ); - - let mut record = sender_key_store - .load_sender_key(sender_key_name) - .await? - .ok_or_else(|| { - SignalProtocolError::NoSenderKeyState(format!( - "no sender key record for group {} sender {}", - sender_key_name.group_id(), - sender_key_name.sender_id() - )) - })?; - - let sender_key_state = record - .sender_key_state_mut() - .map_err(|e| anyhow!("Invalid SenderKey session: {:?}", e))?; - - let sender_chain_key = sender_key_state - .sender_chain_key() - .ok_or_else(|| anyhow!("Invalid SenderKey session: missing chain key"))?; - - let message_keys = sender_chain_key.sender_message_key(); - - let mut ciphertext = Vec::new(); - aes_256_cbc_encrypt_into( - plaintext, - message_keys.cipher_key(), - message_keys.iv(), - &mut ciphertext, - ) - .map_err(|_| anyhow!("AES encryption failed"))?; - - let signing_key = sender_key_state - .signing_key_private() - .map_err(|e| anyhow!("Invalid SenderKey session: missing signing key: {:?}", e))?; - - let skm = SenderKeyMessage::new( - SENDERKEY_MESSAGE_CURRENT_VERSION, - sender_key_state.chain_id(), - message_keys.iteration(), - ciphertext.into_boxed_slice(), - csprng, - &signing_key, - )?; - - sender_key_state.set_sender_chain_key(sender_chain_key.next()?); - - sender_key_store - .store_sender_key(sender_key_name, record) - .await?; - - Ok(skm) -} - -pub struct SignalStores<'a, S, I, P, SP> { - pub sender_key_store: &'a mut (dyn crate::libsignal::protocol::SenderKeyStore + Send + Sync), - pub session_store: &'a mut S, - pub identity_store: &'a mut I, - pub prekey_store: &'a mut P, - pub signed_prekey_store: &'a SP, -} - -/// Check if an anyhow error is a 406 "not-acceptable" server error (device unregistered). -/// Uses typed downcast to `ServerErrorCode` — the shared error type that the -/// `SendContextResolver` impl wraps server errors in. -pub(crate) fn is_device_unregistered_error(err: &anyhow::Error) -> bool { - crate::request::ServerErrorCode::from_anyhow(err).is_some_and(|e| e.code == 406) -} - -pub struct EncryptResult { - pub participant_nodes: Vec<Node>, - pub includes_prekey_message: bool, - pub encrypted_devices: Vec<Jid>, - /// True if any device returned 406 (unregistered) during prekey fetch. - pub had_unregistered_device: bool, -} - -/// Maximum number of concurrent per-device crypto tasks during group send -/// fan-out. Picked from the `perf-audit` benchmark: speedup plateaus around -/// 16 on Oracle ARM64; 32 gives only ~10% more for double the task overhead. -const ENCRYPT_FANOUT_CONCURRENCY: usize = 16; - -/// Per-task encrypt result, shipped from a spawned task back to the orchestrator. -struct EncryptOneResult { - enc_type: &'static str, - is_prekey: bool, - ciphertext: Vec<u8>, - hide_decrypt_fail: bool, -} - -/// Surfaces a spawned task that didn't deliver its result — either the task -/// itself panicked or the runtime tore it down (e.g., during shutdown). -/// Surfacing this as an Err lets the encrypt fan-out fall through to its -/// existing log+skip path instead of propagating a panic. -#[derive(Debug, thiserror::Error)] -#[error("spawned task did not produce a result (panic or runtime shutdown)")] -struct SpawnCanceled; - -/// Future returned by [`spawn_oneshot`]. Holds the spawned task's -/// [`AbortHandle`] until the result is received, so dropping the future mid- -/// flight (e.g., the outer send was cancelled by a timeout) cancels the -/// in-flight crypto work instead of orphaning it. -struct Spawned<T> { - rx: futures::channel::oneshot::Receiver<T>, - abort: Option<AbortHandle>, -} - -impl<T> Future for Spawned<T> { - type Output = std::result::Result<T, SpawnCanceled>; - - fn poll( - mut self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll<Self::Output> { - match std::pin::Pin::new(&mut self.rx).poll(cx) { - std::task::Poll::Ready(Ok(value)) => { - // Result delivered: disarm so Drop doesn't try to abort an - // already-completed task. - if let Some(handle) = self.abort.take() { - handle.detach(); - } - std::task::Poll::Ready(Ok(value)) - } - std::task::Poll::Ready(Err(_)) => { - if let Some(handle) = self.abort.take() { - handle.detach(); - } - std::task::Poll::Ready(Err(SpawnCanceled)) - } - std::task::Poll::Pending => std::task::Poll::Pending, - } - } -} - -impl<T> Drop for Spawned<T> { - fn drop(&mut self) { - // If the future was dropped before completion, abort the spawned - // task to stop the wasted CPU work. AbortHandle::abort is a no-op - // after the task has already finished, so this is always safe. - if let Some(handle) = self.abort.take() { - handle.abort(); - } - } -} - -/// Spawn `fut` on the runtime and return a future that resolves to its -/// output. Cancellation propagates: dropping the returned future aborts -/// the spawned task. A spawned-task panic surfaces as `Err(SpawnCanceled)` -/// rather than a panic on `rx.await`. -#[cfg(not(target_arch = "wasm32"))] -fn spawn_oneshot<F, T>( - rt: &dyn Runtime, - fut: F, -) -> impl Future<Output = std::result::Result<T, SpawnCanceled>> + Send + 'static -where - F: Future<Output = T> + Send + 'static, - T: Send + 'static, -{ - let (tx, rx) = futures::channel::oneshot::channel(); - let abort = rt.spawn(Box::pin(async move { - let _ = tx.send(fut.await); - })); - Spawned { - rx, - abort: Some(abort), - } -} - -#[cfg(target_arch = "wasm32")] -fn spawn_oneshot<F, T>( - rt: &dyn Runtime, - fut: F, -) -> impl Future<Output = std::result::Result<T, SpawnCanceled>> + 'static -where - F: Future<Output = T> + 'static, - T: 'static, -{ - let (tx, rx) = futures::channel::oneshot::channel(); - let abort = rt.spawn(Box::pin(async move { - let _ = tx.send(fut.await); - })); - Spawned { - rx, - abort: Some(abort), - } -} - -/// Encrypt padded plaintext for each device JID, producing participant `<to>` nodes. -/// -/// Encrypt the plaintext for one device's Signal session. Shared by the -/// single-device fast path and the parallel fan-out so both behave identically. -async fn encrypt_one_device( - plaintext: &[u8], - addr: &ProtocolAddress, - session_store: &mut dyn crate::libsignal::protocol::SessionStore, - identity_store: &mut dyn crate::libsignal::protocol::IdentityKeyStore, - device_jid: Jid, - hide_decrypt_fail: bool, -) -> (Jid, Result<Option<EncryptOneResult>, String>) { - match message_encrypt(plaintext, addr, session_store, identity_store).await { - Ok(encrypted_payload) => { - let Some((enc_type, is_prekey, serialized_bytes)) = - extract_ciphertext(encrypted_payload) - else { - return (device_jid, Ok(None)); - }; - ( - device_jid, - Ok(Some(EncryptOneResult { - enc_type, - is_prekey, - // Box<[u8]> -> Vec<u8> reuses the allocation (no copy). - ciphertext: serialized_bytes.into(), - hide_decrypt_fail, - })), - ) - } - Err(e) => (device_jid, Err(format!("{addr}: {e}"))), - } -} - -/// Append one encrypt result to the fan-out output: a `<to>` participant node on -/// success, a logged skip on failure. -fn push_encrypt_result( - (device_jid, res): (Jid, Result<Option<EncryptOneResult>, String>), - mediatype: Option<&str>, - participant_nodes: &mut Vec<Node>, - encrypted_devices: &mut Vec<Jid>, - includes_prekey_message: &mut bool, -) { - match res { - Ok(Some(one)) => { - *includes_prekey_message |= one.is_prekey; - let mut enc_builder = NodeBuilder::new("enc") - .attr("v", stanza::ENC_VERSION) - .attr("type", one.enc_type); - // `mediatype` is batch-level (same for every device) and originates as - // a `&'static str`, so it's threaded here instead of cloned per result. - if let Some(mt) = mediatype { - enc_builder = enc_builder.attr("mediatype", mt); - } - if one.hide_decrypt_fail { - enc_builder = enc_builder.attr("decrypt-fail", "hide"); - } - let enc_node = enc_builder.bytes(one.ciphertext).build(); - participant_nodes.push( - NodeBuilder::new("to") - .attr("jid", device_jid.clone()) - .children([enc_node]) - .build(), - ); - encrypted_devices.push(device_jid); - } - Ok(None) => {} - Err(msg) => log::warn!("Failed to encrypt for device: {msg}. Skipping."), - } -} - -/// Per-device Signal sessions are independent (different ratchet state per -/// recipient), so this fans the encrypt loop out across tokio tasks bounded -/// by [`ENCRYPT_FANOUT_CONCURRENCY`]. Each task clones the store handles -/// (Arc bumps under the hood); the shared cache provides interior mutability. -/// -/// Callers must hold per-device session locks before calling this function — -/// concurrent ratchet mutations will corrupt Signal session state. -pub async fn encrypt_for_devices<'a, S, I, P, SP>( - runtime: &dyn Runtime, - stores: &mut SignalStores<'a, S, I, P, SP>, - resolver: &dyn SendContextResolver, - devices: &[Jid], - plaintext_to_encrypt: &[u8], - hide_decrypt_fail: bool, - mediatype: Option<&str>, -) -> Result<EncryptResult> -where - S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static, - I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static, - P: crate::libsignal::protocol::PreKeyStore + Send + Sync, - SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync, -{ - // Per-device LID upgrade map: encryption_overrides[i] mirrors devices[i]. - // None = use devices[i] as-is; Some(jid) = use this LID-upgraded version. - // The Vec replaces a HashMap<&Jid, Jid> that paid hash + alloc per insert - // and per get (~666 of each on a large group). Plain Vec<Option<Jid>> is - // direct indexing and contiguous memory. - let mut encryption_overrides: Vec<Option<Jid>> = vec![None; devices.len()]; - // 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 reusable_addr = crate::types::jid::make_reusable_protocol_address(); - - for (idx, device_jid) in devices.iter().enumerate() { - // WhatsApp Web's SignalAddress.toString() normalizes PN → LID before - // creating signal addresses. We do the same: check LID session FIRST. - // This prevents using stale PN sessions when a newer LID session exists. - if device_jid.is_pn() - && let Some(lid_user) = resolver.get_lid_for_phone(&device_jid.user).await - { - // Construct the LID JID with the same device ID - let lid_jid = Jid::lid_device(lid_user, device_jid.device); - lid_jid.reset_protocol_address(&mut reusable_addr); - - if stores.session_store.has_session(&reusable_addr).await? { - log::debug!( - "Using LID session {} for PN {} (LID-first lookup)", - lid_jid, - device_jid - ); - encryption_overrides[idx] = Some(lid_jid); - continue; - } - } - - device_jid.reset_protocol_address(&mut reusable_addr); - if stores.session_store.has_session(&reusable_addr).await? { - continue; - } - - // No session found - need to fetch prekeys and create session. - // Keep device_jid for prekey fetch (server returns bundles keyed by this), - // but normalize to LID for the actual session creation. - if device_jid.is_pn() - && let Some(lid_user) = resolver.get_lid_for_phone(&device_jid.user).await - { - let lid_jid = Jid::lid_device(lid_user, device_jid.device); - log::debug!( - "Will create LID session {} for PN {} (no existing session)", - lid_jid, - device_jid - ); - encryption_overrides[idx] = Some(lid_jid); - } - indices_needing_prekeys.push(idx); - } - - if !indices_needing_prekeys.is_empty() { - log::debug!( - "Fetching prekeys for {} devices without sessions", - indices_needing_prekeys.len() - ); - // Materialize the Jid slice for the resolver call. fetch_prekeys - // wants &[Jid]; same per-device clone count as the previous Vec - // model, just sourced from the indices. - let jids_for_fetch: Vec<Jid> = indices_needing_prekeys - .iter() - .map(|&i| devices[i].clone()) - .collect(); - // 406 on this batch is all-or-nothing — per-device retries just wasted - // N·RTT with the same failure. Mark `had_406` so the caller invalidates - // the users and the next send re-fetches. Matches WA Web's - // `GroupSkmsgJob`: log, continue without those devices. - let prekey_bundles = match resolver - .fetch_prekeys_for_identity_check(&jids_for_fetch) - .await - { - Ok(bundles) => bundles, - Err(e) if is_device_unregistered_error(&e) => { - log::warn!( - "Prekey fetch returned 406 for {} device(s); skipping them this round", - jids_for_fetch.len() - ); - had_406 = true; - std::collections::HashMap::new() - } - Err(e) => return Err(e), - }; - - // Parallel session establishment via process_prekey_bundle. Each - // recipient device has an independent Signal session and an - // independent prekey bundle, so the X3DH derivation runs on a - // separate task per device, bounded at ENCRYPT_FANOUT_CONCURRENCY. - // Spawning goes through `Runtime::spawn` (the platform-agnostic - // abstraction) plus a oneshot channel for result delivery — - // `FuturesUnordered` handles the in-flight window. - let prekey_bundles = std::sync::Arc::new(prekey_bundles); - let total = indices_needing_prekeys.len(); - let mut next_spawn = 0usize; - - let make_session_task = |spawn_idx: usize| { - let idx = indices_needing_prekeys[spawn_idx]; - let device_jid = devices[idx].clone(); - let mut encryption_jid = encryption_overrides[idx] - .clone() - .unwrap_or_else(|| device_jid.clone()); - - // Normalize agent to 0 for LID JIDs to match how pre-key bundles are stored. - // prekeys.rs forces agent=0 for LID; we must match that here. - if encryption_jid.is_lid() { - encryption_jid.agent = 0; - } - - let lookup_jid = device_jid.normalize_for_prekey_bundle(); - let bundles = prekey_bundles.clone(); - let mut session_store = stores.session_store.clone(); - let mut identity_store = stores.identity_store.clone(); - - spawn_oneshot(runtime, async move { - let mut addr = crate::types::jid::make_reusable_protocol_address(); - encryption_jid.reset_protocol_address(&mut addr); - - let Some(bundle) = bundles.get(&lookup_jid) else { - log::warn!( - "No pre-key bundle returned for device {}. This device will be skipped for encryption.", - addr - ); - return Ok::<Option<Jid>, anyhow::Error>(None); - }; - - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - // No UntrustedIdentity recovery: WA Web's isTrustedIdentity is - // unconditional Ok(true) (TOFU), and save_identity inside - // process_prekey_bundle persists rotations transparently. - match process_prekey_bundle( - &addr, - &mut session_store, - &mut identity_store, - bundle, - &mut rng, - UsePQRatchet::No, - ) - .await - { - // Surface a replaced identity so the caller can react - // (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 - )), - } - }) - }; - - let mut in_flight: FuturesUnordered<_> = FuturesUnordered::new(); - while next_spawn < total && in_flight.len() < ENCRYPT_FANOUT_CONCURRENCY { - in_flight.push(make_session_task(next_spawn)); - next_spawn += 1; - } - while let Some(spawn_result) = in_flight.next().await { - match spawn_result { - // Some(jid) => establishing this session replaced a stored - // identity; notify the client so it can react off-path. - Ok(Ok(Some(changed_jid))) => resolver.on_local_identity_change(&changed_jid), - Ok(Ok(None)) => {} - Ok(Err(e)) => return Err(e), - Err(SpawnCanceled) => { - log::warn!( - "Session-establishment task did not deliver a result; skipping device." - ); - } - } - if next_spawn < total { - in_flight.push(make_session_task(next_spawn)); - next_spawn += 1; - } - } - } - - let mut participant_nodes = Vec::with_capacity(devices.len()); - let mut includes_prekey_message = false; - let mut encrypted_devices = Vec::with_capacity(devices.len()); - - // The wire-order of `<to>` participants does not need to match the input - // device order: WA Web's `phash` (computed both client and server side) - // sorts before hashing, as does our `participant_list_hash`. - if devices.len() == 1 { - // Single recipient device: the parallel fan-out is pure overhead here - // (an Arc<[u8]> copy of the plaintext, a spawned task + oneshot channel, - // a FuturesUnordered, and two store clones), with no parallelism to gain. - // Encrypt inline. - let device_jid = devices[0].clone(); - let addr = encryption_overrides[0] - .as_ref() - .unwrap_or(&devices[0]) - .to_protocol_address(); - let res = encrypt_one_device( - plaintext_to_encrypt, - &addr, - &mut *stores.session_store, - &mut *stores.identity_store, - device_jid, - hide_decrypt_fail, - ) - .await; - push_encrypt_result( - res, - mediatype, - &mut participant_nodes, - &mut encrypted_devices, - &mut includes_prekey_message, - ); - } else { - // Parallel encrypt fan-out across tokio tasks bounded by - // ENCRYPT_FANOUT_CONCURRENCY; collected in completion order so the - // fastest encrypts ship first. - let plaintext_arc: std::sync::Arc<[u8]> = std::sync::Arc::from(plaintext_to_encrypt); - - let total = devices.len(); - let mut next_spawn = 0usize; - - let make_encrypt_task = |idx: usize| { - let device_jid = devices[idx].clone(); - // The encryption JID is only needed to build the Signal address, so - // derive it here from a borrow rather than cloning the whole Jid into - // the task (device_jid is still cloned because it's returned). - let addr = encryption_overrides[idx] - .as_ref() - .unwrap_or(&devices[idx]) - .to_protocol_address(); - let plaintext = plaintext_arc.clone(); - let mut session_store = stores.session_store.clone(); - let mut identity_store = stores.identity_store.clone(); - - spawn_oneshot(runtime, async move { - encrypt_one_device( - &plaintext, - &addr, - &mut session_store, - &mut identity_store, - device_jid, - hide_decrypt_fail, - ) - .await - }) - }; - - let mut in_flight: FuturesUnordered<_> = FuturesUnordered::new(); - while next_spawn < total && in_flight.len() < ENCRYPT_FANOUT_CONCURRENCY { - in_flight.push(make_encrypt_task(next_spawn)); - next_spawn += 1; - } - while let Some(spawn_result) = in_flight.next().await { - match spawn_result { - Ok(res) => push_encrypt_result( - res, - mediatype, - &mut participant_nodes, - &mut encrypted_devices, - &mut includes_prekey_message, - ), - Err(SpawnCanceled) => { - log::warn!("Encrypt task did not deliver a result; skipping device."); - } - } - - if next_spawn < total { - in_flight.push(make_encrypt_task(next_spawn)); - next_spawn += 1; - } - } - } - - Ok(EncryptResult { - participant_nodes, - includes_prekey_message, - encrypted_devices, - had_unregistered_device: had_406, - }) -} - -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) - || own_lid - .is_some_and(|lid| device_jid.is_same_user_as(lid) && device_jid.device == lid.device) -} - -fn partition_dm_devices( - all_devices: Vec<Jid>, - own_jid: &Jid, - own_lid: Option<&Jid>, -) -> (Vec<Jid>, Vec<Jid>) { - let mut recipient_devices = Vec::with_capacity(all_devices.len()); - let mut own_other_devices = Vec::with_capacity(4); - - for device_jid in all_devices { - if is_exact_dm_sender_device(&device_jid, own_jid, own_lid) { - continue; - } - - if device_jid.matches_user_or_lid(own_jid, own_lid) { - own_other_devices.push(device_jid); - } else { - recipient_devices.push(device_jid); - } - } - - (recipient_devices, own_other_devices) -} - -/// Result of `prepare_dm_stanza` — carries the stanza node and the -/// locally computed phash for server ACK validation. -pub struct PreparedDmStanza { - pub node: Node, - /// Locally computed phash from the sent device set. Not sent on the - /// wire (WA Web only sends phash for groups). Used by the caller to - /// compare against the server's ACK phash for device-list drift detection. - pub phash: Option<String>, - /// `MessageContextInfo.message_secret` generated for this stanza so the - /// caller can persist it for later addon (msmsg/poll/edit) decryption. - /// `None` when the message had no reporting token (no secret was used). - pub message_secret: Option<[u8; crate::reporting_token::MESSAGE_SECRET_SIZE]>, -} - -#[allow(clippy::too_many_arguments)] -pub async fn prepare_dm_stanza< - 'a, - S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static, - I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static, - P: crate::libsignal::protocol::PreKeyStore + Send + Sync, - SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync, ->( - runtime: &dyn Runtime, - stores: &mut SignalStores<'a, S, I, P, SP>, - 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>, -) -> Result<PreparedDmStanza> { - // sender is the author's own jid, remote is the chat jid (WAWebReportingTokenUtils: - // getSender vs e.to). Both previously used to_jid, conflating sender with remote. - let reporting_result = generate_reporting_token(message, &request_id, own_jid, &to_jid, None); - - let message_for_encryption = if let Some(ref result) = reporting_result { - prepare_message_with_context(message, &result.message_secret) - } else { - message.clone() - }; - - let recipient_plaintext = MessageUtils::encode_and_pad(&message_for_encryption); - - // Partition first so phash reflects the actual sent set (sender excluded) - let total_devices = all_devices.len(); - let (recipient_devices, own_other_devices) = - partition_dm_devices(all_devices, own_jid, own_lid); - - let phash = MessageUtils::participant_list_hash( - recipient_devices.iter().chain(own_other_devices.iter()), - ) - .ok(); - - let dsm = crate::messages::wrap_device_sent(message_for_encryption, to_jid.to_string()); - - let own_devices_plaintext = MessageUtils::encode_and_pad(&dsm); - - 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 mediatype = media_type_from_message(message); - - // NOTE: WA Web has a bare-<enc> fast path for single primary device - // (WAWebSendMsgCreateFanoutStanza). Not implemented here because - // encrypt_for_devices always wraps in <to jid=...> nodes; - // a bare-enc mode would require refactoring the encryption layer. - // The <participants> form is accepted by the server regardless. - - if !recipient_devices.is_empty() { - let result = encrypt_for_devices( - runtime, - stores, - resolver, - &recipient_devices, - &recipient_plaintext, - hide_decrypt_fail, - mediatype, - ) - .await?; - participant_nodes.extend(result.participant_nodes); - includes_prekey_message = includes_prekey_message || result.includes_prekey_message; - } - - if !own_other_devices.is_empty() { - let result = encrypt_for_devices( - runtime, - stores, - resolver, - &own_other_devices, - &own_devices_plaintext, - hide_decrypt_fail, - mediatype, - ) - .await?; - participant_nodes.extend(result.participant_nodes); - includes_prekey_message = includes_prekey_message || result.includes_prekey_message; - } - - // All per-device encrypts failed: an empty <participants> would silently - // drop the message. WA Web's encryptAndSendUserMsg rejects here too. - let attempted_devices = recipient_devices.len() + own_other_devices.len(); - if participant_nodes.is_empty() && attempted_devices > 0 { - return Err(anyhow!( - "encryption failed for all {attempted_devices} recipient device(s)" - )); - } - - let mut message_content_nodes = vec![ - NodeBuilder::new("participants") - .children(participant_nodes) - .build(), - ]; - - if includes_prekey_message && let Some(acc) = account { - let device_identity_bytes = acc.encode_to_vec(); - message_content_nodes.push( - NodeBuilder::new("device-identity") - .bytes(device_identity_bytes) - .build(), - ); - } - - // Add reporting token node if we generated one - if let Some(ref result) = reporting_result { - message_content_nodes.push(build_reporting_node(result)); - } - - // Add any extra stanza nodes provided by the caller - message_content_nodes.extend(extra_stanza_nodes.iter().cloned()); - - let stanza_type = stanza_type_from_message(message); - - let mut stanza_builder = NodeBuilder::new("message") - .attr("to", to_jid) - .attr("id", request_id) - .attr("type", stanza_type); - - if let Some(edit_attr) = edit - && edit_attr != crate::types::message::EditAttribute::Empty - { - stanza_builder = stanza_builder.attr("edit", edit_attr.to_string_val()); - } - - let stanza = stanza_builder.children(message_content_nodes).build(); - - Ok(PreparedDmStanza { - node: stanza, - phash, - message_secret: reporting_result.map(|r| r.message_secret), - }) -} - -/// Returns true if `message_encrypt` on `signal_address` would produce -/// a pkmsg (no session yet, or session with un-acked pre-key still -/// pending). Used before `message_encrypt` to fail-fast when `account` -/// is None — pkmsg without `<device-identity>` reproduces the linked -/// device deadlock. -/// -/// `SessionStore::load_session` is take-semantics in production -/// (`SessionAdapter` → `SignalStoreCache::get_session` marks the slot -/// `CheckedOut`); the loaded record is put back via `store_session` -/// so the subsequent `message_encrypt` finds the slot Present. -async fn pkmsg_would_be_emitted<S>( - session_store: &mut S, - signal_address: &ProtocolAddress, -) -> Result<bool> -where - S: crate::libsignal::protocol::SessionStore, -{ - let loaded = session_store.load_session(signal_address).await?; - // Conservative read: treat any failure to interrogate the session as - // "would be pkmsg" so the caller bails. Silently treating Err as false - // would let message_encrypt run with a corrupt session and potentially - // burn the sender chain. - let needs_pkmsg = match &loaded { - None => true, - Some(record) => match record.session_state() { - None => true, - Some(state) => match state.unacknowledged_pre_key_message_items() { - Ok(Some(_)) => true, - Ok(None) => false, - Err(_) => true, - }, - }, - }; - if let Some(record) = loaded { - session_store - .store_session(signal_address, record) - .await - .map_err(|e| anyhow!("restoring checked-out session after pre-flight: {e}"))?; - } - Ok(needs_pkmsg) -} - -#[allow(clippy::too_many_arguments)] -pub async fn prepare_peer_stanza<S, I>( - session_store: &mut S, - identity_store: &mut I, - transport_jid: Jid, - signal_address: &ProtocolAddress, - message: &wa::Message, - request_id: String, - account: Option<&wa::AdvSignedDeviceIdentity>, -) -> Result<Node> -where - S: crate::libsignal::protocol::SessionStore, - I: crate::libsignal::protocol::IdentityKeyStore, -{ - let options = peer_message_options_from_message(message); - prepare_peer_stanza_with_options( - session_store, - identity_store, - transport_jid, - signal_address, - message, - request_id, - account, - options, - ) - .await -} - -#[allow(clippy::too_many_arguments)] -pub async fn prepare_peer_stanza_with_options<S, I>( - session_store: &mut S, - identity_store: &mut I, - transport_jid: Jid, - signal_address: &ProtocolAddress, - message: &wa::Message, - request_id: String, - account: Option<&wa::AdvSignedDeviceIdentity>, - options: PeerMessageOptions, -) -> Result<Node> -where - S: crate::libsignal::protocol::SessionStore, - I: crate::libsignal::protocol::IdentityKeyStore, -{ - let plaintext = MessageUtils::encode_and_pad(message); - - if account.is_none() && pkmsg_would_be_emitted(session_store, signal_address).await? { - bail!( - "peer pkmsg requires <device-identity> (account is None); \ - refusing before message_encrypt to avoid advancing the sender chain" - ); - } - - let encrypted_message = - message_encrypt(&plaintext, signal_address, session_store, identity_store).await?; - - let (enc_type, is_prekey, serialized_bytes) = extract_ciphertext(encrypted_message) - .ok_or_else(|| anyhow!("Unexpected peer encryption message type"))?; - - let enc_node = NodeBuilder::new("enc") - .attrs([("v", "2"), ("type", enc_type)]) - .bytes(serialized_bytes) - .build(); - - let meta_node = NodeBuilder::new("meta").attr("appdata", "default").build(); - - let mut children = vec![meta_node, enc_node]; - if is_prekey { - // Defense in depth: pre-flight should have caught this, but a corrupt - // session that triggers a fresh pkmsg mid-call would slip past. - let account = account.ok_or_else(|| { - anyhow!("peer pkmsg without <device-identity> (unreachable via pre-flight)") - })?; - children.push( - NodeBuilder::new("device-identity") - .bytes(account.encode_to_vec()) - .build(), - ); - } - - let mut stanza_builder = NodeBuilder::new("message") - .attr("to", transport_jid) - .attr("id", request_id) - .attr("type", stanza::MSG_TYPE_TEXT) - .attr("category", "peer") - .attr("push_priority", options.push_priority().as_str()); - if let Some(privacy_sensitive) = options.privacy_sensitive() { - stanza_builder = stanza_builder.attr("privacy_sensitive", privacy_sensitive.as_str()); - } - - Ok(stanza_builder.children(children).build()) -} - -/// 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. -#[allow(clippy::too_many_arguments)] -pub async fn prepare_dm_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>, -) -> 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!( - "DM 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 DM retry"))?; - - let hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message); - 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); - } - if hide_decrypt_fail { - enc_builder = enc_builder.attr("decrypt-fail", "hide"); - } - let enc_node = enc_builder.bytes(serialized).build(); - - let mut children = vec![enc_node]; - if is_prekey { - // Defense in depth: pre-flight should have caught this, but a corrupt - // session that triggers a fresh pkmsg mid-call would slip past. - let acc = account.ok_or_else(|| { - anyhow!("DM retry pkmsg without <device-identity> (unreachable via pre-flight)") - })?; - children.push( - NodeBuilder::new("device-identity") - .bytes(acc.encode_to_vec()) - .build(), - ); - } - - let mut stanza_builder = NodeBuilder::new("message") - .attr("to", to_jid) - .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. - 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()) -} - -/// 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. -#[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]; - - if is_prekey { - // Defense in depth: pre-flight should have caught this, but a corrupt - // session that triggers a fresh pkmsg mid-call would slip past. - let acc = account.ok_or_else(|| { - anyhow!("group retry pkmsg without <device-identity> (unreachable via pre-flight)") - })?; - children.push( - NodeBuilder::new("device-identity") - .bytes(acc.encode_to_vec()) - .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()) -} - -/// Result of `prepare_group_stanza` — carries the stanza node and the exact -/// device list used for SKDM distribution, so callers can persist sender key -/// tracking without re-resolving devices. -pub struct PreparedGroupStanza { - pub node: Node, - /// Full SKDM distribution target set, marked `has_key=true` after the - /// server ACK. Mirrors WA Web `markHasSenderKey(x, M)` which marks the - /// whole target set `M`, not only the devices that encrypted successfully: - /// devices that failed (406 / no bundle) are marked too so they are not - /// re-targeted on every send (the retry-receipt path repairs any that are - /// actually alive and keyless via `mark_forget_sender_key`). - pub skdm_devices: Vec<Jid>, - /// Users whose device registry should be invalidated because their - /// devices returned 406 (unregistered) during SKDM prekey fetch. - /// Empty when no 406 occurred. - pub stale_device_users: Vec<String>, - /// Generated `MessageContextInfo.message_secret`; populated when the - /// reporting token was produced for this send. - pub message_secret: Option<[u8; crate::reporting_token::MESSAGE_SECRET_SIZE]>, - /// The identity we addressed this group send under (LID for LID-mode - /// groups, PN for PN-mode). Used to key the persisted `messageSecret` - /// so msmsg bot replies referencing this msg_id hit the same row that - /// `<meta target_sender_jid>` echoes back at lookup time. - pub sender_identity: Jid, -} - -#[allow(clippy::too_many_arguments)] -pub async fn prepare_group_stanza< - 'a, - S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static, - I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static, - P: crate::libsignal::protocol::PreKeyStore + Send + Sync, - SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync, ->( - runtime: &dyn Runtime, - stores: &mut SignalStores<'a, S, I, P, SP>, - 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<Vec<Jid>>, - edit: Option<crate::types::message::EditAttribute>, - extra_stanza_nodes: &[Node], -) -> Result<PreparedGroupStanza> { - 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"), - }; - - // Generate reporting token if the message type supports it - // For groups, both sender_jid and remote_jid are the group JID (to_jid) per Baileys implementation - let reporting_result = generate_reporting_token(message, &request_id, &to_jid, &to_jid, None); - - // Prepare message with MessageContextInfo containing the message secret - let message_for_encryption = if let Some(ref result) = reporting_result { - prepare_message_with_context(message, &result.message_secret) - } else { - message.clone() - }; - - let own_base_jid = own_sending_jid.to_non_ad(); - - let mut message_children: Vec<Node> = Vec::new(); - let mut includes_prekey_message = false; - let mut phash_for_stanza: Option<String> = None; - let mut skdm_encrypted_devices: Vec<Jid> = Vec::new(); - - // Build the chain name once and hold its lock across SKDM creation + the - // skmsg encrypt, so concurrent same-(group, sender) sends can't split the - // key between the SKDM and the skmsg (nor reuse a chain iteration). - let sender_key_name = make_sender_key_name(&to_jid, &own_sending_jid.to_protocol_address()); - let chain_lock = stores - .sender_key_store - .sender_key_lock(&sender_key_name) - .await; - let _chain_guard = chain_lock.lock().await; - - // Determine if we need to distribute SKDM and to which devices - let distribution_list: Option<Vec<Jid>> = if let Some(target_devices) = skdm_target_devices { - // Use the specific list of devices that need SKDM - if target_devices.is_empty() { - None - } else { - log::debug!( - "SKDM distribution to {} specific devices for group {}", - target_devices.len(), - to_jid - ); - Some(target_devices) - } - } else if force_skdm_distribution { - // Resolve all devices for all participants (legacy behavior) - // For LID groups, use phone numbers for device queries (LID usync may not work for own JID) - // For PN groups, use JIDs directly - let mut jids_to_resolve: Vec<Jid> = group_info - .participants - .iter() - .map(|jid| { - let base_jid = jid.to_non_ad(); - // If this is a LID JID and we have a phone number mapping, use it for device query - if base_jid.is_lid() - && let Some(phone_jid) = group_info.phone_jid_for_lid_user(&base_jid.user) - { - log::debug!( - "Using phone number {} for LID {} device query", - phone_jid, - base_jid - ); - return phone_jid.to_non_ad(); - } - base_jid - }) - .collect(); - - // Determine what user to check for — use the PN user when own is LID - // and we have a mapping. Keeping this as a borrow avoids allocating a - // throwaway Jid when own is already in the list. - let own_pn_mapping = if own_base_jid.is_lid() { - group_info.phone_jid_for_lid_user(&own_base_jid.user) - } else { - None - }; - let own_check_user = own_pn_mapping - .map(|pn| pn.user.as_str()) - .unwrap_or(own_base_jid.user.as_str()); - - if !jids_to_resolve.iter().any(|p| p.user == own_check_user) { - jids_to_resolve.push(match own_pn_mapping { - Some(pn) => pn.to_non_ad(), - None => own_base_jid.clone(), - }); - } - - crate::types::jid::sort_dedup_by_user(&mut jids_to_resolve); - - log::debug!( - "Resolving devices for {} participants", - jids_to_resolve.len() - ); - - let mut resolved_list = resolver.resolve_devices(&jids_to_resolve).await?; - - // For LID groups, convert phone-based device JIDs back to LID format - // This is necessary because WhatsApp Web expects LID addressing in SKDM <to> nodes - if group_info.addressing_mode == crate::types::message::AddressingMode::Lid { - resolved_list = resolved_list - .into_iter() - .map(|device_jid| group_info.phone_device_jid_into_lid(device_jid)) - .collect(); - log::debug!( - "Converted {} devices to LID addressing for group {}", - resolved_list.len(), - to_jid - ); - } - - // Dedup AFTER LID conversion to avoid duplicates when both phone and LID - // queries return the same user (e.g., 559980000003:33 and 100000037037034:33 - // both convert to 100000037037034:33@lid). - // Key on (user, server, agent, device) — excludes `integrator` which is not - // part of the wire JID identity used in <to jid> and phash. - crate::types::jid::sort_dedup_by_device(&mut resolved_list); - - // Filter devices for SKDM distribution: - // - Exclude the exact sending device (own_sending_jid) - we already have our own sender key - // - Keep ALL other devices including our own other devices (phone, other companions) - // because they need the SKDM to decrypt messages we send from this device - // - Exclude hosted/Cloud API devices (device ID 99 or @hosted server) - they don't - // participate in group E2EE, only in 1:1 chats - 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 - }); - log::debug!( - "Filtered SKDM devices from {} to {} (excluded sender {}:{} and hosted devices)", - before_filter, - resolved_list.len(), - own_user, - own_device - ); - - log::debug!( - "SKDM distribution list for {} resolved to {} devices", - to_jid, - resolved_list.len(), - ); - - Some(resolved_list) - } else { - None - }; - - // 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. - if to_jid.is_group() { - // Warm/partial sends pass the complete set in `all_devices_for_phash`; - // the cold `force_skdm` path leaves it None and `distribution_list` - // already holds the full resolved set. - if let Some(src) = all_devices_for_phash - .as_deref() - .or(distribution_list.as_deref()) - { - let phash_set = build_group_phash_set(src, &own_sending_jid); - match MessageUtils::participant_list_hash(&phash_set) { - Ok(phash) => phash_for_stanza = Some(phash), - Err(e) => log::warn!("Failed to compute group phash for {}: {:?}", to_jid, e), - } - } - } else if let Some(ref distribution_list) = distribution_list { - match MessageUtils::participant_list_hash(distribution_list) { - Ok(phash) => phash_for_stanza = Some(phash), - Err(e) => log::warn!("Failed to compute phash for {}: {:?}", to_jid, e), - } - } - - let mut had_unregistered_devices = false; - - if let Some(ref distribution_list) = distribution_list { - let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group( - stores.sender_key_store, - &sender_key_name, - ) - .await?; - - let skdm_wrapper_msg = wa::Message { - sender_key_distribution_message: Some(wa::message::SenderKeyDistributionMessage { - group_id: Some(to_jid.to_string()), - axolotl_sender_key_distribution_message: Some(axolotl_skdm_bytes), - }), - ..Default::default() - }; - let skdm_plaintext_to_encrypt = MessageUtils::encode_and_pad(&skdm_wrapper_msg); - - // WA Web's GroupSkmsgJob wraps ensureE2ESessions in try/catch — logs error - // but does NOT rethrow. SKDM distribution failure must not prevent the group - // message from being sent. Only successfully encrypted devices are tracked. - // 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( - runtime, - stores, - resolver, - distribution_list, - &skdm_plaintext_to_encrypt, - skdm_hide_decrypt_fail, - None, - ) - .await - { - Ok(result) => { - includes_prekey_message = includes_prekey_message || result.includes_prekey_message; - if result.had_unregistered_device { - had_unregistered_devices = true; - } - skdm_encrypted_devices = result.encrypted_devices; - - if !result.participant_nodes.is_empty() { - message_children.push( - NodeBuilder::new("participants") - .children(result.participant_nodes) - .build(), - ); - if includes_prekey_message && let Some(acc) = account { - message_children.push( - NodeBuilder::new("device-identity") - .bytes(acc.encode_to_vec()) - .build(), - ); - } - } - } - Err(e) => { - log::warn!( - "SKDM distribution failed for group {}, continuing without it: {e}", - to_jid - ); - if is_device_unregistered_error(&e) { - had_unregistered_devices = true; - } - } - } - } - - let plaintext = MessageUtils::encode_and_pad(&message_for_encryption); - let skmsg = encrypt_group_message( - stores.sender_key_store, - &sender_key_name, - &plaintext, - &mut rand::make_rng::<rand::rngs::StdRng>(), - ) - .await?; - - 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 mut enc_builder = NodeBuilder::new("enc") - .attr("v", stanza::ENC_VERSION) - .attr("type", stanza::ENC_TYPE_SKMSG); - if let Some(mt) = mediatype { - enc_builder = enc_builder.attr("mediatype", mt); - } - enc_builder = enc_builder.bytes(skmsg_ciphertext); - if hide_decrypt_fail { - enc_builder = enc_builder.attr("decrypt-fail", "hide"); - } - let content_node = enc_builder.build(); - - let stanza_type = stanza_type_from_message(message); - let mut stanza_builder = NodeBuilder::new("message") - .attr("to", to_jid) - .attr("id", request_id) - .attr("type", stanza_type); - - // WA Web always sets addressing_mode for groups (MsgCreateDeviceStanza.js:131-135) - stanza_builder = stanza_builder.attr("addressing_mode", group_info.addressing_mode.as_str()); - - if let Some(edit_attr) = &edit - && *edit_attr != crate::types::message::EditAttribute::Empty - { - stanza_builder = stanza_builder.attr("edit", edit_attr.to_string_val()); - } - // NOTE: WhatsApp Web does NOT include participant attribute on initial admin revoke send - // The participant attribute only appears on retry/fanout messages - - message_children.push(content_node); - - // Add reporting token node if we generated one - if let Some(ref result) = reporting_result { - message_children.push(build_reporting_node(result)); - } - - // Add phash if we distributed keys in this message - if let Some(phash) = phash_for_stanza { - stanza_builder = stanza_builder.attr("phash", phash); - } - - // Add any extra stanza nodes provided by the caller - message_children.extend(extra_stanza_nodes.iter().cloned()); - - let stanza = stanza_builder.children(message_children).build(); - - let stale_users = if had_unregistered_devices { - collect_stale_device_users( - distribution_list.as_deref(), - &skdm_encrypted_devices, - group_info, - ) - } else { - Vec::new() - }; - - Ok(PreparedGroupStanza { - node: stanza, - // Mark the full target set (matches WA Web `markHasSenderKey(x, M)`), not - // just `skdm_encrypted_devices`. `stale_users` above already used the - // encrypted subset to find which devices to re-resolve. - skdm_devices: distribution_list.unwrap_or_default(), - stale_device_users: stale_users, - message_secret: reporting_result.map(|r| r.message_secret), - sender_identity: own_sending_jid, - }) -} - -/// Build the device set hashed into a group `phash`, matching WA Web -/// `phashV2([].concat(A, [B]))`: every participant device (`A`) plus the -/// sending device `B`. `devices` is the resolved set (recipients); the sending -/// device is excluded from it (we never SKDM ourselves) so it is appended here. -/// Hosted devices don't take part in group E2EE and are dropped, mirroring the -/// SKDM distribution filter. `participant_list_hash` sorts before hashing, so -/// order here is irrelevant. -pub(crate) fn build_group_phash_set(devices: &[Jid], own_sending_jid: &Jid) -> Vec<Jid> { - let mut set: Vec<Jid> = devices.iter().filter(|d| !d.is_hosted()).cloned().collect(); - if !set - .iter() - .any(|d| d.user == own_sending_jid.user && d.device == own_sending_jid.device) - { - set.push(own_sending_jid.clone()); - } - crate::types::jid::sort_dedup_by_device(&mut set); - set -} - -/// Collect users whose devices failed SKDM so the caller can invalidate their -/// registry entries. In LID-mode groups, both the LID and PN aliases are -/// emitted when the group knows the mapping — `invalidate_device_cache` needs -/// both to clean up zombie records that were stored under whichever alias -/// `update_device_list` canonicalised to at the time of the write. -pub(crate) fn collect_stale_device_users( - distribution_list: Option<&[Jid]>, - skdm_encrypted_devices: &[Jid], - group_info: &GroupInfo, -) -> Vec<String> { - let Some(dist) = distribution_list else { - return Vec::new(); - }; - let is_lid_mode = group_info.addressing_mode == crate::types::message::AddressingMode::Lid; - let encrypted_set: HashSet<&Jid> = skdm_encrypted_devices.iter().collect(); - let mut user_set: HashSet<String> = HashSet::new(); - for d in dist { - if encrypted_set.contains(d) { - continue; - } - user_set.insert(d.user.to_string()); - if is_lid_mode - && d.is_lid() - && let Some(pn_jid) = group_info.phone_jid_for_lid_user(&d.user) - && pn_jid.is_pn() - { - user_set.insert(pn_jid.user.to_string()); - } - } - user_set.into_iter().collect() -} - -/// Caller must hold `SenderKeyStore::sender_key_lock` for `sender_key_name` -/// across this creation + the matching skmsg encrypt (see `encrypt_group_message`). -pub async fn create_sender_key_distribution_message_for_group( - store: &mut (dyn SenderKeyStore + Send + Sync), - sender_key_name: &SenderKeyName, -) -> Result<Vec<u8>> { - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - - let skdm = crate::libsignal::protocol::create_sender_key_distribution_message( - sender_key_name, - store, - &mut rng, - ) - .await?; - - Ok(skdm.into_serialized().into_vec()) -} - -/// Ensure the status stanza has a `<participants>` node listing all recipient -/// user JIDs. WhatsApp Web's `participantList` uses bare USER JIDs (not -/// device JIDs) -- `<to jid="user@s.whatsapp.net"/>` -- to tell the server -/// which users should receive the skmsg. The SKDM distribution list -/// (already in `<participants>`) uses device JIDs with `<enc>` children. -/// -/// This is a pure function (no runtime or client dependencies). -pub fn ensure_status_participants( - mut stanza: Node, - group_info: &crate::client::context::GroupInfo, -) -> Node { - use wacore_binary::NodeContent; - use wacore_binary::builder::NodeBuilder; - - // Build bare <to jid="USER_JID"/> entries for each participant. - // WhatsApp Web uses USER_JID (not DEVICE_JID) for the participantList. - let bare_to_nodes: Vec<Node> = group_info - .participants - .iter() - .map(|jid| NodeBuilder::new("to").attr("jid", jid.to_non_ad()).build()) - .collect(); - - // Check if <participants> already exists in the stanza children - let children = match &mut stanza.content { - Some(NodeContent::Nodes(nodes)) => nodes, - _ => { - stanza.content = Some(NodeContent::Nodes(vec![])); - match &mut stanza.content { - Some(NodeContent::Nodes(nodes)) => nodes, - _ => unreachable!(), - } - } - }; - - if let Some(participants_node) = children.iter_mut().find(|n| n.tag == "participants") { - // <participants> already exists (from SKDM distribution). - // Add bare <to> user JID entries for users whose devices are NOT - // already represented by SKDM device-level entries. - let existing_users: std::collections::HashSet<wacore_binary::CompactString> = - participants_node - .children() - .unwrap_or_default() - .iter() - .filter_map(|n| n.attrs.get("jid").and_then(|v| v.to_jid()).map(|j| j.user)) - .collect(); - - let new_to_nodes: Vec<Node> = bare_to_nodes - .into_iter() - .filter(|n| { - n.attrs - .get("jid") - .and_then(|v| v.to_jid()) - .is_some_and(|j| !existing_users.contains(&j.user)) - }) - .collect(); - - if !new_to_nodes.is_empty() { - match &mut participants_node.content { - Some(NodeContent::Nodes(nodes)) => nodes.extend(new_to_nodes), - _ => { - participants_node.content = Some(NodeContent::Nodes(new_to_nodes)); - } - } - } - } else { - // No <participants> node — create one with bare <to> entries. - let participants_node = NodeBuilder::new("participants") - .children(bare_to_nodes) - .build(); - children.insert(0, participants_node); - } - - stanza -} - -/// True when a `status@broadcast` message should carry the -/// `<meta status_setting="..."/>` child. Only applies to actual status posts: -/// reactions (handled server-side as addons) and revokes must omit it, per -/// `WAWebEncryptAndSendStatusMsg` vs `WAWebSendReactionMsgAction`. -/// -/// Descends `ephemeral_message` / `device_sent_message` / view-once wrappers -/// before classifying (same as `stanza_type_from_message`), so a reaction -/// nested inside a wrapper cannot slip past and re-trigger 479. -pub fn status_carries_privacy_meta(message: &wa::Message) -> bool { - let msg = unwrap_message(message); - let is_revoke = msg - .protocol_message - .as_ref() - .is_some_and(|pm| pm.r#type == Some(wa::message::protocol_message::Type::Revoke as i32)); - let is_reaction = msg.reaction_message.is_some() || msg.enc_reaction_message.is_some(); - !is_revoke && !is_reaction -} - -/// 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 -/// entries; an empty result means "nothing to send to"). -/// -/// Pure function: no allocations besides the returned `Vec` and (when needed) -/// the own-LID push. Dedup is a linear Vec scan — status lists stay small -/// enough that a HashSet is not worth its allocation. -pub fn assemble_status_participants<I>(resolved: I, own_lid: &Jid) -> anyhow::Result<Vec<Jid>> -where - I: IntoIterator<Item = Option<Jid>>, -{ - let iter = resolved.into_iter(); - let (lower, _upper) = iter.size_hint(); - let mut out: Vec<Jid> = Vec::with_capacity(lower.saturating_add(1)); - for jid in iter.flatten() { - if !out.iter().any(|r| r.user == jid.user) { - out.push(jid); - } - } - if out.is_empty() { - anyhow::bail!("No valid status recipients after LID resolution"); - } - if !out.iter().any(|r| r.user == own_lid.user) { - out.push(own_lid.to_non_ad()); - } - Ok(out) -} - -/// Build a `Message.ProtocolMessage` for `GROUP_MEMBER_LABEL_CHANGE`. -/// -/// Sent via the standard E2EE fanout, not an IQ. Empty `label` clears. -/// `ts_secs` is unix seconds, matching WA Web's `unixTime()`. -pub fn build_member_label_message(label: String, ts_secs: i64) -> wa::Message { - wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage { - r#type: Some(wa::message::protocol_message::Type::GroupMemberLabelChange as i32), - member_label: Some(wa::MemberLabel { - label: Some(label), - label_timestamp: Some(ts_secs), - }), - ..Default::default() - })), - ..Default::default() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::client::context::{GroupInfo, SendContextResolver}; - use crate::libsignal::protocol::{IdentityKeyPair, KeyPair, PreKeyBundle}; - use std::collections::HashMap; - use wacore_binary::Jid; - - mod assemble_status_participants { - use super::*; - - fn lid(u: &str) -> Jid { - u.parse().expect("parse LID jid") - } - - #[test] - fn dedup_keeps_first_entry_per_user_and_anchors_own() { - let own = lid("99999999999999@lid"); - let out = assemble_status_participants( - vec![ - Some(lid("111@lid")), - Some(lid("222@lid")), - Some(lid("111@lid")), - Some(lid("333@lid")), - ], - &own, - ) - .expect("should succeed"); - let users: Vec<&str> = out.iter().map(|j| j.user.as_str()).collect(); - assert_eq!(users, ["111", "222", "333", "99999999999999"]); - } - - #[test] - fn skips_none_entries_matching_wa_web_compactmap() { - // Unresolvable recipients arrive as `None` and must be silently - // dropped — mirrors WA Web's `compactMap(list, toUserLid)`. - let own = lid("me@lid"); - let out = assemble_status_participants( - vec![None, Some(lid("111@lid")), None, Some(lid("222@lid"))], - &own, - ) - .expect("should succeed"); - let users: Vec<&str> = out.iter().map(|j| j.user.as_str()).collect(); - assert_eq!(users, ["111", "222", "me"]); - } - - #[test] - fn does_not_duplicate_own_when_already_in_list() { - let own = lid("me@lid"); - let out = - assemble_status_participants(vec![Some(lid("111@lid")), Some(lid("me@lid"))], &own) - .expect("should succeed"); - let users: Vec<&str> = out.iter().map(|j| j.user.as_str()).collect(); - assert_eq!(users, ["111", "me"]); - } - - #[test] - fn errors_when_every_recipient_is_unresolvable() { - // Regression guard for the original bug: a single LID-only - // contact used to hard-abort the send with - // `No PN mapping for LID ...`. The new contract is softer — - // individual unresolvable entries are dropped — but we still - // refuse to send when the entire list came back empty, rather - // than silently broadcasting to own devices only. - let own = lid("me@lid"); - let err = assemble_status_participants(vec![None, None, None], &own) - .expect_err("all-None list must error"); - assert!(err.to_string().contains("No valid status recipients")); - } - - #[test] - fn errors_when_list_is_empty() { - let own = lid("me@lid"); - let err = assemble_status_participants(Vec::<Option<Jid>>::new(), &own) - .expect_err("empty list must error"); - assert!(err.to_string().contains("No valid status recipients")); - } - - #[test] - fn strips_device_suffix_from_own_lid() { - // Snapshot lid from the device store carries a device id; the - // participant list uses bare USER JIDs. - let own: Jid = "me:5@lid".parse().unwrap(); - let out = assemble_status_participants(vec![Some(lid("111@lid"))], &own) - .expect("should succeed"); - let me = out - .iter() - .find(|j| j.user.as_str() == "me") - .expect("own LID should be present"); - assert_eq!(me.device, 0, "own LID should be non-ad (device=0)"); - } - } - - mod peer_message_options { - use super::*; - use crate::types::message::{PrivacySensitiveType, PushPriority}; - - fn pdo_message_raw(request_type: i32) -> wa::Message { - wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage { - r#type: Some( - wa::message::protocol_message::Type::PeerDataOperationRequestMessage as i32, - ), - peer_data_operation_request_message: Some( - wa::message::PeerDataOperationRequestMessage { - peer_data_operation_request_type: Some(request_type), - ..Default::default() - }, - ), - ..Default::default() - })), - ..Default::default() - } - } - - fn pdo_message(request_type: wa::message::PeerDataOperationRequestType) -> wa::Message { - pdo_message_raw(request_type as i32) - } - - #[test] - fn pdo_priority_map_matches_wa_web_non_message_requests() { - use wa::message::PeerDataOperationRequestType as PdoType; - - let high_force_cases = [ - (PdoType::GenerateLinkPreview, PushPriority::HighForce, None), - ( - PdoType::PlaceholderMessageResend, - PushPriority::HighForce, - None, - ), - ( - PdoType::HistorySyncOnDemand, - PushPriority::HighForce, - Some(PrivacySensitiveType::OnDemand), - ), - ( - PdoType::CompanionCanonicalUserNonceFetch, - PushPriority::HighForce, - None, - ), - ]; - - for (request_type, push_priority, privacy_sensitive) in high_force_cases { - let options = peer_message_options_from_message(&pdo_message(request_type)); - assert_eq!(options.push_priority(), push_priority, "{request_type:?}"); - assert_eq!( - options.privacy_sensitive(), - privacy_sensitive, - "{request_type:?}" - ); - } - - let default_cases = [ - PdoType::UploadSticker, - PdoType::SendRecentStickerBootstrap, - PdoType::WaffleLinkingNonceFetch, - PdoType::FullHistorySyncOnDemand, - PdoType::CompanionMetaNonceFetch, - PdoType::CompanionSyncdSnapshotFatalRecovery, - PdoType::HistorySyncChunkRetry, - PdoType::GalaxyFlowAction, - PdoType::BusinessBroadcastInsightsDeliveredTo, - PdoType::BusinessBroadcastInsightsRefresh, - ]; - - for request_type in default_cases { - let options = peer_message_options_from_message(&pdo_message(request_type)); - assert_eq!( - options.push_priority(), - PushPriority::High, - "{request_type:?}" - ); - assert_eq!(options.privacy_sensitive(), None, "{request_type:?}"); - } - } - - #[test] - fn non_pdo_and_unknown_pdo_keep_peer_defaults() { - let app_state_key_request = wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage { - r#type: Some( - wa::message::protocol_message::Type::AppStateSyncKeyRequest as i32, - ), - app_state_sync_key_request: Some(wa::message::AppStateSyncKeyRequest { - key_ids: Vec::new(), - }), - ..Default::default() - })), - ..Default::default() - }; - - for msg in [app_state_key_request, pdo_message_raw(99)] { - let options = peer_message_options_from_message(&msg); - assert_eq!(options.push_priority(), PushPriority::High); - assert_eq!(options.privacy_sensitive(), None); - } - } - } - - mod status_carries_privacy_meta { - use super::*; - - #[test] - fn true_for_text_post() { - let msg = wa::Message { - extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { - text: Some("hi".into()), - ..Default::default() - })), - ..Default::default() - }; - assert!(status_carries_privacy_meta(&msg)); - } - - #[test] - fn true_for_image_post() { - let msg = wa::Message { - image_message: Some(Box::new(wa::message::ImageMessage::default())), - ..Default::default() - }; - assert!(status_carries_privacy_meta(&msg)); - } - - #[test] - fn false_for_reaction() { - let msg = wa::Message { - reaction_message: Some(wa::message::ReactionMessage { - text: Some("💚".into()), - ..Default::default() - }), - ..Default::default() - }; - assert!( - !status_carries_privacy_meta(&msg), - "reactions must omit <meta status_setting> (479 SmaxInvalid otherwise)" - ); - } - - #[test] - fn false_for_enc_reaction() { - let msg = wa::Message { - enc_reaction_message: Some(wa::message::EncReactionMessage::default()), - ..Default::default() - }; - assert!(!status_carries_privacy_meta(&msg)); - } - - #[test] - fn false_for_revoke() { - let msg = wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage { - r#type: Some(wa::message::protocol_message::Type::Revoke as i32), - ..Default::default() - })), - ..Default::default() - }; - assert!(!status_carries_privacy_meta(&msg)); - } - - #[test] - fn true_for_non_revoke_protocol_message() { - // Other ProtocolMessage types (e.g., EphemeralSettings) aren't - // reactions and aren't revokes — treat as posts for now. - let msg = wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage { - r#type: Some(wa::message::protocol_message::Type::EphemeralSetting as i32), - ..Default::default() - })), - ..Default::default() - }; - assert!(status_carries_privacy_meta(&msg)); - } - - #[test] - fn false_for_reaction_inside_ephemeral_wrapper() { - let inner = wa::Message { - reaction_message: Some(wa::message::ReactionMessage::default()), - ..Default::default() - }; - let msg = wa::Message { - ephemeral_message: Some(Box::new(wa::message::FutureProofMessage { - message: Some(Box::new(inner)), - })), - ..Default::default() - }; - assert!(!status_carries_privacy_meta(&msg)); - } - - #[test] - fn false_for_revoke_inside_device_sent_wrapper() { - let inner = wa::Message { - protocol_message: Some(Box::new(wa::message::ProtocolMessage { - r#type: Some(wa::message::protocol_message::Type::Revoke as i32), - ..Default::default() - })), - ..Default::default() - }; - let msg = wa::Message { - device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { - destination_jid: Some(String::new()), - message: Some(Box::new(inner)), - ..Default::default() - })), - ..Default::default() - }; - assert!(!status_carries_privacy_meta(&msg)); - } - } - - #[test] - fn build_member_label_message_sets_fields() { - let msg = build_member_label_message("VIP".to_string(), 1_766_847_151); - let pm = msg.protocol_message.as_ref().expect("protocol_message set"); - assert_eq!( - pm.r#type, - Some(wa::message::protocol_message::Type::GroupMemberLabelChange as i32) - ); - let ml = pm.member_label.as_ref().expect("member_label set"); - assert_eq!(ml.label.as_deref(), Some("VIP")); - assert_eq!(ml.label_timestamp, Some(1_766_847_151)); - assert!( - pm.key.is_none(), - "MessageKey must NOT be set (WA Web parity)" - ); - } - - #[test] - fn build_member_label_message_clear_uses_empty_string() { - let msg = build_member_label_message(String::new(), 1); - let ml = msg - .protocol_message - .as_ref() - .unwrap() - .member_label - .as_ref() - .unwrap(); - assert_eq!(ml.label.as_deref(), Some("")); - } - - #[test] - fn build_member_label_message_preserves_unicode() { - let msg = build_member_label_message("🚀 BOT".to_string(), 2); - let ml = msg - .protocol_message - .as_ref() - .unwrap() - .member_label - .as_ref() - .unwrap(); - assert_eq!(ml.label.as_deref(), Some("🚀 BOT")); - } - - /// Mock implementation of SendContextResolver for testing - struct MockSendContextResolver { - /// Pre-key bundles to return: JID -> Option<PreKeyBundle> - prekey_bundles: HashMap<Jid, Option<PreKeyBundle>>, - /// Devices to return from resolve_devices - devices: Vec<Jid>, - /// Phone number to LID mappings for testing LID session lookup - phone_to_lid: HashMap<String, String>, - /// JIDs reported via `on_local_identity_change` (send-path detection). - identity_changes: std::sync::Mutex<Vec<Jid>>, - } - - impl MockSendContextResolver { - fn new() -> Self { - Self { - prekey_bundles: HashMap::new(), - devices: Vec::new(), - phone_to_lid: HashMap::new(), - identity_changes: std::sync::Mutex::new(Vec::new()), - } - } - - fn captured_identity_changes(&self) -> Vec<Jid> { - self.identity_changes.lock().unwrap().clone() - } - - fn with_missing_bundle(mut self, jid: Jid) -> Self { - self.prekey_bundles.insert(jid, None); - self - } - - fn with_bundle(mut self, jid: Jid, bundle: PreKeyBundle) -> Self { - self.prekey_bundles.insert(jid, Some(bundle)); - self - } - - fn with_devices(mut self, devices: Vec<Jid>) -> Self { - self.devices = devices; - self - } - - fn with_phone_to_lid(mut self, phone: &str, lid: &str) -> Self { - self.phone_to_lid.insert(phone.to_string(), lid.to_string()); - self - } - } - - #[async_trait::async_trait] - impl SendContextResolver for MockSendContextResolver { - async fn resolve_devices(&self, _jids: &[Jid]) -> Result<Vec<Jid>> { - Ok(self.devices.clone()) - } - - async fn fetch_prekeys(&self, jids: &[Jid]) -> Result<HashMap<Jid, PreKeyBundle>> { - let mut result = HashMap::new(); - for jid in jids { - if let Some(bundle_opt) = self.prekey_bundles.get(jid) - && let Some(bundle) = bundle_opt - { - result.insert(jid.clone(), bundle.clone()); - } - } - Ok(result) - } - - async fn fetch_prekeys_for_identity_check( - &self, - jids: &[Jid], - ) -> Result<HashMap<Jid, PreKeyBundle>> { - let mut result = HashMap::new(); - for jid in jids { - if let Some(bundle_opt) = self.prekey_bundles.get(jid) - && let Some(bundle) = bundle_opt - { - result.insert(jid.clone(), bundle.clone()); - } - // If None, we intentionally omit it from the result (simulating server not returning it) - } - Ok(result) - } - - async fn resolve_group_info(&self, _jid: &Jid) -> Result<std::sync::Arc<GroupInfo>> { - unimplemented!("resolve_group_info not needed for send.rs tests") - } - - async fn get_lid_for_phone( - &self, - phone_user: &str, - ) -> Option<wacore_binary::CompactString> { - self.phone_to_lid.get(phone_user).map(|s| s.as_str().into()) - } - - fn on_local_identity_change(&self, jid: &Jid) { - self.identity_changes.lock().unwrap().push(jid.clone()); - } - } - - /// Test case: Missing pre-key bundle for a single device skips gracefully - /// - /// When sending to multiple devices, if some don't have pre-key bundles (e.g., Cloud API), - /// we should skip them instead of failing the entire message. - #[test] - fn test_missing_prekey_bundle_skips_device() { - let device_with_bundle: Jid = "1234567890:0@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - let device_without_bundle: Jid = "1234567890:1@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - let cloud_api: Jid = "1234567890:99@hosted" - .parse() - .expect("test JID should be valid"); - - let bundle = create_mock_bundle(); - - let resolver = MockSendContextResolver::new() - .with_bundle(device_with_bundle.clone(), bundle) - .with_missing_bundle(device_without_bundle.clone()) - .with_missing_bundle(cloud_api.clone()) - .with_devices(vec![ - device_with_bundle.clone(), - device_without_bundle.clone(), - cloud_api.clone(), - ]); - - // Check that the resolver correctly returns only available bundles - assert_eq!( - resolver.prekey_bundles.len(), - 3, - "Resolver should have 3 entries" - ); - - // Verify device_with_bundle has a Some(bundle) - assert!( - resolver.prekey_bundles[&device_with_bundle].is_some(), - "device_with_bundle should have a Some entry" - ); - - // Verify others have None - assert!( - resolver.prekey_bundles[&device_without_bundle].is_none(), - "device_without_bundle should have None" - ); - assert!( - resolver.prekey_bundles[&cloud_api].is_none(), - "cloud_api should have None" - ); - - println!("✅ Missing pre-key bundle skips device gracefully"); - } - - /// Test case: All devices missing pre-key bundles - /// - /// If all devices are unavailable, the batch should still complete without panic. - #[test] - fn test_all_devices_missing_prekey_bundles() { - let device1: Jid = "1234567890:0@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - let device2: Jid = "1234567890:1@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - let device3: Jid = "9876543210:0@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - - let resolver = MockSendContextResolver::new() - .with_missing_bundle(device1.clone()) - .with_missing_bundle(device2.clone()) - .with_missing_bundle(device3.clone()) - .with_devices(vec![device1.clone(), device2.clone(), device3.clone()]); - - // All entries should be None - assert!(resolver.prekey_bundles[&device1].is_none()); - assert!(resolver.prekey_bundles[&device2].is_none()); - assert!(resolver.prekey_bundles[&device3].is_none()); - - println!("✅ All devices missing bundles handled gracefully"); - } - - /// Test case: Large group with mixed device availability - /// - /// In real-world scenarios, large groups may have some unavailable devices. - /// The encryption should proceed for available devices and skip unavailable ones. - #[test] - fn test_large_group_with_mixed_device_availability() { - let mut all_devices = Vec::new(); - - for i in 0..10u16 { - let device_jid = Jid::pn_device("1234567890", i); - all_devices.push(device_jid); - } - - let mut resolver = MockSendContextResolver::new().with_devices(all_devices.clone()); - - // Add bundles for devices 0-6, mark 7-9 as missing - for i in 0..10u16 { - let device_jid = Jid::pn_device("1234567890", i); - - if i < 7 { - resolver = resolver.with_bundle(device_jid, create_mock_bundle()); - } else { - resolver = resolver.with_missing_bundle(device_jid); - } - } - - // Verify bundle availability - let available_count = resolver - .prekey_bundles - .values() - .filter(|v| v.is_some()) - .count(); - - assert_eq!(available_count, 7, "Should have 7 available devices"); - assert_eq!( - resolver.prekey_bundles.len(), - 10, - "Should have 10 total entries" - ); - - println!("✅ Large group with 7 available, 3 unavailable devices"); - } - - /// Test case: Cloud API / HOSTED device without pre-key - /// - /// # Context: What are HOSTED devices? - /// - /// HOSTED devices (Cloud API / Meta Business API) are WhatsApp Business accounts - /// that use Meta's server-side infrastructure instead of traditional E2EE. - /// - /// ## Identification: - /// - Device ID 99 (`:99`) on any server - /// - Server `@hosted` or `@hosted.lid` - /// - /// ## Behavior: - /// - They do NOT have Signal protocol prekey bundles - /// - For 1:1 chats: included in device list, but prekey fetch fails gracefully - /// - For groups: proactively filtered out before SKDM distribution - /// - /// This test verifies that when a hosted device is included in the device list - /// (which would happen for 1:1 chats), the missing prekey is handled gracefully. - #[test] - fn test_cloud_api_device_without_prekey() { - let regular_device: Jid = "1234567890:0@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - let cloud_api: Jid = "1234567890:99@hosted" - .parse() - .expect("test JID should be valid"); - - // Verify the cloud_api device is detected as hosted - assert!( - cloud_api.is_hosted(), - "Device with :99@hosted should be detected as hosted" - ); - assert!( - !regular_device.is_hosted(), - "Regular device should NOT be detected as hosted" - ); - - let resolver = MockSendContextResolver::new() - .with_bundle(regular_device.clone(), create_mock_bundle()) - .with_missing_bundle(cloud_api.clone()) - .with_devices(vec![regular_device.clone(), cloud_api.clone()]); - - assert!( - resolver.prekey_bundles[&regular_device].is_some(), - "Regular device should have a bundle" - ); - assert!( - resolver.prekey_bundles[&cloud_api].is_none(), - "Cloud API device should not have a bundle (they don't use Signal protocol)" - ); - - println!("✅ Cloud API device has no prekey bundle (expected behavior)"); - } - - /// Test case: HOSTED devices are filtered from group SKDM distribution - /// - /// # Why filter hosted devices from groups? - /// - /// WhatsApp Web explicitly excludes hosted devices from group message fanout. - /// From the JS code (`getFanOutList`): - /// ```javascript - /// var isHosted = e.id === 99 || e.isHosted === true; - /// var includeInFanout = !isHosted || isOneToOneChat; - /// ``` - /// - /// ## Reasons: - /// 1. Hosted devices don't use Signal protocol - they can't process SKDM - /// 2. Including them causes unnecessary prekey fetch failures - /// 3. Group encryption is handled differently for Cloud API businesses - /// - /// This test verifies that `is_hosted()` correctly identifies devices that - /// should be filtered from group SKDM distribution. - #[test] - fn test_hosted_devices_filtered_from_group_skdm() { - // Simulate devices returned from usync for a group - let devices: Vec<Jid> = vec![ - // Regular devices - should receive SKDM - "5511999887766:0@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), // Primary phone - "5511999887766:33@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), // WhatsApp Web companion - "5521988776655:0@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), // Another participant - "100000012345678:33@lid" - .parse() - .expect("test JID should be valid"), // LID companion device - // HOSTED devices - should be EXCLUDED from group SKDM - "5531977665544:99@s.whatsapp.net" - .parse() - .expect("test JID should be valid"), // Cloud API on regular server - "100000087654321:99@lid" - .parse() - .expect("test JID should be valid"), // Cloud API on LID server - "5541966554433:0@hosted" - .parse() - .expect("test JID should be valid"), // Explicit @hosted server - ]; - - // This is the filtering logic used in prepare_group_stanza - let filtered_for_skdm: Vec<Jid> = - devices.into_iter().filter(|jid| !jid.is_hosted()).collect(); - - assert_eq!( - filtered_for_skdm.len(), - 4, - "Should have 4 devices after filtering out hosted devices" - ); - - // Verify all remaining devices are NOT hosted - for jid in &filtered_for_skdm { - assert!( - !jid.is_hosted(), - "Filtered list should not contain hosted device: {}", - jid - ); - } - - // Verify specific devices are included/excluded by checking struct fields - // (Device ID 0 is not serialized in the string representation) - let has_primary_phone = filtered_for_skdm - .iter() - .any(|j| j.user == "5511999887766" && j.device == 0 && j.server == "s.whatsapp.net"); - let has_companion = filtered_for_skdm - .iter() - .any(|j| j.user == "5511999887766" && j.device == 33 && j.server == "s.whatsapp.net"); - let has_cloud_api = filtered_for_skdm - .iter() - .any(|j| j.user == "5531977665544" && j.device == 99); - let has_hosted_server = filtered_for_skdm.iter().any(|j| j.server == "hosted"); - - assert!(has_primary_phone, "Primary phone should be included"); - assert!(has_companion, "WhatsApp Web companion should be included"); - assert!( - !has_cloud_api, - "Cloud API device (ID 99) should be excluded" - ); - assert!( - !has_hosted_server, - "@hosted server device should be excluded" - ); - - println!("✅ Hosted devices correctly filtered from group SKDM distribution"); - } - - /// Test case: Device recovery between retries - /// - /// If a device was temporarily unavailable, a retry should succeed. - #[test] - fn test_device_recovery_between_requests() { - let device: Jid = "1234567890:0@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - - // First attempt: device unavailable - let resolver_first = MockSendContextResolver::new().with_missing_bundle(device.clone()); - - assert!( - resolver_first.prekey_bundles[&device].is_none(), - "First attempt: device should be unavailable" - ); - - // Second attempt: device recovered - let resolver_second = - MockSendContextResolver::new().with_bundle(device.clone(), create_mock_bundle()); - - assert!( - resolver_second.prekey_bundles[&device].is_some(), - "Second attempt: device should be available" - ); - - println!("✅ Device recovery between retries works correctly"); - } - - /// Helper function to create a mock PreKeyBundle with valid types - fn create_mock_bundle() -> PreKeyBundle { - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - let identity_pair = IdentityKeyPair::generate(&mut rng); - let signed_prekey_pair = KeyPair::generate(&mut rng); - let prekey_pair = KeyPair::generate(&mut rng); - - PreKeyBundle::new( - 1, // registration_id - 1u32.into(), // device_id - Some((1u32.into(), prekey_pair.public_key)), // pre_key - 2u32.into(), // signed_pre_key_id - signed_prekey_pair.public_key, - vec![0u8; 64], - *identity_pair.identity_key(), - ) - .expect("Failed to create PreKeyBundle") - } - - // These tests validate the fix for the LID-PN session mismatch issue. - // When a message is received with sender_lid, the session is stored under the LID address. - // When sending a reply using the phone number, we must reuse the existing LID session - // instead of creating a new PN session, otherwise subsequent messages will fail with - // MAC verification errors. - - /// Test that phone_to_lid mapping returns the cached LID mapping. - /// - /// This verifies the MockSendContextResolver correctly stores phone-to-LID - /// mappings used for LID session lookup. - #[test] - fn test_mock_resolver_phone_to_lid_mapping() { - let phone = "559980000001"; - let lid = "100000012345678"; - - let resolver = MockSendContextResolver::new().with_phone_to_lid(phone, lid); - - // Access the HashMap directly (synchronous) - let result = resolver.phone_to_lid.get(phone).cloned(); - - assert!(result.is_some(), "Should return LID for known phone"); - assert_eq!( - result.expect("known phone should return LID"), - lid, - "Should return correct LID" - ); - - // Unknown phone should return None - let unknown = resolver.phone_to_lid.get("999999999").cloned(); - assert!(unknown.is_none(), "Should return None for unknown phone"); - - println!("✅ MockSendContextResolver phone_to_lid mapping works correctly"); - } - - /// Test that the resolver correctly maps phone numbers to LIDs. - /// - /// This is a building block for the session lookup logic. - #[test] - fn test_phone_to_lid_mapping_multiple_users() { - let resolver = MockSendContextResolver::new() - .with_phone_to_lid("559980000001", "100000012345678") - .with_phone_to_lid("559980000002", "100000024691356") - .with_phone_to_lid("559980000003", "100000037037034"); - - // Verify all mappings using direct HashMap access - let lid1 = resolver.phone_to_lid.get("559980000001").cloned(); - let lid2 = resolver.phone_to_lid.get("559980000002").cloned(); - let lid3 = resolver.phone_to_lid.get("559980000003").cloned(); - - assert_eq!( - lid1.expect("phone 1 should have LID mapping"), - "100000012345678" - ); - assert_eq!( - lid2.expect("phone 2 should have LID mapping"), - "100000024691356" - ); - assert_eq!( - lid3.expect("phone 3 should have LID mapping"), - "100000037037034" - ); - - println!("✅ Multiple phone-to-LID mappings work correctly"); - } - - /// Test the scenario that caused the original bug: - /// - Session exists under LID address (from receiving a message with sender_lid) - /// - Send to PN address should reuse the LID session, not create a new one - /// - /// This test verifies the logic flow, though full integration testing - /// requires the actual encrypt_for_devices function with real sessions. - #[test] - fn test_lid_session_lookup_scenario() { - // Scenario setup: - // - Received message from 559980000001@s.whatsapp.net with sender_lid=100000012345678@lid - // - Session was stored under 100000012345678.0 - // - Now sending reply to 559980000001@s.whatsapp.net - // - Should look up LID and check for session under 100000012345678.0 - - let phone = "559980000001"; - let lid = "100000012345678"; - let device_id = 0u16; - - let resolver = MockSendContextResolver::new().with_phone_to_lid(phone, lid); - - // Simulate the device JID we're trying to send to (PN format) - let pn_device_jid = Jid::pn_device(phone, device_id); - - // Step 1: Look up LID for the phone number (using direct HashMap access) - let lid_user = resolver - .phone_to_lid - .get(pn_device_jid.user.as_str()) - .cloned(); - assert!(lid_user.is_some(), "Should find LID for phone"); - let lid_user = lid_user.expect("phone should have LID mapping"); - - // Step 2: Construct the LID JID with same device ID - let lid_jid = Jid::lid_device(lid_user.clone(), pn_device_jid.device); - - // Step 3: Verify the LID JID is correctly constructed - assert_eq!(lid_jid.user, lid, "LID user should match"); - assert_eq!(lid_jid.server, "lid", "Server should be 'lid'"); - assert_eq!(lid_jid.device, device_id, "Device ID should be preserved"); - - // Step 4: Convert to protocol addresses and verify they're different - use crate::types::jid::JidExt; - let pn_address = pn_device_jid.to_protocol_address(); - let lid_address = lid_jid.to_protocol_address(); - - assert_ne!( - pn_address.name(), - lid_address.name(), - "PN and LID addresses should have different names" - ); - assert_eq!( - pn_address.device_id(), - lid_address.device_id(), - "Device IDs should match" - ); - - println!("✅ LID session lookup scenario works correctly:"); - println!(" - PN JID: {} -> Address: {}", pn_device_jid, pn_address); - println!(" - LID JID: {} -> Address: {}", lid_jid, lid_address); - println!(" - Would check for session under LID address first"); - } - - /// Test that companion device IDs are preserved in LID JID construction. - /// - /// WhatsApp Web uses device ID 33, and this must be preserved when - /// constructing the LID JID for session lookup. - #[test] - fn test_lid_jid_preserves_companion_device_id() { - let phone = "559980000001"; - let lid = "100000012345678"; - let companion_device_id = 33u16; // WhatsApp Web device ID - - let resolver = MockSendContextResolver::new().with_phone_to_lid(phone, lid); - - // Simulate sending to a companion device (WhatsApp Web) - let pn_device_jid = Jid::pn_device(phone, companion_device_id); - - // Look up LID using direct HashMap access - let lid_user = resolver - .phone_to_lid - .get(pn_device_jid.user.as_str()) - .cloned(); - - // Construct LID JID - let lid_jid = Jid::lid_device( - lid_user.expect("phone should have LID mapping for companion test"), - pn_device_jid.device, - ); - - assert_eq!( - lid_jid.device, companion_device_id, - "Device ID 33 should be preserved" - ); - assert_eq!(lid_jid.to_string(), "100000012345678:33@lid"); - - println!("✅ Companion device ID (33) correctly preserved in LID JID"); - } - - /// Test that LID lookup only applies to s.whatsapp.net JIDs. - /// - /// LID JIDs (@lid) and group JIDs (@g.us) should not trigger LID lookup. - #[test] - fn test_lid_lookup_only_for_pn_jids() { - let _resolver = - MockSendContextResolver::new().with_phone_to_lid("559980000001", "100000012345678"); - - // These JIDs should NOT trigger LID lookup - let lid_jid: Jid = "100000012345678:0@lid" - .parse() - .expect("test JID should be valid"); - let group_jid: Jid = "120363123456789012@g.us" - .parse() - .expect("test JID should be valid"); - - // Only s.whatsapp.net JIDs should be looked up - assert_ne!( - lid_jid.server, "s.whatsapp.net", - "LID JID should not be s.whatsapp.net" - ); - assert_ne!( - group_jid.server, "s.whatsapp.net", - "Group JID should not be s.whatsapp.net" - ); - - // PN JID should be eligible for lookup - let pn_jid: Jid = "559980000001:0@s.whatsapp.net" - .parse() - .expect("test JID should be valid"); - assert_eq!( - pn_jid.server, "s.whatsapp.net", - "PN JID should be s.whatsapp.net" - ); - - println!("✅ LID lookup correctly limited to s.whatsapp.net JIDs"); - } - - /// Test case: Regression test for self-encryption bug. - /// - /// The sender's own device (e.g. device 79) must be excluded from the encryption list - /// to prevent "SESSION BASE KEY CHANGED" warnings caused by establishing a session with oneself. - #[test] - fn test_dm_encryption_excludes_sender_device() { - // Setup: - // - Own user: 123456789 - // - Specific own device (Sender): 79 - // - Other own device: 0 - // - Recipient: 987654321 - - let own_user = "123456789"; - let own_device_id = 79; - - // Own JID (Sender) - let own_jid = Jid::lid_device(own_user.to_string(), own_device_id); - - // Simulate devices returned by resolver.resolve_devices() - // This includes: - // 1. The sender's own device (should be excluded) - // 2. Another device of the sender (should be in own_other_devices) - // 3. The recipient's device (should be in recipient_devices) - let all_devices: Vec<Jid> = vec![ - Jid::lid_device(own_user.to_string(), own_device_id), // Sender (79) - Jid::lid_device(own_user.to_string(), 0), // Other own device (0) - Jid::lid_device("987654321".to_string(), 0), // Recipient - ]; - - let (recipient_devices, own_other_devices) = - partition_dm_devices(all_devices, &own_jid, None); - - // Verifications - - // 1. Sender device (79) should NOT be in either list - let sender_in_own = own_other_devices.iter().any(|d| d.device == own_device_id); - let sender_in_recipient = recipient_devices.iter().any(|d| d.device == own_device_id); - - assert!( - !sender_in_own, - "Sender device (79) should be excluded from own_other_devices" - ); - assert!( - !sender_in_recipient, - "Sender device (79) should be excluded from recipient_devices" - ); - - // 2. Other own device (0) MUST be in own_other_devices - let other_own_present = own_other_devices - .iter() - .any(|d| d.device == 0 && d.user == own_user); - assert!( - other_own_present, - "Other own device (0) should be included in own_other_devices" - ); - - // 3. Recipient MUST be in recipient_devices - let recipient_present = recipient_devices.iter().any(|d| d.user == "987654321"); - assert!( - recipient_present, - "Recipient should be included in recipient_devices" - ); - - println!("✅ Self-encryption regression test passed: Sender device correctly excluded."); - } - - #[test] - fn test_dm_encryption_treats_own_lid_devices_as_self() { - let own_pn = Jid::pn_device("559980000001".to_string(), 18); - let own_lid = Jid::lid_device("123456789012345".to_string(), 18); - - let all_devices = vec![ - Jid::lid_device("123456789012345".to_string(), 18), // Exact sender device via LID - Jid::lid_device("123456789012345".to_string(), 0), // Other own device via LID - Jid::lid_device("987654321012345".to_string(), 0), // Recipient - ]; - - let (recipient_devices, own_other_devices) = - partition_dm_devices(all_devices, &own_pn, Some(&own_lid)); - - assert!( - !own_other_devices - .iter() - .any(|d| d.user == own_lid.user && d.device == 18), - "Exact sender LID device should be excluded from own_other_devices" - ); - assert!( - !recipient_devices - .iter() - .any(|d| d.user == own_lid.user && d.device == 18), - "Exact sender LID device should be excluded from recipient_devices" - ); - assert!( - own_other_devices - .iter() - .any(|d| d.user == own_lid.user && d.device == 0), - "Other own LID devices should be routed through DSM as own_other_devices" - ); - assert!( - recipient_devices - .iter() - .any(|d| d.user == "987654321012345" && d.device == 0), - "Non-self devices must remain in recipient_devices" - ); - } - - /// Test case: LID Prekey Lookup Normalization - /// - /// Verifies that when looking up pre-key bundles for LID JIDs, the lookup key - /// is normalized (agent=0) to match how the bundles are stored in the map. - /// - /// This validates the fix for "No pre-key bundle returned" when the requested JID - /// has non-standard agent/server fields but the bundle is stored under the normalized key. - #[test] - fn test_lid_prekey_lookup_normalization() { - // 1. Define JIDs - // The JID we request (simulating what comes from resolve_devices or elsewhere) - // Let's pretend it has agent=1 to simulate a mismatch - let mut requested_jid = Jid::lid_device("123456789".to_string(), 0); - requested_jid.agent = 1; - - // The normalized JID (how it's stored in the bundle map) - let normalized_jid = Jid::lid_device("123456789".to_string(), 0); // agent=0 by default - - // 2. Setup Resolver - // Store the bundle under the NORMALIZED key (agent=0) - let resolver = MockSendContextResolver::new() - .with_bundle(normalized_jid.clone(), create_mock_bundle()) - .with_devices(vec![requested_jid.clone()]); - - // 3. Verify Mock Setup - // Ensure bundle is accessible via normalized key but NOT via requested (raw) key - // This confirms our test condition is valid (that implicit lookup would fail) - assert!( - resolver.prekey_bundles.contains_key(&normalized_jid), - "Setup: bundle should exist for normalized key" - ); - assert!( - !resolver.prekey_bundles.contains_key(&requested_jid), - "Setup: bundle should NOT exist for requested raw key" - ); - - // 4. Test logic mirroring `encrypt_for_devices` - let mut jid_to_encryption_jid = HashMap::new(); - // Assume direct mapping for simplicity - jid_to_encryption_jid.insert(requested_jid.clone(), requested_jid.clone()); - - // Get the bundles map (mocks `fetch_prekeys_for_identity_check`) - // The mock implementation returns the map as-is filtered by keys. - // HOWEVER, `fetch_prekeys` usually takes a list. - // In `encrypt_for_devices`, we call: - // let prekey_bundles = resolver.fetch_prekeys_for_identity_check(&[requested_jid]).await?; - - // Let's simulate what `fetch_prekeys_for_identity_check` would return. - // Our mock implementation `fetch_prekeys` logic: - // if let Some(bundle_opt) = self.prekey_bundles.get(jid) - - // Wait, if the mock follows exact HashMap lookup, `fetch_prekeys(&[requested_jid])` - // will return EMPTY because `requested_jid` is not in `prekey_bundles`. - // The REAL `fetch_prekeys` (in `client.rs` -> `prekeys.rs`) sends an IQ to the server, - // and the server response is parsed. The parsing logic (in `prekeys.rs`) normalizes the key. - // So the HashMap returned by `fetch_prekeys` will contain NORMALIZED keys. - - // So for this test to be accurate, we must simulate that `fetch_prekeys` returned a map - // where the key is NORMALIZED, even if we asked for `requested_jid`? - // Actually, `PreKeyFetchSpec` asks for JIDs. The response contains JIDs. - // If we ask for `agent=1`, does the server return `agent=1`? - // The logs showed: - // parsed: `...:82@lid` (agent=0 probably, or just not printed?) - // lookup: `...` (failed) - - // The critical part is that the `HashMap` returned by `resolver.fetch_prekeys` - // definitely contains the bundle under some key. - // If `prekeys.rs` normalizes it, it's under the normalized key. - // The `encrypt_for_devices` logic has: - // `match prekey_bundles.get(device_jid)` - // where `device_jid` is the one from the loop (requested_jid). - - // If `fetch_prekeys` returns a map with `normalized_jid`, and we lookup `requested_jid`, it fails. - // My fix was to normalize `requested_jid` before lookup. - - // So I need to construct the `prekey_bundles` map manually here to simulate the return from fetch. - let mut prekey_bundles = HashMap::new(); - prekey_bundles.insert(normalized_jid.clone(), create_mock_bundle()); - - // Now test the logic: - let device_jid = &requested_jid; - - // -- Logic from fix -- - // Use centralized normalization logic - let lookup_jid = device_jid.normalize_for_prekey_bundle(); - - // Fix: Use the normalized device_jid to lookup the bundle - let bundle = prekey_bundles.get(&lookup_jid); - // -------------------- - - assert!(bundle.is_some(), "Should find bundle after normalization"); - - // Verify it would have failed without normalization - let raw_lookup = prekey_bundles.get(device_jid); - assert!( - raw_lookup.is_none(), - "Should NOT find bundle without normalization" - ); - - println!("✅ LID Prekey Lookup Normalization passed"); - } - - mod group_retry { - use super::*; - use crate::libsignal::protocol::{ - Direction, IdentityChange, IdentityKey, IdentityKeyPair, IdentityKeyStore, KeyPair, - PreKeyBundle, ProtocolAddress, SessionStore, process_prekey_bundle, - }; - use crate::types::message::AddressingMode; - use std::collections::HashMap; - use wacore_binary::NodeContent; - - struct MemSessionStore(HashMap<ProtocolAddress, Vec<u8>>); - impl MemSessionStore { - fn new() -> Self { - Self(HashMap::new()) - } - } - #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] - #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] - impl SessionStore for MemSessionStore { - async fn load_session( - &self, - a: &ProtocolAddress, - ) -> crate::libsignal::protocol::error::Result< - Option<crate::libsignal::protocol::SessionRecord>, - > { - Ok(self - .0 - .get(a) - .and_then(|b| crate::libsignal::protocol::SessionRecord::deserialize(b).ok())) - } - async fn has_session( - &self, - a: &ProtocolAddress, - ) -> crate::libsignal::protocol::error::Result<bool> { - Ok(self.0.contains_key(a)) - } - async fn store_session( - &mut self, - a: &ProtocolAddress, - r: crate::libsignal::protocol::SessionRecord, - ) -> crate::libsignal::protocol::error::Result<()> { - self.0.insert(a.clone(), r.serialize()?); - Ok(()) - } - } - - struct MemIdentityStore { - pair: IdentityKeyPair, - reg_id: u32, - known: HashMap<ProtocolAddress, IdentityKey>, - } - #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] - #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] - impl IdentityKeyStore for MemIdentityStore { - async fn get_identity_key_pair( - &self, - ) -> crate::libsignal::protocol::error::Result<IdentityKeyPair> { - Ok(self.pair.clone()) - } - async fn get_local_registration_id( - &self, - ) -> crate::libsignal::protocol::error::Result<u32> { - Ok(self.reg_id) - } - async fn save_identity( - &mut self, - a: &ProtocolAddress, - id: &IdentityKey, - ) -> crate::libsignal::protocol::error::Result<IdentityChange> { - self.known.insert(a.clone(), *id); - Ok(IdentityChange::from_changed(false)) - } - async fn is_trusted_identity( - &self, - _: &ProtocolAddress, - _: &IdentityKey, - _: Direction, - ) -> crate::libsignal::protocol::error::Result<bool> { - Ok(true) - } - async fn get_identity( - &self, - a: &ProtocolAddress, - ) -> crate::libsignal::protocol::error::Result<Option<IdentityKey>> { - Ok(self.known.get(a).copied()) - } - } - - async fn setup_session() -> (MemSessionStore, MemIdentityStore, Jid) { - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - let sender = IdentityKeyPair::generate(&mut rng); - let receiver = IdentityKeyPair::generate(&mut rng); - let spk = KeyPair::generate(&mut rng); - let opk = KeyPair::generate(&mut rng); - let sig = receiver - .private_key() - .calculate_signature(&spk.public_key.serialize(), &mut rng) - .unwrap(); - let bundle = PreKeyBundle::new( - 1, - 1u32.into(), - Some((1u32.into(), opk.public_key)), - 1u32.into(), - spk.public_key, - sig.to_vec(), - *receiver.identity_key(), - ) - .unwrap(); - let jid: Jid = "559911112222@s.whatsapp.net".parse().unwrap(); - let addr = jid.to_protocol_address(); - let mut ss = MemSessionStore::new(); - let mut is = MemIdentityStore { - pair: sender, - reg_id: 42, - known: HashMap::new(), - }; - process_prekey_bundle( - &addr, - &mut ss, - &mut is, - &bundle, - &mut rand::make_rng::<rand::rngs::StdRng>(), - crate::libsignal::protocol::UsePQRatchet::No, - ) - .await - .unwrap(); - (ss, is, jid) - } - - #[tokio::test] - async fn group_retry_pkmsg_with_account_emits_device_identity() { - let (mut ss, mut is, jid) = setup_session().await; - 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( - &mut ss, - &mut is, - group.clone(), - p.clone(), - p.clone(), - &wa::Message::default(), - "3EB0ABC".into(), - 1, - Some(&account), - AddressingMode::Pn, - None, - ) - .await - .unwrap(); - - assert_eq!(n.tag, "message"); - let mut a = n.attrs(); - assert_eq!(a.optional_string("to").unwrap().as_ref(), group.to_string()); - assert_eq!( - a.optional_string("participant").unwrap().as_ref(), - p.to_string() - ); - // Default (empty) message falls through to "media" per WA Web's typeAttributeFromProtobuf - assert_eq!( - a.optional_string("type").unwrap().as_ref(), - stanza::MSG_TYPE_MEDIA - ); - assert!(a.optional_string("category").is_none()); - assert_eq!(a.optional_string("addressing_mode").unwrap().as_ref(), "pn"); - let enc = n.get_optional_child("enc").unwrap(); - let mut ea = enc.attrs(); - assert_eq!( - ea.optional_string("v").unwrap().as_ref(), - stanza::ENC_VERSION - ); - assert_eq!( - ea.optional_string("type").unwrap().as_ref(), - stanza::ENC_TYPE_PKMSG - ); - assert_eq!(ea.optional_string("count").unwrap().as_ref(), "1"); - assert!(matches!(&enc.content, Some(NodeContent::Bytes(_)))); - assert!( - n.get_optional_child("device-identity").is_some(), - "pkmsg group retry with account must include <device-identity>" - ); - } - - /// Symmetric to peer/dm pre-flights: refuse group retry pkmsg when - /// account is missing rather than silently dropping device-identity. - #[tokio::test] - async fn group_retry_pkmsg_preflight_errors_when_account_missing() { - let (mut ss, mut is, jid) = setup_session().await; - let group: Jid = "120363098765432100@g.us".parse().unwrap(); - let p: Jid = jid.to_string().parse().unwrap(); - - let before = ss - .load_session(&p.to_protocol_address()) - .await - .unwrap() - .expect("pre-condition: session present") - .serialize() - .expect("serialize before"); - - let result = prepare_group_retry_stanza( - &mut ss, - &mut is, - group, - p.clone(), - p.clone(), - &wa::Message::default(), - "grp-retry-no-account".into(), - 1, - None, - AddressingMode::Pn, - None, - ) - .await; - let err = result.expect_err("group retry pkmsg must reject missing account"); - assert!( - err.to_string().contains("device-identity"), - "error must name <device-identity>; got: {err}" - ); - - let after = ss - .load_session(&p.to_protocol_address()) - .await - .unwrap() - .expect("session still present") - .serialize() - .expect("serialize after"); - assert_eq!( - before, after, - "group retry pre-flight must leave the session byte-identical" - ); - } - - /// Pins the WAWebSendMsgCreateDeviceStanza retry shape: `<enc>` - /// directly under `<message>` plus a `recipient` attribute. - /// Pre-fix this regressed to the fanout shape and the server - /// rejected every retry with 479. - #[tokio::test] - async fn dm_retry_emits_enc_directly_under_message_with_recipient() { - let (mut ss, mut is, jid) = setup_session().await; - // Distinct values so a swapped-args regression (e.g. `recipient = - // to_jid`) fails the assertions below instead of silently passing. - let to: Jid = "559922223333:5@s.whatsapp.net".parse().unwrap(); - 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( - &mut ss, - &mut is, - to.clone(), - Some(recipient.clone()), - requester, - &wa::Message::default(), - "dm-retry-format-1".into(), - 1, - Some(&account), - None, - ) - .await - .unwrap(); - - assert_eq!(n.tag, "message"); - // <enc> is a direct child — no <participants> wrapper. - assert!( - n.get_optional_child("participants").is_none(), - "DM retry must not wrap <enc> in <participants> \ - (matches WAWebSendMsgCreateDeviceStanza)" - ); - assert!( - n.get_optional_child("enc").is_some(), - "<enc> must be a direct child of <message>" - ); - assert_eq!( - n.attrs().optional_string("to").unwrap().as_ref(), - to.to_string(), - "`to` should target the requesting device verbatim" - ); - assert_eq!( - n.attrs().optional_string("recipient").unwrap().as_ref(), - recipient.to_string(), - "`recipient` should mirror the original message's recipient \ - (forwarded from the retry receipt's `recipient` attr)" - ); - } - - #[tokio::test] - async fn dm_retry_pkmsg_targets_single_device() { - let (mut ss, mut is, jid) = setup_session().await; - let to: Jid = "559922223333@s.whatsapp.net".parse().unwrap(); - let encryption = jid.clone(); - let account = pkmsg_account_proto(); - - let n = prepare_dm_retry_stanza( - &mut ss, - &mut is, - to.clone(), - Some(to.clone()), - encryption, - &wa::Message::default(), - "dm-retry-1".into(), - 1, - Some(&account), - None, - ) - .await - .unwrap(); - - assert_eq!(n.tag, "message"); - let mut attrs = n.attrs(); - assert_eq!( - attrs.optional_string("to").unwrap().as_ref(), - to.to_string() - ); - assert_eq!( - attrs.optional_string("recipient").unwrap().as_ref(), - to.to_string() - ); - assert_eq!(attrs.optional_string("id").unwrap().as_ref(), "dm-retry-1"); - assert_eq!( - attrs.optional_string("type").unwrap().as_ref(), - stanza::MSG_TYPE_MEDIA - ); - assert!(attrs.optional_string("participant").is_none()); - assert!(attrs.optional_string("addressing_mode").is_none()); - - // `<enc>` is a direct child of `<message>` (no `<participants>` wrapper). - assert!(n.get_optional_child("participants").is_none()); - let enc = n.get_optional_child("enc").unwrap(); - let mut enc_attrs = enc.attrs(); - assert_eq!( - enc_attrs.optional_string("type").unwrap().as_ref(), - stanza::ENC_TYPE_PKMSG - ); - assert_eq!(enc_attrs.optional_string("count").unwrap().as_ref(), "1"); - assert!( - n.get_optional_child("device-identity").is_some(), - "pkmsg DM retry with account must include <device-identity>" - ); - } - - #[tokio::test] - async fn dm_retry_pkmsg_with_account_has_device_identity() { - let (mut ss, mut is, jid) = setup_session().await; - let to: Jid = "559922223333@s.whatsapp.net".parse().unwrap(); - let acc = wa::AdvSignedDeviceIdentity { - details: Some(b"t".to_vec()), - ..Default::default() - }; - - let n = prepare_dm_retry_stanza( - &mut ss, - &mut is, - to.clone(), - Some(to), - jid, - &wa::Message::default(), - "dm-retry-2".into(), - 2, - Some(&acc), - None, - ) - .await - .unwrap(); - - let enc = n.get_optional_child("enc").unwrap(); - assert_eq!( - enc.attrs().optional_string("type").unwrap().as_ref(), - stanza::ENC_TYPE_PKMSG - ); - assert_eq!(enc.attrs().optional_string("count").unwrap().as_ref(), "2"); - assert!(n.get_optional_child("device-identity").is_some()); - } - - #[tokio::test] - async fn pkmsg_with_account_has_device_identity() { - let (mut ss, mut is, jid) = setup_session().await; - let group: Jid = "120363098765432100@g.us".parse().unwrap(); - let p: Jid = jid.to_string().parse().unwrap(); - let acc = wa::AdvSignedDeviceIdentity { - details: Some(b"t".to_vec()), - ..Default::default() - }; - let n = prepare_group_retry_stanza( - &mut ss, - &mut is, - group, - p.clone(), - p, - &wa::Message::default(), - "id2".into(), - 2, - Some(&acc), - AddressingMode::Pn, - None, - ) - .await - .unwrap(); - assert_eq!( - n.get_optional_child("enc") - .unwrap() - .attrs() - .optional_string("type") - .unwrap() - .as_ref(), - stanza::ENC_TYPE_PKMSG - ); - assert!(n.get_optional_child("device-identity").is_some()); - assert_eq!( - n.attrs() - .optional_string("addressing_mode") - .unwrap() - .as_ref(), - "pn" - ); - } - - #[tokio::test] - async fn lid_addressing_mode() { - let (mut ss, mut is, jid) = setup_session().await; - 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( - &mut ss, - &mut is, - group, - p.clone(), - p, - &wa::Message::default(), - "m2".into(), - 3, - Some(&wa::AdvSignedDeviceIdentity::default()), - AddressingMode::Lid, - None, - ) - .await - .unwrap(); - let mut ea = n.get_optional_child("enc").unwrap().attrs(); - assert_eq!(ea.optional_string("count").unwrap().as_ref(), "3"); - assert_eq!( - n.attrs() - .optional_string("addressing_mode") - .unwrap() - .as_ref(), - "lid" - ); - } - - #[tokio::test] - async fn group_retry_preserves_edit_attribute() { - let (mut ss, mut is, jid) = setup_session().await; - 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( - &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), - ) - .await - .unwrap(); - assert_eq!(n.attrs().optional_string("edit").unwrap().as_ref(), "8"); - } - - #[tokio::test] - async fn dm_retry_preserves_edit_attribute() { - 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( - &mut ss, - &mut is, - to.clone(), - Some(to), - jid, - &wa::Message::default(), - "edit-1".into(), - 1, - Some(&account), - Some(crate::types::message::EditAttribute::MessageEdit), - ) - .await - .unwrap(); - assert_eq!(n.attrs().optional_string("edit").unwrap().as_ref(), "1"); - } - - #[tokio::test] - async fn retry_without_edit_omits_attribute() { - let (mut ss, mut is, jid) = setup_session().await; - 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( - &mut ss, - &mut is, - group, - p.clone(), - p, - &wa::Message::default(), - "plain-1".into(), - 1, - Some(&account), - AddressingMode::Lid, - None, - ) - .await - .unwrap(); - assert!(n.attrs().optional_string("edit").is_none()); - } - - // Peer pkmsg layout: `[<meta appdata="default"/>, <enc>, <device-identity>]`. - // Without `<device-identity>` the phone XMPP-acks but its Signal - // layer skips session promotion. Mirrors whatsmeow's - // `preparePeerMessageNode`. - - fn pkmsg_account_proto() -> wa::AdvSignedDeviceIdentity { - // Opaque placeholder bytes — the assertions only check that - // the element carries non-empty content. - wa::AdvSignedDeviceIdentity { - details: Some(vec![0u8; 32]), - account_signature_key: Some(vec![0u8; 32]), - account_signature: Some(vec![0u8; 64]), - device_signature: Some(vec![0u8; 64]), - } - } - - async fn build_peer_stanza( - account: Option<&wa::AdvSignedDeviceIdentity>, - ) -> wacore_binary::Node { - build_peer_stanza_with_options(account, PeerMessageOptions::default()).await - } - - async fn build_peer_stanza_with_options( - account: Option<&wa::AdvSignedDeviceIdentity>, - options: PeerMessageOptions, - ) -> wacore_binary::Node { - let (mut ss, mut is, jid) = setup_session().await; - let addr = jid.to_protocol_address(); - prepare_peer_stanza_with_options( - &mut ss, - &mut is, - jid.clone(), - &addr, - &wa::Message::default(), - "peer-test-1".into(), - account, - options, - ) - .await - .expect("peer stanza builds") - } - - #[tokio::test] - async fn peer_pkmsg_includes_meta_and_device_identity() { - let account = pkmsg_account_proto(); - let n = build_peer_stanza(Some(&account)).await; - - assert_eq!(n.tag, "message"); - assert_eq!( - n.attrs().optional_string("category").unwrap().as_ref(), - "peer" - ); - assert_eq!( - n.attrs().optional_string("push_priority").unwrap().as_ref(), - "high" - ); - assert!(n.attrs().optional_string("privacy_sensitive").is_none()); - - let children = n.children().expect("peer message has children"); - let tags: Vec<&str> = children.iter().map(|c| c.tag.as_ref()).collect(); - // Layout matches whatsmeow's preparePeerMessageNode for pkmsg: - // [<meta>, <enc>, <device-identity>]. - assert_eq!( - tags, - vec!["meta", "enc", "device-identity"], - "peer pkmsg children order/identity must match whatsmeow" - ); - - let meta = n.get_optional_child("meta").expect("meta present"); - assert_eq!( - meta.attrs().optional_string("appdata").unwrap().as_ref(), - "default", - "<meta appdata=\"default\"/> is what the phone uses to route the peer payload" - ); - - let enc = n.get_optional_child("enc").expect("enc present"); - assert_eq!( - enc.attrs().optional_string("type").unwrap().as_ref(), - "pkmsg", - "fresh session must produce pkmsg, not msg" - ); - - let device_identity = n - .get_optional_child("device-identity") - .expect("device-identity present"); - match &device_identity.content { - Some(NodeContent::Bytes(b)) => assert!( - !b.is_empty(), - "device-identity content must be the proto-encoded \ - AdvSignedDeviceIdentity, not empty" - ), - other => panic!("device-identity must carry bytes, got {other:?}"), - } - } - - #[tokio::test] - async fn peer_stanza_carries_high_force_and_privacy_attrs() { - let account = pkmsg_account_proto(); - let n = build_peer_stanza_with_options( - Some(&account), - PeerMessageOptions::high_force_on_demand(), - ) - .await; - - assert_eq!( - n.attrs().optional_string("push_priority").unwrap().as_ref(), - "high_force" - ); - assert_eq!( - n.attrs() - .optional_string("privacy_sensitive") - .unwrap() - .as_ref(), - "1" - ); - } - - #[tokio::test] - async fn peer_pkmsg_errors_when_account_missing_without_ratchet_advance() { - // Pkmsg without <device-identity> would reproduce the deadlock — - // refuse AND prove the session is byte-identical after the failed - // call so the next retry has the same ratchet position. - let (mut ss, mut is, jid) = setup_session().await; - let addr = jid.to_protocol_address(); - - let before = ss - .load_session(&addr) - .await - .unwrap() - .expect("pre-condition: session loaded") - .serialize() - .expect("serialize before"); - - let result = prepare_peer_stanza( - &mut ss, - &mut is, - jid.clone(), - &addr, - &wa::Message::default(), - "peer-test-no-account".into(), - None, - ) - .await; - let err = result.expect_err("pkmsg path must reject missing account"); - assert!( - err.to_string().contains("device-identity"), - "error must name the missing element; got: {err}" - ); - - let after = ss - .load_session(&addr) - .await - .unwrap() - .expect("session still present after failed call") - .serialize() - .expect("serialize after"); - assert_eq!( - before, after, - "session record must be byte-identical after a failed prepare — \ - any difference means a ratchet step was committed for a stanza we couldn't ship" - ); - } - - /// Pre-flight check: when no session exists and account is None, - /// `prepare_peer_stanza` must refuse before `message_encrypt` runs, - /// otherwise the sender chain is persisted for a stanza we cannot ship - /// (CodeRabbit-flagged ratchet-burn-on-fail-fast). - #[tokio::test] - async fn peer_pkmsg_preflight_no_ratchet_burn_without_session() { - let jid: Jid = "559911112222@s.whatsapp.net".parse().unwrap(); - let addr = jid.to_protocol_address(); - let mut ss = MemSessionStore::new(); - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - let mut is = MemIdentityStore { - pair: IdentityKeyPair::generate(&mut rng), - reg_id: 42, - known: HashMap::new(), - }; - - assert!( - !ss.has_session(&addr).await.unwrap(), - "precondition: store has no session for this address" - ); - - let result = prepare_peer_stanza( - &mut ss, - &mut is, - jid.clone(), - &addr, - &wa::Message::default(), - "peer-preflight-1".into(), - None, - ) - .await; - let err = result.expect_err("must refuse before message_encrypt"); - assert!( - err.to_string().contains("device-identity"), - "error must name <device-identity>; got: {err}" - ); - assert!( - !ss.has_session(&addr).await.unwrap(), - "pre-flight must NOT advance/persist a session — the ratchet \ - must remain unburned for the retry attempt" - ); - } - - /// Symmetric to peer_pkmsg_preflight: prepare_dm_retry_stanza must - /// also refuse to ship pkmsg without <device-identity>, otherwise - /// message_encrypt would advance the sender chain for a stanza the - /// peer's Signal layer cannot promote. - #[tokio::test] - async fn dm_retry_pkmsg_preflight_errors_when_account_missing() { - let (mut ss, mut is, jid) = setup_session().await; - let addr = jid.to_protocol_address(); - - let before = ss - .load_session(&addr) - .await - .unwrap() - .expect("pre-condition: session present") - .serialize() - .expect("serialize before"); - - let to: Jid = "559922223333@s.whatsapp.net".parse().unwrap(); - let result = prepare_dm_retry_stanza( - &mut ss, - &mut is, - to.clone(), - Some(to), - jid.clone(), - &wa::Message::default(), - "dm-retry-no-account".into(), - 1, - None, - None, - ) - .await; - let err = result.expect_err("DM retry pkmsg path must reject missing account"); - assert!( - err.to_string().contains("device-identity"), - "error must name <device-identity>; got: {err}" - ); - - let after = ss - .load_session(&addr) - .await - .unwrap() - .expect("session still present") - .serialize() - .expect("serialize after"); - assert_eq!( - before, after, - "DM retry pre-flight must leave the session byte-identical" - ); - } - - /// Production's SessionAdapter::load_session has TAKE semantics - /// (SignalStoreCache marks the slot CheckedOut until store_session - /// puts the record back). If the pre-flight only loads without - /// restoring, the slot stays stranded and message_encrypt sees no - /// session. The mock here mirrors that contract via interior - /// mutability (Mutex) on the &self load_session. - #[tokio::test] - async fn preflight_restores_session_with_take_store_semantics() { - use std::collections::{HashMap, HashSet}; - use std::sync::Mutex; - - struct TakeStore { - inner: Mutex<TakeInner>, - } - struct TakeInner { - present: HashMap<ProtocolAddress, Vec<u8>>, - taken: HashSet<ProtocolAddress>, - } - impl TakeStore { - fn from(ss: &MemSessionStore) -> Self { - Self { - inner: Mutex::new(TakeInner { - present: ss.0.clone(), - taken: HashSet::new(), - }), - } - } - fn is_present(&self, addr: &ProtocolAddress) -> bool { - let g = self.inner.lock().unwrap(); - g.present.contains_key(addr) && !g.taken.contains(addr) - } - } - #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] - #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] - impl SessionStore for TakeStore { - async fn load_session( - &self, - a: &ProtocolAddress, - ) -> crate::libsignal::protocol::error::Result< - Option<crate::libsignal::protocol::SessionRecord>, - > { - let mut g = self.inner.lock().unwrap(); - if g.taken.contains(a) { - return Ok(None); - } - let rec = g.present.get(a).and_then(|b| { - crate::libsignal::protocol::SessionRecord::deserialize(b).ok() - }); - if rec.is_some() { - g.taken.insert(a.clone()); - } - Ok(rec) - } - async fn has_session( - &self, - a: &ProtocolAddress, - ) -> crate::libsignal::protocol::error::Result<bool> { - let g = self.inner.lock().unwrap(); - Ok(g.present.contains_key(a) && !g.taken.contains(a)) - } - async fn store_session( - &mut self, - a: &ProtocolAddress, - r: crate::libsignal::protocol::SessionRecord, - ) -> crate::libsignal::protocol::error::Result<()> { - let mut g = self.inner.lock().unwrap(); - g.present.insert(a.clone(), r.serialize()?); - g.taken.remove(a); - Ok(()) - } - } - - let (mem_ss, mut is, jid) = setup_session().await; - let mut ss = TakeStore::from(&mem_ss); - let addr = jid.to_protocol_address(); - - // setup_session leaves pending_pre_key set, so account=None - // would bail. Use Some(account) — pre-flight still runs - // load+restore because it's gated on account.is_none() at the - // call site; switch to account=None and we want the assertion - // to verify that the BAIL path also restores the slot. - assert!( - ss.is_present(&addr), - "precondition: session is Present before pre-flight" - ); - - // Drive the bail path: account=None + session has pending_pre_key - // → pre-flight bails. Even on bail, the loaded record must be - // put back so a retry with Some(account) doesn't see a stranded slot. - let bail = prepare_peer_stanza( - &mut ss, - &mut is, - jid.clone(), - &addr, - &wa::Message::default(), - "preflight-take-bail".into(), - None, - ) - .await; - bail.expect_err("must bail with account=None on a pending-pkmsg session"); - assert!( - ss.is_present(&addr), - "pre-flight bail path must still restore the checked-out session" - ); - - // And the pass path: with Some(account), the pre-flight still - // does load+restore, then message_encrypt runs successfully. - let account = pkmsg_account_proto(); - let ok = prepare_peer_stanza( - &mut ss, - &mut is, - jid.clone(), - &addr, - &wa::Message::default(), - "preflight-take-pass".into(), - Some(&account), - ) - .await; - ok.expect("peer stanza builds with Some(account)"); - assert!( - ss.is_present(&addr), - "session must be Present after a successful encrypt+store" - ); - } - } - - mod decrypt_fail { - use super::*; - - #[test] - fn regular_message() { - let msg = wa::Message { - conversation: Some("hi".into()), - ..Default::default() - }; - assert!(!should_hide_decrypt_fail(&msg)); - } - - #[test] - fn reaction() { - let msg = wa::Message { - reaction_message: Some(Default::default()), - ..Default::default() - }; - assert!(should_hide_decrypt_fail(&msg)); - } - - #[test] - fn pin() { - let msg = wa::Message { - pin_in_chat_message: Some(Default::default()), - ..Default::default() - }; - assert!(should_hide_decrypt_fail(&msg)); - } - - #[test] - fn poll_vote() { - let msg = wa::Message { - poll_update_message: Some(wa::message::PollUpdateMessage { - vote: Some(Default::default()), - ..Default::default() - }), - ..Default::default() - }; - assert!(should_hide_decrypt_fail(&msg)); - } - - #[test] - fn poll_update_without_vote() { - let msg = wa::Message { - poll_update_message: Some(Default::default()), - ..Default::default() - }; - assert!(!should_hide_decrypt_fail(&msg)); - } - - #[test] - fn reaction_inside_ephemeral_wrapper() { - let msg = wa::Message { - ephemeral_message: Some(Box::new(wa::message::FutureProofMessage { - message: Some(Box::new(wa::Message { - reaction_message: Some(Default::default()), - ..Default::default() - })), - })), - ..Default::default() - }; - assert!(should_hide_decrypt_fail(&msg)); - } - - #[test] - fn conditional_reveal() { - let msg = wa::Message { - conditional_reveal_message: Some(Default::default()), - ..Default::default() - }; - assert!(should_hide_decrypt_fail(&msg)); - } - - #[test] - fn poll_add_option_edit() { - use wa::message::secret_encrypted_message::SecretEncType; - let msg = wa::Message { - secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { - secret_enc_type: Some(SecretEncType::PollAddOption as i32), - ..Default::default() - }), - ..Default::default() - }; - assert!(should_hide_decrypt_fail(&msg)); - } - } - - mod decrypt_fail_for_send { - use super::*; - use crate::types::message::EditAttribute; - - fn plain() -> wa::Message { - wa::Message { - conversation: Some("hi".into()), - ..Default::default() - } - } - - #[test] - fn sender_revoke_is_not_hidden() { - assert!(!should_hide_decrypt_fail_for_send( - Some(&EditAttribute::SenderRevoke), - &plain() - )); - } - - #[test] - fn admin_revoke_is_not_hidden() { - assert!(!should_hide_decrypt_fail_for_send( - Some(&EditAttribute::AdminRevoke), - &plain() - )); - } - - #[test] - fn message_edit_is_hidden() { - assert!(should_hide_decrypt_fail_for_send( - Some(&EditAttribute::MessageEdit), - &plain() - )); - } - - #[test] - fn revoke_does_not_block_content_based_hide() { - // A reaction still hides on its own merits even under a revoke edit. - let msg = wa::Message { - reaction_message: Some(Default::default()), - ..Default::default() - }; - assert!(should_hide_decrypt_fail_for_send( - Some(&EditAttribute::SenderRevoke), - &msg - )); - } - } - - mod stanza_type { - use super::*; - use wa::message::secret_encrypted_message::SecretEncType; - - fn secret(enc: SecretEncType) -> wa::Message { - wa::Message { - secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { - secret_enc_type: Some(enc as i32), - ..Default::default() - }), - ..Default::default() - } - } - - #[test] - fn poll_add_option_edit_is_poll() { - assert_eq!( - stanza_type_from_message(&secret(SecretEncType::PollAddOption)), - stanza::MSG_TYPE_POLL - ); - } - - #[test] - fn poll_edit_is_poll() { - assert_eq!( - stanza_type_from_message(&secret(SecretEncType::PollEdit)), - stanza::MSG_TYPE_POLL - ); - } - - #[test] - fn album_is_text() { - let msg = wa::Message { - album_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&msg), stanza::MSG_TYPE_TEXT); - } - - // Helpers for wrapper tests. WA Web's typeAttributeFromProtobuf unwraps - // FutureProofMessage wrappers (via getUnwrappedProtobufMessage) and then - // classifies the inner message. - fn fpm(inner: wa::Message) -> Box<wa::message::FutureProofMessage> { - Box::new(wa::message::FutureProofMessage { - message: Some(Box::new(inner)), - }) - } - fn text_inner() -> wa::Message { - wa::Message { - conversation: Some("hi".to_string()), - ..Default::default() - } - } - fn image_inner() -> wa::Message { - wa::Message { - image_message: Some(Box::default()), - ..Default::default() - } - } - - #[test] - fn group_status_v2_classifies_by_inner() { - let txt = wa::Message { - group_status_message_v2: Some(fpm(text_inner())), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&txt), stanza::MSG_TYPE_TEXT); - - // Regression guard: forcing this wrapper to "text" dropped the - // mediatype and silently dropped the stanza. WA Web unwraps it and - // sends type="media" mediatype="image". - let img = wa::Message { - group_status_message_v2: Some(fpm(image_inner())), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&img), stanza::MSG_TYPE_MEDIA); - assert_eq!(media_type_from_message(&img), Some("image")); - } - - #[test] - fn group_status_v2_empty_is_media() { - // An empty wrapper is not one of WA Web's four re-checked wrappers - // (ephemeral/groupMentioned/botInvoke/deviceSent), so it falls through - // to the media default in both WA Web and here. - let m = wa::Message { - group_status_message_v2: Some(Box::default()), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&m), stanza::MSG_TYPE_MEDIA); - } - - #[test] - fn payment_family_is_text() { - // Payment family classifies as text; the media default would be dropped. - let cases = [ - wa::Message { - request_payment_message: Some(Box::default()), - ..Default::default() - }, - wa::Message { - send_payment_message: Some(Box::default()), - ..Default::default() - }, - wa::Message { - decline_payment_request_message: Some(Default::default()), - ..Default::default() - }, - wa::Message { - cancel_payment_request_message: Some(Default::default()), - ..Default::default() - }, - wa::Message { - payment_invite_message: Some(Default::default()), - ..Default::default() - }, - ]; - for m in cases { - assert_eq!(media_type_from_message(&m), None); - assert_eq!(stanza_type_from_message(&m), stanza::MSG_TYPE_TEXT); - } - } - - #[test] - fn backfilled_wrappers_classify_by_inner() { - let spoiler = wa::Message { - spoiler_message: Some(fpm(text_inner())), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&spoiler), stanza::MSG_TYPE_TEXT); - - let status_mention = wa::Message { - status_mention_message: Some(fpm(image_inner())), - ..Default::default() - }; - assert_eq!( - stanza_type_from_message(&status_mention), - stanza::MSG_TYPE_MEDIA - ); - assert_eq!(media_type_from_message(&status_mention), Some("image")); - - let question = wa::Message { - question_message: Some(fpm(text_inner())), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&question), stanza::MSG_TYPE_TEXT); - - let group_status_v1 = wa::Message { - group_status_message: Some(fpm(text_inner())), - ..Default::default() - }; - assert_eq!( - stanza_type_from_message(&group_status_v1), - stanza::MSG_TYPE_TEXT - ); - } - - #[test] - fn nested_wrappers_reach_innermost() { - // ephemeral { viewOnceV2 { image } } -> media + mediatype. - let inner = wa::Message { - view_once_message_v2: Some(fpm(image_inner())), - ..Default::default() - }; - let m = wa::Message { - ephemeral_message: Some(fpm(inner)), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&m), stanza::MSG_TYPE_MEDIA); - assert_eq!(media_type_from_message(&m), Some("image")); - } - - #[test] - fn preserved_classifier_branches() { - let r = wa::Message { - reaction_message: Some(Default::default()), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&r), stanza::MSG_TYPE_REACTION); - - let ev = wa::Message { - event_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&ev), stanza::MSG_TYPE_EVENT); - - let poll = wa::Message { - poll_creation_message_v3: Some(Box::default()), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&poll), stanza::MSG_TYPE_POLL); - - assert_eq!( - stanza_type_from_message(&text_inner()), - stanza::MSG_TYPE_TEXT - ); - assert_eq!( - stanza_type_from_message(&image_inner()), - stanza::MSG_TYPE_MEDIA - ); - - let proto = wa::Message { - protocol_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&proto), stanza::MSG_TYPE_TEXT); - - let url = wa::Message { - extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { - matched_text: Some("https://example.com".to_string()), - ..Default::default() - })), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&url), stanza::MSG_TYPE_MEDIA); - } - - #[test] - fn interactive_and_list_types_get_their_mediatype() { - // WA Web's mediaTypeFromProtobuf maps these to concrete mediatypes; - // omitting the attribute makes the server drop the type="media" stanza. - let list = wa::Message { - list_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!(stanza_type_from_message(&list), stanza::MSG_TYPE_MEDIA); - assert_eq!(media_type_from_message(&list), Some("list")); - - let list_response = wa::Message { - list_response_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!( - media_type_from_message(&list_response), - Some("list_response") - ); - - let buttons_response = wa::Message { - buttons_response_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!( - media_type_from_message(&buttons_response), - Some("buttons_response") - ); - - let order = wa::Message { - order_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!(media_type_from_message(&order), Some("order")); - - let product = wa::Message { - product_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!(media_type_from_message(&product), Some("product")); - - let interactive_response = wa::Message { - interactive_response_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!( - media_type_from_message(&interactive_response), - Some("native_flow_response") - ); - - let history_bundle = wa::Message { - message_history_bundle: Some(Box::default()), - ..Default::default() - }; - assert_eq!( - media_type_from_message(&history_bundle), - Some("group_history") - ); - } - - #[test] - fn buttons_message_has_no_mediatype() { - // WA Web maps buttonsMessage to EncMediaType.Button, but its string - // mapper has no Button case (returns null/DROP_ATTR), so the attribute - // is omitted. Adding a "buttons" mediatype would diverge from WA Web. - let buttons = wa::Message { - buttons_message: Some(Box::default()), - ..Default::default() - }; - assert_eq!(media_type_from_message(&buttons), None); - } - - #[test] - fn ephemeral_wrapped_list_reaches_list_mediatype() { - let m = wa::Message { - ephemeral_message: Some(fpm(wa::Message { - list_message: Some(Box::default()), - ..Default::default() - })), - ..Default::default() - }; - assert_eq!(media_type_from_message(&m), Some("list")); - } - - #[test] - fn top_level_lottie_sticker_is_terminal_sticker() { - // WA Web's mediaTypeFromProtobuf treats a top-level lottieStickerMessage - // as a terminal "sticker" and does NOT recurse into it, unlike the - // stanza-type path which unwraps it. - let lottie = wa::Message { - lottie_sticker_message: Some(fpm(image_inner())), - ..Default::default() - }; - assert_eq!(media_type_from_message(&lottie), Some("sticker")); - } - } - - #[cfg(test)] - mod device_unregistered_tests { - use super::is_device_unregistered_error; - use crate::request::ServerErrorCode; - - #[test] - fn detects_406_server_error_code() { - let err = anyhow::Error::new(ServerErrorCode { - code: 406, - text: "not-acceptable".to_string(), - }); - assert!(is_device_unregistered_error(&err)); - } - - #[test] - fn rejects_non_406_server_error() { - let err = anyhow::Error::new(ServerErrorCode { - code: 404, - text: "not-found".to_string(), - }); - assert!(!is_device_unregistered_error(&err)); - } - - #[test] - fn rejects_unrelated_error() { - let err = anyhow::anyhow!("some random error"); - assert!(!is_device_unregistered_error(&err)); - } - - #[test] - fn rejects_wacore_iq_error_without_server_error_code_wrapper() { - // wacore::IqError::ServerError is NOT the same as ServerErrorCode. - // This simulates the old bug: if someone wraps wacore IqError directly - // without the ServerErrorCode wrapper, the check should not match. - let err = anyhow::Error::new(crate::request::IqError::ServerError { - code: 406, - text: "not-acceptable".to_string(), - }); - // This would only match if we also checked IqError (we don't — we use ServerErrorCode) - // The SendContextResolver impl is responsible for wrapping in ServerErrorCode - assert!(!is_device_unregistered_error(&err)); - } - } - - mod collect_stale_device_users { - use super::super::collect_stale_device_users; - use crate::client::context::GroupInfo; - use crate::types::message::AddressingMode; - use std::collections::{HashMap, HashSet}; - use wacore_binary::{CompactString, Jid}; - - fn lid_device(user: &str, dev: u16) -> Jid { - Jid::lid_device(user.to_string(), dev) - } - - fn pn_user(user: &str) -> Jid { - Jid::pn(user) - } - - fn group_info_lid(mapping: &[(&str, &str)]) -> GroupInfo { - let mut info = GroupInfo::new(Vec::new(), AddressingMode::Lid); - if !mapping.is_empty() { - let mut map: HashMap<CompactString, Jid> = HashMap::new(); - for (lid_user, pn) in mapping { - map.insert(CompactString::from(*lid_user), pn_user(pn)); - } - info.set_lid_to_pn_map(map); - } - info - } - - #[test] - fn emits_lid_and_pn_alias_when_mapping_known() { - let info = group_info_lid(&[("100000000000001", "15550000001")]); - let dist = vec![lid_device("100000000000001", 5)]; - let out = collect_stale_device_users(Some(&dist), &[], &info); - let set: HashSet<String> = out.into_iter().collect(); - assert!(set.contains("100000000000001")); - assert!(set.contains("15550000001")); - assert_eq!(set.len(), 2); - } - - #[test] - fn emits_only_lid_when_mapping_unknown() { - let info = group_info_lid(&[]); - let dist = vec![lid_device("100000000000002", 7)]; - let out = collect_stale_device_users(Some(&dist), &[], &info); - assert_eq!(out, vec!["100000000000002".to_string()]); - } - - #[test] - fn dedups_multiple_devices_of_same_user() { - let info = group_info_lid(&[("100000000000003", "15550000003")]); - let dist = vec![ - lid_device("100000000000003", 1), - lid_device("100000000000003", 2), - lid_device("100000000000003", 3), - ]; - let out = collect_stale_device_users(Some(&dist), &[], &info); - let set: HashSet<String> = out.into_iter().collect(); - assert_eq!(set.len(), 2); - assert!(set.contains("100000000000003")); - assert!(set.contains("15550000003")); - } - - #[test] - fn skips_successfully_encrypted_devices() { - let info = group_info_lid(&[]); - let encrypted = lid_device("100000000000004", 5); - let dist = vec![encrypted.clone(), lid_device("100000000000005", 5)]; - let encrypted_set = vec![encrypted]; - let out = collect_stale_device_users(Some(&dist), &encrypted_set, &info); - assert_eq!(out, vec!["100000000000005".to_string()]); - } - - #[test] - fn pn_mode_group_does_not_emit_alias() { - // In PN-mode groups the distribution list is already PN-form, so - // there's no LID↔PN duality to emit. - let mut info = GroupInfo::new(Vec::new(), AddressingMode::Pn); - let mut map: HashMap<CompactString, Jid> = HashMap::new(); - map.insert( - CompactString::from("100000000000006"), - pn_user("15550000006"), - ); - info.set_lid_to_pn_map(map); - let dist = vec![Jid::pn_device("15550000006", 3)]; - let out = collect_stale_device_users(Some(&dist), &[], &info); - assert_eq!(out, vec!["15550000006".to_string()]); - } - - #[test] - fn skips_non_pn_alias() { - // If phone_jid_for_lid_user returns a JID whose server isn't PN - // (malformed/adversarial server response), do not emit it. - let mut info = GroupInfo::new(Vec::new(), AddressingMode::Lid); - let mut map: HashMap<CompactString, Jid> = HashMap::new(); - map.insert( - CompactString::from("100000000000007"), - Jid::lid("100000000000099"), - ); - info.set_lid_to_pn_map(map); - let dist = vec![lid_device("100000000000007", 5)]; - let out = collect_stale_device_users(Some(&dist), &[], &info); - assert_eq!(out, vec!["100000000000007".to_string()]); - } - - #[test] - fn empty_distribution_list_yields_empty() { - let info = group_info_lid(&[]); - let out = collect_stale_device_users(None, &[], &info); - assert!(out.is_empty()); - let out = collect_stale_device_users(Some(&[]), &[], &info); - assert!(out.is_empty()); - } - } - - /// Item 2 — WA Web `markHasSenderKey(x, M)`: a key-distributing group send - /// marks the FULL SKDM target set `has_key=true`, not only the devices that - /// encrypted successfully. A device whose SKDM encryption fails (no session - /// and no bundle, mimicking a 406) must still land in - /// `PreparedGroupStanza.skdm_devices`, so the next send does not re-target - /// it every time (the fan-out storm); the retry-receipt path repairs any - /// device that is actually alive and keyless. - mod mark_full_distribution_list { - use super::*; - use crate::libsignal::protocol::{ - Direction, IdentityChange, IdentityKey, IdentityKeyStore, PreKeyId, PreKeyRecord, - PreKeyStore, ProtocolAddress, SenderKeyRecord, SenderKeyStore, SessionStore, - SignedPreKeyId, SignedPreKeyRecord, SignedPreKeyStore, UsePQRatchet, - process_prekey_bundle, - }; - use crate::libsignal::store::sender_key_name::SenderKeyName; - use crate::runtime::{AbortHandle, Runtime}; - use crate::types::jid::JidExt; - use crate::types::message::AddressingMode; - use std::future::Future; - use std::pin::Pin; - use std::time::Duration; - - type SigResult<T> = crate::libsignal::protocol::error::Result<T>; - - #[derive(Clone, Default)] - struct MemSessionStore(HashMap<ProtocolAddress, Vec<u8>>); - #[async_trait::async_trait] - impl SessionStore for MemSessionStore { - async fn load_session( - &self, - a: &ProtocolAddress, - ) -> SigResult<Option<crate::libsignal::protocol::SessionRecord>> { - Ok(self - .0 - .get(a) - .and_then(|b| crate::libsignal::protocol::SessionRecord::deserialize(b).ok())) - } - async fn has_session(&self, a: &ProtocolAddress) -> SigResult<bool> { - Ok(self.0.contains_key(a)) - } - async fn store_session( - &mut self, - a: &ProtocolAddress, - r: crate::libsignal::protocol::SessionRecord, - ) -> SigResult<()> { - self.0.insert(a.clone(), r.serialize()?); - Ok(()) - } - } - - #[derive(Clone)] - struct MemIdentityStore { - pair: IdentityKeyPair, - reg_id: u32, - known: HashMap<ProtocolAddress, IdentityKey>, - } - #[async_trait::async_trait] - impl IdentityKeyStore for MemIdentityStore { - async fn get_identity_key_pair(&self) -> SigResult<IdentityKeyPair> { - Ok(self.pair.clone()) - } - async fn get_local_registration_id(&self) -> SigResult<u32> { - Ok(self.reg_id) - } - async fn save_identity( - &mut self, - a: &ProtocolAddress, - id: &IdentityKey, - ) -> SigResult<IdentityChange> { - self.known.insert(a.clone(), *id); - Ok(IdentityChange::from_changed(false)) - } - async fn is_trusted_identity( - &self, - _: &ProtocolAddress, - _: &IdentityKey, - _: Direction, - ) -> SigResult<bool> { - Ok(true) - } - async fn get_identity(&self, a: &ProtocolAddress) -> SigResult<Option<IdentityKey>> { - Ok(self.known.get(a).copied()) - } - } - - #[derive(Default)] - struct MemSenderKeyStore(HashMap<SenderKeyName, SenderKeyRecord>); - #[async_trait::async_trait] - impl SenderKeyStore for MemSenderKeyStore { - async fn store_sender_key( - &mut self, - n: &SenderKeyName, - r: SenderKeyRecord, - ) -> SigResult<()> { - self.0.insert(n.clone(), r); - Ok(()) - } - async fn load_sender_key( - &self, - n: &SenderKeyName, - ) -> SigResult<Option<SenderKeyRecord>> { - Ok(self.0.get(n).cloned()) - } - } - - // Outgoing group encryption never consumes our own prekeys, and device B - // has no bundle (so no session is established for it) — these are never - // called; present only to satisfy the generic bounds. - struct UnusedPreKeyStore; - #[async_trait::async_trait] - impl PreKeyStore for UnusedPreKeyStore { - async fn get_pre_key(&self, _: PreKeyId) -> SigResult<PreKeyRecord> { - unreachable!("prekey store not used in outgoing group encrypt") - } - async fn save_pre_key(&mut self, _: PreKeyId, _: &PreKeyRecord) -> SigResult<()> { - unreachable!() - } - async fn remove_pre_key(&mut self, _: PreKeyId) -> SigResult<()> { - unreachable!() - } - } - struct UnusedSignedPreKeyStore; - #[async_trait::async_trait] - impl SignedPreKeyStore for UnusedSignedPreKeyStore { - async fn get_signed_pre_key(&self, _: SignedPreKeyId) -> SigResult<SignedPreKeyRecord> { - unreachable!("signed prekey store not used in outgoing group encrypt") - } - async fn save_signed_pre_key( - &mut self, - _: SignedPreKeyId, - _: &SignedPreKeyRecord, - ) -> SigResult<()> { - unreachable!() - } - } - - struct TokioTestRuntime; - #[async_trait::async_trait] - impl Runtime for TokioTestRuntime { - fn spawn( - &self, - future: Pin<Box<dyn Future<Output = ()> + Send + 'static>>, - ) -> AbortHandle { - let handle = tokio::spawn(future); - AbortHandle::new(move || handle.abort()) - } - fn sleep(&self, _d: Duration) -> Pin<Box<dyn Future<Output = ()> + Send>> { - // Not exercised on the send path; wacore dev-deps omit tokio's - // "time" feature, so resolve immediately rather than time out. - Box::pin(async {}) - } - fn spawn_blocking( - &self, - f: Box<dyn FnOnce() + Send + 'static>, - ) -> Pin<Box<dyn Future<Output = ()> + Send>> { - Box::pin(async move { - let _ = tokio::task::spawn_blocking(f).await; - }) - } - fn yield_now(&self) -> Option<Pin<Box<dyn Future<Output = ()> + Send>>> { - None - } - } - - // Establish a real Signal session for `a` so its SKDM encrypts; the - // returned identity store is the sender's (knows `a` after X3DH). - async fn established_stores(a: &Jid) -> (MemSessionStore, MemIdentityStore) { - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - let sender = IdentityKeyPair::generate(&mut rng); - let receiver = IdentityKeyPair::generate(&mut rng); - let spk = KeyPair::generate(&mut rng); - let opk = KeyPair::generate(&mut rng); - let sig = receiver - .private_key() - .calculate_signature(&spk.public_key.serialize(), &mut rng) - .unwrap(); - let bundle = PreKeyBundle::new( - 1, - 1u32.into(), - Some((1u32.into(), opk.public_key)), - 1u32.into(), - spk.public_key, - sig.to_vec(), - *receiver.identity_key(), - ) - .unwrap(); - let mut ss = MemSessionStore::default(); - let mut is = MemIdentityStore { - pair: sender, - reg_id: 42, - known: HashMap::new(), - }; - process_prekey_bundle( - &a.to_protocol_address(), - &mut ss, - &mut is, - &bundle, - &mut rng, - UsePQRatchet::No, - ) - .await - .unwrap(); - (ss, is) - } - - #[tokio::test] - async fn failed_device_is_still_marked_has_key() { - let group: Jid = "120363000000000001@g.us".parse().unwrap(); - let own_jid: Jid = "559900000000@s.whatsapp.net".parse().unwrap(); - let own_lid: Jid = "100000000000000@lid".parse().unwrap(); - // A has a session (encrypts ok); B has neither session nor bundle, - // mimicking a device that 406'd / has no key material. - let a: Jid = "559911112222:0@s.whatsapp.net".parse().unwrap(); - let b: Jid = "559933334444:0@s.whatsapp.net".parse().unwrap(); - - let (mut ss, mut is) = established_stores(&a).await; - let mut sks = MemSenderKeyStore::default(); - let mut pks = UnusedPreKeyStore; - let spks = UnusedSignedPreKeyStore; - let mut stores = SignalStores { - sender_key_store: &mut sks, - session_store: &mut ss, - identity_store: &mut is, - prekey_store: &mut pks, - signed_prekey_store: &spks, - }; - - // Empty resolver: no LID overrides; B's prekey fetch returns nothing - // → B is dropped by the encrypt fan-out (not in encrypted_devices). - let resolver = MockSendContextResolver::new(); - let rt = TokioTestRuntime; - - let group_info = GroupInfo::new( - vec![own_jid.to_non_ad(), a.to_non_ad(), b.to_non_ad()], - AddressingMode::Pn, - ); - let msg = wa::Message { - conversation: Some("hi".into()), - ..Default::default() - }; - - let prepared = prepare_group_stanza( - &rt, - &mut stores, - &resolver, - &group_info, - &own_jid, - &own_lid, - None, - group, - &msg, - "TESTREQID".into(), - false, - Some(vec![a.clone(), b.clone()]), - None, - None, - &[], - ) - .await - .expect("prepare_group_stanza should succeed even when a device fails to encrypt"); - - let marked: std::collections::HashSet<String> = prepared - .skdm_devices - .iter() - .map(|j| j.to_string()) - .collect(); - - assert!( - marked.contains(&a.to_string()), - "device that encrypted must be marked" - ); - assert!( - marked.contains(&b.to_string()), - "device whose SKDM encryption FAILED must still be marked has_key \ - (WA Web markHasSenderKey(x, M) marks the full target set → no re-fanout storm)" - ); - assert_eq!( - prepared.skdm_devices.len(), - 2, - "exactly the full distribution list (A + B), not just the encrypted subset" - ); - - // A key-distributing send must carry a phash (computed over the list). - assert!( - prepared.node.attrs().optional_string("phash").is_some(), - "a key-distributing group send must carry a phash" - ); - } - } - - /// Item 3 — phash device-set construction. The set hashed is the full - /// recipient list PLUS the sending device (which is never in the recipient - /// list, since we don't SKDM ourselves), matching WA Web - /// `phashV2([].concat(A, [B]))`. - /// - /// This was confirmed against a real WA Web capture sent to the production - /// server: the recipient `<to>` set plus the sending device reproduced the - /// exact `phash` on the wire, while the recipient set alone did not — so the - /// sending device is part of the hash. Raw identifiers are not committed - /// (PII); the vectors below are fictitious but exercise the same logic. - mod group_phash_golden { - use super::*; - - #[test] - fn phash_set_includes_sending_device() { - // Fictitious group: a few users with bare (device 0) + companion - // devices. The self user appears as a companion (device 0) in the - // recipient list; its SENDING device (24) is excluded, mirroring a - // real send (we never SKDM ourselves). - let recipients: Vec<Jid> = [ - "100000000000001@lid", - "100000000000001:5@lid", - "100000000000002@lid", - "100000000000003@lid", - "100000000000003:12@lid", - "100000000000099@lid", - ] - .iter() - .map(|s| s.parse().expect("valid LID jid")) - .collect(); - - let own_sending: Jid = "100000000000099:24@lid".parse().unwrap(); - assert!( - !recipients - .iter() - .any(|j: &Jid| j.user == "100000000000099" && j.device == 24), - "the sending device must not already be in the recipient list" - ); - - let set = build_group_phash_set(&recipients, &own_sending); - assert_eq!(set.len(), 7, "6 recipients + the sending device"); - - // Dropping the sending device changes the hash, proving it is part - // of the hashed set (WA Web `[].concat(A, [B])`). - let with_self = MessageUtils::participant_list_hash(&set).unwrap(); - let without_self = MessageUtils::participant_list_hash(&recipients).unwrap(); - assert_ne!(with_self, without_self); - - // Deterministic standard-base64 vectors (regression guard). - assert_eq!(without_self, "2:rZoSAdIV"); - assert_eq!(with_self, "2:sti8OtHX"); - } - - #[test] - fn phash_set_drops_hosted_devices() { - // Hosted (Cloud API) devices don't take part in group E2EE and must - // not enter the phash, mirroring the SKDM distribution filter. - let with_hosted: Vec<Jid> = ["100000000000001@lid", "100000000000002:99@hosted"] - .iter() - .map(|s| s.parse().expect("valid jid")) - .collect(); - let without_hosted: Vec<Jid> = ["100000000000001@lid"] - .iter() - .map(|s| s.parse().expect("valid jid")) - .collect(); - let own: Jid = "100000000000099:24@lid".parse().unwrap(); - - assert_eq!( - build_group_phash_set(&with_hosted, &own), - build_group_phash_set(&without_hosted, &own), - "hosted devices must not affect the phash set" - ); - } - } - - mod local_identity_change_on_send { - use super::*; - use crate::libsignal::protocol::{ - Direction, IdentityChange, IdentityKey, IdentityKeyStore, PreKeyId, PreKeyRecord, - PreKeyStore, ProtocolAddress, SenderKeyRecord, SessionRecord, SessionStore, - SignedPreKeyId, SignedPreKeyRecord, SignedPreKeyStore, - }; - use crate::runtime::{AbortHandle, Runtime}; - use crate::types::jid::JidExt; - use std::future::Future; - use std::pin::Pin; - use std::time::Duration; - - type SigResult<T> = crate::libsignal::protocol::error::Result<T>; - - #[derive(Clone, Default)] - struct MemSessionStore(HashMap<ProtocolAddress, Vec<u8>>); - #[async_trait::async_trait] - impl SessionStore for MemSessionStore { - async fn load_session(&self, a: &ProtocolAddress) -> SigResult<Option<SessionRecord>> { - Ok(self - .0 - .get(a) - .and_then(|b| SessionRecord::deserialize(b).ok())) - } - async fn has_session(&self, a: &ProtocolAddress) -> SigResult<bool> { - Ok(self.0.contains_key(a)) - } - async fn store_session( - &mut self, - a: &ProtocolAddress, - r: SessionRecord, - ) -> SigResult<()> { - self.0.insert(a.clone(), r.serialize()?); - Ok(()) - } - } - - /// Identity store that reports the real change (unlike the hardcoded - /// stub elsewhere), so a pre-seeded stale key surfaces as ReplacedExisting. - #[derive(Clone)] - struct MemIdentityStore { - pair: IdentityKeyPair, - known: HashMap<ProtocolAddress, IdentityKey>, - } - #[async_trait::async_trait] - impl IdentityKeyStore for MemIdentityStore { - async fn get_identity_key_pair(&self) -> SigResult<IdentityKeyPair> { - Ok(self.pair.clone()) - } - async fn get_local_registration_id(&self) -> SigResult<u32> { - Ok(42) - } - async fn save_identity( - &mut self, - a: &ProtocolAddress, - id: &IdentityKey, - ) -> SigResult<IdentityChange> { - let changed = self.known.get(a).is_some_and(|k| k != id); - self.known.insert(a.clone(), *id); - Ok(IdentityChange::from_changed(changed)) - } - async fn is_trusted_identity( - &self, - _: &ProtocolAddress, - _: &IdentityKey, - _: Direction, - ) -> SigResult<bool> { - Ok(true) - } - async fn get_identity(&self, a: &ProtocolAddress) -> SigResult<Option<IdentityKey>> { - Ok(self.known.get(a).copied()) - } - } - - struct UnusedPreKeyStore; - #[async_trait::async_trait] - impl PreKeyStore for UnusedPreKeyStore { - async fn get_pre_key(&self, _: PreKeyId) -> SigResult<PreKeyRecord> { - unreachable!() - } - async fn save_pre_key(&mut self, _: PreKeyId, _: &PreKeyRecord) -> SigResult<()> { - unreachable!() - } - async fn remove_pre_key(&mut self, _: PreKeyId) -> SigResult<()> { - unreachable!() - } - } - struct UnusedSignedPreKeyStore; - #[async_trait::async_trait] - impl SignedPreKeyStore for UnusedSignedPreKeyStore { - async fn get_signed_pre_key(&self, _: SignedPreKeyId) -> SigResult<SignedPreKeyRecord> { - unreachable!() - } - async fn save_signed_pre_key( - &mut self, - _: SignedPreKeyId, - _: &SignedPreKeyRecord, - ) -> SigResult<()> { - unreachable!() - } - } - #[derive(Default)] - struct MemSenderKeyStore( - HashMap<crate::libsignal::store::sender_key_name::SenderKeyName, SenderKeyRecord>, - ); - #[async_trait::async_trait] - impl SenderKeyStore for MemSenderKeyStore { - async fn store_sender_key( - &mut self, - n: &crate::libsignal::store::sender_key_name::SenderKeyName, - r: SenderKeyRecord, - ) -> SigResult<()> { - self.0.insert(n.clone(), r); - Ok(()) - } - async fn load_sender_key( - &self, - n: &crate::libsignal::store::sender_key_name::SenderKeyName, - ) -> SigResult<Option<SenderKeyRecord>> { - Ok(self.0.get(n).cloned()) - } - } - - struct TokioTestRuntime; - #[async_trait::async_trait] - impl Runtime for TokioTestRuntime { - fn spawn( - &self, - future: Pin<Box<dyn Future<Output = ()> + Send + 'static>>, - ) -> AbortHandle { - let handle = tokio::spawn(future); - AbortHandle::new(move || handle.abort()) - } - fn sleep(&self, _d: Duration) -> Pin<Box<dyn Future<Output = ()> + Send>> { - Box::pin(async {}) - } - fn spawn_blocking( - &self, - f: Box<dyn FnOnce() + Send + 'static>, - ) -> Pin<Box<dyn Future<Output = ()> + Send>> { - Box::pin(async move { - let _ = tokio::task::spawn_blocking(f).await; - }) - } - fn yield_now(&self) -> Option<Pin<Box<dyn Future<Output = ()> + Send>>> { - None - } - } - - /// The send path must report a replaced identity via the resolver when - /// establishing a session whose bundle carries a new identity key for an - /// address we already knew (peer reinstall). Mirrors WA Web saveIdentity - /// -> handleNewIdentity firing during outbound session setup. - #[tokio::test] - async fn encrypt_for_devices_reports_replaced_identity() { - let mut rng = rand::make_rng::<rand::rngs::StdRng>(); - - // Receiver device D with a valid signed bundle. - let device: Jid = "5511777777777:0@s.whatsapp.net".parse().unwrap(); - let receiver = IdentityKeyPair::generate(&mut rng); - let spk = KeyPair::generate(&mut rng); - let opk = KeyPair::generate(&mut rng); - let sig = receiver - .private_key() - .calculate_signature(&spk.public_key.serialize(), &mut rng) - .unwrap(); - let bundle = PreKeyBundle::new( - 1, - 1u32.into(), - Some((1u32.into(), opk.public_key)), - 1u32.into(), - spk.public_key, - sig.to_vec(), - *receiver.identity_key(), - ) - .unwrap(); - - // Local stores: no session for D + a STALE identity pre-seeded for D's - // address, so establishing the session reports ReplacedExisting. - let sender = IdentityKeyPair::generate(&mut rng); - let stale = *IdentityKeyPair::generate(&mut rng).identity_key(); - let mut known = HashMap::new(); - known.insert(device.to_protocol_address(), stale); - - let mut session_store = MemSessionStore::default(); - let mut identity_store = MemIdentityStore { - pair: sender, - known, - }; - let mut prekey_store = UnusedPreKeyStore; - let signed_prekey_store = UnusedSignedPreKeyStore; - let mut sender_key_store = MemSenderKeyStore::default(); - - let mut stores = SignalStores { - sender_key_store: &mut sender_key_store, - session_store: &mut session_store, - identity_store: &mut identity_store, - prekey_store: &mut prekey_store, - signed_prekey_store: &signed_prekey_store, - }; - - let resolver = MockSendContextResolver::new() - .with_bundle(device.clone(), bundle) - .with_devices(vec![device.clone()]); - let rt = TokioTestRuntime; - - encrypt_for_devices( - &rt, - &mut stores, - &resolver, - std::slice::from_ref(&device), - b"hello", - false, - None, - ) - .await - .expect("encrypt_for_devices"); +mod classify; +mod dm; +mod encrypt; +mod group; +mod peer; +mod status; + +pub use classify::*; +#[cfg(test)] +pub(crate) use dm::partition_dm_devices; +pub(crate) use dm::pkmsg_would_be_emitted; +pub use dm::*; +pub use encrypt::*; +pub use group::*; +pub use peer::*; +pub use status::*; - assert_eq!( - resolver.captured_identity_changes(), - vec![device], - "replaced identity on the send path must be reported via the resolver" - ); - } - } -} +#[cfg(test)] +mod tests; diff --git a/wacore/src/send/classify.rs b/wacore/src/send/classify.rs new file mode 100644 index 000000000..3668db87b --- /dev/null +++ b/wacore/src/send/classify.rs @@ -0,0 +1,316 @@ +//! Message classification: stanza/media type, ciphertext extraction, decrypt-fail gating. + +use super::*; + +/// Extract (enc_type, is_prekey, serialized) from a CiphertextMessage. +pub fn extract_ciphertext(msg: CiphertextMessage) -> Option<(&'static str, bool, Box<[u8]>)> { + match msg { + CiphertextMessage::SignalMessage(m) => { + Some((stanza::ENC_TYPE_MSG, false, m.into_serialized())) + } + CiphertextMessage::PreKeySignalMessage(m) => { + Some((stanza::ENC_TYPE_PKMSG, true, m.into_serialized())) + } + _ => None, + } +} + +/// Unwrap wrapper message types to reach the inner message. +/// Matches WA Web's getUnwrappedProtobufMessage. Does not unwrap +/// `edited_message`; that field is itself a signal callers may need. +pub(crate) fn unwrap_message(msg: &wa::Message) -> &wa::Message { + macro_rules! try_unwrap { + ($($field:ident),+ $(,)?) => { + $( + if let Some(ref w) = msg.$field { + if let Some(ref inner) = w.message { + return unwrap_message(inner); + } + } + )+ + }; + } + try_unwrap!( + ephemeral_message, + view_once_message, + view_once_message_v2, + view_once_message_v2_extension, + document_with_caption_message, + group_mentioned_message, + bot_invoke_message, + associated_child_message, + poll_creation_option_image_message, + // Remaining FutureProofMessage wrappers from WA Web's + // getUnwrappedProtobufMessage list; classify by the inner message. + event_cover_image, + group_status_message, + group_status_message_v2, + group_status_mention_message, + status_add_yours, + status_mention_message, + question_message, + question_reply_message, + spoiler_message, + lottie_sticker_message, + limit_sharing_message, + newsletter_admin_profile_message, + newsletter_admin_profile_message_v2, + poll_creation_message_v4, + ); + if let Some(ref dsm) = msg.device_sent_message + && let Some(ref inner) = dsm.message + { + return unwrap_message(inner); + } + msg +} + +/// Matches WAWebE2EProtoUtils.typeAttributeFromProtobuf. +pub fn stanza_type_from_message(msg: &wa::Message) -> &'static str { + let msg = unwrap_message(msg); + + if msg.reaction_message.is_some() || msg.enc_reaction_message.is_some() { + return stanza::MSG_TYPE_REACTION; + } + if msg.event_message.is_some() || msg.enc_event_response_message.is_some() { + return stanza::MSG_TYPE_EVENT; + } + if let Some(ref sec) = msg.secret_encrypted_message { + use wa::message::secret_encrypted_message::SecretEncType; + match SecretEncType::try_from(sec.secret_enc_type.unwrap_or(0)) { + Ok(SecretEncType::EventEdit) => return stanza::MSG_TYPE_EVENT, + Ok(SecretEncType::MessageEdit) => return stanza::MSG_TYPE_TEXT, + Ok(SecretEncType::PollEdit | SecretEncType::PollAddOption) => { + return stanza::MSG_TYPE_POLL; + } + _ => {} + } + } + if msg.poll_creation_message.is_some() + || msg.poll_creation_message_v2.is_some() + || msg.poll_creation_message_v3.is_some() + || msg.poll_creation_message_v5.is_some() + || msg.poll_update_message.is_some() + { + return stanza::MSG_TYPE_POLL; + } + if msg.conversation.is_some() + || msg.protocol_message.is_some() + || msg.keep_in_chat_message.is_some() + || msg.edited_message.is_some() + || msg.pin_in_chat_message.is_some() + || msg.interactive_message.is_some() + || msg.template_button_reply_message.is_some() + || msg.request_phone_number_message.is_some() + || msg.enc_comment_message.is_some() + || msg.newsletter_admin_invite_message.is_some() + || msg.newsletter_follower_invite_message_v2.is_some() + || msg.message_history_notice.is_some() + || msg.album_message.is_some() + // Payment family. WA Web's typeAttributeFromProtobuf leaves these at the media + // default, but media-without-mediatype is dropped by the server (so is a bare + // "pay" stanza); text is what delivers and renders on Android. + || msg.request_payment_message.is_some() + || msg.send_payment_message.is_some() + || msg.payment_invite_message.is_some() + || msg.decline_payment_request_message.is_some() + || msg.cancel_payment_request_message.is_some() + { + return stanza::MSG_TYPE_TEXT; + } + // pollResultSnapshotMessage maps to "text" by default in WA Web + // (gated behind isPollResultSnapshotPollTypeEnvelopeEnabled for "poll") + if msg.poll_result_snapshot_message.is_some() || msg.poll_result_snapshot_message_v3.is_some() { + return stanza::MSG_TYPE_TEXT; + } + if let Some(ref ext) = msg.extended_text_message { + if ext + .matched_text + .as_ref() + .is_some_and(|t| !t.trim().is_empty()) + { + return stanza::MSG_TYPE_MEDIA; + } + return stanza::MSG_TYPE_TEXT; + } + stanza::MSG_TYPE_MEDIA +} + +pub fn peer_message_options_from_message(msg: &wa::Message) -> PeerMessageOptions { + use wa::message::PeerDataOperationRequestType as PdoType; + + // WAWebSendNonMessageDataRequest's A/F helpers gate rollout flags we do + // not model; use the default-on wire shape for supported peer PDO flows. + let request_type = unwrap_message(msg) + .protocol_message + .as_deref() + .and_then(|pm| pm.peer_data_operation_request_message.as_ref()) + .and_then(|pdo| pdo.peer_data_operation_request_type) + .and_then(|raw| PdoType::try_from(raw).ok()); + + match request_type { + Some(PdoType::HistorySyncOnDemand) => PeerMessageOptions::high_force_on_demand(), + Some( + PdoType::GenerateLinkPreview + | PdoType::PlaceholderMessageResend + | PdoType::CompanionCanonicalUserNonceFetch, + ) => PeerMessageOptions::high_force(), + _ => PeerMessageOptions::default(), + } +} + +/// Matches WAWebBackendJobsCommon.mediaTypeFromProtobuf + encodeMaybeMediaType. +/// Returns `None` when the attribute should be omitted. +pub fn media_type_from_message(msg: &wa::Message) -> Option<&'static str> { + // WA Web's mediaTypeFromProtobuf treats a top-level lottieStickerMessage as a + // terminal "sticker" and does NOT recurse into it (unlike typeAttributeFromProtobuf, + // which unwraps it via getUnwrappedProtobufMessage). Check before the shared unwrap. + if msg.lottie_sticker_message.is_some() { + return Some("sticker"); + } + + let msg = unwrap_message(msg); + + if msg.image_message.is_some() { + return Some("image"); + } + if let Some(ref vid) = msg.video_message { + return if vid.gif_playback == Some(true) { + Some("gif") + } else { + Some("video") + }; + } + if msg.ptv_message.is_some() { + return Some("ptv"); + } + if let Some(ref audio) = msg.audio_message { + return if audio.ptt == Some(true) { + Some("ptt") + } else { + Some("audio") + }; + } + if msg.document_message.is_some() { + return Some("document"); + } + if msg.sticker_message.is_some() { + return Some("sticker"); + } + if msg.sticker_pack_message.is_some() { + return Some("sticker_pack"); + } + if let Some(ref loc) = msg.location_message { + return if loc.is_live == Some(true) { + Some("livelocation") + } else { + Some("location") + }; + } + if msg.live_location_message.is_some() { + return Some("livelocation"); + } + if msg.contact_message.is_some() { + return Some("vcard"); + } + if msg.contacts_array_message.is_some() { + return Some("contact_array"); + } + if let Some(ref ext) = msg.extended_text_message + && ext + .matched_text + .as_ref() + .is_some_and(|t| !t.trim().is_empty()) + { + return Some("url"); + } + if msg.group_invite_message.is_some() { + return Some("url"); + } + // Interactive / business message families. WA Web's mediaTypeFromProtobuf maps + // each to a concrete mediatype; without it the server drops the type="media" + // stanza. buttonsMessage is intentionally absent: WA Web maps it to + // EncMediaType.Button, which its string mapper drops (no attribute). + if msg.list_message.is_some() { + return Some("list"); + } + if msg.list_response_message.is_some() { + return Some("list_response"); + } + if msg.buttons_response_message.is_some() { + return Some("buttons_response"); + } + if msg.order_message.is_some() { + return Some("order"); + } + if msg.product_message.is_some() { + return Some("product"); + } + if msg.interactive_response_message.is_some() { + return Some("native_flow_response"); + } + if msg.message_history_bundle.is_some() { + return Some("group_history"); + } + None +} + +/// Canonical rule for `decrypt-fail="hide"` on outgoing `<enc>` nodes. +/// Shared by DM fanout, group SKDM and group SKMSG so the three paths can't drift. +/// Both revoke kinds are excluded: WA Web never hides REVOKE, and the server +/// drops revoke stanzas carrying the hide attribute. +pub fn should_hide_decrypt_fail_for_send( + edit: Option<&crate::types::message::EditAttribute>, + msg: &wa::Message, +) -> bool { + use crate::types::message::EditAttribute; + edit.is_some_and(|e| { + *e != EditAttribute::Empty + && *e != EditAttribute::AdminRevoke + && *e != EditAttribute::SenderRevoke + }) || should_hide_decrypt_fail(msg) +} + +/// Infrastructure messages get decrypt-fail="hide" so recipients don't see +/// "waiting for this message" placeholders for things like reactions or pin changes. +pub fn should_hide_decrypt_fail(msg: &wa::Message) -> bool { + let msg = unwrap_message(msg); + + use wa::message::protocol_message::Type as ProtocolType; + use wa::message::secret_encrypted_message::SecretEncType; + + msg.reaction_message.is_some() + || msg.enc_reaction_message.is_some() + || msg.pin_in_chat_message.is_some() + || msg.edited_message.is_some() + || msg.keep_in_chat_message.is_some() + || msg.enc_event_response_message.is_some() + || msg + .poll_update_message + .as_ref() + .is_some_and(|p| p.vote.is_some()) + || msg.message_history_notice.is_some() + || msg.conditional_reveal_message.is_some() + || msg.secret_encrypted_message.as_ref().is_some_and(|s| { + matches!( + SecretEncType::try_from(s.secret_enc_type.unwrap_or(0)), + Ok(SecretEncType::EventEdit + | SecretEncType::PollEdit + | SecretEncType::PollAddOption) + ) + }) + || msg + .bot_invoke_message + .as_ref() + .and_then(|b| b.message.as_ref()) + .and_then(|m| m.protocol_message.as_ref()) + .is_some_and(|p| p.r#type == Some(ProtocolType::RequestWelcomeMessage as i32)) + || msg.protocol_message.as_ref().is_some_and(|p| { + matches!( + p.r#type, + Some(t) if t == ProtocolType::EphemeralSyncResponse as i32 + || t == ProtocolType::RequestWelcomeMessage as i32 + || t == ProtocolType::GroupMemberLabelChange as i32 + ) || p.edited_message.is_some() + }) +} diff --git a/wacore/src/send/dm.rs b/wacore/src/send/dm.rs new file mode 100644 index 000000000..216a80823 --- /dev/null +++ b/wacore/src/send/dm.rs @@ -0,0 +1,317 @@ +//! 1:1 (DM) stanza preparation and DM retry stanzas. + +use super::*; + +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) + || own_lid + .is_some_and(|lid| device_jid.is_same_user_as(lid) && device_jid.device == lid.device) +} + +pub(crate) fn partition_dm_devices( + all_devices: Vec<Jid>, + own_jid: &Jid, + own_lid: Option<&Jid>, +) -> (Vec<Jid>, Vec<Jid>) { + let mut recipient_devices = Vec::with_capacity(all_devices.len()); + let mut own_other_devices = Vec::with_capacity(4); + + for device_jid in all_devices { + if is_exact_dm_sender_device(&device_jid, own_jid, own_lid) { + continue; + } + + if device_jid.matches_user_or_lid(own_jid, own_lid) { + own_other_devices.push(device_jid); + } else { + recipient_devices.push(device_jid); + } + } + + (recipient_devices, own_other_devices) +} + +/// Result of `prepare_dm_stanza` — carries the stanza node and the +/// locally computed phash for server ACK validation. +pub struct PreparedDmStanza { + pub node: Node, + /// Locally computed phash from the sent device set. Not sent on the + /// wire (WA Web only sends phash for groups). Used by the caller to + /// compare against the server's ACK phash for device-list drift detection. + pub phash: Option<String>, + /// `MessageContextInfo.message_secret` generated for this stanza so the + /// caller can persist it for later addon (msmsg/poll/edit) decryption. + /// `None` when the message had no reporting token (no secret was used). + pub message_secret: Option<[u8; crate::reporting_token::MESSAGE_SECRET_SIZE]>, +} + +#[allow(clippy::too_many_arguments)] +pub async fn prepare_dm_stanza< + 'a, + S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static, + I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static, + P: crate::libsignal::protocol::PreKeyStore + Send + Sync, + SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync, +>( + runtime: &dyn Runtime, + stores: &mut SignalStores<'a, S, I, P, SP>, + 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>, +) -> Result<PreparedDmStanza> { + // sender is the author's own jid, remote is the chat jid (WAWebReportingTokenUtils: + // getSender vs e.to). Both previously used to_jid, conflating sender with remote. + let reporting_result = generate_reporting_token(message, &request_id, own_jid, &to_jid, None); + + let message_for_encryption = if let Some(ref result) = reporting_result { + prepare_message_with_context(message, &result.message_secret) + } else { + message.clone() + }; + + let recipient_plaintext = MessageUtils::encode_and_pad(&message_for_encryption); + + // Partition first so phash reflects the actual sent set (sender excluded) + let total_devices = all_devices.len(); + let (recipient_devices, own_other_devices) = + partition_dm_devices(all_devices, own_jid, own_lid); + + let phash = MessageUtils::participant_list_hash( + recipient_devices.iter().chain(own_other_devices.iter()), + ) + .ok(); + + let dsm = crate::messages::wrap_device_sent(message_for_encryption, to_jid.to_string()); + + let own_devices_plaintext = MessageUtils::encode_and_pad(&dsm); + + 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 mediatype = media_type_from_message(message); + + // NOTE: WA Web has a bare-<enc> fast path for single primary device + // (WAWebSendMsgCreateFanoutStanza). Not implemented here because + // encrypt_for_devices always wraps in <to jid=...> nodes; + // a bare-enc mode would require refactoring the encryption layer. + // The <participants> form is accepted by the server regardless. + + if !recipient_devices.is_empty() { + let result = encrypt_for_devices( + runtime, + stores, + resolver, + &recipient_devices, + &recipient_plaintext, + hide_decrypt_fail, + mediatype, + ) + .await?; + participant_nodes.extend(result.participant_nodes); + includes_prekey_message = includes_prekey_message || result.includes_prekey_message; + } + + if !own_other_devices.is_empty() { + let result = encrypt_for_devices( + runtime, + stores, + resolver, + &own_other_devices, + &own_devices_plaintext, + hide_decrypt_fail, + mediatype, + ) + .await?; + participant_nodes.extend(result.participant_nodes); + includes_prekey_message = includes_prekey_message || result.includes_prekey_message; + } + + // All per-device encrypts failed: an empty <participants> would silently + // drop the message. WA Web's encryptAndSendUserMsg rejects here too. + let attempted_devices = recipient_devices.len() + own_other_devices.len(); + if participant_nodes.is_empty() && attempted_devices > 0 { + return Err(anyhow!( + "encryption failed for all {attempted_devices} recipient device(s)" + )); + } + + let mut message_content_nodes = vec![ + NodeBuilder::new("participants") + .children(participant_nodes) + .build(), + ]; + + if includes_prekey_message && let Some(acc) = account { + let device_identity_bytes = acc.encode_to_vec(); + message_content_nodes.push( + NodeBuilder::new("device-identity") + .bytes(device_identity_bytes) + .build(), + ); + } + + // Add reporting token node if we generated one + if let Some(ref result) = reporting_result { + message_content_nodes.push(build_reporting_node(result)); + } + + // Add any extra stanza nodes provided by the caller + message_content_nodes.extend(extra_stanza_nodes.iter().cloned()); + + let stanza_type = stanza_type_from_message(message); + + let mut stanza_builder = NodeBuilder::new("message") + .attr("to", to_jid) + .attr("id", request_id) + .attr("type", stanza_type); + + if let Some(edit_attr) = edit + && edit_attr != crate::types::message::EditAttribute::Empty + { + stanza_builder = stanza_builder.attr("edit", edit_attr.to_string_val()); + } + + let stanza = stanza_builder.children(message_content_nodes).build(); + + Ok(PreparedDmStanza { + node: stanza, + phash, + message_secret: reporting_result.map(|r| r.message_secret), + }) +} + +/// Returns true if `message_encrypt` on `signal_address` would produce +/// a pkmsg (no session yet, or session with un-acked pre-key still +/// pending). Used before `message_encrypt` to fail-fast when `account` +/// is None — pkmsg without `<device-identity>` reproduces the linked +/// device deadlock. +/// +/// `SessionStore::load_session` is take-semantics in production +/// (`SessionAdapter` → `SignalStoreCache::get_session` marks the slot +/// `CheckedOut`); the loaded record is put back via `store_session` +/// so the subsequent `message_encrypt` finds the slot Present. +pub(crate) async fn pkmsg_would_be_emitted<S>( + session_store: &mut S, + signal_address: &ProtocolAddress, +) -> Result<bool> +where + S: crate::libsignal::protocol::SessionStore, +{ + let loaded = session_store.load_session(signal_address).await?; + // Conservative read: treat any failure to interrogate the session as + // "would be pkmsg" so the caller bails. Silently treating Err as false + // would let message_encrypt run with a corrupt session and potentially + // burn the sender chain. + let needs_pkmsg = match &loaded { + None => true, + Some(record) => match record.session_state() { + None => true, + Some(state) => match state.unacknowledged_pre_key_message_items() { + Ok(Some(_)) => true, + Ok(None) => false, + Err(_) => true, + }, + }, + }; + if let Some(record) = loaded { + session_store + .store_session(signal_address, record) + .await + .map_err(|e| anyhow!("restoring checked-out session after pre-flight: {e}"))?; + } + Ok(needs_pkmsg) +} + +/// 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. +#[allow(clippy::too_many_arguments)] +pub async fn prepare_dm_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>, +) -> 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!( + "DM 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 DM retry"))?; + + let hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message); + 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); + } + if hide_decrypt_fail { + enc_builder = enc_builder.attr("decrypt-fail", "hide"); + } + let enc_node = enc_builder.bytes(serialized).build(); + + let mut children = vec![enc_node]; + if is_prekey { + // Defense in depth: pre-flight should have caught this, but a corrupt + // session that triggers a fresh pkmsg mid-call would slip past. + let acc = account.ok_or_else(|| { + anyhow!("DM retry pkmsg without <device-identity> (unreachable via pre-flight)") + })?; + children.push( + NodeBuilder::new("device-identity") + .bytes(acc.encode_to_vec()) + .build(), + ); + } + + let mut stanza_builder = NodeBuilder::new("message") + .attr("to", to_jid) + .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. + 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()) +} diff --git a/wacore/src/send/encrypt.rs b/wacore/src/send/encrypt.rs new file mode 100644 index 000000000..44136cc35 --- /dev/null +++ b/wacore/src/send/encrypt.rs @@ -0,0 +1,583 @@ +//! Per-device Signal encryption fanout and the bounded spawn helper. + +use super::*; + +/// Caller must hold `SenderKeyStore::sender_key_lock` for `sender_key_name` +/// across the surrounding SKDM creation + this encrypt, so a concurrent send +/// can't split the key between the SKDM and the skmsg. +pub async fn encrypt_group_message<S, R>( + sender_key_store: &mut S, + sender_key_name: &SenderKeyName, + plaintext: &[u8], + csprng: &mut R, +) -> Result<SenderKeyMessage> +where + S: SenderKeyStore + ?Sized, + R: Rng + CryptoRng, +{ + log::debug!( + "Attempting to load sender key for group {} sender {}", + sender_key_name.group_id(), + sender_key_name.sender_id() + ); + + let mut record = sender_key_store + .load_sender_key(sender_key_name) + .await? + .ok_or_else(|| { + SignalProtocolError::NoSenderKeyState(format!( + "no sender key record for group {} sender {}", + sender_key_name.group_id(), + sender_key_name.sender_id() + )) + })?; + + let sender_key_state = record + .sender_key_state_mut() + .map_err(|e| anyhow!("Invalid SenderKey session: {:?}", e))?; + + let sender_chain_key = sender_key_state + .sender_chain_key() + .ok_or_else(|| anyhow!("Invalid SenderKey session: missing chain key"))?; + + let message_keys = sender_chain_key.sender_message_key(); + + let mut ciphertext = Vec::new(); + aes_256_cbc_encrypt_into( + plaintext, + message_keys.cipher_key(), + message_keys.iv(), + &mut ciphertext, + ) + .map_err(|_| anyhow!("AES encryption failed"))?; + + let signing_key = sender_key_state + .signing_key_private() + .map_err(|e| anyhow!("Invalid SenderKey session: missing signing key: {:?}", e))?; + + let skm = SenderKeyMessage::new( + SENDERKEY_MESSAGE_CURRENT_VERSION, + sender_key_state.chain_id(), + message_keys.iteration(), + ciphertext.into_boxed_slice(), + csprng, + &signing_key, + )?; + + sender_key_state.set_sender_chain_key(sender_chain_key.next()?); + + sender_key_store + .store_sender_key(sender_key_name, record) + .await?; + + Ok(skm) +} + +pub struct SignalStores<'a, S, I, P, SP> { + pub sender_key_store: &'a mut (dyn crate::libsignal::protocol::SenderKeyStore + Send + Sync), + pub session_store: &'a mut S, + pub identity_store: &'a mut I, + pub prekey_store: &'a mut P, + pub signed_prekey_store: &'a SP, +} + +/// Check if an anyhow error is a 406 "not-acceptable" server error (device unregistered). +/// Uses typed downcast to `ServerErrorCode` — the shared error type that the +/// `SendContextResolver` impl wraps server errors in. +pub(crate) fn is_device_unregistered_error(err: &anyhow::Error) -> bool { + crate::request::ServerErrorCode::from_anyhow(err).is_some_and(|e| e.code == 406) +} + +pub struct EncryptResult { + pub participant_nodes: Vec<Node>, + pub includes_prekey_message: bool, + pub encrypted_devices: Vec<Jid>, + /// True if any device returned 406 (unregistered) during prekey fetch. + pub had_unregistered_device: bool, +} + +/// Maximum number of concurrent per-device crypto tasks during group send +/// fan-out. Picked from the `perf-audit` benchmark: speedup plateaus around +/// 16 on Oracle ARM64; 32 gives only ~10% more for double the task overhead. +const ENCRYPT_FANOUT_CONCURRENCY: usize = 16; + +/// Per-task encrypt result, shipped from a spawned task back to the orchestrator. +struct EncryptOneResult { + enc_type: &'static str, + is_prekey: bool, + ciphertext: Vec<u8>, + hide_decrypt_fail: bool, +} + +/// Surfaces a spawned task that didn't deliver its result — either the task +/// itself panicked or the runtime tore it down (e.g., during shutdown). +/// Surfacing this as an Err lets the encrypt fan-out fall through to its +/// existing log+skip path instead of propagating a panic. +#[derive(Debug, thiserror::Error)] +#[error("spawned task did not produce a result (panic or runtime shutdown)")] +struct SpawnCanceled; + +/// Future returned by [`spawn_oneshot`]. Holds the spawned task's +/// [`AbortHandle`] until the result is received, so dropping the future mid- +/// flight (e.g., the outer send was cancelled by a timeout) cancels the +/// in-flight crypto work instead of orphaning it. +struct Spawned<T> { + rx: futures::channel::oneshot::Receiver<T>, + abort: Option<AbortHandle>, +} + +impl<T> Future for Spawned<T> { + type Output = std::result::Result<T, SpawnCanceled>; + + fn poll( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll<Self::Output> { + match std::pin::Pin::new(&mut self.rx).poll(cx) { + std::task::Poll::Ready(Ok(value)) => { + // Result delivered: disarm so Drop doesn't try to abort an + // already-completed task. + if let Some(handle) = self.abort.take() { + handle.detach(); + } + std::task::Poll::Ready(Ok(value)) + } + std::task::Poll::Ready(Err(_)) => { + if let Some(handle) = self.abort.take() { + handle.detach(); + } + std::task::Poll::Ready(Err(SpawnCanceled)) + } + std::task::Poll::Pending => std::task::Poll::Pending, + } + } +} + +impl<T> Drop for Spawned<T> { + fn drop(&mut self) { + // If the future was dropped before completion, abort the spawned + // task to stop the wasted CPU work. AbortHandle::abort is a no-op + // after the task has already finished, so this is always safe. + if let Some(handle) = self.abort.take() { + handle.abort(); + } + } +} + +/// Spawn `fut` on the runtime and return a future that resolves to its +/// output. Cancellation propagates: dropping the returned future aborts +/// the spawned task. A spawned-task panic surfaces as `Err(SpawnCanceled)` +/// rather than a panic on `rx.await`. +#[cfg(not(target_arch = "wasm32"))] +fn spawn_oneshot<F, T>( + rt: &dyn Runtime, + fut: F, +) -> impl Future<Output = std::result::Result<T, SpawnCanceled>> + Send + 'static +where + F: Future<Output = T> + Send + 'static, + T: Send + 'static, +{ + let (tx, rx) = futures::channel::oneshot::channel(); + let abort = rt.spawn(Box::pin(async move { + let _ = tx.send(fut.await); + })); + Spawned { + rx, + abort: Some(abort), + } +} + +#[cfg(target_arch = "wasm32")] +fn spawn_oneshot<F, T>( + rt: &dyn Runtime, + fut: F, +) -> impl Future<Output = std::result::Result<T, SpawnCanceled>> + 'static +where + F: Future<Output = T> + 'static, + T: 'static, +{ + let (tx, rx) = futures::channel::oneshot::channel(); + let abort = rt.spawn(Box::pin(async move { + let _ = tx.send(fut.await); + })); + Spawned { + rx, + abort: Some(abort), + } +} + +/// Encrypt padded plaintext for each device JID, producing participant `<to>` nodes. +/// +/// Encrypt the plaintext for one device's Signal session. Shared by the +/// single-device fast path and the parallel fan-out so both behave identically. +async fn encrypt_one_device( + plaintext: &[u8], + addr: &ProtocolAddress, + session_store: &mut dyn crate::libsignal::protocol::SessionStore, + identity_store: &mut dyn crate::libsignal::protocol::IdentityKeyStore, + device_jid: Jid, + hide_decrypt_fail: bool, +) -> (Jid, Result<Option<EncryptOneResult>, String>) { + match message_encrypt(plaintext, addr, session_store, identity_store).await { + Ok(encrypted_payload) => { + let Some((enc_type, is_prekey, serialized_bytes)) = + extract_ciphertext(encrypted_payload) + else { + return (device_jid, Ok(None)); + }; + ( + device_jid, + Ok(Some(EncryptOneResult { + enc_type, + is_prekey, + // Box<[u8]> -> Vec<u8> reuses the allocation (no copy). + ciphertext: serialized_bytes.into(), + hide_decrypt_fail, + })), + ) + } + Err(e) => (device_jid, Err(format!("{addr}: {e}"))), + } +} + +/// Append one encrypt result to the fan-out output: a `<to>` participant node on +/// success, a logged skip on failure. +fn push_encrypt_result( + (device_jid, res): (Jid, Result<Option<EncryptOneResult>, String>), + mediatype: Option<&str>, + participant_nodes: &mut Vec<Node>, + encrypted_devices: &mut Vec<Jid>, + includes_prekey_message: &mut bool, +) { + match res { + Ok(Some(one)) => { + *includes_prekey_message |= one.is_prekey; + let mut enc_builder = NodeBuilder::new("enc") + .attr("v", stanza::ENC_VERSION) + .attr("type", one.enc_type); + // `mediatype` is batch-level (same for every device) and originates as + // a `&'static str`, so it's threaded here instead of cloned per result. + if let Some(mt) = mediatype { + enc_builder = enc_builder.attr("mediatype", mt); + } + if one.hide_decrypt_fail { + enc_builder = enc_builder.attr("decrypt-fail", "hide"); + } + let enc_node = enc_builder.bytes(one.ciphertext).build(); + participant_nodes.push( + NodeBuilder::new("to") + .attr("jid", device_jid.clone()) + .children([enc_node]) + .build(), + ); + encrypted_devices.push(device_jid); + } + Ok(None) => {} + Err(msg) => log::warn!("Failed to encrypt for device: {msg}. Skipping."), + } +} + +/// Per-device Signal sessions are independent (different ratchet state per +/// recipient), so this fans the encrypt loop out across tokio tasks bounded +/// by [`ENCRYPT_FANOUT_CONCURRENCY`]. Each task clones the store handles +/// (Arc bumps under the hood); the shared cache provides interior mutability. +/// +/// Callers must hold per-device session locks before calling this function — +/// concurrent ratchet mutations will corrupt Signal session state. +pub async fn encrypt_for_devices<'a, S, I, P, SP>( + runtime: &dyn Runtime, + stores: &mut SignalStores<'a, S, I, P, SP>, + resolver: &dyn SendContextResolver, + devices: &[Jid], + plaintext_to_encrypt: &[u8], + hide_decrypt_fail: bool, + mediatype: Option<&str>, +) -> Result<EncryptResult> +where + S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static, + I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static, + P: crate::libsignal::protocol::PreKeyStore + Send + Sync, + SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync, +{ + // Per-device LID upgrade map: encryption_overrides[i] mirrors devices[i]. + // None = use devices[i] as-is; Some(jid) = use this LID-upgraded version. + // The Vec replaces a HashMap<&Jid, Jid> that paid hash + alloc per insert + // and per get (~666 of each on a large group). Plain Vec<Option<Jid>> is + // direct indexing and contiguous memory. + let mut encryption_overrides: Vec<Option<Jid>> = vec![None; devices.len()]; + // 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 reusable_addr = crate::types::jid::make_reusable_protocol_address(); + + for (idx, device_jid) in devices.iter().enumerate() { + // WhatsApp Web's SignalAddress.toString() normalizes PN → LID before + // creating signal addresses. We do the same: check LID session FIRST. + // This prevents using stale PN sessions when a newer LID session exists. + if device_jid.is_pn() + && let Some(lid_user) = resolver.get_lid_for_phone(&device_jid.user).await + { + // Construct the LID JID with the same device ID + let lid_jid = Jid::lid_device(lid_user, device_jid.device); + lid_jid.reset_protocol_address(&mut reusable_addr); + + if stores.session_store.has_session(&reusable_addr).await? { + log::debug!( + "Using LID session {} for PN {} (LID-first lookup)", + lid_jid, + device_jid + ); + encryption_overrides[idx] = Some(lid_jid); + continue; + } + } + + device_jid.reset_protocol_address(&mut reusable_addr); + if stores.session_store.has_session(&reusable_addr).await? { + continue; + } + + // No session found - need to fetch prekeys and create session. + // Keep device_jid for prekey fetch (server returns bundles keyed by this), + // but normalize to LID for the actual session creation. + if device_jid.is_pn() + && let Some(lid_user) = resolver.get_lid_for_phone(&device_jid.user).await + { + let lid_jid = Jid::lid_device(lid_user, device_jid.device); + log::debug!( + "Will create LID session {} for PN {} (no existing session)", + lid_jid, + device_jid + ); + encryption_overrides[idx] = Some(lid_jid); + } + indices_needing_prekeys.push(idx); + } + + if !indices_needing_prekeys.is_empty() { + log::debug!( + "Fetching prekeys for {} devices without sessions", + indices_needing_prekeys.len() + ); + // Materialize the Jid slice for the resolver call. fetch_prekeys + // wants &[Jid]; same per-device clone count as the previous Vec + // model, just sourced from the indices. + let jids_for_fetch: Vec<Jid> = indices_needing_prekeys + .iter() + .map(|&i| devices[i].clone()) + .collect(); + // 406 on this batch is all-or-nothing — per-device retries just wasted + // N·RTT with the same failure. Mark `had_406` so the caller invalidates + // the users and the next send re-fetches. Matches WA Web's + // `GroupSkmsgJob`: log, continue without those devices. + let prekey_bundles = match resolver + .fetch_prekeys_for_identity_check(&jids_for_fetch) + .await + { + Ok(bundles) => bundles, + Err(e) if is_device_unregistered_error(&e) => { + log::warn!( + "Prekey fetch returned 406 for {} device(s); skipping them this round", + jids_for_fetch.len() + ); + had_406 = true; + std::collections::HashMap::new() + } + Err(e) => return Err(e), + }; + + // Parallel session establishment via process_prekey_bundle. Each + // recipient device has an independent Signal session and an + // independent prekey bundle, so the X3DH derivation runs on a + // separate task per device, bounded at ENCRYPT_FANOUT_CONCURRENCY. + // Spawning goes through `Runtime::spawn` (the platform-agnostic + // abstraction) plus a oneshot channel for result delivery — + // `FuturesUnordered` handles the in-flight window. + let prekey_bundles = std::sync::Arc::new(prekey_bundles); + let total = indices_needing_prekeys.len(); + let mut next_spawn = 0usize; + + let make_session_task = |spawn_idx: usize| { + let idx = indices_needing_prekeys[spawn_idx]; + let device_jid = devices[idx].clone(); + let mut encryption_jid = encryption_overrides[idx] + .clone() + .unwrap_or_else(|| device_jid.clone()); + + // Normalize agent to 0 for LID JIDs to match how pre-key bundles are stored. + // prekeys.rs forces agent=0 for LID; we must match that here. + if encryption_jid.is_lid() { + encryption_jid.agent = 0; + } + + let lookup_jid = device_jid.normalize_for_prekey_bundle(); + let bundles = prekey_bundles.clone(); + let mut session_store = stores.session_store.clone(); + let mut identity_store = stores.identity_store.clone(); + + spawn_oneshot(runtime, async move { + let mut addr = crate::types::jid::make_reusable_protocol_address(); + encryption_jid.reset_protocol_address(&mut addr); + + let Some(bundle) = bundles.get(&lookup_jid) else { + log::warn!( + "No pre-key bundle returned for device {}. This device will be skipped for encryption.", + addr + ); + return Ok::<Option<Jid>, anyhow::Error>(None); + }; + + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + // No UntrustedIdentity recovery: WA Web's isTrustedIdentity is + // unconditional Ok(true) (TOFU), and save_identity inside + // process_prekey_bundle persists rotations transparently. + match process_prekey_bundle( + &addr, + &mut session_store, + &mut identity_store, + bundle, + &mut rng, + UsePQRatchet::No, + ) + .await + { + // Surface a replaced identity so the caller can react + // (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 + )), + } + }) + }; + + let mut in_flight: FuturesUnordered<_> = FuturesUnordered::new(); + while next_spawn < total && in_flight.len() < ENCRYPT_FANOUT_CONCURRENCY { + in_flight.push(make_session_task(next_spawn)); + next_spawn += 1; + } + while let Some(spawn_result) = in_flight.next().await { + match spawn_result { + // Some(jid) => establishing this session replaced a stored + // identity; notify the client so it can react off-path. + Ok(Ok(Some(changed_jid))) => resolver.on_local_identity_change(&changed_jid), + Ok(Ok(None)) => {} + Ok(Err(e)) => return Err(e), + Err(SpawnCanceled) => { + log::warn!( + "Session-establishment task did not deliver a result; skipping device." + ); + } + } + if next_spawn < total { + in_flight.push(make_session_task(next_spawn)); + next_spawn += 1; + } + } + } + + let mut participant_nodes = Vec::with_capacity(devices.len()); + let mut includes_prekey_message = false; + let mut encrypted_devices = Vec::with_capacity(devices.len()); + + // The wire-order of `<to>` participants does not need to match the input + // device order: WA Web's `phash` (computed both client and server side) + // sorts before hashing, as does our `participant_list_hash`. + if devices.len() == 1 { + // Single recipient device: the parallel fan-out is pure overhead here + // (an Arc<[u8]> copy of the plaintext, a spawned task + oneshot channel, + // a FuturesUnordered, and two store clones), with no parallelism to gain. + // Encrypt inline. + let device_jid = devices[0].clone(); + let addr = encryption_overrides[0] + .as_ref() + .unwrap_or(&devices[0]) + .to_protocol_address(); + let res = encrypt_one_device( + plaintext_to_encrypt, + &addr, + &mut *stores.session_store, + &mut *stores.identity_store, + device_jid, + hide_decrypt_fail, + ) + .await; + push_encrypt_result( + res, + mediatype, + &mut participant_nodes, + &mut encrypted_devices, + &mut includes_prekey_message, + ); + } else { + // Parallel encrypt fan-out across tokio tasks bounded by + // ENCRYPT_FANOUT_CONCURRENCY; collected in completion order so the + // fastest encrypts ship first. + let plaintext_arc: std::sync::Arc<[u8]> = std::sync::Arc::from(plaintext_to_encrypt); + + let total = devices.len(); + let mut next_spawn = 0usize; + + let make_encrypt_task = |idx: usize| { + let device_jid = devices[idx].clone(); + // The encryption JID is only needed to build the Signal address, so + // derive it here from a borrow rather than cloning the whole Jid into + // the task (device_jid is still cloned because it's returned). + let addr = encryption_overrides[idx] + .as_ref() + .unwrap_or(&devices[idx]) + .to_protocol_address(); + let plaintext = plaintext_arc.clone(); + let mut session_store = stores.session_store.clone(); + let mut identity_store = stores.identity_store.clone(); + + spawn_oneshot(runtime, async move { + encrypt_one_device( + &plaintext, + &addr, + &mut session_store, + &mut identity_store, + device_jid, + hide_decrypt_fail, + ) + .await + }) + }; + + let mut in_flight: FuturesUnordered<_> = FuturesUnordered::new(); + while next_spawn < total && in_flight.len() < ENCRYPT_FANOUT_CONCURRENCY { + in_flight.push(make_encrypt_task(next_spawn)); + next_spawn += 1; + } + while let Some(spawn_result) = in_flight.next().await { + match spawn_result { + Ok(res) => push_encrypt_result( + res, + mediatype, + &mut participant_nodes, + &mut encrypted_devices, + &mut includes_prekey_message, + ), + Err(SpawnCanceled) => { + log::warn!("Encrypt task did not deliver a result; skipping device."); + } + } + + if next_spawn < total { + in_flight.push(make_encrypt_task(next_spawn)); + next_spawn += 1; + } + } + } + + Ok(EncryptResult { + participant_nodes, + includes_prekey_message, + encrypted_devices, + had_unregistered_device: had_406, + }) +} diff --git a/wacore/src/send/group.rs b/wacore/src/send/group.rs new file mode 100644 index 000000000..70201568a --- /dev/null +++ b/wacore/src/send/group.rs @@ -0,0 +1,562 @@ +//! Group stanza preparation, phash/stale-device helpers and sender-key distribution. + +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. +#[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]; + + if is_prekey { + // Defense in depth: pre-flight should have caught this, but a corrupt + // session that triggers a fresh pkmsg mid-call would slip past. + let acc = account.ok_or_else(|| { + anyhow!("group retry pkmsg without <device-identity> (unreachable via pre-flight)") + })?; + children.push( + NodeBuilder::new("device-identity") + .bytes(acc.encode_to_vec()) + .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()) +} + +/// Result of `prepare_group_stanza` — carries the stanza node and the exact +/// device list used for SKDM distribution, so callers can persist sender key +/// tracking without re-resolving devices. +pub struct PreparedGroupStanza { + pub node: Node, + /// Full SKDM distribution target set, marked `has_key=true` after the + /// server ACK. Mirrors WA Web `markHasSenderKey(x, M)` which marks the + /// whole target set `M`, not only the devices that encrypted successfully: + /// devices that failed (406 / no bundle) are marked too so they are not + /// re-targeted on every send (the retry-receipt path repairs any that are + /// actually alive and keyless via `mark_forget_sender_key`). + pub skdm_devices: Vec<Jid>, + /// Users whose device registry should be invalidated because their + /// devices returned 406 (unregistered) during SKDM prekey fetch. + /// Empty when no 406 occurred. + pub stale_device_users: Vec<String>, + /// Generated `MessageContextInfo.message_secret`; populated when the + /// reporting token was produced for this send. + pub message_secret: Option<[u8; crate::reporting_token::MESSAGE_SECRET_SIZE]>, + /// The identity we addressed this group send under (LID for LID-mode + /// groups, PN for PN-mode). Used to key the persisted `messageSecret` + /// so msmsg bot replies referencing this msg_id hit the same row that + /// `<meta target_sender_jid>` echoes back at lookup time. + pub sender_identity: Jid, +} + +#[allow(clippy::too_many_arguments)] +pub async fn prepare_group_stanza< + 'a, + S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static, + I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static, + P: crate::libsignal::protocol::PreKeyStore + Send + Sync, + SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync, +>( + runtime: &dyn Runtime, + stores: &mut SignalStores<'a, S, I, P, SP>, + 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<Vec<Jid>>, + edit: Option<crate::types::message::EditAttribute>, + extra_stanza_nodes: &[Node], +) -> Result<PreparedGroupStanza> { + 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"), + }; + + // Generate reporting token if the message type supports it + // For groups, both sender_jid and remote_jid are the group JID (to_jid) per Baileys implementation + let reporting_result = generate_reporting_token(message, &request_id, &to_jid, &to_jid, None); + + // Prepare message with MessageContextInfo containing the message secret + let message_for_encryption = if let Some(ref result) = reporting_result { + prepare_message_with_context(message, &result.message_secret) + } else { + message.clone() + }; + + let own_base_jid = own_sending_jid.to_non_ad(); + + let mut message_children: Vec<Node> = Vec::new(); + let mut includes_prekey_message = false; + let mut phash_for_stanza: Option<String> = None; + let mut skdm_encrypted_devices: Vec<Jid> = Vec::new(); + + // Build the chain name once and hold its lock across SKDM creation + the + // skmsg encrypt, so concurrent same-(group, sender) sends can't split the + // key between the SKDM and the skmsg (nor reuse a chain iteration). + let sender_key_name = make_sender_key_name(&to_jid, &own_sending_jid.to_protocol_address()); + let chain_lock = stores + .sender_key_store + .sender_key_lock(&sender_key_name) + .await; + let _chain_guard = chain_lock.lock().await; + + // Determine if we need to distribute SKDM and to which devices + let distribution_list: Option<Vec<Jid>> = if let Some(target_devices) = skdm_target_devices { + // Use the specific list of devices that need SKDM + if target_devices.is_empty() { + None + } else { + log::debug!( + "SKDM distribution to {} specific devices for group {}", + target_devices.len(), + to_jid + ); + Some(target_devices) + } + } else if force_skdm_distribution { + // Resolve all devices for all participants (legacy behavior) + // For LID groups, use phone numbers for device queries (LID usync may not work for own JID) + // For PN groups, use JIDs directly + let mut jids_to_resolve: Vec<Jid> = group_info + .participants + .iter() + .map(|jid| { + let base_jid = jid.to_non_ad(); + // If this is a LID JID and we have a phone number mapping, use it for device query + if base_jid.is_lid() + && let Some(phone_jid) = group_info.phone_jid_for_lid_user(&base_jid.user) + { + log::debug!( + "Using phone number {} for LID {} device query", + phone_jid, + base_jid + ); + return phone_jid.to_non_ad(); + } + base_jid + }) + .collect(); + + // Determine what user to check for — use the PN user when own is LID + // and we have a mapping. Keeping this as a borrow avoids allocating a + // throwaway Jid when own is already in the list. + let own_pn_mapping = if own_base_jid.is_lid() { + group_info.phone_jid_for_lid_user(&own_base_jid.user) + } else { + None + }; + let own_check_user = own_pn_mapping + .map(|pn| pn.user.as_str()) + .unwrap_or(own_base_jid.user.as_str()); + + if !jids_to_resolve.iter().any(|p| p.user == own_check_user) { + jids_to_resolve.push(match own_pn_mapping { + Some(pn) => pn.to_non_ad(), + None => own_base_jid.clone(), + }); + } + + crate::types::jid::sort_dedup_by_user(&mut jids_to_resolve); + + log::debug!( + "Resolving devices for {} participants", + jids_to_resolve.len() + ); + + let mut resolved_list = resolver.resolve_devices(&jids_to_resolve).await?; + + // For LID groups, convert phone-based device JIDs back to LID format + // This is necessary because WhatsApp Web expects LID addressing in SKDM <to> nodes + if group_info.addressing_mode == crate::types::message::AddressingMode::Lid { + resolved_list = resolved_list + .into_iter() + .map(|device_jid| group_info.phone_device_jid_into_lid(device_jid)) + .collect(); + log::debug!( + "Converted {} devices to LID addressing for group {}", + resolved_list.len(), + to_jid + ); + } + + // Dedup AFTER LID conversion to avoid duplicates when both phone and LID + // queries return the same user (e.g., 559980000003:33 and 100000037037034:33 + // both convert to 100000037037034:33@lid). + // Key on (user, server, agent, device) — excludes `integrator` which is not + // part of the wire JID identity used in <to jid> and phash. + crate::types::jid::sort_dedup_by_device(&mut resolved_list); + + // Filter devices for SKDM distribution: + // - Exclude the exact sending device (own_sending_jid) - we already have our own sender key + // - Keep ALL other devices including our own other devices (phone, other companions) + // because they need the SKDM to decrypt messages we send from this device + // - Exclude hosted/Cloud API devices (device ID 99 or @hosted server) - they don't + // participate in group E2EE, only in 1:1 chats + 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 + }); + log::debug!( + "Filtered SKDM devices from {} to {} (excluded sender {}:{} and hosted devices)", + before_filter, + resolved_list.len(), + own_user, + own_device + ); + + log::debug!( + "SKDM distribution list for {} resolved to {} devices", + to_jid, + resolved_list.len(), + ); + + Some(resolved_list) + } else { + None + }; + + // 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. + if to_jid.is_group() { + // Warm/partial sends pass the complete set in `all_devices_for_phash`; + // the cold `force_skdm` path leaves it None and `distribution_list` + // already holds the full resolved set. + if let Some(src) = all_devices_for_phash + .as_deref() + .or(distribution_list.as_deref()) + { + let phash_set = build_group_phash_set(src, &own_sending_jid); + match MessageUtils::participant_list_hash(&phash_set) { + Ok(phash) => phash_for_stanza = Some(phash), + Err(e) => log::warn!("Failed to compute group phash for {}: {:?}", to_jid, e), + } + } + } else if let Some(ref distribution_list) = distribution_list { + match MessageUtils::participant_list_hash(distribution_list) { + Ok(phash) => phash_for_stanza = Some(phash), + Err(e) => log::warn!("Failed to compute phash for {}: {:?}", to_jid, e), + } + } + + let mut had_unregistered_devices = false; + + if let Some(ref distribution_list) = distribution_list { + let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group( + stores.sender_key_store, + &sender_key_name, + ) + .await?; + + let skdm_wrapper_msg = wa::Message { + sender_key_distribution_message: Some(wa::message::SenderKeyDistributionMessage { + group_id: Some(to_jid.to_string()), + axolotl_sender_key_distribution_message: Some(axolotl_skdm_bytes), + }), + ..Default::default() + }; + let skdm_plaintext_to_encrypt = MessageUtils::encode_and_pad(&skdm_wrapper_msg); + + // WA Web's GroupSkmsgJob wraps ensureE2ESessions in try/catch — logs error + // but does NOT rethrow. SKDM distribution failure must not prevent the group + // message from being sent. Only successfully encrypted devices are tracked. + // 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( + runtime, + stores, + resolver, + distribution_list, + &skdm_plaintext_to_encrypt, + skdm_hide_decrypt_fail, + None, + ) + .await + { + Ok(result) => { + includes_prekey_message = includes_prekey_message || result.includes_prekey_message; + if result.had_unregistered_device { + had_unregistered_devices = true; + } + skdm_encrypted_devices = result.encrypted_devices; + + if !result.participant_nodes.is_empty() { + message_children.push( + NodeBuilder::new("participants") + .children(result.participant_nodes) + .build(), + ); + if includes_prekey_message && let Some(acc) = account { + message_children.push( + NodeBuilder::new("device-identity") + .bytes(acc.encode_to_vec()) + .build(), + ); + } + } + } + Err(e) => { + log::warn!( + "SKDM distribution failed for group {}, continuing without it: {e}", + to_jid + ); + if is_device_unregistered_error(&e) { + had_unregistered_devices = true; + } + } + } + } + + let plaintext = MessageUtils::encode_and_pad(&message_for_encryption); + let skmsg = encrypt_group_message( + stores.sender_key_store, + &sender_key_name, + &plaintext, + &mut rand::make_rng::<rand::rngs::StdRng>(), + ) + .await?; + + 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 mut enc_builder = NodeBuilder::new("enc") + .attr("v", stanza::ENC_VERSION) + .attr("type", stanza::ENC_TYPE_SKMSG); + if let Some(mt) = mediatype { + enc_builder = enc_builder.attr("mediatype", mt); + } + enc_builder = enc_builder.bytes(skmsg_ciphertext); + if hide_decrypt_fail { + enc_builder = enc_builder.attr("decrypt-fail", "hide"); + } + let content_node = enc_builder.build(); + + let stanza_type = stanza_type_from_message(message); + let mut stanza_builder = NodeBuilder::new("message") + .attr("to", to_jid) + .attr("id", request_id) + .attr("type", stanza_type); + + // WA Web always sets addressing_mode for groups (MsgCreateDeviceStanza.js:131-135) + stanza_builder = stanza_builder.attr("addressing_mode", group_info.addressing_mode.as_str()); + + if let Some(edit_attr) = &edit + && *edit_attr != crate::types::message::EditAttribute::Empty + { + stanza_builder = stanza_builder.attr("edit", edit_attr.to_string_val()); + } + // NOTE: WhatsApp Web does NOT include participant attribute on initial admin revoke send + // The participant attribute only appears on retry/fanout messages + + message_children.push(content_node); + + // Add reporting token node if we generated one + if let Some(ref result) = reporting_result { + message_children.push(build_reporting_node(result)); + } + + // Add phash if we distributed keys in this message + if let Some(phash) = phash_for_stanza { + stanza_builder = stanza_builder.attr("phash", phash); + } + + // Add any extra stanza nodes provided by the caller + message_children.extend(extra_stanza_nodes.iter().cloned()); + + let stanza = stanza_builder.children(message_children).build(); + + let stale_users = if had_unregistered_devices { + collect_stale_device_users( + distribution_list.as_deref(), + &skdm_encrypted_devices, + group_info, + ) + } else { + Vec::new() + }; + + Ok(PreparedGroupStanza { + node: stanza, + // Mark the full target set (matches WA Web `markHasSenderKey(x, M)`), not + // just `skdm_encrypted_devices`. `stale_users` above already used the + // encrypted subset to find which devices to re-resolve. + skdm_devices: distribution_list.unwrap_or_default(), + stale_device_users: stale_users, + message_secret: reporting_result.map(|r| r.message_secret), + sender_identity: own_sending_jid, + }) +} + +/// Build the device set hashed into a group `phash`, matching WA Web +/// `phashV2([].concat(A, [B]))`: every participant device (`A`) plus the +/// sending device `B`. `devices` is the resolved set (recipients); the sending +/// device is excluded from it (we never SKDM ourselves) so it is appended here. +/// Hosted devices don't take part in group E2EE and are dropped, mirroring the +/// SKDM distribution filter. `participant_list_hash` sorts before hashing, so +/// order here is irrelevant. +pub(crate) fn build_group_phash_set(devices: &[Jid], own_sending_jid: &Jid) -> Vec<Jid> { + let mut set: Vec<Jid> = devices.iter().filter(|d| !d.is_hosted()).cloned().collect(); + if !set + .iter() + .any(|d| d.user == own_sending_jid.user && d.device == own_sending_jid.device) + { + set.push(own_sending_jid.clone()); + } + crate::types::jid::sort_dedup_by_device(&mut set); + set +} + +/// Collect users whose devices failed SKDM so the caller can invalidate their +/// registry entries. In LID-mode groups, both the LID and PN aliases are +/// emitted when the group knows the mapping — `invalidate_device_cache` needs +/// both to clean up zombie records that were stored under whichever alias +/// `update_device_list` canonicalised to at the time of the write. +pub(crate) fn collect_stale_device_users( + distribution_list: Option<&[Jid]>, + skdm_encrypted_devices: &[Jid], + group_info: &GroupInfo, +) -> Vec<String> { + let Some(dist) = distribution_list else { + return Vec::new(); + }; + let is_lid_mode = group_info.addressing_mode == crate::types::message::AddressingMode::Lid; + let encrypted_set: HashSet<&Jid> = skdm_encrypted_devices.iter().collect(); + let mut user_set: HashSet<String> = HashSet::new(); + for d in dist { + if encrypted_set.contains(d) { + continue; + } + user_set.insert(d.user.to_string()); + if is_lid_mode + && d.is_lid() + && let Some(pn_jid) = group_info.phone_jid_for_lid_user(&d.user) + && pn_jid.is_pn() + { + user_set.insert(pn_jid.user.to_string()); + } + } + user_set.into_iter().collect() +} + +/// Caller must hold `SenderKeyStore::sender_key_lock` for `sender_key_name` +/// across this creation + the matching skmsg encrypt (see `encrypt_group_message`). +pub async fn create_sender_key_distribution_message_for_group( + store: &mut (dyn SenderKeyStore + Send + Sync), + sender_key_name: &SenderKeyName, +) -> Result<Vec<u8>> { + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + + let skdm = crate::libsignal::protocol::create_sender_key_distribution_message( + sender_key_name, + store, + &mut rng, + ) + .await?; + + Ok(skdm.into_serialized().into_vec()) +} + +/// Build a `Message.ProtocolMessage` for `GROUP_MEMBER_LABEL_CHANGE`. +/// +/// Sent via the standard E2EE fanout, not an IQ. Empty `label` clears. +/// `ts_secs` is unix seconds, matching WA Web's `unixTime()`. +pub fn build_member_label_message(label: String, ts_secs: i64) -> wa::Message { + wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::GroupMemberLabelChange as i32), + member_label: Some(wa::MemberLabel { + label: Some(label), + label_timestamp: Some(ts_secs), + }), + ..Default::default() + })), + ..Default::default() + } +} diff --git a/wacore/src/send/peer.rs b/wacore/src/send/peer.rs new file mode 100644 index 000000000..8ca5ec381 --- /dev/null +++ b/wacore/src/send/peer.rs @@ -0,0 +1,95 @@ +//! Peer (own-device) stanza preparation. + +use super::*; + +#[allow(clippy::too_many_arguments)] +pub async fn prepare_peer_stanza<S, I>( + session_store: &mut S, + identity_store: &mut I, + transport_jid: Jid, + signal_address: &ProtocolAddress, + message: &wa::Message, + request_id: String, + account: Option<&wa::AdvSignedDeviceIdentity>, +) -> Result<Node> +where + S: crate::libsignal::protocol::SessionStore, + I: crate::libsignal::protocol::IdentityKeyStore, +{ + let options = peer_message_options_from_message(message); + prepare_peer_stanza_with_options( + session_store, + identity_store, + transport_jid, + signal_address, + message, + request_id, + account, + options, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +pub async fn prepare_peer_stanza_with_options<S, I>( + session_store: &mut S, + identity_store: &mut I, + transport_jid: Jid, + signal_address: &ProtocolAddress, + message: &wa::Message, + request_id: String, + account: Option<&wa::AdvSignedDeviceIdentity>, + options: PeerMessageOptions, +) -> Result<Node> +where + S: crate::libsignal::protocol::SessionStore, + I: crate::libsignal::protocol::IdentityKeyStore, +{ + let plaintext = MessageUtils::encode_and_pad(message); + + if account.is_none() && pkmsg_would_be_emitted(session_store, signal_address).await? { + bail!( + "peer pkmsg requires <device-identity> (account is None); \ + refusing before message_encrypt to avoid advancing the sender chain" + ); + } + + let encrypted_message = + message_encrypt(&plaintext, signal_address, session_store, identity_store).await?; + + let (enc_type, is_prekey, serialized_bytes) = extract_ciphertext(encrypted_message) + .ok_or_else(|| anyhow!("Unexpected peer encryption message type"))?; + + let enc_node = NodeBuilder::new("enc") + .attrs([("v", "2"), ("type", enc_type)]) + .bytes(serialized_bytes) + .build(); + + let meta_node = NodeBuilder::new("meta").attr("appdata", "default").build(); + + let mut children = vec![meta_node, enc_node]; + if is_prekey { + // Defense in depth: pre-flight should have caught this, but a corrupt + // session that triggers a fresh pkmsg mid-call would slip past. + let account = account.ok_or_else(|| { + anyhow!("peer pkmsg without <device-identity> (unreachable via pre-flight)") + })?; + children.push( + NodeBuilder::new("device-identity") + .bytes(account.encode_to_vec()) + .build(), + ); + } + + let mut stanza_builder = NodeBuilder::new("message") + .attr("to", transport_jid) + .attr("id", request_id) + .attr("type", stanza::MSG_TYPE_TEXT) + .attr("category", "peer") + .attr("push_priority", options.push_priority().as_str()); + if let Some(privacy_sensitive) = options.privacy_sensitive() { + stanza_builder = stanza_builder.attr("privacy_sensitive", privacy_sensitive.as_str()); + } + + Ok(stanza_builder.children(children).build()) +} diff --git a/wacore/src/send/status.rs b/wacore/src/send/status.rs new file mode 100644 index 000000000..64b2d0a0a --- /dev/null +++ b/wacore/src/send/status.rs @@ -0,0 +1,125 @@ +//! Status-broadcast participant assembly and privacy metadata. + +use super::*; + +/// Ensure the status stanza has a `<participants>` node listing all recipient +/// user JIDs. WhatsApp Web's `participantList` uses bare USER JIDs (not +/// device JIDs) -- `<to jid="user@s.whatsapp.net"/>` -- to tell the server +/// which users should receive the skmsg. The SKDM distribution list +/// (already in `<participants>`) uses device JIDs with `<enc>` children. +/// +/// This is a pure function (no runtime or client dependencies). +pub fn ensure_status_participants( + mut stanza: Node, + group_info: &crate::client::context::GroupInfo, +) -> Node { + use wacore_binary::NodeContent; + use wacore_binary::builder::NodeBuilder; + + // Build bare <to jid="USER_JID"/> entries for each participant. + // WhatsApp Web uses USER_JID (not DEVICE_JID) for the participantList. + let bare_to_nodes: Vec<Node> = group_info + .participants + .iter() + .map(|jid| NodeBuilder::new("to").attr("jid", jid.to_non_ad()).build()) + .collect(); + + // Check if <participants> already exists in the stanza children + let children = match &mut stanza.content { + Some(NodeContent::Nodes(nodes)) => nodes, + _ => { + stanza.content = Some(NodeContent::Nodes(vec![])); + match &mut stanza.content { + Some(NodeContent::Nodes(nodes)) => nodes, + _ => unreachable!(), + } + } + }; + + if let Some(participants_node) = children.iter_mut().find(|n| n.tag == "participants") { + // <participants> already exists (from SKDM distribution). + // Add bare <to> user JID entries for users whose devices are NOT + // already represented by SKDM device-level entries. + let existing_users: std::collections::HashSet<wacore_binary::CompactString> = + participants_node + .children() + .unwrap_or_default() + .iter() + .filter_map(|n| n.attrs.get("jid").and_then(|v| v.to_jid()).map(|j| j.user)) + .collect(); + + let new_to_nodes: Vec<Node> = bare_to_nodes + .into_iter() + .filter(|n| { + n.attrs + .get("jid") + .and_then(|v| v.to_jid()) + .is_some_and(|j| !existing_users.contains(&j.user)) + }) + .collect(); + + if !new_to_nodes.is_empty() { + match &mut participants_node.content { + Some(NodeContent::Nodes(nodes)) => nodes.extend(new_to_nodes), + _ => { + participants_node.content = Some(NodeContent::Nodes(new_to_nodes)); + } + } + } + } else { + // No <participants> node — create one with bare <to> entries. + let participants_node = NodeBuilder::new("participants") + .children(bare_to_nodes) + .build(); + children.insert(0, participants_node); + } + + stanza +} + +/// True when a `status@broadcast` message should carry the +/// `<meta status_setting="..."/>` child. Only applies to actual status posts: +/// reactions (handled server-side as addons) and revokes must omit it, per +/// `WAWebEncryptAndSendStatusMsg` vs `WAWebSendReactionMsgAction`. +/// +/// Descends `ephemeral_message` / `device_sent_message` / view-once wrappers +/// before classifying (same as `stanza_type_from_message`), so a reaction +/// nested inside a wrapper cannot slip past and re-trigger 479. +pub fn status_carries_privacy_meta(message: &wa::Message) -> bool { + let msg = unwrap_message(message); + let is_revoke = msg + .protocol_message + .as_ref() + .is_some_and(|pm| pm.r#type == Some(wa::message::protocol_message::Type::Revoke as i32)); + let is_reaction = msg.reaction_message.is_some() || msg.enc_reaction_message.is_some(); + !is_revoke && !is_reaction +} + +/// 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 +/// entries; an empty result means "nothing to send to"). +/// +/// Pure function: no allocations besides the returned `Vec` and (when needed) +/// the own-LID push. Dedup is a linear Vec scan — status lists stay small +/// enough that a HashSet is not worth its allocation. +pub fn assemble_status_participants<I>(resolved: I, own_lid: &Jid) -> anyhow::Result<Vec<Jid>> +where + I: IntoIterator<Item = Option<Jid>>, +{ + let iter = resolved.into_iter(); + let (lower, _upper) = iter.size_hint(); + let mut out: Vec<Jid> = Vec::with_capacity(lower.saturating_add(1)); + for jid in iter.flatten() { + if !out.iter().any(|r| r.user == jid.user) { + out.push(jid); + } + } + if out.is_empty() { + anyhow::bail!("No valid status recipients after LID resolution"); + } + if !out.iter().any(|r| r.user == own_lid.user) { + out.push(own_lid.to_non_ad()); + } + Ok(out) +} diff --git a/wacore/src/send/tests.rs b/wacore/src/send/tests.rs new file mode 100644 index 000000000..076ce39a6 --- /dev/null +++ b/wacore/src/send/tests.rs @@ -0,0 +1,3256 @@ +//! Tests for stanza preparation and encryption fanout. + +use super::*; +use crate::client::context::{GroupInfo, SendContextResolver}; +use crate::libsignal::protocol::{IdentityKeyPair, KeyPair, PreKeyBundle}; +use std::collections::HashMap; +use wacore_binary::Jid; + +mod assemble_status_participants { + use super::*; + + fn lid(u: &str) -> Jid { + u.parse().expect("parse LID jid") + } + + #[test] + fn dedup_keeps_first_entry_per_user_and_anchors_own() { + let own = lid("99999999999999@lid"); + let out = assemble_status_participants( + vec![ + Some(lid("111@lid")), + Some(lid("222@lid")), + Some(lid("111@lid")), + Some(lid("333@lid")), + ], + &own, + ) + .expect("should succeed"); + let users: Vec<&str> = out.iter().map(|j| j.user.as_str()).collect(); + assert_eq!(users, ["111", "222", "333", "99999999999999"]); + } + + #[test] + fn skips_none_entries_matching_wa_web_compactmap() { + // Unresolvable recipients arrive as `None` and must be silently + // dropped — mirrors WA Web's `compactMap(list, toUserLid)`. + let own = lid("me@lid"); + let out = assemble_status_participants( + vec![None, Some(lid("111@lid")), None, Some(lid("222@lid"))], + &own, + ) + .expect("should succeed"); + let users: Vec<&str> = out.iter().map(|j| j.user.as_str()).collect(); + assert_eq!(users, ["111", "222", "me"]); + } + + #[test] + fn does_not_duplicate_own_when_already_in_list() { + let own = lid("me@lid"); + let out = + assemble_status_participants(vec![Some(lid("111@lid")), Some(lid("me@lid"))], &own) + .expect("should succeed"); + let users: Vec<&str> = out.iter().map(|j| j.user.as_str()).collect(); + assert_eq!(users, ["111", "me"]); + } + + #[test] + fn errors_when_every_recipient_is_unresolvable() { + // Regression guard for the original bug: a single LID-only + // contact used to hard-abort the send with + // `No PN mapping for LID ...`. The new contract is softer — + // individual unresolvable entries are dropped — but we still + // refuse to send when the entire list came back empty, rather + // than silently broadcasting to own devices only. + let own = lid("me@lid"); + let err = assemble_status_participants(vec![None, None, None], &own) + .expect_err("all-None list must error"); + assert!(err.to_string().contains("No valid status recipients")); + } + + #[test] + fn errors_when_list_is_empty() { + let own = lid("me@lid"); + let err = assemble_status_participants(Vec::<Option<Jid>>::new(), &own) + .expect_err("empty list must error"); + assert!(err.to_string().contains("No valid status recipients")); + } + + #[test] + fn strips_device_suffix_from_own_lid() { + // Snapshot lid from the device store carries a device id; the + // participant list uses bare USER JIDs. + let own: Jid = "me:5@lid".parse().unwrap(); + let out = + assemble_status_participants(vec![Some(lid("111@lid"))], &own).expect("should succeed"); + let me = out + .iter() + .find(|j| j.user.as_str() == "me") + .expect("own LID should be present"); + assert_eq!(me.device, 0, "own LID should be non-ad (device=0)"); + } +} + +mod peer_message_options { + use super::*; + use crate::types::message::{PrivacySensitiveType, PushPriority}; + + fn pdo_message_raw(request_type: i32) -> wa::Message { + wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + r#type: Some( + wa::message::protocol_message::Type::PeerDataOperationRequestMessage as i32, + ), + peer_data_operation_request_message: Some( + wa::message::PeerDataOperationRequestMessage { + peer_data_operation_request_type: Some(request_type), + ..Default::default() + }, + ), + ..Default::default() + })), + ..Default::default() + } + } + + fn pdo_message(request_type: wa::message::PeerDataOperationRequestType) -> wa::Message { + pdo_message_raw(request_type as i32) + } + + #[test] + fn pdo_priority_map_matches_wa_web_non_message_requests() { + use wa::message::PeerDataOperationRequestType as PdoType; + + let high_force_cases = [ + (PdoType::GenerateLinkPreview, PushPriority::HighForce, None), + ( + PdoType::PlaceholderMessageResend, + PushPriority::HighForce, + None, + ), + ( + PdoType::HistorySyncOnDemand, + PushPriority::HighForce, + Some(PrivacySensitiveType::OnDemand), + ), + ( + PdoType::CompanionCanonicalUserNonceFetch, + PushPriority::HighForce, + None, + ), + ]; + + for (request_type, push_priority, privacy_sensitive) in high_force_cases { + let options = peer_message_options_from_message(&pdo_message(request_type)); + assert_eq!(options.push_priority(), push_priority, "{request_type:?}"); + assert_eq!( + options.privacy_sensitive(), + privacy_sensitive, + "{request_type:?}" + ); + } + + let default_cases = [ + PdoType::UploadSticker, + PdoType::SendRecentStickerBootstrap, + PdoType::WaffleLinkingNonceFetch, + PdoType::FullHistorySyncOnDemand, + PdoType::CompanionMetaNonceFetch, + PdoType::CompanionSyncdSnapshotFatalRecovery, + PdoType::HistorySyncChunkRetry, + PdoType::GalaxyFlowAction, + PdoType::BusinessBroadcastInsightsDeliveredTo, + PdoType::BusinessBroadcastInsightsRefresh, + ]; + + for request_type in default_cases { + let options = peer_message_options_from_message(&pdo_message(request_type)); + assert_eq!( + options.push_priority(), + PushPriority::High, + "{request_type:?}" + ); + assert_eq!(options.privacy_sensitive(), None, "{request_type:?}"); + } + } + + #[test] + fn non_pdo_and_unknown_pdo_keep_peer_defaults() { + let app_state_key_request = wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::AppStateSyncKeyRequest as i32), + app_state_sync_key_request: Some(wa::message::AppStateSyncKeyRequest { + key_ids: Vec::new(), + }), + ..Default::default() + })), + ..Default::default() + }; + + for msg in [app_state_key_request, pdo_message_raw(99)] { + let options = peer_message_options_from_message(&msg); + assert_eq!(options.push_priority(), PushPriority::High); + assert_eq!(options.privacy_sensitive(), None); + } + } +} + +mod status_carries_privacy_meta { + use super::*; + + #[test] + fn true_for_text_post() { + let msg = wa::Message { + extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { + text: Some("hi".into()), + ..Default::default() + })), + ..Default::default() + }; + assert!(status_carries_privacy_meta(&msg)); + } + + #[test] + fn true_for_image_post() { + let msg = wa::Message { + image_message: Some(Box::new(wa::message::ImageMessage::default())), + ..Default::default() + }; + assert!(status_carries_privacy_meta(&msg)); + } + + #[test] + fn false_for_reaction() { + let msg = wa::Message { + reaction_message: Some(wa::message::ReactionMessage { + text: Some("💚".into()), + ..Default::default() + }), + ..Default::default() + }; + assert!( + !status_carries_privacy_meta(&msg), + "reactions must omit <meta status_setting> (479 SmaxInvalid otherwise)" + ); + } + + #[test] + fn false_for_enc_reaction() { + let msg = wa::Message { + enc_reaction_message: Some(wa::message::EncReactionMessage::default()), + ..Default::default() + }; + assert!(!status_carries_privacy_meta(&msg)); + } + + #[test] + fn false_for_revoke() { + let msg = wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::Revoke as i32), + ..Default::default() + })), + ..Default::default() + }; + assert!(!status_carries_privacy_meta(&msg)); + } + + #[test] + fn true_for_non_revoke_protocol_message() { + // Other ProtocolMessage types (e.g., EphemeralSettings) aren't + // reactions and aren't revokes — treat as posts for now. + let msg = wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::EphemeralSetting as i32), + ..Default::default() + })), + ..Default::default() + }; + assert!(status_carries_privacy_meta(&msg)); + } + + #[test] + fn false_for_reaction_inside_ephemeral_wrapper() { + let inner = wa::Message { + reaction_message: Some(wa::message::ReactionMessage::default()), + ..Default::default() + }; + let msg = wa::Message { + ephemeral_message: Some(Box::new(wa::message::FutureProofMessage { + message: Some(Box::new(inner)), + })), + ..Default::default() + }; + assert!(!status_carries_privacy_meta(&msg)); + } + + #[test] + fn false_for_revoke_inside_device_sent_wrapper() { + let inner = wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + r#type: Some(wa::message::protocol_message::Type::Revoke as i32), + ..Default::default() + })), + ..Default::default() + }; + let msg = wa::Message { + device_sent_message: Some(Box::new(wa::message::DeviceSentMessage { + destination_jid: Some(String::new()), + message: Some(Box::new(inner)), + ..Default::default() + })), + ..Default::default() + }; + assert!(!status_carries_privacy_meta(&msg)); + } +} + +#[test] +fn build_member_label_message_sets_fields() { + let msg = build_member_label_message("VIP".to_string(), 1_766_847_151); + let pm = msg.protocol_message.as_ref().expect("protocol_message set"); + assert_eq!( + pm.r#type, + Some(wa::message::protocol_message::Type::GroupMemberLabelChange as i32) + ); + let ml = pm.member_label.as_ref().expect("member_label set"); + assert_eq!(ml.label.as_deref(), Some("VIP")); + assert_eq!(ml.label_timestamp, Some(1_766_847_151)); + assert!( + pm.key.is_none(), + "MessageKey must NOT be set (WA Web parity)" + ); +} + +#[test] +fn build_member_label_message_clear_uses_empty_string() { + let msg = build_member_label_message(String::new(), 1); + let ml = msg + .protocol_message + .as_ref() + .unwrap() + .member_label + .as_ref() + .unwrap(); + assert_eq!(ml.label.as_deref(), Some("")); +} + +#[test] +fn build_member_label_message_preserves_unicode() { + let msg = build_member_label_message("🚀 BOT".to_string(), 2); + let ml = msg + .protocol_message + .as_ref() + .unwrap() + .member_label + .as_ref() + .unwrap(); + assert_eq!(ml.label.as_deref(), Some("🚀 BOT")); +} + +/// Mock implementation of SendContextResolver for testing +struct MockSendContextResolver { + /// Pre-key bundles to return: JID -> Option<PreKeyBundle> + prekey_bundles: HashMap<Jid, Option<PreKeyBundle>>, + /// Devices to return from resolve_devices + devices: Vec<Jid>, + /// Phone number to LID mappings for testing LID session lookup + phone_to_lid: HashMap<String, String>, + /// JIDs reported via `on_local_identity_change` (send-path detection). + identity_changes: std::sync::Mutex<Vec<Jid>>, +} + +impl MockSendContextResolver { + fn new() -> Self { + Self { + prekey_bundles: HashMap::new(), + devices: Vec::new(), + phone_to_lid: HashMap::new(), + identity_changes: std::sync::Mutex::new(Vec::new()), + } + } + + fn captured_identity_changes(&self) -> Vec<Jid> { + self.identity_changes.lock().unwrap().clone() + } + + fn with_missing_bundle(mut self, jid: Jid) -> Self { + self.prekey_bundles.insert(jid, None); + self + } + + fn with_bundle(mut self, jid: Jid, bundle: PreKeyBundle) -> Self { + self.prekey_bundles.insert(jid, Some(bundle)); + self + } + + fn with_devices(mut self, devices: Vec<Jid>) -> Self { + self.devices = devices; + self + } + + fn with_phone_to_lid(mut self, phone: &str, lid: &str) -> Self { + self.phone_to_lid.insert(phone.to_string(), lid.to_string()); + self + } +} + +#[async_trait::async_trait] +impl SendContextResolver for MockSendContextResolver { + async fn resolve_devices(&self, _jids: &[Jid]) -> Result<Vec<Jid>> { + Ok(self.devices.clone()) + } + + async fn fetch_prekeys(&self, jids: &[Jid]) -> Result<HashMap<Jid, PreKeyBundle>> { + let mut result = HashMap::new(); + for jid in jids { + if let Some(bundle_opt) = self.prekey_bundles.get(jid) + && let Some(bundle) = bundle_opt + { + result.insert(jid.clone(), bundle.clone()); + } + } + Ok(result) + } + + async fn fetch_prekeys_for_identity_check( + &self, + jids: &[Jid], + ) -> Result<HashMap<Jid, PreKeyBundle>> { + let mut result = HashMap::new(); + for jid in jids { + if let Some(bundle_opt) = self.prekey_bundles.get(jid) + && let Some(bundle) = bundle_opt + { + result.insert(jid.clone(), bundle.clone()); + } + // If None, we intentionally omit it from the result (simulating server not returning it) + } + Ok(result) + } + + async fn resolve_group_info(&self, _jid: &Jid) -> Result<std::sync::Arc<GroupInfo>> { + unimplemented!("resolve_group_info not needed for send.rs tests") + } + + async fn get_lid_for_phone(&self, phone_user: &str) -> Option<wacore_binary::CompactString> { + self.phone_to_lid.get(phone_user).map(|s| s.as_str().into()) + } + + fn on_local_identity_change(&self, jid: &Jid) { + self.identity_changes.lock().unwrap().push(jid.clone()); + } +} + +/// Test case: Missing pre-key bundle for a single device skips gracefully +/// +/// When sending to multiple devices, if some don't have pre-key bundles (e.g., Cloud API), +/// we should skip them instead of failing the entire message. +#[test] +fn test_missing_prekey_bundle_skips_device() { + let device_with_bundle: Jid = "1234567890:0@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let device_without_bundle: Jid = "1234567890:1@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let cloud_api: Jid = "1234567890:99@hosted" + .parse() + .expect("test JID should be valid"); + + let bundle = create_mock_bundle(); + + let resolver = MockSendContextResolver::new() + .with_bundle(device_with_bundle.clone(), bundle) + .with_missing_bundle(device_without_bundle.clone()) + .with_missing_bundle(cloud_api.clone()) + .with_devices(vec![ + device_with_bundle.clone(), + device_without_bundle.clone(), + cloud_api.clone(), + ]); + + // Check that the resolver correctly returns only available bundles + assert_eq!( + resolver.prekey_bundles.len(), + 3, + "Resolver should have 3 entries" + ); + + // Verify device_with_bundle has a Some(bundle) + assert!( + resolver.prekey_bundles[&device_with_bundle].is_some(), + "device_with_bundle should have a Some entry" + ); + + // Verify others have None + assert!( + resolver.prekey_bundles[&device_without_bundle].is_none(), + "device_without_bundle should have None" + ); + assert!( + resolver.prekey_bundles[&cloud_api].is_none(), + "cloud_api should have None" + ); + + println!("✅ Missing pre-key bundle skips device gracefully"); +} + +/// Test case: All devices missing pre-key bundles +/// +/// If all devices are unavailable, the batch should still complete without panic. +#[test] +fn test_all_devices_missing_prekey_bundles() { + let device1: Jid = "1234567890:0@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let device2: Jid = "1234567890:1@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let device3: Jid = "9876543210:0@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + + let resolver = MockSendContextResolver::new() + .with_missing_bundle(device1.clone()) + .with_missing_bundle(device2.clone()) + .with_missing_bundle(device3.clone()) + .with_devices(vec![device1.clone(), device2.clone(), device3.clone()]); + + // All entries should be None + assert!(resolver.prekey_bundles[&device1].is_none()); + assert!(resolver.prekey_bundles[&device2].is_none()); + assert!(resolver.prekey_bundles[&device3].is_none()); + + println!("✅ All devices missing bundles handled gracefully"); +} + +/// Test case: Large group with mixed device availability +/// +/// In real-world scenarios, large groups may have some unavailable devices. +/// The encryption should proceed for available devices and skip unavailable ones. +#[test] +fn test_large_group_with_mixed_device_availability() { + let mut all_devices = Vec::new(); + + for i in 0..10u16 { + let device_jid = Jid::pn_device("1234567890", i); + all_devices.push(device_jid); + } + + let mut resolver = MockSendContextResolver::new().with_devices(all_devices.clone()); + + // Add bundles for devices 0-6, mark 7-9 as missing + for i in 0..10u16 { + let device_jid = Jid::pn_device("1234567890", i); + + if i < 7 { + resolver = resolver.with_bundle(device_jid, create_mock_bundle()); + } else { + resolver = resolver.with_missing_bundle(device_jid); + } + } + + // Verify bundle availability + let available_count = resolver + .prekey_bundles + .values() + .filter(|v| v.is_some()) + .count(); + + assert_eq!(available_count, 7, "Should have 7 available devices"); + assert_eq!( + resolver.prekey_bundles.len(), + 10, + "Should have 10 total entries" + ); + + println!("✅ Large group with 7 available, 3 unavailable devices"); +} + +/// Test case: Cloud API / HOSTED device without pre-key +/// +/// # Context: What are HOSTED devices? +/// +/// HOSTED devices (Cloud API / Meta Business API) are WhatsApp Business accounts +/// that use Meta's server-side infrastructure instead of traditional E2EE. +/// +/// ## Identification: +/// - Device ID 99 (`:99`) on any server +/// - Server `@hosted` or `@hosted.lid` +/// +/// ## Behavior: +/// - They do NOT have Signal protocol prekey bundles +/// - For 1:1 chats: included in device list, but prekey fetch fails gracefully +/// - For groups: proactively filtered out before SKDM distribution +/// +/// This test verifies that when a hosted device is included in the device list +/// (which would happen for 1:1 chats), the missing prekey is handled gracefully. +#[test] +fn test_cloud_api_device_without_prekey() { + let regular_device: Jid = "1234567890:0@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + let cloud_api: Jid = "1234567890:99@hosted" + .parse() + .expect("test JID should be valid"); + + // Verify the cloud_api device is detected as hosted + assert!( + cloud_api.is_hosted(), + "Device with :99@hosted should be detected as hosted" + ); + assert!( + !regular_device.is_hosted(), + "Regular device should NOT be detected as hosted" + ); + + let resolver = MockSendContextResolver::new() + .with_bundle(regular_device.clone(), create_mock_bundle()) + .with_missing_bundle(cloud_api.clone()) + .with_devices(vec![regular_device.clone(), cloud_api.clone()]); + + assert!( + resolver.prekey_bundles[&regular_device].is_some(), + "Regular device should have a bundle" + ); + assert!( + resolver.prekey_bundles[&cloud_api].is_none(), + "Cloud API device should not have a bundle (they don't use Signal protocol)" + ); + + println!("✅ Cloud API device has no prekey bundle (expected behavior)"); +} + +/// Test case: HOSTED devices are filtered from group SKDM distribution +/// +/// # Why filter hosted devices from groups? +/// +/// WhatsApp Web explicitly excludes hosted devices from group message fanout. +/// From the JS code (`getFanOutList`): +/// ```javascript +/// var isHosted = e.id === 99 || e.isHosted === true; +/// var includeInFanout = !isHosted || isOneToOneChat; +/// ``` +/// +/// ## Reasons: +/// 1. Hosted devices don't use Signal protocol - they can't process SKDM +/// 2. Including them causes unnecessary prekey fetch failures +/// 3. Group encryption is handled differently for Cloud API businesses +/// +/// This test verifies that `is_hosted()` correctly identifies devices that +/// should be filtered from group SKDM distribution. +#[test] +fn test_hosted_devices_filtered_from_group_skdm() { + // Simulate devices returned from usync for a group + let devices: Vec<Jid> = vec![ + // Regular devices - should receive SKDM + "5511999887766:0@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), // Primary phone + "5511999887766:33@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), // WhatsApp Web companion + "5521988776655:0@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), // Another participant + "100000012345678:33@lid" + .parse() + .expect("test JID should be valid"), // LID companion device + // HOSTED devices - should be EXCLUDED from group SKDM + "5531977665544:99@s.whatsapp.net" + .parse() + .expect("test JID should be valid"), // Cloud API on regular server + "100000087654321:99@lid" + .parse() + .expect("test JID should be valid"), // Cloud API on LID server + "5541966554433:0@hosted" + .parse() + .expect("test JID should be valid"), // Explicit @hosted server + ]; + + // This is the filtering logic used in prepare_group_stanza + let filtered_for_skdm: Vec<Jid> = devices.into_iter().filter(|jid| !jid.is_hosted()).collect(); + + assert_eq!( + filtered_for_skdm.len(), + 4, + "Should have 4 devices after filtering out hosted devices" + ); + + // Verify all remaining devices are NOT hosted + for jid in &filtered_for_skdm { + assert!( + !jid.is_hosted(), + "Filtered list should not contain hosted device: {}", + jid + ); + } + + // Verify specific devices are included/excluded by checking struct fields + // (Device ID 0 is not serialized in the string representation) + let has_primary_phone = filtered_for_skdm + .iter() + .any(|j| j.user == "5511999887766" && j.device == 0 && j.server == "s.whatsapp.net"); + let has_companion = filtered_for_skdm + .iter() + .any(|j| j.user == "5511999887766" && j.device == 33 && j.server == "s.whatsapp.net"); + let has_cloud_api = filtered_for_skdm + .iter() + .any(|j| j.user == "5531977665544" && j.device == 99); + let has_hosted_server = filtered_for_skdm.iter().any(|j| j.server == "hosted"); + + assert!(has_primary_phone, "Primary phone should be included"); + assert!(has_companion, "WhatsApp Web companion should be included"); + assert!( + !has_cloud_api, + "Cloud API device (ID 99) should be excluded" + ); + assert!( + !has_hosted_server, + "@hosted server device should be excluded" + ); + + println!("✅ Hosted devices correctly filtered from group SKDM distribution"); +} + +/// Test case: Device recovery between retries +/// +/// If a device was temporarily unavailable, a retry should succeed. +#[test] +fn test_device_recovery_between_requests() { + let device: Jid = "1234567890:0@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + + // First attempt: device unavailable + let resolver_first = MockSendContextResolver::new().with_missing_bundle(device.clone()); + + assert!( + resolver_first.prekey_bundles[&device].is_none(), + "First attempt: device should be unavailable" + ); + + // Second attempt: device recovered + let resolver_second = + MockSendContextResolver::new().with_bundle(device.clone(), create_mock_bundle()); + + assert!( + resolver_second.prekey_bundles[&device].is_some(), + "Second attempt: device should be available" + ); + + println!("✅ Device recovery between retries works correctly"); +} + +/// Helper function to create a mock PreKeyBundle with valid types +fn create_mock_bundle() -> PreKeyBundle { + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + let identity_pair = IdentityKeyPair::generate(&mut rng); + let signed_prekey_pair = KeyPair::generate(&mut rng); + let prekey_pair = KeyPair::generate(&mut rng); + + PreKeyBundle::new( + 1, // registration_id + 1u32.into(), // device_id + Some((1u32.into(), prekey_pair.public_key)), // pre_key + 2u32.into(), // signed_pre_key_id + signed_prekey_pair.public_key, + vec![0u8; 64], + *identity_pair.identity_key(), + ) + .expect("Failed to create PreKeyBundle") +} + +// These tests validate the fix for the LID-PN session mismatch issue. +// When a message is received with sender_lid, the session is stored under the LID address. +// When sending a reply using the phone number, we must reuse the existing LID session +// instead of creating a new PN session, otherwise subsequent messages will fail with +// MAC verification errors. + +/// Test that phone_to_lid mapping returns the cached LID mapping. +/// +/// This verifies the MockSendContextResolver correctly stores phone-to-LID +/// mappings used for LID session lookup. +#[test] +fn test_mock_resolver_phone_to_lid_mapping() { + let phone = "559980000001"; + let lid = "100000012345678"; + + let resolver = MockSendContextResolver::new().with_phone_to_lid(phone, lid); + + // Access the HashMap directly (synchronous) + let result = resolver.phone_to_lid.get(phone).cloned(); + + assert!(result.is_some(), "Should return LID for known phone"); + assert_eq!( + result.expect("known phone should return LID"), + lid, + "Should return correct LID" + ); + + // Unknown phone should return None + let unknown = resolver.phone_to_lid.get("999999999").cloned(); + assert!(unknown.is_none(), "Should return None for unknown phone"); + + println!("✅ MockSendContextResolver phone_to_lid mapping works correctly"); +} + +/// Test that the resolver correctly maps phone numbers to LIDs. +/// +/// This is a building block for the session lookup logic. +#[test] +fn test_phone_to_lid_mapping_multiple_users() { + let resolver = MockSendContextResolver::new() + .with_phone_to_lid("559980000001", "100000012345678") + .with_phone_to_lid("559980000002", "100000024691356") + .with_phone_to_lid("559980000003", "100000037037034"); + + // Verify all mappings using direct HashMap access + let lid1 = resolver.phone_to_lid.get("559980000001").cloned(); + let lid2 = resolver.phone_to_lid.get("559980000002").cloned(); + let lid3 = resolver.phone_to_lid.get("559980000003").cloned(); + + assert_eq!( + lid1.expect("phone 1 should have LID mapping"), + "100000012345678" + ); + assert_eq!( + lid2.expect("phone 2 should have LID mapping"), + "100000024691356" + ); + assert_eq!( + lid3.expect("phone 3 should have LID mapping"), + "100000037037034" + ); + + println!("✅ Multiple phone-to-LID mappings work correctly"); +} + +/// Test the scenario that caused the original bug: +/// - Session exists under LID address (from receiving a message with sender_lid) +/// - Send to PN address should reuse the LID session, not create a new one +/// +/// This test verifies the logic flow, though full integration testing +/// requires the actual encrypt_for_devices function with real sessions. +#[test] +fn test_lid_session_lookup_scenario() { + // Scenario setup: + // - Received message from 559980000001@s.whatsapp.net with sender_lid=100000012345678@lid + // - Session was stored under 100000012345678.0 + // - Now sending reply to 559980000001@s.whatsapp.net + // - Should look up LID and check for session under 100000012345678.0 + + let phone = "559980000001"; + let lid = "100000012345678"; + let device_id = 0u16; + + let resolver = MockSendContextResolver::new().with_phone_to_lid(phone, lid); + + // Simulate the device JID we're trying to send to (PN format) + let pn_device_jid = Jid::pn_device(phone, device_id); + + // Step 1: Look up LID for the phone number (using direct HashMap access) + let lid_user = resolver + .phone_to_lid + .get(pn_device_jid.user.as_str()) + .cloned(); + assert!(lid_user.is_some(), "Should find LID for phone"); + let lid_user = lid_user.expect("phone should have LID mapping"); + + // Step 2: Construct the LID JID with same device ID + let lid_jid = Jid::lid_device(lid_user.clone(), pn_device_jid.device); + + // Step 3: Verify the LID JID is correctly constructed + assert_eq!(lid_jid.user, lid, "LID user should match"); + assert_eq!(lid_jid.server, "lid", "Server should be 'lid'"); + assert_eq!(lid_jid.device, device_id, "Device ID should be preserved"); + + // Step 4: Convert to protocol addresses and verify they're different + use crate::types::jid::JidExt; + let pn_address = pn_device_jid.to_protocol_address(); + let lid_address = lid_jid.to_protocol_address(); + + assert_ne!( + pn_address.name(), + lid_address.name(), + "PN and LID addresses should have different names" + ); + assert_eq!( + pn_address.device_id(), + lid_address.device_id(), + "Device IDs should match" + ); + + println!("✅ LID session lookup scenario works correctly:"); + println!(" - PN JID: {} -> Address: {}", pn_device_jid, pn_address); + println!(" - LID JID: {} -> Address: {}", lid_jid, lid_address); + println!(" - Would check for session under LID address first"); +} + +/// Test that companion device IDs are preserved in LID JID construction. +/// +/// WhatsApp Web uses device ID 33, and this must be preserved when +/// constructing the LID JID for session lookup. +#[test] +fn test_lid_jid_preserves_companion_device_id() { + let phone = "559980000001"; + let lid = "100000012345678"; + let companion_device_id = 33u16; // WhatsApp Web device ID + + let resolver = MockSendContextResolver::new().with_phone_to_lid(phone, lid); + + // Simulate sending to a companion device (WhatsApp Web) + let pn_device_jid = Jid::pn_device(phone, companion_device_id); + + // Look up LID using direct HashMap access + let lid_user = resolver + .phone_to_lid + .get(pn_device_jid.user.as_str()) + .cloned(); + + // Construct LID JID + let lid_jid = Jid::lid_device( + lid_user.expect("phone should have LID mapping for companion test"), + pn_device_jid.device, + ); + + assert_eq!( + lid_jid.device, companion_device_id, + "Device ID 33 should be preserved" + ); + assert_eq!(lid_jid.to_string(), "100000012345678:33@lid"); + + println!("✅ Companion device ID (33) correctly preserved in LID JID"); +} + +/// Test that LID lookup only applies to s.whatsapp.net JIDs. +/// +/// LID JIDs (@lid) and group JIDs (@g.us) should not trigger LID lookup. +#[test] +fn test_lid_lookup_only_for_pn_jids() { + let _resolver = + MockSendContextResolver::new().with_phone_to_lid("559980000001", "100000012345678"); + + // These JIDs should NOT trigger LID lookup + let lid_jid: Jid = "100000012345678:0@lid" + .parse() + .expect("test JID should be valid"); + let group_jid: Jid = "120363123456789012@g.us" + .parse() + .expect("test JID should be valid"); + + // Only s.whatsapp.net JIDs should be looked up + assert_ne!( + lid_jid.server, "s.whatsapp.net", + "LID JID should not be s.whatsapp.net" + ); + assert_ne!( + group_jid.server, "s.whatsapp.net", + "Group JID should not be s.whatsapp.net" + ); + + // PN JID should be eligible for lookup + let pn_jid: Jid = "559980000001:0@s.whatsapp.net" + .parse() + .expect("test JID should be valid"); + assert_eq!( + pn_jid.server, "s.whatsapp.net", + "PN JID should be s.whatsapp.net" + ); + + println!("✅ LID lookup correctly limited to s.whatsapp.net JIDs"); +} + +/// Test case: Regression test for self-encryption bug. +/// +/// The sender's own device (e.g. device 79) must be excluded from the encryption list +/// to prevent "SESSION BASE KEY CHANGED" warnings caused by establishing a session with oneself. +#[test] +fn test_dm_encryption_excludes_sender_device() { + // Setup: + // - Own user: 123456789 + // - Specific own device (Sender): 79 + // - Other own device: 0 + // - Recipient: 987654321 + + let own_user = "123456789"; + let own_device_id = 79; + + // Own JID (Sender) + let own_jid = Jid::lid_device(own_user.to_string(), own_device_id); + + // Simulate devices returned by resolver.resolve_devices() + // This includes: + // 1. The sender's own device (should be excluded) + // 2. Another device of the sender (should be in own_other_devices) + // 3. The recipient's device (should be in recipient_devices) + let all_devices: Vec<Jid> = vec![ + Jid::lid_device(own_user.to_string(), own_device_id), // Sender (79) + Jid::lid_device(own_user.to_string(), 0), // Other own device (0) + Jid::lid_device("987654321".to_string(), 0), // Recipient + ]; + + let (recipient_devices, own_other_devices) = partition_dm_devices(all_devices, &own_jid, None); + + // Verifications + + // 1. Sender device (79) should NOT be in either list + let sender_in_own = own_other_devices.iter().any(|d| d.device == own_device_id); + let sender_in_recipient = recipient_devices.iter().any(|d| d.device == own_device_id); + + assert!( + !sender_in_own, + "Sender device (79) should be excluded from own_other_devices" + ); + assert!( + !sender_in_recipient, + "Sender device (79) should be excluded from recipient_devices" + ); + + // 2. Other own device (0) MUST be in own_other_devices + let other_own_present = own_other_devices + .iter() + .any(|d| d.device == 0 && d.user == own_user); + assert!( + other_own_present, + "Other own device (0) should be included in own_other_devices" + ); + + // 3. Recipient MUST be in recipient_devices + let recipient_present = recipient_devices.iter().any(|d| d.user == "987654321"); + assert!( + recipient_present, + "Recipient should be included in recipient_devices" + ); + + println!("✅ Self-encryption regression test passed: Sender device correctly excluded."); +} + +#[test] +fn test_dm_encryption_treats_own_lid_devices_as_self() { + let own_pn = Jid::pn_device("559980000001".to_string(), 18); + let own_lid = Jid::lid_device("123456789012345".to_string(), 18); + + let all_devices = vec![ + Jid::lid_device("123456789012345".to_string(), 18), // Exact sender device via LID + Jid::lid_device("123456789012345".to_string(), 0), // Other own device via LID + Jid::lid_device("987654321012345".to_string(), 0), // Recipient + ]; + + let (recipient_devices, own_other_devices) = + partition_dm_devices(all_devices, &own_pn, Some(&own_lid)); + + assert!( + !own_other_devices + .iter() + .any(|d| d.user == own_lid.user && d.device == 18), + "Exact sender LID device should be excluded from own_other_devices" + ); + assert!( + !recipient_devices + .iter() + .any(|d| d.user == own_lid.user && d.device == 18), + "Exact sender LID device should be excluded from recipient_devices" + ); + assert!( + own_other_devices + .iter() + .any(|d| d.user == own_lid.user && d.device == 0), + "Other own LID devices should be routed through DSM as own_other_devices" + ); + assert!( + recipient_devices + .iter() + .any(|d| d.user == "987654321012345" && d.device == 0), + "Non-self devices must remain in recipient_devices" + ); +} + +/// Test case: LID Prekey Lookup Normalization +/// +/// Verifies that when looking up pre-key bundles for LID JIDs, the lookup key +/// is normalized (agent=0) to match how the bundles are stored in the map. +/// +/// This validates the fix for "No pre-key bundle returned" when the requested JID +/// has non-standard agent/server fields but the bundle is stored under the normalized key. +#[test] +fn test_lid_prekey_lookup_normalization() { + // 1. Define JIDs + // The JID we request (simulating what comes from resolve_devices or elsewhere) + // Let's pretend it has agent=1 to simulate a mismatch + let mut requested_jid = Jid::lid_device("123456789".to_string(), 0); + requested_jid.agent = 1; + + // The normalized JID (how it's stored in the bundle map) + let normalized_jid = Jid::lid_device("123456789".to_string(), 0); // agent=0 by default + + // 2. Setup Resolver + // Store the bundle under the NORMALIZED key (agent=0) + let resolver = MockSendContextResolver::new() + .with_bundle(normalized_jid.clone(), create_mock_bundle()) + .with_devices(vec![requested_jid.clone()]); + + // 3. Verify Mock Setup + // Ensure bundle is accessible via normalized key but NOT via requested (raw) key + // This confirms our test condition is valid (that implicit lookup would fail) + assert!( + resolver.prekey_bundles.contains_key(&normalized_jid), + "Setup: bundle should exist for normalized key" + ); + assert!( + !resolver.prekey_bundles.contains_key(&requested_jid), + "Setup: bundle should NOT exist for requested raw key" + ); + + // 4. Test logic mirroring `encrypt_for_devices` + let mut jid_to_encryption_jid = HashMap::new(); + // Assume direct mapping for simplicity + jid_to_encryption_jid.insert(requested_jid.clone(), requested_jid.clone()); + + // Get the bundles map (mocks `fetch_prekeys_for_identity_check`) + // The mock implementation returns the map as-is filtered by keys. + // HOWEVER, `fetch_prekeys` usually takes a list. + // In `encrypt_for_devices`, we call: + // let prekey_bundles = resolver.fetch_prekeys_for_identity_check(&[requested_jid]).await?; + + // Let's simulate what `fetch_prekeys_for_identity_check` would return. + // Our mock implementation `fetch_prekeys` logic: + // if let Some(bundle_opt) = self.prekey_bundles.get(jid) + + // Wait, if the mock follows exact HashMap lookup, `fetch_prekeys(&[requested_jid])` + // will return EMPTY because `requested_jid` is not in `prekey_bundles`. + // The REAL `fetch_prekeys` (in `client.rs` -> `prekeys.rs`) sends an IQ to the server, + // and the server response is parsed. The parsing logic (in `prekeys.rs`) normalizes the key. + // So the HashMap returned by `fetch_prekeys` will contain NORMALIZED keys. + + // So for this test to be accurate, we must simulate that `fetch_prekeys` returned a map + // where the key is NORMALIZED, even if we asked for `requested_jid`? + // Actually, `PreKeyFetchSpec` asks for JIDs. The response contains JIDs. + // If we ask for `agent=1`, does the server return `agent=1`? + // The logs showed: + // parsed: `...:82@lid` (agent=0 probably, or just not printed?) + // lookup: `...` (failed) + + // The critical part is that the `HashMap` returned by `resolver.fetch_prekeys` + // definitely contains the bundle under some key. + // If `prekeys.rs` normalizes it, it's under the normalized key. + // The `encrypt_for_devices` logic has: + // `match prekey_bundles.get(device_jid)` + // where `device_jid` is the one from the loop (requested_jid). + + // If `fetch_prekeys` returns a map with `normalized_jid`, and we lookup `requested_jid`, it fails. + // My fix was to normalize `requested_jid` before lookup. + + // So I need to construct the `prekey_bundles` map manually here to simulate the return from fetch. + let mut prekey_bundles = HashMap::new(); + prekey_bundles.insert(normalized_jid.clone(), create_mock_bundle()); + + // Now test the logic: + let device_jid = &requested_jid; + + // -- Logic from fix -- + // Use centralized normalization logic + let lookup_jid = device_jid.normalize_for_prekey_bundle(); + + // Fix: Use the normalized device_jid to lookup the bundle + let bundle = prekey_bundles.get(&lookup_jid); + // -------------------- + + assert!(bundle.is_some(), "Should find bundle after normalization"); + + // Verify it would have failed without normalization + let raw_lookup = prekey_bundles.get(device_jid); + assert!( + raw_lookup.is_none(), + "Should NOT find bundle without normalization" + ); + + println!("✅ LID Prekey Lookup Normalization passed"); +} + +mod group_retry { + use super::*; + use crate::libsignal::protocol::{ + Direction, IdentityChange, IdentityKey, IdentityKeyPair, IdentityKeyStore, KeyPair, + PreKeyBundle, ProtocolAddress, SessionStore, process_prekey_bundle, + }; + use crate::types::message::AddressingMode; + use std::collections::HashMap; + use wacore_binary::NodeContent; + + struct MemSessionStore(HashMap<ProtocolAddress, Vec<u8>>); + impl MemSessionStore { + fn new() -> Self { + Self(HashMap::new()) + } + } + #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] + #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] + impl SessionStore for MemSessionStore { + async fn load_session( + &self, + a: &ProtocolAddress, + ) -> crate::libsignal::protocol::error::Result< + Option<crate::libsignal::protocol::SessionRecord>, + > { + Ok(self + .0 + .get(a) + .and_then(|b| crate::libsignal::protocol::SessionRecord::deserialize(b).ok())) + } + async fn has_session( + &self, + a: &ProtocolAddress, + ) -> crate::libsignal::protocol::error::Result<bool> { + Ok(self.0.contains_key(a)) + } + async fn store_session( + &mut self, + a: &ProtocolAddress, + r: crate::libsignal::protocol::SessionRecord, + ) -> crate::libsignal::protocol::error::Result<()> { + self.0.insert(a.clone(), r.serialize()?); + Ok(()) + } + } + + struct MemIdentityStore { + pair: IdentityKeyPair, + reg_id: u32, + known: HashMap<ProtocolAddress, IdentityKey>, + } + #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] + #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] + impl IdentityKeyStore for MemIdentityStore { + async fn get_identity_key_pair( + &self, + ) -> crate::libsignal::protocol::error::Result<IdentityKeyPair> { + Ok(self.pair.clone()) + } + async fn get_local_registration_id( + &self, + ) -> crate::libsignal::protocol::error::Result<u32> { + Ok(self.reg_id) + } + async fn save_identity( + &mut self, + a: &ProtocolAddress, + id: &IdentityKey, + ) -> crate::libsignal::protocol::error::Result<IdentityChange> { + self.known.insert(a.clone(), *id); + Ok(IdentityChange::from_changed(false)) + } + async fn is_trusted_identity( + &self, + _: &ProtocolAddress, + _: &IdentityKey, + _: Direction, + ) -> crate::libsignal::protocol::error::Result<bool> { + Ok(true) + } + async fn get_identity( + &self, + a: &ProtocolAddress, + ) -> crate::libsignal::protocol::error::Result<Option<IdentityKey>> { + Ok(self.known.get(a).copied()) + } + } + + async fn setup_session() -> (MemSessionStore, MemIdentityStore, Jid) { + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + let sender = IdentityKeyPair::generate(&mut rng); + let receiver = IdentityKeyPair::generate(&mut rng); + let spk = KeyPair::generate(&mut rng); + let opk = KeyPair::generate(&mut rng); + let sig = receiver + .private_key() + .calculate_signature(&spk.public_key.serialize(), &mut rng) + .unwrap(); + let bundle = PreKeyBundle::new( + 1, + 1u32.into(), + Some((1u32.into(), opk.public_key)), + 1u32.into(), + spk.public_key, + sig.to_vec(), + *receiver.identity_key(), + ) + .unwrap(); + let jid: Jid = "559911112222@s.whatsapp.net".parse().unwrap(); + let addr = jid.to_protocol_address(); + let mut ss = MemSessionStore::new(); + let mut is = MemIdentityStore { + pair: sender, + reg_id: 42, + known: HashMap::new(), + }; + process_prekey_bundle( + &addr, + &mut ss, + &mut is, + &bundle, + &mut rand::make_rng::<rand::rngs::StdRng>(), + crate::libsignal::protocol::UsePQRatchet::No, + ) + .await + .unwrap(); + (ss, is, jid) + } + + #[tokio::test] + async fn group_retry_pkmsg_with_account_emits_device_identity() { + let (mut ss, mut is, jid) = setup_session().await; + 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( + &mut ss, + &mut is, + group.clone(), + p.clone(), + p.clone(), + &wa::Message::default(), + "3EB0ABC".into(), + 1, + Some(&account), + AddressingMode::Pn, + None, + ) + .await + .unwrap(); + + assert_eq!(n.tag, "message"); + let mut a = n.attrs(); + assert_eq!(a.optional_string("to").unwrap().as_ref(), group.to_string()); + assert_eq!( + a.optional_string("participant").unwrap().as_ref(), + p.to_string() + ); + // Default (empty) message falls through to "media" per WA Web's typeAttributeFromProtobuf + assert_eq!( + a.optional_string("type").unwrap().as_ref(), + stanza::MSG_TYPE_MEDIA + ); + assert!(a.optional_string("category").is_none()); + assert_eq!(a.optional_string("addressing_mode").unwrap().as_ref(), "pn"); + let enc = n.get_optional_child("enc").unwrap(); + let mut ea = enc.attrs(); + assert_eq!( + ea.optional_string("v").unwrap().as_ref(), + stanza::ENC_VERSION + ); + assert_eq!( + ea.optional_string("type").unwrap().as_ref(), + stanza::ENC_TYPE_PKMSG + ); + assert_eq!(ea.optional_string("count").unwrap().as_ref(), "1"); + assert!(matches!(&enc.content, Some(NodeContent::Bytes(_)))); + assert!( + n.get_optional_child("device-identity").is_some(), + "pkmsg group retry with account must include <device-identity>" + ); + } + + /// Symmetric to peer/dm pre-flights: refuse group retry pkmsg when + /// account is missing rather than silently dropping device-identity. + #[tokio::test] + async fn group_retry_pkmsg_preflight_errors_when_account_missing() { + let (mut ss, mut is, jid) = setup_session().await; + let group: Jid = "120363098765432100@g.us".parse().unwrap(); + let p: Jid = jid.to_string().parse().unwrap(); + + let before = ss + .load_session(&p.to_protocol_address()) + .await + .unwrap() + .expect("pre-condition: session present") + .serialize() + .expect("serialize before"); + + let result = prepare_group_retry_stanza( + &mut ss, + &mut is, + group, + p.clone(), + p.clone(), + &wa::Message::default(), + "grp-retry-no-account".into(), + 1, + None, + AddressingMode::Pn, + None, + ) + .await; + let err = result.expect_err("group retry pkmsg must reject missing account"); + assert!( + err.to_string().contains("device-identity"), + "error must name <device-identity>; got: {err}" + ); + + let after = ss + .load_session(&p.to_protocol_address()) + .await + .unwrap() + .expect("session still present") + .serialize() + .expect("serialize after"); + assert_eq!( + before, after, + "group retry pre-flight must leave the session byte-identical" + ); + } + + /// Pins the WAWebSendMsgCreateDeviceStanza retry shape: `<enc>` + /// directly under `<message>` plus a `recipient` attribute. + /// Pre-fix this regressed to the fanout shape and the server + /// rejected every retry with 479. + #[tokio::test] + async fn dm_retry_emits_enc_directly_under_message_with_recipient() { + let (mut ss, mut is, jid) = setup_session().await; + // Distinct values so a swapped-args regression (e.g. `recipient = + // to_jid`) fails the assertions below instead of silently passing. + let to: Jid = "559922223333:5@s.whatsapp.net".parse().unwrap(); + 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( + &mut ss, + &mut is, + to.clone(), + Some(recipient.clone()), + requester, + &wa::Message::default(), + "dm-retry-format-1".into(), + 1, + Some(&account), + None, + ) + .await + .unwrap(); + + assert_eq!(n.tag, "message"); + // <enc> is a direct child — no <participants> wrapper. + assert!( + n.get_optional_child("participants").is_none(), + "DM retry must not wrap <enc> in <participants> \ + (matches WAWebSendMsgCreateDeviceStanza)" + ); + assert!( + n.get_optional_child("enc").is_some(), + "<enc> must be a direct child of <message>" + ); + assert_eq!( + n.attrs().optional_string("to").unwrap().as_ref(), + to.to_string(), + "`to` should target the requesting device verbatim" + ); + assert_eq!( + n.attrs().optional_string("recipient").unwrap().as_ref(), + recipient.to_string(), + "`recipient` should mirror the original message's recipient \ + (forwarded from the retry receipt's `recipient` attr)" + ); + } + + #[tokio::test] + async fn dm_retry_pkmsg_targets_single_device() { + let (mut ss, mut is, jid) = setup_session().await; + let to: Jid = "559922223333@s.whatsapp.net".parse().unwrap(); + let encryption = jid.clone(); + let account = pkmsg_account_proto(); + + let n = prepare_dm_retry_stanza( + &mut ss, + &mut is, + to.clone(), + Some(to.clone()), + encryption, + &wa::Message::default(), + "dm-retry-1".into(), + 1, + Some(&account), + None, + ) + .await + .unwrap(); + + assert_eq!(n.tag, "message"); + let mut attrs = n.attrs(); + assert_eq!( + attrs.optional_string("to").unwrap().as_ref(), + to.to_string() + ); + assert_eq!( + attrs.optional_string("recipient").unwrap().as_ref(), + to.to_string() + ); + assert_eq!(attrs.optional_string("id").unwrap().as_ref(), "dm-retry-1"); + assert_eq!( + attrs.optional_string("type").unwrap().as_ref(), + stanza::MSG_TYPE_MEDIA + ); + assert!(attrs.optional_string("participant").is_none()); + assert!(attrs.optional_string("addressing_mode").is_none()); + + // `<enc>` is a direct child of `<message>` (no `<participants>` wrapper). + assert!(n.get_optional_child("participants").is_none()); + let enc = n.get_optional_child("enc").unwrap(); + let mut enc_attrs = enc.attrs(); + assert_eq!( + enc_attrs.optional_string("type").unwrap().as_ref(), + stanza::ENC_TYPE_PKMSG + ); + assert_eq!(enc_attrs.optional_string("count").unwrap().as_ref(), "1"); + assert!( + n.get_optional_child("device-identity").is_some(), + "pkmsg DM retry with account must include <device-identity>" + ); + } + + #[tokio::test] + async fn dm_retry_pkmsg_with_account_has_device_identity() { + let (mut ss, mut is, jid) = setup_session().await; + let to: Jid = "559922223333@s.whatsapp.net".parse().unwrap(); + let acc = wa::AdvSignedDeviceIdentity { + details: Some(b"t".to_vec()), + ..Default::default() + }; + + let n = prepare_dm_retry_stanza( + &mut ss, + &mut is, + to.clone(), + Some(to), + jid, + &wa::Message::default(), + "dm-retry-2".into(), + 2, + Some(&acc), + None, + ) + .await + .unwrap(); + + let enc = n.get_optional_child("enc").unwrap(); + assert_eq!( + enc.attrs().optional_string("type").unwrap().as_ref(), + stanza::ENC_TYPE_PKMSG + ); + assert_eq!(enc.attrs().optional_string("count").unwrap().as_ref(), "2"); + assert!(n.get_optional_child("device-identity").is_some()); + } + + #[tokio::test] + async fn pkmsg_with_account_has_device_identity() { + let (mut ss, mut is, jid) = setup_session().await; + let group: Jid = "120363098765432100@g.us".parse().unwrap(); + let p: Jid = jid.to_string().parse().unwrap(); + let acc = wa::AdvSignedDeviceIdentity { + details: Some(b"t".to_vec()), + ..Default::default() + }; + let n = prepare_group_retry_stanza( + &mut ss, + &mut is, + group, + p.clone(), + p, + &wa::Message::default(), + "id2".into(), + 2, + Some(&acc), + AddressingMode::Pn, + None, + ) + .await + .unwrap(); + assert_eq!( + n.get_optional_child("enc") + .unwrap() + .attrs() + .optional_string("type") + .unwrap() + .as_ref(), + stanza::ENC_TYPE_PKMSG + ); + assert!(n.get_optional_child("device-identity").is_some()); + assert_eq!( + n.attrs() + .optional_string("addressing_mode") + .unwrap() + .as_ref(), + "pn" + ); + } + + #[tokio::test] + async fn lid_addressing_mode() { + let (mut ss, mut is, jid) = setup_session().await; + 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( + &mut ss, + &mut is, + group, + p.clone(), + p, + &wa::Message::default(), + "m2".into(), + 3, + Some(&wa::AdvSignedDeviceIdentity::default()), + AddressingMode::Lid, + None, + ) + .await + .unwrap(); + let mut ea = n.get_optional_child("enc").unwrap().attrs(); + assert_eq!(ea.optional_string("count").unwrap().as_ref(), "3"); + assert_eq!( + n.attrs() + .optional_string("addressing_mode") + .unwrap() + .as_ref(), + "lid" + ); + } + + #[tokio::test] + async fn group_retry_preserves_edit_attribute() { + let (mut ss, mut is, jid) = setup_session().await; + 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( + &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), + ) + .await + .unwrap(); + assert_eq!(n.attrs().optional_string("edit").unwrap().as_ref(), "8"); + } + + #[tokio::test] + async fn dm_retry_preserves_edit_attribute() { + 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( + &mut ss, + &mut is, + to.clone(), + Some(to), + jid, + &wa::Message::default(), + "edit-1".into(), + 1, + Some(&account), + Some(crate::types::message::EditAttribute::MessageEdit), + ) + .await + .unwrap(); + assert_eq!(n.attrs().optional_string("edit").unwrap().as_ref(), "1"); + } + + #[tokio::test] + async fn retry_without_edit_omits_attribute() { + let (mut ss, mut is, jid) = setup_session().await; + 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( + &mut ss, + &mut is, + group, + p.clone(), + p, + &wa::Message::default(), + "plain-1".into(), + 1, + Some(&account), + AddressingMode::Lid, + None, + ) + .await + .unwrap(); + assert!(n.attrs().optional_string("edit").is_none()); + } + + // Peer pkmsg layout: `[<meta appdata="default"/>, <enc>, <device-identity>]`. + // Without `<device-identity>` the phone XMPP-acks but its Signal + // layer skips session promotion. Mirrors whatsmeow's + // `preparePeerMessageNode`. + + fn pkmsg_account_proto() -> wa::AdvSignedDeviceIdentity { + // Opaque placeholder bytes — the assertions only check that + // the element carries non-empty content. + wa::AdvSignedDeviceIdentity { + details: Some(vec![0u8; 32]), + account_signature_key: Some(vec![0u8; 32]), + account_signature: Some(vec![0u8; 64]), + device_signature: Some(vec![0u8; 64]), + } + } + + async fn build_peer_stanza( + account: Option<&wa::AdvSignedDeviceIdentity>, + ) -> wacore_binary::Node { + build_peer_stanza_with_options(account, PeerMessageOptions::default()).await + } + + async fn build_peer_stanza_with_options( + account: Option<&wa::AdvSignedDeviceIdentity>, + options: PeerMessageOptions, + ) -> wacore_binary::Node { + let (mut ss, mut is, jid) = setup_session().await; + let addr = jid.to_protocol_address(); + prepare_peer_stanza_with_options( + &mut ss, + &mut is, + jid.clone(), + &addr, + &wa::Message::default(), + "peer-test-1".into(), + account, + options, + ) + .await + .expect("peer stanza builds") + } + + #[tokio::test] + async fn peer_pkmsg_includes_meta_and_device_identity() { + let account = pkmsg_account_proto(); + let n = build_peer_stanza(Some(&account)).await; + + assert_eq!(n.tag, "message"); + assert_eq!( + n.attrs().optional_string("category").unwrap().as_ref(), + "peer" + ); + assert_eq!( + n.attrs().optional_string("push_priority").unwrap().as_ref(), + "high" + ); + assert!(n.attrs().optional_string("privacy_sensitive").is_none()); + + let children = n.children().expect("peer message has children"); + let tags: Vec<&str> = children.iter().map(|c| c.tag.as_ref()).collect(); + // Layout matches whatsmeow's preparePeerMessageNode for pkmsg: + // [<meta>, <enc>, <device-identity>]. + assert_eq!( + tags, + vec!["meta", "enc", "device-identity"], + "peer pkmsg children order/identity must match whatsmeow" + ); + + let meta = n.get_optional_child("meta").expect("meta present"); + assert_eq!( + meta.attrs().optional_string("appdata").unwrap().as_ref(), + "default", + "<meta appdata=\"default\"/> is what the phone uses to route the peer payload" + ); + + let enc = n.get_optional_child("enc").expect("enc present"); + assert_eq!( + enc.attrs().optional_string("type").unwrap().as_ref(), + "pkmsg", + "fresh session must produce pkmsg, not msg" + ); + + let device_identity = n + .get_optional_child("device-identity") + .expect("device-identity present"); + match &device_identity.content { + Some(NodeContent::Bytes(b)) => assert!( + !b.is_empty(), + "device-identity content must be the proto-encoded \ + AdvSignedDeviceIdentity, not empty" + ), + other => panic!("device-identity must carry bytes, got {other:?}"), + } + } + + #[tokio::test] + async fn peer_stanza_carries_high_force_and_privacy_attrs() { + let account = pkmsg_account_proto(); + let n = build_peer_stanza_with_options( + Some(&account), + PeerMessageOptions::high_force_on_demand(), + ) + .await; + + assert_eq!( + n.attrs().optional_string("push_priority").unwrap().as_ref(), + "high_force" + ); + assert_eq!( + n.attrs() + .optional_string("privacy_sensitive") + .unwrap() + .as_ref(), + "1" + ); + } + + #[tokio::test] + async fn peer_pkmsg_errors_when_account_missing_without_ratchet_advance() { + // Pkmsg without <device-identity> would reproduce the deadlock — + // refuse AND prove the session is byte-identical after the failed + // call so the next retry has the same ratchet position. + let (mut ss, mut is, jid) = setup_session().await; + let addr = jid.to_protocol_address(); + + let before = ss + .load_session(&addr) + .await + .unwrap() + .expect("pre-condition: session loaded") + .serialize() + .expect("serialize before"); + + let result = prepare_peer_stanza( + &mut ss, + &mut is, + jid.clone(), + &addr, + &wa::Message::default(), + "peer-test-no-account".into(), + None, + ) + .await; + let err = result.expect_err("pkmsg path must reject missing account"); + assert!( + err.to_string().contains("device-identity"), + "error must name the missing element; got: {err}" + ); + + let after = ss + .load_session(&addr) + .await + .unwrap() + .expect("session still present after failed call") + .serialize() + .expect("serialize after"); + assert_eq!( + before, after, + "session record must be byte-identical after a failed prepare — \ + any difference means a ratchet step was committed for a stanza we couldn't ship" + ); + } + + /// Pre-flight check: when no session exists and account is None, + /// `prepare_peer_stanza` must refuse before `message_encrypt` runs, + /// otherwise the sender chain is persisted for a stanza we cannot ship + /// (CodeRabbit-flagged ratchet-burn-on-fail-fast). + #[tokio::test] + async fn peer_pkmsg_preflight_no_ratchet_burn_without_session() { + let jid: Jid = "559911112222@s.whatsapp.net".parse().unwrap(); + let addr = jid.to_protocol_address(); + let mut ss = MemSessionStore::new(); + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + let mut is = MemIdentityStore { + pair: IdentityKeyPair::generate(&mut rng), + reg_id: 42, + known: HashMap::new(), + }; + + assert!( + !ss.has_session(&addr).await.unwrap(), + "precondition: store has no session for this address" + ); + + let result = prepare_peer_stanza( + &mut ss, + &mut is, + jid.clone(), + &addr, + &wa::Message::default(), + "peer-preflight-1".into(), + None, + ) + .await; + let err = result.expect_err("must refuse before message_encrypt"); + assert!( + err.to_string().contains("device-identity"), + "error must name <device-identity>; got: {err}" + ); + assert!( + !ss.has_session(&addr).await.unwrap(), + "pre-flight must NOT advance/persist a session — the ratchet \ + must remain unburned for the retry attempt" + ); + } + + /// Symmetric to peer_pkmsg_preflight: prepare_dm_retry_stanza must + /// also refuse to ship pkmsg without <device-identity>, otherwise + /// message_encrypt would advance the sender chain for a stanza the + /// peer's Signal layer cannot promote. + #[tokio::test] + async fn dm_retry_pkmsg_preflight_errors_when_account_missing() { + let (mut ss, mut is, jid) = setup_session().await; + let addr = jid.to_protocol_address(); + + let before = ss + .load_session(&addr) + .await + .unwrap() + .expect("pre-condition: session present") + .serialize() + .expect("serialize before"); + + let to: Jid = "559922223333@s.whatsapp.net".parse().unwrap(); + let result = prepare_dm_retry_stanza( + &mut ss, + &mut is, + to.clone(), + Some(to), + jid.clone(), + &wa::Message::default(), + "dm-retry-no-account".into(), + 1, + None, + None, + ) + .await; + let err = result.expect_err("DM retry pkmsg path must reject missing account"); + assert!( + err.to_string().contains("device-identity"), + "error must name <device-identity>; got: {err}" + ); + + let after = ss + .load_session(&addr) + .await + .unwrap() + .expect("session still present") + .serialize() + .expect("serialize after"); + assert_eq!( + before, after, + "DM retry pre-flight must leave the session byte-identical" + ); + } + + /// Production's SessionAdapter::load_session has TAKE semantics + /// (SignalStoreCache marks the slot CheckedOut until store_session + /// puts the record back). If the pre-flight only loads without + /// restoring, the slot stays stranded and message_encrypt sees no + /// session. The mock here mirrors that contract via interior + /// mutability (Mutex) on the &self load_session. + #[tokio::test] + async fn preflight_restores_session_with_take_store_semantics() { + use std::collections::{HashMap, HashSet}; + use std::sync::Mutex; + + struct TakeStore { + inner: Mutex<TakeInner>, + } + struct TakeInner { + present: HashMap<ProtocolAddress, Vec<u8>>, + taken: HashSet<ProtocolAddress>, + } + impl TakeStore { + fn from(ss: &MemSessionStore) -> Self { + Self { + inner: Mutex::new(TakeInner { + present: ss.0.clone(), + taken: HashSet::new(), + }), + } + } + fn is_present(&self, addr: &ProtocolAddress) -> bool { + let g = self.inner.lock().unwrap(); + g.present.contains_key(addr) && !g.taken.contains(addr) + } + } + #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] + #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] + impl SessionStore for TakeStore { + async fn load_session( + &self, + a: &ProtocolAddress, + ) -> crate::libsignal::protocol::error::Result< + Option<crate::libsignal::protocol::SessionRecord>, + > { + let mut g = self.inner.lock().unwrap(); + if g.taken.contains(a) { + return Ok(None); + } + let rec = g + .present + .get(a) + .and_then(|b| crate::libsignal::protocol::SessionRecord::deserialize(b).ok()); + if rec.is_some() { + g.taken.insert(a.clone()); + } + Ok(rec) + } + async fn has_session( + &self, + a: &ProtocolAddress, + ) -> crate::libsignal::protocol::error::Result<bool> { + let g = self.inner.lock().unwrap(); + Ok(g.present.contains_key(a) && !g.taken.contains(a)) + } + async fn store_session( + &mut self, + a: &ProtocolAddress, + r: crate::libsignal::protocol::SessionRecord, + ) -> crate::libsignal::protocol::error::Result<()> { + let mut g = self.inner.lock().unwrap(); + g.present.insert(a.clone(), r.serialize()?); + g.taken.remove(a); + Ok(()) + } + } + + let (mem_ss, mut is, jid) = setup_session().await; + let mut ss = TakeStore::from(&mem_ss); + let addr = jid.to_protocol_address(); + + // setup_session leaves pending_pre_key set, so account=None + // would bail. Use Some(account) — pre-flight still runs + // load+restore because it's gated on account.is_none() at the + // call site; switch to account=None and we want the assertion + // to verify that the BAIL path also restores the slot. + assert!( + ss.is_present(&addr), + "precondition: session is Present before pre-flight" + ); + + // Drive the bail path: account=None + session has pending_pre_key + // → pre-flight bails. Even on bail, the loaded record must be + // put back so a retry with Some(account) doesn't see a stranded slot. + let bail = prepare_peer_stanza( + &mut ss, + &mut is, + jid.clone(), + &addr, + &wa::Message::default(), + "preflight-take-bail".into(), + None, + ) + .await; + bail.expect_err("must bail with account=None on a pending-pkmsg session"); + assert!( + ss.is_present(&addr), + "pre-flight bail path must still restore the checked-out session" + ); + + // And the pass path: with Some(account), the pre-flight still + // does load+restore, then message_encrypt runs successfully. + let account = pkmsg_account_proto(); + let ok = prepare_peer_stanza( + &mut ss, + &mut is, + jid.clone(), + &addr, + &wa::Message::default(), + "preflight-take-pass".into(), + Some(&account), + ) + .await; + ok.expect("peer stanza builds with Some(account)"); + assert!( + ss.is_present(&addr), + "session must be Present after a successful encrypt+store" + ); + } +} + +mod decrypt_fail { + use super::*; + + #[test] + fn regular_message() { + let msg = wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }; + assert!(!should_hide_decrypt_fail(&msg)); + } + + #[test] + fn reaction() { + let msg = wa::Message { + reaction_message: Some(Default::default()), + ..Default::default() + }; + assert!(should_hide_decrypt_fail(&msg)); + } + + #[test] + fn pin() { + let msg = wa::Message { + pin_in_chat_message: Some(Default::default()), + ..Default::default() + }; + assert!(should_hide_decrypt_fail(&msg)); + } + + #[test] + fn poll_vote() { + let msg = wa::Message { + poll_update_message: Some(wa::message::PollUpdateMessage { + vote: Some(Default::default()), + ..Default::default() + }), + ..Default::default() + }; + assert!(should_hide_decrypt_fail(&msg)); + } + + #[test] + fn poll_update_without_vote() { + let msg = wa::Message { + poll_update_message: Some(Default::default()), + ..Default::default() + }; + assert!(!should_hide_decrypt_fail(&msg)); + } + + #[test] + fn reaction_inside_ephemeral_wrapper() { + let msg = wa::Message { + ephemeral_message: Some(Box::new(wa::message::FutureProofMessage { + message: Some(Box::new(wa::Message { + reaction_message: Some(Default::default()), + ..Default::default() + })), + })), + ..Default::default() + }; + assert!(should_hide_decrypt_fail(&msg)); + } + + #[test] + fn conditional_reveal() { + let msg = wa::Message { + conditional_reveal_message: Some(Default::default()), + ..Default::default() + }; + assert!(should_hide_decrypt_fail(&msg)); + } + + #[test] + fn poll_add_option_edit() { + use wa::message::secret_encrypted_message::SecretEncType; + let msg = wa::Message { + secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { + secret_enc_type: Some(SecretEncType::PollAddOption as i32), + ..Default::default() + }), + ..Default::default() + }; + assert!(should_hide_decrypt_fail(&msg)); + } +} + +mod decrypt_fail_for_send { + use super::*; + use crate::types::message::EditAttribute; + + fn plain() -> wa::Message { + wa::Message { + conversation: Some("hi".into()), + ..Default::default() + } + } + + #[test] + fn sender_revoke_is_not_hidden() { + assert!(!should_hide_decrypt_fail_for_send( + Some(&EditAttribute::SenderRevoke), + &plain() + )); + } + + #[test] + fn admin_revoke_is_not_hidden() { + assert!(!should_hide_decrypt_fail_for_send( + Some(&EditAttribute::AdminRevoke), + &plain() + )); + } + + #[test] + fn message_edit_is_hidden() { + assert!(should_hide_decrypt_fail_for_send( + Some(&EditAttribute::MessageEdit), + &plain() + )); + } + + #[test] + fn revoke_does_not_block_content_based_hide() { + // A reaction still hides on its own merits even under a revoke edit. + let msg = wa::Message { + reaction_message: Some(Default::default()), + ..Default::default() + }; + assert!(should_hide_decrypt_fail_for_send( + Some(&EditAttribute::SenderRevoke), + &msg + )); + } +} + +mod stanza_type { + use super::*; + use wa::message::secret_encrypted_message::SecretEncType; + + fn secret(enc: SecretEncType) -> wa::Message { + wa::Message { + secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { + secret_enc_type: Some(enc as i32), + ..Default::default() + }), + ..Default::default() + } + } + + #[test] + fn poll_add_option_edit_is_poll() { + assert_eq!( + stanza_type_from_message(&secret(SecretEncType::PollAddOption)), + stanza::MSG_TYPE_POLL + ); + } + + #[test] + fn poll_edit_is_poll() { + assert_eq!( + stanza_type_from_message(&secret(SecretEncType::PollEdit)), + stanza::MSG_TYPE_POLL + ); + } + + #[test] + fn album_is_text() { + let msg = wa::Message { + album_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&msg), stanza::MSG_TYPE_TEXT); + } + + // Helpers for wrapper tests. WA Web's typeAttributeFromProtobuf unwraps + // FutureProofMessage wrappers (via getUnwrappedProtobufMessage) and then + // classifies the inner message. + fn fpm(inner: wa::Message) -> Box<wa::message::FutureProofMessage> { + Box::new(wa::message::FutureProofMessage { + message: Some(Box::new(inner)), + }) + } + fn text_inner() -> wa::Message { + wa::Message { + conversation: Some("hi".to_string()), + ..Default::default() + } + } + fn image_inner() -> wa::Message { + wa::Message { + image_message: Some(Box::default()), + ..Default::default() + } + } + + #[test] + fn group_status_v2_classifies_by_inner() { + let txt = wa::Message { + group_status_message_v2: Some(fpm(text_inner())), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&txt), stanza::MSG_TYPE_TEXT); + + // Regression guard: forcing this wrapper to "text" dropped the + // mediatype and silently dropped the stanza. WA Web unwraps it and + // sends type="media" mediatype="image". + let img = wa::Message { + group_status_message_v2: Some(fpm(image_inner())), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&img), stanza::MSG_TYPE_MEDIA); + assert_eq!(media_type_from_message(&img), Some("image")); + } + + #[test] + fn group_status_v2_empty_is_media() { + // An empty wrapper is not one of WA Web's four re-checked wrappers + // (ephemeral/groupMentioned/botInvoke/deviceSent), so it falls through + // to the media default in both WA Web and here. + let m = wa::Message { + group_status_message_v2: Some(Box::default()), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&m), stanza::MSG_TYPE_MEDIA); + } + + #[test] + fn payment_family_is_text() { + // Payment family classifies as text; the media default would be dropped. + let cases = [ + wa::Message { + request_payment_message: Some(Box::default()), + ..Default::default() + }, + wa::Message { + send_payment_message: Some(Box::default()), + ..Default::default() + }, + wa::Message { + decline_payment_request_message: Some(Default::default()), + ..Default::default() + }, + wa::Message { + cancel_payment_request_message: Some(Default::default()), + ..Default::default() + }, + wa::Message { + payment_invite_message: Some(Default::default()), + ..Default::default() + }, + ]; + for m in cases { + assert_eq!(media_type_from_message(&m), None); + assert_eq!(stanza_type_from_message(&m), stanza::MSG_TYPE_TEXT); + } + } + + #[test] + fn backfilled_wrappers_classify_by_inner() { + let spoiler = wa::Message { + spoiler_message: Some(fpm(text_inner())), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&spoiler), stanza::MSG_TYPE_TEXT); + + let status_mention = wa::Message { + status_mention_message: Some(fpm(image_inner())), + ..Default::default() + }; + assert_eq!( + stanza_type_from_message(&status_mention), + stanza::MSG_TYPE_MEDIA + ); + assert_eq!(media_type_from_message(&status_mention), Some("image")); + + let question = wa::Message { + question_message: Some(fpm(text_inner())), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&question), stanza::MSG_TYPE_TEXT); + + let group_status_v1 = wa::Message { + group_status_message: Some(fpm(text_inner())), + ..Default::default() + }; + assert_eq!( + stanza_type_from_message(&group_status_v1), + stanza::MSG_TYPE_TEXT + ); + } + + #[test] + fn nested_wrappers_reach_innermost() { + // ephemeral { viewOnceV2 { image } } -> media + mediatype. + let inner = wa::Message { + view_once_message_v2: Some(fpm(image_inner())), + ..Default::default() + }; + let m = wa::Message { + ephemeral_message: Some(fpm(inner)), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&m), stanza::MSG_TYPE_MEDIA); + assert_eq!(media_type_from_message(&m), Some("image")); + } + + #[test] + fn preserved_classifier_branches() { + let r = wa::Message { + reaction_message: Some(Default::default()), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&r), stanza::MSG_TYPE_REACTION); + + let ev = wa::Message { + event_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&ev), stanza::MSG_TYPE_EVENT); + + let poll = wa::Message { + poll_creation_message_v3: Some(Box::default()), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&poll), stanza::MSG_TYPE_POLL); + + assert_eq!( + stanza_type_from_message(&text_inner()), + stanza::MSG_TYPE_TEXT + ); + assert_eq!( + stanza_type_from_message(&image_inner()), + stanza::MSG_TYPE_MEDIA + ); + + let proto = wa::Message { + protocol_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&proto), stanza::MSG_TYPE_TEXT); + + let url = wa::Message { + extended_text_message: Some(Box::new(wa::message::ExtendedTextMessage { + matched_text: Some("https://example.com".to_string()), + ..Default::default() + })), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&url), stanza::MSG_TYPE_MEDIA); + } + + #[test] + fn interactive_and_list_types_get_their_mediatype() { + // WA Web's mediaTypeFromProtobuf maps these to concrete mediatypes; + // omitting the attribute makes the server drop the type="media" stanza. + let list = wa::Message { + list_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!(stanza_type_from_message(&list), stanza::MSG_TYPE_MEDIA); + assert_eq!(media_type_from_message(&list), Some("list")); + + let list_response = wa::Message { + list_response_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!( + media_type_from_message(&list_response), + Some("list_response") + ); + + let buttons_response = wa::Message { + buttons_response_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!( + media_type_from_message(&buttons_response), + Some("buttons_response") + ); + + let order = wa::Message { + order_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!(media_type_from_message(&order), Some("order")); + + let product = wa::Message { + product_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!(media_type_from_message(&product), Some("product")); + + let interactive_response = wa::Message { + interactive_response_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!( + media_type_from_message(&interactive_response), + Some("native_flow_response") + ); + + let history_bundle = wa::Message { + message_history_bundle: Some(Box::default()), + ..Default::default() + }; + assert_eq!( + media_type_from_message(&history_bundle), + Some("group_history") + ); + } + + #[test] + fn buttons_message_has_no_mediatype() { + // WA Web maps buttonsMessage to EncMediaType.Button, but its string + // mapper has no Button case (returns null/DROP_ATTR), so the attribute + // is omitted. Adding a "buttons" mediatype would diverge from WA Web. + let buttons = wa::Message { + buttons_message: Some(Box::default()), + ..Default::default() + }; + assert_eq!(media_type_from_message(&buttons), None); + } + + #[test] + fn ephemeral_wrapped_list_reaches_list_mediatype() { + let m = wa::Message { + ephemeral_message: Some(fpm(wa::Message { + list_message: Some(Box::default()), + ..Default::default() + })), + ..Default::default() + }; + assert_eq!(media_type_from_message(&m), Some("list")); + } + + #[test] + fn top_level_lottie_sticker_is_terminal_sticker() { + // WA Web's mediaTypeFromProtobuf treats a top-level lottieStickerMessage + // as a terminal "sticker" and does NOT recurse into it, unlike the + // stanza-type path which unwraps it. + let lottie = wa::Message { + lottie_sticker_message: Some(fpm(image_inner())), + ..Default::default() + }; + assert_eq!(media_type_from_message(&lottie), Some("sticker")); + } +} + +#[cfg(test)] +mod device_unregistered_tests { + use super::is_device_unregistered_error; + use crate::request::ServerErrorCode; + + #[test] + fn detects_406_server_error_code() { + let err = anyhow::Error::new(ServerErrorCode { + code: 406, + text: "not-acceptable".to_string(), + }); + assert!(is_device_unregistered_error(&err)); + } + + #[test] + fn rejects_non_406_server_error() { + let err = anyhow::Error::new(ServerErrorCode { + code: 404, + text: "not-found".to_string(), + }); + assert!(!is_device_unregistered_error(&err)); + } + + #[test] + fn rejects_unrelated_error() { + let err = anyhow::anyhow!("some random error"); + assert!(!is_device_unregistered_error(&err)); + } + + #[test] + fn rejects_wacore_iq_error_without_server_error_code_wrapper() { + // wacore::IqError::ServerError is NOT the same as ServerErrorCode. + // This simulates the old bug: if someone wraps wacore IqError directly + // without the ServerErrorCode wrapper, the check should not match. + let err = anyhow::Error::new(crate::request::IqError::ServerError { + code: 406, + text: "not-acceptable".to_string(), + }); + // This would only match if we also checked IqError (we don't — we use ServerErrorCode) + // The SendContextResolver impl is responsible for wrapping in ServerErrorCode + assert!(!is_device_unregistered_error(&err)); + } +} + +mod collect_stale_device_users { + use super::super::collect_stale_device_users; + use crate::client::context::GroupInfo; + use crate::types::message::AddressingMode; + use std::collections::{HashMap, HashSet}; + use wacore_binary::{CompactString, Jid}; + + fn lid_device(user: &str, dev: u16) -> Jid { + Jid::lid_device(user.to_string(), dev) + } + + fn pn_user(user: &str) -> Jid { + Jid::pn(user) + } + + fn group_info_lid(mapping: &[(&str, &str)]) -> GroupInfo { + let mut info = GroupInfo::new(Vec::new(), AddressingMode::Lid); + if !mapping.is_empty() { + let mut map: HashMap<CompactString, Jid> = HashMap::new(); + for (lid_user, pn) in mapping { + map.insert(CompactString::from(*lid_user), pn_user(pn)); + } + info.set_lid_to_pn_map(map); + } + info + } + + #[test] + fn emits_lid_and_pn_alias_when_mapping_known() { + let info = group_info_lid(&[("100000000000001", "15550000001")]); + let dist = vec![lid_device("100000000000001", 5)]; + let out = collect_stale_device_users(Some(&dist), &[], &info); + let set: HashSet<String> = out.into_iter().collect(); + assert!(set.contains("100000000000001")); + assert!(set.contains("15550000001")); + assert_eq!(set.len(), 2); + } + + #[test] + fn emits_only_lid_when_mapping_unknown() { + let info = group_info_lid(&[]); + let dist = vec![lid_device("100000000000002", 7)]; + let out = collect_stale_device_users(Some(&dist), &[], &info); + assert_eq!(out, vec!["100000000000002".to_string()]); + } + + #[test] + fn dedups_multiple_devices_of_same_user() { + let info = group_info_lid(&[("100000000000003", "15550000003")]); + let dist = vec![ + lid_device("100000000000003", 1), + lid_device("100000000000003", 2), + lid_device("100000000000003", 3), + ]; + let out = collect_stale_device_users(Some(&dist), &[], &info); + let set: HashSet<String> = out.into_iter().collect(); + assert_eq!(set.len(), 2); + assert!(set.contains("100000000000003")); + assert!(set.contains("15550000003")); + } + + #[test] + fn skips_successfully_encrypted_devices() { + let info = group_info_lid(&[]); + let encrypted = lid_device("100000000000004", 5); + let dist = vec![encrypted.clone(), lid_device("100000000000005", 5)]; + let encrypted_set = vec![encrypted]; + let out = collect_stale_device_users(Some(&dist), &encrypted_set, &info); + assert_eq!(out, vec!["100000000000005".to_string()]); + } + + #[test] + fn pn_mode_group_does_not_emit_alias() { + // In PN-mode groups the distribution list is already PN-form, so + // there's no LID↔PN duality to emit. + let mut info = GroupInfo::new(Vec::new(), AddressingMode::Pn); + let mut map: HashMap<CompactString, Jid> = HashMap::new(); + map.insert( + CompactString::from("100000000000006"), + pn_user("15550000006"), + ); + info.set_lid_to_pn_map(map); + let dist = vec![Jid::pn_device("15550000006", 3)]; + let out = collect_stale_device_users(Some(&dist), &[], &info); + assert_eq!(out, vec!["15550000006".to_string()]); + } + + #[test] + fn skips_non_pn_alias() { + // If phone_jid_for_lid_user returns a JID whose server isn't PN + // (malformed/adversarial server response), do not emit it. + let mut info = GroupInfo::new(Vec::new(), AddressingMode::Lid); + let mut map: HashMap<CompactString, Jid> = HashMap::new(); + map.insert( + CompactString::from("100000000000007"), + Jid::lid("100000000000099"), + ); + info.set_lid_to_pn_map(map); + let dist = vec![lid_device("100000000000007", 5)]; + let out = collect_stale_device_users(Some(&dist), &[], &info); + assert_eq!(out, vec!["100000000000007".to_string()]); + } + + #[test] + fn empty_distribution_list_yields_empty() { + let info = group_info_lid(&[]); + let out = collect_stale_device_users(None, &[], &info); + assert!(out.is_empty()); + let out = collect_stale_device_users(Some(&[]), &[], &info); + assert!(out.is_empty()); + } +} + +/// Item 2 — WA Web `markHasSenderKey(x, M)`: a key-distributing group send +/// marks the FULL SKDM target set `has_key=true`, not only the devices that +/// encrypted successfully. A device whose SKDM encryption fails (no session +/// and no bundle, mimicking a 406) must still land in +/// `PreparedGroupStanza.skdm_devices`, so the next send does not re-target +/// it every time (the fan-out storm); the retry-receipt path repairs any +/// device that is actually alive and keyless. +mod mark_full_distribution_list { + use super::*; + use crate::libsignal::protocol::{ + Direction, IdentityChange, IdentityKey, IdentityKeyStore, PreKeyId, PreKeyRecord, + PreKeyStore, ProtocolAddress, SenderKeyRecord, SenderKeyStore, SessionStore, + SignedPreKeyId, SignedPreKeyRecord, SignedPreKeyStore, UsePQRatchet, process_prekey_bundle, + }; + use crate::libsignal::store::sender_key_name::SenderKeyName; + use crate::runtime::{AbortHandle, Runtime}; + use crate::types::jid::JidExt; + use crate::types::message::AddressingMode; + use std::future::Future; + use std::pin::Pin; + use std::time::Duration; + + type SigResult<T> = crate::libsignal::protocol::error::Result<T>; + + #[derive(Clone, Default)] + struct MemSessionStore(HashMap<ProtocolAddress, Vec<u8>>); + #[async_trait::async_trait] + impl SessionStore for MemSessionStore { + async fn load_session( + &self, + a: &ProtocolAddress, + ) -> SigResult<Option<crate::libsignal::protocol::SessionRecord>> { + Ok(self + .0 + .get(a) + .and_then(|b| crate::libsignal::protocol::SessionRecord::deserialize(b).ok())) + } + async fn has_session(&self, a: &ProtocolAddress) -> SigResult<bool> { + Ok(self.0.contains_key(a)) + } + async fn store_session( + &mut self, + a: &ProtocolAddress, + r: crate::libsignal::protocol::SessionRecord, + ) -> SigResult<()> { + self.0.insert(a.clone(), r.serialize()?); + Ok(()) + } + } + + #[derive(Clone)] + struct MemIdentityStore { + pair: IdentityKeyPair, + reg_id: u32, + known: HashMap<ProtocolAddress, IdentityKey>, + } + #[async_trait::async_trait] + impl IdentityKeyStore for MemIdentityStore { + async fn get_identity_key_pair(&self) -> SigResult<IdentityKeyPair> { + Ok(self.pair.clone()) + } + async fn get_local_registration_id(&self) -> SigResult<u32> { + Ok(self.reg_id) + } + async fn save_identity( + &mut self, + a: &ProtocolAddress, + id: &IdentityKey, + ) -> SigResult<IdentityChange> { + self.known.insert(a.clone(), *id); + Ok(IdentityChange::from_changed(false)) + } + async fn is_trusted_identity( + &self, + _: &ProtocolAddress, + _: &IdentityKey, + _: Direction, + ) -> SigResult<bool> { + Ok(true) + } + async fn get_identity(&self, a: &ProtocolAddress) -> SigResult<Option<IdentityKey>> { + Ok(self.known.get(a).copied()) + } + } + + #[derive(Default)] + struct MemSenderKeyStore(HashMap<SenderKeyName, SenderKeyRecord>); + #[async_trait::async_trait] + impl SenderKeyStore for MemSenderKeyStore { + async fn store_sender_key( + &mut self, + n: &SenderKeyName, + r: SenderKeyRecord, + ) -> SigResult<()> { + self.0.insert(n.clone(), r); + Ok(()) + } + async fn load_sender_key(&self, n: &SenderKeyName) -> SigResult<Option<SenderKeyRecord>> { + Ok(self.0.get(n).cloned()) + } + } + + // Outgoing group encryption never consumes our own prekeys, and device B + // has no bundle (so no session is established for it) — these are never + // called; present only to satisfy the generic bounds. + struct UnusedPreKeyStore; + #[async_trait::async_trait] + impl PreKeyStore for UnusedPreKeyStore { + async fn get_pre_key(&self, _: PreKeyId) -> SigResult<PreKeyRecord> { + unreachable!("prekey store not used in outgoing group encrypt") + } + async fn save_pre_key(&mut self, _: PreKeyId, _: &PreKeyRecord) -> SigResult<()> { + unreachable!() + } + async fn remove_pre_key(&mut self, _: PreKeyId) -> SigResult<()> { + unreachable!() + } + } + struct UnusedSignedPreKeyStore; + #[async_trait::async_trait] + impl SignedPreKeyStore for UnusedSignedPreKeyStore { + async fn get_signed_pre_key(&self, _: SignedPreKeyId) -> SigResult<SignedPreKeyRecord> { + unreachable!("signed prekey store not used in outgoing group encrypt") + } + async fn save_signed_pre_key( + &mut self, + _: SignedPreKeyId, + _: &SignedPreKeyRecord, + ) -> SigResult<()> { + unreachable!() + } + } + + struct TokioTestRuntime; + #[async_trait::async_trait] + impl Runtime for TokioTestRuntime { + fn spawn(&self, future: Pin<Box<dyn Future<Output = ()> + Send + 'static>>) -> AbortHandle { + let handle = tokio::spawn(future); + AbortHandle::new(move || handle.abort()) + } + fn sleep(&self, _d: Duration) -> Pin<Box<dyn Future<Output = ()> + Send>> { + // Not exercised on the send path; wacore dev-deps omit tokio's + // "time" feature, so resolve immediately rather than time out. + Box::pin(async {}) + } + fn spawn_blocking( + &self, + f: Box<dyn FnOnce() + Send + 'static>, + ) -> Pin<Box<dyn Future<Output = ()> + Send>> { + Box::pin(async move { + let _ = tokio::task::spawn_blocking(f).await; + }) + } + fn yield_now(&self) -> Option<Pin<Box<dyn Future<Output = ()> + Send>>> { + None + } + } + + // Establish a real Signal session for `a` so its SKDM encrypts; the + // returned identity store is the sender's (knows `a` after X3DH). + async fn established_stores(a: &Jid) -> (MemSessionStore, MemIdentityStore) { + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + let sender = IdentityKeyPair::generate(&mut rng); + let receiver = IdentityKeyPair::generate(&mut rng); + let spk = KeyPair::generate(&mut rng); + let opk = KeyPair::generate(&mut rng); + let sig = receiver + .private_key() + .calculate_signature(&spk.public_key.serialize(), &mut rng) + .unwrap(); + let bundle = PreKeyBundle::new( + 1, + 1u32.into(), + Some((1u32.into(), opk.public_key)), + 1u32.into(), + spk.public_key, + sig.to_vec(), + *receiver.identity_key(), + ) + .unwrap(); + let mut ss = MemSessionStore::default(); + let mut is = MemIdentityStore { + pair: sender, + reg_id: 42, + known: HashMap::new(), + }; + process_prekey_bundle( + &a.to_protocol_address(), + &mut ss, + &mut is, + &bundle, + &mut rng, + UsePQRatchet::No, + ) + .await + .unwrap(); + (ss, is) + } + + #[tokio::test] + async fn failed_device_is_still_marked_has_key() { + let group: Jid = "120363000000000001@g.us".parse().unwrap(); + let own_jid: Jid = "559900000000@s.whatsapp.net".parse().unwrap(); + let own_lid: Jid = "100000000000000@lid".parse().unwrap(); + // A has a session (encrypts ok); B has neither session nor bundle, + // mimicking a device that 406'd / has no key material. + let a: Jid = "559911112222:0@s.whatsapp.net".parse().unwrap(); + let b: Jid = "559933334444:0@s.whatsapp.net".parse().unwrap(); + + let (mut ss, mut is) = established_stores(&a).await; + let mut sks = MemSenderKeyStore::default(); + let mut pks = UnusedPreKeyStore; + let spks = UnusedSignedPreKeyStore; + let mut stores = SignalStores { + sender_key_store: &mut sks, + session_store: &mut ss, + identity_store: &mut is, + prekey_store: &mut pks, + signed_prekey_store: &spks, + }; + + // Empty resolver: no LID overrides; B's prekey fetch returns nothing + // → B is dropped by the encrypt fan-out (not in encrypted_devices). + let resolver = MockSendContextResolver::new(); + let rt = TokioTestRuntime; + + let group_info = GroupInfo::new( + vec![own_jid.to_non_ad(), a.to_non_ad(), b.to_non_ad()], + AddressingMode::Pn, + ); + let msg = wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }; + + let prepared = prepare_group_stanza( + &rt, + &mut stores, + &resolver, + &group_info, + &own_jid, + &own_lid, + None, + group, + &msg, + "TESTREQID".into(), + false, + Some(vec![a.clone(), b.clone()]), + None, + None, + &[], + ) + .await + .expect("prepare_group_stanza should succeed even when a device fails to encrypt"); + + let marked: std::collections::HashSet<String> = prepared + .skdm_devices + .iter() + .map(|j| j.to_string()) + .collect(); + + assert!( + marked.contains(&a.to_string()), + "device that encrypted must be marked" + ); + assert!( + marked.contains(&b.to_string()), + "device whose SKDM encryption FAILED must still be marked has_key \ + (WA Web markHasSenderKey(x, M) marks the full target set → no re-fanout storm)" + ); + assert_eq!( + prepared.skdm_devices.len(), + 2, + "exactly the full distribution list (A + B), not just the encrypted subset" + ); + + // A key-distributing send must carry a phash (computed over the list). + assert!( + prepared.node.attrs().optional_string("phash").is_some(), + "a key-distributing group send must carry a phash" + ); + } +} + +/// Item 3 — phash device-set construction. The set hashed is the full +/// recipient list PLUS the sending device (which is never in the recipient +/// list, since we don't SKDM ourselves), matching WA Web +/// `phashV2([].concat(A, [B]))`. +/// +/// This was confirmed against a real WA Web capture sent to the production +/// server: the recipient `<to>` set plus the sending device reproduced the +/// exact `phash` on the wire, while the recipient set alone did not — so the +/// sending device is part of the hash. Raw identifiers are not committed +/// (PII); the vectors below are fictitious but exercise the same logic. +mod group_phash_golden { + use super::*; + + #[test] + fn phash_set_includes_sending_device() { + // Fictitious group: a few users with bare (device 0) + companion + // devices. The self user appears as a companion (device 0) in the + // recipient list; its SENDING device (24) is excluded, mirroring a + // real send (we never SKDM ourselves). + let recipients: Vec<Jid> = [ + "100000000000001@lid", + "100000000000001:5@lid", + "100000000000002@lid", + "100000000000003@lid", + "100000000000003:12@lid", + "100000000000099@lid", + ] + .iter() + .map(|s| s.parse().expect("valid LID jid")) + .collect(); + + let own_sending: Jid = "100000000000099:24@lid".parse().unwrap(); + assert!( + !recipients + .iter() + .any(|j: &Jid| j.user == "100000000000099" && j.device == 24), + "the sending device must not already be in the recipient list" + ); + + let set = build_group_phash_set(&recipients, &own_sending); + assert_eq!(set.len(), 7, "6 recipients + the sending device"); + + // Dropping the sending device changes the hash, proving it is part + // of the hashed set (WA Web `[].concat(A, [B])`). + let with_self = MessageUtils::participant_list_hash(&set).unwrap(); + let without_self = MessageUtils::participant_list_hash(&recipients).unwrap(); + assert_ne!(with_self, without_self); + + // Deterministic standard-base64 vectors (regression guard). + assert_eq!(without_self, "2:rZoSAdIV"); + assert_eq!(with_self, "2:sti8OtHX"); + } + + #[test] + fn phash_set_drops_hosted_devices() { + // Hosted (Cloud API) devices don't take part in group E2EE and must + // not enter the phash, mirroring the SKDM distribution filter. + let with_hosted: Vec<Jid> = ["100000000000001@lid", "100000000000002:99@hosted"] + .iter() + .map(|s| s.parse().expect("valid jid")) + .collect(); + let without_hosted: Vec<Jid> = ["100000000000001@lid"] + .iter() + .map(|s| s.parse().expect("valid jid")) + .collect(); + let own: Jid = "100000000000099:24@lid".parse().unwrap(); + + assert_eq!( + build_group_phash_set(&with_hosted, &own), + build_group_phash_set(&without_hosted, &own), + "hosted devices must not affect the phash set" + ); + } +} + +mod local_identity_change_on_send { + use super::*; + use crate::libsignal::protocol::{ + Direction, IdentityChange, IdentityKey, IdentityKeyStore, PreKeyId, PreKeyRecord, + PreKeyStore, ProtocolAddress, SenderKeyRecord, SessionRecord, SessionStore, SignedPreKeyId, + SignedPreKeyRecord, SignedPreKeyStore, + }; + use crate::runtime::{AbortHandle, Runtime}; + use crate::types::jid::JidExt; + use std::future::Future; + use std::pin::Pin; + use std::time::Duration; + + type SigResult<T> = crate::libsignal::protocol::error::Result<T>; + + #[derive(Clone, Default)] + struct MemSessionStore(HashMap<ProtocolAddress, Vec<u8>>); + #[async_trait::async_trait] + impl SessionStore for MemSessionStore { + async fn load_session(&self, a: &ProtocolAddress) -> SigResult<Option<SessionRecord>> { + Ok(self + .0 + .get(a) + .and_then(|b| SessionRecord::deserialize(b).ok())) + } + async fn has_session(&self, a: &ProtocolAddress) -> SigResult<bool> { + Ok(self.0.contains_key(a)) + } + async fn store_session(&mut self, a: &ProtocolAddress, r: SessionRecord) -> SigResult<()> { + self.0.insert(a.clone(), r.serialize()?); + Ok(()) + } + } + + /// Identity store that reports the real change (unlike the hardcoded + /// stub elsewhere), so a pre-seeded stale key surfaces as ReplacedExisting. + #[derive(Clone)] + struct MemIdentityStore { + pair: IdentityKeyPair, + known: HashMap<ProtocolAddress, IdentityKey>, + } + #[async_trait::async_trait] + impl IdentityKeyStore for MemIdentityStore { + async fn get_identity_key_pair(&self) -> SigResult<IdentityKeyPair> { + Ok(self.pair.clone()) + } + async fn get_local_registration_id(&self) -> SigResult<u32> { + Ok(42) + } + async fn save_identity( + &mut self, + a: &ProtocolAddress, + id: &IdentityKey, + ) -> SigResult<IdentityChange> { + let changed = self.known.get(a).is_some_and(|k| k != id); + self.known.insert(a.clone(), *id); + Ok(IdentityChange::from_changed(changed)) + } + async fn is_trusted_identity( + &self, + _: &ProtocolAddress, + _: &IdentityKey, + _: Direction, + ) -> SigResult<bool> { + Ok(true) + } + async fn get_identity(&self, a: &ProtocolAddress) -> SigResult<Option<IdentityKey>> { + Ok(self.known.get(a).copied()) + } + } + + struct UnusedPreKeyStore; + #[async_trait::async_trait] + impl PreKeyStore for UnusedPreKeyStore { + async fn get_pre_key(&self, _: PreKeyId) -> SigResult<PreKeyRecord> { + unreachable!() + } + async fn save_pre_key(&mut self, _: PreKeyId, _: &PreKeyRecord) -> SigResult<()> { + unreachable!() + } + async fn remove_pre_key(&mut self, _: PreKeyId) -> SigResult<()> { + unreachable!() + } + } + struct UnusedSignedPreKeyStore; + #[async_trait::async_trait] + impl SignedPreKeyStore for UnusedSignedPreKeyStore { + async fn get_signed_pre_key(&self, _: SignedPreKeyId) -> SigResult<SignedPreKeyRecord> { + unreachable!() + } + async fn save_signed_pre_key( + &mut self, + _: SignedPreKeyId, + _: &SignedPreKeyRecord, + ) -> SigResult<()> { + unreachable!() + } + } + #[derive(Default)] + struct MemSenderKeyStore( + HashMap<crate::libsignal::store::sender_key_name::SenderKeyName, SenderKeyRecord>, + ); + #[async_trait::async_trait] + impl SenderKeyStore for MemSenderKeyStore { + async fn store_sender_key( + &mut self, + n: &crate::libsignal::store::sender_key_name::SenderKeyName, + r: SenderKeyRecord, + ) -> SigResult<()> { + self.0.insert(n.clone(), r); + Ok(()) + } + async fn load_sender_key( + &self, + n: &crate::libsignal::store::sender_key_name::SenderKeyName, + ) -> SigResult<Option<SenderKeyRecord>> { + Ok(self.0.get(n).cloned()) + } + } + + struct TokioTestRuntime; + #[async_trait::async_trait] + impl Runtime for TokioTestRuntime { + fn spawn(&self, future: Pin<Box<dyn Future<Output = ()> + Send + 'static>>) -> AbortHandle { + let handle = tokio::spawn(future); + AbortHandle::new(move || handle.abort()) + } + fn sleep(&self, _d: Duration) -> Pin<Box<dyn Future<Output = ()> + Send>> { + Box::pin(async {}) + } + fn spawn_blocking( + &self, + f: Box<dyn FnOnce() + Send + 'static>, + ) -> Pin<Box<dyn Future<Output = ()> + Send>> { + Box::pin(async move { + let _ = tokio::task::spawn_blocking(f).await; + }) + } + fn yield_now(&self) -> Option<Pin<Box<dyn Future<Output = ()> + Send>>> { + None + } + } + + /// The send path must report a replaced identity via the resolver when + /// establishing a session whose bundle carries a new identity key for an + /// address we already knew (peer reinstall). Mirrors WA Web saveIdentity + /// -> handleNewIdentity firing during outbound session setup. + #[tokio::test] + async fn encrypt_for_devices_reports_replaced_identity() { + let mut rng = rand::make_rng::<rand::rngs::StdRng>(); + + // Receiver device D with a valid signed bundle. + let device: Jid = "5511777777777:0@s.whatsapp.net".parse().unwrap(); + let receiver = IdentityKeyPair::generate(&mut rng); + let spk = KeyPair::generate(&mut rng); + let opk = KeyPair::generate(&mut rng); + let sig = receiver + .private_key() + .calculate_signature(&spk.public_key.serialize(), &mut rng) + .unwrap(); + let bundle = PreKeyBundle::new( + 1, + 1u32.into(), + Some((1u32.into(), opk.public_key)), + 1u32.into(), + spk.public_key, + sig.to_vec(), + *receiver.identity_key(), + ) + .unwrap(); + + // Local stores: no session for D + a STALE identity pre-seeded for D's + // address, so establishing the session reports ReplacedExisting. + let sender = IdentityKeyPair::generate(&mut rng); + let stale = *IdentityKeyPair::generate(&mut rng).identity_key(); + let mut known = HashMap::new(); + known.insert(device.to_protocol_address(), stale); + + let mut session_store = MemSessionStore::default(); + let mut identity_store = MemIdentityStore { + pair: sender, + known, + }; + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + + let mut stores = SignalStores { + sender_key_store: &mut sender_key_store, + session_store: &mut session_store, + identity_store: &mut identity_store, + prekey_store: &mut prekey_store, + signed_prekey_store: &signed_prekey_store, + }; + + let resolver = MockSendContextResolver::new() + .with_bundle(device.clone(), bundle) + .with_devices(vec![device.clone()]); + let rt = TokioTestRuntime; + + encrypt_for_devices( + &rt, + &mut stores, + &resolver, + std::slice::from_ref(&device), + b"hello", + false, + None, + ) + .await + .expect("encrypt_for_devices"); + + assert_eq!( + resolver.captured_identity_changes(), + vec![device], + "replaced identity on the send path must be reported via the resolver" + ); + } +}