1 files changed, 186 insertions, 0 deletions
diff --git a/src/util/windows.rs b/src/util/windows.rs
new file mode 100644
index 0000000..3b1339e
--- /dev/null
+++ b/src/util/windows.rs
@@ -0,0 +1,186 @@
+use windows::{
+    core::{
+        HSTRING,
+        PWSTR,
+    },
+    Win32::{
+        Foundation,
+        System::Services,
+    },
+};
+
+lazy_static::lazy_static! {
+    static ref PSQL_REGEX: regex::Regex = regex::Regex::new(r#"^(?:postgresql|psql|postgres)-.*([\d\.]*)"#).unwrap();
+}
+
+pub unsafe fn ensure_postgres_started() -> windows::core::Result<bool> {
+    let sc_manager = Services::OpenSCManagerW(
+        None,
+        None,
+        (Foundation::GENERIC_READ | Foundation::GENERIC_WRITE).0,
+    )?;
+
+    let Some(svc) = get_psql_service(sc_manager)? else {
+        return Ok(false);
+    };
+
+    match svc.status.dwCurrentState {
+        Services::SERVICE_RUNNING => {
+            log::info!("postgres was already running, done");
+            return Ok(false);
+        },
+
+        Services::SERVICE_START_PENDING => {
+            log::info!("postgres was already in startup, done");
+            return Ok(false);
+        },
+
+        _ => {
+            log::info!("postgres is not running: starting it")
+        },
+    }
+
+    let service =
+        Services::OpenServiceW(sc_manager, &HSTRING::from(&svc.name), Services::SERVICE_START)?;
+
+    Services::StartServiceW(service, None)?;
+
+    log::info!("started postgres service");
+
+    Services::CloseServiceHandle(service)?;
+    Services::CloseServiceHandle(sc_manager)?;
+
+    Ok(true)
+}
+
+pub unsafe fn shutdown_postgres() -> windows::core::Result<()> {
+    let sc_manager = Services::OpenSCManagerW(
+        None,
+        None,
+        (Foundation::GENERIC_READ | Foundation::GENERIC_WRITE).0,
+    )?;
+
+    let Some(svc) = get_psql_service(sc_manager)? else {
+        log::warn!("wasn't able to find postgres service");
+        return Ok(());
+    };
+
+    if svc.status.dwCurrentState != Services::SERVICE_RUNNING {
+        log::warn!("postgres wasn't in running state, not attempting to shut it down");
+        return Ok(());
+    }
+
+    let service =
+        Services::OpenServiceW(sc_manager, &HSTRING::from(&svc.name), Services::SERVICE_STOP)?;
+
+    let mut params = Services::SERVICE_CONTROL_STATUS_REASON_PARAMSW {
+        dwReason:      Services::SERVICE_STOP_REASON_FLAG_PLANNED
+            | Services::SERVICE_STOP_REASON_MAJOR_NONE
+            | Services::SERVICE_STOP_REASON_MINOR_NONE,
+        pszComment:    PWSTR::null(),
+        ServiceStatus: Services::SERVICE_STATUS_PROCESS::default(),
+    };
+
+    Services::ControlServiceExW(
+        service,
+        Services::SERVICE_CONTROL_STOP,
+        Services::SERVICE_CONTROL_STATUS_REASON_INFO,
+        &mut params as *mut _ as _,
+    )?;
+
+    log::info!("stopped postgres service");
+
+    Services::CloseServiceHandle(service)?;
+    Services::CloseServiceHandle(sc_manager)?;
+
+    Ok(())
+}
+
+pub unsafe fn get_psql_service(
+    sc_manager: Services::SC_HANDLE,
+) -> windows::core::Result<Option<ServiceStatus>> {
+    let services = list_services(sc_manager)?;
+
+    let mut psql_services =
+        services.into_iter().filter(|status| PSQL_REGEX.is_match(&status.name)).collect::<Vec<_>>();
+
+    match psql_services.len() {
+        1 => {},
+
+        0 => {
+            log::warn!("unable to identify an installed postgres service");
+            return Ok(None);
+        },
+
+        other => {
+            log::warn!("unable to identify postgres service: {other} potential matches found");
+            return Ok(None);
+        },
+    }
+
+    let svc = psql_services.remove(0);
+    log::debug!("identified postgres service: {}", svc.name);
+
+    Ok(Some(svc))
+}
+
+const RESULT_SIZE: usize = size_of::<Services::ENUM_SERVICE_STATUS_PROCESSW>();
+
+pub struct ServiceStatus {
+    pub display_name: String,
+    pub name:         String,
+    pub status:       Services::SERVICE_STATUS_PROCESS,
+}
+
+pub unsafe fn list_services(
+    sc_manager: Services::SC_HANDLE,
+) -> windows::core::Result<Vec<ServiceStatus>> {
+    let mut bytes_needed: u32 = 0;
+    let mut services_returned: u32 = 0;
+    let mut resume = 0;
+
+    let mut svcs = vec![];
+
+    loop {
+        let mut data = [0u8; 256 * 1024];
+
+        let result = Services::EnumServicesStatusExW(
+            sc_manager,
+            Services::SC_ENUM_PROCESS_INFO,
+            Services::SERVICE_WIN32,
+            Services::SERVICE_STATE_ALL,
+            Some(&mut data),
+            &mut bytes_needed as *mut _,
+            &mut services_returned as *mut _,
+            Some(&mut resume),
+            None,
+        );
+
+        if let Err(ref e) = result
+            && e.code() != Foundation::ERROR_MORE_DATA.to_hresult()
+        {
+            return result.map(|_| unreachable!());
+        }
+
+        let data = bytemuck::cast_slice::<u8, [u8; RESULT_SIZE]>(
+            &data[..services_returned as usize * RESULT_SIZE],
+        );
+
+        let elems = data
+            .into_iter()
+            .map(|x| core::mem::transmute::<_, Services::ENUM_SERVICE_STATUS_PROCESSW>(*x))
+            .map(|x| ServiceStatus {
+                display_name: x.lpDisplayName.to_string().unwrap(),
+                name:         x.lpServiceName.to_string().unwrap(),
+                status:       x.ServiceStatusProcess,
+            });
+
+        svcs.extend(elems);
+
+        if result.is_ok() {
+            break;
+        }
+    }
+
+    Ok(svcs)
+}