libsoliton/soliton/tests/integration_kex_ratchet.rs

#![allow(deprecated)] // Tests exercise from_bytes directly for parser coverage.

use soliton::constants;
use soliton::identity::{
    GeneratedIdentity, IdentityPublicKey, IdentitySecretKey, generate_identity,
};
use soliton::kex::{
    PreKeyBundle, build_first_message_aad, initiate_session, receive_session, sign_prekey,
    verify_bundle,
};
use soliton::primitives::xwing;
use soliton::ratchet::RatchetState;

/// Generate a full identity + signed pre-key bundle for testing.
fn setup_peer() -> (
    IdentityPublicKey,
    IdentitySecretKey,
    xwing::PublicKey,
    xwing::SecretKey,
    PreKeyBundle,
) {
    let GeneratedIdentity {
        public_key: ik_pk,
        secret_key: ik_sk,
        ..
    } = generate_identity().unwrap();
    let (spk_pk, spk_sk) = xwing::keygen().unwrap();
    let spk_sig = sign_prekey(&ik_sk, &spk_pk).unwrap();
    let bundle = PreKeyBundle {
        ik_pub: IdentityPublicKey::from_bytes(ik_pk.as_bytes().to_vec()).unwrap(),
        crypto_version: constants::CRYPTO_VERSION.to_string(),
        spk_pub: spk_pk.clone(),
        spk_id: 1,
        spk_sig,
        opk_pub: None,
        opk_id: None,
    };
    (ik_pk, ik_sk, spk_pk, spk_sk, bundle)
}

/// Run a full KEX and return initialized ratchet states + fingerprints.
fn do_kex(with_opk: bool) -> (RatchetState, RatchetState, [u8; 32], [u8; 32]) {
    let (alice_ik_pk, alice_ik_sk, _alice_spk_pk, _alice_spk_sk, _alice_bundle) = setup_peer();
    let (bob_ik_pk, bob_ik_sk, _bob_spk_pk, bob_spk_sk, mut bob_bundle) = setup_peer();

    let opk_sk = if with_opk {
        let (pk, sk) = xwing::keygen().unwrap();
        bob_bundle.opk_pub = Some(pk);
        bob_bundle.opk_id = Some(42);
        Some(sk)
    } else {
        None
    };

    let fp_a = alice_ik_pk.fingerprint_raw();
    let fp_b = bob_ik_pk.fingerprint_raw();

    // Alice verifies Bob's bundle and initiates session.
    let verified = verify_bundle(bob_bundle, &bob_ik_pk).unwrap();
    let mut initiated = initiate_session(&alice_ik_pk, &alice_ik_sk, &verified).unwrap();

    // Build first message AAD.
    let aad = build_first_message_aad(&fp_a, &fp_b, &initiated.session_init).unwrap();

    // Alice encrypts first message, consuming the initial chain key.
    let (first_ct, alice_ck) =
        RatchetState::encrypt_first_message(initiated.take_initial_chain_key(), b"hello bob", &aad)
            .unwrap();

    // Bob receives the session — sender_sig proves Alice initiated (not an impersonator).
    let mut received = receive_session(
        &bob_ik_pk,
        &bob_ik_sk,
        &alice_ik_pk,
        &initiated.session_init,
        &initiated.sender_sig,
        &bob_spk_sk,
        opk_sk.as_ref(),
    )
    .unwrap();

    // Bob decrypts the first message, consuming the initial chain key.
    let (first_pt, bob_ck) =
        RatchetState::decrypt_first_message(received.take_initial_chain_key(), &first_ct, &aad)
            .unwrap();
    assert_eq!(&*first_pt, b"hello bob");

    // Chain keys must match after first message.
    assert_eq!(*alice_ck, *bob_ck);

    // Extract keys via take_ methods — each replaces the internal value with
    // zeros, enforcing single-use. ek_pk/peer_ek are pub (non-secret).
    let ek_pk = xwing::PublicKey::from_bytes(initiated.ek_pk.as_bytes().to_vec()).unwrap();
    let ek_sk = xwing::SecretKey::from_bytes(initiated.ek_sk().as_bytes().to_vec()).unwrap();
    let peer_ek = xwing::PublicKey::from_bytes(received.peer_ek.as_bytes().to_vec()).unwrap();
    let alice = RatchetState::init_alice(
        *initiated.take_root_key(),
        *alice_ck,
        fp_a,
        fp_b,
        ek_pk,
        ek_sk,
    )
    .unwrap();
    let bob =
        RatchetState::init_bob(*received.take_root_key(), *bob_ck, fp_b, fp_a, peer_ek).unwrap();

    (alice, bob, fp_a, fp_b)
}

/// Helper: Alice sends to Bob.
fn send_a_to_b(alice: &mut RatchetState, bob: &mut RatchetState, msg: &[u8]) -> Vec<u8> {
    let enc = alice.encrypt(msg).unwrap();
    let pt = bob.decrypt(&enc.header, &enc.ciphertext).unwrap();
    pt.to_vec()
}

/// Helper: Bob sends to Alice.
fn send_b_to_a(alice: &mut RatchetState, bob: &mut RatchetState, msg: &[u8]) -> Vec<u8> {
    let enc = bob.encrypt(msg).unwrap();
    let pt = alice.decrypt(&enc.header, &enc.ciphertext).unwrap();
    pt.to_vec()
}

#[test]
fn full_session_lifecycle() {
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Exchange a few messages.
    let pt = send_a_to_b(&mut alice, &mut bob, b"msg1");
    assert_eq!(pt, b"msg1");
    let pt = send_b_to_a(&mut alice, &mut bob, b"msg2");
    assert_eq!(pt, b"msg2");

    // Serialize both sides.
    let alice_bytes = alice.to_bytes().unwrap().0;
    let bob_bytes = bob.to_bytes().unwrap().0;

    // Deserialize.
    let mut alice2 = RatchetState::from_bytes(&alice_bytes).unwrap();
    let mut bob2 = RatchetState::from_bytes(&bob_bytes).unwrap();

    // Continue conversation after deserialization.
    let pt = send_a_to_b(&mut alice2, &mut bob2, b"post-restore");
    assert_eq!(pt, b"post-restore");
}

#[test]
fn basic_session_with_opk_enabled() {
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);
    let pt = send_a_to_b(&mut alice, &mut bob, b"with opk");
    assert_eq!(pt, b"with opk");
    let pt = send_b_to_a(&mut alice, &mut bob, b"reply");
    assert_eq!(pt, b"reply");
}

#[test]
fn full_session_without_opk() {
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(false);
    let pt = send_a_to_b(&mut alice, &mut bob, b"no opk");
    assert_eq!(pt, b"no opk");
    let pt = send_b_to_a(&mut alice, &mut bob, b"reply");
    assert_eq!(pt, b"reply");
}

#[test]
fn bidirectional_ratchet() {
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);
    // Multiple ratchet epochs: direction changes trigger DH ratchet steps.
    for i in 0..10u32 {
        if i % 2 == 0 {
            let msg = format!("a-to-b #{i}");
            let pt = send_a_to_b(&mut alice, &mut bob, msg.as_bytes());
            assert_eq!(pt, msg.as_bytes());
        } else {
            let msg = format!("b-to-a #{i}");
            let pt = send_b_to_a(&mut alice, &mut bob, msg.as_bytes());
            assert_eq!(pt, msg.as_bytes());
        }
    }
}

#[test]
fn out_of_order_across_ratchet_steps() {
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Alice sends 3 messages in the same epoch.
    let enc0 = alice.encrypt(b"msg0").unwrap();
    let enc1 = alice.encrypt(b"msg1").unwrap();
    let enc2 = alice.encrypt(b"msg2").unwrap();

    // Bob receives msg2 first (skipping 0 and 1).
    let pt2 = bob.decrypt(&enc2.header, &enc2.ciphertext).unwrap();
    assert_eq!(&*pt2, b"msg2");

    // Bob sends back (triggers ratchet step on both sides when Alice receives).
    let enc_b = bob.encrypt(b"bob reply").unwrap();
    let pt_b = alice.decrypt(&enc_b.header, &enc_b.ciphertext).unwrap();
    assert_eq!(&*pt_b, b"bob reply");

    // Bob can still decrypt the earlier messages via prev_recv_epoch_key grace period.
    let pt0 = bob.decrypt(&enc0.header, &enc0.ciphertext).unwrap();
    assert_eq!(&*pt0, b"msg0");
    let pt1 = bob.decrypt(&enc1.header, &enc1.ciphertext).unwrap();
    assert_eq!(&*pt1, b"msg1");
}

#[test]
fn serialize_with_out_of_order_then_consume() {
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Alice sends 3 messages.
    let enc0 = alice.encrypt(b"msg0").unwrap();
    let enc1 = alice.encrypt(b"msg1").unwrap();
    let enc2 = alice.encrypt(b"msg2").unwrap();

    // Bob receives only msg2 (out of order).
    bob.decrypt(&enc2.header, &enc2.ciphertext).unwrap();

    // Serialize and deserialize Bob.
    let bob_bytes = bob.to_bytes().unwrap().0;
    let mut bob2 = RatchetState::from_bytes(&bob_bytes).unwrap();

    // Consume the out-of-order messages via counter-mode derivation.
    let pt0 = bob2.decrypt(&enc0.header, &enc0.ciphertext).unwrap();
    assert_eq!(&*pt0, b"msg0");
    let pt1 = bob2.decrypt(&enc1.header, &enc1.ciphertext).unwrap();
    assert_eq!(&*pt1, b"msg1");
}

#[test]
fn session_reset_reestablishment() {
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);
    send_a_to_b(&mut alice, &mut bob, b"before reset");

    // Reset both sides.
    alice.reset();
    bob.reset();

    // Old states are unusable.
    assert!(alice.encrypt(b"test").is_err());
    assert!(bob.encrypt(b"test").is_err());

    // Establish a completely new session.
    let (mut alice2, mut bob2, _fp_a2, _fp_b2) = do_kex(false);
    let pt = send_a_to_b(&mut alice2, &mut bob2, b"new session");
    assert_eq!(pt, b"new session");
}

#[test]
fn long_conversation_stress() {
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);
    for i in 0..120u32 {
        let pt = if i % 5 < 3 {
            send_a_to_b(&mut alice, &mut bob, &i.to_be_bytes())
        } else {
            send_b_to_a(&mut alice, &mut bob, &i.to_be_bytes())
        };
        // Serialize mid-conversation.
        if i == 60 {
            let a_bytes = alice.to_bytes().unwrap().0;
            let b_bytes = bob.to_bytes().unwrap().0;
            alice = RatchetState::from_bytes(&a_bytes).unwrap();
            bob = RatchetState::from_bytes(&b_bytes).unwrap();
        }
        // Verify plaintext content for the first message after the serialization
        // round-trip to confirm state survives serialization without corrupting
        // message content (not just AEAD authentication).
        if i == 61 {
            assert_eq!(
                pt,
                (61u32).to_be_bytes(),
                "plaintext must survive serialization round-trip"
            );
        }
    }
}

#[test]
fn ratchet_derive_call_keys_delegates() {
    let (alice, bob, fp_a, fp_b) = do_kex(true);
    let kem_ss: [u8; 32] = soliton::primitives::random::random_array();
    let call_id: [u8; 16] = soliton::primitives::random::random_array();

    // Cross-party agreement: alice's send == bob's recv and vice versa.
    let alice_keys = alice.derive_call_keys(&kem_ss, &call_id).unwrap();
    let bob_keys = bob.derive_call_keys(&kem_ss, &call_id).unwrap();
    assert_eq!(alice_keys.send_key(), bob_keys.recv_key());
    assert_eq!(alice_keys.recv_key(), bob_keys.send_key());

    // White-box: verify RatchetState::derive_call_keys truly delegates to
    // call::derive_call_keys with self.root_key and stored fingerprints.
    // Both parties share the same root key after KEX, so verify with Alice's side.
    #[allow(deprecated)]
    let root_key = alice.root_key_bytes();
    let direct =
        soliton::call::derive_call_keys(root_key, &kem_ss, &call_id, &fp_a, &fp_b).unwrap();
    assert_eq!(alice_keys.send_key(), direct.send_key());
    assert_eq!(alice_keys.recv_key(), direct.recv_key());
}

#[test]
fn ratchet_mismatched_fingerprints_returns_aead_failed() {
    // Fingerprints are bound at init time. If Alice and Bob are initialized
    // with inconsistent fingerprint orderings, the AAD won't match and
    // decryption fails with AeadFailed.
    use soliton::error::Error;

    let (ek_pk, ek_sk) = xwing::keygen().unwrap();
    let rk = [0x11u8; 32];
    let ck = [0x22u8; 32];
    let fp_a = [0xAAu8; 32];
    let fp_b = [0xBBu8; 32];

    // Alice: local=fp_a, remote=fp_b (correct)
    let mut alice = RatchetState::init_alice(rk, ck, fp_a, fp_b, ek_pk.clone(), ek_sk).unwrap();
    // Bob initialized with WRONG fingerprint ordering (swapped local/remote)
    let mut bad_bob = RatchetState::init_bob(rk, ck, fp_a, fp_b, ek_pk).unwrap();

    let enc = alice.encrypt(b"secret").unwrap();

    // Bad Bob's AAD uses remote_fp=fp_b as sender, local_fp=fp_a as recipient,
    // producing AAD = fp_b || fp_a. Alice used fp_a || fp_b → mismatch.
    let result = bad_bob.decrypt(&enc.header, &enc.ciphertext);
    assert!(
        matches!(result, Err(Error::AeadFailed)),
        "mismatched fingerprint ordering must return AeadFailed, got: {:?}",
        result,
    );
}

#[test]
fn reflection_attack_rejected() {
    // Abstract.md Theorem 12: A→B ciphertext reflected back to A as B→A must
    // fail. The fingerprint ordering in AAD differs between encrypt (local_fp
    // || remote_fp) and decrypt (remote_fp || local_fp), so the AEAD tag
    // cannot authenticate. In practice, the reflection may fail earlier (e.g.,
    // InvalidData from KEM decapsulation of a ciphertext encapsulated to the
    // wrong key) — the exact error variant depends on ratchet state. The
    // security property is that decryption never succeeds.

    let (mut alice, _bob, _fp_a, _fp_b) = do_kex(true);
    let enc = alice.encrypt(b"hello bob").unwrap();

    // Alice tries to decrypt her own ciphertext (reflection).
    let result = alice.decrypt(&enc.header, &enc.ciphertext);
    assert!(
        result.is_err(),
        "reflected ciphertext must be rejected, got: Ok(...)",
    );
}

#[test]
fn bob_sends_before_receiving() {
    // Bob comes online and sends before any of Alice's messages arrive.
    // This exercises init_bob → immediate encrypt with ratchet_pending=true
    // and recv_ratchet_pk = pk_EK from KEX (the realistic network-delay scenario).
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Alice sends messages Bob hasn't received yet.
    let enc_a0 = alice.encrypt(b"alice msg 0").unwrap();
    let enc_a1 = alice.encrypt(b"alice msg 1").unwrap();

    // Bob sends first — triggers KEM ratchet from his pending state.
    let pt = send_b_to_a(&mut alice, &mut bob, b"bob first");
    assert_eq!(pt, b"bob first");

    // Alice's messages still decrypt: both sides start counters at 1 (post-KEX),
    // so enc_a0 has n=1 matching Bob's recv_count=1 — no skipping needed.
    let pt0 = bob.decrypt(&enc_a0.header, &enc_a0.ciphertext).unwrap();
    assert_eq!(&*pt0, b"alice msg 0");
    let pt1 = bob.decrypt(&enc_a1.header, &enc_a1.ciphertext).unwrap();
    assert_eq!(&*pt1, b"alice msg 1");

    // Conversation continues normally.
    let pt = send_a_to_b(&mut alice, &mut bob, b"after catch-up");
    assert_eq!(pt, b"after catch-up");
}

#[test]
fn ratchet_new_epoch_without_kem_ct_returns_invalid_data() {
    use soliton::error::Error;
    use soliton::ratchet::RatchetHeader;

    // Alice's initial state has recv_ratchet_pk = None. Any incoming ratchet_pk
    // is treated as a new epoch (KEM ratchet step). Without a KEM ciphertext,
    // the ratchet step cannot proceed → InvalidData.
    let (ek_pk, ek_sk) = xwing::keygen().unwrap();
    let rk = [0x11u8; 32];
    let ck = [0x22u8; 32];
    let fp_a = [0xAAu8; 32];
    let fp_b = [0xBBu8; 32];
    let mut alice = RatchetState::init_alice(rk, ck, fp_a, fp_b, ek_pk, ek_sk).unwrap();

    let dummy_pk = xwing::PublicKey::from_bytes(vec![0u8; 1216]).unwrap();
    let header = RatchetHeader {
        ratchet_pk: dummy_pk,
        kem_ct: None,
        n: 0,
        pn: 1,
    };

    let result = alice.decrypt(&header, &[]);
    assert!(
        matches!(result, Err(Error::InvalidData)),
        "new-epoch message without kem_ct must return InvalidData, got: {:?}",
        result,
    );
}

// ── Security property tests (RT-540, RT-541, RT-542) ──────────────────

#[test]
fn forward_secrecy_old_epoch_keys_unrecoverable() {
    // RT-540: After a KEM ratchet step, message keys from the old epoch
    // must be unrecoverable. Verify by capturing a ciphertext, advancing
    // the ratchet, then confirming the old ciphertext cannot be replayed
    // against the new state.
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Epoch 0: Alice sends, Bob receives.
    let enc0 = alice.encrypt(b"epoch0 msg").unwrap();
    bob.decrypt(&enc0.header, &enc0.ciphertext).unwrap();

    // Trigger KEM ratchet: Bob replies, Alice receives → new epoch.
    let enc_b = bob.encrypt(b"reply").unwrap();
    alice.decrypt(&enc_b.header, &enc_b.ciphertext).unwrap();

    // Another ratchet step: Alice sends again → epoch advances again.
    let enc_a2 = alice.encrypt(b"epoch2 msg").unwrap();
    bob.decrypt(&enc_a2.header, &enc_a2.ciphertext).unwrap();

    // Bob replies again → second ratchet step. prev_recv_epoch_key from
    // epoch 0 is now gone (only one-epoch grace period).
    let enc_b2 = bob.encrypt(b"reply2").unwrap();
    alice.decrypt(&enc_b2.header, &enc_b2.ciphertext).unwrap();

    // Serialize Bob to get a snapshot of current state.
    let bob_bytes = bob.to_bytes().unwrap().0;
    let mut bob_replay = RatchetState::from_bytes(&bob_bytes).unwrap();

    // Attempt to decrypt the epoch-0 message — must fail. The epoch-0
    // epoch key has been overwritten by two subsequent KEM ratchet steps.
    // The replayed message's ratchet_pk matches prev_recv_ratchet_pk
    // (one-epoch grace period), so AEAD succeeds with the retained key.
    // Counter n=0 is already in prev_recv_seen → DuplicateMessage.
    let result = bob_replay.decrypt(&enc0.header, &enc0.ciphertext);
    assert!(
        matches!(result, Err(soliton::error::Error::DuplicateMessage)),
        "old epoch replay must be caught as duplicate"
    );
}

#[test]
fn cross_session_isolation() {
    // RT-541: Two independent KEX sessions from the same identities must
    // produce completely independent key material.
    let (mut alice1, mut bob1, _fp_a1, _fp_b1) = do_kex(true);
    let (mut alice2, mut bob2, _fp_a2, _fp_b2) = do_kex(true);

    // Encrypt the same plaintext in both sessions.
    let enc1 = alice1.encrypt(b"same plaintext").unwrap();
    let enc2 = alice2.encrypt(b"same plaintext").unwrap();

    // Ciphertexts must differ (different keys, different nonces from
    // independent ratchet states).
    assert_ne!(
        enc1.ciphertext, enc2.ciphertext,
        "two sessions must produce different ciphertexts for the same plaintext"
    );

    // Cross-decrypt must fail: session 1's ciphertext cannot decrypt
    // under session 2's state.
    let result = bob2.decrypt(&enc1.header, &enc1.ciphertext);
    assert!(result.is_err(), "cross-session decryption must fail");

    // Verify session 1 still works independently.
    let pt1 = bob1.decrypt(&enc1.header, &enc1.ciphertext).unwrap();
    assert_eq!(&*pt1, b"same plaintext");
}

#[test]
fn break_in_recovery_via_kem_ratchet() {
    // RT-542: After a KEM ratchet step heals a compromised state, the
    // attacker's snapshot becomes useless for future messages.
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Exchange initial messages to establish a working session.
    send_a_to_b(&mut alice, &mut bob, b"setup1");
    send_b_to_a(&mut alice, &mut bob, b"setup2");

    // Attacker captures Bob's state (simulating key compromise).
    let compromised_bob_bytes = bob.to_bytes().unwrap().0;

    // Legitimate conversation continues — triggers KEM ratchet steps.
    let mut bob = RatchetState::from_bytes(&compromised_bob_bytes).unwrap();
    send_a_to_b(&mut alice, &mut bob, b"heal1");
    send_b_to_a(&mut alice, &mut bob, b"heal2");
    // One more round to push the compromised epoch out of grace period.
    send_a_to_b(&mut alice, &mut bob, b"heal3");
    send_b_to_a(&mut alice, &mut bob, b"heal4");

    // New message after healing ratchet steps.
    let enc_post_heal = alice.encrypt(b"secret after healing").unwrap();
    bob.decrypt(&enc_post_heal.header, &enc_post_heal.ciphertext)
        .unwrap();

    // Attacker's compromised snapshot cannot decrypt the post-heal message.
    let mut attacker_bob = RatchetState::from_bytes(&compromised_bob_bytes).unwrap();
    let result = attacker_bob.decrypt(&enc_post_heal.header, &enc_post_heal.ciphertext);
    assert!(
        result.is_err(),
        "compromised state must not decrypt messages after KEM ratchet healing"
    );
}

// ── Ratchet property tests (RT-546, RT-547, RT-548, RT-557, RT-558) ───

#[test]
fn replay_rejected_across_serialization_boundary() {
    // RT-546: Replay detection must survive serialization round-trip.
    use soliton::error::Error;

    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);
    let enc = alice.encrypt(b"unique msg").unwrap();

    // Bob decrypts once.
    bob.decrypt(&enc.header, &enc.ciphertext).unwrap();

    // Serialize and deserialize Bob.
    let bob_bytes = bob.to_bytes().unwrap().0;
    let mut bob2 = RatchetState::from_bytes(&bob_bytes).unwrap();

    // Replay the same message — must be rejected as duplicate.
    let result = bob2.decrypt(&enc.header, &enc.ciphertext);
    assert!(
        matches!(result, Err(Error::DuplicateMessage)),
        "replay after serialization must return DuplicateMessage, got: {:?}",
        result,
    );
}

#[test]
fn recv_seen_persists_across_serialization() {
    // RT-547: recv_seen set must survive serialization round-trip.
    // Send messages out of order, serialize, then verify the gaps are
    // still tracked and late messages still decrypt.
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    let enc0 = alice.encrypt(b"msg0").unwrap();
    let enc1 = alice.encrypt(b"msg1").unwrap();
    let enc2 = alice.encrypt(b"msg2").unwrap();

    // Bob receives 0 and 2, skipping 1.
    bob.decrypt(&enc0.header, &enc0.ciphertext).unwrap();
    bob.decrypt(&enc2.header, &enc2.ciphertext).unwrap();

    // Serialize/deserialize.
    let bob_bytes = bob.to_bytes().unwrap().0;
    let mut bob2 = RatchetState::from_bytes(&bob_bytes).unwrap();

    // msg1 (the gap) must still be decryptable.
    let pt1 = bob2.decrypt(&enc1.header, &enc1.ciphertext).unwrap();
    assert_eq!(&*pt1, b"msg1");

    // msg0 and msg2 must be rejected as duplicates.
    use soliton::error::Error;
    assert!(matches!(
        bob2.decrypt(&enc0.header, &enc0.ciphertext),
        Err(Error::DuplicateMessage)
    ));
    assert!(matches!(
        bob2.decrypt(&enc2.header, &enc2.ciphertext),
        Err(Error::DuplicateMessage)
    ));
}

#[test]
fn prev_recv_epoch_key_grace_period() {
    // RT-548: Messages from the previous epoch decrypt after one KEM
    // ratchet step (grace period) but fail after two (grace period expired).
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Alice sends messages in epoch 0.
    let enc_old = alice.encrypt(b"old epoch msg").unwrap();

    // Trigger one KEM ratchet step: Bob replies → new epoch.
    let enc_b = bob.encrypt(b"trigger ratchet").unwrap();
    alice.decrypt(&enc_b.header, &enc_b.ciphertext).unwrap();

    // Bob has NOT received enc_old yet. One ratchet step has occurred.
    // Grace period: old-epoch messages still decrypt.
    let pt = bob.decrypt(&enc_old.header, &enc_old.ciphertext).unwrap();
    assert_eq!(&*pt, b"old epoch msg");

    // Now trigger a SECOND ratchet step.
    let (mut alice2, mut bob2, _, _) = do_kex(true);
    let enc_old2 = alice2.encrypt(b"old epoch 2").unwrap();

    // First ratchet step.
    let enc_b2 = bob2.encrypt(b"ratchet1").unwrap();
    alice2.decrypt(&enc_b2.header, &enc_b2.ciphertext).unwrap();
    let enc_a3 = alice2.encrypt(b"ratchet1 reply").unwrap();
    bob2.decrypt(&enc_a3.header, &enc_a3.ciphertext).unwrap();

    // Second ratchet step.
    let enc_b3 = bob2.encrypt(b"ratchet2").unwrap();
    alice2.decrypt(&enc_b3.header, &enc_b3.ciphertext).unwrap();
    let enc_a4 = alice2.encrypt(b"ratchet2 reply").unwrap();
    bob2.decrypt(&enc_a4.header, &enc_a4.ciphertext).unwrap();

    // Now try to decrypt the epoch-0 message — grace period expired.
    let result = bob2.decrypt(&enc_old2.header, &enc_old2.ciphertext);
    assert!(
        result.is_err(),
        "old epoch message must fail after two ratchet steps (grace period expired), got: {:?}",
        result,
    );
}

#[test]
fn high_counter_encrypt_decrypt_roundtrip() {
    // RT-557: Verify many messages can be sent and received within a single
    // epoch, including out-of-order delivery at various counter values.
    // The exact u32::MAX-1 edge case is tested by the unit test
    // `encrypt_at_send_count_max_minus_one_succeeds` which has direct
    // field access. This integration test verifies the contract end-to-end.
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Send many messages in one direction (same epoch, counter climbs).
    let mut pending = Vec::new();
    for i in 0..50u32 {
        let msg = format!("msg-{i}");
        pending.push((alice.encrypt(msg.as_bytes()).unwrap(), msg));
    }

    // Deliver out of order: last first, then remaining.
    let (last_enc, last_msg) = pending.pop().unwrap();
    let pt = bob.decrypt(&last_enc.header, &last_enc.ciphertext).unwrap();
    assert_eq!(&*pt, last_msg.as_bytes());

    // Deliver rest in forward order (counter-mode key derivation for each).
    for (enc, msg) in &pending {
        let pt = bob.decrypt(&enc.header, &enc.ciphertext).unwrap();
        assert_eq!(&*pt, msg.as_bytes());
    }
}

#[test]
fn multi_epoch_stress_test() {
    // RT-558: Exercise >3 consecutive KEM ratchet steps in both directions.
    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    for epoch in 0..8u32 {
        let msg_ab = format!("a→b epoch {epoch}");
        let pt = send_a_to_b(&mut alice, &mut bob, msg_ab.as_bytes());
        assert_eq!(pt, msg_ab.as_bytes());

        let msg_ba = format!("b→a epoch {epoch}");
        let pt = send_b_to_a(&mut alice, &mut bob, msg_ba.as_bytes());
        assert_eq!(pt, msg_ba.as_bytes());
    }

    // Verify serialization still works after many epochs.
    let alice_bytes = alice.to_bytes().unwrap().0;
    let bob_bytes = bob.to_bytes().unwrap().0;
    let mut alice2 = RatchetState::from_bytes(&alice_bytes).unwrap();
    let mut bob2 = RatchetState::from_bytes(&bob_bytes).unwrap();

    let pt = send_a_to_b(&mut alice2, &mut bob2, b"post-multi-epoch");
    assert_eq!(pt, b"post-multi-epoch");
}

#[test]
fn epoch_isolation_survives_serialization() {
    // RT-985: Verify that epoch isolation holds across a serialization boundary.
    // Serialize Bob at epoch N, advance to epoch N+2, deserialize the epoch-N
    // snapshot, and verify it cannot decrypt epoch-N+2 messages.

    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Epoch 0: exchange messages, then serialize Bob.
    send_a_to_b(&mut alice, &mut bob, b"epoch0");
    // to_bytes consumes self — save the snapshot, then reload Bob to continue.
    let bob_epoch0_bytes = bob.to_bytes().unwrap().0;
    let mut bob = RatchetState::from_bytes(&bob_epoch0_bytes).unwrap();

    // Advance to epoch 1: Bob sends → Alice receives → Alice sends → Bob receives.
    let enc_b = bob.encrypt(b"trigger ratchet 1").unwrap();
    alice.decrypt(&enc_b.header, &enc_b.ciphertext).unwrap();
    let enc_a = alice.encrypt(b"epoch1 reply").unwrap();
    bob.decrypt(&enc_a.header, &enc_a.ciphertext).unwrap();

    // Advance to epoch 2: Bob sends → Alice receives → Alice sends → Bob receives.
    let enc_b2 = bob.encrypt(b"trigger ratchet 2").unwrap();
    alice.decrypt(&enc_b2.header, &enc_b2.ciphertext).unwrap();
    let enc_a2 = alice.encrypt(b"epoch2 msg").unwrap();
    bob.decrypt(&enc_a2.header, &enc_a2.ciphertext).unwrap();

    // Deserialize the epoch-0 snapshot and try to decrypt the epoch-2 message.
    let mut bob_old = RatchetState::from_bytes(&bob_epoch0_bytes).unwrap();
    let result = bob_old.decrypt(&enc_a2.header, &enc_a2.ciphertext);
    assert!(
        result.is_err(),
        "epoch-0 snapshot must not decrypt epoch-2 message, got: {:?}",
        result,
    );
}

#[test]
fn prev_recv_seen_survives_serialization() {
    // RT-990: Verify prev_recv_seen duplicate detection survives serialization.
    // A previous-epoch message that was already decrypted must be rejected as
    // duplicate after serialize/deserialize.
    use soliton::error::Error;

    let (mut alice, mut bob, _fp_a, _fp_b) = do_kex(true);

    // Alice sends two messages in epoch 0.
    let enc0 = alice.encrypt(b"prev-epoch msg0").unwrap();
    let enc1 = alice.encrypt(b"prev-epoch msg1").unwrap();

    // Bob decrypts msg0 in epoch 0 (msg1 held back).
    bob.decrypt(&enc0.header, &enc0.ciphertext).unwrap();

    // Trigger Bob's receive-side KEM ratchet: Bob sends → Alice receives →
    // Alice sends with new ratchet_pk → Bob receives (triggers KEM ratchet,
    // moving recv_seen → prev_recv_seen).
    let enc_b = bob.encrypt(b"trigger ratchet").unwrap();
    alice.decrypt(&enc_b.header, &enc_b.ciphertext).unwrap();
    let enc_a_new = alice.encrypt(b"new epoch reply").unwrap();
    bob.decrypt(&enc_a_new.header, &enc_a_new.ciphertext)
        .unwrap();

    // Bob now decrypts msg1 from the previous epoch (grace period).
    // msg1's counter is added to prev_recv_seen.
    bob.decrypt(&enc1.header, &enc1.ciphertext).unwrap();

    // Serialize/deserialize Bob — prev_recv_seen should contain msg0 and msg1.
    let bob_bytes = bob.to_bytes().unwrap().0;
    let mut bob2 = RatchetState::from_bytes(&bob_bytes).unwrap();

    // Replay msg0 from the previous epoch — must be rejected as duplicate.
    let result = bob2.decrypt(&enc0.header, &enc0.ciphertext);
    assert!(
        matches!(result, Err(Error::DuplicateMessage)),
        "prev-epoch replay after serialization must return DuplicateMessage, got: {:?}",
        result,
    );

    // Replay msg1 from the previous epoch — also duplicate.
    let result = bob2.decrypt(&enc1.header, &enc1.ciphertext);
    assert!(
        matches!(result, Err(Error::DuplicateMessage)),
        "prev-epoch replay after serialization must return DuplicateMessage, got: {:?}",
        result,
    );
}