diff --git a/Cargo.lock b/Cargo.lock index 598e4fc..7c1ec3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -713,6 +713,7 @@ dependencies = [ "libc", "socket2", "thiserror 2.0.19", + "tokio", ] [[package]] @@ -1716,6 +1717,19 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" +[[package]] +name = "tokio" +version = "1.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +dependencies = [ + "libc", + "mio", + "pin-project-lite", + "socket2", + "windows-sys", +] + [[package]] name = "typenum" version = "1.20.1" diff --git a/crates/irtt-client/Cargo.toml b/crates/irtt-client/Cargo.toml index e9ef818..10342e4 100644 --- a/crates/irtt-client/Cargo.toml +++ b/crates/irtt-client/Cargo.toml @@ -16,10 +16,12 @@ irtt-proto.workspace = true libc = { version = "0.2.189", optional = true } socket2 = { version = "0.6.5", features = ["all"] } thiserror = "2.0.19" +tokio = { version = "1", default-features = false, features = ["rt", "net", "time", "sync"], optional = true} [features] default = [] ancillary = ["dep:libc"] +tokio = ["dep:tokio"] [package.metadata.docs.rs] all-features = true diff --git a/crates/irtt-client/src/async_client.rs b/crates/irtt-client/src/async_client.rs new file mode 100644 index 0000000..59fc34c --- /dev/null +++ b/crates/irtt-client/src/async_client.rs @@ -0,0 +1,741 @@ +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_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> { + 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..ec4df19 --- /dev/null +++ b/crates/irtt-client/src/async_client/tests.rs @@ -0,0 +1,1099 @@ +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_options::tokio_socket_traffic_class, Client, RunMode, SocketConfig}; + +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() + } +} + +async fn opened_client(addr: SocketAddr, dscp: u8) -> AsyncClient { + let mut client = AsyncClient::connect(config(addr, None, dscp)) + .await + .unwrap(); + client.open().await.unwrap(); + client +} + +fn inject_send(client: &AsyncClient, result: InjectedSend) { + client.test_hooks.sends.borrow_mut().push_back(result); +} + +fn socket_dscp(client: &AsyncClient) -> u32 { + tokio_socket_traffic_class(&client.socket, client.remote).unwrap() >> 2 +} + +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 open_session( + socket: &UdpSocket, + tx: &mpsc::Sender>, + key: Option<&[u8]>, +) -> (Params, SocketAddr) { + let (packet, peer) = recv_packet(socket, tx); + let request = decode_open_request(&packet, key).unwrap(); + send_open_reply( + socket, + peer, + request.params.clone(), + key, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + (request.params, peer) +} + +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| { + open_session(&socket, &tx, key.as_deref()); + 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 (params, peer) = open_session(&socket, &tx, None); + let (probe_packet, _) = recv_packet(&socket, &tx); + let probe = decode_echo_request(&probe_packet, ¶ms, None).unwrap(); + socket + .send_to( + &echo_reply( + ¶ms, + probe.sequence, + TOKEN, + flags::FLAG_REPLY | flags::FLAG_CLOSE, + None, + ), + peer, + ) + .unwrap(); + }) +} + +fn start_echo_close_server(key: Option>) -> TestServer { + start_server(move |socket, tx| { + let (params, peer) = open_session(&socket, &tx, key.as_deref()); + + let (probe_packet, _) = recv_packet(&socket, &tx); + let probe = decode_echo_request(&probe_packet, ¶ms, key.as_deref()).unwrap(); + socket + .send_to( + &echo_reply( + ¶ms, + probe.sequence, + TOKEN, + flags::FLAG_REPLY, + key.as_deref(), + ), + peer, + ) + .unwrap(); + + let (close_packet, _) = recv_packet(&socket, &tx); + assert_eq!( + decode_close_request(&close_packet, key.as_deref()) + .unwrap() + .token, + TOKEN + ); + }) +} + +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 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(); + inject_send(&client, 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| { + open_session(&socket, &tx, None); + 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| { + open_session(&socket, &tx, None); + }); + + 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 = opened_client(server.addr, 0).await; + 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 = opened_client(server.addr, 0).await; + 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), + }, + ]); + inject_send(&client, 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 = opened_client(server.addr, 0).await; + 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(); + inject_send(&client, 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 = opened_client(server.addr, 0).await; + 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 (params, peer) = open_session(&socket, &tx, None); + let (probe_packet, _) = recv_packet(&socket, &tx); + let probe = decode_echo_request(&probe_packet, ¶ms, None).unwrap(); + socket + .send_to( + &echo_reply(¶ms, probe.sequence, TOKEN, flags::FLAG_REPLY, None), + peer, + ) + .unwrap(); + }); + + runtime().block_on(async { + let mut client = opened_client(server.addr, 0).await; + client.send_probe().await.unwrap(); + client.test_hooks.receive_would_block.set(2); + + assert!(matches!( + client.recv().await.unwrap().as_slice(), + [ClientEvent::EchoReply { 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 = opened_client(server.addr, 46).await; + 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!(socket_dscp(&client), 46); + }); + 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| { + open_session(&socket, &tx, None); + 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 = opened_client(close_server.addr, 0).await; + 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 = opened_client(server.addr, 46).await; + let deadline = client.next_send_deadline(); + let attempts = client.test_hooks.send_attempts.get(); + assert_eq!(socket_dscp(&client), 46); + + 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!(socket_dscp(&client), 46); + + inject_send(&client, 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!(socket_dscp(&client), 46); + + 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 = opened_client(server.addr, 46).await; + let deadline = client.next_send_deadline(); + + inject_send(&client, 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!(socket_dscp(&client), 46); + + client.test_hooks.fail_dscp_restore.set(true); + inject_send(&client, 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 = opened_client(server.addr, 46).await; + client.test_hooks.fail_dscp_restore.set(true); + inject_send(&client, 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 successful_short_close_commits_before_length_error() { + let server = start_open_server(None, 0); + runtime().block_on(async { + let mut client = opened_client(server.addr, 0).await; + 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(); + inject_send(&client, 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) + )); + }); + assert_eq!(server.finish().len(), 1); +} + +#[test] +fn successful_close_uses_post_acceptance_timestamp() { + let server = start_open_server(None, 1); + runtime().block_on(async { + let mut client = opened_client(server.addr, 0).await; + 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)); + + assert!(matches!( + client.close().await.unwrap().as_slice(), + [ClientEvent::SessionClosed { at, .. }] if *at == close_sent_at + )); + }); + assert_eq!(server.finish().len(), 2); +} + +#[test] +fn blocking_and_async_hmac_dscp_lifecycle_are_semantically_equivalent() { + let key = b"async-equivalence-key".to_vec(); + let blocking_server = start_echo_close_server(Some(key.clone())); + let mut blocking = + Client::connect(config(blocking_server.addr, Some(key.clone()), 46)).unwrap(); + let blocking_open = blocking.open().unwrap(); + let blocking_sent = blocking.send_probe().unwrap(); + let blocking_reply = blocking.recv_once().unwrap(); + let blocking_close = blocking.close().unwrap(); + assert_eq!(blocking_server.finish().len(), 3); + + let async_server = start_echo_close_server(Some(key.clone())); + let (async_open, async_sent, async_reply, async_close) = runtime().block_on(async { + let mut client = AsyncClient::connect(config(async_server.addr, Some(key), 46)) + .await + .unwrap(); + let opened = client.open().await.unwrap(); + assert_eq!(socket_dscp(&client), 46); + let sent = client.send_probe().await.unwrap(); + let reply = client.recv().await.unwrap(); + let closed = client.close().await.unwrap(); + (opened, sent, reply, closed) + }); + assert_eq!(async_server.finish().len(), 3); + + assert_eq!( + open_negotiated(&blocking_open), + open_negotiated(&async_open) + ); + assert_eq!(open_negotiated(&async_open).params.dscp, 46); + assert_matching_event_shape(&blocking_sent[0], &async_sent[0]); + assert_matching_event_shape(&blocking_reply[0], &async_reply[0]); + assert_matching_event_shape(&blocking_close[0], &async_close[0]); +} + +#[test] +fn blocking_and_async_no_test_are_semantically_equivalent() { + let blocking_no_test_server = start_no_test_server(); + let mut blocking_no_test_config = config(blocking_no_test_server.addr, None, 0); + blocking_no_test_config.run_mode = RunMode::NoTest; + let mut blocking_no_test = Client::connect(blocking_no_test_config).unwrap(); + let blocking_no_test_open = blocking_no_test.open().unwrap(); + assert!(blocking_no_test.is_run_complete()); + assert_eq!(blocking_no_test_server.finish().len(), 1); + + let async_no_test_server = start_no_test_server(); + let mut async_no_test_config = config(async_no_test_server.addr, None, 0); + async_no_test_config.run_mode = RunMode::NoTest; + let async_no_test_open = runtime().block_on(async { + let mut client = AsyncClient::connect(async_no_test_config).await.unwrap(); + let opened = client.open().await.unwrap(); + assert!(client.is_run_complete()); + opened + }); + assert_eq!(async_no_test_server.finish().len(), 1); + assert!(matches!( + blocking_no_test_open, + OpenOutcome::NoTestCompleted { .. } + )); + assert!(matches!( + async_no_test_open, + OpenOutcome::NoTestCompleted { .. } + )); + assert_eq!( + open_negotiated(&blocking_no_test_open), + open_negotiated(&async_no_test_open) + ); +} + +#[test] +fn blocking_and_async_peer_close_are_semantically_equivalent() { + let blocking_peer_close_server = start_peer_close_server(); + let mut blocking_peer_close = + Client::connect(config(blocking_peer_close_server.addr, None, 0)).unwrap(); + blocking_peer_close.open().unwrap(); + blocking_peer_close.send_probe().unwrap(); + let blocking_peer_close_events = blocking_peer_close.recv_once().unwrap(); + assert!(blocking_peer_close.is_peer_closed()); + assert_eq!(blocking_peer_close_server.finish().len(), 2); + + let async_peer_close_server = start_peer_close_server(); + let async_peer_close_events = runtime().block_on(async { + let mut client = AsyncClient::connect(config(async_peer_close_server.addr, None, 0)) + .await + .unwrap(); + client.open().await.unwrap(); + client.send_probe().await.unwrap(); + let events = client.recv().await.unwrap(); + assert!(client.is_peer_closed()); + events + }); + assert_eq!(async_peer_close_server.finish().len(), 2); + assert_eq!( + blocking_peer_close_events.len(), + async_peer_close_events.len() + ); + for (blocking, asynchronous) in blocking_peer_close_events + .iter() + .zip(&async_peer_close_events) + { + assert_matching_event_shape(blocking, asynchronous); + } +} + +fn open_negotiated(outcome: &OpenOutcome) -> &crate::NegotiatedParams { + match outcome { + OpenOutcome::Started { negotiated, .. } + | OpenOutcome::NoTestCompleted { negotiated, .. } => negotiated, + } +} + +fn assert_matching_event_shape(blocking: &ClientEvent, asynchronous: &ClientEvent) { + assert_eq!(event_shape(blocking), event_shape(asynchronous)); +} + +fn event_shape(event: &ClientEvent) -> (&'static str, u64, usize) { + match event { + ClientEvent::EchoSent { seq, bytes, .. } => ("sent", u64::from(*seq), *bytes), + ClientEvent::EchoReply { seq, bytes, .. } => ("reply", u64::from(*seq), *bytes), + ClientEvent::SessionClosed { token, .. } => ("closed", *token, 0), + other => panic!("unexpected equivalence event: {other:?}"), + } +} diff --git a/crates/irtt-client/src/error.rs b/crates/irtt-client/src/error.rs index ee03b2a..4423122 100644 --- a/crates/irtt-client/src/error.rs +++ b/crates/irtt-client/src/error.rs @@ -12,6 +12,10 @@ use thiserror::Error; #[derive(Debug, Error)] #[non_exhaustive] pub enum ClientError { + /// A Tokio-backed client operation was polled without a current runtime. + #[cfg(feature = "tokio")] + #[error("Tokio AsyncClient requires a current Tokio runtime")] + NoTokioRuntime, /// The configured server address could not be resolved to a socket address. /// /// Check the host, port, address-family restrictions, and name service diff --git a/crates/irtt-client/src/lib.rs b/crates/irtt-client/src/lib.rs index 0f04c19..e0971c3 100644 --- a/crates/irtt-client/src/lib.rs +++ b/crates/irtt-client/src/lib.rs @@ -3,6 +3,8 @@ //! `irtt-client` opens IRTT-compatible sessions, sends echo probes, receives //! replies, classifies loss/late/duplicate packets, and emits [`ClientEvent`] //! values for callers to consume directly or aggregate with `irtt-stats`. +//! With the `tokio` feature enabled, `AsyncClient` provides the same low-level +//! lifecycle for callers that own a Tokio runtime and drive readiness directly. //! //! Timing values intentionally preserve signed measurement semantics. When //! server timing is available, [`RttSample::effective`] is adjusted for server @@ -51,6 +53,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 +68,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, diff --git a/crates/irtt-client/src/receive.rs b/crates/irtt-client/src/receive.rs index f239d42..ca5fbd1 100644 --- a/crates/irtt-client/src/receive.rs +++ b/crates/irtt-client/src/receive.rs @@ -62,6 +62,32 @@ pub(crate) fn recv_datagram( linux::recv_datagram(socket, buf) } +#[cfg(all( + feature = "tokio", + not(all(target_os = "linux", feature = "ancillary")) +))] +pub(crate) fn try_recv_tokio_datagram( + socket: &tokio::net::UdpSocket, + buf: &mut [u8], +) -> Result { + let len = socket.try_recv(buf)?; + let received_at = ClientTimestamp::now(); + + Ok(ReceivedDatagram { + len, + received_at, + meta: ReceiveMeta::default(), + }) +} + +#[cfg(all(feature = "tokio", target_os = "linux", feature = "ancillary"))] +pub(crate) fn try_recv_tokio_datagram( + socket: &tokio::net::UdpSocket, + buf: &mut [u8], +) -> Result { + linux::try_recv_tokio_datagram(socket, buf) +} + #[cfg(all(target_os = "linux", feature = "ancillary"))] pub(crate) fn recv_datagram_from( socket: &UdpSocket, @@ -133,4 +159,29 @@ mod tests { assert!(datagram.received_at.mono >= before.mono); assert!(datagram.received_at.mono <= after.mono); } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_try_receive_preserves_fallback_metadata_and_would_block() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_io() + .build() + .unwrap(); + + runtime.block_on(async { + let (sender, receiver) = connected_loopback_pair(); + receiver.set_nonblocking(true).unwrap(); + let receiver = tokio::net::UdpSocket::from_std(receiver).unwrap(); + sender.send(b"tokio").unwrap(); + receiver.readable().await.unwrap(); + + let mut buf = [0_u8; 16]; + let datagram = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap(); + assert_eq!(&buf[..datagram.len], b"tokio"); + assert_eq!(datagram.meta, ReceiveMeta::default()); + + let error = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::WouldBlock); + }); + } } diff --git a/crates/irtt-client/src/receive/linux.rs b/crates/irtt-client/src/receive/linux.rs index ca0c3a4..3509058 100644 --- a/crates/irtt-client/src/receive/linux.rs +++ b/crates/irtt-client/src/receive/linux.rs @@ -2,7 +2,7 @@ use std::{ io, mem, net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, UdpSocket}, - os::fd::AsRawFd, + os::fd::{AsRawFd, RawFd}, ptr, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -49,6 +49,24 @@ fn enable_kernel_rx_timestamps(socket: &UdpSocket) -> io::Result<()> { pub(crate) fn recv_datagram( socket: &UdpSocket, buf: &mut [u8], +) -> Result { + recv_datagram_fd(socket.as_raw_fd(), buf, 0) +} + +#[cfg(feature = "tokio")] +pub(crate) fn try_recv_tokio_datagram( + socket: &tokio::net::UdpSocket, + buf: &mut [u8], +) -> Result { + socket.try_io(tokio::io::Interest::READABLE, || { + recv_datagram_fd(socket.as_raw_fd(), buf, libc::MSG_DONTWAIT) + }) +} + +fn recv_datagram_fd( + socket_fd: RawFd, + buf: &mut [u8], + flags: libc::c_int, ) -> Result { let mut control = ControlBuffer::new(); let mut iov = libc::iovec { @@ -64,8 +82,8 @@ pub(crate) fn recv_datagram( let len = unsafe { // SAFETY: `msg` points to one writable iovec backed by `buf`, and the // control buffer is writable and lives until `recvmsg` returns. The - // socket file descriptor is borrowed from a valid `UdpSocket`. - libc::recvmsg(socket.as_raw_fd(), &mut msg, 0) + // socket file descriptor is borrowed from a valid UDP socket. + libc::recvmsg(socket_fd, &mut msg, flags) }; if len < 0 { return Err(io::Error::last_os_error()); @@ -366,6 +384,52 @@ mod tests { assert!(datagram.received_at.mono <= after.mono); } + #[cfg(feature = "tokio")] + #[test] + fn tokio_recvmsg_would_block_clears_readiness() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_io() + .enable_time() + .build() + .unwrap(); + + runtime.block_on(async { + let (sender, receiver) = connected_ipv4_loopback_pair(); + configure_receive_metadata(&receiver, sender.local_addr().unwrap()).unwrap(); + receiver.set_nonblocking(true).unwrap(); + let receiver = tokio::net::UdpSocket::from_std(receiver).unwrap(); + let mut buf = [0_u8; 16]; + + sender.send(b"first").unwrap(); + tokio::time::timeout(Duration::from_secs(1), receiver.readable()) + .await + .expect("initial readability timed out") + .unwrap(); + + let first = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap(); + assert_eq!(&buf[..first.len], b"first"); + + let error = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::WouldBlock); + + assert!( + tokio::time::timeout(Duration::from_millis(20), receiver.readable()) + .await + .is_err(), + "readable completed from stale readiness after recvmsg WouldBlock" + ); + + sender.send(b"second").unwrap(); + tokio::time::timeout(Duration::from_secs(1), receiver.readable()) + .await + .expect("second readability timed out") + .unwrap(); + + let second = super::try_recv_tokio_datagram(&receiver, &mut buf).unwrap(); + assert_eq!(&buf[..second.len], b"second"); + }); + } + #[test] fn kernel_rx_timestamp_metadata_is_observed_when_kernel_provides_it() { let (sender, receiver) = connected_ipv4_loopback_pair(); diff --git a/crates/irtt-client/src/socket.rs b/crates/irtt-client/src/socket.rs index 767e825..d3ba12b 100644 --- a/crates/irtt-client/src/socket.rs +++ b/crates/irtt-client/src/socket.rs @@ -28,18 +28,57 @@ pub(crate) fn validate_open_timeouts(timeouts: &[Duration]) -> Result<(), Client } pub(crate) fn resolve_remote(config: &ClientConfig) -> Result { + #[cfg(all(test, feature = "tokio"))] + SYNC_RESOLVER_CALLS.with(|calls| calls.set(calls.get() + 1)); + let addr = normalize_server_addr(&config.server_addr); let mut addrs = addr .to_socket_addrs() .map_err(|_| ClientError::Resolve { addr: addr.clone() })?; addrs - .find(|addr| { - (!config.socket_config.ipv4_only || addr.is_ipv4()) - && (!config.socket_config.ipv6_only || addr.is_ipv6()) - }) + .find(|addr| address_family_allowed(config, *addr)) .ok_or(ClientError::Resolve { addr }) } +#[cfg(feature = "tokio")] +pub(crate) async fn resolve_remote_tokio(config: &ClientConfig) -> Result { + let addr = normalize_server_addr(&config.server_addr); + if let Ok(remote) = addr.parse::() { + return address_family_allowed(config, remote) + .then_some(remote) + .ok_or(ClientError::Resolve { addr }); + } + + #[cfg(all(test, feature = "tokio"))] + TOKIO_DNS_LOOKUPS.with(|calls| calls.set(calls.get() + 1)); + + let mut addrs = tokio::net::lookup_host(&addr) + .await + .map_err(|_| ClientError::Resolve { addr: addr.clone() })?; + addrs + .find(|remote| address_family_allowed(config, *remote)) + .ok_or_else(|| ClientError::Resolve { addr: addr.clone() }) +} + +fn address_family_allowed(config: &ClientConfig, remote: SocketAddr) -> bool { + (!config.socket_config.ipv4_only || remote.is_ipv4()) + && (!config.socket_config.ipv6_only || remote.is_ipv6()) +} + +#[cfg(all(test, feature = "tokio"))] +std::thread_local! { + static SYNC_RESOLVER_CALLS: std::cell::Cell = const { std::cell::Cell::new(0) }; + static TOKIO_DNS_LOOKUPS: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +#[cfg(all(test, feature = "tokio"))] +pub(crate) fn resolution_call_counts() -> (usize, usize) { + ( + SYNC_RESOLVER_CALLS.with(std::cell::Cell::get), + TOKIO_DNS_LOOKUPS.with(std::cell::Cell::get), + ) +} + pub(crate) fn normalize_server_addr(addr: &str) -> String { if addr.parse::().is_ok() { return addr.to_owned(); @@ -66,16 +105,16 @@ pub(crate) fn connect_udp_socket( config: &SocketConfig, remote: SocketAddr, ) -> Result { - let domain = if remote.is_ipv4() { - Domain::IPV4 - } else { - Domain::IPV6 - }; - let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?; + let socket = create_connected_udp_socket(config, remote)?; + socket.set_read_timeout(config.recv_timeout)?; + Ok(socket) +} - if config.ipv6_only && remote.is_ipv6() { - socket.set_only_v6(true)?; - } +fn create_connected_udp_socket( + config: &SocketConfig, + remote: SocketAddr, +) -> Result { + let socket = create_prebind_udp_socket(config, remote)?; let bind_addr = config.bind_addr.unwrap_or_else(|| { if remote.is_ipv4() { SocketAddr::from(([0, 0, 0, 0], 0)) @@ -95,10 +134,36 @@ pub(crate) fn connect_udp_socket( if let Some(ttl) = config.ttl { apply_ttl_to_socket(&socket, remote, ttl)?; } - socket.set_read_timeout(config.recv_timeout)?; Ok(socket) } +fn create_prebind_udp_socket( + config: &SocketConfig, + remote: SocketAddr, +) -> Result { + let domain = if remote.is_ipv4() { + Domain::IPV4 + } else { + Domain::IPV6 + }; + let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?; + + if config.ipv6_only && remote.is_ipv6() { + socket.set_only_v6(true)?; + } + Ok(socket) +} + +#[cfg(feature = "tokio")] +pub(crate) fn connect_tokio_udp_socket( + config: &SocketConfig, + remote: SocketAddr, +) -> Result { + let socket = create_connected_udp_socket(config, remote)?; + socket.set_nonblocking(true)?; + Ok(tokio::net::UdpSocket::from_std(socket)?) +} + pub(crate) fn bind_unconnected_udp_socket( config: &SocketConfig, family_remote: SocketAddr, @@ -151,4 +216,192 @@ mod tests { assert_eq!(normalize_server_addr("[::1]"), "[::1]:2112"); assert_eq!(normalize_server_addr("[::1]:1234"), "[::1]:1234"); } + + #[test] + fn address_family_filtering_remains_unchanged() { + let ipv4 = SocketAddr::from(([127, 0, 0, 1], DEFAULT_PORT)); + let ipv6 = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], DEFAULT_PORT)); + let mut config = ClientConfig::default(); + + assert!(address_family_allowed(&config, ipv4)); + assert!(address_family_allowed(&config, ipv6)); + + config.socket_config.ipv4_only = true; + assert!(address_family_allowed(&config, ipv4)); + assert!(!address_family_allowed(&config, ipv6)); + + config.socket_config.ipv4_only = false; + config.socket_config.ipv6_only = true; + assert!(!address_family_allowed(&config, ipv4)); + assert!(address_family_allowed(&config, ipv6)); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_ipv6_only_is_applied_before_bind() { + let config = SocketConfig { + ipv6_only: true, + ..SocketConfig::default() + }; + let remote = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], DEFAULT_PORT)); + + let socket = create_prebind_udp_socket(&config, remote).unwrap(); + + assert!(socket.only_v6().unwrap()); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_ipv6_socket_preserves_v6only_and_explicit_bind() { + let bind_addr = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], 0)); + let peer = match UdpSocket::bind(bind_addr) { + Ok(peer) => peer, + Err(error) if ipv6_loopback_unavailable(&error) => { + eprintln!("skipping IPv6 Tokio socket test: IPv6 loopback unavailable: {error}"); + return; + } + Err(error) => panic!("{error}"), + }; + let remote = peer.local_addr().unwrap(); + let config = SocketConfig { + bind_addr: Some(bind_addr), + ipv6_only: true, + ..SocketConfig::default() + }; + + tokio_runtime().block_on(async { + let socket = connect_tokio_udp_socket(&config, remote).unwrap(); + + assert!(socket2::SockRef::from(&socket).only_v6().unwrap()); + assert!(socket.local_addr().unwrap().is_ipv6()); + assert_eq!(socket.local_addr().unwrap().ip(), bind_addr.ip()); + assert_eq!(socket.peer_addr().unwrap(), remote); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_ipv4_socket_preserves_family_bind_peer_and_ttl() { + let peer = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap(); + let remote = peer.local_addr().unwrap(); + let bind_addr = SocketAddr::from(([127, 0, 0, 1], 0)); + let config = SocketConfig { + bind_addr: Some(bind_addr), + ttl: Some(64), + ..SocketConfig::default() + }; + + tokio_runtime().block_on(async { + let socket = connect_tokio_udp_socket(&config, remote).unwrap(); + + assert!(socket.local_addr().unwrap().is_ipv4()); + assert_eq!(socket.local_addr().unwrap().ip(), bind_addr.ip()); + assert_eq!(socket.peer_addr().unwrap(), remote); + assert_eq!(socket2::SockRef::from(&socket).ttl_v4().unwrap(), 64); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_filters_families_without_using_blocking_resolution() { + let mut config = ClientConfig { + server_addr: "localhost:2112".to_owned(), + ..ClientConfig::default() + }; + config.socket_config.ipv4_only = true; + + tokio_runtime().block_on(async { + let before = resolution_call_counts(); + let remote = resolve_remote_tokio(&config).await.unwrap(); + let after = resolution_call_counts(); + + assert!(remote.is_ipv4()); + assert_eq!(after.0, before.0); + assert_eq!(after.1, before.1 + 1); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_selects_ipv6_for_ipv6_only_hostname() { + let mut config = ClientConfig { + server_addr: "localhost:2112".to_owned(), + ..ClientConfig::default() + }; + config.socket_config.ipv6_only = true; + + tokio_runtime().block_on(async { + let remote = resolve_remote_tokio(&config).await.unwrap(); + assert!(remote.is_ipv6()); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_literal_path_bypasses_all_name_resolution() { + let config = ClientConfig { + server_addr: "127.0.0.1:2112".to_owned(), + ..ClientConfig::default() + }; + + tokio_runtime().block_on(async { + let before = resolution_call_counts(); + let remote = resolve_remote_tokio(&config).await.unwrap(); + let after = resolution_call_counts(); + + assert_eq!(remote, SocketAddr::from(([127, 0, 0, 1], 2112))); + assert_eq!(after, before); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_rejects_disallowed_literal_family() { + let mut config = ClientConfig { + server_addr: "127.0.0.1:2112".to_owned(), + ..ClientConfig::default() + }; + config.socket_config.ipv6_only = true; + + tokio_runtime().block_on(async { + assert!(matches!( + resolve_remote_tokio(&config).await, + Err(ClientError::Resolve { .. }) + )); + }); + } + + #[cfg(feature = "tokio")] + #[test] + fn tokio_resolver_rejects_mixed_family_restrictions() { + let mut config = ClientConfig { + server_addr: "127.0.0.1:2112".to_owned(), + ..ClientConfig::default() + }; + config.socket_config.ipv4_only = true; + config.socket_config.ipv6_only = true; + + tokio_runtime().block_on(async { + assert!(matches!( + resolve_remote_tokio(&config).await, + Err(ClientError::Resolve { .. }) + )); + }); + } + + #[cfg(feature = "tokio")] + fn tokio_runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_current_thread() + .enable_io() + .build() + .unwrap() + } + + #[cfg(feature = "tokio")] + fn ipv6_loopback_unavailable(error: &std::io::Error) -> bool { + matches!( + error.kind(), + std::io::ErrorKind::AddrNotAvailable | std::io::ErrorKind::Unsupported + ) + } } diff --git a/crates/irtt-client/src/socket_options.rs b/crates/irtt-client/src/socket_options.rs index 283e87e..bc8a1d0 100644 --- a/crates/irtt-client/src/socket_options.rs +++ b/crates/irtt-client/src/socket_options.rs @@ -3,6 +3,8 @@ use std::{ net::{SocketAddr, UdpSocket}, }; +use socket2::SockRef; + use crate::{ config::{MAX_DSCP_CODEPOINT, MAX_TTL}, error::ClientError, @@ -33,6 +35,24 @@ pub(crate) fn clear_dscp_on_socket( set_socket_traffic_class(socket, remote, 0, "clear DSCP before close") } +#[cfg(feature = "tokio")] +pub(crate) fn apply_dscp_to_tokio_socket( + socket: &tokio::net::UdpSocket, + remote: SocketAddr, + dscp_codepoint: u8, +) -> Result<(), ClientError> { + let traffic_class = dscp_codepoint_to_traffic_class(dscp_codepoint)?; + set_tokio_socket_traffic_class(socket, remote, traffic_class, "set negotiated DSCP") +} + +#[cfg(feature = "tokio")] +pub(crate) fn clear_dscp_on_tokio_socket( + socket: &tokio::net::UdpSocket, + remote: SocketAddr, +) -> Result<(), ClientError> { + set_tokio_socket_traffic_class(socket, remote, 0, "clear DSCP before close") +} + pub(crate) fn validate_ttl(ttl: u32) -> Result<(), ClientError> { if ttl == 0 || ttl > MAX_TTL { return Err(ClientError::InvalidConfig { @@ -67,6 +87,20 @@ pub(crate) fn socket_traffic_class( }) } +#[cfg(all(test, feature = "tokio"))] +pub(crate) fn tokio_socket_traffic_class( + socket: &tokio::net::UdpSocket, + remote: SocketAddr, +) -> Result { + get_socket_traffic_class_ref(SockRef::from(socket), remote).map_err(|source| { + ClientError::SocketOption { + operation: "read DSCP socket option", + remote, + source, + } + }) +} + #[cfg(test)] pub(crate) fn socket_ttl(socket: &UdpSocket, remote: SocketAddr) -> Result { get_socket_ttl(socket, remote).map_err(|source| ClientError::SocketOption { @@ -77,19 +111,24 @@ pub(crate) fn socket_ttl(socket: &UdpSocket, remote: SocketAddr) -> Result io::Result<()> { + set_socket_ttl_ref(SockRef::from(socket), remote, ttl) +} + +fn set_socket_ttl_ref(socket: SockRef<'_>, remote: SocketAddr, ttl: u32) -> io::Result<()> { if remote.is_ipv4() { - socket2::SockRef::from(socket).set_ttl_v4(ttl) + socket.set_ttl_v4(ttl) } else { - socket2::SockRef::from(socket).set_unicast_hops_v6(ttl) + socket.set_unicast_hops_v6(ttl) } } #[cfg(test)] fn get_socket_ttl(socket: &UdpSocket, remote: SocketAddr) -> io::Result { + let socket = SockRef::from(socket); if remote.is_ipv4() { - socket2::SockRef::from(socket).ttl_v4() + socket.ttl_v4() } else { - socket2::SockRef::from(socket).unicast_hops_v6() + socket.unicast_hops_v6() } } @@ -112,6 +151,30 @@ fn set_socket_traffic_class_io( socket: &UdpSocket, remote: SocketAddr, traffic_class: u32, +) -> io::Result<()> { + set_socket_traffic_class_ref(SockRef::from(socket), remote, traffic_class) +} + +#[cfg(feature = "tokio")] +fn set_tokio_socket_traffic_class( + socket: &tokio::net::UdpSocket, + remote: SocketAddr, + traffic_class: u32, + operation: &'static str, +) -> Result<(), ClientError> { + set_socket_traffic_class_ref(SockRef::from(socket), remote, traffic_class).map_err(|source| { + ClientError::SocketOption { + operation, + remote, + source, + } + }) +} + +fn set_socket_traffic_class_ref( + socket: SockRef<'_>, + remote: SocketAddr, + traffic_class: u32, ) -> io::Result<()> { if remote.is_ipv4() { set_ipv4_traffic_class(socket, traffic_class) @@ -127,8 +190,8 @@ fn set_socket_traffic_class_io( target_os = "illumos", target_os = "haiku", )))] -fn set_ipv4_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result<()> { - socket2::SockRef::from(socket).set_tos_v4(traffic_class) +fn set_ipv4_traffic_class(socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> { + socket.set_tos_v4(traffic_class) } #[cfg(any( @@ -138,7 +201,7 @@ fn set_ipv4_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result< target_os = "illumos", target_os = "haiku", ))] -fn set_ipv4_traffic_class(_socket: &UdpSocket, traffic_class: u32) -> io::Result<()> { +fn set_ipv4_traffic_class(_socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> { unsupported_traffic_class(traffic_class, "IPv4 DSCP socket options") } @@ -153,8 +216,8 @@ fn set_ipv4_traffic_class(_socket: &UdpSocket, traffic_class: u32) -> io::Result target_os = "openbsd", target_os = "cygwin", ))] -fn set_ipv6_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result<()> { - socket2::SockRef::from(socket).set_tclass_v6(traffic_class) +fn set_ipv6_traffic_class(socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> { + socket.set_tclass_v6(traffic_class) } #[cfg(not(any( @@ -168,7 +231,7 @@ fn set_ipv6_traffic_class(socket: &UdpSocket, traffic_class: u32) -> io::Result< target_os = "openbsd", target_os = "cygwin", )))] -fn set_ipv6_traffic_class(_socket: &UdpSocket, traffic_class: u32) -> io::Result<()> { +fn set_ipv6_traffic_class(_socket: SockRef<'_>, traffic_class: u32) -> io::Result<()> { unsupported_traffic_class(traffic_class, "IPv6 DSCP socket options") } @@ -211,8 +274,8 @@ fn unsupported_traffic_class(traffic_class: u32, feature: &'static str) -> io::R target_os = "haiku", )) ))] -fn get_ipv4_traffic_class(socket: &UdpSocket) -> io::Result { - socket2::SockRef::from(socket).tos_v4() +fn get_ipv4_traffic_class(socket: SockRef<'_>) -> io::Result { + socket.tos_v4() } #[cfg(all( @@ -225,7 +288,7 @@ fn get_ipv4_traffic_class(socket: &UdpSocket) -> io::Result { target_os = "haiku", ) ))] -fn get_ipv4_traffic_class(_socket: &UdpSocket) -> io::Result { +fn get_ipv4_traffic_class(_socket: SockRef<'_>) -> io::Result { unsupported_readback("IPv4 DSCP socket option readback") } @@ -243,8 +306,8 @@ fn get_ipv4_traffic_class(_socket: &UdpSocket) -> io::Result { target_os = "cygwin", ) ))] -fn get_ipv6_traffic_class(socket: &UdpSocket) -> io::Result { - socket2::SockRef::from(socket).tclass_v6() +fn get_ipv6_traffic_class(socket: SockRef<'_>) -> io::Result { + socket.tclass_v6() } #[cfg(all( @@ -261,12 +324,17 @@ fn get_ipv6_traffic_class(socket: &UdpSocket) -> io::Result { target_os = "cygwin", )) ))] -fn get_ipv6_traffic_class(_socket: &UdpSocket) -> io::Result { +fn get_ipv6_traffic_class(_socket: SockRef<'_>) -> io::Result { unsupported_readback("IPv6 DSCP socket option readback") } #[cfg(test)] fn get_socket_traffic_class(socket: &UdpSocket, remote: SocketAddr) -> io::Result { + get_socket_traffic_class_ref(SockRef::from(socket), remote) +} + +#[cfg(test)] +fn get_socket_traffic_class_ref(socket: SockRef<'_>, remote: SocketAddr) -> io::Result { if remote.is_ipv4() { get_ipv4_traffic_class(socket) } else { @@ -395,6 +463,40 @@ mod tests { assert_eq!(socket_traffic_class(&socket, remote).unwrap(), 0); } + #[cfg(all( + feature = "tokio", + not(any( + target_os = "fuchsia", + target_os = "redox", + target_os = "solaris", + target_os = "illumos", + target_os = "haiku", + )) + ))] + #[test] + fn tokio_ipv4_socket_option_sets_and_clears_traffic_class() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_io() + .build() + .unwrap(); + + runtime.block_on(async { + let remote = SocketAddr::from(([127, 0, 0, 1], 9)); + let socket = UdpSocket::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap(); + socket.connect(remote).unwrap(); + socket.set_nonblocking(true).unwrap(); + let socket = tokio::net::UdpSocket::from_std(socket).unwrap(); + + apply_dscp_to_tokio_socket(&socket, remote, 46).unwrap(); + let traffic_class = tokio_socket_traffic_class(&socket, remote).unwrap(); + assert_eq!(traffic_class & 0xfc, 184); + assert_eq!(traffic_class & 0b11, 0); + + clear_dscp_on_tokio_socket(&socket, remote).unwrap(); + assert_eq!(tokio_socket_traffic_class(&socket, remote).unwrap(), 0); + }); + } + #[test] #[cfg(any( target_os = "android",