refactor(client): transact blocking session opening

crates/irtt-client/src/client.rs

@@ -4,6 +4,9 @@ use std::{
     time::{Duration, Instant},
 };
 
+#[cfg(test)]
+use std::cell::Cell;
+
 #[cfg(test)]
 use irtt_proto::{flags, Params, TimestampFields, PROTOCOL_VERSION};
 
@@ -12,7 +15,10 @@ use crate::{
     error::ClientError,
     event::{ClientEvent, OpenOutcome},
     receive::recv_datagram,
-    session::machine::{recv_buffer_size, ProbeSent, SessionMachine, MAX_OPEN_PACKET_SIZE},
+    session::machine::{
+        recv_buffer_size, OpenDatagramDisposition, PreparedOpenAcceptance, ProbeSent,
+        SessionMachine, MAX_OPEN_PACKET_SIZE,
+    },
     socket::{connect_udp_socket, resolve_remote, validate_open_timeouts},
     socket_options::{apply_dscp_to_socket, clear_dscp_on_socket},
     timing::ClientTimestamp,
@@ -28,6 +34,34 @@ enum ProbeScheduleMode {
     Managed { scheduled_at: Instant },
 }
 
+#[derive(Debug)]
+struct PreparedClientOpen {
+    machine: PreparedOpenAcceptance,
+    schedule: Option<ProbeSchedule>,
+    recv_buffer_len: Option<usize>,
+    negotiated_dscp: Option<u8>,
+    previous_dscp: Option<u8>,
+    post_open_recv_timeout: Option<Duration>,
+}
+
+#[derive(Debug)]
+struct PreparedClientOpenFailure {
+    primary: ClientError,
+    machine: PreparedOpenAcceptance,
+}
+
+#[cfg(test)]
+#[derive(Debug, Default)]
+struct ClientTestHooks {
+    fail_open_dscp: Cell<bool>,
+    fail_open_timeout_restore: Cell<bool>,
+    fail_cleanup_send: Cell<bool>,
+    recv_buffer_len_override: Cell<Option<usize>>,
+    fail_dscp_restore: Cell<bool>,
+    last_restored_read_timeout: Cell<Option<Duration>>,
+    prepared_active_session_before_dscp: Cell<bool>,
+}
+
 #[cfg(test)]
 #[derive(Debug, Clone, Copy)]
 pub(crate) struct ProbeSendTimestamps {
@@ -60,6 +94,9 @@ pub struct Client {
     socket: UdpSocket,
     remote: SocketAddr,
     recv_buffer: Vec<u8>,
+    applied_dscp: Option<u8>,
+    #[cfg(test)]
+    test_hooks: ClientTestHooks,
     #[cfg(test)]
     probe_reported_len: Option<usize>,
     #[cfg(test)]
@@ -86,6 +123,9 @@ impl Client {
             socket,
             remote,
             recv_buffer: vec![0_u8; recv_buffer_size(false, None)?],
+            applied_dscp: None,
+            #[cfg(test)]
+            test_hooks: ClientTestHooks::default(),
             #[cfg(test)]
             probe_reported_len: None,
             #[cfg(test)]
@@ -99,66 +139,24 @@ impl Client {
     ///
     /// On success, returns the negotiated open outcome and transitions the
     /// client into either an open probe session or completed no-test state.
-    /// Open attempts use [`ClientConfig::open_timeouts`].
+    /// Open attempts use [`ClientConfig::open_timeouts`]. Malformed, unrelated,
+    /// or unauthenticated datagrams are ignored until the current attempt's
+    /// absolute deadline, so one attempt may consume several datagrams without
+    /// retransmitting. Silence or ignored traffic eventually produces
+    /// [`ClientError::OpenTimeout`], while authenticated incompatibility remains
+    /// terminal.
+    ///
+    /// When a trusted reply allocates a token but later negotiation or socket
+    /// preparation fails, the client sends a best-effort cleanup close and
+    /// preserves the original failure. Failed opening never leaves the session
+    /// machine open.
     pub fn open(&mut self) -> Result<OpenOutcome, ClientError> {
-        let outcome = (|| -> Result<OpenOutcome, ClientError> {
-            let packet = self.runtime.open_packet()?;
-            let mut buf = [0_u8; MAX_OPEN_PACKET_SIZE];
-
-            for timeout in &self.runtime.config().open_timeouts {
-                self.socket.set_read_timeout(Some(*timeout))?;
-                self.socket.send(&packet)?;
-
-                match self.socket.recv(&mut buf) {
-                    Ok(size) => {
-                        let reply = self.runtime.decode_open_reply(&buf[..size])?;
-                        let opened_at = ClientTimestamp::now();
-                        let remote = self.remote;
-                        let has_hmac = self.runtime.has_hmac();
-                        let socket = &self.socket;
-                        let recv_buffer = &mut self.recv_buffer;
-                        let mut schedule = None;
-
-                        let outcome =
-                            self.runtime
-                                .accept_open_reply(reply, opened_at, |negotiated| {
-                                    let size = recv_buffer_size(has_hmac, Some(negotiated))?;
-                                    let negotiated_dscp = u8::try_from(negotiated.params.dscp)
-                                        .map_err(|_| ClientError::InvalidConfig {
-                                            reason: "negotiated dscp must be in range 0..=63"
-                                                .to_owned(),
-                                        })?;
-                                    apply_dscp_to_socket(socket, remote, negotiated_dscp)?;
-                                    recv_buffer.resize(size, 0);
-                                    schedule =
-                                        Some(ProbeSchedule::new(opened_at.mono, negotiated)?);
-                                    Ok(())
-                                })?;
-                        self.schedule = schedule;
-                        return Ok(outcome);
-                    }
-                    Err(err)
-                        if matches!(
-                            err.kind(),
-                            io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
-                        ) => {}
-                    Err(err) => return Err(ClientError::Socket(err)),
-                }
-            }
-
-            Err(ClientError::OpenTimeout)
-        })();
-
-        let restore = self
-            .socket
-            .set_read_timeout(self.runtime.config().socket_config.recv_timeout);
-
-        match (outcome, restore) {
-            (Ok(outcome), Ok(())) => Ok(outcome),
-            (Ok(_), Err(source)) => Err(ClientError::ReadTimeoutRestore { source }),
-            (Err(err), Ok(())) => Err(err),
-            (Err(err), Err(_)) => Err(err),
+        let result = self.open_transaction();
+        if result.is_err() {
+            let _ =
+                self.restore_open_read_timeout(self.runtime.config().socket_config.recv_timeout);
         }
+        result
     }
 
     /// Send a close request and emit a [`ClientEvent::SessionClosed`] event.
@@ -417,6 +415,231 @@ impl Client {
         ));
         Ok(events)
     }
+
+    fn open_transaction(&mut self) -> Result<OpenOutcome, ClientError> {
+        let request = self.runtime.prepare_open_request()?;
+        let mut buf = [0_u8; MAX_OPEN_PACKET_SIZE];
+        let attempt_count = self.runtime.config().open_timeouts.len();
+
+        for attempt in 0..attempt_count {
+            let timeout = self.runtime.config().open_timeouts[attempt];
+            let deadline = Instant::now()
+                .checked_add(timeout)
+                .ok_or(ClientError::DurationOverflow)?;
+            self.socket.set_read_timeout(Some(timeout))?;
+            self.socket.send(&request.bytes)?;
+
+            loop {
+                let now = Instant::now();
+                let Some(remaining) = deadline.checked_duration_since(now) else {
+                    break;
+                };
+                if remaining.is_zero() {
+                    break;
+                }
+                self.socket.set_read_timeout(Some(remaining))?;
+
+                let datagram = match recv_datagram(&self.socket, &mut buf) {
+                    Ok(datagram) => datagram,
+                    Err(err)
+                        if matches!(
+                            err.kind(),
+                            io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
+                        ) =>
+                    {
+                        if Instant::now() >= deadline {
+                            break;
+                        }
+                        continue;
+                    }
+                    Err(err) => return Err(ClientError::Socket(err)),
+                };
+                if datagram.received_at.mono > deadline {
+                    break;
+                }
+
+                let reply = match self.runtime.inspect_open_datagram(&buf[..datagram.len])? {
+                    OpenDatagramDisposition::Ignore => continue,
+                    OpenDatagramDisposition::Trusted(reply) => reply,
+                };
+                let machine = match self
+                    .runtime
+                    .prepare_open_acceptance(reply, datagram.received_at)
+                {
+                    Ok(machine) => machine,
+                    Err(failure) => {
+                        self.send_cleanup_close_best_effort(failure.cleanup_close.as_deref());
+                        return Err(failure.primary);
+                    }
+                };
+                let prepared = match self.prepare_client_open(machine, datagram.received_at) {
+                    Ok(prepared) => prepared,
+                    Err(failure) => {
+                        self.send_cleanup_close_best_effort(failure.machine.cleanup_close_packet());
+                        return Err(failure.primary);
+                    }
+                };
+                return self.apply_prepared_open(prepared);
+            }
+        }
+
+        Err(ClientError::OpenTimeout)
+    }
+
+    fn prepare_client_open(
+        &self,
+        machine: PreparedOpenAcceptance,
+        opened_at: ClientTimestamp,
+    ) -> Result<PreparedClientOpen, Box<PreparedClientOpenFailure>> {
+        let Some(negotiated) = machine.normal_negotiated() else {
+            return Ok(PreparedClientOpen {
+                machine,
+                schedule: None,
+                recv_buffer_len: None,
+                negotiated_dscp: None,
+                previous_dscp: self.applied_dscp,
+                post_open_recv_timeout: self.runtime.config().socket_config.recv_timeout,
+            });
+        };
+        let schedule = match ProbeSchedule::new(opened_at.mono, negotiated) {
+            Ok(schedule) => schedule,
+            Err(primary) => return Err(Box::new(PreparedClientOpenFailure { primary, machine })),
+        };
+        let recv_buffer_len = match recv_buffer_size(self.runtime.has_hmac(), Some(negotiated)) {
+            Ok(size) => size,
+            Err(primary) => return Err(Box::new(PreparedClientOpenFailure { primary, machine })),
+        };
+        #[cfg(test)]
+        let recv_buffer_len = self
+            .test_hooks
+            .recv_buffer_len_override
+            .get()
+            .unwrap_or(recv_buffer_len);
+        let negotiated_dscp = match u8::try_from(negotiated.params.dscp) {
+            Ok(dscp) => dscp,
+            Err(_) => {
+                return Err(Box::new(PreparedClientOpenFailure {
+                    primary: ClientError::InvalidConfig {
+                        reason: "negotiated dscp must be in range 0..=63".to_owned(),
+                    },
+                    machine,
+                }));
+            }
+        };
+
+        Ok(PreparedClientOpen {
+            machine,
+            schedule: Some(schedule),
+            recv_buffer_len: Some(recv_buffer_len),
+            negotiated_dscp: Some(negotiated_dscp),
+            previous_dscp: self.applied_dscp,
+            post_open_recv_timeout: self.runtime.config().socket_config.recv_timeout,
+        })
+    }
+
+    fn apply_prepared_open(
+        &mut self,
+        prepared: PreparedClientOpen,
+    ) -> Result<OpenOutcome, ClientError> {
+        let PreparedClientOpen {
+            machine,
+            schedule,
+            recv_buffer_len,
+            negotiated_dscp,
+            previous_dscp,
+            post_open_recv_timeout,
+        } = prepared;
+
+        if let (Some(recv_buffer_len), Some(negotiated_dscp)) = (recv_buffer_len, negotiated_dscp) {
+            let previous_len = self.recv_buffer.len();
+            let additional = recv_buffer_len.saturating_sub(previous_len);
+            if let Err(source) = self.recv_buffer.try_reserve(additional) {
+                self.send_cleanup_close_best_effort(machine.cleanup_close_packet());
+                return Err(ClientError::AllocationFailed {
+                    operation: "negotiated receive buffer",
+                    source,
+                });
+            }
+            self.recv_buffer.resize(recv_buffer_len, 0);
+
+            #[cfg(test)]
+            self.test_hooks
+                .prepared_active_session_before_dscp
+                .set(machine.has_prepared_active_session());
+            if let Err(primary) = self.apply_open_dscp(negotiated_dscp) {
+                self.recv_buffer.truncate(previous_len);
+                self.restore_dscp_best_effort(previous_dscp);
+                self.send_cleanup_close_best_effort(machine.cleanup_close_packet());
+                return Err(primary);
+            }
+            if let Err(source) = self.restore_open_read_timeout(post_open_recv_timeout) {
+                let primary = ClientError::ReadTimeoutRestore { source };
+                self.recv_buffer.truncate(previous_len);
+                self.restore_dscp_best_effort(previous_dscp);
+                self.send_cleanup_close_best_effort(machine.cleanup_close_packet());
+                return Err(primary);
+            }
+
+            let outcome = self.runtime.commit_open(machine);
+            self.schedule = schedule;
+            self.applied_dscp = Some(negotiated_dscp);
+            Ok(outcome)
+        } else {
+            debug_assert!(schedule.is_none());
+            if let Err(source) = self.restore_open_read_timeout(post_open_recv_timeout) {
+                return Err(ClientError::ReadTimeoutRestore { source });
+            }
+            let outcome = self.runtime.commit_open(machine);
+            self.schedule = None;
+            self.applied_dscp = None;
+            Ok(outcome)
+        }
+    }
+
+    fn restore_dscp_best_effort(&self, previous_dscp: Option<u8>) {
+        #[cfg(test)]
+        if self.test_hooks.fail_dscp_restore.replace(false) {
+            return;
+        }
+        let _ = match previous_dscp {
+            Some(dscp) => apply_dscp_to_socket(&self.socket, self.remote, dscp),
+            None => clear_dscp_on_socket(&self.socket, self.remote),
+        };
+    }
+
+    fn send_cleanup_close_best_effort(&self, packet: Option<&[u8]>) {
+        if let Some(packet) = packet {
+            #[cfg(test)]
+            if self.test_hooks.fail_cleanup_send.replace(false) {
+                return;
+            }
+            let _ = self.socket.send(packet);
+        }
+    }
+
+    fn apply_open_dscp(&self, dscp: u8) -> Result<(), ClientError> {
+        #[cfg(test)]
+        if self.test_hooks.fail_open_dscp.replace(false) {
+            return Err(ClientError::SocketOption {
+                operation: "set negotiated DSCP",
+                remote: self.remote,
+                source: io::Error::other("injected negotiated DSCP failure"),
+            });
+        }
+        apply_dscp_to_socket(&self.socket, self.remote, dscp)
+    }
+
+    fn restore_open_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
+        #[cfg(test)]
+        self.test_hooks.last_restored_read_timeout.set(timeout);
+        #[cfg(test)]
+        if self.test_hooks.fail_open_timeout_restore.replace(false) {
+            return Err(io::Error::other(
+                "injected configured read timeout restoration failure",
+            ));
+        }
+        self.socket.set_read_timeout(timeout)
+    }
 }
 
 pub(crate) fn echo_sent_event(