diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index 746b11b..f9526e4 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -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 = HashMap::new(); let mut pending_socks_replies: HashSet = HashSet::new(); + let mut opened_streams: HashSet = HashSet::new(); + let mut stream_send_credit: HashMap = HashMap::new(); + let mut stream_pending_data: HashMap>> = 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::(&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, + opened_streams: &HashSet, + stream_send_credit: &mut HashMap, + stream_pending_data: &mut HashMap>>, +) -> 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, + stream_pending_data: &mut HashMap>>, +) -> 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, 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, diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 2f5f90e..8317585 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -15,11 +15,11 @@ use dosh::protocol::{ self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind, ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, - StreamClose, StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, TicketAttachBody, - TicketAttachEnvelope, TicketAttachOkEnvelope, + StreamClose, StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, + TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; use dosh::pty::{PtyHandle, PtyOutput, spawn_pty_session}; -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::net::SocketAddr; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -27,6 +27,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-server")] struct Args { @@ -191,6 +193,9 @@ struct ClientState { last_screen: Option, allowed_forwardings: Vec, stream_writers: HashMap>>, + opened_streams: HashSet, + stream_send_credit: HashMap, + stream_pending_data: HashMap>>, } #[derive(Clone)] @@ -384,9 +389,12 @@ async fn handle_packet( PacketKind::Ack => handle_ack(state, &packet).await, PacketKind::Detach => handle_detach(state, &packet).await, PacketKind::StreamOpen => handle_stream_open(state, socket, peer, &packet).await, - PacketKind::StreamOpenOk => handle_stream_open_ok(state, peer, &packet).await, + PacketKind::StreamOpenOk => handle_stream_open_ok(state, socket, peer, &packet).await, PacketKind::StreamOpenReject => handle_stream_open_reject(state, peer, &packet).await, - PacketKind::StreamData => handle_stream_data(state, peer, &packet).await, + PacketKind::StreamData => handle_stream_data(state, socket, peer, &packet).await, + PacketKind::StreamWindowAdjust => { + handle_stream_window_adjust(state, socket, peer, &packet).await + } PacketKind::StreamClose => handle_stream_close(state, &packet).await, _ => Ok(()), } @@ -623,6 +631,9 @@ async fn handle_native_user_auth( last_screen: Some(screen), allowed_forwardings: req.auth.requested_forwardings.clone(), stream_writers: HashMap::new(), + opened_streams: HashSet::new(), + stream_send_credit: HashMap::new(), + stream_pending_data: HashMap::new(), }, ); Ok(( @@ -750,6 +761,9 @@ async fn handle_bootstrap_attach( last_screen: Some(screen), allowed_forwardings: Vec::new(), stream_writers: HashMap::new(), + opened_streams: HashSet::new(), + stream_send_credit: HashMap::new(), + stream_pending_data: HashMap::new(), }, ); ( @@ -862,6 +876,9 @@ async fn handle_ticket_attach( last_screen: Some(screen), allowed_forwardings: Vec::new(), stream_writers: HashMap::new(), + opened_streams: HashSet::new(), + stream_send_credit: HashMap::new(), + stream_pending_data: HashMap::new(), }, ); (client_id, output_seq, snapshot) @@ -1189,6 +1206,10 @@ async fn handle_stream_open( .get_mut(&packet.header.conn_id) .ok_or_else(|| anyhow!("unknown client"))?; client.stream_writers.insert(open.stream_id, writer_tx); + client.opened_streams.insert(open.stream_id); + client + .stream_send_credit + .insert(open.stream_id, STREAM_INITIAL_WINDOW); } tokio::spawn(async move { @@ -1235,6 +1256,7 @@ async fn handle_stream_open( async fn handle_stream_data( state: &Arc>, + socket: &Arc, peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { @@ -1259,34 +1281,51 @@ async fn handle_stream_data( client.stream_writers.get(&data.stream_id).cloned() }; if let Some(writer) = writer { + let len = data.bytes.len(); let _ = writer.send(data.bytes).await; + send_stream_window_adjust_to_client( + state, + socket, + packet.header.conn_id, + data.stream_id, + len, + ) + .await?; } Ok(()) } async fn handle_stream_open_ok( state: &Arc>, + socket: &Arc, peer: SocketAddr, packet: &protocol::Packet, ) -> Result<()> { let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; - let _ok: StreamOpenOk = protocol::from_body(&body)?; - let mut locked = state.lock().expect("server state poisoned"); - let session = locked - .sessions - .get_mut(&session_name) - .ok_or_else(|| anyhow!("unknown session"))?; - let client = session - .clients - .get_mut(&packet.header.conn_id) - .ok_or_else(|| anyhow!("unknown client"))?; - if !client.replay.accept(packet.header.seq) { - return Ok(()); + let ok: StreamOpenOk = protocol::from_body(&body)?; + { + let mut locked = state.lock().expect("server state poisoned"); + let session = locked + .sessions + .get_mut(&session_name) + .ok_or_else(|| anyhow!("unknown session"))?; + let client = session + .clients + .get_mut(&packet.header.conn_id) + .ok_or_else(|| anyhow!("unknown client"))?; + if !client.replay.accept(packet.header.seq) { + return Ok(()); + } + client.endpoint = peer; + client.last_seen = Instant::now(); + client.opened_streams.insert(ok.stream_id); + client + .stream_send_credit + .entry(ok.stream_id) + .or_insert(STREAM_INITIAL_WINDOW); } - client.endpoint = peer; - client.last_seen = Instant::now(); - Ok(()) + flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, ok.stream_id).await } async fn handle_stream_open_reject( @@ -1312,9 +1351,46 @@ async fn handle_stream_open_reject( client.endpoint = peer; client.last_seen = Instant::now(); client.stream_writers.remove(&reject.stream_id); + client.opened_streams.remove(&reject.stream_id); + client.stream_send_credit.remove(&reject.stream_id); + client.stream_pending_data.remove(&reject.stream_id); Ok(()) } +async fn handle_stream_window_adjust( + state: &Arc>, + socket: &Arc, + peer: SocketAddr, + packet: &protocol::Packet, +) -> Result<()> { + let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; + let adjust: StreamWindowAdjust = protocol::from_body(&body)?; + { + let mut locked = state.lock().expect("server state poisoned"); + let session = locked + .sessions + .get_mut(&session_name) + .ok_or_else(|| anyhow!("unknown session"))?; + let client = session + .clients + .get_mut(&packet.header.conn_id) + .ok_or_else(|| anyhow!("unknown client"))?; + if !client.replay.accept(packet.header.seq) { + return Ok(()); + } + client.endpoint = peer; + client.last_seen = Instant::now(); + add_stream_credit( + &mut client.stream_send_credit, + adjust.stream_id, + adjust.bytes, + ); + } + flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, adjust.stream_id) + .await +} + async fn handle_stream_close( state: &Arc>, packet: &protocol::Packet, @@ -1334,6 +1410,9 @@ async fn handle_stream_close( } client.last_seen = Instant::now(); client.stream_writers.remove(&close.stream_id); + client.opened_streams.remove(&close.stream_id); + client.stream_send_credit.remove(&close.stream_id); + client.stream_pending_data.remove(&close.stream_id); Ok(()) } @@ -1397,6 +1476,9 @@ async fn start_remote_forwards( && let Some(client) = session.clients.get_mut(&client_id) { client.stream_writers.insert(stream_id, writer_tx); + client + .stream_send_credit + .insert(stream_id, STREAM_INITIAL_WINDOW); } else { break; } @@ -1535,8 +1617,144 @@ async fn send_stream_data_to_client( stream_id: u64, bytes: Vec, ) -> Result<()> { - let body = protocol::to_body(&StreamData { stream_id, bytes })?; - send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamData, body).await + queue_or_send_stream_data_to_client(state, socket, client_id, stream_id, bytes).await +} + +async fn queue_or_send_stream_data_to_client( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + bytes: Vec, +) -> Result<()> { + let send = { + let mut locked = state.lock().expect("server state poisoned"); + let mut send = None; + for session in locked.sessions.values_mut() { + if let Some(client) = session.clients.get_mut(&client_id) { + if !client.opened_streams.contains(&stream_id) + || client + .stream_send_credit + .get(&stream_id) + .copied() + .unwrap_or(0) + < bytes.len() + || client + .stream_pending_data + .get(&stream_id) + .is_some_and(|pending| !pending.is_empty()) + { + client + .stream_pending_data + .entry(stream_id) + .or_default() + .push_back(bytes); + return Ok(()); + } + *client.stream_send_credit.entry(stream_id).or_default() -= bytes.len(); + client.send_seq += 1; + let body = protocol::to_body(&StreamData { stream_id, bytes })?; + let packet = protocol::encode_encrypted( + PacketKind::StreamData, + client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + send = Some((client.endpoint, packet)); + break; + } + } + send + }; + if let Some((endpoint, packet)) = send { + socket.send_to(&packet, endpoint).await?; + } + Ok(()) +} + +async fn flush_stream_pending_data_to_client( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, +) -> Result<()> { + loop { + let send = { + let mut locked = state.lock().expect("server state poisoned"); + let mut send = None; + for session in locked.sessions.values_mut() { + if let Some(client) = session.clients.get_mut(&client_id) { + let Some(pending) = client.stream_pending_data.get_mut(&stream_id) else { + return Ok(()); + }; + let Some(bytes) = pending.front() else { + client.stream_pending_data.remove(&stream_id); + return Ok(()); + }; + let credit = client + .stream_send_credit + .get(&stream_id) + .copied() + .unwrap_or(0); + if credit < bytes.len() { + return Ok(()); + } + let bytes = pending.pop_front().expect("pending front exists"); + if pending.is_empty() { + client.stream_pending_data.remove(&stream_id); + } + *client.stream_send_credit.entry(stream_id).or_default() -= bytes.len(); + client.send_seq += 1; + let body = protocol::to_body(&StreamData { stream_id, bytes })?; + let packet = protocol::encode_encrypted( + PacketKind::StreamData, + client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + send = Some((client.endpoint, packet)); + break; + } + } + send + }; + let Some((endpoint, packet)) = send else { + return Ok(()); + }; + socket.send_to(&packet, endpoint).await?; + } +} + +fn add_stream_credit(stream_send_credit: &mut HashMap, 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_to_client( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + 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_to_client( + state, + socket, + client_id, + PacketKind::StreamWindowAdjust, + body, + ) + .await } async fn send_stream_close_to_client( @@ -1550,6 +1768,9 @@ async fn send_stream_close_to_client( for session in locked.sessions.values_mut() { if let Some(client) = session.clients.get_mut(&client_id) { client.stream_writers.remove(&stream_id); + client.opened_streams.remove(&stream_id); + client.stream_send_credit.remove(&stream_id); + client.stream_pending_data.remove(&stream_id); break; } } @@ -1851,4 +2072,83 @@ mod tests { .contains("invalid environment variable name") ); } + + #[tokio::test] + async fn stream_data_waits_for_open_and_credit() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + let client_id = [7u8; 16]; + let session_key = [9u8; 32]; + let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap()); + let endpoint = receiver.local_addr().unwrap(); + state.sessions.insert( + "test".to_string(), + Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + clients: HashMap::from([( + client_id, + ClientState { + endpoint, + mode: "forward-only".to_string(), + session_key, + last_acked: 0, + replay: ReplayWindow::default(), + send_seq: 1, + cols: 80, + rows: 24, + last_seen: Instant::now(), + pending: VecDeque::new(), + last_screen: None, + allowed_forwardings: Vec::new(), + stream_writers: HashMap::new(), + opened_streams: HashSet::new(), + stream_send_credit: HashMap::new(), + stream_pending_data: HashMap::new(), + }, + )]), + output_seq: 0, + recent: VecDeque::new(), + }, + ); + let state = Arc::new(Mutex::new(state)); + + queue_or_send_stream_data_to_client(&state, &sender, client_id, 42, b"hello".to_vec()) + .await + .unwrap(); + { + let locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert_eq!(client.stream_pending_data[&42].len(), 1); + } + + { + let mut locked = state.lock().unwrap(); + let client = locked + .sessions + .get_mut("test") + .unwrap() + .clients + .get_mut(&client_id) + .unwrap(); + client.opened_streams.insert(42); + client.stream_send_credit.insert(42, STREAM_INITIAL_WINDOW); + } + flush_stream_pending_data_to_client(&state, &sender, client_id, 42) + .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::StreamData); + let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); + let data: StreamData = protocol::from_body(&plain).unwrap(); + assert_eq!(data.stream_id, 42); + assert_eq!(data.bytes, b"hello"); + } }