libsoliton 0.1.3

//! LO-KEX session initiation and reception (§5).
//!
//! KEM-based key agreement for establishing DM sessions.
//! Produces a root key and chain key for initializing LO-Ratchet.

use crate::constants;
use crate::error::{Error, Result};
use crate::identity::{self, HybridSignature, IdentityPublicKey, IdentitySecretKey};
use crate::primitives::{hkdf, xwing};
use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};

/// A pre-key bundle published to the user's home DM relay (§5.3).
pub struct PreKeyBundle {
    /// Identity public key (3200 bytes).
    pub ik_pub: IdentityPublicKey,
    /// Crypto version string (e.g. "lo-crypto-v1").
    pub crypto_version: String,
    /// Signed pre-key public key (X-Wing, 1216 bytes).
    pub spk_pub: xwing::PublicKey,
    /// Signed pre-key ID.
    pub spk_id: u32,
    /// Hybrid signature over the raw SPK public key bytes.
    pub spk_sig: HybridSignature,
    /// Optional one-time pre-key public key (X-Wing, 1216 bytes).
    pub opk_pub: Option<xwing::PublicKey>,
    /// Optional one-time pre-key ID.
    pub opk_id: Option<u32>,
}

/// A pre-key bundle that has passed verification (§5.4 Step 1).
///
/// Produced by [`verify_bundle`] and required by [`initiate_session`].
/// This newtype enforces at the type level that bundle verification
/// cannot be skipped before session initiation.
pub struct VerifiedBundle(PreKeyBundle);

impl std::ops::Deref for VerifiedBundle {
    type Target = PreKeyBundle;
    fn deref(&self) -> &PreKeyBundle {
        &self.0
    }
}

/// Session init message sent from Alice to Bob (§5.4 Step 5).
pub struct SessionInit {
    /// Crypto version string.
    pub crypto_version: String,
    /// SHA3-256(Alice.IK_pub) — 32 bytes raw.
    pub sender_ik_fingerprint: [u8; 32],
    /// SHA3-256(Bob.IK_pub) — 32 bytes raw. Binds this session init to a
    /// specific recipient; covered by Alice's HybridSig (§5.4 Step 6).
    pub recipient_ik_fingerprint: [u8; 32],
    /// Alice's ephemeral X-Wing public key (becomes initial ratchet key).
    pub sender_ek: xwing::PublicKey,
    /// X-Wing ciphertext encapsulated to Bob's IK.
    pub ct_ik: xwing::Ciphertext,
    /// X-Wing ciphertext encapsulated to Bob's SPK.
    pub ct_spk: xwing::Ciphertext,
    /// Bob's signed pre-key ID.
    pub spk_id: u32,
    /// X-Wing ciphertext encapsulated to Bob's OPK (if available).
    pub ct_opk: Option<xwing::Ciphertext>,
    /// Bob's one-time pre-key ID (if used).
    pub opk_id: Option<u32>,
}

/// Result of session initiation (Alice's side).
///
/// Secret key fields (`root_key`, `initial_chain_key`, `ek_sk`) are private
/// to enforce single-use semantics. `ZeroizeOnDrop` generates a custom `Drop`
/// impl, which prevents partial moves out of the struct — without `take_`
/// methods, callers would need to `.clone()` secret fields (since
/// `Zeroizing<[u8; 32]>` is `Clone`), defeating the ownership-transfer
/// guarantee documented on `encrypt_first_message`.
///
/// Use [`take_root_key`](Self::take_root_key) and
/// [`take_initial_chain_key`](Self::take_initial_chain_key) to extract keys.
/// Each replaces the internal value with zeros, preventing accidental reuse.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct InitiatedSession {
    /// The session init message to send to Bob.
    #[zeroize(skip)]
    pub session_init: SessionInit,
    /// Root key for LO-Ratchet initialization.
    root_key: Zeroizing<[u8; 32]>,
    /// Initial chain key — must be consumed via `take_initial_chain_key()`
    /// and passed to `encrypt_first_message`.
    initial_chain_key: Zeroizing<[u8; 32]>,
    /// Alice's ephemeral X-Wing public key (EK), becomes send_ratchet in ratchet.
    #[zeroize(skip)]
    pub ek_pk: xwing::PublicKey,
    /// Alice's ephemeral X-Wing secret key. Implements `ZeroizeOnDrop`.
    #[zeroize(skip)]
    ek_sk: xwing::SecretKey,
    /// Alice's hybrid signature (Ed25519 + ML-DSA-65) over the encoded SessionInit.
    ///
    /// Must be sent to Bob alongside the session init fields so Bob can verify
    /// that the session was initiated by the holder of Alice's identity secret key.
    /// Bob passes this to [`receive_session`] as `sender_sig`.
    #[zeroize(skip)]
    pub sender_sig: HybridSignature,
    /// Whether a one-time pre-key (OPK) was used in this session.
    ///
    /// When `false`, forward secrecy between sessions established during the same
    /// SPK rotation period depends solely on Alice's ephemeral key. Applications
    /// should treat this as a signal to prioritize OPK replenishment on the server.
    #[zeroize(skip)]
    pub opk_used: bool,
}

impl InitiatedSession {
    /// Extract the root key, replacing the internal copy with zeros.
    ///
    /// Returns the 32-byte root key for `RatchetState::init_alice`. A second
    /// call returns zeros, which `init_alice` rejects as `InvalidData`.
    pub fn take_root_key(&mut self) -> Zeroizing<[u8; 32]> {
        std::mem::replace(&mut self.root_key, Zeroizing::new([0u8; 32]))
    }

    /// Extract the initial chain key, replacing the internal copy with zeros.
    ///
    /// Pass the returned value to `encrypt_first_message`. A second call
    /// returns zeros — `encrypt_first_message` will derive a zero-based
    /// message key, and the resulting ratchet will fail at `init_alice`
    /// (all-zero epoch key is rejected).
    pub fn take_initial_chain_key(&mut self) -> Zeroizing<[u8; 32]> {
        std::mem::replace(&mut self.initial_chain_key, Zeroizing::new([0u8; 32]))
    }

    /// Reference to Alice's ephemeral X-Wing secret key.
    ///
    /// The secret key is owned by the struct and zeroized on drop. Callers
    /// that need the raw bytes (e.g., CAPI serialization) should copy from
    /// this reference and zeroize their copy when done.
    pub fn ek_sk(&self) -> &xwing::SecretKey {
        &self.ek_sk
    }

    /// Raw pointer to the internal root_key bytes (test-utils only).
    #[cfg(all(feature = "test-utils", debug_assertions))]
    #[deprecated(note = "test-utils only — do not call in production code")]
    pub fn root_key_ptr(&self) -> *const u8 {
        self.root_key.as_ptr()
    }

    /// Raw pointer to the internal initial_chain_key bytes (test-utils only).
    #[cfg(all(feature = "test-utils", debug_assertions))]
    #[deprecated(note = "test-utils only — do not call in production code")]
    pub fn initial_chain_key_ptr(&self) -> *const u8 {
        self.initial_chain_key.as_ptr()
    }
}

/// Result of session reception (Bob's side).
///
/// Secret key fields are private — see [`InitiatedSession`] for rationale.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct ReceivedSession {
    /// Root key for LO-Ratchet initialization.
    root_key: Zeroizing<[u8; 32]>,
    /// Initial chain key — must be consumed via `take_initial_chain_key()`
    /// and passed to `decrypt_first_message`.
    initial_chain_key: Zeroizing<[u8; 32]>,
    /// Alice's ephemeral public key (recv_ratchet_pk for Bob).
    #[zeroize(skip)]
    pub peer_ek: xwing::PublicKey,
}

impl ReceivedSession {
    /// Extract the root key, replacing the internal copy with zeros.
    ///
    /// Returns the 32-byte root key for `RatchetState::init_bob`. A second
    /// call returns zeros, which `init_bob` rejects as `InvalidData`.
    pub fn take_root_key(&mut self) -> Zeroizing<[u8; 32]> {
        std::mem::replace(&mut self.root_key, Zeroizing::new([0u8; 32]))
    }

    /// Extract the initial chain key, replacing the internal copy with zeros.
    ///
    /// Pass the returned value to `decrypt_first_message`. A second call
    /// returns zeros — `decrypt_first_message` will derive a zero-based
    /// message key, and AEAD decryption will fail.
    pub fn take_initial_chain_key(&mut self) -> Zeroizing<[u8; 32]> {
        std::mem::replace(&mut self.initial_chain_key, Zeroizing::new([0u8; 32]))
    }

    /// Raw pointer to the internal root_key bytes (test-utils only).
    #[cfg(all(feature = "test-utils", debug_assertions))]
    #[deprecated(note = "test-utils only — do not call in production code")]
    pub fn root_key_ptr(&self) -> *const u8 {
        self.root_key.as_ptr()
    }

    /// Raw pointer to the internal initial_chain_key bytes (test-utils only).
    #[cfg(all(feature = "test-utils", debug_assertions))]
    #[deprecated(note = "test-utils only — do not call in production code")]
    pub fn initial_chain_key_ptr(&self) -> *const u8 {
        self.initial_chain_key.as_ptr()
    }
}

#[must_use = "bundle verification result must be checked"]
/// Verify a pre-key bundle (§5.4 Step 1).
///
/// Checks that:
/// 1. The bundle's IK matches the known IK for this peer.
/// 2. The SPK signature is valid (hybrid verify over raw SPK bytes).
/// 3. The crypto version is supported.
///
/// On success, returns a [`VerifiedBundle`] that can be passed to
/// [`initiate_session`]. This enforces at the type level that
/// verification cannot be skipped.
///
/// # Errors
///
/// - [`Error::InvalidData`] — OPK fields partially present (structural violation).
/// - [`Error::BundleVerificationFailed`] — IK mismatch, unsupported crypto version,
///   or invalid SPK signature. All non-structural failures return this single error
///   to prevent an attacker from iteratively probing which check failed.
pub fn verify_bundle(bundle: PreKeyBundle, known_ik: &IdentityPublicKey) -> Result<VerifiedBundle> {
    // OPK fields must be both present or both absent — partial OPK would
    // produce a SessionInit that the receiver cannot decode.
    // Cheap structural check first — rejects malformed bundles before IK comparison.
    if bundle.opk_pub.is_some() != bundle.opk_id.is_some() {
        return Err(Error::InvalidData);
    }

    // Reject bundles whose IK differs from the trusted key — a mismatch
    // indicates a MITM substitution or a corrupted bundle.
    // Uses the type's PartialEq impl (constant-time via ct_eq) for
    // consistency, even though both values are public.
    if bundle.ik_pub != *known_ik {
        return Err(Error::BundleVerificationFailed);
    }

    // Check crypto version before expensive signature verification.
    // crypto_version is intentionally excluded from the SPK signature (§5.3):
    // downgrade protection relies on this hard-fail rejection (§14.12), not
    // signature binding. Tampering with crypto_version in a relayed bundle
    // causes rejection regardless — the attacker gains no advantage over
    // dropping the bundle entirely. These two design decisions are coupled:
    // removing the hard-fail policy here would require adding crypto_version
    // to the SPK signature to maintain downgrade resistance.
    //
    // Returns BundleVerificationFailed (not UnsupportedCryptoVersion) to
    // avoid exposing which check failed — an attacker iteratively probing
    // bundles could otherwise narrow down which field is wrong.
    if bundle.crypto_version != constants::CRYPTO_VERSION {
        return Err(Error::BundleVerificationFailed);
    }

    // Authenticate the SPK against the IK — a forged or relayed SPK would
    // let an attacker substitute their own key in the X3DH exchange.
    // Domain-separated: "lo-spk-sig-v1" ‖ SPK_pub (matches sign_prekey).
    let mut msg =
        Vec::with_capacity(constants::SPK_SIG_LABEL.len() + bundle.spk_pub.as_bytes().len());
    msg.extend_from_slice(constants::SPK_SIG_LABEL);
    msg.extend_from_slice(bundle.spk_pub.as_bytes());
    // Map signature verification failure to BundleVerificationFailed for
    // consistency — all non-structural failures return the same error.
    identity::hybrid_verify(&bundle.ik_pub, &msg, &bundle.spk_sig)
        .map_err(|_| Error::BundleVerificationFailed)?;

    Ok(VerifiedBundle(bundle))
}

/// Initiate a session from Alice to Bob (§5.4).
///
/// Requires a [`VerifiedBundle`] (produced by [`verify_bundle`]),
/// ensuring the bundle's IK, SPK signature, and crypto version have
/// been validated before any KEM operations are performed.
///
/// Alice's identity secret key (`alice_ik_sk`) is used to sign the encoded
/// `SessionInit`, proving to Bob that the session was initiated by the holder
/// of Alice's identity key (not merely someone who knows her public key).
/// Bob verifies this signature in [`receive_session`].
///
/// Returns an [`InitiatedSession`] containing:
/// - `session_init`: the wire-format session init to send to Bob.
/// - `root_key` / `initial_chain_key`: session keys for ratchet initialization.
/// - `ek_pk` / `ek_sk`: Alice's ephemeral keypair (EK becomes the send ratchet key).
/// - `sender_sig`: Alice's signature over the encoded SessionInit — must be sent to Bob.
/// - `opk_used`: whether a one-time pre-key was consumed.
///
/// # Security
///
/// When `opk_used` is `false`, forward secrecy between sessions established during
/// the same SPK rotation period depends solely on Alice's ephemeral key. Applications
/// should treat `opk_used == false` as a signal to replenish OPKs on the server.
#[must_use = "session establishment result contains key material that must not be discarded"]
pub fn initiate_session(
    alice_ik_pk: &IdentityPublicKey,
    alice_ik_sk: &IdentitySecretKey,
    bundle: &VerifiedBundle,
) -> Result<InitiatedSession> {
    // Ephemeral keypair: per-session randomness ensures forward secrecy —
    // compromise of long-term keys cannot decrypt past sessions.
    let (ek_pk, ek_sk) = xwing::keygen()?;

    // KEM encapsulations produce independent shared secrets that are combined
    // via HKDF below. Each targets a different key to provide layered security.
    //
    // IdentityPublicKey::from_bytes validates LO_PUBLIC_KEY_SIZE (3200 bytes) at
    // construction; VerifiedBundle already holds a validated IdentityPublicKey.
    // xwing_pk() slices &self.0[..XWING_PUBLIC_KEY_SIZE] — always exactly 1216
    // bytes regardless of sub-key content. xwing::PublicKey::from_bytes_unchecked
    // requires exactly XWING_PUBLIC_KEY_SIZE bytes and performs no further
    // structural validation; sub-key structural validity (X25519 point validity,
    // ML-KEM-768 modulus checks) is intentionally deferred to use-site. If the
    // X-Wing sub-key bytes are malformed, encapsulate returns a pseudorandom
    // ciphertext (X-Wing implicit rejection) and the derived session key diverges
    // from Bob's — AeadFailed at decrypt_first_message, not undefined behavior.
    // from_bytes_unchecked is safe here because the size invariant is enforced by
    // xwing_pk().
    let bob_xwing_pk = xwing::PublicKey::from_bytes_unchecked(bundle.ik_pub.xwing_pk().to_vec());
    let (ct_ik, mut ss_ik) = xwing::encapsulate(&bob_xwing_pk)?;

    // SPK encapsulation binds the session to Bob's medium-term signing identity.
    let (ct_spk, mut ss_spk) = xwing::encapsulate(&bundle.spk_pub)?;

    // OPK provides replay protection — each OPK is used once and discarded.
    let (ct_opk, mut ss_opk) = if let Some(ref opk_pub) = bundle.opk_pub {
        let (ct, ss) = xwing::encapsulate(opk_pub)?;
        (Some(ct), Some(ss))
    } else {
        (None, None)
    };

    // HKDF combines independent shared secrets into a single session key,
    // ensuring compromise of any one KEM doesn't break the session.
    // Pre-allocate for max IKM: 3 × 32-byte shared secrets (ik + spk + opk).
    let mut ikm = Zeroizing::new(Vec::with_capacity(96));
    ikm.extend_from_slice(ss_ik.as_bytes());
    ikm.extend_from_slice(ss_spk.as_bytes());
    if let Some(ref ss) = ss_opk {
        ikm.extend_from_slice(ss.as_bytes());
    }

    let info = build_kex_info(alice_ik_pk, &bundle.ik_pub, &ek_pk)?;

    let mut session_key = Zeroizing::new([0u8; 64]);
    // Output size (64) is compile-time constant — InvalidLength structurally unreachable.
    hkdf::hkdf_sha3_256(&constants::HKDF_ZERO_SALT, &ikm, &info, &mut *session_key)?;
    ikm.zeroize();

    // Split the 64-byte HKDF output: first half seeds the ratchet root chain,
    // second half becomes the initial send/receive chain key.
    let mut root_key = Zeroizing::new([0u8; 32]);
    let mut chain_key = Zeroizing::new([0u8; 32]);
    root_key.copy_from_slice(&session_key[..32]);
    chain_key.copy_from_slice(&session_key[32..64]);
    session_key.zeroize(); // earliest safe point — both halves extracted above

    // Shared secrets were copied into ikm but backing buffers persist —
    // Zeroizing<Vec> on ikm doesn't cover these originals.
    ss_ik.0.zeroize();
    ss_spk.0.zeroize();
    if let Some(ref mut ss) = ss_opk {
        ss.0.zeroize();
    }

    let sender_ik_fingerprint = alice_ik_pk.fingerprint_raw();

    let session_init = SessionInit {
        crypto_version: constants::CRYPTO_VERSION.to_string(),
        sender_ik_fingerprint,
        recipient_ik_fingerprint: bundle.ik_pub.fingerprint_raw(),
        sender_ek: ek_pk.clone(),
        ct_ik,
        ct_spk,
        spk_id: bundle.spk_id,
        ct_opk,
        opk_id: bundle.opk_id,
    };

    // Sign the encoded SessionInit to prove Alice possesses sk_IK_A.
    // Without this, anyone who knows pk_IK_A could impersonate Alice by
    // encapsulating to Bob's public keys and placing pk_IK_A in the info field.
    // Domain-separated with INITIATOR_SIG_LABEL to prevent cross-context reuse.
    let si_encoded = encode_session_init(&session_init)?;
    let mut sign_msg = Vec::with_capacity(constants::INITIATOR_SIG_LABEL.len() + si_encoded.len());
    sign_msg.extend_from_slice(constants::INITIATOR_SIG_LABEL);
    sign_msg.extend_from_slice(&si_encoded);
    let sender_sig = identity::hybrid_sign(alice_ik_sk, &sign_msg)?;

    let opk_used = session_init.ct_opk.is_some();
    Ok(InitiatedSession {
        session_init,
        root_key,
        initial_chain_key: chain_key,
        ek_pk,
        ek_sk,
        sender_sig,
        opk_used,
    })
}

/// Receive a session init from Alice (Bob's side, §5.5).
///
/// Bob decapsulates and derives the same session key.
///
/// `sender_sig` is Alice's hybrid signature over the encoded `SessionInit`,
/// produced by [`initiate_session`]. It is verified against `alice_ik_pk`
/// before any KEM operations are performed, proving that the session was
/// initiated by the holder of Alice's identity secret key.
///
/// `spk_sk` must correspond to the SPK identified by `si.spk_id`.
/// `opk_sk` must correspond to the OPK identified by `si.opk_id` (if present).
/// Wrong-key selection is undetectable at this layer — ML-KEM uses implicit
/// rejection (returns a pseudorandom shared secret on decapsulation with the
/// wrong key). The resulting session keys will be incorrect, causing the first
/// ratchet message to fail with `AeadFailed`.
///
/// Returns a [`ReceivedSession`] containing `root_key`, `initial_chain_key`
/// (must be advanced through `decrypt_first_message` before ratchet init),
/// and `peer_ek` (Alice's ephemeral public key).
#[must_use = "session establishment result contains key material that must not be discarded"]
pub fn receive_session(
    bob_ik_pk: &IdentityPublicKey,
    bob_ik_sk: &IdentitySecretKey,
    alice_ik_pk: &IdentityPublicKey,
    si: &SessionInit,
    sender_sig: &HybridSignature,
    spk_sk: &xwing::SecretKey,
    opk_sk: Option<&xwing::SecretKey>,
) -> Result<ReceivedSession> {
    // Reject incompatible protocol versions before expensive KEM operations.
    // Returns UnsupportedCryptoVersion (not collapsed to InvalidData) because
    // crypto_version is a cleartext field: Bob reads it before any KEM or
    // signature operation, so the caller already knows which version they sent.
    // There is no oracle risk — distinguishing a version mismatch from other
    // failures reveals nothing an active attacker cannot derive from the
    // plaintext field they constructed. Compare verify_bundle, which collapses
    // all failures to BundleVerificationFailed because a bundle prober could
    // use distinct error codes to enumerate which field is wrong.
    if si.crypto_version != constants::CRYPTO_VERSION {
        return Err(Error::UnsupportedCryptoVersion);
    }

    // Verify sender fingerprint matches Alice's IK — prevents impersonation.
    // InvalidData (not BundleVerificationFailed) — no pre-key bundle is involved
    // in receive_session; this is a sender-identity mismatch.
    // Constant-time: fingerprints arrive in untrusted wire data before signature
    // verification. Variable-time comparison would let an attacker probe the
    // expected fingerprint byte-by-byte via response timing.
    use subtle::ConstantTimeEq;
    let expected_fp = alice_ik_pk.fingerprint_raw();
    if expected_fp.ct_ne(&si.sender_ik_fingerprint).into() {
        return Err(Error::InvalidData);
    }

    // Verify recipient fingerprint matches Bob's IK — detects misdirected sessions.
    // Constant-time for same reason as sender fingerprint above.
    let expected_recipient_fp = bob_ik_pk.fingerprint_raw();
    if expected_recipient_fp
        .ct_ne(&si.recipient_ik_fingerprint)
        .into()
    {
        return Err(Error::InvalidData);
    }

    // Verify Alice's identity signature over the encoded SessionInit.
    // This proves the session was initiated by the holder of sk_IK_A, not
    // merely someone who knows pk_IK_A (which is public). Without this check,
    // any party could impersonate Alice by encapsulating to Bob's public keys
    // and placing pk_IK_A in the HKDF info field.
    // §5.5 Step 2: encode_session_init enforces OPK co-presence (ct_opk
    // and opk_id must be both present or both absent) — returns InvalidData
    // on mismatch. This runs before hybrid_verify (Step 3), matching spec
    // ordering: structural validation before cryptographic verification.
    let si_encoded = encode_session_init(si)?;
    let mut verify_msg =
        Vec::with_capacity(constants::INITIATOR_SIG_LABEL.len() + si_encoded.len());
    verify_msg.extend_from_slice(constants::INITIATOR_SIG_LABEL);
    verify_msg.extend_from_slice(&si_encoded);
    identity::hybrid_verify(alice_ik_pk, &verify_msg, sender_sig)?;

    // Co-presence guards: reject mismatched ct_opk / opk_sk before any KEM
    // operations — avoids wasting two decapsulations on malformed input.
    if si.ct_opk.is_none() && opk_sk.is_some() {
        return Err(Error::InvalidData);
    }
    if si.ct_opk.is_some() && opk_sk.is_none() {
        return Err(Error::InvalidData);
    }

    // Mirror Alice's encapsulations — each decapsulation recovers the same
    // shared secret Alice derived, so the HKDF output matches.

    // IdentitySecretKey::from_bytes validates LO_SECRET_KEY_SIZE (2496 bytes) at
    // construction. xwing_sk() slices &self.0[..XWING_SECRET_KEY_SIZE] — always
    // exactly 2432 bytes regardless of sub-key content. xwing::SecretKey::
    // from_bytes_unchecked requires exactly XWING_SECRET_KEY_SIZE bytes and
    // performs no further structural validation; ML-KEM-768 NTT-domain validity
    // is deferred to use-site. If the X-Wing sub-key bytes are malformed,
    // decapsulate returns a pseudorandom shared secret (ML-KEM implicit rejection)
    // and the derived session key diverges from Alice's — AeadFailed at
    // decrypt_first_message, not undefined behavior. from_bytes_unchecked is safe
    // here because the size invariant is enforced by xwing_sk().
    let bob_xwing_sk = xwing::SecretKey::from_bytes_unchecked(bob_ik_sk.xwing_sk().to_vec());
    let mut ss_ik = xwing::decapsulate(&bob_xwing_sk, &si.ct_ik)?;

    // SPK shared secret — binds session to Bob's medium-term signing identity.
    let mut ss_spk = xwing::decapsulate(spk_sk, &si.ct_spk)?;

    // OPK shared secret — provides replay protection (each OPK is used once).
    // Co-presence validated above (ct_opk.is_some() → opk_sk.is_some()).
    let mut ss_opk = if let Some(ref ct_opk) = si.ct_opk {
        Some(xwing::decapsulate(opk_sk.ok_or(Error::Internal)?, ct_opk)?)
    } else {
        None
    };

    // HKDF input must concatenate shared secrets in the same ik ‖ spk ‖ opk
    // order as Alice — HKDF output is order-dependent. This ordering is a
    // documentation-only guarantee; there is no type-level enforcement. The
    // three shared secrets are all xwing::SharedSecret values with no ordering
    // annotation. Reordering the decapsulations above or the extend_from_slice
    // calls below would compile and run silently, producing a different session
    // key with no error until AeadFailed at decrypt_first_message. The
    // session_agreement_with_opk integration test is the only runtime guard
    // against an order-breaking refactor — any change to the decapsulation
    // sequence or IKM concatenation MUST verify that test still passes and
    // confirms matching keys between initiator and responder.
    // Pre-allocate for max IKM: 3 × 32-byte shared secrets (ik + spk + opk).
    let mut ikm = Zeroizing::new(Vec::with_capacity(96));
    ikm.extend_from_slice(ss_ik.as_bytes());
    ikm.extend_from_slice(ss_spk.as_bytes());
    if let Some(ref ss) = ss_opk {
        ikm.extend_from_slice(ss.as_bytes());
    }

    let info = build_kex_info(alice_ik_pk, bob_ik_pk, &si.sender_ek)?;

    let mut session_key = Zeroizing::new([0u8; 64]);
    // Output size (64) is compile-time constant — InvalidLength structurally unreachable.
    hkdf::hkdf_sha3_256(&constants::HKDF_ZERO_SALT, &ikm, &info, &mut *session_key)?;
    ikm.zeroize();

    // Split the 64-byte HKDF output: first half seeds the ratchet root chain,
    // second half becomes the initial send/receive chain key.
    let mut root_key = Zeroizing::new([0u8; 32]);
    let mut chain_key = Zeroizing::new([0u8; 32]);
    root_key.copy_from_slice(&session_key[..32]);
    chain_key.copy_from_slice(&session_key[32..64]);
    session_key.zeroize(); // earliest safe point — both halves extracted above

    // Shared secrets were copied into ikm but backing buffers persist —
    // Zeroizing<Vec> on ikm doesn't cover these originals.
    ss_ik.0.zeroize();
    ss_spk.0.zeroize();
    if let Some(ref mut ss) = ss_opk {
        ss.0.zeroize();
    }

    Ok(ReceivedSession {
        root_key,
        initial_chain_key: chain_key,
        peer_ek: si.sender_ek.clone(),
    })
}

/// Build the HKDF info parameter for LO-KEX (§5.4 Step 4).
///
/// info = "lo-kex-v1" || len(crypto_version) || crypto_version ||
///        len(alice_ik) || alice_ik || len(bob_ik) || bob_ik || len(ek) || ek
///
/// `crypto_version` is included to prevent cross-version domain separation
/// collision: if a future "lo-crypto-v2" reuses the same HKDF label, the
/// version binding ensures sessions under different protocol versions
/// derive independent key material.
fn build_kex_info(
    alice_ik: &IdentityPublicKey,
    bob_ik: &IdentityPublicKey,
    ek: &xwing::PublicKey,
) -> Result<Vec<u8>> {
    let cv = constants::CRYPTO_VERSION.as_bytes();
    // Pre-allocate: KEX_HKDF_INFO_PFX(9) + crypto_version(2+len) +
    // alice_ik(2+LO_PUBLIC_KEY_SIZE) + bob_ik(2+LO_PUBLIC_KEY_SIZE) +
    // ek(2+XWING_PUBLIC_KEY_SIZE).
    let mut info = Vec::with_capacity(
        constants::KEX_HKDF_INFO_PFX.len()
            + 2
            + cv.len()
            + 2
            + constants::LO_PUBLIC_KEY_SIZE
            + 2
            + constants::LO_PUBLIC_KEY_SIZE
            + 2
            + constants::XWING_PUBLIC_KEY_SIZE,
    );
    info.extend_from_slice(constants::KEX_HKDF_INFO_PFX);

    // Bind the protocol version into the HKDF info to prevent cross-version
    // key derivation collisions.
    push_u16_prefixed(&mut info, cv)?;
    push_u16_prefixed(&mut info, alice_ik.as_bytes())?;
    push_u16_prefixed(&mut info, bob_ik.as_bytes())?;
    push_u16_prefixed(&mut info, ek.as_bytes())?;

    Ok(info)
}

/// Check that a byte slice fits in a u16 length prefix.
///
/// Used for variable-length fields (e.g. `crypto_version`) where an oversized
/// value is a recoverable input error. Fixed-size protocol fields (`ct_ik`,
/// `ct_spk`, `ct_opk`) use `assert!` instead because their sizes are
/// compile-time constants — a violation there indicates a library bug.
///
/// Returns `Ok(())` if `data.len() <= u16::MAX`, or `InvalidLength` otherwise.
fn require_u16_len(data: &[u8]) -> Result<()> {
    if data.len() > u16::MAX as usize {
        // `expected` here is the maximum allowed length (u16 ceiling), not an
        // exact size requirement — the wire format uses u16 length prefixes.
        return Err(Error::InvalidLength {
            expected: u16::MAX as usize,
            got: data.len(),
        });
    }
    Ok(())
}

/// Append a u16-length-prefixed field: BE16(len) || data.
///
/// Must only be called with compile-time-constant-sized protocol fields
/// (X-Wing keys, ciphertexts). Variable-length fields (e.g. `crypto_version`)
/// must go through `require_u16_len` first.
fn push_u16_prefixed(buf: &mut Vec<u8>, data: &[u8]) -> Result<()> {
    let len = u16::try_from(data.len()).map_err(|_| Error::Internal)?;
    buf.extend_from_slice(&len.to_be_bytes());
    buf.extend_from_slice(data);
    Ok(())
}

/// Encode a SessionInit into deterministic binary (§7.4).
///
/// Used both as the signed message for identity authentication (§5.4 Step 6,
/// §5.5 Step 2) and as the AAD component for first-message encryption.
///
/// Wire layout (all lengths are 2-byte big-endian u16):
/// ```text
/// len(crypto_version) || crypto_version
/// sender_ik_fingerprint (32 bytes, fixed)
/// recipient_ik_fingerprint (32 bytes, fixed)
/// sender_ek (1216 bytes, fixed)
/// len(ct_ik) || ct_ik (1120 bytes)
/// len(ct_spk) || ct_spk (1120 bytes)
/// spk_id (4 bytes, big-endian u32)
/// has_opk (1 byte: 0x00 or 0x01)
/// [if has_opk] len(ct_opk) || ct_opk (1120 bytes) || opk_id (4 bytes, big-endian u32)
/// ```
pub fn encode_session_init(si: &SessionInit) -> Result<Vec<u8>> {
    // Pre-allocate: 2+12 + 32 + 32 + 1216 + 2+1120 + 2+1120 + 4 + 1 + optionally 2+1120+4 = 4669 with OPK.
    // 5120 = 5 × 1024, rounded up from 4669 for allocator alignment.
    let mut buf = Vec::with_capacity(5120);

    // len(crypto_version) || crypto_version
    let cv = si.crypto_version.as_bytes();
    require_u16_len(cv)?;
    push_u16_prefixed(&mut buf, cv)?;

    // sender_ik_fingerprint_raw (32 bytes, fixed — no length prefix)
    buf.extend_from_slice(&si.sender_ik_fingerprint);

    // recipient_ik_fingerprint (32 bytes, fixed — no length prefix)
    buf.extend_from_slice(&si.recipient_ik_fingerprint);

    // sender_ek (1216 bytes, fixed — no length prefix)
    buf.extend_from_slice(si.sender_ek.as_bytes());

    // len(ct_ik) || ct_ik, len(ct_spk) || ct_spk — defense-in-depth: X-Wing
    // ciphertexts are always XWING_CIPHERTEXT_SIZE bytes, so the error inside
    // push_u16_prefixed is an invariant check, not a user-input guard.
    push_u16_prefixed(&mut buf, si.ct_ik.as_bytes())?;
    push_u16_prefixed(&mut buf, si.ct_spk.as_bytes())?;

    buf.extend_from_slice(&si.spk_id.to_be_bytes());

    // Partial OPK produces undecodable wire format — decoder branches on has_opk.
    if si.ct_opk.is_some() != si.opk_id.is_some() {
        return Err(Error::InvalidData);
    }

    if let (Some(ct_opk), Some(opk_id)) = (&si.ct_opk, si.opk_id) {
        buf.push(0x01);
        push_u16_prefixed(&mut buf, ct_opk.as_bytes())?;
        buf.extend_from_slice(&opk_id.to_be_bytes());
    } else {
        buf.push(0x00);
    }

    Ok(buf)
}

/// Decode a SessionInit from the binary produced by [`encode_session_init`].
///
/// Parses the wire format in exact field order. Returns `Err` on any structural
/// anomaly so that callers never operate on partially-decoded state:
///
/// - `InvalidData`  — input too short to contain the next field (truncation returns
///   `InvalidData`, not `InvalidLength`, to avoid leaking internal parser offsets);
///   invalid UTF-8 in `crypto_version`; `has_opk` not `0x00`/`0x01`;
///   trailing bytes after the last field.
///
/// `xwing::PublicKey::from_bytes` and `xwing::Ciphertext::from_bytes` errors
/// (wrong fixed size) are propagated directly.
pub fn decode_session_init(data: &[u8]) -> Result<SessionInit> {
    // Arithmetic safety: `cur + N` uses plain addition (not checked_add) because
    // the maximum cursor position is bounded by validated ciphertext sizes
    // (3 × XWING_CIPHERTEXT_SIZE = 3360) plus one u16-prefixed crypto_version
    // (up to 65535) plus fixed-size fields (~71K worst case), well under
    // usize::MAX on all targets. Unlike `from_bytes` (which parses
    // variable-count entries), this parser has a fixed number of fields with
    // either validated-constant or u16-bounded lengths.
    //
    // Truncation returns InvalidData (not InvalidLength) to avoid leaking
    // internal parser offsets through error payloads.
    let mut cur = 0usize;

    // len(crypto_version) || crypto_version
    if data.len() < cur + 2 {
        return Err(Error::InvalidData);
    }
    let ver_len = u16::from_be_bytes([data[cur], data[cur + 1]]) as usize;
    cur += 2;
    // Specification.md §5.3: parsers MUST reject crypto_version longer than 64 bytes.
    if ver_len > 64 {
        return Err(Error::InvalidLength {
            expected: 64,
            got: ver_len,
        });
    }
    if data.len() < cur + ver_len {
        return Err(Error::InvalidData);
    }
    let crypto_version =
        String::from_utf8(data[cur..cur + ver_len].to_vec()).map_err(|_| Error::InvalidData)?;
    cur += ver_len;

    // Reject unknown protocol versions at decode time — a standalone
    // decode_session_init call (via CAPI) must not return an unvalidated
    // version that the caller might act on.
    if crypto_version != constants::CRYPTO_VERSION {
        return Err(Error::UnsupportedCryptoVersion);
    }

    // sender_ik_fingerprint (32 bytes, fixed)
    if data.len() < cur + 32 {
        return Err(Error::InvalidData);
    }
    let sender_ik_fingerprint: [u8; 32] = data[cur..cur + 32]
        .try_into()
        .map_err(|_| Error::Internal)?;
    cur += 32;

    // recipient_ik_fingerprint (32 bytes, fixed)
    if data.len() < cur + 32 {
        return Err(Error::InvalidData);
    }
    let recipient_ik_fingerprint: [u8; 32] = data[cur..cur + 32]
        .try_into()
        .map_err(|_| Error::Internal)?;
    cur += 32;

    // sender_ek (XWING_PUBLIC_KEY_SIZE bytes, fixed — no length prefix)
    if data.len() < cur + constants::XWING_PUBLIC_KEY_SIZE {
        return Err(Error::InvalidData);
    }
    let sender_ek =
        xwing::PublicKey::from_bytes(data[cur..cur + constants::XWING_PUBLIC_KEY_SIZE].to_vec())?;
    cur += constants::XWING_PUBLIC_KEY_SIZE;

    // len(ct_ik) || ct_ik
    if data.len() < cur + 2 {
        return Err(Error::InvalidData);
    }
    let ct_ik_len = u16::from_be_bytes([data[cur], data[cur + 1]]) as usize;
    // Validate before allocation — wire-format ciphertext length must match
    // X-Wing's fixed size. InvalidData (not InvalidLength) because this is a
    // wire-format violation, not a caller-size error.
    if ct_ik_len != constants::XWING_CIPHERTEXT_SIZE {
        return Err(Error::InvalidData);
    }
    cur += 2;
    if data.len() < cur + ct_ik_len {
        return Err(Error::InvalidData);
    }
    let ct_ik = xwing::Ciphertext::from_bytes(data[cur..cur + ct_ik_len].to_vec())?;
    cur += ct_ik_len;

    // len(ct_spk) || ct_spk
    if data.len() < cur + 2 {
        return Err(Error::InvalidData);
    }
    let ct_spk_len = u16::from_be_bytes([data[cur], data[cur + 1]]) as usize;
    if ct_spk_len != constants::XWING_CIPHERTEXT_SIZE {
        return Err(Error::InvalidData);
    }
    cur += 2;
    if data.len() < cur + ct_spk_len {
        return Err(Error::InvalidData);
    }
    let ct_spk = xwing::Ciphertext::from_bytes(data[cur..cur + ct_spk_len].to_vec())?;
    cur += ct_spk_len;

    // spk_id (4 bytes, BE u32)
    if data.len() < cur + 4 {
        return Err(Error::InvalidData);
    }
    let spk_id = u32::from_be_bytes(data[cur..cur + 4].try_into().map_err(|_| Error::Internal)?);
    cur += 4;

    // has_opk (1 byte: 0x00 or 0x01 — strict)
    if data.len() < cur + 1 {
        return Err(Error::InvalidData);
    }
    let has_opk_byte = data[cur];
    cur += 1;

    let (ct_opk, opk_id) = match has_opk_byte {
        0x00 => (None, None),
        0x01 => {
            // len(ct_opk) || ct_opk
            if data.len() < cur + 2 {
                return Err(Error::InvalidData);
            }
            let ct_opk_len = u16::from_be_bytes([data[cur], data[cur + 1]]) as usize;
            if ct_opk_len != constants::XWING_CIPHERTEXT_SIZE {
                return Err(Error::InvalidData);
            }
            cur += 2;
            if data.len() < cur + ct_opk_len {
                return Err(Error::InvalidData);
            }
            let ct = xwing::Ciphertext::from_bytes(data[cur..cur + ct_opk_len].to_vec())?;
            cur += ct_opk_len;

            // opk_id (4 bytes, BE u32)
            if data.len() < cur + 4 {
                return Err(Error::InvalidData);
            }
            let id =
                u32::from_be_bytes(data[cur..cur + 4].try_into().map_err(|_| Error::Internal)?);
            cur += 4;

            (Some(ct), Some(id))
        }
        _ => return Err(Error::InvalidData),
    };

    // Reject trailing bytes — a strict parser prevents format malleability.
    if cur != data.len() {
        return Err(Error::InvalidData);
    }

    Ok(SessionInit {
        crypto_version,
        sender_ik_fingerprint,
        recipient_ik_fingerprint,
        sender_ek,
        ct_ik,
        ct_spk,
        spk_id,
        ct_opk,
        opk_id,
    })
}

/// Build AAD for the first message in a session (§7.3).
///
/// # Security
///
/// The AAD binds sender/recipient identity fingerprints and the full
/// `SessionInit` to the AEAD encryption, preventing cross-session replay
/// and cross-party confusion.
pub fn build_first_message_aad(
    sender_fingerprint: &[u8; 32],
    recipient_fingerprint: &[u8; 32],
    si: &SessionInit,
) -> Result<Vec<u8>> {
    let si_bytes = encode_session_init(si)?;
    build_first_message_aad_from_encoded(sender_fingerprint, recipient_fingerprint, &si_bytes)
}

/// Build AAD from pre-encoded session init bytes (§7.3).
///
/// Same as [`build_first_message_aad`] but accepts the already-encoded
/// `SessionInit` bytes. Used by the CAPI where the caller provides the
/// encoded form directly.
///
/// Returns `Err(InvalidData)` if `si_encoded` is empty — an empty encoding
/// would strip session binding from the AAD, weakening the AEAD guarantee.
pub fn build_first_message_aad_from_encoded(
    sender_fingerprint: &[u8; 32],
    recipient_fingerprint: &[u8; 32],
    si_encoded: &[u8],
) -> Result<Vec<u8>> {
    if si_encoded.is_empty() {
        return Err(Error::InvalidData);
    }
    let mut aad = Vec::with_capacity(constants::DM_AAD.len() + 32 + 32 + si_encoded.len());
    aad.extend_from_slice(constants::DM_AAD);
    aad.extend_from_slice(sender_fingerprint);
    aad.extend_from_slice(recipient_fingerprint);
    aad.extend_from_slice(si_encoded);
    Ok(aad)
}

/// Sign a pre-key's public key with the identity key (§5.3).
///
/// Produces a hybrid signature (Ed25519 + ML-DSA-65) over the domain-separated
/// message `"lo-spk-sig-v1" ‖ SPK_pub`. The verifier calls [`verify_bundle`]
/// to check this signature.
pub fn sign_prekey(
    ik_sk: &IdentitySecretKey,
    spk_pub: &xwing::PublicKey,
) -> Result<HybridSignature> {
    let mut msg = Vec::with_capacity(constants::SPK_SIG_LABEL.len() + spk_pub.as_bytes().len());
    msg.extend_from_slice(constants::SPK_SIG_LABEL);
    msg.extend_from_slice(spk_pub.as_bytes());
    identity::hybrid_sign(ik_sk, &msg)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::Error;
    use crate::identity::{GeneratedIdentity, generate_identity};
    use crate::ratchet::RatchetState;
    use hex_literal::hex;

    /// Generate a valid pre-key bundle for Bob.
    fn make_bundle(
        bob_pk: &IdentityPublicKey,
        bob_sk: &IdentitySecretKey,
        with_opk: bool,
    ) -> (PreKeyBundle, xwing::SecretKey, Option<xwing::SecretKey>) {
        let (spk_pub, spk_sk) = xwing::keygen().expect("keygen");
        let spk_sig = sign_prekey(bob_sk, &spk_pub).expect("sign_prekey");
        let (opk_pub, opk_sk, opk_id) = if with_opk {
            let (pub_key, sec_key) = xwing::keygen().expect("keygen (opk)");
            (Some(pub_key), Some(sec_key), Some(1))
        } else {
            (None, None, None)
        };
        let bundle = PreKeyBundle {
            ik_pub: IdentityPublicKey(bob_pk.as_bytes().to_vec()),
            crypto_version: constants::CRYPTO_VERSION.to_string(),
            spk_pub,
            spk_id: 42,
            spk_sig,
            opk_pub,
            opk_id,
        };
        (bundle, spk_sk, opk_sk)
    }

    #[test]
    fn verify_bundle_valid() {
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, _, _) = make_bundle(&bob_pk, &bob_sk, false);
        assert!(verify_bundle(bundle, &bob_pk).is_ok());
    }

    #[test]
    fn verify_bundle_ik_mismatch() {
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: other_pk,
            ..
        } = generate_identity().unwrap();
        let (bundle, _, _) = make_bundle(&bob_pk, &bob_sk, false);
        assert!(matches!(
            verify_bundle(bundle, &other_pk),
            Err(Error::BundleVerificationFailed)
        ));
    }

    #[test]
    fn verify_bundle_bad_spk_sig() {
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (mut bundle, _, _) = make_bundle(&bob_pk, &bob_sk, false);
        // Tamper with the SPK signature.
        let mut bad_sig = bundle.spk_sig.as_bytes().to_vec();
        bad_sig[0] ^= 0xFF;
        bundle.spk_sig = HybridSignature::from_bytes(bad_sig).unwrap();
        // Returns BundleVerificationFailed (not VerificationFailed) — all
        // non-structural failures are collapsed to prevent error oracle.
        assert!(matches!(
            verify_bundle(bundle, &bob_pk),
            Err(Error::BundleVerificationFailed)
        ));
    }

    #[test]
    fn verify_bundle_wrong_crypto_version() {
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (mut bundle, _, _) = make_bundle(&bob_pk, &bob_sk, false);
        bundle.crypto_version = "lo-crypto-v999".to_string();
        // Returns BundleVerificationFailed (not UnsupportedCryptoVersion) to
        // prevent iterative probing of which check failed.
        assert!(matches!(
            verify_bundle(bundle, &bob_pk),
            Err(Error::BundleVerificationFailed)
        ));
    }

    #[test]
    fn verify_bundle_partial_opk() {
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (mut bundle, _, _) = make_bundle(&bob_pk, &bob_sk, true);
        // Make opk_id None while opk_pub is Some.
        bundle.opk_id = None;
        assert!(matches!(
            verify_bundle(bundle, &bob_pk),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn session_agreement_with_opk() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, opk_sk) = make_bundle(&bob_pk, &bob_sk, true);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let mut initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let mut received = receive_session(
            &bob_pk,
            &bob_sk,
            &alice_pk,
            &initiated.session_init,
            &initiated.sender_sig,
            &spk_sk,
            opk_sk.as_ref(),
        )
        .unwrap();
        assert_eq!(*initiated.take_root_key(), *received.take_root_key());
        assert_eq!(
            *initiated.take_initial_chain_key(),
            *received.take_initial_chain_key()
        );
    }

    #[test]
    fn session_agreement_without_opk() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let mut initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let mut received = receive_session(
            &bob_pk,
            &bob_sk,
            &alice_pk,
            &initiated.session_init,
            &initiated.sender_sig,
            &spk_sk,
            None,
        )
        .unwrap();
        assert_eq!(*initiated.take_root_key(), *received.take_root_key());
        assert_eq!(
            *initiated.take_initial_chain_key(),
            *received.take_initial_chain_key()
        );
    }

    #[test]
    fn kex_derivation_kat() {
        // Known-answer test for the LO-KEX HKDF derivation step (§5.4 Step 4).
        //
        // Fixed synthetic inputs; expected output computed independently by
        // Python's hmac/hashlib (HKDF-SHA3-256 per RFC 5869):
        //
        //   salt = [0x00; 32]              (HKDF_ZERO_SALT)
        //   ikm  = [0xAA; 32] ‖ [0xBB; 32] (ss_ik ‖ ss_spk, no OPK)
        //   info = "lo-kex-v1"
        //          ‖ BE16(3200) ‖ [0x01; 3200]  (alice IK)
        //          ‖ BE16(3200) ‖ [0x02; 3200]  (bob IK)
        //          ‖ BE16(1216) ‖ [0x03; 1216]  (ephemeral key)
        //
        // Verifies: info string construction, HKDF parameters (salt, L=64),
        // and output split (root_key=[0..32], chain_key=[32..64]).
        // IKM concatenation order (ss_ik ‖ ss_spk ‖ [ss_opk]) is covered by
        // session_agreement_with_opk / session_agreement_without_opk, which
        // would fail if Alice and Bob assembled IKM asymmetrically.
        let alice_ik =
            IdentityPublicKey::from_bytes(vec![0x01u8; constants::LO_PUBLIC_KEY_SIZE]).unwrap();
        let bob_ik =
            IdentityPublicKey::from_bytes(vec![0x02u8; constants::LO_PUBLIC_KEY_SIZE]).unwrap();
        let ek =
            xwing::PublicKey::from_bytes(vec![0x03u8; constants::XWING_PUBLIC_KEY_SIZE]).unwrap();

        let mut ikm = [0u8; 64];
        ikm[..32].fill(0xAA); // ss_ik
        ikm[32..].fill(0xBB); // ss_spk

        let info = build_kex_info(&alice_ik, &bob_ik, &ek).unwrap();
        let mut out = [0u8; 64];
        hkdf::hkdf_sha3_256(&constants::HKDF_ZERO_SALT, &ikm, &info, &mut out).unwrap();

        // KAT vectors updated after RT-182: crypto_version is now bound into
        // the HKDF info string, preventing cross-version domain separation
        // collision.
        assert_eq!(
            out[..32],
            hex!("b90ead32626ae4b0be4864ba8ad2fad4b570c4021d4f20b83c79207dd2a91074"),
            "root_key derivation mismatch — info construction or HKDF parameters changed"
        );
        assert_eq!(
            out[32..],
            hex!("965624c6263ca5663765bc97f179891ea3e0ee73010b037de182abb3b123d0bf"),
            "chain_key derivation mismatch — info construction or HKDF parameters changed"
        );
    }

    #[test]
    fn receive_session_fingerprint_mismatch() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: fake_alice_pk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        // Bob thinks the sender is fake_alice, not alice — fingerprint check fires first.
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &fake_alice_pk,
                &initiated.session_init,
                &initiated.sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn receive_session_wrong_recipient_fingerprint() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: carol_pk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let mut initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        // Tamper: replace recipient fingerprint with Carol's. Bob's pre-signature
        // check (si.recipient_ik_fingerprint != SHA3-256(bob_pk)) fires immediately.
        initiated.session_init.recipient_ik_fingerprint = carol_pk.fingerprint_raw();
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &initiated.session_init,
                &initiated.sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn receive_session_wrong_crypto_version() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let mut initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        initiated.session_init.crypto_version = "lo-crypto-v999".to_string();
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &initiated.session_init,
                &initiated.sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::UnsupportedCryptoVersion)
        ));
    }

    #[test]
    fn receive_session_opk_sk_without_ct_opk() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        // Provide an opk_sk when ct_opk is absent.
        let (_, extra_sk) = xwing::keygen().unwrap();
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &initiated.session_init,
                &initiated.sender_sig,
                &spk_sk,
                Some(&extra_sk)
            ),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn receive_session_ct_opk_without_opk_sk() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, true);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        // ct_opk is present (with OPK) but we don't provide opk_sk.
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &initiated.session_init,
                &initiated.sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn receive_session_partial_opk_fields() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, opk_sk) = make_bundle(&bob_pk, &bob_sk, true);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let mut initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        // Make ct_opk present but opk_id absent.
        initiated.session_init.opk_id = None;
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &initiated.session_init,
                &initiated.sender_sig,
                &spk_sk,
                opk_sk.as_ref()
            ),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn receive_session_bad_sig() {
        // Verify that tampered signature is rejected.
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        // Tamper with the signature.
        let mut bad_sig_bytes = initiated.sender_sig.as_bytes().to_vec();
        bad_sig_bytes[0] ^= 0xFF;
        let bad_sig = HybridSignature::from_bytes(bad_sig_bytes).unwrap();
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &initiated.session_init,
                &bad_sig,
                &spk_sk,
                None
            ),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn receive_session_tampered_ct_ik() {
        // Flip one byte in ct_ik of the SessionInit. The sender_sig was computed
        // over the original SessionInit encoding, so receive_session must reject
        // the tampered session with VerificationFailed — the signature no longer
        // matches the altered content.
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        // InitiatedSession implements Drop, so fields cannot be moved out.
        // Clone the data we need before constructing the tampered SessionInit.
        let sender_sig = initiated.sender_sig.clone();
        let si = &initiated.session_init;
        // Flip the first byte of ct_ik — changes the encoded SessionInit that
        // the signature covers.
        let mut ct_ik_bytes = si.ct_ik.as_bytes().to_vec();
        ct_ik_bytes[0] ^= 0xFF;
        let tampered = SessionInit {
            crypto_version: si.crypto_version.clone(),
            sender_ik_fingerprint: si.sender_ik_fingerprint,
            recipient_ik_fingerprint: si.recipient_ik_fingerprint,
            sender_ek: si.sender_ek.clone(),
            ct_ik: xwing::Ciphertext::from_bytes(ct_ik_bytes).unwrap(),
            ct_spk: si.ct_spk.clone(),
            spk_id: si.spk_id,
            ct_opk: si.ct_opk.clone(),
            opk_id: si.opk_id,
        };
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &tampered,
                &sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn receive_session_tampered_sender_ek() {
        // Flip a byte in sender_ek — changes the encoded SessionInit that the
        // signature covers → VerificationFailed.
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let sender_sig = initiated.sender_sig.clone();
        let si = &initiated.session_init;
        let mut ek_bytes = si.sender_ek.as_bytes().to_vec();
        ek_bytes[0] ^= 0xFF;
        let tampered = SessionInit {
            crypto_version: si.crypto_version.clone(),
            sender_ik_fingerprint: si.sender_ik_fingerprint,
            recipient_ik_fingerprint: si.recipient_ik_fingerprint,
            sender_ek: xwing::PublicKey::from_bytes(ek_bytes).unwrap(),
            ct_ik: si.ct_ik.clone(),
            ct_spk: si.ct_spk.clone(),
            spk_id: si.spk_id,
            ct_opk: si.ct_opk.clone(),
            opk_id: si.opk_id,
        };
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &tampered,
                &sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn receive_session_tampered_ct_spk() {
        // Flip a byte in ct_spk — changes the encoded SessionInit that the
        // signature covers → VerificationFailed.
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let sender_sig = initiated.sender_sig.clone();
        let si = &initiated.session_init;
        let mut ct_spk_bytes = si.ct_spk.as_bytes().to_vec();
        ct_spk_bytes[0] ^= 0xFF;
        let tampered = SessionInit {
            crypto_version: si.crypto_version.clone(),
            sender_ik_fingerprint: si.sender_ik_fingerprint,
            recipient_ik_fingerprint: si.recipient_ik_fingerprint,
            sender_ek: si.sender_ek.clone(),
            ct_ik: si.ct_ik.clone(),
            ct_spk: xwing::Ciphertext::from_bytes(ct_spk_bytes).unwrap(),
            spk_id: si.spk_id,
            ct_opk: si.ct_opk.clone(),
            opk_id: si.opk_id,
        };
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &tampered,
                &sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn receive_session_tampered_sender_ik_fingerprint() {
        // Flip a byte in sender_ik_fingerprint — the fingerprint check fires
        // (SHA3-256(alice_pk) ≠ si.sender_ik_fingerprint) before signature
        // verification → InvalidData.
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let sender_sig = initiated.sender_sig.clone();
        let si = &initiated.session_init;
        let mut bad_fp = si.sender_ik_fingerprint;
        bad_fp[0] ^= 0xFF;
        let tampered = SessionInit {
            crypto_version: si.crypto_version.clone(),
            sender_ik_fingerprint: bad_fp,
            recipient_ik_fingerprint: si.recipient_ik_fingerprint,
            sender_ek: si.sender_ek.clone(),
            ct_ik: si.ct_ik.clone(),
            ct_spk: si.ct_spk.clone(),
            spk_id: si.spk_id,
            ct_opk: si.ct_opk.clone(),
            opk_id: si.opk_id,
        };
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &tampered,
                &sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn receive_session_tampered_crypto_version() {
        // Replace crypto_version with an unknown string — the version check fires
        // before signature verification → UnsupportedCryptoVersion.
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let sender_sig = initiated.sender_sig.clone();
        let si = &initiated.session_init;
        let tampered = SessionInit {
            crypto_version: "lo-crypto-v99".to_string(),
            sender_ik_fingerprint: si.sender_ik_fingerprint,
            recipient_ik_fingerprint: si.recipient_ik_fingerprint,
            sender_ek: si.sender_ek.clone(),
            ct_ik: si.ct_ik.clone(),
            ct_spk: si.ct_spk.clone(),
            spk_id: si.spk_id,
            ct_opk: si.ct_opk.clone(),
            opk_id: si.opk_id,
        };
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &tampered,
                &sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::UnsupportedCryptoVersion)
        ));
    }

    #[test]
    fn receive_session_wrong_signer() {
        // Eve initiates with her own keys but Bob is told the sender is Alice.
        // Eve's sender_ik_fingerprint is SHA3-256(eve_pk), not SHA3-256(alice_pk),
        // so the fingerprint check fires first and returns InvalidData — before
        // signature verification is reached.
        let GeneratedIdentity {
            public_key: alice_pk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: eve_pk,
            secret_key: eve_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb_eve = verify_bundle(bundle, &bob_pk).unwrap();
        let eve_initiated = initiate_session(&eve_pk, &eve_sk, &vb_eve).unwrap();
        assert!(matches!(
            receive_session(
                &bob_pk,
                &bob_sk,
                &alice_pk,
                &eve_initiated.session_init,
                &eve_initiated.sender_sig,
                &spk_sk,
                None
            ),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn zero_knowledge_impersonation_rejected() {
        // The exact zero-knowledge impersonation attack: Eve constructs a SessionInit
        // with Alice's fingerprint and valid KEM ciphertexts to Bob's keys — but signs
        // with her own sk because she doesn't have sk_IK_A.
        // The fingerprint check passes (Alice's fingerprint is present), so signature
        // verification is reached and must return VerificationFailed.
        let GeneratedIdentity {
            public_key: alice_pk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            secret_key: eve_sk, ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();

        // Eve generates her own ephemeral key and encapsulates to Bob's keys.
        let (ek_pk, _ek_sk) = xwing::keygen().unwrap();
        let bob_xwing_pk = xwing::PublicKey::from_bytes_unchecked(bob_pk.xwing_pk().to_vec());
        let (ct_ik, _ss) = xwing::encapsulate(&bob_xwing_pk).unwrap();
        let (ct_spk, _ss) = xwing::encapsulate(&vb.spk_pub).unwrap();

        // Eve places Alice's fingerprint in the SessionInit — sender fingerprint check passes.
        // Eve also places Bob's fingerprint correctly — recipient fingerprint check passes.
        // Only the signature will catch Eve's impersonation.
        let si = SessionInit {
            crypto_version: constants::CRYPTO_VERSION.to_string(),
            sender_ik_fingerprint: alice_pk.fingerprint_raw(),
            recipient_ik_fingerprint: bob_pk.fingerprint_raw(),
            sender_ek: ek_pk,
            ct_ik,
            ct_spk,
            spk_id: vb.spk_id,
            ct_opk: None,
            opk_id: None,
        };

        // Eve signs with her own sk (she doesn't have sk_IK_A).
        let si_encoded = encode_session_init(&si).unwrap();
        let mut sign_msg =
            Vec::with_capacity(constants::INITIATOR_SIG_LABEL.len() + si_encoded.len());
        sign_msg.extend_from_slice(constants::INITIATOR_SIG_LABEL);
        sign_msg.extend_from_slice(&si_encoded);
        let eve_sig = identity::hybrid_sign(&eve_sk, &sign_msg).unwrap();

        // Bob rejects: Eve's signature doesn't verify under alice_pk.
        assert!(matches!(
            receive_session(&bob_pk, &bob_sk, &alice_pk, &si, &eve_sig, &spk_sk, None),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn encode_session_init_deterministic() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, _, _) = make_bundle(&bob_pk, &bob_sk, true);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let enc1 = encode_session_init(&initiated.session_init).unwrap();
        let enc2 = encode_session_init(&initiated.session_init).unwrap();
        assert_eq!(enc1, enc2);
    }

    #[test]
    fn encode_session_init_rejects_oversized_crypto_version() {
        // crypto_version exceeding u16::MAX cannot be encoded (u16 length prefix).
        let si = SessionInit {
            crypto_version: "x".repeat(u16::MAX as usize + 1),
            sender_ik_fingerprint: [0u8; 32],
            recipient_ik_fingerprint: [0u8; 32],
            sender_ek: xwing::PublicKey::from_bytes(vec![0u8; constants::XWING_PUBLIC_KEY_SIZE])
                .unwrap(),
            ct_ik: xwing::Ciphertext::from_bytes(vec![0u8; constants::XWING_CIPHERTEXT_SIZE])
                .unwrap(),
            ct_spk: xwing::Ciphertext::from_bytes(vec![0u8; constants::XWING_CIPHERTEXT_SIZE])
                .unwrap(),
            spk_id: 0,
            ct_opk: None,
            opk_id: None,
        };
        assert!(matches!(
            encode_session_init(&si),
            Err(Error::InvalidLength { .. })
        ));
    }

    #[test]
    fn build_first_message_aad_structure() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, _, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let fp_a = alice_pk.fingerprint_raw();
        let fp_b = bob_pk.fingerprint_raw();
        let aad = build_first_message_aad(&fp_a, &fp_b, &initiated.session_init).unwrap();
        let si_bytes = encode_session_init(&initiated.session_init).unwrap();
        // AAD = "lo-dm-v1" || fp_A || fp_B || Encode(SI)
        let mut expected = Vec::new();
        expected.extend_from_slice(constants::DM_AAD);
        expected.extend_from_slice(&fp_a);
        expected.extend_from_slice(&fp_b);
        expected.extend_from_slice(&si_bytes);
        assert_eq!(aad, expected);
    }

    #[test]
    #[allow(clippy::cast_possible_truncation)]
    fn kex_info_field_structure() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, _, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        // Rebuild info independently.
        // Stable reference — deliberately uses literal b"lo-kex-v1" and
        // b"lo-crypto-v1", not constants, so a constant change is caught.
        let info = build_kex_info(&alice_pk, &bob_pk, &initiated.ek_pk).unwrap();
        let cv = b"lo-crypto-v1";
        let mut expected = Vec::new();
        expected.extend_from_slice(b"lo-kex-v1");
        expected.extend_from_slice(&(cv.len() as u16).to_be_bytes());
        expected.extend_from_slice(cv);
        expected.extend_from_slice(&(alice_pk.as_bytes().len() as u16).to_be_bytes());
        expected.extend_from_slice(alice_pk.as_bytes());
        expected.extend_from_slice(&(bob_pk.as_bytes().len() as u16).to_be_bytes());
        expected.extend_from_slice(bob_pk.as_bytes());
        expected.extend_from_slice(&(initiated.ek_pk.as_bytes().len() as u16).to_be_bytes());
        expected.extend_from_slice(initiated.ek_pk.as_bytes());
        assert_eq!(info, expected);
    }

    #[test]
    fn sign_prekey_verify() {
        let GeneratedIdentity {
            public_key: pk,
            secret_key: sk,
            ..
        } = generate_identity().unwrap();
        let (spk_pub, _) = xwing::keygen().unwrap();
        let sig = sign_prekey(&sk, &spk_pub).unwrap();
        // Verify with the same domain-separated message that sign_prekey uses.
        let mut msg = Vec::with_capacity(constants::SPK_SIG_LABEL.len() + spk_pub.as_bytes().len());
        msg.extend_from_slice(constants::SPK_SIG_LABEL);
        msg.extend_from_slice(spk_pub.as_bytes());
        assert!(identity::hybrid_verify(&pk, &msg, &sig).is_ok());
    }

    #[test]
    fn first_message_round_trip() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let mut initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let fp_a = alice_pk.fingerprint_raw();
        let fp_b = bob_pk.fingerprint_raw();
        let aad = build_first_message_aad(&fp_a, &fp_b, &initiated.session_init).unwrap();
        let (payload, ck_alice) = RatchetState::encrypt_first_message(
            initiated.take_initial_chain_key(),
            b"first message",
            &aad,
        )
        .unwrap();
        // Bob receives and decrypts.
        let mut received = receive_session(
            &bob_pk,
            &bob_sk,
            &alice_pk,
            &initiated.session_init,
            &initiated.sender_sig,
            &spk_sk,
            None,
        )
        .unwrap();
        let (pt, ck_bob) =
            RatchetState::decrypt_first_message(received.take_initial_chain_key(), &payload, &aad)
                .unwrap();
        assert_eq!(&*pt, b"first message");
        assert_eq!(*ck_alice, *ck_bob);
    }

    // ── decode_session_init tests ──────────────────────────────────────────────

    /// Build an all-zeros (no keygen) SessionInit for MIRI-safe roundtrip tests.
    ///
    /// Uses fixed-size zeroed byte slices for xwing types — `from_bytes` is a
    /// size check only, so no PQ operations are performed.
    fn make_zero_session_init(with_opk: bool) -> SessionInit {
        SessionInit {
            crypto_version: constants::CRYPTO_VERSION.to_string(),
            sender_ik_fingerprint: [0x01u8; 32],
            recipient_ik_fingerprint: [0x02u8; 32],
            sender_ek: xwing::PublicKey::from_bytes(vec![0x03u8; constants::XWING_PUBLIC_KEY_SIZE])
                .unwrap(),
            ct_ik: xwing::Ciphertext::from_bytes(vec![0x04u8; constants::XWING_CIPHERTEXT_SIZE])
                .unwrap(),
            ct_spk: xwing::Ciphertext::from_bytes(vec![0x05u8; constants::XWING_CIPHERTEXT_SIZE])
                .unwrap(),
            spk_id: 0xDEAD_BEEF,
            ct_opk: if with_opk {
                Some(
                    xwing::Ciphertext::from_bytes(vec![0x06u8; constants::XWING_CIPHERTEXT_SIZE])
                        .unwrap(),
                )
            } else {
                None
            },
            opk_id: if with_opk { Some(0x1234_5678) } else { None },
        }
    }

    #[test]
    fn decode_session_init_roundtrip_without_opk() {
        let si = make_zero_session_init(false);
        let encoded = encode_session_init(&si).unwrap();
        let decoded = decode_session_init(&encoded).unwrap();
        // Field-level assertions.
        assert_eq!(decoded.crypto_version, si.crypto_version);
        assert_eq!(decoded.sender_ik_fingerprint, si.sender_ik_fingerprint);
        assert_eq!(
            decoded.recipient_ik_fingerprint,
            si.recipient_ik_fingerprint
        );
        assert_eq!(decoded.sender_ek.as_bytes(), si.sender_ek.as_bytes());
        assert_eq!(decoded.ct_ik.as_bytes(), si.ct_ik.as_bytes());
        assert_eq!(decoded.ct_spk.as_bytes(), si.ct_spk.as_bytes());
        assert_eq!(decoded.spk_id, si.spk_id);
        assert!(decoded.ct_opk.is_none());
        assert!(decoded.opk_id.is_none());
        // Byte-level: re-encoding the decoded struct must produce identical bytes.
        let re_encoded = encode_session_init(&decoded).unwrap();
        assert_eq!(encoded, re_encoded);
    }

    #[test]
    fn decode_session_init_roundtrip_with_opk() {
        let si = make_zero_session_init(true);
        let encoded = encode_session_init(&si).unwrap();
        let decoded = decode_session_init(&encoded).unwrap();
        assert_eq!(decoded.crypto_version, si.crypto_version);
        assert_eq!(decoded.sender_ik_fingerprint, si.sender_ik_fingerprint);
        assert_eq!(
            decoded.recipient_ik_fingerprint,
            si.recipient_ik_fingerprint
        );
        assert_eq!(decoded.sender_ek.as_bytes(), si.sender_ek.as_bytes());
        assert_eq!(decoded.ct_ik.as_bytes(), si.ct_ik.as_bytes());
        assert_eq!(decoded.ct_spk.as_bytes(), si.ct_spk.as_bytes());
        assert_eq!(decoded.spk_id, si.spk_id);
        assert_eq!(
            decoded.ct_opk.as_ref().unwrap().as_bytes(),
            si.ct_opk.as_ref().unwrap().as_bytes()
        );
        assert_eq!(decoded.opk_id, si.opk_id);
        let re_encoded = encode_session_init(&decoded).unwrap();
        assert_eq!(encoded, re_encoded);
    }

    #[test]
    fn decode_session_init_empty_returns_invalid_data() {
        assert!(matches!(decode_session_init(&[]), Err(Error::InvalidData)));
    }

    #[test]
    fn decode_session_init_truncated_mid_field_returns_invalid_data() {
        let encoded = encode_session_init(&make_zero_session_init(false)).unwrap();
        // Truncate at several representative offsets covering different fields:
        // within crypto_version length prefix (1), within sender_fp (10),
        // within sender_ek (100), within ct_ik (1500), past spk_id but before has_opk.
        for cut in [1, 10, 100, 1500, encoded.len() - 1] {
            assert!(
                matches!(
                    decode_session_init(&encoded[..cut]),
                    Err(Error::InvalidData)
                ),
                "expected InvalidData when truncated to {cut} bytes"
            );
        }
    }

    #[test]
    fn decode_session_init_trailing_byte_returns_invalid_data() {
        let encoded = encode_session_init(&make_zero_session_init(false)).unwrap();
        let mut with_trail = encoded;
        with_trail.push(0x00);
        assert!(matches!(
            decode_session_init(&with_trail),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn decode_session_init_invalid_utf8_version_returns_invalid_data() {
        // Build a wire buffer with a 1-byte crypto_version set to 0xFF (not valid UTF-8).
        let mut buf: Vec<u8> = Vec::new();
        buf.extend_from_slice(&1u16.to_be_bytes()); // len = 1
        buf.push(0xFF); // invalid UTF-8
        // Fill the rest with zeros up to the minimum required length so the parser
        // reaches the UTF-8 decode and doesn't fail with InvalidLength first.
        buf.extend_from_slice(&[0x01u8; 32]); // sender_fp
        buf.extend_from_slice(&[0x02u8; 32]); // recipient_fp
        buf.extend_from_slice(&[0x03u8; constants::XWING_PUBLIC_KEY_SIZE]); // sender_ek
        // ct_ik with length prefix
        buf.extend_from_slice(
            &u16::try_from(constants::XWING_CIPHERTEXT_SIZE)
                .unwrap()
                .to_be_bytes(),
        );
        buf.extend_from_slice(&[0x04u8; constants::XWING_CIPHERTEXT_SIZE]);
        // ct_spk with length prefix
        buf.extend_from_slice(
            &u16::try_from(constants::XWING_CIPHERTEXT_SIZE)
                .unwrap()
                .to_be_bytes(),
        );
        buf.extend_from_slice(&[0x05u8; constants::XWING_CIPHERTEXT_SIZE]);
        buf.extend_from_slice(&0u32.to_be_bytes()); // spk_id
        buf.push(0x00); // has_opk = false
        assert!(matches!(decode_session_init(&buf), Err(Error::InvalidData)));
    }

    #[test]
    fn decode_session_init_bad_has_opk_byte_returns_invalid_data() {
        // Encode a valid no-OPK SessionInit, then replace has_opk with 0x02.
        let encoded = encode_session_init(&make_zero_session_init(false)).unwrap();
        let mut bad = encoded;
        // has_opk is the last byte in a no-OPK encoding.
        *bad.last_mut().unwrap() = 0x02;
        assert!(matches!(decode_session_init(&bad), Err(Error::InvalidData)));
    }

    #[test]
    fn first_message_wrong_aad() {
        let GeneratedIdentity {
            public_key: alice_pk,
            secret_key: alice_sk,
            ..
        } = generate_identity().unwrap();
        let GeneratedIdentity {
            public_key: bob_pk,
            secret_key: bob_sk,
            ..
        } = generate_identity().unwrap();
        let (bundle, spk_sk, _) = make_bundle(&bob_pk, &bob_sk, false);
        let vb = verify_bundle(bundle, &bob_pk).unwrap();
        let mut initiated = initiate_session(&alice_pk, &alice_sk, &vb).unwrap();
        let fp_a = alice_pk.fingerprint_raw();
        let fp_b = bob_pk.fingerprint_raw();
        let aad = build_first_message_aad(&fp_a, &fp_b, &initiated.session_init).unwrap();
        let (payload, _) =
            RatchetState::encrypt_first_message(initiated.take_initial_chain_key(), b"first", &aad)
                .unwrap();
        let mut received = receive_session(
            &bob_pk,
            &bob_sk,
            &alice_pk,
            &initiated.session_init,
            &initiated.sender_sig,
            &spk_sk,
            None,
        )
        .unwrap();
        // Decrypt with wrong AAD.
        assert!(matches!(
            RatchetState::decrypt_first_message(
                received.take_initial_chain_key(),
                &payload,
                b"wrong aad"
            ),
            Err(Error::AeadFailed)
        ));
    }
}