From 36a62018e58c23ad21a8910a0e06ef3e44825d80 Mon Sep 17 00:00:00 2001 From: Erica Marigold Date: Thu, 30 Jan 2025 06:57:02 +0000 Subject: [PATCH] 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 | 63 +++++++++++++++++++++++++++++++---------------------- src/ssh.rs | 34 ++++++++++++++++++++--------- 2 files changed, 61 insertions(+), 36 deletions(-) diff --git a/src/main.rs b/src/main.rs index b3513a3..348832b 100644 --- a/src/main.rs +++ b/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,38 +18,42 @@ mod cli; mod components; mod config; mod errors; +mod keycode; mod logging; mod ssh; mod tui; -mod keycode; -const SOCKET_ADDR: LazyLock = LazyLock::new(|| SocketAddr::from(([127, 0, 0, 1], 2222))); -pub static SSH_CONFIG: OnceLock> = 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!("../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 } diff --git a/src/ssh.rs b/src/ssh.rs index 05b133c..9fa78e4 100644 --- a/src/ssh.rs +++ b/src/ssh.rs @@ -5,7 +5,7 @@ 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)] @@ -54,7 +55,10 @@ impl TermWriter { impl Write for TermWriter { #[instrument(skip(self, buf), level = "debug")] fn write(&mut self, buf: &[u8]) -> std::io::Result { - tracing::trace!("Writing {} bytes into SSH terminal writer buffer", buf.len()); + tracing::trace!( + "Writing {} bytes into SSH terminal writer buffer", + buf.len() + ); self.inner.extend(buf); Ok(buf.len()) } @@ -72,14 +76,12 @@ pub struct SshSession { tui: Arc>>, } -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)), @@ -87,10 +89,22 @@ impl SshSession { } } - async fn run_app(app: Arc>, writer: Arc>, tui: Arc>>, session: &Handle, channel_id: ChannelId) -> eyre::Result<()> { + async fn run_app( + app: Arc>, + writer: Arc>, + tui: Arc>>, + 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")) + 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")) } } @@ -126,7 +140,7 @@ impl Handler for SshSession { Err(err) => { tracing::error!("Session exited with error: {err}"); let _ = session_handle.channel_failure(channel_id).await; - }, + } } }); @@ -149,7 +163,7 @@ impl Handler for SshSession { ) -> Result<(), Self::Error> { 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}"));