From bf50b38ce25520f82fd273023a33722470929ebf Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Mon, 13 Jul 2026 00:20:52 -0400 Subject: [PATCH] Support stream EOF half closes --- src/bin/dosh-client.rs | 246 ++++++++++++++++++++++++++++++++++++++--- src/bin/dosh-server.rs | 225 +++++++++++++++++++++++++++++++++++-- src/transport.rs | 54 ++++++++- 3 files changed, 499 insertions(+), 26 deletions(-) diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index 4aa74dd..e2c66ac 100644 --- a/src/bin/dosh-client.rs +++ b/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, - TicketAttachEnvelope, TicketAttachOkEnvelope, + StreamData, StreamEof, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, + TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; use dosh::ssh_agent; use dosh::transport::{ @@ -218,6 +218,9 @@ enum ForwardEvent { stream_id: u64, bytes: Vec, }, + Eof { + stream_id: u64, + }, Close { stream_id: u64, }, @@ -5900,6 +5903,8 @@ async fn run_terminal( let mut stream_retransmit_srtt: Option = None; let mut stream_next_recv_offset: HashMap = HashMap::new(); let mut stream_recv_pending: HashMap>> = HashMap::new(); + let mut stream_local_eof_sent: HashSet = HashSet::new(); + let mut stream_remote_eof_received: HashSet = HashSet::new(); let mut stream_retired: HashSet = HashSet::new(); let mut stream_retired_order: VecDeque = VecDeque::new(); let mut pending_user_input: VecDeque = 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 { - stream_id: open.stream_id, + .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 { - stream_id: open.stream_id, - }).await; + let _ = forward_tx + .send(if graceful_eof { + ForwardEvent::Eof { + stream_id: open.stream_id, + } + } 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::(&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>, stream_next_recv_offset: &mut HashMap, stream_recv_pending: &mut HashMap>>, + stream_local_eof_sent: &mut HashSet, + stream_remote_eof_received: &mut HashSet, ) { 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, + stream_remote_eof_received: &HashSet, +) -> 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>, stream_next_recv_offset: &HashMap, stream_recv_pending: &HashMap>>, + stream_local_eof_sent: &HashSet, + stream_remote_eof_received: &HashSet, ) -> 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>, stream_next_recv_offset: &mut HashMap, stream_recv_pending: &mut HashMap>>, + stream_local_eof_sent: &mut HashSet, + stream_remote_eof_received: &mut HashSet, stream_retired: &mut HashSet, stream_retired_order: &mut VecDeque, ) { @@ -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(); diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 9f007bd..ae65627 100644 --- a/src/bin/dosh-server.rs +++ b/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, - TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, + StreamClose, StreamData, StreamEof, StreamOpen, StreamOpenOk, StreamOpenReject, + 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>, stream_next_recv_offset: HashMap, stream_recv_pending: HashMap>>, + stream_local_eof_sent: HashSet, + stream_remote_eof_received: HashSet, stream_retired: HashSet, stream_retired_order: VecDeque, stream_retransmit_srtt: Option, @@ -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>, + 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>, 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 _ = - 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 + }; }); } }); @@ -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>, + socket: &Arc, + 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>, socket: &Arc, @@ -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, diff --git a/src/transport.rs b/src/transport.rs index 6e6d1dc..0d96ab0 100644 --- a/src/transport.rs +++ b/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, }, + 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 { + 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());