fix(prekeys): act on the device a rejected fetch names

src/client/context_impl.rs

@@ -18,13 +18,14 @@ impl SendContextResolver for Client {
         &self,
         jids: &[Jid],
     ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error> {
-        self.fetch_pre_keys(jids, None).await
+        // The fan-out has its own batch-level handling; it wants the bundles only.
+        self.fetch_pre_keys(jids, None).await.map(|o| o.bundles)
     }
 
     async fn fetch_prekeys_for_identity_check(
         &self,
         jids: &[Jid],
-    ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error> {
+    ) -> Result<wacore::prekeys::PreKeyFetchOutcome, anyhow::Error> {
         self.fetch_pre_keys(jids, Some(PreKeyFetchReason::Identity))
             .await
             .map_err(|e| {

src/client/sessions.rs

@@ -323,9 +323,13 @@ impl Client {
 /// `wacore::request::ServerErrorCode` to cross the crate boundary. A downcast
 /// to either one alone silently answers `false` for the other, and the failure
 /// mode of that is the send failing exactly as it did before.
+/// The `<error code>` the server attaches to a device it no longer knows.
+const UNREGISTERED_DEVICE_CODE: u16 = 406;
+
 fn is_device_unregistered(err: &anyhow::Error) -> bool {
     use crate::error::ErrorChainExt;
-    err.server_rejection().is_some_and(|r| r.code == 406)
+    err.server_rejection()
+        .is_some_and(|r| r.code == UNREGISTERED_DEVICE_CODE)
 }
 
 /// The distinct users named by `jids`, in first-seen order.
@@ -457,6 +461,30 @@ impl Client {
             Err(e) => return Err(e),
         };
 
+        // The server named these individually, which is the per-device signal a
+        // batch-wide failure cannot give: refresh exactly their device lists and
+        // leave the rest of the batch alone. The send continues, because the
+        // devices that did come back with a bundle are unaffected and skipping
+        // them would deliver to fewer devices for a reason that only concerns
+        // the named ones.
+        if !prekey_bundles.rejected.is_empty() {
+            let rejected: Vec<Jid> = prekey_bundles
+                .rejected
+                .iter()
+                .filter(|device| device.code == UNREGISTERED_DEVICE_CODE)
+                .map(|device| device.jid.clone())
+                .collect();
+            if !rejected.is_empty() {
+                log::debug!(
+                    "prekey fetch rejected {} of {} device(s) as unregistered; \
+                     refreshing their device lists",
+                    rejected.len(),
+                    jids.len()
+                );
+                self.invalidate_device_caches_for(&rejected).await;
+            }
+        }
+
         let mut adapter = self.signal_adapter().await;
         let mut rng = rand::make_rng::<StdRng>();
 
@@ -465,7 +493,10 @@ impl Client {
         let mut failed_count = 0;
 
         for jid in jids {
-            if let Some(bundle) = prekey_bundles.get(&jid.normalize_for_prekey_bundle()) {
+            if let Some(bundle) = prekey_bundles
+                .bundles
+                .get(&jid.normalize_for_prekey_bundle())
+            {
                 match self
                     .install_prekey_bundle_cached(jid, bundle, &mut adapter, &mut rng)
                     .await

src/prekeys.rs

@@ -13,7 +13,7 @@ use std::sync::atomic::Ordering;
 use wacore::iq::prekeys::{
     DigestKeyBundleSpec, PreKeyCountSpec, PreKeyFetchReason, PreKeyFetchSpec, PreKeyUploadSpec,
 };
-use wacore::libsignal::protocol::{KeyPair, PreKeyBundle, PublicKey};
+use wacore::libsignal::protocol::{KeyPair, PublicKey};
 use wacore::libsignal::store::record_helpers::encode_pre_key_record_to;
 use wacore::store::commands::DeviceCommand;
 use wacore_binary::Jid;
@@ -147,7 +147,7 @@ impl Client {
         &self,
         jids: &[Jid],
         reason: Option<PreKeyFetchReason>,
-    ) -> Result<std::collections::HashMap<Jid, PreKeyBundle>, anyhow::Error> {
+    ) -> Result<wacore::prekeys::PreKeyFetchOutcome, anyhow::Error> {
         let spec = match reason {
             Some(r) => PreKeyFetchSpec::with_reason(jids.to_vec(), r),
             None => PreKeyFetchSpec::new(jids.to_vec()),
@@ -161,13 +161,13 @@ impl Client {
         // identity as the fallback in validateADVwithIdentityKey).
         let spec = spec.with_account_identities(self.collect_account_identities(jids).await);
 
-        let bundles = self.execute(spec).await?;
+        let outcome = self.execute(spec).await?;
 
-        for jid in bundles.keys() {
+        for jid in outcome.bundles.keys() {
             log::debug!("Successfully parsed pre-key bundle for {}", jid.observe());
         }
 
-        Ok(bundles)
+        Ok(outcome)
     }
 
     /// Load, for each companion JID, its account (device 0) identity key from the

wacore/benches/send_receive_benchmark.rs

@@ -441,8 +441,8 @@ impl SendContextResolver for MockResolver {
     async fn fetch_prekeys_for_identity_check(
         &self,
         _: &[Jid],
-    ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error> {
-        Ok(HashMap::new())
+    ) -> Result<wacore::prekeys::PreKeyFetchOutcome, anyhow::Error> {
+        Ok(wacore::prekeys::PreKeyFetchOutcome::default())
     }
     async fn resolve_group_info(
         &self,

wacore/src/client/context.rs

@@ -220,10 +220,13 @@ pub trait SendContextResolver: crate::sync_marker::MaybeSendSync {
         jids: &[Jid],
     ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error>;
 
+    /// Returns the bundles alongside the devices the server rejected by name,
+    /// so a per-device rejection is not flattened into "no bundle" before the
+    /// fan-out can tell the two apart.
     async fn fetch_prekeys_for_identity_check(
         &self,
         jids: &[Jid],
-    ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error>;
+    ) -> Result<crate::prekeys::PreKeyFetchOutcome, anyhow::Error>;
 
     async fn resolve_group_info(&self, jid: &Jid) -> Result<Arc<GroupInfo>, anyhow::Error>;
 

wacore/src/iq/prekeys.rs

@@ -144,7 +144,7 @@ impl PreKeyFetchSpec {
 }
 
 impl IqSpec for PreKeyFetchSpec {
-    type Response = std::collections::HashMap<Jid, PreKeyBundle>;
+    type Response = crate::prekeys::PreKeyFetchOutcome;
 
     fn build_iq(&self) -> InfoQuery<'static> {
         let content = PreKeyUtils::build_fetch_prekeys_request(

wacore/src/prekeys.rs

@@ -5,6 +5,27 @@ use wacore_binary::Jid;
 use wacore_binary::builder::NodeBuilder;
 use wacore_binary::{Node, NodeRef};
 
+/// A device the server refused to hand a bundle for, named individually.
+#[derive(Debug, Clone)]
+pub struct RejectedDevice {
+    pub jid: Jid,
+    /// The `<error code>` the server attached; `406` means unregistered.
+    pub code: u16,
+}
+
+/// What one prekey fetch produced: the bundles it did return, and the devices
+/// it named as rejected.
+///
+/// Kept apart because they call for opposite responses. A missing bundle is
+/// ambiguous and the device is simply skipped this round; a rejected one is the
+/// server telling us which cached device is gone, which is the only per-device
+/// signal a batch-wide failure cannot give.
+#[derive(Default)]
+pub struct PreKeyFetchOutcome {
+    pub bundles: HashMap<Jid, PreKeyBundle>,
+    pub rejected: Vec<RejectedDevice>,
+}
+
 pub struct PreKeyUtils;
 
 /// Compute SHA-1 digest of a key bundle for validation against server.
@@ -175,13 +196,16 @@ impl PreKeyUtils {
     pub fn parse_prekeys_response(
         resp_node: &NodeRef<'_>,
         account_identities: &HashMap<Jid, [u8; 32]>,
-    ) -> Result<HashMap<Jid, PreKeyBundle>, anyhow::Error> {
+    ) -> Result<PreKeyFetchOutcome, anyhow::Error> {
         let list_node = resp_node
             .get_optional_child("list")
             .ok_or_else(|| anyhow::anyhow!("<list> not found in pre-key response"))?;
 
         let children = list_node.children().unwrap_or_default();
-        let mut bundles = HashMap::with_capacity(children.len());
+        let mut outcome = PreKeyFetchOutcome {
+            bundles: HashMap::with_capacity(children.len()),
+            rejected: Vec::new(),
+        };
         for user_node_ref in children {
             if user_node_ref.tag != "user" {
                 continue;
@@ -201,6 +225,19 @@ impl PreKeyUtils {
                 jid.user = CompactString::from(user_base);
                 jid.device = device;
             }
+            // A rejected device answers with an `<error>` in place of its key
+            // material, which is how the server names the one device a
+            // batch-wide failure cannot. Distinguished from a malformed bundle
+            // because only this one means "this device is gone": the caller
+            // refreshes that device list, and the rest of the batch is
+            // unaffected. Mirrors WA Web's `FetchKeyBundlesUserError` arm,
+            // which collects per-item errors alongside the bundles.
+            if let Some(error_node) = user_node_ref.get_optional_child("error") {
+                let code = error_node.attrs().optional_u64("code").unwrap_or(0) as u16;
+                log::debug!("prekey fetch rejected {} with code {code}", jid.observe());
+                outcome.rejected.push(RejectedDevice { jid, code });
+                continue;
+            }
             let account_identity = account_identities.get(&jid);
             let bundle =
                 match Self::node_to_pre_key_bundle_ref(&jid, user_node_ref, account_identity) {
@@ -210,10 +247,10 @@ impl PreKeyUtils {
                         continue;
                     }
                 };
-            bundles.insert(jid, bundle);
+            outcome.bundles.insert(jid, bundle);
         }
 
-        Ok(bundles)
+        Ok(outcome)
     }
 
     fn node_to_pre_key_bundle_ref(
@@ -508,7 +545,8 @@ mod tests {
             .build();
 
         let bundles = PreKeyUtils::parse_prekeys_response(&response.as_node_ref(), &HashMap::new())
-            .expect("parse bundles");
+            .expect("parse bundles")
+            .bundles;
         assert!(bundles.contains_key(&base_jid));
         assert!(!bundles.contains_key(&raw_jid));
 
@@ -518,6 +556,53 @@ mod tests {
         assert_eq!(parsed_jid.agent, 0);
     }
 
+    /// The server names a rejected device inside its own `<user>`, which is the
+    /// per-device signal a batch-wide `406` cannot give. The rest of the batch
+    /// must come back intact, or one absent device would cost the send every
+    /// other device in it.
+    #[test]
+    fn a_rejected_user_is_named_without_costing_the_rest_of_the_batch() {
+        let good_jid = Jid::lid_device("100000012345678", 1);
+        let gone_jid = Jid::lid_device("100000087654321", 2);
+
+        let good =
+            PreKeyBundleUserNode::from_bundle(good_jid.clone(), &create_mock_bundle(1), None)
+                .expect("build bundle node")
+                .into_node();
+
+        // What the server sends in place of key material for a device it no
+        // longer knows: an <error> where the bundle would be.
+        let gone = NodeBuilder::new("user")
+            .attr("jid", NodeValue::Jid(gone_jid.clone()))
+            .children([NodeBuilder::new("error")
+                .attr("code", "406")
+                .attr("text", "not-acceptable")
+                .build()])
+            .build();
+
+        let response = NodeBuilder::new("iq")
+            .children([NodeBuilder::new("list").children([good, gone]).build()])
+            .build();
+
+        let outcome = PreKeyUtils::parse_prekeys_response(&response.as_node_ref(), &HashMap::new())
+            .expect("a rejected user must not fail the whole parse");
+
+        assert!(
+            outcome.bundles.contains_key(&good_jid),
+            "the healthy device still gets its bundle"
+        );
+        assert!(
+            !outcome.bundles.contains_key(&gone_jid),
+            "the rejected device has no bundle to give"
+        );
+        assert_eq!(outcome.rejected.len(), 1);
+        assert_eq!(outcome.rejected[0].jid, gone_jid);
+        assert_eq!(
+            outcome.rejected[0].code, 406,
+            "the code is what separates an unregistered device from another refusal"
+        );
+    }
+
     fn parse_one(jid: Jid, device_identity: Option<Vec<u8>>) -> HashMap<Jid, PreKeyBundle> {
         let bundle = create_mock_bundle(jid.device as u32);
         let user_node = PreKeyBundleUserNode::from_bundle(jid, &bundle, device_identity)
@@ -528,6 +613,7 @@ mod tests {
             .build();
         PreKeyUtils::parse_prekeys_response(&response.as_node_ref(), &HashMap::new())
             .expect("parse bundles")
+            .bundles
     }
 
     #[test]

wacore/src/send/encrypt.rs

@@ -77,8 +77,12 @@ pub struct SignalStores<'a> {
 /// 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.
+/// The `<error code>` the server attaches to a device it no longer knows.
+pub(crate) const UNREGISTERED_DEVICE_CODE: u16 = 406;
+
 pub(crate) fn is_device_unregistered_error(err: &anyhow::Error) -> bool {
-    crate::request::ServerErrorCode::from_anyhow(err).is_some_and(|e| e.code == 406)
+    crate::request::ServerErrorCode::from_anyhow(err)
+        .is_some_and(|e| e.code == UNREGISTERED_DEVICE_CODE)
 }
 
 pub struct EncryptResult {
@@ -565,15 +569,33 @@ pub async fn ensure_sessions_for_devices(
             .iter()
             .map(|&i| devices[i].clone())
             .collect();
-        // 406 on this batch is all-or-nothing — per-device retries just wasted
+        // A batch-wide 406 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.
+        //
+        // A per-device rejection is the better-informed case: the server names
+        // the device in its own `<user>`, so it sets the same flag without
+        // condemning the rest of the batch.
         let prekey_bundles = match resolver
             .fetch_prekeys_for_identity_check(&jids_for_fetch)
             .await
         {
-            Ok(bundles) => bundles,
+            Ok(outcome) => {
+                if outcome
+                    .rejected
+                    .iter()
+                    .any(|device| device.code == UNREGISTERED_DEVICE_CODE)
+                {
+                    log::debug!(
+                        "prekey fetch rejected {} of {} device(s) by name",
+                        outcome.rejected.len(),
+                        jids_for_fetch.len()
+                    );
+                    had_406 = true;
+                }
+                outcome.bundles
+            }
             Err(e) if is_device_unregistered_error(&e) => {
                 // No server prekeys for these devices this round; the next send
                 // re-fetches. Debug, not warn — a batch 406 would otherwise flood the log.