Add reliable stream retransmission

tests/hostile_network.rs

@@ -20,7 +20,7 @@
 //! tests are reproducible and fast.
 
 use std::fs;
-use std::io::Read;
+use std::io::{Read, Write};
 use std::net::{SocketAddr, TcpListener, UdpSocket};
 use std::process::{Child, Command, Stdio};
 use std::sync::Arc;
@@ -38,7 +38,7 @@ use dosh::native::{self, ForwardingKind, ForwardingRequest, NativeAuthOk, Native
 use dosh::protocol::{
     self, AttachOk, AttachReject, CLIENT_TO_SERVER, Frame, Header, Input, NativeAuthOkBody,
     NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind, ResumeRequest,
-    SERVER_TO_CLIENT, StreamData, StreamOpen, StreamOpenOk,
+    SERVER_TO_CLIENT, StreamData, StreamOpen, StreamOpenOk, StreamWindowAdjust,
 };
 use ed25519_dalek::{SigningKey, VerifyingKey};
 use rand::rngs::StdRng;
@@ -627,6 +627,25 @@ fn wait_for_stream_open_ok(
     None
 }
 
+fn wait_for_stream_data(
+    socket: &UdpSocket,
+    key: &[u8; 32],
+    stream_id: u64,
+    millis: u64,
+) -> Option<(Header, StreamData)> {
+    let deadline = Instant::now() + Duration::from_millis(millis);
+    while Instant::now() < deadline {
+        if let Some((header, data)) =
+            recv_stream_packet::<StreamData>(socket, key, PacketKind::StreamData)
+        {
+            if data.stream_id == stream_id {
+                return Some((header, data));
+            }
+        }
+    }
+    None
+}
+
 fn start_tcp_collector() -> (u16, Receiver<Vec<u8>>) {
     let listener = TcpListener::bind("127.0.0.1:0").unwrap();
     let port = listener.local_addr().unwrap().port();
@@ -656,6 +675,34 @@ fn start_tcp_collector() -> (u16, Receiver<Vec<u8>>) {
     (port, rx)
 }
 
+fn start_tcp_echo() -> u16 {
+    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+    let port = listener.local_addr().unwrap().port();
+    thread::spawn(move || {
+        let Ok((mut stream, _)) = listener.accept() else {
+            return;
+        };
+        stream
+            .set_read_timeout(Some(Duration::from_millis(100)))
+            .unwrap();
+        let deadline = Instant::now() + Duration::from_secs(5);
+        let mut buf = [0u8; 1024];
+        while Instant::now() < deadline {
+            match stream.read(&mut buf) {
+                Ok(0) => return,
+                Ok(n) => {
+                    let _ = stream.write_all(&buf[..n]);
+                }
+                Err(ref err)
+                    if err.kind() == std::io::ErrorKind::WouldBlock
+                        || err.kind() == std::io::ErrorKind::TimedOut => {}
+                Err(_) => return,
+            }
+        }
+    });
+    port
+}
+
 fn collect_tcp(rx: &Receiver<Vec<u8>>, millis: u64) -> Vec<u8> {
     let deadline = Instant::now() + Duration::from_millis(millis);
     let mut out = Vec::new();
@@ -895,6 +942,7 @@ fn forwarded_stream_data_survives_reorder_and_rejects_replay() {
         open_ok.seq,
         protocol::to_body(&StreamData {
             stream_id,
+            offset: 0,
             bytes: b"A".to_vec(),
         })
         .unwrap(),
@@ -909,6 +957,7 @@ fn forwarded_stream_data_survives_reorder_and_rejects_replay() {
         open_ok.seq,
         protocol::to_body(&StreamData {
             stream_id,
+            offset: 1,
             bytes: b"B".to_vec(),
         })
         .unwrap(),
@@ -916,8 +965,8 @@ fn forwarded_stream_data_survives_reorder_and_rejects_replay() {
 
     let mut seen = collect_tcp(&tcp_rx, 1500);
     assert!(
-        seen.contains(&b'A') && seen.contains(&b'B'),
-        "reordered stream bytes did not both reach TCP target: {:?}",
+        seen.starts_with(b"AB"),
+        "reordered stream bytes were not delivered once in-order to TCP target: {:?}",
         String::from_utf8_lossy(&seen)
     );
 
@@ -933,6 +982,7 @@ fn forwarded_stream_data_survives_reorder_and_rejects_replay() {
         CLIENT_TO_SERVER,
         &protocol::to_body(&StreamData {
             stream_id,
+            offset: 2,
             bytes: b"X".to_vec(),
         })
         .unwrap(),
@@ -959,6 +1009,94 @@ fn forwarded_stream_data_survives_reorder_and_rejects_replay() {
     );
 }
 
+#[test]
+fn forwarded_stream_data_recovers_after_server_to_client_loss() {
+    let dir = tempfile::tempdir().unwrap();
+    let port = free_udp_port();
+    let config = write_server_config(&dir, port);
+    let signing_key = SigningKey::from_bytes(&[43u8; 32]);
+    authorize_ed25519_key(&dir, &signing_key);
+    let target_port = start_tcp_echo();
+    let mut server = start_server(&dir, &config);
+    let relay = Relay::spawn(port, 0xA11CEu64);
+
+    let attached = native_attach_through_relay(
+        &relay,
+        &signing_key,
+        vec![ForwardingRequest {
+            kind: ForwardingKind::Local,
+            bind_host: Some("127.0.0.1".to_string()),
+            listen_port: 0,
+            target_host: Some("127.0.0.1".to_string()),
+            target_port: Some(target_port),
+        }],
+    );
+    let stream_id = 55;
+    let open = StreamOpen {
+        stream_id,
+        target_host: "127.0.0.1".to_string(),
+        target_port,
+    };
+    send_stream_packet(
+        &attached.socket,
+        &relay,
+        &attached.ok,
+        PacketKind::StreamOpen,
+        3,
+        attached.ok.initial_seq,
+        protocol::to_body(&open).unwrap(),
+    );
+    let open_ok =
+        wait_for_stream_open_ok(&attached.socket, &attached.ok.session_key, stream_id, 3000)
+            .expect("stream did not open");
+
+    // Drop the first server->client echo and at least one retransmit tick. Once
+    // the link clears, the unacked stream chunk must be re-encrypted with a new
+    // packet sequence and delivered at the same stream offset.
+    relay.set_drop_s2c(100);
+    send_stream_packet(
+        &attached.socket,
+        &relay,
+        &attached.ok,
+        PacketKind::StreamData,
+        4,
+        open_ok.seq,
+        protocol::to_body(&StreamData {
+            stream_id,
+            offset: 0,
+            bytes: b"PING".to_vec(),
+        })
+        .unwrap(),
+    );
+    thread::sleep(Duration::from_millis(300));
+    relay.set_drop_s2c(0);
+
+    let (_header, data) =
+        wait_for_stream_data(&attached.socket, &attached.ok.session_key, stream_id, 3000)
+            .expect("lost server->client stream data was not retransmitted");
+    assert_eq!(data.offset, 0);
+    assert_eq!(data.bytes, b"PING");
+
+    send_stream_packet(
+        &attached.socket,
+        &relay,
+        &attached.ok,
+        PacketKind::StreamWindowAdjust,
+        5,
+        open_ok.seq,
+        protocol::to_body(&StreamWindowAdjust {
+            stream_id,
+            received_offset: 4,
+            bytes: 4,
+        })
+        .unwrap(),
+    );
+
+    drop(relay);
+    let _ = server.kill();
+    let _ = server.wait();
+}
+
 #[test]
 fn stale_packets_after_resume_are_ignored_not_fatal() {
     let dir = tempfile::tempdir().unwrap();