From 357e1356ea0e1ff17e7f8a28b472d2278925caf2 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 13:03:51 +0700 Subject: [PATCH] fix(client): transact blocking session close --- crates/irtt-client/src/client.rs | 84 ++++++- crates/irtt-client/src/client/tests/close.rs | 251 ++++++++++++++++++- crates/irtt-client/src/client/tests/dscp.rs | 129 ++++++++++ 3 files changed, 454 insertions(+), 10 deletions(-) diff --git a/crates/irtt-client/src/client.rs b/crates/irtt-client/src/client.rs index 592e3ce..1d1c85f 100644 --- a/crates/irtt-client/src/client.rs +++ b/crates/irtt-client/src/client.rs @@ -55,6 +55,11 @@ struct PreparedClientOpenFailure { struct ClientTestHooks { fail_open_dscp: Cell, fail_open_timeout_restore: Cell, + fail_close_event_reserve: Cell, + fail_close_dscp_clear: Cell, + fail_close_send: Cell, + close_reported_len: Cell>, + close_send_attempts: Cell, fail_cleanup_send: Cell, recv_buffer_len_override: Cell>, fail_dscp_restore: Cell, @@ -162,16 +167,44 @@ impl Client { /// Send a close request and emit a [`ClientEvent::SessionClosed`] event. /// /// The close event means the client has sent its close packet and stopped - /// tracking the session locally; it is not a server acknowledgement. + /// tracking the session locally; it is not a server acknowledgement. If the + /// send fails, the negotiated DSCP is restored best-effort and the session + /// and probe schedule remain open. pub fn close(&mut self) -> Result, ClientError> { - let socket = &self.socket; - let remote = self.remote; - let events = self.runtime.close_with(|packet| { - clear_dscp_on_socket(socket, remote)?; - socket.send(packet)?; - Ok(()) - })?; + let prepared = self.runtime.prepare_close()?; + let mut events = Vec::new(); + #[cfg(test)] + if self.test_hooks.fail_close_event_reserve.replace(false) { + events + .try_reserve(usize::MAX) + .map_err(|source| ClientError::AllocationFailed { + operation: "close event result", + source, + })?; + } + events + .try_reserve(1) + .map_err(|source| ClientError::AllocationFailed { + operation: "close event result", + source, + })?; + let previous_dscp = self.applied_dscp; + let expected_bytes = prepared.bytes.len(); + + self.clear_close_dscp()?; + let bytes = match self.send_close_datagram(prepared.bytes) { + Ok(bytes) => bytes, + Err(err) => { + self.restore_dscp_best_effort(previous_dscp); + return Err(ClientError::Socket(err)); + } + }; + let event = self.runtime.commit_local_close(prepared.commit); self.schedule = None; + self.applied_dscp = None; + + validate_datagram_length(expected_bytes, bytes)?; + events.push(event); Ok(events) } @@ -244,8 +277,11 @@ impl Client { datagram.received_at, datagram.meta, )?; - if self.runtime.is_terminal() { + if self.runtime.is_peer_closed() { self.schedule = None; + if self.clear_close_dscp().is_ok() { + self.applied_dscp = None; + } } Ok(events) } @@ -607,6 +643,18 @@ impl Client { }; } + fn clear_close_dscp(&self) -> Result<(), ClientError> { + #[cfg(test)] + if self.test_hooks.fail_close_dscp_clear.replace(false) { + return Err(ClientError::SocketOption { + operation: "clear negotiated DSCP", + remote: self.remote, + source: io::Error::other("injected negotiated DSCP clear failure"), + }); + } + clear_dscp_on_socket(&self.socket, self.remote) + } + fn send_cleanup_close_best_effort(&self, packet: Option<&[u8]>) { if let Some(packet) = packet { #[cfg(test)] @@ -640,6 +688,24 @@ impl Client { } self.socket.set_read_timeout(timeout) } + + fn send_close_datagram(&self, packet: &[u8]) -> io::Result { + #[cfg(test)] + { + self.test_hooks + .close_send_attempts + .set(self.test_hooks.close_send_attempts.get() + 1); + if self.test_hooks.fail_close_send.replace(false) { + return Err(io::Error::other("injected close send failure")); + } + } + let bytes = self.socket.send(packet)?; + #[cfg(test)] + if let Some(reported) = self.test_hooks.close_reported_len.take() { + return Ok(reported); + } + Ok(bytes) + } } pub(crate) fn echo_sent_event( diff --git a/crates/irtt-client/src/client/tests/close.rs b/crates/irtt-client/src/client/tests/close.rs index f5c7f23..656692f 100644 --- a/crates/irtt-client/src/client/tests/close.rs +++ b/crates/irtt-client/src/client/tests/close.rs @@ -30,7 +30,13 @@ fn close_sends_one_close_packet_with_negotiated_token() { let mut client = Client::connect(default_test_config(server.addr)).unwrap(); assert_open_started(client.open().unwrap()); let events = client.close().unwrap(); - assert_eq!(events.len(), 1); + assert!(matches!( + events.as_slice(), + [ClientEvent::SessionClosed { token: TOKEN, .. }] + )); + assert!(!client.runtime.is_open()); + assert!(client.schedule.is_none()); + assert_eq!(client.applied_dscp, None); let packets: Vec<_> = server.rx.iter().take(2).collect(); let close = &packets[1]; assert_eq!(close[3], flags::FLAG_CLOSE); @@ -38,6 +44,138 @@ fn close_sends_one_close_packet_with_negotiated_token() { server.join(); } +#[test] +fn close_event_reservation_failure_precedes_dscp_clear_and_send() { + let params = default_params(); + let server = start_fake_server(move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + socket + .send_to( + &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms, None), + peer, + ) + .unwrap(); + let _ = recv_request(&socket, &tx); + }); + let mut client = Client::connect(default_test_config(server.addr)).unwrap(); + assert_open_started(client.open().unwrap()); + client.test_hooks.fail_close_event_reserve.set(true); + client.test_hooks.fail_close_dscp_clear.set(true); + + assert!(matches!( + client.close(), + Err(ClientError::AllocationFailed { + operation: "close event result", + .. + }) + )); + assert_eq!(client.test_hooks.close_send_attempts.get(), 0); + assert!(client.test_hooks.fail_close_dscp_clear.get()); + assert!(client.runtime.is_open()); + assert!(client.schedule.is_some()); + assert_eq!(client.applied_dscp, Some(0)); + + client.test_hooks.fail_close_dscp_clear.set(false); + client.close().unwrap(); + assert_eq!(client.test_hooks.close_send_attempts.get(), 1); + server.join(); +} + +#[test] +fn close_dscp_clear_failure_precedes_send_and_preserves_open_state() { + let params = default_params(); + let server = start_fake_server(move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + socket + .send_to( + &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms, None), + peer, + ) + .unwrap(); + let _ = recv_request(&socket, &tx); + }); + let mut client = Client::connect(default_test_config(server.addr)).unwrap(); + assert_open_started(client.open().unwrap()); + client.test_hooks.fail_close_dscp_clear.set(true); + + assert!(matches!( + client.close(), + Err(ClientError::SocketOption { .. }) + )); + assert_eq!(client.test_hooks.close_send_attempts.get(), 0); + assert!(client.runtime.is_open()); + assert!(client.schedule.is_some()); + assert_eq!(client.applied_dscp, Some(0)); + + client.close().unwrap(); + assert_eq!(client.test_hooks.close_send_attempts.get(), 1); + server.join(); +} + +#[test] +fn close_send_failure_leaves_machine_and_schedule_open_for_retry() { + let params = default_params(); + let server = start_fake_server(move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + socket + .send_to( + &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms, None), + peer, + ) + .unwrap(); + let _ = recv_request(&socket, &tx); + }); + let mut client = Client::connect(default_test_config(server.addr)).unwrap(); + assert_open_started(client.open().unwrap()); + client.test_hooks.fail_close_send.set(true); + + assert!(matches!(client.close(), Err(ClientError::Socket(_)))); + assert!(client.runtime.is_open()); + assert!(client.schedule.is_some()); + assert_eq!(client.applied_dscp, Some(0)); + assert_eq!(client.test_hooks.close_send_attempts.get(), 1); + + assert!(matches!( + client.close().unwrap().as_slice(), + [ClientEvent::SessionClosed { token: TOKEN, .. }] + )); + assert_eq!(client.test_hooks.close_send_attempts.get(), 2); + assert!(matches!(client.close(), Err(ClientError::AlreadyClosed))); + server.join(); +} + +#[test] +fn short_successful_close_commits_before_length_mismatch() { + let params = default_params(); + let server = start_fake_server(move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + socket + .send_to( + &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms, None), + peer, + ) + .unwrap(); + let _ = recv_request(&socket, &tx); + }); + let mut client = Client::connect(default_test_config(server.addr)).unwrap(); + assert_open_started(client.open().unwrap()); + let expected = client.runtime.prepare_close().unwrap().bytes.len(); + client.test_hooks.close_reported_len.set(Some(expected - 1)); + + assert!(matches!( + client.close(), + Err(ClientError::DatagramLengthMismatch { + expected: error_expected, + actual, + }) if error_expected == expected && actual + 1 == expected + )); + assert!(!client.runtime.is_open()); + assert!(client.schedule.is_none()); + assert_eq!(client.applied_dscp, None); + assert!(matches!(client.close(), Err(ClientError::AlreadyClosed))); + server.join(); +} + #[test] fn send_probe_fails_after_close() { let params = default_params(); @@ -132,6 +270,117 @@ fn close_flagged_echo_reply_emits_reply_then_closes_without_sending_close() { server.join(); } +#[test] +fn peer_close_dscp_cleanup_failure_preserves_events_and_closed_state() { + let params = default_params(); + let server = start_fake_server({ + let params = params.clone(); + move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + socket + .send_to( + &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms, None), + peer, + ) + .unwrap(); + let (request, _) = recv_request(&socket, &tx); + let seq = u32::from_le_bytes(request[12..16].try_into().unwrap()); + socket + .send_to( + &echo_reply_packet_with_flags( + TOKEN, + seq, + ¶ms, + &TimestampFields::default(), + None, + FLAG_REPLY | flags::FLAG_CLOSE, + ), + peer, + ) + .unwrap(); + socket + .set_read_timeout(Some(Duration::from_millis(250))) + .unwrap(); + while recv_request_timeout(&socket, &tx).is_some() {} + } + }); + let mut client = Client::connect(ClientConfig { + socket_config: crate::SocketConfig { + recv_timeout: Some(Duration::from_millis(200)), + ..Default::default() + }, + ..default_test_config(server.addr) + }) + .unwrap(); + assert_open_started(client.open().unwrap()); + client.send_probe().unwrap(); + client.test_hooks.fail_close_dscp_clear.set(true); + + assert!(matches!( + client.recv_once().unwrap().as_slice(), + [ + ClientEvent::EchoReply { .. }, + ClientEvent::SessionClosed { token: TOKEN, .. } + ] + )); + assert!(client.is_peer_closed()); + assert!(client.schedule.is_none()); + assert_eq!(client.applied_dscp, Some(0)); + assert_eq!(client.test_hooks.close_send_attempts.get(), 0); + let open = server.rx.recv_timeout(Duration::from_millis(100)).unwrap(); + let echo = server.rx.recv_timeout(Duration::from_millis(100)).unwrap(); + assert_ne!(open[3] & FLAG_OPEN, 0); + assert_eq!(echo[3] & flags::FLAG_CLOSE, 0); + assert!(server.rx.recv_timeout(Duration::from_millis(400)).is_err()); + server.join(); +} + +#[test] +fn recv_available_stops_after_peer_close() { + let params = default_params(); + let server = start_fake_server({ + let params = params.clone(); + move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + socket + .send_to( + &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms, None), + peer, + ) + .unwrap(); + let (request, _) = recv_request(&socket, &tx); + let seq = u32::from_le_bytes(request[12..16].try_into().unwrap()); + let close = echo_reply_packet_with_flags( + TOKEN, + seq, + ¶ms, + &TimestampFields::default(), + None, + FLAG_REPLY | flags::FLAG_CLOSE, + ); + socket.send_to(&close, peer).unwrap(); + socket.send_to(&close, peer).unwrap(); + } + }); + let mut client = Client::connect(default_test_config(server.addr)).unwrap(); + assert_open_started(client.open().unwrap()); + client.send_probe().unwrap(); + + let events = client + .recv_available(RecvBudget { max_packets: 8 }) + .unwrap(); + + assert!(matches!( + events.as_slice(), + [ + ClientEvent::EchoReply { .. }, + ClientEvent::SessionClosed { token: TOKEN, .. } + ] + )); + assert!(client.is_peer_closed()); + server.join(); +} + #[test] fn normal_echo_reply_does_not_close_session() { let params = default_params(); diff --git a/crates/irtt-client/src/client/tests/dscp.rs b/crates/irtt-client/src/client/tests/dscp.rs index 21956c8..c98e876 100644 --- a/crates/irtt-client/src/client/tests/dscp.rs +++ b/crates/irtt-client/src/client/tests/dscp.rs @@ -73,3 +73,132 @@ fn normal_open_uses_negotiated_dscp_not_requested_dscp() { ); server.join(); } + +#[test] +#[cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + target_os = "haiku", +)))] +fn failed_close_send_restores_negotiated_dscp_and_keeps_session_open() { + let mut params = default_params(); + params.dscp = 46; + let server = start_fake_server(move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + socket + .send_to( + &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms, None), + peer, + ) + .unwrap(); + let _ = recv_request(&socket, &tx); + }); + let mut config = default_test_config(server.addr); + config.dscp = 46; + let mut client = Client::connect(config).unwrap(); + assert_open_started(client.open().unwrap()); + client.test_hooks.fail_close_send.set(true); + + let error = client.close().unwrap_err(); + + assert!(matches!(error, ClientError::Socket(_))); + assert!(client.runtime.is_open()); + assert!(client.schedule.is_some()); + assert_eq!(client.applied_dscp, Some(46)); + assert_eq!( + socket_traffic_class(&client.socket, client.remote).unwrap() & 0xfc, + 184 + ); + client.close().unwrap(); + server.join(); +} + +#[test] +fn failed_close_keeps_send_error_primary_when_dscp_restoration_also_fails() { + let params = default_params(); + let server = start_fake_server(move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + socket + .send_to( + &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms, None), + peer, + ) + .unwrap(); + let _ = recv_request(&socket, &tx); + }); + let mut client = Client::connect(default_test_config(server.addr)).unwrap(); + assert_open_started(client.open().unwrap()); + client.test_hooks.fail_close_send.set(true); + client.test_hooks.fail_dscp_restore.set(true); + + let error = client.close().unwrap_err(); + + assert!(matches!(error, ClientError::Socket(_))); + assert!(error.to_string().contains("injected close send failure")); + assert!(client.runtime.is_open()); + assert!(client.schedule.is_some()); + assert_eq!(client.applied_dscp, Some(0)); + client.close().unwrap(); + server.join(); +} + +#[test] +#[cfg(not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + target_os = "haiku", +)))] +fn authenticated_peer_close_clears_schedule_and_negotiated_dscp() { + let mut params = default_params(); + params.dscp = 46; + let server = start_fake_server(move |socket, tx| { + let (_, peer) = recv_request(&socket, &tx); + socket + .send_to( + &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, ¶ms, None), + peer, + ) + .unwrap(); + let (request, _) = recv_request(&socket, &tx); + let seq = u32::from_le_bytes(request[12..16].try_into().unwrap()); + socket + .send_to( + &echo_reply_packet_with_flags( + TOKEN, + seq, + ¶ms, + &TimestampFields::default(), + None, + FLAG_REPLY | flags::FLAG_CLOSE, + ), + peer, + ) + .unwrap(); + }); + let mut config = default_test_config(server.addr); + config.dscp = 46; + let mut client = Client::connect(config).unwrap(); + assert_open_started(client.open().unwrap()); + client.send_probe().unwrap(); + + let events = client.recv_once().unwrap(); + + assert!(matches!( + events.as_slice(), + [ + ClientEvent::EchoReply { .. }, + ClientEvent::SessionClosed { .. } + ] + )); + assert!(client.schedule.is_none()); + assert_eq!(client.applied_dscp, None); + assert_eq!( + socket_traffic_class(&client.socket, client.remote).unwrap(), + 0 + ); + server.join(); +}