From abbcae618f95289021b6ad623dc1a366ae653087 Mon Sep 17 00:00:00 2001 From: Douile Date: Wed, 22 Nov 2023 02:20:03 +0000 Subject: [PATCH] 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;