wacore 0.7.0

use crate::libsignal::protocol::{IdentityKey, PreKeyBundle, PreKeyId, PublicKey, SignedPreKeyId};
use std::collections::HashMap;
use wacore_binary::CompactString;
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.
///
/// Matches WA Web's `validateLocalKeyBundle` hash computation:
/// SHA-1(identity_pub_key || signed_prekey_pub || signed_prekey_signature || prekey_pub_1 || ...)
pub fn compute_key_bundle_digest(
    identity_pub_key: &[u8],
    signed_prekey_pub: &[u8],
    signed_prekey_signature: &[u8],
    prekey_pubkeys: &[&[u8]],
) -> Vec<u8> {
    use sha1::Digest;
    let mut hasher = sha1::Sha1::new();
    hasher.update(identity_pub_key);
    hasher.update(signed_prekey_pub);
    hasher.update(signed_prekey_signature);
    for pk in prekey_pubkeys {
        hasher.update(pk);
    }
    hasher.finalize().to_vec()
}

/// Extract the `publicKey` field (tag 2) from a protobuf-encoded
/// PreKeyRecordStructure without a full prost decode.
///
/// Validates the record framing end-to-end: a malformed varint, a truncated
/// length-delimited/fixed field, or an unsupported wire type (e.g. the
/// deprecated group types) yields `None` — the same records a full
/// `PreKeyRecordStructure::decode` rejects. This keeps callers that skip on
/// `None` (the pre-key upload and the digestKey check) from admitting a record
/// this device could not later decode. Uses last-one-wins semantics for a
/// repeated `publicKey` field, per the protobuf spec.
pub fn extract_prekey_public_key(record: &[u8]) -> Option<&[u8]> {
    let mut pos = 0;
    let mut result: Option<&[u8]> = None;
    while pos < record.len() {
        let (tag_byte, consumed) = decode_varint(&record[pos..])?;
        pos += consumed;
        let field_number = (tag_byte >> 3) as u32;
        let wire_type = (tag_byte & 0x7) as u32;
        match wire_type {
            // varint
            0 => {
                let (_, c) = decode_varint(&record[pos..])?;
                pos += c;
            }
            // length-delimited
            2 => {
                let (len, c) = decode_varint(&record[pos..])?;
                pos += c;
                let len = len as usize;
                if pos + len > record.len() {
                    return None;
                }
                if field_number == 2 {
                    result = Some(&record[pos..pos + len]);
                }
                pos += len;
            }
            // fixed64
            1 => {
                if pos + 8 > record.len() {
                    return None;
                }
                pos += 8;
            }
            // fixed32
            5 => {
                if pos + 4 > record.len() {
                    return None;
                }
                pos += 4;
            }
            // Unsupported / invalid wire type (groups, reserved): reject the record.
            _ => return None,
        }
    }
    result
}

fn decode_varint(buf: &[u8]) -> Option<(u64, usize)> {
    let mut result: u64 = 0;
    for (i, &byte) in buf.iter().enumerate().take(10) {
        // The 10th byte (i==9) carries the highest bits; only the low bit
        // is valid payload (64 - 9*7 = 1). Reject if more bits are set.
        if i == 9 && (byte & 0x7F) > 1 {
            return None;
        }
        result |= ((byte & 0x7F) as u64) << (i * 7);
        if byte & 0x80 == 0 {
            return Some((result, i + 1));
        }
    }
    None
}

impl PreKeyUtils {
    pub fn build_fetch_prekeys_request(jids: &[Jid], reason: Option<&str>) -> Node {
        let user_nodes = jids.iter().map(|jid| {
            let mut user_builder = NodeBuilder::new("user").attr("jid", jid);
            if let Some(r) = reason {
                user_builder = user_builder.attr("reason", r);
            }
            user_builder.build()
        });

        NodeBuilder::new("key").children(user_nodes).build()
    }

    pub fn build_upload_prekeys_request<'a>(
        registration_id: u32,
        identity_key_bytes: &[u8],
        signed_pre_key_id: u32,
        signed_pre_key_public_bytes: &[u8],
        signed_pre_key_signature: &[u8],
        pre_keys: impl IntoIterator<Item = (u32, &'a [u8])>,
    ) -> Vec<Node> {
        let pre_keys = pre_keys.into_iter();
        let (lower, upper) = pre_keys.size_hint();
        let mut pre_key_nodes = Vec::with_capacity(upper.unwrap_or(lower));
        for (pre_key_id, public_bytes) in pre_keys {
            let node = NodeBuilder::new("key")
                .children([
                    NodeBuilder::new("id")
                        .bytes(pre_key_id.to_be_bytes()[1..].to_vec())
                        .build(),
                    NodeBuilder::new("value").bytes(public_bytes).build(),
                ])
                .build();
            pre_key_nodes.push(node);
        }

        let signed_pre_key_node = NodeBuilder::new("skey")
            .children([
                NodeBuilder::new("id")
                    .bytes(signed_pre_key_id.to_be_bytes()[1..].to_vec())
                    .build(),
                NodeBuilder::new("value")
                    .bytes(signed_pre_key_public_bytes)
                    .build(),
                NodeBuilder::new("signature")
                    .bytes(signed_pre_key_signature)
                    .build(),
            ])
            .build();

        vec![
            NodeBuilder::new("registration")
                .bytes(registration_id.to_be_bytes().to_vec())
                .build(),
            NodeBuilder::new("type").bytes(vec![5u8]).build(),
            NodeBuilder::new("identity")
                .bytes(identity_key_bytes)
                .build(),
            NodeBuilder::new("list").children(pre_key_nodes).build(),
            signed_pre_key_node,
        ]
    }

    /// Parse a prekey-fetch response into bundles.
    ///
    /// `account_identities` maps a (normalized) companion JID to its account's
    /// primary (device 0) identity key, used as the ADV `account_signature_key`
    /// fallback when the server omits that field from `<device-identity>` (the
    /// normal case for a contact's companion device). The caller pre-loads these
    /// from the identity store; pass an empty map when no fallback is available.
    pub fn parse_prekeys_response(
        resp_node: &NodeRef<'_>,
        account_identities: &HashMap<Jid, [u8; 32]>,
    ) -> 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 outcome = PreKeyFetchOutcome {
            bundles: HashMap::with_capacity(children.len()),
            rejected: Vec::new(),
        };
        for user_node_ref in children {
            if user_node_ref.tag != "user" {
                continue;
            }
            // A `<user>` whose jid is missing or unparseable names no device, and
            // the default would name device 0 of an empty user -- an address a
            // rejection must never be recorded against.
            let Some(named) = user_node_ref.attrs().optional_jid("jid") else {
                log::warn!("prekey response carried a <user> with no usable jid; skipping");
                continue;
            };
            let mut jid = named;
            if jid.device == 0
                && matches!(
                    jid.server,
                    wacore_binary::Server::Pn | wacore_binary::Server::Lid
                )
                && let Some((user_base, device_str)) = jid.user.split_once(':')
                && let Ok(device) = device_str.parse::<u16>()
            {
                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") {
                // Narrowed, not truncated: `65942 as u16` is 406, so a cast
                // would let an out-of-range code arrive as the one value that
                // makes us refresh a device list. A code we cannot read names
                // nothing we can act on, so the entry is dropped rather than
                // guessed at.
                match error_node
                    .attrs()
                    .optional_u64("code")
                    .and_then(|raw| u16::try_from(raw).ok())
                {
                    Some(code) => {
                        log::debug!("prekey fetch rejected {} with code {code}", jid.observe());
                        outcome.rejected.push(RejectedDevice { jid, code });
                    }
                    None => log::warn!(
                        "prekey fetch rejected {} with an unreadable code; ignoring",
                        jid.observe()
                    ),
                }
                continue;
            }
            let account_identity = account_identities.get(&jid);
            let bundle =
                match Self::node_to_pre_key_bundle_ref(&jid, user_node_ref, account_identity) {
                    Ok(b) => b,
                    Err(e) => {
                        log::warn!("Failed to parse prekey bundle for {}: {}", jid, e);
                        continue;
                    }
                };
            outcome.bundles.insert(jid, bundle);
        }

        Ok(outcome)
    }

    fn node_to_pre_key_bundle_ref(
        jid: &Jid,
        node: &NodeRef<'_>,
        account_identity: Option<&[u8; 32]>,
    ) -> Result<PreKeyBundle, anyhow::Error> {
        use crate::xml::DisplayableNodeRef;
        use wacore_binary::NodeContentRef;

        fn extract_bytes_ref(node: Option<&NodeRef<'_>>) -> Result<Vec<u8>, anyhow::Error> {
            match node.and_then(|n| n.content.as_ref()) {
                Some(NodeContentRef::Bytes(b)) => Ok(b.to_vec()),
                _ => Err(anyhow::anyhow!("Expected bytes in node content")),
            }
        }

        if let Some(error_node) = node.get_optional_child("error") {
            return Err(anyhow::anyhow!(
                "Error getting prekeys: {}",
                DisplayableNodeRef(error_node)
            ));
        }

        let reg_id_bytes = extract_bytes_ref(node.get_optional_child("registration"))?;
        if reg_id_bytes.len() != 4 {
            return Err(anyhow::anyhow!("Invalid registration ID length"));
        }
        let registration_id = u32::from_be_bytes([
            reg_id_bytes[0],
            reg_id_bytes[1],
            reg_id_bytes[2],
            reg_id_bytes[3],
        ]);

        let keys_node = node.get_optional_child("keys").unwrap_or(node);

        let identity_key_bytes = extract_bytes_ref(keys_node.get_optional_child("identity"))?;
        let identity_key_array: [u8; 32] =
            identity_key_bytes.try_into().map_err(|v: Vec<u8>| {
                anyhow::anyhow!("Invalid identity key length: got {}, expected 32", v.len())
            })?;
        let identity_key =
            IdentityKey::new(PublicKey::from_djb_public_key_bytes(&identity_key_array)?);

        let mut pre_key_tuple = None;
        if let Some(pre_key_node) = keys_node.get_optional_child("key")
            && let Some((id, key_bytes)) = Self::node_to_pre_key_ref(pre_key_node)?
        {
            let pre_key_id: PreKeyId = id.into();
            let pre_key_public = PublicKey::from_djb_public_key_bytes(&key_bytes)?;
            pre_key_tuple = Some((pre_key_id, pre_key_public));
        }

        let signed_pre_key_node = keys_node
            .get_optional_child("skey")
            .ok_or(anyhow::anyhow!("Missing signed prekey"))?;
        let (signed_pre_key_id_u32, signed_pre_key_public_bytes, signed_pre_key_signature) =
            Self::node_to_signed_pre_key_ref(signed_pre_key_node)?;

        let signed_pre_key_id: SignedPreKeyId = signed_pre_key_id_u32.into();
        let signed_pre_key_public =
            PublicKey::from_djb_public_key_bytes(&signed_pre_key_public_bytes)?;

        let bundle = PreKeyBundle::new(
            registration_id,
            (jid.device as u32).into(),
            pre_key_tuple,
            signed_pre_key_id,
            signed_pre_key_public,
            signed_pre_key_signature.to_vec(),
            identity_key,
        )?;

        // Companion devices (device != 0) carry a <device-identity> that ADV-binds
        // the fetched identity key to the account, so a relay can't substitute a
        // fabricated identity. Matches WA Web SessionApi.createSignalSession. The
        // account key is the in-blob `account_signature_key` or, when the server
        // omits it (the normal case for a contact's companion), the account's
        // primary identity from the store (`account_identity`). Invalid → reject
        // the bundle; a missing device-identity or an unverifiable-for-lack-of-key
        // chain is logged but not fatal (WA Web throws here, but our store-set is
        // best-effort and a relay could already strip <device-identity> entirely).
        if jid.device != 0 {
            let device_identity = node
                .get_optional_child("device-identity")
                .or_else(|| keys_node.get_optional_child("device-identity"))
                .and_then(|n| match n.content.as_ref() {
                    Some(NodeContentRef::Bytes(b)) => Some(b),
                    _ => None,
                });
            match device_identity {
                Some(di) => match crate::adv::validate_adv_with_identity_key(
                    di,
                    &identity_key_array,
                    account_identity,
                ) {
                    crate::adv::AdvValidation::Valid => {}
                    crate::adv::AdvValidation::Invalid => {
                        return Err(anyhow::anyhow!(
                            "device-identity ADV validation failed for companion {jid}"
                        ));
                    }
                    crate::adv::AdvValidation::NoAccountKey => log::debug!(
                        "prekey bundle for companion {jid} omits account_signature_key and no stored account identity; proceeding without ADV validation"
                    ),
                },
                None => log::warn!(
                    "prekey bundle for companion {jid} omits <device-identity>; proceeding without ADV validation"
                ),
            }
        }

        Ok(bundle)
    }

    fn node_to_pre_key_ref(node: &NodeRef<'_>) -> Result<Option<(u32, [u8; 32])>, anyhow::Error> {
        use wacore_binary::NodeContentRef;

        let id_content = node
            .get_optional_child("id")
            .and_then(|n| n.content.as_ref());

        let id = match id_content {
            Some(NodeContentRef::Bytes(b)) if !b.is_empty() => {
                if b.len() == 3 {
                    Ok(u32::from_be_bytes([0, b[0], b[1], b[2]]))
                } else if let Ok(s) = std::str::from_utf8(b.as_ref()) {
                    let trimmed_s = s.trim();
                    if trimmed_s.is_empty() {
                        Err(anyhow::anyhow!("ID content is only whitespace"))
                    } else {
                        u32::from_str_radix(trimmed_s, 16).map_err(|e| e.into())
                    }
                } else {
                    Err(anyhow::anyhow!("ID is not valid UTF-8 hex or 3-byte int"))
                }
            }
            _ => Err(anyhow::anyhow!("Missing or empty pre-key ID content")),
        };

        let id = match id {
            Ok(val) => val,
            Err(_e) => return Ok(None),
        };

        let value_bytes = node
            .get_optional_child("value")
            .and_then(|n| n.content.as_ref())
            .and_then(|c| {
                if let NodeContentRef::Bytes(b) = c {
                    Some(b.to_vec())
                } else {
                    None
                }
            })
            .ok_or(anyhow::anyhow!("Missing pre-key value"))?;
        if value_bytes.len() != 32 {
            return Err(anyhow::anyhow!("Invalid pre-key value length"));
        }

        let mut value_arr = [0u8; 32];
        value_arr.copy_from_slice(&value_bytes);
        Ok(Some((id, value_arr)))
    }

    fn node_to_signed_pre_key_ref(
        node: &NodeRef<'_>,
    ) -> Result<(u32, [u8; 32], [u8; 64]), anyhow::Error> {
        use wacore_binary::NodeContentRef;

        let (id, public_key_bytes) = match Self::node_to_pre_key_ref(node)? {
            Some((id, key)) => (id, key),
            None => return Err(anyhow::anyhow!("Signed pre-key is missing ID or value")),
        };
        let signature_bytes = node
            .get_optional_child("signature")
            .and_then(|n| n.content.as_ref())
            .and_then(|c| {
                if let NodeContentRef::Bytes(b) = c {
                    Some(b.to_vec())
                } else {
                    None
                }
            })
            .ok_or(anyhow::anyhow!("Missing signed pre-key signature"))?;
        if signature_bytes.len() != 64 {
            return Err(anyhow::anyhow!("Invalid signature length"));
        }

        let mut sig_arr = [0u8; 64];
        sig_arr.copy_from_slice(&signature_bytes);
        Ok((id, public_key_bytes, sig_arr))
    }
}

#[cfg(test)]
#[allow(clippy::disallowed_methods)]
mod tests {
    use super::*;
    use crate::iq::prekeys::PreKeyBundleUserNode;
    use crate::libsignal::protocol::{IdentityKeyPair, KeyPair};
    use crate::protocol::ProtocolNode;

    use wacore_binary::NodeValue;

    #[test]
    fn extract_prekey_public_key_matches_full_decode_validation() {
        use buffa::Message;
        let public_key = vec![0x05u8; 33];
        let record = waproto::whatsapp::PreKeyRecordStructure {
            id: Some(1),
            public_key: Some(public_key.clone()),
            private_key: Some(vec![0x09u8; 32]),
        }
        .encode_to_vec();

        // Well-formed record: the extractor returns the public key.
        assert_eq!(
            extract_prekey_public_key(&record),
            Some(public_key.as_slice())
        );

        // Truncating into the trailing private_key field leaves a valid publicKey
        // earlier in the buffer but a malformed tail. A full prost decode rejects
        // it; the extractor must agree (return None) so the upload path never ships
        // a record the consume path's full decode would later reject.
        let truncated = &record[..record.len() - 1];
        assert!(waproto::whatsapp::PreKeyRecordStructure::decode_from_slice(truncated).is_err());
        assert_eq!(extract_prekey_public_key(truncated), None);

        // A malformed varint in a trailing field (10 continuation bytes that never
        // terminate): a full decode errors on the bad varint, and the extractor
        // agrees even though a valid publicKey precedes it.
        let mut bad_varint = record.clone();
        bad_varint.push(0x08); // field 1, wire type 0 (varint)
        bad_varint.extend_from_slice(&[0xFF; 10]);
        assert!(
            waproto::whatsapp::PreKeyRecordStructure::decode_from_slice(&bad_varint[..]).is_err()
        );
        assert_eq!(extract_prekey_public_key(&bad_varint), None);

        // An unsupported wire type (3 = start group): prost rejects the dangling
        // group, and the extractor rejects every wire type it cannot frame.
        let mut bad_wire = record.clone();
        bad_wire.push(0x0B); // field 1, wire type 3 (start group)
        assert!(
            waproto::whatsapp::PreKeyRecordStructure::decode_from_slice(&bad_wire[..]).is_err()
        );
        assert_eq!(extract_prekey_public_key(&bad_wire), None);
    }

    fn create_mock_bundle(device_id: u32) -> PreKeyBundle {
        let mut rng = rand::make_rng::<rand::rngs::StdRng>();
        let identity_pair = IdentityKeyPair::generate(&mut rng);
        let signed_prekey_pair = KeyPair::generate(&mut rng);
        let prekey_pair = KeyPair::generate(&mut rng);

        PreKeyBundle::new(
            1,
            device_id.into(),
            Some((1u32.into(), prekey_pair.public_key)),
            2u32.into(),
            signed_prekey_pair.public_key,
            vec![0u8; 64],
            *identity_pair.identity_key(),
        )
        .expect("Failed to create PreKeyBundle")
    }

    #[test]
    fn test_parse_prekeys_response_normalizes_lid_device_jid() {
        let base_jid = Jid::lid_device("100000012345678", 33);
        let bundle = create_mock_bundle(33);
        let mut user_node = PreKeyBundleUserNode::from_bundle(base_jid.clone(), &bundle, None)
            .expect("build bundle node")
            .into_node();

        let raw_jid = Jid {
            user: "100000012345678:33".into(),
            server: wacore_binary::Server::Lid,
            agent: 1,
            device: 0,
            integrator: 0,
        };
        user_node
            .attrs
            .insert("jid".to_string(), NodeValue::Jid(raw_jid.clone()));

        let response = NodeBuilder::new("iq")
            .children([NodeBuilder::new("list").children([user_node]).build()])
            .build();

        let bundles = PreKeyUtils::parse_prekeys_response(&response.as_node_ref(), &HashMap::new())
            .expect("parse bundles")
            .bundles;
        assert!(bundles.contains_key(&base_jid));
        assert!(!bundles.contains_key(&raw_jid));

        let parsed_jid = bundles.keys().next().expect("parsed jid");
        assert_eq!(parsed_jid.user, base_jid.user);
        assert_eq!(parsed_jid.device, base_jid.device);
        // The server's stray agent byte rides along on the parsed key rather than
        // being scrubbed off it. It is inert on a LID, so it never reaches
        // identity — which is what the two `contains_key` assertions above prove,
        // and what the scrubbing used to be needed for.
        assert_eq!(parsed_jid.agent, 1);
    }

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

    /// `65942 as u16` is exactly 406, so a truncating cast turns an
    /// out-of-range code into the one value that makes the caller refresh a
    /// device list. The rejection is dropped instead of being invented.
    #[test]
    fn an_out_of_range_error_code_is_not_narrowed_into_a_406() {
        let gone_jid = Jid::lid_device("100000087654321", 2);
        assert_eq!(65942_u64 as u16, 406, "the cast this guards against");

        let gone = NodeBuilder::new("user")
            .attr("jid", NodeValue::Jid(gone_jid.clone()))
            .children([NodeBuilder::new("error").attr("code", "65942").build()])
            .build();
        let response = NodeBuilder::new("iq")
            .children([NodeBuilder::new("list").children([gone]).build()])
            .build();

        let outcome = PreKeyUtils::parse_prekeys_response(&response.as_node_ref(), &HashMap::new())
            .expect("parse");

        assert!(
            outcome.rejected.is_empty(),
            "an unreadable code must not be recorded as a rejection"
        );
        assert!(outcome.bundles.is_empty());
    }

    /// A `<user>` with no usable jid names no device. Defaulting it would record
    /// the rejection against device 0 of an empty user and refresh whatever
    /// that resolves to.
    #[test]
    fn a_user_without_a_usable_jid_is_skipped_rather_than_defaulted() {
        for attrs in [None, Some("not a jid")] {
            let mut builder = NodeBuilder::new("user");
            if let Some(raw) = attrs {
                builder = builder.attr("jid", raw);
            }
            let user = builder
                .children([NodeBuilder::new("error").attr("code", "406").build()])
                .build();
            let response = NodeBuilder::new("iq")
                .children([NodeBuilder::new("list").children([user]).build()])
                .build();

            let outcome =
                PreKeyUtils::parse_prekeys_response(&response.as_node_ref(), &HashMap::new())
                    .expect("parse");

            assert!(
                outcome.rejected.is_empty(),
                "a rejection with no device to name must not be recorded ({attrs:?})"
            );
        }
    }

    /// The code travels intact for values that do fit, since it is what decides
    /// whether the caller acts at all.
    #[test]
    fn a_rejection_keeps_the_code_the_server_sent() {
        for code in [400_u16, 406, 503] {
            let jid = Jid::lid_device("100000011112222", 3);
            let user = NodeBuilder::new("user")
                .attr("jid", NodeValue::Jid(jid.clone()))
                .children([NodeBuilder::new("error")
                    .attr("code", code.to_string())
                    .build()])
                .build();
            let response = NodeBuilder::new("iq")
                .children([NodeBuilder::new("list").children([user]).build()])
                .build();

            let outcome =
                PreKeyUtils::parse_prekeys_response(&response.as_node_ref(), &HashMap::new())
                    .expect("parse");

            assert_eq!(outcome.rejected.len(), 1);
            assert_eq!(outcome.rejected[0].code, code);
        }
    }

    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)
            .expect("build bundle node")
            .into_node();
        let response = NodeBuilder::new("iq")
            .children([NodeBuilder::new("list").children([user_node]).build()])
            .build();
        PreKeyUtils::parse_prekeys_response(&response.as_node_ref(), &HashMap::new())
            .expect("parse bundles")
            .bundles
    }

    #[test]
    fn parse_skips_companion_bundle_with_invalid_device_identity() {
        let companion = Jid::lid_device("100000012345678", 33);
        let bundles = parse_one(companion.clone(), Some(vec![0xDE, 0xAD, 0xBE, 0xEF]));
        assert!(
            !bundles.contains_key(&companion),
            "companion bundle with an unverifiable device-identity must be dropped"
        );
    }

    // A trimmed device-identity (no account_signature_key) with no stored
    // account-identity fallback is unverifiable, not forged: the bundle must be
    // KEPT (proceed without ADV), otherwise the companion stops receiving. This
    // is the regression #772 introduced (it dropped the bundle outright).
    #[test]
    fn parse_keeps_companion_bundle_when_account_key_omitted_and_no_fallback() {
        use crate::libsignal::protocol::KeyPair;
        let mut rng = rand::make_rng::<rand::rngs::StdRng>();
        let account = KeyPair::generate(&mut rng);
        let device = KeyPair::generate(&mut rng);
        let companion = Jid::lid_device("100000012345678", 33);
        // device-identity without account_signature_key, signatures over the
        // bundle's identity key (create_mock_bundle's identity is unrelated, but
        // the NoAccountKey path short-circuits before signature checks anyway).
        let di = trimmed_device_identity(&account, &device, b"details");
        let bundles = parse_one(companion.clone(), Some(di));
        assert!(
            bundles.contains_key(&companion),
            "trimmed device-identity with no fallback must be kept, not dropped"
        );
    }

    #[test]
    fn parse_keeps_primary_bundle_ignoring_device_identity() {
        // device 0 is the account's primary: ADV validation does not apply, so a
        // junk device-identity must not get it dropped.
        let primary = Jid::lid_device("100000012345678", 0);
        let bundles = parse_one(primary.clone(), Some(vec![0xDE, 0xAD, 0xBE, 0xEF]));
        assert!(bundles.contains_key(&primary));
    }

    /// Encode an ADVSignedDeviceIdentity without the account_signature_key field.
    fn trimmed_device_identity(account: &KeyPair, device: &KeyPair, details: &[u8]) -> Vec<u8> {
        use buffa::Message;
        let mut rng = rand::make_rng::<rand::rngs::StdRng>();
        let identity = device.public_key.public_key_bytes();
        let account_key = account.public_key.public_key_bytes();
        let account_sig = account
            .private_key
            .calculate_signature(&[&[6u8, 0][..], details, identity].concat(), &mut rng)
            .unwrap()
            .to_vec();
        let device_sig = device
            .private_key
            .calculate_signature(
                &[&[6u8, 1][..], details, identity, account_key].concat(),
                &mut rng,
            )
            .unwrap()
            .to_vec();
        waproto::whatsapp::ADVSignedDeviceIdentity {
            details: Some(details.to_vec()),
            account_signature_key: None,
            account_signature: Some(account_sig),
            device_signature: Some(device_sig),
        }
        .encode_to_vec()
    }
}