rust-api-template/src/cli.rs

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<PathBuf>,

    /// Subcommands for specific actions
    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(Debug, Subcommand)]
enum Commands {
    Serve {
        /// lists test values
        #[arg(short, long)]
        port: Option<u16>,
    },
    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<DaemonStatus, Report<AppError>> {
        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<UnboundedReceiver<()>>)),
}

async fn start_service(
    config: &Configuration,
    shutdown_signal: Option<UnboundedReceiver<()>>,
) -> Result<DaemonStatus, Report<AppError>> {
    // 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<AppError>> {
        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<AppError>> {
        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<AppError>> {
        // Generate the windows service boilerplate.
        // The boilerplate contains the low-level service entry function (ffi_service_main) that parses
        // incoming service arguments into Vec<OsString> and passes them to user defined service
        // entry (my_service_main).
        define_windows_service!(ffi_service_main, service_main);

        fn run_service() -> Result<(), Report<AppError>> {
            // 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<OsString>) {
            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)
    }
}