Error handling and better structure

src/errors.rs

@@ -3,13 +3,21 @@ use std::fmt::Formatter;
 
 #[derive(Debug, Clone)]
 pub enum GDError {
-    PacketOverflow
+    PacketOverflow(String),
+    PacketUnderflow(String),
+    PacketBad(String),
+    PacketSend(String),
+    PacketReceive(String)
 }
 
 impl fmt::Display for GDError {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
-            GDError::PacketOverflow => write!(f, "Packet overflow!")
+            GDError::PacketOverflow(details) => write!(f, "Packet overflow: {details}"),
+            GDError::PacketUnderflow(details) => write!(f, "Packet underflow: {details}"),
+            GDError::PacketBad(details) => write!(f, "Packet bad: {details}"),
+            GDError::PacketSend(details) => write!(f, "Couldn't send a packet: {details}"),
+            GDError::PacketReceive(details) => write!(f, "Couldn't receive a packet: {details}")
         }
     }
 }

src/games/tf2.rs

@@ -8,6 +8,6 @@ impl TF2 {
         ValveProtocol::query(App::TF2, address, match port {
             None => 27015,
             Some(port) => port
-        })
+        }, true, true)
     }
 }

src/protocols/valve.rs

@@ -54,7 +54,9 @@ pub struct ExtraData {
 }
 
 pub enum Request {
-    A2sInfo(Option<[u8; 4]>)
+    INFO,
+    PLAYER,
+    RULES
 }
 
 #[derive(PartialEq)]
@@ -78,49 +80,55 @@ impl ValveProtocol {
         }
     }
 
-    pub fn do_request(&self, kind: Request, data_packet: Option<&[u8]>) -> bool {
+    fn send(&self, data: &[u8]) -> Result<(), GDError> {
-        let default: Vec<u8> = vec![0xFF, 0xFF, 0xFF, 0xFF, 0x54, 0x53, 0x6F, 0x75, 0x72, 0x63, 0x65, 0x20,
+        self.socket.send_to(&data, &self.complete_address).map_err(|e| GDError::PacketSend(e.to_string()))?;
-            0x45, 0x6E, 0x67, 0x69, 0x6E, 0x65, 0x20, 0x51, 0x75, 0x65, 0x72, 0x79, 0x00];
+        Ok(())
-
-        let request_kind_packet = match kind {
-            Request::A2sInfo(challenge) => match challenge {
-                None => default,
-                Some(value) => concat_u8_arrays(&default, &value)
-            }
-        };
-
-        let mut packet = request_kind_packet;
-        match data_packet {
-            None => (),
-            Some(data) => packet.extend_from_slice(data)
-        }
-
-        match self.socket.send_to(&packet, &self.complete_address) {
-            Err(_) => false,
-            Ok(_) => true
-        }
     }
 
-    pub fn receive(&self, buffer_size: usize) -> Vec<u8> {
+    fn receive(&self, buffer_size: usize) -> Result<Vec<u8>, GDError> {
         let mut buffer: Vec<u8> = vec![0; buffer_size];
-        let (amt, _) = self.socket.recv_from(&mut buffer.as_mut_slice()).unwrap();
+        let (amt, _) = self.socket.recv_from(&mut buffer.as_mut_slice()).map_err(|e| GDError::PacketReceive(e.to_string()))?;
-        buffer[..amt].to_vec()
+        Ok(buffer[..amt].to_vec())
+    }
+
+    pub fn do_request(&self, kind: Request) -> Result<Vec<u8>, GDError> {
+        let info_initial_packet = vec![0xFF, 0xFF, 0xFF, 0xFF, 0x54, 0x53, 0x6F, 0x75, 0x72, 0x63, 0x65, 0x20, 0x45, 0x6E, 0x67, 0x69, 0x6E, 0x65, 0x20, 0x51, 0x75, 0x65, 0x72, 0x79, 0x00];
+        let player_initial_packet = vec![0xFF, 0xFF, 0xFF, 0xFF, 0x55];
+        let rules_initial_packet = vec![0xFF, 0xFF, 0xFF, 0xFF, 0x56];
+
+        let no_challenge: [u8; 4] = [0xFF, 0xFF, 0xFF, 0xFF];
+
+        let request_initial_packet = match kind {
+            Request::INFO => info_initial_packet,
+            Request::PLAYER => concat_u8_arrays(&player_initial_packet, &no_challenge),
+            Request::RULES => concat_u8_arrays(&rules_initial_packet, &no_challenge)
+        };
+
+        self.send(&request_initial_packet)?;
+        let buffer = self.receive(DEFAULT_PACKET_SIZE)?;
+
+        if buffer.len() < 9 {
+            return Err(GDError::PacketOverflow("Any Valve Protocol response can't be under 9 bytes long.".to_string()));
+        }
+
+        if buffer[4] != 41 { //'A'
+            return Ok(buffer);
+        }
+
+        let challenge: [u8; 4] = [buffer[5], buffer[6], buffer[7], buffer[8]];
+        self.send(&concat_u8_arrays(&request_initial_packet, &challenge))?;
+
+        Ok(self.receive(DEFAULT_PACKET_SIZE)?)
     }
 }
 
 impl ValveProtocol {
-    pub(crate) fn query(app: App, address: &str, port: u16) -> Result<Response, GDError> {
+    pub(crate) fn query(app: App, address: &str, port: u16, gather_players: bool, gather_rules: bool) -> Result<Response, GDError> {
         let client = ValveProtocol::new(address, port);
 
-        client.do_request(Request::A2sInfo(None), None);
+        let buf = client.do_request(Request::INFO)?;
-        let mut buf = client.receive(DEFAULT_PACKET_SIZE);
         let mut pos = 4;
 
-        if buffer::get_u8(&buf, &mut pos)? == 0x41 {
-            client.do_request(Request::A2sInfo(Some([buf[pos], buf[pos + 1], buf[pos + 2], buf[pos + 3]])), None);
-            buf = client.receive(DEFAULT_PACKET_SIZE);
-        }
-
         Ok(Response {
             protocol: buffer::get_u8(&buf, &mut pos)?,
             name: buffer::get_string(&buf, &mut pos)?,
@@ -131,15 +139,15 @@ impl ValveProtocol {
             players: buffer::get_u8(&buf, &mut pos)?,
             max_players: buffer::get_u8(&buf, &mut pos)?,
             bots: buffer::get_u8(&buf, &mut pos)?,
-            server_type: match buffer::get_u8(&buf, &mut pos)? as char {
+            server_type: match buffer::get_u8(&buf, &mut pos)? {
-                'd' => Server::Dedicated,
+                100 => Server::Dedicated, //'d'
-                'l' => Server::NonDedicated,
+                108 => Server::NonDedicated, //'l'
-                _ => Server::SourceTV
+                _ => Server::SourceTV //'p'
             },
-            environment_type: match buffer::get_u8(&buf, &mut pos)? as char {
+            environment_type: match buffer::get_u8(&buf, &mut pos)? {
-                'l' => Environment::Linux,
+                100 => Environment::Linux, //'l'
-                'w' => Environment::Windows,
+                119 => Environment::Windows, //'w'
-                _ => Environment::Mac
+                _ => Environment::Mac //'m' or 'o'
             },
             has_password: buffer::get_u8(&buf, &mut pos)? == 1,
             vac_secured: buffer::get_u8(&buf, &mut pos)? == 1,

src/utils.rs

@@ -22,7 +22,7 @@ pub mod buffer {
 
     pub fn get_u8(buf: &[u8], pos: &mut usize) -> Result<u8, GDError> {
         if buf.len() <= *pos {
-            return Err(GDError::PacketOverflow);
+            return Err(GDError::PacketUnderflow("Unexpectedly short packet.".to_string()));
         }
 
         let value = buf[*pos];
@@ -32,7 +32,7 @@ pub mod buffer {
 
     pub fn get_u16_le(buf: &[u8], pos: &mut usize) -> Result<u16, GDError> {
         if buf.len() <= *pos + 1 {
-            return Err(GDError::PacketOverflow);
+            return Err(GDError::PacketUnderflow("Unexpectedly short packet.".to_string()));
         }
 
         let value = combine_two_u8_le(buf[*pos + 1], buf[*pos]);
@@ -42,7 +42,7 @@ pub mod buffer {
 
     pub fn get_u64_le(buf: &[u8], pos: &mut usize) -> Result<u64, GDError> {
         if buf.len() <= *pos + 7 {
-            return Err(GDError::PacketOverflow);
+            return Err(GDError::PacketUnderflow("Unexpectedly short packet.".to_string()));
         }
 
         let value = combine_eight_u8_le(buf[*pos + 7], buf[*pos + 6], buf[*pos + 5], buf[*pos + 4], buf[*pos + 3], buf[*pos + 2], buf[*pos + 1], buf[*pos]);
@@ -52,7 +52,7 @@ pub mod buffer {
 
     pub fn get_string(buf: &[u8], pos: &mut usize) -> Result<String, GDError> {
         let sub_buf = &buf[*pos..];
-        let first_null_position = sub_buf.iter().position(|&x| x == 0).ok_or(GDError::PacketOverflow)?;
+        let first_null_position = sub_buf.iter().position(|&x| x == 0).ok_or(GDError::PacketBad("Unexpectedly formatted packet.".to_string()))?;
         let value = std::str::from_utf8(&sub_buf[..first_null_position]).unwrap().to_string();
         *pos += value.len() + 1;
         Ok(value)