feat(protocol): updated the driver interface to the updated driver pr…

…otocol

driver/src/handler/init.rs

@@ -52,7 +52,8 @@ pub fn handler_init(command: &mut DriverCommandInitialize) -> anyhow::Result<()>
 
     command.result = InitializeResult::Success;
     command.driver_version = driver_version();
-    command.driver_features = DriverFeature::ProcessModules |
+    command.driver_features = DriverFeature::ProcessList |
+        DriverFeature::ProcessModules |
         DriverFeature::MemoryRead |
         DriverFeature::MemoryWrite |
         feature_mouse |

driver/src/handler/memory_read.rs

@@ -2,7 +2,7 @@ use alloc::vec::Vec;
 
 use kapi::Process;
 use valthrun_driver_protocol::{
-    command::DriverCommandProcessMemoryRead,
+    command::DriverCommandMemoryRead,
     types::MemoryAccessResult,
 };
 use winapi::{
@@ -70,7 +70,7 @@ fn read_memory_physical(ctx: &mut ReadContext) -> bool {
     pmem::read_process_memory(ctx.process, ctx.target_address, ctx.read_buffer).is_ok()
 }
 
-pub fn handler_read(command: &mut DriverCommandProcessMemoryRead) -> anyhow::Result<()> {
+pub fn handler_read(command: &mut DriverCommandMemoryRead) -> anyhow::Result<()> {
     let out_buffer = unsafe { core::slice::from_raw_parts_mut(command.buffer, command.count) };
     if !seh::probe_write(
         out_buffer as *const _ as *const () as u64,

driver/src/handler/memory_write.rs

@@ -1,7 +1,7 @@
 use kapi::Process;
 use obfstr::obfstr;
 use valthrun_driver_protocol::{
-    command::DriverCommandProcessMemoryWrite,
+    command::DriverCommandMemoryWrite,
     types::MemoryAccessResult,
 };
 use winapi::{
@@ -61,7 +61,7 @@ fn write_memory_mm(ctx: &WriteContext) -> bool {
     }
 }
 
-pub fn handler_write(command: &mut DriverCommandProcessMemoryWrite) -> anyhow::Result<()> {
+pub fn handler_write(command: &mut DriverCommandMemoryWrite) -> anyhow::Result<()> {
     let buffer = unsafe {
         if !seh::probe_read(command.buffer as u64, command.count, 0x01) {
             anyhow::bail!("{}", obfstr!("output buffer is not writeable"))

driver/src/handler/module.rs

@@ -1,4 +1,3 @@
-use alloc::vec;
 use core::{
     mem::size_of,
     str,
@@ -8,93 +7,88 @@ use kapi::Process;
 use obfstr::obfstr;
 use valthrun_driver_protocol::{
     command::{
+        DriverCommandProcessList,
         DriverCommandProcessModules,
-        ProcessModulesResult,
     },
     types::{
-        ProcessFilter,
+        ProcessId,
+        ProcessInfo,
         ProcessModuleInfo,
     },
 };
 
 use crate::util::kprocess;
 
-pub fn handler_get_modules(command: &mut DriverCommandProcessModules) -> anyhow::Result<()> {
-    let module_buffer = unsafe {
+pub fn handler_get_processes(command: &mut DriverCommandProcessList) -> anyhow::Result<()> {
+    let buffer = unsafe {
         if !seh::probe_write(
-            command.module_buffer as u64,
-            command.module_buffer_length * size_of::<ProcessModuleInfo>(),
+            command.buffer as u64,
+            command.buffer_capacity * size_of::<ProcessInfo>(),
             0x01,
         ) {
             anyhow::bail!("{}", obfstr!("response buffer not writeable"));
         }
 
-        core::slice::from_raw_parts_mut(command.module_buffer, command.module_buffer_length)
+        core::slice::from_raw_parts_mut(command.buffer, command.buffer_capacity)
     };
 
-    let process_candidates = match command.target_process {
-        ProcessFilter::None => {
-            command.result = ProcessModulesResult::ProcessUnknown;
-            return Ok(());
-        }
-        ProcessFilter::Id { id } => {
-            Process::by_id(id as i32)
-                .map(|p| vec![p])
-                .unwrap_or_default()
-        }
-        ProcessFilter::ImageBaseName { name, name_length } => {
-            let name = unsafe {
-                if !seh::probe_read(name as u64, name_length, 0x01) {
-                    anyhow::bail!("{}", obfstr!("name buffer not readable"));
-                }
+    command.process_count = 0;
+    kprocess::iter(|process| {
+        if let Some(output) = buffer.get_mut(command.process_count) {
+            output.process_id = process.get_id() as ProcessId;
+            output.directory_table_base = process.get_directory_table_base();
 
-                core::slice::from_raw_parts(name, name_length)
-            };
-            let target_name =
-                str::from_utf8(name).map_err(|_| anyhow::anyhow!("invalid name (not utf-8)"))?;
-            kprocess::find_processes_by_name(target_name)?
-        }
-    };
+            {
+                let image_base_name = process.get_image_file_name().unwrap_or_default();
+                let copy_length = image_base_name.len().min(output.image_base_name.len());
+                output.image_base_name[0..copy_length]
+                    .copy_from_slice(&image_base_name.as_bytes()[0..copy_length]);
 
-    let process = match process_candidates.len() {
-        0 => {
-            command.result = ProcessModulesResult::ProcessUnknown;
-            return Ok(());
+                if copy_length < output.image_base_name.len() - 1 {
+                    output.image_base_name[copy_length] = 0x00;
+                }
+            }
         }
-        1 => process_candidates.first().unwrap(),
-        _count => {
-            command.result = ProcessModulesResult::ProcessUbiquitous;
-            return Ok(());
+
+        command.process_count += 1;
+    });
+
+    Ok(())
+}
+
+pub fn handler_get_modules(command: &mut DriverCommandProcessModules) -> anyhow::Result<()> {
+    let buffer = unsafe {
+        if !seh::probe_write(
+            command.buffer as u64,
+            command.buffer_capacity * size_of::<ProcessModuleInfo>(),
+            0x01,
+        ) {
+            anyhow::bail!("{}", obfstr!("response buffer not writeable"));
         }
+
+        core::slice::from_raw_parts_mut(command.buffer, command.buffer_capacity)
     };
 
-    command.process_id = process.get_id() as u32;
-    log::trace!(
-        "Found process id {}. PEP at {:X}",
-        command.process_id,
-        process.eprocess() as u64
-    );
+    let Some(process) = Process::by_id(command.process_id as i32) else {
+        command.process_unknown = true;
+        return Ok(());
+    };
 
     let modules = {
         let attached_process = process.attach();
         attached_process.get_modules()
     };
 
-    command.module_count = modules.len();
-    if modules.len() > module_buffer.len() {
-        command.result = ProcessModulesResult::BufferTooSmall;
-        return Ok(());
-    }
-
-    for index in 0..modules.len() {
-        let output = &mut module_buffer[index];
-        let input = &modules[index];
+    command.process_unknown = false;
+    for module in &modules {
+        if let Some(output) = buffer.get_mut(command.module_count) {
+            output.base_dll_name.copy_from_slice(&module.base_dll_name);
+            output.base_address = module.base_address as u64;
+            output.module_size = module.module_size as u64;
+        }
 
-        output.base_dll_name.copy_from_slice(&input.base_dll_name);
-        output.base_address = input.base_address as u64;
-        output.module_size = input.module_size as u64;
+        command.module_count += 1;
     }
 
-    command.result = ProcessModulesResult::Success;
     Ok(())
 }

driver/src/util/kprocess.rs

@@ -4,6 +4,31 @@ use kapi::Process;
 
 use crate::offsets::get_nt_offsets;
 
+pub fn iter(mut consumer: impl FnMut(&Process)) {
+    #[allow(non_snake_case)]
+    let PsGetNextProcess = get_nt_offsets().PsGetNextProcess;
+
+    let mut current_peprocess = core::ptr::null_mut();
+    loop {
+        current_peprocess = unsafe { PsGetNextProcess(current_peprocess) };
+        if current_peprocess.is_null() {
+            break;
+        }
+
+        let process = Process::from_raw(current_peprocess, false);
+
+        // let active_threads = unsafe {
+        //     current_peprocess
+        //         /* The ActiveThreads comes after the thread list head. Thread list head has a size of 0x10. */
+        //         .byte_offset(EPROCESS_ThreadListHead as isize + 0x10)
+        //         .cast::<u32>()
+        //         .read_volatile()
+        // };
+
+        consumer(&process);
+    }
+}
+
 pub fn find_processes_by_name(target_name: &str) -> anyhow::Result<Vec<Process>> {
     #[allow(non_snake_case)]
     let PsGetNextProcess = get_nt_offsets().PsGetNextProcess;