quicproquo/crates/quicnprotochat-client/tests/noise_transport.rs

//! M1 integration test: Noise_XX handshake + Ping/Pong round-trip.
//!
//! Both the server-side and client-side logic run in the same Tokio runtime
//! using `tokio::spawn`. The test verifies:
//!
//! 1. The Noise_XX handshake completes from both sides.
//! 2. A Ping sent by the client arrives as a Ping on the server side.
//! 3. The server's Pong arrives correctly on the client side.
//! 4. Mutual authentication: each peer's observed remote static key matches the
//!    other peer's actual public key (the core security property of XX).

use std::sync::Arc;

use tokio::net::TcpListener;

use quicnprotochat_core::{handshake_initiator, handshake_responder, NoiseKeypair};
use quicnprotochat_proto::{MsgType, ParsedEnvelope};

/// Completes a full Noise_XX handshake and Ping/Pong exchange, then verifies
/// mutual authentication by comparing observed vs. actual static public keys.
#[tokio::test]
async fn noise_xx_ping_pong_round_trip() {
    let server_keypair = Arc::new(NoiseKeypair::generate());
    let client_keypair = NoiseKeypair::generate();

    // Bind the listener *before* spawning so the port is ready when the client
    // calls connect — no sleep or retry needed.
    let listener = TcpListener::bind("127.0.0.1:0")
        .await
        .expect("failed to bind test listener");
    let server_addr = listener.local_addr().expect("failed to get local addr");

    // ── Server task ───────────────────────────────────────────────────────────
    //
    // Handles exactly one connection: completes the handshake, asserts that it
    // receives a Ping, sends a Pong, then returns the client's observed key.
    let server_kp = Arc::clone(&server_keypair);
    let server_task = tokio::spawn(async move {
        let (stream, _peer) = listener.accept().await.expect("server accept failed");

        let mut transport = handshake_responder(stream, &server_kp)
            .await
            .expect("server Noise_XX handshake failed");

        let env = transport
            .recv_envelope()
            .await
            .expect("server recv_envelope failed");

        match env.msg_type {
            MsgType::Ping => {}
            _ => panic!("server expected Ping, received a different message type"),
        }

        transport
            .send_envelope(&ParsedEnvelope {
                msg_type: MsgType::Pong,
                group_id: vec![],
                sender_id: vec![],
                payload: vec![],
                timestamp_ms: 0,
            })
            .await
            .expect("server send_envelope failed");

        // Return the client's public key as authenticated by the server.
        transport
            .remote_static_public_key()
            .expect("server: no remote static key after completed XX handshake")
            .to_vec()
    });

    // ── Client side ───────────────────────────────────────────────────────────
    let stream = tokio::net::TcpStream::connect(server_addr)
        .await
        .expect("client connect failed");

    let mut transport = handshake_initiator(stream, &client_keypair)
        .await
        .expect("client Noise_XX handshake failed");

    // Capture the server's public key as authenticated by the client.
    let server_key_seen_by_client = transport
        .remote_static_public_key()
        .expect("client: no remote static key after completed XX handshake")
        .to_vec();

    transport
        .send_envelope(&ParsedEnvelope {
            msg_type: MsgType::Ping,
            group_id: vec![],
            sender_id: vec![],
            payload: vec![],
            timestamp_ms: 1_700_000_000_000,
        })
        .await
        .expect("client send_envelope failed");

    let pong = tokio::time::timeout(std::time::Duration::from_secs(5), transport.recv_envelope())
        .await
        .expect("timed out waiting for Pong — server task likely panicked")
        .expect("client recv_envelope failed");

    match pong.msg_type {
        MsgType::Pong => {}
        _ => panic!("client expected Pong, received a different message type"),
    }

    // ── Mutual authentication assertions ──────────────────────────────────────
    let client_key_seen_by_server = server_task
        .await
        .expect("server task panicked — see output above");

    // The server authenticated the client's static public key correctly.
    assert_eq!(
        client_key_seen_by_server,
        client_keypair.public_bytes().to_vec(),
        "server's authenticated view of client key does not match client's actual public key"
    );

    // The client authenticated the server's static public key correctly.
    assert_eq!(
        server_key_seen_by_client,
        server_keypair.public_bytes().to_vec(),
        "client's authenticated view of server key does not match server's actual public key"
    );
}

/// A second independent connection on the same server must also succeed,
/// confirming that the server keypair reuse across connections is correct.
#[tokio::test]
async fn two_sequential_connections_both_authenticate() {
    let server_keypair = Arc::new(NoiseKeypair::generate());

    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind failed");
    let server_addr = listener.local_addr().expect("local_addr failed");

    let server_kp = Arc::clone(&server_keypair);
    tokio::spawn(async move {
        for _ in 0..2_u8 {
            let (stream, _) = listener.accept().await.expect("accept failed");
            let kp = Arc::clone(&server_kp);
            tokio::spawn(async move {
                let mut t = handshake_responder(stream, &kp)
                    .await
                    .expect("server handshake failed");
                let env = t.recv_envelope().await.expect("recv failed");
                match env.msg_type {
                    MsgType::Ping => {}
                    _ => panic!("expected Ping"),
                }
                t.send_envelope(&ParsedEnvelope {
                    msg_type: MsgType::Pong,
                    group_id: vec![],
                    sender_id: vec![],
                    payload: vec![],
                    timestamp_ms: 0,
                })
                .await
                .expect("server send failed");
            });
        }
    });

    for _ in 0..2_u8 {
        let kp = NoiseKeypair::generate();
        let stream = tokio::net::TcpStream::connect(server_addr)
            .await
            .expect("connect failed");
        let mut t = handshake_initiator(stream, &kp)
            .await
            .expect("client handshake failed");

        t.send_envelope(&ParsedEnvelope {
            msg_type: MsgType::Ping,
            group_id: vec![],
            sender_id: vec![],
            payload: vec![],
            timestamp_ms: 0,
        })
        .await
        .expect("client send failed");

        let pong = tokio::time::timeout(std::time::Duration::from_secs(5), t.recv_envelope())
            .await
            .expect("timeout")
            .expect("recv failed");

        match pong.msg_type {
            MsgType::Pong => {}
            _ => panic!("expected Pong"),
        }

        // Each client sees the *same* server public key (key reuse across connections).
        let seen = t
            .remote_static_public_key()
            .expect("no remote key")
            .to_vec();
        assert_eq!(seen, server_keypair.public_bytes().to_vec());
    }
}