smux/
smux.rs

1use std::{
2    fmt, fs,
3    io::{self, BufWriter, Read, Write},
4    path::PathBuf,
5    sync::{Arc, Mutex, RwLock},
6    time::Duration,
7};
8
9use bytes::Bytes;
10use crossterm::{
11    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
12    execute,
13    terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
14};
15use portable_pty::{native_pty_system, CommandBuilder, MasterPty, PtySize};
16use ratatui::{
17    backend::CrosstermBackend,
18    layout::{Alignment, Constraint, Direction, Layout, Rect},
19    style::{Color, Modifier, Style},
20    widgets::{Block, Borders, Paragraph},
21    Terminal,
22};
23use tokio::{
24    sync::mpsc::{channel, Sender},
25    task::spawn_blocking,
26};
27use tracing::Level;
28use tracing_subscriber::FmtSubscriber;
29use tui_term::widget::{Cursor, PseudoTerminal};
30
31#[derive(Debug, Clone, Copy)]
32struct Size {
33    cols: u16,
34    rows: u16,
35}
36
37#[tokio::main]
38async fn main() -> io::Result<()> {
39    init_panic_hook();
40    let (mut terminal, mut size) = setup_terminal().unwrap();
41
42    let cwd = std::env::current_dir().unwrap();
43    let mut cmd = CommandBuilder::new_default_prog();
44    cmd.cwd(cwd);
45
46    let mut panes: Vec<PtyPane> = Vec::new();
47    let mut active_pane: Option<usize> = None;
48
49    // Add a default pane
50    let pane_size = calc_pane_size(size, 1);
51    open_new_pane(&mut panes, &mut active_pane, &cmd, pane_size)?;
52
53    loop {
54        terminal.draw(|f| {
55            let chunks = Layout::default()
56                .direction(Direction::Vertical)
57                .margin(1)
58                .constraints([Constraint::Percentage(100), Constraint::Min(1)].as_ref())
59                .split(f.area());
60
61            let pane_height = if panes.is_empty() {
62                chunks[0].height
63            } else {
64                (chunks[0].height.saturating_sub(1)) / panes.len() as u16
65            };
66
67            for (index, pane) in panes.iter().enumerate() {
68                let block = Block::default()
69                    .borders(Borders::ALL)
70                    .style(Style::default().add_modifier(Modifier::BOLD));
71                let mut cursor = Cursor::default();
72                let block = if Some(index) == active_pane {
73                    block.style(
74                        Style::default()
75                            .add_modifier(Modifier::BOLD)
76                            .fg(Color::LightMagenta),
77                    )
78                } else {
79                    cursor.hide();
80                    block
81                };
82                let parser = pane.parser.read().unwrap();
83                let screen = parser.screen();
84                let pseudo_term = PseudoTerminal::new(screen).block(block).cursor(cursor);
85                let pane_chunk = Rect {
86                    x: chunks[0].x,
87                    y: chunks[0].y + (index as u16 * pane_height), /* Adjust the y coordinate for
88                                                                    * each pane */
89                    width: chunks[0].width,
90                    height: pane_height, // Use the calculated pane height directly
91                };
92                f.render_widget(pseudo_term, pane_chunk);
93            }
94
95            let explanation =
96                "Ctrl+n to open a new pane | Ctrl+x to close the active pane | Ctrl+q to quit";
97            let explanation = Paragraph::new(explanation)
98                .style(Style::default().add_modifier(Modifier::BOLD | Modifier::REVERSED))
99                .alignment(Alignment::Center);
100            f.render_widget(explanation, chunks[1]);
101        })?;
102
103        if event::poll(Duration::from_millis(10))? {
104            tracing::info!("Terminal Size: {:?}", terminal.size());
105            match event::read()? {
106                Event::Key(key) => match key.code {
107                    KeyCode::Char('q') if key.modifiers.contains(KeyModifiers::CONTROL) => {
108                        cleanup_terminal(&mut terminal).unwrap();
109                        return Ok(());
110                    }
111                    KeyCode::Char('n') if key.modifiers.contains(KeyModifiers::CONTROL) => {
112                        let pane_size = calc_pane_size(size, panes.len() + 1);
113                        tracing::info!("Opened new pane with size: {size:?}");
114                        resize_all_panes(&mut panes, pane_size);
115                        open_new_pane(&mut panes, &mut active_pane, &cmd, pane_size)?;
116                    }
117                    KeyCode::Char('x') if key.modifiers.contains(KeyModifiers::CONTROL) => {
118                        close_active_pane(&mut panes, &mut active_pane).await?;
119                        resize_all_panes(&mut panes, pane_size);
120                    }
121                    KeyCode::Char('k') if key.modifiers.contains(KeyModifiers::CONTROL) => {
122                        if let Some(pane) = active_pane {
123                            active_pane = Some(pane.saturating_sub(1));
124                        }
125                    }
126                    KeyCode::Char('j') if key.modifiers.contains(KeyModifiers::CONTROL) => {
127                        if let Some(pane) = active_pane {
128                            if pane < panes.len() - 1 {
129                                active_pane = Some(pane.saturating_add(1));
130                            }
131                        }
132                    }
133                    _ => {
134                        if let Some(index) = active_pane {
135                            if handle_pane_key_event(&mut panes[index], &key).await {
136                                continue;
137                            }
138                        }
139                    }
140                },
141                Event::Resize(cols, rows) => {
142                    tracing::info!("Resized to: rows: {} cols: {}", rows, cols);
143                    size.rows = rows;
144                    size.cols = cols;
145                    let pane_size = calc_pane_size(size, panes.len());
146                    resize_all_panes(&mut panes, pane_size);
147                }
148                _ => {}
149            }
150        }
151    }
152}
153
154fn calc_pane_size(mut size: Size, nr_panes: usize) -> Size {
155    size.rows -= 2;
156    size.rows /= nr_panes as u16;
157    size
158}
159
160fn resize_all_panes(panes: &mut Vec<PtyPane>, size: Size) {
161    for pane in panes.iter() {
162        pane.resize(size);
163    }
164}
165
166struct PtyPane {
167    parser: Arc<RwLock<vt100::Parser>>,
168    sender: Sender<Bytes>,
169    master_pty: Box<dyn MasterPty>,
170}
171
172impl PtyPane {
173    fn new(size: Size, cmd: CommandBuilder) -> io::Result<Self> {
174        let pty_system = native_pty_system();
175        let pty_pair = pty_system
176            .openpty(PtySize {
177                rows: size.rows - 4,
178                cols: size.cols - 4,
179                pixel_width: 0,
180                pixel_height: 0,
181            })
182            .unwrap();
183        let parser = Arc::new(RwLock::new(vt100::Parser::new(
184            size.rows - 4,
185            size.cols - 4,
186            0,
187        )));
188
189        spawn_blocking(move || {
190            let mut child = pty_pair.slave.spawn_command(cmd).unwrap();
191            let _ = child.wait();
192            drop(pty_pair.slave);
193        });
194
195        {
196            let mut reader = pty_pair.master.try_clone_reader().unwrap();
197            let parser = parser.clone();
198            tokio::spawn(async move {
199                let mut processed_buf = Vec::new();
200                let mut buf = [0u8; 8192];
201
202                loop {
203                    let size = reader.read(&mut buf).unwrap();
204                    if size == 0 {
205                        break;
206                    }
207                    if size > 0 {
208                        processed_buf.extend_from_slice(&buf[..size]);
209                        let mut parser = parser.write().unwrap();
210                        parser.process(&processed_buf);
211
212                        // Clear the processed portion of the buffer
213                        processed_buf.clear();
214                    }
215                }
216            });
217        }
218
219        let (tx, mut rx) = channel::<Bytes>(32);
220
221        let mut writer = BufWriter::new(pty_pair.master.take_writer().unwrap());
222        // writer is moved into the tokio task below
223        tokio::spawn(async move {
224            while let Some(bytes) = rx.recv().await {
225                writer.write_all(&bytes).unwrap();
226                writer.flush().unwrap();
227            }
228        });
229
230        Ok(Self {
231            parser,
232            sender: tx,
233            master_pty: pty_pair.master,
234        })
235    }
236
237    fn resize(&self, size: Size) {
238        self.parser
239            .write()
240            .unwrap()
241            .set_size(size.rows - 4, size.cols - 4);
242        self.master_pty
243            .resize(PtySize {
244                rows: size.rows - 4,
245                cols: size.cols - 4,
246                pixel_width: 0,
247                pixel_height: 0,
248            })
249            .unwrap();
250    }
251}
252
253async fn handle_pane_key_event(pane: &mut PtyPane, key: &KeyEvent) -> bool {
254    let input_bytes = match key.code {
255        KeyCode::Char(ch) => {
256            let mut send = vec![ch as u8];
257            let upper = ch.to_ascii_uppercase();
258            if key.modifiers == KeyModifiers::CONTROL {
259                match upper {
260                    'N' => {
261                        // Ignore Ctrl+n within a pane
262                        return true;
263                    }
264                    'X' => {
265                        // Close the pane
266                        return false;
267                    }
268                    // https://github.com/fyne-io/terminal/blob/master/input.go
269                    // https://gist.github.com/ConnerWill/d4b6c776b509add763e17f9f113fd25b
270                    '2' | '@' | ' ' => send = vec![0],
271                    '3' | '[' => send = vec![27],
272                    '4' | '\\' => send = vec![28],
273                    '5' | ']' => send = vec![29],
274                    '6' | '^' => send = vec![30],
275                    '7' | '-' | '_' => send = vec![31],
276                    char if ('A'..='_').contains(&char) => {
277                        // Since A == 65,
278                        // we can safely subtract 64 to get
279                        // the corresponding control character
280                        let ascii_val = char as u8;
281                        let ascii_to_send = ascii_val - 64;
282                        send = vec![ascii_to_send];
283                    }
284                    _ => {}
285                }
286            }
287            send
288        }
289        #[cfg(unix)]
290        KeyCode::Enter => vec![b'\n'],
291        #[cfg(windows)]
292        KeyCode::Enter => vec![b'\r', b'\n'],
293        KeyCode::Backspace => vec![8],
294        KeyCode::Left => vec![27, 91, 68],
295        KeyCode::Right => vec![27, 91, 67],
296        KeyCode::Up => vec![27, 91, 65],
297        KeyCode::Down => vec![27, 91, 66],
298        KeyCode::Tab => vec![9],
299        KeyCode::Home => vec![27, 91, 72],
300        KeyCode::End => vec![27, 91, 70],
301        KeyCode::PageUp => vec![27, 91, 53, 126],
302        KeyCode::PageDown => vec![27, 91, 54, 126],
303        KeyCode::BackTab => vec![27, 91, 90],
304        KeyCode::Delete => vec![27, 91, 51, 126],
305        KeyCode::Insert => vec![27, 91, 50, 126],
306        KeyCode::Esc => vec![27],
307        _ => return true,
308    };
309
310    pane.sender.send(Bytes::from(input_bytes)).await.ok();
311    true
312}
313
314fn open_new_pane(
315    panes: &mut Vec<PtyPane>,
316    active_pane: &mut Option<usize>,
317    cmd: &CommandBuilder,
318    size: Size,
319) -> io::Result<()> {
320    let new_pane = PtyPane::new(size, cmd.clone())?;
321    let new_pane_index = panes.len();
322    panes.push(new_pane);
323    *active_pane = Some(new_pane_index);
324    Ok(())
325}
326
327async fn close_active_pane(
328    panes: &mut Vec<PtyPane>,
329    active_pane: &mut Option<usize>,
330) -> io::Result<()> {
331    if let Some(active_index) = active_pane {
332        let _pane = panes.remove(*active_index);
333        // TODO: shutdown pane correctly
334        if !panes.is_empty() {
335            let remaining_panes = panes.len();
336            let new_active_index = *active_index % remaining_panes;
337            *active_pane = Some(new_active_index);
338        }
339    }
340    Ok(())
341}
342
343fn setup_terminal() -> io::Result<(Terminal<CrosstermBackend<BufWriter<io::Stdout>>>, Size)> {
344    enable_raw_mode()?;
345    let stdout = io::stdout();
346    let backend = CrosstermBackend::new(BufWriter::new(stdout));
347    let mut terminal = Terminal::new(backend)?;
348    let initial_size = terminal.size()?;
349    let size = Size {
350        rows: initial_size.height,
351        cols: initial_size.width,
352    };
353    execute!(terminal.backend_mut(), EnterAlternateScreen)?;
354    Ok((terminal, size))
355}
356
357fn cleanup_terminal(
358    terminal: &mut Terminal<CrosstermBackend<BufWriter<io::Stdout>>>,
359) -> io::Result<()> {
360    execute!(terminal.backend_mut(), LeaveAlternateScreen)?;
361    disable_raw_mode()?;
362    terminal.show_cursor()?;
363    terminal.clear()?;
364    Ok(())
365}
366
367fn init_panic_hook() {
368    let log_file = Some(PathBuf::from("/tmp/tui-term/smux.log"));
369    let log_file = match log_file {
370        Some(path) => {
371            if let Some(parent) = path.parent() {
372                let _ = fs::create_dir_all(parent);
373            }
374            Some(fs::File::create(path).unwrap())
375        }
376        None => None,
377    };
378
379    let subscriber = FmtSubscriber::builder()
380        // all spans/events with a level higher than TRACE (e.g, debug, info, warn, etc.)
381        // will be written to output path.
382        .with_max_level(Level::TRACE)
383        .with_writer(Mutex::new(log_file.unwrap()))
384        .with_thread_ids(true)
385        .with_ansi(true)
386        .with_line_number(true);
387
388    let subscriber = subscriber.finish();
389    tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed");
390
391    // Set the panic hook to log panic information before panicking
392    std::panic::set_hook(Box::new(|panic| {
393        let original_hook = std::panic::take_hook();
394        tracing::error!("Panic Error: {}", panic);
395        crossterm::terminal::disable_raw_mode().expect("Could not disable raw mode");
396        crossterm::execute!(std::io::stdout(), crossterm::terminal::LeaveAlternateScreen)
397            .expect("Could not leave the alternate screen");
398
399        original_hook(panic);
400    }));
401    tracing::debug!("Set panic hook")
402}
403
404impl fmt::Debug for PtyPane {
405    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
406        let parser = self.parser.read().unwrap();
407        let screen = parser.screen();
408
409        f.debug_struct("PtyPane")
410            .field("screen", screen)
411            .field("title:", &screen.title())
412            .field("icon_name:", &screen.icon_name())
413            .finish()
414    }
415}
smux/smux.rs

smux/
smux.rs
