image-help/src/lib.rs

#![allow(dead_code)]

#[cfg(test)]
mod tests;

use std::{
    ffi::CString,
    ptr::{null_mut, NonNull},
    sync::{Arc, RwLock},
};

use lazy_static::lazy_static;
use log::error;
use utils::{string::ToWide, win32::Win32Error};
use winapi::{
    shared::ntdef::{HANDLE, ULONG},
    um::{
        dbghelp::{
            self, SYMOPT_DEBUG, SYMOPT_DEFERRED_LOADS, SYMOPT_FAVOR_COMPRESSED, SYMOPT_UNDNAME,
        },
        libloaderapi::{GetProcAddress, LoadLibraryW},
        processthreadsapi::GetCurrentProcess,
    },
};

pub mod error {
    use std::ffi::NulError;

    #[derive(Debug, thiserror::Error)]
    pub enum Error {
        #[error(transparent)]
        ImageHelpError(#[from] ImageHelpError),
        #[error(transparent)]
        NullError(#[from] NulError),
        #[error(transparent)]
        Win32Error(#[from] utils::win32::error::Win32Error),
    }

    #[derive(Debug, thiserror::Error)]
    pub enum ImageHelpError {
        #[error("Module could not be found")]
        ModuleNotFound,
    }
}

pub struct ImageHelp {
    process_handle: Arc<RwLock<HANDLE>>,
}

unsafe impl Send for ImageHelp {}
unsafe impl Sync for ImageHelp {}

impl ImageHelp {
    pub fn init() -> Self {
        let process_handle = unsafe { GetCurrentProcess() };

        unsafe {
            dbghelp::SymInitializeW(process_handle, std::ptr::null_mut(), 0);

            dbghelp::SymSetOptions(
                SYMOPT_DEFERRED_LOADS | SYMOPT_DEBUG | SYMOPT_UNDNAME | SYMOPT_FAVOR_COMPRESSED,
            );
        }

        Self {
            process_handle: Arc::new(RwLock::new(process_handle)),
        }
    }

    pub fn get() -> &'static Self {
        &IMAGE_HELP
    }

    fn add_sym_options(&self, new_options: u32) {
        unsafe {
            let options = dbghelp::SymGetOptions();
            dbghelp::SymSetOptions(options | new_options);
        }
    }

    fn remove_sym_options(&self, new_options: u32) {
        unsafe {
            let options = dbghelp::SymGetOptions();
            dbghelp::SymSetOptions(options & !new_options);
        }
    }

    pub fn find_function(
        &self,
        module_name: &str,
        function_name: &str,
    ) -> Result<*const u8, error::Error> {
        let process_handle = self.process_handle.read().unwrap();

        let module_name_w = module_name.to_wide_null();

        let module_handle = unsafe {
            let handle = LoadLibraryW(module_name_w.as_ptr());

            if handle == std::ptr::null_mut() {
                error!("Module handle {} is NULL.", module_name);
                Err(error::ImageHelpError::ModuleNotFound)
            } else {
                Ok(handle)
            }
        }?;

        log::debug!("trying GetProcAddress");

        let proc_addr = unsafe {
            let func_cname = CString::new(function_name)?;
            let proc_adr = GetProcAddress(module_handle, func_cname.as_ptr());

            NonNull::new(proc_adr)
        };

        if let Some(proc_addr) = proc_addr {
            log::debug!("Found address for symbol with GetProcAddress");
            return Ok(proc_addr.as_ptr() as *const _);
        }

        log::debug!("SymLoadModuleExW..");

        match utils::win32::error::true_or_last_error(|| unsafe {
            dbghelp::SymLoadModuleExW(
                GetCurrentProcess(),
                null_mut(),
                module_name_w.as_ptr(),
                null_mut(),
                module_handle as u64,
                0,
                null_mut(),
                0,
            ) != 0
        }) {
            Ok(_) => Ok(()),
            Err(Win32Error { error_code: 0 }) => {
                log::debug!("Symbols already loaded.");
                Ok(())
            }
            Err(err) => Err(err),
        }?;

        let try_find_symbol = || {
            let mut sym_info =
                unsafe { std::mem::MaybeUninit::<dbghelp::SYMBOL_INFOW>::zeroed().assume_init() };
            sym_info.SizeOfStruct = std::mem::size_of::<dbghelp::SYMBOL_INFOW>() as ULONG;

            let symbol_name_utf16 = function_name.to_wide_null();
            sym_info.MaxNameLen = 255;

            log::debug!("SymFromNameW..");

            utils::win32::error::true_or_last_error(|| unsafe {
                dbghelp::SymFromNameW(*process_handle, symbol_name_utf16.as_ptr(), &mut sym_info)
                    != 0
            })?;

            Ok(sym_info.Address as *const u8)
        };

        self.add_sym_options(SYMOPT_UNDNAME);

        log::debug!("trying to find symbol with SYMOPT_UNDNAME");
        try_find_symbol().or_else(|_| {
            self.remove_sym_options(SYMOPT_UNDNAME);

            log::debug!("trying to find symbol without SYMOPT_UNDNAME");
            try_find_symbol()
        })
    }
}

lazy_static! {
    pub static ref IMAGE_HELP: ImageHelp = ImageHelp::init();
}