wgcloak/tests/tls_tcp.rs

use std::{collections::HashMap, time::Duration};

use tokio::io::{duplex, split};
use wgcloak::transport::tls_tcp;
use wgcloak::transport::tls_tcp::{
    client_authenticate, server_authenticate, AuthChallenge, FrameCodec, FrameError, FramedReader,
    FramedWriter,
};

#[test]
fn codec_handles_split_and_coalesced_frames() {
    let codec = FrameCodec::new(32).expect("valid maximum");
    let mut wire = codec.encode(b"first").unwrap();
    wire.extend_from_slice(&codec.encode(b"second").unwrap());
    let mut decoder = FrameCodec::new(32).unwrap();
    assert!(decoder.feed(&wire[..1]).unwrap().is_empty());
    assert!(decoder.feed(&wire[1..4]).unwrap().is_empty());
    assert_eq!(
        decoder.feed(&wire[4..]).unwrap(),
        vec![b"first".to_vec(), b"second".to_vec()]
    );
    decoder.finish().unwrap();
}

#[test]
fn codec_rejects_oversized_before_payload() {
    let mut decoder = FrameCodec::new(4).unwrap();
    let error = decoder.feed(&[0, 5]).unwrap_err();
    assert_eq!(
        error,
        FrameError::Oversized {
            length: 5,
            maximum: 4
        }
    );
    assert!(decoder.finish().is_ok());
}

#[test]
fn codec_rejects_maximum_above_u16_wire_limit() {
    assert_eq!(
        FrameCodec::new(65_536).unwrap_err(),
        FrameError::MaxTooLarge
    );
}

#[test]
fn codec_reports_truncated_header_or_payload() {
    let mut decoder = FrameCodec::new(32).unwrap();
    decoder.feed(&[0]).unwrap();
    assert_eq!(decoder.finish(), Err(FrameError::Truncated));
    let mut decoder = FrameCodec::new(32).unwrap();
    decoder.feed(&[0, 3, b'x']).unwrap();
    assert_eq!(decoder.finish(), Err(FrameError::Truncated));
}

#[test]
fn auth_response_is_bound_to_challenge_and_secret() {
    let challenge = AuthChallenge::random();
    let nonce = [9u8; 16];
    let message = tls_tcp::encode_auth_response(&challenge, "laptop", &nonce, &[7u8; 32]).unwrap();
    let mut clients = HashMap::new();
    clients.insert("laptop".to_owned(), vec![7u8; 32]);
    // Exercise the full asynchronous exchange below; this assertion also
    // ensures the generated challenge is not the all-zero sentinel.
    assert_ne!(challenge.bytes, [0u8; 32]);
    assert_eq!(message[0], 0x11);
    assert_eq!(clients["laptop"].len(), 32);
}

#[tokio::test]
async fn tls_stream_authentication_exchange_uses_challenge_response() {
    let (left, right) = duplex(4096);
    let (mut server_reader, mut server_writer) = split(left);
    let (mut client_reader, mut client_writer) = split(right);
    let mut clients = HashMap::new();
    clients.insert("laptop".to_owned(), vec![7u8; 32]);
    let server_clients = clients.clone();

    let server = tokio::spawn(async move {
        server_authenticate(
            &mut server_reader,
            &mut server_writer,
            &server_clients,
            1024,
            Duration::from_secs(1),
        )
        .await
    });
    let client = tokio::spawn(async move {
        client_authenticate(
            &mut client_reader,
            &mut client_writer,
            "laptop",
            &[7u8; 32],
            1024,
            Duration::from_secs(1),
        )
        .await
    });
    assert!(client.await.unwrap().is_ok());
    let authenticated = server.await.unwrap().unwrap();
    assert_eq!(authenticated.id, "laptop");
}

#[tokio::test]
async fn tls_authentication_accepts_validated_maximum_datagram_size() {
    let (left, right) = duplex(4096);
    let (mut server_reader, mut server_writer) = split(left);
    let (mut client_reader, mut client_writer) = split(right);
    let mut clients = HashMap::new();
    clients.insert("laptop".to_owned(), vec![7u8; 32]);
    let server_clients = clients.clone();

    let server = tokio::spawn(async move {
        server_authenticate(
            &mut server_reader,
            &mut server_writer,
            &server_clients,
            65_503 + 512,
            Duration::from_secs(1),
        )
        .await
    });
    let client = tokio::spawn(async move {
        client_authenticate(
            &mut client_reader,
            &mut client_writer,
            "laptop",
            &[7u8; 32],
            65_503 + 512,
            Duration::from_secs(1),
        )
        .await
    });
    assert!(client.await.unwrap().is_ok());
    assert_eq!(server.await.unwrap().unwrap().id, "laptop");
}

#[tokio::test]
async fn bad_secret_is_rejected_without_authenticating() {
    let (left, right) = duplex(4096);
    let (mut server_reader, mut server_writer) = split(left);
    let (mut client_reader, mut client_writer) = split(right);
    let mut clients = HashMap::new();
    clients.insert("laptop".to_owned(), vec![7u8; 32]);
    let server_clients = clients.clone();
    let server = tokio::spawn(async move {
        server_authenticate(
            &mut server_reader,
            &mut server_writer,
            &server_clients,
            1024,
            Duration::from_secs(1),
        )
        .await
    });
    let client = tokio::spawn(async move {
        client_authenticate(
            &mut client_reader,
            &mut client_writer,
            "laptop",
            &[8u8; 32],
            1024,
            Duration::from_secs(1),
        )
        .await
    });
    assert!(matches!(
        client.await.unwrap(),
        Err(tls_tcp::AuthError::Rejected)
    ));
    assert!(matches!(
        server.await.unwrap(),
        Err(tls_tcp::AuthError::InvalidCredentials)
    ));
}

#[tokio::test]
async fn framed_session_preserves_empty_datagrams_and_heartbeats() {
    let (left, right) = duplex(4096);
    let (mut reader, _unused_writer) = split(left);
    let (unused_reader, writer) = split(right);
    drop(unused_reader);
    let mut writer = FramedWriter::new(writer, 1200);
    let mut reader = FramedReader::new(&mut reader, 1200, Some(Duration::from_secs(1)));
    // Write frames using the writer half and read them from the other side.
    // The opposite halves are intentionally dropped only after construction;
    // duplex keeps the active directions independent.
    let send = async {
        writer.send_keepalive().await.unwrap();
        writer.send_datagram(b"").await.unwrap();
        writer.send_datagram(b"payload").await.unwrap();
    };
    tokio::pin!(send);
    send.await;
    assert_eq!(reader.recv_datagram().await.unwrap(), b"");
    assert_eq!(reader.recv_datagram().await.unwrap(), b"payload");
}