Harden embedded transport under UDP churn

src/bin/dosh-client.rs

@@ -36,6 +36,7 @@ use dosh::transport::{
     DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig,
     TransportEvent,
 };
+use dosh::udp::is_transient_udp_send_error;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 use std::collections::BTreeMap;
@@ -3534,52 +3535,6 @@ async fn send_terminal_udp(socket: &UdpSocket, packet: &[u8], addr: SocketAddr)
     }
 }
 
-fn is_transient_udp_send_error(err: &std::io::Error) -> bool {
-    matches!(
-        err.kind(),
-        std::io::ErrorKind::Interrupted
-            | std::io::ErrorKind::TimedOut
-            | std::io::ErrorKind::WouldBlock
-    ) || err.raw_os_error().is_some_and(is_transient_udp_os_error)
-}
-
-#[cfg(unix)]
-fn is_transient_udp_os_error(code: i32) -> bool {
-    matches!(
-        code,
-        libc::EADDRNOTAVAIL
-            | libc::ECONNREFUSED
-            | libc::ECONNRESET
-            | libc::EHOSTDOWN
-            | libc::EHOSTUNREACH
-            | libc::ENETDOWN
-            | libc::ENETRESET
-            | libc::ENETUNREACH
-            | libc::ETIMEDOUT
-    )
-}
-
-#[cfg(windows)]
-fn is_transient_udp_os_error(code: i32) -> bool {
-    matches!(
-        code,
-        10049  // WSAEADDRNOTAVAIL
-            | 10050 // WSAENETDOWN
-            | 10051 // WSAENETUNREACH
-            | 10052 // WSAENETRESET
-            | 10054 // WSAECONNRESET
-            | 10060 // WSAETIMEDOUT
-            | 10061 // WSAECONNREFUSED
-            | 10064 // WSAEHOSTDOWN
-            | 10065 // WSAEHOSTUNREACH
-    )
-}
-
-#[cfg(not(any(unix, windows)))]
-fn is_transient_udp_os_error(_code: i32) -> bool {
-    false
-}
-
 #[allow(clippy::too_many_arguments)]
 async fn try_native_auth(
     socket: &UdpSocket,
@@ -7278,23 +7233,23 @@ mod tests {
         PredictMode, Predictor, RESTART_STATUS_SCRIPT, RemoteForward, STARTUP_INPUT_HOLD,
         SshConfig, StartupGateMode, StatusAction, auth_allows, cache_key, cache_server_prefix,
         clear_cached_credentials, ensure_tui_safe_status_overlay, input_matches_escape,
-        is_local_status_target, is_resume_response_for_client, is_transient_udp_os_error,
+        is_local_status_target, is_resume_response_for_client, latest_release_download_url,
-        latest_release_download_url, load_first_native_identity_with_prompt, parse_dynamic_forward,
+        load_first_native_identity_with_prompt, parse_dynamic_forward, parse_escape_key,
-        parse_escape_key, parse_local_forward, parse_remote_forward, parse_ssh_config,
+        parse_local_forward, parse_remote_forward, parse_ssh_config, post_submit_hold_duration,
-        post_submit_hold_duration, queue_pending_user_input, raw_contains_host_table,
+        queue_pending_user_input, raw_contains_host_table, recv_response_until, refresh_live_addr,
-        recv_response_until, refresh_live_addr, release_tag_download_url,
+        release_tag_download_url, release_tag_from_effective_url, release_version_from_tag,
-        release_tag_from_effective_url, release_version_from_tag, render_status_clear,
+        render_status_clear, render_status_overlay, requested_env, resolved_startup_command,
-        render_status_overlay, requested_env, resolved_startup_command, retransmit_stream_opens,
+        retransmit_stream_opens, rewrite_forward_command, selected_predict_mode, selected_udp_host,
-        rewrite_forward_command, selected_predict_mode, selected_udp_host, server_version_mismatch,
+        server_version_mismatch, should_flush_terminal_input_after_contact,
-        should_flush_terminal_input_after_contact, should_hold_during_startup_gate,
+        should_hold_during_startup_gate, should_hold_post_submit_input, split_after_command_submit,
-        should_hold_post_submit_input, split_after_command_submit, ssh_destination_host,
+        ssh_destination_host, ssh_username, ssh_with_user, startup_command, status_ssh_target,
-        ssh_username, ssh_with_user, startup_command, status_ssh_target, strip_stale_mouse_reports,
+        strip_stale_mouse_reports, toml_bare_key_or_quoted, update_check_requested,
-        toml_bare_key_or_quoted, update_check_requested, update_version_status,
+        update_version_status, upsert_managed_block, valid_forward_host, vscode_safe_alias,
-        upsert_managed_block, valid_forward_host, vscode_safe_alias,
     };
     use dosh::config::{ClientConfig, CommandExtension, HostConfig};
     use dosh::native::EnvVar;
     use dosh::protocol::{self, Frame, PacketKind};
+    use dosh::udp::is_transient_udp_send_error;
     use ed25519_dalek::{SigningKey, VerifyingKey};
     use std::collections::{HashMap, VecDeque};
     use std::fs;
@@ -8524,18 +8479,34 @@ mod tests {
     fn transient_udp_send_errors_are_not_terminal_fatal() {
         #[cfg(unix)]
         {
-            assert!(is_transient_udp_os_error(libc::ENETUNREACH));
+            assert!(is_transient_udp_send_error(
-            assert!(is_transient_udp_os_error(libc::EHOSTUNREACH));
+                &std::io::Error::from_raw_os_error(libc::ENETUNREACH)
-            assert!(is_transient_udp_os_error(libc::EADDRNOTAVAIL));
+            ));
-            assert!(!is_transient_udp_os_error(libc::EINVAL));
+            assert!(is_transient_udp_send_error(
+                &std::io::Error::from_raw_os_error(libc::EHOSTUNREACH)
+            ));
+            assert!(is_transient_udp_send_error(
+                &std::io::Error::from_raw_os_error(libc::EADDRNOTAVAIL)
+            ));
+            assert!(!is_transient_udp_send_error(
+                &std::io::Error::from_raw_os_error(libc::EINVAL)
+            ));
         }
 
         #[cfg(windows)]
         {
-            assert!(is_transient_udp_os_error(10051)); // WSAENETUNREACH
+            assert!(is_transient_udp_send_error(
-            assert!(is_transient_udp_os_error(10065)); // WSAEHOSTUNREACH
+                &std::io::Error::from_raw_os_error(10051)
-            assert!(is_transient_udp_os_error(10049)); // WSAEADDRNOTAVAIL
+            )); // WSAENETUNREACH
-            assert!(!is_transient_udp_os_error(10022)); // WSAEINVAL
+            assert!(is_transient_udp_send_error(
+                &std::io::Error::from_raw_os_error(10065)
+            )); // WSAEHOSTUNREACH
+            assert!(is_transient_udp_send_error(
+                &std::io::Error::from_raw_os_error(10049)
+            )); // WSAEADDRNOTAVAIL
+            assert!(!is_transient_udp_send_error(
+                &std::io::Error::from_raw_os_error(10022)
+            )); // WSAEINVAL
         }
     }
 

src/bin/dosh-server.rs

@@ -28,6 +28,7 @@ use dosh::protocol::{
     TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
 };
 use dosh::pty::{PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session};
+use dosh::udp::is_transient_udp_send_error;
 use sha2::{Digest, Sha256};
 use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use std::fs;
@@ -1713,52 +1714,6 @@ async fn send_udp(socket: &UdpSocket, packet: &[u8], peer: SocketAddr) -> Result
     }
 }
 
-fn is_transient_udp_send_error(err: &std::io::Error) -> bool {
-    matches!(
-        err.kind(),
-        std::io::ErrorKind::Interrupted
-            | std::io::ErrorKind::TimedOut
-            | std::io::ErrorKind::WouldBlock
-    ) || err.raw_os_error().is_some_and(is_transient_udp_os_error)
-}
-
-#[cfg(unix)]
-fn is_transient_udp_os_error(code: i32) -> bool {
-    matches!(
-        code,
-        libc::EADDRNOTAVAIL
-            | libc::ECONNREFUSED
-            | libc::ECONNRESET
-            | libc::EHOSTDOWN
-            | libc::EHOSTUNREACH
-            | libc::ENETDOWN
-            | libc::ENETRESET
-            | libc::ENETUNREACH
-            | libc::ETIMEDOUT
-    )
-}
-
-#[cfg(windows)]
-fn is_transient_udp_os_error(code: i32) -> bool {
-    matches!(
-        code,
-        10049  // WSAEADDRNOTAVAIL
-            | 10050 // WSAENETDOWN
-            | 10051 // WSAENETUNREACH
-            | 10052 // WSAENETRESET
-            | 10054 // WSAECONNRESET
-            | 10060 // WSAETIMEDOUT
-            | 10061 // WSAECONNREFUSED
-            | 10064 // WSAEHOSTDOWN
-            | 10065 // WSAEHOSTUNREACH
-    )
-}
-
-#[cfg(not(any(unix, windows)))]
-fn is_transient_udp_os_error(_code: i32) -> bool {
-    false
-}
-
 async fn handle_ack(
     state: &Arc<Mutex<ServerState>>,
     peer: SocketAddr,

src/lib.rs

@@ -23,3 +23,4 @@ pub mod pty;
 pub mod server;
 pub mod ssh_agent;
 pub mod transport;
+pub mod udp;

src/transport.rs

@@ -19,6 +19,7 @@ use crate::protocol::{
     self, CLIENT_TO_SERVER, PacketKind, ReplayWindow, SERVER_TO_CLIENT, StreamClose, StreamData,
     StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
 };
+use crate::udp::is_transient_udp_send_error;
 use anyhow::{Result, anyhow, bail};
 use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use std::net::SocketAddr;
@@ -348,8 +349,8 @@ impl StreamMux {
         &mut self,
         adjust: StreamWindowAdjust,
     ) -> Result<Vec<OutgoingStreamPacket>> {
-        self.ack_data(adjust.stream_id, adjust.received_offset);
+        let acked_bytes = self.ack_data(adjust.stream_id, adjust.received_offset);
-        self.add_credit(adjust.stream_id, adjust.bytes as usize);
+        self.add_credit(adjust.stream_id, acked_bytes);
         self.flush_pending_data(adjust.stream_id)
     }
 
@@ -556,9 +557,9 @@ impl StreamMux {
         (chunks, consumed, *expected)
     }
 
-    fn ack_data(&mut self, stream_id: u64, received_offset: u64) {
+    fn ack_data(&mut self, stream_id: u64, received_offset: u64) -> usize {
         let Some(sent) = self.sent_data.get_mut(&stream_id) else {
-            return;
+            return 0;
         };
         let acked_offsets = sent
             .iter()
@@ -567,12 +568,16 @@ impl StreamMux {
                 (end <= received_offset).then_some(*offset)
             })
             .collect::<Vec<_>>();
+        let mut acked_bytes = 0usize;
         for offset in acked_offsets {
-            sent.remove(&offset);
+            if let Some(chunk) = sent.remove(&offset) {
+                acked_bytes = acked_bytes.saturating_add(chunk.bytes.len());
+            }
         }
         if sent.is_empty() {
             self.sent_data.remove(&stream_id);
         }
+        acked_bytes
     }
 
     fn add_credit(&mut self, stream_id: u64, bytes: usize) {
@@ -815,7 +820,12 @@ impl DoshTransport {
 
     async fn send_kind(&mut self, kind: PacketKind, body: &[u8]) -> Result<()> {
         let encoded = self.encode_kind(kind, body)?;
-        self.socket.send_to(&encoded, self.peer_addr).await?;
+        if let Err(err) = self.socket.send_to(&encoded, self.peer_addr).await {
+            if is_transient_udp_send_error(&err) {
+                return Ok(());
+            }
+            return Err(err.into());
+        }
         Ok(())
     }
 
@@ -980,6 +990,30 @@ mod tests {
         assert_eq!(data.bytes, b"!");
     }
 
+    #[test]
+    fn cumulative_window_adjust_restores_credit_even_if_delta_was_lost() {
+        let mut mux = StreamMux::new(TransportConfig {
+            initial_window: 5,
+            retransmit_after: Duration::from_secs(1),
+            ..TransportConfig::default()
+        });
+        mux.open_stream(1, "@dosh-test", 0).unwrap();
+        mux.handle_open_ok(StreamOpenOk { stream_id: 1 }).unwrap();
+        assert_eq!(mux.send_data(1, b"hello".to_vec()).unwrap().len(), 1);
+        assert_eq!(mux.send_credit[&1], 0);
+
+        let flushed = mux
+            .handle_window_adjust(StreamWindowAdjust {
+                stream_id: 1,
+                received_offset: 5,
+                bytes: 0,
+            })
+            .unwrap();
+        assert!(flushed.is_empty());
+        assert_eq!(mux.send_credit[&1], 5);
+        assert_eq!(mux.send_data(1, b"!".to_vec()).unwrap().len(), 1);
+    }
+
     #[test]
     fn handle_packet_decodes_and_dispatches_stream_packets() {
         let mut mux = StreamMux::new(TransportConfig::default());

src/udp.rs

@@ -0,0 +1,98 @@
+//! UDP error classification shared by Dosh clients, servers, and embedders.
+
+pub fn is_transient_udp_send_error(err: &std::io::Error) -> bool {
+    matches!(
+        err.kind(),
+        std::io::ErrorKind::Interrupted
+            | std::io::ErrorKind::TimedOut
+            | std::io::ErrorKind::WouldBlock
+    ) || err.raw_os_error().is_some_and(is_transient_udp_os_error)
+}
+
+#[cfg(unix)]
+fn is_transient_udp_os_error(code: i32) -> bool {
+    matches!(
+        code,
+        libc::EADDRNOTAVAIL
+            | libc::ECONNREFUSED
+            | libc::ECONNRESET
+            | libc::EHOSTDOWN
+            | libc::EHOSTUNREACH
+            | libc::ENETDOWN
+            | libc::ENETRESET
+            | libc::ENETUNREACH
+            | libc::ETIMEDOUT
+    )
+}
+
+#[cfg(windows)]
+fn is_transient_udp_os_error(code: i32) -> bool {
+    matches!(
+        code,
+        10049  // WSAEADDRNOTAVAIL
+            | 10050 // WSAENETDOWN
+            | 10051 // WSAENETUNREACH
+            | 10052 // WSAENETRESET
+            | 10054 // WSAECONNRESET
+            | 10060 // WSAETIMEDOUT
+            | 10061 // WSAECONNREFUSED
+            | 10064 // WSAEHOSTDOWN
+            | 10065 // WSAEHOSTUNREACH
+    )
+}
+
+#[cfg(not(any(unix, windows)))]
+fn is_transient_udp_os_error(_code: i32) -> bool {
+    false
+}
+
+#[cfg(test)]
+mod tests {
+    use super::is_transient_udp_send_error;
+
+    #[test]
+    fn classifies_portable_transient_send_errors() {
+        for kind in [
+            std::io::ErrorKind::Interrupted,
+            std::io::ErrorKind::TimedOut,
+            std::io::ErrorKind::WouldBlock,
+        ] {
+            assert!(is_transient_udp_send_error(&std::io::Error::from(kind)));
+        }
+        assert!(!is_transient_udp_send_error(&std::io::Error::from(
+            std::io::ErrorKind::PermissionDenied,
+        )));
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn classifies_unix_network_churn_as_transient() {
+        for code in [
+            libc::EADDRNOTAVAIL,
+            libc::ECONNREFUSED,
+            libc::ECONNRESET,
+            libc::EHOSTDOWN,
+            libc::EHOSTUNREACH,
+            libc::ENETDOWN,
+            libc::ENETRESET,
+            libc::ENETUNREACH,
+            libc::ETIMEDOUT,
+        ] {
+            assert!(is_transient_udp_send_error(
+                &std::io::Error::from_raw_os_error(code)
+            ));
+        }
+    }
+
+    #[cfg(windows)]
+    #[test]
+    fn classifies_windows_network_churn_as_transient() {
+        for code in [
+            10049, 10050, 10051, 10052, 10054, 10060, 10061, 10064, 10065,
+        ] {
+            assert!(is_transient_udp_send_error(
+                &std::io::Error::from_raw_os_error(code)
+            ));
+        }
+    }
+}