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    // The TUI runs in-process and streams nothing to a recorder, so a policy rule that demands
88    // session recording must not reach it — the connection is refused by `ChannelServer`'s
89    // fail-closed gate instead.
90    const RECORDS_SESSION: bool = false;
91
92    async fn new(
93        rt: tokio::runtime::Handle,
94        channel_id: ChannelId,
95        session: Handle,
96        dev: Arc<Device>,
97        // The TUI demo handler ignores the policy-mapped local user; it runs purely in-process.
98        _ctx: &crate::ssh::ChannelContext,
99    ) -> Result<Self, Self::Error> {
100        let mut term = Self {
101            term: make_term(rt, session.clone(), channel_id)?,
102            dev,
103            channel_id,
104            session,
105            io: Default::default(),
106        };
107        term.refresh()?;
108
109        Ok(term)
110    }
111
112    async fn handle_event(&mut self, event: &ChannelEvent) -> Result<(), Self::Error> {
113        match event {
114            ChannelEvent::Data(d) => {
115                self.io
116                    .input(
117                        d,
118                        Env {
119                            dev: &self.dev,
120                            channel_id: self.channel_id,
121                            session: &self.session,
122                        },
123                    )
124                    .await;
125
126                self.draw()?;
127            }
128            ChannelEvent::Resize { width, height } => {
129                self.term.resize(Rect::new(0, 0, *width, *height))?;
130                self.draw()?;
131            }
132            ChannelEvent::Eof
133            | ChannelEvent::Signal(Sig::ABRT | Sig::QUIT | Sig::TERM | Sig::KILL | Sig::INT) => {
134                tracing::debug!(?event, channel_id = %self.channel_id, "close channel");
135
136                if self.session.close(self.channel_id).await.is_err() {
137                    tracing::error!("session already shut down");
138
139                    return Err(std::io::ErrorKind::BrokenPipe.into());
140                }
141            }
142            ChannelEvent::Signal(sig) => {
143                tracing::debug!(?sig, "unhandled signal");
144            }
145            ChannelEvent::Close => {
146                self.term.clear()?;
147            }
148        }
149
150        Ok(())
151    }
152}
153
154fn make_term(
155    rt: tokio::runtime::Handle,
156    session_handle: Handle,
157    channel_id: ChannelId,
158) -> Result<Terminal<Backend>, <Backend as ratatui::backend::Backend>::Error> {
159    let terminal_handle = ChannelWrite::new(rt, session_handle, channel_id);
160    let backend = CrosstermBackend::new(terminal_handle);
161
162    let options = TerminalOptions {
163        viewport: Viewport::Fixed(Rect::default()),
164    };
165
166    Terminal::with_options(backend, options)
167}