Implement process and module enumeration

src/server/commands_request.rs

@@ -46,3 +46,71 @@ impl CERequest for Process32FirstRequest {
         }
     }
 }
+
+#[derive(Debug)]
+pub struct Process32NextRequest {
+    pub handle: usize,
+}
+
+impl CERequest for Process32NextRequest {
+    type Response = Process32Response;
+
+    const ID: Command = CMD_PROCESS32NEXT;
+
+    fn read(buf: &mut dyn Buf) -> Self {
+        Self {
+            handle: read_usize(buf),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct CloseHandleRequest {
+    pub handle: usize,
+}
+
+impl CERequest for CloseHandleRequest {
+    type Response = I32Response;
+
+    const ID: Command = CMD_CLOSEHANDLE;
+
+    fn read(buf: &mut dyn Buf) -> Self {
+        Self {
+            handle: read_usize(buf),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct Module32FirstRequest {
+    pub handle: usize,
+}
+
+impl CERequest for Module32FirstRequest {
+    type Response = Module32Response;
+
+    const ID: Command = CMD_MODULE32FIRST;
+
+    fn read(buf: &mut dyn Buf) -> Self {
+        Self {
+            handle: read_usize(buf),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct Module32NextRequest {
+    pub handle: usize,
+}
+
+impl CERequest for Module32NextRequest {
+    type Response = Module32Response;
+
+    const ID: Command = CMD_MODULE32NEXT;
+
+    fn read(buf: &mut dyn Buf) -> Self {
+        Self {
+            handle: read_usize(buf),
+        }
+    }
+}

src/server/commands_response.rs

@@ -26,7 +26,7 @@ impl CEResponse for Process32Response {
     fn serialize(self, writer: &mut dyn BufMut) {
         if let Some(entry) = self.entry {
             writer.put_u32_le(1u32);
-            writer.put_i32(entry.pid);
+            writer.put_i32_le(entry.pid);
             write_i32_prefixed_string(writer, entry.process_name)
         }
         else {
@@ -36,3 +36,36 @@ impl CEResponse for Process32Response {
         }
     }
 }
+
+#[derive(Debug)]
+pub struct Module32Response {
+    pub entry: Option<CeModuleEntry>,
+}
+
+impl CEResponse for Module32Response {
+    fn serialize(self, writer: &mut dyn BufMut) {
+        if let Some(entry) = self.entry {
+            writer.put_u32_le(1u32);
+            writer.put_i64_le(entry.module_base);
+            writer.put_i32_le(entry.module_size);
+            write_i32_prefixed_string(writer, entry.module_name)
+        }
+        else {
+            writer.put_u32_le(0u32);
+            writer.put_i64_le(0); // base
+            writer.put_i32_le(0); // size
+            writer.put_i32_le(0); // module name length
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct I32Response {
+    pub response: i32,
+}
+
+impl CEResponse for I32Response {
+    fn serialize(self, writer: &mut dyn BufMut) {
+        writer.put_i32_le(self.response);
+    }
+}

src/server/handler.rs

@@ -5,7 +5,11 @@ pub trait HandlerFactory<T: FullHandler> {
 
 pub trait FullHandler
     : Handler<CreateToolHelp32SnapshotRequest>
-    + Handler<Process32FirstRequest> {
+    + Handler<Process32FirstRequest>
+    + Handler<Process32NextRequest>
+    + Handler<Module32FirstRequest>
+    + Handler<Module32NextRequest>
+    + Handler<CloseHandleRequest> {
     fn create() -> Self;
 }
 

src/server/handlers/windows.rs

@@ -1,5 +1,13 @@
-use crate::server::{ce_common::CeProcessEntry, commands_request::*, commands_response::*, ce_common::cstring_to_string, handler::*};
-use winapi::um::tlhelp32::{CreateToolhelp32Snapshot, PROCESSENTRY32, Process32First};
+use crate::server::{ce_common::*, commands_request::*, commands_response::*, handler::*};
+use winapi::um::{
+    handleapi::CloseHandle,
+    tlhelp32::{
+        CreateToolhelp32Snapshot,
+        Process32First, Process32Next, Module32First, Module32Next,
+        LPPROCESSENTRY32, LPMODULEENTRY32
+    },
+    winnt::HANDLE
+};
 
 pub struct WindowsHandler;
 
@@ -22,22 +30,85 @@ impl Handler<CreateToolHelp32SnapshotRequest> for WindowsHandler {
 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);
+            get_process_response(req.handle, |x, y| Process32First(x, y))
+        }
+    }
+}
 
-            if response != 0 {
-                Process32Response {
-                    entry: Some(
-                        CeProcessEntry {
-                            pid: entry.th32ProcessID as i32,
-                            process_name: cstring_to_string(std::mem::transmute(&entry.szExeFile[..])),
-                        }
-                    )
+impl Handler<Process32NextRequest> for WindowsHandler {
+    fn handle(&self, req: Process32NextRequest) -> Process32Response {
+        unsafe {
+            get_process_response(req.handle, |x, y| Process32Next(x, y))
+        }
+    }
+}
+
+unsafe fn get_process_response<F>(handle: usize, func: F) -> Process32Response
+    where F : FnOnce(HANDLE, LPPROCESSENTRY32) -> i32 {
+    let mut entry = std::mem::MaybeUninit::uninit().assume_init();
+    let response = func(std::mem::transmute(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
+            )
+        }
+    } else {
+        Process32Response {
+            entry: None
+        }
+    }
+}
+
+impl Handler<Module32FirstRequest> for WindowsHandler {
+    fn handle(&self, req: Module32FirstRequest) -> Module32Response {
+        unsafe {
+            get_module_response(req.handle, |x, y| Module32First(x, y))
+        }
+    }
+}
+
+impl Handler<Module32NextRequest> for WindowsHandler {
+    fn handle(&self, req: Module32NextRequest) -> Module32Response {
+        unsafe {
+            get_module_response(req.handle, |x, y| Module32Next(x, y))
+        }
+    }
+}
+
+unsafe fn get_module_response<F>(handle: usize, func: F) -> Module32Response
+    where F : FnOnce(HANDLE, LPMODULEENTRY32) -> i32 {
+    let mut entry = std::mem::MaybeUninit::uninit().assume_init();
+    let response = func(std::mem::transmute(handle), &mut entry);
+
+    if response != 0 {
+        Module32Response {
+            entry: Some(
+                CeModuleEntry {
+                    module_base: entry.modBaseAddr as i64,
+                    module_size: entry.modBaseSize as i32,
+                    module_name: cstring_to_string(std::mem::transmute(&entry.szModule[..]))
                 }
+            )
+        }
+    } else {
+        Module32Response {
+            entry: None
+        }
+    }
+}
+
+impl Handler<CloseHandleRequest> for WindowsHandler {
+    fn handle(&self, req: CloseHandleRequest) -> I32Response {
+        unsafe {
+            let response = CloseHandle(std::mem::transmute(req.handle));
+
+            I32Response {
+                response
             }
         }
     }

src/server/server.rs

@@ -26,12 +26,15 @@ pub struct CheatEngineConnection(TcpStream);
 
 // generics werent powerful enough
 macro_rules! gen_request_dispatch {
-    ($bytes: ident for $handler: ident, $($request: ty),*) => {
+    ($bytes: ident for $handler: ident, $($request: ty,)*) => {
         match $bytes.get_u8() {
             $(
                 <$request>::ID => gen_request_dispatch!($request, $bytes, $handler),
             )*
-            byte => todo!("Unimplemented packet byte {}", byte),
+            byte => {
+                error!("Unimplemented packet byte {} (data: {:?})", byte, $bytes);
+                todo!("Unimplemented packet byte {}", byte);
+            },
         }
     };
     ($request: ty, $bytes: ident, $handler: ident) => {
@@ -83,7 +86,11 @@ impl CheatEngineConnection {
         gen_request_dispatch!(
             bytes for handler,
             CreateToolHelp32SnapshotRequest,
-            Process32FirstRequest
+            Process32FirstRequest,
+            Process32NextRequest,
+            Module32FirstRequest,
+            Module32NextRequest,
+            CloseHandleRequest,
         )
     }
 }