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_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 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"); }