From 8d6ad1f8122d92702bb854fd1e1d56c9b2bcbfba Mon Sep 17 00:00:00 2001 From: KiraCoding <38864051+KiraCoding@users.noreply.github.com> Date: Fri, 13 Sep 2024 06:47:18 +0200 Subject: [PATCH] wip --- Cargo.lock | 69 ++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 3 +++ src/base.rs | 13 +++++----- src/program.rs | 64 ++++++++++++++++++++-------------------------- src/section.rs | 2 +- 5 files changed, 107 insertions(+), 44 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d8db507..2ce2088 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -27,6 +27,21 @@ version = "0.8.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" +[[package]] +name = "dataview" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50eb3a329e19d78c3a3dfa4ec5a51ecb84fa3a20c06edad04be25356018218f9" +dependencies = [ + "derive_pod", +] + +[[package]] +name = "derive_pod" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2ea6706d74fca54e15f1d40b5cf7fe7f764aaec61352a9fcec58fe27e042fc8" + [[package]] name = "either" version = "1.13.0" @@ -37,10 +52,42 @@ checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" name = "inka" version = "0.1.0" dependencies = [ + "pelite", "rayon", "windows", ] +[[package]] +name = "libc" +version = "0.2.158" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439" + +[[package]] +name = "no-std-compat" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" + +[[package]] +name = "pelite" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88dccf4bd32294364aeb7bd55d749604450e9db54605887551f21baea7617685" +dependencies = [ + "dataview", + "libc", + "no-std-compat", + "pelite-macros", + "winapi", +] + +[[package]] +name = "pelite-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a7cf3f8ecebb0f4895f4892a8be0a0dc81b498f9d56735cb769dc31bf00815b" + [[package]] name = "proc-macro2" version = "1.0.86" @@ -96,6 +143,28 @@ version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows" version = "0.58.0" diff --git a/Cargo.toml b/Cargo.toml index 4b94ad7..07c527f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,9 @@ description = "todo" [dependencies] rayon = "1.10.0" +[target.'cfg(windows)'.dependencies.pelite] +version = "0.10.0" + [target.'cfg(windows)'.dependencies.windows] version = "0.58.0" features = [ diff --git a/src/base.rs b/src/base.rs index 1b61623..c61d2df 100644 --- a/src/base.rs +++ b/src/base.rs @@ -14,18 +14,18 @@ impl Base { pub fn program() -> Self { let raw_base = unsafe { GetModuleHandleW(PCWSTR::null()).unwrap_unchecked().0.cast() }; - // SAFETY: `raw_base` is a valid and non-null. + // SAFETY: `raw_base` is valid and non-null. unsafe { Self::new_unchecked(raw_base) } } #[inline] - pub fn as_ptr(&self) -> NonNull { - self.ptr + pub const unsafe fn add(&self, count: usize) -> NonNull { + unsafe { self.ptr.add(count) } } #[inline] - pub unsafe fn add(&self, count: usize) -> NonNull { - unsafe { self.ptr.add(count) } + pub(crate) const fn as_nonnull(&self) -> NonNull { + self.ptr } /// Creates a new `Base`. @@ -33,7 +33,8 @@ impl Base { /// # Safety /// /// `ptr` must be non-null. - pub(crate) unsafe fn new_unchecked(ptr: *mut u8) -> Self { + #[inline] + pub(crate) const unsafe fn new_unchecked(ptr: *mut u8) -> Self { // SAFETY: Caller ensures that `ptr` is non-null. let ptr = unsafe { NonNull::new_unchecked(ptr) }; diff --git a/src/program.rs b/src/program.rs index dde71ca..86971f7 100644 --- a/src/program.rs +++ b/src/program.rs @@ -3,11 +3,10 @@ use core::ops::Index; use core::ptr::NonNull; use core::slice::{from_raw_parts, SliceIndex}; use core::str::from_utf8_unchecked; +use pelite::pe::{Pe, PeView}; use rayon::iter::IndexedParallelIterator; use rayon::slice::ParallelSlice; use std::sync::LazyLock; -use windows::Win32::System::Diagnostics::Debug::{IMAGE_NT_HEADERS64, IMAGE_SECTION_HEADER}; -use windows::Win32::System::SystemServices::IMAGE_DOS_HEADER; static PROGRAM: LazyLock = LazyLock::new(Program::init); @@ -44,7 +43,7 @@ impl Program { #[inline] pub fn as_slice(&self) -> &[u8] { // SAFETY: todo!() - unsafe { from_raw_parts(self.base.as_ptr().as_ptr(), self.len) } + unsafe { from_raw_parts(self.base.as_nonnull().as_ptr(), self.len) } } /// Returns `true` if the program contains the byte pattern. @@ -110,7 +109,7 @@ impl Program { /// Returns [`None`] if the pattern doesn’t match. /// /// # Examples - /// + /// /// ``` /// use inka::program; /// @@ -121,13 +120,13 @@ impl Program { /// 0x5f, 0x00, 0xf1, 0x10, /// 0x80, 0x5e, 0x5f, 0xbf, /// ]; - /// + /// /// let pattern = &[ /// 0x7c, 0x73, 0xe1, 0x3d, /// 0x1a, 0x7d, 0xb3, 0x00, /// 0xd2, /// ]; - /// + /// /// let ptr = program() /// .rfind(pattern) /// .unwrap(); @@ -146,43 +145,34 @@ impl Program { fn init() -> Self { let base = Base::program(); - let dos_header = base.as_ptr().as_ptr() as *const IMAGE_DOS_HEADER; - let nt_headers64: &IMAGE_NT_HEADERS64 = - unsafe { &*(base.add((*dos_header).e_lfanew as usize).as_ptr().cast()) }; - - let len = nt_headers64.OptionalHeader.SizeOfImage as usize; - - let section_header_ptr = unsafe { - (nt_headers64 as *const IMAGE_NT_HEADERS64).add(1) as *const IMAGE_SECTION_HEADER - }; + { + let pe = unsafe { PeView::module(base.as_nonnull().as_ptr()) }; - let sections = (0..nt_headers64.FileHeader.NumberOfSections) - .map(|index| unsafe { &*section_header_ptr.add(index as usize) }) - .map(|section| { - let name = { - let name_len = section - .Name - .iter() - .position(|&char| char == 0) - .unwrap_or(section.Name.len()); + let len = pe.nt_headers().OptionalHeader.SizeOfImage as usize; - unsafe { from_utf8_unchecked(§ion.Name[..name_len]) } - }; + let sections = pe + .section_headers() + .iter() + .map(|section| { + let name = section.name().unwrap(); - let section_base = unsafe { - Base::new_unchecked(base.add(section.VirtualAddress as usize).as_ptr().cast()) - }; + let base = unsafe { + Base::new_unchecked( + base.add(section.VirtualAddress as usize).as_ptr().cast(), + ) + }; - let len = unsafe { section.Misc.VirtualSize as usize }; + let len = section.VirtualSize as usize; - Section::new(name, section_base, len) - }) - .collect(); + Section::new(name, base, len) + }) + .collect(); - Self { - base, - len, - sections, + Self { + base, + len, + sections, + } } } } diff --git a/src/section.rs b/src/section.rs index 894c64e..8d20da5 100644 --- a/src/section.rs +++ b/src/section.rs @@ -34,7 +34,7 @@ impl Section { #[inline] pub fn as_slice(&self) -> &[u8] { - unsafe { from_raw_parts(self.base.as_ptr().as_ptr(), self.len) } + unsafe { from_raw_parts(self.base.as_nonnull().as_ptr(), self.len) } } pub fn contains(&self, pattern: &[u8]) -> bool {