master: reuse server challenges

Author: numas13

master/src/master_server.rs

@@ -314,21 +314,14 @@ impl<Addr: AddrExt> MasterServer<Addr> {
 
         match p {
             server::Packet::Challenge(p) => {
-                let master_challenge = self.add_challenge(from);
-                let resp = master::ChallengeResponse::new(master_challenge, p.server_challenge);
+                let challenge = self.server_challenge_add(from);
+                let resp = master::ChallengeResponse::new(challenge, p.server_challenge);
                 trace!("{from}: send {resp:?}");
                 let mut buf = [0; 32];
                 let packet = resp.encode(&mut buf)?;
                 self.sock.send_to(packet, from)?;
             }
             server::Packet::ServerAdd(p) => {
-                let challenge = match self.challenges.get(&from) {
-                    Some(e) => e.value,
-                    None => {
-                        trace!("{from}: challenge does not exists");
-                        return Ok(());
-                    }
-                };
                 if p.version < self.cfg.server.min_version {
                     let min = self.cfg.server.min_version;
                     warn!(
@@ -337,6 +330,10 @@ impl<Addr: AddrExt> MasterServer<Addr> {
                     );
                     return Ok(());
                 }
+                let Some(challenge) = self.server_challenge_get(&from) else {
+                    trace!("{from}: challenge does not exist");
+                    return Ok(());
+                };
                 if p.challenge != challenge {
                     warn!(
                         "{from}: expected challenge {challenge} but received {}",
@@ -606,10 +603,15 @@ impl<Addr: AddrExt> MasterServer<Addr> {
         Err(Error::UndefinedPacket)
     }
 
-    fn add_challenge(&mut self, addr: Addr) -> u32 {
-        let x = self.rng.u32(..);
-        self.challenges.insert(addr, x);
-        x
+    fn server_challenge_add(&mut self, addr: Addr) -> u32 {
+        self.challenges
+            .entry(addr)
+            .or_insert_with(|| Timed::new(self.rng.u32(..)))
+            .value
+    }
+
+    fn server_challenge_get(&self, addr: &Addr) -> Option<u32> {
+        self.challenges.get(addr).map(|i| i.value)
     }
 
     fn admin_challenge_add(&mut self, addr: Addr) -> (u32, u32) {

master/tests/master.rs

@@ -33,12 +33,21 @@ struct Test {
 }
 
 impl Test {
-    fn new(cfg: &Config) -> Test {
-        let mut test = Test {
+    fn new() -> Test {
+        Test {
             master_addr: UNSPECIFIED.into(),
             server_addr: UNSPECIFIED.into(),
-        };
+        }
+    }
+
+    fn with_master(cfg: &Config) -> Test {
+        let mut test = Self::new();
         test.create_master(cfg);
+        test
+    }
+
+    fn with_master_and_server(cfg: &Config) -> Test {
+        let mut test = Self::with_master(cfg);
         test.add_server(cfg);
         test
     }
@@ -87,7 +96,7 @@ fn server_add() {
     log::set_max_level(log::LevelFilter::Trace);
 
     let cfg = Config::default();
-    let test = Test::new(&cfg);
+    let test = Test::with_master_and_server(&cfg);
     let mut buf = [0; 1024];
     let sock = UdpSocket::bind(UNSPECIFIED).unwrap();
     let game_key = Some(0xbeefdead);
@@ -115,6 +124,30 @@ fn server_add() {
     assert_eq!(servers[0].port(), test.server_addr.port());
 }
 
+#[test]
+fn server_reuse_challenge() {
+    log::set_logger(&LOGGER).ok();
+    log::set_max_level(log::LevelFilter::Trace);
+
+    let test = Test::with_master(&Config::default());
+    let sock = UdpSocket::bind(UNSPECIFIED).unwrap();
+    let mut challenges = [0; 3];
+    for (i, j) in challenges.iter_mut().enumerate() {
+        let challenge = Some(i as u32);
+        let packet = server::Challenge::new(challenge);
+        let mut buf = [0; 512];
+        let p = packet.encode(&mut buf).unwrap();
+        sock.send_to(p, test.master_addr).unwrap();
+
+        let (len, _) = sock.recv_from(&mut buf).unwrap();
+        let resp = master::ChallengeResponse::decode(&buf[..len]).unwrap();
+        assert_eq!(resp.server_challenge, challenge);
+        *j = resp.master_challenge;
+    }
+
+    assert!(challenges.iter().all(|i| *i == challenges[0]));
+}
+
 #[test]
 fn client_rate_limit() {
     log::set_logger(&LOGGER).ok();
@@ -124,7 +157,7 @@ fn client_rate_limit() {
     cfg.master.server.client_rate_limit = 10;
     info!("client rate limit {}", cfg.master.server.client_rate_limit);
 
-    let test = Test::new(&cfg);
+    let test = Test::with_master_and_server(&cfg);
     let sock = UdpSocket::bind(UNSPECIFIED).unwrap();
     sock.set_read_timeout(Some(Duration::from_millis(100)))
         .unwrap();