capture: Add packet capture infrastructure

crates/lib/Cargo.toml

@@ -23,6 +23,7 @@ services = []
 game_defs = ["dep:phf", "games"]
 serde = ["dep:serde", "serde/derive"]
 clap = ["dep:clap"]
+packet_capture = ["dep:pcap-file", "dep:pnet_packet"]
 
 [dependencies]
 byteorder = "1.5"
@@ -37,6 +38,9 @@ phf = { version = "0.11", optional = true, features = ["macros"] }
 
 clap = { version = "4.1.11", optional = true, features = ["derive"] }
 
+pcap-file = { version = "2.0", optional = true }
+pnet_packet = { version = "0.34", optional = true }
+
 # Examples
 [[example]]
 name = "minecraft"

crates/lib/src/capture.rs

@@ -0,0 +1,453 @@
+use std::io::Write;
+use std::net::{IpAddr, SocketAddr};
+
+use crate::GDResult;
+
+use pcap_file::pcapng::{blocks::enhanced_packet::EnhancedPacketOption, PcapNgBlock};
+use pnet_packet::{
+    ethernet::{EtherType, MutableEthernetPacket},
+    ip::{IpNextHeaderProtocol, IpNextHeaderProtocols},
+    ipv4::MutableIpv4Packet,
+    ipv6::MutableIpv6Packet,
+    tcp::{MutableTcpPacket, TcpFlags},
+    udp::MutableUdpPacket,
+    PacketSize,
+};
+
+/// Info about a packet we have sent or recieved.
+#[derive(Clone, Debug, PartialEq)]
+pub struct PacketInfo<'a> {
+    pub direction: PacketDirection,
+    pub protocol: PacketProtocol,
+    pub remote_address: &'a SocketAddr,
+    pub local_address: &'a SocketAddr,
+}
+
+/// The direction of a packet.
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum PacketDirection {
+    /// The packet is coming from us, destined for a server.
+    Send,
+    /// A server has sent this packet to us.
+    Receive,
+}
+
+/// The protocol of a packet.
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum PacketProtocol {
+    TCP,
+    UDP,
+}
+
+/// Trait for objects that can write packet captures.
+pub trait CaptureWriter {
+    fn write(&mut self, packet: &PacketInfo, data: &[u8]) -> crate::GDResult<()>;
+    fn new_connect(&mut self, packet: &PacketInfo) -> crate::GDResult<()>;
+    // TODO: Tcp FIN when socket ends
+}
+
+// Packet size constants
+const PACKET_SIZE: usize = 5012;
+const HEADER_SIZE_ETHERNET: usize = 14;
+const HEADER_SIZE_IP4: usize = 20;
+const HEADER_SIZE_IP6: usize = 40;
+const HEADER_SIZE_UDP: usize = 4;
+
+/// A writer that does nothing
+struct NullWriter;
+impl CaptureWriter for NullWriter {
+    fn write(&mut self, _: &PacketInfo, _: &[u8]) -> GDResult<()> { Ok(()) }
+    fn new_connect(&mut self, _: &PacketInfo) -> GDResult<()> { Ok(()) }
+}
+
+/// Writer that writes to pcap file
+struct PcapWriter<W: Write> {
+    writer: pcap_file::pcapng::PcapNgWriter<W>,
+    start_time: std::time::Instant,
+    send_seq: u32,
+    rec_seq: u32,
+    has_sent_handshake: bool,
+    stream_count: u32,
+}
+impl<W: Write> PcapWriter<W> {
+    fn new(writer: pcap_file::pcapng::PcapNgWriter<W>) -> Self {
+        Self {
+            writer,
+            start_time: std::time::Instant::now(),
+            send_seq: 0,
+            rec_seq: 0,
+            has_sent_handshake: false,
+            stream_count: 0,
+        }
+    }
+}
+
+impl<W: Write> CaptureWriter for PcapWriter<W> {
+    fn write(&mut self, info: &PacketInfo, data: &[u8]) -> GDResult<()> {
+        self.write_transport_packet(info, data);
+
+        Ok(())
+    }
+
+    fn new_connect(&mut self, packet: &PacketInfo) -> GDResult<()> {
+        match packet.protocol {
+            PacketProtocol::TCP => {
+                self.write_tcp_handshake(packet);
+            }
+            PacketProtocol::UDP => {}
+        }
+
+        self.stream_count = self.stream_count.wrapping_add(1);
+
+        Ok(())
+    }
+}
+
+impl<W: Write> PcapWriter<W> {
+    /// Encode the transport layer packet with a payload and write it.
+    fn write_transport_packet(&mut self, info: &PacketInfo, payload: &[u8]) {
+        let mut buf = vec![0; PACKET_SIZE - usize::max(HEADER_SIZE_IP4, HEADER_SIZE_IP6) - HEADER_SIZE_ETHERNET];
+
+        let (source_port, dest_port) = match info.direction {
+            PacketDirection::Send => (info.local_address.port(), info.remote_address.port()),
+            PacketDirection::Receive => (info.remote_address.port(), info.local_address.port()),
+        };
+
+        match info.protocol {
+            PacketProtocol::TCP => {
+                let buf_size = {
+                    let mut tcp = MutableTcpPacket::new(&mut buf).unwrap();
+                    tcp.set_source(source_port);
+                    tcp.set_destination(dest_port);
+                    tcp.set_payload(payload);
+                    tcp.set_data_offset(5);
+                    tcp.set_window(43440);
+                    match info.direction {
+                        PacketDirection::Send => {
+                            tcp.set_sequence(self.send_seq);
+                            tcp.set_acknowledgement(self.rec_seq);
+
+                            self.send_seq = self.send_seq.wrapping_add(payload.len() as u32);
+                        }
+                        PacketDirection::Receive => {
+                            tcp.set_sequence(self.rec_seq);
+                            tcp.set_acknowledgement(self.send_seq);
+
+                            self.rec_seq = self.rec_seq.wrapping_add(payload.len() as u32);
+                        }
+                    }
+                    tcp.set_flags(TcpFlags::PSH | TcpFlags::ACK);
+
+                    tcp.packet_size()
+                };
+
+                self.write_transport_payload(
+                    info,
+                    IpNextHeaderProtocols::Tcp,
+                    &buf[.. buf_size + payload.len()],
+                    vec![],
+                );
+
+                let mut info = info.clone();
+                let buf_size = {
+                    let mut tcp = MutableTcpPacket::new(&mut buf).unwrap();
+                    tcp.set_source(dest_port);
+                    tcp.set_destination(source_port);
+                    tcp.set_data_offset(5);
+                    tcp.set_window(43440);
+                    match &info.direction {
+                        PacketDirection::Send => {
+                            tcp.set_sequence(self.rec_seq);
+                            tcp.set_acknowledgement(self.send_seq);
+
+                            info.direction = PacketDirection::Receive;
+                        }
+                        PacketDirection::Receive => {
+                            tcp.set_sequence(self.send_seq);
+                            tcp.set_acknowledgement(self.rec_seq);
+
+                            info.direction = PacketDirection::Send;
+                        }
+                    }
+                    tcp.set_flags(TcpFlags::ACK);
+
+                    tcp.packet_size()
+                };
+
+                self.write_transport_payload(
+                    &info,
+                    IpNextHeaderProtocols::Tcp,
+                    &buf[.. buf_size],
+                    vec![EnhancedPacketOption::Comment("Generated TCP ack".into())],
+                );
+            }
+            PacketProtocol::UDP => {
+                let buf_size = {
+                    let mut udp = MutableUdpPacket::new(&mut buf).unwrap();
+                    udp.set_source(source_port);
+                    udp.set_destination(dest_port);
+                    udp.set_length((payload.len() + HEADER_SIZE_UDP) as u16);
+                    udp.set_payload(payload);
+
+                    udp.packet_size()
+                };
+
+                self.write_transport_payload(
+                    info,
+                    IpNextHeaderProtocols::Udp,
+                    &buf[.. buf_size + payload.len()],
+                    vec![],
+                );
+            }
+        }
+    }
+
+    /// Encode a network layer (IP) packet with a payload.
+    fn encode_ip_packet(
+        &self,
+        buf: &mut [u8],
+        info: &PacketInfo,
+        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 == PacketDirection::Send {
+                    (local_address, remote_address)
+                } else {
+                    (remote_address, local_address)
+                };
+
+                let header_size = HEADER_SIZE_IP4 + (32 / 8);
+
+                let mut ip = MutableIpv4Packet::new(buf).unwrap();
+                ip.set_version(4);
+                ip.set_total_length((payload.len() + header_size) as u16);
+                ip.set_next_level_protocol(protocol);
+                // https://en.wikipedia.org/wiki/Internet_Protocol_version_4#Total_Length
+
+                ip.set_header_length((header_size / 4) as u8);
+                ip.set_source(source);
+                ip.set_destination(destination);
+                ip.set_payload(payload);
+                ip.set_ttl(64);
+                ip.set_flags(pnet_packet::ipv4::Ipv4Flags::DontFragment);
+
+                let mut options_writer =
+                    pnet_packet::ipv4::MutableIpv4OptionPacket::new(ip.get_options_raw_mut()).unwrap();
+                options_writer.set_copied(1);
+                options_writer.set_class(0);
+                options_writer.set_number(pnet_packet::ipv4::Ipv4OptionNumbers::SID);
+                options_writer.set_length(&[4]);
+                options_writer.set_data(&(self.stream_count as u16).to_be_bytes());
+
+                ip.set_checksum(pnet_packet::ipv4::checksum(&ip.to_immutable()));
+
+                (ip.packet_size(), pnet_packet::ethernet::EtherTypes::Ipv4)
+            }
+            (IpAddr::V6(local_address), IpAddr::V6(remote_address)) => {
+                let (source, destination) = match info.direction {
+                    PacketDirection::Send => (local_address, remote_address),
+                    PacketDirection::Receive => (remote_address, local_address),
+                };
+
+                let mut ip = MutableIpv6Packet::new(buf).unwrap();
+                ip.set_version(6);
+                ip.set_payload_length(payload.len() as u16);
+                ip.set_next_header(protocol);
+                ip.set_source(source);
+                ip.set_destination(destination);
+                ip.set_hop_limit(64);
+                ip.set_payload(payload);
+                ip.set_flow_label(self.stream_count);
+
+                (ip.packet_size(), pnet_packet::ethernet::EtherTypes::Ipv6)
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    /// Encode a physical layer (ethernet) packet with a payload.
+    fn encode_ethernet_packet(
+        &self,
+        buf: &mut [u8],
+        ethertype: pnet_packet::ethernet::EtherType,
+        payload: &[u8],
+    ) -> usize {
+        let mut ethernet = MutableEthernetPacket::new(buf).unwrap();
+        ethernet.set_ethertype(ethertype);
+        ethernet.set_payload(payload);
+
+        ethernet.packet_size()
+    }
+
+    /// Write a TCP handshake.
+    fn write_tcp_handshake(&mut self, info: &PacketInfo) {
+        let (source_port, dest_port) = (info.local_address.port(), info.remote_address.port());
+
+        let mut info = info.clone();
+        info.direction = PacketDirection::Send;
+        let mut buf = vec![0; PACKET_SIZE];
+        // Add a generated comment to all packets
+        let options = vec![
+            pcap_file::pcapng::blocks::enhanced_packet::EnhancedPacketOption::Comment("Generated TCP handshake".into()),
+        ];
+
+        // SYN
+        let buf_size = {
+            let mut tcp = MutableTcpPacket::new(&mut buf).unwrap();
+            self.send_seq = 500;
+            tcp.set_sequence(self.send_seq);
+            tcp.set_flags(TcpFlags::SYN);
+            tcp.set_source(source_port);
+            tcp.set_destination(dest_port);
+            tcp.set_window(43440);
+            tcp.set_data_offset(5);
+
+            tcp.packet_size()
+        };
+        self.write_transport_payload(
+            &info,
+            IpNextHeaderProtocols::Tcp,
+            &buf[.. buf_size],
+            options.clone(),
+        );
+
+        // SYN + ACK
+        info.direction = PacketDirection::Receive;
+        let buf_size = {
+            let mut tcp = MutableTcpPacket::new(&mut buf).unwrap();
+            self.send_seq = self.send_seq.wrapping_add(1);
+            tcp.set_acknowledgement(self.send_seq);
+            self.rec_seq = 1000;
+            tcp.set_sequence(self.rec_seq);
+            tcp.set_flags(TcpFlags::SYN | TcpFlags::ACK);
+            tcp.set_source(dest_port);
+            tcp.set_destination(source_port);
+            tcp.set_window(43440);
+            tcp.set_data_offset(5);
+
+            tcp.packet_size()
+        };
+        self.write_transport_payload(
+            &info,
+            IpNextHeaderProtocols::Tcp,
+            &buf[.. buf_size],
+            options.clone(),
+        );
+
+        // ACK
+        info.direction = PacketDirection::Send;
+        let buf_size = {
+            let mut tcp = MutableTcpPacket::new(&mut buf).unwrap();
+            tcp.set_sequence(self.send_seq);
+            self.rec_seq = self.rec_seq.wrapping_add(1);
+            tcp.set_acknowledgement(self.rec_seq);
+            tcp.set_flags(TcpFlags::ACK);
+            tcp.set_source(source_port);
+            tcp.set_destination(dest_port);
+            tcp.set_window(43440);
+            tcp.set_data_offset(5);
+
+            tcp.packet_size()
+        };
+        self.write_transport_payload(
+            &info,
+            IpNextHeaderProtocols::Tcp,
+            &buf[.. buf_size],
+            options,
+        );
+
+        self.has_sent_handshake = true;
+    }
+
+    /// Take a transport layer packet as a buffer and write it after encoding
+    /// all the layers under it.
+    fn write_transport_payload(
+        &mut self,
+        info: &PacketInfo,
+        protocol: IpNextHeaderProtocol,
+        payload: &[u8],
+        options: Vec<pcap_file::pcapng::blocks::enhanced_packet::EnhancedPacketOption>,
+    ) {
+        let mut network_packet = vec![0; PACKET_SIZE - HEADER_SIZE_ETHERNET];
+        let (network_size, ethertype) = self.encode_ip_packet(&mut network_packet, info, protocol, payload);
+        let network_size = network_size + payload.len();
+        network_packet.truncate(network_size);
+
+        let mut physical_packet = vec![0; PACKET_SIZE];
+        let physical_size =
+            self.encode_ethernet_packet(&mut physical_packet, ethertype, &network_packet) + network_size;
+
+        physical_packet.truncate(physical_size);
+
+        self.writer
+            .write_block(
+                &pcap_file::pcapng::blocks::enhanced_packet::EnhancedPacketBlock {
+                    original_len: physical_size as u32,
+                    data: physical_packet.into(),
+                    interface_id: 0,
+                    timestamp: self.start_time.elapsed(),
+                    options,
+                }
+                .into_block(),
+            )
+            .unwrap();
+    }
+}
+
+/// Setup the static capture into a file or to nowhere.
+///
+/// This leaks the writer.
+///
+/// # Panics
+/// - If this is called more than once (OnceLock used internally).
+///
+/// # Safety
+/// The safety of this function has not been evaluated yet, and
+/// testing has only been done with limited CLI use cases.
+pub unsafe fn simple_setup_capture(file_name: Option<String>) {
+    let writer: Box<dyn CaptureWriter + Send + Sync> = if let Some(file_name) = file_name {
+        let file = std::fs::OpenOptions::new()
+            .create_new(true)
+            .write(true)
+            .open(file_name)
+            .unwrap();
+        let mut writer = pcap_file::pcapng::PcapNgWriter::new(file).unwrap();
+
+        // Write headers
+        writer
+            .write_block(
+                &pcap_file::pcapng::blocks::interface_description::InterfaceDescriptionBlock {
+                    linktype: pcap_file::DataLink::ETHERNET,
+                    snaplen: 0xFFFF,
+                    options: vec![],
+                }
+                .into_block(),
+            )
+            .unwrap();
+
+        let writer = PcapWriter::new(writer);
+        Box::new(writer)
+    } else {
+        Box::new(NullWriter)
+    };
+    setup_capture(writer);
+}
+
+/// Set a capture writer to handle packet send/recieve data.
+///
+/// This leaks the writer.
+///
+/// # Panics
+/// - If this is called more than once (OnceLock used internally).
+///
+/// # Safety
+/// The safety of this function has not been evaluated yet, and
+/// testing has only been done with limited CLI use cases.
+pub unsafe fn setup_capture(writer: Box<dyn CaptureWriter + Send + Sync>) {
+    // TODO: safety
+    unsafe {
+        crate::socket::capture::set_writer(writer);
+    }
+}

crates/lib/src/lib.rs

@@ -45,6 +45,9 @@ mod buffer;
 mod socket;
 mod utils;
 
+#[cfg(feature = "packet_capture")]
+pub mod capture;
+
 pub use errors::*;
 #[cfg(feature = "games")]
 pub use games::*;

crates/lib/src/socket.rs

@@ -30,12 +30,12 @@ pub trait Socket {
     fn local_addr(&self) -> std::io::Result<SocketAddr>;
 }
 
-pub struct TcpSocket {
+pub struct TcpSocketImpl {
     socket: net::TcpStream,
     address: SocketAddr,
 }
 
-impl Socket for TcpSocket {
+impl Socket for TcpSocketImpl {
     fn new(address: &SocketAddr, timeout_settings: &Option<TimeoutSettings>) -> GDResult<Self> {
         let socket = if let Some(timeout) = TimeoutSettings::get_connect_or_default(timeout_settings) {
             net::TcpStream::connect_timeout(address, timeout)
@@ -79,12 +79,12 @@ impl Socket for TcpSocket {
     fn local_addr(&self) -> std::io::Result<SocketAddr> { self.socket.local_addr() }
 }
 
-pub struct UdpSocket {
+pub struct UdpSocketImpl {
     socket: net::UdpSocket,
     address: SocketAddr,
 }
 
-impl Socket for UdpSocket {
+impl Socket for UdpSocketImpl {
     fn new(address: &SocketAddr, timeout_settings: &Option<TimeoutSettings>) -> GDResult<Self> {
         let socket = net::UdpSocket::bind("0.0.0.0:0").map_err(|e| SocketBind.context(e))?;
 
@@ -128,6 +128,135 @@ impl Socket for UdpSocket {
     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, sync::OnceLock};
+
+    use super::{Socket, TcpSocketImpl, UdpSocketImpl};
+
+    use crate::{
+        capture::{CaptureWriter, PacketDirection, PacketInfo, PacketProtocol},
+        protocols::types::TimeoutSettings,
+        GDResult,
+    };
+
+    static mut WRITER_CELL: OnceLock<Box<dyn CaptureWriter + Send + Sync>> = OnceLock::new();
+
+    /// Set a GLOBAL capture writer that handles writing all packet data.
+    ///
+    /// This is unsafe because the writer is a mutable static, the caller
+    /// is responsible for ensuring that the writer type is able to be stored
+    /// as such.
+    pub(crate) unsafe fn set_writer(writer: Box<dyn CaptureWriter + Send + Sync>) {
+        if WRITER_CELL.set(writer).is_err() {
+            panic!("Should only set writer once");
+        }
+    }
+
+    pub trait PacketProtocolProvider {
+        fn protocol() -> PacketProtocol;
+    }
+    pub struct PacketProtocolTCP;
+    impl PacketProtocolProvider for PacketProtocolTCP {
+        fn protocol() -> PacketProtocol { PacketProtocol::TCP }
+    }
+
+    pub struct PacketProtocolUDP;
+    impl PacketProtocolProvider for PacketProtocolUDP {
+        fn protocol() -> PacketProtocol { PacketProtocol::UDP }
+    }
+
+    /// A socket that allows capturing
+    #[derive(Clone, Debug)]
+    pub struct WrappedCaptureSocket<I, P> {
+        inner: I,
+        remote_address: SocketAddr,
+        _protocol: PhantomData<P>,
+    }
+
+    impl<I: Socket, P: PacketProtocolProvider> Socket for WrappedCaptureSocket<I, P> {
+        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 = PacketInfo {
+                direction: PacketDirection::Send,
+                protocol: P::protocol(),
+                remote_address: address,
+                local_address: &v.local_addr().unwrap(),
+            };
+            // TODO: Safety
+            unsafe {
+                if let Some(writer) = WRITER_CELL.get_mut() {
+                    writer.new_connect(&info)?;
+                }
+            }
+
+            Ok(v)
+        }
+
+        fn send(&mut self, data: &[u8]) -> crate::GDResult<()> {
+            let info = PacketInfo {
+                direction: PacketDirection::Send,
+                protocol: P::protocol(),
+                remote_address: &self.remote_address,
+                local_address: &self.local_addr().unwrap(),
+            };
+            // TODO: Safety
+            unsafe {
+                if let Some(writer) = WRITER_CELL.get_mut() {
+                    writer.write(&info, data)?;
+                }
+            }
+            self.inner.send(data)
+        }
+
+        fn receive(&mut self, size: Option<usize>) -> crate::GDResult<Vec<u8>> {
+            let data = self.inner.receive(size)?;
+            let info = PacketInfo {
+                direction: PacketDirection::Receive,
+                protocol: P::protocol(),
+                remote_address: &self.remote_address,
+                local_address: &self.local_addr().unwrap(),
+            };
+            // TODO: Safety
+            unsafe {
+                if let Some(writer) = WRITER_CELL.get_mut() {
+                    writer.write(&info, &data)?;
+                }
+            }
+            Ok(data)
+        }
+
+        fn apply_timeout(
+            &self,
+            timeout_settings: &Option<crate::protocols::types::TimeoutSettings>,
+        ) -> crate::GDResult<()> {
+            self.inner.apply_timeout(timeout_settings)
+        }
+        fn port(&self) -> u16 { self.inner.port() }
+        fn local_addr(&self) -> std::io::Result<SocketAddr> { self.inner.local_addr() }
+    }
+
+    pub type CapturedUdpSocket = WrappedCaptureSocket<UdpSocketImpl, PacketProtocolUDP>;
+    pub type CapturedTcpSocket = WrappedCaptureSocket<TcpSocketImpl, PacketProtocolTCP>;
+}
+
+#[cfg(not(feature = "packet_capture"))]
+pub type UdpSocket = UdpSocketImpl;
+#[cfg(not(feature = "packet_capture"))]
+pub type TcpSocket = TcpSocketImpl;
+
+#[cfg(feature = "packet_capture")]
+pub type UdpSocket = capture::CapturedUdpSocket;
+#[cfg(feature = "packet_capture")]
+pub type TcpSocket = capture::CapturedTcpSocket;
+
 #[cfg(test)]
 mod tests {
     use std::thread;