Harden embedded transport under UDP churn

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);
-        self.add_credit(adjust.stream_id, adjust.bytes as usize);
+        let acked_bytes = self.ack_data(adjust.stream_id, adjust.received_offset);
+        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());