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(); !local_forwards.is_empty() || !remote_forwards.is_empty() || !dynamic_forwards.is_empty();
let predict = args.predict || host.predict.unwrap_or(config.predict); let predict = args.predict || host.predict.unwrap_or(config.predict);
let session = select_session(args.session.as_deref(), args.new); 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" "view-only"
} else { } else {
"read-write" "read-write"
@@ -1945,7 +1947,7 @@ async fn run_terminal(
if bytes == [0x1d] { if bytes == [0x1d] {
break; break;
} }
if cred.mode != "view-only" { if cred.mode != "view-only" && cred.mode != "forward-only" {
predictor.observe_input(&bytes)?; predictor.observe_input(&bytes)?;
send_input(&socket, addr, &cred, &mut send_seq, bytes).await?; send_input(&socket, addr, &cred, &mut send_seq, bytes).await?;
} }
@@ -1958,7 +1960,7 @@ async fn run_terminal(
&& (cols, rows) != last_size && (cols, rows) != last_size
{ {
last_size = (cols, rows); 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?; 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"); let mut locked = state.lock().expect("server state poisoned");
for session in config.prewarm_sessions.clone() { 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 { struct Session {
pty: PtyHandle, pty: Option<PtyHandle>,
parser: vt100::Parser, parser: vt100::Parser,
clients: HashMap<[u8; 16], ClientState>, clients: HashMap<[u8; 16], ClientState>,
output_seq: u64, output_seq: u64,
@@ -225,17 +225,30 @@ impl ServerState {
} }
} }
fn ensure_session(&mut self, name: &str, cols: u16, rows: u16) -> Result<()> { fn ensure_session(&mut self, name: &str, cols: u16, rows: u16, mode: &str) -> Result<()> {
if self.sessions.contains_key(name) { 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(()); return Ok(());
} }
let pty = spawn_pty_session( let pty = if mode_uses_pty(mode) {
name.to_string(), Some(spawn_pty_session(
&self.config.shell, name.to_string(),
cols.max(1), &self.config.shell,
rows.max(1), cols.max(1),
self.pty_tx.clone(), rows.max(1),
)?; self.pty_tx.clone(),
)?)
} else {
None
};
self.sessions.insert( self.sessions.insert(
name.to_string(), name.to_string(),
Session { 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( async fn handle_packet(
state: &Arc<Mutex<ServerState>>, state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>, socket: &Arc<UdpSocket>,
@@ -471,13 +492,10 @@ async fn handle_native_user_auth(
let mode = pending.client.requested_mode.clone(); let mode = pending.client.requested_mode.clone();
let cols = pending.client.terminal_size.0; let cols = pending.client.terminal_size.0;
let rows = pending.client.terminal_size.1; let rows = pending.client.terminal_size.1;
if !locked.sessions.contains_key(&session_name) { if !locked.sessions.contains_key(&session_name) && !locked.config.create_on_attach {
if locked.config.create_on_attach { return Err(anyhow!("session does not exist"));
locked.ensure_session(&session_name, cols, rows)?;
} else {
return Err(anyhow!("session does not exist"));
}
} }
locked.ensure_session(&session_name, cols, rows, &mode)?;
let session_key = pending.session_key; let session_key = pending.session_key;
let key_id = protocol::session_key_id(&session_key); let key_id = protocol::session_key_id(&session_key);
let attach_ticket_psk = crypto::random_32(); let attach_ticket_psk = crypto::random_32();
@@ -497,8 +515,12 @@ async fn handle_native_user_auth(
.sessions .sessions
.get_mut(&session_name) .get_mut(&session_name)
.expect("session exists"); .expect("session exists");
if mode != "view-only" { if mode_allows_terminal_updates(&mode) {
session.pty.resize(cols, rows)?; session
.pty
.as_ref()
.context("terminal session has no pty")?
.resize(cols, rows)?;
session.parser.set_size(rows, cols); session.parser.set_size(rows, cols);
} }
let client_id = crypto::random_16(); let client_id = crypto::random_16();
@@ -605,19 +627,26 @@ async fn handle_bootstrap_attach(
if !verify_bootstrap(&req.bootstrap, &locked.secret)? { if !verify_bootstrap(&req.bootstrap, &locked.secret)? {
return send_reject(socket, peer, "invalid or expired bootstrap").await; return send_reject(socket, peer, "invalid or expired bootstrap").await;
} }
if !locked.sessions.contains_key(&req.bootstrap.session) { if !locked.sessions.contains_key(&req.bootstrap.session) && !locked.config.create_on_attach
if locked.config.create_on_attach { {
locked.ensure_session(&req.bootstrap.session, req.cols, req.rows)?; return send_reject(socket, peer, "session does not exist").await;
} else {
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 let session = locked
.sessions .sessions
.get_mut(&req.bootstrap.session) .get_mut(&req.bootstrap.session)
.expect("session exists"); .expect("session exists");
if req.bootstrap.mode != "view-only" { if mode_allows_terminal_updates(&req.bootstrap.mode) {
session.pty.resize(req.cols, req.rows)?; session
.pty
.as_ref()
.context("terminal session has no pty")?
.resize(req.cols, req.rows)?;
session.parser.set_size(req.rows, req.cols); session.parser.set_size(req.rows, req.cols);
} }
let client_id = crypto::random_16(); 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 session_key_id = protocol::session_key_id(&session_key);
let (client_id, output_seq, snapshot) = { let (client_id, output_seq, snapshot) = {
let mut locked = state.lock().expect("server state poisoned"); let mut locked = state.lock().expect("server state poisoned");
if !locked.sessions.contains_key(&req.session) { if !locked.sessions.contains_key(&req.session) && !locked.config.create_on_attach {
if locked.config.create_on_attach { return send_reject(socket, peer, "session does not exist").await;
locked.ensure_session(&req.session, req.cols, req.rows)?;
} else {
return send_reject(socket, peer, "session does not exist").await;
}
} }
locked.ensure_session(&req.session, req.cols, req.rows, &req.mode)?;
let session = locked let session = locked
.sessions .sessions
.get_mut(&req.session) .get_mut(&req.session)
.expect("session exists"); .expect("session exists");
if req.mode != "view-only" { if mode_allows_terminal_updates(&req.mode) {
session.pty.resize(req.cols, req.rows)?; session
.pty
.as_ref()
.context("terminal session has no pty")?
.resize(req.cols, req.rows)?;
session.parser.set_size(req.rows, req.cols); session.parser.set_size(req.rows, req.cols);
} }
let client_id = crypto::random_16(); let client_id = crypto::random_16();
@@ -830,8 +860,12 @@ async fn handle_resume(
client.rows = req.rows; client.rows = req.rows;
client.last_seen = Instant::now(); client.last_seen = Instant::now();
client.send_seq += 1; client.send_seq += 1;
if client.mode != "view-only" { if mode_allows_terminal_updates(&client.mode) {
session.pty.resize(req.cols, req.rows)?; session
.pty
.as_ref()
.context("terminal session has no pty")?
.resize(req.cols, req.rows)?;
session.parser.set_size(req.rows, req.cols); session.parser.set_size(req.rows, req.cols);
} }
let snapshot = screen_snapshot(session.parser.screen()); let snapshot = screen_snapshot(session.parser.screen());
@@ -883,10 +917,14 @@ async fn handle_input(
client.endpoint = peer; client.endpoint = peer;
} }
client.last_seen = Instant::now(); client.last_seen = Instant::now();
if client.mode == "view-only" { if !mode_allows_terminal_updates(&client.mode) {
return Ok(()); return Ok(());
} }
session.pty.write_all(&input.bytes)?; session
.pty
.as_ref()
.context("terminal session has no pty")?
.write_all(&input.bytes)?;
Ok(()) Ok(())
} }
@@ -910,11 +948,15 @@ async fn handle_resize(
if !client.replay.accept(packet.header.seq) { if !client.replay.accept(packet.header.seq) {
return Ok(()); return Ok(());
} }
if client.mode != "view-only" { if mode_allows_terminal_updates(&client.mode) {
client.endpoint = peer; client.endpoint = peer;
client.cols = resize.cols; client.cols = resize.cols;
client.rows = resize.rows; 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); session.parser.set_size(resize.rows, resize.cols);
} }
Ok(()) 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")?; let (cols, rows) = raw.split_once('x').context("size must be COLSxROWS")?;
Ok((cols.parse()?, rows.parse()?)) 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());
}
}