Harden reconnect and stream open reliability

src/bin/dosh-server.rs

@@ -394,6 +394,7 @@ struct ClientState {
     opened_streams: HashSet<u64>,
     stream_send_credit: HashMap<u64, usize>,
     stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>>,
+    stream_pending_opens: HashMap<u64, PendingStreamOpen>,
     stream_next_send_offset: HashMap<u64, u64>,
     stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
     stream_next_recv_offset: HashMap<u64, u64>,
@@ -425,6 +426,14 @@ struct PendingStreamChunk {
     attempts: u32,
 }
 
+#[derive(Clone)]
+struct PendingStreamOpen {
+    target_host: String,
+    target_port: u16,
+    last_sent: Instant,
+    attempts: u32,
+}
+
 struct PendingNativeAuth {
     client: dosh::native::NativeClientHello,
     server: NativeServerHello,
@@ -1150,6 +1159,7 @@ async fn handle_native_user_auth(
                     opened_streams: HashSet::new(),
                     stream_send_credit: HashMap::new(),
                     stream_pending_data: HashMap::new(),
+                    stream_pending_opens: HashMap::new(),
                     stream_next_send_offset: HashMap::new(),
                     stream_sent_data: HashMap::new(),
                     stream_next_recv_offset: HashMap::new(),
@@ -1293,6 +1303,7 @@ async fn handle_bootstrap_attach(
                     opened_streams: HashSet::new(),
                     stream_send_credit: HashMap::new(),
                     stream_pending_data: HashMap::new(),
+                    stream_pending_opens: HashMap::new(),
                     stream_next_send_offset: HashMap::new(),
                     stream_sent_data: HashMap::new(),
                     stream_next_recv_offset: HashMap::new(),
@@ -1428,6 +1439,7 @@ async fn handle_ticket_attach(
                     opened_streams: HashSet::new(),
                     stream_send_credit: HashMap::new(),
                     stream_pending_data: HashMap::new(),
+                    stream_pending_opens: HashMap::new(),
                     stream_next_send_offset: HashMap::new(),
                     stream_sent_data: HashMap::new(),
                     stream_next_recv_offset: HashMap::new(),
@@ -1503,7 +1515,10 @@ async fn handle_resume(
 ) -> Result<()> {
     let (key, session_name) = match find_client_decrypt_key(state, &packet.header) {
         Ok(found) => found,
-        Err(_) => return send_reject(socket, peer, "unknown client").await,
+        Err(_) => {
+            return send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client")
+                .await;
+        }
     };
     let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
     let req: ResumeRequest = protocol::from_body(&body)?;
@@ -1743,7 +1758,11 @@ async fn handle_stream_open(
     let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
     let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
     let open: StreamOpen = protocol::from_body(&body)?;
-    let allowed = {
+    enum StreamOpenCheck {
+        Duplicate,
+        New(Result<()>),
+    }
+    let check = {
         let mut locked = state.lock().expect("server state poisoned");
         let config = locked.config.clone();
         let session = locked
@@ -1759,17 +1778,27 @@ async fn handle_stream_open(
         }
         client.endpoint = peer;
         client.last_seen = Instant::now();
-        stream_open_allowed(&config, client, &open)
+        if client.opened_streams.contains(&open.stream_id) {
+            StreamOpenCheck::Duplicate
+        } else {
+            StreamOpenCheck::New(stream_open_allowed(&config, client, &open))
+        }
     };
-    if let Err(err) = allowed {
-        return send_stream_open_reject(
-            state,
-            socket,
-            packet.header.conn_id,
-            open.stream_id,
-            err.to_string(),
-        )
-        .await;
+    match check {
+        StreamOpenCheck::Duplicate => {
+            return send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await;
+        }
+        StreamOpenCheck::New(Ok(())) => {}
+        StreamOpenCheck::New(Err(err)) => {
+            return send_stream_open_reject(
+                state,
+                socket,
+                packet.header.conn_id,
+                open.stream_id,
+                err.to_string(),
+            )
+            .await;
+        }
     }
 
     let target = format!("{}:{}", open.target_host, open.target_port);
@@ -1928,6 +1957,7 @@ async fn handle_stream_open_ok(
         }
         client.endpoint = peer;
         client.last_seen = Instant::now();
+        client.stream_pending_opens.remove(&ok.stream_id);
         client.opened_streams.insert(ok.stream_id);
         client
             .stream_send_credit
@@ -1960,6 +1990,7 @@ async fn handle_stream_open_reject(
     client.endpoint = peer;
     client.last_seen = Instant::now();
     client.stream_writers.remove(&reject.stream_id);
+    client.stream_pending_opens.remove(&reject.stream_id);
     client.opened_streams.remove(&reject.stream_id);
     client.stream_send_credit.remove(&reject.stream_id);
     client.stream_pending_data.remove(&reject.stream_id);
@@ -2018,6 +2049,7 @@ async fn handle_stream_close(
     client.endpoint = peer;
     client.last_seen = Instant::now();
     client.stream_writers.remove(&close.stream_id);
+    client.stream_pending_opens.remove(&close.stream_id);
     client.opened_streams.remove(&close.stream_id);
     client.stream_send_credit.remove(&close.stream_id);
     client.stream_pending_data.remove(&close.stream_id);
@@ -2331,10 +2363,40 @@ async fn send_stream_open_to_client(
 ) -> Result<()> {
     let body = protocol::to_body(&StreamOpen {
         stream_id,
-        target_host,
+        target_host: target_host.clone(),
         target_port,
     })?;
-    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpen, body).await
+    let send = {
+        let mut locked = state.lock().expect("server state poisoned");
+        let mut found = None;
+        if let Some(client) = locked.client_mut(&client_id) {
+            client.send_seq += 1;
+            let packet = protocol::encode_encrypted(
+                PacketKind::StreamOpen,
+                client_id,
+                client.send_seq,
+                client.last_acked,
+                &client.session_key,
+                SERVER_TO_CLIENT,
+                &body,
+            )?;
+            client.stream_pending_opens.insert(
+                stream_id,
+                PendingStreamOpen {
+                    target_host,
+                    target_port,
+                    last_sent: Instant::now(),
+                    attempts: 1,
+                },
+            );
+            found = Some((client.endpoint, packet));
+        }
+        found
+    };
+    if let Some((endpoint, packet)) = send {
+        socket.send_to(&packet, endpoint).await?;
+    }
+    Ok(())
 }
 
 async fn send_stream_open_reject(
@@ -2857,6 +2919,38 @@ async fn retransmit_pending(
                         sends.push((client.endpoint, pending.packet.clone()));
                     }
                 }
+                let pending_open_ids: Vec<u64> =
+                    client.stream_pending_opens.keys().copied().collect();
+                let mut retransmit_opens = Vec::new();
+                for stream_id in pending_open_ids {
+                    let Some(open) = client.stream_pending_opens.get_mut(&stream_id) else {
+                        continue;
+                    };
+                    if now.duration_since(open.last_sent) < Duration::from_millis(200) {
+                        continue;
+                    }
+                    open.last_sent = now;
+                    open.attempts = open.attempts.saturating_add(1);
+                    retransmit_opens.push((stream_id, open.target_host.clone(), open.target_port));
+                }
+                for (stream_id, target_host, target_port) in retransmit_opens {
+                    client.send_seq += 1;
+                    let body = protocol::to_body(&StreamOpen {
+                        stream_id,
+                        target_host,
+                        target_port,
+                    })?;
+                    let packet = protocol::encode_encrypted(
+                        PacketKind::StreamOpen,
+                        *client_id,
+                        client.send_seq,
+                        client.last_acked,
+                        &client.session_key,
+                        SERVER_TO_CLIENT,
+                        &body,
+                    )?;
+                    sends.push((client.endpoint, packet));
+                }
                 let stream_ids: Vec<u64> = client.stream_sent_data.keys().copied().collect();
                 let mut retransmit_chunks = Vec::new();
                 for stream_id in stream_ids {
@@ -3299,6 +3393,46 @@ mod tests {
         assert_eq!(ServerConfig::default().client_timeout_secs, 2_592_000);
     }
 
+    #[tokio::test]
+    async fn unknown_resume_reject_keeps_client_id() {
+        let (pty_tx, _rx) = mpsc::unbounded_channel();
+        let state = Arc::new(Mutex::new(ServerState::new(
+            ServerConfig::default(),
+            [0u8; 32],
+            pty_tx,
+        )));
+        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
+        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let client_id = [9u8; 16];
+        let packet = protocol::Packet {
+            header: protocol::Header {
+                kind: PacketKind::ResumeRequest,
+                flags: 1,
+                conn_id: client_id,
+                seq: 77,
+                ack: 0,
+                session_key_id: [8u8; 16],
+                body_len: 0,
+            },
+            body: Vec::new(),
+        };
+
+        handle_resume(&state, &sender, receiver.local_addr().unwrap(), &packet)
+            .await
+            .unwrap();
+
+        let mut buf = [0u8; 1024];
+        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
+            .await
+            .unwrap()
+            .unwrap();
+        let reject = protocol::decode(&buf[..n]).unwrap();
+        assert_eq!(reject.header.kind, PacketKind::AttachReject);
+        assert_eq!(reject.header.conn_id, client_id);
+        let body: AttachReject = protocol::from_body(&reject.body).unwrap();
+        assert_eq!(body.reason, "unknown client");
+    }
+
     fn test_client_state(session_key: [u8; 32]) -> ClientState {
         ClientState {
             endpoint: "127.0.0.1:9".parse().unwrap(),
@@ -3318,6 +3452,7 @@ mod tests {
             opened_streams: HashSet::new(),
             stream_send_credit: HashMap::new(),
             stream_pending_data: HashMap::new(),
+            stream_pending_opens: HashMap::new(),
             stream_next_send_offset: HashMap::new(),
             stream_sent_data: HashMap::new(),
             stream_next_recv_offset: HashMap::new(),
@@ -3389,6 +3524,121 @@ mod tests {
         assert!(locked.lookup_client(&client_id).is_none());
     }
 
+    #[tokio::test]
+    async fn duplicate_stream_open_for_open_stream_resends_ok() {
+        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
+        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
+        let client_id = [11u8; 16];
+        let session_key = [12u8; 32];
+        let stream_id = 99;
+        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
+        let mut client = test_client_state(session_key);
+        client.endpoint = receiver.local_addr().unwrap();
+        client.opened_streams.insert(stream_id);
+        state.sessions.insert(
+            "test".to_string(),
+            Session {
+                pty: None,
+                parser: vt100::Parser::new(24, 80, 100),
+                clients: HashMap::from([(client_id, client)]),
+                output_seq: 0,
+                recent: VecDeque::new(),
+                empty_since: None,
+                holder_control: None,
+                persistent: false,
+                bytes_since_persist: 0,
+                last_persisted_seq: 0,
+            },
+        );
+        state.client_index.insert(client_id, "test".to_string());
+        let state = Arc::new(Mutex::new(state));
+        let raw = protocol::encode_encrypted(
+            PacketKind::StreamOpen,
+            client_id,
+            2,
+            0,
+            &session_key,
+            CLIENT_TO_SERVER,
+            &protocol::to_body(&StreamOpen {
+                stream_id,
+                target_host: "127.0.0.1".to_string(),
+                target_port: 9,
+            })
+            .unwrap(),
+        )
+        .unwrap();
+        let packet = protocol::decode(&raw).unwrap();
+
+        handle_stream_open(&state, &sender, receiver.local_addr().unwrap(), &packet)
+            .await
+            .unwrap();
+
+        let mut buf = [0u8; 2048];
+        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
+            .await
+            .unwrap()
+            .unwrap();
+        let packet = protocol::decode(&buf[..n]).unwrap();
+        assert_eq!(packet.header.kind, PacketKind::StreamOpenOk);
+        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
+        let ok: StreamOpenOk = protocol::from_body(&plain).unwrap();
+        assert_eq!(ok.stream_id, stream_id);
+    }
+
+    #[tokio::test]
+    async fn pending_server_stream_open_is_retransmitted() {
+        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
+        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
+        let client_id = [13u8; 16];
+        let session_key = [14u8; 32];
+        let stream_id = 123;
+        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
+        let mut client = test_client_state(session_key);
+        client.endpoint = receiver.local_addr().unwrap();
+        client.stream_pending_opens.insert(
+            stream_id,
+            PendingStreamOpen {
+                target_host: "127.0.0.1".to_string(),
+                target_port: 8080,
+                last_sent: Instant::now() - Duration::from_millis(250),
+                attempts: 1,
+            },
+        );
+        state.sessions.insert(
+            "test".to_string(),
+            Session {
+                pty: None,
+                parser: vt100::Parser::new(24, 80, 100),
+                clients: HashMap::from([(client_id, client)]),
+                output_seq: 0,
+                recent: VecDeque::new(),
+                empty_since: None,
+                holder_control: None,
+                persistent: false,
+                bytes_since_persist: 0,
+                last_persisted_seq: 0,
+            },
+        );
+        state.client_index.insert(client_id, "test".to_string());
+        let state = Arc::new(Mutex::new(state));
+
+        retransmit_pending(&state, &sender).await.unwrap();
+
+        let mut buf = [0u8; 2048];
+        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
+            .await
+            .unwrap()
+            .unwrap();
+        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: StreamOpen = protocol::from_body(&plain).unwrap();
+        assert_eq!(open.stream_id, stream_id);
+        assert_eq!(open.target_port, 8080);
+    }
+
     #[test]
     fn forward_only_session_does_not_allocate_pty() {
         let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
@@ -3520,6 +3770,7 @@ mod tests {
                         opened_streams: HashSet::new(),
                         stream_send_credit: HashMap::new(),
                         stream_pending_data: HashMap::new(),
+                        stream_pending_opens: HashMap::new(),
                         stream_next_send_offset: HashMap::new(),
                         stream_sent_data: HashMap::new(),
                         stream_next_recv_offset: HashMap::new(),
@@ -3615,6 +3866,7 @@ mod tests {
                         opened_streams: HashSet::from([42]),
                         stream_send_credit: HashMap::from([(42, 0)]),
                         stream_pending_data: HashMap::new(),
+                        stream_pending_opens: HashMap::new(),
                         stream_next_send_offset: HashMap::new(),
                         stream_sent_data: HashMap::new(),
                         stream_next_recv_offset: HashMap::new(),