use std::{ ffi::OsString, iter::FromIterator, }; use tokio::sync::watch; mod diceware; mod download; mod ulid; #[derive (Clone, Copy, Debug, PartialEq)] enum Subcommand { Diceware, Download, PtthServer, PtthFileServer, PtthQuicEndServer, Ulid, } #[tokio::main] async fn main () -> anyhow::Result <()> { use Subcommand::*; tracing_subscriber::fmt::init (); let args = Vec::from_iter (std::env::args_os ()); let (subcommand, args) = parse_args (&args)?; match subcommand { Diceware => { diceware::main (); Ok (()) }, Download => download::main (args).await, PtthServer => ptth_server::executable::main (args).await, PtthFileServer => ptth_file_server::main (args).await, PtthQuicEndServer => { let (shutdown_tx, shutdown_rx) = watch::channel (false); ctrlc::set_handler (move || { shutdown_tx.send (true).expect ("Couldn't forward Ctrl+C signal"); })?; tracing::trace! ("Set Ctrl+C handler"); quic_demo::executable_end_server::main (args, Some (shutdown_rx)).await?; Ok (()) } Ulid => ulid::main (args).await, } } fn parse_subcommand (arg: &str) -> Option { use Subcommand::*; let map = vec! [ ("diceware", Diceware), ("download", Download), ("ptth_server", PtthServer), ("ptth_file_server", PtthFileServer), ("ptth_quic_end_server", PtthQuicEndServer), ("ulid", Ulid), ]; let arg = arg.strip_suffix (".exe").unwrap_or (arg); for (suffix, subcommand) in &map { if arg.ends_with (suffix) { return Some (*subcommand); } } None } fn parse_args (args: &[OsString]) -> anyhow::Result <(Subcommand, &[OsString])> { let arg_0 = match args.get (0) { Some (x) => x, None => anyhow::bail! ("arg 0 must be the exe name"), }; let arg_0 = arg_0.to_str ().ok_or_else (|| anyhow::anyhow! ("arg 0 should be valid UTF-8"))?; match parse_subcommand (arg_0) { Some (x) => return Ok ((x, args)), None => (), } let arg_1 = match args.get (1) { Some (x) => x, None => anyhow::bail! ("arg 1 must be the subcommand if arg 0 is not"), }; let arg_1 = arg_1.to_str ().ok_or_else (|| anyhow::anyhow! ("arg 1 subcommand should be valid UTF-8"))?; match parse_subcommand (arg_1) { Some (x) => return Ok ((x, &args [1..])), None => (), } anyhow::bail! ("Subcommand must be either arg 0 (exe name) or arg 1") } #[cfg (test)] mod tests { use super::*; #[test] fn multi_call () -> anyhow::Result <()> { let negative_cases = vec! [ vec! [], vec! ["invalid_exe_name"], vec! ["ptth_multi_call_server"], vec! ["ptth_server.ex"], vec! ["ptth_multi_call_server", "invalid_subcommand"], ]; for input in &negative_cases { let input: Vec <_> = input.iter ().map (OsString::from).collect (); let actual = parse_args (&input); assert! (actual.is_err ()); } let positive_cases = vec! [ (vec! ["ptth_server.exe"], (Subcommand::PtthServer, vec! ["ptth_server.exe"])), (vec! ["ptth_server"], (Subcommand::PtthServer, vec! ["ptth_server"])), (vec! ["ptth_server", "--help"], (Subcommand::PtthServer, vec! ["ptth_server", "--help"])), (vec! ["ptth_file_server"], (Subcommand::PtthFileServer, vec! ["ptth_file_server"])), (vec! ["ptth_quic_end_server", "--help"], (Subcommand::PtthQuicEndServer, vec! ["ptth_quic_end_server", "--help"])), (vec! ["ptth_multi_call_server", "ptth_server"], (Subcommand::PtthServer, vec! ["ptth_server"])), ( vec! [ "ptth_multi_call_server", "ptth_server", "--help" ], ( Subcommand::PtthServer, vec! [ "ptth_server", "--help" ] ) ), ( vec! [ "invalid_exe_name", "ptth_server", "--help" ], ( Subcommand::PtthServer, vec! [ "ptth_server", "--help" ] ) ), ]; for (input, (expected_subcommand, expected_args)) in &positive_cases { let input: Vec <_> = input.iter ().map (OsString::from).collect (); let (actual_subcommand, actual_args) = parse_args (&input)?; assert_eq! (expected_subcommand, &actual_subcommand); assert_eq! (expected_args, actual_args); } Ok (()) } }