skylink/src/main.rs

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<dyn std::error::Error>> {
    // 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<dyn std::error::Error>> {
    // 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<OsString>) {
    // 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);
}