merge: Douile fork + local (broken af)

crates/lib/src/capture/mod.rs

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

crates/lib/src/capture/packet.rs

@@ -0,0 +1,254 @@
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
+
+/// Size of a standard network packet.
+pub(crate) const PACKET_SIZE: usize = 5012;
+/// Size of an Ethernet header.
+pub(crate) const HEADER_SIZE_ETHERNET: usize = 14;
+/// Size of an IPv4 header.
+pub(crate) const HEADER_SIZE_IP4: usize = 20;
+/// Size of an IPv6 header.
+pub(crate) const HEADER_SIZE_IP6: usize = 40;
+/// Size of a UDP header.
+pub(crate) const HEADER_SIZE_UDP: usize = 4;
+
+/// Represents the direction of a network packet.
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub(crate) enum Direction {
+    /// Packet is outgoing (sent by us).
+    Send,
+    /// Packet is incoming (received by us).
+    Receive,
+}
+
+/// Defines the protocol of a network packet.
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub(crate) enum Protocol {
+    /// Transmission Control Protocol.
+    TCP,
+    /// User Datagram Protocol.
+    UDP,
+}
+
+/// Trait for handling different types of IP addresses (IPv4, IPv6).
+pub trait IpAddress: Sized {
+    /// Creates an instance from a standard `IpAddr`, returning `None` if the types are incompatible.
+    fn from_std(ip: IpAddr) -> Option<Self>;
+}
+
+/// Represents a captured network packet with metadata.
+#[derive(Clone, Debug, PartialEq)]
+pub(crate) struct CapturePacket<'a> {
+    /// Direction of the packet (Send/Receive).
+    pub(crate) direction: Direction,
+    /// Protocol of the packet (TCP/UDP).
+    pub(crate) protocol: Protocol,
+    /// Remote socket address.
+    pub(crate) remote_address: &'a SocketAddr,
+    /// Local socket address.
+    pub(crate) local_address: &'a SocketAddr,
+}
+
+impl CapturePacket<'_> {
+    /// Retrieves the local and remote ports based on the packet's direction.
+    ///
+    /// Returns:
+    /// - (u16, u16): Tuple of (source port, destination port).
+    pub(crate) const fn ports_by_direction(&self) -> (u16, u16) {
+        let (local, remote) = (self.local_address.port(), self.remote_address.port());
+        self.direction.order(local, remote)
+    }
+
+    /// Retrieves the local and remote IP addresses.
+    ///
+    /// Returns:
+    /// - (IpAddr, IpAddr): Tuple of (local IP, remote IP).
+    pub(crate) 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).
+    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.
+    pub(crate) 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"),
+        );
+
+        self.direction.order(local, remote)
+    }
+}
+
+impl Direction {
+    /// Orders two elements (source and destination) based on the packet's direction.
+    ///
+    /// Returns:
+    /// - (T, T): Ordered tuple (source, destination).
+    pub(crate) const fn order<T>(&self, source: T, remote: T) -> (T, T) {
+        match self {
+            Direction::Send => (source, remote),
+            Direction::Receive => (remote, source),
+        }
+    }
+}
+
+/// Implements the `IpAddress` trait for `Ipv4Addr`.
+impl IpAddress for Ipv4Addr {
+    /// Creates an `Ipv4Addr` from a standard `IpAddr`, if it's IPv4.
+    fn from_std(ip: IpAddr) -> Option<Self> {
+        match ip {
+            IpAddr::V4(ipv4) => Some(ipv4),
+            _ => None,
+        }
+    }
+}
+
+/// Implements the `IpAddress` trait for `Ipv6Addr`.
+impl IpAddress for Ipv6Addr {
+    /// Creates an `Ipv6Addr` from a standard `IpAddr`, if it's IPv6.
+    fn from_std(ip: IpAddr) -> Option<Self> {
+        match ip {
+            IpAddr::V6(ipv6) => Some(ipv6),
+            _ => None,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::str::FromStr;
+
+    // Helper function to create a SocketAddr from a string
+    fn socket_addr(addr: &str) -> SocketAddr {
+        SocketAddr::from_str(addr).unwrap()
+    }
+
+    #[test]
+    fn test_ports_by_direction() {
+        let packet_send = CapturePacket {
+            direction: Direction::Send,
+            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::TCP,
+            local_address: &socket_addr("127.0.0.1:8080"),
+            remote_address: &socket_addr("192.168.1.1:80"),
+        };
+
+        assert_eq!(packet_send.ports_by_direction(), (8080, 80));
+        assert_eq!(packet_receive.ports_by_direction(), (80, 8080));
+    }
+
+    #[test]
+    fn test_ip_addr_by_direction_ipv4() {
+        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"),
+        };
+
+        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"),
+        };
+
+        assert_eq!(
+            packet_send.ip_addr_by_direction(),
+            (
+                IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)),
+                IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2))
+            )
+        );
+        assert_eq!(
+            packet_receive.ip_addr_by_direction(),
+            (
+                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))
+            )
+        );
+    }
+
+    #[test]
+    fn test_ip_by_direction_type_specific() {
+        let packet = CapturePacket {
+            direction: Direction::Send,
+            protocol: Protocol::TCP,
+            local_address: &socket_addr("127.0.0.1:8080"),
+            remote_address: &socket_addr("192.168.1.1:80"),
+        };
+
+        let ipv4_result: Result<(Ipv4Addr, Ipv4Addr), _> =
+            std::panic::catch_unwind(|| packet.ipvt_by_direction::<Ipv4Addr>());
+        assert!(ipv4_result.is_ok());
+
+        let ipv6_result: Result<(Ipv6Addr, Ipv6Addr), _> =
+            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();
+    }
+}

crates/lib/src/capture/pcap.rs

@@ -0,0 +1,311 @@
+use pcap_file::pcapng::{
+    blocks::enhanced_packet::{EnhancedPacketBlock, EnhancedPacketOption},
+    PcapNgBlock, PcapNgWriter,
+};
+use pnet_packet::{
+    ethernet::{EtherType, EtherTypes, MutableEthernetPacket},
+    ip::{IpNextHeaderProtocol, IpNextHeaderProtocols},
+    ipv4::MutableIpv4Packet,
+    ipv6::MutableIpv6Packet,
+    tcp::{MutableTcpPacket, TcpFlags},
+    udp::MutableUdpPacket,
+    PacketSize,
+};
+use std::{
+    io::Write,
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
+    time::Instant,
+};
+
+use super::packet::{
+    CapturePacket, Direction, Protocol, HEADER_SIZE_ETHERNET, HEADER_SIZE_IP4, HEADER_SIZE_IP6, HEADER_SIZE_UDP,
+    PACKET_SIZE,
+};
+
+const DEFAULT_TTL: u8 = 64;
+const TCP_WINDOW_SIZE: u16 = 43440;
+const BUFFER_SIZE: usize = PACKET_SIZE
+    - (if HEADER_SIZE_IP4 > HEADER_SIZE_IP6 {
+        HEADER_SIZE_IP4
+    } else {
+        HEADER_SIZE_IP6
+    })
+    - HEADER_SIZE_ETHERNET;
+
+pub(crate) struct Pcap<W: Write> {
+    writer: PcapNgWriter<W>,
+    pub(crate) state: State,
+    buffer: Vec<u8>,
+}
+
+pub(crate) struct State {
+    pub(crate) start_time: Instant,
+    pub(crate) send_seq: u32,
+    pub(crate) rec_seq: u32,
+    pub(crate) has_sent_handshake: bool,
+    pub(crate) has_sent_fin: bool,
+    pub(crate) stream_count: u32,
+}
+
+impl<W: Write> Pcap<W> {
+    pub fn new(writer: PcapNgWriter<W>) -> Self {
+        Self {
+            writer,
+            state: State::default(),
+            buffer: vec![0; BUFFER_SIZE],
+        }
+    }
+
+    pub fn write_transport_packet(&mut self, info: &CapturePacket, payload: &[u8]) -> Result<(), std::io::Error> {
+        let (source_port, dest_port) = info.ports_by_direction();
+
+        match info.protocol {
+            Protocol::TCP => self.handle_tcp(info, payload, source_port, dest_port)?,
+            Protocol::UDP => self.handle_udp(info, payload, source_port, dest_port)?,
+        }
+
+        Ok(())
+    }
+
+    fn handle_tcp(
+        &mut self,
+        info: &CapturePacket,
+        payload: &[u8],
+        source_port: u16,
+        dest_port: u16,
+    ) -> Result<(), std::io::Error> {
+        let buf_size = self.setup_tcp_packet(info, payload, source_port, dest_port)?;
+        self.write_transport_payload(
+            info,
+            IpNextHeaderProtocols::Tcp,
+            &self.buffer[..buf_size + payload.len()],
+            vec![],
+        );
+
+        Ok(())
+    }
+
+    fn setup_tcp_packet(
+        &mut self,
+        info: &CapturePacket,
+        payload: &[u8],
+        source_port: u16,
+        dest_port: u16,
+    ) -> Result<usize, std::io::Error> {
+        let mut tcp = MutableTcpPacket::new(&mut self.buffer)
+            .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::Other, "Failed to create TCP packet"))?;
+        tcp.set_source(source_port);
+        tcp.set_destination(dest_port);
+        tcp.set_payload(payload);
+        tcp.set_data_offset(5);
+        tcp.set_window(TCP_WINDOW_SIZE);
+
+        // Set sequence and acknowledgement numbers
+        match info.direction {
+            Direction::Send => {
+                tcp.set_sequence(self.state.send_seq);
+                tcp.set_acknowledgement(self.state.rec_seq);
+                self.state.send_seq = self.state.send_seq.wrapping_add(payload.len() as u32);
+            }
+            Direction::Receive => {
+                tcp.set_sequence(self.state.rec_seq);
+                tcp.set_acknowledgement(self.state.send_seq);
+                self.state.rec_seq = self.state.rec_seq.wrapping_add(payload.len() as u32);
+            }
+        }
+        tcp.set_flags(TcpFlags::PSH | TcpFlags::ACK);
+
+        Ok(tcp.packet_size())
+    }
+
+    pub fn write_tcp_handshake(&mut self, info: &CapturePacket) {
+        // Initialize sequence numbers for demonstration purposes
+        self.state.send_seq = 500;
+        self.state.rec_seq = 1000;
+
+        // Common setup for TCP handshake packets
+        let mut tcp_handshake_packet =
+            |info: &CapturePacket, direction: Direction, flags: u8| -> Result<(), std::io::Error> {
+                let (source_port, dest_port) = info.ports_by_direction();
+                let adjusted_info = CapturePacket {
+                    direction,
+                    ..info.clone()
+                };
+                self.setup_tcp_packet(&adjusted_info, &[], source_port, dest_port)?;
+                Ok(self.write_transport_payload(
+                    &adjusted_info,
+                    IpNextHeaderProtocols::Tcp,
+                    &self.buffer,
+                    vec![EnhancedPacketOption::Comment(
+                        format!(
+                            "Generated TCP {}",
+                            match flags {
+                                TcpFlags::SYN => "SYN",
+                                TcpFlags::SYN | TcpFlags::ACK => "SYN-ACK",
+                                TcpFlags::ACK => "ACK",
+                            }
+                        )
+                        .into(),
+                    )],
+                ))
+            };
+
+        // Send SYN
+        tcp_handshake_packet(info, Direction::Send, TcpFlags::SYN);
+
+        // Send SYN-ACK
+        self.state.send_seq = self.state.send_seq.wrapping_add(1); // Update sequence number after SYN
+        tcp_handshake_packet(info, Direction::Receive, TcpFlags::SYN | TcpFlags::ACK);
+
+        // Send ACK
+        self.state.rec_seq = self.state.rec_seq.wrapping_add(1); // Update sequence number after SYN-ACK
+        tcp_handshake_packet(info, Direction::Send, TcpFlags::ACK);
+    }
+
+    fn handle_udp(
+        &mut self,
+        info: &CapturePacket,
+        payload: &[u8],
+        source_port: u16,
+        dest_port: u16,
+    ) -> Result<(), std::io::Error> {
+        let mut udp = MutableUdpPacket::new(&mut self.buffer)
+            .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::Other, "Failed to create UDP packet"))?;
+        udp.set_source(source_port);
+        udp.set_destination(dest_port);
+        udp.set_length((payload.len() + HEADER_SIZE_UDP) as u16);
+        udp.set_payload(payload);
+
+        let buf_size = udp.packet_size();
+        self.write_transport_payload(
+            info,
+            IpNextHeaderProtocols::Udp,
+            &self.buffer[..buf_size + payload.len()],
+            vec![],
+        );
+
+        Ok(())
+    }
+
+    fn write_transport_payload(
+        &mut self,
+        info: &CapturePacket,
+        protocol: IpNextHeaderProtocol,
+        payload: &[u8],
+        options: Vec<EnhancedPacketOption>,
+    ) {
+        let network_packet_size = self.encode_ip_packet(info, protocol, payload).unwrap().0;
+        let ethertype = self.encode_ip_packet(info, protocol, payload).unwrap().1;
+        let ethernet_packet_size = self.encode_ethernet_packet(info, ethertype, &self.buffer[..network_packet_size]).unwrap();
+
+        let enhanced_packet_block = EnhancedPacketBlock {
+            original_len: ethernet_packet_size as u32,
+            data: self.buffer[..ethernet_packet_size].to_vec().into(),
+            interface_id: 0,
+            timestamp: self.state.start_time.elapsed(),
+            options,
+        };
+
+        self.writer.write_block(&enhanced_packet_block.into_block());
+    }
+
+    fn encode_ethernet_packet(
+        &mut self,
+        info: &CapturePacket,
+        ethertype: EtherType,
+        payload: &[u8],
+    ) -> Result<usize, std::io::Error> {
+        let mut ethernet_packet = MutableEthernetPacket::new(&mut self.buffer).ok_or_else(|| {
+            std::io::Error::new(
+                std::io::ErrorKind::Other,
+                "Failed to create Ethernet packet",
+            )
+        })?;
+
+        ethernet_packet.set_ethertype(ethertype);
+        ethernet_packet.set_payload(payload);
+
+        Ok(ethernet_packet.packet_size())
+    }
+
+    fn encode_ip_packet(
+        &mut self,
+        info: &CapturePacket,
+        protocol: IpNextHeaderProtocol,
+        payload: &[u8],
+    ) -> Result<(usize, EtherType), std::io::Error> {
+        match info.ip_addr() {
+            (IpAddr::V4(_), IpAddr::V4(_)) => {
+                let (source, destination) = info.ipvt_by_direction();
+
+                let mut ip_packet = MutableIpv4Packet::new(&mut self.buffer)
+                    .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::Other, "Failed to create IPv4 packet"))?;
+                self.set_ipv4_packet_fields(&mut ip_packet, source, destination, payload, protocol);
+                ip_packet.set_checksum(pnet_packet::ipv4::checksum(&ip_packet.to_immutable()));
+
+                Ok((ip_packet.packet_size(), EtherTypes::Ipv4))
+            }
+            (IpAddr::V6(_), IpAddr::V6(_)) => {
+                let (source, destination) = info.ipvt_by_direction();
+
+                let mut ip_packet = MutableIpv6Packet::new(&mut self.buffer)
+                    .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::Other, "Failed to create IPv6 packet"))?;
+                self.set_ipv6_packet_fields(&mut ip_packet, source, destination, payload, protocol);
+
+                Ok((ip_packet.packet_size(), EtherTypes::Ipv6))
+            }
+            _ => Err(std::io::Error::new(
+                std::io::ErrorKind::Other,
+                "Unsupported or mismatched IP address types",
+            )),
+        }
+    }
+
+    fn set_ipv4_packet_fields(
+        &mut self,
+        ip_packet: &mut MutableIpv4Packet,
+        source: Ipv4Addr,
+        destination: Ipv4Addr,
+        payload: &[u8],
+        protocol: IpNextHeaderProtocol,
+    ) {
+        ip_packet.set_version(4);
+        ip_packet.set_header_length(5); // No options
+        ip_packet.set_total_length((payload.len() + HEADER_SIZE_IP4) as u16);
+        ip_packet.set_next_level_protocol(protocol);
+        ip_packet.set_source(source);
+        ip_packet.set_destination(destination);
+        ip_packet.set_ttl(DEFAULT_TTL);
+        ip_packet.set_payload(payload);
+    }
+
+    fn set_ipv6_packet_fields(
+        &mut self,
+        ip_packet: &mut MutableIpv6Packet,
+        source: Ipv6Addr,
+        destination: Ipv6Addr,
+        payload: &[u8],
+        protocol: IpNextHeaderProtocol,
+    ) {
+        ip_packet.set_version(6);
+        ip_packet.set_payload_length(payload.len() as u16);
+        ip_packet.set_next_header(protocol);
+        ip_packet.set_source(source);
+        ip_packet.set_destination(destination);
+        ip_packet.set_hop_limit(DEFAULT_TTL);
+        ip_packet.set_payload(payload);
+    }
+}
+
+impl Default for State {
+    fn default() -> Self {
+        Self {
+            start_time: Instant::now(),
+            send_seq: 0,
+            rec_seq: 0,
+            has_sent_handshake: false,
+            has_sent_fin: false,
+            stream_count: 0,
+        }
+    }
+}

crates/lib/src/capture/writer.rs

@@ -0,0 +1,41 @@
+use std::io::Write;
+
+use crate::{
+    capture::packet::{CapturePacket, Protocol},
+    GDResult,
+};
+
+use super::pcap::Pcap;
+
+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 trait Writer {
+    fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> crate::GDResult<()>;
+    fn new_connect(&mut self, packet: &CapturePacket) -> crate::GDResult<()>;
+}
+
+impl<W: Write> Writer for Pcap<W> {
+    fn write(&mut self, info: &CapturePacket, data: &[u8]) -> GDResult<()> {
+        self.write_transport_packet(info, data);
+
+        Ok(())
+    }
+
+    fn new_connect(&mut self, packet: &CapturePacket) -> GDResult<()> {
+        match packet.protocol {
+            Protocol::TCP => {
+                self.write_tcp_handshake(packet);
+            }
+            Protocol::UDP => {}
+        }
+
+        self.state.stream_count = self.state.stream_count.wrapping_add(1);
+
+        Ok(())
+    }
+}