code renew build ok

src/driver_state.rs

@@ -1,3 +1,5 @@
+use std::error::Error;
+// use std::io::{Error, ErrorKind};
 use std::ffi::c_void;
 use std::mem::{size_of_val};
 
@@ -5,19 +7,21 @@ use winapi::shared::ntdef::{NTSTATUS};
 use winapi::shared::minwindef::{DWORD};
 use winapi::um::winioctl::{
     CTL_CODE, FILE_ANY_ACCESS,
-    METHOD_IN_DIRECT, METHOD_OUT_DIRECT, METHOD_BUFFERED, METHOD_NEITHER
+    METHOD_IN_DIRECT, METHOD_OUT_DIRECT, /* METHOD_BUFFERED, */ METHOD_NEITHER
 };
 
 use crate::pdb_store::{PdbStore};
 use crate::windows::{WindowsFFI, WindowsVersion};
 use crate::ioctl_protocol::{
-    InputData, OffsetData, DerefAddr, ScanRange, HideProcess,
+    InputData, OffsetData, DerefAddr, ScanPoolData, /* HideProcess, */
-    OutputData, Nothing
+    /* OutputData, */ Nothing
 };
 
+type BoxResult<T> = Result<T, Box<dyn Error>>;
+
 const SIOCTL_TYPE: DWORD = 40000;
 
-fn to_epoch(filetime: u64) -> u64 {
+pub fn to_epoch(filetime: u64) -> u64 {
     let windows_epoch_diff: u64 = 11644473600000 * 10000;
     if filetime < windows_epoch_diff {
         return 0;
@@ -69,8 +73,9 @@ impl PartialEq for EprocessPoolChunk {
 
 #[allow(dead_code)]
 pub struct DriverState {
+    // TODO: Make private, only call methods of DriverState
     pub pdb_store: PdbStore,
-    windows_ffi: WindowsFFI,
+    pub windows_ffi: WindowsFFI,
 }
 
 impl DriverState {
@@ -79,27 +84,25 @@ impl DriverState {
         windows_ffi.print_version();
         Self {
             pdb_store,
-            windows_ffi,
+            windows_ffi
-            ntosbase: 0u64,
-            nonpaged_range: [0, 0],
-            eprocess_traverse_result: Vec::new(),
-            pool_scan_result: Vec::new()
         }
     }
 
     pub fn startup(&mut self) -> NTSTATUS {
-        self.windows_ffi.load_driver()
+        let s = self.windows_ffi.load_driver();
         let mut input = InputData {
             offset_value: OffsetData::new(&self.pdb_store, self.windows_ffi.short_version)
         };
-        self.windows_ffi.device_io(code, &mut input, &mut Nothing);
+        self.windows_ffi.device_io(DriverAction::SetupOffset.get_code(),
+                                   &mut input, &mut Nothing);
+        s
     }
 
     pub fn shutdown(&self) -> NTSTATUS {
         self.windows_ffi.unload_driver()
     }
 
-    pub fn get_kernel_base(&self) -> Result<u64, io::Error> {
+    pub fn get_kernel_base(&self) -> BoxResult<u64> {
         let mut ntosbase = 0u64;
         self.windows_ffi.device_io(DriverAction::GetKernelBase.get_code(),
                                    &mut Nothing, &mut ntosbase);
@@ -107,16 +110,16 @@ impl DriverState {
         Ok(ntosbase)
     }
 
-    pub fn scan_active_head(&self, ntosbase: u64) -> Result<Vec<EprocessPoolChunk>, io::Error> {
+    pub fn scan_active_head(&self, ntosbase: u64) -> BoxResult<Vec<EprocessPoolChunk>> {
-        let ps_active_head = ntosbase + self.pdb_store.get_offset("PsActiveProcessHead");
+        let ps_active_head = ntosbase + self.pdb_store.get_offset_r("PsActiveProcessHead")?;
-        let flink_offset = self.pdb_store.get_offset("_LIST_ENTRY.Flink");
+        let flink_offset = self.pdb_store.get_offset_r("_LIST_ENTRY.Flink")?;
-        let eprocess_link_offset = self.pdb_store.get_offset("_EPROCESS.ActiveProcessLinks");
+        let eprocess_link_offset = self.pdb_store.get_offset_r("_EPROCESS.ActiveProcessLinks")?;
-        let eprocess_name_offset = self.pdb_store.get_offset("_EPROCESS.ImageFileName");
+        let eprocess_name_offset = self.pdb_store.get_offset_r("_EPROCESS.ImageFileName")?;
 
         let mut ptr = ps_active_head;
         self.deref_addr(ptr + flink_offset, &mut ptr);
 
-        let mut result: Vec<EprocessPoolChunk>;
+        let mut result: Vec<EprocessPoolChunk> = Vec::new();
         while ptr != ps_active_head {
             let mut image_name = [0u8; 15];
             let eprocess = ptr - eprocess_link_offset;
@@ -141,23 +144,33 @@ impl DriverState {
         Ok(result)
     }
 
-    pub fn scan_pool(&self, ntosbase: u64, tag: [u8; 4],
+    pub fn scan_pool<F>(&self, ntosbase: u64, tag: [u8; 4], mut handler: F) -> BoxResult<bool>
-                     handler: FnMut(&DriverState, u64) -> Result<bool, io::Error>
+                        where F: FnMut(&DriverState, u64) -> BoxResult<bool>
-                    ) -> Result<bool, io::Error> {
+    {
-        let range = self.get_nonpaged_range(ntosbase);
+        let code = DriverAction::ScanPoolRemote.get_code();
+        let range = self.get_nonpaged_range(ntosbase)?;
         let start_address = range[0];
         let end_address = range[1];
         let mut ptr = start_address;
         while ptr < end_address {
             let mut input = InputData {
-                scan_range: ScanPoolData::new(&[ptr, end_address], tag)
+                scan_range: ScanPoolData::new(&[ptr, end_address], &tag)
             };
             self.windows_ffi.device_io(code, &mut input, &mut ptr);
             if ptr >= end_address {
                 break;
             }
-            handler(&self, ptr)?;
+            ptr += match handler(&self, ptr) {
-            ptr += pool_header_size;
+                Ok(success) => {
+                    if success {
+                    }
+                    else {
+                    }
+                },
+                Err(e) => println!("Handle error {:?}", e),
+                // found, ptr += chunk size
+                // ptr += pool_header_size;
+            };
         }
         Ok(true)
     }
@@ -188,27 +201,27 @@ impl DriverState {
                                        outptr as *mut c_void, output_len as DWORD);
     }
 
-    pub fn get_unicode_string(&self, unicode_str_addr: u64) -> Result<&str, io::Error> {
+    pub fn get_unicode_string(&self, unicode_str_addr: u64) -> BoxResult<String> {
-        let mut strlen: u16;
+        let mut strlen = 0u16;
-        let mut bufaddr : u64;
+        let mut bufaddr = 0u64;
-        let buffer_ptr = unicode_str_addr + self.pdb_store.get_offset("_UNICODE_STRING.Buffer")?;
+        let buffer_ptr = unicode_str_addr + self.pdb_store.get_offset_r("_UNICODE_STRING.Buffer")?;
 
-        self.defer_addr(unicode_str_addr, &mut strlen);
+        self.deref_addr(unicode_str_addr, &mut strlen);
-        self.defer_addr(buffer_ptr, &mut bufaddr);
+        self.deref_addr(buffer_ptr, &mut bufaddr);
 
         let mut buf = vec![0u8; strlen as usize];
-        dr.deref_addr_ptr(bufaddr, buf.as_mut_ptr(), strlen);
+        self.deref_addr_ptr(bufaddr, buf.as_mut_ptr(), strlen as u64);
 
-        prinln!("unicode string {?}", buf);
+        println!("unicode string {:?}", buf);
 
-        Ok(str::from_utf8(&buf)?)
+        Ok(std::str::from_utf8(&buf)?.to_string())
     }
 
-    pub fn get_nonpaged_range(&self, ntosbase: u64) -> Result<[u64; 2], io::Error> {
+    pub fn get_nonpaged_range(&self, ntosbase: u64) -> BoxResult<[u64; 2]> {
         // TODO: Add support for other Windows version here
         match self.windows_ffi.short_version {
             WindowsVersion::Windows10FastRing => {
-                let mistate = ntosbase + self.pdb_store.get_offset("MiState")?;
+                let mistate = ntosbase + self.pdb_store.get_offset_r("MiState")?;
                 let system_node_ptr = self.pdb_store.addr_decompose(
                                         mistate, "_MI_SYSTEM_INFORMATION.Hardware.SystemNodeNonPagedPool")?;
                 let mut system_node_addr = 0u64;
@@ -217,19 +230,19 @@ impl DriverState {
                 let mut first_va = 0u64;
                 let mut last_va = 0u64;
                 self.deref_addr(
-                    system_node_addr + self.pdb_store.get_offset(
+                    system_node_addr
-                        "_MI_SYSTEM_NODE_NONPAGED_POOL.NonPagedPoolFirstVa")?,
+                    + self.pdb_store.get_offset_r("_MI_SYSTEM_NODE_NONPAGED_POOL.NonPagedPoolFirstVa")?,
                     &mut first_va);
 
                 self.deref_addr(
-                    system_node_addr + self.pdb_store.get_offset(
+                    system_node_addr
-                        "_MI_SYSTEM_NODE_NONPAGED_POOL.NonPagedPoolLastVa")?,
+                    + self.pdb_store.get_offset_r("_MI_SYSTEM_NODE_NONPAGED_POOL.NonPagedPoolLastVa")?,
                     &mut last_va);
 
                 Ok([first_va, last_va])
             }
             _ => {
-                Err("Windows version for nonpaged pool algorithm is not implemented")
+                Err("Windows version for nonpaged pool algorithm is not implemented".into())
             }
         }
     }

src/ioctl_protocol.rs

@@ -75,7 +75,7 @@ impl ScanPoolData{
         Self {
             start: arr[0],
             end: arr[1],
-            tag: tag
+            tag: *tag
         }
     }
 }
@@ -91,7 +91,7 @@ pub struct HideProcess {
 pub union InputData {
     pub offset_value: OffsetData,
     pub deref_addr: DerefAddr,
-    pub scan_range: ScanRange,
+    pub scan_range: ScanPoolData,
     pub hide_process: HideProcess,
 }
 

src/main.rs

@@ -5,50 +5,53 @@ mod windows;
 mod ioctl_protocol;
 mod driver_state;
 
-use chrono::prelude::DateTime;
+use std::error::Error;
-use chrono::Utc;
+use std::str::{from_utf8};
-use chrono::{Local, DateTime};
+// use chrono::prelude::DateTime;
-use std::time::{SystemTime, UNIX_EPOCH, Duration};
+// use chrono::Utc;
+// use chrono::{Local, DateTime};
+// use std::time::{SystemTime, UNIX_EPOCH, Duration};
 
 use pdb_store::parse_pdb;
 use windows::WindowsFFI;
-use driver_state::{DriverState, DriverAction};
+use driver_state::{DriverState, EprocessPoolChunk, to_epoch};
 
-fn to_str_time(time_ms: u64) -> String {
+fn to_str_time(_time_ms: u64) -> String {
-    if time_ms == 0 {
+    // if time_ms == 0 {
-        return "".to_string();
+    //     return "".to_string();
-    }
+    // }
-    let d = UNIX_EPOCH + Duration::from_millis(time_ms);
+    // let d = UNIX_EPOCH + Duration::from_millis(time_ms);
-    let datetime = DateTime::<Utc>::from(d);
+    // let datetime = DateTime::<Utc>::from(d);
-    let timestamp_str = datetime.format("%Y-%m-%d %H:%M:%S.%f").to_string();
+    // let timestamp_str = datetime.format("%Y-%m-%d %H:%M:%S.%f").to_string();
-    timestamp_str
+    // timestamp_str
+    "".to_string()
 }
 
-fn main() -> Result<(), io::Error> {
+fn main() -> Result<(), Box<dyn Error>> {
     // for windows admin require
     // https://github.com/nabijaczleweli/rust-embed-resource
 
     let mut driver = DriverState::new(parse_pdb(), WindowsFFI::new());
     println!("NtLoadDriver()   -> 0x{:x}", driver.startup());
 
-    let ntosbase = driver.get_ntosbase()?;
+    let ntosbase = driver.get_kernel_base()?;
-    let pool_header_size = driver.pdb_store.get_offset("_POOL_HEADER.struct_size")?;
+    let pool_header_size = driver.pdb_store.get_offset_r("_POOL_HEADER.struct_size")?;
 
     let eprocess_tag: [u8; 4] = [80, 114, 111, 99]; // Proc
-    let eprocess_name_offset = driver.pdb_store.get_offset("_EPROCESS.ImageFileName")?;
+    let eprocess_name_offset = driver.pdb_store.get_offset_r("_EPROCESS.ImageFileName")?;
-    let eprocess_create_time_offset = driver.pdb_store.get_offset("_EPROCESS.CreateTime")?;
+    let eprocess_create_time_offset = driver.pdb_store.get_offset_r("_EPROCESS.CreateTime")?;
-    let eprocess_exit_time_offset = driver.pdb_store.get_offset("_EPROCESS.ExitTime")?;
+    let eprocess_exit_time_offset = driver.pdb_store.get_offset_r("_EPROCESS.ExitTime")?;
-    let eprocess_size = driver.pdb_store.get_offset("_EPROCESS.struct_size")?;
+    let eprocess_size = driver.pdb_store.get_offset_r("_EPROCESS.struct_size")?;
 
     let eprocess_scan_head = driver.scan_active_head(ntosbase)?;
-    let mut eprocess_list: Vec<EprocessPoolChunk>;
+    let mut eprocess_list: Vec<EprocessPoolChunk> = Vec::new();
     driver.scan_pool(ntosbase, eprocess_tag, |dr, pool_addr| {
         let mut pool = vec![0u8; pool_header_size as usize];
         dr.deref_addr_ptr(pool_addr, pool.as_mut_ptr(), pool_header_size);
 
-        let pool_size = (pool[2] as u64) * 16u64;
+        let chunk_size = (pool[2] as u64) * 16u64;
         let eprocess_valid_start = pool_addr + pool_header_size;
-        let eprocess_valid_end = pool_addr + pool_size - eprocess_size;
+        let eprocess_valid_end = pool_addr + chunk_size - eprocess_size;
         let mut try_eprocess_ptr = eprocess_valid_start;
 
         let mut create_time = 0u64;
@@ -63,7 +66,7 @@ fn main() -> Result<(), io::Error> {
         }
         let mut image_name = [0u8; 15];
         dr.deref_addr(try_eprocess_ptr + eprocess_name_offset, &mut image_name);
-        let eprocess_name = std::str::from_utf8(&image_name)?
+        let eprocess_name = from_utf8(&image_name)?
                             .to_string()
                             .trim_end_matches(char::from(0))
                             .to_string();
@@ -75,22 +78,22 @@ fn main() -> Result<(), io::Error> {
             exit_time: to_epoch(exit_time)
         });
         Ok(try_eprocess_ptr <= eprocess_valid_end)
-    });
+    })?;
 
     let ethread_tag: [u8; 4] = [84, 104, 114, 101]; // Thre
-    let ethread_create_time_offset = driver.pdb_store.get_offset("_ETHREAD.CreateTime")?;
+    let ethread_create_time_offset = driver.pdb_store.get_offset_r("_ETHREAD.CreateTime")?;
-    let ethread_exit_time_offset = driver.pdb_store.get_offset("_ETHREAD.ExitTime")?;
+    let ethread_exit_time_offset = driver.pdb_store.get_offset_r("_ETHREAD.ExitTime")?;
-    let ethread_threadname_offset = driver.pdb_store.get_offset("_ETHREAD.TheadName")?;
+    let ethread_threadname_offset = driver.pdb_store.get_offset_r("_ETHREAD.TheadName")?;
-    let ethread_size = driver.pdb_store.get_offset("_ETHREAD.struct_size")?;
+    let ethread_size = driver.pdb_store.get_offset_r("_ETHREAD.struct_size")?;
 
-    // let mut ethread_list: Vec<EprocessPoolChunk>;
+    // let mut ethread_list: Vec<EprocessPoolChunk> = Vec::new();
     driver.scan_pool(ntosbase, ethread_tag, |dr, pool_addr| {
         let mut pool = vec![0u8; pool_header_size as usize];
         dr.deref_addr_ptr(pool_addr, pool.as_mut_ptr(), pool_header_size);
 
-        let pool_size = (pool[2] as u64) * 16u64;
+        let chunk_size = (pool[2] as u64) * 16u64;
         let ethread_valid_start = pool_addr + pool_header_size;
-        let ethread_valid_end = pool_addr + pool_size - ethread_size;
+        let ethread_valid_end = pool_addr + chunk_size - ethread_size;
         let mut try_ethread_ptr = ethread_valid_start;
 
         let mut create_time = 0u64;
@@ -105,10 +108,10 @@ fn main() -> Result<(), io::Error> {
         }
         let mut threadname_ptr = 0u64;
         dr.deref_addr(try_ethread_ptr + ethread_threadname_offset, &mut threadname_ptr);
-        let threadname = dr.get_unicode_string(threadname_ptr)?
+        let threadname = dr.get_unicode_string(threadname_ptr)?;
         println!("threadname: {}", threadname);
         Ok(try_ethread_ptr <= ethread_valid_end)
-    });
+    })?;
 
     // for result in &driver.eprocess_traverse_result {
     //     println!("- [{}] 0x{:x} {}",

src/pdb_store.rs

@@ -1,3 +1,4 @@
+use std::error::Error;
 use std::io;
 use std::io::{Read};
 use std::path::Path;
@@ -29,6 +30,10 @@ pub struct PdbStore {
 }
 
 impl PdbStore {
+    pub fn get_offset_r(&self, name: &str) -> Result<u64, Box<dyn Error>> {
+        self.get_offset(name)
+            .ok_or(format!("{} is not found in PDB", name).into())
+    }
     #[allow(dead_code)]
     pub fn get_offset(&self, name: &str) -> Option<u64> {
         if name.contains(".") {
@@ -52,9 +57,9 @@ impl PdbStore {
     }
 
     #[allow(dead_code)]
-    pub fn addr_decompose(&self, addr: u64, full_name: &str) -> Result<u64, String>{
+    pub fn addr_decompose(&self, addr: u64, full_name: &str) -> Result<u64, Box<dyn Error>>{
         if !full_name.contains(".") {
-            return Err("Not decomposable".to_string());
+            return Err("Not decomposable".into());
         }
 
         let mut name_part: Vec<&str> = full_name.split_terminator('.').collect();
@@ -65,7 +70,7 @@ impl PdbStore {
                     Some((memtype, offset)) => {
                         if next.len() != 0 {
                             if memtype.contains("*") {
-                                return Err(format!("Cannot dereference pointer at {} {}", memtype, name_part[1]));
+                                return Err(format!("Cannot dereference pointer at {} {}", memtype, name_part[1]).into());
                             }
                             next.insert(0, memtype);
                             self.addr_decompose(addr + *offset, &next.join("."))
@@ -74,10 +79,10 @@ impl PdbStore {
                             Ok(addr + *offset)
                         }
                     },
-                    None => Err(format!("Not found member {}", name_part[1]))
+                    None => Err(format!("Not found member {}", name_part[1]).into())
                 }
             },
-            None => Err(format!("Struct {} not found", name_part[0]))
+            None => Err(format!("Struct {} not found", name_part[0]).into())
         }
     }