diff --git a/src/client/lid_pn.rs b/src/client/lid_pn.rs index e615c83cd..906d508a1 100644 --- a/src/client/lid_pn.rs +++ b/src/client/lid_pn.rs @@ -83,10 +83,12 @@ impl Client { .await .map_err(|e| anyhow!("persisting LID-PN mapping: {e}"))?; - // If this is a new LID mapping, migrate any existing PN-keyed device registry entries + // If this is a new LID mapping, migrate any existing PN-keyed entries to LID if is_new_mapping { self.migrate_device_registry_on_lid_discovery(phone_number, lid) .await; + self.migrate_signal_sessions_on_lid_discovery(phone_number, lid) + .await; } Ok(()) @@ -166,6 +168,64 @@ impl Client { } } + /// Migrate Signal sessions and identity keys from PN to LID address. + /// WA Web never stores sessions under PN when a LID mapping is known. + pub(crate) async fn migrate_signal_sessions_on_lid_discovery(&self, pn: &str, lid: &str) { + use log::{info, warn}; + use wacore::types::jid::JidExt; + + let backend = self.persistence_manager.backend(); + + for device_id in 0..=99u16 { + let pn_jid = Jid::pn_device(pn.to_string(), device_id); + let lid_jid = Jid::lid_device(lid.to_string(), device_id); + + let pn_proto = pn_jid.to_protocol_address(); + let lid_proto = lid_jid.to_protocol_address(); + let pn_addr_key = pn_proto.as_str(); + let lid_addr_key = lid_proto.as_str(); + + // Migrate session if PN session exists + if let Ok(Some(session_data)) = backend.get_session(pn_addr_key).await { + match backend.get_session(lid_addr_key).await { + Ok(Some(_)) => { + if let Err(e) = backend.delete_session(pn_addr_key).await { + warn!("Failed to delete stale PN session {pn_addr_key}: {e}"); + } + self.signal_cache.delete_session(&pn_proto).await; + info!("Deleted stale PN session {pn_addr_key} (LID exists)"); + } + Ok(None) => { + if let Err(e) = backend.put_session(lid_addr_key, &session_data).await { + warn!("Failed to write LID session {lid_addr_key}: {e}"); + } else { + if let Err(e) = backend.delete_session(pn_addr_key).await { + warn!("Failed to delete PN session {pn_addr_key}: {e}"); + } + self.signal_cache.delete_session(&pn_proto).await; + info!("Migrated session {pn_addr_key} -> {lid_addr_key}"); + } + } + Err(e) => warn!("Failed to check LID session {lid_addr_key}: {e}"), + } + } + + // Migrate identity independently of session (can outlive deleted sessions) + if let Ok(Some(identity_data)) = backend.load_identity(pn_addr_key).await + && let Ok(None) = backend.load_identity(lid_addr_key).await + { + let Ok(key): Result<[u8; 32], _> = identity_data.as_slice().try_into() else { + continue; + }; + if let Err(e) = backend.put_identity(lid_addr_key, key).await { + warn!("Failed to migrate identity {pn_addr_key} -> {lid_addr_key}: {e}"); + } else if let Err(e) = backend.delete_identity(pn_addr_key).await { + warn!("Failed to delete PN identity {pn_addr_key}: {e}"); + } + } + } + } + /// Get the phone number (user part) for a given LID. /// Looks up the LID-PN mapping from the in-memory cache. /// diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 49815a4d8..2f440d9f2 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -259,15 +259,8 @@ impl Client { Ok(success_count) } - /// Establish session with primary phone (device 0) immediately for PDO. - /// - /// Called during login BEFORE offline messages arrive. Checks both PN and LID - /// sessions but does NOT establish PN sessions proactively. The primary phone's - /// PN session will be established via LID pkmsg when needed, which prevents - /// dual-session conflicts where both PN and LID sessions exist for the same user. - /// This matches WhatsApp Web's `prekey_fetch_iq_pnh_lid_enabled: false` behavior. - /// - /// Returns error if session check fails (fail-safe to prevent replacing existing sessions). + /// Log primary phone (device 0) session state at login. + /// Migration is lazy via try_pn_to_lid_migration_decrypt on first message. pub(crate) async fn establish_primary_phone_session_immediate(&self) -> Result<()> { let device_snapshot = self.persistence_manager.get_device_snapshot().await; @@ -276,46 +269,30 @@ impl Client { .clone() .ok_or_else(|| anyhow::Error::from(crate::client::ClientError::NotLoggedIn))?; + let Some(ref own_lid) = device_snapshot.lid else { + log::debug!("No own LID yet, skipping primary phone session check"); + return Ok(()); + }; + + let primary_phone_lid = own_lid.with_device(0); let primary_phone_pn = own_pn.with_device(0); - let primary_phone_lid = device_snapshot.lid.as_ref().map(|lid| lid.with_device(0)); - let pn_session_exists = - self.check_session_exists(&primary_phone_pn) - .await - .map_err(|e| { - anyhow::anyhow!( - "Cannot verify PN session existence for primary phone {}: {}. \ - Refusing to establish session to prevent potential MAC failures.", - primary_phone_pn, - e - ) - })?; - - // Don't proactively establish PN session - matches WhatsApp Web's - // prekey_fetch_iq_pnh_lid_enabled: false behavior. The primary phone will - // establish the session via pkmsg from LID address, which prevents dual-session - // conflicts where both PN and LID sessions exist for the same user. - if pn_session_exists { - log::debug!( - "PN session with primary phone {} already exists", - primary_phone_pn - ); - } else { - log::debug!( - "No PN session with primary phone {} - will be established via LID pkmsg", - primary_phone_pn - ); - } + let lid_exists = self + .check_session_exists(&primary_phone_lid) + .await + .unwrap_or(false); + let pn_exists = self + .check_session_exists(&primary_phone_pn) + .await + .unwrap_or(false); - // Check LID session existence (don't establish - primary phone does that via pkmsg) - if let Some(ref lid_jid) = primary_phone_lid { - match self.check_session_exists(lid_jid).await { - Ok(true) => log::debug!("LID session with {} already exists", lid_jid), - Ok(false) => log::debug!( - "No LID session with {} - established on first message", - lid_jid - ), - Err(e) => log::debug!("Could not check LID session for {}: {}", lid_jid, e), + match (lid_exists, pn_exists) { + (true, _) => log::debug!("LID session with {} exists", primary_phone_lid), + (false, true) => { + log::debug!("PN-only session for own device 0 — will migrate on first message") + } + (false, false) => { + log::debug!("No session with own device 0 — will establish on first message") } } diff --git a/src/message.rs b/src/message.rs index 816e7e95e..72227f547 100644 --- a/src/message.rs +++ b/src/message.rs @@ -269,85 +269,10 @@ impl Client { return; } - // Determine the JID to use for end-to-end decryption. - // ... (previous JID resolution comments) - let sender_encryption_jid = { - let sender = &info.source.sender; - let alt = info.source.sender_alt.as_ref(); - let pn_server = wacore_binary::jid::DEFAULT_USER_SERVER; - let lid_server = wacore_binary::jid::HIDDEN_USER_SERVER; - - if sender.server == lid_server { - // Sender is already LID - use it directly for session lookup. - // Also cache the LID-to-PN mapping if PN alt is available. - if let Some(alt_jid) = alt - && alt_jid.server == pn_server - { - if let Err(err) = self - .add_lid_pn_mapping( - &sender.user, - &alt_jid.user, - crate::lid_pn_cache::LearningSource::PeerLidMessage, - ) - .await - { - warn!( - "Failed to persist LID-to-PN mapping {} -> {}: {err}", - sender.user, alt_jid.user - ); - } - debug!( - "Cached LID-to-PN mapping: {} -> {}", - sender.user, alt_jid.user - ); - } - sender.clone() - } else if sender.server == pn_server { - // ... (PN to LID resolution logic) - if let Some(alt_jid) = alt - && alt_jid.server == lid_server - { - if let Err(err) = self - .add_lid_pn_mapping( - &alt_jid.user, - &sender.user, - crate::lid_pn_cache::LearningSource::PeerPnMessage, - ) - .await - { - warn!( - "Failed to persist PN-to-LID mapping {} -> {}: {err}", - sender.user, alt_jid.user - ); - } - debug!( - "Cached PN-to-LID mapping: {} -> {}", - sender.user, alt_jid.user - ); - - Jid { - user: alt_jid.user.clone(), - server: wacore_binary::jid::cow_server_from_str(lid_server), - device: sender.device, - agent: sender.agent, - integrator: sender.integrator, - } - } else if let Some(lid_user) = self.lid_pn_cache.get_current_lid(&sender.user).await - { - Jid { - user: lid_user.clone(), - server: wacore_binary::jid::cow_server_from_str(lid_server), - device: sender.device, - agent: sender.agent, - integrator: sender.integrator, - } - } else { - sender.clone() - } - } else { - sender.clone() - } - }; + // Warm LID-PN cache before resolution so resolve_encryption_jid() finds the mapping + self.cache_lid_pn_from_message(&info.source.sender, info.source.sender_alt.as_ref()) + .await; + let sender_encryption_jid = self.resolve_encryption_jid(&info.source.sender).await; let has_unavailable = node.get_optional_child("unavailable").is_some(); @@ -944,13 +869,29 @@ impl Client { continue; } - // Handle SessionNotFound gracefully - send retry receipt to request session establishment + // Try PN→LID session migration before sending retry receipt if let SignalProtocolError::SessionNotFound(_) = e { + if self + .try_pn_to_lid_migration_decrypt( + sender_encryption_jid, + &signal_address, + &parsed_message, + &mut adapter, + &mut rng, + &enc_type, + padding_version, + info, + ) + .await + { + any_success = true; + continue; + } + warn!( "[msg:{}] No session found for {} message from {}. Sending retry receipt to request session establishment.", info.id, enc_type, info.source.sender ); - // Send retry receipt so the sender resends with a PreKeySignalMessage dispatched_undecryptable = self.handle_decrypt_failure( info, RetryReason::NoSession, @@ -1228,6 +1169,136 @@ impl Client { Ok(()) } + /// Attempt PN→LID session migration and retry decryption. + /// Returns true if decryption succeeded after migration. + #[allow(clippy::too_many_arguments)] + async fn try_pn_to_lid_migration_decrypt( + self: &Arc, + 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: &MessageInfo, + ) -> bool { + use wacore::libsignal::protocol::{UsePQRatchet, message_decrypt}; + + if !sender_jid.is_lid() { + return false; + } + + let Some(pn) = self.lid_pn_cache.get_phone_number(&sender_jid.user).await else { + return false; + }; + + self.migrate_signal_sessions_on_lid_discovery(&pn, &sender_jid.user) + .await; + + // Reload into cache (replacing negative cache entries) + let backend = self.persistence_manager.backend(); + let addr_key = signal_address.as_str(); + if let Ok(Some(session_data)) = backend.get_session(addr_key).await + && let Ok(record) = + wacore::libsignal::protocol::SessionRecord::deserialize(&session_data) + { + self.signal_cache.put_session(signal_address, record).await; + } + if let Ok(Some(identity_data)) = backend.load_identity(addr_key).await { + self.signal_cache + .put_identity(signal_address, &identity_data) + .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 + { + Ok(padded_plaintext) => { + log::info!( + "[msg:{}] Decrypted after PN→LID session migration for {}", + info.id, + info.source.sender + ); + if let Err(e) = self + .clone() + .handle_decrypted_plaintext(enc_type, &padded_plaintext, padding_version, info) + .await + { + log::warn!( + "[msg:{}] Failed processing plaintext after migration: {e:?}", + info.id + ); + } + true + } + Err(SignalProtocolError::DuplicatedMessage(chain, counter)) => { + log::debug!( + "[msg:{}] Already processed (chain {chain}, counter {counter}) after migration", + info.id + ); + true + } + Err(retry_err) => { + log::warn!( + "[msg:{}] Decryption still failed after PN→LID migration: {retry_err:?}", + info.id + ); + false + } + } + } + + /// Cache LID-PN mapping from message attributes (before resolve_encryption_jid). + async fn cache_lid_pn_from_message(&self, sender: &Jid, alt: Option<&Jid>) { + let pn_server = wacore_binary::jid::DEFAULT_USER_SERVER; + let lid_server = wacore_binary::jid::HIDDEN_USER_SERVER; + + let (lid_user, pn_user, source) = if sender.server == lid_server { + if let Some(alt_jid) = alt + && alt_jid.server == pn_server + { + ( + &sender.user, + &alt_jid.user, + crate::lid_pn_cache::LearningSource::PeerLidMessage, + ) + } else { + return; + } + } else if sender.server == pn_server { + if let Some(alt_jid) = alt + && alt_jid.server == lid_server + { + ( + &alt_jid.user, + &sender.user, + crate::lid_pn_cache::LearningSource::PeerPnMessage, + ) + } else { + return; + } + } else { + return; + }; + + if let Err(err) = self.add_lid_pn_mapping(lid_user, pn_user, source).await { + warn!( + "Failed to cache LID-PN mapping {} <-> {}: {err}", + lid_user, pn_user + ); + } + } + pub(crate) async fn parse_message_info( &self, node: &Node, @@ -2872,10 +2943,15 @@ mod tests { "Should detect self-sent DM from own LID" ); - // 2. sender_alt should be None (peer_recipient_pn is recipient's PN, not sender's) + // 2. sender_alt should be own PN (derived from own_jid, not message attrs) assert!( - info.source.sender_alt.is_none(), - "sender_alt should be None for self-sent DMs (peer_recipient_pn is recipient's PN)" + 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!( @@ -3059,8 +3135,13 @@ mod tests { ); assert!( - info.source.sender_alt.is_none(), - "sender_alt should be None for self-sent messages" + 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!( diff --git a/src/retry.rs b/src/retry.rs index 98a9b6833..6bdd2f491 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -175,8 +175,9 @@ impl Client { .parse::() .unwrap_or_else(|_| receipt.source.sender.clone()); - // Device existence check (matches WhatsApp Web's WAWebApiDeviceList.hasDevice). - // This prevents processing retry receipts from unknown/stale devices. + // Resolved JID for session operations; keep original for stanza addressing + let resolved_jid = self.resolve_encryption_jid(&participant_jid).await; + let sender_device_id = participant_jid.device() as u32; let sender_user = participant_jid.user.clone(); if !self.has_device(&sender_user, sender_device_id).await { @@ -192,7 +193,11 @@ impl Client { let is_peer = device_snapshot .pn .as_ref() - .is_some_and(|our_pn| participant_jid.user == our_pn.user); + .is_some_and(|our_pn| participant_jid.is_same_user_as(our_pn)) + || device_snapshot + .lid + .as_ref() + .is_some_and(|our_lid| participant_jid.is_same_user_as(our_lid)); // Process key bundle to establish a pairwise session for the retry. // Needed for both DMs and groups (group retries use pairwise, not sender key). @@ -200,7 +205,7 @@ impl Client { if !receipt.source.chat.is_status_broadcast() { // Try to process key bundle if present let key_bundle_result = self - .process_retry_key_bundle(node, &participant_jid, is_peer) + .process_retry_key_bundle(node, &resolved_jid, is_peer) .await; if let Err(e) = &key_bundle_result { @@ -213,7 +218,7 @@ impl Client { // session, delete the session to force re-establishment. // This handles the case where the requester reinstalled but didn't include keys. if let Some(received_reg_id) = extract_registration_id_from_node(node) { - let signal_address = participant_jid.to_protocol_address(); + let signal_address = resolved_jid.to_protocol_address(); let device_store = self.persistence_manager.get_device_arc().await; let device_guard = device_store.read().await; @@ -362,7 +367,7 @@ impl Client { // For DMs, handle base key tracking for collision detection (matches WhatsApp Web). // This detects when we haven't regenerated our session despite receiving retry receipts, // which can cause infinite retry loops where both sides are stuck with stale keys. - let signal_address = participant_jid.to_protocol_address(); + let signal_address = resolved_jid.to_protocol_address(); let device_store = self.persistence_manager.get_device_arc().await; // Check for base key collision before deleting the session. @@ -462,7 +467,6 @@ impl Client { if receipt.source.chat.is_group() { // Group retry: pairwise encrypt to failing device only (RetryMsgJob.js:71). // Using sender-key broadcast would resend to ALL participants → duplicates. - let encryption_jid = self.resolve_encryption_jid(&participant_jid).await; let device_snapshot = self.persistence_manager.get_device_snapshot().await; let addressing_mode = cached_group_info @@ -477,7 +481,7 @@ impl Client { &mut store_adapter.identity_store, receipt.source.chat.clone(), participant_jid, - encryption_jid, + resolved_jid.clone(), &original_msg, message_id, retry_count, @@ -548,7 +552,8 @@ impl Client { return Err(anyhow::anyhow!("Invalid registration ID in retry receipt")); } - let signal_address = requester_jid.to_protocol_address(); + let resolved_jid = self.resolve_encryption_jid(requester_jid).await; + let signal_address = resolved_jid.to_protocol_address(); // Check if the registration ID changed (indicates device reinstall). // Read session through cache for consistent state. diff --git a/tests/e2e/tests/lid_sessions.rs b/tests/e2e/tests/lid_sessions.rs new file mode 100644 index 000000000..e2239bb32 --- /dev/null +++ b/tests/e2e/tests/lid_sessions.rs @@ -0,0 +1,587 @@ +//! LID-first Signal session tests. +//! +//! Validates that Signal sessions are always stored under LID addresses +//! (not PN) and that own-device messaging works without NoSession errors. +//! +//! These tests exercise: +//! - Sessions created via the send path are under LID, not PN +//! - PN sessions injected into the DB get migrated to LID +//! - Own device 0 has a LID session after login +//! - LID sessions survive reconnect (DB reload) +//! - No UndecryptableMessage events during normal messaging +//! - Multiple sequential sends don't regress to PN sessions + +use e2e_tests::{TestClient, send_and_expect_text}; +use log::info; +use wacore::store::traits::SignalStore; +use wacore::types::events::Event; + +fn mask_addr(addr: &str) -> String { + if let Some(at) = addr.find('@') { + let user = &addr[..at]; + let rest = &addr[at..]; + if user.len() > 4 { + format!("{}...{}{rest}", &user[..2], &user[user.len() - 2..]) + } else { + format!("{user}{rest}") + } + } else { + addr.to_string() + } +} + +/// Scan backend for sessions matching a user across device IDs 0..=99. +async fn scan_sessions( + backend: &dyn SignalStore, + user: &str, + server: &str, +) -> anyhow::Result> { + let mut results = Vec::new(); + for device_id in 0..=99u16 { + let addr = if device_id == 0 { + format!("{user}@{server}.0") + } else { + format!("{user}:{device_id}@{server}.0") + }; + if backend.get_session(&addr).await?.is_some() { + results.push(addr); + } + } + Ok(results) +} + +/// Assert that ALL sessions for a user are under LID, NONE under PN. +async fn assert_lid_only_sessions( + backend: &dyn SignalStore, + pn_user: &str, + lid_user: &str, + context: &str, +) { + let pn_sessions = scan_sessions(backend, pn_user, "c.us") + .await + .expect("scan PN sessions"); + let lid_sessions = scan_sessions(backend, lid_user, "lid") + .await + .expect("scan LID sessions"); + + assert!( + !lid_sessions.is_empty(), + "[{context}] Expected at least one LID session, found 0. PN count: {}", + pn_sessions.len() + ); + assert!( + pn_sessions.is_empty(), + "[{context}] Expected 0 PN sessions, found {}. LID count: {}", + pn_sessions.len(), + lid_sessions.len() + ); + + info!( + "[{context}] LID-only: {} session(s), first={}", + lid_sessions.len(), + lid_sessions + .first() + .map(|s| mask_addr(s)) + .unwrap_or_default() + ); +} + +/// After a roundtrip between two clients, sessions should be stored +/// exclusively under LID addresses, not PN. Also verifies that +/// both sides have LID-only sessions (bidirectional check). +#[tokio::test] +async fn test_sessions_stored_under_lid_not_pn() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut client_a = TestClient::connect("e2e_lid_sess_a").await?; + let mut client_b = TestClient::connect("e2e_lid_sess_b").await?; + + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + let lid_a = client_a.client.get_lid().await.expect("A should have LID"); + let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + + // Roundtrip to establish sessions in both directions + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "LID session test A->B", + 30, + ) + .await?; + send_and_expect_text( + &client_b.client, + &mut client_a, + &jid_a, + "LID session test B->A", + 30, + ) + .await?; + info!("Roundtrip complete"); + + // Check A's view: B's sessions should be LID-only + let backend_a = client_a.client.persistence_manager().backend(); + assert_lid_only_sessions(&*backend_a, &jid_b.user, &lid_b.user, "A's store for B").await; + + // Check B's view: A's sessions should be LID-only + let backend_b = client_b.client.persistence_manager().backend(); + assert_lid_only_sessions(&*backend_b, &jid_a.user, &lid_a.user, "B's store for A").await; + + // Verify continued delivery (session is functional, not just stored) + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "Post-check A->B", + 30, + ) + .await?; + send_and_expect_text( + &client_b.client, + &mut client_a, + &jid_a, + "Post-check B->A", + 30, + ) + .await?; + info!("Bidirectional post-check delivery confirmed"); + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +/// Multiple sequential sends should never create PN sessions. +/// Regression guard: the first send establishes a session, subsequent sends +/// must reuse the LID session (not accidentally create a PN one). +#[tokio::test] +async fn test_multiple_sends_stay_lid_only() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let client_a = TestClient::connect("e2e_lid_multi_send_a").await?; + let mut client_b = TestClient::connect("e2e_lid_multi_send_b").await?; + + let jid_b = client_b.jid().await; + let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + + // Send 5 messages sequentially + for i in 1..=5 { + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + &format!("Sequential msg {i}"), + 30, + ) + .await?; + } + info!("All 5 sequential messages delivered"); + + let backend_a = client_a.client.persistence_manager().backend(); + + // After 5 sends, still only LID sessions + assert_lid_only_sessions( + &*backend_a, + &jid_b.user, + &lid_b.user, + "After 5 sequential sends", + ) + .await; + + // Count LID sessions — should be exactly 1 primary device session + // (not 5 sessions from 5 sends) + let lid_sessions = scan_sessions(&*backend_a, &lid_b.user, "lid").await?; + info!("LID session count after 5 sends: {}", lid_sessions.len()); + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +/// If a stale PN session exists in the database (simulating a legacy pairing), +/// messaging should still work via the LID session without undecryptable errors. +#[tokio::test] +async fn test_stale_pn_session_does_not_break_lid_messaging() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut client_a = TestClient::connect("e2e_lid_migrate_a").await?; + let mut client_b = TestClient::connect("e2e_lid_migrate_b").await?; + + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + + // First, establish a normal LID session via roundtrip + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "Establish session", + 30, + ) + .await?; + send_and_expect_text( + &client_b.client, + &mut client_a, + &jid_a, + "Reply to establish", + 30, + ) + .await?; + + let backend_a = client_a.client.persistence_manager().backend(); + + // Read the existing LID session data + let lid_addr = format!("{}@lid.0", lid_b.user); + let lid_session_data = backend_a + .get_session(&lid_addr) + .await? + .expect("LID session should exist after roundtrip"); + + // Inject a copy under the PN address (simulates legacy database state) + let pn_addr = format!("{}@c.us.0", jid_b.user); + backend_a.put_session(&pn_addr, &lid_session_data).await?; + info!("Injected stale PN session at {}", mask_addr(&pn_addr)); + + // Verify both exist before migration + assert!( + backend_a.get_session(&pn_addr).await?.is_some(), + "PN session should exist after injection" + ); + assert!( + backend_a.get_session(&lid_addr).await?.is_some(), + "LID session should still exist" + ); + + // Send more messages — should use the LID session, not the stale PN copy + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "Post-inject A->B", + 30, + ) + .await?; + send_and_expect_text( + &client_b.client, + &mut client_a, + &jid_a, + "Post-inject B->A", + 30, + ) + .await?; + info!("Messaging works despite stale PN session in DB"); + + // LID session should still be authoritative + assert!( + backend_a.get_session(&lid_addr).await?.is_some(), + "LID session at {lid_addr} should survive after messaging with stale PN present" + ); + + // No undecryptable events on either side + client_a + .assert_no_event( + 3, + |e| matches!(e, Event::UndecryptableMessage(_)), + "A should have no undecryptable messages with stale PN session", + ) + .await?; + client_b + .assert_no_event( + 3, + |e| matches!(e, Event::UndecryptableMessage(_)), + "B should have no undecryptable messages with stale PN session", + ) + .await?; + info!("No undecryptable events despite stale PN session in DB"); + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +/// After reconnect, sessions loaded from SQLite should still be under LID. +/// Also verifies no PN sessions appear after the cache is cleared and reloaded. +#[tokio::test] +async fn test_lid_session_survives_reconnect() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut client_a = TestClient::connect("e2e_lid_recon_a").await?; + let mut client_b = TestClient::connect("e2e_lid_recon_b").await?; + + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + + // Establish sessions with a roundtrip + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "Pre-reconnect A->B", + 30, + ) + .await?; + send_and_expect_text( + &client_b.client, + &mut client_a, + &jid_a, + "Pre-reconnect B->A", + 30, + ) + .await?; + info!("Sessions established"); + + // Verify LID-only before reconnect + let backend_a = client_a.client.persistence_manager().backend(); + assert_lid_only_sessions(&*backend_a, &jid_b.user, &lid_b.user, "Before reconnect").await; + + // Reconnect A — clears in-memory signal cache, forces DB reload + client_a.reconnect_and_wait().await?; + info!("Client A reconnected (cache cleared)"); + + // Session should still be under LID after reload + let backend_a = client_a.client.persistence_manager().backend(); + assert_lid_only_sessions(&*backend_a, &jid_b.user, &lid_b.user, "After reconnect").await; + + // Verify delivery still works with the reloaded session (no new pkmsg needed) + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "Post-reconnect 1", + 30, + ) + .await?; + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "Post-reconnect 2", + 30, + ) + .await?; + info!("Post-reconnect delivery confirmed (2 messages)"); + + // Final check: still LID-only after post-reconnect sends + assert_lid_only_sessions( + &*backend_a, + &jid_b.user, + &lid_b.user, + "After post-reconnect sends", + ) + .await; + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +/// Own device 0 should have a LID session after login, not PN. +#[tokio::test] +async fn test_own_device_0_has_lid_session_after_login() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let client = TestClient::connect("e2e_lid_dev0").await?; + + let own_pn = client + .client + .get_pn() + .await + .expect("Client should have PN after connect"); + let own_lid = client + .client + .get_lid() + .await + .expect("Client should have LID after connect"); + + let backend = client.client.persistence_manager().backend(); + + // Device 0 is the primary phone — session should exist under LID + let lid_addr = format!("{}@lid.0", own_lid.user); + let lid_session = backend.get_session(&lid_addr).await?; + assert!( + lid_session.is_some(), + "Should have LID session with own device 0 ({lid_addr}) after login" + ); + info!("Own device 0 LID session exists: {}", mask_addr(&lid_addr)); + + // No PN session should exist for own device 0 + let pn_addr = format!("{}@c.us.0", own_pn.user); + let pn_session = backend.get_session(&pn_addr).await?; + assert!( + pn_session.is_none(), + "Should NOT have PN session for own device 0 ({pn_addr}). \ + Own device sessions must be under LID." + ); + info!("Confirmed no PN session for own device 0"); + + client.disconnect().await; + Ok(()) +} + +/// Normal messaging should never produce UndecryptableMessage events. +/// This guards against the NoSession bug where decryption fails silently. +#[tokio::test] +async fn test_no_undecryptable_events_during_messaging() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut client_a = TestClient::connect("e2e_lid_undec_a").await?; + let mut client_b = TestClient::connect("e2e_lid_undec_b").await?; + + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + + // Exchange several messages in both directions + for i in 1..=3 { + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + &format!("A->B msg {i}"), + 30, + ) + .await?; + send_and_expect_text( + &client_b.client, + &mut client_a, + &jid_a, + &format!("B->A msg {i}"), + 30, + ) + .await?; + } + info!("6 messages exchanged successfully"); + + // Neither side should have any undecryptable messages + client_a + .assert_no_event( + 3, + |e| matches!(e, Event::UndecryptableMessage(_)), + "A should have no undecryptable messages", + ) + .await?; + client_b + .assert_no_event( + 3, + |e| matches!(e, Event::UndecryptableMessage(_)), + "B should have no undecryptable messages", + ) + .await?; + info!("No undecryptable events on either side"); + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +/// Reproduces the exact production bug: a session exists under PN address +/// but NOT under LID. When a message arrives addressed from LID, decryption +/// fails with NoSession and an UndecryptableMessage event is dispatched. +/// +/// This simulates a database from an old pairing where the session was stored +/// under PN, but WhatsApp has since migrated to LID-based addressing. +#[tokio::test] +async fn test_pn_only_session_causes_undecryptable_on_lid_lookup() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut client_a = TestClient::connect("e2e_lid_repro_a").await?; + let mut client_b = TestClient::connect("e2e_lid_repro_b").await?; + + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + let lid_b = client_b.client.get_lid().await.expect("B should have LID"); + + // Step 1: Establish sessions via roundtrip + send_and_expect_text(&client_a.client, &mut client_b, &jid_b, "Setup A->B", 30).await?; + send_and_expect_text(&client_b.client, &mut client_a, &jid_a, "Setup B->A", 30).await?; + info!("Sessions established via roundtrip"); + + // Step 2: First reconnect — ensures sessions are flushed to backend + clears cache + client_a.reconnect_and_wait().await?; + info!("First reconnect done"); + + // Step 3: Verify messaging works (forces session load from backend into clean cache) + send_and_expect_text( + &client_b.client, + &mut client_a, + &jid_a, + "Post-reconnect verify", + 30, + ) + .await?; + info!("Verified messaging works after first reconnect"); + + // Step 4: Simulate legacy DB — move session from LID to PN address + let backend_a = client_a.client.persistence_manager().backend(); + let lid_addr = format!("{}@lid.0", lid_b.user); + let pn_addr = format!("{}@c.us.0", jid_b.user); + + let lid_session_data = backend_a + .get_session(&lid_addr) + .await? + .expect("LID session should exist in backend"); + info!("Read LID session ({} bytes)", lid_session_data.len()); + + // Copy to PN address, then delete LID — simulates legacy state + backend_a.put_session(&pn_addr, &lid_session_data).await?; + backend_a.delete_session(&lid_addr).await?; + info!("Moved session to PN (simulating legacy DB)"); + + // Verify backend state is now PN-only + assert!( + backend_a.get_session(&lid_addr).await?.is_none(), + "LID session should be deleted from backend" + ); + assert!( + backend_a.get_session(&pn_addr).await?.is_some(), + "PN session should exist in backend" + ); + + // Step 5: Second reconnect — clears cache, reloads from modified backend + // Now A only has a PN session for B, no LID session + client_a.reconnect_and_wait().await?; + info!("Second reconnect done (cache now has PN-only state)"); + + // Step 6: B sends to A — A should fail to decrypt (looks up LID, finds nothing) + // This is the exact production bug: phone sends from LID, session is under PN + let test_text = "This should trigger NoSession"; + client_b + .client + .send_message(jid_a.clone(), e2e_tests::text_msg(test_text)) + .await?; + info!("B sent message to A (expecting decryption failure on A)"); + + // Message must decrypt successfully (migration should happen on-the-fly) + client_a + .wait_for_text(test_text, 15) + .await + .expect("Message should decrypt after on-the-fly PN->LID migration"); + info!("Message decrypted despite PN-only backend state"); + + // No UndecryptableMessage should have been emitted + client_a + .assert_no_event( + 3, + |e| matches!(e, Event::UndecryptableMessage(_)), + "No undecryptable events after migration", + ) + .await?; + + // Session should now be under LID (migrated from PN) + let backend_a = client_a.client.persistence_manager().backend(); + assert!( + backend_a.get_session(&lid_addr).await?.is_some(), + "Session should be under LID after migration" + ); + assert!( + backend_a.get_session(&pn_addr).await?.is_none(), + "PN session should be cleaned up after migration" + ); + info!("Session correctly under LID after on-the-fly migration"); + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} diff --git a/wacore/src/messages.rs b/wacore/src/messages.rs index 1b1553e0c..0b10d7887 100644 --- a/wacore/src/messages.rs +++ b/wacore/src/messages.rs @@ -194,11 +194,20 @@ pub fn parse_message_info( .as_ref() .map(|r| r.to_non_ad()) .unwrap_or_else(|| from.to_non_ad()); + // Populate sender_alt so LID-PN cache warms from self-messages + let sender_alt = if from.server == jid::HIDDEN_USER_SERVER { + Some(own_jid.clone()) + } else if from.server == jid::DEFAULT_USER_SERVER && own_lid.is_some() { + own_lid.cloned() + } else { + None + }; MessageSource { chat, sender: from.clone(), is_from_me: true, recipient, + sender_alt, ..Default::default() } } else {