Error handling and better structure

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 {
-        let default: Vec<u8> = 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 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
-        }
+    fn send(&self, data: &[u8]) -> Result<(), GDError> {
+        self.socket.send_to(&data, &self.complete_address).map_err(|e| GDError::PacketSend(e.to_string()))?;
+        Ok(())
     }
 
-    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();
-        buffer[..amt].to_vec()
+        let (amt, _) = self.socket.recv_from(&mut buffer.as_mut_slice()).map_err(|e| GDError::PacketReceive(e.to_string()))?;
+        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 mut buf = client.receive(DEFAULT_PACKET_SIZE);
+        let buf = client.do_request(Request::INFO)?;
         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 {
-                'd' => Server::Dedicated,
-                'l' => Server::NonDedicated,
-                _ => Server::SourceTV
+            server_type: match buffer::get_u8(&buf, &mut pos)? {
+                100 => Server::Dedicated, //'d'
+                108 => Server::NonDedicated, //'l'
+                _ => Server::SourceTV //'p'
             },
-            environment_type: match buffer::get_u8(&buf, &mut pos)? as char {
-                'l' => Environment::Linux,
-                'w' => Environment::Windows,
-                _ => Environment::Mac
+            environment_type: match buffer::get_u8(&buf, &mut pos)? {
+                100 => Environment::Linux, //'l'
+                119 => Environment::Windows, //'w'
+                _ => Environment::Mac //'m' or 'o'
             },
             has_password: buffer::get_u8(&buf, &mut pos)? == 1,
             vac_secured: buffer::get_u8(&buf, &mut pos)? == 1,