From 3b1edd8e3dce0bad933249c48781caa989353f91 Mon Sep 17 00:00:00 2001 From: Douile Date: Wed, 22 Nov 2023 01:11:50 +0000 Subject: [PATCH 01/35] socket: Add method to get local address --- crates/lib/src/socket.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index 8ae6bef..e49872b 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -24,7 +24,10 @@ pub trait Socket { fn send(&mut self, data: &[u8]) -> GDResult<()>; fn receive(&mut self, size: Option) -> GDResult>; + /// Get the remote port fn port(&self) -> u16; + /// Get the local SocketAddr + fn local_addr(&self) -> std::io::Result; } pub struct TcpSocket { @@ -73,6 +76,7 @@ impl Socket for TcpSocket { } fn port(&self) -> u16 { self.address.port() } + fn local_addr(&self) -> std::io::Result { self.socket.local_addr() } } pub struct UdpSocket { @@ -121,6 +125,7 @@ impl Socket for UdpSocket { } fn port(&self) -> u16 { self.address.port() } + fn local_addr(&self) -> std::io::Result { self.socket.local_addr() } } #[cfg(test)] From abbcae618f95289021b6ad623dc1a366ae653087 Mon Sep 17 00:00:00 2001 From: Douile Date: Wed, 22 Nov 2023 02:20:03 +0000 Subject: [PATCH 02/35] capture: Add packet capture infrastructure --- crates/lib/Cargo.toml | 4 + crates/lib/src/capture.rs | 453 ++++++++++++++++++++++++++++++++++++++ crates/lib/src/lib.rs | 3 + crates/lib/src/socket.rs | 137 +++++++++++- 4 files changed, 593 insertions(+), 4 deletions(-) create mode 100644 crates/lib/src/capture.rs diff --git a/crates/lib/Cargo.toml b/crates/lib/Cargo.toml index 5e24b5f..6b17d0c 100644 --- a/crates/lib/Cargo.toml +++ b/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" diff --git a/crates/lib/src/capture.rs b/crates/lib/src/capture.rs new file mode 100644 index 0000000..9c7a7fa --- /dev/null +++ b/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 { + writer: pcap_file::pcapng::PcapNgWriter, + start_time: std::time::Instant, + send_seq: u32, + rec_seq: u32, + has_sent_handshake: bool, + stream_count: u32, +} +impl PcapWriter { + fn new(writer: pcap_file::pcapng::PcapNgWriter) -> Self { + Self { + writer, + start_time: std::time::Instant::now(), + send_seq: 0, + rec_seq: 0, + has_sent_handshake: false, + stream_count: 0, + } + } +} + +impl CaptureWriter for PcapWriter { + 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 PcapWriter { + /// 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, + ) { + 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) { + let writer: Box = 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) { + // TODO: safety + unsafe { + crate::socket::capture::set_writer(writer); + } +} diff --git a/crates/lib/src/lib.rs b/crates/lib/src/lib.rs index 7cafba7..38ec19c 100644 --- a/crates/lib/src/lib.rs +++ b/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::*; diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index e49872b..7dfaf19 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -30,12 +30,12 @@ pub trait Socket { fn local_addr(&self) -> std::io::Result; } -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) -> GDResult { 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 { 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) -> GDResult { 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 { 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> = 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) { + 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 { + inner: I, + remote_address: SocketAddr, + _protocol: PhantomData

, + } + + impl Socket for WrappedCaptureSocket { + 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 = 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) -> crate::GDResult> { + 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::GDResult<()> { + self.inner.apply_timeout(timeout_settings) + } + fn port(&self) -> u16 { self.inner.port() } + fn local_addr(&self) -> std::io::Result { self.inner.local_addr() } + } + + pub type CapturedUdpSocket = WrappedCaptureSocket; + pub type CapturedTcpSocket = WrappedCaptureSocket; +} + +#[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; From 8d17ca4e48a1fe9078666028629f865789f4931c Mon Sep 17 00:00:00 2001 From: Douile Date: Wed, 22 Nov 2023 02:23:32 +0000 Subject: [PATCH 03/35] cli: Add option to create a packet capture --- crates/cli/Cargo.toml | 3 ++- crates/cli/src/main.rs | 9 +++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index d74c1d0..dd4ed97 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -11,8 +11,9 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [features] -default = ["json"] +default = ["json", "packet_capture"] json = ["dep:serde", "dep:serde_json", "gamedig/serde"] +packet_capture = ["gamedig/packet_capture"] [dependencies] clap = { version = "4.1.11", features = ["derive"] } diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 526f227..6a6c6d5 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -38,6 +38,10 @@ struct Cli { #[arg(short, long, default_value = "generic")] output_mode: OutputMode, + #[cfg(feature = "packet_capture")] + #[arg(short, long)] + capture: Option, + /// Optional timeout settings for the server query. #[command(flatten, next_help_heading = "Timeouts")] timeout_settings: Option, @@ -177,6 +181,11 @@ fn main() -> Result<()> { // Resolve the IP address let ip = resolve_ip_or_domain(&args.ip, &mut extra_options)?; + #[cfg(feature = "packet_capture")] + unsafe { + gamedig::capture::simple_setup_capture(args.capture.clone()); + } + // Query the server using game definition, parsed IP, and user command line // flags. let result = query_with_timeout_and_extra_settings(game, &ip, args.port, args.timeout_settings, extra_options)?; From 49096e46bb9e4f432580e37c7a6af278691c41be Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 12 Jan 2024 07:57:16 +0000 Subject: [PATCH 04/35] merge: Douile fork + local (broken af) --- crates/cli/src/main.rs | 4 +- crates/lib/Cargo.toml | 5 +- crates/lib/src/capture.rs | 453 ------------------------------- crates/lib/src/capture/mod.rs | 43 +++ crates/lib/src/capture/packet.rs | 254 +++++++++++++++++ crates/lib/src/capture/pcap.rs | 311 +++++++++++++++++++++ crates/lib/src/capture/writer.rs | 41 +++ crates/lib/src/lib.rs | 2 +- crates/lib/src/socket.rs | 258 ++++++++++++++---- 9 files changed, 852 insertions(+), 519 deletions(-) delete mode 100644 crates/lib/src/capture.rs create mode 100644 crates/lib/src/capture/mod.rs create mode 100644 crates/lib/src/capture/packet.rs create mode 100644 crates/lib/src/capture/pcap.rs create mode 100644 crates/lib/src/capture/writer.rs diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 6a6c6d5..788bba7 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -182,9 +182,7 @@ fn main() -> Result<()> { let ip = resolve_ip_or_domain(&args.ip, &mut extra_options)?; #[cfg(feature = "packet_capture")] - unsafe { - gamedig::capture::simple_setup_capture(args.capture.clone()); - } + gamedig::capture::setup_capture(args.capture); // Query the server using game definition, parsed IP, and user command line // flags. diff --git a/crates/lib/Cargo.toml b/crates/lib/Cargo.toml index 6b17d0c..8cc7cdb 100644 --- a/crates/lib/Cargo.toml +++ b/crates/lib/Cargo.toml @@ -17,13 +17,13 @@ rust-version = "1.65.0" categories = ["parser-implementations", "parsing", "network-programming", "encoding"] [features] -default = ["games", "services", "game_defs"] +default = ["games", "services", "game_defs", "packet_capture",] games = [] services = [] game_defs = ["dep:phf", "games"] serde = ["dep:serde", "serde/derive"] clap = ["dep:clap"] -packet_capture = ["dep:pcap-file", "dep:pnet_packet"] +packet_capture = ["dep:pcap-file", "dep:pnet_packet", "dep:lazy_static"] [dependencies] byteorder = "1.5" @@ -40,6 +40,7 @@ clap = { version = "4.1.11", optional = true, features = ["derive"] } pcap-file = { version = "2.0", optional = true } pnet_packet = { version = "0.34", optional = true } +lazy_static = { version = "1.4", optional = true } # Examples [[example]] diff --git a/crates/lib/src/capture.rs b/crates/lib/src/capture.rs deleted file mode 100644 index 9c7a7fa..0000000 --- a/crates/lib/src/capture.rs +++ /dev/null @@ -1,453 +0,0 @@ -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 { - writer: pcap_file::pcapng::PcapNgWriter, - start_time: std::time::Instant, - send_seq: u32, - rec_seq: u32, - has_sent_handshake: bool, - stream_count: u32, -} -impl PcapWriter { - fn new(writer: pcap_file::pcapng::PcapNgWriter) -> Self { - Self { - writer, - start_time: std::time::Instant::now(), - send_seq: 0, - rec_seq: 0, - has_sent_handshake: false, - stream_count: 0, - } - } -} - -impl CaptureWriter for PcapWriter { - 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 PcapWriter { - /// 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, - ) { - 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) { - let writer: Box = 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) { - // TODO: safety - unsafe { - crate::socket::capture::set_writer(writer); - } -} diff --git a/crates/lib/src/capture/mod.rs b/crates/lib/src/capture/mod.rs new file mode 100644 index 0000000..b659dca --- /dev/null +++ b/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) { + 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) { + crate::socket::capture::set_writer(writer); +} diff --git a/crates/lib/src/capture/packet.rs b/crates/lib/src/capture/packet.rs new file mode 100644 index 0000000..4f726fc --- /dev/null +++ b/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; +} + +/// 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(&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(&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 { + 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 { + 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::()); + assert!(ipv4_result.is_ok()); + + let ipv6_result: Result<(Ipv6Addr, Ipv6Addr), _> = + std::panic::catch_unwind(|| packet.ipvt_by_direction::()); + 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(); + } +} diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs new file mode 100644 index 0000000..054c8e3 --- /dev/null +++ b/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 { + writer: PcapNgWriter, + pub(crate) state: State, + buffer: Vec, +} + +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 Pcap { + pub fn new(writer: PcapNgWriter) -> 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 { + 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, + ) { + 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 { + 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, + } + } +} diff --git a/crates/lib/src/capture/writer.rs b/crates/lib/src/capture/writer.rs new file mode 100644 index 0000000..012b700 --- /dev/null +++ b/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>> = 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 Writer for Pcap { + 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(()) + } +} diff --git a/crates/lib/src/lib.rs b/crates/lib/src/lib.rs index 38ec19c..aeb3be2 100644 --- a/crates/lib/src/lib.rs +++ b/crates/lib/src/lib.rs @@ -46,7 +46,7 @@ mod socket; mod utils; #[cfg(feature = "packet_capture")] -pub mod capture; +pub(crate) mod capture; pub use errors::*; #[cfg(feature = "games")] diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index 7dfaf19..f0fe3d3 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -12,26 +12,65 @@ use std::{ const DEFAULT_PACKET_SIZE: usize = 1024; +/// A trait defining the basic functionalities of a network socket. pub trait Socket { - /// Create a new socket and connect to the remote address (if required). + /// Create a new socket and connect to the remote address. /// - /// Calls [Self::apply_timeout] with the given timeout settings. + /// # Arguments + /// * `address` - The address to connect the socket to. + /// * `timeout_settings` - Optional timeout settings for the socket. + /// + /// # Returns + /// A result containing the socket instance or an error. fn new(address: &SocketAddr, timeout_settings: &Option) -> GDResult - where Self: Sized; + where + Self: Sized; + /// Apply read and write timeouts to the socket. + /// + /// # 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) -> GDResult<()>; + /// Send data over the socket. + /// + /// # Arguments + /// * `data` - Data to be sent. + /// + /// # Returns + /// A result indicating success or error in sending data. fn send(&mut self, data: &[u8]) -> GDResult<()>; + + /// Receive data from the socket. + /// + /// # Arguments + /// * `size` - Optional size of data to receive. + /// + /// # Returns + /// A result containing received data or an error. fn receive(&mut self, size: Option) -> GDResult>; - /// Get the remote port + /// Get the remote port of the socket. + /// + /// # Returns + /// The port number. fn port(&self) -> u16; - /// Get the local SocketAddr + + /// Get the local SocketAddr. + /// + /// # Returns + /// The local SocketAddr. fn local_addr(&self) -> std::io::Result; } +/// Implementation of a TCP socket. pub struct TcpSocketImpl { + /// The underlying TCP socket stream. socket: net::TcpStream, + /// The address of the remote host. address: SocketAddr, } @@ -75,12 +114,19 @@ impl Socket for TcpSocketImpl { Ok(buf) } - fn port(&self) -> u16 { self.address.port() } - fn local_addr(&self) -> std::io::Result { self.socket.local_addr() } + fn port(&self) -> u16 { + self.address.port() + } + fn local_addr(&self) -> std::io::Result { + self.socket.local_addr() + } } +/// Implementation of a UDP socket. pub struct UdpSocketImpl { + /// The underlying UDP socket. socket: net::UdpSocket, + /// The address of the remote host. address: SocketAddr, } @@ -121,53 +167,78 @@ impl Socket for UdpSocketImpl { .recv_from(&mut buf) .map_err(|e| PacketReceive.context(e))?; - Ok(buf[.. number_of_bytes_received].to_vec()) + Ok(buf[..number_of_bytes_received].to_vec()) } - fn port(&self) -> u16 { self.address.port() } - fn local_addr(&self) -> std::io::Result { self.socket.local_addr() } + fn port(&self) -> u16 { + self.address.port() + } + 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, sync::OnceLock}; + use std::{marker::PhantomData, net::SocketAddr}; use super::{Socket, TcpSocketImpl, UdpSocketImpl}; use crate::{ - capture::{CaptureWriter, PacketDirection, PacketInfo, PacketProtocol}, + capture::{ + packet::CapturePacket, + packet::{Direction, Protocol}, + writer::{Writer, CAPTURE_WRITER}, + }, protocols::types::TimeoutSettings, GDResult, }; - static mut WRITER_CELL: OnceLock> = OnceLock::new(); - - /// Set a GLOBAL capture writer that handles writing all packet data. + /// Sets a global capture writer for handling 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) { - if WRITER_CELL.set(writer).is_err() { - panic!("Should only set writer once"); + /// # 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 trait ProtocolProvider { + /// Returns the protocol used by the provider. + fn protocol() -> Protocol; + } + + /// Represents the TCP protocol provider. + pub struct ProtocolTCP; + impl ProtocolProvider for ProtocolTCP { + fn protocol() -> Protocol { + Protocol::TCP } } - pub trait PacketProtocolProvider { - fn protocol() -> PacketProtocol; - } - pub struct PacketProtocolTCP; - impl PacketProtocolProvider for PacketProtocolTCP { - fn protocol() -> PacketProtocol { PacketProtocol::TCP } + /// Represents the UDP protocol provider. + pub struct ProtocolUDP; + impl ProtocolProvider for ProtocolUDP { + fn protocol() -> Protocol { + Protocol::UDP + } } - pub struct PacketProtocolUDP; - impl PacketProtocolProvider for PacketProtocolUDP { - fn protocol() -> PacketProtocol { PacketProtocol::UDP } - } - - /// A socket that allows capturing + /// A socket wrapper that allows capturing packets. + /// + /// # Type parameters + /// * `I` - The inner socket type. + /// * `P` - The protocol provider. #[derive(Clone, Debug)] pub struct WrappedCaptureSocket { inner: I, @@ -175,76 +246,143 @@ pub mod capture { _protocol: PhantomData

, } - impl Socket for WrappedCaptureSocket { + 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 { + where + Self: Sized, + { let v = Self { inner: I::new(address, timeout_settings)?, remote_address: *address, _protocol: PhantomData, }; - let info = PacketInfo { - direction: PacketDirection::Send, + let info = CapturePacket { + direction: Direction::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)?; - } + + 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]) -> crate::GDResult<()> { - let info = PacketInfo { - direction: PacketDirection::Send, + let info = CapturePacket { + direction: Direction::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)?; - } + + 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 = PacketInfo { - direction: PacketDirection::Receive, + let info = CapturePacket { + direction: Direction::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)?; - } + + 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) } - fn port(&self) -> u16 { self.inner.port() } - fn local_addr(&self) -> std::io::Result { self.inner.local_addr() } + + /// 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() + } } - pub type CapturedUdpSocket = WrappedCaptureSocket; - pub type CapturedTcpSocket = WrappedCaptureSocket; + /// 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 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 type CapturedTcpSocket = WrappedCaptureSocket; } #[cfg(not(feature = "packet_capture"))] From 3d47180e85db58d5b6dcd08dc8c7f448a504c4bf Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Sat, 13 Jan 2024 03:22:50 +0000 Subject: [PATCH 05/35] refactor: backport pcap impl --- crates/cli/src/main.rs | 2 +- crates/lib/src/capture/mod.rs | 2 +- crates/lib/src/capture/packet.rs | 4 + crates/lib/src/capture/pcap.rs | 483 +++++++++++++++++-------------- crates/lib/src/capture/writer.rs | 2 +- crates/lib/src/lib.rs | 2 +- crates/lib/src/socket.rs | 12 +- 7 files changed, 272 insertions(+), 235 deletions(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 788bba7..969e5e3 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -182,7 +182,7 @@ fn main() -> Result<()> { let ip = resolve_ip_or_domain(&args.ip, &mut extra_options)?; #[cfg(feature = "packet_capture")] - gamedig::capture::setup_capture(args.capture); + gamedig::capture::setup_capture(args.capture.clone()); // Query the server using game definition, parsed IP, and user command line // flags. diff --git a/crates/lib/src/capture/mod.rs b/crates/lib/src/capture/mod.rs index b659dca..6cfd8ad 100644 --- a/crates/lib/src/capture/mod.rs +++ b/crates/lib/src/capture/mod.rs @@ -18,7 +18,7 @@ pub fn setup_capture(file_name: Option) { let mut pcap_writer = pcap_file::pcapng::PcapNgWriter::new(file).unwrap(); // Write headers - pcap_writer.write_block( + let _ = pcap_writer.write_block( &pcap_file::pcapng::blocks::interface_description::InterfaceDescriptionBlock { linktype: pcap_file::DataLink::ETHERNET, snaplen: 0xFFFF, diff --git a/crates/lib/src/capture/packet.rs b/crates/lib/src/capture/packet.rs index 4f726fc..8ad2917 100644 --- a/crates/lib/src/capture/packet.rs +++ b/crates/lib/src/capture/packet.rs @@ -53,6 +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) { let (local, remote) = (self.local_address.port(), self.remote_address.port()); self.direction.order(local, remote) @@ -62,6 +63,7 @@ impl CapturePacket<'_> { /// /// Returns: /// - (IpAddr, IpAddr): Tuple of (local IP, remote IP). + #[allow(dead_code)] pub(crate) fn ip_addr(&self) -> (IpAddr, IpAddr) { let (local, remote) = (self.local_address.ip(), self.remote_address.ip()); (local, remote) @@ -71,6 +73,7 @@ impl CapturePacket<'_> { /// /// 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) @@ -82,6 +85,7 @@ impl CapturePacket<'_> { /// /// Returns: /// - (T, T): Tuple of (source IP, destination IP) of the specified type in order. + #[allow(dead_code)] pub(crate) fn ipvt_by_direction(&self) -> (T, T) { let (local, remote) = ( T::from_std(self.local_address.ip()).expect("Incorrect IP type for local address"), diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs index 054c8e3..be65968 100644 --- a/crates/lib/src/capture/pcap.rs +++ b/crates/lib/src/capture/pcap.rs @@ -1,9 +1,6 @@ -use pcap_file::pcapng::{ - blocks::enhanced_packet::{EnhancedPacketBlock, EnhancedPacketOption}, - PcapNgBlock, PcapNgWriter, -}; +use pcap_file::pcapng::{blocks::enhanced_packet::EnhancedPacketOption, PcapNgBlock, PcapNgWriter}; use pnet_packet::{ - ethernet::{EtherType, EtherTypes, MutableEthernetPacket}, + ethernet::{EtherType, MutableEthernetPacket}, ip::{IpNextHeaderProtocol, IpNextHeaderProtocols}, ipv4::MutableIpv4Packet, ipv6::MutableIpv6Packet, @@ -11,19 +8,14 @@ use pnet_packet::{ udp::MutableUdpPacket, PacketSize, }; -use std::{ - io::Write, - net::{IpAddr, Ipv4Addr, Ipv6Addr}, - time::Instant, -}; +use std::{io::Write, net::IpAddr, 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 @@ -35,7 +27,6 @@ const BUFFER_SIZE: usize = PACKET_SIZE pub(crate) struct Pcap { writer: PcapNgWriter, pub(crate) state: State, - buffer: Vec, } pub(crate) struct State { @@ -43,257 +34,300 @@ pub(crate) struct State { 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 Pcap { - pub fn new(writer: PcapNgWriter) -> Self { + pub(crate) fn new(writer: PcapNgWriter) -> 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(); + 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()), + }; 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)?, - } + Protocol::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 { + Direction::Send => { + tcp.set_sequence(self.state.send_seq); + tcp.set_acknowledgement(self.state.rec_seq); - Ok(()) - } + 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); - 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![], - ); + self.state.rec_seq = self.state.rec_seq.wrapping_add(payload.len() as u32); + } + } + tcp.set_flags(TcpFlags::PSH | TcpFlags::ACK); - Ok(()) - } - - fn setup_tcp_packet( - &mut self, - info: &CapturePacket, - payload: &[u8], - source_port: u16, - dest_port: u16, - ) -> Result { - 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() + tcp.packet_size() }; - self.setup_tcp_packet(&adjusted_info, &[], source_port, dest_port)?; - Ok(self.write_transport_payload( - &adjusted_info, + + self.write_transport_payload( + 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(), - )], - )) - }; + &buf[..buf_size + payload.len()], + vec![], + ); - // Send SYN - tcp_handshake_packet(info, Direction::Send, TcpFlags::SYN); + 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 { + Direction::Send => { + tcp.set_sequence(self.state.rec_seq); + tcp.set_acknowledgement(self.state.send_seq); - // 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); + info.direction = Direction::Receive; + } + Direction::Receive => { + tcp.set_sequence(self.state.send_seq); + tcp.set_acknowledgement(self.state.rec_seq); - // 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); + info.direction = Direction::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())], + ); + } + Protocol::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![], + ); + } + } } - fn handle_udp( - &mut self, + /// Encode a network layer (IP) packet with a payload. + fn encode_ip_packet( + &self, + buf: &mut [u8], info: &CapturePacket, + protocol: IpNextHeaderProtocol, 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); + ) -> (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) + }; - let buf_size = udp.packet_size(); + 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.state.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 { + Direction::Send => (local_address, remote_address), + Direction::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.state.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. + pub(crate) fn write_tcp_handshake(&mut self, info: &CapturePacket) { + let (source_port, dest_port) = (info.local_address.port(), info.remote_address.port()); + + let mut info = info.clone(); + info.direction = Direction::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.state.send_seq = 500; + tcp.set_sequence(self.state.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::Udp, - &self.buffer[..buf_size + payload.len()], - vec![], + &info, + IpNextHeaderProtocols::Tcp, + &buf[..buf_size], + options.clone(), ); - Ok(()) + // SYN + ACK + info.direction = Direction::Receive; + let buf_size = { + let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); + self.state.send_seq = self.state.send_seq.wrapping_add(1); + tcp.set_acknowledgement(self.state.send_seq); + self.state.rec_seq = 1000; + tcp.set_sequence(self.state.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 = Direction::Send; + let buf_size = { + let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); + tcp.set_sequence(self.state.send_seq); + self.state.rec_seq = self.state.rec_seq.wrapping_add(1); + tcp.set_acknowledgement(self.state.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.state.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: &CapturePacket, protocol: IpNextHeaderProtocol, payload: &[u8], - options: Vec, + options: Vec, ) { - 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 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 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, - }; + let mut physical_packet = vec![0; PACKET_SIZE]; + let physical_size = + self.encode_ethernet_packet(&mut physical_packet, ethertype, &network_packet) + network_size; - self.writer.write_block(&enhanced_packet_block.into_block()); - } + physical_packet.truncate(physical_size); - fn encode_ethernet_packet( - &mut self, - info: &CapturePacket, - ethertype: EtherType, - payload: &[u8], - ) -> Result { - 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", + 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.state.start_time.elapsed(), + options, + } + .into_block(), ) - })?; - - 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); + .unwrap(); } } @@ -304,7 +338,6 @@ impl Default for State { send_seq: 0, rec_seq: 0, has_sent_handshake: false, - has_sent_fin: false, stream_count: 0, } } diff --git a/crates/lib/src/capture/writer.rs b/crates/lib/src/capture/writer.rs index 012b700..476853c 100644 --- a/crates/lib/src/capture/writer.rs +++ b/crates/lib/src/capture/writer.rs @@ -14,7 +14,7 @@ lazy_static! { pub(crate) static ref CAPTURE_WRITER: Mutex>> = Mutex::new(None); } -pub trait Writer { +pub(crate) trait Writer { fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> crate::GDResult<()>; fn new_connect(&mut self, packet: &CapturePacket) -> crate::GDResult<()>; } diff --git a/crates/lib/src/lib.rs b/crates/lib/src/lib.rs index aeb3be2..38ec19c 100644 --- a/crates/lib/src/lib.rs +++ b/crates/lib/src/lib.rs @@ -46,7 +46,7 @@ mod socket; mod utils; #[cfg(feature = "packet_capture")] -pub(crate) mod capture; +pub mod capture; pub use errors::*; #[cfg(feature = "games")] diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index f0fe3d3..f380c89 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -213,13 +213,13 @@ pub mod capture { } /// A trait representing a provider of a network protocol. - pub trait ProtocolProvider { + pub(crate) trait ProtocolProvider { /// Returns the protocol used by the provider. fn protocol() -> Protocol; } /// Represents the TCP protocol provider. - pub struct ProtocolTCP; + pub(crate) struct ProtocolTCP; impl ProtocolProvider for ProtocolTCP { fn protocol() -> Protocol { Protocol::TCP @@ -375,14 +375,14 @@ pub mod capture { /// /// This type captures and processes UDP packets, wrapping around standard UDP socket /// functionalities with additional packet capture capabilities. - pub type CapturedUdpSocket = WrappedCaptureSocket; + 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 type CapturedTcpSocket = WrappedCaptureSocket; + pub(crate) type CapturedTcpSocket = WrappedCaptureSocket; } #[cfg(not(feature = "packet_capture"))] @@ -391,9 +391,9 @@ pub type UdpSocket = UdpSocketImpl; pub type TcpSocket = TcpSocketImpl; #[cfg(feature = "packet_capture")] -pub type UdpSocket = capture::CapturedUdpSocket; +pub(crate) type UdpSocket = capture::CapturedUdpSocket; #[cfg(feature = "packet_capture")] -pub type TcpSocket = capture::CapturedTcpSocket; +pub(crate) type TcpSocket = capture::CapturedTcpSocket; #[cfg(test)] mod tests { From 0543cabce26210bc8eb50bc4e3012a93f570c15f Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Sat, 13 Jan 2024 09:18:56 +0000 Subject: [PATCH 06/35] refactor: clean up --- crates/cli/src/main.rs | 140 ++++++++++++++++++++----------- crates/lib/src/capture/mod.rs | 24 +++--- crates/lib/src/capture/packet.rs | 94 ++++----------------- crates/lib/src/capture/pcap.rs | 23 ++--- crates/lib/src/capture/writer.rs | 18 ++-- crates/lib/src/socket.rs | 9 +- 6 files changed, 137 insertions(+), 171 deletions(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 969e5e3..df7078a 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -1,6 +1,9 @@ -use std::net::{IpAddr, ToSocketAddrs}; +use std::{ + net::{IpAddr, ToSocketAddrs}, + path::PathBuf, +}; -use clap::{Parser, ValueEnum}; +use clap::{Parser, Subcommand, ValueEnum}; use gamedig::{ games::*, protocols::types::{CommonResponse, ExtraRequestSettings, TimeoutSettings}, @@ -13,42 +16,68 @@ use self::error::{Error, Result}; // NOTE: For some reason without setting long_about here the doc comment for // ExtraRequestSettings gets set as the about for the CLI. #[derive(Debug, Parser)] -#[command(author, version, about, long_about = None)] +#[command(author, version, long_about = None, about = r" + + _____ _____ _ _____ _ _____ + / ____| | __ \(_) / ____| | |_ _| +| | __ __ _ _ __ ___ ___| | | |_ __ _ | | | | | | +| | |_ |/ _` | '_ ` _ \ / _ \ | | | |/ _` | | | | | | | +| |__| | (_| | | | | | | __/ |__| | | (_| | | |____| |____ _| |_ + \_____|\__,_|_| |_| |_|\___|_____/|_|\__, | \_____|______|_____| + __/ | + |___/ + + A command line interface for querying game servers. + Copyright (C) 2022 - Present GameDig Organization & Contributors + Licensed under the MIT license +")] struct Cli { - /// Unique identifier of the game for which server information is being - /// queried. - #[arg(short, long)] - game: String, + #[command(subcommand)] + action: Action, +} - /// Hostname or IP address of the server. - #[arg(short, long)] - ip: String, +#[derive(Subcommand, Debug)] +enum Action { + /// Query game server information + Query { + /// Unique identifier of the game for which server information is being + /// queried. + #[arg(short, long)] + game: String, - /// Optional query port number for the server. If not provided the default - /// port for the game is used. - #[arg(short, long)] - port: Option, + /// Hostname or IP address of the server. + #[arg(short, long)] + ip: String, - /// Flag indicating if the output should be in JSON format. - #[cfg(feature = "json")] - #[arg(short, long)] - json: bool, + /// Optional query port number for the server. If not provided the default + /// port for the game is used. + #[arg(short, long)] + port: Option, - /// Which response variant to use when outputting. - #[arg(short, long, default_value = "generic")] - output_mode: OutputMode, + /// Flag indicating if the output should be in JSON format for programmatic use. + #[cfg(feature = "json")] + #[arg(short, long)] + json: bool, - #[cfg(feature = "packet_capture")] - #[arg(short, long)] - capture: Option, + /// Which response variant to use when outputting + #[arg(short, long, default_value = "generic")] + output_mode: OutputMode, - /// Optional timeout settings for the server query. - #[command(flatten, next_help_heading = "Timeouts")] - timeout_settings: Option, + /// Optional file path for packet capture file writer + #[cfg(feature = "packet_capture")] + #[arg(short, long)] + capture: Option, - /// Optional extra settings for the server query. - #[command(flatten, next_help_heading = "Query options")] - extra_options: Option, + /// Optional timeout settings for the server query + #[command(flatten, next_help_heading = "Timeouts")] + timeout_settings: Option, + + /// Optional extra settings for the server query + #[command(flatten, next_help_heading = "Query options")] + extra_options: Option, + }, + /// Display the MIT License information + License, } #[derive(Clone, Debug, PartialEq, Eq, ValueEnum)] @@ -139,12 +168,12 @@ fn set_hostname_if_missing(host: &str, extra_options: &mut Option output_result_json(result.as_json()), + OutputMode::Generic if json => output_result_json(result.as_json()), #[cfg(feature = "json")] - OutputMode::ProtocolSpecific if args.json => output_result_json(result.as_original()), + OutputMode::ProtocolSpecific if json => output_result_json(result.as_original()), OutputMode::Generic => output_result_debug(result.as_json()), OutputMode::ProtocolSpecific => output_result_debug(result.as_original()), @@ -169,27 +198,36 @@ fn output_result_json(result: R) { } fn main() -> Result<()> { - // Parse the command line arguments let args = Cli::parse(); - // Retrieve the game based on the provided ID - let game = find_game(&args.game)?; + match args.action { + Action::Query { + game, + ip, + port, + json, + output_mode, + capture, + timeout_settings, + extra_options, + } => { + // Process the query command + let game = find_game(&game)?; + let mut extra_options = extra_options; + let ip = resolve_ip_or_domain(&ip, &mut extra_options)?; - // Extract extra options for use in setup - let mut extra_options = args.extra_options.clone(); + #[cfg(feature = "packet_capture")] + gamedig::capture::setup_capture(capture); - // Resolve the IP address - let ip = resolve_ip_or_domain(&args.ip, &mut extra_options)?; - - #[cfg(feature = "packet_capture")] - gamedig::capture::setup_capture(args.capture.clone()); - - // Query the server using game definition, parsed IP, and user command line - // flags. - let result = query_with_timeout_and_extra_settings(game, &ip, args.port, args.timeout_settings, extra_options)?; - - // Output the result in the specified format - output_result(&args, result.as_ref()); + let result = query_with_timeout_and_extra_settings(game, &ip, port, timeout_settings, extra_options)?; + output_result(output_mode, json, result.as_ref()); + } + Action::License => { + // Bake the license into the binary + // so we don't have to ship it separately + println!("{}", include_str!("../../../LICENSE.md")); + } + } Ok(()) } diff --git a/crates/lib/src/capture/mod.rs b/crates/lib/src/capture/mod.rs index 6cfd8ad..13bf7af 100644 --- a/crates/lib/src/capture/mod.rs +++ b/crates/lib/src/capture/mod.rs @@ -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) { - if let Some(file_name) = file_name { +pub fn setup_capture(file_path: Option) { + 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) { 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) { /// 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) { crate::socket::capture::set_writer(writer); } diff --git a/crates/lib/src/capture/packet.rs b/crates/lib/src/capture/packet.rs index 8ad2917..83910d0 100644 --- a/crates/lib/src/capture/packet.rs +++ b/crates/lib/src/capture/packet.rs @@ -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; } @@ -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(&self) -> (T, T) { + pub(super) fn ipvt_by_direction(&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(&self, source: T, remote: T) -> (T, T) { + pub(self) const fn order(&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::()); 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(); - } } diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs index be65968..e6cd5b0 100644 --- a/crates/lib/src/capture/pcap.rs +++ b/crates/lib/src/capture/pcap.rs @@ -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 Pcap { 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 Pcap { 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 Pcap { (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); diff --git a/crates/lib/src/capture/writer.rs b/crates/lib/src/capture/writer.rs index 476853c..1951e7e 100644 --- a/crates/lib/src/capture/writer.rs +++ b/crates/lib/src/capture/writer.rs @@ -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>> = 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 Writer for Pcap { diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index f380c89..0ee26f6 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -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 { + pub(crate) struct WrappedCaptureSocket { inner: I, remote_address: SocketAddr, _protocol: PhantomData

, @@ -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(), From 6aee5ebb760ca441850961dfe47a144959cc6677 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Sat, 13 Jan 2024 10:10:31 +0000 Subject: [PATCH 07/35] chore: format --- crates/cli/src/main.rs | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index df7078a..acd68fa 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -13,10 +13,7 @@ mod error; use self::error::{Error, Result}; -// NOTE: For some reason without setting long_about here the doc comment for -// ExtraRequestSettings gets set as the about for the CLI. -#[derive(Debug, Parser)] -#[command(author, version, long_about = None, about = r" +const GAMEDIG_HEADER: &str = r" _____ _____ _ _____ _ _____ / ____| | __ \(_) / ____| | |_ _| @@ -30,7 +27,12 @@ use self::error::{Error, Result}; A command line interface for querying game servers. Copyright (C) 2022 - Present GameDig Organization & Contributors Licensed under the MIT license -")] +"; + +// NOTE: For some reason without setting long_about here the doc comment for +// ExtraRequestSettings gets set as the about for the CLI. +#[derive(Debug, Parser)] +#[command(author, version, about = GAMEDIG_HEADER, long_about = None)] struct Cli { #[command(subcommand)] action: Action, From 6cf6800bffd4830c4fbc360083d19d6e44326250 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Sat, 13 Jan 2024 11:21:41 +0000 Subject: [PATCH 08/35] chore: add source link --- crates/cli/src/main.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index acd68fa..b3eb937 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -78,6 +78,9 @@ enum Action { #[command(flatten, next_help_heading = "Query options")] extra_options: Option, }, + + /// Check out the source code + Source, /// Display the MIT License information License, } @@ -224,6 +227,13 @@ fn main() -> Result<()> { let result = query_with_timeout_and_extra_settings(game, &ip, port, timeout_settings, extra_options)?; output_result(output_mode, json, result.as_ref()); } + Action::Source => { + println!("{}", GAMEDIG_HEADER); + + // Print the source code location + println!("Source code: https://github.com/gamedig/rust-gamedig\n"); + println!("Be sure to leave a star if you like the project :)\n"); + } Action::License => { // Bake the license into the binary // so we don't have to ship it separately From 61ecbab312ab3f920f7dafb7342116e23c8df3b3 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Mon, 15 Jan 2024 15:51:21 +0000 Subject: [PATCH 09/35] refactor: clean up and add more features --- crates/cli/Cargo.toml | 45 +++++++++++--- crates/cli/src/main.rs | 130 +++++++++++++++++++++++++++++++++-------- 2 files changed, 142 insertions(+), 33 deletions(-) diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index dd4ed97..7610004 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -11,15 +11,42 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [features] -default = ["json", "packet_capture"] -json = ["dep:serde", "dep:serde_json", "gamedig/serde"] +default = ["packet_capture", "bson", "json", "xml", "browser"] + +# Tools packet_capture = ["gamedig/packet_capture"] -[dependencies] -clap = { version = "4.1.11", features = ["derive"] } -gamedig = { version = "*", path = "../lib", features = ["clap"] } -thiserror = "1.0.43" +# Output formats +bson = ["dep:serde", "dep:bson", "dep:hex", "gamedig/serde"] +json = ["dep:serde", "dep:serde_json", "gamedig/serde"] +xml = ["dep:serde", "dep:serde-xml-rs", "gamedig/serde"] -# JSON dependencies -serde = { version = "1", optional = true } -serde_json = { version = "1", optional = true } +# Misc +browser = ["dep:webbrowser"] + +[dependencies] +# Core Dependencies +thiserror = "1.0.43" +clap = { version = "4.1.11", default-features = false, features = ["derive"] } +gamedig = { version = "*", path = "../lib", default-features = false, features = [ + "clap", + "games", + "game_defs", +] } + +# Feature Dependencies +# Serialization / Deserialization +serde = { version = "1", optional = true, default-features = false } + +# BSON +bson = { version = "2.8.1", optional = true, default-features = false } +hex = { version = "0.4.3", optional = true, default-features = false } + +# JSON +serde_json = { version = "1", optional = true, default-features = false } + +# XML +serde-xml-rs = { version = "0.6.0", optional = true, default-features = false } + +# Browser +webbrowser = { version = "0.8.12", optional = true, default-features = false } diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index b3eb937..84a4b1f 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -25,7 +25,7 @@ const GAMEDIG_HEADER: &str = r" |___/ A command line interface for querying game servers. - Copyright (C) 2022 - Present GameDig Organization & Contributors + Copyright (C) 2022 - 2023 GameDig Organization & Contributors Licensed under the MIT license "; @@ -56,10 +56,9 @@ enum Action { #[arg(short, long)] port: Option, - /// Flag indicating if the output should be in JSON format for programmatic use. - #[cfg(feature = "json")] - #[arg(short, long)] - json: bool, + /// Specifies the output format + #[arg(short, long, default_value = "debug", value_enum)] + format: OutputFormat, /// Which response variant to use when outputting #[arg(short, long, default_value = "generic")] @@ -95,6 +94,19 @@ enum OutputMode { ProtocolSpecific, } +#[derive(Clone, Debug, PartialEq, Eq, ValueEnum)] +enum OutputFormat { + Debug, + #[cfg(feature = "json")] + JsonPretty, + #[cfg(feature = "json")] + Json, + #[cfg(feature = "xml")] + Xml, + #[cfg(feature = "bson")] + Bson, +} + /// Attempt to find a game from the [library game definitions](GAMES) based on /// its unique identifier. /// @@ -122,13 +134,13 @@ fn find_game(game_id: &str) -> Result<&'static Game> { /// # Returns /// * `Result` - On sucess returns a resolved IP address; on failure /// returns an [Error::InvalidHostname] error. -fn resolve_ip_or_domain(host: &str, extra_options: &mut Option) -> Result { - if let Ok(parsed_ip) = host.parse() { +fn resolve_ip_or_domain>(host: T, extra_options: &mut Option) -> Result { + let host_str = host.as_ref(); + if let Ok(parsed_ip) = host_str.parse() { Ok(parsed_ip) } else { - set_hostname_if_missing(host, extra_options); - - resolve_domain(host) + set_hostname_if_missing(host_str, extra_options); + resolve_domain(host_str) } } @@ -173,15 +185,36 @@ fn set_hostname_if_missing(host: &str, extra_options: &mut Option(output_mode: OutputMode, format: OutputFormat, result: &T) { + match format { + OutputFormat::Debug => match output_mode { + OutputMode::Generic => output_result_debug(result.as_json()), + OutputMode::ProtocolSpecific => output_result_debug(result.as_original()), + }, #[cfg(feature = "json")] - OutputMode::Generic if json => output_result_json(result.as_json()), + OutputFormat::JsonPretty => match output_mode { + OutputMode::Generic => output_result_json_pretty(result.as_json()), + OutputMode::ProtocolSpecific => output_result_json_pretty(result.as_original()), + }, #[cfg(feature = "json")] - OutputMode::ProtocolSpecific if json => output_result_json(result.as_original()), - - OutputMode::Generic => output_result_debug(result.as_json()), - OutputMode::ProtocolSpecific => output_result_debug(result.as_original()), + OutputFormat::Json => match output_mode { + OutputMode::Generic => output_result_json(result.as_json()), + OutputMode::ProtocolSpecific => output_result_json(result.as_original()), + }, + #[cfg(feature = "xml")] + OutputFormat::Xml => match output_mode { + OutputMode::Generic => output_result_xml(result.as_json()), + //BUG: In this case we get a writer error with all serde write methods some reason with xml 0.6.0 + //BUG: gamedig-cli.exe query -g -i -p -f xml --output-mode protocol-specific + //BUG: Writer: emitter error: document start event has already been emitted + //BUG: With xml 0.5.1 we get unsupported operation: 'serialize_unit_variant' + OutputMode::ProtocolSpecific => panic!("XML format is not supported for protocol specific output"), + }, + #[cfg(feature = "bson")] + OutputFormat::Bson => match output_mode { + OutputMode::Generic => output_result_bson(result.as_json()), + OutputMode::ProtocolSpecific => output_result_bson(result.as_original()), + }, } } @@ -198,8 +231,33 @@ fn output_result_debug(result: R) { /// # Arguments /// * `result` - A serde serializable result. #[cfg(feature = "json")] -fn output_result_json(result: R) { - serde_json::to_writer_pretty(std::io::stdout(), &result).unwrap(); +fn output_result_json(result: T) { + println!("{}", serde_json::to_string(&result).unwrap()); +} + +#[cfg(feature = "json")] +fn output_result_json_pretty(result: T) { + println!("{}", serde_json::to_string_pretty(&result).unwrap()); +} + +/// Output the result as an XML object. +/// # Arguments +/// * `result` - A serde serializable result. +fn output_result_xml(result: T) { + println!("{}", serde_xml_rs::to_string(&result).unwrap()); +} + +#[cfg(feature = "bson")] +fn output_result_bson(result: T) { + let bson = bson::to_bson(&result).unwrap(); + + if let bson::Bson::Document(document) = bson { + let bytes = bson::to_vec(&document).unwrap(); + + println!("{}", hex::encode(bytes)); + } else { + panic!("Failed to convert result to BSON"); + } } fn main() -> Result<()> { @@ -210,7 +268,7 @@ fn main() -> Result<()> { game, ip, port, - json, + format, output_mode, capture, timeout_settings, @@ -225,14 +283,38 @@ fn main() -> Result<()> { gamedig::capture::setup_capture(capture); let result = query_with_timeout_and_extra_settings(game, &ip, port, timeout_settings, extra_options)?; - output_result(output_mode, json, result.as_ref()); + output_result(output_mode, format, result.as_ref()); } Action::Source => { println!("{}", GAMEDIG_HEADER); - // Print the source code location - println!("Source code: https://github.com/gamedig/rust-gamedig\n"); - println!("Be sure to leave a star if you like the project :)\n"); + #[cfg(feature = "browser")] + { + // Directly offering to open the URL + println!("\nWould you like to open the GitHub repository in your default browser? [Y/n]"); + + let mut choice = String::new(); + std::io::stdin().read_line(&mut choice).unwrap(); + if choice.trim().eq_ignore_ascii_case("Y") { + if webbrowser::open("https://github.com/gamedig/rust-gamedig").is_ok() { + println!("Opening GitHub repository in default browser..."); + } else { + println!("Failed to open GitHub repository in default browser."); + println!("Please use the following URL: https://github.com/gamedig/rust-gamedig"); + } + } else { + println!("Not to worry, you can always open the repository manually"); + println!("by visiting the following URL: https://github.com/gamedig/rust-gamedig"); + } + } + + #[cfg(not(feature = "browser"))] + { + println!("\nYou can find the source code for this project at the following URL:"); + println!("https://github.com/gamedig/rust-gamedig"); + } + + println!("\nBe sure to leave a star if you like the project :)"); } Action::License => { // Bake the license into the binary From c30f28741f2c815d1ab903d2747377a650b52e57 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Mon, 15 Jan 2024 16:28:56 +0000 Subject: [PATCH 10/35] impl: tcp fin --- crates/lib/src/capture/pcap.rs | 48 ++++++++++++++++++++++++++++++-- crates/lib/src/capture/writer.rs | 12 +++++++- crates/lib/src/socket.rs | 22 ++++++++++++++- 3 files changed, 77 insertions(+), 5 deletions(-) diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs index e6cd5b0..0eb8b21 100644 --- a/crates/lib/src/capture/pcap.rs +++ b/crates/lib/src/capture/pcap.rs @@ -114,7 +114,7 @@ impl Pcap { &info, IpNextHeaderProtocols::Tcp, &buf[..buf_size], - vec![EnhancedPacketOption::Comment("Generated TCP ack".into())], + vec![EnhancedPacketOption::Comment("Generated TCP ACK".into())], ); } Protocol::UDP => { @@ -285,8 +285,50 @@ impl Pcap { self.state.has_sent_handshake = true; } - /// Take a transport layer packet as a buffer and write it after encoding - /// all the layers under it. + pub(crate) fn send_tcp_fin(&mut self, info: &CapturePacket) { + let mut buf = vec![0; BUFFER_SIZE]; + let (source_port, dest_port) = info.ports_by_direction(); + + let buf_size = { + let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); + tcp.set_source(source_port); + tcp.set_destination(dest_port); + tcp.set_data_offset(5); + tcp.set_window(43440); + + match info.direction { + Direction::Send => { + tcp.set_sequence(self.state.send_seq); + tcp.set_acknowledgement(self.state.rec_seq); + } + Direction::Receive => { + tcp.set_sequence(self.state.rec_seq); + tcp.set_acknowledgement(self.state.send_seq); + } + } + + tcp.set_flags(TcpFlags::FIN | TcpFlags::ACK); + tcp.packet_size() + }; + + self.write_transport_payload( + info, + IpNextHeaderProtocols::Tcp, + &buf[..buf_size], + vec![EnhancedPacketOption::Comment("Generated TCP FIN".into())], + ); + + // Update sequence number + match info.direction { + Direction::Send => { + self.state.send_seq = self.state.send_seq.wrapping_add(1); + } + Direction::Receive => { + self.state.rec_seq = self.state.rec_seq.wrapping_add(1); + } + } + } + fn write_transport_payload( &mut self, info: &CapturePacket, diff --git a/crates/lib/src/capture/writer.rs b/crates/lib/src/capture/writer.rs index 1951e7e..4d63655 100644 --- a/crates/lib/src/capture/writer.rs +++ b/crates/lib/src/capture/writer.rs @@ -14,7 +14,7 @@ lazy_static! { pub(crate) trait Writer { fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> GDResult<()>; fn new_connect(&mut self, packet: &CapturePacket) -> GDResult<()>; - //TODO: Close connection + fn close_connection(&mut self, packet: &CapturePacket) -> GDResult<()>; } impl Writer for Pcap { @@ -36,4 +36,14 @@ impl Writer for Pcap { Ok(()) } + + fn close_connection(&mut self, packet: &CapturePacket) -> GDResult<()> { + match packet.protocol { + Protocol::TCP => { + self.send_tcp_fin(packet); + } + Protocol::UDP => {} + } + Ok(()) + } } diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index 0ee26f6..0ba35ee 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -239,7 +239,7 @@ pub mod capture { /// * `I` - The inner socket type. /// * `P` - The protocol provider. #[derive(Clone, Debug)] - pub(crate) struct WrappedCaptureSocket { + pub(crate) struct WrappedCaptureSocket { inner: I, remote_address: SocketAddr, _protocol: PhantomData

, @@ -369,6 +369,26 @@ pub mod capture { } } + // 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. /// From 89d4ddeac772ad0b00e0599ea0f1f7d192154379 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Tue, 16 Jan 2024 18:41:10 +0000 Subject: [PATCH 11/35] fix: bson output --- crates/cli/Cargo.toml | 4 +++- crates/cli/src/main.rs | 40 ++++++++++++++++++++++++++++++++++------ 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index 7610004..4b5a0d9 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -17,7 +17,7 @@ default = ["packet_capture", "bson", "json", "xml", "browser"] packet_capture = ["gamedig/packet_capture"] # Output formats -bson = ["dep:serde", "dep:bson", "dep:hex", "gamedig/serde"] +bson = ["dep:serde", "dep:bson", "dep:hex", "dep:base64", "gamedig/serde"] json = ["dep:serde", "dep:serde_json", "gamedig/serde"] xml = ["dep:serde", "dep:serde-xml-rs", "gamedig/serde"] @@ -40,6 +40,7 @@ serde = { version = "1", optional = true, default-features = false } # BSON bson = { version = "2.8.1", optional = true, default-features = false } +base64 = { version = "0.21.7", optional = true, default-features = false, features = ["std"]} hex = { version = "0.4.3", optional = true, default-features = false } # JSON @@ -50,3 +51,4 @@ serde-xml-rs = { version = "0.6.0", optional = true, default-features = false } # Browser webbrowser = { version = "0.8.12", optional = true, default-features = false } + diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 84a4b1f..b106813 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -96,15 +96,23 @@ enum OutputMode { #[derive(Clone, Debug, PartialEq, Eq, ValueEnum)] enum OutputFormat { + /// Human readable structured output Debug, + /// RFC 8259 #[cfg(feature = "json")] JsonPretty, + /// RFC 8259 #[cfg(feature = "json")] Json, + /// Parser tries to be mostly XML 1.1 (RFC 7303) compliant #[cfg(feature = "xml")] Xml, + /// RFC 4648 section 8 #[cfg(feature = "bson")] - Bson, + BsonHex, + /// RFC 4648 section 4 + #[cfg(feature = "bson")] + BsonBase64, } /// Attempt to find a game from the [library game definitions](GAMES) based on @@ -211,9 +219,14 @@ fn output_result(output_mode: OutputMode, format: Ou OutputMode::ProtocolSpecific => panic!("XML format is not supported for protocol specific output"), }, #[cfg(feature = "bson")] - OutputFormat::Bson => match output_mode { - OutputMode::Generic => output_result_bson(result.as_json()), - OutputMode::ProtocolSpecific => output_result_bson(result.as_original()), + OutputFormat::BsonHex => match output_mode { + OutputMode::Generic => output_result_bson_hex(result.as_json()), + OutputMode::ProtocolSpecific => output_result_bson_hex(result.as_original()), + }, + #[cfg(feature = "bson")] + OutputFormat::BsonBase64 => match output_mode { + OutputMode::Generic => output_result_bson_base64(result.as_json()), + OutputMode::ProtocolSpecific => output_result_bson_base64(result.as_original()), }, } } @@ -248,7 +261,7 @@ fn output_result_xml(result: T) { } #[cfg(feature = "bson")] -fn output_result_bson(result: T) { +fn output_result_bson_hex(result: T) { let bson = bson::to_bson(&result).unwrap(); if let bson::Bson::Document(document) = bson { @@ -256,7 +269,22 @@ fn output_result_bson(result: T) { println!("{}", hex::encode(bytes)); } else { - panic!("Failed to convert result to BSON"); + panic!("Failed to convert result to BSON Hex"); + } +} + +#[cfg(feature = "bson")] +fn output_result_bson_base64(result: T) { + use base64::Engine; + + let bson = bson::to_bson(&result).unwrap(); + + if let bson::Bson::Document(document) = bson { + let bytes = bson::to_vec(&document).unwrap(); + + println!("{}", base64::prelude::BASE64_STANDARD.encode(&bytes)); + } else { + panic!("Failed to convert result to BSON Base64"); } } From 15004f3dae0ea3fa8efa647b2d65f68ea8ac2823 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Tue, 16 Jan 2024 18:41:29 +0000 Subject: [PATCH 12/35] chore: optimize further --- Cargo.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 4dd9c76..d25d3e8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,8 @@ resolver = "2" opt-level = 3 debug = false rpath = true -lto = true +lto = 'fat' +codegen-units = 1 [profile.release.package."*"] opt-level = 3 From 1991c9f5eb7ed12ec908a53dcaf11a27c9482b69 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Thu, 18 Jan 2024 17:48:11 +0000 Subject: [PATCH 13/35] chore: add missing cli fn docs --- crates/cli/src/main.rs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 1201674..5185c76 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -150,7 +150,6 @@ fn resolve_ip_or_domain>(host: T, extra_options: &mut Option(result: T) { println!("{}", serde_json::to_string(&result).unwrap()); } +/// Output the result as a pretty printed JSON object. +/// +/// # Arguments +/// * `result` - A serde serializable result. #[cfg(feature = "json")] fn output_result_json_pretty(result: T) { println!("{}", serde_json::to_string_pretty(&result).unwrap()); @@ -261,6 +264,10 @@ fn output_result_xml(result: T) { println!("{}", serde_xml_rs::to_string(&result).unwrap()); } +/// Output the result as a BSON object encoded as a hex string. +/// +/// # Arguments +/// * `result` - A serde serializable result. #[cfg(feature = "bson")] fn output_result_bson_hex(result: T) { let bson = bson::to_bson(&result).unwrap(); @@ -274,6 +281,10 @@ fn output_result_bson_hex(result: T) { } } +/// Output the result as a BSON object encoded as a base64 string. +/// +/// # Arguments +/// * `result` - A serde serializable result. #[cfg(feature = "bson")] fn output_result_bson_base64(result: T) { use base64::Engine; From b49525543d8b2eaa411e70cd0fdf43a5f78a1493 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Thu, 18 Jan 2024 17:56:58 +0000 Subject: [PATCH 14/35] chore: format --- crates/cli/src/main.rs | 72 +++++++++++++++++++------------- crates/lib/src/capture/mod.rs | 4 +- crates/lib/src/capture/packet.rs | 19 +++++---- crates/lib/src/capture/pcap.rs | 27 ++++++++---- crates/lib/src/socket.rs | 72 +++++++++++++------------------- 5 files changed, 101 insertions(+), 93 deletions(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 5185c76..b12c82e 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -51,8 +51,8 @@ enum Action { #[arg(short, long)] ip: String, - /// Optional query port number for the server. If not provided the default - /// port for the game is used. + /// Optional query port number for the server. If not provided the + /// default port for the game is used. #[arg(short, long)] port: Option, @@ -195,39 +195,51 @@ fn set_hostname_if_missing(host: &str, extra_options: &mut Option(output_mode: OutputMode, format: OutputFormat, result: &T) { match format { - OutputFormat::Debug => match output_mode { - OutputMode::Generic => output_result_debug(result.as_json()), - OutputMode::ProtocolSpecific => output_result_debug(result.as_original()), - }, + OutputFormat::Debug => { + match output_mode { + OutputMode::Generic => output_result_debug(result.as_json()), + OutputMode::ProtocolSpecific => output_result_debug(result.as_original()), + } + } #[cfg(feature = "json")] - OutputFormat::JsonPretty => match output_mode { - OutputMode::Generic => output_result_json_pretty(result.as_json()), - OutputMode::ProtocolSpecific => output_result_json_pretty(result.as_original()), - }, + OutputFormat::JsonPretty => { + match output_mode { + OutputMode::Generic => output_result_json_pretty(result.as_json()), + OutputMode::ProtocolSpecific => output_result_json_pretty(result.as_original()), + } + } #[cfg(feature = "json")] - OutputFormat::Json => match output_mode { - OutputMode::Generic => output_result_json(result.as_json()), - OutputMode::ProtocolSpecific => output_result_json(result.as_original()), - }, + OutputFormat::Json => { + match output_mode { + OutputMode::Generic => output_result_json(result.as_json()), + OutputMode::ProtocolSpecific => output_result_json(result.as_original()), + } + } #[cfg(feature = "xml")] - OutputFormat::Xml => match output_mode { - OutputMode::Generic => output_result_xml(result.as_json()), - //BUG: In this case we get a writer error with all serde write methods some reason with xml 0.6.0 - //BUG: gamedig-cli.exe query -g -i -p -f xml --output-mode protocol-specific - //BUG: Writer: emitter error: document start event has already been emitted - //BUG: With xml 0.5.1 we get unsupported operation: 'serialize_unit_variant' - OutputMode::ProtocolSpecific => panic!("XML format is not supported for protocol specific output"), - }, + OutputFormat::Xml => { + match output_mode { + OutputMode::Generic => output_result_xml(result.as_json()), + // BUG: In this case we get a writer error with all serde write methods some reason with xml 0.6.0 + // BUG: gamedig-cli.exe query -g -i -p -f xml --output-mode protocol-specific + // BUG: Writer: emitter error: document start event has already been emitted + // BUG: With xml 0.5.1 we get unsupported operation: 'serialize_unit_variant' + OutputMode::ProtocolSpecific => panic!("XML format is not supported for protocol specific output"), + } + } #[cfg(feature = "bson")] - OutputFormat::BsonHex => match output_mode { - OutputMode::Generic => output_result_bson_hex(result.as_json()), - OutputMode::ProtocolSpecific => output_result_bson_hex(result.as_original()), - }, + OutputFormat::BsonHex => { + match output_mode { + OutputMode::Generic => output_result_bson_hex(result.as_json()), + OutputMode::ProtocolSpecific => output_result_bson_hex(result.as_original()), + } + } #[cfg(feature = "bson")] - OutputFormat::BsonBase64 => match output_mode { - OutputMode::Generic => output_result_bson_base64(result.as_json()), - OutputMode::ProtocolSpecific => output_result_bson_base64(result.as_original()), - }, + OutputFormat::BsonBase64 => { + match output_mode { + OutputMode::Generic => output_result_bson_base64(result.as_json()), + OutputMode::ProtocolSpecific => output_result_bson_base64(result.as_original()), + } + } } } diff --git a/crates/lib/src/capture/mod.rs b/crates/lib/src/capture/mod.rs index 13bf7af..a1995f8 100644 --- a/crates/lib/src/capture/mod.rs +++ b/crates/lib/src/capture/mod.rs @@ -38,6 +38,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::socket::capture::set_writer(writer); } diff --git a/crates/lib/src/capture/packet.rs b/crates/lib/src/capture/packet.rs index 83910d0..71e8136 100644 --- a/crates/lib/src/capture/packet.rs +++ b/crates/lib/src/capture/packet.rs @@ -31,7 +31,8 @@ pub(crate) enum Protocol { /// Trait for handling different types of IP addresses (IPv4, IPv6). pub(crate) trait IpAddress: Sized { - /// Creates an instance from a standard `IpAddr`, returning `None` if the types are incompatible. + /// Creates an instance from a standard `IpAddr`, returning `None` if the + /// types are incompatible. fn from_std(ip: IpAddr) -> Option; } @@ -67,12 +68,15 @@ impl CapturePacket<'_> { (local, remote) } - /// Retrieves IP addresses of a specific type (IPv4 or IPv6) based on the packet's direction. + /// 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. + /// 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. + /// - (T, T): Tuple of (source IP, destination IP) of the specified type in + /// order. pub(super) fn ipvt_by_direction(&self) -> (T, T) { let (local, remote) = ( T::from_std(self.local_address.ip()).expect("Incorrect IP type for local address"), @@ -84,7 +88,8 @@ impl CapturePacket<'_> { } impl Direction { - /// Orders two elements (source and destination) based on the packet's direction. + /// Orders two elements (source and destination) based on the packet's + /// direction. /// /// Returns: /// - (T, T): Ordered tuple (source, destination). @@ -124,9 +129,7 @@ mod tests { use std::str::FromStr; // Helper function to create a SocketAddr from a string - fn socket_addr(addr: &str) -> SocketAddr { - SocketAddr::from_str(addr).unwrap() - } + fn socket_addr(addr: &str) -> SocketAddr { SocketAddr::from_str(addr).unwrap() } #[test] fn test_ports_by_direction() { diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs index 0eb8b21..0613619 100644 --- a/crates/lib/src/capture/pcap.rs +++ b/crates/lib/src/capture/pcap.rs @@ -11,7 +11,13 @@ use pnet_packet::{ use std::{io::Write, net::IpAddr, time::Instant}; use super::packet::{ - CapturePacket, Direction, Protocol, HEADER_SIZE_ETHERNET, HEADER_SIZE_IP4, HEADER_SIZE_IP6, HEADER_SIZE_UDP, + CapturePacket, + Direction, + Protocol, + HEADER_SIZE_ETHERNET, + HEADER_SIZE_IP4, + HEADER_SIZE_IP6, + HEADER_SIZE_UDP, PACKET_SIZE, }; @@ -80,7 +86,7 @@ impl Pcap { self.write_transport_payload( info, IpNextHeaderProtocols::Tcp, - &buf[..buf_size + payload.len()], + &buf[.. buf_size + payload.len()], vec![], ); @@ -113,7 +119,7 @@ impl Pcap { self.write_transport_payload( &info, IpNextHeaderProtocols::Tcp, - &buf[..buf_size], + &buf[.. buf_size], vec![EnhancedPacketOption::Comment("Generated TCP ACK".into())], ); } @@ -131,7 +137,7 @@ impl Pcap { self.write_transport_payload( info, IpNextHeaderProtocols::Udp, - &buf[..buf_size + payload.len()], + &buf[.. buf_size + payload.len()], vec![], ); } @@ -238,7 +244,7 @@ impl Pcap { self.write_transport_payload( &info, IpNextHeaderProtocols::Tcp, - &buf[..buf_size], + &buf[.. buf_size], options.clone(), ); @@ -261,7 +267,7 @@ impl Pcap { self.write_transport_payload( &info, IpNextHeaderProtocols::Tcp, - &buf[..buf_size], + &buf[.. buf_size], options.clone(), ); @@ -280,7 +286,12 @@ impl Pcap { tcp.packet_size() }; - self.write_transport_payload(&info, IpNextHeaderProtocols::Tcp, &buf[..buf_size], options); + self.write_transport_payload( + &info, + IpNextHeaderProtocols::Tcp, + &buf[.. buf_size], + options, + ); self.state.has_sent_handshake = true; } @@ -314,7 +325,7 @@ impl Pcap { self.write_transport_payload( info, IpNextHeaderProtocols::Tcp, - &buf[..buf_size], + &buf[.. buf_size], vec![EnhancedPacketOption::Comment("Generated TCP FIN".into())], ); diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index 9dd1e91..4548720 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -22,8 +22,7 @@ pub trait Socket { /// # Returns /// A result containing the socket instance or an error. fn new(address: &SocketAddr, timeout_settings: &Option) -> GDResult - where - Self: Sized; + where Self: Sized; /// Apply read and write timeouts to the socket. /// @@ -112,12 +111,8 @@ impl Socket for TcpSocketImpl { Ok(buf) } - fn port(&self) -> u16 { - self.address.port() - } - fn local_addr(&self) -> std::io::Result { - self.socket.local_addr() - } + fn port(&self) -> u16 { self.address.port() } + fn local_addr(&self) -> std::io::Result { self.socket.local_addr() } } /// Implementation of a UDP socket. @@ -165,15 +160,11 @@ impl Socket for UdpSocketImpl { .recv_from(&mut buf) .map_err(|e| PacketReceive.context(e))?; - Ok(buf[..number_of_bytes_received].to_vec()) + Ok(buf[.. number_of_bytes_received].to_vec()) } - fn port(&self) -> u16 { - self.address.port() - } - fn local_addr(&self) -> std::io::Result { - self.socket.local_addr() - } + fn port(&self) -> u16 { self.address.port() } + fn local_addr(&self) -> std::io::Result { self.socket.local_addr() } } /// Things used for capturing packets. @@ -219,17 +210,13 @@ pub mod capture { /// Represents the TCP protocol provider. pub(crate) struct ProtocolTCP; impl ProtocolProvider for ProtocolTCP { - fn protocol() -> Protocol { - Protocol::TCP - } + fn protocol() -> Protocol { Protocol::TCP } } /// Represents the UDP protocol provider. pub(crate) struct ProtocolUDP; impl ProtocolProvider for ProtocolUDP { - fn protocol() -> Protocol { - Protocol::UDP - } + fn protocol() -> Protocol { Protocol::UDP } } /// A socket wrapper that allows capturing packets. @@ -247,8 +234,9 @@ pub mod capture { 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`. + /// 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. @@ -257,9 +245,7 @@ pub mod capture { /// # Returns /// A `GDResult` containing either the wrapped socket or an error. fn new(address: &SocketAddr, timeout_settings: &Option) -> GDResult - where - Self: Sized, - { + where Self: Sized { let v = Self { inner: I::new(address, timeout_settings)?, remote_address: *address, @@ -282,8 +268,8 @@ pub mod capture { /// 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. + /// 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. @@ -307,8 +293,8 @@ pub mod capture { /// 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. + /// 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. @@ -353,9 +339,7 @@ pub mod capture { /// /// # Returns /// The remote port number. - fn port(&self) -> u16 { - self.inner.port() - } + fn port(&self) -> u16 { self.inner.port() } /// Returns the local SocketAddr of the wrapped socket. /// @@ -363,9 +347,7 @@ pub mod capture { /// /// # Returns /// The local SocketAddr. - fn local_addr(&self) -> std::io::Result { - self.inner.local_addr() - } + fn local_addr(&self) -> std::io::Result { self.inner.local_addr() } } // this seems a bad way to do this, but its safe @@ -388,18 +370,20 @@ pub mod capture { } } - /// A specialized `WrappedCaptureSocket` for UDP, using `UdpSocketImpl` as the inner socket - /// and `ProtocolUDP` as the protocol provider. + /// 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. + /// 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. + /// 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. + /// This type captures and processes TCP packets, wrapping around standard + /// TCP socket functionalities with additional packet capture + /// capabilities. pub(crate) type CapturedTcpSocket = WrappedCaptureSocket; } From 962c8564183fd53b93abd4862017d6effe974935 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Thu, 18 Jan 2024 18:08:27 +0000 Subject: [PATCH 15/35] fix: packet support for msrv --- crates/lib/src/capture/packet.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/lib/src/capture/packet.rs b/crates/lib/src/capture/packet.rs index 71e8136..f7c3c3c 100644 --- a/crates/lib/src/capture/packet.rs +++ b/crates/lib/src/capture/packet.rs @@ -54,7 +54,7 @@ impl CapturePacket<'_> { /// /// Returns: /// - (u16, u16): Tuple of (source port, destination port). - pub(super) const fn ports_by_direction(&self) -> (u16, u16) { + pub(super) fn ports_by_direction(&self) -> (u16, u16) { let (local, remote) = (self.local_address.port(), self.remote_address.port()); self.direction.order(local, remote) } From 3eb20b9debfcbca3eb70bbad707964ef623d7f4d Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Thu, 18 Jan 2024 18:13:42 +0000 Subject: [PATCH 16/35] fix: mindustry leak crate-private type --- crates/lib/src/games/mindustry/protocol.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/lib/src/games/mindustry/protocol.rs b/crates/lib/src/games/mindustry/protocol.rs index a42d69a..04e0e05 100644 --- a/crates/lib/src/games/mindustry/protocol.rs +++ b/crates/lib/src/games/mindustry/protocol.rs @@ -16,7 +16,7 @@ pub const MAX_BUFFER_SIZE: usize = 500; /// Send a ping packet. /// /// [Reference](https://github.com/Anuken/Mindustry/blob/a2e5fbdedb2fc1c8d3c157bf344d10ad6d321442/core/src/mindustry/net/ArcNetProvider.java#L248) -pub fn send_ping(socket: &mut UdpSocket) -> GDResult<()> { socket.send(&[-2i8 as u8, 1i8 as u8]) } +pub(crate) fn send_ping(socket: &mut UdpSocket) -> GDResult<()> { socket.send(&[-2i8 as u8, 1i8 as u8]) } /// Parse server data. /// From e615c63ed289da1e78abcb69e851c5b75e0abf5b Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Thu, 18 Jan 2024 18:18:11 +0000 Subject: [PATCH 17/35] fix: cli unused import + undeclared crate --- crates/cli/src/main.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index b12c82e..a2acc96 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -1,7 +1,4 @@ -use std::{ - net::{IpAddr, ToSocketAddrs}, - path::PathBuf, -}; +use std::net::{IpAddr, ToSocketAddrs}; use clap::{Parser, Subcommand, ValueEnum}; use gamedig::{ @@ -67,7 +64,7 @@ enum Action { /// Optional file path for packet capture file writer #[cfg(feature = "packet_capture")] #[arg(short, long)] - capture: Option, + capture: Option, /// Optional timeout settings for the server query #[command(flatten, next_help_heading = "Timeouts")] @@ -272,6 +269,7 @@ fn output_result_json_pretty(result: T) { /// Output the result as an XML object. /// # Arguments /// * `result` - A serde serializable result. +#[cfg(feature = "xml")] fn output_result_xml(result: T) { println!("{}", serde_xml_rs::to_string(&result).unwrap()); } From bedd277027a51d9c84b34a1afa7a94cdd5c961bf Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Thu, 18 Jan 2024 18:37:31 +0000 Subject: [PATCH 18/35] fix: variant `Action::Query` does not have this field --- crates/cli/src/main.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index a2acc96..5594593 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -320,6 +320,7 @@ fn main() -> Result<()> { port, format, output_mode, + #[cfg(feature = "packet_capture")] capture, timeout_settings, extra_options, From b1e42f902361c419b61c960414f65f1613be4d73 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Thu, 18 Jan 2024 18:59:19 +0000 Subject: [PATCH 19/35] refactor: clippy fixs --- crates/cli/src/main.rs | 2 +- crates/lib/src/capture/packet.rs | 16 ++++++++-------- crates/lib/src/capture/pcap.rs | 4 ++-- crates/lib/src/socket.rs | 4 ++-- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 5594593..2bafb77 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -304,7 +304,7 @@ fn output_result_bson_base64(result: T) { if let bson::Bson::Document(document) = bson { let bytes = bson::to_vec(&document).unwrap(); - println!("{}", base64::prelude::BASE64_STANDARD.encode(&bytes)); + println!("{}", base64::prelude::BASE64_STANDARD.encode(bytes)); } else { panic!("Failed to convert result to BSON Base64"); } diff --git a/crates/lib/src/capture/packet.rs b/crates/lib/src/capture/packet.rs index f7c3c3c..ef3ae62 100644 --- a/crates/lib/src/capture/packet.rs +++ b/crates/lib/src/capture/packet.rs @@ -24,9 +24,9 @@ pub(crate) enum Direction { #[derive(Clone, Copy, Debug, PartialEq)] pub(crate) enum Protocol { /// Transmission Control Protocol. - TCP, + Tcp, /// User Datagram Protocol. - UDP, + Udp, } /// Trait for handling different types of IP addresses (IPv4, IPv6). @@ -41,7 +41,7 @@ pub(crate) trait IpAddress: Sized { pub(crate) struct CapturePacket<'a> { /// Direction of the packet (Send/Receive). pub(crate) direction: Direction, - /// Protocol of the packet (TCP/UDP). + /// Protocol of the packet (Tcp/UDP). pub(crate) protocol: Protocol, /// Remote socket address. pub(crate) remote_address: &'a SocketAddr, @@ -135,14 +135,14 @@ mod tests { fn test_ports_by_direction() { let packet_send = CapturePacket { direction: Direction::Send, - protocol: Protocol::TCP, + 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, + protocol: Protocol::Tcp, local_address: &socket_addr("127.0.0.1:8080"), remote_address: &socket_addr("192.168.1.1:80"), }; @@ -155,14 +155,14 @@ mod tests { fn test_ip_addr() { let packet_send = CapturePacket { direction: Direction::Send, - protocol: Protocol::TCP, + 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, + protocol: Protocol::Tcp, local_address: &socket_addr("127.0.0.1:8080"), remote_address: &socket_addr("192.168.1.1:80"), }; @@ -187,7 +187,7 @@ mod tests { fn test_ip_by_direction_type_specific() { let packet = CapturePacket { direction: Direction::Send, - protocol: Protocol::TCP, + protocol: Protocol::Tcp, local_address: &socket_addr("127.0.0.1:8080"), remote_address: &socket_addr("192.168.1.1:80"), }; diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs index 0613619..f041b3a 100644 --- a/crates/lib/src/capture/pcap.rs +++ b/crates/lib/src/capture/pcap.rs @@ -56,7 +56,7 @@ impl Pcap { let (source_port, dest_port) = info.ports_by_direction(); match info.protocol { - Protocol::TCP => { + Protocol::Tcp => { let buf_size = { let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); tcp.set_source(source_port); @@ -123,7 +123,7 @@ impl Pcap { vec![EnhancedPacketOption::Comment("Generated TCP ACK".into())], ); } - Protocol::UDP => { + Protocol::Udp => { let buf_size = { let mut udp = MutableUdpPacket::new(&mut buf).unwrap(); udp.set_source(source_port); diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index 4548720..eeba5eb 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -210,13 +210,13 @@ pub mod capture { /// Represents the TCP protocol provider. pub(crate) struct ProtocolTCP; impl ProtocolProvider for ProtocolTCP { - fn protocol() -> Protocol { Protocol::TCP } + fn protocol() -> Protocol { Protocol::Tcp } } /// Represents the UDP protocol provider. pub(crate) struct ProtocolUDP; impl ProtocolProvider for ProtocolUDP { - fn protocol() -> Protocol { Protocol::UDP } + fn protocol() -> Protocol { Protocol::Udp } } /// A socket wrapper that allows capturing packets. From 5dff511e6f21f54ace830087950ddfeeade3a379 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Thu, 18 Jan 2024 18:59:59 +0000 Subject: [PATCH 20/35] refactor: clippy fix + docs --- crates/lib/src/capture/writer.rs | 45 +++++++++++++++++++++++++++++--- 1 file changed, 41 insertions(+), 4 deletions(-) diff --git a/crates/lib/src/capture/writer.rs b/crates/lib/src/capture/writer.rs index 4d63655..d3a9db0 100644 --- a/crates/lib/src/capture/writer.rs +++ b/crates/lib/src/capture/writer.rs @@ -8,15 +8,52 @@ use crate::GDResult; use lazy_static::lazy_static; lazy_static! { + /// A globally accessible, lazily-initialized static writer instance. + /// This writer is intended for capturing and recording network packets. + /// The writer is wrapped in a Mutex to ensure thread-safe access and modification. pub(crate) static ref CAPTURE_WRITER: Mutex>> = Mutex::new(None); } +/// Trait defining the functionality for a writer that handles network packet +/// captures. This trait includes methods for writing packet data, handling new +/// connections, and closing connections. pub(crate) trait Writer { + /// Writes a given packet's data to an underlying storage or stream. + /// + /// # Arguments + /// * `packet` - Reference to the packet being captured. + /// * `data` - The raw byte data associated with the packet. + /// + /// # Returns + /// A `GDResult` indicating the success or failure of the write operation. fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> GDResult<()>; + + /// Handles the creation of a new connection, potentially logging or + /// initializing resources. + /// + /// # Arguments + /// * `packet` - Reference to the packet indicating a new connection. + /// + /// # Returns + /// A `GDResult` indicating the success or failure of handling the new + /// connection. fn new_connect(&mut self, packet: &CapturePacket) -> GDResult<()>; + + /// Closes a connection, handling any necessary cleanup or finalization. + /// + /// # Arguments + /// * `packet` - Reference to the packet indicating the closure of a + /// connection. + /// + /// # Returns + /// A `GDResult` indicating the success or failure of the connection closure + /// operation. fn close_connection(&mut self, packet: &CapturePacket) -> GDResult<()>; } +/// Implementation of the `Writer` trait for the `Pcap` struct. +/// This implementation enables writing, connection handling, and closure +/// specific to PCAP (Packet Capture) format. impl Writer for Pcap { fn write(&mut self, info: &CapturePacket, data: &[u8]) -> GDResult<()> { self.write_transport_packet(info, data); @@ -26,10 +63,10 @@ impl Writer for Pcap { fn new_connect(&mut self, packet: &CapturePacket) -> GDResult<()> { match packet.protocol { - Protocol::TCP => { + Protocol::Tcp => { self.write_tcp_handshake(packet); } - Protocol::UDP => {} + Protocol::Udp => {} } self.state.stream_count = self.state.stream_count.wrapping_add(1); @@ -39,10 +76,10 @@ impl Writer for Pcap { fn close_connection(&mut self, packet: &CapturePacket) -> GDResult<()> { match packet.protocol { - Protocol::TCP => { + Protocol::Tcp => { self.send_tcp_fin(packet); } - Protocol::UDP => {} + Protocol::Udp => {} } Ok(()) } From 2deb1df4aed0d182eae4b36675cbe036672fbb78 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 19 Jan 2024 04:30:20 +0000 Subject: [PATCH 21/35] impl: cli release workflow (x86_64, win, dar, linux, wasm32-wasi) --- .github/workflows/release.yml | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 .github/workflows/release.yml diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..8c1d8fb --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,31 @@ +name: Release + +on: + release: + types: [created] + +jobs: + release: + name: Release ${{ matrix.target }} + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - target: x86_64-pc-windows-gnu + archive: zip + - target: x86_64-unknown-linux-musl + archive: tar.gz tar.xz tar.zst + - target: x86_64-apple-darwin + archive: zip + - target: wasm32-wasi + archive: zip tar.gz + steps: + - uses: actions/checkout@v4 + - name: Compile and release + uses: rust-build/rust-build.action@v1.4.4 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + RUSTTARGET: ${{ matrix.target }} + ARCHIVE_TYPES: ${{ matrix.archive }} From bba9f5f11b056f2ad8b10463a5030631e88496ea Mon Sep 17 00:00:00 2001 From: Douile Date: Mon, 22 Jan 2024 22:02:34 +0000 Subject: [PATCH 22/35] cli: Improve capture help string --- crates/cli/src/main.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 2bafb77..fc44180 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -62,6 +62,13 @@ enum Action { output_mode: OutputMode, /// Optional file path for packet capture file writer + /// + /// When set a PCAP file will be written to the location. This file can + /// be read with a tool like wireshark. The PCAP contains a log of the + /// TCP and UDP data sent/recieved by the gamedig library, it does not + /// contain an accurate representation of the real packets sent on the + /// wire as some information has to be hallucinated in order for it to + /// display nicely. #[cfg(feature = "packet_capture")] #[arg(short, long)] capture: Option, From 8f381f733c8299add06e8421b047d6d423e8c76b Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 2 Feb 2024 00:45:17 +0000 Subject: [PATCH 23/35] chore: change cli header year --- crates/cli/src/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index fc44180..66c50b6 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -22,7 +22,7 @@ const GAMEDIG_HEADER: &str = r" |___/ A command line interface for querying game servers. - Copyright (C) 2022 - 2023 GameDig Organization & Contributors + Copyright (C) 2022 - 2024 GameDig Organization & Contributors Licensed under the MIT license "; From 5365845bb56ca6a0eab75d8326e3c01a73a60496 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 2 Feb 2024 01:44:44 +0000 Subject: [PATCH 24/35] fix: cli xml `LastElementNameNotAvailable` error case --- crates/cli/Cargo.toml | 4 +-- crates/cli/src/error.rs | 6 ++++ crates/cli/src/main.rs | 62 ++++++++++++++++++++++++++++++++++++++++- 3 files changed, 69 insertions(+), 3 deletions(-) diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index 4b5a0d9..c2e1444 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -19,7 +19,7 @@ packet_capture = ["gamedig/packet_capture"] # Output formats bson = ["dep:serde", "dep:bson", "dep:hex", "dep:base64", "gamedig/serde"] json = ["dep:serde", "dep:serde_json", "gamedig/serde"] -xml = ["dep:serde", "dep:serde-xml-rs", "gamedig/serde"] +xml = ["dep:serde", "dep:quick-xml", "gamedig/serde"] # Misc browser = ["dep:webbrowser"] @@ -47,7 +47,7 @@ hex = { version = "0.4.3", optional = true, default-features = false } serde_json = { version = "1", optional = true, default-features = false } # XML -serde-xml-rs = { version = "0.6.0", optional = true, default-features = false } +quick-xml = { version = "0.31.0", optional = true, default-features = false } # Browser webbrowser = { version = "0.8.12", optional = true, default-features = false } diff --git a/crates/cli/src/error.rs b/crates/cli/src/error.rs index d5db795..396abc0 100644 --- a/crates/cli/src/error.rs +++ b/crates/cli/src/error.rs @@ -11,6 +11,12 @@ pub enum Error { #[error("Gamedig Error: {0}")] Gamedig(#[from] gamedig::errors::GDError), + #[error("Serde Error: {0}")] + Serde(#[from] serde_json::Error), + + #[error("Xml Error: {0}")] + Xml(#[from] quick_xml::Error), + #[error("Unknown Game: {0}")] UnknownGame(String), diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 66c50b6..bce0fda 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -278,7 +278,67 @@ fn output_result_json_pretty(result: T) { /// * `result` - A serde serializable result. #[cfg(feature = "xml")] fn output_result_xml(result: T) { - println!("{}", serde_xml_rs::to_string(&result).unwrap()); + use quick_xml::events::{BytesEnd, BytesStart, Event, BytesText}; + use quick_xml::Writer; + use serde_json::Value; + use std::io::Cursor; + + // Serialize the struct to a JSON Value first + let json = serde_json::to_value(result).expect("Failed to serialize struct to JSON"); + + // Create a buffer and a writer for XML output + let buffer = Cursor::new(Vec::new()); + let mut writer = Writer::new(buffer); + + // Recursive function to convert JSON to XML + fn json_to_xml( + writer: &mut Writer, + key: Option<&str>, + value: &Value, + ) -> Result<()> { + match value { + Value::Object(obj) => { + if let Some(key) = key { + writer.write_event(Event::Start(BytesStart::new(key)))?; + } + for (k, v) in obj { + json_to_xml(writer, Some(k), v)?; + } + if let Some(key) = key { + writer.write_event(Event::End(BytesEnd::new(key)))?; + } + }, + Value::Array(arr) => { + for v in arr { + json_to_xml(writer, key.or(Some("item")), v)?; + } + }, + _ => { + if let Some(key) = key { + writer.write_event(Event::Start(BytesStart::new(key)))?; + } + let text_string = match value { + Value::String(s) => s.to_string(), + _ => value.to_string().trim_matches('"').to_string(), + }; + let text = text_string.as_str(); + writer.write_event(Event::Text(BytesText::new(text)))?; + if let Some(key) = key { + writer.write_event(Event::End(BytesEnd::new(key)))?; + } + } + } + Ok(()) + } + + writer.write_event(Event::Start(BytesStart::new("data"))).expect("Failed to write start tag"); + json_to_xml(&mut writer, None, &json).expect("Failed to convert JSON to XML"); + writer.write_event(Event::End(BytesEnd::new("data"))).expect("Failed to write end tag"); + + let xml_bytes = writer.into_inner().into_inner(); + let xml_string = String::from_utf8(xml_bytes).expect("Failed to convert XML bytes to string"); + + println!("{}", xml_string); } /// Output the result as a BSON object encoded as a hex string. From df51521a79100b83e6395279c2dfb38859656783 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 2 Feb 2024 01:58:48 +0000 Subject: [PATCH 25/35] fix: remove bug panic due to xml protocol format now being supported --- crates/cli/src/main.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index bce0fda..4104c03 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -223,11 +223,7 @@ fn output_result(output_mode: OutputMode, format: Ou OutputFormat::Xml => { match output_mode { OutputMode::Generic => output_result_xml(result.as_json()), - // BUG: In this case we get a writer error with all serde write methods some reason with xml 0.6.0 - // BUG: gamedig-cli.exe query -g -i -p -f xml --output-mode protocol-specific - // BUG: Writer: emitter error: document start event has already been emitted - // BUG: With xml 0.5.1 we get unsupported operation: 'serialize_unit_variant' - OutputMode::ProtocolSpecific => panic!("XML format is not supported for protocol specific output"), + OutputMode::ProtocolSpecific => output_result_xml(result.as_original()), } } #[cfg(feature = "bson")] From 422cb57efa4eedb0b517ff0a94371efec4308e78 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 2 Feb 2024 02:00:15 +0000 Subject: [PATCH 26/35] chore: update cli default features --- crates/cli/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index c2e1444..08f8e2f 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -11,7 +11,7 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [features] -default = ["packet_capture", "bson", "json", "xml", "browser"] +default = ["json", "bson", "xml", "browser"] # Tools packet_capture = ["gamedig/packet_capture"] From 5310200181b71eadaea6da4840b79d2f55ebf535 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 2 Feb 2024 02:01:25 +0000 Subject: [PATCH 27/35] fix: cli xml feat not having json dep --- crates/cli/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index 08f8e2f..501dab3 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -19,7 +19,7 @@ packet_capture = ["gamedig/packet_capture"] # Output formats bson = ["dep:serde", "dep:bson", "dep:hex", "dep:base64", "gamedig/serde"] json = ["dep:serde", "dep:serde_json", "gamedig/serde"] -xml = ["dep:serde", "dep:quick-xml", "gamedig/serde"] +xml = ["dep:serde", "dep:serde_json", "dep:quick-xml", "gamedig/serde"] # Misc browser = ["dep:webbrowser"] From 744230455ca4041f90a5c6176d22b1a5009b2d03 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 2 Feb 2024 02:34:02 +0000 Subject: [PATCH 28/35] revert: lib default feat to exclude packet capture --- crates/lib/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/lib/Cargo.toml b/crates/lib/Cargo.toml index 15e370f..44592e1 100644 --- a/crates/lib/Cargo.toml +++ b/crates/lib/Cargo.toml @@ -17,7 +17,7 @@ rust-version = "1.65.0" categories = ["parser-implementations", "parsing", "network-programming", "encoding"] [features] -default = ["games", "services", "game_defs", "packet_capture",] +default = ["games", "services", "game_defs"] games = [] services = [] game_defs = ["dep:phf", "games"] From 4675b24ff3d80c089e260e7c9e858dab52283cfe Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 2 Feb 2024 20:57:07 +0000 Subject: [PATCH 29/35] refactor: better xml support and error prop --- crates/cli/src/error.rs | 6 ++ crates/cli/src/main.rs | 139 ++++++++++++++++++++++++++-------------- 2 files changed, 97 insertions(+), 48 deletions(-) diff --git a/crates/cli/src/error.rs b/crates/cli/src/error.rs index 396abc0..1700827 100644 --- a/crates/cli/src/error.rs +++ b/crates/cli/src/error.rs @@ -11,9 +11,15 @@ pub enum Error { #[error("Gamedig Error: {0}")] Gamedig(#[from] gamedig::errors::GDError), + #[cfg(any(feature = "json", feature = "xml"))] #[error("Serde Error: {0}")] Serde(#[from] serde_json::Error), + #[cfg(feature = "bson")] + #[error("Bson Error: {0}")] + Bson(#[from] bson::ser::Error), + + #[cfg(feature = "xml")] #[error("Xml Error: {0}")] Xml(#[from] quick_xml::Error), diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 4104c03..ed19a3e 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -203,42 +203,42 @@ fn output_result(output_mode: OutputMode, format: Ou match output_mode { OutputMode::Generic => output_result_debug(result.as_json()), OutputMode::ProtocolSpecific => output_result_debug(result.as_original()), - } + }; } #[cfg(feature = "json")] OutputFormat::JsonPretty => { - match output_mode { + let _ = match output_mode { OutputMode::Generic => output_result_json_pretty(result.as_json()), OutputMode::ProtocolSpecific => output_result_json_pretty(result.as_original()), - } + }; } #[cfg(feature = "json")] OutputFormat::Json => { - match output_mode { + let _ = match output_mode { OutputMode::Generic => output_result_json(result.as_json()), OutputMode::ProtocolSpecific => output_result_json(result.as_original()), - } + }; } #[cfg(feature = "xml")] OutputFormat::Xml => { - match output_mode { + let _ = match output_mode { OutputMode::Generic => output_result_xml(result.as_json()), OutputMode::ProtocolSpecific => output_result_xml(result.as_original()), - } + }; } #[cfg(feature = "bson")] OutputFormat::BsonHex => { - match output_mode { + let _ = match output_mode { OutputMode::Generic => output_result_bson_hex(result.as_json()), OutputMode::ProtocolSpecific => output_result_bson_hex(result.as_original()), - } + }; } #[cfg(feature = "bson")] OutputFormat::BsonBase64 => { - match output_mode { + let _ = match output_mode { OutputMode::Generic => output_result_bson_base64(result.as_json()), OutputMode::ProtocolSpecific => output_result_bson_base64(result.as_original()), - } + }; } } } @@ -256,8 +256,10 @@ fn output_result_debug(result: R) { /// # Arguments /// * `result` - A serde serializable result. #[cfg(feature = "json")] -fn output_result_json(result: T) { - println!("{}", serde_json::to_string(&result).unwrap()); +fn output_result_json(result: T) -> Result<()> { + println!("{}", serde_json::to_string(&result)?); + + Ok(()) } /// Output the result as a pretty printed JSON object. @@ -265,61 +267,92 @@ fn output_result_json(result: T) { /// # Arguments /// * `result` - A serde serializable result. #[cfg(feature = "json")] -fn output_result_json_pretty(result: T) { - println!("{}", serde_json::to_string_pretty(&result).unwrap()); +fn output_result_json_pretty(result: T) -> Result<()> { + println!("{}", serde_json::to_string_pretty(&result)?); + + Ok(()) } /// Output the result as an XML object. /// # Arguments /// * `result` - A serde serializable result. #[cfg(feature = "xml")] -fn output_result_xml(result: T) { - use quick_xml::events::{BytesEnd, BytesStart, Event, BytesText}; - use quick_xml::Writer; +fn output_result_xml(result: T) -> Result<()> { + use quick_xml::{ + events::{BytesEnd, BytesStart, BytesText, Event}, + Writer, + }; use serde_json::Value; - use std::io::Cursor; - // Serialize the struct to a JSON Value first - let json = serde_json::to_value(result).expect("Failed to serialize struct to JSON"); + // Serialize the input `result` of generic type `T` into a JSON value. + // This step converts the Rust data structure into a JSON format, + // which will then be used to generate the corresponding XML. + let json = serde_json::to_value(result)?; - // Create a buffer and a writer for XML output - let buffer = Cursor::new(Vec::new()); - let mut writer = Writer::new(buffer); + // Initialize the XML writer with a new, empty vector to store the XML data. + let mut writer = Writer::new(Vec::new()); - // Recursive function to convert JSON to XML - fn json_to_xml( - writer: &mut Writer, - key: Option<&str>, - value: &Value, - ) -> Result<()> { + // Define a recursive function `json_to_xml` to convert the JSON value into XML + // format. The function takes a mutable reference to the XML writer, an + // optional key as a string slice, and a reference to the JSON value to be + // converted. + fn json_to_xml(writer: &mut Writer, key: Option<&str>, value: &Value) -> Result<()> { match value { + // If the JSON value is an object, iterate through its properties, + // creating XML elements with corresponding keys and values. Value::Object(obj) => { if let Some(key) = key { + // Start an XML element for the object. writer.write_event(Event::Start(BytesStart::new(key)))?; } + for (k, v) in obj { + // Recursively process each property of the object. json_to_xml(writer, Some(k), v)?; } + if let Some(key) = key { + // Close the XML element for the object. writer.write_event(Event::End(BytesEnd::new(key)))?; } - }, + } + + // If the JSON value is an array, iterate through its elements, + // creating XML elements for each item. Value::Array(arr) => { for v in arr { + // Use "item" as the default key for array elements without keys. json_to_xml(writer, key.or(Some("item")), v)?; } - }, + } + + // If the JSON value is null, create an empty XML element. + Value::Null => { + if let Some(key) = key { + writer.write_event(Event::Empty(BytesStart::new(key)))?; + } + } + + // For all other JSON value types (String, Number, Bool), + // convert the value to a string and create an XML element with the text content. + // Note: We handle null strings here as well, as they are treated as a string type. _ => { if let Some(key) = key { + // Start the XML element with the given key. writer.write_event(Event::Start(BytesStart::new(key)))?; } + + // Convert the JSON value to a string, trimming quotes for non-string values. let text_string = match value { Value::String(s) => s.to_string(), - _ => value.to_string().trim_matches('"').to_string(), + _ => value.to_string().trim_matches('"').to_string(), }; - let text = text_string.as_str(); - writer.write_event(Event::Text(BytesText::new(text)))?; + + // Create a text node with the converted string value. + writer.write_event(Event::Text(BytesText::new(&text_string)))?; + if let Some(key) = key { + // Close the XML element. writer.write_event(Event::End(BytesEnd::new(key)))?; } } @@ -327,14 +360,20 @@ fn output_result_xml(result: T) { Ok(()) } - writer.write_event(Event::Start(BytesStart::new("data"))).expect("Failed to write start tag"); - json_to_xml(&mut writer, None, &json).expect("Failed to convert JSON to XML"); - writer.write_event(Event::End(BytesEnd::new("data"))).expect("Failed to write end tag"); + // Start the root XML element named "data". + writer.write_event(Event::Start(BytesStart::new("data")))?; + // Convert the top-level JSON value to XML. + json_to_xml(&mut writer, None, &json)?; + // Close the root XML element. + writer.write_event(Event::End(BytesEnd::new("data")))?; - let xml_bytes = writer.into_inner().into_inner(); - let xml_string = String::from_utf8(xml_bytes).expect("Failed to convert XML bytes to string"); + // Convert the XML data stored in the writer to a UTF-8 string. + let xml_bytes = writer.into_inner(); + let xml_string = String::from_utf8(xml_bytes).expect("Failed to convert XML bytes to UTF-8 string"); println!("{}", xml_string); + + Ok(()) } /// Output the result as a BSON object encoded as a hex string. @@ -342,15 +381,17 @@ fn output_result_xml(result: T) { /// # Arguments /// * `result` - A serde serializable result. #[cfg(feature = "bson")] -fn output_result_bson_hex(result: T) { - let bson = bson::to_bson(&result).unwrap(); +fn output_result_bson_hex(result: T) -> Result<()> { + let bson = bson::to_bson(&result)?; if let bson::Bson::Document(document) = bson { - let bytes = bson::to_vec(&document).unwrap(); + let bytes = bson::to_vec(&document)?; println!("{}", hex::encode(bytes)); + + Ok(()) } else { - panic!("Failed to convert result to BSON Hex"); + panic!("Failed to convert result to BSON Hex (BSON_DOCUMENT_UNAVAILABLE)"); } } @@ -359,17 +400,19 @@ fn output_result_bson_hex(result: T) { /// # Arguments /// * `result` - A serde serializable result. #[cfg(feature = "bson")] -fn output_result_bson_base64(result: T) { +fn output_result_bson_base64(result: T) -> Result<()> { use base64::Engine; - let bson = bson::to_bson(&result).unwrap(); + let bson = bson::to_bson(&result)?; if let bson::Bson::Document(document) = bson { - let bytes = bson::to_vec(&document).unwrap(); + let bytes = bson::to_vec(&document)?; println!("{}", base64::prelude::BASE64_STANDARD.encode(bytes)); + + Ok(()) } else { - panic!("Failed to convert result to BSON Base64"); + panic!("Failed to convert result to BSON Base64 (BSON_DOCUMENT_UNAVAILABLE)"); } } From 48aa5115c024d3973aad0f015e197f972ca70fb4 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Fri, 2 Feb 2024 21:10:41 +0000 Subject: [PATCH 30/35] fix: add xml utf8 and semver declaration --- crates/cli/src/main.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index ed19a3e..f20be84 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -279,7 +279,7 @@ fn output_result_json_pretty(result: T) -> Result<()> { #[cfg(feature = "xml")] fn output_result_xml(result: T) -> Result<()> { use quick_xml::{ - events::{BytesEnd, BytesStart, BytesText, Event}, + events::{BytesDecl, BytesEnd, BytesStart, BytesText, Event}, Writer, }; use serde_json::Value; @@ -292,6 +292,9 @@ fn output_result_xml(result: T) -> Result<()> { // Initialize the XML writer with a new, empty vector to store the XML data. let mut writer = Writer::new(Vec::new()); + // Write the XML 1.1 declaration + writer.write_event(Event::Decl(BytesDecl::new("1.1", Some("utf-8"), None)))?; + // Define a recursive function `json_to_xml` to convert the JSON value into XML // format. The function takes a mutable reference to the XML writer, an // optional key as a string slice, and a reference to the JSON value to be From c34392a3da46aa6f0f3036045afd7e3cda4f4e54 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Wed, 7 Feb 2024 15:26:17 +0000 Subject: [PATCH 31/35] refactor: pcap packet buffer size --- crates/lib/src/capture/pcap.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs index f041b3a..005196d 100644 --- a/crates/lib/src/capture/pcap.rs +++ b/crates/lib/src/capture/pcap.rs @@ -21,13 +21,7 @@ use super::packet::{ PACKET_SIZE, }; -const BUFFER_SIZE: usize = PACKET_SIZE - - (if HEADER_SIZE_IP4 > HEADER_SIZE_IP6 { - HEADER_SIZE_IP4 - } else { - HEADER_SIZE_IP6 - }) - - HEADER_SIZE_ETHERNET; +const BUFFER_SIZE: usize = PACKET_SIZE - HEADER_SIZE_IP6 - HEADER_SIZE_ETHERNET; pub(crate) struct Pcap { writer: PcapNgWriter, From 9a2b953fff13cf6a940f08bff7a8f0094ca56af1 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Wed, 7 Feb 2024 15:28:41 +0000 Subject: [PATCH 32/35] refactor: capture socket and move to capture dir --- crates/lib/src/capture/mod.rs | 3 +- crates/lib/src/capture/socket.rs | 214 +++++++++++++++++++++++++++++ crates/lib/src/socket.rs | 226 +------------------------------ 3 files changed, 219 insertions(+), 224 deletions(-) create mode 100644 crates/lib/src/capture/socket.rs 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 { From 7b37e7122166767c47139a0c72d199f13368d133 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Wed, 7 Feb 2024 15:29:13 +0000 Subject: [PATCH 33/35] chore: remove unused else --- crates/lib/src/capture/mod.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/crates/lib/src/capture/mod.rs b/crates/lib/src/capture/mod.rs index 1e8581a..583ebd8 100644 --- a/crates/lib/src/capture/mod.rs +++ b/crates/lib/src/capture/mod.rs @@ -29,9 +29,6 @@ pub fn setup_capture(file_path: Option) { let writer = Box::new(Pcap::new(pcap_writer)); attach(writer) - } else { - // If no file path is provided - // Do nothing } } From 730c938ad243d8cd1cfb50a84c1a4617869b4d77 Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Wed, 7 Feb 2024 15:30:20 +0000 Subject: [PATCH 34/35] chore: format --- crates/lib/src/socket.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index 77d4c40..6161592 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -168,7 +168,7 @@ impl Socket for UdpSocketImpl { } #[cfg(not(feature = "packet_capture"))] -pub type UdpSocket = UdpSocketImpl; +pub type UdpSocket = UdpSocketImpl; #[cfg(not(feature = "packet_capture"))] pub type TcpSocket = TcpSocketImpl; From 7369dbab19c47d28e3646482cd15622baad2786d Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Wed, 7 Feb 2024 16:10:04 +0000 Subject: [PATCH 35/35] refactor: use slices where possable in pcap --- crates/lib/src/capture/pcap.rs | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs index 005196d..acb653f 100644 --- a/crates/lib/src/capture/pcap.rs +++ b/crates/lib/src/capture/pcap.rs @@ -45,14 +45,15 @@ impl Pcap { } pub(crate) fn write_transport_packet(&mut self, info: &CapturePacket, payload: &[u8]) { - let mut buf = vec![0; BUFFER_SIZE]; + let mut buffer_array: [u8; BUFFER_SIZE] = [0; BUFFER_SIZE]; + let buf: &mut [u8] = &mut buffer_array[..]; let (source_port, dest_port) = info.ports_by_direction(); match info.protocol { Protocol::Tcp => { let buf_size = { - let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); + let mut tcp = MutableTcpPacket::new(buf).unwrap(); tcp.set_source(source_port); tcp.set_destination(dest_port); tcp.set_payload(payload); @@ -86,7 +87,7 @@ impl Pcap { let mut info = info.clone(); let buf_size = { - let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); + let mut tcp = MutableTcpPacket::new(buf).unwrap(); tcp.set_source(dest_port); tcp.set_destination(source_port); tcp.set_data_offset(5); @@ -119,7 +120,7 @@ impl Pcap { } Protocol::Udp => { let buf_size = { - let mut udp = MutableUdpPacket::new(&mut buf).unwrap(); + let mut udp = MutableUdpPacket::new(buf).unwrap(); udp.set_source(source_port); udp.set_destination(dest_port); udp.set_length((payload.len() + HEADER_SIZE_UDP) as u16); @@ -216,7 +217,8 @@ impl Pcap { let mut info = info.clone(); info.direction = Direction::Send; - let mut buf = vec![0; PACKET_SIZE]; + let mut buffer_array: [u8; BUFFER_SIZE] = [0; BUFFER_SIZE]; + let buf: &mut [u8] = &mut buffer_array[..]; // Add a generated comment to all packets let options = vec![ pcap_file::pcapng::blocks::enhanced_packet::EnhancedPacketOption::Comment("Generated TCP handshake".into()), @@ -224,7 +226,7 @@ impl Pcap { // SYN let buf_size = { - let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); + let mut tcp = MutableTcpPacket::new(buf).unwrap(); self.state.send_seq = 500; tcp.set_sequence(self.state.send_seq); tcp.set_flags(TcpFlags::SYN); @@ -245,7 +247,7 @@ impl Pcap { // SYN + ACK info.direction = Direction::Receive; let buf_size = { - let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); + let mut tcp = MutableTcpPacket::new(buf).unwrap(); self.state.send_seq = self.state.send_seq.wrapping_add(1); tcp.set_acknowledgement(self.state.send_seq); self.state.rec_seq = 1000; @@ -268,7 +270,7 @@ impl Pcap { // ACK info.direction = Direction::Send; let buf_size = { - let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); + let mut tcp = MutableTcpPacket::new(buf).unwrap(); tcp.set_sequence(self.state.send_seq); self.state.rec_seq = self.state.rec_seq.wrapping_add(1); tcp.set_acknowledgement(self.state.rec_seq); @@ -291,11 +293,12 @@ impl Pcap { } pub(crate) fn send_tcp_fin(&mut self, info: &CapturePacket) { - let mut buf = vec![0; BUFFER_SIZE]; + let mut buffer_array: [u8; BUFFER_SIZE] = [0; BUFFER_SIZE]; + let buf: &mut [u8] = &mut buffer_array[..]; let (source_port, dest_port) = info.ports_by_direction(); let buf_size = { - let mut tcp = MutableTcpPacket::new(&mut buf).unwrap(); + let mut tcp = MutableTcpPacket::new(buf).unwrap(); tcp.set_source(source_port); tcp.set_destination(dest_port); tcp.set_data_offset(5);