refactor: backport pcap impl

This commit is contained in:
Cain 2024-01-13 03:22:50 +00:00 committed by Douile
parent 49096e46bb
commit 3d47180e85
No known key found for this signature in database
GPG key ID: E048586A5FF6585C
7 changed files with 272 additions and 235 deletions

View file

@ -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.

View file

@ -18,7 +18,7 @@ pub fn setup_capture(file_name: Option<String>) {
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,

View file

@ -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<T: IpAddress>(&self) -> (T, T) {
let (local, remote) = (
T::from_std(self.local_address.ip()).expect("Incorrect IP type for local address"),

View file

@ -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<W: Write> {
writer: PcapNgWriter<W>,
pub(crate) state: State,
buffer: Vec<u8>,
}
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<W: Write> Pcap<W> {
pub fn new(writer: PcapNgWriter<W>) -> Self {
pub(crate) fn new(writer: PcapNgWriter<W>) -> Self {
Self {
writer,
state: State::default(),
buffer: vec![0; BUFFER_SIZE],
}
}
pub fn write_transport_packet(&mut self, info: &CapturePacket, payload: &[u8]) -> Result<(), std::io::Error> {
let (source_port, dest_port) = info.ports_by_direction();
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<usize, std::io::Error> {
let mut tcp = MutableTcpPacket::new(&mut self.buffer)
.ok_or_else(|| std::io::Error::new(std::io::ErrorKind::Other, "Failed to create TCP packet"))?;
tcp.set_source(source_port);
tcp.set_destination(dest_port);
tcp.set_payload(payload);
tcp.set_data_offset(5);
tcp.set_window(TCP_WINDOW_SIZE);
// Set sequence and acknowledgement numbers
match info.direction {
Direction::Send => {
tcp.set_sequence(self.state.send_seq);
tcp.set_acknowledgement(self.state.rec_seq);
self.state.send_seq = self.state.send_seq.wrapping_add(payload.len() as u32);
}
Direction::Receive => {
tcp.set_sequence(self.state.rec_seq);
tcp.set_acknowledgement(self.state.send_seq);
self.state.rec_seq = self.state.rec_seq.wrapping_add(payload.len() as u32);
}
}
tcp.set_flags(TcpFlags::PSH | TcpFlags::ACK);
Ok(tcp.packet_size())
}
pub fn write_tcp_handshake(&mut self, info: &CapturePacket) {
// Initialize sequence numbers for demonstration purposes
self.state.send_seq = 500;
self.state.rec_seq = 1000;
// Common setup for TCP handshake packets
let mut tcp_handshake_packet =
|info: &CapturePacket, direction: Direction, flags: u8| -> Result<(), std::io::Error> {
let (source_port, dest_port) = info.ports_by_direction();
let adjusted_info = CapturePacket {
direction,
..info.clone()
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<EnhancedPacketOption>,
options: Vec<pcap_file::pcapng::blocks::enhanced_packet::EnhancedPacketOption>,
) {
let network_packet_size = self.encode_ip_packet(info, protocol, payload).unwrap().0;
let ethertype = self.encode_ip_packet(info, protocol, payload).unwrap().1;
let ethernet_packet_size = self.encode_ethernet_packet(info, ethertype, &self.buffer[..network_packet_size]).unwrap();
let 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<usize, std::io::Error> {
let mut ethernet_packet = MutableEthernetPacket::new(&mut self.buffer).ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::Other,
"Failed to create Ethernet packet",
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,
}
}

View file

@ -14,7 +14,7 @@ lazy_static! {
pub(crate) static ref CAPTURE_WRITER: Mutex<Option<Box<dyn Writer + Send + Sync>>> = 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<()>;
}

View file

@ -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")]

View file

@ -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<UdpSocketImpl, ProtocolUDP>;
pub(crate) type CapturedUdpSocket = WrappedCaptureSocket<UdpSocketImpl, ProtocolUDP>;
/// 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<TcpSocketImpl, ProtocolTCP>;
pub(crate) type CapturedTcpSocket = WrappedCaptureSocket<TcpSocketImpl, ProtocolTCP>;
}
#[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 {