Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions 58 Packages Projects Releases 51 Wiki Activity

Support stream EOF half closes
ci / test (push) Waiting to run
Details
ci / fuzz-smoke (push) Waiting to run
Details
ci / windows-client (push) Waiting to run
Details
ci / package-release (linux-x86_64, ubuntu-latest) (push) Waiting to run
Details
ci / package-release (macos-aarch64, macos-14) (push) Waiting to run
Details
ci / package-release (macos-x86_64, macos-13) (push) Waiting to run
Details
ci / package-release (windows-x86_64, windows-latest) (push) Waiting to run
Details
ci / publish-gitea-release (push) Blocked by required conditions
Details
ci / remote-bench (push) Waiting to run
Details

Browse Source
This commit is contained in:
DuProcess
2026-07-13 00:20:52 -04:00
parent dccd648ac9
commit bf50b38ce2
3 changed files with 499 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bin/dosh-client.rs
+232 -14
View File
@@ -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<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 {
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::<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
+215 -10
View File
@@ -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<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 _ =
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<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
+52 -2
View File
@@ -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());