From 7d47be4c40ac8f9a9c33ae8c78e44d20c2c46ef6 Mon Sep 17 00:00:00 2001 From: Nathan Perry Date: Fri, 16 Aug 2024 03:09:00 -0400 Subject: windows: windows_util -> util::windows --- src/bot.rs | 14 ++-- src/lib.rs | 3 - src/util/mod.rs | 9 ++- src/util/windows.rs | 186 ++++++++++++++++++++++++++++++++++++++++++++++++++++ src/windows_util.rs | 186 ---------------------------------------------------- 5 files changed, 199 insertions(+), 199 deletions(-) create mode 100644 src/util/windows.rs delete mode 100644 src/windows_util.rs diff --git a/src/bot.rs b/src/bot.rs index de31c73..0c44349 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -258,11 +258,11 @@ fn on_err(err: FrameworkError) -> BoxFuture<()> { } else { return; } - }, + } Err(e) => { error!("processing unrecognized message: {e}"); "BANIC".to_string() - }, + } } } else { if let Err(e) = commands::meme::invoke::_meme( @@ -270,7 +270,7 @@ fn on_err(err: FrameworkError) -> BoxFuture<()> { msg_content, Default::default(), ) - .await + .await { error!("producing meme for unrecognized: {e}"); "BANIC".to_string() @@ -278,7 +278,7 @@ fn on_err(err: FrameworkError) -> BoxFuture<()> { return; } } - }, + } _ => "BANIC".to_string(), }; @@ -377,7 +377,7 @@ fn check(ctx: PoiseContext) -> BoxFuture> { }) } -fn before_handle<'fut>(ctx: PoiseContext<'fut>) -> Pin + Send + 'fut>> { +fn before_handle<'fut>(ctx: PoiseContext<'fut>) -> Pin + Send + 'fut>> { debug!( "got command '{}' from user '{}' ({})", ctx.command().name, @@ -397,7 +397,7 @@ fn after_handle(ctx: PoiseContext) -> BoxFuture<()> { pub async fn run() -> anyhow::Result<()> { #[cfg(all(windows, feature = "windows_autostart_postgres"))] let started_pg = - tokio::task::spawn_blocking(|| unsafe { crate::windows_util::ensure_postgres_started() }) + tokio::task::spawn_blocking(|| unsafe { util::windows::ensure_postgres_started() }) .await??; let token = &CONFIG.discord.auth.token; @@ -431,7 +431,7 @@ pub async fn run() -> anyhow::Result<()> { if started_pg { log::info!("we started postgres, stopping it before shutdown"); - tokio::task::spawn_blocking(|| crate::windows_util::shutdown_postgres()).await??; + tokio::task::spawn_blocking(|| util::windows::shutdown_postgres()).await??; } } diff --git a/src/lib.rs b/src/lib.rs index ac56087..1f03290 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,9 +21,6 @@ pub mod config; pub mod log_setup; pub mod util; -#[cfg(windows)] -pub mod windows_util; - pub use crate::{ config::*, util::*, diff --git a/src/util/mod.rs b/src/util/mod.rs index fb4b552..d214957 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -37,6 +37,9 @@ use crate::{ mod rest_vec; +#[cfg(windows)] +pub mod windows; + pub use rest_vec::*; pub async fn currently_playing(ctx: PoiseContext<'_>) -> bool { @@ -66,9 +69,9 @@ pub async fn users_listening(ctx: &Context) -> anyhow::Result { pub fn msg(ctx: poise::Context<'_, U, E>) -> Option<&Message> { match ctx { poise::Context::Prefix(poise::PrefixContext { - msg, - .. - }) => Some(msg), + msg, + .. + }) => Some(msg), _ => None, } } 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 { + 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> { + let services = list_services(sc_manager)?; + + let mut psql_services = + services.into_iter().filter(|status| PSQL_REGEX.is_match(&status.name)).collect::>(); + + 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::(); + +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> { + 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::( + &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) +} diff --git a/src/windows_util.rs b/src/windows_util.rs deleted file mode 100644 index 3b1339e..0000000 --- a/src/windows_util.rs +++ /dev/null @@ -1,186 +0,0 @@ -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 { - 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> { - let services = list_services(sc_manager)?; - - let mut psql_services = - services.into_iter().filter(|status| PSQL_REGEX.is_match(&status.name)).collect::>(); - - 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::(); - -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> { - 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::( - &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) -} -- cgit v1.3.1