ptth/crates/ptth_multi_call_server/src/main.rs

use std::{
	ffi::OsString,
	iter::FromIterator,
};

use tokio::sync::watch;

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 => {
			ptth_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");
			ptth_quic::executable_end_server::main (args, Some (shutdown_rx)).await?;
			
			Ok (())
		}
		Ulid => ulid::main (args).await,
	}
}

fn parse_subcommand (arg: &str) -> Option <Subcommand>
{
	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 (())
	}
}