Implement kernel driver enumeration via NtQuerySystemInformation and remove module lookup panic

Author: ko1N

Cargo.toml

@@ -33,6 +33,7 @@ procfs = { version = "=0.15.1", features = ["backtrace"] }
 
 [target.'cfg(target_os = "windows")'.dependencies]
 windows = { version = "0.61", features = [
+	"Wdk_System_SystemInformation",
 	"Wdk_System_Threading",
 	"Win32_Foundation",
 	"Win32_Security",

src/windows/mod.rs

@@ -2,10 +2,11 @@ use memflow::os::process::*;
 use memflow::prelude::v1::*;
 
 use windows::core::PCSTR;
+use windows::Wdk::System::SystemInformation::{NtQuerySystemInformation, SYSTEM_INFORMATION_CLASS};
 use windows::Win32::Foundation::{
-    CloseHandle, HANDLE, NTSTATUS, STATUS_ACCESS_DENIED, STATUS_INSUFFICIENT_RESOURCES,
-    STATUS_INVALID_HANDLE, STATUS_INVALID_INFO_CLASS, STATUS_INVALID_PARAMETER,
-    STATUS_NOT_IMPLEMENTED, STATUS_NO_MEMORY, STATUS_PRIVILEGE_NOT_HELD,
+    CloseHandle, HANDLE, NTSTATUS, STATUS_ACCESS_DENIED, STATUS_INFO_LENGTH_MISMATCH,
+    STATUS_INSUFFICIENT_RESOURCES, STATUS_INVALID_HANDLE, STATUS_INVALID_INFO_CLASS,
+    STATUS_INVALID_PARAMETER, STATUS_NOT_IMPLEMENTED, STATUS_NO_MEMORY, STATUS_PRIVILEGE_NOT_HELD,
 };
 
 use windows::Win32::System::Diagnostics::ToolHelp::{
@@ -19,21 +20,43 @@ use windows::Win32::Security::{
     TOKEN_ADJUST_PRIVILEGES, TOKEN_PRIVILEGES,
 };
 
-use core::mem::{size_of, MaybeUninit};
+use core::mem::{align_of, size_of, MaybeUninit};
 use core::ptr;
 use std::ffi::OsString;
 use std::os::windows::ffi::OsStringExt;
 
-pub mod mem;
-use mem::ProcessVirtualMemory;
+const SYSTEM_MODULE_INFORMATION_CLASS: SYSTEM_INFORMATION_CLASS = SYSTEM_INFORMATION_CLASS(11);
+const MODULE_QUERY_INITIAL_CAPACITY: usize = 256 * 1024;
+const MODULE_QUERY_MAX_RETRIES: usize = 6;
 
+pub mod mem;
 pub mod process;
 pub use process::WindowsProcess;
 
 pub mod keyboard;
 pub use keyboard::{WindowsKeyboard, WindowsKeyboardState};
 
-struct KernelModule {}
+struct KernelModule {
+    base: Address,
+    size: umem,
+    path: ReprCString,
+    name: ReprCString,
+}
+
+#[repr(C)]
+#[allow(non_snake_case)]
+struct RtlProcessModuleInformation {
+    Section: *mut core::ffi::c_void,
+    MappedBase: *mut core::ffi::c_void,
+    ImageBase: *mut core::ffi::c_void,
+    ImageSize: u32,
+    Flags: u32,
+    LoadOrderIndex: u16,
+    InitOrderIndex: u16,
+    LoadCount: u16,
+    OffsetToFileName: u16,
+    FullPathName: [u8; 256],
+}
 
 pub(crate) struct Handle(HANDLE);
 
@@ -65,17 +88,50 @@ pub fn conv_err(_err: windows::core::Error) -> Error {
 
 pub(crate) fn conv_ntstatus(status: NTSTATUS) -> Error {
     let kind = match status {
-        STATUS_ACCESS_DENIED | STATUS_PRIVILEGE_NOT_HELD => ErrorKind::Permissions,
+        STATUS_ACCESS_DENIED | STATUS_PRIVILEGE_NOT_HELD => ErrorKind::NotSupported,
         STATUS_INVALID_PARAMETER => ErrorKind::ArgValidation,
         STATUS_NOT_IMPLEMENTED | STATUS_INVALID_INFO_CLASS => ErrorKind::NotSupported,
-        STATUS_INSUFFICIENT_RESOURCES | STATUS_NO_MEMORY => ErrorKind::OutOfMemory,
+        STATUS_INSUFFICIENT_RESOURCES | STATUS_NO_MEMORY => ErrorKind::Unknown,
         STATUS_INVALID_HANDLE => ErrorKind::ProcessNotFound,
         _ => ErrorKind::Unknown,
     };
 
     Error(ErrorOrigin::OsLayer, kind)
 }
 
+fn split_basename(path: &str) -> &str {
+    path.rsplit_once('\\')
+        .or_else(|| path.rsplit_once('/'))
+        .map(|(_, name)| name)
+        .unwrap_or(path)
+}
+
+fn parse_module_path(info: &RtlProcessModuleInformation) -> (ReprCString, ReprCString) {
+    let path_end = info
+        .FullPathName
+        .iter()
+        .position(|&c| c == 0)
+        .unwrap_or(info.FullPathName.len());
+
+    let path_raw = &info.FullPathName[..path_end];
+    let path = String::from_utf8_lossy(path_raw);
+    let path = path.as_ref();
+
+    let name = if (info.OffsetToFileName as usize) < path_end {
+        let name_raw = &info.FullPathName[info.OffsetToFileName as usize..path_end];
+        let parsed = String::from_utf8_lossy(name_raw);
+        if parsed.is_empty() {
+            split_basename(path).to_owned()
+        } else {
+            parsed.into_owned()
+        }
+    } else {
+        split_basename(path).to_owned()
+    };
+
+    (path.into(), name.into())
+}
+
 unsafe fn enable_debug_privilege() -> Result<()> {
     let process = GetCurrentProcess();
     let mut token = HANDLE::default();
@@ -129,6 +185,85 @@ impl WindowsOs {
 
         Ok(Default::default())
     }
+
+    fn query_kernel_modules() -> Result<Vec<KernelModule>> {
+        let mut buffer = vec![0u8; MODULE_QUERY_INITIAL_CAPACITY];
+
+        for _ in 0..MODULE_QUERY_MAX_RETRIES {
+            let mut return_len = 0u32;
+            let status = unsafe {
+                NtQuerySystemInformation(
+                    SYSTEM_MODULE_INFORMATION_CLASS,
+                    buffer.as_mut_ptr().cast(),
+                    buffer.len() as u32,
+                    &mut return_len,
+                )
+            };
+
+            if status == STATUS_INFO_LENGTH_MISMATCH {
+                let table_off = align_of::<RtlProcessModuleInformation>();
+                let next_len = usize::max(
+                    buffer.len().saturating_mul(2),
+                    (return_len as usize).saturating_add(table_off),
+                );
+                buffer.resize(next_len, 0);
+                continue;
+            }
+
+            if status.is_err() {
+                return Err(conv_ntstatus(status));
+            }
+
+            let valid_len = if return_len as usize > buffer.len() {
+                buffer.len()
+            } else {
+                return_len as usize
+            };
+
+            let data = &buffer[..valid_len];
+            let table_off = align_of::<RtlProcessModuleInformation>();
+
+            if data.len() < table_off {
+                return Err(Error(ErrorOrigin::OsLayer, ErrorKind::Unknown));
+            }
+
+            let mut header_buf = [0u8; core::mem::size_of::<u32>()];
+            header_buf.copy_from_slice(&data[..core::mem::size_of::<u32>()]);
+            let header_number = u32::from_ne_bytes(header_buf);
+            let entry_size = core::mem::size_of::<RtlProcessModuleInformation>();
+            let entries_len = header_number as usize;
+            let table_len = entries_len
+                .checked_mul(entry_size)
+                .ok_or(Error(ErrorOrigin::OsLayer, ErrorKind::Unknown))?;
+            let table_end = table_off
+                .checked_add(table_len)
+                .ok_or(Error(ErrorOrigin::OsLayer, ErrorKind::Unknown))?;
+
+            if table_end > data.len() {
+                return Err(Error(ErrorOrigin::OsLayer, ErrorKind::ArgValidation));
+            }
+
+            let mut parsed = Vec::with_capacity(entries_len);
+            for idx in 0..entries_len {
+                let module_off = table_off + idx * entry_size;
+                let module = unsafe {
+                    (data.as_ptr().add(module_off) as *const RtlProcessModuleInformation)
+                        .read_unaligned()
+                };
+                let (path, name) = parse_module_path(&module);
+                parsed.push(KernelModule {
+                    base: Address::from(module.ImageBase as umem),
+                    size: module.ImageSize as umem,
+                    path,
+                    name,
+                });
+            }
+
+            return Ok(parsed);
+        }
+
+        Err(Error(ErrorOrigin::OsLayer, ErrorKind::Unknown))
+    }
 }
 
 impl Clone for WindowsOs {
@@ -244,6 +379,8 @@ impl Os for WindowsOs {
     /// # Arguments
     /// * `callback` - where to pass each matching module to. This is an opaque callback.
     fn module_address_list_callback(&mut self, mut callback: AddressCallback) -> Result<()> {
+        self.cached_modules = Self::query_kernel_modules()?;
+
         (0..self.cached_modules.len())
             .map(Address::from)
             .take_while(|a| callback.call(*a))
@@ -256,29 +393,19 @@ impl Os for WindowsOs {
     ///
     /// # Arguments
     /// * `address` - address where module's information resides in
-    fn module_by_address(&mut self, _address: Address) -> Result<ModuleInfo> {
-        /*self.cached_modules
-        .iter()
-        .skip(address.to_umem() as usize)
-        .next()
-        .map(|km| ModuleInfo {
-            address,
-            size: km.size as umem,
-            base: Address::NULL,
-            name: km
-                .name
-                .split("/")
-                .last()
-                .or(Some(""))
-                .map(ReprCString::from)
-                .unwrap(),
-            arch: self.info.arch,
-            path: km.name.clone().into(),
-            parent_process: Address::INVALID,
-        })
-        .ok_or(Error(ErrorOrigin::OsLayer, ErrorKind::NotFound))*/
-
-        todo!()
+    fn module_by_address(&mut self, address: Address) -> Result<ModuleInfo> {
+        self.cached_modules
+            .get(address.to_umem() as usize)
+            .map(|km| ModuleInfo {
+                address,
+                size: km.size,
+                base: km.base,
+                name: km.name.clone(),
+                arch: self.info.arch,
+                path: km.path.clone(),
+                parent_process: Address::INVALID,
+            })
+            .ok_or(Error(ErrorOrigin::OsLayer, ErrorKind::NotFound))
     }
 
     /// Retrieves address of the primary module structure of the process

src/windows/process.rs

@@ -3,7 +3,8 @@ use memflow::os::process::*;
 use memflow::prelude::v1::*;
 use memflow::types::gap_remover::GapRemover;
 
-use super::{conv_err, conv_ntstatus, ProcessVirtualMemory};
+use super::{conv_err, conv_ntstatus};
+use crate::windows::mem::ProcessVirtualMemory;
 
 use windows::Wdk::System::Threading::{NtQueryInformationProcess, ProcessBasicInformation};
 use windows::Win32::Foundation::{HINSTANCE, HMODULE, STILL_ACTIVE};