Avoid PTY allocation for forward-only mode

src/bin/dosh-client.rs

@@ -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?;
                     }
                 }

src/bin/dosh-server.rs

@@ -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());
+    }
+}