malwaredb 0.3.2

// SPDX-License-Identifier: Apache-2.0

//! Windows service installation and implementation of [ServiceMain](https://learn.microsoft.com/en-us/windows/win32/services/service-servicemain-function)
//! Based on <https://github.com/devinstewart/rust_windows_service_axum/blob/34bec1df321d8653b60693d5a068aeef5f0dd85b/src/main.rs>

use crate::cli::config::Config;
use crate::cli::service::Installer;

use std::path::PathBuf;
use std::process::ExitCode;
use std::{
    ffi::OsString,
    io::{Error, ErrorKind::Other},
};

use anyhow::Result;
use clap::Parser;
use tokio::runtime::Runtime;
use tokio::sync::mpsc;
use windows_service::{
    define_windows_service,
    service::{
        ServiceAccess, ServiceControl, ServiceControlAccept, ServiceErrorControl, ServiceExitCode,
        ServiceInfo, ServiceStartType, ServiceState, ServiceStatus, ServiceType,
    },
    service_control_handler::{self, ServiceControlHandlerResult, ServiceStatusHandle},
    service_manager::{ServiceManager, ServiceManagerAccess},
    Error::Winapi,
};

define_windows_service!(ffi_service_main, service_main);

const SERVICE_NAME: &str = "MalwareDB";
const SERVICE_TYPE: ServiceType = ServiceType::OWN_PROCESS;

/// Install Malware DB as a Windows service
#[derive(Parser, Debug, Clone, PartialEq)]
pub struct Install {
    /// Use a specific configuration file path
    #[arg(value_hint = clap::ValueHint::FilePath)]
    pub config: Option<PathBuf>,

    /// Remove the service
    #[arg(default_value_t = false)]
    pub uninstall: bool,
}

impl Installer for Install {
    fn do_install(&self) -> Result<ExitCode> {
        if self.uninstall {
            let manager =
                ServiceManager::local_computer(None::<&str>, ServiceManagerAccess::CONNECT)?;
            let service_access =
                ServiceAccess::QUERY_STATUS | ServiceAccess::STOP | ServiceAccess::DELETE;
            let service = manager.open_service(SERVICE_NAME, service_access)?;
            service.delete()?;
            if service.query_status()?.current_state != ServiceState::Stopped {
                // If the service cannot be stopped, it will be deleted when the system restarts.
                service.stop()?;
            }
            println!("Service uninstalled successfully");
            return Ok(ExitCode::SUCCESS);
        }

        let manager =
            ServiceManager::local_computer(None::<&str>, ServiceManagerAccess::CREATE_SERVICE)?;

        let config_file_arg = if let Some(path) = &self.config {
            if Config::from_file(path).is_err() {
                eprintln!(
                    "The provided configuration file {} is invalid, not creating the Windows Service", path.display()
                );
                return Ok(ExitCode::FAILURE);
            }
            vec![OsString::from(path)]
        } else {
            if Config::from_found_files().is_err() {
                eprintln!("A default configuration file could not be found or is invalid, not creating the Windows Service");
                return Ok(ExitCode::FAILURE);
            }
            vec![]
        };

        let service_info = ServiceInfo {
            name: SERVICE_NAME.into(),
            display_name: "MalwareDB Service".into(),
            service_type: ServiceType::OWN_PROCESS,
            start_type: ServiceStartType::AutoStart,
            error_control: ServiceErrorControl::Normal,
            executable_path: std::env::current_exe()?,
            launch_arguments: config_file_arg,
            dependencies: vec![],
            account_name: None, // Local system account
            account_password: None,
        };

        manager.create_service(&service_info, ServiceAccess::CHANGE_CONFIG)?;

        Ok(ExitCode::SUCCESS)
    }
}

fn notify_stop_and_return_error(
    status_handle: ServiceStatusHandle,
    error: String,
) -> windows_service::Result<()> {
    match status_handle.set_service_status(ServiceStatus {
        service_type: SERVICE_TYPE,
        current_state: ServiceState::Stopped,
        controls_accepted: ServiceControlAccept::empty(),
        exit_code: ServiceExitCode::Win32(1),
        checkpoint: 0,
        wait_hint: std::time::Duration::default(),
        process_id: None,
    }) {
        Ok(()) => (),
        Err(e) => {
            return Err(Winapi(Error::new(
                Other,
                format!("On Error: {error}\nFailed to set service status: {e:?}"),
            )))
        }
    }

    Err(Winapi(Error::new(Other, error)))
}

/// Service entry point
#[allow(dead_code)]
#[allow(clippy::needless_pass_by_value)]
fn service_main(arguments: Vec<std::ffi::OsString>) {
    crate::init_logger(None);

    let cfg = if arguments.is_empty() {
        Config::from_found_files().expect("No config files found")
    } else {
        let path = PathBuf::from(arguments.first().unwrap());
        Config::from_file(&path).expect("Unable to parse provided config file")
    };

    if let Err(_e) = run_service(cfg) {
        // TODO
    }
}

fn run_service(cfg: Config) -> windows_service::Result<()> {
    let (tx, rx) = mpsc::channel(1);

    let status_handle =
        service_control_handler::register(
            SERVICE_NAME,
            move |control_event| match control_event {
                ServiceControl::Interrogate => ServiceControlHandlerResult::NoError,
                ServiceControl::Stop => {
                    match tx.blocking_send(()) {
                        Ok(()) => (),
                        Err(_) => return ServiceControlHandlerResult::Other(1),
                    }
                    ServiceControlHandlerResult::NoError
                }
                ServiceControl::UserEvent(code) => {
                    if code.to_raw() == 130 {
                        match tx.blocking_send(()) {
                            Ok(()) => (),
                            Err(_) => return ServiceControlHandlerResult::Other(1),
                        }
                    }
                    ServiceControlHandlerResult::NoError
                }
                _ => ServiceControlHandlerResult::NotImplemented,
            },
        )?;

    match status_handle.set_service_status(ServiceStatus {
        service_type: SERVICE_TYPE,
        current_state: ServiceState::Running,
        controls_accepted: ServiceControlAccept::STOP,
        exit_code: ServiceExitCode::Win32(0),
        checkpoint: 0,
        wait_hint: std::time::Duration::default(),
        process_id: None,
    }) {
        Ok(()) => (),
        Err(e) => {
            return notify_stop_and_return_error(
                status_handle,
                format!("Failed to set service status: {e:?}"),
            )
        }
    }

    let rt = match Runtime::new() {
        Ok(rt) => rt,
        Err(e) => {
            return notify_stop_and_return_error(
                status_handle,
                format!("Failed to create runtime: {e:?}"),
            )
        }
    };

    rt.block_on(async {
        let state = cfg.into_state().await.expect("Failed to initialize server");
        state.serve(Some(rx)).await.expect("Failed to start server");
    });

    Ok(())
}