From 9beefd5d6db4468af0587ae29f890e9563f606b3 Mon Sep 17 00:00:00 2001
From: Denis Drakhnia <numas13@gmail.com>
Date: Sat, 7 Oct 2023 20:35:38 +0300
Subject: [PATCH] Reorg missing challenge error

---
 src/client.rs        |  2 +-
 src/master_server.rs |  6 ++++++
 src/server_info.rs   | 13 ++++---------
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 1b23029..49e7cbf 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -38,7 +38,7 @@ impl<'a> Deref for Filter<'a> {
 #[derive(Debug)]
 pub enum Packet<'a> {
     Challenge(Option<u32>),
-    ServerAdd(u32, ServerInfo<&'a str>),
+    ServerAdd(Option<u32>, ServerInfo<&'a str>),
     ServerRemove,
     QueryServers(Region, Filter<'a>),
 }
diff --git a/src/master_server.rs b/src/master_server.rs
index 7df677b..45a522f 100644
--- a/src/master_server.rs
+++ b/src/master_server.rs
@@ -39,6 +39,8 @@ pub enum Error {
     BindSocket(io::Error),
     #[error("Failed to decode packet: {0}")]
     ClientPacket(#[from] crate::client::Error),
+    #[error("Missing challenge in ServerInfo")]
+    MissingChallenge,
     #[error(transparent)]
     Io(#[from] io::Error),
 }
@@ -129,6 +131,10 @@ impl MasterServer {
                 self.remove_outdated_challenges();
             }
             Packet::ServerAdd(challenge, info) => {
+                let challenge = match challenge {
+                    Some(c) => c,
+                    None => return Err(Error::MissingChallenge),
+                };
                 let entry = match self.challenges.get(&from) {
                     Some(e) => e,
                     None => {
diff --git a/src/server_info.rs b/src/server_info.rs
index 235fe61..b8b14d2 100644
--- a/src/server_info.rs
+++ b/src/server_info.rs
@@ -13,8 +13,6 @@ use crate::parser::{Error as ParserError, ParseValue, Parser};
 pub enum Error {
     #[error("Invalid region")]
     InvalidRegion,
-    #[error("Missing challenge in ServerInfo")]
-    MissingChallenge,
     #[error(transparent)]
     Parser(#[from] ParserError),
 }
@@ -182,7 +180,7 @@ impl<'a, T> ServerInfo<T>
 where
     T: 'a + Default + ParseValue<'a, Err = ParserError>,
 {
-    pub fn from_bytes(src: &'a [u8]) -> Result<(u32, Self, &'a [u8]), Error> {
+    pub fn from_bytes(src: &'a [u8]) -> Result<(Option<u32>, Self, &'a [u8]), Error> {
         let mut parser = Parser::new(src);
         let (challenge, info) = parser.parse()?;
         let tail = match parser.end() {
@@ -193,7 +191,7 @@ where
     }
 }
 
-impl<'a, T> ParseValue<'a> for (u32, ServerInfo<T>)
+impl<'a, T> ParseValue<'a> for (Option<u32>, ServerInfo<T>)
 where
     T: 'a + Default + ParseValue<'a, Err = ParserError>,
 {
@@ -239,10 +237,7 @@ where
             }
         }
 
-        match challenge {
-            Some(challenge) => Ok((challenge, info)),
-            None => Err(Error::MissingChallenge),
-        }
+        Ok((challenge, info))
     }
 }
 
@@ -313,7 +308,7 @@ mod tests {
         assert_eq!(
             ServerInfo::from_bytes(&buf[..]),
             Ok((
-                12345678,
+                Some(12345678),
                 ServerInfo::<&str> {
                     protocol: 47,
                     players: 16,