Merge pull request #10 from TotemArts/cleanup

Clean up some UDK slice handling

README.md

@@ -12,7 +12,6 @@ Installation simply involves copying our DLL into the game's binary directory. N
    * `udk_offsets.rs` - Constants describing important offsets in the UDK binary
    * `udk_xaudio.rs` - UDK XAudio FFI and detours
    * `xaudio27.rs` - XAudio2.7 -> 2.9 compatibility layer
- * `winbindings/` - Windows API Rust bindings.
 
 ## Loading the extensions
 When the system loads the UDK, it will load all DLL dependencies alongside the UDK before executing any game code.

src/lib.rs

@@ -19,6 +19,9 @@ const UDK_KNOWN_HASH: [u8; 32] = [
     0x6C, 0x1E, 0x8F, 0x9E, 0xF0, 0x70, 0x40, 0xB8, 0xF9, 0x96, 0x73, 0x8A, 0x00, 0xFB, 0x90, 0x07,
 ];
 
+use std::ops::Range;
+use std::sync::OnceLock;
+
 use anyhow::Context;
 use sha2::{Digest, Sha256};
 
@@ -36,16 +39,15 @@ use windows::Win32::{
     },
 };
 
-/// Cached slice of UDK.exe. This is only touched once upon init, and
-/// never written again.
-// FIXME: The slice is actually unsafe to access; sections of memory may be unmapped!
-// We should use a raw pointer slice instead (if ergonomics permit doing so).
-static mut UDK_SLICE: Option<&'static [u8]> = None;
+/// Cached memory range for UDK.exe
+static UDK_RANGE: OnceLock<Range<usize>> = OnceLock::new();
+
+/// Return the base pointer for UDK.exe
+pub fn get_udk_ptr() -> *const u8 {
+    let range = UDK_RANGE.get().unwrap();
 
-/// Return a slice of UDK.exe
-pub fn get_udk_slice() -> &'static [u8] {
-    // SAFETY: This is only touched once in DllMain.
-    unsafe { UDK_SLICE.unwrap() }
+    // TODO: Once Rust gets better raw slice support, we should return a `*const [u8]` instead.
+    range.start as *const u8
 }
 
 /// Wrapped version of the Win32 OutputDebugString.
@@ -93,20 +95,14 @@ fn get_module_information(process: HANDLE, module: HINSTANCE) -> windows::core::
     }
 }
 
-/// Create a raw slice from a MODULEINFO structure.
-fn get_module_slice(info: &MODULEINFO) -> *const [u8] {
-    core::ptr::slice_from_raw_parts(info.lpBaseOfDll as *const u8, info.SizeOfImage as usize)
-}
-
 /// Called upon DLL attach. This function verifies the UDK and initializes
 /// hooks if the UDK matches our known hash.
 fn dll_attach() -> anyhow::Result<()> {
     let process = unsafe { GetCurrentProcess() };
     let module = unsafe { GetModuleHandleA(None) }.context("failed to get module handle")?;
 
-    let exe_slice = get_module_slice(
-        &get_module_information(process, module.into()).expect("Failed to get module information for UDK"),
-    );
+    let exe_info = get_module_information(process, module.into())
+        .context("Failed to get module information for UDK")?;
 
     // Now that we're attached, let's hash the UDK executable.
     // If the hash does not match what we think it should be, do not attach detours.
@@ -127,9 +123,12 @@ fn dll_attach() -> anyhow::Result<()> {
     }
 
     // Cache the UDK slice.
-    unsafe {
-        UDK_SLICE = Some(exe_slice.as_ref().unwrap());
-    }
+    UDK_RANGE
+        .set(Range {
+            start: exe_info.lpBaseOfDll as usize,
+            end: exe_info.lpBaseOfDll as usize + exe_info.SizeOfImage as usize,
+        })
+        .unwrap();
 
     // Initialize detours.
     udk_xaudio::init()?;

src/udk_log.rs

@@ -1,5 +1,5 @@
 //! This module contains functionality relevant to UDK logging.
-use crate::get_udk_slice;
+use crate::get_udk_ptr;
 
 use crate::udk_offsets::{DEBUG_FN_OFFSET, DEBUG_LOG_OFFSET};
 
@@ -19,9 +19,9 @@ pub enum LogType {
 
 /// Log a message via the UDK logging framework.
 pub fn log(typ: LogType, msg: &str) {
-    let udk_slice = get_udk_slice();
-    let log_obj = unsafe { udk_slice.as_ptr().add(DEBUG_LOG_OFFSET) };
-    let log_fn: UDKLogFn = unsafe { std::mem::transmute(udk_slice.as_ptr().add(DEBUG_FN_OFFSET)) };
+    let udk_slice = get_udk_ptr();
+    let log_obj = unsafe { udk_slice.add(DEBUG_LOG_OFFSET) };
+    let log_fn: UDKLogFn = unsafe { std::mem::transmute(udk_slice.add(DEBUG_FN_OFFSET)) };
 
     // Convert the UTF-8 Rust string into an OS wide string.
     let wmsg: widestring::U16CString = widestring::WideCString::from_str(format!("TotemArts Extensions: {}", msg)).unwrap();

src/udk_xaudio.rs

@@ -2,7 +2,7 @@
 use anyhow::Context;
 use retour::static_detour;
 
-use crate::get_udk_slice;
+use crate::get_udk_ptr;
 use crate::udk_log::{log, LogType};
 use crate::udk_offsets::{UDK_CREATEFX_PTR_OFFSET, UDK_XAUDIO2CREATE_OFFSET};
 use crate::xaudio27::{IXAudio27, XAudio27Wrapper};
@@ -98,13 +98,13 @@ fn xaudio2create_hook(xaudio2_out: *mut IXAudio27, _flags: u32, _processor: u32)
 }
 
 pub fn init() -> anyhow::Result<()> {
-    let udk = get_udk_slice();
+    let udk = get_udk_ptr();
 
     // SAFETY: This is only safe if the UDK binary matches what we expect.
     unsafe {
         XAudio2CreateHook
             .initialize(
-                std::mem::transmute(udk.as_ptr().add(UDK_XAUDIO2CREATE_OFFSET)),
+                std::mem::transmute(udk.add(UDK_XAUDIO2CREATE_OFFSET)),
                 xaudio2create_hook,
             )
             .context("Failed to setup InitializeHardware hook")?;
@@ -113,14 +113,14 @@ pub fn init() -> anyhow::Result<()> {
 
         // Enable RW access to the CreateFX pointer.
         let _guard = region::protect_with_handle(
-            udk.as_ptr().add(UDK_CREATEFX_PTR_OFFSET),
+            udk.add(UDK_CREATEFX_PTR_OFFSET),
             8,
             region::Protection::READ_WRITE,
         )
         .context("failed to adjust memory protection for CreateFX")?;
 
         // Overwrite xapofx!CreateFX pointer with our hook.
-        (udk.as_ptr().add(UDK_CREATEFX_PTR_OFFSET) as *mut usize).write(createfx_hook as usize);
+        (udk.add(UDK_CREATEFX_PTR_OFFSET) as *mut usize).write(createfx_hook as usize);
     }
 
     Ok(())