refactor: accept CLI args as options and minor restructuring

* Made constructed `App` instances use CLI args for `frame_rate` and
  `tick_rate`
* Moved lazy constants from `LazyLock` to `lazy_static!`
* Slightly restructure SSH config initialization
* Also apply rustfmt formatting to some files

src/main.rs

@@ -1,7 +1,14 @@
-use std::{net::SocketAddr, sync::{Arc, LazyLock, OnceLock}};
+use std::{net::SocketAddr, sync::Arc};
 
+use clap::Parser as _;
+use cli::Cli;
 use color_eyre::{eyre::eyre, Result};
-use russh::{keys::PrivateKey, server::{Config, Server}, MethodSet};
+use lazy_static::lazy_static;
+use russh::{
+    keys::PrivateKey,
+    server::{Config, Server},
+    MethodSet,
+};
 use ssh::SshServer;
 use tokio::net::TcpListener;
 
@@ -11,39 +18,42 @@ mod cli;
 mod components;
 mod config;
 mod errors;
+mod keycode;
 mod logging;
 mod ssh;
 mod tui;
-mod keycode;
 
-const SOCKET_ADDR: LazyLock<SocketAddr> = LazyLock::new(|| SocketAddr::from(([127, 0, 0, 1], 2222)));
-pub static SSH_CONFIG: OnceLock<Arc<Config>> = OnceLock::new();
+const SSH_KEYS: &[&[u8]] = &[
+    include_bytes!("../rsa.pem"),
+    include_bytes!("../ed25519.pem"),
+];
+lazy_static! {
+    pub(crate) static ref OPTIONS: Cli = Cli::parse();
+    pub(crate) static ref SOCKET_ADDR: SocketAddr = SocketAddr::from(([127, 0, 0, 1], 2222));
+}
 
 #[tokio::main]
 async fn main() -> Result<()> {
     crate::errors::init()?;
     crate::logging::init()?;
 
-    SSH_CONFIG.get_or_init(|| {
-        tracing::debug!("setting up ssh config");
-
-        let mut conf = Config::default();
-        conf.methods = MethodSet::NONE;
-        conf.keys = vec![
-            PrivateKey::from_openssh(include_bytes!("../rsa.pem")).unwrap(),
-            // PrivateKey::from_openssh(include_bytes!("../ecdsa.pem")).unwrap(),
-            PrivateKey::from_openssh(include_bytes!("../ed25519.pem")).unwrap()
-        ];
-        Arc::new(conf)
-    });
-
-    // let args = Cli::parse();
-    // let mut app = App::new(args.tick_rate, args.frame_rate)?;
-    // app.run().await?;
-    
-    tracing::info!("attemping to listen on {}", *SOCKET_ADDR);
-    SshServer::default().run_on_socket(
-        Arc::clone(SSH_CONFIG.get().unwrap()),
-        &TcpListener::bind(*SOCKET_ADDR).await?,
-    ).await.map_err(|err| eyre!(err))
+    let config = ssh_config();
+    tracing::info!("Attempting to listen on {}", *SOCKET_ADDR);
+    SshServer::default()
+        .run_on_socket(Arc::new(config), &TcpListener::bind(*SOCKET_ADDR).await?)
+        .await
+        .map_err(|err| eyre!(err))
+}
+
+fn ssh_config() -> Config {
+    let mut conf = Config::default();
+    conf.methods = MethodSet::NONE;
+    conf.keys = SSH_KEYS
+        .to_vec()
+        .iter()
+        .filter_map(|pem| PrivateKey::from_openssh(pem).ok())
+        .collect();
+    
+    tracing::trace!("SSH config: {:#?}", conf);
+    conf
 }

src/ssh.rs

@@ -1,11 +1,11 @@
 use std::{io::Write, net::SocketAddr, sync::Arc};
 
 use async_trait::async_trait;
-use color_eyre::eyre::eyre;
+use color_eyre::eyre::{self, eyre};
 use ratatui::prelude::CrosstermBackend;
 use russh::{
     server::{Auth, Handle, Handler, Msg, Server, Session},
-    Channel, ChannelId, CryptoVec, Pty, 
+    Channel, ChannelId, CryptoVec, Pty,
 };
 use tokio::{
     runtime::Handle as TokioHandle,
@@ -16,6 +16,7 @@ use tracing::instrument;
 use crate::{
     app::App,
     tui::{Terminal, Tui},
+    OPTIONS,
 };
 
 #[derive(Debug)]
@@ -27,7 +28,9 @@ pub struct TermWriter {
 }
 
 impl TermWriter {
+    #[instrument(skip(session, channel), level = "trace", fields(channel_id = %channel.id()))]
     fn new(session: &mut Session, channel: Channel<Msg>) -> Self {
+        tracing::trace!("Acquiring new SSH writer");
         Self {
             session: session.handle(),
             channel,
@@ -52,12 +55,17 @@ impl TermWriter {
 impl Write for TermWriter {
     #[instrument(skip(self, buf), level = "debug")]
     fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        tracing::trace!(
+            "Writing {} bytes into SSH terminal writer buffer",
+            buf.len()
+        );
         self.inner.extend(buf);
         Ok(buf.len())
     }
 
     #[instrument(skip(self), level = "trace")]
     fn flush(&mut self) -> std::io::Result<()> {
+        tracing::trace!("Flushing SSH terminal writer buffer");
         tokio::task::block_in_place(|| self.flush_inner())
     }
 }
@@ -68,32 +76,48 @@ pub struct SshSession {
     tui: Arc<RwLock<Option<Tui>>>,
 }
 
-unsafe impl Send for SshSession {}
-
 impl SshSession {
     pub fn new() -> Self {
         let (ssh_tx, ssh_rx) = mpsc::unbounded_channel();
 
         Self {
-            app: App::new(10f64, 60f64, ssh_rx)
+            app: App::new(OPTIONS.tick_rate, OPTIONS.frame_rate, ssh_rx)
                 .ok()
                 .map(|app| Arc::new(Mutex::new(app))),
             tui: Arc::new(RwLock::new(None)),
             ssh_tx,
         }
     }
+
+    async fn run_app(
+        app: Arc<Mutex<App>>,
+        writer: Arc<Mutex<Terminal>>,
+        tui: Arc<RwLock<Option<Tui>>>,
+        session: &Handle,
+        channel_id: ChannelId,
+    ) -> eyre::Result<()> {
+        app.lock_owned().await.run(writer, tui).await?;
+        session
+            .close(channel_id)
+            .await
+            .map_err(|_| eyre!("failed to close session"))?;
+        session
+            .exit_status_request(channel_id, 0)
+            .await
+            .map_err(|_| eyre!("failed to send session exit status"))
+    }
 }
 
 #[async_trait]
 impl Handler for SshSession {
-    type Error = color_eyre::eyre::Error;
+    type Error = eyre::Error;
 
     #[instrument(skip(self), span = "user_login", fields(method = "none"))]
     async fn auth_none(&mut self, user: &str) -> Result<Auth, Self::Error> {
         Ok(Auth::Accept)
     }
 
-    #[instrument(skip(self), span = "channel_establish", level = "trace")]
+    #[instrument(skip(self, session, channel), span = "channel_establish", fields(channel_id = %channel.id()))]
     async fn channel_open_session(
         &mut self,
         channel: Channel<Msg>,
@@ -108,10 +132,16 @@ impl Handler for SshSession {
             let writer = Arc::new(Mutex::new(term));
             let tui = Arc::clone(&self.tui);
 
+            tracing::info!("Serving app to open session");
             tokio::task::spawn(async move {
-                inner_app.lock_owned().await.run(writer, tui).await.unwrap();
-                session_handle.close(channel_id).await.unwrap();
-                session_handle.exit_status_request(channel_id, 0).await.unwrap();
+                let res = Self::run_app(inner_app, writer, tui, &session_handle, channel_id).await;
+                match res {
+                    Ok(()) => tracing::info!("session exited"),
+                    Err(err) => {
+                        tracing::error!("Session exited with error: {err}");
+                        let _ = session_handle.channel_failure(channel_id).await;
+                    }
+                }
             });
 
             return Ok(true);
@@ -119,7 +149,7 @@ impl Handler for SshSession {
 
         Err(eyre!("Failed to initialize App for session"))
     }
-    #[instrument(skip(self, session), level = "trace")]
+    #[instrument(skip_all, fields(channel_id = %channel_id))]
     async fn pty_request(
         &mut self,
         channel_id: ChannelId,
@@ -128,12 +158,12 @@ impl Handler for SshSession {
         row_height: u32,
         pix_width: u32,
         pix_height: u32,
-        modes: &[(Pty, u32)],
+        _modes: &[(Pty, u32)],
         session: &mut Session,
     ) -> Result<(), Self::Error> {
-        tracing::info!("Received pty request from channel {channel_id}; terminal: {term}");
+        tracing::info!("PTY requested by terminal: {term}");
         tracing::debug!("dims: {col_width} * {row_height}, pixel: {pix_width} * {pix_height}");
-        
+
         if !term.contains("xterm") {
             session.channel_failure(channel_id)?;
             return Err(eyre!("Unsupported terminal type: {term}"));

src/tui/mod.rs

@@ -243,6 +243,6 @@ impl Tui {
 
 impl Drop for Tui {
     fn drop(&mut self) {
-        self.exit().unwrap();
+        let _ = self.exit().inspect_err(|err| error!("Failed to exit Tui: {err}"));
     }
 }