diff --git a/CHANGELOG.md b/CHANGELOG.md index d993eb3f..e183583b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Roughly double bssh-server single-connection SFTP write throughput by fixing channel fragmentation and the serial write path** (#187). Three independent bottlenecks identified in the issue's code audit are addressed. First, `build_russh_config` now advertises a 65535-byte `maximum_packet_size` (the russh cap; library default 32768) and an 8 MiB `window_size` (library default 2 MiB), both configurable via `server.maximum_packet_size` / `server.window_size` (env: `BSSH_MAX_PACKET_SIZE` / `BSSH_WINDOW_SIZE`), so a 256 KiB SFTP write is no longer chopped into 8 CHANNEL_DATA packets that each pay the full cipher/copy/scheduling cost. Second, the SFTP server loop in `bssh-russh-sftp` is restructured from a strict read-process-write-flush cycle into a reader task with a bounded read-ahead queue plus an in-order processor: responses are flushed once per burst instead of once per request, and consecutive queued `SSH_FXP_WRITE` requests to the same handle at sequential offsets are coalesced into a single handler call (bounded by `max_write_coalesce_len`, default 256 KiB) while every merged request id still receives its own status reply, preserving SFTP response ordering and error semantics. The bssh write/read handlers additionally track each open file's cursor and elide the per-chunk `seek` for sequential transfers (append handles are exempt since O_APPEND ignores the cursor). Third, the workspace gains a tuned `[profile.release]` (`lto = "fat"`, `codegen-units = 1`) and documented `RUSTFLAGS` target-CPU guidance for self-builds. Local before/after benchmark (loopback, OpenSSH sftp client, 1 GiB random file, 3-run averages, 20-core Linux box): upload 71 MiB/s to 437 MiB/s (6.2x), download 73 MiB/s to 282 MiB/s (3.9x); an ablation with the code changes but library-default channel sizing lands at 333 / 291 MiB/s, attributing most of the gain to the pipelined write path and the rest of the upload gain to the larger packets. Numbers on the reporter's NIC-limited Xeon will differ; the methodology is documented in the PR. Remaining plan items (per-packet copies inside russh, AES-CTR internals) live in upstream russh since #212 removed the fork and are out of scope here. ### Added +- **Preserve an outer TUI's enhanced-keyboard state across interactive PTY sessions** (#236). Before the local input task starts, bssh performs DA-delimited, 100 ms-per-screen bounded queries for raw Kitty progressive-enhancement flags on both main and alternate screens, plus xterm `modifyOtherKeys` and the screen initially selected according to DEC private mode 1049. It uses non-clearing DEC mode 47 transitions to inspect the hidden screen, then normalizes both screens to the legacy keyboard baseline without popping the outer application's Kitty stacks, so the remote shell is isolated from the outer TUI's protocol. Strictly matched replies before each transaction's first primary DA are consumed while every post-DA, unrelated, malformed, partial, or concurrently typed byte is retained in order and forwarded through the normal PTY input path. Teardown resets leaked stacks and restores each screen's captured Kitty flags plus the global xterm value. Fragmented remote 1049 transitions are tracked only after the bytes survive filtering and are successfully written; the streaming parser skips OSC/DCS/APC/PM/SOS payloads with their protocol-specific terminators and validates CSI parameters so escape-looking data is not mistaken for an executed transition. A session that began on main can therefore safely use `1049l` to restore the remote-saved cursor after a disconnect; when a session began in an alternate screen, cleanup intentionally uses non-clearing mode 47 instead of a destructive `1049h`, because the terminal protocol cannot restore the 1049 bit and its single cursor-save slot without clearing or overwriting outer state. Failed or incomplete queries retain #234's ordinary-shell baseline fallback, hidden-screen setup failures make a best-effort return to the initial screen, the same guaranteed snapshot is reused by normal, forced, and repeated cleanup paths, and an atomic RAII ownership token covers both raw and deferred-raw guards. - **Make server-side SSH compression configurable instead of hard-disabled** (#220). The #215 workaround forced the server to advertise only `none` compression for every deployment, with no way for operators to opt back in. A new `server.compression` option (YAML `server.compression: true`, env `BSSH_COMPRESSION`, builder `ServerConfigBuilder::compression`) now controls whether `build_russh_config` advertises russh's default compression list (`none`, `zlib`, `zlib@openssh.com`) or only `none`. The default stays off, so behavior is unchanged unless explicitly enabled, and enabling it logs a warning: russh's delayed-zlib (`zlib@openssh.com`) transport still desyncs a few packets after compression activates post-auth (reproduced on russh 0.61.1 and 0.62.1), dropping clients that negotiate it mid-session. Both settings are covered by unit tests, and the option plus the desync caveat are documented in the man page and config docs. ### Fixed diff --git a/docs/architecture/interactive-mode.md b/docs/architecture/interactive-mode.md index a489f72f..ae393d72 100644 --- a/docs/architecture/interactive-mode.md +++ b/docs/architecture/interactive-mode.md @@ -193,7 +193,18 @@ Run this check in Ghostty and at least one other Kitty-keyboard-compatible termi 3. At the restored local prompt, verify that Space, Enter, arrow keys, and modified keys behave normally and do not print CSI-u fragments such as `;1:3u`, `:1C`, or `:3C`. 4. Repeat with a Kitty-keyboard-aware full-screen application on the remote host and terminate the SSH transport before the application exits cleanly. -The cleanup deliberately restores the ordinary shell baseline. Exact preservation of an enhanced keyboard mode owned by an outer local TUI is not currently available because crossterm exposes protocol support detection but not the active flag value, while issuing a separate `/dev/tty` query would race bssh's PTY input reader. +Before starting the PTY input task, bssh uses a DA-delimited query transaction of at most 100 ms per screen to capture the Kitty enhancement flags independently for the main and alternate screens, together with the global xterm `modifyOtherKeys` value and the screen initially selected according to DEC private mode 1049. It uses DEC mode 47 to inspect the hidden buffer without the clear/initialize behavior of mode 1049, sets both screens to the legacy keyboard baseline without popping the outer application's Kitty stacks, and returns to the original screen before remote input starts. Only strict replies before the first primary DA in each transaction are consumed; post-DA bytes and user input, malformed replies, or unrelated escape sequences are retained in byte order and forwarded to the remote session. A failed hidden-screen write or flush makes a best-effort mode-47 return before capture falls back. + +Teardown resets leaked stacks on both buffers and restores both captured Kitty values plus `modifyOtherKeys`. It also tracks fragmented remote `1049h`/`1049l` transitions, but commits a transition only after the bytes survive response filtering and `write_all` succeeds. Its streaming parser skips 7-bit OSC, DCS, APC, PM, and SOS payloads through fragmented terminators—BEL is accepted only for OSC compatibility, while every string type accepts `ESC \` ST—and accepts only validated private CSI parameters, preventing embedded escape-looking data from changing cleanup state. Raw 8-bit C1 controls are not interpreted because their byte values can also occur as UTF-8 continuation bytes. If bssh began on main and the connection dies in a remote-owned 1049 alternate screen, cleanup uses `1049l` to restore the remote-saved cursor before restoring both screens' keyboard state. The reverse correction is intentionally non-destructive: if bssh began inside an outer alternate screen and the remote reset 1049, cleanup selects the original buffer with mode 47 instead of forcing `1049h`. Generic terminal protocols expose only one cursor-save slot, and `1049h` clears the alternate buffer, so the 1049 mode bit, saved cursor, and existing outer contents cannot all be reconstructed exactly in that case. Keyboard state and screen contents take priority over claiming an impossible exact mode-bit restoration. + +If a DA transaction or the screen-state query is unsupported, malformed, or times out, cleanup falls back to the ordinary-shell baseline from #234 rather than claiming partial preservation. The cleanup snapshot is published with poison recovery and reused by normal, forced, and repeated cleanup paths. A single-owner RAII token covers raw and deferred-raw guards alike so they cannot race the process-global snapshot, while a deferred guard that never enters raw mode performs no terminal cleanup when dropped. + +To validate outer-TUI preservation in Ghostty and Kitty or WezTerm: + +1. Start a Kitty-keyboard-aware TUI that can launch a child command while remaining in its alternate screen, and record its active enhancement flags with `printf '\033[?u'`. +2. Launch `bssh` from that TUI, enable different remote flags with `printf '\033[=31u'`, and exit normally. +3. Verify the outer TUI receives the same keyboard encoding it used before launching bssh and that its alternate-screen contents were not cleared by an unnecessary `1049l`/`1049h` round trip. +4. Repeat after abrupt transport loss and local `~.`. From an ordinary main-screen shell, disconnect after remote `printf '\033[?1049h'` and verify bssh safely returns with `1049l`; from an outer alternate-screen TUI, use remote `printf '\033[?1049l'` and verify cleanup returns with mode 47 without clearing the outer buffer. ### Performance Characteristics diff --git a/src/pty/mod.rs b/src/pty/mod.rs index e962befa..fdaec89e 100644 --- a/src/pty/mod.rs +++ b/src/pty/mod.rs @@ -29,6 +29,11 @@ use tokio::time::Duration; pub mod session; pub mod terminal; +mod terminal_protocol; +#[cfg(test)] +mod terminal_protocol_boundary_tests; +mod terminal_protocol_restore; +mod terminal_screen; pub use session::PtySession; pub use terminal::{TerminalState, TerminalStateGuard, force_terminal_cleanup}; diff --git a/src/pty/session/mod.rs b/src/pty/session/mod.rs index b0504c00..36c0fc28 100644 --- a/src/pty/session/mod.rs +++ b/src/pty/session/mod.rs @@ -18,7 +18,9 @@ mod constants; mod escape_filter; mod input; mod local_escape; -mod raw_input; +mod output_delivery; +pub(crate) mod raw_input; +mod raw_input_task; mod session_manager; mod terminal_modes; diff --git a/src/pty/session/output_delivery.rs b/src/pty/session/output_delivery.rs new file mode 100644 index 00000000..4c614f27 --- /dev/null +++ b/src/pty/session/output_delivery.rs @@ -0,0 +1,110 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Ordered filtering, terminal delivery, and observation of remote output. + +use std::io::{self, Write}; + +use super::escape_filter::EscapeSequenceFilter; + +/// Filter remote bytes, write the surviving bytes, then commit only bytes the +/// terminal actually accepted to the protocol observer. +pub(super) fn deliver_filtered_output( + filter: &mut EscapeSequenceFilter, + writer: &mut impl Write, + input: &[u8], + commit: impl FnOnce(&[u8]), +) -> io::Result<()> { + let output = filter.filter(input); + if output.is_empty() { + return Ok(()); + } + + writer.write_all(&output)?; + commit(&output); + let _ = writer.flush(); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pty::terminal_screen::ScreenModeTracker; + + struct FailedWriter; + + impl Write for FailedWriter { + fn write(&mut self, _buffer: &[u8]) -> io::Result { + Err(io::Error::other("injected delivery failure")) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + #[test] + fn failed_delivery_does_not_commit_screen_transition() { + let mut filter = EscapeSequenceFilter::new(); + let mut tracker = ScreenModeTracker::new(); + let mut committed = false; + + let result = + deliver_filtered_output(&mut filter, &mut FailedWriter, b"\x1b[?1049h", |output| { + committed = true; + tracker.observe(output); + }); + + assert!(result.is_err()); + assert!(!committed); + } + + #[test] + fn filtered_payload_does_not_commit_screen_transition() { + let mut filter = EscapeSequenceFilter::new(); + let mut tracker = ScreenModeTracker::new(); + let mut output = Vec::new(); + let mut committed = false; + + deliver_filtered_output( + &mut filter, + &mut output, + b"\x1b]10;payload\x1b[?1049h\x07", + |delivered| { + committed = true; + tracker.observe(delivered); + }, + ) + .unwrap(); + + assert!(output.is_empty()); + assert!(!committed); + } + + #[test] + fn successful_delivery_commits_screen_transition() { + let mut filter = EscapeSequenceFilter::new(); + let mut tracker = ScreenModeTracker::new(); + let mut output = Vec::new(); + let mut observed = None; + + deliver_filtered_output(&mut filter, &mut output, b"\x1b[?1049h", |delivered| { + observed = tracker.observe(delivered); + }) + .unwrap(); + + assert_eq!(output, b"\x1b[?1049h"); + assert_eq!(observed, Some(true)); + } +} diff --git a/src/pty/session/raw_input_task.rs b/src/pty/session/raw_input_task.rs new file mode 100644 index 00000000..b628fa3a --- /dev/null +++ b/src/pty/session/raw_input_task.rs @@ -0,0 +1,90 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Blocking local-terminal input forwarding for an active PTY session. + +use std::time::Duration; + +use tokio::sync::{mpsc, watch}; + +use super::constants::INPUT_POLL_TIMEOUT_MS; +use super::local_escape::{LocalAction, LocalEscapeDetector}; +use super::raw_input::RawInputReader; +use crate::pty::PtyMessage; + +pub(crate) fn run( + input_tx: mpsc::Sender, + cancel_rx: watch::Receiver, + pending_input: Vec, +) { + let mut reader = RawInputReader::new(); + let mut buffer = [0_u8; 1024]; + let mut escape_detector = LocalEscapeDetector::new(); + + if !pending_input.is_empty() && !forward_input(&input_tx, &mut escape_detector, &pending_input) + { + return; + } + + loop { + if *cancel_rx.borrow() { + break; + } + + match reader.poll(Duration::from_millis(INPUT_POLL_TIMEOUT_MS)) { + Ok(true) => match reader.read(&mut buffer) { + Ok(0) => { + tracing::debug!("EOF received on stdin"); + break; + } + Ok(read) => { + if !forward_input(&input_tx, &mut escape_detector, &buffer[..read]) { + break; + } + } + Err(error) => { + let _ = input_tx.try_send(PtyMessage::Error(format!("Input error: {error}"))); + break; + } + }, + Ok(false) => continue, + Err(error) => { + let _ = input_tx.try_send(PtyMessage::Error(format!("Poll error: {error}"))); + break; + } + } + } +} + +fn forward_input( + input_tx: &mpsc::Sender, + detector: &mut LocalEscapeDetector, + bytes: &[u8], +) -> bool { + if let Some(action) = detector.process(bytes) { + match action { + LocalAction::Disconnect => { + tracing::debug!("Disconnect escape sequence detected"); + let _ = input_tx.try_send(PtyMessage::Terminate); + false + } + LocalAction::Passthrough(data) => { + input_tx.try_send(PtyMessage::LocalInput(data)).is_ok() + } + } + } else { + let data = smallvec::SmallVec::from_slice(bytes); + input_tx.try_send(PtyMessage::LocalInput(data)).is_ok() + } +} diff --git a/src/pty/session/session_manager.rs b/src/pty/session/session_manager.rs index e1b5de34..ae216992 100644 --- a/src/pty/session/session_manager.rs +++ b/src/pty/session/session_manager.rs @@ -16,8 +16,8 @@ use super::constants::*; use super::escape_filter::EscapeSequenceFilter; -use super::local_escape::{LocalAction, LocalEscapeDetector}; -use super::raw_input::RawInputReader; +use super::output_delivery::deliver_filtered_output; +use super::raw_input_task; use super::terminal_modes::configure_terminal_modes; use crate::pty::{ PtyConfig, PtyMessage, PtyState, @@ -156,7 +156,9 @@ impl PtySession { } // Set up terminal state guard - self.terminal_guard = Some(TerminalStateGuard::new()?); + let mut terminal_guard = TerminalStateGuard::new()?; + let pending_input = terminal_guard.take_pending_input(); + self.terminal_guard = Some(terminal_guard); // Enable mouse support if requested if self.config.enable_mouse { @@ -226,70 +228,7 @@ impl PtySession { // NOTE: TerminalStateGuard has already called enable_raw_mode() at this point, // so stdin.read() will return raw bytes without line buffering let input_task = tokio::task::spawn_blocking(move || { - let mut reader = RawInputReader::new(); - let mut buffer = [0u8; 1024]; - let mut escape_detector = LocalEscapeDetector::new(); - - loop { - if *cancel_for_input.borrow() { - break; - } - - let poll_timeout = Duration::from_millis(INPUT_POLL_TIMEOUT_MS); - - match reader.poll(poll_timeout) { - Ok(true) => { - match reader.read(&mut buffer) { - Ok(0) => { - // EOF - user closed stdin - tracing::debug!("EOF received on stdin"); - break; - } - Ok(n) => { - // Check for local escape sequences (e.g., ~. for disconnect) - if let Some(action) = escape_detector.process(&buffer[..n]) { - match action { - LocalAction::Disconnect => { - tracing::debug!("Disconnect escape sequence detected"); - let _ = input_tx.try_send(PtyMessage::Terminate); - break; - } - LocalAction::Passthrough(data) => { - // Send filtered data - if input_tx - .try_send(PtyMessage::LocalInput(data)) - .is_err() - { - break; - } - } - } - } else { - // Pass raw bytes through as-is - // This includes arrow keys, function keys, terminal responses, etc. - let data = smallvec::SmallVec::from_slice(&buffer[..n]); - if input_tx.try_send(PtyMessage::LocalInput(data)).is_err() { - break; - } - } - } - Err(e) => { - let _ = input_tx - .try_send(PtyMessage::Error(format!("Input error: {e}"))); - break; - } - } - } - Ok(false) => { - // Timeout - continue polling - continue; - } - Err(e) => { - let _ = input_tx.try_send(PtyMessage::Error(format!("Poll error: {e}"))); - break; - } - } - } + raw_input_task::run(input_tx, cancel_for_input, pending_input); }); // We'll integrate channel reading into the main loop since russh Channel doesn't clone @@ -317,27 +256,37 @@ impl PtySession { // Filter terminal escape sequence responses before display // This prevents raw XTGETTCAP, DA1/DA2/DA3 responses from appearing // on screen when running applications like Neovim - let filtered_data = self.escape_filter.filter(data); - if !filtered_data.is_empty() { - if let Err(e) = io::stdout().write_all(&filtered_data) { - tracing::error!("Failed to write to stdout: {e}"); - should_terminate = true; - } else { - let _ = io::stdout().flush(); - } + let terminal_guard = &mut self.terminal_guard; + if let Err(e) = deliver_filtered_output( + &mut self.escape_filter, + &mut io::stdout(), + data, + |delivered| { + if let Some(guard) = terminal_guard.as_mut() { + guard.observe_remote_output(delivered); + } + }, + ) { + tracing::error!("Failed to write to stdout: {e}"); + should_terminate = true; } } Some(ChannelMsg::ExtendedData { ref data, ext }) => { if ext == 1 { // stderr - also filter escape sequences - let filtered_data = self.escape_filter.filter(data); - if !filtered_data.is_empty() { - if let Err(e) = io::stdout().write_all(&filtered_data) { - tracing::error!("Failed to write stderr to stdout: {e}"); - should_terminate = true; - } else { - let _ = io::stdout().flush(); - } + let terminal_guard = &mut self.terminal_guard; + if let Err(e) = deliver_filtered_output( + &mut self.escape_filter, + &mut io::stdout(), + data, + |delivered| { + if let Some(guard) = terminal_guard.as_mut() { + guard.observe_remote_output(delivered); + } + }, + ) { + tracing::error!("Failed to write stderr to stdout: {e}"); + should_terminate = true; } } } @@ -381,14 +330,19 @@ impl PtySession { // Apply escape filter for consistency with SSH channel data // This path may receive data from other sources that could // contain terminal responses that shouldn't be displayed - let filtered_data = self.escape_filter.filter(&data); - if !filtered_data.is_empty() { - if let Err(e) = io::stdout().write_all(&filtered_data) { - tracing::error!("Failed to write to stdout: {e}"); - should_terminate = true; - } else { - let _ = io::stdout().flush(); - } + let terminal_guard = &mut self.terminal_guard; + if let Err(e) = deliver_filtered_output( + &mut self.escape_filter, + &mut io::stdout(), + &data, + |delivered| { + if let Some(guard) = terminal_guard.as_mut() { + guard.observe_remote_output(delivered); + } + }, + ) { + tracing::error!("Failed to write to stdout: {e}"); + should_terminate = true; } } Some(PtyMessage::Resize { width, height }) => { diff --git a/src/pty/terminal.rs b/src/pty/terminal.rs index dcc27b7c..2106d2f8 100644 --- a/src/pty/terminal.rs +++ b/src/pty/terminal.rs @@ -14,9 +14,8 @@ //! Terminal state management for PTY sessions. -use std::io::Write; use std::sync::{ - Arc, LazyLock, Mutex, + Arc, LazyLock, Mutex, MutexGuard, PoisonError, TryLockError, atomic::{AtomicBool, Ordering}, }; @@ -27,32 +26,43 @@ use crossterm::{ terminal::{disable_raw_mode, enable_raw_mode}, }; +use super::terminal_protocol::{ProtocolState, capture_protocol_state, write_terminal_cleanup}; +use super::terminal_screen::ScreenModeTracker; + /// Global terminal cleanup synchronization /// Ensures only one cleanup attempt happens even with multiple guards static TERMINAL_MUTEX: LazyLock> = LazyLock::new(|| Mutex::new(())); +static SAVED_PROTOCOL_STATE: LazyLock>> = + LazyLock::new(|| Mutex::new(None)); static RAW_MODE_ACTIVE: AtomicBool = AtomicBool::new(false); +static TERMINAL_OWNER_ACTIVE: AtomicBool = AtomicBool::new(false); -/// Best-effort reset for terminal modes that a remote PTY application can leak. -/// -/// Kitty keyboard mode stacks are independent for the main and alternate screens. -/// Reset the current screen first, leave the alternate screen, and then reset the -/// main screen. `CSI < 999 u` performs a large bounded pop of unbalanced pushes, -/// `CSI = 0 u` also handles implementations without a stack, and -/// `CSI > 4 ; 0 m` disables xterm's `modifyOtherKeys` fallback. -/// -/// This intentionally restores the ordinary shell baseline. Crossterm does not -/// expose the current enhancement flags, and querying `/dev/tty` here would race -/// the PTY input reader. An outer local TUI that starts bssh with enhanced keyboard -/// reporting enabled may therefore need to re-enable its mode after bssh exits. -const TERMINAL_CLEANUP_SEQUENCE: &[u8] = b"\x1b[?2004l\x1b[<999u\x1b[=0u\x1b[>4;0m\x1b[?1000l\x1b[?1002l\x1b[?1003l\x1b[?1006l\x1b[?1015l\x1b[?1049l\x1b[<999u\x1b[=0u\x1b[>4;0m\x1b[?25h"; +fn lock_recover(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} -/// Emit the shared terminal reset sequence to an injectable writer. -/// -/// Cleanup must never replace the original session error, so write and flush -/// failures are deliberately ignored. -fn write_terminal_cleanup(writer: &mut impl Write) { - let _ = writer.write_all(TERMINAL_CLEANUP_SEQUENCE); - let _ = writer.flush(); +fn publish_protocol_state(state: &ProtocolState) { + *lock_recover(&SAVED_PROTOCOL_STATE) = Some(state.clone()); +} + +struct TerminalOwner; + +impl TerminalOwner { + fn acquire() -> Result { + if TERMINAL_OWNER_ACTIVE + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_err() + { + anyhow::bail!("Another terminal state guard is already active"); + } + Ok(Self) + } +} + +impl Drop for TerminalOwner { + fn drop(&mut self) { + TERMINAL_OWNER_ACTIVE.store(false, Ordering::SeqCst); + } } /// Terminal state information that needs to be preserved and restored @@ -86,13 +96,18 @@ impl Default for TerminalState { pub struct TerminalStateGuard { saved_state: TerminalState, is_raw_mode_active: Arc, - // Simplified cleanup - just track if we need cleanup - _needs_cleanup: bool, + needs_cleanup: AtomicBool, + protocol_cleanup_needed: bool, + pending_input: Vec, + protocol_state: ProtocolState, + screen_mode_tracker: ScreenModeTracker, + _owner: TerminalOwner, } impl TerminalStateGuard { /// Create a new terminal state guard and enter raw mode pub fn new() -> Result { + let owner = TerminalOwner::acquire()?; let saved_state = Self::save_terminal_state()?; let is_raw_mode_active = Arc::new(AtomicBool::new(false)); @@ -104,6 +119,14 @@ impl TerminalStateGuard { is_raw_mode_active.store(true, Ordering::Relaxed); } + // This bounded query runs before the PTY input task exists, so terminal + // replies cannot race normal input forwarding. + let capture = capture_protocol_state(&mut std::io::stdout()); + let protocol_state = capture.state.clone(); + + // Publish the snapshot before any later fallible setup so the caller's + // forced cleanup can preserve it if construction returns an error. + publish_protocol_state(&protocol_state); // Enable bracketed paste mode execute!(std::io::stdout(), EnableBracketedPaste) .with_context(|| "Failed to enable bracketed paste mode")?; @@ -111,19 +134,30 @@ impl TerminalStateGuard { Ok(Self { saved_state, is_raw_mode_active, - _needs_cleanup: true, + needs_cleanup: AtomicBool::new(true), + protocol_cleanup_needed: true, + pending_input: capture.pending_input, + protocol_state, + screen_mode_tracker: ScreenModeTracker::new(), + _owner: owner, }) } /// Create a terminal state guard without entering raw mode pub fn new_without_raw_mode() -> Result { + let owner = TerminalOwner::acquire()?; let saved_state = Self::save_terminal_state()?; let is_raw_mode_active = Arc::new(AtomicBool::new(false)); Ok(Self { saved_state, is_raw_mode_active, - _needs_cleanup: false, + needs_cleanup: AtomicBool::new(false), + protocol_cleanup_needed: false, + pending_input: Vec::new(), + protocol_state: ProtocolState::default(), + screen_mode_tracker: ScreenModeTracker::new(), + _owner: owner, }) } @@ -135,6 +169,7 @@ impl TerminalStateGuard { RAW_MODE_ACTIVE.store(true, Ordering::SeqCst); self.is_raw_mode_active.store(true, Ordering::Relaxed); } + self.needs_cleanup.store(true, Ordering::Release); Ok(()) } @@ -146,6 +181,9 @@ impl TerminalStateGuard { RAW_MODE_ACTIVE.store(false, Ordering::SeqCst); self.is_raw_mode_active.store(false, Ordering::Relaxed); } + if !self.protocol_cleanup_needed { + self.needs_cleanup.store(false, Ordering::Release); + } Ok(()) } @@ -159,6 +197,20 @@ impl TerminalStateGuard { &self.saved_state } + /// Take bytes read during the pre-input terminal query that were not replies. + pub(crate) fn take_pending_input(&mut self) -> Vec { + std::mem::take(&mut self.pending_input) + } + + /// Observe remote DEC 1049 transitions so safe teardown can leave a + /// remote-owned alternate screen when bssh originally started on main. + pub(crate) fn observe_remote_output(&mut self, output: &[u8]) { + if let Some(current) = self.screen_mode_tracker.observe(output) { + self.protocol_state.current_1049 = Some(current); + publish_protocol_state(&self.protocol_state); + } + } + /// Save current terminal state fn save_terminal_state() -> Result { let size = if let Some((terminal_size::Width(w), terminal_size::Height(h))) = @@ -186,7 +238,9 @@ impl TerminalStateGuard { // Reset every terminal mode owned by the remote PTY through the same path // used by panic and last-resort cleanup. - write_terminal_cleanup(&mut std::io::stdout()); + if self.protocol_cleanup_needed { + write_terminal_cleanup(&mut std::io::stdout(), &self.protocol_state); + } // Exit raw mode if it's globally active if RAW_MODE_ACTIVE.load(Ordering::SeqCst) { @@ -201,6 +255,7 @@ impl TerminalStateGuard { if self.is_raw_mode_active.load(Ordering::Relaxed) { self.is_raw_mode_active.store(false, Ordering::Relaxed); } + self.needs_cleanup.store(false, Ordering::Release); Ok(()) } @@ -208,7 +263,9 @@ impl TerminalStateGuard { impl Drop for TerminalStateGuard { fn drop(&mut self) { - if let Err(e) = self.restore_terminal_state() { + if self.needs_cleanup.load(Ordering::Acquire) + && let Err(e) = self.restore_terminal_state() + { eprintln!("Warning: Failed to restore terminal state: {e}"); } } @@ -241,7 +298,13 @@ pub fn force_terminal_cleanup() { // operations below are individually safe. let _guard = TERMINAL_MUTEX.try_lock().ok(); - write_terminal_cleanup(&mut std::io::stdout()); + let state = match SAVED_PROTOCOL_STATE.try_lock() { + Ok(saved) => saved.clone(), + Err(TryLockError::Poisoned(error)) => error.into_inner().clone(), + Err(TryLockError::WouldBlock) => None, + } + .unwrap_or_default(); + write_terminal_cleanup(&mut std::io::stdout(), &state); if RAW_MODE_ACTIVE.load(Ordering::SeqCst) { let _ = disable_raw_mode(); @@ -335,79 +398,8 @@ impl TerminalOps { mod tests { use super::*; - // ----------------------------------------------------------------------- - // force_terminal_cleanup tests - // - // Terminal-state mutation (raw mode, alternate screen) cannot be unit- - // tested safely inside `cargo test` because: - // • the test runner does not allocate a real TTY — crossterm operations - // that require a TTY (e.g. enable_raw_mode) will fail or behave - // unexpectedly; - // • the escape-sequence writes go to cargo's captured stdout, which is - // harmless but not observable in a meaningful way; - // • the global statics (TERMINAL_MUTEX, RAW_MODE_ACTIVE) are shared - // across all tests in the same process, so tests that mutate them - // must be carefully ordered or marked #[serial]. - // - // What *can* be reliably tested here: - // 1. Exact cleanup bytes and ordering through an injectable writer. - // 2. Idempotency: calling force_terminal_cleanup() twice does not panic. - // 3. Poisoned-mutex resilience: the `try_lock().ok()` pattern correctly - // yields None (rather than panicking) when the mutex is poisoned. - // Verified using a local Mutex so the global TERMINAL_MUTEX is never - // poisoned, keeping other tests unaffected. - // 4. Held-mutex resilience: `try_lock()` returns WouldBlock (not a - // deadlock) when a lock is already held. Verified using a local Mutex - // for the same isolation reason. - // ----------------------------------------------------------------------- - - #[test] - fn test_terminal_cleanup_writes_exact_mode_reset_sequence() { - let mut output = Vec::new(); - - write_terminal_cleanup(&mut output); - - assert_eq!( - output, - b"\x1b[?2004l\ - \x1b[<999u\x1b[=0u\x1b[>4;0m\ - \x1b[?1000l\x1b[?1002l\x1b[?1003l\x1b[?1006l\x1b[?1015l\ - \x1b[?1049l\ - \x1b[<999u\x1b[=0u\x1b[>4;0m\ - \x1b[?25h" - ); - } - - #[test] - fn test_terminal_cleanup_writer_is_idempotent() { - let mut output = Vec::new(); - - write_terminal_cleanup(&mut output); - write_terminal_cleanup(&mut output); - - assert_eq!(output.len(), TERMINAL_CLEANUP_SEQUENCE.len() * 2); - assert_eq!( - &output[..TERMINAL_CLEANUP_SEQUENCE.len()], - &output[TERMINAL_CLEANUP_SEQUENCE.len()..] - ); - } - - #[test] - fn test_terminal_cleanup_ignores_writer_failures() { - struct FailingWriter; - - impl Write for FailingWriter { - fn write(&mut self, _buffer: &[u8]) -> std::io::Result { - Err(std::io::Error::other("intentional write failure")) - } - - fn flush(&mut self) -> std::io::Result<()> { - Err(std::io::Error::other("intentional flush failure")) - } - } - - write_terminal_cleanup(&mut FailingWriter); - } + // TTY mutations are not observable under the test harness, so these tests + // cover idempotency and synchronization behavior without enabling raw mode. /// Calling force_terminal_cleanup() twice in succession must not panic. /// @@ -470,4 +462,33 @@ mod tests { assert!(!state.was_mouse_enabled); assert_eq!(state.size, (80, 24)); } + + #[test] + #[serial_test::serial(terminal_owner)] + fn test_terminal_owner_applies_to_non_raw_guard_and_releases_on_drop() { + let guard = TerminalStateGuard::new_without_raw_mode().unwrap(); + assert!(!guard.needs_cleanup.load(Ordering::Acquire)); + assert!(TerminalStateGuard::new_without_raw_mode().is_err()); + + drop(guard); + + assert!(TerminalOwner::acquire().is_ok()); + } + + #[test] + fn test_lock_recover_allows_guaranteed_publication_after_poison() { + let saved = Mutex::new(None); + let _ = std::panic::catch_unwind(|| { + let _guard = saved.lock().unwrap(); + panic!("intentional poison"); + }); + let expected = ProtocolState { + main_kitty_flags: Some(7), + ..ProtocolState::default() + }; + + *lock_recover(&saved) = Some(expected.clone()); + + assert_eq!(*lock_recover(&saved), Some(expected)); + } } diff --git a/src/pty/terminal_protocol.rs b/src/pty/terminal_protocol.rs new file mode 100644 index 00000000..ac8d6f34 --- /dev/null +++ b/src/pty/terminal_protocol.rs @@ -0,0 +1,317 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Bounded capture and restoration of local enhanced-keyboard terminal modes. + +use std::io::Write; +use std::time::{Duration, Instant}; + +use super::session::raw_input::RawInputReader; +pub(crate) use super::terminal_protocol_restore::write_terminal_cleanup; + +const INITIAL_QUERY: &[u8] = b"\x1b[?u\x1b[?4m\x1b[?1049$p\x1b[c"; +const HIDDEN_SCREEN_QUERY: &[u8] = b"\x1b[?u\x1b[c"; +const QUERY_TIMEOUT: Duration = Duration::from_millis(100); +const MAX_QUERY_INPUT: usize = 4096; +const MAX_CSI_SEQUENCE: usize = 64; + +const SELECT_MAIN_SCREEN: &[u8] = b"\x1b[?47l"; +const SELECT_ALTERNATE_SCREEN: &[u8] = b"\x1b[?47h"; +const ENTRY_KEYBOARD_BASELINE: &[u8] = b"\x1b[=0u\x1b[>4;0m"; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct ProtocolState { + /// Raw Kitty flags for the main screen, including flags unknown to crossterm. + pub(crate) main_kitty_flags: Option, + /// Raw Kitty flags for the alternate screen. + pub(crate) alternate_kitty_flags: Option, + /// xterm's `modifyOtherKeys` resource value. + pub(crate) modify_other_keys: Option, + /// Screen selected when bssh took control, as reported by DEC mode 1049. + /// + /// This is a buffer selector, not a promise that the 1049 mode bit and its + /// single cursor-save slot can both be restored without clearing content. + pub(crate) initial_screen_alternate: Option, + /// Last observed DECSET/DECRST 1049 state from the remote session. + pub(crate) current_1049: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ModifyOtherKeys { + Value(u32), + Disabled, +} + +#[derive(Debug, Default, PartialEq, Eq)] +pub(crate) struct ProtocolCapture { + pub(crate) state: ProtocolState, + pub(crate) pending_input: Vec, +} + +#[derive(Debug, Default, PartialEq, Eq)] +pub(crate) struct QueryReplies { + kitty_flags: Option, + modify_other_keys: Option, + pub(crate) alternate_screen: Option, + pub(crate) pending_input: Vec, + pub(crate) completed: bool, +} + +trait QueryInput { + fn poll(&self, timeout: Duration) -> std::io::Result; + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result; +} + +impl QueryInput for RawInputReader { + fn poll(&self, timeout: Duration) -> std::io::Result { + RawInputReader::poll(self, timeout) + } + + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + RawInputReader::read(self, buffer) + } +} + +/// Capture terminal protocol state before the PTY input task starts. +/// +/// The primary-device-attributes request is a sentinel: terminals answer it even +/// when they do not implement the preceding optional queries. Bytes that are not +/// strict matches for one of our replies are returned for normal PTY forwarding. +pub(crate) fn capture_protocol_state(writer: &mut impl Write) -> ProtocolCapture { + let mut reader = RawInputReader::new(); + capture_all_screens(writer, &mut reader, QUERY_TIMEOUT) +} + +fn capture_all_screens( + writer: &mut impl Write, + reader: &mut impl QueryInput, + timeout: Duration, +) -> ProtocolCapture { + let first = capture_transaction(writer, reader, timeout, INITIAL_QUERY); + let mut pending_input = first.pending_input; + let Some(initial_alternate) = first.alternate_screen else { + let _ = writer.write_all(ENTRY_KEYBOARD_BASELINE); + let _ = writer.flush(); + return ProtocolCapture { + state: ProtocolState::default(), + pending_input, + }; + }; + if !first.completed { + let _ = writer.write_all(ENTRY_KEYBOARD_BASELINE); + let _ = writer.flush(); + return ProtocolCapture { + state: ProtocolState::default(), + pending_input, + }; + } + + let _ = writer.write_all(ENTRY_KEYBOARD_BASELINE); + let hidden_switch = if initial_alternate { + SELECT_MAIN_SCREEN + } else { + SELECT_ALTERNATE_SCREEN + }; + let return_switch = if initial_alternate { + SELECT_ALTERNATE_SCREEN + } else { + SELECT_MAIN_SCREEN + }; + if writer + .write_all(hidden_switch) + .and_then(|()| writer.flush()) + .is_err() + { + return_to_initial_screen(writer, return_switch); + return ProtocolCapture { + state: ProtocolState::default(), + pending_input, + }; + } + + let hidden = capture_transaction(writer, reader, timeout, HIDDEN_SCREEN_QUERY); + pending_input.extend_from_slice(&hidden.pending_input); + let _ = writer.write_all(ENTRY_KEYBOARD_BASELINE); + return_to_initial_screen(writer, return_switch); + + if !hidden.completed { + return ProtocolCapture { + state: ProtocolState::default(), + pending_input, + }; + } + + let (main_kitty_flags, alternate_kitty_flags) = if initial_alternate { + (hidden.kitty_flags, first.kitty_flags) + } else { + (first.kitty_flags, hidden.kitty_flags) + }; + ProtocolCapture { + state: ProtocolState { + main_kitty_flags, + alternate_kitty_flags, + modify_other_keys: first.modify_other_keys, + initial_screen_alternate: Some(initial_alternate), + current_1049: Some(initial_alternate), + }, + pending_input, + } +} + +fn return_to_initial_screen(writer: &mut impl Write, return_switch: &[u8]) { + let _ = writer.write_all(return_switch); + let _ = writer.write_all(ENTRY_KEYBOARD_BASELINE); + let _ = writer.flush(); +} + +fn capture_transaction( + writer: &mut impl Write, + reader: &mut impl QueryInput, + timeout: Duration, + query: &[u8], +) -> QueryReplies { + if writer + .write_all(query) + .and_then(|()| writer.flush()) + .is_err() + { + return QueryReplies::default(); + } + + let deadline = Instant::now() + timeout; + let mut input = Vec::new(); + let mut buffer = [0_u8; 256]; + + while input.len() < MAX_QUERY_INPUT { + let Some(remaining) = deadline.checked_duration_since(Instant::now()) else { + break; + }; + match reader.poll(remaining) { + Ok(true) => {} + Ok(false) | Err(_) => break, + } + + let read_limit = buffer.len().min(MAX_QUERY_INPUT - input.len()); + match reader.read(&mut buffer[..read_limit]) { + Ok(0) | Err(_) => break, + Ok(read) => input.extend_from_slice(&buffer[..read]), + } + + if contains_primary_device_attributes(&input) { + break; + } + } + + parse_query_input(&input) +} + +fn contains_primary_device_attributes(input: &[u8]) -> bool { + csi_sequences(input).any(is_primary_device_attributes) +} + +pub(crate) fn parse_query_input(input: &[u8]) -> QueryReplies { + let Some(sentinel_end) = csi_ranges(input) + .find_map(|(start, end)| is_primary_device_attributes(&input[start..end]).then_some(end)) + else { + return QueryReplies { + pending_input: input.to_vec(), + ..QueryReplies::default() + }; + }; + + let mut capture = QueryReplies::default(); + let mut cursor = 0; + + for (start, end) in csi_ranges(&input[..sentinel_end]) { + capture + .pending_input + .extend_from_slice(&input[cursor..start]); + let sequence = &input[start..end]; + + if let Some(flags) = parse_decimal_reply(sequence, b"\x1b[?", b"u") { + capture.kitty_flags = Some(flags); + } else if let Some(value) = parse_decimal_reply(sequence, b"\x1b[>4;", b"m") { + capture.modify_other_keys = Some(ModifyOtherKeys::Value(value)); + } else if sequence == b"\x1b[>4n" { + capture.modify_other_keys = Some(ModifyOtherKeys::Disabled); + } else if let Some(mode) = parse_decimal_reply(sequence, b"\x1b[?1049;", b"$y") { + match mode { + 1 | 3 => capture.alternate_screen = Some(true), + 2 | 4 => capture.alternate_screen = Some(false), + _ => capture.pending_input.extend_from_slice(sequence), + } + } else if is_primary_device_attributes(sequence) { + capture.completed = true; + } else { + capture.pending_input.extend_from_slice(sequence); + } + + cursor = end; + } + capture.pending_input.extend_from_slice(&input[cursor..]); + capture +} + +fn parse_decimal_reply(sequence: &[u8], prefix: &[u8], suffix: &[u8]) -> Option { + let digits = sequence.strip_prefix(prefix)?.strip_suffix(suffix)?; + if digits.is_empty() || digits.len() > 10 || !digits.iter().all(u8::is_ascii_digit) { + return None; + } + std::str::from_utf8(digits).ok()?.parse().ok() +} + +fn is_primary_device_attributes(sequence: &[u8]) -> bool { + let Some(parameters) = sequence + .strip_prefix(b"\x1b[?") + .and_then(|value| value.strip_suffix(b"c")) + else { + return false; + }; + !parameters.is_empty() + && parameters + .iter() + .all(|byte| byte.is_ascii_digit() || *byte == b';') +} + +fn csi_sequences(input: &[u8]) -> impl Iterator { + csi_ranges(input).map(|(start, end)| &input[start..end]) +} + +fn csi_ranges(input: &[u8]) -> impl Iterator { + let mut offset = 0; + std::iter::from_fn(move || { + while offset + 1 < input.len() { + if input[offset] != b'\x1b' || input[offset + 1] != b'[' { + offset += 1; + continue; + } + + let start = offset; + let limit = input.len().min(start + MAX_CSI_SEQUENCE); + offset += 2; + while offset < limit { + if (0x40..=0x7e).contains(&input[offset]) { + offset += 1; + return Some((start, offset)); + } + offset += 1; + } + } + None + }) +} + +#[cfg(test)] +#[path = "terminal_protocol_tests.rs"] +mod tests; diff --git a/src/pty/terminal_protocol_boundary_tests.rs b/src/pty/terminal_protocol_boundary_tests.rs new file mode 100644 index 00000000..de86ed15 --- /dev/null +++ b/src/pty/terminal_protocol_boundary_tests.rs @@ -0,0 +1,46 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::terminal_protocol::parse_query_input; + +#[test] +fn preserves_reply_shaped_and_user_bytes_after_da_sentinel() { + let input = b"\x1b[?7u\x1b[?1049;2$y\x1b[?1;2c\x1b[?31uuser"; + + let capture = parse_query_input(input); + + assert!(capture.completed); + assert_eq!(capture.pending_input, b"\x1b[?31uuser"); +} + +#[test] +fn missing_da_discards_state_and_preserves_every_captured_byte() { + let input = b"typed\x1b[?31u\x1b[>4;2m\x1b[?1049;1$y"; + + let capture = parse_query_input(input); + + assert!(!capture.completed); + assert_eq!(capture.pending_input, input); +} + +#[test] +fn invalid_decrqm_mode_is_preserved_as_malformed_input() { + let input = b"\x1b[?1049;9$y\x1b[?1;2c"; + + let capture = parse_query_input(input); + + assert!(capture.completed); + assert!(capture.alternate_screen.is_none()); + assert_eq!(capture.pending_input, b"\x1b[?1049;9$y"); +} diff --git a/src/pty/terminal_protocol_restore.rs b/src/pty/terminal_protocol_restore.rs new file mode 100644 index 00000000..c826d352 --- /dev/null +++ b/src/pty/terminal_protocol_restore.rs @@ -0,0 +1,211 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::io::Write; + +use super::terminal_protocol::{ModifyOtherKeys, ProtocolState}; + +const COMMON_RESET_PREFIX: &[u8] = + b"\x1b[?2004l\x1b[<999u\x1b[=0u\x1b[>4;0m\x1b[?1000l\x1b[?1002l\x1b[?1003l\x1b[?1006l\x1b[?1015l"; +const SELECT_MAIN_SCREEN: &[u8] = b"\x1b[?47l"; +const SELECT_ALTERNATE_SCREEN: &[u8] = b"\x1b[?47h"; +const LEAVE_1049_SCREEN: &[u8] = b"\x1b[?1049l"; +const KEYBOARD_RESET: &[u8] = b"\x1b[<999u\x1b[=0u\x1b[>4;0m"; +const SHOW_CURSOR: &[u8] = b"\x1b[?25h"; + +/// Reset remote-owned modes on both screens and restore both captured states. +pub(crate) fn write_terminal_cleanup(writer: &mut impl Write, state: &ProtocolState) { + let _ = (|| -> std::io::Result<()> { + writer.write_all(COMMON_RESET_PREFIX)?; + let Some(initial_alternate) = state.initial_screen_alternate else { + writer.write_all(b"\x1b[?1049l")?; + writer.write_all(KEYBOARD_RESET)?; + writer.write_all(SHOW_CURSOR)?; + return writer.flush(); + }; + + // If bssh started on the main screen and the remote died inside a + // 1049 alternate screen, DECRST 1049 is both safe and useful: it + // restores the main screen and the cursor saved by the remote. + // + // The inverse is deliberately not attempted. DECSET 1049 clears the + // alternate buffer and overwrites the terminal's single saved-cursor + // slot, so forcing it when bssh started in an alternate screen would + // destroy exactly the outer state this cleanup is preserving. Mode 47 + // below selects that buffer without clearing it. + if !initial_alternate && state.current_1049 == Some(true) { + writer.write_all(LEAVE_1049_SCREEN)?; + } + + writer.write_all(screen_select(!initial_alternate))?; + writer.write_all(KEYBOARD_RESET)?; + write_kitty_restore(writer, kitty_flags(state, !initial_alternate))?; + + writer.write_all(screen_select(initial_alternate))?; + writer.write_all(KEYBOARD_RESET)?; + write_kitty_restore(writer, kitty_flags(state, initial_alternate))?; + write_modify_other_keys_restore(writer, state.modify_other_keys)?; + writer.write_all(SHOW_CURSOR)?; + writer.flush() + })(); +} + +fn screen_select(alternate: bool) -> &'static [u8] { + if alternate { + SELECT_ALTERNATE_SCREEN + } else { + SELECT_MAIN_SCREEN + } +} + +fn kitty_flags(state: &ProtocolState, alternate: bool) -> Option { + if alternate { + state.alternate_kitty_flags + } else { + state.main_kitty_flags + } +} + +fn write_kitty_restore(writer: &mut impl Write, flags: Option) -> std::io::Result<()> { + if let Some(flags) = flags { + write!(writer, "\x1b[={flags}u")?; + } + Ok(()) +} + +fn write_modify_other_keys_restore( + writer: &mut impl Write, + value: Option, +) -> std::io::Result<()> { + match value { + Some(ModifyOtherKeys::Value(value)) => write!(writer, "\x1b[>4;{value}m"), + Some(ModifyOtherKeys::Disabled) => writer.write_all(b"\x1b[>4n"), + None => Ok(()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn restores_both_screens_then_returns_to_initial_main() { + let state = ProtocolState { + main_kitty_flags: Some(31), + alternate_kitty_flags: Some(15), + modify_other_keys: Some(ModifyOtherKeys::Value(2)), + initial_screen_alternate: Some(false), + current_1049: Some(false), + }; + let mut output = Vec::new(); + + write_terminal_cleanup(&mut output, &state); + + assert_eq!( + output, + b"\x1b[?2004l\x1b[<999u\x1b[=0u\x1b[>4;0m\ + \x1b[?1000l\x1b[?1002l\x1b[?1003l\x1b[?1006l\x1b[?1015l\ + \x1b[?47h\x1b[<999u\x1b[=0u\x1b[>4;0m\x1b[=15u\ + \x1b[?47l\x1b[<999u\x1b[=0u\x1b[>4;0m\ + \x1b[=31u\x1b[>4;2m\x1b[?25h" + ); + } + + #[test] + fn restores_both_screens_then_returns_to_initial_alternate() { + let state = ProtocolState { + main_kitty_flags: Some(3), + alternate_kitty_flags: Some(7), + modify_other_keys: Some(ModifyOtherKeys::Value(1)), + initial_screen_alternate: Some(true), + current_1049: Some(true), + }; + let mut output = Vec::new(); + + write_terminal_cleanup(&mut output, &state); + + assert_eq!( + output, + b"\x1b[?2004l\x1b[<999u\x1b[=0u\x1b[>4;0m\ + \x1b[?1000l\x1b[?1002l\x1b[?1003l\x1b[?1006l\x1b[?1015l\ + \x1b[?47l\x1b[<999u\x1b[=0u\x1b[>4;0m\x1b[=3u\ + \x1b[?47h\x1b[<999u\x1b[=0u\x1b[>4;0m\ + \x1b[=7u\x1b[>4;1m\x1b[?25h" + ); + } + + #[test] + fn repeated_cleanup_restores_same_effective_state() { + let state = ProtocolState { + main_kitty_flags: Some(31), + alternate_kitty_flags: Some(15), + modify_other_keys: None, + initial_screen_alternate: Some(false), + current_1049: Some(false), + }; + let mut output = Vec::new(); + + write_terminal_cleanup(&mut output, &state); + write_terminal_cleanup(&mut output, &state); + + assert_eq!( + output + .windows(b"\x1b[=31u".len()) + .filter(|window| *window == b"\x1b[=31u") + .count(), + 2 + ); + } + + #[test] + fn leaves_remote_1049_screen_when_initial_screen_was_main() { + let state = ProtocolState { + main_kitty_flags: Some(3), + alternate_kitty_flags: Some(7), + modify_other_keys: None, + initial_screen_alternate: Some(false), + current_1049: Some(true), + }; + let mut output = Vec::new(); + + write_terminal_cleanup(&mut output, &state); + + let leave = output + .windows(LEAVE_1049_SCREEN.len()) + .position(|window| window == LEAVE_1049_SCREEN) + .expect("cleanup should leave the remote's 1049 screen"); + let select_hidden = output + .windows(SELECT_ALTERNATE_SCREEN.len()) + .position(|window| window == SELECT_ALTERNATE_SCREEN) + .expect("cleanup should still restore hidden-screen keyboard state"); + assert!(leave < select_hidden); + } + + #[test] + fn does_not_clear_outer_alternate_screen_to_force_1049_bit() { + let state = ProtocolState { + main_kitty_flags: Some(3), + alternate_kitty_flags: Some(7), + modify_other_keys: None, + initial_screen_alternate: Some(true), + current_1049: Some(false), + }; + let mut output = Vec::new(); + + write_terminal_cleanup(&mut output, &state); + + assert!(!output.windows(8).any(|window| window == b"\x1b[?1049h")); + assert!(output.ends_with(b"\x1b[?25h")); + } +} diff --git a/src/pty/terminal_protocol_tests.rs b/src/pty/terminal_protocol_tests.rs new file mode 100644 index 00000000..e08af2ef --- /dev/null +++ b/src/pty/terminal_protocol_tests.rs @@ -0,0 +1,218 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::{collections::VecDeque, io::Write}; + +use super::*; + +#[derive(Default)] +struct FakeInput { + chunks: VecDeque>, +} + +impl FakeInput { + fn new(chunks: &[&[u8]]) -> Self { + Self { + chunks: chunks.iter().map(|chunk| chunk.to_vec()).collect(), + } + } +} + +impl QueryInput for FakeInput { + fn poll(&self, _timeout: Duration) -> std::io::Result { + Ok(!self.chunks.is_empty()) + } + + fn read(&mut self, buffer: &mut [u8]) -> std::io::Result { + let Some(chunk) = self.chunks.pop_front() else { + return Ok(0); + }; + let read = chunk.len().min(buffer.len()); + buffer[..read].copy_from_slice(&chunk[..read]); + Ok(read) + } +} + +struct FailOnceWriter { + output: Vec, + writes: usize, + fail_at: usize, +} + +impl Write for FailOnceWriter { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + self.writes += 1; + if self.writes == self.fail_at { + return Err(std::io::Error::other("injected write failure")); + } + self.output.extend_from_slice(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +struct FailOnceFlushWriter { + output: Vec, + flushes: usize, + fail_at: usize, +} + +impl Write for FailOnceFlushWriter { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + self.output.extend_from_slice(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.flushes += 1; + if self.flushes == self.fail_at { + return Err(std::io::Error::other("injected flush failure")); + } + Ok(()) + } +} + +struct NoReplyInput; + +impl QueryInput for NoReplyInput { + fn poll(&self, timeout: Duration) -> std::io::Result { + std::thread::sleep(timeout); + Ok(false) + } + + fn read(&mut self, _buffer: &mut [u8]) -> std::io::Result { + unreachable!("poll never reports readable input") + } +} + +#[test] +fn captures_fragmented_replies_and_preserves_user_input() { + let mut output = Vec::new(); + let mut input = FakeInput::new(&[ + b"typed", + b"\x1b[?3", + b"1u\x1b[>4;2m\x1b[?1049;1$y\x1b[?1;2c", + b"\x1b[?5u\x1b[?1;2c", + ]); + + let capture = capture_all_screens(&mut output, &mut input, Duration::from_secs(1)); + + assert_eq!( + output, + b"\x1b[?u\x1b[?4m\x1b[?1049$p\x1b[c\ + \x1b[=0u\x1b[>4;0m\x1b[?47l\ + \x1b[?u\x1b[c\ + \x1b[=0u\x1b[>4;0m\x1b[?47h\ + \x1b[=0u\x1b[>4;0m" + ); + assert_eq!( + capture.state, + ProtocolState { + main_kitty_flags: Some(5), + alternate_kitty_flags: Some(31), + modify_other_keys: Some(ModifyOtherKeys::Value(2)), + initial_screen_alternate: Some(true), + current_1049: Some(true), + } + ); + assert_eq!(capture.pending_input, b"typed"); +} + +#[test] +fn timeout_and_unsupported_terminal_fall_back_without_losing_input() { + let mut output = Vec::new(); + let mut input = FakeInput::new(&[b"abc\x1b[?1;2c"]); + + let capture = capture_all_screens(&mut output, &mut input, Duration::from_secs(1)); + + assert_eq!(capture.state, ProtocolState::default()); + assert_eq!(capture.pending_input, b"abc"); +} + +#[test] +fn malformed_and_unrelated_sequences_are_preserved() { + let input = b"\x1b[?99999999999u\x1b[?x u\x1b[Aplain"; + + let capture = parse_query_input(input); + + assert!(!capture.completed); + assert_eq!(capture.pending_input, input); +} + +#[test] +fn hidden_switch_write_failure_best_effort_returns_to_initial_screen() { + let mut output = FailOnceWriter { + output: Vec::new(), + writes: 0, + fail_at: 3, + }; + let mut input = FakeInput::new(&[b"\x1b[?7u\x1b[?1049;2$y\x1b[?1;2c"]); + + let capture = capture_all_screens(&mut output, &mut input, Duration::from_secs(1)); + + assert_eq!(capture.state, ProtocolState::default()); + assert!(output.output.ends_with(b"\x1b[?47l\x1b[=0u\x1b[>4;0m")); +} + +#[test] +fn hidden_switch_flush_failure_best_effort_returns_to_initial_screen() { + let mut output = FailOnceFlushWriter { + output: Vec::new(), + flushes: 0, + fail_at: 3, + }; + let mut input = FakeInput::new(&[b"\x1b[?7u\x1b[?1049;2$y\x1b[?1;2c"]); + + let capture = capture_all_screens(&mut output, &mut input, Duration::from_secs(1)); + + assert_eq!(capture.state, ProtocolState::default()); + assert!(output.output.ends_with(b"\x1b[?47l\x1b[=0u\x1b[>4;0m")); +} + +#[test] +fn no_da_reply_waits_for_the_bounded_timeout() { + let timeout = Duration::from_millis(15); + let started = Instant::now(); + let mut output = Vec::new(); + let mut input = NoReplyInput; + + let capture = capture_all_screens(&mut output, &mut input, timeout); + + assert!(started.elapsed() >= Duration::from_millis(10)); + assert_eq!(capture.state, ProtocolState::default()); + assert!(capture.pending_input.is_empty()); +} + +#[test] +fn captures_xtqmodkeys_disabled_reply() { + let capture = parse_query_input(b"\x1b[>4n\x1b[?1;2c"); + + assert!(capture.completed); + assert_eq!(capture.modify_other_keys, Some(ModifyOtherKeys::Disabled)); + assert!(capture.pending_input.is_empty()); +} + +#[test] +fn accepts_permanently_set_and_reset_decrqm_statuses() { + let set = parse_query_input(b"\x1b[?1049;3$y\x1b[?1;2c"); + let reset = parse_query_input(b"\x1b[?1049;4$y\x1b[?1;2c"); + + assert_eq!(set.alternate_screen, Some(true)); + assert_eq!(reset.alternate_screen, Some(false)); + assert!(set.pending_input.is_empty()); + assert!(reset.pending_input.is_empty()); +} diff --git a/src/pty/terminal_screen.rs b/src/pty/terminal_screen.rs new file mode 100644 index 00000000..90651a56 --- /dev/null +++ b/src/pty/terminal_screen.rs @@ -0,0 +1,232 @@ +// Copyright 2025 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Tracks delivered DECSET/DECRST 1049 transitions across fragmented output. + +const MAX_CSI_SEQUENCE: usize = 64; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ParserState { + Ground, + Escape, + Csi, + String(StringKind), + StringEscape(StringKind), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StringKind { + Osc, + Other, +} + +#[derive(Debug)] +pub(crate) struct ScreenModeTracker { + state: ParserState, + csi: Vec, +} + +impl ScreenModeTracker { + pub(crate) fn new() -> Self { + Self { + state: ParserState::Ground, + csi: Vec::with_capacity(16), + } + } + + /// Return the last DEC 1049 state explicitly selected by delivered output. + /// + /// OSC payloads are skipped until BEL or ST; DCS, APC, PM, and SOS payloads + /// require ST. This prevents escape-looking payload data from being + /// mistaken for an executed CSI command. + pub(crate) fn observe(&mut self, input: &[u8]) -> Option { + let mut current = None; + + for &byte in input { + match self.state { + ParserState::Ground => self.consume_ground(byte), + ParserState::Escape => self.consume_escape(byte), + ParserState::Csi => { + if let Some(is_set) = self.consume_csi(byte) { + current = Some(is_set); + } + } + ParserState::String(kind) => self.consume_string(kind, byte), + ParserState::StringEscape(kind) => self.consume_string_escape(kind, byte), + } + } + + current + } + + fn consume_ground(&mut self, byte: u8) { + if byte == b'\x1b' { + self.state = ParserState::Escape; + } + } + + fn consume_escape(&mut self, byte: u8) { + match byte { + b'[' => self.start_csi(), + b']' => self.state = ParserState::String(StringKind::Osc), + // DCS, SOS, PM, and APC. + b'P' | b'X' | b'^' | b'_' => { + self.state = ParserState::String(StringKind::Other); + } + b'\x1b' => {} + _ => self.state = ParserState::Ground, + } + } + + fn start_csi(&mut self) { + self.csi.clear(); + self.state = ParserState::Csi; + } + + fn consume_csi(&mut self, byte: u8) -> Option { + match byte { + // CAN and SUB cancel a control sequence. + 0x18 | 0x1a => { + self.csi.clear(); + self.state = ParserState::Ground; + return None; + } + b'\x1b' => { + self.csi.clear(); + self.state = ParserState::Escape; + return None; + } + _ => self.csi.push(byte), + } + + if self.csi.len() > MAX_CSI_SEQUENCE { + self.csi.clear(); + self.state = ParserState::Ground; + return None; + } + if !(0x40..=0x7e).contains(&byte) { + return None; + } + + let selection = parse_1049_selection(&self.csi); + self.csi.clear(); + self.state = ParserState::Ground; + selection + } + + fn consume_string(&mut self, kind: StringKind, byte: u8) { + match byte { + b'\x07' if kind == StringKind::Osc => self.state = ParserState::Ground, + b'\x1b' => self.state = ParserState::StringEscape(kind), + _ => {} + } + } + + fn consume_string_escape(&mut self, kind: StringKind, byte: u8) { + match byte { + b'\\' => self.state = ParserState::Ground, + b'\x07' if kind == StringKind::Osc => self.state = ParserState::Ground, + b'\x1b' => {} + _ => self.state = ParserState::String(kind), + } + } +} + +fn parse_1049_selection(sequence: &[u8]) -> Option { + let (&final_byte, body) = sequence.split_last()?; + let is_set = match final_byte { + b'h' => true, + b'l' => false, + _ => return None, + }; + let parameters = body.strip_prefix(b"?")?; + if parameters.is_empty() + || !parameters + .iter() + .all(|byte| byte.is_ascii_digit() || *byte == b';') + { + return None; + } + + parameters + .split(|byte| *byte == b';') + .any(|mode| mode == b"1049") + .then_some(is_set) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tracks_fragmented_1049_transitions_and_last_value() { + let mut tracker = ScreenModeTracker::new(); + + assert_eq!(tracker.observe(b"text\x1b[?10"), None); + assert_eq!(tracker.observe(b"49h\x1b[?25;1049l"), Some(false)); + assert_eq!(tracker.observe(b"\x1b[?47h"), None); + } + + #[test] + fn skips_osc_dcs_and_tmux_payload_false_positives() { + let mut tracker = ScreenModeTracker::new(); + + assert_eq!(tracker.observe(b"\x1b]0;title\x1b[?1049h\x07"), None); + assert_eq!(tracker.observe(b"\x1bPpayload\x1b[?1049l\x1b\\"), None); + assert_eq!(tracker.observe(b"\x1bPtmux;\x1b\x1b[?1049h\x1b\\"), None); + assert_eq!(tracker.observe(b"\x1b[?1049h"), Some(true)); + } + + #[test] + fn non_osc_strings_ignore_bel_and_embedded_1049_until_fragmented_st() { + let mut tracker = ScreenModeTracker::new(); + + for introducer in [ + b"\x1bP".as_slice(), + b"\x1b_".as_slice(), + b"\x1b^".as_slice(), + b"\x1bX".as_slice(), + ] { + assert_eq!(tracker.observe(introducer), None); + assert_eq!(tracker.observe(b"payload\x07\x1b[?1049h\x1b"), None); + assert_eq!(tracker.observe(b"\\"), None); + } + assert_eq!(tracker.observe(b"\x1b[?1049l"), Some(false)); + } + + #[test] + fn osc_bel_terminates_before_following_1049() { + let mut tracker = ScreenModeTracker::new(); + + assert_eq!(tracker.observe(b"\x1b]0;title\x07\x1b[?1049h"), Some(true)); + } + + #[test] + fn rejects_invalid_csi_parameters() { + let mut tracker = ScreenModeTracker::new(); + + assert_eq!(tracker.observe(b"\x1b[?1049:1h"), None); + assert_eq!(tracker.observe(b"\x1b[?1049$h"), None); + assert_eq!(tracker.observe(b"\x1b[?x;1049h"), None); + assert_eq!(tracker.observe(b"\x1b[?25;1049h"), Some(true)); + } + + #[test] + fn does_not_treat_utf8_c1_continuation_as_csi() { + let mut tracker = ScreenModeTracker::new(); + + assert_eq!(tracker.observe(b"\xc2\x9b?1049h"), None); + assert_eq!(tracker.observe(b"\x9b?1049h"), None); + } +}