Add process module to talpid-windows

8 files changed, 174 insertions, 203 deletions

diff --git a/Cargo.lock b/Cargo.lock
index a60c970639..50f68e0cd9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1863,6 +1863,7 @@ dependencies = [
  "talpid-platform-metadata",
  "talpid-time",
  "talpid-types",
+ "talpid-windows",
  "tokio",
  "tokio-stream",
  "winapi",
diff --git a/mullvad-daemon/Cargo.toml b/mullvad-daemon/Cargo.toml
index cc89500abd..352db7d9f9 100644
--- a/mullvad-daemon/Cargo.toml
+++ b/mullvad-daemon/Cargo.toml
@@ -60,6 +60,7 @@ ctrlc = "3.0"
 windows-service = "0.6.0"
 winapi = { version = "0.3", features = ["winnt", "excpt"] }
 dirs = "5.0.1"
+talpid-windows = { path = "../talpid-windows" }
 
 [target.'cfg(windows)'.dependencies.windows-sys]
 workspace = true
diff --git a/mullvad-daemon/src/exception_logging/win.rs b/mullvad-daemon/src/exception_logging/win.rs
index 906a8533fc..85f16019b0 100644
--- a/mullvad-daemon/src/exception_logging/win.rs
+++ b/mullvad-daemon/src/exception_logging/win.rs
@@ -1,27 +1,25 @@
 use mullvad_paths::log_dir;
 use std::{
     borrow::Cow,
-    ffi::{c_char, c_void, CStr},
+    ffi::c_void,
     fmt::Write,
-    fs, io, mem,
+    fs, io,
     os::windows::io::AsRawHandle,
     path::{Path, PathBuf},
     ptr,
 };
 use talpid_types::ErrorExt;
+use talpid_windows::process::{ModuleEntry, ProcessSnapshot};
 use winapi::{
     um::winnt::{CONTEXT_CONTROL, CONTEXT_INTEGER, CONTEXT_SEGMENTS},
     vc::excpt::EXCEPTION_EXECUTE_HANDLER,
 };
 use windows_sys::Win32::{
-    Foundation::{CloseHandle, BOOL, ERROR_NO_MORE_FILES, HANDLE, INVALID_HANDLE_VALUE},
+    Foundation::{BOOL, HANDLE},
     System::{
         Diagnostics::{
             Debug::{SetUnhandledExceptionFilter, CONTEXT, EXCEPTION_POINTERS, EXCEPTION_RECORD},
-            ToolHelp::{
-                CreateToolhelp32Snapshot, Module32First, Module32Next, MODULEENTRY32,
-                TH32CS_SNAPMODULE,
-            },
+            ToolHelp::TH32CS_SNAPMODULE,
         },
         Threading::{GetCurrentProcess, GetCurrentProcessId, GetCurrentThreadId},
     },
@@ -291,7 +289,7 @@ fn get_context_info(context: &CONTEXT) -> String {
 }
 
 /// Return module info for the current process and given memory address.
-fn find_address_module(address: *mut c_void) -> io::Result<Option<ModuleInfo>> {
+fn find_address_module(address: *mut c_void) -> io::Result<Option<ModuleEntry>> {
     let snap = ProcessSnapshot::new(TH32CS_SNAPMODULE, 0)?;
 
     for module in snap.modules() {
@@ -306,85 +304,3 @@ fn find_address_module(address: *mut c_void) -> io::Result<Option<ModuleInfo>> {
 
     Ok(None)
 }
-
-struct ModuleInfo {
-    name: String,
-    base_address: *const u8,
-    size: usize,
-}
-
-struct ProcessSnapshot {
-    handle: HANDLE,
-}
-
-impl ProcessSnapshot {
-    fn new(flags: u32, process_id: u32) -> io::Result<ProcessSnapshot> {
-        let snap = unsafe { CreateToolhelp32Snapshot(flags, process_id) };
-
-        if snap == INVALID_HANDLE_VALUE {
-            Err(io::Error::last_os_error())
-        } else {
-            Ok(ProcessSnapshot { handle: snap })
-        }
-    }
-
-    fn handle(&self) -> HANDLE {
-        self.handle
-    }
-
-    fn modules(&self) -> ProcessSnapshotModules<'_> {
-        let mut entry: MODULEENTRY32 = unsafe { mem::zeroed() };
-        entry.dwSize = mem::size_of::<MODULEENTRY32>() as u32;
-
-        ProcessSnapshotModules {
-            snapshot: self,
-            iter_started: false,
-            temp_entry: entry,
-        }
-    }
-}
-
-impl Drop for ProcessSnapshot {
-    fn drop(&mut self) {
-        unsafe {
-            CloseHandle(self.handle);
-        }
-    }
-}
-
-struct ProcessSnapshotModules<'a> {
-    snapshot: &'a ProcessSnapshot,
-    iter_started: bool,
-    temp_entry: MODULEENTRY32,
-}
-
-impl Iterator for ProcessSnapshotModules<'_> {
-    type Item = io::Result<ModuleInfo>;
-
-    fn next(&mut self) -> Option<io::Result<ModuleInfo>> {
-        if self.iter_started {
-            if unsafe { Module32Next(self.snapshot.handle(), &mut self.temp_entry) } == 0 {
-                let last_error = io::Error::last_os_error();
-
-                return if last_error.raw_os_error().unwrap() as u32 == ERROR_NO_MORE_FILES {
-                    None
-                } else {
-                    Some(Err(last_error))
-                };
-            }
-        } else {
-            if unsafe { Module32First(self.snapshot.handle(), &mut self.temp_entry) } == 0 {
-                return Some(Err(io::Error::last_os_error()));
-            }
-            self.iter_started = true;
-        }
-
-        let cstr_ref = &self.temp_entry.szModule[0];
-        let cstr = unsafe { CStr::from_ptr(cstr_ref as *const u8 as *const c_char) };
-        Some(Ok(ModuleInfo {
-            name: cstr.to_string_lossy().into_owned(),
-            base_address: self.temp_entry.modBaseAddr,
-            size: self.temp_entry.modBaseSize as usize,
-        }))
-    }
-}
diff --git a/talpid-core/src/split_tunnel/windows/driver.rs b/talpid-core/src/split_tunnel/windows/driver.rs
index eea79f188f..3f5a4e6a6d 100644
--- a/talpid-core/src/split_tunnel/windows/driver.rs
+++ b/talpid-core/src/split_tunnel/windows/driver.rs
@@ -1,6 +1,6 @@
 use super::windows::{
     get_device_path, get_process_creation_time, get_process_device_path, open_process,
-    ProcessAccess, ProcessSnapshot,
+    ProcessAccess,
 };
 use bitflags::bitflags;
 use memoffset::offset_of;
@@ -22,7 +22,7 @@ use std::{
     time::Duration,
 };
 use talpid_types::ErrorExt;
-use talpid_windows::{io::Overlapped, sync::Event};
+use talpid_windows::{io::Overlapped, process::ProcessSnapshot, sync::Event};
 use windows_sys::Win32::{
     Foundation::{
         ERROR_ACCESS_DENIED, ERROR_FILE_NOT_FOUND, ERROR_INVALID_PARAMETER, ERROR_IO_PENDING,
@@ -486,7 +486,7 @@ fn build_process_tree() -> io::Result<Vec<ProcessInfo>> {
     let mut process_info = HashMap::new();
 
     let snap = ProcessSnapshot::new(TH32CS_SNAPPROCESS, 0)?;
-    for entry in snap.entries() {
+    for entry in snap.processes() {
         let entry = entry?;
 
         let process = match open_process(ProcessAccess::QueryLimitedInformation, false, entry.pid) {
diff --git a/talpid-core/src/split_tunnel/windows/windows.rs b/talpid-core/src/split_tunnel/windows/windows.rs
index 423426c385..c20a837089 100644
--- a/talpid-core/src/split_tunnel/windows/windows.rs
+++ b/talpid-core/src/split_tunnel/windows/windows.rs
@@ -1,6 +1,3 @@
-// TODO: The snapshot code could be combined with the mostly-identical code in
-//       the windows_exception_logging module.
-
 use std::{
     ffi::{OsStr, OsString},
     fs, io, iter, mem,
@@ -12,99 +9,15 @@ use std::{
     ptr,
 };
 use windows_sys::Win32::{
-    Foundation::{
-        CloseHandle, ERROR_INSUFFICIENT_BUFFER, ERROR_NO_MORE_FILES, FILETIME, HANDLE,
-        INVALID_HANDLE_VALUE,
-    },
+    Foundation::{CloseHandle, ERROR_INSUFFICIENT_BUFFER, FILETIME, HANDLE},
     Storage::FileSystem::{GetFinalPathNameByHandleW, QueryDosDeviceW},
     System::{
-        Diagnostics::ToolHelp::{
-            CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W,
-        },
         ProcessStatus::GetProcessImageFileNameW,
         Threading::{GetProcessTimes, OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION},
         WindowsProgramming::VOLUME_NAME_NT,
     },
 };
 
-pub struct ProcessSnapshot {
-    handle: HANDLE,
-}
-
-impl ProcessSnapshot {
-    pub fn new(flags: u32, process_id: u32) -> io::Result<ProcessSnapshot> {
-        let snap = unsafe { CreateToolhelp32Snapshot(flags, process_id) };
-
-        if snap == INVALID_HANDLE_VALUE {
-            Err(io::Error::last_os_error())
-        } else {
-            Ok(ProcessSnapshot { handle: snap })
-        }
-    }
-
-    pub fn handle(&self) -> HANDLE {
-        self.handle
-    }
-
-    pub fn entries(&self) -> ProcessSnapshotEntries<'_> {
-        let mut entry: PROCESSENTRY32W = unsafe { mem::zeroed() };
-        entry.dwSize = mem::size_of::<PROCESSENTRY32W>() as u32;
-
-        ProcessSnapshotEntries {
-            snapshot: self,
-            iter_started: false,
-            temp_entry: entry,
-        }
-    }
-}
-
-impl Drop for ProcessSnapshot {
-    fn drop(&mut self) {
-        unsafe {
-            CloseHandle(self.handle);
-        }
-    }
-}
-
-pub struct ProcessEntry {
-    pub pid: u32,
-    pub parent_pid: u32,
-}
-
-pub struct ProcessSnapshotEntries<'a> {
-    snapshot: &'a ProcessSnapshot,
-    iter_started: bool,
-    temp_entry: PROCESSENTRY32W,
-}
-
-impl Iterator for ProcessSnapshotEntries<'_> {
-    type Item = io::Result<ProcessEntry>;
-
-    fn next(&mut self) -> Option<io::Result<ProcessEntry>> {
-        if self.iter_started {
-            if unsafe { Process32NextW(self.snapshot.handle(), &mut self.temp_entry) } == 0 {
-                let last_error = io::Error::last_os_error();
-
-                return if last_error.raw_os_error().unwrap() as u32 == ERROR_NO_MORE_FILES {
-                    None
-                } else {
-                    Some(Err(last_error))
-                };
-            }
-        } else {
-            if unsafe { Process32FirstW(self.snapshot.handle(), &mut self.temp_entry) } == 0 {
-                return Some(Err(io::Error::last_os_error()));
-            }
-            self.iter_started = true;
-        }
-
-        Some(Ok(ProcessEntry {
-            pid: self.temp_entry.th32ProcessID,
-            parent_pid: self.temp_entry.th32ParentProcessID,
-        }))
-    }
-}
-
 /// Obtains a device path without resolving links or mount points.
 pub fn get_device_path<T: AsRef<Path>>(path: T) -> Result<OsString, io::Error> {
     // Preferentially, use GetFinalPathNameByHandleW. If the file does not exist
diff --git a/talpid-windows/src/lib.rs b/talpid-windows/src/lib.rs
index 865c3e8a81..755f6b2e30 100644
--- a/talpid-windows/src/lib.rs
+++ b/talpid-windows/src/lib.rs
@@ -12,3 +12,6 @@ pub mod net;
 
 /// Synchronization
 pub mod sync;
+
+/// Processes
+pub mod process;
diff --git a/talpid-windows/src/process.rs b/talpid-windows/src/process.rs
new file mode 100644
index 0000000000..ecddbe09f2
--- /dev/null
+++ b/talpid-windows/src/process.rs
@@ -0,0 +1,157 @@
+use std::{
+    ffi::{c_char, CStr},
+    io, mem,
+};
+use windows_sys::Win32::{
+    Foundation::{CloseHandle, ERROR_NO_MORE_FILES, HANDLE, INVALID_HANDLE_VALUE},
+    System::Diagnostics::ToolHelp::{
+        CreateToolhelp32Snapshot, Module32First, Module32Next, Process32FirstW, Process32NextW,
+        MODULEENTRY32, PROCESSENTRY32W,
+    },
+};
+
+/// A snapshot of process modules, threads, and heaps
+pub struct ProcessSnapshot {
+    handle: HANDLE,
+}
+
+impl ProcessSnapshot {
+    /// Create a new process snapshot using `CreateToolhelp32Snapshot`
+    pub fn new(flags: u32, process_id: u32) -> io::Result<ProcessSnapshot> {
+        let snap = unsafe { CreateToolhelp32Snapshot(flags, process_id) };
+
+        if snap == INVALID_HANDLE_VALUE {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(ProcessSnapshot { handle: snap })
+        }
+    }
+
+    /// Return the raw handle
+    pub fn as_raw(&self) -> HANDLE {
+        self.handle
+    }
+
+    /// Return an iterator over the modules in the snapshot
+    pub fn modules(&self) -> ProcessSnapshotModules<'_> {
+        let mut entry: MODULEENTRY32 = unsafe { mem::zeroed() };
+        entry.dwSize = mem::size_of::<MODULEENTRY32>() as u32;
+
+        ProcessSnapshotModules {
+            snapshot: self,
+            iter_started: false,
+            temp_entry: entry,
+        }
+    }
+
+    /// Return an iterator over the processes in the snapshot
+    pub fn processes(&self) -> ProcessSnapshotEntries<'_> {
+        let mut entry: PROCESSENTRY32W = unsafe { mem::zeroed() };
+        entry.dwSize = mem::size_of::<PROCESSENTRY32W>() as u32;
+
+        ProcessSnapshotEntries {
+            snapshot: self,
+            iter_started: false,
+            temp_entry: entry,
+        }
+    }
+}
+
+impl Drop for ProcessSnapshot {
+    fn drop(&mut self) {
+        unsafe {
+            CloseHandle(self.handle);
+        }
+    }
+}
+
+/// Description of a snapshot module entry. See `MODULEENTRY32`
+pub struct ModuleEntry {
+    /// Module name
+    pub name: String,
+    /// Module base address (in the owning process)
+    pub base_address: *const u8,
+    /// Size of the module (in bytes)
+    pub size: usize,
+}
+
+/// Module iterator for [ProcessSnapshot]
+pub struct ProcessSnapshotModules<'a> {
+    snapshot: &'a ProcessSnapshot,
+    iter_started: bool,
+    temp_entry: MODULEENTRY32,
+}
+
+impl Iterator for ProcessSnapshotModules<'_> {
+    type Item = io::Result<ModuleEntry>;
+
+    fn next(&mut self) -> Option<io::Result<ModuleEntry>> {
+        if self.iter_started {
+            if unsafe { Module32Next(self.snapshot.as_raw(), &mut self.temp_entry) } == 0 {
+                let last_error = io::Error::last_os_error();
+
+                return if last_error.raw_os_error().unwrap() as u32 == ERROR_NO_MORE_FILES {
+                    None
+                } else {
+                    Some(Err(last_error))
+                };
+            }
+        } else {
+            if unsafe { Module32First(self.snapshot.as_raw(), &mut self.temp_entry) } == 0 {
+                return Some(Err(io::Error::last_os_error()));
+            }
+            self.iter_started = true;
+        }
+
+        let cstr_ref = &self.temp_entry.szModule[0];
+        let cstr = unsafe { CStr::from_ptr(cstr_ref as *const u8 as *const c_char) };
+        Some(Ok(ModuleEntry {
+            name: cstr.to_string_lossy().into_owned(),
+            base_address: self.temp_entry.modBaseAddr,
+            size: self.temp_entry.modBaseSize as usize,
+        }))
+    }
+}
+
+/// Description of a snapshot process entry. See `PROCESSENTRY32W`
+pub struct ProcessEntry {
+    /// Process identifier
+    pub pid: u32,
+    /// Process identifier of the parent process
+    pub parent_pid: u32,
+}
+
+/// Process iterator for [ProcessSnapshot]
+pub struct ProcessSnapshotEntries<'a> {
+    snapshot: &'a ProcessSnapshot,
+    iter_started: bool,
+    temp_entry: PROCESSENTRY32W,
+}
+
+impl Iterator for ProcessSnapshotEntries<'_> {
+    type Item = io::Result<ProcessEntry>;
+
+    fn next(&mut self) -> Option<io::Result<ProcessEntry>> {
+        if self.iter_started {
+            if unsafe { Process32NextW(self.snapshot.as_raw(), &mut self.temp_entry) } == 0 {
+                let last_error = io::Error::last_os_error();
+
+                return if last_error.raw_os_error().unwrap() as u32 == ERROR_NO_MORE_FILES {
+                    None
+                } else {
+                    Some(Err(last_error))
+                };
+            }
+        } else {
+            if unsafe { Process32FirstW(self.snapshot.as_raw(), &mut self.temp_entry) } == 0 {
+                return Some(Err(io::Error::last_os_error()));
+            }
+            self.iter_started = true;
+        }
+
+        Some(Ok(ProcessEntry {
+            pid: self.temp_entry.th32ProcessID,
+            parent_pid: self.temp_entry.th32ParentProcessID,
+        }))
+    }
+}
diff --git a/talpid-windows/src/sync.rs b/talpid-windows/src/sync.rs
index 7b4ed59be3..202c96524d 100644
--- a/talpid-windows/src/sync.rs
+++ b/talpid-windows/src/sync.rs
@@ -1,7 +1,7 @@
 use std::{io, ptr};
 use windows_sys::Win32::{
-    Foundation::{CloseHandle, DuplicateHandle, BOOL, DUPLICATE_SAME_ACCESS, HANDLE},
-    System::Threading::{CreateEventW, GetCurrentProcess, SetEvent},
+    Foundation::{CloseHandle, BOOL, HANDLE},
+    System::Threading::{CreateEventW, SetEvent},
 };
 
 /// Windows event object
@@ -39,26 +39,6 @@ impl Event {
     pub fn as_raw(&self) -> HANDLE {
         self.0
     }
-
-    /// Duplicate the event object with `DuplicateHandle()`
-    pub fn duplicate(&self) -> io::Result<Event> {
-        let mut new_event = 0;
-        let status = unsafe {
-            DuplicateHandle(
-                GetCurrentProcess(),
-                self.0,
-                GetCurrentProcess(),
-                &mut new_event,
-                0,
-                0,
-                DUPLICATE_SAME_ACCESS,
-            )
-        };
-        if status == 0 {
-            return Err(io::Error::last_os_error());
-        }
-        Ok(Event(new_event))
-    }
 }
 
 impl Drop for Event {