Merge branch 'main' into feat/cli

crates/lib/src/protocols/unreal2/mod.rs

@@ -0,0 +1,58 @@
+/// The implementation.
+pub mod protocol;
+/// All types used by the implementation.
+pub mod types;
+
+pub use protocol::*;
+pub use types::*;
+
+/// Generate a module containing a query function for a valve game.
+///
+/// * `mod_name` - The name to be given to the game module (see ID naming
+///   conventions in CONTRIBUTING.md).
+/// * `pretty_name` - The full name of the game, will be used as the
+///   documentation for the created module.
+/// * `default_port` - Passed through to [game_query_fn].
+#[cfg(feature = "games")]
+macro_rules! game_query_mod {
+    ($mod_name: ident, $pretty_name: expr, $default_port: literal) => {
+        #[doc = $pretty_name]
+        pub mod $mod_name {
+            crate::protocols::unreal2::game_query_fn!($default_port);
+        }
+    };
+}
+
+#[cfg(feature = "games")]
+pub(crate) use game_query_mod;
+
+// Allow generating doc comments:
+// https://users.rust-lang.org/t/macros-filling-text-in-comments/20473
+/// Generate a query function for a valve game.
+///
+/// * `default_port` - The default port the game uses.
+#[cfg(feature = "games")]
+macro_rules! game_query_fn {
+    ($default_port: literal) => {
+        crate::protocols::unreal2::game_query_fn! {@gen $default_port, concat!(
+        "Make a Unreal2 query for with default timeout settings and default extra request settings.\n\n",
+        "If port is `None`, then the default port (", stringify!($default_port), ") will be used.")}
+    };
+
+    (@gen $default_port: literal, $doc: expr) => {
+        #[doc = $doc]
+        pub fn query(
+            address: &std::net::IpAddr,
+            port: Option<u16>,
+        ) -> crate::GDResult<crate::protocols::unreal2::Response> {
+            crate::protocols::unreal2::query(
+                &std::net::SocketAddr::new(*address, port.unwrap_or($default_port)),
+                &crate::protocols::unreal2::GatheringSettings::default(),
+                None,
+            )
+        }
+    };
+}
+
+#[cfg(feature = "games")]
+pub(crate) use game_query_fn;

crates/lib/src/protocols/unreal2/protocol.rs

@@ -0,0 +1,308 @@
+use crate::buffer::{Buffer, StringDecoder};
+use crate::errors::GDErrorKind::PacketBad;
+use crate::protocols::types::TimeoutSettings;
+use crate::socket::{Socket, UdpSocket};
+use crate::utils::retry_on_timeout;
+use crate::GDResult;
+
+use super::{GatheringSettings, MutatorsAndRules, PacketKind, Players, Response, ServerInfo};
+
+use std::net::SocketAddr;
+
+use byteorder::{ByteOrder, LittleEndian};
+use encoding_rs::{UTF_16LE, WINDOWS_1252};
+
+/// Response packets don't seem to exceed 500 bytes, set to 1024 just to be
+/// safe.
+const PACKET_SIZE: usize = 1024;
+
+/// Default amount of players to pre-allocate if numplayers was not included in
+/// server info response.
+const DEFAULT_PLAYER_PREALLOCATION: usize = 10;
+
+/// Maximum amount of players to pre-allocate: if the server specifies a number
+/// larger than this in serverinfo we don't allocate that many.
+const MAXIMUM_PLAYER_PREALLOCATION: usize = 50;
+
+/// The Unreal2 protocol implementation.
+pub(crate) struct Unreal2Protocol {
+    socket: UdpSocket,
+    retry_count: usize,
+}
+
+impl Unreal2Protocol {
+    pub fn new(address: &SocketAddr, timeout_settings: Option<TimeoutSettings>) -> GDResult<Self> {
+        let socket = UdpSocket::new(address)?;
+        let retry_count = timeout_settings
+            .as_ref()
+            .map(|t| t.get_retries())
+            .unwrap_or_else(|| TimeoutSettings::default().get_retries());
+        socket.apply_timeout(&timeout_settings)?;
+
+        Ok(Self {
+            socket,
+            retry_count,
+        })
+    }
+
+    /// Send a request packet and recieve the first response (with retries).
+    fn get_request_data(&mut self, packet_type: PacketKind) -> GDResult<Vec<u8>> {
+        retry_on_timeout(self.retry_count, move || {
+            self.get_request_data_impl(packet_type)
+        })
+    }
+
+    /// Send a request packet
+    fn get_request_data_impl(&mut self, packet_type: PacketKind) -> GDResult<Vec<u8>> {
+        let request = [0x79, 0, 0, 0, packet_type as u8];
+        self.socket.send(&request)?;
+
+        let data = self.socket.receive(Some(PACKET_SIZE))?;
+
+        Ok(data)
+    }
+
+    /// Consume the header part of a response packet, validate that the packet
+    /// type matches what is expected.
+    fn consume_response_headers<B: ByteOrder>(
+        buffer: &mut Buffer<B>,
+        expected_packet_type: PacketKind,
+    ) -> GDResult<()> {
+        // Skip header
+        buffer.move_cursor(4)?;
+
+        let packet_type: u8 = buffer.read()?;
+
+        let packet_type: PacketKind = packet_type.try_into()?;
+
+        if packet_type != expected_packet_type {
+            Err(PacketBad.context(format!(
+                "Packet response ({:?}) didn't match request ({:?}) packet type",
+                packet_type, expected_packet_type
+            )))
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Send server info query.
+    pub fn query_server_info(&mut self) -> GDResult<ServerInfo> {
+        let data = self.get_request_data(PacketKind::ServerInfo)?;
+        let mut buffer = Buffer::<LittleEndian>::new(&data);
+        // TODO: Maybe put consume headers in individual packet parse methods
+        Self::consume_response_headers(&mut buffer, PacketKind::ServerInfo)?;
+        ServerInfo::parse(&mut buffer)
+    }
+
+    /// Send mutators and rules query.
+    pub fn query_mutators_and_rules(&mut self) -> GDResult<MutatorsAndRules> {
+        // This is a required packet so we validate that we get at least one response.
+        // However there can be many packets in response to a single request so
+        // we greedily handle packets until we get a timeout (or any receive
+        // error).
+
+        let mut mutators_and_rules = MutatorsAndRules::default();
+        {
+            let data = self.get_request_data(PacketKind::MutatorsAndRules)?;
+            let mut buffer = Buffer::<LittleEndian>::new(&data);
+            // TODO: Maybe put consume headers in individual packet parse methods
+            Self::consume_response_headers(&mut buffer, PacketKind::MutatorsAndRules)?;
+            mutators_and_rules.parse(&mut buffer)?
+        };
+
+        // We could receive multiple packets in response
+        while let Ok(data) = self.socket.receive(Some(PACKET_SIZE)) {
+            let mut buffer = Buffer::<LittleEndian>::new(&data);
+
+            let r = Self::consume_response_headers(&mut buffer, PacketKind::MutatorsAndRules);
+            if r.is_err() {
+                println!("{:?}", r);
+                break;
+            }
+
+            mutators_and_rules.parse(&mut buffer)?;
+        }
+
+        Ok(mutators_and_rules)
+    }
+
+    /// Send players query.
+    pub fn query_players(&mut self, server_info: Option<&ServerInfo>) -> GDResult<Players> {
+        // Pre-allocate the player arrays, but don't over allocate memory if the server
+        // specifies an insane number of players.
+        let num_players: Option<usize> = server_info.and_then(|i| i.num_players.try_into().ok());
+
+        let mut players = Players::with_capacity(
+            num_players
+                .unwrap_or(DEFAULT_PLAYER_PREALLOCATION)
+                .min(MAXIMUM_PLAYER_PREALLOCATION),
+        );
+
+        // Fetch first players packet (with retries)
+        let mut players_data = self.get_request_data(PacketKind::Players);
+        // Players are non required so if we don't get any responses we continue to
+        // return
+        while let Ok(data) = players_data {
+            let mut buffer = Buffer::<LittleEndian>::new(&data);
+
+            Self::consume_response_headers(&mut buffer, PacketKind::Players)?;
+
+            players.parse(&mut buffer)?;
+
+            if let Some(num_players) = num_players {
+                if players.total_len() >= num_players {
+                    // If we have already received the amount of players specified in server info
+                    // then we don't need to wait for more player packets to time out.
+                    break;
+                }
+            }
+
+            // Receive next packet
+            players_data = self.socket.receive(Some(PACKET_SIZE));
+        }
+
+        Ok(players)
+    }
+
+    /// Make a full server query.
+    pub fn query(&mut self, gather_settings: &GatheringSettings) -> GDResult<Response> {
+        // Fetch the server info, this can only handle one response packet
+        let server_info = self.query_server_info()?;
+
+        let mutators_and_rules = if gather_settings.mutators_and_rules {
+            self.query_mutators_and_rules()?
+        } else {
+            MutatorsAndRules::default()
+        };
+
+        let players = if gather_settings.players {
+            self.query_players(Some(&server_info))?
+        } else {
+            Players::with_capacity(0)
+        };
+
+        // TODO: Handle extra info parsing when we detect certain game types (or maybe
+        // include that in gather settings).
+
+        Ok(Response {
+            server_info,
+            mutators_and_rules,
+            players,
+        })
+    }
+}
+
+/// Unreal 2 string decoder
+pub struct Unreal2StringDecoder;
+impl StringDecoder for Unreal2StringDecoder {
+    type Delimiter = [u8; 1];
+
+    const DELIMITER: Self::Delimiter = [0x00];
+
+    fn decode_string(data: &[u8], cursor: &mut usize, delimiter: Self::Delimiter) -> GDResult<String> {
+        let mut ucs2 = false;
+        let mut length: usize = (*data
+            .first()
+            .ok_or(PacketBad.context("Tried to decode string without length"))?)
+        .into();
+
+        let mut start = 0;
+
+        // Check if it is a UCS-2 string
+        if length >= 0x80 {
+            ucs2 = true;
+
+            length = (length & 0x7f) * 2;
+
+            start += 1;
+
+            // For UCS-2 strings, some unreal 2 games randomly insert an extra 0x01 here,
+            // not included in the length. Skip it if present (hopefully this never happens
+            // legitimately)
+            if let Some(1) = data[start ..].first() {
+                start += 1;
+            }
+        }
+
+        // If UCS2 the first byte is the masked length of the string
+        let result = if ucs2 {
+            let string_data = &data[start .. start + length];
+            if string_data.len() != length {
+                return Err(PacketBad.context("Not enough data in buffer to read string"));
+            }
+
+            // When node decodes UCS2 it uses the UFT16LE encoding.
+            // https://github.com/nodejs/node/blob/2aaa21f9f684484edb54be30589c4af0b923cdef/lib/buffer.js#L637-L645
+            let (result, _, invalid_sequences) = UTF_16LE.decode(string_data);
+
+            if invalid_sequences {
+                return Err(PacketBad.context("UTF-8 string contained invalid character(s)"));
+            }
+
+            result
+        } else {
+            // Else the string is null-delimited latin1
+
+            // TODO: Replace this with delimiter finder helper
+            let position = data
+            // Create an iterator over the data.
+                .iter()
+                // Find the position of the delimiter
+                .position(|&b| b == delimiter.as_ref()[0])
+                // If the delimiter is not found, use the whole data slice.
+                .unwrap_or(data.len());
+
+            length = position + 1;
+
+            // Decode as latin1
+            let (result, _, invalid_sequences) = WINDOWS_1252.decode(&data[0 .. position]);
+
+            if invalid_sequences {
+                return Err(PacketBad.context("latin1 string contained invalid character(s)"));
+            }
+
+            result
+        };
+
+        // Strip color encodings
+        // TODO: Improve efficiency
+        // TODO: There might be a nicer way to do this once string patterns are stable
+        //       https://github.com/rust-lang/rust/issues/27721
+
+        // After '0x1b' skip 3 characters (including the '0x1b')
+        let mut char_skip = 0usize;
+        let result: String = result
+            .chars()
+            .filter(|c: &char| {
+                if '\x1b'.eq(c) {
+                    char_skip = 4;
+                    return false;
+                }
+                char_skip = char_skip.saturating_sub(1);
+
+                char_skip == 0
+            })
+            .collect();
+
+        // Remove all characters between 0x00 and 0x1a
+        let result = result.replace(|c: char| c > '\x00' && c <= '\x1a', "");
+
+        *cursor += start + length;
+
+        // Strip delimiter that wasn't included in length
+        Ok(result.trim_matches('\0').to_string())
+    }
+}
+
+///  Make an unreal2 query.
+pub fn query(
+    address: &SocketAddr,
+    gather_settings: &GatheringSettings,
+    timeout_settings: Option<TimeoutSettings>,
+) -> GDResult<Response> {
+    let mut client = Unreal2Protocol::new(address, timeout_settings)?;
+
+    client.query(gather_settings)
+}
+
+// TODO: Add tests

crates/lib/src/protocols/unreal2/types.rs

@@ -0,0 +1,246 @@
+use crate::buffer::Buffer;
+use crate::errors::GDErrorKind::PacketBad;
+use crate::protocols::types::{CommonPlayer, CommonResponse, ExtraRequestSettings, GenericPlayer};
+use crate::protocols::GenericResponse;
+use crate::{GDError, GDResult};
+
+use super::Unreal2StringDecoder;
+
+use std::collections::{HashMap, HashSet};
+
+use byteorder::ByteOrder;
+
+/// Unreal 2 packet types.
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+#[repr(u8)]
+pub enum PacketKind {
+    ServerInfo = 0,
+    MutatorsAndRules = 1,
+    Players = 2,
+}
+
+impl TryFrom<u8> for PacketKind {
+    type Error = GDError;
+    fn try_from(value: u8) -> GDResult<Self> {
+        match value {
+            0 => Ok(Self::ServerInfo),
+            1 => Ok(Self::MutatorsAndRules),
+            2 => Ok(Self::Players),
+            _ => Err(PacketBad.context("Unknown packet type")),
+        }
+    }
+}
+
+/// Unreal 2 server info.
+#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct ServerInfo {
+    pub server_id: u32,
+    pub ip: String,
+    pub game_port: u32,
+    pub query_port: u32,
+    pub name: String,
+    pub map: String,
+    pub game_type: String,
+    pub num_players: u32,
+    pub max_players: u32,
+}
+
+impl ServerInfo {
+    pub fn parse<B: ByteOrder>(buffer: &mut Buffer<B>) -> GDResult<Self> {
+        Ok(ServerInfo {
+            server_id: buffer.read()?,
+            ip: buffer.read_string::<Unreal2StringDecoder>(None)?,
+            game_port: buffer.read()?,
+            query_port: buffer.read()?,
+            name: buffer.read_string::<Unreal2StringDecoder>(None)?,
+            map: buffer.read_string::<Unreal2StringDecoder>(None)?,
+            game_type: buffer.read_string::<Unreal2StringDecoder>(None)?,
+            num_players: buffer.read()?,
+            max_players: buffer.read()?,
+        })
+    }
+}
+
+/// Unreal 2 mutators and rules.
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct MutatorsAndRules {
+    pub mutators: HashSet<String>,
+    pub rules: HashMap<String, Vec<String>>,
+}
+
+impl MutatorsAndRules {
+    pub fn parse<B: ByteOrder>(&mut self, buffer: &mut Buffer<B>) -> GDResult<()> {
+        while buffer.remaining_length() > 0 {
+            let key = buffer.read_string::<Unreal2StringDecoder>(None)?;
+            let value = buffer.read_string::<Unreal2StringDecoder>(None).ok();
+
+            if key.eq_ignore_ascii_case("mutator") {
+                if let Some(value) = value {
+                    self.mutators.insert(value);
+                }
+            } else {
+                let rule_vec = self.rules.get_mut(&key);
+
+                let rule_vec = if let Some(rule_vec) = rule_vec {
+                    rule_vec
+                } else {
+                    self.rules.insert(key.clone(), Vec::default());
+                    self.rules
+                        .get_mut(&key)
+                        .expect("Value should be in HashMap after we inserted")
+                };
+
+                if let Some(value) = value {
+                    rule_vec.push(value);
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+/// Unreal 2 players and bots.
+#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct Players {
+    /// List of players returned by server (without 0 ping).
+    pub players: Vec<Player>,
+    /// List of bots returned by server (players with 0 ping).
+    pub bots: Vec<Player>,
+}
+
+impl Players {
+    /// Pre-allocate the vectors inside the players struct based on the provided
+    /// capacity.
+    pub fn with_capacity(capacity: usize) -> Self {
+        Players {
+            players: Vec::with_capacity(capacity),
+            // Allocate half as many bots as we don't expect there to be as many
+            bots: Vec::with_capacity(capacity / 2),
+        }
+    }
+
+    /// Parse a raw buffer of players into the current struct.
+    pub fn parse<B: ByteOrder>(&mut self, buffer: &mut Buffer<B>) -> GDResult<()> {
+        while buffer.remaining_length() > 0 {
+            let player = Player {
+                id: buffer.read()?,
+                name: buffer.read_string::<Unreal2StringDecoder>(None)?,
+                ping: buffer.read()?,
+                score: buffer.read()?,
+                stats_id: buffer.read()?,
+            };
+
+            // If ping is 0 the player is a bot
+            if player.ping == 0 {
+                self.bots.push(player);
+            } else {
+                self.players.push(player);
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Length of both players and bots.
+    pub fn total_len(&self) -> usize { self.players.len() + self.bots.len() }
+}
+
+/// Unreal 2 player info.
+#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct Player {
+    pub id: u32,
+    pub name: String,
+    pub ping: u32,
+    pub score: i32,
+    pub stats_id: u32,
+}
+
+impl CommonPlayer for Player {
+    fn name(&self) -> &str { &self.name }
+
+    fn score(&self) -> Option<i32> { Some(self.score) }
+
+    fn as_original(&self) -> GenericPlayer { GenericPlayer::Unreal2(self) }
+}
+
+/// Unreal 2 response.
+#[derive(Clone, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct Response {
+    pub server_info: ServerInfo,
+    pub mutators_and_rules: MutatorsAndRules,
+    pub players: Players,
+}
+
+impl CommonResponse for Response {
+    fn map(&self) -> Option<&str> { Some(&self.server_info.map) }
+
+    fn name(&self) -> Option<&str> { Some(&self.server_info.name) }
+
+    fn game_mode(&self) -> Option<&str> { Some(&self.server_info.game_type) }
+
+    fn players_online(&self) -> u32 { self.server_info.num_players }
+
+    fn players_maximum(&self) -> u32 { self.server_info.max_players }
+
+    fn players(&self) -> Option<Vec<&dyn crate::protocols::types::CommonPlayer>> {
+        Some(
+            self.players
+                .players
+                .iter()
+                .map(|player| player as _)
+                .collect(),
+        )
+    }
+
+    fn as_original(&self) -> GenericResponse { GenericResponse::Unreal2(self) }
+}
+
+/// What data to gather, purely used only with the query function.
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
+pub struct GatheringSettings {
+    pub players: bool,
+    pub mutators_and_rules: bool,
+}
+
+impl GatheringSettings {
+    /// Default values are true for both the players and the rules.
+    pub const fn default() -> Self {
+        Self {
+            players: true,
+            mutators_and_rules: true,
+        }
+    }
+
+    pub const fn into_extra(self) -> ExtraRequestSettings {
+        ExtraRequestSettings {
+            hostname: None,
+            protocol_version: None,
+            gather_players: Some(self.players),
+            gather_rules: Some(self.mutators_and_rules),
+            check_app_id: None,
+        }
+    }
+}
+
+impl Default for GatheringSettings {
+    fn default() -> Self { GatheringSettings::default() }
+}
+
+impl From<ExtraRequestSettings> for GatheringSettings {
+    fn from(value: ExtraRequestSettings) -> Self {
+        let default = Self::default();
+        Self {
+            players: value.gather_players.unwrap_or(default.players),
+            mutators_and_rules: value.gather_rules.unwrap_or(default.mutators_and_rules),
+        }
+    }
+}
+
+// TODO: Add tests