Drive SDK maintenance while receiving

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<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 (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<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(_) => 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<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();