use std::{collections::HashSet, path::PathBuf}; use clap::{command, Parser, Subcommand}; use error_stack::{Report, ResultExt}; use sqlx::PgPool; use tokio::sync::mpsc::UnboundedReceiver; use tokio_util::sync::CancellationToken; use crate::{ api::{ self, routes::api_keys::{models::ApiKey, sql::create_api_key}, }, config::Configuration, errors::AppError, APP_NAME, ROOT_PATH, }; #[derive(Parser, Debug)] #[command(name = APP_NAME, version = "1.0", about = "A tool with subcommands")] pub struct Cli { #[arg(short, long, value_name = "FILE", help = "Configuration file location")] config: Option, /// Subcommands for specific actions #[command(subcommand)] command: Option, } #[derive(Debug, Subcommand)] enum Commands { Serve { /// lists test values #[arg(short, long)] port: Option, }, Migrate, #[command(subcommand)] Service(ServiceCommands), #[command(subcommand)] Create(CreateCommands), } #[derive(Debug, Subcommand)] enum ServiceCommands { Install, Remove, Run, } #[derive(Debug, Subcommand)] enum CreateCommands { ApiKey { /// name of API Key #[arg(short, long)] name: String, /// name of API Key #[arg(short, long, default_value = "true")] requires_auth: bool, }, } impl Cli { pub fn config(&self) -> PathBuf { match self.config.as_ref() { Some(path) => path.clone(), None => ROOT_PATH.with_file_name("config.toml"), } } pub async fn handle( &self, config: &mut Configuration, ) -> Result> { match &self.command { Some(command) => match command { Commands::Serve { port } => { if let Some(port) = port { config.port = *port; } start_service(config, None).await } Commands::Migrate => { let pool = PgPool::connect(&config.database_url) .await .change_context(AppError)?; sqlx::migrate!("src/migrations") .run(&pool) .await .change_context(AppError)?; Ok(DaemonStatus::NotRunning) } Commands::Service(service_commands) => match service_commands { ServiceCommands::Install => { #[cfg(windows)] windows::install_service()?; // Print success message println!("Succssfully installed service {APP_NAME}"); Ok(DaemonStatus::NotRunning) } ServiceCommands::Remove => { #[cfg(windows)] windows::uninstall_service()?; // Print success message println!("Succssfully removed service {APP_NAME}"); Ok(DaemonStatus::NotRunning) } ServiceCommands::Run => { #[cfg(windows)] { windows::run()?; Ok(DaemonStatus::NotRunning) } #[cfg(not(windows))] { Ok(DaemonStatus::NotRunning) } } }, Commands::Create(create_commands) => match create_commands { CreateCommands::ApiKey { name, requires_auth, } => { // create API Key let (key_secret, key) = ApiKey::create(name, *requires_auth, HashSet::new(), config) .change_context(AppError)?; // Add API Key to Database let pool = PgPool::connect(&config.database_url) .await .change_context(AppError)?; create_api_key(&pool, &key).await.change_context(AppError)?; // print API key secret to console println!("Created API Key: {}.{key_secret}", key.id); Ok(DaemonStatus::NotRunning) } }, }, None => start_service(config, None).await, } } } pub enum DaemonStatus { NotRunning, Running((CancellationToken, Option>)), } async fn start_service( config: &Configuration, shutdown_signal: Option>, ) -> Result> { // create cancellation token let cancellation_token = CancellationToken::new(); api::start(config, cancellation_token.clone()).await?; Ok(DaemonStatus::Running((cancellation_token, shutdown_signal))) } #[cfg(windows)] mod windows { use error_stack::{Report, ResultExt}; use std::{ffi::OsString, thread, time::Duration}; use tokio_util::task::TaskTracker; use tracing::{error, info}; use windows_service::{ define_windows_service, service::{ ServiceAccess, ServiceControl, ServiceControlAccept, ServiceErrorControl, ServiceExitCode, ServiceInfo, ServiceStartType, ServiceState, ServiceStatus, ServiceType, }, service_control_handler::{self, ServiceControlHandlerResult}, service_dispatcher, service_manager::{ServiceManager, ServiceManagerAccess}, }; use crate::{ cli::{start_service, DaemonStatus}, APP_NAME, ROOT_PATH, }; use crate::{config::Configuration, errors::AppError}; const SERVICE_NAME: &str = "GenericApiService"; const SERVICE_DISPLAY_NAME: &str = APP_NAME; pub fn install_service() -> Result<(), Report> { let manager_access = ServiceManagerAccess::CONNECT | ServiceManagerAccess::CREATE_SERVICE; let service_manager = ServiceManager::local_computer(None::<&str>, manager_access) .change_context(AppError)?; // This example installs the service defined in `examples/ping_service.rs`. // In the real world code you would set the executable path to point to your own binary // that implements windows service. let service_binary_path = ::std::env::current_exe() .change_context(AppError)? .with_file_name(format!("{}.exe", env!("CARGO_PKG_NAME"))); let service_info = ServiceInfo { name: SERVICE_NAME.into(), display_name: SERVICE_DISPLAY_NAME.into(), service_type: ServiceType::OWN_PROCESS, start_type: ServiceStartType::AutoStart, error_control: ServiceErrorControl::Normal, executable_path: service_binary_path, launch_arguments: vec!["service".into(), "run".into()], dependencies: vec![], account_name: None, // run as System account_password: None, }; let service = service_manager .create_service(&service_info, ServiceAccess::CHANGE_CONFIG) .change_context(AppError)?; service .set_description(format!("{APP_NAME}")) .change_context(AppError)?; Ok(()) } pub fn uninstall_service() -> Result<(), Report> { let manager_access = ServiceManagerAccess::CONNECT; let service_manager = ServiceManager::local_computer(None::<&str>, manager_access) .change_context(AppError)?; let service_access = ServiceAccess::QUERY_STATUS | ServiceAccess::STOP | ServiceAccess::DELETE; let service = service_manager .open_service(SERVICE_NAME, service_access) .change_context(AppError)?; let service_status = service.query_status().change_context(AppError)?; if service_status.current_state != ServiceState::Stopped { service.stop().change_context(AppError)?; // Wait for service to stop thread::sleep(Duration::from_secs(1)); } service.delete().change_context(AppError)?; Ok(()) } pub fn run() -> Result<(), Report> { // Generate the windows service boilerplate. // The boilerplate contains the low-level service entry function (ffi_service_main) that parses // incoming service arguments into Vec and passes them to user defined service // entry (my_service_main). define_windows_service!(ffi_service_main, service_main); fn run_service() -> Result<(), Report> { // Create a channel to be able to poll a stop event from the service worker loop. let (shutdown_tx, shutdown_rx) = tokio::sync::mpsc::unbounded_channel(); // Define system service event handler that will be receiving service events. let event_handler = move |control_event| -> ServiceControlHandlerResult { match control_event { // Notifies a service to report its current status information to the service // control manager. Always return NoError even if not implemented. ServiceControl::Interrogate => ServiceControlHandlerResult::NoError, // Handle stop ServiceControl::Stop => { let _ = shutdown_tx.send(()); ServiceControlHandlerResult::NoError } _ => ServiceControlHandlerResult::NotImplemented, } }; // Register system service event handler. // The returned status handle should be used to report service status changes to the system. let status_handle = service_control_handler::register(SERVICE_NAME, event_handler) .change_context(AppError)?; // Tell the system that service is running status_handle .set_service_status(ServiceStatus { service_type: ServiceType::OWN_PROCESS, current_state: ServiceState::Running, controls_accepted: ServiceControlAccept::STOP, exit_code: ServiceExitCode::Win32(0), checkpoint: 0, wait_hint: Duration::default(), process_id: None, }) .change_context(AppError)?; // load config file let config = Configuration::new(ROOT_PATH.with_file_name("config.toml"))?; config.check()?; // run service - blocking thread let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .unwrap(); rt.block_on(async move { // create task tracker let tracker = TaskTracker::new(); match start_service(&config, Some(shutdown_rx)).await { Ok(DaemonStatus::Running((cancellation_token, Some(mut shutdown_rx)))) => { // wait for shutdown signal shutdown_rx.recv().await; // send shutdown signal to application and wait cancellation_token.cancel(); // close task tracker tracker.close(); // Wait for everything to finish. tracker.wait().await; tokio::time::sleep(Duration::from_millis(200)).await; info!("{APP_NAME} gracefully shut down."); } Err(e) => error!("{e:?}"), _ => unreachable!(), } }); // Tell the system that service has stopped. status_handle .set_service_status(ServiceStatus { service_type: ServiceType::OWN_PROCESS, current_state: ServiceState::Stopped, controls_accepted: ServiceControlAccept::empty(), exit_code: ServiceExitCode::Win32(0), checkpoint: 0, wait_hint: Duration::default(), process_id: None, }) .change_context(AppError)?; Ok(()) } // Service entry function which is called on background thread by the system with service // parameters. There is no stdout or stderr at this point so make sure to configure the log // output to file if needed. pub fn service_main(_arguments: Vec) { if let Err(e) = run_service() { // Handle the error, by logging or something. error!("Error executing windows service: {e}") } } // Register generated `ffi_service_main` with the system and start the service, blocking // this thread until the service is stopped. service_dispatcher::start(SERVICE_NAME, ffi_service_main).change_context(AppError) } }