Add process module to talpid-windows

2 files changed, 4 insertions, 91 deletions

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