Turn dummy handler into proxy for winapi calls

Cargo.lock

@@ -40,6 +40,7 @@ dependencies = [
  "log",
  "pretty_env_logger",
  "tokio",
+ "winapi",
 ]
 
 [[package]]

Cargo.toml

@@ -11,3 +11,9 @@ bytes = "0.6"
 log = "0.4.11"
 pretty_env_logger = "0.4.0"
 tokio = {version = "0.3.3", features = ["net", "rt-multi-thread", "io-util"]}
+winapi = {version = "0.3.9", features = ["tlhelp32"], optional = true}
+
+
+[features]
+default = ["windows"]
+windows = ["winapi"]

src/main.rs

@@ -8,7 +8,7 @@ fn main() {
 
     let rt = tokio::runtime::Runtime::new().unwrap();
     rt.block_on(async {
-        server::run::<server::TestHandler>().await.unwrap();
+        server::run::<server::WindowsHandler>().await.unwrap();
     })
 }
 

src/server/ce_common.rs

@@ -0,0 +1,51 @@
+use bytes::{Buf, BufMut};
+
+#[derive(Debug)]
+pub struct CeProcessEntry {
+    pub pid: i32,
+    pub process_name: String,
+}
+
+#[derive(Debug)]
+pub struct CeModuleEntry {
+    pub module_base: i64,
+    pub module_size: i32,
+    pub module_name: String,
+}
+
+// NOTE: all handles seem to be 32-bit
+pub fn read_usize(buf: &mut dyn Buf) -> usize {
+    match std::mem::size_of::<usize>() {
+        4 | 8 => buf.get_u32_le() as usize,
+        _ => unreachable!(),
+    }
+}
+
+pub fn write_usize(buf: &mut dyn BufMut, value: usize) {
+    match std::mem::size_of::<usize>() {
+        4 | 8 => buf.put_u32_le(value as u32),
+        _ => unreachable!(),
+    }
+}
+
+pub fn write_i32_prefixed_string(buf: &mut dyn BufMut, value: String) {
+    let bytes = value.into_bytes();
+
+    buf.put_i32_le(bytes.len() as i32);
+    buf.put_slice(&bytes[..]);
+}
+
+pub fn cstring_to_string(with_zeroes: &[u8]) -> String {
+    let len = with_zeroes.iter().position(|i| *i == 0).unwrap_or(with_zeroes.len());
+    let without_zeroes = &with_zeroes[0..len];
+    String::from_utf8(without_zeroes.to_vec()).unwrap()
+}
+
+#[test]
+fn test_cstring_to_string() {
+    assert_eq!(String::from(""), cstring_to_string(b""));
+    assert_eq!(String::from(""), cstring_to_string(b"\0"));
+    assert_eq!(String::from("abc"), cstring_to_string(b"abc"));
+    assert_eq!(String::from("abc"), cstring_to_string(b"abc\0"));
+    assert_eq!(String::from("abc"), cstring_to_string(b"abc\0\0abc"));
+}

src/server/commands_request.rs

@@ -2,8 +2,7 @@ use std::fmt::Debug;
 
 use bytes::Buf;
 
-use super::{command::*, commands_response::*};
-use super::command::Command;
+use super::{command::*, commands_response::*, ce_common::*};
 
 pub trait CERequest : Debug + Send {
     type Response: CEResponse + Debug + Send;
@@ -19,7 +18,7 @@ pub struct CreateToolHelp32SnapshotRequest {
 }
 
 impl CERequest for CreateToolHelp32SnapshotRequest {
-    type Response = CreateToolHelp32SnapshotResponse;
+    type Response = HandleResponse;
 
     const ID: Command = CMD_CREATETOOLHELP32SNAPSHOT;
 
@@ -33,17 +32,17 @@ impl CERequest for CreateToolHelp32SnapshotRequest {
 
 #[derive(Debug)]
 pub struct Process32FirstRequest {
-    pub handle: u32,
+    pub handle: usize,
 }
 
 impl CERequest for Process32FirstRequest {
-    type Response = Process32FirstResponse;
+    type Response = Process32Response;
 
     const ID: Command = CMD_PROCESS32FIRST;
 
     fn read(buf: &mut dyn Buf) -> Self {
         Self {
-            handle: buf.get_u32_le(),
+            handle: read_usize(buf),
         }
     }
 }

src/server/commands_response.rs

@@ -1,35 +1,38 @@
 use bytes::BufMut;
 
+use super::ce_common::*;
+
 pub trait CEResponse {
     fn serialize(self, writer: &mut dyn BufMut);
 }
 
 #[derive(Debug)]
-pub struct CreateToolHelp32SnapshotResponse {
-    pub handle: u32,
+pub struct HandleResponse {
+    pub handle: usize,
 }
 
-impl CEResponse for CreateToolHelp32SnapshotResponse {
+impl CEResponse for HandleResponse {
     fn serialize(self, writer: &mut dyn BufMut) {
-        writer.put_u32_le(self.handle);
+        write_usize(writer, self.handle);
     }
 }
 
 #[derive(Debug)]
-pub struct Process32FirstResponse {
-    pub result: bool,
-    // TODO: process entries!!
+pub struct Process32Response {
+    pub entry: Option<CeProcessEntry>,
 }
 
-impl CEResponse for Process32FirstResponse {
+impl CEResponse for Process32Response {
     fn serialize(self, writer: &mut dyn BufMut) {
-        writer.put_u32_le(self.result as u32);
-
-        if self.result {
-            todo!()
-        } else {
-            writer.put_u32_le(0);
-            writer.put_u32_le(0);
+        if let Some(entry) = self.entry {
+            writer.put_u32_le(1u32);
+            writer.put_i32(entry.pid);
+            write_i32_prefixed_string(writer, entry.process_name)
+        }
+        else {
+            writer.put_u32_le(0u32);
+            writer.put_i32_le(0); // pid
+            writer.put_i32_le(0); // process name length
         }
     }
 }

src/server/handler.rs

@@ -1,6 +1,4 @@
-use log::warn;
-
-use super::{commands_response::*, commands_request::*};
+use super::commands_request::*;
 
 pub trait HandlerFactory<T: FullHandler> {
 }
@@ -14,27 +12,3 @@ pub trait FullHandler
 pub trait Handler<TReq: CERequest> {
     fn handle(&self, req: TReq) -> TReq::Response;
 }
-
-pub struct TestHandler;
-
-impl FullHandler for TestHandler {
-    fn create() -> TestHandler { TestHandler }
-}
-
-impl Handler<CreateToolHelp32SnapshotRequest> for TestHandler {
-    fn handle(&self, _req: CreateToolHelp32SnapshotRequest) -> CreateToolHelp32SnapshotResponse {
-        warn!("Stubbed CreateToolHelp32SnapshotResponse::process"); // TODO
-        CreateToolHelp32SnapshotResponse {
-            handle: 0
-        }
-    }
-}
-
-impl Handler<Process32FirstRequest> for TestHandler {
-    fn handle(&self, _req: Process32FirstRequest) -> Process32FirstResponse {
-        warn!("Stubbed Process32FirstResponse::process"); // TODO
-        Process32FirstResponse {
-            result: false,
-        }
-    }
-}

src/server/handlers/mod.rs

@@ -0,0 +1,4 @@
+mod windows;
+
+#[cfg(windows)]
+pub use windows::WindowsHandler;

src/server/handlers/windows.rs

@@ -0,0 +1,44 @@
+use crate::server::{ce_common::CeProcessEntry, commands_request::*, commands_response::*, ce_common::cstring_to_string, handler::*};
+use winapi::um::tlhelp32::{CreateToolhelp32Snapshot, PROCESSENTRY32, Process32First};
+
+pub struct WindowsHandler;
+
+impl FullHandler for WindowsHandler {
+    fn create() -> WindowsHandler { WindowsHandler }
+}
+
+impl Handler<CreateToolHelp32SnapshotRequest> for WindowsHandler {
+    fn handle(&self, req: CreateToolHelp32SnapshotRequest) -> HandleResponse {
+        unsafe {
+            let ret = CreateToolhelp32Snapshot(req.snapshot_flags, req.process_id);
+
+            HandleResponse {
+                handle: ret as usize
+            }
+        }
+    }
+}
+
+impl Handler<Process32FirstRequest> for WindowsHandler {
+    fn handle(&self, req: Process32FirstRequest) -> Process32Response {
+        unsafe {
+            let mut entry = std::mem::MaybeUninit::uninit().assume_init();
+            let response = Process32First(std::mem::transmute(req.handle), &mut entry);
+
+            if response != 0 {
+                Process32Response {
+                    entry: Some(
+                        CeProcessEntry {
+                            pid: entry.th32ProcessID as i32,
+                            process_name: cstring_to_string(std::mem::transmute(&entry.szExeFile[..])),
+                        }
+                    )
+                }
+            } else {
+                Process32Response {
+                    entry: None
+                }
+            }
+        }
+    }
+}

src/server/mod.rs

@@ -3,6 +3,8 @@ mod command;
 mod commands_request;
 mod commands_response;
 mod handler;
+mod handlers;
+mod ce_common;
 
 pub use server::run;
-pub use handler::TestHandler;
+pub use handlers::WindowsHandler;