Windows-Gotify-Client-Rust/src/windows/startup.rs

//! HKCU Run startup integration (AGENTS.md §27).
//!
//! Writes (or removes) a value under
//! `HKCU\Software\Microsoft\Windows\CurrentVersion\Run` whose data is
//! the path of the current executable. That makes the user logged-in
//! start the client automatically; the user can toggle it from
//! Settings → Startup → "Start with Windows" (Phase 6 C2 wires that
//! toggle to `set_enabled`).

use windows_sys::Win32::System::Registry::{
    RegCloseKey, RegDeleteValueW, RegOpenKeyExW, RegQueryValueExW, RegSetValueExW,
    HKEY, HKEY_CURRENT_USER, KEY_READ, KEY_WRITE, REG_SZ,
};
use windows_sys::Win32::Foundation::ERROR_SUCCESS;

/// `windows_sys::core::Error` doesn't exist in 0.59 — wrap a `u32`
/// win32 error in a small private tuple struct so callers can still
/// `?` on it without pulling in a heavier error-handling crate.
#[derive(Debug)]
pub struct Win32Error(pub u32);

impl std::fmt::Display for Win32Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "win32 error 0x{:08x}", self.0)
    }
}
impl std::error::Error for Win32Error {}

/// Translates a `RegOpenKeyExW`/`RegQueryValueExW` outcome into
/// either `Some(path)` / `None` (absent value) or `Err(Win32Error)`.
pub type RegResult<T> = Result<T, Win32Error>;

const RUN_KEY_PATH: &str = "Software\\Microsoft\\Windows\\CurrentVersion\\Run\0";
const RUN_VALUE_NAME: &str = "WindowsGotifyClientRust\0";

/// Returns the path of the current executable as a UTF-16 string
/// suitable for `RegSetValueExW`. Falls back to `<unknown>` if the OS
/// refuses to give us the path; the startup write will then write
/// `<unknown>` which Windows will refuse to run on next login — that
/// fails closed (won't auto-start) rather than fails open.
fn current_exe_path() -> String {
    std::env::current_exe()
        .ok()
        .and_then(|p| p.to_str().map(|s| s.to_string()))
        .unwrap_or_else(|| "<unknown>".into())
}

/// Is "Start with Windows" currently enabled? (Reads the registry,
/// not the in-memory config, so we never desync.)
pub fn is_enabled() -> bool {
    match query_value() {
        Ok(Some(_)) => true,
        _ => false,
    }
}

/// Add (or update) the auto-start entry.
pub fn set_enabled(enabled: bool) -> bool {
    if !enabled {
        return remove_value();
    }
    let path = current_exe_path();
    set_value(&path)
}

/// Look up the existing Run value (returns `Ok(None)` if absent).
fn query_value() -> RegResult<Option<String>> {
    let path_w: Vec<u16> = RUN_KEY_PATH.encode_utf16().collect();
    let name_w: Vec<u16> = RUN_VALUE_NAME.encode_utf16().collect();

    let mut key: HKEY = std::ptr::null_mut();
    let open = unsafe {
        RegOpenKeyExW(
            HKEY_CURRENT_USER,
            path_w.as_ptr(),
            0,
            KEY_READ,
            &mut key,
        )
    };
    if open != ERROR_SUCCESS {
        return Err(Win32Error(open));
    }

    // First call to learn the required buffer size.
    let mut size: u32 = 0;
    let mut kind: u32 = 0;
    let probe = unsafe {
        RegQueryValueExW(
            key,
            name_w.as_ptr(),
            std::ptr::null(),
            &mut kind,
            std::ptr::null_mut(),
            &mut size,
        )
    };

    let result = if probe == ERROR_SUCCESS && (kind == REG_SZ || kind == REG_EXPAND_SZ) && size > 0 {
        // Allocate the buffer and re-query. The buffer includes the
        // null terminator (size is in BYTES, not characters).
        let mut buf = vec![0u8; size as usize];
        let query = unsafe {
            RegQueryValueExW(
                key,
                name_w.as_ptr(),
                std::ptr::null(),
                &mut kind,
                buf.as_mut_ptr(),
                &mut size,
            )
        };
        if query == ERROR_SUCCESS {
            // Truncate at first null terminator; convert wide to UTF-8.
            let wide: &[u16] = unsafe {
                std::slice::from_raw_parts(buf.as_ptr() as *const u16, size as usize / 2)
            };
            let s = wide
                .iter()
                .take_while(|&&c| c != 0)
                .map(|&c| char::from_u32(c as u32).unwrap_or('?'))
                .collect::<String>();
            Ok(Some(s))
        } else {
            Err(Win32Error(query))
        }
    } else if probe == ERROR_SUCCESS {
        Ok(Some(String::new()))
    } else {
        // ERROR_FILE_NOT_FOUND (2) → not enabled.
        Ok(None)
    };

    unsafe { RegCloseKey(key) };
    result
}

const REG_EXPAND_SZ: u32 = 2;

fn set_value(path: &str) -> bool {
    let path_w: Vec<u16> = RUN_KEY_PATH.encode_utf16().collect();
    let name_w: Vec<u16> = RUN_VALUE_NAME.encode_utf16().collect();

    let mut key: HKEY = std::ptr::null_mut();
    let open = unsafe {
        RegOpenKeyExW(
            HKEY_CURRENT_USER,
            path_w.as_ptr(),
            0,
            KEY_WRITE,
            &mut key,
        )
    };
    if open != ERROR_SUCCESS {
        tracing::warn!(error = open, "startup: RegOpenKeyExW failed");
        return false;
    }

    // The registry stores REG_SZ in UTF-16 with a trailing null. Both
    // the bytes (excluding null) and the wide chars (including null)
    // are needed: `cbData` is in BYTES.
    let mut data_w: Vec<u16> = path.encode_utf16().collect();
    data_w.push(0);
    let byte_len = (data_w.len() * std::mem::size_of::<u16>()) as u32;

    let set = unsafe {
        RegSetValueExW(
            key,
            name_w.as_ptr(),
            0,
            REG_SZ,
            data_w.as_ptr() as *const u8,
            byte_len,
        )
    };
    unsafe { RegCloseKey(key) };

    if set != ERROR_SUCCESS {
        tracing::warn!(error = set, "startup: RegSetValueExW failed");
        return false;
    }
    tracing::info!(path = %path, "startup: HKCU Run entry written");
    true
}

fn remove_value() -> bool {
    let path_w: Vec<u16> = RUN_KEY_PATH.encode_utf16().collect();
    let name_w: Vec<u16> = RUN_VALUE_NAME.encode_utf16().collect();

    let mut key: HKEY = std::ptr::null_mut();
    let open = unsafe {
        RegOpenKeyExW(
            HKEY_CURRENT_USER,
            path_w.as_ptr(),
            0,
            KEY_WRITE,
            &mut key,
        )
    };
    if open != ERROR_SUCCESS {
        // Key missing → already removed; treat as success.
        return open == 2;
    }
    let del = unsafe { RegDeleteValueW(key, name_w.as_ptr()) };
    unsafe { RegCloseKey(key) };
    del == ERROR_SUCCESS
}