//! `test_replay_protection` — a Data record that was already delivered, replayed verbatim, must be
//! rejected by the receiver's sliding replay window.
//!
//! Topology: the client and server each talk to one end of their own duplex. A relay task in the
//! middle forwards bytes between the two. For the client->server direction the relay parses whole
//! frames (using the crate's public framing helpers) so that, once the client has sent its data
//! packet, the relay can forward that exact record to the server a SECOND time — a verbatim replay.

mod common;

use aura_proto::frame::{read_frame, write_frame, MsgType, RawFrame};
use aura_proto::{client_handshake, server_handshake, Frame, ProtoError};
use bytes::Bytes;
use tokio::io::{split, AsyncWriteExt};
use tokio::sync::oneshot;

#[tokio::test]
async fn test_replay_protection() {
    let pki = common::mint_pki("vpn.aura.example", "client-alpha");
    let client_cfg = pki.client_config();
    let server_cfg = pki.server_config();

    // Two duplexes with a relay in the middle.
    let (client_io, relay_a) = tokio::io::duplex(64 * 1024);
    let (relay_b, server_io) = tokio::io::duplex(64 * 1024);

    let (c_read, c_write) = split(client_io);
    let (s_read, s_write) = split(server_io);
    let (ra_read, ra_write) = split(relay_a); // faces the client
    let (rb_read, rb_write) = split(relay_b); // faces the server

    // Signal so the relay forwards the replay only after the server has consumed the genuine copy.
    let (genuine_done_tx, genuine_done_rx) = oneshot::channel::<()>();

    // ---- Relay: client -> server, with a one-shot verbatim replay of the first Data record ----
    let relay_c2s = tokio::spawn(async move {
        let mut ra_read = ra_read;
        let mut rb_write = rb_write;
        let mut genuine_done = Some(genuine_done_rx);
        let mut replayed = false;
        loop {
            let frame: RawFrame = match read_frame(&mut ra_read).await {
                Ok(f) => f,
                Err(_) => break, // EOF when the client side closes
            };
            // Forward the frame unchanged.
            write_frame(&mut rb_write, frame.msg_type, &frame.payload)
                .await
                .expect("relay forward c->s");

            // On the first Data record, wait until the server has accepted it, then replay it once.
            if frame.msg_type == MsgType::Data && !replayed {
                replayed = true;
                if let Some(rx) = genuine_done.take() {
                    let _ = rx.await; // server signals after it accepted the genuine record
                }
                write_frame(&mut rb_write, frame.msg_type, &frame.payload)
                    .await
                    .expect("relay replay c->s");
                rb_write.flush().await.expect("flush replay");
            }
        }
    });

    // ---- Relay: server -> client (straight byte copy) ----
    let relay_s2c = tokio::spawn(async move {
        let mut rb_read = rb_read;
        let mut ra_write = ra_write;
        let _ = tokio::io::copy(&mut rb_read, &mut ra_write).await;
        let _ = ra_write.shutdown().await;
    });

    // ---- Client: handshake, then send exactly one Data frame ----
    let client = tokio::spawn(async move {
        let mut sess = client_handshake(c_read, c_write, &client_cfg)
            .await
            .expect("client handshake");
        sess.send_frame(Frame::Data {
            stream_id: 9,
            payload: Bytes::from_static(b"the one and only payload"),
        })
        .await
        .expect("client send");
        // Keep the session (and thus the transport) alive until the test signals completion.
        sess
    });

    // ---- Server: handshake, accept the genuine record, then expect the replay to be rejected ----
    let server = tokio::spawn(async move {
        let mut sess = server_handshake(s_read, s_write, &server_cfg)
            .await
            .expect("server handshake");

        // 1) Genuine record is accepted.
        let first = sess.recv_frame().await.expect("genuine recv");
        match first {
            Frame::Data { stream_id, payload } => {
                assert_eq!(stream_id, 9);
                assert_eq!(&payload[..], b"the one and only payload");
            }
            other => panic!("expected Data, got {other:?}"),
        }

        // Tell the relay it may now inject the verbatim replay.
        genuine_done_tx.send(()).expect("signal genuine done");

        // 2) The replayed record must be rejected by the replay window.
        let replay_result = sess.recv_frame().await;
        assert!(
            matches!(replay_result, Err(ProtoError::Replay(_))),
            "expected ProtoError::Replay, got {replay_result:?}"
        );
        sess
    });

    let (_client_sess, server_outcome) = tokio::join!(client, server);
    server_outcome.expect("server task");
    drop(_client_sess); // closes the client side -> relays drain and exit

    let _ = relay_c2s.await;
    let _ = relay_s2c.await;
}