diff --git a/crates/lib/src/capture/mod.rs b/crates/lib/src/capture/mod.rs index a1995f8..1e8581a 100644 --- a/crates/lib/src/capture/mod.rs +++ b/crates/lib/src/capture/mod.rs @@ -1,5 +1,6 @@ pub(crate) mod packet; mod pcap; +pub(crate) mod socket; pub(crate) mod writer; use self::{pcap::Pcap, writer::Writer}; @@ -38,4 +39,4 @@ pub fn setup_capture(file_path: Option) { /// /// # Errors /// Returns an Error if the writer is already set. -fn attach(writer: Box) { crate::socket::capture::set_writer(writer); } +fn attach(writer: Box) { crate::capture::socket::set_writer(writer); } diff --git a/crates/lib/src/capture/socket.rs b/crates/lib/src/capture/socket.rs new file mode 100644 index 0000000..546e9b8 --- /dev/null +++ b/crates/lib/src/capture/socket.rs @@ -0,0 +1,214 @@ +use std::{marker::PhantomData, net::SocketAddr}; + +use crate::{ + capture::{ + packet::CapturePacket, + packet::{Direction, Protocol}, + writer::{Writer, CAPTURE_WRITER}, + }, + protocols::types::TimeoutSettings, + socket::{Socket, TcpSocketImpl, UdpSocketImpl}, + GDResult, +}; + +/// Sets a global capture writer for handling all packet data. +/// +/// # Panics +/// Panics if a capture writer is already set. +/// +/// # Arguments +/// * `writer` - A boxed writer that implements the `Writer` trait. +pub(crate) fn set_writer(writer: Box) { + let mut lock = CAPTURE_WRITER.lock().unwrap(); + + if lock.is_some() { + panic!("Capture writer already set"); + } + + *lock = Some(writer); +} + +/// A trait representing a provider of a network protocol. +pub(crate) trait ProtocolProvider { + /// Returns the protocol used by the provider. + fn protocol() -> Protocol; +} + +/// Represents the TCP protocol provider. +pub(crate) struct ProtocolTCP; +impl ProtocolProvider for ProtocolTCP { + fn protocol() -> Protocol { Protocol::Tcp } +} + +/// Represents the UDP protocol provider. +pub(crate) struct ProtocolUDP; +impl ProtocolProvider for ProtocolUDP { + fn protocol() -> Protocol { Protocol::Udp } +} + +/// A socket wrapper that allows capturing packets. +/// +/// # Type parameters +/// * `I` - The inner socket type. +/// * `P` - The protocol provider. +#[derive(Clone, Debug)] +pub(crate) struct WrappedCaptureSocket { + inner: I, + remote_address: SocketAddr, + _protocol: PhantomData

, +} + +impl Socket for WrappedCaptureSocket { + /// Creates a new wrapped socket for capturing packets. + /// + /// Initializes a new socket of type `I`, wrapping it to enable packet + /// capturing. Capturing is protocol-specific, as indicated by + /// the `ProtocolProvider`. + /// + /// # Arguments + /// * `address` - The address to connect the socket to. + /// * `timeout_settings` - Optional timeout settings for the socket. + /// + /// # Returns + /// A `GDResult` containing either the wrapped socket or an error. + fn new(address: &SocketAddr, timeout_settings: &Option) -> GDResult + where Self: Sized { + let v = Self { + inner: I::new(address, timeout_settings)?, + remote_address: *address, + _protocol: PhantomData, + }; + + let info = CapturePacket { + direction: Direction::Send, + protocol: P::protocol(), + remote_address: address, + local_address: &v.local_addr().unwrap(), + }; + + if let Some(writer) = CAPTURE_WRITER.lock().unwrap().as_mut() { + writer.new_connect(&info)?; + } + + Ok(v) + } + + /// Sends data over the socket and captures the packet. + /// + /// The method sends data using the inner socket and captures the sent + /// packet if a capture writer is set. + /// + /// # Arguments + /// * `data` - Data to be sent. + /// + /// # Returns + /// A result indicating success or error in sending data. + fn send(&mut self, data: &[u8]) -> GDResult<()> { + let info = CapturePacket { + direction: Direction::Send, + protocol: P::protocol(), + remote_address: &self.remote_address, + local_address: &self.local_addr().unwrap(), + }; + + if let Some(writer) = CAPTURE_WRITER.lock().unwrap().as_mut() { + writer.write(&info, data)?; + } + + self.inner.send(data) + } + + /// Receives data from the socket and captures the packet. + /// + /// The method receives data using the inner socket and captures the + /// incoming packet if a capture writer is set. + /// + /// # Arguments + /// * `size` - Optional size of data to receive. + /// + /// # Returns + /// A result containing received data or an error. + fn receive(&mut self, size: Option) -> crate::GDResult> { + let data = self.inner.receive(size)?; + let info = CapturePacket { + direction: Direction::Receive, + protocol: P::protocol(), + remote_address: &self.remote_address, + local_address: &self.local_addr().unwrap(), + }; + + if let Some(writer) = CAPTURE_WRITER.lock().unwrap().as_mut() { + writer.write(&info, &data)?; + } + + Ok(data) + } + + /// Applies timeout settings to the wrapped socket. + /// + /// Delegates the operation to the inner socket implementation. + /// + /// # Arguments + /// * `timeout_settings` - Optional timeout settings to apply. + /// + /// # Returns + /// A result indicating success or error in applying timeouts. + fn apply_timeout( + &self, + timeout_settings: &Option, + ) -> crate::GDResult<()> { + self.inner.apply_timeout(timeout_settings) + } + + /// Returns the remote port of the wrapped socket. + /// + /// Delegates the operation to the inner socket implementation. + /// + /// # Returns + /// The remote port number. + fn port(&self) -> u16 { self.inner.port() } + + /// Returns the local SocketAddr of the wrapped socket. + /// + /// Delegates the operation to the inner socket implementation. + /// + /// # Returns + /// The local SocketAddr. + fn local_addr(&self) -> std::io::Result { self.inner.local_addr() } +} + +// this seems a bad way to do this, but its safe +impl Drop for WrappedCaptureSocket { + fn drop(&mut self) { + // Construct the CapturePacket info + let info = CapturePacket { + direction: Direction::Send, + protocol: P::protocol(), + remote_address: &self.remote_address, + local_address: &self + .local_addr() + .unwrap_or_else(|_| SocketAddr::new(std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED), 0)), + }; + + // If a capture writer is set, close the connection and capture the packet. + if let Some(writer) = CAPTURE_WRITER.lock().unwrap().as_mut() { + let _ = writer.close_connection(&info); + } + } +} + +/// A specialized `WrappedCaptureSocket` for UDP, using `UdpSocketImpl` as +/// the inner socket and `ProtocolUDP` as the protocol provider. +/// +/// This type captures and processes UDP packets, wrapping around standard +/// UDP socket functionalities with additional packet capture +/// capabilities. +pub(crate) type CapturedUdpSocket = WrappedCaptureSocket; + +/// A specialized `WrappedCaptureSocket` for TCP, using `TcpSocketImpl` as +/// the inner socket and `ProtocolTCP` as the protocol provider. +/// +/// This type captures and processes TCP packets, wrapping around standard +/// TCP socket functionalities with additional packet capture +/// capabilities. +pub(crate) type CapturedTcpSocket = WrappedCaptureSocket; diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index eeba5eb..77d4c40 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -167,235 +167,15 @@ impl Socket for UdpSocketImpl { fn local_addr(&self) -> std::io::Result { self.socket.local_addr() } } -/// Things used for capturing packets. -#[cfg(feature = "packet_capture")] -pub mod capture { - use std::{marker::PhantomData, net::SocketAddr}; - - use super::{Socket, TcpSocketImpl, UdpSocketImpl}; - - use crate::{ - capture::{ - packet::CapturePacket, - packet::{Direction, Protocol}, - writer::{Writer, CAPTURE_WRITER}, - }, - protocols::types::TimeoutSettings, - GDResult, - }; - - /// Sets a global capture writer for handling all packet data. - /// - /// # Panics - /// Panics if a capture writer is already set. - /// - /// # Arguments - /// * `writer` - A boxed writer that implements the `Writer` trait. - pub(crate) fn set_writer(writer: Box) { - let mut lock = CAPTURE_WRITER.lock().unwrap(); - - if lock.is_some() { - panic!("Capture writer already set"); - } - - *lock = Some(writer); - } - - /// A trait representing a provider of a network protocol. - pub(crate) trait ProtocolProvider { - /// Returns the protocol used by the provider. - fn protocol() -> Protocol; - } - - /// Represents the TCP protocol provider. - pub(crate) struct ProtocolTCP; - impl ProtocolProvider for ProtocolTCP { - fn protocol() -> Protocol { Protocol::Tcp } - } - - /// Represents the UDP protocol provider. - pub(crate) struct ProtocolUDP; - impl ProtocolProvider for ProtocolUDP { - fn protocol() -> Protocol { Protocol::Udp } - } - - /// A socket wrapper that allows capturing packets. - /// - /// # Type parameters - /// * `I` - The inner socket type. - /// * `P` - The protocol provider. - #[derive(Clone, Debug)] - pub(crate) struct WrappedCaptureSocket { - inner: I, - remote_address: SocketAddr, - _protocol: PhantomData

, - } - - impl Socket for WrappedCaptureSocket { - /// Creates a new wrapped socket for capturing packets. - /// - /// Initializes a new socket of type `I`, wrapping it to enable packet - /// capturing. Capturing is protocol-specific, as indicated by - /// the `ProtocolProvider`. - /// - /// # Arguments - /// * `address` - The address to connect the socket to. - /// * `timeout_settings` - Optional timeout settings for the socket. - /// - /// # Returns - /// A `GDResult` containing either the wrapped socket or an error. - fn new(address: &SocketAddr, timeout_settings: &Option) -> GDResult - where Self: Sized { - let v = Self { - inner: I::new(address, timeout_settings)?, - remote_address: *address, - _protocol: PhantomData, - }; - - let info = CapturePacket { - direction: Direction::Send, - protocol: P::protocol(), - remote_address: address, - local_address: &v.local_addr().unwrap(), - }; - - if let Some(writer) = CAPTURE_WRITER.lock().unwrap().as_mut() { - writer.new_connect(&info)?; - } - - Ok(v) - } - - /// Sends data over the socket and captures the packet. - /// - /// The method sends data using the inner socket and captures the sent - /// packet if a capture writer is set. - /// - /// # Arguments - /// * `data` - Data to be sent. - /// - /// # Returns - /// A result indicating success or error in sending data. - fn send(&mut self, data: &[u8]) -> GDResult<()> { - let info = CapturePacket { - direction: Direction::Send, - protocol: P::protocol(), - remote_address: &self.remote_address, - local_address: &self.local_addr().unwrap(), - }; - - if let Some(writer) = CAPTURE_WRITER.lock().unwrap().as_mut() { - writer.write(&info, data)?; - } - - self.inner.send(data) - } - - /// Receives data from the socket and captures the packet. - /// - /// The method receives data using the inner socket and captures the - /// incoming packet if a capture writer is set. - /// - /// # Arguments - /// * `size` - Optional size of data to receive. - /// - /// # Returns - /// A result containing received data or an error. - fn receive(&mut self, size: Option) -> crate::GDResult> { - let data = self.inner.receive(size)?; - let info = CapturePacket { - direction: Direction::Receive, - protocol: P::protocol(), - remote_address: &self.remote_address, - local_address: &self.local_addr().unwrap(), - }; - - if let Some(writer) = CAPTURE_WRITER.lock().unwrap().as_mut() { - writer.write(&info, &data)?; - } - - Ok(data) - } - - /// Applies timeout settings to the wrapped socket. - /// - /// Delegates the operation to the inner socket implementation. - /// - /// # Arguments - /// * `timeout_settings` - Optional timeout settings to apply. - /// - /// # Returns - /// A result indicating success or error in applying timeouts. - fn apply_timeout( - &self, - timeout_settings: &Option, - ) -> crate::GDResult<()> { - self.inner.apply_timeout(timeout_settings) - } - - /// Returns the remote port of the wrapped socket. - /// - /// Delegates the operation to the inner socket implementation. - /// - /// # Returns - /// The remote port number. - fn port(&self) -> u16 { self.inner.port() } - - /// Returns the local SocketAddr of the wrapped socket. - /// - /// Delegates the operation to the inner socket implementation. - /// - /// # Returns - /// The local SocketAddr. - fn local_addr(&self) -> std::io::Result { self.inner.local_addr() } - } - - // this seems a bad way to do this, but its safe - impl Drop for WrappedCaptureSocket { - fn drop(&mut self) { - // Construct the CapturePacket info - let info = CapturePacket { - direction: Direction::Send, - protocol: P::protocol(), - remote_address: &self.remote_address, - local_address: &self - .local_addr() - .unwrap_or_else(|_| SocketAddr::new(std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED), 0)), - }; - - // If a capture writer is set, close the connection and capture the packet. - if let Some(writer) = CAPTURE_WRITER.lock().unwrap().as_mut() { - let _ = writer.close_connection(&info); - } - } - } - - /// A specialized `WrappedCaptureSocket` for UDP, using `UdpSocketImpl` as - /// the inner socket and `ProtocolUDP` as the protocol provider. - /// - /// This type captures and processes UDP packets, wrapping around standard - /// UDP socket functionalities with additional packet capture - /// capabilities. - pub(crate) type CapturedUdpSocket = WrappedCaptureSocket; - - /// A specialized `WrappedCaptureSocket` for TCP, using `TcpSocketImpl` as - /// the inner socket and `ProtocolTCP` as the protocol provider. - /// - /// This type captures and processes TCP packets, wrapping around standard - /// TCP socket functionalities with additional packet capture - /// capabilities. - pub(crate) type CapturedTcpSocket = WrappedCaptureSocket; -} - #[cfg(not(feature = "packet_capture"))] -pub type UdpSocket = UdpSocketImpl; +pub type UdpSocket = UdpSocketImpl; #[cfg(not(feature = "packet_capture"))] pub type TcpSocket = TcpSocketImpl; #[cfg(feature = "packet_capture")] -pub(crate) type UdpSocket = capture::CapturedUdpSocket; +pub(crate) type UdpSocket = crate::capture::socket::CapturedUdpSocket; #[cfg(feature = "packet_capture")] -pub(crate) type TcpSocket = capture::CapturedTcpSocket; +pub(crate) type TcpSocket = crate::capture::socket::CapturedTcpSocket; #[cfg(test)] mod tests {