Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions 46 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();