Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions 58 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-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,