From 8af8c4856bd356c8fde2619bf86d1830d17d13ef Mon Sep 17 00:00:00 2001
From: a2x <45197573+a2x@users.noreply.github.com>
Date: Thu, 28 Sep 2023 19:48:27 +1000
Subject: [PATCH] Check if `AddressOfNameOrdinals[i]` is valid

---
 src/error.rs         |  6 ++++++
 src/remote/module.rs | 47 +++++++++++++++++++++++++++++++++++++-------
 2 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index 2c9c3da..f690b4b 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -2,6 +2,7 @@ use std::io;
 
 #[derive(Debug)]
 pub enum Error {
+    BufferSizeMismatch(usize, usize),
     InvalidMagic(u32),
     IOError(io::Error),
     ModuleNotFound,
@@ -33,6 +34,11 @@ impl From<windows::core::Error> for Error {
 impl std::fmt::Display for Error {
     fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
+            Self::BufferSizeMismatch(expected, actual) => write!(
+                fmt,
+                "Buffer size mismatch: expected {}, got {}",
+                expected, actual
+            ),
             Self::InvalidMagic(magic) => write!(fmt, "Invalid magic: {:#X}", magic),
             Self::IOError(err) => write!(fmt, "IO error: {}", err),
             Self::ModuleNotFound => write!(fmt, "Module not found"),
diff --git a/src/remote/module.rs b/src/remote/module.rs
index 94f24e6..d5e731a 100644
--- a/src/remote/module.rs
+++ b/src/remote/module.rs
@@ -1,4 +1,5 @@
 use std::ffi::CStr;
+use std::mem;
 use std::slice;
 
 use windows::Win32::System::Diagnostics::Debug::*;
@@ -25,6 +26,7 @@ pub struct Section {
 pub struct Module<'a> {
     address: usize,
     nt_headers: &'a IMAGE_NT_HEADERS64,
+    size: u32,
     exports: Vec<Export>,
     sections: Vec<Section>,
 }
@@ -35,6 +37,13 @@ impl<'a> Module<'a> {
 
         process.read_memory_raw(address, headers.as_mut_ptr() as *mut _, headers.len())?;
 
+        if headers.len() < mem::size_of::<IMAGE_DOS_HEADER>() {
+            return Err(Error::BufferSizeMismatch(
+                mem::size_of::<IMAGE_DOS_HEADER>(),
+                headers.len(),
+            ));
+        }
+
         let dos_header = unsafe { &*(headers.as_ptr() as *const IMAGE_DOS_HEADER) };
 
         if dos_header.e_magic != IMAGE_DOS_SIGNATURE {
@@ -49,12 +58,15 @@ impl<'a> Module<'a> {
             return Err(Error::InvalidMagic(nt_headers.Signature));
         }
 
-        let exports = unsafe { Self::parse_exports(process, address, nt_headers)? };
+        let size = nt_headers.OptionalHeader.SizeOfImage;
+
+        let exports = unsafe { Self::parse_exports(process, address, size, nt_headers)? };
         let sections = unsafe { Self::parse_sections(nt_headers) };
 
         Ok(Self {
             address,
             nt_headers,
+            size,
             exports,
             sections,
         })
@@ -87,12 +99,13 @@ impl<'a> Module<'a> {
 
     #[inline]
     pub fn size(&self) -> u32 {
-        self.nt_headers.OptionalHeader.SizeOfImage
+        self.size
     }
 
     unsafe fn parse_exports(
         process: &Process,
         address: usize,
+        size: u32,
         nt_headers: &IMAGE_NT_HEADERS64,
     ) -> Result<Vec<Export>> {
         let export_data_directory =
@@ -109,6 +122,13 @@ impl<'a> Module<'a> {
             buffer.len(),
         )?;
 
+        if buffer.len() < mem::size_of::<IMAGE_EXPORT_DIRECTORY>() {
+            return Err(Error::BufferSizeMismatch(
+                mem::size_of::<IMAGE_EXPORT_DIRECTORY>(),
+                buffer.len(),
+            ));
+        }
+
         let export_directory = &*(buffer.as_ptr() as *const IMAGE_EXPORT_DIRECTORY);
 
         let delta =
@@ -121,7 +141,18 @@ impl<'a> Module<'a> {
         let mut exports: Vec<Export> = Vec::with_capacity(export_directory.NumberOfNames as usize);
 
         for i in 0..export_directory.NumberOfNames {
+            let target = ordinal_table as usize + i as usize * mem::size_of::<u16>();
+
+            if target > address + size as usize || target < ordinal_table as usize {
+                continue;
+            }
+
             let function_ordinal = *ordinal_table.offset(i as isize);
+
+            if function_ordinal as usize > export_directory.NumberOfFunctions as usize {
+                continue;
+            }
+
             let function_va = address + *function_table.offset(function_ordinal as isize) as usize;
 
             // Skip forwarded exports.
@@ -129,11 +160,13 @@ impl<'a> Module<'a> {
                 continue;
             }
 
-            let name = CStr::from_ptr(
-                delta.wrapping_add(*name_table.offset(i as isize) as usize) as *const i8
-            )
-            .to_string_lossy()
-            .into_owned();
+            let name_ptr = delta.wrapping_add(*name_table.offset(i as isize) as usize) as *const i8;
+
+            if name_ptr.is_null() {
+                continue;
+            }
+
+            let name = CStr::from_ptr(name_ptr).to_string_lossy().into_owned();
 
             exports.push(Export {
                 name,