Support stream EOF half closes

src/bin/dosh-client.rs

@@ -28,8 +28,8 @@ use dosh::protocol::{
     self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
     NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody,
     NativeUserAuthBody, PacketKind, Rekey, Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose,
-    StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, TicketAttachBody,
+    StreamData, StreamEof, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
-    TicketAttachEnvelope, TicketAttachOkEnvelope,
+    TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
 };
 use dosh::ssh_agent;
 use dosh::transport::{
@@ -218,6 +218,9 @@ enum ForwardEvent {
         stream_id: u64,
         bytes: Vec<u8>,
     },
+    Eof {
+        stream_id: u64,
+    },
     Close {
         stream_id: u64,
     },
@@ -5900,6 +5903,8 @@ async fn run_terminal(
     let mut stream_retransmit_srtt: Option<Duration> = None;
     let mut stream_next_recv_offset: HashMap<u64, u64> = HashMap::new();
     let mut stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>> = HashMap::new();
+    let mut stream_local_eof_sent: HashSet<u64> = HashSet::new();
+    let mut stream_remote_eof_received: HashSet<u64> = HashSet::new();
     let mut stream_retired: HashSet<u64> = HashSet::new();
     let mut stream_retired_order: VecDeque<u64> = VecDeque::new();
     let mut pending_user_input: VecDeque<PendingUserInput> = VecDeque::new();
@@ -6653,9 +6658,13 @@ async fn run_terminal(
                                         let forward_tx = remote_forward_tx.clone();
                                         tokio::spawn(async move {
                                             let mut buf = [0u8; 16 * 1024];
+                                            let mut graceful_eof = false;
                                             loop {
                                                 match reader.read(&mut buf).await {
-                                                    Ok(0) => break,
+                                                    Ok(0) => {
+                                                        graceful_eof = true;
+                                                        break;
+                                                    }
                                                     Ok(n) => {
                                                         if forward_tx
                                                             .send(ForwardEvent::Data {
@@ -6672,8 +6681,14 @@ async fn run_terminal(
                                                 }
                                             }
                                             let _ = forward_tx
-                                                .send(ForwardEvent::Close {
+                                                .send(if graceful_eof {
+                                                    ForwardEvent::Eof {
                                                         stream_id: open.stream_id,
+                                                    }
+                                                } else {
+                                                    ForwardEvent::Close {
+                                                        stream_id: open.stream_id,
+                                                    }
                                                 })
                                                 .await;
                                         });
@@ -6705,9 +6720,13 @@ async fn run_terminal(
                                 let forward_tx = remote_forward_tx.clone();
                                 tokio::spawn(async move {
                                     let mut buf = [0u8; 16 * 1024];
+                                    let mut graceful_eof = false;
                                     loop {
                                         match reader.read(&mut buf).await {
-                                            Ok(0) => break,
+                                            Ok(0) => {
+                                                graceful_eof = true;
+                                                break;
+                                            }
                                             Ok(n) => {
                                                 if forward_tx
                                                     .send(ForwardEvent::Data {
@@ -6723,9 +6742,17 @@ async fn run_terminal(
                                             Err(_) => break,
                                         }
                                     }
-                                    let _ = forward_tx.send(ForwardEvent::Close {
+                                    let _ = forward_tx
+                                        .send(if graceful_eof {
+                                            ForwardEvent::Eof {
                                                 stream_id: open.stream_id,
-                                    }).await;
+                                            }
+                                        } else {
+                                            ForwardEvent::Close {
+                                                stream_id: open.stream_id,
+                                            }
+                                        })
+                                        .await;
                                 });
                             }
                             Err(err) => {
@@ -6778,6 +6805,10 @@ async fn run_terminal(
                             &mut stream_next_send_offset,
                             &mut stream_sent_data,
                         ).await?;
+                        if stream_local_eof_sent.contains(&ok.stream_id) {
+                            send_stream_eof(&socket, addr, &cred, &mut send_seq, ok.stream_id)
+                                .await?;
+                        }
                     }
                     PacketKind::StreamOpenReject => {
                         let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
@@ -6803,6 +6834,8 @@ async fn run_terminal(
                             &mut stream_sent_data,
                             &mut stream_next_recv_offset,
                             &mut stream_recv_pending,
+                            &mut stream_local_eof_sent,
+                            &mut stream_remote_eof_received,
                             &mut stream_retired,
                             &mut stream_retired_order,
                         );
@@ -6876,6 +6909,51 @@ async fn run_terminal(
                             &mut stream_sent_data,
                         ).await?;
                     }
+                    PacketKind::StreamEof => {
+                        let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
+                            continue;
+                        };
+                        let Ok(eof) = protocol::from_body::<StreamEof>(&plain) else {
+                            continue;
+                        };
+                        last_packet_at = Instant::now();
+                        if !opened_streams.contains(&eof.stream_id) {
+                            continue;
+                        }
+                        stream_remote_eof_received.insert(eof.stream_id);
+                        if let Some(writer) = stream_writers.get_mut(&eof.stream_id) {
+                            let _ = writer.shutdown().await;
+                        }
+                        #[cfg(unix)]
+                        if let Some(writer) = agent_writers.get_mut(&eof.stream_id) {
+                            let _ = writer.shutdown().await;
+                        }
+                        if stream_eof_complete(
+                            eof.stream_id,
+                            &stream_local_eof_sent,
+                            &stream_remote_eof_received,
+                        ) {
+                            pending_socks_replies.remove(&eof.stream_id);
+                            retire_stream_state(
+                                eof.stream_id,
+                                &mut opened_streams,
+                                &mut stream_send_credit,
+                                &mut stream_pending_data,
+                                &mut stream_pending_opens,
+                                &mut stream_next_send_offset,
+                                &mut stream_sent_data,
+                                &mut stream_next_recv_offset,
+                                &mut stream_recv_pending,
+                                &mut stream_local_eof_sent,
+                                &mut stream_remote_eof_received,
+                                &mut stream_retired,
+                                &mut stream_retired_order,
+                            );
+                            stream_writers.remove(&eof.stream_id);
+                            #[cfg(unix)]
+                            agent_writers.remove(&eof.stream_id);
+                        }
+                    }
                     PacketKind::StreamClose => {
                         let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
                             continue;
@@ -6895,6 +6973,8 @@ async fn run_terminal(
                             &mut stream_sent_data,
                             &mut stream_next_recv_offset,
                             &mut stream_recv_pending,
+                            &mut stream_local_eof_sent,
+                            &mut stream_remote_eof_received,
                             &mut stream_retired,
                             &mut stream_retired_order,
                         );
@@ -6947,6 +7027,37 @@ async fn run_terminal(
                             &mut stream_sent_data,
                         ).await?;
                     }
+                    Some(ForwardEvent::Eof { stream_id }) => {
+                        stream_local_eof_sent.insert(stream_id);
+                        if opened_streams.contains(&stream_id) {
+                            send_stream_eof(&socket, addr, &cred, &mut send_seq, stream_id).await?;
+                        }
+                        if stream_eof_complete(
+                            stream_id,
+                            &stream_local_eof_sent,
+                            &stream_remote_eof_received,
+                        ) {
+                            pending_socks_replies.remove(&stream_id);
+                            retire_stream_state(
+                                stream_id,
+                                &mut opened_streams,
+                                &mut stream_send_credit,
+                                &mut stream_pending_data,
+                                &mut stream_pending_opens,
+                                &mut stream_next_send_offset,
+                                &mut stream_sent_data,
+                                &mut stream_next_recv_offset,
+                                &mut stream_recv_pending,
+                                &mut stream_local_eof_sent,
+                                &mut stream_remote_eof_received,
+                                &mut stream_retired,
+                                &mut stream_retired_order,
+                            );
+                            stream_writers.remove(&stream_id);
+                            #[cfg(unix)]
+                            agent_writers.remove(&stream_id);
+                        }
+                    }
                     Some(ForwardEvent::Close { stream_id }) => {
                         pending_socks_replies.remove(&stream_id);
                         retire_stream_state(
@@ -6959,6 +7070,8 @@ async fn run_terminal(
                             &mut stream_sent_data,
                             &mut stream_next_recv_offset,
                             &mut stream_recv_pending,
+                            &mut stream_local_eof_sent,
+                            &mut stream_remote_eof_received,
                             &mut stream_retired,
                             &mut stream_retired_order,
                         );
@@ -7166,9 +7279,13 @@ async fn start_local_forwards(
                 let forward_tx = forward_tx.clone();
                 tokio::spawn(async move {
                     let mut buf = [0u8; 16 * 1024];
+                    let mut graceful_eof = false;
                     loop {
                         match reader.read(&mut buf).await {
-                            Ok(0) => break,
+                            Ok(0) => {
+                                graceful_eof = true;
+                                break;
+                            }
                             Ok(n) => {
                                 if forward_tx
                                     .send(ForwardEvent::Data {
@@ -7184,7 +7301,13 @@ async fn start_local_forwards(
                             Err(_) => break,
                         }
                     }
-                    let _ = forward_tx.send(ForwardEvent::Close { stream_id }).await;
+                    let _ = forward_tx
+                        .send(if graceful_eof {
+                            ForwardEvent::Eof { stream_id }
+                        } else {
+                            ForwardEvent::Close { stream_id }
+                        })
+                        .await;
                 });
             }
         });
@@ -7287,9 +7410,13 @@ async fn handle_socks5_connect(
         .await?;
     tokio::spawn(async move {
         let mut buf = [0u8; 16 * 1024];
+        let mut graceful_eof = false;
         loop {
             match reader.read(&mut buf).await {
-                Ok(0) => break,
+                Ok(0) => {
+                    graceful_eof = true;
+                    break;
+                }
                 Ok(n) => {
                     if forward_tx
                         .send(ForwardEvent::Data {
@@ -7305,7 +7432,13 @@ async fn handle_socks5_connect(
                 Err(_) => break,
             }
         }
-        let _ = forward_tx.send(ForwardEvent::Close { stream_id }).await;
+        let _ = forward_tx
+            .send(if graceful_eof {
+                ForwardEvent::Eof { stream_id }
+            } else {
+                ForwardEvent::Close { stream_id }
+            })
+            .await;
     });
     Ok(())
 }
@@ -8145,6 +8278,8 @@ fn cleanup_stream_state(
     stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
     stream_next_recv_offset: &mut HashMap<u64, u64>,
     stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
+    stream_local_eof_sent: &mut HashSet<u64>,
+    stream_remote_eof_received: &mut HashSet<u64>,
 ) {
     opened_streams.remove(&stream_id);
     stream_send_credit.remove(&stream_id);
@@ -8154,6 +8289,16 @@ fn cleanup_stream_state(
     stream_sent_data.remove(&stream_id);
     stream_next_recv_offset.remove(&stream_id);
     stream_recv_pending.remove(&stream_id);
+    stream_local_eof_sent.remove(&stream_id);
+    stream_remote_eof_received.remove(&stream_id);
+}
+
+fn stream_eof_complete(
+    stream_id: u64,
+    stream_local_eof_sent: &HashSet<u64>,
+    stream_remote_eof_received: &HashSet<u64>,
+) -> bool {
+    stream_local_eof_sent.contains(&stream_id) && stream_remote_eof_received.contains(&stream_id)
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -8167,6 +8312,8 @@ fn stream_state_exists(
     stream_sent_data: &HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
     stream_next_recv_offset: &HashMap<u64, u64>,
     stream_recv_pending: &HashMap<u64, BTreeMap<u64, Vec<u8>>>,
+    stream_local_eof_sent: &HashSet<u64>,
+    stream_remote_eof_received: &HashSet<u64>,
 ) -> bool {
     opened_streams.contains(&stream_id)
         || stream_send_credit.contains_key(&stream_id)
@@ -8176,6 +8323,8 @@ fn stream_state_exists(
         || stream_sent_data.contains_key(&stream_id)
         || stream_next_recv_offset.contains_key(&stream_id)
         || stream_recv_pending.contains_key(&stream_id)
+        || stream_local_eof_sent.contains(&stream_id)
+        || stream_remote_eof_received.contains(&stream_id)
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -8189,6 +8338,8 @@ fn retire_stream_state(
     stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
     stream_next_recv_offset: &mut HashMap<u64, u64>,
     stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
+    stream_local_eof_sent: &mut HashSet<u64>,
+    stream_remote_eof_received: &mut HashSet<u64>,
     stream_retired: &mut HashSet<u64>,
     stream_retired_order: &mut VecDeque<u64>,
 ) {
@@ -8202,6 +8353,8 @@ fn retire_stream_state(
         stream_sent_data,
         stream_next_recv_offset,
         stream_recv_pending,
+        stream_local_eof_sent,
+        stream_remote_eof_received,
     );
     cleanup_stream_state(
         stream_id,
@@ -8213,6 +8366,8 @@ fn retire_stream_state(
         stream_sent_data,
         stream_next_recv_offset,
         stream_recv_pending,
+        stream_local_eof_sent,
+        stream_remote_eof_received,
     );
     if !had_state || STREAM_RETIRED_TOMBSTONES == 0 {
         return;
@@ -8307,6 +8462,17 @@ async fn send_stream_close(
     send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamClose, &body).await
 }
 
+async fn send_stream_eof(
+    socket: &UdpSocket,
+    addr: SocketAddr,
+    cred: &CachedCredential,
+    send_seq: &mut u64,
+    stream_id: u64,
+) -> Result<()> {
+    let body = protocol::to_body(&StreamEof { stream_id })?;
+    send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamEof, &body).await
+}
+
 async fn send_stream_packet(
     socket: &UdpSocket,
     addr: SocketAddr,
@@ -9603,7 +9769,7 @@ mod tests {
         release_tag_from_effective_url, release_version_from_tag, render_status_clear,
         render_status_overlay, requested_env, resolved_startup_command, retire_stream_state,
         retransmit_stream_opens, rewrite_forward_command, sanitize_trace_name,
-        selected_predict_mode, selected_udp_host, server_version_mismatch,
+        selected_predict_mode, selected_udp_host, send_stream_eof, server_version_mismatch,
         should_flush_terminal_input_after_contact, should_health_log_client_start,
         should_hold_during_startup_gate, should_hold_post_submit_input,
         should_reconnect_before_input_for_local_sleep, should_repaint_idle_terminal,
@@ -9730,6 +9896,8 @@ mod tests {
             (stream_id, BTreeMap::from([(77, vec![5])])),
             (other_stream_id, BTreeMap::from([(88, vec![6])])),
         ]);
+        let mut stream_local_eof_sent = HashSet::from([stream_id, other_stream_id]);
+        let mut stream_remote_eof_received = HashSet::from([stream_id, other_stream_id]);
 
         cleanup_stream_state(
             stream_id,
@@ -9741,6 +9909,8 @@ mod tests {
             &mut stream_sent_data,
             &mut stream_next_recv_offset,
             &mut stream_recv_pending,
+            &mut stream_local_eof_sent,
+            &mut stream_remote_eof_received,
         );
 
         assert!(!opened_streams.contains(&stream_id));
@@ -9751,6 +9921,8 @@ mod tests {
         assert!(!stream_sent_data.contains_key(&stream_id));
         assert!(!stream_next_recv_offset.contains_key(&stream_id));
         assert!(!stream_recv_pending.contains_key(&stream_id));
+        assert!(!stream_local_eof_sent.contains(&stream_id));
+        assert!(!stream_remote_eof_received.contains(&stream_id));
 
         assert!(opened_streams.contains(&other_stream_id));
         assert!(stream_send_credit.contains_key(&other_stream_id));
@@ -9760,6 +9932,8 @@ mod tests {
         assert!(stream_sent_data.contains_key(&other_stream_id));
         assert!(stream_next_recv_offset.contains_key(&other_stream_id));
         assert!(stream_recv_pending.contains_key(&other_stream_id));
+        assert!(stream_local_eof_sent.contains(&other_stream_id));
+        assert!(stream_remote_eof_received.contains(&other_stream_id));
     }
 
     #[test]
@@ -9773,6 +9947,8 @@ mod tests {
         let mut stream_sent_data = HashMap::new();
         let mut stream_next_recv_offset = HashMap::new();
         let mut stream_recv_pending = HashMap::new();
+        let mut stream_local_eof_sent = HashSet::new();
+        let mut stream_remote_eof_received = HashSet::new();
         let mut stream_retired = HashSet::new();
         let mut stream_retired_order = VecDeque::new();
 
@@ -9786,6 +9962,8 @@ mod tests {
             &mut stream_sent_data,
             &mut stream_next_recv_offset,
             &mut stream_recv_pending,
+            &mut stream_local_eof_sent,
+            &mut stream_remote_eof_received,
             &mut stream_retired,
             &mut stream_retired_order,
         );
@@ -9799,6 +9977,8 @@ mod tests {
             &mut stream_sent_data,
             &mut stream_next_recv_offset,
             &mut stream_recv_pending,
+            &mut stream_local_eof_sent,
+            &mut stream_remote_eof_received,
             &mut stream_retired,
             &mut stream_retired_order,
         );
@@ -10922,6 +11102,44 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn client_stream_eof_is_sent_without_stream_close() {
+        let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let sender = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let addr = receiver.local_addr().unwrap();
+        let cred = CachedCredential {
+            server: "host".to_string(),
+            session: "term".to_string(),
+            mode: "forward-only".to_string(),
+            udp_host: "127.0.0.1".to_string(),
+            udp_port: addr.port(),
+            client_id: [1; 16],
+            session_key: [2; 32],
+            session_key_id: protocol::session_key_id(&[2; 32]),
+            attach_ticket: Vec::new(),
+            attach_ticket_psk: [3; 32],
+            last_rendered_seq: 0,
+        };
+        let mut send_seq = 10;
+
+        send_stream_eof(&sender, addr, &cred, &mut send_seq, 44)
+            .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::StreamEof);
+        let plain =
+            protocol::decrypt_body(&packet, &cred.session_key, protocol::CLIENT_TO_SERVER).unwrap();
+        let eof: protocol::StreamEof = protocol::from_body(&plain).unwrap();
+        assert_eq!(eof.stream_id, 44);
+        assert_eq!(send_seq, 11);
+    }
+
     #[tokio::test]
     async fn client_stream_retransmit_uses_observed_rtt() {
         let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();

src/bin/dosh-server.rs

@@ -24,8 +24,8 @@ use dosh::protocol::{
     self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
     NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody,
     NativeUserAuthBody, PacketKind, ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT,
-    StreamClose, StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
+    StreamClose, StreamData, StreamEof, StreamOpen, StreamOpenOk, StreamOpenReject,
-    TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
+    StreamWindowAdjust, TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
 };
 use dosh::pty::{PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session};
 use dosh::udp::{is_transient_udp_error, is_transient_udp_send_error};
@@ -433,6 +433,8 @@ struct ClientState {
     stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
     stream_next_recv_offset: HashMap<u64, u64>,
     stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
+    stream_local_eof_sent: HashSet<u64>,
+    stream_remote_eof_received: HashSet<u64>,
     stream_retired: HashSet<u64>,
     stream_retired_order: VecDeque<u64>,
     stream_retransmit_srtt: Option<Duration>,
@@ -939,6 +941,7 @@ async fn handle_packet(
         PacketKind::StreamWindowAdjust => {
             handle_stream_window_adjust(state, socket, peer, &packet).await
         }
+        PacketKind::StreamEof => handle_stream_eof(state, peer, &packet).await,
         PacketKind::StreamClose => handle_stream_close(state, peer, &packet).await,
         _ => Ok(()),
     }
@@ -1268,6 +1271,8 @@ async fn handle_native_user_auth(
                     stream_sent_data: HashMap::new(),
                     stream_next_recv_offset: HashMap::new(),
                     stream_recv_pending: HashMap::new(),
+                    stream_local_eof_sent: HashSet::new(),
+                    stream_remote_eof_received: HashSet::new(),
                     stream_retired: HashSet::new(),
                     stream_retired_order: VecDeque::new(),
                     stream_retransmit_srtt: None,
@@ -1418,6 +1423,8 @@ async fn handle_bootstrap_attach(
                     stream_sent_data: HashMap::new(),
                     stream_next_recv_offset: HashMap::new(),
                     stream_recv_pending: HashMap::new(),
+                    stream_local_eof_sent: HashSet::new(),
+                    stream_remote_eof_received: HashSet::new(),
                     stream_retired: HashSet::new(),
                     stream_retired_order: VecDeque::new(),
                     stream_retransmit_srtt: None,
@@ -1568,6 +1575,8 @@ async fn handle_ticket_attach(
                     stream_sent_data: HashMap::new(),
                     stream_next_recv_offset: HashMap::new(),
                     stream_recv_pending: HashMap::new(),
+                    stream_local_eof_sent: HashSet::new(),
+                    stream_remote_eof_received: HashSet::new(),
                     stream_retired: HashSet::new(),
                     stream_retired_order: VecDeque::new(),
                     stream_retransmit_srtt: None,
@@ -2152,6 +2161,7 @@ async fn handle_stream_open(
                 break;
             }
         }
+        let _ = writer.shutdown().await;
     });
 
     send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;
@@ -2162,9 +2172,13 @@ async fn handle_stream_open(
     let stream_id = open.stream_id;
     tokio::spawn(async move {
         let mut buf = [0u8; 16 * 1024];
+        let mut graceful_eof = false;
         loop {
             match reader.read(&mut buf).await {
-                Ok(0) => break,
+                Ok(0) => {
+                    graceful_eof = true;
+                    break;
+                }
                 Ok(n) => {
                     if let Err(err) = send_stream_data_to_client(
                         &state,
@@ -2182,7 +2196,11 @@ async fn handle_stream_open(
                 Err(_) => break,
             }
         }
-        let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
+        let _ = if graceful_eof {
+            send_stream_eof_to_client(&state, &socket, client_id, stream_id).await
+        } else {
+            send_stream_close_to_client(&state, &socket, client_id, stream_id).await
+        };
     });
 
     Ok(())
@@ -2305,7 +2323,27 @@ async fn handle_stream_open_ok(
             .entry(ok.stream_id)
             .or_insert(STREAM_INITIAL_WINDOW);
     }
-    flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, ok.stream_id).await
+    flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, ok.stream_id).await?;
+    let should_send_eof = {
+        let mut locked = state.lock().expect("server state poisoned");
+        locked
+            .client_mut(&packet.header.conn_id)
+            .is_some_and(|client| client.stream_local_eof_sent.contains(&ok.stream_id))
+    };
+    if should_send_eof {
+        let body = protocol::to_body(&StreamEof {
+            stream_id: ok.stream_id,
+        })?;
+        send_stream_packet_to_client(
+            state,
+            socket,
+            packet.header.conn_id,
+            PacketKind::StreamEof,
+            body,
+        )
+        .await?;
+    }
+    Ok(())
 }
 
 async fn handle_stream_open_reject(
@@ -2367,6 +2405,37 @@ async fn handle_stream_window_adjust(
         .await
 }
 
+async fn handle_stream_eof(
+    state: &Arc<Mutex<ServerState>>,
+    peer: SocketAddr,
+    packet: &protocol::Packet,
+) -> Result<()> {
+    let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
+    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
+    let eof: StreamEof = protocol::from_body(&body)?;
+    let mut locked = state.lock().expect("server state poisoned");
+    let Some(session) = locked.sessions.get_mut(&session_name) else {
+        return Ok(());
+    };
+    let Some(client) = session.clients.get_mut(&packet.header.conn_id) else {
+        return Ok(());
+    };
+    if !client.replay.accept(packet.header.seq) {
+        return Ok(());
+    }
+    client.endpoint = peer;
+    client.last_seen = Instant::now();
+    if !client.opened_streams.contains(&eof.stream_id) {
+        return Ok(());
+    }
+    client.stream_remote_eof_received.insert(eof.stream_id);
+    client.stream_writers.remove(&eof.stream_id);
+    if client_stream_eof_complete(client, eof.stream_id) {
+        retire_client_stream(client, eof.stream_id);
+    }
+    Ok(())
+}
+
 async fn handle_stream_close(
     state: &Arc<Mutex<ServerState>>,
     peer: SocketAddr,
@@ -2462,6 +2531,7 @@ async fn start_remote_forwards(
                             break;
                         }
                     }
+                    let _ = writer.shutdown().await;
                 });
                 if let Err(err) = send_stream_open_to_client(
                     &state,
@@ -2480,9 +2550,13 @@ async fn start_remote_forwards(
                 let socket = Arc::clone(&socket);
                 tokio::spawn(async move {
                     let mut buf = [0u8; 16 * 1024];
+                    let mut graceful_eof = false;
                     loop {
                         match reader.read(&mut buf).await {
-                            Ok(0) => break,
+                            Ok(0) => {
+                                graceful_eof = true;
+                                break;
+                            }
                             Ok(n) => {
                                 if let Err(err) = send_stream_data_to_client(
                                     &state,
@@ -2500,8 +2574,11 @@ async fn start_remote_forwards(
                             Err(_) => break,
                         }
                     }
-                    let _ =
+                    let _ = if graceful_eof {
-                        send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
+                        send_stream_eof_to_client(&state, &socket, client_id, stream_id).await
+                    } else {
+                        send_stream_close_to_client(&state, &socket, client_id, stream_id).await
+                    };
                 });
             }
         });
@@ -2588,6 +2665,7 @@ fn spawn_agent_forward(
                         break;
                     }
                 }
+                let _ = writer.shutdown().await;
             });
             if let Err(err) = send_stream_open_to_client(
                 &state,
@@ -2606,9 +2684,13 @@ fn spawn_agent_forward(
             let socket = Arc::clone(&socket);
             tokio::spawn(async move {
                 let mut buf = [0u8; 16 * 1024];
+                let mut graceful_eof = false;
                 loop {
                     match reader.read(&mut buf).await {
-                        Ok(0) => break,
+                        Ok(0) => {
+                            graceful_eof = true;
+                            break;
+                        }
                         Ok(n) => {
                             if let Err(err) = send_stream_data_to_client(
                                 &state,
@@ -2626,7 +2708,11 @@ fn spawn_agent_forward(
                         Err(_) => break,
                     }
                 }
-                let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
+                let _ = if graceful_eof {
+                    send_stream_eof_to_client(&state, &socket, client_id, stream_id).await
+                } else {
+                    send_stream_close_to_client(&state, &socket, client_id, stream_id).await
+                };
             });
         }
         let _ = std::fs::remove_file(&socket_path);
@@ -3543,6 +3629,8 @@ fn cleanup_client_stream(client: &mut ClientState, stream_id: u64) {
     client.stream_sent_data.remove(&stream_id);
     client.stream_next_recv_offset.remove(&stream_id);
     client.stream_recv_pending.remove(&stream_id);
+    client.stream_local_eof_sent.remove(&stream_id);
+    client.stream_remote_eof_received.remove(&stream_id);
 }
 
 fn client_has_stream_state(client: &ClientState, stream_id: u64) -> bool {
@@ -3555,6 +3643,13 @@ fn client_has_stream_state(client: &ClientState, stream_id: u64) -> bool {
         || client.stream_sent_data.contains_key(&stream_id)
         || client.stream_next_recv_offset.contains_key(&stream_id)
         || client.stream_recv_pending.contains_key(&stream_id)
+        || client.stream_local_eof_sent.contains(&stream_id)
+        || client.stream_remote_eof_received.contains(&stream_id)
+}
+
+fn client_stream_eof_complete(client: &ClientState, stream_id: u64) -> bool {
+    client.stream_local_eof_sent.contains(&stream_id)
+        && client.stream_remote_eof_received.contains(&stream_id)
 }
 
 fn retire_client_stream(client: &mut ClientState, stream_id: u64) {
@@ -3650,6 +3745,22 @@ async fn send_stream_close_to_client(
     send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamClose, body).await
 }
 
+async fn send_stream_eof_to_client(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    client_id: [u8; 16],
+    stream_id: u64,
+) -> Result<()> {
+    {
+        let mut locked = state.lock().expect("server state poisoned");
+        if let Some(client) = locked.client_mut(&client_id) {
+            client.stream_local_eof_sent.insert(stream_id);
+        }
+    }
+    let body = protocol::to_body(&StreamEof { stream_id })?;
+    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamEof, body).await
+}
+
 async fn send_stream_packet_to_client(
     state: &Arc<Mutex<ServerState>>,
     socket: &Arc<UdpSocket>,
@@ -4513,6 +4624,8 @@ mod tests {
             stream_sent_data: HashMap::new(),
             stream_next_recv_offset: HashMap::new(),
             stream_recv_pending: HashMap::new(),
+            stream_local_eof_sent: HashSet::new(),
+            stream_remote_eof_received: HashSet::new(),
             stream_retired: HashSet::new(),
             stream_retired_order: VecDeque::new(),
             stream_retransmit_srtt: None,
@@ -5177,6 +5290,8 @@ mod tests {
                         stream_sent_data: HashMap::new(),
                         stream_next_recv_offset: HashMap::new(),
                         stream_recv_pending: HashMap::new(),
+                        stream_local_eof_sent: HashSet::new(),
+                        stream_remote_eof_received: HashSet::new(),
                         stream_retired: HashSet::new(),
                         stream_retired_order: VecDeque::new(),
                         stream_retransmit_srtt: None,
@@ -5238,6 +5353,92 @@ mod tests {
         assert_eq!(data.bytes, b"hello");
     }
 
+    #[tokio::test]
+    async fn stream_eof_to_client_does_not_retire_stream() {
+        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 receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
+        let endpoint = receiver.local_addr().unwrap();
+        state.sessions.insert(
+            "test".to_string(),
+            Session {
+                pty: None,
+                parser: vt100::Parser::new(24, 80, 100),
+                restored_screen: None,
+                clients: HashMap::from([(
+                    client_id,
+                    ClientState {
+                        endpoint,
+                        mode: "forward-only".to_string(),
+                        session_key,
+                        last_acked: 0,
+                        replay: ReplayWindow::default(),
+                        send_seq: 1,
+                        cols: 80,
+                        rows: 24,
+                        last_seen: Instant::now(),
+                        pending: VecDeque::new(),
+                        allowed_forwardings: Vec::new(),
+                        stream_writers: HashMap::new(),
+                        opened_streams: HashSet::from([42]),
+                        stream_send_credit: HashMap::from([(42, STREAM_INITIAL_WINDOW)]),
+                        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(),
+                        stream_recv_pending: HashMap::new(),
+                        stream_local_eof_sent: HashSet::new(),
+                        stream_remote_eof_received: HashSet::new(),
+                        stream_retired: HashSet::new(),
+                        stream_retired_order: VecDeque::new(),
+                        stream_retransmit_srtt: None,
+                        epoch: 0,
+                        epoch_started: Instant::now(),
+                        epoch_packets: 0,
+                        previous_session_key: None,
+                    },
+                )]),
+                output_seq: 0,
+                recent: VecDeque::new(),
+                empty_since: None,
+                holder_control: None,
+                persistent: false,
+                bytes_since_persist: 0,
+                last_persisted_seq: 0,
+                last_screen_persist_at: Instant::now(),
+            },
+        );
+        state.client_index.insert(client_id, "test".to_string());
+        let state = Arc::new(Mutex::new(state));
+
+        send_stream_eof_to_client(&state, &sender, client_id, 42)
+            .await
+            .unwrap();
+
+        {
+            let locked = state.lock().unwrap();
+            let client = locked.sessions["test"].clients.get(&client_id).unwrap();
+            assert!(client.opened_streams.contains(&42));
+            assert!(client.stream_local_eof_sent.contains(&42));
+            assert!(!client.stream_retired.contains(&42));
+        }
+
+        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::StreamEof);
+        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
+        let eof: StreamEof = protocol::from_body(&plain).unwrap();
+        assert_eq!(eof.stream_id, 42);
+    }
+
     #[tokio::test]
     async fn server_stream_send_splits_large_writes() {
         let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
@@ -5276,6 +5477,8 @@ mod tests {
                         stream_sent_data: HashMap::new(),
                         stream_next_recv_offset: HashMap::new(),
                         stream_recv_pending: HashMap::new(),
+                        stream_local_eof_sent: HashSet::new(),
+                        stream_remote_eof_received: HashSet::new(),
                         stream_retired: HashSet::new(),
                         stream_retired_order: VecDeque::new(),
                         stream_retransmit_srtt: None,
@@ -5378,6 +5581,8 @@ mod tests {
                         stream_sent_data: HashMap::new(),
                         stream_next_recv_offset: HashMap::new(),
                         stream_recv_pending: HashMap::new(),
+                        stream_local_eof_sent: HashSet::new(),
+                        stream_remote_eof_received: HashSet::new(),
                         stream_retired: HashSet::new(),
                         stream_retired_order: VecDeque::new(),
                         stream_retransmit_srtt: None,

src/transport.rs

@@ -17,7 +17,7 @@
 
 use crate::protocol::{
     self, CLIENT_TO_SERVER, PacketKind, ReplayWindow, SERVER_TO_CLIENT, StreamClose, StreamData,
-    StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
+    StreamEof, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
 };
 use crate::udp::{is_transient_udp_error, is_transient_udp_send_error};
 use anyhow::{Result, anyhow, bail};
@@ -144,6 +144,9 @@ pub enum TransportEvent {
         stream_id: u64,
         flushed: Vec<OutgoingStreamPacket>,
     },
+    Eof {
+        stream_id: u64,
+    },
     Close {
         stream_id: u64,
     },
@@ -437,6 +440,17 @@ impl StreamMux {
         encode_packet(PacketKind::StreamClose, &StreamClose { stream_id })
     }
 
+    pub fn eof_stream(&self, stream_id: u64) -> Result<OutgoingStreamPacket> {
+        if !self.has_stream_state(stream_id) {
+            bail!("stream {stream_id} is not open");
+        }
+        encode_packet(PacketKind::StreamEof, &StreamEof { stream_id })
+    }
+
+    pub fn handle_eof(&self, eof: StreamEof) -> bool {
+        self.has_stream_state(eof.stream_id)
+    }
+
     pub fn handle_close(&mut self, close: StreamClose) -> bool {
         if !self.has_stream_state(close.stream_id) {
             return false;
@@ -507,6 +521,15 @@ impl StreamMux {
                     Ok(TransportEvent::Ignored { stream_id })
                 }
             }
+            PacketKind::StreamEof => {
+                let eof: StreamEof = protocol::from_body(body)?;
+                let stream_id = eof.stream_id;
+                if self.handle_eof(eof) {
+                    Ok(TransportEvent::Eof { stream_id })
+                } else {
+                    Ok(TransportEvent::Ignored { stream_id })
+                }
+            }
             _ => bail!("packet kind {kind:?} is not a stream transport packet"),
         }
     }
@@ -1004,6 +1027,11 @@ impl DoshTransport {
         Ok(())
     }
 
+    pub async fn eof(&mut self, stream_id: u64) -> Result<()> {
+        let packet = self.mux.eof_stream(stream_id)?;
+        self.send_outgoing(packet).await
+    }
+
     pub async fn close(&mut self, stream_id: u64) -> Result<()> {
         let packet = self.mux.close_stream(stream_id)?;
         self.send_outgoing(packet).await
@@ -1098,7 +1126,8 @@ impl DoshTransport {
             | PacketKind::StreamOpenReject
             | PacketKind::StreamData
             | PacketKind::StreamWindowAdjust
-            | PacketKind::StreamClose => {
+            | PacketKind::StreamClose
+            | PacketKind::StreamEof => {
                 let event = self.mux.handle_packet(kind, plain)?;
                 self.send_followups(&event).await?;
                 Ok(SessionEvent::Stream(event))
@@ -1123,6 +1152,7 @@ impl DoshTransport {
             }
             TransportEvent::Open(_)
             | TransportEvent::OpenReject { .. }
+            | TransportEvent::Eof { .. }
             | TransportEvent::Close { .. }
             | TransportEvent::Ignored { .. } => {}
         }
@@ -1724,6 +1754,26 @@ mod tests {
         assert!(mux.send_data(99, b"lost".to_vec()).is_err());
     }
 
+    #[test]
+    fn stream_eof_is_delivered_without_closing_stream_state() {
+        let mut mux = StreamMux::new(TransportConfig::default());
+        mux.open_stream(1, "@dosh-test", 0).unwrap();
+        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
+            .unwrap()
+            .unwrap();
+
+        let eof = mux.eof_stream(1).unwrap();
+        assert_eq!(eof.kind, PacketKind::StreamEof);
+        assert!(mux.is_open(1));
+
+        let body = protocol::to_body(&StreamEof { stream_id: 1 }).unwrap();
+        assert_eq!(
+            mux.handle_packet(PacketKind::StreamEof, &body).unwrap(),
+            TransportEvent::Eof { stream_id: 1 }
+        );
+        assert!(mux.is_open(1));
+    }
+
     #[test]
     fn rejecting_incoming_stream_retires_local_state() {
         let mut mux = StreamMux::new(TransportConfig::default());