winkit 0.1.0

Thin checked wrappers over the Win32 an installer needs: elevation, ACLs, services, the Restart Manager, the registry and shortcuts.
use std::{
    mem, ptr,
    time::{Duration, Instant},
};

use windows_sys::Win32::System::Services::{
    CloseServiceHandle, ControlService, OpenSCManagerW, OpenServiceW, QueryServiceStatus,
    SC_MANAGER_CONNECT, SERVICE_CONTROL_STOP, SERVICE_QUERY_STATUS, SERVICE_RUNNING, SERVICE_START,
    SERVICE_STATUS, SERVICE_STOP, SERVICE_STOP_PENDING, SERVICE_STOPPED, StartServiceW,
};

use crate::{
    error::{Error, Result},
    wide::wide,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ServiceState {
    Running,
    Stopped,
    Stopping,
    Other(u32),
}

/// Reads a service's current state, or reports that it is not installed.
pub fn service_state(name: &str) -> Result<ServiceState> {
    let service = Service::open(name, SERVICE_QUERY_STATUS)?;
    Ok(state_of(service.status()?))
}

/// Stops a service and waits for it to actually be stopped.
///
/// This is the part `sc stop` cannot do: it returns as soon as the control is
/// accepted, so anything built on it has to guess at a sleep afterwards.
/// Waiting on the real status is both quicker when the service stops fast and
/// correct when it does not.
///
/// A service that is not installed is not an error; there is nothing to stop.
pub fn stop_service(name: &str, timeout: Duration) -> Result<()> {
    let service = match Service::open(name, SERVICE_STOP | SERVICE_QUERY_STATUS) {
        Ok(service) => service,
        Err(error) if error.is_not_found() => return Ok(()),
        Err(error) => return Err(error),
    };

    if state_of(service.status()?) == ServiceState::Stopped {
        return Ok(());
    }

    let mut status: SERVICE_STATUS = unsafe { mem::zeroed() };
    if unsafe { ControlService(service.handle, SERVICE_CONTROL_STOP, &mut status) } == 0 {
        // Losing the race with something else that stopped it is still success.
        let error = Error::last("ControlService");
        if state_of(service.status()?) != ServiceState::Stopped {
            return Err(error);
        }
        return Ok(());
    }

    let deadline = Instant::now() + timeout;
    loop {
        match state_of(service.status()?) {
            ServiceState::Stopped => return Ok(()),
            _ if Instant::now() >= deadline => {
                return Err(Error::saying(
                    "stop_service",
                    format!("{name} did not stop within {} seconds", timeout.as_secs()),
                ));
            }
            // Polling rather than waiting on an event because the service
            // control manager offers no waitable handle for this.
            _ => std::thread::sleep(Duration::from_millis(100)),
        }
    }
}

/// Starts a service, treating "already running" as success.
pub fn start_service(name: &str) -> Result<()> {
    let service = Service::open(name, SERVICE_START | SERVICE_QUERY_STATUS)?;
    if state_of(service.status()?) == ServiceState::Running {
        return Ok(());
    }
    if unsafe { StartServiceW(service.handle, 0, ptr::null()) } == 0 {
        return Err(Error::last("StartServiceW"));
    }
    Ok(())
}

fn state_of(status: SERVICE_STATUS) -> ServiceState {
    match status.dwCurrentState {
        SERVICE_RUNNING => ServiceState::Running,
        SERVICE_STOPPED => ServiceState::Stopped,
        SERVICE_STOP_PENDING => ServiceState::Stopping,
        other => ServiceState::Other(other),
    }
}

/// An open service handle that closes itself. The manager it came from is
/// held alongside it and dropped after it, because closing the manager first
/// would invalidate the service handle.
struct Service {
    handle: windows_sys::Win32::System::Services::SC_HANDLE,
    _manager: Manager,
}

struct Manager(windows_sys::Win32::System::Services::SC_HANDLE);

impl Service {
    fn open(name: &str, access: u32) -> Result<Self> {
        let manager = unsafe { OpenSCManagerW(ptr::null(), ptr::null(), SC_MANAGER_CONNECT) };
        if manager.is_null() {
            return Err(Error::last("OpenSCManagerW"));
        }
        let manager = Manager(manager);

        let name = wide(name);
        let service = unsafe { OpenServiceW(manager.0, name.as_ptr(), access) };
        if service.is_null() {
            return Err(Error::last("OpenServiceW"));
        }
        Ok(Self {
            handle: service,
            _manager: manager,
        })
    }

    fn status(&self) -> Result<SERVICE_STATUS> {
        let mut status: SERVICE_STATUS = unsafe { mem::zeroed() };
        if unsafe { QueryServiceStatus(self.handle, &mut status) } == 0 {
            return Err(Error::last("QueryServiceStatus"));
        }
        Ok(status)
    }
}

impl Drop for Service {
    fn drop(&mut self) {
        unsafe { CloseServiceHandle(self.handle) };
    }
}

impl Drop for Manager {
    fn drop(&mut self) {
        unsafe { CloseServiceHandle(self.0) };
    }
}