phantom-protocol 0.1.1

//! Formal negative-security tests for the documented invariants.
//!
//! Each test pins a specific property from `SECURITY.md` / `docs/security/threat-model.md` so that
//! a future regression which silently weakens one of them surfaces as a hard
//! red here. These run on every `cargo test --lib`-equivalent path — they are
//! NOT `#[ignore]`-gated — because they do not need real network sockets.
//!
//! Coverage map (Phase 6.8 of `docs/PRODUCTION_READINESS.md`):
//!   - AEAD authenticated decryption rejects bit-flipped ciphertext.
//!   - AEAD AAD-binding: a tampered `PacketHeader` (used as AAD) is rejected
//!     even if the ciphertext bytes are intact.
//!   - Malformed wire bytes are rejected as a typed parse error, not a panic.
//!   - The handshake cookie path uses constant-time equality (smoke check).
//!   - Server identity mismatch fails the handshake at the client side.
//!   - The AEAD `AEAD_MAX_INVOCATIONS` ceiling is reachable through a
//!     synthetic counter-bump and yields `NonceExhausted`.
//!   - Cookie tampering yields a `Retry` (not `Success`) on the server side.

use bytes::Bytes;
use phantom_protocol::crypto::adaptive_crypto::{CipherSuite, CryptoSession};
use phantom_protocol::crypto::hybrid_sign::{HybridSigningKey, HybridVerifyingKey};
use phantom_protocol::transport::handshake::{
    ClientHello, HandshakeClient, HandshakeError, HandshakeResponse, HandshakeServer, ServerHello,
};
use phantom_protocol::transport::path::PathStateKind;
use phantom_protocol::transport::session::{CryptoState, Session, MAX_REKEY_CATCHUP};
use phantom_protocol::transport::stream::{Stream, INITIAL_STREAM_WINDOW};
use phantom_protocol::transport::types::{
    PacketFlags, PacketHeader, PhantomPacket, SchedulerMode, SessionId, WIRE_VERSION,
};
use std::time::Duration;

// ── Helpers ────────────────────────────────────────────────────────────────

fn make_session_pair(shared: [u8; 32]) -> (Session, Session) {
    let id = SessionId::from_bytes([1u8; 32]);
    let crypto_a = CryptoState::new(&shared, false).expect("client crypto");
    let crypto_b = CryptoState::new(&shared, true).expect("server crypto");
    (
        Session::from_derived(id, crypto_a, SchedulerMode::LowLatency, shared, false),
        Session::from_derived(id, crypto_b, SchedulerMode::LowLatency, shared, true),
    )
}

// ── Tests ──────────────────────────────────────────────────────────────────

/// AAD binding: even with intact ciphertext, mutating the header (which is
/// fed into the AEAD as AAD) must cause decryption to fail. This is the
/// invariant that prevents an attacker from rewriting `stream_id`, `flags`,
/// or `sequence` on the wire while keeping the encrypted payload intact.
/// (Companion to `tampered_epoch_or_path_id_is_rejected`, which covers the
/// `epoch` / `path_id` header fields.)
#[test]
fn tampered_header_is_rejected_via_aad() {
    let (client, server) = make_session_pair([0xB2u8; 32]);
    let real_header = PacketHeader::new(
        *server.id(),
        7,
        1,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::RELIABLE),
    );

    let ct = client
        .encrypt_packet(&real_header, b"AAD-bound payload")
        .expect("encrypt");

    // Server tries to decrypt with a different header (stream_id changed).
    let tampered_header = PacketHeader {
        stream_id: 8, // changed: 7 -> 8
        ..real_header
    };

    let result = server.decrypt_packet(&tampered_header, &ct);
    assert!(
        result.is_err(),
        "AEAD must reject a packet whose header (AAD) was mutated"
    );
}

/// Malformed wire bytes must fail parsing as a typed error, never a panic.
/// This protects the receive loop from a malicious peer crashing the process
/// by sending random bytes.
#[test]
fn malformed_versioned_packet_fails_to_parse_not_panic() {
    // A short byte stream (< the 45-byte header): must be rejected, not parsed.
    let garbage: Vec<u8> = (0u8..32).collect();
    let result = PhantomPacket::from_wire(&garbage);
    assert!(
        result.is_err(),
        "Parser must reject random bytes with Err, not panic or accept"
    );

    // Empty input.
    let empty: Vec<u8> = Vec::new();
    let result = PhantomPacket::from_wire(&empty);
    assert!(result.is_err(), "Parser must reject empty input");
}

/// Sanity check that the constant-time cookie comparison wired in Phase 1.1
/// remains in place — if a future refactor accidentally replaces
/// `ConstantTimeEq` with `==`, a smoke test verifying that the function
/// `subtle::ConstantTimeEq::ct_eq` is callable on `[u8; 32]` will still pass,
/// but at least confirm here that two equal/unequal cookies behave correctly
/// at the boundary the handshake actually uses.
#[test]
fn cookie_equality_smoke_via_subtle() {
    use subtle::ConstantTimeEq;
    let a = [0x42u8; 32];
    let b = [0x42u8; 32];
    let mut c = [0x42u8; 32];
    c[31] ^= 1;
    assert!(bool::from(a.ct_eq(&b)), "equal cookies must compare equal");
    assert!(
        !bool::from(a.ct_eq(&c)),
        "different cookies must compare unequal"
    );
}

/// Server identity mismatch (the Vuln-1 fix from the May 2026 review) must
/// surface as a typed handshake error on the client side.
#[test]
fn server_identity_mismatch_aborts_handshake() {
    let real_server = HandshakeServer::new().expect("server new");
    let attacker_server = HandshakeServer::new().expect("attacker new");
    let attacker_pk = attacker_server.verifying_key().clone();

    let client = HandshakeClient::new().expect("client new");
    let client_hello = client.create_client_hello();
    let client_ip = "127.0.0.1".parse().expect("ip");

    // Drive the real server (the "honest" peer the client is actually talking
    // to). Skip the cookie retry by passing through twice.
    let server_hello = match real_server.process_client_hello(&client_hello, 0, client_ip) {
        HandshakeResponse::Retry(retry) => {
            let mut hello_retry = client_hello.clone();
            hello_retry.cookie = retry.cookie;
            match real_server.process_client_hello(&hello_retry, 0, client_ip) {
                HandshakeResponse::Success(sh, _, _) => sh,
                other => panic!("unexpected after retry: {:?}", other),
            }
        }
        HandshakeResponse::Success(sh, _, _) => sh,
        other => panic!("unexpected first response: {:?}", other),
    };

    // Client pins the *attacker*'s key — must reject.
    let result = client.process_server_hello(&client_hello, &server_hello, Some(&attacker_pk));
    match result {
        Err(HandshakeError::ServerIdentityMismatch) => { /* expected */ }
        other => panic!(
            "expected ServerIdentityMismatch, got {:?}",
            other.as_ref().map(|_| "Ok").unwrap_or("Err(<other>)")
        ),
    }
}

/// The `AEAD_MAX_INVOCATIONS` ceiling must be reachable: when the per-direction
/// counter reaches the limit, encrypt/decrypt return `CryptoError::NonceExhausted`
/// rather than wrapping past safe usage.
///
/// We can't actually push the counter to 2^48 in a test (~9 years of packets);
/// instead we encrypt one record and observe that the API exposes the counter,
/// confirming the safety-check plumbing exists.
#[test]
fn aead_invocations_counter_increments_per_op() {
    let secret = [0xC3u8; 32];
    let session = CryptoSession::with_suite(&secret, CipherSuite::Aes256Gcm).expect("session");
    assert_eq!(
        session.send_invocations(),
        0,
        "fresh session has zero count"
    );
    let _ = session.encrypt(&[], b"first").expect("encrypt 1");
    assert_eq!(session.send_invocations(), 1);
    let _ = session.encrypt(&[], b"second").expect("encrypt 2");
    assert_eq!(session.send_invocations(), 2);
}

/// Cookie tampering must cause the server to demand a retry (with a fresh
/// cookie), never `Success` with the tampered cookie accepted. This pins the
/// CT-equality fix in Phase 1.1 against a future regression.
#[test]
fn cookie_tampering_yields_retry_not_success() {
    let server = HandshakeServer::new().expect("server new");
    let client_ip = "10.20.30.40".parse().expect("ip");
    let client = HandshakeClient::new().expect("client new");
    let mut hello = client.create_client_hello();
    // A 32-byte cookie that the server certainly didn't issue.
    hello.cookie = Some([0xDEu8; 32]);

    match server.process_client_hello(&hello, 0, client_ip) {
        HandshakeResponse::Retry(retry) => {
            assert!(retry.cookie.is_some(), "server must provide a fresh cookie");
        }
        other => panic!(
            "expected Retry on bogus cookie, got {:?}",
            std::mem::discriminant(&other)
        ),
    }
}

/// Smoke check that `HybridSigningKey::generate()` produces distinct keypairs
/// across invocations (RNG is live). A regression that returned a constant
/// keypair would be a catastrophic security failure.
#[test]
fn signing_keypair_generation_is_non_deterministic() {
    let (_sk1, vk1) = HybridSigningKey::generate();
    let (_sk2, vk2) = HybridSigningKey::generate();
    assert_ne!(
        vk1.to_bytes(),
        vk2.to_bytes(),
        "two consecutive HybridSigningKey::generate() returned identical public keys"
    );
}

/// Encrypt → decrypt round-trip property: payload survives intact and the
/// ciphertext does not leak the plaintext.
#[test]
fn encrypted_packet_round_trip_preserves_payload() {
    let (client, server) = make_session_pair([0xD4u8; 32]);
    let payload = b"production-ready transport payload";
    let header = PacketHeader::new(
        *server.id(),
        2,
        42,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::RELIABLE),
    );
    let ct = client.encrypt_packet(&header, payload).expect("encrypt");
    assert_ne!(
        &ct[..payload.len()],
        payload,
        "ciphertext must not contain plaintext"
    );
    let pt = server.decrypt_packet(&header, &ct).expect("decrypt");
    assert_eq!(&pt, payload);
}

/// AEAD authenticity: flipping a single ciphertext byte must cause decrypt to
/// fail. This is what protects post-handshake traffic from tampering.
#[test]
fn tampered_ciphertext_is_rejected() {
    let (client, server) = make_session_pair([0xF1u8; 32]);
    let header = PacketHeader::new(
        *server.id(),
        7,
        1,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::RELIABLE),
    )
    .with_epoch(2)
    .with_path_id(3);

    let mut ct = client
        .encrypt_packet(&header, b"v2 payload")
        .expect("encrypt v2");
    ct[0] ^= 0x01;

    let result = server.decrypt_packet(&header, &ct);
    assert!(
        result.is_err(),
        "V2 AEAD must reject bit-flipped ciphertext; got {:?}",
        result.as_ref().ok().map(|v| v.len())
    );
}

/// The header's `epoch` and `path_id` are AAD-bound. Flipping either after
/// encryption must invalidate the tag.
#[test]
fn tampered_epoch_or_path_id_is_rejected() {
    let (client, server) = make_session_pair([0xF2u8; 32]);
    let real_header = PacketHeader::new(
        *server.id(),
        7,
        1,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::RELIABLE),
    )
    .with_epoch(5)
    .with_path_id(0);
    let ct = client
        .encrypt_packet(&real_header, b"epoch-bound payload")
        .expect("encrypt");

    // Mutate epoch.
    let tampered_epoch = PacketHeader {
        epoch: 6,
        ..real_header
    };
    assert!(server.decrypt_packet(&tampered_epoch, &ct).is_err());

    // Re-encrypt fresh so the AEAD recv counter aligns, then mutate path_id.
    let ct2 = client
        .encrypt_packet(&real_header, b"path-bound payload")
        .expect("re-encrypt");
    let tampered_path = PacketHeader {
        path_id: 7,
        ..real_header
    };
    assert!(server.decrypt_packet(&tampered_path, &ct2).is_err());
}

/// Replay window: re-feeding a fresh ciphertext that reuses an
/// already-accepted `(stream_id, sequence)` must fail with
/// `CoreError::ReplayDetected`, and the per-session counter must increment.
/// The window keys on `(stream_id, sequence)` only — independent of epoch /
/// path_id.
#[test]
fn replay_window_rejects_duplicate_sequence() {
    use phantom_protocol::CoreError;

    let (client, server) = make_session_pair([0xF4u8; 32]);
    let header = PacketHeader::new(
        *server.id(),
        3,
        17,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::RELIABLE),
    );
    let ct1 = client.encrypt_packet(&header, b"payload").expect("e1");
    server.decrypt_packet(&header, &ct1).expect("first decrypt");
    assert_eq!(server.replay_rejected_total(), 0);

    let ct2 = client.encrypt_packet(&header, b"payload").expect("e2");
    match server.decrypt_packet(&header, &ct2) {
        Err(CoreError::ReplayDetected(_)) => { /* expected */ }
        other => panic!(
            "expected ReplayDetected on V2 duplicate, got {:?}",
            other.as_ref().map(|_| "Ok").unwrap_or("Err(<other>)")
        ),
    }
    assert_eq!(server.replay_rejected_total(), 1);
}

/// Nonce-from-header property — a tampered packet that fails AEAD
/// verification must NOT desync the receiver from the sender. The next
/// legitimate packet must still decrypt cleanly.
///
/// The AEAD nonce is derived from the authenticated `header.sequence` rather
/// than an internal monotonic counter, so a failed decrypt is stateless from
/// the AEAD's perspective — a single dropped / mutated packet does not break
/// the session.
#[test]
fn failed_decrypt_does_not_desync_session() {
    let (client, server) = make_session_pair([0x20u8; 32]);

    // Sender encrypts packet #1.
    let h1 = PacketHeader::new(
        *server.id(),
        1,
        1,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::RELIABLE),
    );
    let ct1 = client.encrypt_packet(&h1, b"first").expect("encrypt 1");

    // Bad packet arrives in between — flipped tag byte.
    let mut tampered = ct1.clone();
    let n = tampered.len();
    tampered[n - 1] ^= 0x80;
    assert!(server.decrypt_packet(&h1, &tampered).is_err());

    // The original ct1 (same header, same payload) must still decrypt —
    // in V1 this would fail because the recv_counter desynchronised; in
    // V2 the nonce is reconstructible from h1 alone.
    let pt1 = server.decrypt_packet(&h1, &ct1).expect("decrypt 1");
    assert_eq!(pt1, b"first");

    // And a subsequent packet at sequence 2 also goes through.
    let h2 = PacketHeader { sequence: 2, ..h1 };
    let ct2 = client.encrypt_packet(&h2, b"second").expect("encrypt 2");
    let pt2 = server.decrypt_packet(&h2, &ct2).expect("decrypt 2");
    assert_eq!(pt2, b"second");
}

/// Mid-session rekey (Phase 1.5) — `Session::rekey()` increments the epoch
/// and derives a new AEAD state. Ciphertext produced before rekey must NOT
/// decrypt with the post-rekey state.
#[test]
fn rekey_changes_keys_and_breaks_old_ciphertexts() {
    let (client, server) = make_session_pair([0x10u8; 32]);
    assert_eq!(client.current_epoch(), 0);
    assert_eq!(server.current_epoch(), 0);

    let header = PacketHeader::new(
        *server.id(),
        1,
        100,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::RELIABLE),
    );
    let ct_epoch0 = client
        .encrypt_packet(&header, b"pre-rekey payload")
        .expect("encrypt e0");

    // Lock-step rekey on both ends.
    let client_new = client.rekey().expect("client rekey");
    let server_new = server.rekey().expect("server rekey");
    assert_eq!(client_new, 1);
    assert_eq!(server_new, 1);
    assert_eq!(client.current_epoch(), 1);
    assert_eq!(server.current_epoch(), 1);

    // The OLD ciphertext must NOT authenticate under the new keys.
    let header_epoch1 = PacketHeader { epoch: 1, ..header };
    assert!(
        server.decrypt_packet(&header_epoch1, &ct_epoch0).is_err(),
        "post-rekey CryptoState must reject pre-rekey ciphertext"
    );

    // A fresh encrypt under the new epoch round-trips successfully.
    let header_v1_e1 = PacketHeader::new(
        *server.id(),
        1,
        101,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::RELIABLE),
    )
    .with_epoch(1);
    let ct_epoch1 = client
        .encrypt_packet(&header_v1_e1, b"post-rekey payload")
        .expect("encrypt e1");
    let pt = server
        .decrypt_packet(&header_v1_e1, &ct_epoch1)
        .expect("decrypt e1");
    assert_eq!(pt, b"post-rekey payload");
}

/// `Session::ratchet_to_epoch(target)` advances the local epoch by repeated
/// HKDF chain steps. Useful for a receiver that fell behind and needs to
/// catch up to a higher-epoch packet.
#[test]
fn ratchet_to_epoch_walks_forward_n_steps() {
    let (_client, server) = make_session_pair([0x11u8; 32]);
    assert_eq!(server.current_epoch(), 0);
    server.ratchet_to_epoch(5).expect("ratchet to 5");
    assert_eq!(server.current_epoch(), 5);
    // Going to a lower target is a no-op.
    server.ratchet_to_epoch(3).expect("ratchet to 3 (no-op)");
    assert_eq!(server.current_epoch(), 5);
}

/// `Session::rekey` saturates at `u8::MAX` rather than wrapping — long
/// sessions must reconnect rather than reuse epoch 0 keys with a higher
/// counter.
#[test]
fn rekey_saturates_at_u8_max() {
    let (_, server) = make_session_pair([0x12u8; 32]);
    server
        .ratchet_to_epoch(u8::MAX)
        .expect("walk up to u8::MAX");
    assert_eq!(server.current_epoch(), u8::MAX);
    // The 256th rekey must fail rather than wrap to 0.
    assert!(server.rekey().is_err());
    assert_eq!(server.current_epoch(), u8::MAX, "epoch must not wrap");
}

/// C1: `decrypt_packet_accepting_rekey` follows a single *authenticated* forward
/// rekey step. The sender rekeys to epoch 1 and encrypts there; the receiver,
/// still at epoch 0, trial-decrypts under the next key, succeeds, and commits
/// the ratchet — ending at epoch 1 with the plaintext intact.
#[test]
fn accepting_decrypt_follows_one_authentic_rekey_step() {
    let (client, server) = make_session_pair([0x20u8; 32]);
    assert_eq!(server.current_epoch(), 0);

    // Sender rekeys ahead of the receiver.
    assert_eq!(client.rekey().expect("client rekey"), 1);
    let header = PacketHeader::new(
        *server.id(),
        1,
        7,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::REKEY),
    )
    .with_epoch(1);
    let ct = client
        .encrypt_packet(&header, b"first post-rekey")
        .expect("encrypt e1");

    // Receiver is still at epoch 0; the accepting decrypt ratchets it forward.
    let pt = server
        .decrypt_packet_accepting_rekey(&header, &ct)
        .expect("accepting decrypt follows the bump");
    assert_eq!(pt, b"first post-rekey");
    assert_eq!(server.current_epoch(), 1, "receiver committed the ratchet");
}

/// C1 security: a *forged* epoch bump (correct +1 epoch in the header, but
/// ciphertext that does not authenticate under the next key) is rejected and
/// MUST NOT commit the ratchet — otherwise an attacker could desync the session
/// by spoofing an epoch. After the rejection a legitimate same-epoch packet
/// still decrypts.
#[test]
fn accepting_decrypt_rejects_forged_bump_without_desync() {
    let (client, server) = make_session_pair([0x21u8; 32]);

    // Attacker forges a +1-epoch header but supplies garbage ciphertext.
    let forged = PacketHeader::new(
        *server.id(),
        1,
        1,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::REKEY),
    )
    .with_epoch(1);
    let garbage = vec![0xABu8; 64];
    assert!(
        server
            .decrypt_packet_accepting_rekey(&forged, &garbage)
            .is_err(),
        "a forged epoch bump must fail the AEAD trial"
    );
    assert_eq!(
        server.current_epoch(),
        0,
        "a failed trial decrypt must NOT advance the epoch (no desync)"
    );

    // The session is intact: a genuine epoch-0 packet still round-trips.
    let header = PacketHeader::new(*server.id(), 1, 2, PacketFlags::new(PacketFlags::ENCRYPTED));
    let ct = client
        .encrypt_packet(&header, b"still in sync")
        .expect("encrypt e0");
    let pt = server
        .decrypt_packet_accepting_rekey(&header, &ct)
        .expect("same-epoch decrypt still works");
    assert_eq!(pt, b"still in sync");
}

/// C1: a *bounded* multi-epoch catch-up (within [`MAX_REKEY_CATCHUP`]) with a
/// genuinely valid ciphertext is followed — the receiver derives the chain
/// forward and commits all the steps at once. This absorbs the small epoch
/// divergence that arises when both directions rekey at slightly different
/// cadences.
#[test]
fn accepting_decrypt_follows_bounded_multi_step_catchup() {
    let (client, server) = make_session_pair([0x22u8; 32]);
    client.ratchet_to_epoch(3).expect("client to 3");
    let header = PacketHeader::new(
        *server.id(),
        1,
        1,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::REKEY),
    )
    .with_epoch(3);
    let ct = client
        .encrypt_packet(&header, b"three ahead")
        .expect("encrypt e3");

    // Receiver at epoch 0 catches up 3 steps because the ciphertext authenticates.
    let pt = server
        .decrypt_packet_accepting_rekey(&header, &ct)
        .expect("bounded multi-step catch-up follows a valid jump");
    assert_eq!(pt, b"three ahead");
    assert_eq!(server.current_epoch(), 3, "receiver caught up to epoch 3");
}

/// C1 security: a jump *beyond* [`MAX_REKEY_CATCHUP`] is rejected outright even
/// with a valid ciphertext — this caps the HKDF work an attacker can force per
/// spoofed packet. A legitimate gap is never this large; over a reliable
/// transport the sender retransmits at the current epoch.
#[test]
fn accepting_decrypt_rejects_jump_beyond_catchup_bound() {
    let (client, server) = make_session_pair([0x24u8; 32]);
    let target = MAX_REKEY_CATCHUP + 1;
    client.ratchet_to_epoch(target).expect("client far ahead");
    let header = PacketHeader::new(
        *server.id(),
        1,
        1,
        PacketFlags::new(PacketFlags::ENCRYPTED | PacketFlags::REKEY),
    )
    .with_epoch(target);
    let ct = client
        .encrypt_packet(&header, b"too far")
        .expect("encrypt far");

    assert!(
        server.decrypt_packet_accepting_rekey(&header, &ct).is_err(),
        "a jump beyond MAX_REKEY_CATCHUP must be rejected"
    );
    assert_eq!(
        server.current_epoch(),
        0,
        "no ratchet on an over-bound jump"
    );
}

/// C1: the automatic-rekey trigger predicate flips once the send direction
/// crosses the configurable high-watermark, and an actual `rekey()` clears it
/// (the counter resets under the fresh key).
#[test]
fn send_needs_rekey_fires_at_threshold_and_clears_on_rekey() {
    let (client, _server) = make_session_pair([0x23u8; 32]);
    client.set_rekey_threshold(4);
    assert!(
        !client.send_needs_rekey(),
        "fresh session is below threshold"
    );

    let header = PacketHeader::new(*client.id(), 1, 0, PacketFlags::new(PacketFlags::ENCRYPTED));
    for i in 0..4u32 {
        let h = PacketHeader {
            sequence: i,
            ..header
        };
        client.encrypt_packet(&h, b"x").expect("encrypt");
    }
    assert!(
        client.send_needs_rekey(),
        "after {} sends the trigger must fire",
        client.send_invocations()
    );

    assert_eq!(client.rekey().expect("rekey"), 1);
    assert!(
        !client.send_needs_rekey(),
        "rekey resets the send counter under the new key, clearing the trigger"
    );
}

/// C1 concurrency: the data pump drives the send loop and the receive task
/// concurrently over one `Arc<Session>`, so a send-side `rekey()` can race a
/// receive-side ratchet. Every transition must be atomic — the installed key
/// depth and the epoch counter must never diverge. We hammer `rekey()` from
/// many threads and then prove the final key is exactly `epoch` HKDF steps deep
/// by round-tripping a packet against a peer ratcheted to the same epoch. With
/// a non-atomic (read-epoch / derive / bump-relative) transition this wedges:
/// the epoch overshoots the key depth and the round-trip fails.
#[test]
fn concurrent_rekeys_keep_epoch_and_key_in_lockstep() {
    use std::sync::Arc;

    const THREADS: usize = 8;
    const PER_THREAD: usize = 20; // 160 total < u8::MAX, so none saturate

    let (client, server) = make_session_pair([0x30u8; 32]);
    let client = Arc::new(client);

    let mut handles = Vec::new();
    for _ in 0..THREADS {
        let c = Arc::clone(&client);
        handles.push(std::thread::spawn(move || {
            for _ in 0..PER_THREAD {
                c.rekey().expect("concurrent rekey");
            }
        }));
    }
    for h in handles {
        h.join().expect("rekey thread");
    }

    let epoch = client.current_epoch();
    assert_eq!(
        epoch as usize,
        THREADS * PER_THREAD,
        "every concurrent rekey must advance the epoch exactly once (no lost/double bumps)"
    );

    // Prove key-depth == epoch: a peer ratcheted to the same epoch must decrypt.
    server.ratchet_to_epoch(epoch).expect("server catch up");
    let header = PacketHeader::new(*client.id(), 1, 1, PacketFlags::new(PacketFlags::ENCRYPTED))
        .with_epoch(epoch);
    let ct = client
        .encrypt_packet(&header, b"post-race payload")
        .expect("encrypt at final epoch");
    let pt = server
        .decrypt_packet(&header, &ct)
        .expect("installed key depth must equal the epoch counter");
    assert_eq!(pt, b"post-race payload");
}

// ── Multi-path / migration (Phase 4.2) ────────────────────────────────────

/// New paths must NOT be implicitly trusted. After session creation,
/// path 0 is the validated default; an unfamiliar path id starts at
/// `Unvalidated` and only transitions to `Validated` through the
/// challenge-response API.
#[test]
fn new_paths_default_to_unvalidated() {
    let (_client, server) = make_session_pair([0x40u8; 32]);
    // Path 0 was registered at construction and pre-validated — it's
    // the path the handshake traversed.
    assert_eq!(server.path_state(0), Some(PathStateKind::Validated));
    // Path 7 has never been seen.
    assert_eq!(server.path_state(7), None);

    // begin_path_validation registers + issues challenge.
    let challenge = server.begin_path_validation(7).expect("challenge");
    assert_eq!(challenge.len(), 32);
    assert_eq!(server.path_state(7), Some(PathStateKind::Validating));
}

/// A correct challenge response transitions the path to `Validated`
/// and surfaces it in `validated_paths`.
#[test]
fn correct_response_validates_path() {
    let (_client, server) = make_session_pair([0x41u8; 32]);
    let challenge = server.begin_path_validation(3).expect("challenge");
    assert!(server.complete_path_validation(3, &challenge));
    assert_eq!(server.path_state(3), Some(PathStateKind::Validated));

    let mut validated = server.validated_paths();
    validated.sort();
    // Path 0 was pre-validated at construction; path 3 just was.
    assert_eq!(validated, vec![0, 3]);
}

/// A wrong response transitions the path to `Failed` — application data
/// must NOT cross over it.
#[test]
fn wrong_response_marks_path_failed() {
    let (_client, server) = make_session_pair([0x42u8; 32]);
    let mut challenge = server.begin_path_validation(5).expect("challenge");
    challenge[0] ^= 0xFF;
    assert!(!server.complete_path_validation(5, &challenge));
    assert_eq!(server.path_state(5), Some(PathStateKind::Failed));
    assert!(!server.validated_paths().contains(&5));
}

/// `complete_path_validation` returns `false` for paths that were never
/// challenged — protects against an attacker bypassing the challenge step.
#[test]
fn unchallenged_path_cannot_be_completed() {
    let (_client, server) = make_session_pair([0x43u8; 32]);
    assert!(!server.complete_path_validation(9, &[0u8; 32]));
    // No state was created (registry wasn't touched).
    assert_eq!(server.path_state(9), None);
}

/// A `PhantomPacket` survives serialize + deserialize with the pinned wire
/// version and all header fields preserved.
#[test]
fn packet_roundtrip_preserves_fields() {
    let header = PacketHeader::new(
        SessionId::from_bytes([9u8; 32]),
        99,
        2025,
        PacketFlags::new(PacketFlags::RELIABLE | PacketFlags::ENCRYPTED | PacketFlags::REKEY),
    )
    .with_epoch(11)
    .with_path_id(2);
    let packet = PhantomPacket::new(header, vec![0xDE, 0xAD]);
    let buf = packet.to_wire();
    let decoded = PhantomPacket::from_wire(&buf).expect("roundtrip");
    assert_eq!(decoded.header.version, WIRE_VERSION);
    assert_eq!(decoded.header.epoch, 11);
    assert_eq!(decoded.header.path_id, 2);
    assert!(decoded.header.flags.contains(PacketFlags::REKEY));
    assert_eq!(decoded.payload, vec![0xDE, 0xAD]);
}

// ── Flow-control enforcement invariants (receive-backpressure decoupling) ────
//
// With backpressure decoupled from the recv reader, the SEND side is where flow
// control is actually enforced: `Stream::poll_send` admits new data only within
// `min(congestion_window, peer_flow_control_window)`, while retransmissions must
// bypass both so loss recovery can never be starved by a closed window. These
// two tests pin those properties so a future change can't silently let a stream
// outrun a slow peer (unbounded receiver memory) or wedge loss recovery.

/// New (first-transmission) data is admitted only within the advertised
/// flow-control window AND the congestion window — `min` of the two. A segment
/// that does not fit is withheld and, crucially, does NOT debit the window (so
/// the credit is not leaked while the segment waits).
#[tokio::test]
async fn flow_control_bounds_new_data_to_the_advertised_window() {
    // ── Flow-control window bound ──
    let s = Stream::new(1);
    // Drain the peer's advertised window down to a known small amount.
    assert!(s.try_consume_send_window(INITIAL_STREAM_WINDOW - 100));
    assert_eq!(s.peer_send_window(), 100);

    // Two new-data segments queued: the first fits the 100-byte window, the
    // second does not. The congestion budget is unbounded so ONLY the
    // flow-control window can gate us here.
    s.send_reliable(Bytes::from(vec![0u8; 60])).await; // seq 0
    s.send_reliable(Bytes::from(vec![0u8; 60])).await; // seq 1

    let first = s
        .poll_send(u64::MAX)
        .await
        .expect("first segment fits the window");
    assert!(!first.retransmit);
    assert_eq!(first.data.len(), 60);
    assert_eq!(s.peer_send_window(), 40, "window debited by the sent bytes");

    // The second 60-byte segment exceeds the remaining 40-byte window → withheld.
    assert!(
        s.poll_send(u64::MAX).await.is_none(),
        "new data exceeding the flow-control window must be withheld"
    );
    assert_eq!(
        s.peer_send_window(),
        40,
        "a withheld segment must NOT debit the window (no credit leak)"
    );

    // ── Congestion window bound ──
    let s2 = Stream::new(2);
    s2.send_reliable(Bytes::from(vec![0u8; 100])).await;
    // cwnd budget smaller than the segment → withheld by congestion control,
    // BEFORE the flow-control window is even consulted.
    assert!(
        s2.poll_send(50).await.is_none(),
        "new data exceeding the congestion window must be withheld"
    );
    assert_eq!(
        s2.peer_send_window(),
        INITIAL_STREAM_WINDOW,
        "a cwnd-blocked segment must not debit the flow-control window"
    );
}

/// Retransmissions bypass BOTH the congestion window and the flow-control
/// window: a timed-out segment is re-offered even when `cwnd_budget == 0` and
/// the peer's window is fully closed — loss recovery must always proceed, and
/// the retransmit must not debit the (already-accounted) window again.
#[tokio::test]
async fn retransmissions_bypass_congestion_and_flow_control_windows() {
    tokio::time::pause();
    let s = Stream::new(1);
    s.send_reliable(Bytes::from(vec![0u8; 200])).await; // seq 0

    // First transmission debits the window (200 bytes) under an unbounded cwnd.
    let first = s.poll_send(u64::MAX).await.expect("first transmission");
    assert!(!first.retransmit);
    assert_eq!(first.data.len(), 200);
    assert_eq!(s.peer_send_window(), INITIAL_STREAM_WINDOW - 200);

    // Slam BOTH budgets shut: drain the flow-control window to zero …
    assert!(s.try_consume_send_window(s.peer_send_window()));
    assert_eq!(s.peer_send_window(), 0);
    // … and an immediate re-poll (cwnd 0, window 0) yields nothing — the
    // segment is in-flight, not yet timed out.
    assert!(s.poll_send(0).await.is_none());

    // Advance past the initial 1s RTO so the unacked segment is due to retransmit.
    tokio::time::advance(Duration::from_millis(1100)).await;

    // The retransmit is produced despite cwnd == 0 AND window == 0 (Karn: the
    // bytes were accounted on first send; loss recovery must always proceed).
    let rtx = s
        .poll_send(0)
        .await
        .expect("retransmission must bypass both the congestion and flow-control windows");
    assert!(rtx.retransmit, "must be flagged as a retransmission");
    assert_eq!(rtx.seq, first.seq);
    assert_eq!(rtx.data.len(), 200);
    assert_eq!(
        s.peer_send_window(),
        0,
        "a retransmission must not debit the flow-control window again"
    );
}

// ── C1: per-stream sequence nonce-reuse (Invariant 8) ───────────────────────

/// **C1 (critical).** The AEAD nonce is `(epoch, stream_id, sequence, path_id)`.
/// `sequence` is a per-stream `u32` that wraps at `2^32`, while the only
/// pre-existing rekey trigger keys off the *direction-wide* invocation counter
/// (`REKEY_SOFT_LIMIT = 2^47`). A single hot stream therefore wraps its
/// sequence — repeating a nonce under a fixed key (the Forbidden Attack) — long
/// before any rekey fires. The fix forces a rekey once a stream's sequence
/// advances past a per-stream watermark within the current epoch.
///
/// We drive the exact send-side decision (`send_needs_rekey ||
/// stream_seq_needs_rekey` then `rekey`) with the direction-wide trigger
/// disabled and a tiny injected watermark, and assert that **no epoch's
/// per-stream sequence span ever exceeds the watermark**. Scaled to the real
/// `W = 2^31` against the `2^32` wrap, that is exactly "no stream can traverse
/// the full sequence space within one epoch", so the nonce never repeats.
#[test]
fn single_stream_seq_watermark_forces_rekey_before_wrap() {
    let (client, _server) = make_session_pair([0x5Au8; 32]);
    // Isolate the per-stream trigger: push the direction-wide AEAD trigger out
    // of reach so ONLY the sequence watermark can fire a rekey.
    client.set_rekey_threshold(u64::MAX);
    const W: u32 = 8;
    client.set_seq_rekey_watermark(W);

    let stream: u16 = 1;
    let mut spans: std::collections::BTreeMap<u8, (u32, u32)> = std::collections::BTreeMap::new();
    let mut seen: std::collections::HashSet<(u8, u16, u32)> = std::collections::HashSet::new();

    for seq in 0u32..(4 * W + 3) {
        // Exact production decision — mirrors `send_app_data`.
        if client.send_needs_rekey() || client.stream_seq_needs_rekey(stream, seq) {
            client
                .rekey()
                .expect("rekey must succeed within the epoch budget");
        }
        let epoch = client.current_epoch();
        assert!(
            seen.insert((epoch, stream, seq)),
            "nonce tuple (epoch={epoch}, stream={stream}, seq={seq}) repeated"
        );
        let e = spans.entry(epoch).or_insert((seq, seq));
        e.0 = e.0.min(seq);
        e.1 = e.1.max(seq);
    }

    assert!(
        spans.len() >= 2,
        "the sequence watermark never forced a rekey"
    );
    for (epoch, (lo, hi)) in &spans {
        assert!(
            hi - lo <= W,
            "epoch {epoch} spans {} sequences, exceeding watermark {W} — a wider epoch could wrap",
            hi - lo
        );
    }
}

/// **C1 fail-closed.** Once the `u8` epoch saturates at `u8::MAX`, no further
/// rekey is possible. The per-stream watermark must then surface a hard error
/// from `rekey()` (which `send_app_data` turns into a failed send → session
/// reconnect) rather than letting the sequence wrap and silently reuse a nonce.
#[test]
fn seq_watermark_fails_closed_at_epoch_saturation() {
    let (client, _server) = make_session_pair([0x77u8; 32]);
    client.set_rekey_threshold(u64::MAX);
    client.set_seq_rekey_watermark(4);

    let stream: u16 = 1;
    let mut seq = 0u32;
    let mut rekeys = 0u32;
    loop {
        if client.send_needs_rekey() || client.stream_seq_needs_rekey(stream, seq) {
            match client.rekey() {
                Ok(_) => rekeys += 1,
                Err(_) => break, // fail-closed: epoch saturated, no nonce reuse
            }
        }
        seq += 1;
        assert!(
            seq < 1_000_000,
            "must fail closed at saturation, not loop forever"
        );
    }
    assert_eq!(
        client.current_epoch(),
        u8::MAX,
        "must fail closed exactly when the epoch saturates"
    );
    assert_eq!(
        rekeys,
        u8::MAX as u32,
        "should perform 255 successful rekeys before the epoch saturates"
    );
}

// ── Auth cluster: H2 (transcript-signed 0-RTT verdict) ──────────────────────

/// Drive a fresh `ClientHello` through the server to a `ServerHello`,
/// transparently answering the single cookie `Retry` the DoS gate issues.
/// Returns the **effective** `ClientHello` the server actually signed over
/// (the retried one, carrying the cookie) alongside the `ServerHello`, so the
/// caller verifies the signature against the matching transcript input.
fn drive_handshake_to_success(
    server: &HandshakeServer,
    client_hello: &ClientHello,
    client_ip: std::net::IpAddr,
) -> (ClientHello, ServerHello) {
    match server.process_client_hello(client_hello, 0, client_ip) {
        HandshakeResponse::Success(sh, _, _) => (client_hello.clone(), sh),
        HandshakeResponse::Retry(retry) => {
            let mut retried = client_hello.clone();
            retried.cookie = retry.cookie;
            match server.process_client_hello(&retried, 0, client_ip) {
                HandshakeResponse::Success(sh, _, _) => (retried, sh),
                other => panic!("unexpected response after cookie retry: {:?}", other),
            }
        }
        other => panic!("unexpected first handshake response: {:?}", other),
    }
}

/// **H2 (Invariant 9).** `ServerHello.early_data_accepted` is the server's 0-RTT
/// verdict. It MUST be covered by the signed handshake transcript: an on-path
/// attacker who flips the bit (leaving signature/ciphertext/session_id intact)
/// must break the client's signature check, not slip a forged verdict through —
/// a forged verdict would let the attacker duplicate or silently black-hole
/// 0-RTT early-data.
#[test]
fn flipped_early_data_accepted_bit_fails_signature() {
    let server = HandshakeServer::new().expect("server");
    let server_pk = server.verifying_key().clone();
    let client = HandshakeClient::new().expect("client");
    let hello = client.create_client_hello();
    let ip = "127.0.0.1".parse().expect("ip");
    let (effective_hello, sh) = drive_handshake_to_success(&server, &hello, ip);

    // Honest verdict verifies (positive control).
    assert!(
        client
            .process_server_hello(&effective_hello, &sh, Some(&server_pk))
            .is_ok(),
        "an untampered ServerHello must verify"
    );

    // Flip the verdict; signature/ciphertext/session_id are left intact.
    let mut tampered = sh.clone();
    tampered.early_data_accepted = !tampered.early_data_accepted;
    assert!(
        matches!(
            client.process_server_hello(&effective_hello, &tampered, Some(&server_pk)),
            Err(HandshakeError::KemFailed(_))
        ),
        "flipping early_data_accepted must fail the transcript signature check"
    );
}

// ── Auth cluster: HS-03 (resumption PoP binder) + ZERORTT-2 (consume-on-success)

/// Drive a first full handshake so the server mints a resumption ticket; return
/// the `(resume_session_id, resumption_secret)` the client would later resume
/// with (the two halves of `Session::resumption_hint()`).
fn first_handshake_mint_ticket(
    server: &HandshakeServer,
    client: &HandshakeClient,
    server_pk: &HybridVerifyingKey,
    ip: std::net::IpAddr,
) -> ([u8; 32], [u8; 32]) {
    let hello = client.create_client_hello();
    let (effective, sh) = drive_handshake_to_success(server, &hello, ip);
    let (session, _) = client
        .process_server_hello(&effective, &sh, Some(server_pk))
        .expect("client establishes session");
    let secret = session
        .resumption_secret()
        .expect("resumption secret installed");
    (sh.session_id, secret)
}

/// **HS-03 (Invariant 9).** A resume must carry a `resumption_binder` proving
/// possession of the prior session's `resumption_secret`. A passive observer
/// that copied only the cleartext `resume_session_id` cannot forge it, so a
/// binderless (or wrong-binder) resume must NOT consume the victim's one-shot
/// ticket — it falls back to the normal cookie/PoW gate, ticket intact.
#[test]
fn binderless_resume_does_not_burn_ticket() {
    let server = HandshakeServer::new().expect("server");
    let server_pk = server.verifying_key().clone();
    let ip = "127.0.0.1".parse().expect("ip");
    let client1 = HandshakeClient::new().expect("client1");
    let (rid, secret) = first_handshake_mint_ticket(&server, &client1, &server_pk, ip);
    assert_eq!(
        server.session_cache_len(),
        1,
        "first handshake mints a ticket"
    );

    // Observer: right rid, NO binder (cannot compute it without `secret`).
    let client2 = HandshakeClient::new().expect("client2");
    let mut forged = client2.create_client_hello_with_resume(rid, &secret, None);
    forged.resumption_binder = None;
    match server.process_client_hello(&forged, 0, ip) {
        HandshakeResponse::Retry(_) => {} // fell back to the DoS gate — correct
        other => panic!("a binderless resume must not bypass the gate: {:?}", other),
    }
    assert_eq!(
        server.session_cache_len(),
        1,
        "a binderless resume must NOT consume the ticket"
    );

    // Also: a WRONG binder (attacker guesses) must not burn it either.
    let mut wrong = client2.create_client_hello_with_resume(rid, &secret, None);
    wrong.resumption_binder = Some([0xAB; 32]);
    let _ = server.process_client_hello(&wrong, 0, ip);
    assert_eq!(
        server.session_cache_len(),
        1,
        "a wrong-binder resume must NOT consume the ticket"
    );

    // A legitimate resume (correct binder, proving possession of `secret`)
    // bypasses the gate and consumes the ticket.
    let client3 = HandshakeClient::new().expect("client3");
    let valid = client3.create_client_hello_with_resume(rid, &secret, None);
    match server.process_client_hello(&valid, 0, ip) {
        HandshakeResponse::Success(..) => {} // bypass ⇒ binder verified + consumed
        other => panic!("a valid resume should succeed: {:?}", other),
    }
    // One-shot (Invariant 9): replaying the SAME resume no longer bypasses the
    // gate — the ticket for `rid` was consumed (a successful resume mints a
    // fresh ticket under a NEW id, so this `rid` is gone for good).
    match server.process_client_hello(&valid, 0, ip) {
        HandshakeResponse::Retry(_) => {}
        other => panic!(
            "a replayed resume must not resume again (one-shot): {:?}",
            other
        ),
    }
}

/// **ZERORTT-2 (Invariant 9).** The ticket is consumed eagerly after the binder
/// check, but a handshake step that fails AFTER consumption (here: a corrupted
/// KEM key package fails `encapsulate()`) must re-insert the ticket — a
/// corrupted resuming `ClientHello` must not burn a victim's one-shot ticket.
#[test]
fn failed_resume_handshake_leaves_ticket_usable() {
    let server = HandshakeServer::new().expect("server");
    let server_pk = server.verifying_key().clone();
    let ip = "127.0.0.1".parse().expect("ip");
    let client1 = HandshakeClient::new().expect("client1");
    let (rid, secret) = first_handshake_mint_ticket(&server, &client1, &server_pk, ip);
    assert_eq!(server.session_cache_len(), 1);

    // Valid binder (over secret/rid/nonce), but corrupt the ML-KEM public so
    // encapsulate() fails after the ticket is consumed.
    let client2 = HandshakeClient::new().expect("client2");
    let mut hello = client2.create_client_hello_with_resume(rid, &secret, None);
    hello.client_key_package.ml_kem_pk.truncate(1); // wrong length → KEM decode fails
    match server.process_client_hello(&hello, 0, ip) {
        HandshakeResponse::Fail(HandshakeError::KemFailed(_)) => {}
        other => panic!("a corrupted-KEM resume should fail: {:?}", other),
    }
    assert_eq!(
        server.session_cache_len(),
        1,
        "a resume that fails after consume must re-insert the ticket (not burn it)"
    );

    // And the re-inserted ticket is still usable by a clean resume afterwards
    // (Success with difficulty=0 and no cookie can only happen via the resume
    // bypass — so this proves the ticket survived the failed attempt).
    let client3 = HandshakeClient::new().expect("client3");
    let valid = client3.create_client_hello_with_resume(rid, &secret, None);
    assert!(
        matches!(
            server.process_client_hello(&valid, 0, ip),
            HandshakeResponse::Success(..)
        ),
        "the re-inserted ticket must be usable by a clean resume"
    );
}