Add stream window flow control

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<u64, tokio::net::tcp::OwnedWriteHalf> = HashMap::new();
     let mut pending_socks_replies: HashSet<u64> = HashSet::new();
+    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();
     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::<StreamWindowAdjust>(&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<u8>,
+    opened_streams: &HashSet<u64>,
+    stream_send_credit: &mut HashMap<u64, usize>,
+    stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
+) -> 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<u64, usize>,
+    stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
+) -> 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<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(
+    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,

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,19 +1281,30 @@ 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 ok: StreamOpenOk = protocol::from_body(&body)?;
+    {
         let mut locked = state.lock().expect("server state poisoned");
         let session = locked
             .sessions
@@ -1286,7 +1319,13 @@ async fn handle_stream_open_ok(
         }
         client.endpoint = peer;
         client.last_seen = Instant::now();
-    Ok(())
+        client.opened_streams.insert(ok.stream_id);
+        client
+            .stream_send_credit
+            .entry(ok.stream_id)
+            .or_insert(STREAM_INITIAL_WINDOW);
+    }
+    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<()> {
+    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 })?;
-    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamData, body).await
+                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");
+    }
 }