tribufu/rust-gamedig
1
0
Fork 0
mirror of https://github.com/tribufu/rust-gamedig synced 2026-06-01 09:42:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor: clean up

Browse source
This commit is contained in:
Cain 2024-01-13 09:18:56 +00:00 • committed by Douile
parent 3d47180e85
commit 0543cabce2
No known key found for this signature in database
GPG key ID: E048586A5FF6585C
6 changed files with 137 additions and 171 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1,25 +1,24 @@
pub(crate) mod packet;
mod pcap;
pub mod writer;
pub(crate) mod writer;
use pcap_file::pcapng::PcapNgBlock;
use writer::Writer;
use self::{pcap::Pcap, writer::Writer};
use pcap_file::pcapng::{blocks::interface_description::InterfaceDescriptionBlock, PcapNgBlock, PcapNgWriter};
use std::path::PathBuf;
use self::pcap::Pcap;
pub fn setup_capture(file_name: Option<String>) {
if let Some(file_name) = file_name {
pub fn setup_capture(file_path: Option<PathBuf>) {
if let Some(file_path) = file_path {
let file = std::fs::OpenOptions::new()
.create_new(true)
.write(true)
.open(file_name)
.open(file_path.with_extension("pcap"))
.unwrap();
let mut pcap_writer = pcap_file::pcapng::PcapNgWriter::new(file).unwrap();
let mut pcap_writer = PcapNgWriter::new(file).unwrap();
// Write headers
let _ = pcap_writer.write_block(
&pcap_file::pcapng::blocks::interface_description::InterfaceDescriptionBlock {
let _ = pcap_writer.write_block(
&InterfaceDescriptionBlock {
linktype: pcap_file::DataLink::ETHERNET,
snaplen: 0xFFFF,
options: vec![],
@ -30,6 +29,7 @@ pub fn setup_capture(file_name: Option<String>) {
let writer = Box::new(Pcap::new(pcap_writer));
attach(writer)
} else {
// If no file path is provided
// Do nothing
}
}
@ -37,7 +37,7 @@ pub fn setup_capture(file_name: Option<String>) {
/// Attaches a writer to the capture module.
///
/// # Errors
/// Returns an `io::Error` if the writer is already set.
/// Returns an Error if the writer is already set.
fn attach(writer: Box<dyn Writer + Send + Sync>) {
crate::socket::capture::set_writer(writer);
}

94
crates/lib/src/capture/packet.rs
View file

@ -30,7 +30,7 @@ pub(crate) enum Protocol {
}
/// Trait for handling different types of IP addresses (IPv4, IPv6).
pub trait IpAddress: Sized {
pub(crate) trait IpAddress: Sized {
/// Creates an instance from a standard `IpAddr`, returning `None` if the types are incompatible.
fn from_std(ip: IpAddr) -> Option<Self>;
}
@ -53,8 +53,7 @@ impl CapturePacket<'_> {
///
/// Returns:
/// - (u16, u16): Tuple of (source port, destination port).
#[allow(dead_code)]
pub(crate) const fn ports_by_direction(&self) -> (u16, u16) {
pub(super) const fn ports_by_direction(&self) -> (u16, u16) {
let (local, remote) = (self.local_address.port(), self.remote_address.port());
self.direction.order(local, remote)
}
@ -63,30 +62,18 @@ impl CapturePacket<'_> {
///
/// Returns:
/// - (IpAddr, IpAddr): Tuple of (local IP, remote IP).
#[allow(dead_code)]
pub(crate) fn ip_addr(&self) -> (IpAddr, IpAddr) {
pub(super) fn ip_addr(&self) -> (IpAddr, IpAddr) {
let (local, remote) = (self.local_address.ip(), self.remote_address.ip());
(local, remote)
}
/// Retrieves IP addresses based on the packet's direction.
///
/// Returns:
/// - (IpAddr, IpAddr): Tuple of (source IP, destination IP).
#[allow(dead_code)]
pub(crate) fn ip_addr_by_direction(&self) -> (IpAddr, IpAddr) {
let (local, remote) = self.ip_addr();
self.direction.order(local, remote)
}
/// Retrieves IP addresses of a specific type (IPv4 or IPv6) based on the packet's direction.
///
/// Panics if the IP type of the addresses does not match the requested type.
///
/// Returns:
/// - (T, T): Tuple of (source IP, destination IP) of the specified type in order.
#[allow(dead_code)]
pub(crate) fn ipvt_by_direction<T: IpAddress>(&self) -> (T, T) {
pub(super) fn ipvt_by_direction<T: IpAddress>(&self) -> (T, T) {
let (local, remote) = (
T::from_std(self.local_address.ip()).expect("Incorrect IP type for local address"),
T::from_std(self.remote_address.ip()).expect("Incorrect IP type for remote address"),
@ -101,7 +88,7 @@ impl Direction {
///
/// Returns:
/// - (T, T): Ordered tuple (source, destination).
pub(crate) const fn order<T>(&self, source: T, remote: T) -> (T, T) {
pub(self) const fn order<T>(&self, source: T, remote: T) -> (T, T) {
match self {
Direction::Send => (source, remote),
Direction::Receive => (remote, source),
@ -162,65 +149,33 @@ mod tests {
}
#[test]
fn test_ip_addr_by_direction_ipv4() {
fn test_ip_addr() {
let packet_send = CapturePacket {
direction: Direction::Send,
protocol: Protocol::UDP,
local_address: &socket_addr("10.0.0.1:3000"),
remote_address: &socket_addr("10.0.0.2:3001"),
protocol: Protocol::TCP,
local_address: &socket_addr("127.0.0.1:8080"),
remote_address: &socket_addr("192.168.1.1:80"),
};
let packet_receive = CapturePacket {
direction: Direction::Receive,
protocol: Protocol::UDP,
local_address: &socket_addr("10.0.0.1:3000"),
remote_address: &socket_addr("10.0.0.2:3001"),
protocol: Protocol::TCP,
local_address: &socket_addr("127.0.0.1:8080"),
remote_address: &socket_addr("192.168.1.1:80"),
};
assert_eq!(
packet_send.ip_addr_by_direction(),
packet_send.ip_addr(),
(
IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)),
IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2))
IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))
)
);
assert_eq!(
packet_receive.ip_addr_by_direction(),
packet_receive.ip_addr(),
(
IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)),
IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))
)
);
}
#[test]
fn test_ip_addr_by_direction_ipv6() {
let packet_send = CapturePacket {
direction: Direction::Send,
protocol: Protocol::UDP,
local_address: &socket_addr("[::1]:3000"),
remote_address: &socket_addr("[::2]:3001"),
};
let packet_receive = CapturePacket {
direction: Direction::Receive,
protocol: Protocol::UDP,
local_address: &socket_addr("[::1]:3000"),
remote_address: &socket_addr("[::2]:3001"),
};
assert_eq!(
packet_send.ip_addr_by_direction(),
(
IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)),
IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2))
)
);
assert_eq!(
packet_receive.ip_addr_by_direction(),
(
IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2)),
IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1))
IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))
)
);
}
@ -242,17 +197,4 @@ mod tests {
std::panic::catch_unwind(|| packet.ipvt_by_direction::<Ipv6Addr>());
assert!(ipv6_result.is_err());
}
#[test]
#[should_panic(expected = "Local and remote IP addresses must be of the same version")]
fn test_mismatched_ip_version_panic() {
let packet = CapturePacket {
direction: Direction::Send,
protocol: Protocol::UDP,
local_address: &socket_addr("127.0.0.1:8080"), // IPv4
remote_address: &socket_addr("[::1]:80"), // IPv6
};
packet.ip_addr_by_direction();
}
}

23
crates/lib/src/capture/pcap.rs
View file

@ -15,7 +15,6 @@ use super::packet::{
PACKET_SIZE,
};
const BUFFER_SIZE: usize = PACKET_SIZE
- (if HEADER_SIZE_IP4 > HEADER_SIZE_IP6 {
HEADER_SIZE_IP4
@ -48,10 +47,7 @@ impl<W: Write> Pcap<W> {
pub(crate) fn write_transport_packet(&mut self, info: &CapturePacket, payload: &[u8]) {
let mut buf = vec![0; BUFFER_SIZE];
let (source_port, dest_port) = match info.direction {
Direction::Send => (info.local_address.port(), info.remote_address.port()),
Direction::Receive => (info.remote_address.port(), info.local_address.port()),
};
let (source_port, dest_port) = info.ports_by_direction();
match info.protocol {
Protocol::TCP => {
@ -150,13 +146,9 @@ impl<W: Write> Pcap<W> {
protocol: IpNextHeaderProtocol,
payload: &[u8],
) -> (usize, EtherType) {
match (info.local_address.ip(), info.remote_address.ip()) {
(IpAddr::V4(local_address), IpAddr::V4(remote_address)) => {
let (source, destination) = if info.direction == Direction::Send {
(local_address, remote_address)
} else {
(remote_address, local_address)
};
match info.ip_addr() {
(IpAddr::V4(_), IpAddr::V4(_)) => {
let (source, destination) = info.ipvt_by_direction();
let header_size = HEADER_SIZE_IP4 + (32 / 8);
@ -185,11 +177,8 @@ impl<W: Write> Pcap<W> {
(ip.packet_size(), pnet_packet::ethernet::EtherTypes::Ipv4)
}
(IpAddr::V6(local_address), IpAddr::V6(remote_address)) => {
let (source, destination) = match info.direction {
Direction::Send => (local_address, remote_address),
Direction::Receive => (remote_address, local_address),
};
(IpAddr::V6(_), IpAddr::V6(_)) => {
let (source, destination) = info.ipvt_by_direction();
let mut ip = MutableIpv6Packet::new(buf).unwrap();
ip.set_version(6);

18
crates/lib/src/capture/writer.rs
View file

@ -1,22 +1,20 @@
use std::io::Write;
use std::{io::Write, sync::Mutex};
use crate::{
capture::packet::{CapturePacket, Protocol},
GDResult,
use super::{
packet::{CapturePacket, Protocol},
pcap::Pcap,
};
use super::pcap::Pcap;
use crate::GDResult;
use lazy_static::lazy_static;
use std::sync::Mutex;
lazy_static! {
pub(crate) static ref CAPTURE_WRITER: Mutex<Option<Box<dyn Writer + Send + Sync>>> = Mutex::new(None);
}
pub(crate) trait Writer {
fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> crate::GDResult<()>;
fn new_connect(&mut self, packet: &CapturePacket) -> crate::GDResult<()>;
fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> GDResult<()>;
fn new_connect(&mut self, packet: &CapturePacket) -> GDResult<()>;
//TODO: Close connection
}
impl<W: Write> Writer for Pcap<W> {

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

@ -4,10 +4,9 @@ use crate::{
GDResult,
};
use std::net::SocketAddr;
use std::{
io::{Read, Write},
net,
net::{self, SocketAddr},
};
const DEFAULT_PACKET_SIZE: usize = 1024;
@ -227,7 +226,7 @@ pub mod capture {
}
/// Represents the UDP protocol provider.
pub struct ProtocolUDP;
pub(crate) struct ProtocolUDP;
impl ProtocolProvider for ProtocolUDP {
fn protocol() -> Protocol {
Protocol::UDP
@ -240,7 +239,7 @@ pub mod capture {
/// * `I` - The inner socket type.
/// * `P` - The protocol provider.
#[derive(Clone, Debug)]
pub struct WrappedCaptureSocket<I, P> {
pub(crate) struct WrappedCaptureSocket<I, P> {
inner: I,
remote_address: SocketAddr,
_protocol: PhantomData<P>,
@ -292,7 +291,7 @@ pub mod capture {
///
/// # Returns
/// A result indicating success or error in sending data.
fn send(&mut self, data: &[u8]) -> crate::GDResult<()> {
fn send(&mut self, data: &[u8]) -> GDResult<()> {
let info = CapturePacket {
direction: Direction::Send,
protocol: P::protocol(),