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(
src/bin/dosh-server.rs
+216 -17
View File
@@ -20,7 +20,7 @@ use dosh::protocol::{
TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
};
use dosh::pty::{PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session};
use std::collections::{HashMap, HashSet, VecDeque};
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::net::{IpAddr, SocketAddr};
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::UnixStream as StdUnixStream;
@@ -375,6 +375,10 @@ struct ClientState {
opened_streams: HashSet<u64>,
stream_send_credit: HashMap<u64, usize>,
stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>>,
stream_next_send_offset: HashMap<u64, u64>,
stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
stream_next_recv_offset: HashMap<u64, u64>,
stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
/// Current transport key epoch (0 = original handshake key). Bumped on rekey.
epoch: u64,
/// When the current epoch began, for the wall-clock rekey trigger.
@@ -394,6 +398,14 @@ struct PendingFrame {
attempts: u8,
}
#[derive(Clone)]
struct PendingStreamChunk {
offset: u64,
bytes: Vec<u8>,
last_sent: Instant,
attempts: u32,
}
struct PendingNativeAuth {
client: dosh::native::NativeClientHello,
server: NativeServerHello,
@@ -1115,6 +1127,10 @@ async fn handle_native_user_auth(
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
stream_recv_pending: HashMap::new(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
@@ -1260,6 +1276,10 @@ async fn handle_bootstrap_attach(
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
stream_recv_pending: HashMap::new(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
@@ -1379,6 +1399,10 @@ async fn handle_ticket_attach(
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
stream_recv_pending: HashMap::new(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
@@ -1745,6 +1769,14 @@ async fn handle_stream_open(
client
.stream_send_credit
.insert(open.stream_id, STREAM_INITIAL_WINDOW);
client
.stream_next_send_offset
.entry(open.stream_id)
.or_insert(0);
client
.stream_next_recv_offset
.entry(open.stream_id)
.or_insert(0);
}
tokio::spawn(async move {
@@ -1798,7 +1830,8 @@ async fn handle_stream_data(
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let data: StreamData = protocol::from_body(&body)?;
let writer = {
let stream_id = data.stream_id;
let (writer, writes, consumed, received_offset) = {
let mut locked = state.lock().expect("server state poisoned");
let session = locked
.sessions
@@ -1813,17 +1846,27 @@ async fn handle_stream_data(
}
client.endpoint = peer;
client.last_seen = Instant::now();
client.stream_writers.get(&data.stream_id).cloned()
let writer = client.stream_writers.get(&data.stream_id).cloned();
let (writes, consumed, received_offset) = accept_stream_data(client, data);
(writer, writes, consumed, received_offset)
};
let len = data.bytes.len();
if let Some(writer) = writer {
let _ = writer.send(data.bytes).await;
for bytes in writes {
let _ = writer.send(bytes).await;
}
}
// Always return flow-control credit, even when the stream's writer is already
// gone (closed/unknown). The peer debited its send window for these bytes, so
// skipping the adjust would wedge the stream at a permanent credit deficit.
send_stream_window_adjust_to_client(state, socket, packet.header.conn_id, data.stream_id, len)
.await?;
send_stream_window_adjust_to_client(
state,
socket,
packet.header.conn_id,
stream_id,
consumed,
received_offset,
)
.await?;
Ok(())
}
@@ -1913,11 +1956,7 @@ async fn handle_stream_window_adjust(
}
client.endpoint = peer;
client.last_seen = Instant::now();
add_stream_credit(
&mut client.stream_send_credit,
adjust.stream_id,
adjust.bytes,
);
ack_stream_data(client, adjust.stream_id, adjust.received_offset);
}
flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, adjust.stream_id)
.await
@@ -1948,6 +1987,10 @@ async fn handle_stream_close(
client.opened_streams.remove(&close.stream_id);
client.stream_send_credit.remove(&close.stream_id);
client.stream_pending_data.remove(&close.stream_id);
client.stream_next_send_offset.remove(&close.stream_id);
client.stream_sent_data.remove(&close.stream_id);
client.stream_next_recv_offset.remove(&close.stream_id);
client.stream_recv_pending.remove(&close.stream_id);
Ok(())
}
@@ -2312,8 +2355,16 @@ async fn queue_or_send_stream_data_to_client(
return Ok(());
}
*client.stream_send_credit.entry(stream_id).or_default() -= bytes.len();
let offset = *client.stream_next_send_offset.entry(stream_id).or_default();
client
.stream_next_send_offset
.insert(stream_id, offset.saturating_add(bytes.len() as u64));
client.send_seq += 1;
let body = protocol::to_body(&StreamData { stream_id, bytes })?;
let body = protocol::to_body(&StreamData {
stream_id,
offset,
bytes: bytes.clone(),
})?;
let packet = protocol::encode_encrypted(
PacketKind::StreamData,
client_id,
@@ -2323,6 +2374,19 @@ async fn queue_or_send_stream_data_to_client(
SERVER_TO_CLIENT,
&body,
)?;
client
.stream_sent_data
.entry(stream_id)
.or_default()
.insert(
offset,
PendingStreamChunk {
offset,
bytes,
last_sent: Instant::now(),
attempts: 1,
},
);
send = Some((client.endpoint, packet));
}
send
@@ -2364,8 +2428,16 @@ async fn flush_stream_pending_data_to_client(
client.stream_pending_data.remove(&stream_id);
}
*client.stream_send_credit.entry(stream_id).or_default() -= bytes.len();
let offset = *client.stream_next_send_offset.entry(stream_id).or_default();
client
.stream_next_send_offset
.insert(stream_id, offset.saturating_add(bytes.len() as u64));
client.send_seq += 1;
let body = protocol::to_body(&StreamData { stream_id, bytes })?;
let body = protocol::to_body(&StreamData {
stream_id,
offset,
bytes: bytes.clone(),
})?;
let packet = protocol::encode_encrypted(
PacketKind::StreamData,
client_id,
@@ -2375,6 +2447,19 @@ async fn flush_stream_pending_data_to_client(
SERVER_TO_CLIENT,
&body,
)?;
client
.stream_sent_data
.entry(stream_id)
.or_default()
.insert(
offset,
PendingStreamChunk {
offset,
bytes,
last_sent: Instant::now(),
attempts: 1,
},
);
send = Some((client.endpoint, packet));
}
send
@@ -2386,9 +2471,72 @@ async fn flush_stream_pending_data_to_client(
}
}
fn add_stream_credit(stream_send_credit: &mut HashMap<u64, usize>, stream_id: u64, bytes: u32) {
fn accept_stream_data(client: &mut ClientState, data: StreamData) -> (Vec<Vec<u8>>, usize, u64) {
let stream_id = data.stream_id;
let expected = client.stream_next_recv_offset.entry(stream_id).or_insert(0);
if data.offset < *expected {
return (Vec::new(), 0, *expected);
}
if data.offset > *expected {
client
.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) = client
.stream_recv_pending
.entry(stream_id)
.or_default()
.remove(expected)
{
consumed += bytes.len();
*expected = expected.saturating_add(bytes.len() as u64);
writes.push(bytes);
}
if client
.stream_recv_pending
.get(&stream_id)
.is_some_and(BTreeMap::is_empty)
{
client.stream_recv_pending.remove(&stream_id);
}
(writes, consumed, *expected)
}
fn ack_stream_data(client: &mut ClientState, stream_id: u64, received_offset: u64) {
let Some(sent) = client.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() {
client.stream_sent_data.remove(&stream_id);
}
add_stream_credit(&mut client.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_to_client(
@@ -2397,9 +2545,11 @@ async fn send_stream_window_adjust_to_client(
client_id: [u8; 16],
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_to_client(
@@ -2425,6 +2575,10 @@ async fn send_stream_close_to_client(
client.opened_streams.remove(&stream_id);
client.stream_send_credit.remove(&stream_id);
client.stream_pending_data.remove(&stream_id);
client.stream_next_send_offset.remove(&stream_id);
client.stream_sent_data.remove(&stream_id);
client.stream_next_recv_offset.remove(&stream_id);
client.stream_recv_pending.remove(&stream_id);
}
}
let body = protocol::to_body(&StreamClose { stream_id })?;
@@ -2642,7 +2796,7 @@ async fn retransmit_pending(
let now = Instant::now();
let mut sends = Vec::new();
for session in locked.sessions.values_mut() {
for client in session.clients.values_mut() {
for (client_id, client) in session.clients.iter_mut() {
for pending in client.pending.iter_mut() {
if pending.output_seq <= client.last_acked {
continue;
@@ -2655,6 +2809,39 @@ async fn retransmit_pending(
sends.push((client.endpoint, pending.packet.clone()));
}
}
let stream_ids: Vec<u64> = client.stream_sent_data.keys().copied().collect();
let mut retransmit_chunks = Vec::new();
for stream_id in stream_ids {
let Some(chunks) = client.stream_sent_data.get_mut(&stream_id) else {
continue;
};
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_chunks.push((stream_id, chunk.offset, chunk.bytes.clone()));
}
}
for (stream_id, offset, bytes) in retransmit_chunks {
client.send_seq += 1;
let body = protocol::to_body(&StreamData {
stream_id,
offset,
bytes,
})?;
let packet = protocol::encode_encrypted(
PacketKind::StreamData,
*client_id,
client.send_seq,
client.last_acked,
&client.session_key,
SERVER_TO_CLIENT,
&body,
)?;
sends.push((client.endpoint, packet));
}
}
}
sends
@@ -3015,6 +3202,10 @@ mod tests {
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
stream_recv_pending: HashMap::new(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
@@ -3211,6 +3402,10 @@ mod tests {
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
stream_recv_pending: HashMap::new(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
@@ -3300,6 +3495,10 @@ mod tests {
opened_streams: HashSet::from([42]),
stream_send_credit: HashMap::from([(42, 0)]),
stream_pending_data: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
stream_recv_pending: HashMap::new(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
src/protocol.rs
+5 -1
View File
@@ -38,12 +38,14 @@ pub fn is_implicit_session_name(name: &str) -> bool {
}
pub const MAGIC: &[u8; 4] = b"DOSH";
// v4: added reliable stream offsets/acks to `StreamData` and
// `StreamWindowAdjust`.
// v3: added `ForwardingKind::Agent` (SSH-agent forwarding). The new variant rides
// inside `NativeUserAuth.requested_forwardings`, so a pre-agent peer would fail to
// deserialize it; bumping the wire version makes such a peer answer with a clear
// version-mismatch reject instead. Existing variants' bincode discriminants are
// unchanged, so the bump is purely a compatibility gate.
pub const VERSION: u8 = 3;
pub const VERSION: u8 = 4;
pub const HEADER_LEN: usize = 58;
/// Stable, user-facing reason string the server puts in an `AttachReject` when a
@@ -416,12 +418,14 @@ pub struct StreamOpenReject {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamData {
pub stream_id: u64,
pub offset: u64,
pub bytes: Vec<u8>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamWindowAdjust {
pub stream_id: u64,
pub received_offset: u64,
pub bytes: u32,
}