crates/inject/src/main.rs at 8cd54197dd2feb41bf26a5990dd293b4b4b00988

kretrod/lodestone Public
lodestone Classic menu, and ship as one file 9b869ac · 3 days ago History
374 lines (349 loc) · 13.5 KB Raw
//! Loader: maps `lodestone-client.dll` into the running game.
//!
//! The classic route — write the DLL path into the target, then run
//! `LoadLibraryW` on it in a remote thread. kernel32 sits at the same base in
//! every process on a given boot, so the address we resolve locally is the
//! address the target will call.

#![cfg(windows)]

use std::ffi::{c_void, OsStr};
use std::os::windows::ffi::OsStrExt;

use windows_sys::Win32::Foundation::{CloseHandle, HANDLE, MAX_PATH};
use windows_sys::Win32::System::Diagnostics::Debug::WriteProcessMemory;
use windows_sys::Win32::System::Diagnostics::ToolHelp::{
    CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W, TH32CS_SNAPPROCESS,
};
use windows_sys::Win32::System::LibraryLoader::{GetModuleHandleA, GetProcAddress};
use windows_sys::Win32::System::Memory::{
    VirtualAllocEx, VirtualFreeEx, MEM_COMMIT, MEM_RELEASE, MEM_RESERVE, PAGE_READWRITE,
};
use windows_sys::Win32::System::ProcessStatus::{
    EnumProcessModulesEx, GetModuleBaseNameW, LIST_MODULES_ALL,
};
use windows_sys::Win32::System::Threading::{
    CreateRemoteThread, GetExitCodeThread, OpenProcess, WaitForSingleObject, INFINITE,
    PROCESS_ALL_ACCESS,
};

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.iter().any(|a| a == "--eject") {
        match eject(args.iter().any(|a| a == "--force")) {
            Ok(m) => println!("{m}"),
            Err(e) => {
                eprintln!("error: {e}");
                std::process::exit(1);
            }
        }
        return;
    }
    // An explicit path wins, so a freshly built DLL can be loaded without
    // rebuilding the loader around it.
    let dll = match args.iter().find(|a| !a.starts_with("--")) {
        Some(path) => Ok(path.clone()),
        None => extract_client().or_else(|e| {
            let fallback = default_dll_path();
            if std::path::Path::new(&fallback).exists() {
                Ok(fallback)
            } else {
                Err(e)
            }
        }),
    };
    let dll = match dll {
        Ok(d) => d,
        Err(e) => {
            eprintln!("error: {e}");
            std::process::exit(1);
        }
    };

    match run(&dll) {
        Ok(msg) => println!("{msg}"),
        Err(e) => {
            eprintln!("error: {e}");
            std::process::exit(1);
        }
    }
}

/// Ask the client to unload.
///
/// The graceful path drops a marker file the client polls for, so it can pull
/// its hooks out in the right order from its own thread. `--force` instead
/// calls FreeLibrary from a remote thread, which is only safe when the client
/// never got as far as hooking anything.
/// Beside the loader, so the pair works from wherever they are put.
fn marker() -> std::path::PathBuf {
    std::env::current_exe()
        .ok()
        .and_then(|p| p.parent().map(|d| d.join("unload")))
        .unwrap_or_else(|| std::path::PathBuf::from("unload"))
}

fn eject(force: bool) -> Result<String, String> {
    let pid = find_game()?;
    if !force {
        // The client stops itself and stays resident, so there is no module to
        // watch for: drop the marker, give it time to act, take it away again.
        std::fs::write(marker(), "1")
            .map_err(|e| format!("could not write the unload marker: {e}"))?;
        std::thread::sleep(std::time::Duration::from_millis(1200));
        let _ = std::fs::remove_file(marker());
        return Ok(format!("asked the client in pid {pid} to stop"));
    }

    // SAFETY: standard remote-call sequence, every handle checked.
    unsafe {
        let proc = OpenProcess(PROCESS_ALL_ACCESS, 0, pid);
        if proc.is_null() {
            return Err(format!("OpenProcess({pid}): {}", std::io::Error::last_os_error()));
        }
        let module = module_handle(proc, "lodestone-client.dll");
        let Some(module) = module else {
            CloseHandle(proc);
            return Ok(format!("the client is not loaded in pid {pid}"));
        };
        let k32 = GetModuleHandleA(c"kernel32.dll".as_ptr() as *const u8);
        let free = GetProcAddress(k32, c"FreeLibrary".as_ptr() as *const u8)
            .ok_or("kernel32!FreeLibrary not found")?;
        let thread = CreateRemoteThread(
            proc,
            std::ptr::null(),
            0,
            Some(std::mem::transmute::<
                unsafe extern "system" fn() -> isize,
                unsafe extern "system" fn(*mut c_void) -> u32,
            >(free)),
            module,
            0,
            std::ptr::null_mut(),
        );
        if thread.is_null() {
            CloseHandle(proc);
            return Err(format!("CreateRemoteThread: {}", std::io::Error::last_os_error()));
        }
        WaitForSingleObject(thread, INFINITE);
        CloseHandle(thread);
        CloseHandle(proc);
        Ok(format!("forced the client out of pid {pid}"))
    }
}

/// The base address of a loaded module in the target.
fn module_handle(proc: HANDLE, needle: &str) -> Option<*mut c_void> {
    let mut handles = vec![0usize; 2048];
    let mut needed: u32 = 0;
    // SAFETY: buffer sized in bytes for the call.
    let ok = unsafe {
        EnumProcessModulesEx(
            proc,
            handles.as_mut_ptr() as *mut _,
            (handles.len() * std::mem::size_of::<usize>()) as u32,
            &mut needed,
            LIST_MODULES_ALL,
        )
    };
    if ok == 0 {
        return None;
    }
    let n = (needed as usize / std::mem::size_of::<usize>()).min(handles.len());
    for &h in &handles[..n] {
        let mut buf = [0u16; MAX_PATH as usize];
        // SAFETY: h came from the enumeration.
        let len =
            unsafe { GetModuleBaseNameW(proc, h as *mut _, buf.as_mut_ptr(), buf.len() as u32) };
        if len > 0 && String::from_utf16_lossy(&buf[..len as usize]).eq_ignore_ascii_case(needle) {
            return Some(h as *mut c_void);
        }
    }
    None
}

/// The client, built into this binary by build.rs. Empty when the DLL was not
/// built at the time the loader was compiled.
const EMBEDDED_CLIENT: &[u8] = include_bytes!(env!("LODESTONE_CLIENT_DLL"));

/// Write the embedded client out so it can be loaded.
///
/// Under a fresh name each run: the client never unmaps itself, so reusing a
/// name would only ever load the copy already in the process.
fn extract_client() -> Result<String, String> {
    if EMBEDDED_CLIENT.is_empty() {
        return Err("this build has no client embedded; pass the DLL path".into());
    }
    let dir = std::env::current_exe()
        .ok()
        .and_then(|p| p.parent().map(|d| d.to_path_buf()))
        .unwrap_or_else(std::env::temp_dir);
    let stamp = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);
    let path = dir.join(format!("lodestone-client-{stamp}.dll"));
    if let Err(e) = std::fs::write(&path, EMBEDDED_CLIENT) {
        // A read-only folder is a normal thing to be dropped into.
        let fallback = std::env::temp_dir().join(format!("lodestone-client-{stamp}.dll"));
        std::fs::write(&fallback, EMBEDDED_CLIENT)
            .map_err(|e2| format!("could not write the client: {e} / {e2}"))?;
        return Ok(fallback.to_string_lossy().into_owned());
    }
    Ok(path.to_string_lossy().into_owned())
}

fn default_dll_path() -> String {
    // Next to the injector by default, so the pair can be copied anywhere.
    std::env::current_exe()
        .ok()
        .and_then(|p| p.parent().map(|d| d.join("lodestone-client.dll")))
        .map(|p| p.to_string_lossy().into_owned())
        .unwrap_or_else(|| "lodestone-client.dll".into())
}

fn run(dll: &str) -> Result<String, String> {
    let path = std::fs::canonicalize(dll)
        .map_err(|e| format!("{dll}: {e}"))?
        .to_string_lossy()
        // canonicalize hands back a \\?\ extended path; LoadLibraryW takes it,
        // but the plain form is what shows up in module lists.
        .trim_start_matches("\\\\?\\")
        .to_string();

    let pid = find_game()?;
    // SAFETY: pid came from the snapshot above; the handle is closed below.
    let proc = unsafe { OpenProcess(PROCESS_ALL_ACCESS, 0, pid) };
    if proc.is_null() {
        return Err(format!(
            "OpenProcess({pid}): {} (run as administrator?)",
            std::io::Error::last_os_error()
        ));
    }
    let result = inject(proc, pid, &path);
    // SAFETY: proc is a live handle from OpenProcess.
    unsafe { CloseHandle(proc) };
    result
}

fn inject(proc: HANDLE, pid: u32, path: &str) -> Result<String, String> {
    let file_name = path
        .rsplit('\\')
        .next()
        .unwrap_or("lodestone-client.dll")
        .to_string();
    if let Some(name) = loaded_module(proc, &file_name) {
        return Ok(format!("{name} is already loaded in pid {pid}"));
    }

    let wide: Vec<u16> = OsStr::new(path).encode_wide().chain(Some(0)).collect();
    let bytes = wide.len() * 2;

    // SAFETY: standard remote-allocate / write / call sequence; every handle
    // and pointer is checked before use and the allocation is freed at the end.
    unsafe {
        let remote = VirtualAllocEx(
            proc,
            std::ptr::null(),
            bytes,
            MEM_COMMIT | MEM_RESERVE,
            PAGE_READWRITE,
        );
        if remote.is_null() {
            return Err(format!("VirtualAllocEx: {}", std::io::Error::last_os_error()));
        }

        let mut written = 0usize;
        if WriteProcessMemory(proc, remote, wide.as_ptr() as *const c_void, bytes, &mut written) == 0
        {
            VirtualFreeEx(proc, remote, 0, MEM_RELEASE);
            return Err(format!("WriteProcessMemory: {}", std::io::Error::last_os_error()));
        }

        let k32 = GetModuleHandleA(c"kernel32.dll".as_ptr() as *const u8);
        let load_library = GetProcAddress(k32, c"LoadLibraryW".as_ptr() as *const u8)
            .ok_or("kernel32!LoadLibraryW not found")?;

        let thread = CreateRemoteThread(
            proc,
            std::ptr::null(),
            0,
            Some(std::mem::transmute::<
                unsafe extern "system" fn() -> isize,
                unsafe extern "system" fn(*mut c_void) -> u32,
            >(load_library)),
            remote,
            0,
            std::ptr::null_mut(),
        );
        if thread.is_null() {
            VirtualFreeEx(proc, remote, 0, MEM_RELEASE);
            return Err(format!("CreateRemoteThread: {}", std::io::Error::last_os_error()));
        }

        WaitForSingleObject(thread, INFINITE);
        CloseHandle(thread);
        VirtualFreeEx(proc, remote, 0, MEM_RELEASE);

        // The thread's exit code is only the low 32 bits of the returned
        // HMODULE, so a module that happens to land on a 4 GB boundary would
        // look like a failure. Ask the module list instead.
        match module_handle(proc, &file_name) {
            Some(base) => Ok(format!("injected into pid {pid} ({file_name} at {base:p})")),
            None => Err(format!(
                "{file_name} did not load into pid {pid} — wrong architecture, or a \
                 dependency is missing"
            )),
        }
    }
}

fn loaded_module(proc: HANDLE, needle: &str) -> Option<String> {
    let mut handles = vec![0usize; 2048];
    let mut needed: u32 = 0;
    // SAFETY: buffer is sized in bytes for the call.
    let ok = unsafe {
        EnumProcessModulesEx(
            proc,
            handles.as_mut_ptr() as *mut _,
            (handles.len() * std::mem::size_of::<usize>()) as u32,
            &mut needed,
            LIST_MODULES_ALL,
        )
    };
    if ok == 0 {
        return None;
    }
    let n = (needed as usize / std::mem::size_of::<usize>()).min(handles.len());
    for &h in &handles[..n] {
        let mut buf = [0u16; MAX_PATH as usize];
        // SAFETY: h is a module handle from the enumeration above.
        let len =
            unsafe { GetModuleBaseNameW(proc, h as *mut _, buf.as_mut_ptr(), buf.len() as u32) };
        if len > 0 {
            let name = String::from_utf16_lossy(&buf[..len as usize]);
            if name.eq_ignore_ascii_case(needle) {
                return Some(name);
            }
        }
    }
    None
}

/// The JVM with GLFW loaded is the game.
fn find_game() -> Result<u32, String> {
    let mut candidates = Vec::new();
    // SAFETY: snapshot handle is checked and closed.
    unsafe {
        let snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if snap.is_null() {
            return Err("CreateToolhelp32Snapshot failed".into());
        }
        let mut e: PROCESSENTRY32W = std::mem::zeroed();
        e.dwSize = std::mem::size_of::<PROCESSENTRY32W>() as u32;
        let mut ok = Process32FirstW(snap, &mut e);
        while ok != 0 {
            let end = e.szExeFile.iter().position(|&c| c == 0).unwrap_or(0);
            let name = String::from_utf16_lossy(&e.szExeFile[..end]).to_ascii_lowercase();
            if name == "javaw.exe" || name == "java.exe" {
                candidates.push(e.th32ProcessID);
            }
            ok = Process32NextW(snap, &mut e);
        }
        CloseHandle(snap);
    }

    let mut games = Vec::new();
    for pid in candidates {
        // SAFETY: pid from the snapshot; handle closed right after the check.
        let h = unsafe { OpenProcess(PROCESS_ALL_ACCESS, 0, pid) };
        if h.is_null() {
            continue;
        }
        if loaded_module(h, "glfw.dll").is_some() {
            games.push(pid);
        }
        // SAFETY: h is a live handle.
        unsafe { CloseHandle(h) };
    }
    match games.len() {
        1 => Ok(games[0]),
        0 => Err("Minecraft is not running".into()),
        n => Err(format!("{n} Minecraft processes are running")),
    }
}