Avoid PTY allocation for forward-only mode
ci / test (push) Has been cancelled
ci / remote-bench (push) Has been cancelled

This commit is contained in:
DuProcess
2026-06-13 22:18:15 -04:00
parent 7fec3592f5
commit 76e2e980f1
2 changed files with 105 additions and 44 deletions
+5 -3
View File
@@ -186,7 +186,9 @@ async fn main() -> Result<()> {
!local_forwards.is_empty() || !remote_forwards.is_empty() || !dynamic_forwards.is_empty();
let predict = args.predict || host.predict.unwrap_or(config.predict);
let session = select_session(args.session.as_deref(), args.new);
let mode = if args.view_only || config.view_only {
let mode = if args.forward_only {
"forward-only"
} else if args.view_only || config.view_only {
"view-only"
} else {
"read-write"
@@ -1945,7 +1947,7 @@ async fn run_terminal(
if bytes == [0x1d] {
break;
}
if cred.mode != "view-only" {
if cred.mode != "view-only" && cred.mode != "forward-only" {
predictor.observe_input(&bytes)?;
send_input(&socket, addr, &cred, &mut send_seq, bytes).await?;
}
@@ -1958,7 +1960,7 @@ async fn run_terminal(
&& (cols, rows) != last_size
{
last_size = (cols, rows);
if cred.mode != "view-only" {
if cred.mode != "view-only" && cred.mode != "forward-only" {
send_resize(&socket, addr, &cred, &mut send_seq, cols, rows).await?;
}
}
+100 -41
View File
@@ -115,7 +115,7 @@ async fn serve(config_path: Option<std::path::PathBuf>) -> Result<()> {
{
let mut locked = state.lock().expect("server state poisoned");
for session in config.prewarm_sessions.clone() {
locked.ensure_session(&session, 80, 24)?;
locked.ensure_session(&session, 80, 24, "read-write")?;
}
}
@@ -169,7 +169,7 @@ struct ServerState {
}
struct Session {
pty: PtyHandle,
pty: Option<PtyHandle>,
parser: vt100::Parser,
clients: HashMap<[u8; 16], ClientState>,
output_seq: u64,
@@ -225,17 +225,30 @@ impl ServerState {
}
}
fn ensure_session(&mut self, name: &str, cols: u16, rows: u16) -> Result<()> {
if self.sessions.contains_key(name) {
fn ensure_session(&mut self, name: &str, cols: u16, rows: u16, mode: &str) -> Result<()> {
if let Some(session) = self.sessions.get_mut(name) {
if session.pty.is_none() && mode_uses_pty(mode) {
session.pty = Some(spawn_pty_session(
name.to_string(),
&self.config.shell,
cols.max(1),
rows.max(1),
self.pty_tx.clone(),
)?);
}
return Ok(());
}
let pty = spawn_pty_session(
name.to_string(),
&self.config.shell,
cols.max(1),
rows.max(1),
self.pty_tx.clone(),
)?;
let pty = if mode_uses_pty(mode) {
Some(spawn_pty_session(
name.to_string(),
&self.config.shell,
cols.max(1),
rows.max(1),
self.pty_tx.clone(),
)?)
} else {
None
};
self.sessions.insert(
name.to_string(),
Session {
@@ -250,6 +263,14 @@ impl ServerState {
}
}
fn mode_uses_pty(mode: &str) -> bool {
mode != "forward-only"
}
fn mode_allows_terminal_updates(mode: &str) -> bool {
mode != "view-only" && mode != "forward-only"
}
async fn handle_packet(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
@@ -471,13 +492,10 @@ async fn handle_native_user_auth(
let mode = pending.client.requested_mode.clone();
let cols = pending.client.terminal_size.0;
let rows = pending.client.terminal_size.1;
if !locked.sessions.contains_key(&session_name) {
if locked.config.create_on_attach {
locked.ensure_session(&session_name, cols, rows)?;
} else {
return Err(anyhow!("session does not exist"));
}
if !locked.sessions.contains_key(&session_name) && !locked.config.create_on_attach {
return Err(anyhow!("session does not exist"));
}
locked.ensure_session(&session_name, cols, rows, &mode)?;
let session_key = pending.session_key;
let key_id = protocol::session_key_id(&session_key);
let attach_ticket_psk = crypto::random_32();
@@ -497,8 +515,12 @@ async fn handle_native_user_auth(
.sessions
.get_mut(&session_name)
.expect("session exists");
if mode != "view-only" {
session.pty.resize(cols, rows)?;
if mode_allows_terminal_updates(&mode) {
session
.pty
.as_ref()
.context("terminal session has no pty")?
.resize(cols, rows)?;
session.parser.set_size(rows, cols);
}
let client_id = crypto::random_16();
@@ -605,19 +627,26 @@ async fn handle_bootstrap_attach(
if !verify_bootstrap(&req.bootstrap, &locked.secret)? {
return send_reject(socket, peer, "invalid or expired bootstrap").await;
}
if !locked.sessions.contains_key(&req.bootstrap.session) {
if locked.config.create_on_attach {
locked.ensure_session(&req.bootstrap.session, req.cols, req.rows)?;
} else {
return send_reject(socket, peer, "session does not exist").await;
}
if !locked.sessions.contains_key(&req.bootstrap.session) && !locked.config.create_on_attach
{
return send_reject(socket, peer, "session does not exist").await;
}
locked.ensure_session(
&req.bootstrap.session,
req.cols,
req.rows,
&req.bootstrap.mode,
)?;
let session = locked
.sessions
.get_mut(&req.bootstrap.session)
.expect("session exists");
if req.bootstrap.mode != "view-only" {
session.pty.resize(req.cols, req.rows)?;
if mode_allows_terminal_updates(&req.bootstrap.mode) {
session
.pty
.as_ref()
.context("terminal session has no pty")?
.resize(req.cols, req.rows)?;
session.parser.set_size(req.rows, req.cols);
}
let client_id = crypto::random_16();
@@ -710,19 +739,20 @@ async fn handle_ticket_attach(
let session_key_id = protocol::session_key_id(&session_key);
let (client_id, output_seq, snapshot) = {
let mut locked = state.lock().expect("server state poisoned");
if !locked.sessions.contains_key(&req.session) {
if locked.config.create_on_attach {
locked.ensure_session(&req.session, req.cols, req.rows)?;
} else {
return send_reject(socket, peer, "session does not exist").await;
}
if !locked.sessions.contains_key(&req.session) && !locked.config.create_on_attach {
return send_reject(socket, peer, "session does not exist").await;
}
locked.ensure_session(&req.session, req.cols, req.rows, &req.mode)?;
let session = locked
.sessions
.get_mut(&req.session)
.expect("session exists");
if req.mode != "view-only" {
session.pty.resize(req.cols, req.rows)?;
if mode_allows_terminal_updates(&req.mode) {
session
.pty
.as_ref()
.context("terminal session has no pty")?
.resize(req.cols, req.rows)?;
session.parser.set_size(req.rows, req.cols);
}
let client_id = crypto::random_16();
@@ -830,8 +860,12 @@ async fn handle_resume(
client.rows = req.rows;
client.last_seen = Instant::now();
client.send_seq += 1;
if client.mode != "view-only" {
session.pty.resize(req.cols, req.rows)?;
if mode_allows_terminal_updates(&client.mode) {
session
.pty
.as_ref()
.context("terminal session has no pty")?
.resize(req.cols, req.rows)?;
session.parser.set_size(req.rows, req.cols);
}
let snapshot = screen_snapshot(session.parser.screen());
@@ -883,10 +917,14 @@ async fn handle_input(
client.endpoint = peer;
}
client.last_seen = Instant::now();
if client.mode == "view-only" {
if !mode_allows_terminal_updates(&client.mode) {
return Ok(());
}
session.pty.write_all(&input.bytes)?;
session
.pty
.as_ref()
.context("terminal session has no pty")?
.write_all(&input.bytes)?;
Ok(())
}
@@ -910,11 +948,15 @@ async fn handle_resize(
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
if client.mode != "view-only" {
if mode_allows_terminal_updates(&client.mode) {
client.endpoint = peer;
client.cols = resize.cols;
client.rows = resize.rows;
session.pty.resize(resize.cols, resize.rows)?;
session
.pty
.as_ref()
.context("terminal session has no pty")?
.resize(resize.cols, resize.rows)?;
session.parser.set_size(resize.rows, resize.cols);
}
Ok(())
@@ -1656,3 +1698,20 @@ fn parse_size(raw: &str) -> Result<(u16, u16)> {
let (cols, rows) = raw.split_once('x').context("size must be COLSxROWS")?;
Ok((cols.parse()?, rows.parse()?))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn forward_only_session_does_not_allocate_pty() {
let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
state
.ensure_session("forward", 80, 24, "forward-only")
.unwrap();
assert!(state.sessions["forward"].pty.is_none());
}
}