Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 61 additions & 1 deletion src/client/lid_pn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(())
Expand Down Expand Up @@ -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}");
Comment on lines +189 to +193

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Serialize PN→LID session migration with session locks

This migration path reads/writes sessions directly through backend and then calls signal_cache.delete_session without acquiring the per-address session_lock_for lock that the send/decrypt/retry paths rely on. If a message is concurrently advancing the PN session in cache, migration can copy an older backend record to LID and then drop the newer PN state before it is flushed, which can roll back ratchet state and trigger follow-up BadMac/SessionNotFound failures for active chats.

Useful? React with 👍 / 👎.

}
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.
///
Expand Down
69 changes: 23 additions & 46 deletions src/client/sessions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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(());
};
Comment on lines +272 to +275

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Don't silently no-op when the own LID hasn't been populated yet.

get_device_snapshot() can legitimately return lid = None during login before appstate/offline sync fills it in. Returning Ok(()) here means the proactive own-device migration/prekey fetch never runs on that path, so device 0 can remain PN-only until some later flow repairs it. Please defer or reschedule this step once the own LID is available instead of skipping it.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/client/sessions.rs` around lines 272 - 275, The check in the block that
inspects device_snapshot.lid (the let Some(ref own_lid) = device_snapshot.lid
else { ... return Ok(()); }) must not silently no-op; instead, when
get_device_snapshot() yields lid == None, enqueue or re-schedule the
primary-phone migration/prekey fetch work to run once the own LID becomes
available (e.g., push a task to your existing retry/worker queue or register a
callback/listener that triggers the same primary-phone session check when the
appstate/offline sync populates lid). Update the logic around
get_device_snapshot(), own_lid, and the primary phone session check so that the
function returns a pending/rescheduled outcome rather than Ok(()) when lid is
missing, and ensure the re-scheduled action uses the same code path that runs
when own_lid is present.


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")
}
}

Expand Down
Loading
Loading