Drive SDK maintenance while receiving

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;

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,33 +169,52 @@ 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]) {
+            if let Some(event) = self.handle_datagram(peer, &buf[..n]).await? {
+                return Ok(event);
+            }
+        }
+    }
+
+    async fn handle_datagram(
+        &mut self,
+        peer: SocketAddr,
+        datagram: &[u8],
+    ) -> Result<Option<DoshServerEvent>> {
+        let packet = match protocol::decode(datagram) {
             Ok(packet) => packet,
-                Err(_) => continue,
+            Err(_) => return Ok(None),
         };
         if packet.header.kind == PacketKind::NativeClientHello {
             self.handle_client_hello(peer, packet.body).await?;
-                return Ok(DoshServerEvent::Ignored);
+            return Ok(Some(DoshServerEvent::Ignored));
         }
         if packet.header.kind == PacketKind::NativeUserAuth {
-                return self.handle_user_auth(peer, &packet).await;
+            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(&buf[..n], peer).await?;
-                return Ok(DoshServerEvent::Session {
+            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?;
-            return Ok(DoshServerEvent::Ignored);
-        }
+        Ok(Some(DoshServerEvent::Ignored))
     }
 
     pub async fn accept_stream(&mut self, conn_id: [u8; 16], stream_id: u64) -> Result<()> {
@@ -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();

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,36 +1180,52 @@ impl DoshTransport {
     pub async fn recv(&mut self) -> Result<SessionEvent> {
         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 packet = match protocol::decode(&buf[..n]) {
-                Ok(packet) => packet,
-                Err(_) => continue,
-            };
-            if packet.header.conn_id != self.conn_id {
+            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 plain = match protocol::decrypt_body(
-                &packet,
-                &self.session_key,
-                self.role.recv_direction(),
-            ) {
+            };
+            if let Some(event) = self.accept_received_datagram(&buf[..n], peer).await? {
+                return Ok(event);
+            }
+        }
+    }
+
+    async fn accept_received_datagram(
+        &mut self,
+        datagram: &[u8],
+        peer: SocketAddr,
+    ) -> Result<Option<SessionEvent>> {
+        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(_) => continue,
+                Err(_) => return Ok(None),
             };
         if !self.replay.accept(packet.header.seq) {
-                continue;
+            return Ok(None);
         }
         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;
-        }
+        self.handle_decrypted_packet(packet.header.kind, &plain)
+            .await
+            .map(Some)
     }
 
     pub async fn handle_datagram(
@@ -1214,26 +1233,10 @@ impl DoshTransport {
         datagram: &[u8],
         peer: SocketAddr,
     ) -> Result<SessionEvent> {
-        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();