Bound forwarded stream buffering
ci / test (push) Has been cancelled
ci / remote-bench (push) Has been cancelled

This commit is contained in:
DuProcess
2026-06-13 15:43:42 -04:00
parent 6a0f12c59e
commit 64b00783fd
2 changed files with 19 additions and 14 deletions
+11 -6
View File
@@ -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 {
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;
});
}
});
+4 -4
View File
@@ -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)