311 lines
11 KiB
Rust
311 lines
11 KiB
Rust
|
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,
|
||
|
|
Uninstall,
|
||
|
|
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 => {
|
||
|
|
// TODO do things
|
||
|
|
Ok(DaemonStatus::NotRunning)
|
||
|
|
}
|
||
|
|
ServiceCommands::Uninstall => {
|
||
|
|
// TODO do things
|
||
|
|
Ok(DaemonStatus::NotRunning)
|
||
|
|
}
|
||
|
|
ServiceCommands::Run => {
|
||
|
|
// TODO do things
|
||
|
|
// create shutdown signals
|
||
|
|
// let (shutdown_send, mut shutdown_recv) = mpsc::unbounded_channel();
|
||
|
|
// Register generated `ffi_service_main` with the system and start the service, blocking
|
||
|
|
// this thread until the service is stopped.
|
||
|
|
#[cfg(windows)]
|
||
|
|
let _ = service_dispatcher::start(env!("CARGO_PKG_NAME"), ffi_service_main);
|
||
|
|
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 windows_service::{
|
||
|
|
service::{
|
||
|
|
ServiceAccess, ServiceErrorControl, ServiceInfo, ServiceStartType, ServiceState,
|
||
|
|
ServiceType,
|
||
|
|
},
|
||
|
|
service_manager::{ServiceManager, ServiceManagerAccess},
|
||
|
|
};
|
||
|
|
|
||
|
|
use crate::errors::AppError;
|
||
|
|
use crate::APP_NAME;
|
||
|
|
|
||
|
|
pub fn install() -> 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: OsString::from(env!("CARGO_PKG_NAME")),
|
||
|
|
display_name: OsString::from(&format!("{APP_NAME}")),
|
||
|
|
service_type: ServiceType::OWN_PROCESS,
|
||
|
|
start_type: ServiceStartType::AutoStart,
|
||
|
|
error_control: ServiceErrorControl::Normal,
|
||
|
|
executable_path: service_binary_path,
|
||
|
|
launch_arguments: vec!["service run".into()],
|
||
|
|
dependencies: vec![],
|
||
|
|
account_name: None, // run as System
|
||
|
|
account_password: None,
|
||
|
|
};
|
||
|
|
let service = service_manager
|
||
|
|
.create_service(&service_info, ServiceAccess::CHANGE_CONFIG)
|
||
|
|
.into_report()
|
||
|
|
.change_context(AppError)?;
|
||
|
|
|
||
|
|
service
|
||
|
|
.set_description(format!("{APP_NAME}"))
|
||
|
|
.into_report()
|
||
|
|
.change_context(AppError)?;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn remove() -> 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(env!("CARGO_PKG_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().into_report().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.
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
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::channel(1);
|
||
|
|
|
||
|
|
// 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 => {
|
||
|
|
shutdown_tx.try_send(()).unwrap();
|
||
|
|
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(env!("CARGO_PKG_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)?;
|
||
|
|
|
||
|
|
// run service - blocking thread
|
||
|
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||
|
|
.enable_all()
|
||
|
|
.build()
|
||
|
|
.unwrap();
|
||
|
|
|
||
|
|
rt.block_on(async move {
|
||
|
|
if let Err(e) = execute(shutdown_rx).await {
|
||
|
|
error!("{e:?}");
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
// 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,
|
||
|
|
})
|
||
|
|
.into_report()
|
||
|
|
.change_context(ServiceError::Starting)?;
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|