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

Add stream window flow control
ci / test (push) Has been cancelled
Details
ci / remote-bench (push) Has been cancelled
Details

Browse Source
This commit is contained in:
DuProcess
2026-06-14 09:08:19 -04:00
parent dd9c4daeda
commit fb5c25ea0b
2 changed files with 476 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bin/dosh-client.rs
+154 -3
View File
@@ -18,14 +18,15 @@ use dosh::protocol::{
self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody,
NativeUserAuthBody, PacketKind, Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose,
StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, TicketAttachBody, TicketAttachEnvelope,
TicketAttachOkEnvelope,
StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, TicketAttachBody,
TicketAttachEnvelope, TicketAttachOkEnvelope,
};
use dosh::ssh_agent;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use std::fs;
use std::io::{Read, Write};
use std::net::{
@@ -41,6 +42,8 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use tokio::sync::mpsc;
const STREAM_INITIAL_WINDOW: usize = 1024 * 1024;
#[derive(Debug, Parser)]
#[command(name = "dosh-client")]
struct Args {
@@ -2085,6 +2088,9 @@ async fn run_terminal(
signal_background_ready();
let mut stream_writers: HashMap<u64, tokio::net::tcp::OwnedWriteHalf> = HashMap::new();
let mut pending_socks_replies: HashSet<u64> = HashSet::new();
let mut opened_streams: HashSet<u64> = HashSet::new();
let mut stream_send_credit: HashMap<u64, usize> = HashMap::new();
let mut stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>> = HashMap::new();
if let Some(frame) = first_frame {
if !forward_only {
render_frame(&frame)?;
@@ -2238,6 +2244,8 @@ async fn run_terminal(
Ok(stream) => {
let (mut reader, writer) = stream.into_split();
stream_writers.insert(open.stream_id, writer);
opened_streams.insert(open.stream_id);
stream_send_credit.insert(open.stream_id, STREAM_INITIAL_WINDOW);
send_stream_open_ok(&socket, addr, &cred, &mut send_seq, open.stream_id).await?;
let forward_tx = remote_forward_tx.clone();
tokio::spawn(async move {
@@ -2286,11 +2294,22 @@ async fn run_terminal(
continue;
};
last_packet_at = Instant::now();
opened_streams.insert(ok.stream_id);
stream_send_credit.entry(ok.stream_id).or_insert(STREAM_INITIAL_WINDOW);
if pending_socks_replies.remove(&ok.stream_id)
&& let Some(writer) = stream_writers.get_mut(&ok.stream_id)
{
let _ = writer.write_all(&[5, 0, 0, 1, 0, 0, 0, 0, 0, 0]).await;
}
flush_stream_pending_data(
&socket,
addr,
&cred,
&mut send_seq,
ok.stream_id,
&mut stream_send_credit,
&mut stream_pending_data,
).await?;
}
PacketKind::StreamOpenReject => {
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
@@ -2305,6 +2324,9 @@ async fn run_terminal(
{
let _ = writer.write_all(&[5, 5, 0, 1, 0, 0, 0, 0, 0, 0]).await;
}
opened_streams.remove(&reject.stream_id);
stream_send_credit.remove(&reject.stream_id);
stream_pending_data.remove(&reject.stream_id);
stream_writers.remove(&reject.stream_id);
}
PacketKind::StreamData => {
@@ -2317,8 +2339,35 @@ async fn run_terminal(
last_packet_at = Instant::now();
if let Some(writer) = stream_writers.get_mut(&data.stream_id) {
let _ = writer.write_all(&data.bytes).await;
send_stream_window_adjust(
&socket,
addr,
&cred,
&mut send_seq,
data.stream_id,
data.bytes.len(),
).await?;
}
}
PacketKind::StreamWindowAdjust => {
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
continue;
};
let Ok(adjust) = protocol::from_body::<StreamWindowAdjust>(&plain) else {
continue;
};
last_packet_at = Instant::now();
add_stream_credit(&mut stream_send_credit, adjust.stream_id, adjust.bytes);
flush_stream_pending_data(
&socket,
addr,
&cred,
&mut send_seq,
adjust.stream_id,
&mut stream_send_credit,
&mut stream_pending_data,
).await?;
}
PacketKind::StreamClose => {
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
continue;
@@ -2328,6 +2377,9 @@ async fn run_terminal(
};
last_packet_at = Instant::now();
pending_socks_replies.remove(&close.stream_id);
opened_streams.remove(&close.stream_id);
stream_send_credit.remove(&close.stream_id);
stream_pending_data.remove(&close.stream_id);
stream_writers.remove(&close.stream_id);
}
_ => {}
@@ -2337,6 +2389,7 @@ async fn run_terminal(
match forward_event {
Some(ForwardEvent::Open { stream_id, target_host, target_port, writer, socks5_reply }) => {
stream_writers.insert(stream_id, writer);
stream_send_credit.insert(stream_id, STREAM_INITIAL_WINDOW);
if socks5_reply {
pending_socks_replies.insert(stream_id);
}
@@ -2352,10 +2405,23 @@ async fn run_terminal(
.await?;
}
Some(ForwardEvent::Data { stream_id, bytes }) => {
send_stream_data(&socket, addr, &cred, &mut send_seq, stream_id, bytes).await?;
queue_or_send_stream_data(
&socket,
addr,
&cred,
&mut send_seq,
stream_id,
bytes,
&opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
).await?;
}
Some(ForwardEvent::Close { stream_id }) => {
pending_socks_replies.remove(&stream_id);
opened_streams.remove(&stream_id);
stream_send_credit.remove(&stream_id);
stream_pending_data.remove(&stream_id);
stream_writers.remove(&stream_id);
send_stream_close(&socket, addr, &cred, &mut send_seq, stream_id).await?;
}
@@ -2710,6 +2776,91 @@ async fn send_stream_data(
send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamData, &body).await
}
async fn queue_or_send_stream_data(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
bytes: Vec<u8>,
opened_streams: &HashSet<u64>,
stream_send_credit: &mut HashMap<u64, usize>,
stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
) -> Result<()> {
if !opened_streams.contains(&stream_id)
|| stream_send_credit.get(&stream_id).copied().unwrap_or(0) < bytes.len()
|| stream_pending_data
.get(&stream_id)
.is_some_and(|pending| !pending.is_empty())
{
stream_pending_data
.entry(stream_id)
.or_default()
.push_back(bytes);
return Ok(());
}
*stream_send_credit.entry(stream_id).or_default() -= bytes.len();
send_stream_data(socket, addr, cred, send_seq, stream_id, bytes).await
}
async fn flush_stream_pending_data(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
stream_send_credit: &mut HashMap<u64, usize>,
stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
) -> Result<()> {
let Some(pending) = stream_pending_data.get_mut(&stream_id) else {
return Ok(());
};
loop {
let Some(bytes) = pending.front() else {
break;
};
let credit = stream_send_credit.get(&stream_id).copied().unwrap_or(0);
if credit < bytes.len() {
break;
}
let bytes = pending.pop_front().expect("pending front exists");
*stream_send_credit.entry(stream_id).or_default() -= bytes.len();
send_stream_data(socket, addr, cred, send_seq, stream_id, bytes).await?;
}
if pending.is_empty() {
stream_pending_data.remove(&stream_id);
}
Ok(())
}
fn add_stream_credit(stream_send_credit: &mut HashMap<u64, usize>, stream_id: u64, bytes: u32) {
let credit = stream_send_credit.entry(stream_id).or_default();
*credit = credit.saturating_add(bytes as usize);
}
async fn send_stream_window_adjust(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
bytes: usize,
) -> Result<()> {
let body = protocol::to_body(&StreamWindowAdjust {
stream_id,
bytes: bytes.min(u32::MAX as usize) as u32,
})?;
send_stream_packet(
socket,
addr,
cred,
send_seq,
PacketKind::StreamWindowAdjust,
&body,
)
.await
}
async fn send_stream_close(
socket: &UdpSocket,
addr: SocketAddr,