mod auth; mod daemon; mod protocol; mod session; use std::path::PathBuf; use tokio::signal::unix::{signal, SignalKind}; use tokio::sync::broadcast; use auth::AuthToken; use daemon::Daemon; const VERSION: &str = "0.1.0"; // --------------------------------------------------------------------------- // CLI argument parsing — no clap needed for 3 flags. // --------------------------------------------------------------------------- struct Cli { socket_dir: Option, default_shell: Option, verbose: bool, } impl Cli { fn parse() -> Result { let mut args = std::env::args().skip(1).peekable(); let mut socket_dir = None; let mut default_shell = None; let mut verbose = false; while let Some(arg) = args.next() { match arg.as_str() { "--socket-dir" => { socket_dir = Some(PathBuf::from( args.next().ok_or("--socket-dir requires a value")?, )); } "--shell" => { default_shell = Some( args.next().ok_or("--shell requires a value")?, ); } "--verbose" | "-v" => { verbose = true; } "--help" | "-h" => { print_usage(); std::process::exit(0); } other => { return Err(format!("unknown argument: {other}")); } } } Ok(Self { socket_dir, default_shell, verbose, }) } } fn print_usage() { eprintln!( "USAGE: agor-ptyd [OPTIONS]\n\ \n\ OPTIONS:\n\ --socket-dir Socket directory\n\ (default: /run/user/$UID/agor or ~/.local/share/agor/run)\n\ --shell Default shell (default: $SHELL or /bin/bash)\n\ --verbose Enable debug logging\n\ --help Show this message" ); } // --------------------------------------------------------------------------- // Socket directory resolution // --------------------------------------------------------------------------- fn resolve_socket_dir(override_path: Option) -> Result { if let Some(p) = override_path { return Ok(p); } // Prefer XDG runtime dir. if let Ok(uid_str) = std::env::var("UID").or_else(|_| { // UID is not always exported; fall back to getuid(). Ok::<_, std::env::VarError>(unsafe { libc_getuid() }.to_string()) }) { let xdg = PathBuf::from(format!("/run/user/{uid_str}/agor")); if xdg.parent().map(|p| p.exists()).unwrap_or(false) { return Ok(xdg); } } // Fallback: ~/.local/share/agor/run let home = std::env::var("HOME").map_err(|_| "HOME not set".to_string())?; Ok(PathBuf::from(home).join(".local/share/agor/run")) } #[cfg(target_os = "linux")] unsafe fn libc_getuid() -> u32 { // Safety: getuid() is always safe. extern "C" { fn getuid() -> u32; } getuid() } #[cfg(not(target_os = "linux"))] unsafe fn libc_getuid() -> u32 { 0 } // --------------------------------------------------------------------------- // Default shell resolution // --------------------------------------------------------------------------- fn resolve_shell(override_shell: Option) -> String { override_shell .or_else(|| std::env::var("SHELL").ok()) .filter(|s| !s.is_empty()) .unwrap_or_else(|| "/bin/bash".into()) } // --------------------------------------------------------------------------- // Entry point // --------------------------------------------------------------------------- #[tokio::main] async fn main() { let cli = match Cli::parse() { Ok(c) => c, Err(e) => { eprintln!("error: {e}"); print_usage(); std::process::exit(1); } }; // Initialise logging. let log_level = if cli.verbose { "debug" } else { "info" }; env_logger::Builder::from_env( env_logger::Env::default().default_filter_or(log_level), ) .init(); log::info!("agor-ptyd v{VERSION} starting"); let socket_dir = match resolve_socket_dir(cli.socket_dir) { Ok(d) => d, Err(e) => { log::error!("cannot resolve socket directory: {e}"); std::process::exit(1); } }; // Ensure the directory exists. if let Err(e) = std::fs::create_dir_all(&socket_dir) { log::error!("cannot create socket directory {socket_dir:?}: {e}"); std::process::exit(1); } let socket_path = socket_dir.join("ptyd.sock"); let shell = resolve_shell(cli.default_shell); // Generate and persist auth token. let token = match AuthToken::generate_and_persist(&socket_dir) { Ok(t) => t, Err(e) => { log::error!("token generation failed: {e}"); std::process::exit(1); } }; // Shutdown broadcast channel — one sender, N receivers. let (shutdown_tx, shutdown_rx) = broadcast::channel::<()>(1); // Signal handlers for SIGTERM and SIGINT. let shutdown_tx_sigterm = shutdown_tx.clone(); let shutdown_tx_sigint = shutdown_tx.clone(); tokio::spawn(async move { let mut sigterm = signal(SignalKind::terminate()).expect("SIGTERM handler"); sigterm.recv().await; log::info!("SIGTERM received"); let _ = shutdown_tx_sigterm.send(()); }); tokio::spawn(async move { let mut sigint = signal(SignalKind::interrupt()).expect("SIGINT handler"); sigint.recv().await; log::info!("SIGINT received"); let _ = shutdown_tx_sigint.send(()); }); let daemon = Daemon::new(socket_path, token, shell); if let Err(e) = daemon.run(shutdown_rx).await { log::error!("daemon exited with error: {e}"); std::process::exit(1); } log::info!("agor-ptyd shut down cleanly"); }