tribufu/rust-gamedig
1
0
Fork 0
mirror of https://github.com/tribufu/rust-gamedig synced 2026-05-06 07:17:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor: capture socket and move to capture dir

Browse source
This commit is contained in:
Cain 2024-02-07 15:28:41 +00:00
parent c34392a3da
commit 9a2b953fff
3 changed files with 219 additions and 224 deletions
Download patch file Download diff file Expand all files Collapse all files

3
crates/lib/src/capture/mod.rs
View file

@ -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<PathBuf>) {
///
/// # Errors
/// Returns an Error if the writer is already set.
fn attach(writer: Box<dyn Writer + Send + Sync>) { crate::socket::capture::set_writer(writer); }
fn attach(writer: Box<dyn Writer + Send + Sync>) { crate::capture::socket::set_writer(writer); }

214
crates/lib/src/capture/socket.rs Normal file
View file

@ -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<dyn Writer + Send + Sync>) {
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<I: Socket, P: ProtocolProvider> {
inner: I,
remote_address: SocketAddr,
_protocol: PhantomData<P>,
}
impl<I: Socket, P: ProtocolProvider> Socket for WrappedCaptureSocket<I, P> {
/// 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<TimeoutSettings>) -> GDResult<Self>
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<usize>) -> crate::GDResult<Vec<u8>> {
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::protocols::types::TimeoutSettings>,
) -> 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<SocketAddr> { self.inner.local_addr() }
}
// this seems a bad way to do this, but its safe
impl<I: Socket, P: ProtocolProvider> Drop for WrappedCaptureSocket<I, P> {
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<UdpSocketImpl, ProtocolUDP>;
/// 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<TcpSocketImpl, ProtocolTCP>;

226
crates/lib/src/socket.rs
View file

@ -167,235 +167,15 @@ impl Socket for UdpSocketImpl {
fn local_addr(&self) -> std::io::Result<SocketAddr> { 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<dyn Writer + Send + Sync>) {
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<I: Socket, P: ProtocolProvider> {
inner: I,
remote_address: SocketAddr,
_protocol: PhantomData<P>,
}
impl<I: Socket, P: ProtocolProvider> Socket for WrappedCaptureSocket<I, P> {
/// 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<TimeoutSettings>) -> GDResult<Self>
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<usize>) -> crate::GDResult<Vec<u8>> {
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::protocols::types::TimeoutSettings>,
) -> 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<SocketAddr> { self.inner.local_addr() }
}
// this seems a bad way to do this, but its safe
impl<I: Socket, P: ProtocolProvider> Drop for WrappedCaptureSocket<I, P> {
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<UdpSocketImpl, ProtocolUDP>;
/// 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<TcpSocketImpl, ProtocolTCP>;
}
#[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 {