From 64b00783fdcacef028db848e3b2eb038637d3003 Mon Sep 17 00:00:00 2001
From: DuProcess <273172371+DuProcess@users.noreply.github.com>
Date: Sat, 13 Jun 2026 15:43:42 -0400
Subject: [PATCH] Bound forwarded stream buffering

---
 src/bin/dosh-client.rs | 25 +++++++++++++++----------
 src/bin/dosh-server.rs |  8 ++++----
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs
index 89e4a2c..7e058cc 100644
--- a/src/bin/dosh-client.rs
+++ b/src/bin/dosh-client.rs
@@ -1449,7 +1449,7 @@ async fn run_terminal(
     let mut last_size = size().unwrap_or((80, 24));
     let mut frame_buffer = FrameBuffer::default();
     let mut predictor = Predictor::new(predict && cred.mode != "view-only" && !forward_only);
-    let (forward_tx, mut forward_rx) = mpsc::unbounded_channel::<ForwardEvent>();
+    let (forward_tx, mut forward_rx) = mpsc::channel::<ForwardEvent>(1024);
     let _forward_keepalive = if local_forwards.is_empty() {
         Some(forward_tx.clone())
     } else {
@@ -1502,6 +1502,7 @@ async fn run_terminal(
     let mut recv_buf = vec![0u8; 65535];
     loop {
         tokio::select! {
+            biased;
             stdin_msg = stdin_rx.recv() => {
                 match stdin_msg {
                     Some(bytes) => {
@@ -1624,6 +1625,7 @@ async fn run_terminal(
                                                         stream_id: open.stream_id,
                                                         bytes: buf[..n].to_vec(),
                                                     })
+                                                    .await
                                                     .is_err()
                                                 {
                                                     return;
@@ -1634,7 +1636,7 @@ async fn run_terminal(
                                     }
                                     let _ = forward_tx.send(ForwardEvent::Close {
                                         stream_id: open.stream_id,
-                                    });
+                                    }).await;
                                 });
                             }
                             Err(err) => {
@@ -1761,7 +1763,7 @@ async fn run_terminal(
 
 async fn start_local_forwards(
     local_forwards: &[LocalForward],
-    forward_tx: mpsc::UnboundedSender<ForwardEvent>,
+    forward_tx: mpsc::Sender<ForwardEvent>,
     stream_ids: Arc<AtomicU64>,
 ) -> Result<()> {
     for forward in local_forwards {
@@ -1779,12 +1781,14 @@ async fn start_local_forwards(
                 };
                 let stream_id = stream_ids.fetch_add(1, Ordering::Relaxed);
                 let (mut reader, writer) = stream.into_split();
-                let _ = forward_tx.send(ForwardEvent::Open {
-                    stream_id,
-                    target_host: forward.target_host.clone(),
-                    target_port: forward.target_port,
-                    writer,
-                });
+                let _ = forward_tx
+                    .send(ForwardEvent::Open {
+                        stream_id,
+                        target_host: forward.target_host.clone(),
+                        target_port: forward.target_port,
+                        writer,
+                    })
+                    .await;
                 let forward_tx = forward_tx.clone();
                 tokio::spawn(async move {
                     let mut buf = [0u8; 16 * 1024];
@@ -1797,6 +1801,7 @@ async fn start_local_forwards(
                                         stream_id,
                                         bytes: buf[..n].to_vec(),
                                     })
+                                    .await
                                     .is_err()
                                 {
                                     return;
@@ -1805,7 +1810,7 @@ async fn start_local_forwards(
                             Err(_) => break,
                         }
                     }
-                    let _ = forward_tx.send(ForwardEvent::Close { stream_id });
+                    let _ = forward_tx.send(ForwardEvent::Close { stream_id }).await;
                 });
             }
         });
diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs
index 581b93f..a9075ed 100644
--- a/src/bin/dosh-server.rs
+++ b/src/bin/dosh-server.rs
@@ -189,7 +189,7 @@ struct ClientState {
     pending: VecDeque<PendingFrame>,
     last_screen: Option<vt100::Screen>,
     allowed_forwardings: Vec<ForwardingRequest>,
-    stream_writers: HashMap<u64, mpsc::UnboundedSender<Vec<u8>>>,
+    stream_writers: HashMap<u64, mpsc::Sender<Vec<u8>>>,
 }
 
 #[derive(Clone)]
@@ -1014,7 +1014,7 @@ async fn handle_stream_open(
         }
     };
     let (mut reader, mut writer) = stream.into_split();
-    let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::<Vec<u8>>();
+    let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
     {
         let mut locked = state.lock().expect("server state poisoned");
         let session = locked
@@ -1096,7 +1096,7 @@ async fn handle_stream_data(
         client.stream_writers.get(&data.stream_id).cloned()
     };
     if let Some(writer) = writer {
-        let _ = writer.send(data.bytes);
+        let _ = writer.send(data.bytes).await;
     }
     Ok(())
 }
@@ -1227,7 +1227,7 @@ async fn start_remote_forwards(
                 };
                 let stream_id = allocate_server_stream_id(&state);
                 let (mut reader, mut writer) = stream.into_split();
-                let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::<Vec<u8>>();
+                let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
                 {
                     let mut locked = state.lock().expect("server state poisoned");
                     if let Some(session) = locked.sessions.get_mut(&session_name)