Add stream window flow control

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<vt100::Screen>,
     allowed_forwardings: Vec<ForwardingRequest>,
     stream_writers: HashMap<u64, mpsc::Sender<Vec<u8>>>,
+    opened_streams: HashSet<u64>,
+    stream_send_credit: HashMap<u64, usize>,
+    stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>>,
 }
 
 #[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<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
     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<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
     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<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    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<Mutex<ServerState>>,
     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<u8>,
 ) -> 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<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    client_id: [u8; 16],
+    stream_id: u64,
+    bytes: Vec<u8>,
+) -> 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<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    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<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_to_client(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    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");
+    }
 }