Refactor handle shutdown

Author: Rubensei

CHANGELOG.md

@@ -12,6 +12,10 @@ and this project adheres to
 
 - Add `Drop` implementation for `WinDivert` to automatically close the handle
   when it goes out of scope.
+- Add `ShutdownHandle` struct to improve multithreaded handle shutdown
+  ergonomics.
+- Add `WinDivert::shutdown_handle(&self)` to create a `ShutdownHandle` for the
+  current instance.
 
 ### Changed
 
@@ -21,8 +25,14 @@ and this project adheres to
 - Bump `thiserror` to 2.0
 - Changed `WinDivert::shutdown()` method to use a shared reference instead of a
   mutable reference. (#16)
-- Changed `WinDivert::close()` method to be consuming instead of using a mutable
-  reference. (#15)
+- Changed `WinDivert::close()` method to be consuming, and remove it's `action`
+  parameter.
+
+### Removed
+
+- Removed `CloseAction`
+- Removed `Windivert::shutdown()` method in favor of
+  `Windivert::shutdown_handle()`.
 
 ## [Unreleased-sys]
 

windivert/src/core/winapi/mutex.rs

@@ -17,7 +17,7 @@ impl InstallMutex {
         Ok(Self { handle })
     }
 
-    pub fn lock(&mut self) -> Result<InstallMutexGuard, windows::core::Error> {
+    pub fn lock(&mut self) -> Result<InstallMutexGuard<'_>, windows::core::Error> {
         unsafe {
             match WaitForSingleObject(self.handle, INFINITE) {
                 WAIT_ABANDONED | WAIT_OBJECT_0 => {}

windivert/src/divert/mod.rs

@@ -4,6 +4,7 @@ use std::{
     marker::PhantomData,
     mem::MaybeUninit,
     path::Path,
+    sync::{Arc, Weak},
 };
 
 use windows::{
@@ -44,11 +45,10 @@ mod socket;
 #[non_exhaustive]
 #[derive(Debug)]
 pub struct WinDivert<L: layer::WinDivertLayerTrait> {
-    handle: HANDLE,
+    handle: Arc<HANDLE>,
     tls_index: TlsIndex,
     core: SysWrapper,
     _layer: PhantomData<L>,
-    _is_closed: bool,
 }
 
 unsafe impl<L: layer::WinDivertLayerTrait> Send for WinDivert<L> {}
@@ -74,11 +74,10 @@ impl<L: layer::WinDivertLayerTrait> WinDivert<L> {
             Err(open_err.into())
         } else {
             Ok(Self {
-                handle,
+                handle: Arc::new(handle),
                 tls_index: windivert_tls_idx,
                 core: sys_wrapper,
                 _layer: PhantomData::<L>,
-                _is_closed: false,
             })
         }
     }
@@ -348,24 +347,14 @@ impl<L: layer::WinDivertLayerTrait> WinDivert<L> {
     }
 
     /// Handle close function.
-    pub fn close(mut self, action: CloseAction) -> Result<(), WinDivertError> {
-        self.inner_close(action)
-    }
+    pub fn close(self) {}
 
-    /// Handle close function (internally, non-consuming).
-    fn inner_close(&mut self, action: CloseAction) -> Result<(), WinDivertError> {
-        self._is_closed = true;
-        unsafe { BOOL(WinDivertClose(self.handle.0)) }.ok()?;
-        match action {
-            CloseAction::Uninstall => WinDivert::uninstall(),
-            CloseAction::Nothing => Ok(()),
+    /// Returns a new ShutdownHandle that can be used to remotely shutdown WinDivert handle
+    pub fn shutdown_handle(&self) -> ShutdownHandle {
+        ShutdownHandle {
+            handle: Arc::downgrade(&self.handle),
         }
     }
-
-    /// Shutdown function.
-    pub fn shutdown(&self, mode: WinDivertShutdownMode) -> Result<(), WinDivertError> {
-        Ok(unsafe { BOOL(WinDivertShutdown(self.handle.0, mode)) }.ok()?)
-    }
 }
 
 /// Utility methods for WinDivert.
@@ -404,27 +393,24 @@ impl WinDivert<()> {
 
 impl<L: layer::WinDivertLayerTrait> Drop for WinDivert<L> {
     fn drop(&mut self) {
-        if !self._is_closed {
-            // SAFETY: Internal close should only fail if:
-            //   * Handle is closed: Checked
-            //   * Handle is invalid: Impossible with current API
-            //   * Permission issues: Impossible due to admin required for open
-            // It's safe to ignore the return value
-            let _ = self.inner_close(CloseAction::Nothing);
-        }
+        let _ = unsafe { BOOL(WinDivertClose(self.handle.0)) }.ok();
     }
 }
 
-/// Action parameter for  [`WinDivert::close()`](`fn@WinDivert::close`)
-pub enum CloseAction {
-    /// Close the handle and try to uninstall the WinDivert driver.
-    Uninstall,
-    /// Close the handle without uninstalling the driver.
-    Nothing,
+/// Struct that allows remote signaling the shutdown on the associated handle
+#[derive(Debug)]
+pub struct ShutdownHandle {
+    handle: Weak<HANDLE>,
 }
 
-impl Default for CloseAction {
-    fn default() -> Self {
-        Self::Nothing
+impl ShutdownHandle {
+    /// Shuts down the associated handle
+    /// This will prevent any further send, as well as stopping ongoing recv
+    /// Ongoing recv operations might not end immediately, all queued events/packets from the driver to the associated handle must be exhausted before reaching the `WinDivertRecvError::NoData` result.
+    pub fn shutdown(&self) -> Result<(), WinDivertError> {
+        if let Some(handle) = self.handle.upgrade() {
+            unsafe { BOOL(WinDivertShutdown(handle.0, WinDivertShutdownMode::Both)) }.ok()?;
+        };
+        Ok(())
     }
 }

windivert/src/divert/network.rs

@@ -125,11 +125,10 @@ mod tests {
 
     fn setup_divert(sys_wrapper: SysWrapper) -> WinDivert<NetworkLayer> {
         WinDivert {
-            handle: HANDLE(1usize as *mut c_void),
+            handle: Arc::new(HANDLE(1usize as *mut c_void)),
             tls_index: TlsIndex::alloc_tls().unwrap(),
             core: sys_wrapper,
             _layer: PhantomData::<NetworkLayer>,
-            _is_closed: false
         }
     }
 

windivert/src/lib.rs

@@ -16,43 +16,38 @@ let Ok(divert) = WinDivert::network("ip and tcp.DstPort == 443", 0, Default::def
     panic!("Failed to create WinDivert");
 };
 
-let divert = std::sync::Arc::new(divert);
+let shutdown_handle = divert.shutdown_handle();
 
-let divert_shared = divert.clone();
 let handle = std::thread::spawn(move || {
     // Do something in the background
     let mut buffer = [0u8; 1500];
 
     loop {
-        match divert_shared.recv(&mut buffer) {
+        match divert.recv(&mut buffer) {
             Ok(packet) => {
                 // In capture mode the packet is captured and not calling `send()` with it will prevent it from reaching the destination.
-                divert_shared.send(&packet).expect("Failed to send packet");
-            },
+                divert.send(&packet).expect("Failed to send packet");
+            }
             Err(WinDivertError::Recv(WinDivertRecvError::NoData)) => {
                 // Handle was shutdown, and there is no more pending data to receive
                 break;
-            },
+            }
             Err(e) => {
                 // Other errors
                 eprintln!("Error receiving packet: {}", e);
             }
         }
     }
+    // The handle is implicitly closed once `divert` is dropped
 });
 
 std::thread::sleep(std::time::Duration::from_secs(10));
 
-divert
-    .shutdown(WinDivertShutdownMode::Both)
+shutdown_handle
+    .shutdown()
     .expect("Failed to shutdown WinDivert");
 
 handle.join().unwrap();
-
-std::sync::Arc::try_unwrap(divert)
-    .expect("Thread already finished, no references remaining")
-    .close(CloseAction::Nothing)
-    .expect("Failed to close WinDivert");
 ```
 */
 

windivert/src/utils.rs

@@ -1,6 +1,6 @@
 use std::borrow::Cow;
 
-pub(crate) fn prepare_internet_slice_data(slice: &[u8]) -> (&[u8], Cow<[u8]>) {
+pub(crate) fn prepare_internet_slice_data(slice: &[u8]) -> (&[u8], Cow<'_, [u8]>) {
     let headers = etherparse::SlicedPacket::from_ip(slice)
         .expect("WinDivert can't capture anything below ip");
     let offset = match headers.net.unwrap() {