From 97dc1b3f8d776c79f299c219f493ab1cfdabb733 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Sat, 20 Jun 2026 22:59:43 +0700 Subject: [PATCH] test(client): add tests for peer-initiated session closure --- crates/irtt-cli/src/shared/client/session.rs | 190 +++++++++++++++++++ crates/irtt-client/src/managed/runner.rs | 84 +++++++- 2 files changed, 273 insertions(+), 1 deletion(-) diff --git a/crates/irtt-cli/src/shared/client/session.rs b/crates/irtt-cli/src/shared/client/session.rs index fdc47e9..5dbaa85 100644 --- a/crates/irtt-cli/src/shared/client/session.rs +++ b/crates/irtt-cli/src/shared/client/session.rs @@ -157,6 +157,29 @@ fn open_event(outcome: &OpenOutcome) -> &ClientEvent { #[cfg(test)] mod tests { use super::*; + use irtt_proto::{ + echo_packet_len, + flags::{self, FLAG_OPEN, FLAG_REPLY}, + layout::PacketLayout, + Clock, Params, ReceivedStats, StampAt, TimestampFields, MAGIC, PROTOCOL_VERSION, + }; + use std::{ + net::{SocketAddr, UdpSocket}, + thread::JoinHandle, + }; + + const TOKEN: u64 = 0x1234_5678_90ab_cdef; + + struct FakeServer { + addr: SocketAddr, + done: JoinHandle<()>, + } + + impl FakeServer { + fn join(self) { + self.done.join().unwrap(); + } + } #[test] fn final_drain_uses_capped_probe_timeout() { @@ -177,4 +200,171 @@ mod tests { assert!(should_drain_final(true, true)); assert!(!should_drain_final(true, false)); } + + #[test] + fn peer_close_followed_by_session_cleanup_is_successful() { + let params = test_params(None, Duration::from_millis(10)); + let server = start_peer_close_server(params); + let mut session = ClientSession::connect(config(server.addr, None), true).unwrap(); + let shutdown_requested = AtomicBool::new(false); + + assert!(matches!( + session.open().unwrap().as_slice(), + [ClientEvent::SessionStarted { .. }] + )); + + let events = session.step(&shutdown_requested).unwrap(); + assert!(events + .iter() + .any(|event| matches!(event, ClientEvent::EchoReply { .. }))); + assert!(events + .iter() + .any(|event| matches!(event, ClientEvent::SessionClosed { .. }))); + assert!(!session.should_continue(&shutdown_requested)); + + session.drain_final(|_| {}).unwrap(); + assert!(session.poll_timeouts().unwrap().is_empty()); + assert!(session.close().unwrap().is_empty()); + server.join(); + } + + fn config(addr: SocketAddr, duration: Option) -> ClientConfig { + ClientConfig { + server_addr: addr.to_string(), + duration, + interval: Duration::from_millis(10), + open_timeouts: vec![Duration::from_millis(200)], + probe_timeout: Duration::from_millis(20), + socket_config: irtt_client::SocketConfig { + recv_timeout: Some(Duration::from_millis(200)), + ..Default::default() + }, + ..ClientConfig::default() + } + } + + fn test_params(duration: Option, interval: Duration) -> Params { + Params { + protocol_version: PROTOCOL_VERSION, + duration_ns: duration.map_or(0, test_duration_ns_i64), + interval_ns: test_duration_ns_i64(interval), + length: 0, + received_stats: ReceivedStats::Both, + stamp_at: StampAt::Both, + clock: Clock::Both, + dscp: 0, + server_fill: None, + } + } + + fn test_duration_ns_i64(duration: Duration) -> i64 { + i64::try_from(duration.as_nanos()).expect("test duration fits i64 nanoseconds") + } + + fn start_peer_close_server(params: Params) -> FakeServer { + let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + let addr = socket.local_addr().unwrap(); + let done = thread::spawn(move || { + let (_, peer) = recv_request(&socket); + socket + .send_to(&open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms), peer) + .unwrap(); + + let (packet, peer) = recv_request(&socket); + let seq = u32::from_le_bytes(packet[12..16].try_into().unwrap()); + socket + .send_to( + &echo_reply_packet( + TOKEN, + seq, + ¶ms, + &TimestampFields::default(), + FLAG_REPLY | flags::FLAG_CLOSE, + ), + peer, + ) + .unwrap(); + + socket + .set_read_timeout(Some(Duration::from_millis(250))) + .unwrap(); + while let Some((packet, _)) = recv_request_timeout(&socket) { + assert_eq!( + packet[3] & flags::FLAG_CLOSE, + 0, + "session cleanup must not send a close after peer close" + ); + } + }); + FakeServer { addr, done } + } + + fn recv_request(socket: &UdpSocket) -> (Vec, SocketAddr) { + let mut buf = [0_u8; 2048]; + let (size, peer) = socket.recv_from(&mut buf).unwrap(); + (buf[..size].to_vec(), peer) + } + + fn recv_request_timeout(socket: &UdpSocket) -> Option<(Vec, SocketAddr)> { + let mut buf = [0_u8; 2048]; + socket + .recv_from(&mut buf) + .ok() + .map(|(size, peer)| (buf[..size].to_vec(), peer)) + } + + fn open_reply(flags: u8, token: u64, params: &Params) -> Vec { + let mut packet = Vec::new(); + packet.extend_from_slice(&MAGIC); + packet.push(flags); + packet.extend_from_slice(&token.to_le_bytes()); + packet.extend_from_slice(¶ms.encode()); + packet + } + + fn echo_reply_packet( + token: u64, + seq: u32, + params: &Params, + timestamps: &TimestampFields, + flags: u8, + ) -> Vec { + let layout = PacketLayout::echo(false, params); + let packet_len = echo_packet_len(false, params) + .expect("session test params must have a non-negative packet length"); + let mut packet = Vec::with_capacity(packet_len); + + packet.extend_from_slice(&MAGIC); + packet.push(flags); + packet.extend_from_slice(&token.to_le_bytes()); + packet.extend_from_slice(&seq.to_le_bytes()); + + if layout.recv_count { + packet.extend_from_slice(&42_u32.to_le_bytes()); + } + if layout.recv_window { + packet.extend_from_slice(&0_u64.to_le_bytes()); + } + if layout.recv_wall { + packet.extend_from_slice(×tamps.recv_wall.unwrap_or(0).to_le_bytes()); + } + if layout.recv_mono { + packet.extend_from_slice(×tamps.recv_mono.unwrap_or(0).to_le_bytes()); + } + if layout.midpoint_wall { + packet.extend_from_slice(×tamps.midpoint_wall.unwrap_or(0).to_le_bytes()); + } + if layout.midpoint_mono { + packet.extend_from_slice(×tamps.midpoint_mono.unwrap_or(0).to_le_bytes()); + } + if layout.send_wall { + packet.extend_from_slice(×tamps.send_wall.unwrap_or(0).to_le_bytes()); + } + if layout.send_mono { + packet.extend_from_slice(×tamps.send_mono.unwrap_or(0).to_le_bytes()); + } + + packet.resize(packet_len, 0); + packet + } } diff --git a/crates/irtt-client/src/managed/runner.rs b/crates/irtt-client/src/managed/runner.rs index bf80863..f24ab12 100644 --- a/crates/irtt-client/src/managed/runner.rs +++ b/crates/irtt-client/src/managed/runner.rs @@ -496,6 +496,53 @@ mod tests { FakeServer { addr, done } } + fn start_peer_close_server(params: Params) -> FakeServer { + let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + let addr = socket.local_addr().unwrap(); + let done = thread::spawn(move || { + let (_, peer) = recv_request(&socket); + socket + .send_to(&open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms), peer) + .unwrap(); + socket + .set_read_timeout(Some(Duration::from_millis(500))) + .unwrap(); + + let Some((packet, peer)) = recv_request_timeout(&socket) else { + return; + }; + let seq = u32::from_le_bytes(packet[12..16].try_into().unwrap()); + let ts = TimestampFields { + recv_wall: Some(1_000_000_000), + recv_mono: Some(100_000), + send_wall: Some(1_000_000_000), + send_mono: Some(100_000), + ..Default::default() + }; + socket + .send_to( + &echo_reply_packet_with_flags( + TOKEN, + seq, + ¶ms, + &ts, + FLAG_REPLY | flags::FLAG_CLOSE, + ), + peer, + ) + .unwrap(); + + while let Some((packet, _)) = recv_request_timeout(&socket) { + assert_eq!( + packet[3] & flags::FLAG_CLOSE, + 0, + "managed cleanup must not send a close after peer close" + ); + } + }); + FakeServer { addr, done } + } + fn recv_request(socket: &UdpSocket) -> (Vec, SocketAddr) { let mut buf = [0_u8; 2048]; let (size, peer) = socket.recv_from(&mut buf).unwrap(); @@ -524,13 +571,23 @@ mod tests { seq: u32, params: &Params, timestamps: &TimestampFields, + ) -> Vec { + echo_reply_packet_with_flags(token, seq, params, timestamps, FLAG_REPLY) + } + + fn echo_reply_packet_with_flags( + token: u64, + seq: u32, + params: &Params, + timestamps: &TimestampFields, + flags: u8, ) -> Vec { let layout = PacketLayout::echo(false, params); let packet_len = test_echo_packet_len(false, params); let mut packet = Vec::with_capacity(packet_len); packet.extend_from_slice(&MAGIC); - packet.push(FLAG_REPLY); + packet.push(flags); packet.extend_from_slice(&token.to_le_bytes()); packet.extend_from_slice(&seq.to_le_bytes()); @@ -682,6 +739,31 @@ mod tests { assert!(late_before_close < close); } + #[test] + fn peer_close_during_managed_run_is_successful() { + let server = start_peer_close_server(test_params(None, Duration::from_millis(10))); + let (session, sub) = ManagedClient::start_with_subscription( + config(server.addr, None), + SubscriberConfig { + capacity: 16, + overflow: SubscriberOverflow::DropNewest, + }, + ) + .unwrap(); + + let events = collect_until_closed(&sub); + let outcome = session.join().unwrap(); + server.join(); + + assert_eq!(outcome.end_reason, SessionEndReason::TestComplete); + assert!(events + .iter() + .any(|event| matches!(event, ClientEvent::EchoReply { .. }))); + assert!(events + .iter() + .any(|event| matches!(event, ClientEvent::SessionClosed { .. }))); + } + #[test] fn continuous_managed_run_can_be_stopped_cleanly() { let server = start_echo_server(test_params(None, Duration::from_millis(10)));