diff --git a/src/transport.rs b/src/transport.rs index 0d96ab0..969cd46 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -9,7 +9,8 @@ //! 1. authenticate and maintain a Dosh session using [`crate::protocol`], //! 2. encrypt each [`OutgoingStreamPacket::body`] as the given packet kind, //! 3. pass received stream packet bodies back into [`StreamMux`], -//! 4. call [`StreamMux::tick`] periodically to retransmit lost stream data. +//! 4. call [`StreamMux::tick`] periodically to retransmit lost stream data and +//! idempotent stream control packets. //! //! That split is intentional: Dosh roaming/reconnect stays at the authenticated //! session layer, while this module provides the reusable fast persistent byte @@ -33,6 +34,7 @@ pub const ADAPTIVE_RETRANSMIT_PAD: Duration = Duration::from_millis(10); pub const ADAPTIVE_RETRANSMIT_MIN: Duration = Duration::from_millis(10); pub const DEFAULT_KEEPALIVE_AFTER: Duration = Duration::from_secs(2); pub const DEFAULT_RETIRED_STREAM_TOMBSTONES: usize = 16 * 1024; +pub const STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS: u32 = 8; pub const SERVICE_TARGET_PREFIX: &str = "@dosh-"; pub const MAX_STREAM_DATA_BYTES: usize = 60 * 1024; @@ -171,6 +173,20 @@ struct PendingStreamOpen { attempts: u32, } +#[derive(Debug, Clone)] +struct PendingStreamControl { + last_sent: Instant, + attempts: u32, +} + +#[derive(Debug, Clone)] +struct PendingWindowAdjust { + received_offset: u64, + bytes: usize, + last_sent: Instant, + attempts: u32, +} + #[derive(Debug, Clone)] pub struct StreamMux { config: TransportConfig, @@ -182,6 +198,9 @@ pub struct StreamMux { next_send_offset: HashMap, next_recv_offset: HashMap, recv_pending: HashMap>>, + pending_eofs: HashMap, + pending_closes: HashMap, + pending_window_adjusts: HashMap, retired_streams: HashSet, retired_stream_order: VecDeque, srtt: Option, @@ -240,6 +259,9 @@ impl StreamMux { next_send_offset: HashMap::new(), next_recv_offset: HashMap::new(), recv_pending: HashMap::new(), + pending_eofs: HashMap::new(), + pending_closes: HashMap::new(), + pending_window_adjusts: HashMap::new(), retired_streams: HashSet::new(), retired_stream_order: VecDeque::new(), srtt: None, @@ -411,6 +433,17 @@ impl StreamMux { bytes: consumed.min(u32::MAX as usize) as u32, }, )?; + if consumed > 0 { + self.pending_window_adjusts.insert( + stream_id, + PendingWindowAdjust { + received_offset, + bytes: consumed, + last_sent: Instant::now(), + attempts: 1, + }, + ); + } Ok(Some(IncomingStreamData { stream_id, chunks, @@ -437,13 +470,27 @@ impl StreamMux { bail!("stream {stream_id} is not open"); } self.retire_stream(stream_id); + self.pending_closes.insert( + stream_id, + PendingStreamControl { + last_sent: Instant::now(), + attempts: 1, + }, + ); encode_packet(PacketKind::StreamClose, &StreamClose { stream_id }) } - pub fn eof_stream(&self, stream_id: u64) -> Result { + pub fn eof_stream(&mut self, stream_id: u64) -> Result { if !self.has_stream_state(stream_id) { bail!("stream {stream_id} is not open"); } + self.pending_eofs.insert( + stream_id, + PendingStreamControl { + last_sent: Instant::now(), + attempts: 1, + }, + ); encode_packet(PacketKind::StreamEof, &StreamEof { stream_id }) } @@ -452,7 +499,9 @@ impl StreamMux { } pub fn handle_close(&mut self, close: StreamClose) -> bool { - if !self.has_stream_state(close.stream_id) { + let had_state = self.has_stream_state(close.stream_id); + self.pending_closes.remove(&close.stream_id); + if !had_state { return false; } self.retire_stream(close.stream_id); @@ -579,6 +628,79 @@ impl StreamMux { }, )?); } + + let mut retransmit_eofs = Vec::new(); + for (stream_id, pending) in self.pending_eofs.iter_mut() { + if !self.opened_streams.contains(stream_id) { + continue; + } + if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS { + continue; + } + if now.duration_since(pending.last_sent) < retransmit_after { + continue; + } + pending.last_sent = now; + pending.attempts = pending.attempts.saturating_add(1); + retransmit_eofs.push(*stream_id); + } + let opened_streams = self.opened_streams.clone(); + self.pending_eofs.retain(|stream_id, pending| { + opened_streams.contains(stream_id) + && pending.attempts < STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS + }); + for stream_id in retransmit_eofs { + out.push(encode_packet( + PacketKind::StreamEof, + &StreamEof { stream_id }, + )?); + } + + let mut retransmit_closes = Vec::new(); + for (stream_id, pending) in self.pending_closes.iter_mut() { + if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS { + continue; + } + if now.duration_since(pending.last_sent) < retransmit_after { + continue; + } + pending.last_sent = now; + pending.attempts = pending.attempts.saturating_add(1); + retransmit_closes.push(*stream_id); + } + self.pending_closes + .retain(|_, pending| pending.attempts < STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS); + for stream_id in retransmit_closes { + out.push(encode_packet( + PacketKind::StreamClose, + &StreamClose { stream_id }, + )?); + } + + let mut retransmit_adjusts = Vec::new(); + for (stream_id, pending) in self.pending_window_adjusts.iter_mut() { + if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS { + continue; + } + if now.duration_since(pending.last_sent) < retransmit_after { + continue; + } + pending.last_sent = now; + pending.attempts = pending.attempts.saturating_add(1); + retransmit_adjusts.push((*stream_id, pending.received_offset, pending.bytes)); + } + self.pending_window_adjusts + .retain(|_, pending| pending.attempts < STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS); + for (stream_id, received_offset, bytes) in retransmit_adjusts { + out.push(encode_packet( + PacketKind::StreamWindowAdjust, + &StreamWindowAdjust { + stream_id, + received_offset, + bytes: bytes.min(u32::MAX as usize) as u32, + }, + )?); + } Ok(out) } @@ -738,6 +860,8 @@ impl StreamMux { || self.next_send_offset.contains_key(&stream_id) || self.next_recv_offset.contains_key(&stream_id) || self.recv_pending.contains_key(&stream_id) + || self.pending_eofs.contains_key(&stream_id) + || self.pending_window_adjusts.contains_key(&stream_id) } fn retire_stream(&mut self, stream_id: u64) { @@ -764,6 +888,8 @@ impl StreamMux { self.next_send_offset.remove(&stream_id); self.next_recv_offset.remove(&stream_id); self.recv_pending.remove(&stream_id); + self.pending_eofs.remove(&stream_id); + self.pending_window_adjusts.remove(&stream_id); } } @@ -1691,6 +1817,160 @@ mod tests { assert_eq!(mux.send_data(1, b"!".to_vec()).unwrap().len(), 1); } + #[test] + fn tick_retransmits_pending_stream_eof() { + let mut mux = StreamMux::new(TransportConfig { + retransmit_after: Duration::from_millis(200), + ..TransportConfig::default() + }); + mux.accept_open(StreamOpen { + stream_id: 1, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .unwrap(); + + let eof = mux.eof_stream(1).unwrap(); + assert_eq!(eof.kind, PacketKind::StreamEof); + mux.pending_eofs.get_mut(&1).unwrap().last_sent = + Instant::now() - Duration::from_millis(250); + + let retransmits = mux.tick().unwrap(); + assert_eq!(retransmits.len(), 1); + assert_eq!(retransmits[0].kind, PacketKind::StreamEof); + let eof: StreamEof = decode(&retransmits[0]); + assert_eq!(eof.stream_id, 1); + assert_eq!(mux.pending_eofs[&1].attempts, 2); + } + + #[test] + fn tick_retransmits_pending_stream_close() { + let mut mux = StreamMux::new(TransportConfig { + retransmit_after: Duration::from_millis(200), + ..TransportConfig::default() + }); + mux.accept_open(StreamOpen { + stream_id: 1, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .unwrap(); + + let close = mux.close_stream(1).unwrap(); + assert_eq!(close.kind, PacketKind::StreamClose); + mux.pending_closes.get_mut(&1).unwrap().last_sent = + Instant::now() - Duration::from_millis(250); + + let retransmits = mux.tick().unwrap(); + assert_eq!(retransmits.len(), 1); + assert_eq!(retransmits[0].kind, PacketKind::StreamClose); + let close: StreamClose = decode(&retransmits[0]); + assert_eq!(close.stream_id, 1); + assert_eq!(mux.pending_closes[&1].attempts, 2); + } + + #[test] + fn tick_retransmits_pending_window_adjust() { + let mut mux = StreamMux::new(TransportConfig { + retransmit_after: Duration::from_millis(200), + ..TransportConfig::default() + }); + mux.accept_open(StreamOpen { + stream_id: 1, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .unwrap(); + + let incoming = mux + .handle_data(StreamData { + stream_id: 1, + offset: 0, + bytes: b"hello".to_vec(), + }) + .unwrap() + .unwrap(); + assert_eq!(incoming.consumed, 5); + assert_eq!(incoming.received_offset, 5); + assert_eq!(incoming.window_adjust.kind, PacketKind::StreamWindowAdjust); + mux.pending_window_adjusts.get_mut(&1).unwrap().last_sent = + Instant::now() - Duration::from_millis(250); + + let retransmits = mux.tick().unwrap(); + assert_eq!(retransmits.len(), 1); + assert_eq!(retransmits[0].kind, PacketKind::StreamWindowAdjust); + let adjust: StreamWindowAdjust = decode(&retransmits[0]); + assert_eq!(adjust.stream_id, 1); + assert_eq!(adjust.received_offset, 5); + assert_eq!(adjust.bytes, 5); + assert_eq!(mux.pending_window_adjusts[&1].attempts, 2); + } + + #[test] + fn stream_control_retransmits_expire_after_cap() { + let mut mux = StreamMux::new(TransportConfig { + retransmit_after: Duration::from_millis(200), + ..TransportConfig::default() + }); + mux.accept_open(StreamOpen { + stream_id: 1, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .unwrap(); + mux.close_stream(1).unwrap(); + let pending = mux.pending_closes.get_mut(&1).unwrap(); + pending.last_sent = Instant::now() - Duration::from_millis(250); + pending.attempts = STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS; + + assert!(mux.tick().unwrap().is_empty()); + assert!(!mux.pending_closes.contains_key(&1)); + } + + #[test] + fn cleanup_stream_removes_reliable_controls() { + let mut mux = StreamMux::new(TransportConfig::default()); + mux.accept_open(StreamOpen { + stream_id: 1, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .unwrap(); + mux.eof_stream(1).unwrap(); + mux.handle_data(StreamData { + stream_id: 1, + offset: 0, + bytes: b"hello".to_vec(), + }) + .unwrap() + .unwrap(); + + mux.retire_stream(1); + + assert!(!mux.pending_eofs.contains_key(&1)); + assert!(!mux.pending_window_adjusts.contains_key(&1)); + assert!(!mux.is_open(1)); + } + + #[test] + fn peer_close_after_local_close_stops_close_retransmits() { + let mut mux = StreamMux::new(TransportConfig { + retransmit_after: Duration::from_millis(200), + ..TransportConfig::default() + }); + mux.accept_open(StreamOpen { + stream_id: 1, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .unwrap(); + mux.close_stream(1).unwrap(); + + assert!(!mux.handle_close(StreamClose { stream_id: 1 })); + assert!(!mux.pending_closes.contains_key(&1)); + assert!(mux.tick().unwrap().is_empty()); + } + #[test] fn handle_packet_decodes_and_dispatches_stream_packets() { let mut mux = StreamMux::new(TransportConfig::default());