smux/
smux.rs

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

smux/
smux.rs
