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

Add reliable stream retransmission

Browse Source
This commit is contained in:
DuProcess
2026-06-19 16:00:51 -04:00
parent d0d6f59cdf
commit 90e53f4b68
8 changed files with 618 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bin/dosh-client.rs
+200 -13
View File
@@ -179,6 +179,14 @@ enum ForwardEvent {
},
}
#[derive(Clone)]
struct PendingStreamChunk {
offset: u64,
bytes: Vec<u8>,
last_sent: Instant,
attempts: u32,
}
#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
let args = Args::parse();
@@ -2135,6 +2143,7 @@ async fn run_terminal(
let mut last_packet_at = Instant::now();
let mut status_tick = tokio::time::interval(Duration::from_secs(1));
let mut resize_tick = tokio::time::interval(Duration::from_millis(250));
let mut stream_retransmit_tick = tokio::time::interval(Duration::from_millis(200));
let mut last_size = terminal_size();
// React to terminal resize the instant it happens via SIGWINCH (mosh-style),
// instead of waiting up to one `resize_tick`. The 250ms poll below stays as a
@@ -2174,6 +2183,10 @@ async fn run_terminal(
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();
let mut stream_next_send_offset: HashMap<u64, u64> = HashMap::new();
let mut stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>> = HashMap::new();
let mut stream_next_recv_offset: HashMap<u64, u64> = HashMap::new();
let mut stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>> = HashMap::new();
if let Some(frame) = first_frame {
if !forward_only {
render_frame(&frame)?;
@@ -2398,6 +2411,8 @@ async fn run_terminal(
agent_writers.insert(open.stream_id, writer);
opened_streams.insert(open.stream_id);
stream_send_credit.insert(open.stream_id, STREAM_INITIAL_WINDOW);
stream_next_send_offset.entry(open.stream_id).or_insert(0);
stream_next_recv_offset.entry(open.stream_id).or_insert(0);
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 {
@@ -2440,6 +2455,8 @@ async fn run_terminal(
stream_writers.insert(open.stream_id, writer);
opened_streams.insert(open.stream_id);
stream_send_credit.insert(open.stream_id, STREAM_INITIAL_WINDOW);
stream_next_send_offset.entry(open.stream_id).or_insert(0);
stream_next_recv_offset.entry(open.stream_id).or_insert(0);
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 {
@@ -2490,6 +2507,8 @@ async fn run_terminal(
last_packet_at = Instant::now();
opened_streams.insert(ok.stream_id);
stream_send_credit.entry(ok.stream_id).or_insert(STREAM_INITIAL_WINDOW);
stream_next_send_offset.entry(ok.stream_id).or_insert(0);
stream_next_recv_offset.entry(ok.stream_id).or_insert(0);
if pending_socks_replies.remove(&ok.stream_id)
&& let Some(writer) = stream_writers.get_mut(&ok.stream_id)
{
@@ -2503,6 +2522,8 @@ async fn run_terminal(
ok.stream_id,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_next_send_offset,
&mut stream_sent_data,
).await?;
}
PacketKind::StreamOpenReject => {
@@ -2521,6 +2542,10 @@ async fn run_terminal(
opened_streams.remove(&reject.stream_id);
stream_send_credit.remove(&reject.stream_id);
stream_pending_data.remove(&reject.stream_id);
stream_next_send_offset.remove(&reject.stream_id);
stream_sent_data.remove(&reject.stream_id);
stream_next_recv_offset.remove(&reject.stream_id);
stream_recv_pending.remove(&reject.stream_id);
stream_writers.remove(&reject.stream_id);
}
PacketKind::StreamData => {
@@ -2531,11 +2556,17 @@ async fn run_terminal(
continue;
};
last_packet_at = Instant::now();
let len = data.bytes.len();
if let Some(writer) = stream_writers.get_mut(&data.stream_id) {
let _ = writer.write_all(&data.bytes).await;
} else if let Some(writer) = agent_writers.get_mut(&data.stream_id) {
let _ = writer.write_all(&data.bytes).await;
let stream_id = data.stream_id;
let (writes, consumed, received_offset) =
accept_stream_data(&mut stream_next_recv_offset, &mut stream_recv_pending, data);
if let Some(writer) = stream_writers.get_mut(&stream_id) {
for bytes in &writes {
let _ = writer.write_all(bytes).await;
}
} else if let Some(writer) = agent_writers.get_mut(&stream_id) {
for bytes in &writes {
let _ = writer.write_all(bytes).await;
}
}
// Always return flow-control credit so a stream whose local
// writer is already gone cannot wedge the server's send window.
@@ -2544,8 +2575,9 @@ async fn run_terminal(
addr,
&cred,
&mut send_seq,
data.stream_id,
len,
stream_id,
consumed,
received_offset,
).await?;
}
PacketKind::StreamWindowAdjust => {
@@ -2556,7 +2588,12 @@ async fn run_terminal(
continue;
};
last_packet_at = Instant::now();
add_stream_credit(&mut stream_send_credit, adjust.stream_id, adjust.bytes);
ack_stream_data(
&mut stream_sent_data,
&mut stream_send_credit,
adjust.stream_id,
adjust.received_offset,
);
flush_stream_pending_data(
&socket,
addr,
@@ -2565,6 +2602,8 @@ async fn run_terminal(
adjust.stream_id,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_next_send_offset,
&mut stream_sent_data,
).await?;
}
PacketKind::StreamClose => {
@@ -2579,6 +2618,10 @@ async fn run_terminal(
opened_streams.remove(&close.stream_id);
stream_send_credit.remove(&close.stream_id);
stream_pending_data.remove(&close.stream_id);
stream_next_send_offset.remove(&close.stream_id);
stream_sent_data.remove(&close.stream_id);
stream_next_recv_offset.remove(&close.stream_id);
stream_recv_pending.remove(&close.stream_id);
stream_writers.remove(&close.stream_id);
agent_writers.remove(&close.stream_id);
}
@@ -2590,6 +2633,8 @@ async fn run_terminal(
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);
stream_next_send_offset.entry(stream_id).or_insert(0);
stream_next_recv_offset.entry(stream_id).or_insert(0);
if socks5_reply {
pending_socks_replies.insert(stream_id);
}
@@ -2615,6 +2660,8 @@ async fn run_terminal(
&opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_next_send_offset,
&mut stream_sent_data,
).await?;
}
Some(ForwardEvent::Close { stream_id }) => {
@@ -2622,6 +2669,10 @@ async fn run_terminal(
opened_streams.remove(&stream_id);
stream_send_credit.remove(&stream_id);
stream_pending_data.remove(&stream_id);
stream_next_send_offset.remove(&stream_id);
stream_sent_data.remove(&stream_id);
stream_next_recv_offset.remove(&stream_id);
stream_recv_pending.remove(&stream_id);
stream_writers.remove(&stream_id);
agent_writers.remove(&stream_id);
send_stream_close(&socket, addr, &cred, &mut send_seq, stream_id).await?;
@@ -2676,6 +2727,15 @@ async fn run_terminal(
disconnect_status.apply(action)?;
}
}
_ = stream_retransmit_tick.tick() => {
retransmit_stream_data(
&socket,
addr,
&cred,
&mut send_seq,
&mut stream_sent_data,
).await?;
}
}
}
// Leave the terminal clean on exit: erase any lingering status line.
@@ -2986,9 +3046,14 @@ async fn send_stream_data(
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
offset: u64,
bytes: Vec<u8>,
) -> Result<()> {
let body = protocol::to_body(&StreamData { stream_id, bytes })?;
let body = protocol::to_body(&StreamData {
stream_id,
offset,
bytes,
})?;
send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamData, &body).await
}
@@ -3002,6 +3067,8 @@ async fn queue_or_send_stream_data(
opened_streams: &HashSet<u64>,
stream_send_credit: &mut HashMap<u64, usize>,
stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
stream_next_send_offset: &mut HashMap<u64, u64>,
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
) -> Result<()> {
if !opened_streams.contains(&stream_id)
|| stream_send_credit.get(&stream_id).copied().unwrap_or(0) < bytes.len()
@@ -3016,7 +3083,18 @@ async fn queue_or_send_stream_data(
return Ok(());
}
*stream_send_credit.entry(stream_id).or_default() -= bytes.len();
send_stream_data(socket, addr, cred, send_seq, stream_id, bytes).await
let offset = *stream_next_send_offset.entry(stream_id).or_default();
stream_next_send_offset.insert(stream_id, offset.saturating_add(bytes.len() as u64));
stream_sent_data.entry(stream_id).or_default().insert(
offset,
PendingStreamChunk {
offset,
bytes: bytes.clone(),
last_sent: Instant::now(),
attempts: 1,
},
);
send_stream_data(socket, addr, cred, send_seq, stream_id, offset, bytes).await
}
async fn flush_stream_pending_data(
@@ -3027,6 +3105,8 @@ async fn flush_stream_pending_data(
stream_id: u64,
stream_send_credit: &mut HashMap<u64, usize>,
stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
stream_next_send_offset: &mut HashMap<u64, u64>,
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
) -> Result<()> {
let Some(pending) = stream_pending_data.get_mut(&stream_id) else {
return Ok(());
@@ -3041,7 +3121,18 @@ async fn flush_stream_pending_data(
}
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?;
let offset = *stream_next_send_offset.entry(stream_id).or_default();
stream_next_send_offset.insert(stream_id, offset.saturating_add(bytes.len() as u64));
stream_sent_data.entry(stream_id).or_default().insert(
offset,
PendingStreamChunk {
offset,
bytes: bytes.clone(),
last_sent: Instant::now(),
attempts: 1,
},
);
send_stream_data(socket, addr, cred, send_seq, stream_id, offset, bytes).await?;
}
if pending.is_empty() {
stream_pending_data.remove(&stream_id);
@@ -3049,9 +3140,103 @@ async fn flush_stream_pending_data(
Ok(())
}
fn add_stream_credit(stream_send_credit: &mut HashMap<u64, usize>, stream_id: u64, bytes: u32) {
async fn retransmit_stream_data(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
) -> Result<()> {
let now = Instant::now();
let mut retransmit = Vec::new();
for (stream_id, chunks) in stream_sent_data.iter_mut() {
for chunk in chunks.values_mut() {
if now.duration_since(chunk.last_sent) < Duration::from_millis(200) {
continue;
}
chunk.last_sent = now;
chunk.attempts = chunk.attempts.saturating_add(1);
retransmit.push((*stream_id, chunk.offset, chunk.bytes.clone()));
}
}
for (stream_id, offset, bytes) in retransmit {
send_stream_data(socket, addr, cred, send_seq, stream_id, offset, bytes).await?;
}
Ok(())
}
fn accept_stream_data(
stream_next_recv_offset: &mut HashMap<u64, u64>,
stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
data: StreamData,
) -> (Vec<Vec<u8>>, usize, u64) {
let stream_id = data.stream_id;
let expected = stream_next_recv_offset.entry(stream_id).or_insert(0);
if data.offset < *expected {
return (Vec::new(), 0, *expected);
}
if data.offset > *expected {
stream_recv_pending
.entry(stream_id)
.or_default()
.entry(data.offset)
.or_insert(data.bytes);
return (Vec::new(), 0, *expected);
}
let mut writes = vec![data.bytes];
let mut consumed = writes[0].len();
*expected = expected.saturating_add(consumed as u64);
while let Some(bytes) = stream_recv_pending
.entry(stream_id)
.or_default()
.remove(expected)
{
consumed = consumed.saturating_add(bytes.len());
*expected = expected.saturating_add(bytes.len() as u64);
writes.push(bytes);
}
if stream_recv_pending
.get(&stream_id)
.is_some_and(BTreeMap::is_empty)
{
stream_recv_pending.remove(&stream_id);
}
(writes, consumed, *expected)
}
fn ack_stream_data(
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
stream_send_credit: &mut HashMap<u64, usize>,
stream_id: u64,
received_offset: u64,
) {
let Some(sent) = stream_sent_data.get_mut(&stream_id) else {
return;
};
let acked_offsets: Vec<u64> = sent
.iter()
.filter_map(|(offset, chunk)| {
let end = chunk.offset.saturating_add(chunk.bytes.len() as u64);
(end <= received_offset).then_some(*offset)
})
.collect();
let mut acked_bytes = 0usize;
for offset in acked_offsets {
if let Some(chunk) = sent.remove(&offset) {
acked_bytes = acked_bytes.saturating_add(chunk.bytes.len());
}
}
if sent.is_empty() {
stream_sent_data.remove(&stream_id);
}
add_stream_credit(stream_send_credit, stream_id, acked_bytes);
}
fn add_stream_credit(stream_send_credit: &mut HashMap<u64, usize>, stream_id: u64, bytes: usize) {
let credit = stream_send_credit.entry(stream_id).or_default();
*credit = credit.saturating_add(bytes as usize);
*credit = credit.saturating_add(bytes).min(STREAM_INITIAL_WINDOW);
}
async fn send_stream_window_adjust(
@@ -3061,9 +3246,11 @@ async fn send_stream_window_adjust(
send_seq: &mut u64,
stream_id: u64,
bytes: usize,
received_offset: u64,
) -> Result<()> {
let body = protocol::to_body(&StreamWindowAdjust {
stream_id,
received_offset,
bytes: bytes.min(u32::MAX as usize) as u32,
})?;
send_stream_packet(