From 0125966d2a782996dae7ff7d5e038f6c39a916b8 Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Mon, 13 Jul 2026 00:58:54 -0400 Subject: [PATCH] Drive SDK maintenance while receiving --- src/lib.rs | 4 +- src/server.rs | 150 +++++++++++++++++++++++++++++++++-------- src/transport.rs | 170 ++++++++++++++++++++++++++++++++++------------- 3 files changed, 250 insertions(+), 74 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 14360cd..18cf28a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,7 +5,9 @@ //! handle native authentication, host key checks, roaming, keepalives, reliable //! streams, and service routing. Use [`transport::DoshTransport`] only when an //! application already owns session establishment and wants direct access to the -//! encrypted stream transport. +//! encrypted stream transport. The SDK receive loops drive retransmission and +//! keepalives while waiting, so simple applications can call `recv().await` +//! without building their own maintenance loop. pub mod auth; pub mod build_info; diff --git a/src/server.rs b/src/server.rs index 3ae3fe7..e73f7b3 100644 --- a/src/server.rs +++ b/src/server.rs @@ -10,8 +10,8 @@ use crate::protocol::{ NativeServerHelloBody, NativeUserAuthBody, PacketKind, SERVER_TO_CLIENT, }; use crate::transport::{ - DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig, - service_name_from_target, + ADAPTIVE_RETRANSMIT_MIN, DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, + TransportConfig, service_name_from_target, }; use crate::udp::{is_transient_udp_error, is_transient_udp_send_error}; use anyhow::{Context, Result, anyhow, bail}; @@ -169,35 +169,54 @@ impl DoshServer { let mut buf = vec![0u8; 65535]; loop { self.expire_pending(); - let (n, peer) = match self.socket.recv_from(&mut buf).await { - Ok(value) => value, - Err(err) if is_transient_udp_error(&err) => continue, - Err(err) => return Err(err.into()), + let (n, peer) = match tokio::time::timeout( + ADAPTIVE_RETRANSMIT_MIN, + self.socket.recv_from(&mut buf), + ) + .await + { + Ok(Ok(value)) => value, + Ok(Err(err)) if is_transient_udp_error(&err) => continue, + Ok(Err(err)) => return Err(err.into()), + Err(_) => { + self.maintenance().await?; + continue; + } }; - let packet = match protocol::decode(&buf[..n]) { - Ok(packet) => packet, - Err(_) => continue, - }; - if packet.header.kind == PacketKind::NativeClientHello { - self.handle_client_hello(peer, packet.body).await?; - return Ok(DoshServerEvent::Ignored); + if let Some(event) = self.handle_datagram(peer, &buf[..n]).await? { + return Ok(event); } - if packet.header.kind == PacketKind::NativeUserAuth { - return self.handle_user_auth(peer, &packet).await; - } - if let Some(transport) = self.transports.get_mut(&packet.header.conn_id) { - let event = transport.handle_datagram(&buf[..n], peer).await?; - return Ok(DoshServerEvent::Session { - conn_id: packet.header.conn_id, - event, - }); - } - self.send_reject(peer, packet.header.conn_id, "unknown Dosh connection") - .await?; - return Ok(DoshServerEvent::Ignored); } } + async fn handle_datagram( + &mut self, + peer: SocketAddr, + datagram: &[u8], + ) -> Result> { + let packet = match protocol::decode(datagram) { + Ok(packet) => packet, + Err(_) => return Ok(None), + }; + if packet.header.kind == PacketKind::NativeClientHello { + self.handle_client_hello(peer, packet.body).await?; + return Ok(Some(DoshServerEvent::Ignored)); + } + if packet.header.kind == PacketKind::NativeUserAuth { + return self.handle_user_auth(peer, &packet).await.map(Some); + } + if let Some(transport) = self.transports.get_mut(&packet.header.conn_id) { + let event = transport.handle_datagram(datagram, peer).await?; + return Ok(Some(DoshServerEvent::Session { + conn_id: packet.header.conn_id, + event, + })); + } + self.send_reject(peer, packet.header.conn_id, "unknown Dosh connection") + .await?; + Ok(Some(DoshServerEvent::Ignored)) + } + pub async fn accept_stream(&mut self, conn_id: [u8; 16], stream_id: u64) -> Result<()> { self.transport_mut(&conn_id) .ok_or_else(|| anyhow!("unknown Dosh connection"))? @@ -620,6 +639,85 @@ mod tests { } } + #[tokio::test] + async fn sdk_server_recv_drives_transport_maintenance_while_idle() { + let dir = tempfile::tempdir().unwrap(); + let host_key = dir.path().join("host_key"); + let authorized_keys = dir.path().join("authorized_keys"); + let signing = SigningKey::from_bytes(&[94u8; 32]); + let keypair = ssh_key::private::Ed25519Keypair::from(&signing); + let private = + ssh_key::PrivateKey::new(ssh_key::private::KeypairData::from(keypair), "").unwrap(); + std::fs::write( + &authorized_keys, + format!("{}\n", private.public_key().to_openssh().unwrap()), + ) + .unwrap(); + + let server_config = ServerConfig { + host_key: host_key.to_string_lossy().to_string(), + authorized_keys: vec![authorized_keys.to_string_lossy().to_string()], + ..ServerConfig::default() + }; + let server_config = DoshServerConfig::new(server_config) + .bind_addr("127.0.0.1:0".parse().unwrap()) + .service("echo") + .unwrap() + .require_current_user(false) + .transport(TransportConfig { + retransmit_after: Duration::from_millis(20), + keepalive_after: Duration::from_secs(60), + ..TransportConfig::default() + }); + let mut server = DoshServer::bind(server_config).await.unwrap(); + let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let client_addr = client_socket.local_addr().unwrap(); + let conn_id = [95u8; 16]; + let session_key = [96u8; 32]; + let transport = DoshTransport::new( + Arc::clone(&server.socket), + SessionTransportConfig { + role: SessionRole::Server, + conn_id, + session_key, + peer_addr: client_addr, + initial_send_seq: 1, + initial_ack: 0, + stream: server.config.transport.clone(), + }, + ); + server.transports.insert(conn_id, transport); + server + .transport_mut(&conn_id) + .unwrap() + .open_service("echo") + .await + .unwrap(); + + let mut buf = [0u8; 65535]; + tokio::time::timeout( + Duration::from_millis(200), + client_socket.recv_from(&mut buf), + ) + .await + .unwrap() + .unwrap(); + + let recv_task = tokio::spawn(async move { server.recv().await }); + let (n, _) = + tokio::time::timeout(Duration::from_secs(1), client_socket.recv_from(&mut buf)) + .await + .unwrap() + .unwrap(); + recv_task.abort(); + + let packet = protocol::decode(&buf[..n]).unwrap(); + assert_eq!(packet.header.kind, PacketKind::StreamOpen); + let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); + let open: protocol::StreamOpen = protocol::from_body(&plain).unwrap(); + assert_eq!(open.target_host, "@dosh-echo"); + } + #[tokio::test] async fn bad_native_auth_does_not_consume_pending_challenge() { let dir = tempfile::tempdir().unwrap(); diff --git a/src/transport.rs b/src/transport.rs index 969cd46..afe5f18 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -14,7 +14,10 @@ //! //! That split is intentional: Dosh roaming/reconnect stays at the authenticated //! session layer, while this module provides the reusable fast persistent byte -//! streams that higher-level app protocols can embed. +//! streams that higher-level app protocols can embed. [`DoshTransport::recv`] +//! calls [`DoshTransport::maintenance`] while it is idle, so embedders using the +//! session wrapper get retransmission and keepalives in a simple `recv().await` +//! loop. use crate::protocol::{ self, CLIENT_TO_SERVER, PacketKind, ReplayWindow, SERVER_TO_CLIENT, StreamClose, StreamData, @@ -1177,63 +1180,63 @@ impl DoshTransport { pub async fn recv(&mut self) -> Result { let mut buf = vec![0u8; 65535]; loop { - let (n, peer) = match self.socket.recv_from(&mut buf).await { - Ok(value) => value, - Err(err) if transport_recv_error_is_retryable(&err) => continue, - Err(err) => return Err(err.into()), + let (n, peer) = match tokio::time::timeout( + ADAPTIVE_RETRANSMIT_MIN, + self.socket.recv_from(&mut buf), + ) + .await + { + Ok(Ok(value)) => value, + Ok(Err(err)) if transport_recv_error_is_retryable(&err) => continue, + Ok(Err(err)) => return Err(err.into()), + Err(_) => { + self.maintenance().await?; + continue; + } }; - let packet = match protocol::decode(&buf[..n]) { - Ok(packet) => packet, - Err(_) => continue, - }; - if packet.header.conn_id != self.conn_id { - continue; + if let Some(event) = self.accept_received_datagram(&buf[..n], peer).await? { + return Ok(event); } - let plain = match protocol::decrypt_body( - &packet, - &self.session_key, - self.role.recv_direction(), - ) { - Ok(plain) => plain, - Err(_) => continue, - }; - if !self.replay.accept(packet.header.seq) { - continue; - } - self.peer_addr = peer; - self.ack_seq = self.ack_seq.max(packet.header.seq); - self.last_contact = Instant::now(); - return self - .handle_decrypted_packet(packet.header.kind, &plain) - .await; } } + async fn accept_received_datagram( + &mut self, + datagram: &[u8], + peer: SocketAddr, + ) -> Result> { + let packet = match protocol::decode(datagram) { + Ok(packet) => packet, + Err(_) => return Ok(None), + }; + if packet.header.conn_id != self.conn_id { + return Ok(None); + } + let plain = + match protocol::decrypt_body(&packet, &self.session_key, self.role.recv_direction()) { + Ok(plain) => plain, + Err(_) => return Ok(None), + }; + if !self.replay.accept(packet.header.seq) { + return Ok(None); + } + self.peer_addr = peer; + self.ack_seq = self.ack_seq.max(packet.header.seq); + self.last_contact = Instant::now(); + self.handle_decrypted_packet(packet.header.kind, &plain) + .await + .map(Some) + } + pub async fn handle_datagram( &mut self, datagram: &[u8], peer: SocketAddr, ) -> Result { - let packet = match protocol::decode(datagram) { - Ok(packet) => packet, - Err(_) => return Ok(SessionEvent::Ignored), - }; - if packet.header.conn_id != self.conn_id { - return Ok(SessionEvent::Ignored); + match self.accept_received_datagram(datagram, peer).await? { + Some(event) => Ok(event), + None => Ok(SessionEvent::Ignored), } - let plain = - match protocol::decrypt_body(&packet, &self.session_key, self.role.recv_direction()) { - Ok(plain) => plain, - Err(_) => return Ok(SessionEvent::Ignored), - }; - if !self.replay.accept(packet.header.seq) { - return Ok(SessionEvent::Ignored); - } - self.peer_addr = peer; - self.ack_seq = self.ack_seq.max(packet.header.seq); - self.last_contact = Instant::now(); - self.handle_decrypted_packet(packet.header.kind, &plain) - .await } async fn handle_decrypted_packet( @@ -2346,6 +2349,79 @@ mod tests { } } + #[tokio::test] + async fn session_transport_recv_drives_retransmit_while_waiting() { + let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let client_addr = client_socket.local_addr().unwrap(); + let server_addr = server_socket.local_addr().unwrap(); + let key = [48u8; 32]; + let conn_id = [12u8; 16]; + let stream_config = TransportConfig { + retransmit_after: Duration::from_millis(20), + keepalive_after: Duration::from_secs(60), + ..TransportConfig::default() + }; + let mut client = DoshTransport::new_owned( + client_socket, + SessionTransportConfig { + role: SessionRole::Client, + conn_id, + session_key: key, + peer_addr: server_addr, + initial_send_seq: 1, + initial_ack: 0, + stream: stream_config.clone(), + }, + ); + + let stream_id = client.open_service("app").await.unwrap(); + + let mut dropped = [0u8; 65535]; + tokio::time::timeout( + Duration::from_millis(200), + server_socket.recv_from(&mut dropped), + ) + .await + .unwrap() + .unwrap(); + + let mut server = DoshTransport::new_owned( + server_socket, + SessionTransportConfig { + role: SessionRole::Server, + conn_id, + session_key: key, + peer_addr: client_addr, + initial_send_seq: 1, + initial_ack: 0, + stream: stream_config, + }, + ); + + let client_wait = tokio::spawn(async move { client.recv().await }); + match tokio::time::timeout(Duration::from_secs(1), server.recv()) + .await + .unwrap() + .unwrap() + { + SessionEvent::Stream(TransportEvent::Open(open)) => { + assert_eq!(open.stream_id, stream_id); + server.accept_stream(open.stream_id).await.unwrap(); + } + other => panic!("unexpected server event after lost open: {other:?}"), + } + + assert!(matches!( + tokio::time::timeout(Duration::from_secs(1), client_wait) + .await + .unwrap() + .unwrap() + .unwrap(), + SessionEvent::Stream(TransportEvent::OpenOk { stream_id: ok, .. }) if ok == stream_id + )); + } + #[tokio::test] async fn session_transport_reject_stream_retires_server_state() { let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();