unreal-sdk/pdb-helper/src/lib.rs

use std::{ptr::NonNull, sync::Mutex};

use anyhow::Context;
use once_cell::sync::OnceCell;
use pdb::{DataSymbol, FallibleIterator, ProcedureReferenceSymbol, ProcedureSymbol, SymbolData};
use windows::{core::PCWSTR, Win32::System::LibraryLoader::GetModuleHandleW};

#[derive(Debug, Clone, Copy)]
pub struct Globals {
    pub names: NonNull<()>,
    pub objects: NonNull<()>,
    pub process_event: usize,
}

static PDB: OnceCell<Mutex<PdbCtx>> = OnceCell::new();

pub fn get_pdb() -> std::sync::MutexGuard<'static, PdbCtx> {
    PDB.get_or_init(|| Mutex::new(PdbCtx::new().expect("failed to create PDB helper.")))
        .lock()
        .unwrap()
}

#[derive(Debug)]
pub struct PdbCtx {
    base: usize,
    pdb: pdb::PDB<'static, std::fs::File>,
}

// i dont care i need this to be safe and its stored as a dyn so any Send trait is removed
unsafe impl Send for PdbCtx {}

impl PdbCtx {
    pub fn new() -> anyhow::Result<Self> {
        let base = unsafe {
            GetModuleHandleW(PCWSTR::null())
                .context("could not get process base")?
                .0 as usize
        };

        let pdb_path = std::env::current_exe()
            .context("current exe")?
            .with_file_name("ShooterGame.pdb");
        let pdb_file = std::fs::File::open(&pdb_path).context("pdb file open")?;

        let pdb = pdb::PDB::open(pdb_file).context("parse pdb file")?;

        Ok(Self { base, pdb })
    }

    /// returns the address of the function
    pub fn find_function(&mut self, function_name: &str) -> anyhow::Result<usize> {
        let globals = self.pdb.global_symbols().context("global symbols")?;
        let mut iter = globals.iter();

        while let Some(symbol) = iter.next()? {
            match symbol.parse() {
                Ok(pdb::SymbolData::ProcedureReference(data))
                    if data.name.map(|name| name.to_string() == function_name) == Some(true) =>
                {
                    log::trace!("{data:?}");
                    let dbg_info = self.pdb.debug_information()?;
                    let mods = dbg_info
                        .modules()?
                        .nth(data.module.context("module idx none")?)
                        .context("module not found")?
                        .context("module found but none")?;
                    let a = self
                        .pdb
                        .module_info(&mods)
                        .context("module info not found")?
                        .context("module info found but none")?;
                    let sym = a
                        .symbols_at(data.symbol_index)
                        .context("symbol not found")?
                        .next()
                        .context("symbol failed to read")?
                        .context("symbol read but none")?
                        .parse()
                        .context("symbol failed to parse")?;

                    log::debug!("actual symbol: {sym:?}");

                    if let SymbolData::Procedure(ProcedureSymbol { offset, .. }) = sym {
                        // asdf
                        let get_names = offset
                            .to_rva(&self.pdb.address_map()?)
                            .context("rva none")?;
                        return Ok(self.base + get_names.0 as usize);
                    } else {
                        return Err(anyhow::anyhow!("no proceduresymbol found"));
                    }
                }
                _ => {}
            }
        }
        return Err(anyhow::anyhow!("no procedureref found"));
    }

    pub fn parse_names(&mut self, data: ProcedureReferenceSymbol) -> anyhow::Result<usize> {
        log::trace!("{data:?}");
        let dbg_info = self.pdb.debug_information()?;
        let mods = dbg_info
            .modules()?
            .nth(data.module.context("module idx none")?)
            .context("module not found")?
            .context("module found but none")?;
        let a = self
            .pdb
            .module_info(&mods)
            .context("module info not found")?
            .context("module info found but none")?;
        let sym = a
            .symbols_at(data.symbol_index)
            .context("symbol not found")?
            .next()
            .context("symbol failed to read")?
            .context("symbol read but none")?
            .parse()
            .context("symbol failed to parse")?;

        log::debug!("actual symbol: {sym:?}");

        if let SymbolData::Procedure(ProcedureSymbol { offset, .. }) = sym {
            // asdf
            let get_names = offset
                .to_rva(&self.pdb.address_map()?)
                .context("rva none")?;
            let get_names = unsafe {
                core::mem::transmute::<_, unsafe extern "win64" fn() -> *mut ()>(
                    self.base + get_names.0 as usize,
                )
            };

            let names = unsafe { get_names() };

            log::debug!("names: {:?}", names);

            Ok(names as usize)
        } else {
            Err(anyhow::anyhow!("no proceduresymbol found"))
        }
    }

    pub fn parse_objects(&mut self, data: DataSymbol) -> anyhow::Result<usize> {
        let get_names = data
            .offset
            .to_rva(&self.pdb.address_map()?)
            .context("rva none")?;
        let objects = self.base + get_names.0 as usize;

        Ok(objects)
    }

    pub fn find_ue_globals(&mut self) -> anyhow::Result<Globals> {
        let globals = self.pdb.global_symbols().context("global symbols")?;
        let mut iter = globals.iter();

        let mut names = None;
        let mut objects = None;

        while let Some(symbol) = iter.next()? {
            match symbol.parse() {
                Ok(pdb::SymbolData::Data(data)) if data.name.to_string() == "GUObjectArray" => {
                    objects = Some(self.parse_objects(data)?);
                }
                Ok(pdb::SymbolData::ProcedureReference(data))
                    if data.name.map(|name| name.to_string() == "FName::GetNames")
                        == Some(true) =>
                {
                    names = Some(self.parse_names(data)?);
                }
                _ => {}
            }
        }

        let process_event = self
            .find_function("UObject::ProcessEvent")
            .context("could not find UObject::ProcessEvent!")?;

        let names = names.context("could not find names!")?;
        let objects = objects.context("could not find objects!")?;

        Ok(Globals {
            names: NonNull::new(names as _).context("GNames was null!")?,
            objects: NonNull::new(objects as _).context("GObjects was null!")?,
            process_event,
        })
    }
}