bssh/pty/
terminal.rs

1// Copyright 2025 Lablup Inc. and Jeongkyu Shin
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Terminal state management for PTY sessions.
16
17use anyhow::{Context, Result};
18use crossterm::terminal::{disable_raw_mode, enable_raw_mode};
19use once_cell::sync::Lazy;
20use std::sync::{
21    atomic::{AtomicBool, Ordering},
22    Arc, Mutex,
23};
24
25/// Global terminal cleanup synchronization
26/// Ensures only one cleanup attempt happens even with multiple guards
27static TERMINAL_MUTEX: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
28static RAW_MODE_ACTIVE: AtomicBool = AtomicBool::new(false);
29
30/// Terminal state information that needs to be preserved and restored
31#[derive(Debug, Clone)]
32pub struct TerminalState {
33    /// Whether raw mode was enabled before we took control
34    pub was_raw_mode: bool,
35    /// Terminal size when state was saved
36    pub size: (u32, u32),
37    /// Whether alternate screen buffer was in use
38    pub was_alternate_screen: bool,
39    /// Whether mouse reporting was enabled
40    pub was_mouse_enabled: bool,
41}
42
43impl Default for TerminalState {
44    fn default() -> Self {
45        Self {
46            was_raw_mode: false,
47            size: (80, 24),
48            was_alternate_screen: false,
49            was_mouse_enabled: false,
50        }
51    }
52}
53
54/// RAII guard for terminal state management
55///
56/// This ensures that terminal state is properly restored even if
57/// the PTY session is interrupted or fails.
58pub struct TerminalStateGuard {
59    saved_state: TerminalState,
60    is_raw_mode_active: Arc<AtomicBool>,
61    // Simplified cleanup - just track if we need cleanup
62    _needs_cleanup: bool,
63}
64
65impl TerminalStateGuard {
66    /// Create a new terminal state guard and enter raw mode
67    pub fn new() -> Result<Self> {
68        let saved_state = Self::save_terminal_state()?;
69        let is_raw_mode_active = Arc::new(AtomicBool::new(false));
70
71        // Enter raw mode with global synchronization
72        let _guard = TERMINAL_MUTEX.lock().unwrap();
73        if !RAW_MODE_ACTIVE.load(Ordering::SeqCst) {
74            enable_raw_mode().with_context(|| "Failed to enable raw mode")?;
75            RAW_MODE_ACTIVE.store(true, Ordering::SeqCst);
76            is_raw_mode_active.store(true, Ordering::Relaxed);
77        }
78
79        Ok(Self {
80            saved_state,
81            is_raw_mode_active,
82            _needs_cleanup: true,
83        })
84    }
85
86    /// Create a terminal state guard without entering raw mode
87    pub fn new_without_raw_mode() -> Result<Self> {
88        let saved_state = Self::save_terminal_state()?;
89        let is_raw_mode_active = Arc::new(AtomicBool::new(false));
90
91        Ok(Self {
92            saved_state,
93            is_raw_mode_active,
94            _needs_cleanup: false,
95        })
96    }
97
98    /// Manually enter raw mode
99    pub fn enter_raw_mode(&self) -> Result<()> {
100        let _guard = TERMINAL_MUTEX.lock().unwrap();
101        if !RAW_MODE_ACTIVE.load(Ordering::SeqCst) {
102            enable_raw_mode().with_context(|| "Failed to enable raw mode")?;
103            RAW_MODE_ACTIVE.store(true, Ordering::SeqCst);
104            self.is_raw_mode_active.store(true, Ordering::Relaxed);
105        }
106        Ok(())
107    }
108
109    /// Manually exit raw mode
110    pub fn exit_raw_mode(&self) -> Result<()> {
111        let _guard = TERMINAL_MUTEX.lock().unwrap();
112        if RAW_MODE_ACTIVE.load(Ordering::SeqCst) {
113            disable_raw_mode().with_context(|| "Failed to disable raw mode")?;
114            RAW_MODE_ACTIVE.store(false, Ordering::SeqCst);
115            self.is_raw_mode_active.store(false, Ordering::Relaxed);
116        }
117        Ok(())
118    }
119
120    /// Check if raw mode is currently active
121    pub fn is_raw_mode_active(&self) -> bool {
122        self.is_raw_mode_active.load(Ordering::Relaxed)
123    }
124
125    /// Get the saved terminal state
126    pub fn saved_state(&self) -> &TerminalState {
127        &self.saved_state
128    }
129
130    /// Save current terminal state
131    fn save_terminal_state() -> Result<TerminalState> {
132        let size = if let Some((terminal_size::Width(w), terminal_size::Height(h))) =
133            terminal_size::terminal_size()
134        {
135            (u32::from(w), u32::from(h))
136        } else {
137            (80, 24) // Default fallback
138        };
139
140        // TODO: Detect if we're already in raw mode, alternate screen, etc.
141        // For now, assume we're starting from a clean state
142        Ok(TerminalState {
143            was_raw_mode: false,
144            size,
145            was_alternate_screen: false,
146            was_mouse_enabled: false,
147        })
148    }
149
150    /// Restore terminal state to its original condition
151    fn restore_terminal_state(&self) -> Result<()> {
152        // Use global synchronization to prevent race conditions
153        let _guard = TERMINAL_MUTEX.lock().unwrap();
154
155        // Exit raw mode if it's globally active
156        if RAW_MODE_ACTIVE.load(Ordering::SeqCst) {
157            if let Err(e) = disable_raw_mode() {
158                eprintln!("Warning: Failed to disable raw mode during cleanup: {e}");
159            } else {
160                RAW_MODE_ACTIVE.store(false, Ordering::SeqCst);
161            }
162        }
163
164        // Mark our local state as cleaned
165        if self.is_raw_mode_active.load(Ordering::Relaxed) {
166            self.is_raw_mode_active.store(false, Ordering::Relaxed);
167        }
168
169        // TODO: Restore other terminal settings if needed
170        // For now, just exiting raw mode is sufficient
171
172        Ok(())
173    }
174}
175
176impl Drop for TerminalStateGuard {
177    fn drop(&mut self) {
178        if let Err(e) = self.restore_terminal_state() {
179            eprintln!("Warning: Failed to restore terminal state: {e}");
180        }
181    }
182}
183
184/// Force terminal cleanup - can be called from anywhere to ensure terminal is restored
185pub fn force_terminal_cleanup() {
186    let _guard = TERMINAL_MUTEX.lock().unwrap();
187    if RAW_MODE_ACTIVE.load(Ordering::SeqCst) {
188        let _ = disable_raw_mode();
189        RAW_MODE_ACTIVE.store(false, Ordering::SeqCst);
190    }
191}
192
193/// Terminal operations for PTY sessions
194pub struct TerminalOps;
195
196impl TerminalOps {
197    /// Enable mouse support in terminal
198    pub fn enable_mouse() -> Result<()> {
199        use crossterm::event::EnableMouseCapture;
200        use crossterm::execute;
201
202        execute!(std::io::stdout(), EnableMouseCapture)
203            .with_context(|| "Failed to enable mouse capture")?;
204
205        Ok(())
206    }
207
208    /// Disable mouse support in terminal
209    pub fn disable_mouse() -> Result<()> {
210        use crossterm::event::DisableMouseCapture;
211        use crossterm::execute;
212
213        execute!(std::io::stdout(), DisableMouseCapture)
214            .with_context(|| "Failed to disable mouse capture")?;
215
216        Ok(())
217    }
218
219    /// Enable alternate screen buffer
220    pub fn enable_alternate_screen() -> Result<()> {
221        use crossterm::execute;
222        use crossterm::terminal::EnterAlternateScreen;
223
224        execute!(std::io::stdout(), EnterAlternateScreen)
225            .with_context(|| "Failed to enter alternate screen")?;
226
227        Ok(())
228    }
229
230    /// Disable alternate screen buffer
231    pub fn disable_alternate_screen() -> Result<()> {
232        use crossterm::execute;
233        use crossterm::terminal::LeaveAlternateScreen;
234
235        execute!(std::io::stdout(), LeaveAlternateScreen)
236            .with_context(|| "Failed to leave alternate screen")?;
237
238        Ok(())
239    }
240
241    /// Clear the terminal screen
242    pub fn clear_screen() -> Result<()> {
243        use crossterm::execute;
244        use crossterm::terminal::{Clear, ClearType};
245
246        execute!(std::io::stdout(), Clear(ClearType::All))
247            .with_context(|| "Failed to clear screen")?;
248
249        Ok(())
250    }
251
252    /// Move cursor to home position (0, 0)
253    pub fn cursor_home() -> Result<()> {
254        use crossterm::cursor::MoveTo;
255        use crossterm::execute;
256
257        execute!(std::io::stdout(), MoveTo(0, 0))
258            .with_context(|| "Failed to move cursor to home")?;
259
260        Ok(())
261    }
262
263    /// Set terminal title
264    pub fn set_title(title: &str) -> Result<()> {
265        use crossterm::execute;
266        use crossterm::terminal::SetTitle;
267
268        execute!(std::io::stdout(), SetTitle(title))
269            .with_context(|| "Failed to set terminal title")?;
270
271        Ok(())
272    }
273}