Reorganized files

src/io/input.rs

@@ -0,0 +1,34 @@
+use tracing::trace;
+
+#[derive(Debug)]
+pub enum Input {
+    Char(char),
+    Up,
+    Down,
+    Delete,
+    Esc,
+    Enter,
+    Backspace,
+    CtrlP,
+    Other,
+}
+
+impl From<&[u8]> for Input {
+    fn from(value: &[u8]) -> Self {
+        match value {
+            [c] if c.is_ascii_graphic() => Input::Char(*c as char),
+            [27] => Input::Esc,
+            [27, 91, 65] => Input::Up,
+            [27, 91, 66] => Input::Down,
+            [27, 91, 51, 126] => Input::Delete,
+            [13] => Input::Enter,
+            // NOTE: Actual char is DLE, this happens to map to ctrl-p
+            [16] => Input::CtrlP,
+            [127] => Input::Backspace,
+            other => {
+                trace!("{other:?}");
+                Input::Other
+            }
+        }
+    }
+}

src/io/mod.rs

@@ -0,0 +1,7 @@
+mod input;
+mod stats;
+mod terminal_handle;
+
+pub use input::Input;
+pub use stats::{Stats, TrackStats};
+pub use terminal_handle::TerminalHandle;

src/io/stats.rs

@@ -0,0 +1,126 @@
+use std::{
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
+};
+
+use pin_project_lite::pin_project;
+use russh::{ChannelStream, server::Msg};
+
+use crate::helper::Unit;
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+#[derive(Debug, Default)]
+pub struct Stats {
+    connections: AtomicUsize,
+    rx: AtomicUsize,
+    tx: AtomicUsize,
+}
+
+impl Stats {
+    pub fn add_connection(&self) {
+        self.connections.fetch_add(1, Ordering::Relaxed);
+    }
+
+    pub fn add_rx_bytes(&self, n: usize) {
+        self.rx.fetch_add(n, Ordering::Relaxed);
+    }
+
+    pub fn add_tx_bytes(&self, n: usize) {
+        self.tx.fetch_add(n, Ordering::Relaxed);
+    }
+
+    pub fn connections(&self) -> usize {
+        self.connections.load(Ordering::Relaxed)
+    }
+
+    pub fn rx(&self) -> Unit {
+        Unit::new(self.rx.load(Ordering::Relaxed), "B")
+    }
+
+    pub fn tx(&self) -> Unit {
+        Unit::new(self.tx.load(Ordering::Relaxed), "B")
+    }
+}
+
+pin_project! {
+    pub struct TrackStats {
+        #[pin]
+        inner: ChannelStream<Msg>,
+        stats: Arc<Stats>,
+    }
+}
+
+impl TrackStats {
+    pub fn new(inner: ChannelStream<Msg>, stats: Arc<Stats>) -> Self {
+        Self { inner, stats }
+    }
+}
+
+impl hyper::rt::Read for TrackStats {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        mut buf: hyper::rt::ReadBufCursor<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        let project = self.project();
+        let n = unsafe {
+            let mut tbuf = tokio::io::ReadBuf::uninit(buf.as_mut());
+            match tokio::io::AsyncRead::poll_read(project.inner, cx, &mut tbuf) {
+                Poll::Ready(Ok(())) => tbuf.filled().len(),
+                other => return other,
+            }
+        };
+
+        project.stats.add_tx_bytes(n);
+
+        unsafe {
+            buf.advance(n);
+        }
+        Poll::Ready(Ok(()))
+    }
+}
+
+impl hyper::rt::Write for TrackStats {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        let project = self.project();
+        tokio::io::AsyncWrite::poll_write(project.inner, cx, buf).map(|res| {
+            res.inspect(|n| {
+                project.stats.add_rx_bytes(*n);
+            })
+        })
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
+        tokio::io::AsyncWrite::poll_flush(self.project().inner, cx)
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        tokio::io::AsyncWrite::poll_shutdown(self.project().inner, cx)
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        tokio::io::AsyncWrite::is_write_vectored(&self.inner)
+    }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[std::io::IoSlice<'_>],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        let project = self.project();
+        tokio::io::AsyncWrite::poll_write_vectored(project.inner, cx, bufs).map(|res| {
+            res.inspect(|n| {
+                project.stats.add_rx_bytes(*n);
+            })
+        })
+    }
+}

src/io/terminal_handle.rs

@@ -0,0 +1,68 @@
+use crossterm::{
+    execute,
+    terminal::{EnterAlternateScreen, LeaveAlternateScreen},
+};
+use russh::{ChannelId, server::Handle};
+use tokio::sync::mpsc::{UnboundedSender, unbounded_channel};
+use tracing::error;
+
+pub struct TerminalHandle {
+    sender: UnboundedSender<Vec<u8>>,
+    sink: Vec<u8>,
+}
+
+impl TerminalHandle {
+    pub async fn start(handle: Handle, channel_id: ChannelId) -> std::io::Result<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.into()).await;
+
+                if let Err(e) = result {
+                    error!("Failed to send data: {e:?}");
+                };
+            }
+
+            if let Err(e) = handle.close(channel_id).await {
+                error!("Failed to close session: {e:?}");
+            }
+        });
+
+        let mut terminal_handle = Self {
+            sender,
+            sink: Vec::new(),
+        };
+
+        execute!(terminal_handle, EnterAlternateScreen)?;
+
+        Ok(terminal_handle)
+    }
+
+    pub fn leave_alternate_screen(&mut self) -> std::io::Result<()> {
+        execute!(self, LeaveAlternateScreen)
+    }
+}
+
+impl Drop for TerminalHandle {
+    fn drop(&mut self) {
+        self.leave_alternate_screen().ok();
+    }
+}
+
+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 let Err(e) = result {
+            return Err(std::io::Error::new(std::io::ErrorKind::BrokenPipe, e));
+        }
+
+        self.sink.clear();
+        Ok(())
+    }
+}