Skip to main content

tailscale/ssh/
ratatui.rs

1use std::sync::Arc;
2
3use ratatui::{Terminal, TerminalOptions, Viewport, backend::CrosstermBackend, layout::Rect};
4use russh::{ChannelId, Sig, server::Handle};
5
6use crate::{
7    Device,
8    ssh::{ChannelEvent, ChannelHandler, channel_write::ChannelWrite},
9};
10
11type Backend = CrosstermBackend<ChannelWrite>;
12
13/// Terminal environment for [`RatatuiApp`].
14pub trait RatatuiEnv {
15    /// Request that the terminal close.
16    fn close(&self) -> impl Future<Output = ()> + Send;
17
18    /// Get a reference to the Tailscale [`Device`] this is running in.
19    fn tailscale(&self) -> &Device;
20}
21
22/// A [`ratatui`] application designed to be driven by a
23/// [`ChannelServer`][crate::ssh::ChannelServer].
24pub trait RatatuiApp {
25    /// Process new input from the channel.
26    fn input(
27        &mut self,
28        data: &[u8],
29        env: impl RatatuiEnv + Send,
30    ) -> impl Future<Output = ()> + Send;
31
32    /// Render the app to the [`ratatui::Frame`].
33    fn draw(&mut self, frame: &mut ratatui::Frame);
34}
35
36/// A [`ChannelHandler`] that runs a [`RatatuiApp`].
37pub struct RatatuiTerm<Io> {
38    channel_id: ChannelId,
39    session: Handle,
40    term: Terminal<Backend>,
41    dev: Arc<Device>,
42    io: Io,
43}
44
45struct Env<'a> {
46    channel_id: ChannelId,
47    session: &'a Handle,
48    dev: &'a Device,
49}
50
51impl RatatuiEnv for Env<'_> {
52    async fn close(&self) {
53        if self.session.close(self.channel_id).await.is_err() {
54            tracing::error!("channel closed while closing ratatui app");
55        }
56    }
57
58    fn tailscale(&self) -> &Device {
59        self.dev
60    }
61}
62
63impl<Io> RatatuiTerm<Io>
64where
65    Io: RatatuiApp,
66{
67    fn refresh(&mut self) -> std::io::Result<()> {
68        self.term.clear()?;
69        self.draw()?;
70
71        Ok(())
72    }
73
74    fn draw(&mut self) -> std::io::Result<()> {
75        self.term.draw(|frame| self.io.draw(frame))?;
76
77        Ok(())
78    }
79}
80
81impl<Io> ChannelHandler for RatatuiTerm<Io>
82where
83    Io: RatatuiApp + Default + Send,
84{
85    type Error = std::io::Error;
86
87    fn new(
88        rt: tokio::runtime::Handle,
89        channel_id: ChannelId,
90        session: Handle,
91        dev: Arc<Device>,
92        // The TUI demo handler ignores the policy-mapped local user; it runs purely in-process.
93        _accept: &crate::ssh::SshAccept,
94    ) -> Result<Self, Self::Error> {
95        let mut term = Self {
96            term: make_term(rt, session.clone(), channel_id)?,
97            dev,
98            channel_id,
99            session,
100            io: Default::default(),
101        };
102        term.refresh()?;
103
104        Ok(term)
105    }
106
107    async fn handle_event(&mut self, event: &ChannelEvent) -> Result<(), Self::Error> {
108        match event {
109            ChannelEvent::Data(d) => {
110                self.io
111                    .input(
112                        d,
113                        Env {
114                            dev: &self.dev,
115                            channel_id: self.channel_id,
116                            session: &self.session,
117                        },
118                    )
119                    .await;
120
121                self.draw()?;
122            }
123            ChannelEvent::Resize { width, height } => {
124                self.term.resize(Rect::new(0, 0, *width, *height))?;
125                self.draw()?;
126            }
127            ChannelEvent::Eof
128            | ChannelEvent::Signal(Sig::ABRT | Sig::QUIT | Sig::TERM | Sig::KILL | Sig::INT) => {
129                tracing::debug!(?event, channel_id = %self.channel_id, "close channel");
130
131                if self.session.close(self.channel_id).await.is_err() {
132                    tracing::error!("session already shut down");
133
134                    return Err(std::io::ErrorKind::BrokenPipe.into());
135                }
136            }
137            ChannelEvent::Signal(sig) => {
138                tracing::debug!(?sig, "unhandled signal");
139            }
140            ChannelEvent::Close => {
141                self.term.clear()?;
142            }
143        }
144
145        Ok(())
146    }
147}
148
149fn make_term(
150    rt: tokio::runtime::Handle,
151    session_handle: Handle,
152    channel_id: ChannelId,
153) -> Result<Terminal<Backend>, <Backend as ratatui::backend::Backend>::Error> {
154    let terminal_handle = ChannelWrite::new(rt, session_handle, channel_id);
155    let backend = CrosstermBackend::new(terminal_handle);
156
157    let options = TerminalOptions {
158        viewport: Viewport::Fixed(Rect::default()),
159    };
160
161    Terminal::with_options(backend, options)
162}