Require timeout, use async udp operations (tokio).

Author: vikpe

Cargo.toml

@@ -4,7 +4,7 @@ description = "Get server addresses from QuakeWorld master servers."
 keywords = ["masters", "quake", "quakeworld", "servers"]
 repository = "https://github.com/quakeworld/masterstat"
 authors = ["Viktor Persson <viktor.persson@arcsin.se>"]
-version = "0.2.0"
+version = "0.3.0"
 edition = "2021"
 license = "MIT"
 include = [
@@ -21,8 +21,8 @@ include = [
 anyhow = "1.0.82"
 binrw = "0.14.1"
 futures = "0.3.30"
-tinyudp = "0.4.0"
-tokio = { version = "1.37.0", features = ["rt", "sync"] }
+thiserror = "2.0.11"
+tokio = { version = "1.37.0", features = ["macros", "net", "rt-multi-thread", "sync", "time"] }
 
 [dev-dependencies]
 pretty_assertions = "1.4.0"

src/command.rs

@@ -7,6 +7,7 @@ use binrw::BinRead;
 use tokio::sync::Mutex;
 
 use crate::server_address::{RawServerAddress, ServerAddress};
+use crate::tinyudp;
 
 /// Get server addresses from a single master server
 ///
@@ -15,28 +16,33 @@ use crate::server_address::{RawServerAddress, ServerAddress};
 /// ```
 /// use std::time::Duration;
 ///
-/// let master = "master.quakeworld.nu:27000";
-/// let timeout = Some(Duration::from_secs(2));
-/// match masterstat::server_addresses(&master, timeout) {
-///     Ok(addresses) => { println!("found {} server addresses", addresses.len()) },
-///     Err(e) => { eprintln!("error: {}", e); }
+/// async fn test() {
+///     let master = "master.quakeworld.nu:27000";
+///     let timeout = Duration::from_secs(2);
+///     match masterstat::server_addresses(&master, timeout).await {
+///         Ok(result) => { println!("found {} server addresses", result.len()) },
+///         Err(e) => { eprintln!("error: {}", e); }
+///     }
 /// }
 /// ```
-pub fn server_addresses(
+pub async fn server_addresses(
     master_address: &str,
-    timeout: Option<Duration>,
+    timeout: Duration,
 ) -> Result<Vec<ServerAddress>> {
-    const MESSAGE: [u8; 3] = [99, 10, 0];
-    let options = tinyudp::ReadOptions {
-        timeout,
-        buffer_size: 16 * 1024, // 16 kb
-    };
-    let response = tinyudp::send_and_read(master_address, &MESSAGE, &options)?;
-    let server_addresses = parse_servers_response(&response)?;
-    Ok(sorted_and_unique(&server_addresses))
+    const STATUS_MSG: [u8; 3] = [99, 10, 0];
+    let response = tinyudp::send_and_receive(
+        master_address,
+        &STATUS_MSG,
+        tinyudp::Options {
+            timeout,
+            buffer_size: 64 * 1024, // 64 kb
+        },
+    )
+    .await?;
+    parse_servers_response(&response)
 }
 
-/// Get server addresses from many master servers (async, in parallel)
+/// Get server addresses from many master servers (concurrently)
 ///
 /// # Example
 ///
@@ -45,21 +51,22 @@ pub fn server_addresses(
 ///
 /// async fn test() {
 ///     let masters = ["master.quakeworld.nu:27000", "master.quakeservers.net:27000"];
-///     let timeout = Some(Duration::from_secs(2));
-///     let server_addresses = masterstat::server_addresses_from_many(&masters, timeout).await;
+///     let timeout = Duration::from_secs(2);
+///     let result = masterstat::server_addresses_from_many(&masters, timeout).await;
+///     println!("found {} server addresses", result.len());
 /// }
 /// ```
 pub async fn server_addresses_from_many(
     master_addresses: &[impl AsRef<str>],
-    timeout: Option<Duration>,
+    timeout: Duration,
 ) -> Vec<ServerAddress> {
     let mut task_handles = vec![];
     let result_mux = Arc::<Mutex<Vec<ServerAddress>>>::default();
 
     for master_address in master_addresses.iter().map(|a| a.as_ref().to_string()) {
         let result_mux = result_mux.clone();
         let task = tokio::spawn(async move {
-            if let Ok(servers) = server_addresses(&master_address, timeout) {
+            if let Ok(servers) = server_addresses(&master_address, timeout).await {
                 let mut result = result_mux.lock().await;
                 result.extend(servers);
             }
@@ -99,12 +106,33 @@ fn sorted_and_unique(server_addresses: &[ServerAddress]) -> Vec<ServerAddress> {
 
 #[cfg(test)]
 mod tests {
-    use pretty_assertions::assert_eq;
-
     use super::*;
+    // use pretty_assertions::assert_eq;
+
+    #[tokio::test]
+    async fn test_server_addresses() -> Result<()> {
+        let master = "master.quakeservers.net:27000";
+        let timeout = Duration::from_secs(1);
+        let result = server_addresses(master, timeout).await?;
+        assert!(!result.is_empty());
+        Ok(())
+    }
 
-    #[test]
-    fn test_parse_servers_response() -> Result<()> {
+    #[tokio::test]
+    async fn test_server_addresses_from_many() -> Result<()> {
+        let masters = [
+            "master.quakeservers.net:27000",
+            "master.quakeworld.nu:27000",
+            "qwmaster.fodquake.net:27000",
+        ];
+        let timeout = Duration::from_secs(1);
+        let result = server_addresses_from_many(&masters, timeout).await;
+        assert!(result.len() > 500);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_parse_servers_response() -> Result<()> {
         // invalid response header
         {
             let response = [0xff, 0xff];

src/lib.rs

@@ -4,6 +4,7 @@
 
 mod command;
 mod server_address;
+mod tinyudp;
 
 pub use crate::command::server_addresses;
 pub use crate::command::server_addresses_from_many;

src/tinyudp.rs

@@ -0,0 +1,75 @@
+use anyhow::Result;
+use std::net::{Ipv4Addr, SocketAddrV4};
+use std::time::Duration;
+use thiserror::Error;
+use tokio::net::{ToSocketAddrs, UdpSocket};
+
+#[derive(Debug, Error)]
+pub enum TinyudpError {
+    #[error("failed to bind socket: {0}")]
+    BindFailed(#[source] std::io::Error),
+
+    #[error("failed to send message: {0}")]
+    SendFailed(#[source] std::io::Error),
+
+    #[error("failed to receive message: {0}")]
+    ReceiveFailed(#[source] std::io::Error),
+
+    #[error("timeout reached while waiting for response")]
+    TimeoutReached,
+}
+
+pub async fn send_and_receive(
+    target: impl ToSocketAddrs,
+    message: &[u8],
+    options: Options,
+) -> Result<Vec<u8>, TinyudpError> {
+    let socket = bind().await?;
+    socket
+        .send_to(message, target)
+        .await
+        .map_err(TinyudpError::SendFailed)?;
+
+    let mut buffer = vec![0; options.buffer_size];
+    let (bytes_read, _) = tokio::select! {
+        _ = tokio::time::sleep(options.timeout) => Err(TinyudpError::TimeoutReached),
+        res = socket.recv_from(&mut buffer) => res.map_err(TinyudpError::ReceiveFailed),
+    }?;
+
+    let response = buffer[..bytes_read].to_vec();
+    Ok(response)
+}
+
+async fn bind() -> Result<UdpSocket, TinyudpError> {
+    let address = SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0);
+    UdpSocket::bind(address)
+        .await
+        .map_err(TinyudpError::BindFailed)
+}
+
+#[derive(Debug)]
+pub struct Options {
+    pub timeout: Duration,
+    pub buffer_size: usize,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use anyhow::Result;
+
+    #[tokio::test]
+    async fn test_send_and_receive() -> Result<()> {
+        let response = send_and_receive(
+            "quake.se:28501",
+            b"\xff\xff\xff\xffstatus",
+            Options {
+                timeout: Duration::from_secs_f32(0.2),
+                buffer_size: 32 * 1024,
+            },
+        )
+        .await?;
+        assert!(String::from_utf8_lossy(&response).contains("QUAKE.SE KTX"));
+        Ok(())
+    }
+}