tribufu/rust-gamedig
1
0
Fork 0
mirror of https://github.com/tribufu/rust-gamedig synced 2026-06-01 09:42:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor: clean up

Browse source
This commit is contained in:
Cain 2024-01-13 09:18:56 +00:00 • committed by Douile
parent 3d47180e85
commit 0543cabce2
No known key found for this signature in database
GPG key ID: E048586A5FF6585C
6 changed files with 137 additions and 171 deletions
Download patch file Download diff file Expand all files Collapse all files

94
crates/cli/src/main.rs
View file

@ -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::{ use gamedig::{
games::*, games::*,
protocols::types::{CommonResponse, ExtraRequestSettings, TimeoutSettings}, protocols::types::{CommonResponse, ExtraRequestSettings, TimeoutSettings},
@ -13,8 +16,30 @@ use self::error::{Error, Result};
// NOTE: For some reason without setting long_about here the doc comment for // NOTE: For some reason without setting long_about here the doc comment for
// ExtraRequestSettings gets set as the about for the CLI. // ExtraRequestSettings gets set as the about for the CLI.
#[derive(Debug, Parser)] #[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 { struct Cli {
#[command(subcommand)]
action: Action,
}
#[derive(Subcommand, Debug)]
enum Action {
/// Query game server information
Query {
/// Unique identifier of the game for which server information is being /// Unique identifier of the game for which server information is being
/// queried. /// queried.
#[arg(short, long)] #[arg(short, long)]
@ -29,26 +54,30 @@ struct Cli {
#[arg(short, long)] #[arg(short, long)]
port: Option<u16>, port: Option<u16>,
/// Flag indicating if the output should be in JSON format. /// Flag indicating if the output should be in JSON format for programmatic use.
#[cfg(feature = "json")] #[cfg(feature = "json")]
#[arg(short, long)] #[arg(short, long)]
json: bool, json: bool,
/// Which response variant to use when outputting. /// Which response variant to use when outputting
#[arg(short, long, default_value = "generic")] #[arg(short, long, default_value = "generic")]
output_mode: OutputMode, output_mode: OutputMode,
/// Optional file path for packet capture file writer
#[cfg(feature = "packet_capture")] #[cfg(feature = "packet_capture")]
#[arg(short, long)] #[arg(short, long)]
capture: Option<String>, capture: Option<PathBuf>,
/// Optional timeout settings for the server query. /// Optional timeout settings for the server query
#[command(flatten, next_help_heading = "Timeouts")] #[command(flatten, next_help_heading = "Timeouts")]
timeout_settings: Option<TimeoutSettings>, timeout_settings: Option<TimeoutSettings>,
/// Optional extra settings for the server query. /// Optional extra settings for the server query
#[command(flatten, next_help_heading = "Query options")] #[command(flatten, next_help_heading = "Query options")]
extra_options: Option<ExtraRequestSettings>, extra_options: Option<ExtraRequestSettings>,
},
/// Display the MIT License information
License,
} }
#[derive(Clone, Debug, PartialEq, Eq, ValueEnum)] #[derive(Clone, Debug, PartialEq, Eq, ValueEnum)]
@ -139,12 +168,12 @@ fn set_hostname_if_missing(host: &str, extra_options: &mut Option<ExtraRequestSe
/// # Arguments /// # Arguments
/// * `args` - A reference to the command line options. /// * `args` - A reference to the command line options.
/// * `result` - A reference to the result of the query. /// * `result` - A reference to the result of the query.
fn output_result(args: &Cli, result: &dyn CommonResponse) { fn output_result(output: OutputMode, json: bool, result: &dyn CommonResponse) {
match args.output_mode { match output {
#[cfg(feature = "json")] #[cfg(feature = "json")]
OutputMode::Generic if args.json => output_result_json(result.as_json()), OutputMode::Generic if json => output_result_json(result.as_json()),
#[cfg(feature = "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::Generic => output_result_debug(result.as_json()),
OutputMode::ProtocolSpecific => output_result_debug(result.as_original()), OutputMode::ProtocolSpecific => output_result_debug(result.as_original()),
@ -169,27 +198,36 @@ fn output_result_json<R: serde::Serialize>(result: R) {
} }
fn main() -> Result<()> { fn main() -> Result<()> {
// Parse the command line arguments
let args = Cli::parse(); let args = Cli::parse();
// Retrieve the game based on the provided ID match args.action {
let game = find_game(&args.game)?; Action::Query {
game,
// Extract extra options for use in setup ip,
let mut extra_options = args.extra_options.clone(); port,
json,
// Resolve the IP address output_mode,
let ip = resolve_ip_or_domain(&args.ip, &mut extra_options)?; 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)?;
#[cfg(feature = "packet_capture")] #[cfg(feature = "packet_capture")]
gamedig::capture::setup_capture(args.capture.clone()); gamedig::capture::setup_capture(capture);
// Query the server using game definition, parsed IP, and user command line let result = query_with_timeout_and_extra_settings(game, &ip, port, timeout_settings, extra_options)?;
// flags. output_result(output_mode, json, result.as_ref());
let result = query_with_timeout_and_extra_settings(game, &ip, args.port, args.timeout_settings, extra_options)?; }
Action::License => {
// Output the result in the specified format // Bake the license into the binary
output_result(&args, result.as_ref()); // so we don't have to ship it separately
println!("{}", include_str!("../../../LICENSE.md"));
}
}
Ok(()) Ok(())
} }

22
crates/lib/src/capture/mod.rs
View file

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

94
crates/lib/src/capture/packet.rs
View file

@ -30,7 +30,7 @@ pub(crate) enum Protocol {
} }
/// Trait for handling different types of IP addresses (IPv4, IPv6). /// 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. /// Creates an instance from a standard `IpAddr`, returning `None` if the types are incompatible.
fn from_std(ip: IpAddr) -> Option<Self>; fn from_std(ip: IpAddr) -> Option<Self>;
} }
@ -53,8 +53,7 @@ impl CapturePacket<'_> {
/// ///
/// Returns: /// Returns:
/// - (u16, u16): Tuple of (source port, destination port). /// - (u16, u16): Tuple of (source port, destination port).
#[allow(dead_code)] pub(super) const fn ports_by_direction(&self) -> (u16, u16) {
pub(crate) const fn ports_by_direction(&self) -> (u16, u16) {
let (local, remote) = (self.local_address.port(), self.remote_address.port()); let (local, remote) = (self.local_address.port(), self.remote_address.port());
self.direction.order(local, remote) self.direction.order(local, remote)
} }
@ -63,30 +62,18 @@ impl CapturePacket<'_> {
/// ///
/// Returns: /// Returns:
/// - (IpAddr, IpAddr): Tuple of (local IP, remote IP). /// - (IpAddr, IpAddr): Tuple of (local IP, remote IP).
#[allow(dead_code)] pub(super) fn ip_addr(&self) -> (IpAddr, IpAddr) {
pub(crate) fn ip_addr(&self) -> (IpAddr, IpAddr) {
let (local, remote) = (self.local_address.ip(), self.remote_address.ip()); let (local, remote) = (self.local_address.ip(), self.remote_address.ip());
(local, remote) (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. /// 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: /// 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.
#[allow(dead_code)] pub(super) fn ipvt_by_direction<T: IpAddress>(&self) -> (T, T) {
pub(crate) fn ipvt_by_direction<T: IpAddress>(&self) -> (T, T) {
let (local, remote) = ( let (local, remote) = (
T::from_std(self.local_address.ip()).expect("Incorrect IP type for local address"), 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"), T::from_std(self.remote_address.ip()).expect("Incorrect IP type for remote address"),
@ -101,7 +88,7 @@ impl Direction {
/// ///
/// Returns: /// Returns:
/// - (T, T): Ordered tuple (source, destination). /// - (T, T): Ordered tuple (source, destination).
pub(crate) const fn order<T>(&self, source: T, remote: T) -> (T, T) { pub(self) const fn order<T>(&self, source: T, remote: T) -> (T, T) {
match self { match self {
Direction::Send => (source, remote), Direction::Send => (source, remote),
Direction::Receive => (remote, source), Direction::Receive => (remote, source),
@ -162,65 +149,33 @@ mod tests {
} }
#[test] #[test]
fn test_ip_addr_by_direction_ipv4() { fn test_ip_addr() {
let packet_send = CapturePacket { let packet_send = CapturePacket {
direction: Direction::Send, direction: Direction::Send,
protocol: Protocol::UDP, protocol: Protocol::TCP,
local_address: &socket_addr("10.0.0.1:3000"), local_address: &socket_addr("127.0.0.1:8080"),
remote_address: &socket_addr("10.0.0.2:3001"), remote_address: &socket_addr("192.168.1.1:80"),
}; };
let packet_receive = CapturePacket { let packet_receive = CapturePacket {
direction: Direction::Receive, direction: Direction::Receive,
protocol: Protocol::UDP, protocol: Protocol::TCP,
local_address: &socket_addr("10.0.0.1:3000"), local_address: &socket_addr("127.0.0.1:8080"),
remote_address: &socket_addr("10.0.0.2:3001"), remote_address: &socket_addr("192.168.1.1:80"),
}; };
assert_eq!( assert_eq!(
packet_send.ip_addr_by_direction(), packet_send.ip_addr(),
( (
IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)) IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))
) )
); );
assert_eq!( assert_eq!(
packet_receive.ip_addr_by_direction(), packet_receive.ip_addr(),
( (
IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)), IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)) IpAddr::V4(Ipv4Addr::new(192, 168, 1, 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))
) )
); );
} }
@ -242,17 +197,4 @@ mod tests {
std::panic::catch_unwind(|| packet.ipvt_by_direction::<Ipv6Addr>()); std::panic::catch_unwind(|| packet.ipvt_by_direction::<Ipv6Addr>());
assert!(ipv6_result.is_err()); 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();
}
} }

23
crates/lib/src/capture/pcap.rs
View file

@ -15,7 +15,6 @@ use super::packet::{
PACKET_SIZE, PACKET_SIZE,
}; };
const BUFFER_SIZE: usize = PACKET_SIZE const BUFFER_SIZE: usize = PACKET_SIZE
- (if HEADER_SIZE_IP4 > HEADER_SIZE_IP6 { - (if HEADER_SIZE_IP4 > HEADER_SIZE_IP6 {
HEADER_SIZE_IP4 HEADER_SIZE_IP4
@ -48,10 +47,7 @@ impl<W: Write> Pcap<W> {
pub(crate) fn write_transport_packet(&mut self, info: &CapturePacket, payload: &[u8]) { pub(crate) fn write_transport_packet(&mut self, info: &CapturePacket, payload: &[u8]) {
let mut buf = vec![0; BUFFER_SIZE]; let mut buf = vec![0; BUFFER_SIZE];
let (source_port, dest_port) = match info.direction { let (source_port, dest_port) = info.ports_by_direction();
Direction::Send => (info.local_address.port(), info.remote_address.port()),
Direction::Receive => (info.remote_address.port(), info.local_address.port()),
};
match info.protocol { match info.protocol {
Protocol::TCP => { Protocol::TCP => {
@ -150,13 +146,9 @@ impl<W: Write> Pcap<W> {
protocol: IpNextHeaderProtocol, protocol: IpNextHeaderProtocol,
payload: &[u8], payload: &[u8],
) -> (usize, EtherType) { ) -> (usize, EtherType) {
match (info.local_address.ip(), info.remote_address.ip()) { match info.ip_addr() {
(IpAddr::V4(local_address), IpAddr::V4(remote_address)) => { (IpAddr::V4(_), IpAddr::V4(_)) => {
let (source, destination) = if info.direction == Direction::Send { let (source, destination) = info.ipvt_by_direction();
(local_address, remote_address)
} else {
(remote_address, local_address)
};
let header_size = HEADER_SIZE_IP4 + (32 / 8); let header_size = HEADER_SIZE_IP4 + (32 / 8);
@ -185,11 +177,8 @@ impl<W: Write> Pcap<W> {
(ip.packet_size(), pnet_packet::ethernet::EtherTypes::Ipv4) (ip.packet_size(), pnet_packet::ethernet::EtherTypes::Ipv4)
} }
(IpAddr::V6(local_address), IpAddr::V6(remote_address)) => { (IpAddr::V6(_), IpAddr::V6(_)) => {
let (source, destination) = match info.direction { let (source, destination) = info.ipvt_by_direction();
Direction::Send => (local_address, remote_address),
Direction::Receive => (remote_address, local_address),
};
let mut ip = MutableIpv6Packet::new(buf).unwrap(); let mut ip = MutableIpv6Packet::new(buf).unwrap();
ip.set_version(6); ip.set_version(6);

18
crates/lib/src/capture/writer.rs
View file

@ -1,22 +1,20 @@
use std::io::Write; use std::{io::Write, sync::Mutex};
use crate::{ use super::{
capture::packet::{CapturePacket, Protocol}, packet::{CapturePacket, Protocol},
GDResult, pcap::Pcap,
}; };
use crate::GDResult;
use super::pcap::Pcap;
use lazy_static::lazy_static; use lazy_static::lazy_static;
use std::sync::Mutex;
lazy_static! { lazy_static! {
pub(crate) static ref CAPTURE_WRITER: Mutex<Option<Box<dyn Writer + Send + Sync>>> = Mutex::new(None); pub(crate) static ref CAPTURE_WRITER: Mutex<Option<Box<dyn Writer + Send + Sync>>> = Mutex::new(None);
} }
pub(crate) trait Writer { pub(crate) trait Writer {
fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> crate::GDResult<()>; fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> GDResult<()>;
fn new_connect(&mut self, packet: &CapturePacket) -> crate::GDResult<()>; fn new_connect(&mut self, packet: &CapturePacket) -> GDResult<()>;
//TODO: Close connection
} }
impl<W: Write> Writer for Pcap<W> { impl<W: Write> Writer for Pcap<W> {

9
crates/lib/src/socket.rs
View file

@ -4,10 +4,9 @@ use crate::{
GDResult, GDResult,
}; };
use std::net::SocketAddr;
use std::{ use std::{
io::{Read, Write}, io::{Read, Write},
net, net::{self, SocketAddr},
}; };
const DEFAULT_PACKET_SIZE: usize = 1024; const DEFAULT_PACKET_SIZE: usize = 1024;
@ -227,7 +226,7 @@ pub mod capture {
} }
/// Represents the UDP protocol provider. /// Represents the UDP protocol provider.
pub struct ProtocolUDP; pub(crate) struct ProtocolUDP;
impl ProtocolProvider for ProtocolUDP { impl ProtocolProvider for ProtocolUDP {
fn protocol() -> Protocol { fn protocol() -> Protocol {
Protocol::UDP Protocol::UDP
@ -240,7 +239,7 @@ pub mod capture {
/// * `I` - The inner socket type. /// * `I` - The inner socket type.
/// * `P` - The protocol provider. /// * `P` - The protocol provider.
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct WrappedCaptureSocket<I, P> { pub(crate) struct WrappedCaptureSocket<I, P> {
inner: I, inner: I,
remote_address: SocketAddr, remote_address: SocketAddr,
_protocol: PhantomData<P>, _protocol: PhantomData<P>,
@ -292,7 +291,7 @@ pub mod capture {
/// ///
/// # Returns /// # Returns
/// A result indicating success or error in sending data. /// 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 { let info = CapturePacket {
direction: Direction::Send, direction: Direction::Send,
protocol: P::protocol(), protocol: P::protocol(),