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

@@ -361,39 +361,43 @@ 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 pn_key = pn_proto.to_string();
+            let lid_key = lid_proto.to_string();
+            let (first_key, second_key) = if pn_key <= lid_key {
+                (pn_key.clone(), lid_key.clone())
+            } else {
+                (lid_key.clone(), pn_key.clone())
+            };
+            let first_lock = self.session_lock_for(&first_key).await;
+            let _first_guard = first_lock.lock().await;
+            let _second_guard = if first_key == second_key {
+                None
+            } else {
+                let second_lock = self.session_lock_for(&second_key).await;
+                Some(second_lock.lock_arc().await)
+            };
+
+            // 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 +850,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");
+    }
 }

src/message.rs

@@ -1075,13 +1075,27 @@ impl Client {
                         e,
                         SignalProtocolError::BadMac(_) | SignalProtocolError::InvalidMessage(_, _)
                     ) {
+                        // whatsmeow migrates PN sessions before decrypt; a fresh
+                        // LID record can otherwise shadow the sender's PN ratchet.
+                        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;
+                        }
+
                         // WAWebMsgProcessingDecryptionHandler classifies both as
-                        // SignalRetryable -> sendRetryReceipt only, no session ops.
-                        // When the sender resends as pkmsg, process_prekey_bundle
-                        // calls promote_state on the existing record, archiving
-                        // current into previous_sessions[0]. That archived state
-                        // is the only fallback for in-flight messages still on
-                        // the old ratchet (see decrypt_message_with_record).
+                        // SignalRetryable -> sendRetryReceipt only, with no delete.
                         let (reason, label) = if matches!(e, SignalProtocolError::BadMac(_)) {
                             (RetryReason::BadMac, "BadMac")
                         } else {
@@ -2204,7 +2218,7 @@ mod tests {
         IdentityKeyStore as SigIdentityKeyStore, SignalProtocolError,
     };
 
-    #[derive(Default)]
+    #[derive(Default, Clone)]
     struct MemSessionStore(HashMap<ProtocolAddress, SessionRecord>);
 
     #[async_trait]
@@ -2228,6 +2242,7 @@ mod tests {
         }
     }
 
+    #[derive(Clone)]
     struct MemIdentityStore {
         kp: IdentityKeyPair,
         reg_id: u32,
@@ -2270,6 +2285,7 @@ mod tests {
         }
     }
 
+    #[derive(Clone)]
     struct AlicePeer {
         jid: Jid,
         address: ProtocolAddress,
@@ -2439,6 +2455,74 @@ mod tests {
         (success, dups, dispatched, 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(&bob_addr, b"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(&bob_addr, b"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(&bob_addr, b"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");
+    }
+
     /// 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