diff options
Diffstat (limited to 'crates/windivert/src/divert')
| -rw-r--r-- | crates/windivert/src/divert/blocking.rs | 353 | ||||
| -rw-r--r-- | crates/windivert/src/divert/mod.rs | 229 |
2 files changed, 582 insertions, 0 deletions
diff --git a/crates/windivert/src/divert/blocking.rs b/crates/windivert/src/divert/blocking.rs new file mode 100644 index 0000000..c003246 --- /dev/null +++ b/crates/windivert/src/divert/blocking.rs @@ -0,0 +1,353 @@ +use std::borrow::Cow; +use std::{ffi::c_void, mem::MaybeUninit}; + +use crate::address::WinDivertAddress; +use crate::layer; +use crate::prelude::*; +use etherparse::{InternetSlice, SlicedPacket}; +use sys::address::WINDIVERT_ADDRESS; +use windivert_sys as sys; + +const ADDR_SIZE: usize = std::mem::size_of::<WINDIVERT_ADDRESS>(); + +impl<L: layer::WinDivertLayerTrait> WinDivert<L> { + fn internal_recv<'a>( + &self, + buffer: Option<&'a mut [u8]>, + ) -> Result<WinDivertPacket<'a, L>, WinDivertError> { + let mut packet_length = 0; + let mut addr = MaybeUninit::uninit(); + let (buffer_ptr, buffer_len) = if let Some(ref buffer) = buffer { + (buffer.as_ptr(), buffer.len()) + } else { + (std::ptr::null(), 0) + }; + + let res = unsafe { + sys::WinDivertRecv( + self.handle, + buffer_ptr as *mut c_void, + buffer_len as u32, + &mut packet_length, + addr.as_mut_ptr(), + ) + }; + + if res.as_bool() { + Ok(WinDivertPacket { + address: WinDivertAddress::<L>::from_raw(unsafe { addr.assume_init() }), + data: buffer + .map(|b| Cow::Borrowed(&b[..packet_length as usize])) + .unwrap_or_default(), + }) + } else { + let recv_err = WinDivertRecvError::try_from(std::io::Error::last_os_error())?; + Err(recv_err.into()) + } + } + + fn internal_recv_ex<'a>( + &self, + buffer: Option<&'a mut [u8]>, + packet_count: usize, + ) -> Result<(Option<&'a [u8]>, Vec<WINDIVERT_ADDRESS>), WinDivertError> { + let mut packet_length = 0; + + let mut addr_len = (ADDR_SIZE * packet_count) as u32; + let mut addr_buffer: Vec<WINDIVERT_ADDRESS> = + vec![WINDIVERT_ADDRESS::default(); packet_count]; + + let (buffer_ptr, buffer_len) = if let Some(buffer) = &buffer { + (buffer.as_ptr(), buffer.len()) + } else { + (std::ptr::null(), 0) + }; + + let res = unsafe { + sys::WinDivertRecvEx( + self.handle, + buffer_ptr as *mut c_void, + buffer_len as u32, + &mut packet_length, + 0, + addr_buffer.as_mut_ptr(), + &mut addr_len, + std::ptr::null_mut(), + ) + }; + + if res.as_bool() { + addr_buffer.truncate((addr_len / ADDR_SIZE as u32) as usize); + Ok(( + buffer.map(|buffer| &buffer[..packet_length as usize]), + addr_buffer, + )) + } else { + let recv_err = WinDivertRecvError::try_from(std::io::Error::last_os_error())?; + Err(recv_err.into()) + } + } + + fn internal_send(&self, packet: &WinDivertPacket<L>) -> Result<u32, WinDivertError> { + let mut injected_length = 0; + + let res = unsafe { + sys::WinDivertSend( + self.handle, + packet.data.as_ptr() as *const c_void, + packet.data.len() as u32, + &mut injected_length, + packet.address.as_ref(), + ) + }; + + if !res.as_bool() { + return Err(std::io::Error::last_os_error().into()); + } + + Ok(injected_length) + } + + fn internal_send_ex<'data, 'packets, P>(&self, packets: P) -> Result<u32, WinDivertError> + where + P: ExactSizeIterator<Item = &'packets WinDivertPacket<'data, L>>, + 'data: 'packets, + L: 'packets, + { + let packet_count = packets.len(); + let mut injected_length = 0; + let mut packet_buffer: Vec<u8> = Vec::new(); + let mut address_buffer: Vec<WINDIVERT_ADDRESS> = Vec::with_capacity(packet_count); + packets.for_each(|packet: &'packets WinDivertPacket<'data, L>| { + packet_buffer.extend(&packet.data[..]); + address_buffer.push(*packet.address.as_ref()); + }); + + let res = unsafe { + sys::WinDivertSendEx( + self.handle, + packet_buffer.as_ptr() as *const c_void, + packet_buffer.len() as u32, + &mut injected_length, + 0, + address_buffer.as_ptr(), + (std::mem::size_of::<WINDIVERT_ADDRESS>() * packet_count) as u32, + std::ptr::null_mut(), + ) + }; + + if !res.as_bool() { + return Err(std::io::Error::last_os_error().into()); + } + + Ok(injected_length) + } +} + +impl WinDivert<layer::NetworkLayer> { + /// Single packet blocking recv function. + pub fn recv<'a>( + &self, + buffer: Option<&'a mut [u8]>, + ) -> Result<WinDivertPacket<'a, layer::NetworkLayer>, WinDivertError> { + self.internal_recv(buffer) + } + + /// Batched blocking recv function. + pub fn recv_ex<'a>( + &self, + buffer: Option<&'a mut [u8]>, + packet_count: usize, + ) -> Result<Vec<WinDivertPacket<'a, layer::NetworkLayer>>, WinDivertError> { + let (mut buffer, addresses) = self.internal_recv_ex(buffer, packet_count)?; + let mut packets = Vec::with_capacity(addresses.len()); + for addr in addresses.into_iter() { + packets.push(WinDivertPacket { + address: WinDivertAddress::<layer::NetworkLayer>::from_raw(addr), + data: buffer + .map(|inner_buffer| { + let headers = SlicedPacket::from_ip(inner_buffer) + .expect("WinDivert can't capture anything below ip"); + let offset = match headers.ip.unwrap() { + InternetSlice::Ipv4(ip_header, _) => ip_header.total_len() as usize, + InternetSlice::Ipv6(ip6header, _) => { + ip6header.payload_length() as usize + 40 + } + }; + let (data, tail) = inner_buffer.split_at(offset); + buffer = Some(tail); + Cow::Borrowed(data) + }) + .unwrap_or_default(), + }); + } + Ok(packets) + } + + /// Single packet send function. + pub fn send( + &self, + packet: &WinDivertPacket<layer::NetworkLayer>, + ) -> Result<u32, WinDivertError> { + self.internal_send(packet) + } + + /// Batched packet send function. + pub fn send_ex<'data, 'packets, P, I>(&self, packets: P) -> Result<u32, WinDivertError> + where + P: IntoIterator<IntoIter = I>, + I: ExactSizeIterator<Item = &'packets WinDivertPacket<'data, layer::NetworkLayer>>, + 'data: 'packets, + { + self.internal_send_ex(packets.into_iter()) + } +} + +impl WinDivert<layer::ForwardLayer> { + /// Single packet blocking recv function. + pub fn recv<'a>( + &self, + buffer: Option<&'a mut [u8]>, + ) -> Result<WinDivertPacket<'a, layer::ForwardLayer>, WinDivertError> { + self.internal_recv(buffer) + } + + /// Batched blocking recv function. + pub fn recv_ex<'a>( + &self, + buffer: Option<&'a mut [u8]>, + packet_count: usize, + ) -> Result<Vec<WinDivertPacket<'a, layer::NetworkLayer>>, WinDivertError> { + let (mut buffer, addresses) = self.internal_recv_ex(buffer, packet_count)?; + let mut packets = Vec::with_capacity(addresses.len()); + for addr in addresses.into_iter() { + packets.push(WinDivertPacket { + address: WinDivertAddress::<layer::NetworkLayer>::from_raw(addr), + data: buffer + .map(|inner_buffer| { + let headers = SlicedPacket::from_ip(inner_buffer) + .expect("WinDivert can't capture anything below ip"); + let offset = match headers.ip.unwrap() { + InternetSlice::Ipv4(ip_header, _) => ip_header.total_len() as usize, + InternetSlice::Ipv6(ip6header, _) => { + ip6header.payload_length() as usize + 40 + } + }; + let (data, tail) = inner_buffer.split_at(offset); + buffer = Some(tail); + Cow::Borrowed(data) + }) + .unwrap_or_default(), + }); + } + Ok(packets) + } + + /// Single packet send function. + pub fn send( + &self, + packet: &WinDivertPacket<layer::ForwardLayer>, + ) -> Result<u32, WinDivertError> { + self.internal_send(packet) + } + + /// Batched packet send function. + pub fn send_ex<'data, 'packets, P, I>(&self, packets: P) -> Result<u32, WinDivertError> + where + P: IntoIterator<IntoIter = I>, + I: ExactSizeIterator<Item = &'packets WinDivertPacket<'data, layer::ForwardLayer>>, + 'data: 'packets, + { + self.internal_send_ex(packets.into_iter()) + } +} + +impl WinDivert<layer::FlowLayer> { + /// Single packet blocking recv function. + pub fn recv<'a>( + &self, + buffer: Option<&'a mut [u8]>, + ) -> Result<WinDivertPacket<'a, layer::FlowLayer>, WinDivertError> { + self.internal_recv(buffer) + } + + /// Batched blocking recv function. + pub fn recv_ex<'a>( + &self, + packet_count: usize, + ) -> Result<Vec<WinDivertPacket<'a, layer::FlowLayer>>, WinDivertError> { + let (_, addresses) = self.internal_recv_ex(None, packet_count)?; + let mut packets = Vec::with_capacity(addresses.len()); + for addr in addresses.into_iter() { + packets.push(WinDivertPacket::<layer::FlowLayer> { + address: WinDivertAddress::<layer::FlowLayer>::from_raw(addr), + data: Default::default(), + }); + } + Ok(packets) + } +} + +impl WinDivert<layer::SocketLayer> { + /// Single packet blocking recv function. + pub fn recv<'a>( + &self, + buffer: Option<&'a mut [u8]>, + ) -> Result<WinDivertPacket<'a, layer::SocketLayer>, WinDivertError> { + self.internal_recv(buffer) + } + + /// Batched blocking recv function. + pub fn recv_ex<'a>( + &self, + packet_count: usize, + ) -> Result<Vec<WinDivertPacket<'a, layer::SocketLayer>>, WinDivertError> { + let (_, addresses) = self.internal_recv_ex(None, packet_count)?; + let mut packets = Vec::with_capacity(addresses.len()); + for addr in addresses.into_iter() { + packets.push(WinDivertPacket::<layer::SocketLayer> { + address: WinDivertAddress::<layer::SocketLayer>::from_raw(addr), + data: Default::default(), + }); + } + Ok(packets) + } +} + +impl WinDivert<layer::ReflectLayer> { + /// Single packet blocking recv function. + pub fn recv<'a>( + &self, + buffer: Option<&'a mut [u8]>, + ) -> Result<WinDivertPacket<'a, layer::ReflectLayer>, WinDivertError> { + self.internal_recv(buffer) + } + + /// Batched blocking recv function. + pub fn recv_ex<'a>( + &self, + buffer: Option<&'a mut [u8]>, + packet_count: usize, + ) -> Result<Vec<WinDivertPacket<'a, layer::ReflectLayer>>, WinDivertError> { + let (mut buffer, addresses) = self.internal_recv_ex(buffer, packet_count)?; + let mut packets = Vec::with_capacity(addresses.len()); + for addr in addresses.into_iter() { + packets.push(WinDivertPacket { + address: WinDivertAddress::<layer::ReflectLayer>::from_raw(addr), + data: buffer + .map(|inner_buffer| { + let (data, tail) = inner_buffer.split_at( + inner_buffer + .iter() + .position(|&x| x == b'\0') + .expect("CStrings always end in null"), + ); + buffer = Some(tail); + Cow::Borrowed(data) + }) + .unwrap_or_default(), + }); + } + Ok(packets) + } +} diff --git a/crates/windivert/src/divert/mod.rs b/crates/windivert/src/divert/mod.rs new file mode 100644 index 0000000..2bce256 --- /dev/null +++ b/crates/windivert/src/divert/mod.rs @@ -0,0 +1,229 @@ +mod blocking; + +use std::{ + ffi::{c_void, CString}, + marker::PhantomData, + mem::MaybeUninit, +}; + +use crate::layer; +use crate::prelude::*; +use sys::{WinDivertParam, WinDivertShutdownMode}; +use windivert_sys as sys; + +use windows::{ + core::{Error as WinError, Result as WinResult, PCSTR}, + Win32::{ + Foundation::{GetLastError, HANDLE}, + System::{ + Services::{ + CloseServiceHandle, ControlService, OpenSCManagerA, OpenServiceA, + SC_MANAGER_ALL_ACCESS, SERVICE_CONTROL_STOP, SERVICE_STATUS, + }, + Threading::{CreateEventA, TlsAlloc, TlsGetValue, TlsSetValue}, + }, + }, +}; + +/// Main wrapper struct around windivert functionalities. +#[non_exhaustive] +pub struct WinDivert<L: layer::WinDivertLayerTrait> { + handle: HANDLE, + _tls_idx: u32, + _layer: PhantomData<L>, +} + +/// Recv implementations +impl<L: layer::WinDivertLayerTrait> WinDivert<L> { + /// Open a handle using the specified parameters. + fn new( + filter: &str, + layer: WinDivertLayer, + priority: i16, + flags: WinDivertFlags, + ) -> Result<Self, WinDivertError> { + let filter = CString::new(filter)?; + let windivert_tls_idx = unsafe { TlsAlloc() }; + let handle = unsafe { sys::WinDivertOpen(filter.as_ptr(), layer, priority, flags) }; + if handle.is_invalid() { + let open_err = WinDivertOpenError::try_from(std::io::Error::last_os_error())?; + Err(open_err.into()) + } else { + Ok(Self { + handle, + _tls_idx: windivert_tls_idx, + _layer: PhantomData::<L>, + }) + } + } + + pub(crate) fn _get_event(tls_idx: u32) -> Result<HANDLE, WinDivertError> { + let mut event = HANDLE::default(); + unsafe { + event.0 = TlsGetValue(tls_idx) as isize; + if event.is_invalid() { + event = CreateEventA(None, false, false, None)?; + TlsSetValue(tls_idx, Some(event.0 as *mut c_void)); + } + } + Ok(event) + } + + /// Methods that allows to query the driver for parameters. + pub fn get_param(&self, param: WinDivertParam) -> Result<u64, WinDivertError> { + let mut value = 0; + let res = unsafe { sys::WinDivertGetParam(self.handle, param, &mut value) }; + if !res.as_bool() { + return Err(std::io::Error::last_os_error().into()); + } + Ok(value) + } + + /// Method that allows setting driver parameters. + pub fn set_param(&self, param: WinDivertParam, value: u64) -> Result<(), WinDivertError> { + match param { + WinDivertParam::VersionMajor | WinDivertParam::VersionMinor => { + Err(WinDivertError::Parameter(param, value)) + } + _ => unsafe { sys::WinDivertSetParam(self.handle, param, value) } + .ok() + .map_err(|_| std::io::Error::last_os_error().into()), + } + } + + /// Handle close function. + pub fn close(&mut self, action: CloseAction) -> WinResult<()> { + let res = unsafe { sys::WinDivertClose(self.handle) }; + if !res.as_bool() { + return Err(WinError::from(unsafe { GetLastError() })); + } + match action { + CloseAction::Uninstall => WinDivert::uninstall(), + CloseAction::Nothing => Ok(()), + } + } + + /// Shutdown function. + pub fn shutdown(&self, mode: WinDivertShutdownMode) -> WinResult<()> { + let res = unsafe { sys::WinDivertShutdown(self.handle, mode) }; + if !res.as_bool() { + return Err(WinError::from(unsafe { GetLastError() })); + } + Ok(()) + } +} + +impl WinDivert<layer::NetworkLayer> { + /// WinDivert constructor for network layer. + pub fn network( + filter: impl AsRef<str>, + priority: i16, + flags: WinDivertFlags, + ) -> Result<Self, WinDivertError> { + Self::new(filter.as_ref(), WinDivertLayer::Network, priority, flags) + } +} + +impl WinDivert<layer::ForwardLayer> { + /// WinDivert constructor for forward layer. + pub fn forward( + filter: impl AsRef<str>, + priority: i16, + flags: WinDivertFlags, + ) -> Result<Self, WinDivertError> { + Self::new(filter.as_ref(), WinDivertLayer::Forward, priority, flags) + } +} + +impl WinDivert<layer::FlowLayer> { + /// WinDivert constructor for flow layer. + pub fn flow( + filter: &str, + priority: i16, + flags: WinDivertFlags, + ) -> Result<Self, WinDivertError> { + Self::new( + filter, + WinDivertLayer::Flow, + priority, + flags.set_recv_only().set_sniff(), + ) + } +} + +impl WinDivert<layer::SocketLayer> { + /// WinDivert constructor for socket layer. + pub fn socket( + filter: impl AsRef<str>, + priority: i16, + flags: WinDivertFlags, + ) -> Result<Self, WinDivertError> { + Self::new( + filter.as_ref(), + WinDivertLayer::Socket, + priority, + flags.set_recv_only(), + ) + } +} + +impl WinDivert<layer::ReflectLayer> { + /// WinDivert constructor for reflect layer. + pub fn reflect( + filter: impl AsRef<str>, + priority: i16, + flags: WinDivertFlags, + ) -> Result<Self, WinDivertError> { + Self::new( + filter.as_ref(), + WinDivertLayer::Reflect, + priority, + flags.set_recv_only().set_sniff(), + ) + } +} + +impl WinDivert<()> { + /// Maximum number of packets that can be captured/sent in a single batched operation + pub const MAX_BATCH: u8 = windivert_sys::WINDIVERT_BATCH_MAX as u8; + + /// Method that tries to uninstall WinDivert driver. + pub fn uninstall() -> WinResult<()> { + let mut status = MaybeUninit::<SERVICE_STATUS>::uninit(); + unsafe { + let manager = OpenSCManagerA(None, None, SC_MANAGER_ALL_ACCESS)?; + let service = OpenServiceA( + manager, + PCSTR::from_raw(c"WinDivert".as_ptr() as *const u8), + SC_MANAGER_ALL_ACCESS, + )?; + let res = ControlService(service, SERVICE_CONTROL_STOP, status.as_mut_ptr()); + if !res.as_bool() { + return Err(WinError::from(GetLastError())); + } + let res = CloseServiceHandle(service); + if !res.as_bool() { + return Err(WinError::from(GetLastError())); + } + let res = CloseServiceHandle(manager); + if !res.as_bool() { + return Err(WinError::from(GetLastError())); + } + } + 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, +} + +impl Default for CloseAction { + fn default() -> Self { + Self::Nothing + } +} |