ram usage with ssh server and tui

.gitignore

@@ -0,0 +1 @@
+/target

Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "ssh-status"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+anyhow = "1.0.102"
+crossterm = "0.29.0"
+env_logger = "0.11.10"
+log = "0.4.29"
+rand = "0.10.0"
+ratatui = "0.30.0"
+russh = "0.60.0"
+sysinfo = "0.38.4"
+termios = "0.3.3"
+tokio = "1.51.1"

src/app.rs

@@ -0,0 +1,94 @@
+use ratatui::backend::CrosstermBackend;
+use ratatui::layout::Rect;
+use ratatui::style::{Color, Style};
+use ratatui::symbols::line;
+use ratatui::widgets::{Block, Borders, Clear, LineGauge, Paragraph};
+use ratatui::{Terminal, TerminalOptions, Viewport};
+
+use sysinfo::{Components, Disks, Networks, System};
+
+#[derive(Debug, PartialEq)]
+pub enum Action {
+    Continue,
+    Stop,
+}
+
+pub struct App<W: std::io::Write> {
+    pub terminal: Terminal<CrosstermBackend<W>>,
+    pub counter: usize,
+}
+
+impl<W: std::io::Write> App<W> {
+    pub fn new(terminal_handle: W) -> Result<Self, anyhow::Error> {
+        let backend = CrosstermBackend::new(terminal_handle);
+
+        // the correct viewport area will be set when the client request a pty
+        let options = TerminalOptions {
+            viewport: Viewport::Fixed(Rect::default()),
+        };
+
+        let terminal = Terminal::with_options(backend, options)?;
+
+        Ok(Self {
+            terminal,
+            counter: 0,
+        })
+    }
+
+    pub fn input(&mut self, data: &[u8]) -> Action {
+        match data {
+            // Pressing 'q' closes the connection.
+            b"q" => Action::Stop,
+            // Pressing 'c' resets the counter for the app.
+            // Only the client with the id sees the counter reset.
+            b"c" => {
+                self.counter = 0;
+                Action::Continue
+            }
+            _ => Action::Continue,
+        }
+    }
+
+    pub fn update(&mut self) {
+        if (self.counter % 100) == 0 {
+            let mut sys = System::new_all();
+            sys.refresh_all();
+
+            self.terminal
+                .draw(|f| {
+                    let area = f.area();
+                    f.render_widget(Clear, area);
+                    let style = Style::default().fg(Color::Red);
+                    let block = Block::default()
+                        .title("Press 'q' to quit")
+                        .borders(Borders::ALL);
+
+                    let mem_total: f64 = ((sys.total_memory() >> 20) as f64) / 1024.0;
+                    let mem_used: f64 = ((sys.used_memory() >> 20) as f64) / 1024.0;
+
+                    let line_gauge = LineGauge::default()
+                        .block(Block::bordered().title("RAM usage"))
+                        .label(format!("{mem_used:.2}/{mem_total:.2}Go"))
+                        .ratio(mem_used / mem_total)
+                        .filled_symbol("=")
+                        .unfilled_symbol(" ");
+                    f.render_widget(line_gauge.block(block), area);
+                })
+                .unwrap();
+        }
+
+        self.counter += 1;
+    }
+
+    pub fn resize(&mut self, col: u32, row: u32) -> Result<(), anyhow::Error> {
+        let rect = Rect {
+            x: 0,
+            y: 0,
+            width: col as u16,
+            height: row as u16,
+        };
+
+        self.terminal.resize(rect)?;
+        Ok(())
+    }
+}

src/bin/ssh.rs

@@ -0,0 +1,7 @@
+use ssh_status::server;
+
+#[tokio::main]
+async fn main() {
+    let mut server = server::AppServer::new();
+    server.run().await.expect("Failed running server");
+}

src/bin/tui.rs

@@ -0,0 +1,57 @@
+//extern crate termios;
+use ssh_status::app::{Action, App};
+use std::io::{self, Write};
+use std::time::Duration;
+
+use crossterm::cursor;
+use crossterm::event::{DisableMouseCapture, EnableMouseCapture, Event, KeyCode};
+use crossterm::terminal::{EnterAlternateScreen, LeaveAlternateScreen};
+
+fn main() {
+    let mut stdout = io::stdout();
+    crossterm::terminal::enable_raw_mode().unwrap();
+    crossterm::execute!(
+        stdout,
+        EnterAlternateScreen,
+        EnableMouseCapture,
+        cursor::Hide
+    )
+    .unwrap();
+
+    let mut app = App::new(stdout).unwrap();
+    let (w, h) = crossterm::terminal::size().unwrap();
+    app.resize(w as u32, h as u32).unwrap();
+    app.update();
+
+    loop {
+        if crossterm::event::poll(Duration::from_millis(10)).unwrap() {
+            match crossterm::event::read().unwrap() {
+                Event::Key(key) => match key.code {
+                    KeyCode::Char(c) => {
+                        let buffer: [u8; 1] = [c as u8];
+                        match app.input(&buffer) {
+                            Action::Stop => break,
+                            _ => (),
+                        }
+                    }
+                    _ => (),
+                },
+                Event::Resize(w, h) => app.resize(w as u32, h as u32).unwrap(),
+                _ => (),
+            }
+        };
+
+        app.update();
+    }
+
+    let mut stdout = io::stdout();
+    crossterm::terminal::disable_raw_mode().unwrap();
+    stdout.flush().unwrap();
+    crossterm::execute!(
+        stdout,
+        LeaveAlternateScreen,
+        DisableMouseCapture,
+        cursor::Show
+    )
+    .unwrap();
+}

src/lib.rs

@@ -0,0 +1,2 @@
+pub mod app;
+pub mod server;

src/server.rs

@@ -0,0 +1,202 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use russh::keys::{Algorithm, ssh_key::PublicKey};
+use russh::server::*;
+use russh::{Channel, ChannelId, Pty};
+use tokio::sync::Mutex;
+use tokio::sync::mpsc::{UnboundedSender, unbounded_channel};
+
+use crate::app::{Action, App};
+
+#[derive(Clone)]
+pub struct AppServer {
+    clients: Arc<Mutex<HashMap<usize, App<TerminalHandle>>>>,
+    id: usize,
+}
+
+impl AppServer {
+    pub fn new() -> Self {
+        Self {
+            clients: Arc::new(Mutex::new(HashMap::new())),
+            id: 0,
+        }
+    }
+
+    pub async fn run(&mut self) -> Result<(), anyhow::Error> {
+        let clients = self.clients.clone();
+        tokio::spawn(async move {
+            loop {
+                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
+
+                for (_, app) in clients.lock().await.iter_mut() {
+                    app.update();
+                }
+            }
+        });
+
+        let config = Config {
+            inactivity_timeout: Some(std::time::Duration::from_secs(3600)),
+            auth_rejection_time: std::time::Duration::from_secs(3),
+            auth_rejection_time_initial: Some(std::time::Duration::from_secs(0)),
+            keys: vec![
+                russh::keys::PrivateKey::random(&mut rand::rng(), Algorithm::Ed25519).unwrap(),
+            ],
+            nodelay: true,
+            ..Default::default()
+        };
+
+        self.run_on_address(Arc::new(config), ("0.0.0.0", 2222))
+            .await?;
+        Ok(())
+    }
+}
+
+impl Server for AppServer {
+    type Handler = Self;
+    fn new_client(&mut self, _: Option<std::net::SocketAddr>) -> Self {
+        let s = self.clone();
+        self.id += 1;
+        s
+    }
+}
+
+impl Handler for AppServer {
+    type Error = anyhow::Error;
+
+    async fn channel_open_session(
+        &mut self,
+        channel: Channel<Msg>,
+        session: &mut Session,
+    ) -> Result<bool, Self::Error> {
+        let terminal_handle = TerminalHandle::start(session.handle(), channel.id()).await;
+        let app = App::new(terminal_handle)?;
+
+        let mut clients = self.clients.lock().await;
+        clients.insert(self.id, app);
+
+        Ok(true)
+    }
+
+    async fn auth_publickey(&mut self, _: &str, _: &PublicKey) -> Result<Auth, Self::Error> {
+        Ok(Auth::Accept)
+    }
+
+    async fn data(
+        &mut self,
+        channel: ChannelId,
+        data: &[u8],
+        session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        let mut clients = self.clients.lock().await;
+        let app = clients.get_mut(&self.id).unwrap();
+        let action = app.input(data);
+
+        match action {
+            Action::Stop => {
+                clients.remove(&self.id);
+                session.close(channel)?;
+                Ok(())
+            }
+            Action::Continue => Ok(()),
+        }
+    }
+
+    /// The client's window size has changed.
+    async fn window_change_request(
+        &mut self,
+        _: ChannelId,
+        col_width: u32,
+        row_height: u32,
+        _: u32,
+        _: u32,
+        _: &mut Session,
+    ) -> Result<(), Self::Error> {
+        let mut clients = self.clients.lock().await;
+        let app = clients.get_mut(&self.id).unwrap();
+        app.resize(col_width, row_height)?;
+
+        Ok(())
+    }
+
+    /// The client requests a pseudo-terminal with the given
+    /// specifications.
+    ///
+    /// **Note:** Success or failure should be communicated to the client by calling
+    /// `session.channel_success(channel)` or `session.channel_failure(channel)` respectively.
+    async fn pty_request(
+        &mut self,
+        channel: ChannelId,
+        _: &str,
+        col_width: u32,
+        row_height: u32,
+        _: u32,
+        _: u32,
+        _: &[(Pty, u32)],
+        session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        let mut clients = self.clients.lock().await;
+        let app = clients.get_mut(&self.id).unwrap();
+        app.resize(col_width, row_height)?;
+
+        session.channel_success(channel)?;
+
+        Ok(())
+    }
+}
+
+impl Drop for AppServer {
+    fn drop(&mut self) {
+        let id = self.id;
+        let clients = self.clients.clone();
+        tokio::spawn(async move {
+            let mut clients = clients.lock().await;
+            clients.remove(&id);
+        });
+    }
+}
+
+pub struct TerminalHandle {
+    sender: UnboundedSender<Vec<u8>>,
+    // The sink collects the data which is finally sent to sender.
+    sink: Vec<u8>,
+}
+
+impl TerminalHandle {
+    pub async fn start(handle: Handle, channel_id: ChannelId) -> Self {
+        let (sender, mut receiver) = unbounded_channel::<Vec<u8>>();
+        tokio::spawn(async move {
+            while let Some(data) = receiver.recv().await {
+                let result = handle.data(channel_id, data).await;
+                if result.is_err() {
+                    eprintln!("Failed to send data: {result:?}");
+                }
+            }
+        });
+        Self {
+            sender,
+            sink: Vec::new(),
+        }
+    }
+}
+
+// The crossterm backend writes to the terminal handle.
+impl std::io::Write for TerminalHandle {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.sink.extend_from_slice(buf);
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        let result = self.sender.send(self.sink.clone());
+        if result.is_err() {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::BrokenPipe,
+                result.unwrap_err(),
+            ));
+        }
+
+        self.sink.clear();
+        Ok(())
+    }
+}