From 62a9c8899ea0e64e424b03f25637ea7a2d4f9c11 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 15:30:42 +0700 Subject: [PATCH] feat(client): add Tokio async client --- crates/irtt-client/src/async_client.rs | 750 +++++++++++ crates/irtt-client/src/async_client/tests.rs | 1192 ++++++++++++++++++ crates/irtt-client/src/lib.rs | 4 + 3 files changed, 1946 insertions(+) create mode 100644 crates/irtt-client/src/async_client.rs create mode 100644 crates/irtt-client/src/async_client/tests.rs diff --git a/crates/irtt-client/src/async_client.rs b/crates/irtt-client/src/async_client.rs new file mode 100644 index 0000000..afa3bec --- /dev/null +++ b/crates/irtt-client/src/async_client.rs @@ -0,0 +1,750 @@ +use std::{ + io, + net::SocketAddr, + time::{Duration, Instant}, +}; + +#[cfg(test)] +use std::{ + cell::{Cell, RefCell}, + collections::VecDeque, + future, +}; + +use tokio::time; + +use crate::{ + client::{ + echo_sent_event, + schedule::{instant_abs_diff, ProbeSchedule}, + validate_datagram_length, + }, + config::ClientConfig, + error::ClientError, + event::{ClientEvent, OpenOutcome}, + receive::try_recv_tokio_datagram, + session::machine::{ + recv_buffer_size, OpenDatagramDisposition, PreparedOpenAcceptance, PreparedOpenRequest, + PreparedProbe, SessionMachine, MAX_OPEN_PACKET_SIZE, + }, + socket::{connect_tokio_udp_socket, resolve_remote_tokio, validate_open_timeouts}, + socket_options::{apply_dscp_to_tokio_socket, clear_dscp_on_tokio_socket}, + timing::ClientTimestamp, +}; + +#[cfg(test)] +use crate::client::ProbeSendTimestamps; + +#[derive(Debug)] +struct PreparedAsyncOpen { + machine: PreparedOpenAcceptance, + schedule: Option, + recv_buffer_len: Option, + negotiated_dscp: Option, +} + +#[derive(Debug)] +struct PreparedAsyncOpenFailure { + primary: ClientError, + machine: PreparedOpenAcceptance, +} + +#[cfg(test)] +#[derive(Debug)] +enum InjectedSend { + WouldBlock, + Error, + ReportedLength(usize), +} + +#[cfg(test)] +#[derive(Debug, Default)] +struct AsyncClientTestHooks { + sends: RefCell>, + send_attempts: Cell, + receive_would_block: Cell, + fail_open_dscp: Cell, + fail_cleanup_send: Cell, + fail_peer_close_dscp: Cell, + fail_close_dscp_clear: Cell, + fail_dscp_restore: Cell, + pause_open_before_writable: Cell, + pause_open_before_readable: Cell, + pause_probe_before_writable: Cell, + pause_close_before_writable: Cell, + pause_close_after_would_block: Cell, + probe_timestamps: RefCell>, + close_sent_at: Cell>, +} + +#[cfg(test)] +impl AsyncClientTestHooks { + fn try_send(&self, socket: &tokio::net::UdpSocket, packet: &[u8]) -> io::Result { + self.send_attempts.set(self.send_attempts.get() + 1); + if let Some(injected) = self.sends.borrow_mut().pop_front() { + return match injected { + InjectedSend::WouldBlock => Err(io::Error::from(io::ErrorKind::WouldBlock)), + InjectedSend::Error => Err(io::Error::other("injected send failure")), + InjectedSend::ReportedLength(bytes) => Ok(bytes), + }; + } + socket.try_send(packet) + } + + fn take_probe_timestamps(&self) -> Option { + self.probe_timestamps.borrow_mut().pop_front() + } +} + +/// Low-level Tokio IRTT client for one connected UDP target. +/// +/// `AsyncClient` does not construct, own, or store a Tokio runtime. Its async +/// methods are polled by the caller and require a current Tokio runtime with I/O +/// and time enabled. +#[derive(Debug)] +pub struct AsyncClient { + socket: tokio::net::UdpSocket, + machine: SessionMachine, + schedule: Option, + remote: SocketAddr, + recv_buffer: Vec, + applied_dscp: Option, + prepared_open: Option, + prepared_probe: Option, + #[cfg(test)] + test_hooks: AsyncClientTestHooks, +} + +impl AsyncClient { + /// Resolve the configured server and construct one connected Tokio UDP + /// socket. + /// + /// Polling this future without a current Tokio runtime returns + /// [`ClientError::NoTokioRuntime`]. A runtime without enabled I/O or time + /// drivers is outside this type's runtime contract. + pub async fn connect(config: ClientConfig) -> Result { + tokio::runtime::Handle::try_current().map_err(|_| ClientError::NoTokioRuntime)?; + validate_open_timeouts(&config.open_timeouts)?; + let remote = resolve_remote_tokio(&config).await?; + let machine = SessionMachine::new(config.clone(), remote)?; + let prepared_open = machine.prepare_open_request()?; + let socket = connect_tokio_udp_socket(&config.socket_config, remote)?; + + Ok(Self { + socket, + machine, + schedule: None, + remote, + recv_buffer: vec![0_u8; recv_buffer_size(false, None)?], + applied_dscp: None, + prepared_open: Some(prepared_open), + prepared_probe: None, + #[cfg(test)] + test_hooks: AsyncClientTestHooks::default(), + }) + } + + /// Perform the IRTT open exchange. + /// + /// Each configured attempt uses one absolute deadline for one successful + /// request submission and all replies inspected during that attempt. + /// Malformed, unrelated, wrong-direction, and unauthenticated datagrams are + /// ignored without resending the request. + /// + /// Dropping this future before trusted acceptance leaves the local machine + /// connected with no committed schedule or negotiated adapter state. If it + /// is dropped during best-effort post-token cleanup, cleanup may remain + /// incomplete because this low-level client never detaches cleanup work. + pub async fn open(&mut self) -> Result { + if self.prepared_open.is_none() { + let _ = self.machine.prepare_open_request()?; + unreachable!("connected async clients retain an open request"); + } + + let attempt_count = self.machine.config().open_timeouts.len(); + let mut open_buffer = [0_u8; MAX_OPEN_PACKET_SIZE]; + + for attempt in 0..attempt_count { + let timeout = self.machine.config().open_timeouts[attempt]; + let deadline = Instant::now() + .checked_add(timeout) + .ok_or(ClientError::DurationOverflow)?; + if !self.send_open_request(deadline).await? { + continue; + } + + loop { + let Some(datagram) = self.recv_open_datagram(&mut open_buffer, deadline).await? + else { + break; + }; + if datagram.received_at.mono > deadline { + break; + } + + let reply = match self + .machine + .inspect_open_datagram(&open_buffer[..datagram.len])? + { + OpenDatagramDisposition::Ignore => continue, + OpenDatagramDisposition::Trusted(reply) => reply, + }; + let machine = match self + .machine + .prepare_open_acceptance(reply, datagram.received_at) + { + Ok(machine) => machine, + Err(failure) => { + self.send_cleanup_close_best_effort( + failure.cleanup_close.as_deref(), + deadline, + ) + .await; + return Err(failure.primary); + } + }; + let prepared = match self.prepare_async_open(machine, datagram.received_at) { + Ok(prepared) => prepared, + Err(failure) => { + self.send_cleanup_close_best_effort( + failure.machine.cleanup_close_packet(), + deadline, + ) + .await; + return Err(failure.primary); + } + }; + return Ok(self.commit_async_open(prepared)); + } + } + + Err(ClientError::OpenTimeout) + } + + /// Send one caller-paced probe with transactional protocol and schedule + /// commits. + /// + /// A prepared packet is retained across readiness false positives, + /// cancellation, and socket errors so a later call retries the same logical + /// probe without advancing state before kernel acceptance. + pub async fn send_probe(&mut self) -> Result, ClientError> { + self.machine.ensure_open()?; + if self.prepared_probe.is_none() { + self.prepared_probe = self.machine.prepare_probe()?; + } + if self.prepared_probe.is_none() { + return Ok(Vec::new()); + } + + let mut events = Vec::new(); + loop { + #[cfg(test)] + if self.test_hooks.pause_probe_before_writable.replace(false) { + future::pending::<()>().await; + } + self.socket.writable().await?; + + #[cfg(test)] + let test_timestamps = self.test_hooks.take_probe_timestamps(); + #[cfg(not(test))] + let permission_at = Instant::now(); + #[cfg(test)] + let permission_at = test_timestamps + .map(|timestamps| timestamps.permission_at) + .unwrap_or_else(Instant::now); + + let schedule = self + .schedule + .as_mut() + .expect("open sessions always have a probe schedule"); + if !schedule.permit_probe_at(permission_at) { + self.prepared_probe = None; + return Ok(events); + } + let prepared = self + .prepared_probe + .as_ref() + .expect("prepared probe was retained across readiness"); + let machine_preflight = self.machine.preflight_probe_commit(prepared)?; + let schedule_commit = + schedule.preflight_caller_commit(machine_preflight.next_packets_sent)?; + events + .try_reserve(1) + .map_err(|source| ClientError::AllocationFailed { + operation: "probe event result", + source, + })?; + + let expected_bytes = prepared.bytes.len(); + let scheduled_at = schedule_commit.scheduled_at; + #[cfg(not(test))] + let sent_at = ClientTimestamp::now(); + #[cfg(test)] + let sent_at = test_timestamps + .map(|timestamps| timestamps.sent_at) + .unwrap_or_else(ClientTimestamp::now); + let machine_commit = self + .machine + .finalize_probe_commit(machine_preflight, sent_at)?; + #[cfg(not(test))] + let send_call_start = Instant::now(); + #[cfg(test)] + let send_call_start = test_timestamps + .map(|timestamps| timestamps.send_call_start) + .unwrap_or_else(Instant::now); + #[cfg(test)] + let send_result = self.test_hooks.try_send(&self.socket, &prepared.bytes); + #[cfg(not(test))] + let send_result = self.socket.try_send(&prepared.bytes); + let bytes = match send_result { + Ok(bytes) => bytes, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, + Err(error) => return Err(ClientError::Socket(error)), + }; + + #[cfg(not(test))] + let send_finished_at = Instant::now(); + #[cfg(test)] + let send_finished_at = test_timestamps + .map(|timestamps| timestamps.send_finished_at) + .unwrap_or_else(Instant::now); + let sent = self.machine.commit_probe_sent(machine_commit, bytes); + schedule.commit(schedule_commit); + self.prepared_probe = None; + + let send_call = send_finished_at.saturating_duration_since(send_call_start); + let timer_error = instant_abs_diff(sent_at.mono, scheduled_at); + validate_datagram_length(expected_bytes, bytes)?; + events.push(echo_sent_event( + self.remote, + sent, + send_call, + scheduled_at, + timer_error, + )); + return Ok(events); + } + } + + /// Await and classify one complete UDP datagram. + /// + /// State is validated before readiness is awaited. Readiness false positives + /// retry without changing protocol state. Authenticated peer-close events + /// remain authoritative even if best-effort DSCP cleanup fails. + pub async fn recv(&mut self) -> Result, ClientError> { + self.machine.ensure_open()?; + + let datagram = loop { + self.socket.readable().await?; + #[cfg(test)] + if self.test_hooks.receive_would_block.get() > 0 { + self.test_hooks + .receive_would_block + .set(self.test_hooks.receive_would_block.get() - 1); + continue; + } + match try_recv_tokio_datagram(&self.socket, &mut self.recv_buffer) { + Ok(datagram) => break datagram, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => continue, + Err(error) => return Err(ClientError::Socket(error)), + } + }; + + let events = self.machine.process_received_echo_packet( + &self.recv_buffer[..datagram.len], + datagram.received_at, + datagram.meta, + )?; + if self.machine.is_peer_closed() { + self.schedule = None; + self.prepared_probe = None; + if self.clear_peer_close_dscp().is_ok() { + self.applied_dscp = None; + } + } + Ok(events) + } + + /// Poll protocol timeouts using a newly captured monotonic timestamp. + pub fn poll_timeouts(&mut self) -> Result, ClientError> { + self.poll_timeouts_at(Instant::now()) + } + + /// Poll protocol timeouts using the caller's monotonic timestamp. + pub fn poll_timeouts_at(&mut self, now: Instant) -> Result, ClientError> { + self.machine.poll_timeouts_at(now) + } + + /// Send the retained close packet and commit local close exactly once. + /// + /// DSCP is cleared only after write readiness and immediately before the + /// nonblocking send attempt. A `WouldBlock` result restores it before the + /// next await, so dropping this future at a suspension point cannot leave + /// an otherwise-open session with cleared DSCP. + pub async fn close(&mut self) -> Result, ClientError> { + let prepared = self.machine.prepare_close()?; + let mut events = Vec::new(); + 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(); + + loop { + #[cfg(test)] + if self.test_hooks.pause_close_before_writable.replace(false) { + future::pending::<()>().await; + } + self.socket.writable().await?; + self.clear_close_dscp()?; + let mut rollback = DscpRollback::armed( + &self.socket, + self.remote, + previous_dscp, + #[cfg(test)] + &self.test_hooks.fail_dscp_restore, + ); + #[cfg(test)] + let send_result = self.test_hooks.try_send(&self.socket, prepared.bytes); + #[cfg(not(test))] + let send_result = self.socket.try_send(prepared.bytes); + let bytes = match send_result { + Ok(bytes) => bytes, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + rollback.restore()?; + #[cfg(test)] + if self.test_hooks.pause_close_after_would_block.replace(false) { + future::pending::<()>().await; + } + continue; + } + Err(error) => return Err(ClientError::Socket(error)), + }; + + #[cfg(not(test))] + let close_sent_at = ClientTimestamp::now(); + #[cfg(test)] + let close_sent_at = self + .test_hooks + .close_sent_at + .take() + .unwrap_or_else(ClientTimestamp::now); + let event = self + .machine + .commit_local_close(prepared.commit, close_sent_at); + rollback.disarm(); + self.schedule = None; + self.prepared_probe = None; + self.applied_dscp = None; + + validate_datagram_length(expected_bytes, bytes)?; + events.push(event); + return Ok(events); + } + } + + /// Return the next scheduled probe deadline. + pub fn next_send_deadline(&self) -> Option { + if !self.machine.is_open() { + return None; + } + self.schedule.as_ref()?.next_send_deadline() + } + + /// Return the configured local probe timeout. + pub fn probe_timeout(&self) -> Duration { + self.machine.probe_timeout() + } + + /// Return whether the current run has completed. + pub fn is_run_complete(&self) -> bool { + self.machine.is_terminal() + || self + .schedule + .as_ref() + .is_some_and(|schedule| schedule.is_finished() && self.machine.pending_is_empty()) + } + + /// Return whether an authenticated peer close ended the session. + pub fn is_peer_closed(&self) -> bool { + self.machine.is_peer_closed() + } + + async fn send_open_request(&self, deadline: Instant) -> Result { + let request = self + .prepared_open + .as_ref() + .expect("connected clients retain their prepared open request"); + loop { + #[cfg(test)] + if self.test_hooks.pause_open_before_writable.replace(false) { + future::pending::<()>().await; + } + if !wait_for_writable(&self.socket, deadline).await? { + return Ok(false); + } + #[cfg(test)] + let send_result = self.test_hooks.try_send(&self.socket, &request.bytes); + #[cfg(not(test))] + let send_result = self.socket.try_send(&request.bytes); + match send_result { + Ok(bytes) => { + validate_datagram_length(request.bytes.len(), bytes)?; + return Ok(true); + } + Err(error) if error.kind() == io::ErrorKind::WouldBlock => {} + Err(error) => return Err(ClientError::Socket(error)), + } + } + } + + async fn recv_open_datagram( + &self, + buffer: &mut [u8], + deadline: Instant, + ) -> Result, ClientError> { + loop { + #[cfg(test)] + if self.test_hooks.pause_open_before_readable.replace(false) { + future::pending::<()>().await; + } + if !wait_for_readable(&self.socket, deadline).await? { + return Ok(None); + } + match try_recv_tokio_datagram(&self.socket, buffer) { + Ok(datagram) => return Ok(Some(datagram)), + Err(error) if error.kind() == io::ErrorKind::WouldBlock => {} + Err(error) => return Err(ClientError::Socket(error)), + } + } + } + + fn prepare_async_open( + &mut self, + machine: PreparedOpenAcceptance, + opened_at: ClientTimestamp, + ) -> Result> { + let Some(negotiated) = machine.normal_negotiated() else { + return Ok(PreparedAsyncOpen { + machine, + schedule: None, + recv_buffer_len: None, + negotiated_dscp: None, + }); + }; + let schedule = match ProbeSchedule::new(opened_at.mono, negotiated) { + Ok(schedule) => schedule, + Err(primary) => { + return Err(Box::new(PreparedAsyncOpenFailure { primary, machine })); + } + }; + let recv_buffer_len = match recv_buffer_size(self.machine.has_hmac(), Some(negotiated)) { + Ok(size) => size, + Err(primary) => { + return Err(Box::new(PreparedAsyncOpenFailure { primary, machine })); + } + }; + let negotiated_dscp = match u8::try_from(negotiated.params.dscp) { + Ok(dscp) => dscp, + Err(_) => { + return Err(Box::new(PreparedAsyncOpenFailure { + primary: ClientError::InvalidConfig { + reason: "negotiated dscp must be in range 0..=63".to_owned(), + }, + machine, + })); + } + }; + let additional = recv_buffer_len.saturating_sub(self.recv_buffer.len()); + if let Err(source) = self.recv_buffer.try_reserve(additional) { + return Err(Box::new(PreparedAsyncOpenFailure { + primary: ClientError::AllocationFailed { + operation: "negotiated receive buffer", + source, + }, + machine, + })); + } + #[cfg(test)] + let dscp_result = if self.test_hooks.fail_open_dscp.replace(false) { + Err(ClientError::SocketOption { + operation: "set negotiated DSCP", + remote: self.remote, + source: io::Error::other("injected negotiated DSCP failure"), + }) + } else { + apply_dscp_to_tokio_socket(&self.socket, self.remote, negotiated_dscp) + }; + #[cfg(not(test))] + let dscp_result = apply_dscp_to_tokio_socket(&self.socket, self.remote, negotiated_dscp); + if let Err(primary) = dscp_result { + self.restore_dscp_best_effort(self.applied_dscp); + return Err(Box::new(PreparedAsyncOpenFailure { primary, machine })); + } + + Ok(PreparedAsyncOpen { + machine, + schedule: Some(schedule), + recv_buffer_len: Some(recv_buffer_len), + negotiated_dscp: Some(negotiated_dscp), + }) + } + + fn commit_async_open(&mut self, prepared: PreparedAsyncOpen) -> OpenOutcome { + if let Some(recv_buffer_len) = prepared.recv_buffer_len { + self.recv_buffer.resize(recv_buffer_len, 0); + } + let outcome = self.machine.commit_open(prepared.machine); + self.schedule = prepared.schedule; + self.applied_dscp = prepared.negotiated_dscp; + self.prepared_open = None; + outcome + } + + async fn send_cleanup_close_best_effort(&self, packet: Option<&[u8]>, deadline: Instant) { + let Some(packet) = packet else { + return; + }; + loop { + #[cfg(test)] + if self.test_hooks.fail_cleanup_send.replace(false) { + return; + } + #[cfg(test)] + let send_result = self.test_hooks.try_send(&self.socket, packet); + #[cfg(not(test))] + let send_result = self.socket.try_send(packet); + match send_result { + Ok(_) => return, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => {} + Err(_) => return, + } + match wait_for_writable(&self.socket, deadline).await { + Ok(true) => {} + Ok(false) | Err(_) => return, + } + } + } + + fn clear_peer_close_dscp(&self) -> Result<(), ClientError> { + #[cfg(test)] + if self.test_hooks.fail_peer_close_dscp.replace(false) { + return Err(ClientError::SocketOption { + operation: "clear DSCP before close", + remote: self.remote, + source: io::Error::other("injected peer-close DSCP cleanup failure"), + }); + } + clear_dscp_on_tokio_socket(&self.socket, self.remote) + } + + 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 DSCP before close", + remote: self.remote, + source: io::Error::other("injected close DSCP clear failure"), + }); + } + clear_dscp_on_tokio_socket(&self.socket, self.remote) + } + + fn restore_dscp_best_effort(&self, previous_dscp: Option) { + let _ = match previous_dscp { + Some(dscp) => apply_dscp_to_tokio_socket(&self.socket, self.remote, dscp), + None => clear_dscp_on_tokio_socket(&self.socket, self.remote), + }; + } +} + +struct DscpRollback<'a> { + socket: &'a tokio::net::UdpSocket, + remote: SocketAddr, + previous_dscp: Option, + armed: bool, + #[cfg(test)] + fail_restore: &'a Cell, +} + +impl<'a> DscpRollback<'a> { + fn armed( + socket: &'a tokio::net::UdpSocket, + remote: SocketAddr, + previous_dscp: Option, + #[cfg(test)] fail_restore: &'a Cell, + ) -> Self { + Self { + socket, + remote, + previous_dscp, + armed: true, + #[cfg(test)] + fail_restore, + } + } + + fn restore(&mut self) -> Result<(), ClientError> { + if !self.armed { + return Ok(()); + } + #[cfg(test)] + if self.fail_restore.replace(false) { + return Err(ClientError::SocketOption { + operation: "restore negotiated DSCP", + remote: self.remote, + source: io::Error::other("injected DSCP restore failure"), + }); + } + match self.previous_dscp { + Some(dscp) => apply_dscp_to_tokio_socket(self.socket, self.remote, dscp)?, + None => clear_dscp_on_tokio_socket(self.socket, self.remote)?, + } + self.armed = false; + Ok(()) + } + + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for DscpRollback<'_> { + fn drop(&mut self) { + if self.armed { + let _ = self.restore(); + } + } +} + +async fn wait_for_writable( + socket: &tokio::net::UdpSocket, + deadline: Instant, +) -> Result { + match time::timeout_at(deadline.into(), socket.writable()).await { + Ok(result) => { + result?; + Ok(true) + } + Err(_) => Ok(false), + } +} + +async fn wait_for_readable( + socket: &tokio::net::UdpSocket, + deadline: Instant, +) -> Result { + match time::timeout_at(deadline.into(), socket.readable()).await { + Ok(result) => { + result?; + Ok(true) + } + Err(_) => Ok(false), + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/irtt-client/src/async_client/tests.rs b/crates/irtt-client/src/async_client/tests.rs new file mode 100644 index 0000000..f2dee54 --- /dev/null +++ b/crates/irtt-client/src/async_client/tests.rs @@ -0,0 +1,1192 @@ +use std::{ + future::Future, + net::UdpSocket, + pin::Pin, + sync::mpsc, + task::{Context, Poll, Waker}, + thread::{self, JoinHandle}, + time::SystemTime, +}; + +use irtt_proto::{ + decode_close_request, decode_echo_request, decode_open_request, encode_echo_reply, + encode_open_reply, flags, EchoReply, OpenReply, Params, ReceivedStats, StampAt, + TimestampFields, +}; +use tokio::runtime::{Builder, Runtime}; + +use super::*; +use crate::{ + socket::resolution_call_counts, socket_options::tokio_socket_traffic_class, RunMode, + SocketConfig, WarningKind, +}; + +const TOKEN: u64 = 0x1234_5678_90ab_cdef; + +struct TestServer { + addr: SocketAddr, + packets: mpsc::Receiver>, + done: JoinHandle<()>, +} + +impl TestServer { + fn finish(self) -> Vec> { + self.done.join().unwrap(); + self.packets.try_iter().collect() + } +} + +fn runtime() -> Runtime { + Builder::new_current_thread() + .enable_io() + .enable_time() + .build() + .unwrap() +} + +fn config(addr: SocketAddr, key: Option>, dscp: u8) -> ClientConfig { + ClientConfig { + server_addr: addr.to_string(), + received_stats: ReceivedStats::None, + stamp_at: StampAt::None, + dscp, + hmac_key: key, + open_timeouts: vec![Duration::from_millis(200)], + socket_config: SocketConfig { + recv_timeout: Some(Duration::from_millis(200)), + ..SocketConfig::default() + }, + ..ClientConfig::default() + } +} + +fn start_server(handler: F) -> TestServer +where + F: FnOnce(UdpSocket, mpsc::Sender>) + Send + 'static, +{ + let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + socket + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + let addr = socket.local_addr().unwrap(); + let (tx, packets) = mpsc::channel(); + let done = thread::spawn(move || handler(socket, tx)); + TestServer { + addr, + packets, + done, + } +} + +fn recv_packet(socket: &UdpSocket, tx: &mpsc::Sender>) -> (Vec, SocketAddr) { + let mut buffer = [0_u8; 8192]; + let (len, peer) = socket.recv_from(&mut buffer).unwrap(); + let packet = buffer[..len].to_vec(); + tx.send(packet.clone()).unwrap(); + (packet, peer) +} + +fn send_open_reply( + socket: &UdpSocket, + peer: SocketAddr, + params: Params, + key: Option<&[u8]>, + reply_flags: u8, + token: u64, +) { + let reply = encode_open_reply( + &OpenReply { + flags: reply_flags, + token, + params, + }, + key, + ) + .unwrap(); + socket.send_to(&reply, peer).unwrap(); +} + +fn echo_reply( + params: &Params, + sequence: u32, + token: u64, + reply_flags: u8, + key: Option<&[u8]>, +) -> Vec { + encode_echo_reply( + &EchoReply { + flags: reply_flags, + token, + sequence, + recv_count: None, + recv_window: None, + timestamps: TimestampFields::default(), + payload: Vec::new(), + }, + params, + key, + ) + .unwrap() +} + +fn start_open_server(key: Option>, packets_after_open: usize) -> TestServer { + start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, key.as_deref()).unwrap(); + send_open_reply( + &socket, + peer, + request.params, + key.as_deref(), + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + for _ in 0..packets_after_open { + let _ = recv_packet(&socket, &tx); + } + }) +} + +fn start_no_test_server() -> TestServer { + start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, None).unwrap(); + assert!(request.close); + send_open_reply( + &socket, + peer, + request.params, + None, + flags::FLAG_OPEN | flags::FLAG_REPLY | flags::FLAG_CLOSE, + 0, + ); + }) +} + +fn start_peer_close_server() -> TestServer { + start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, None).unwrap(); + send_open_reply( + &socket, + peer, + request.params.clone(), + None, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + let (probe_packet, _) = recv_packet(&socket, &tx); + let probe = decode_echo_request(&probe_packet, &request.params, None).unwrap(); + socket + .send_to( + &echo_reply( + &request.params, + probe.sequence, + TOKEN, + flags::FLAG_REPLY | flags::FLAG_CLOSE, + None, + ), + peer, + ) + .unwrap(); + }) +} + +fn poll_once(future: Pin<&mut F>) -> Poll { + let mut context = Context::from_waker(Waker::noop()); + Future::poll(future, &mut context) +} + +fn poll_recv_once(client: &mut AsyncClient) -> Poll, ClientError>> { + let mut future = Box::pin(client.recv()); + poll_once(Pin::as_mut(&mut future)) +} + +#[test] +fn connect_requires_current_runtime_when_polled() { + let mut future = Box::pin(AsyncClient::connect(ClientConfig::default())); + assert!(matches!( + poll_once(Pin::as_mut(&mut future)), + Poll::Ready(Err(ClientError::NoTokioRuntime)) + )); +} + +#[test] +fn hostname_connect_uses_tokio_lookup_without_synchronous_resolution() { + runtime().block_on(async { + let mut client_config = config(SocketAddr::from(([127, 0, 0, 1], 2112)), None, 0); + client_config.server_addr = "localhost:2112".to_owned(); + client_config.socket_config.ipv4_only = true; + let before = resolution_call_counts(); + + let client = AsyncClient::connect(client_config).await.unwrap(); + + let after = resolution_call_counts(); + assert!(client.remote.is_ipv4()); + assert_eq!(after.0, before.0); + assert_eq!(after.1, before.1 + 1); + }); +} + +#[test] +fn literal_connect_bypasses_all_name_resolution() { + runtime().block_on(async { + let remote = SocketAddr::from(([127, 0, 0, 1], 2112)); + let before = resolution_call_counts(); + + let client = AsyncClient::connect(config(remote, None, 0)).await.unwrap(); + + assert_eq!(client.remote, remote); + assert_eq!(resolution_call_counts(), before); + }); +} + +#[test] +fn recv_before_open_fails_on_first_poll_without_socket_readiness() { + runtime().block_on(async { + let remote = SocketAddr::from(([127, 0, 0, 1], 2112)); + let mut client = AsyncClient::connect(config(remote, None, 0)).await.unwrap(); + + assert!(matches!( + poll_recv_once(&mut client), + Poll::Ready(Err(ClientError::NotOpen)) + )); + }); +} + +#[test] +fn open_ignores_malformed_and_bad_hmac_without_resending() { + let key = b"open-filter-key".to_vec(); + let server_key = key.clone(); + let server = start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, Some(&server_key)).unwrap(); + socket.send_to(&[0_u8], peer).unwrap(); + let mut bad_hmac = encode_open_reply( + &OpenReply { + flags: flags::FLAG_OPEN | flags::FLAG_REPLY, + token: TOKEN, + params: request.params.clone(), + }, + Some(&server_key), + ) + .unwrap(); + bad_hmac[4] ^= 0xff; + socket.send_to(&bad_hmac, peer).unwrap(); + send_open_reply( + &socket, + peer, + request.params, + Some(&server_key), + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + }); + + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, Some(key), 0)) + .await + .unwrap(); + client + .test_hooks + .sends + .borrow_mut() + .push_back(InjectedSend::WouldBlock); + let attempts = client.test_hooks.send_attempts.get(); + assert!(matches!( + client.open().await.unwrap(), + OpenOutcome::Started { .. } + )); + assert_eq!(client.test_hooks.send_attempts.get(), attempts + 2); + assert!(client.prepared_open.is_none()); + assert!(client.schedule.is_some()); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn ignored_open_traffic_uses_one_request_and_deadline_per_attempt() { + let server = start_server(move |socket, tx| { + for _ in 0..2 { + let (_, peer) = recv_packet(&socket, &tx); + for _ in 0..3 { + socket.send_to(&[0_u8], peer).unwrap(); + } + } + }); + let mut client_config = config(server.addr, None, 0); + client_config.open_timeouts = vec![Duration::from_millis(200), Duration::from_millis(200)]; + + runtime().block_on(async { + let mut client = AsyncClient::connect(client_config).await.unwrap(); + assert!(matches!(client.open().await, Err(ClientError::OpenTimeout))); + assert!(client.machine.prepare_open_request().is_ok()); + assert!(client.schedule.is_none()); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn ignored_open_traffic_cannot_extend_the_absolute_attempt_deadline() { + let server = start_server(move |socket, tx| { + let (first_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&first_packet, None).unwrap(); + socket + .set_read_timeout(Some(Duration::from_millis(10))) + .unwrap(); + let started = Instant::now(); + let mut next_noise = started; + let mut buffer = [0_u8; 512]; + + while started.elapsed() < Duration::from_millis(350) { + let now = Instant::now(); + if now >= next_noise { + socket.send_to(&[0_u8], peer).unwrap(); + next_noise = now + Duration::from_millis(25); + } + match socket.recv_from(&mut buffer) { + Ok((len, second_peer)) => { + tx.send(buffer[..len].to_vec()).unwrap(); + send_open_reply( + &socket, + second_peer, + request.params, + None, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + return; + } + Err(error) + if matches!( + error.kind(), + io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut + ) => {} + Err(error) => panic!("{error}"), + } + } + panic!("second open request did not arrive after the first absolute deadline"); + }); + let mut client_config = config(server.addr, None, 0); + client_config.open_timeouts = vec![Duration::from_millis(200), Duration::from_millis(200)]; + + runtime().block_on(async { + let mut client = AsyncClient::connect(client_config).await.unwrap(); + assert!(matches!( + client.open().await.unwrap(), + OpenOutcome::Started { .. } + )); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn authenticated_rejection_is_terminal_without_retry() { + let server = start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, None).unwrap(); + send_open_reply( + &socket, + peer, + request.params, + None, + flags::FLAG_OPEN | flags::FLAG_REPLY | flags::FLAG_CLOSE, + 0, + ); + }); + let mut client_config = config(server.addr, None, 0); + client_config.open_timeouts = vec![Duration::from_millis(200), Duration::from_millis(200)]; + + runtime().block_on(async { + let mut client = AsyncClient::connect(client_config).await.unwrap(); + assert!(matches!( + client.open().await, + Err(ClientError::ServerRejected) + )); + assert!(client.machine.prepare_open_request().is_ok()); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn opening_cancellation_before_submission_is_transactional() { + let server = start_server(move |socket, tx| { + socket + .set_read_timeout(Some(Duration::from_millis(50))) + .unwrap(); + let mut buffer = [0_u8; 512]; + if let Ok((len, _)) = socket.recv_from(&mut buffer) { + tx.send(buffer[..len].to_vec()).unwrap(); + } + }); + + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 0)) + .await + .unwrap(); + client.test_hooks.pause_open_before_writable.set(true); + let mut open = Box::pin(client.open()); + assert!(matches!(poll_once(Pin::as_mut(&mut open)), Poll::Pending)); + drop(open); + + assert!(client.machine.prepare_open_request().is_ok()); + assert!(client.prepared_open.is_some()); + assert!(client.schedule.is_none()); + assert_eq!(client.applied_dscp, None); + }); + assert!(server.finish().is_empty()); +} + +#[test] +fn opening_cancellation_after_submission_can_retry() { + let server = start_server(move |socket, tx| { + let (first, _) = recv_packet(&socket, &tx); + let first = decode_open_request(&first, None).unwrap(); + let (_, peer) = recv_packet(&socket, &tx); + send_open_reply( + &socket, + peer, + first.params, + None, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + }); + + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 0)) + .await + .unwrap(); + client.socket.writable().await.unwrap(); + client.test_hooks.pause_open_before_readable.set(true); + let mut open = Box::pin(client.open()); + assert!(matches!(poll_once(Pin::as_mut(&mut open)), Poll::Pending)); + drop(open); + + assert!(client.machine.prepare_open_request().is_ok()); + assert!(client.prepared_open.is_some()); + assert!(client.schedule.is_none()); + assert_eq!(client.applied_dscp, None); + assert!(matches!( + client.open().await.unwrap(), + OpenOutcome::Started { .. } + )); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn adapter_failure_sends_preencoded_cleanup_without_committing_open() { + let server = start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, None).unwrap(); + send_open_reply( + &socket, + peer, + request.params, + None, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + let (close, _) = recv_packet(&socket, &tx); + assert_eq!(decode_close_request(&close, None).unwrap().token, TOKEN); + }); + + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 46)) + .await + .unwrap(); + client.test_hooks.fail_open_dscp.set(true); + assert!(matches!( + client.open().await, + Err(ClientError::SocketOption { .. }) + )); + assert!(client.machine.prepare_open_request().is_ok()); + assert!(client.schedule.is_none()); + assert_eq!(client.applied_dscp, None); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn cleanup_failure_preserves_primary_open_error() { + let server = start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, None).unwrap(); + send_open_reply( + &socket, + peer, + request.params, + None, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + }); + + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 46)) + .await + .unwrap(); + client.test_hooks.fail_open_dscp.set(true); + client.test_hooks.fail_cleanup_send.set(true); + assert!(matches!( + client.open().await, + Err(ClientError::SocketOption { + operation: "set negotiated DSCP", + .. + }) + )); + assert!(client.machine.prepare_open_request().is_ok()); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn no_test_open_has_no_schedule_or_dscp_state() { + let server = start_no_test_server(); + runtime().block_on(async { + let mut client_config = config(server.addr, None, 46); + client_config.run_mode = RunMode::NoTest; + let mut client = AsyncClient::connect(client_config).await.unwrap(); + + assert!(matches!( + client.open().await.unwrap(), + OpenOutcome::NoTestCompleted { .. } + )); + assert!(client.prepared_open.is_none()); + assert!(client.schedule.is_none()); + assert_eq!(client.applied_dscp, None); + assert!(client.is_run_complete()); + assert!(matches!( + client.close().await, + Err(ClientError::AlreadyCompleted) + )); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn probe_cancellation_would_block_and_error_retain_one_preparation() { + let server = start_open_server(None, 1); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 0)) + .await + .unwrap(); + client.open().await.unwrap(); + let deadline = client.next_send_deadline(); + let attempts = client.test_hooks.send_attempts.get(); + + client.test_hooks.pause_probe_before_writable.set(true); + let mut send = Box::pin(client.send_probe()); + assert!(matches!(poll_once(Pin::as_mut(&mut send)), Poll::Pending)); + drop(send); + + let retained = client.prepared_probe.as_ref().unwrap().bytes.clone(); + let retained_ptr = client.prepared_probe.as_ref().unwrap().bytes.as_ptr(); + assert_eq!(client.machine.packets_sent(), 0); + assert!(client.machine.pending_is_empty()); + assert_eq!(client.next_send_deadline(), deadline); + assert_eq!(client.test_hooks.send_attempts.get(), attempts); + + client + .test_hooks + .sends + .borrow_mut() + .extend([InjectedSend::WouldBlock, InjectedSend::Error]); + assert!(matches!( + client.send_probe().await, + Err(ClientError::Socket(_)) + )); + assert_eq!(client.prepared_probe.as_ref().unwrap().bytes, retained); + assert_eq!( + client.prepared_probe.as_ref().unwrap().bytes.as_ptr(), + retained_ptr + ); + assert_eq!(client.machine.packets_sent(), 0); + assert!(client.machine.pending_is_empty()); + assert_eq!(client.next_send_deadline(), deadline); + + assert!(matches!( + client.send_probe().await.unwrap().as_slice(), + [ClientEvent::EchoSent { seq: 0, .. }] + )); + assert_eq!(client.machine.packets_sent(), 1); + assert!(client.prepared_probe.is_none()); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn probe_false_readiness_resamples_all_timestamps() { + let server = start_open_server(None, 1); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 0)) + .await + .unwrap(); + client.open().await.unwrap(); + let scheduled_at = client.next_send_deadline().unwrap(); + let first_sent = ClientTimestamp { + wall: SystemTime::UNIX_EPOCH + Duration::from_secs(1), + mono: scheduled_at + Duration::from_millis(3), + }; + let second_sent = ClientTimestamp { + wall: SystemTime::UNIX_EPOCH + Duration::from_secs(2), + mono: scheduled_at + Duration::from_millis(11), + }; + client.test_hooks.probe_timestamps.borrow_mut().extend([ + ProbeSendTimestamps { + permission_at: scheduled_at + Duration::from_millis(1), + sent_at: first_sent, + send_call_start: scheduled_at + Duration::from_millis(4), + send_finished_at: scheduled_at + Duration::from_millis(6), + }, + ProbeSendTimestamps { + permission_at: scheduled_at + Duration::from_millis(8), + sent_at: second_sent, + send_call_start: scheduled_at + Duration::from_millis(12), + send_finished_at: scheduled_at + Duration::from_millis(17), + }, + ]); + client + .test_hooks + .sends + .borrow_mut() + .push_back(InjectedSend::WouldBlock); + + let events = client.send_probe().await.unwrap(); + assert!(matches!( + events.as_slice(), + [ClientEvent::EchoSent { + scheduled_at: event_scheduled, + sent_at, + send_call, + timer_error, + .. + }] if *event_scheduled == scheduled_at + && *sent_at == second_sent + && *send_call == Duration::from_millis(5) + && *timer_error == Duration::from_millis(11) + )); + assert!(matches!( + client + .poll_timeouts_at(second_sent.mono + client.probe_timeout()) + .unwrap() + .as_slice(), + [ClientEvent::EchoLoss { sent_at, .. }] if *sent_at == second_sent + )); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn successful_short_probe_commits_before_length_error() { + let server = start_open_server(None, 0); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 0)) + .await + .unwrap(); + client.open().await.unwrap(); + client.test_hooks.pause_probe_before_writable.set(true); + let mut send = Box::pin(client.send_probe()); + assert!(matches!(poll_once(Pin::as_mut(&mut send)), Poll::Pending)); + drop(send); + let probe_len = client.prepared_probe.as_ref().unwrap().bytes.len(); + let deadline = client.next_send_deadline().unwrap(); + client + .test_hooks + .sends + .borrow_mut() + .push_back(InjectedSend::ReportedLength(probe_len - 1)); + + assert!(matches!( + client.send_probe().await, + Err(ClientError::DatagramLengthMismatch { .. }) + )); + assert_eq!(client.machine.packets_sent(), 1); + assert!(!client.machine.pending_is_empty()); + assert!(client.prepared_probe.is_none()); + assert_eq!( + client.next_send_deadline(), + Some(deadline + client.machine.config().interval) + ); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn stale_retained_probe_is_rejected_before_send() { + let server = start_open_server(None, 0); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 0)) + .await + .unwrap(); + client.open().await.unwrap(); + client.test_hooks.pause_probe_before_writable.set(true); + let mut send = Box::pin(client.send_probe()); + assert!(matches!(poll_once(Pin::as_mut(&mut send)), Poll::Pending)); + drop(send); + + let advanced = client.machine.prepare_probe().unwrap().unwrap(); + let preflight = client.machine.preflight_probe_commit(&advanced).unwrap(); + let commit = client + .machine + .finalize_probe_commit(preflight, ClientTimestamp::now()) + .unwrap(); + client + .machine + .commit_probe_sent(commit, advanced.bytes.len()); + let attempts = client.test_hooks.send_attempts.get(); + + assert!(matches!( + client.send_probe().await, + Err(ClientError::StalePreparedProbe { .. }) + )); + assert_eq!(client.test_hooks.send_attempts.get(), attempts); + assert!(client.prepared_probe.is_some()); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn finite_duration_boundary_clears_preparation_without_send() { + let server = start_open_server(None, 0); + runtime().block_on(async { + let mut client_config = config(server.addr, None, 0); + client_config.duration = Some(Duration::from_millis(20)); + let mut client = AsyncClient::connect(client_config).await.unwrap(); + client.open().await.unwrap(); + let start = client.next_send_deadline().unwrap(); + let attempts = client.test_hooks.send_attempts.get(); + client + .test_hooks + .probe_timestamps + .borrow_mut() + .push_back(ProbeSendTimestamps { + permission_at: start + Duration::from_millis(20), + sent_at: ClientTimestamp { + wall: SystemTime::UNIX_EPOCH, + mono: start + Duration::from_millis(21), + }, + send_call_start: start + Duration::from_millis(22), + send_finished_at: start + Duration::from_millis(23), + }); + + assert!(client.send_probe().await.unwrap().is_empty()); + assert!(client.prepared_probe.is_none()); + assert_eq!(client.machine.packets_sent(), 0); + assert_eq!(client.test_hooks.send_attempts.get(), attempts); + assert!(client.is_run_complete()); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn receive_retries_false_readiness_and_shares_packet_classification() { + let server = start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, None).unwrap(); + send_open_reply( + &socket, + peer, + request.params.clone(), + None, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + let (probe_packet, _) = recv_packet(&socket, &tx); + let probe = decode_echo_request(&probe_packet, &request.params, None).unwrap(); + socket.send_to(&[0_u8], peer).unwrap(); + socket + .send_to( + &echo_reply( + &request.params, + probe.sequence, + TOKEN + 1, + flags::FLAG_REPLY, + None, + ), + peer, + ) + .unwrap(); + let normal = echo_reply( + &request.params, + probe.sequence, + TOKEN, + flags::FLAG_REPLY, + None, + ); + socket.send_to(&normal, peer).unwrap(); + socket.send_to(&normal, peer).unwrap(); + socket + .send_to( + &echo_reply( + &request.params, + probe.sequence, + TOKEN, + flags::FLAG_REPLY | flags::FLAG_CLOSE, + None, + ), + peer, + ) + .unwrap(); + }); + + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 46)) + .await + .unwrap(); + client.open().await.unwrap(); + client.send_probe().await.unwrap(); + client.test_hooks.receive_would_block.set(2); + + assert!(matches!( + client.recv().await.unwrap().as_slice(), + [ClientEvent::Warning { + kind: WarningKind::MalformedOrUnrelatedPacket, + .. + }] + )); + assert!(matches!( + client.recv().await.unwrap().as_slice(), + [ClientEvent::Warning { + kind: WarningKind::WrongToken, + .. + }] + )); + assert!(matches!( + client.recv().await.unwrap().as_slice(), + [ClientEvent::EchoReply { packet_meta, .. }] + if *packet_meta == crate::PacketMeta::default() + )); + assert!(matches!( + client.recv().await.unwrap().as_slice(), + [ClientEvent::DuplicateReply { .. }] + )); + assert!(matches!( + client.recv().await.unwrap().as_slice(), + [ + ClientEvent::DuplicateReply { .. }, + ClientEvent::SessionClosed { .. } + ] + )); + assert!(client.is_peer_closed()); + assert!(client.schedule.is_none()); + assert!(client.prepared_probe.is_none()); + assert_eq!(client.applied_dscp, None); + assert_eq!( + tokio_socket_traffic_class(&client.socket, client.remote).unwrap() & 0xfc, + 0 + ); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn timed_out_probe_reply_is_classified_late() { + let server = start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, None).unwrap(); + send_open_reply( + &socket, + peer, + request.params.clone(), + None, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + let (probe_packet, _) = recv_packet(&socket, &tx); + let probe = decode_echo_request(&probe_packet, &request.params, None).unwrap(); + thread::sleep(Duration::from_millis(30)); + socket + .send_to( + &echo_reply( + &request.params, + probe.sequence, + TOKEN, + flags::FLAG_REPLY, + None, + ), + peer, + ) + .unwrap(); + }); + + runtime().block_on(async { + let mut client_config = config(server.addr, None, 0); + client_config.probe_timeout = Duration::from_millis(5); + let mut client = AsyncClient::connect(client_config).await.unwrap(); + client.open().await.unwrap(); + client.send_probe().await.unwrap(); + time::sleep(Duration::from_millis(10)).await; + assert!(matches!( + client.poll_timeouts().unwrap().as_slice(), + [ClientEvent::EchoLoss { seq: 0, .. }] + )); + assert!(matches!( + client.recv().await.unwrap().as_slice(), + [ClientEvent::LateReply { seq: 0, .. }] + )); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn peer_close_dscp_failure_preserves_events_and_logical_ownership() { + let server = start_peer_close_server(); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 46)) + .await + .unwrap(); + client.open().await.unwrap(); + client.send_probe().await.unwrap(); + client.prepared_probe = client.machine.prepare_probe().unwrap(); + client.test_hooks.fail_peer_close_dscp.set(true); + + assert!(matches!( + client.recv().await.unwrap().as_slice(), + [ + ClientEvent::EchoReply { .. }, + ClientEvent::SessionClosed { .. } + ] + )); + assert!(client.is_peer_closed()); + assert!(client.is_run_complete()); + assert!(client.schedule.is_none()); + assert!(client.prepared_probe.is_none()); + assert_eq!(client.applied_dscp, Some(46)); + assert_eq!(client.next_send_deadline(), None); + assert_eq!( + tokio_socket_traffic_class(&client.socket, client.remote).unwrap() & 0xfc, + 46 << 2 + ); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn recv_after_local_close_and_no_test_fails_on_first_poll() { + let close_server = start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, None).unwrap(); + send_open_reply( + &socket, + peer, + request.params, + None, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + let (close_packet, _) = recv_packet(&socket, &tx); + assert_eq!( + decode_close_request(&close_packet, None).unwrap().token, + TOKEN + ); + }); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(close_server.addr, None, 0)) + .await + .unwrap(); + client.open().await.unwrap(); + client.close().await.unwrap(); + assert!(matches!( + poll_recv_once(&mut client), + Poll::Ready(Err(ClientError::AlreadyClosed)) + )); + }); + assert_eq!(close_server.finish().len(), 2); + + let no_test_server = start_no_test_server(); + runtime().block_on(async { + let mut client_config = config(no_test_server.addr, None, 0); + client_config.run_mode = RunMode::NoTest; + let mut client = AsyncClient::connect(client_config).await.unwrap(); + client.open().await.unwrap(); + assert!(matches!( + poll_recv_once(&mut client), + Poll::Ready(Err(ClientError::AlreadyCompleted)) + )); + }); + assert_eq!(no_test_server.finish().len(), 1); +} + +#[test] +fn close_cancellation_and_would_block_restore_dscp_before_suspend() { + let server = start_open_server(None, 1); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 46)) + .await + .unwrap(); + client.open().await.unwrap(); + let deadline = client.next_send_deadline(); + let attempts = client.test_hooks.send_attempts.get(); + assert_eq!( + tokio_socket_traffic_class(&client.socket, client.remote).unwrap() & 0xfc, + 46 << 2 + ); + + client.test_hooks.pause_close_before_writable.set(true); + let mut close = Box::pin(client.close()); + assert!(matches!(poll_once(Pin::as_mut(&mut close)), Poll::Pending)); + drop(close); + assert!(client.machine.is_open()); + assert!(client.schedule.is_some()); + assert_eq!(client.next_send_deadline(), deadline); + assert_eq!(client.applied_dscp, Some(46)); + assert_eq!(client.test_hooks.send_attempts.get(), attempts); + assert_eq!( + tokio_socket_traffic_class(&client.socket, client.remote).unwrap() & 0xfc, + 46 << 2 + ); + + client + .test_hooks + .sends + .borrow_mut() + .push_back(InjectedSend::WouldBlock); + client.test_hooks.pause_close_after_would_block.set(true); + let mut close = Box::pin(client.close()); + assert!(matches!(poll_once(Pin::as_mut(&mut close)), Poll::Pending)); + drop(close); + assert!(client.machine.is_open()); + assert!(client.schedule.is_some()); + assert_eq!(client.applied_dscp, Some(46)); + assert_eq!( + tokio_socket_traffic_class(&client.socket, client.remote).unwrap() & 0xfc, + 46 << 2 + ); + + assert!(matches!( + client.close().await.unwrap().as_slice(), + [ClientEvent::SessionClosed { .. }] + )); + assert!(client.schedule.is_none()); + assert_eq!(client.applied_dscp, None); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn close_errors_preserve_open_state_and_primary_send_error() { + let server = start_open_server(None, 0); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 46)) + .await + .unwrap(); + client.open().await.unwrap(); + let deadline = client.next_send_deadline(); + + client + .test_hooks + .sends + .borrow_mut() + .push_back(InjectedSend::Error); + assert!(matches!(client.close().await, Err(ClientError::Socket(_)))); + assert!(client.machine.is_open()); + assert!(client.schedule.is_some()); + assert_eq!(client.next_send_deadline(), deadline); + assert_eq!(client.applied_dscp, Some(46)); + assert_eq!( + tokio_socket_traffic_class(&client.socket, client.remote).unwrap() & 0xfc, + 46 << 2 + ); + + client.test_hooks.fail_dscp_restore.set(true); + client + .test_hooks + .sends + .borrow_mut() + .push_back(InjectedSend::Error); + assert!(matches!(client.close().await, Err(ClientError::Socket(_)))); + assert!(client.machine.is_open()); + assert!(client.schedule.is_some()); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn close_would_block_restore_failure_stops_without_committing() { + let server = start_open_server(None, 0); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 46)) + .await + .unwrap(); + client.open().await.unwrap(); + client.test_hooks.fail_dscp_restore.set(true); + client + .test_hooks + .sends + .borrow_mut() + .push_back(InjectedSend::WouldBlock); + + assert!(matches!( + client.close().await, + Err(ClientError::SocketOption { + operation: "restore negotiated DSCP", + .. + }) + )); + assert!(client.machine.is_open()); + assert!(client.schedule.is_some()); + assert_eq!(client.applied_dscp, Some(46)); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn close_timestamp_is_sampled_after_success_and_short_send_commits() { + let server = start_open_server(None, 0); + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 0)) + .await + .unwrap(); + client.open().await.unwrap(); + client.test_hooks.pause_probe_before_writable.set(true); + let mut send = Box::pin(client.send_probe()); + assert!(matches!(poll_once(Pin::as_mut(&mut send)), Poll::Pending)); + drop(send); + assert!(client.prepared_probe.is_some()); + + let close_len = client.machine.prepare_close().unwrap().bytes.len(); + let close_sent_at = ClientTimestamp { + wall: SystemTime::UNIX_EPOCH + Duration::from_secs(44), + mono: Instant::now() + Duration::from_secs(1), + }; + client.test_hooks.close_sent_at.set(Some(close_sent_at)); + client + .test_hooks + .sends + .borrow_mut() + .push_back(InjectedSend::ReportedLength(close_len - 1)); + + assert!(matches!( + client.close().await, + Err(ClientError::DatagramLengthMismatch { .. }) + )); + assert!(!client.machine.is_open()); + assert!(client.schedule.is_none()); + assert!(client.prepared_probe.is_none()); + assert_eq!(client.applied_dscp, None); + assert!(matches!( + client.machine.prepare_close(), + Err(ClientError::AlreadyClosed) + )); + + let server = start_open_server(None, 1); + let mut exact = AsyncClient::connect(config(server.addr, None, 0)) + .await + .unwrap(); + exact.open().await.unwrap(); + exact.test_hooks.close_sent_at.set(Some(close_sent_at)); + let events = exact.close().await.unwrap(); + assert!(matches!( + events.as_slice(), + [ClientEvent::SessionClosed { at, .. }] if *at == close_sent_at + )); + assert!(matches!( + exact.close().await, + Err(ClientError::AlreadyClosed) + )); + assert_eq!(server.finish().len(), 2); + }); + assert_eq!(server.finish().len(), 1); +} diff --git a/crates/irtt-client/src/lib.rs b/crates/irtt-client/src/lib.rs index 0f04c19..4ecbc75 100644 --- a/crates/irtt-client/src/lib.rs +++ b/crates/irtt-client/src/lib.rs @@ -51,6 +51,8 @@ )] #![cfg_attr(all(target_os = "linux", feature = "ancillary"), deny(unsafe_code))] +#[cfg(feature = "tokio")] +mod async_client; mod client; mod config; mod error; @@ -64,6 +66,8 @@ mod socket; mod socket_options; mod timing; +#[cfg(feature = "tokio")] +pub use async_client::AsyncClient; pub use client::Client; pub use config::{ ClientAuthConfig, ClientConfig, NegotiationPolicy, RecvBudget, RunMode, SocketConfig,