Add stream window flow control
ci / test (push) Has been cancelled
ci / remote-bench (push) Has been cancelled

This commit is contained in:
DuProcess
2026-06-14 09:08:19 -04:00
parent dd9c4daeda
commit fb5c25ea0b
2 changed files with 476 additions and 25 deletions
+154 -3
View File
@@ -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,
+322 -22
View File
@@ -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");
}
}