fix: self-DM/sibling decryption deadlock (retry shape + session recovery + peer pkmsg identity)

src/client.rs

@@ -423,6 +423,15 @@ pub struct Client {
     /// ran (the count alone can't separate NoSession from BadMac etc.).
     pub(crate) recent_retry_reasons: Cache<String, wacore::protocol::retry::RetryReason>,
 
+    /// Per-peer timestamp of the last forced session recreate via the
+    /// "no keys + retry≥2 + >1h since last" path (whatsmeow parity).
+    /// WA Web's updateLocalSignalSession only deletes on regId mismatch /
+    /// base-key collision — sessions that diverged without either trigger
+    /// stay stuck. This map throttles the fallback so a noisy peer can't
+    /// loop us through prekey fetches.
+    pub(crate) session_recreate_history:
+        Arc<std::sync::Mutex<HashMap<wacore_binary::jid::Jid, wacore::time::Instant>>>,
+
     /// Dispatch-once gate for `UndecryptableMessage`: a server resend of a
     /// failed id re-enters the failure path and would otherwise fire a
     /// duplicate event. Mirrors WA Web's DB-level placeholder uniqueness
@@ -821,6 +830,8 @@ impl Client {
 
             recent_retry_reasons: cache_config.message_retry_counts.build_with_ttl(),
 
+            session_recreate_history: Arc::new(std::sync::Mutex::new(HashMap::new())),
+
             undecryptable_dispatched: cache_config.undecryptable_dispatched.build_with_ttl(),
 
             offline_sync_metrics: Arc::new(OfflineSyncMetrics {

src/client/lid_pn.rs

@@ -349,6 +349,31 @@ impl Client {
     /// from the backend when the cache has unflushed mutations (e.g., after
     /// SKDM encryption ratcheted the session).
     pub(crate) async fn migrate_signal_sessions_on_lid_discovery(&self, pn: &str, lid: &str) {
+        self.migrate_signal_sessions_on_lid_discovery_inner(pn, lid, None)
+            .await;
+    }
+
+    /// Variant called from inside `decrypt_message`, which already holds
+    /// `session_lock_for(<lid_addr>.<device_id>)`. `async_lock::Mutex` is
+    /// not reentrant, so re-acquiring that lock in the migration loop
+    /// deadlocks. `skip_lid_lock_for_device` tells us which LID device
+    /// lock to skip (caller already serializes it).
+    pub(crate) async fn migrate_signal_sessions_on_lid_discovery_with_held_lock(
+        &self,
+        pn: &str,
+        lid: &str,
+        held_device_id: u16,
+    ) {
+        self.migrate_signal_sessions_on_lid_discovery_inner(pn, lid, Some(held_device_id))
+            .await;
+    }
+
+    async fn migrate_signal_sessions_on_lid_discovery_inner(
+        &self,
+        pn: &str,
+        lid: &str,
+        skip_lid_lock_for_device: Option<u16>,
+    ) {
         use log::{info, warn};
         use wacore::types::jid::JidExt;
 
@@ -361,39 +386,51 @@ impl Client {
             let pn_proto = pn_jid.to_protocol_address();
             let lid_proto = lid_jid.to_protocol_address();
 
-            // Migrate session: take from cache (authoritative), write to cache
+            // Read-modify-write of both PN and LID slots must hold the same
+            // per-address locks that encrypt/decrypt take, otherwise a
+            // concurrent message_encrypt on LID can clobber the migrated
+            // session (or read mid-update state). Acquire in stable
+            // (lexicographic) order to avoid deadlocks with operations that
+            // legitimately hold one and not the other.
+            let skip_lid = skip_lid_lock_for_device == Some(device_id);
+            let pn_key = pn_proto.to_string();
+            let lid_key = lid_proto.to_string();
+            let pn_lock = self.session_lock_for(&pn_key).await;
+            // Lock guards must outlive the read-modify-write window. Storing
+            // them as separately-typed Options keeps stable acquisition order
+            // (PN first when both are taken; LID-only acquisition only when
+            // PN ordering would put LID before PN but the caller still holds
+            // PN — not possible here, so PN-first is unconditional and safe).
+            let (_pn_guard_opt, _lid_guard_opt) = if skip_lid {
+                (Some(pn_lock.lock_arc().await), None)
+            } else {
+                let lid_lock = self.session_lock_for(&lid_key).await;
+                if pn_key <= lid_key {
+                    let pn_g = pn_lock.lock_arc().await;
+                    let lid_g = lid_lock.lock_arc().await;
+                    (Some(pn_g), Some(lid_g))
+                } else {
+                    let lid_g = lid_lock.lock_arc().await;
+                    let pn_g = pn_lock.lock_arc().await;
+                    (Some(pn_g), Some(lid_g))
+                }
+            };
+
+            // PN wins on conflict — mirrors whatsmeow's `MigratePNToLID`
+            // (`ON CONFLICT DO UPDATE SET session=excluded.session`).
             if let Ok(Some(session)) = self
                 .signal_cache
                 .get_session(&pn_proto, backend.as_ref())
                 .await
             {
-                match self
-                    .signal_cache
-                    .has_session(&lid_proto, backend.as_ref())
-                    .await
-                {
-                    Ok(true) => {
-                        self.signal_cache.delete_session(&pn_proto).await;
-                        info!("Deleted stale PN session {} (LID exists)", pn_proto);
-                    }
-                    Ok(false) => {
-                        self.signal_cache.put_session(&lid_proto, session).await;
-                        self.signal_cache.delete_session(&pn_proto).await;
-                        info!("Migrated session {} -> {}", pn_proto, lid_proto);
-                    }
-                    Err(e) => {
-                        // Restore the taken PN session to avoid losing it
-                        self.signal_cache.put_session(&pn_proto, session).await;
-                        log::warn!(
-                            "Skipping session migration {} -> {}: {e}",
-                            pn_proto,
-                            lid_proto
-                        );
-                    }
-                }
+                self.signal_cache.put_session(&lid_proto, session).await;
+                self.signal_cache.delete_session(&pn_proto).await;
+                info!(
+                    "Migrated session {} -> {} (PN wins on conflict)",
+                    pn_proto, lid_proto
+                );
             }
 
-            // Migrate identity: same cache-first pattern
             if let Ok(Some(identity_data)) = self
                 .signal_cache
                 .get_identity(&pn_proto, backend.as_ref())
@@ -846,4 +883,168 @@ mod tests {
             "offline batch must not persist to DB"
         );
     }
+
+    /// Produce a SessionRecord blob with a distinctive remote_registration_id
+    /// so we can tell which side of a migration won by parsing the surviving
+    /// session, not by raw-byte comparison.
+    fn tagged_session_blob(remote_regid: u32) -> Vec<u8> {
+        use wacore::libsignal::protocol::{SessionRecord, SessionState};
+        use waproto::whatsapp::SessionStructure;
+
+        let state = SessionState::from_session_structure(SessionStructure {
+            session_version: Some(3),
+            local_identity_public: None,
+            remote_identity_public: None,
+            root_key: None,
+            previous_counter: Some(0),
+            sender_chain: None,
+            receiver_chains: vec![],
+            pending_pre_key: None,
+            remote_registration_id: Some(remote_regid),
+            local_registration_id: Some(0),
+            alice_base_key: Some(vec![]),
+            needs_refresh: None,
+            pending_key_exchange: None,
+        });
+        SessionRecord::new(state)
+            .serialize()
+            .expect("serialize session record")
+    }
+
+    /// Both PN and LID slots hold a session for the same peer; the
+    /// PN one is the working Double Ratchet state, the LID one was
+    /// built freshly by `process_prekey_bundle` and has no link to
+    /// the peer's outbound chain. Migration must keep the PN blob —
+    /// silently dropping it leaves the linked device pinned to the
+    /// fresh stub forever. Reg-id tags identify which side won.
+    #[tokio::test]
+    async fn migration_preserves_working_session_when_both_namespaces_present() {
+        use wacore::libsignal::protocol::SessionRecord;
+        use wacore::types::jid::JidExt as _;
+
+        let client: Arc<Client> = create_test_client().await;
+        let pn = "5500000000000";
+        let lid = "111111111111111";
+
+        client
+            .add_lid_pn_mapping(lid, pn, LearningSource::PeerPnMessage)
+            .await
+            .unwrap();
+
+        let pn_addr = Jid::pn_device(pn.to_string(), 0).to_protocol_address();
+        let lid_addr = Jid::lid_device(lid.to_string(), 0).to_protocol_address();
+
+        // The working session — what Bob's outbound chain is actually
+        // ratcheted against — lives in the PN slot. Tag it with a
+        // distinctive registration id so post-migration we can prove
+        // the surviving session is the SAME blob.
+        const WORKING_REGID: u32 = 0xDEAD_BEEF;
+        const FRESH_REGID: u32 = 0x0BAD_F00D;
+
+        let backend = client.persistence_manager.backend();
+
+        // Seed both slots through signal_cache so the cache holds Present
+        // entries when migrate runs. Raw backend writes alone leave the
+        // cache cold and migrate's `get_session` then races with whatever
+        // populated Absent markers for unknown peers during test bring-up.
+        client
+            .signal_cache
+            .put_session(
+                &pn_addr,
+                SessionRecord::deserialize(&tagged_session_blob(WORKING_REGID))
+                    .expect("seed PN blob deserializes"),
+            )
+            .await;
+        client
+            .signal_cache
+            .put_session(
+                &lid_addr,
+                SessionRecord::deserialize(&tagged_session_blob(FRESH_REGID))
+                    .expect("seed LID blob deserializes"),
+            )
+            .await;
+        client.signal_cache.flush(backend.as_ref()).await.unwrap();
+
+        client
+            .migrate_signal_sessions_on_lid_discovery(pn, lid)
+            .await;
+
+        // PN must be drained — future loads route to LID once the
+        // mapping is known.
+        assert!(
+            backend
+                .get_session(pn_addr.as_str())
+                .await
+                .unwrap()
+                .is_none(),
+            "PN address must be cleared post-migration"
+        );
+
+        let surviving_bytes = backend
+            .get_session(lid_addr.as_str())
+            .await
+            .unwrap()
+            .expect("LID slot must have a session after migration");
+        let record = SessionRecord::deserialize(&surviving_bytes)
+            .expect("surviving session blob must parse");
+        let surviving_regid = record
+            .remote_registration_id()
+            .expect("surviving session must expose its remote reg id");
+
+        assert_eq!(
+            surviving_regid, WORKING_REGID,
+            "LID slot held the FRESH (regid={:#x}) blob — that's the prod \
+             deadlock: the working PN session ({:#x}) got discarded by the \
+             'both exist' branch, leaving us pinned to a session that has no \
+             link to the peer's outbound chain.",
+            surviving_regid, WORKING_REGID
+        );
+    }
+
+    /// Migration must hold the same per-address session locks that
+    /// encrypt/decrypt take. Otherwise a concurrent `message_encrypt`
+    /// on the LID slot can clobber the just-migrated session (or read
+    /// mid-update state). Externally hold the LID lock, kick off
+    /// migration, and assert it blocks until the lock is released.
+    #[tokio::test]
+    async fn migration_blocks_on_per_address_session_lock() {
+        use std::time::Duration;
+        use wacore::types::jid::JidExt as _;
+
+        let client: Arc<Client> = create_test_client().await;
+        let pn = "5500000000000";
+        let lid = "111111111111111";
+        client
+            .add_lid_pn_mapping(lid, pn, LearningSource::PeerPnMessage)
+            .await
+            .unwrap();
+
+        let lid_addr = Jid::lid_device(lid.to_string(), 0).to_protocol_address();
+        let lid_lock = client.session_lock_for(lid_addr.as_str()).await;
+        let held = lid_lock.lock().await;
+
+        let migrate_client = client.clone();
+        let pn_s = pn.to_string();
+        let lid_s = lid.to_string();
+        let mut handle = tokio::spawn(async move {
+            migrate_client
+                .migrate_signal_sessions_on_lid_discovery(&pn_s, &lid_s)
+                .await;
+        });
+
+        let blocked = tokio::time::timeout(Duration::from_millis(200), &mut handle).await;
+        assert!(
+            blocked.is_err(),
+            "migration must block while another holder owns the LID address \
+             session lock — otherwise concurrent encrypt/decrypt races"
+        );
+
+        // Release the lock; migration should now complete so the spawned task
+        // doesn't outlive the test (and contaminate parallel test state).
+        drop(held);
+        tokio::time::timeout(Duration::from_secs(5), handle)
+            .await
+            .expect("migration must complete once the lock is released")
+            .expect("migration task must not panic");
+    }
 }