Sign in Sign up
kretrod/lodestone Public
Branches
master
338 lines (312 loc) · 11.7 KB Raw
//! The only privileged thing Lodestone does: open a process and read/write its
//! memory. Everything above this file is pure interpretation of those bytes.
//!
//! Deliberately *no* injection primitives live here — no VirtualAllocEx, no
//! CreateRemoteThread, no LoadLibrary. A true external never puts a single byte
//! of its own code inside the target.

#![cfg(windows)]

use std::ffi::{c_void, OsString};
use std::os::windows::ffi::OsStringExt;

use windows_sys::Win32::Foundation::{CloseHandle, HANDLE, MAX_PATH};
use windows_sys::Win32::System::Diagnostics::Debug::{ReadProcessMemory, WriteProcessMemory};
use windows_sys::Win32::System::Diagnostics::ToolHelp::{
    CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W, TH32CS_SNAPPROCESS,
};
use windows_sys::Win32::System::ProcessStatus::{
    EnumProcessModulesEx, GetModuleFileNameExW, GetModuleInformation, LIST_MODULES_ALL, MODULEINFO,
};
use windows_sys::Win32::System::Threading::{
    OpenProcess, PROCESS_QUERY_INFORMATION, PROCESS_VM_READ, PROCESS_VM_WRITE,
    PROCESS_VM_OPERATION,
};

pub type Result<T> = std::result::Result<T, String>;

/// Read + write + query. No PROCESS_CREATE_THREAD: we never run code in there.
const ACCESS: u32 =
    PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_VM_OPERATION | PROCESS_QUERY_INFORMATION;

#[derive(Debug, Clone)]
pub struct ProcInfo {
    pub pid: u32,
    pub name: String,
}

/// Snapshot every running process. Used to find `javaw.exe`.
pub fn list_processes() -> Result<Vec<ProcInfo>> {
    let mut out = Vec::new();
    // SAFETY: snapshot handle is checked and closed below.
    unsafe {
        let snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if snap.is_null() {
            return Err(format!("CreateToolhelp32Snapshot: {}", last_err()));
        }
        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 {
            out.push(ProcInfo {
                pid: e.th32ProcessID,
                name: wide_to_string(&e.szExeFile),
            });
            ok = Process32NextW(snap, &mut e);
        }
        CloseHandle(snap);
    }
    Ok(out)
}

#[derive(Debug, Clone)]
pub struct Module {
    pub name: String,
    pub path: String,
    pub base: u64,
    pub size: u32,
}

/// An attached process. Owns the OS handle; closes it on drop.
pub struct Process {
    pid: u32,
    handle: isize,
}

// SAFETY: ReadProcessMemory/WriteProcessMemory are safe to call concurrently on
// one handle, and we never mutate `handle` after construction.
unsafe impl Send for Process {}
unsafe impl Sync for Process {}

impl Drop for Process {
    fn drop(&mut self) {
        // SAFETY: handle came from OpenProcess and is closed exactly once.
        unsafe { CloseHandle(self.handle as HANDLE) };
    }
}

impl Process {
    pub fn attach(pid: u32) -> Result<Self> {
        // SAFETY: valid access mask; null return means failure.
        let h = unsafe { OpenProcess(ACCESS, 0, pid) };
        if h.is_null() {
            return Err(format!("OpenProcess({pid}): {}", last_err()));
        }
        Ok(Self { pid, handle: h as isize })
    }

    pub fn pid(&self) -> u32 {
        self.pid
    }

    /// Read `buf.len()` bytes at `addr`. Partial reads are an error: a short
    /// read here almost always means a stale pointer, and silently zero-filling
    /// would turn that into a confusing wrong answer higher up.
    pub fn read_into(&self, addr: u64, buf: &mut [u8]) -> Result<()> {
        let mut got: usize = 0;
        // SAFETY: buf is a valid local slice; the kernel validates `addr`.
        let ok = unsafe {
            ReadProcessMemory(
                self.handle as HANDLE,
                addr as *const c_void,
                buf.as_mut_ptr() as *mut c_void,
                buf.len(),
                &mut got,
            )
        };
        if ok == 0 || got != buf.len() {
            return Err(format!(
                "read {:#x} ({} bytes): {}",
                addr,
                buf.len(),
                last_err()
            ));
        }
        Ok(())
    }

    pub fn read_bytes(&self, addr: u64, len: usize) -> Result<Vec<u8>> {
        let mut v = vec![0u8; len];
        self.read_into(addr, &mut v)?;
        Ok(v)
    }

    pub fn write_bytes(&self, addr: u64, data: &[u8]) -> Result<()> {
        let mut put: usize = 0;
        // SAFETY: data is a valid local slice; the kernel validates `addr`.
        let ok = unsafe {
            WriteProcessMemory(
                self.handle as HANDLE,
                addr as *const c_void,
                data.as_ptr() as *const c_void,
                data.len(),
                &mut put,
            )
        };
        if ok == 0 || put != data.len() {
            return Err(format!("write {:#x} ({} bytes): {}", addr, data.len(), last_err()));
        }
        Ok(())
    }

    /// Best-effort read that yields `None` instead of an error. The heap walk
    /// hits unmapped pages constantly; those are expected, not exceptional.
    pub fn try_read<const N: usize>(&self, addr: u64) -> Option<[u8; N]> {
        let mut b = [0u8; N];
        self.read_into(addr, &mut b).ok().map(|_| b)
    }

    pub fn u8(&self, a: u64) -> Result<u8> {
        Ok(self.read_bytes(a, 1)?[0])
    }
    pub fn u16(&self, a: u64) -> Result<u16> {
        let mut b = [0u8; 2];
        self.read_into(a, &mut b)?;
        Ok(u16::from_le_bytes(b))
    }
    pub fn u32(&self, a: u64) -> Result<u32> {
        let mut b = [0u8; 4];
        self.read_into(a, &mut b)?;
        Ok(u32::from_le_bytes(b))
    }
    pub fn i32(&self, a: u64) -> Result<i32> {
        Ok(self.u32(a)? as i32)
    }
    pub fn u64(&self, a: u64) -> Result<u64> {
        let mut b = [0u8; 8];
        self.read_into(a, &mut b)?;
        Ok(u64::from_le_bytes(b))
    }
    pub fn f32(&self, a: u64) -> Result<f32> {
        Ok(f32::from_bits(self.u32(a)?))
    }
    pub fn f64(&self, a: u64) -> Result<f64> {
        Ok(f64::from_bits(self.u64(a)?))
    }

    pub fn write_u8(&self, a: u64, v: u8) -> Result<()> {
        self.write_bytes(a, &[v])
    }
    pub fn write_u32(&self, a: u64, v: u32) -> Result<()> {
        self.write_bytes(a, &v.to_le_bytes())
    }
    pub fn write_i32(&self, a: u64, v: i32) -> Result<()> {
        self.write_bytes(a, &v.to_le_bytes())
    }
    pub fn write_u64(&self, a: u64, v: u64) -> Result<()> {
        self.write_bytes(a, &v.to_le_bytes())
    }
    pub fn write_f32(&self, a: u64, v: f32) -> Result<()> {
        self.write_bytes(a, &v.to_le_bytes())
    }
    pub fn write_f64(&self, a: u64, v: f64) -> Result<()> {
        self.write_bytes(a, &v.to_le_bytes())
    }

    /// NUL-terminated ASCII/UTF-8 C string, read in chunks so we don't fault on
    /// a string that sits near the end of a page.
    pub fn cstring(&self, addr: u64, max: usize) -> Result<String> {
        let mut out = Vec::new();
        let mut a = addr;
        while out.len() < max {
            let chunk = self.read_bytes(a, 32)?;
            if let Some(p) = chunk.iter().position(|&c| c == 0) {
                out.extend_from_slice(&chunk[..p]);
                return Ok(String::from_utf8_lossy(&out).into_owned());
            }
            out.extend_from_slice(&chunk);
            a += 32;
        }
        Ok(String::from_utf8_lossy(&out).into_owned())
    }

    pub fn modules(&self) -> Result<Vec<Module>> {
        let mut handles = vec![0usize; 1024];
        let mut needed: u32 = 0;
        // SAFETY: handles is sized in bytes for the call; needed is written back.
        let ok = unsafe {
            EnumProcessModulesEx(
                self.handle as HANDLE,
                handles.as_mut_ptr() as *mut _,
                (handles.len() * std::mem::size_of::<usize>()) as u32,
                &mut needed,
                LIST_MODULES_ALL,
            )
        };
        if ok == 0 {
            return Err(format!("EnumProcessModulesEx: {}", last_err()));
        }
        let count = (needed as usize / std::mem::size_of::<usize>()).min(handles.len());
        let mut out = Vec::with_capacity(count);
        for &h in &handles[..count] {
            let mut namebuf = [0u16; MAX_PATH as usize];
            // SAFETY: h is a module handle from the enumeration above.
            let n = unsafe {
                GetModuleFileNameExW(
                    self.handle as HANDLE,
                    h as *mut _,
                    namebuf.as_mut_ptr(),
                    namebuf.len() as u32,
                )
            };
            let path = if n > 0 {
                String::from_utf16_lossy(&namebuf[..n as usize])
            } else {
                String::new()
            };
            let mut mi: MODULEINFO = unsafe { std::mem::zeroed() };
            // SAFETY: mi is a valid out-param of the documented size.
            let ok = unsafe {
                GetModuleInformation(
                    self.handle as HANDLE,
                    h as *mut _,
                    &mut mi,
                    std::mem::size_of::<MODULEINFO>() as u32,
                )
            };
            if ok == 0 {
                continue;
            }
            let name = path.rsplit('\\').next().unwrap_or("").to_string();
            out.push(Module {
                name,
                path,
                base: mi.lpBaseOfDll as u64,
                size: mi.SizeOfImage,
            });
        }
        Ok(out)
    }

    pub fn module(&self, name: &str) -> Result<Module> {
        self.modules()?
            .into_iter()
            .find(|m| m.name.eq_ignore_ascii_case(name))
            .ok_or_else(|| format!("module {name} not loaded"))
    }

    /// Resolve exported symbols out of a module's PE export directory, read
    /// straight from the target's mapped image. `jvm.dll` publishes the whole
    /// VMStructs table this way, which is how we avoid hardcoding any offset.
    pub fn exports(&self, m: &Module) -> Result<Vec<(String, u64)>> {
        let base = m.base;
        if self.u16(base)? != 0x5A4D {
            return Err("not an MZ image".into());
        }
        let e_lfanew = self.u32(base + 0x3C)? as u64;
        let nt = base + e_lfanew;
        if self.u32(nt)? != 0x0000_4550 {
            return Err("not a PE image".into());
        }
        // COFF header is 20 bytes; the optional header follows.
        let opt = nt + 4 + 20;
        let magic = self.u16(opt)?;
        // PE32+ puts the data directories at +112, PE32 at +96.
        let dir_off = if magic == 0x20B { 112 } else { 96 };
        let export_rva = self.u32(opt + dir_off)? as u64;
        if export_rva == 0 {
            return Ok(Vec::new());
        }
        let ed = base + export_rva;
        let n_names = self.u32(ed + 24)? as usize;
        let func_rva = self.u32(ed + 28)? as u64;
        let name_rva = self.u32(ed + 32)? as u64;
        let ord_rva = self.u32(ed + 36)? as u64;

        // Bulk-read the three parallel arrays instead of 3 reads per symbol.
        let names = self.read_bytes(base + name_rva, n_names * 4)?;
        let ords = self.read_bytes(base + ord_rva, n_names * 2)?;
        let mut out = Vec::with_capacity(n_names);
        for i in 0..n_names {
            let nrva = u32::from_le_bytes(names[i * 4..i * 4 + 4].try_into().unwrap()) as u64;
            let ord = u16::from_le_bytes(ords[i * 2..i * 2 + 2].try_into().unwrap()) as u64;
            let name = self.cstring(base + nrva, 256)?;
            let addr_rva = self.u32(base + func_rva + ord * 4)? as u64;
            out.push((name, base + addr_rva));
        }
        Ok(out)
    }
}

fn wide_to_string(w: &[u16]) -> String {
    let end = w.iter().position(|&c| c == 0).unwrap_or(w.len());
    OsString::from_wide(&w[..end]).to_string_lossy().into_owned()
}

fn last_err() -> String {
    std::io::Error::last_os_error().to_string()
}