Improve querying

Author: Nickid2018

Cargo.toml

@@ -24,4 +24,4 @@ opt-level = "z"
 strip = true
 lto = true
 codegen-units = 1
-panic = "abort"
+panic = "abort"

src/mode/bedrock.rs

@@ -1,4 +1,5 @@
 use crate::analyze::{MotdInfo, StatusPayload};
+use crate::mode::QueryMode::BEDROCK;
 use crate::mode::QueryModeHandler;
 use crate::network::resolve::resolve_addr;
 use crate::network::util;
@@ -9,13 +10,13 @@ use std::io::ErrorKind;
 use std::net::SocketAddr;
 use std::time::Duration;
 use tokio::net::UdpSocket;
+use tokio::task::JoinSet;
 use tokio::time::timeout;
-use crate::mode::QueryMode::BEDROCK;
 
 const MAGIC_HIGH: u64 = 0x00ffff00fefefefeu64;
 const MAGIC_LOW: u64 = 0xfdfdfdfd12345678u64;
 
-async fn single_ip_check(addr: &SocketAddr) -> std::io::Result<StatusPayload> {
+async fn single_ip_check(addr: SocketAddr) -> std::io::Result<StatusPayload> {
     let timeout_time = Duration::from_secs(5);
     let socket = UdpSocket::bind("0.0.0.0:0").await?;
     socket.connect(addr).await?;
@@ -32,7 +33,7 @@ async fn single_ip_check(addr: &SocketAddr) -> std::io::Result<StatusPayload> {
 
     let mut recv_buf = [0u8; 1024];
     let recv = timeout(timeout_time, socket.recv_from(&mut recv_buf)).await??;
-    log::trace!("Received response");
+    log::trace!("Received response from {}", addr);
 
     let mut bytes = BytesMut::from(&recv_buf[..recv.0]);
     if bytes.get_u8() != 0x1C {
@@ -87,18 +88,31 @@ async fn single_ip_check(addr: &SocketAddr) -> std::io::Result<StatusPayload> {
     })
 }
 
-async fn check_bedrock_server(addr_vec: &Vec<SocketAddr>) -> std::io::Result<StatusPayload> {
+async fn safe_ip_check(addr: SocketAddr) -> std::io::Result<StatusPayload> {
+    match single_ip_check(addr).await {
+        Ok(status) => Ok(status),
+        Err(e) => {
+            log::warn!("Failed to check available server ip {}: {}", addr, e);
+            Err(e)
+        }
+    }
+}
+
+async fn check_bedrock_server(addr_vec: Vec<SocketAddr>) -> std::io::Result<StatusPayload> {
+    let mut set = JoinSet::new();
+
     for addr in addr_vec {
-        match single_ip_check(&addr).await {
-            Ok(r) => {
-                return Ok(r);
-            }
-            Err(e) => {
-                log::warn!("Failed to check available server ip {}: {}", addr, e);
-                continue;
-            }
+        set.spawn(safe_ip_check(addr));
+    }
+
+    while let Some(join_res) = set.join_next().await {
+        if let Ok(res) = join_res
+            && res.is_ok()
+        {
+            return res;
         }
     }
+
     Err(std::io::Error::new(ErrorKind::NotFound, "No server found"))
 }
 
@@ -107,8 +121,8 @@ pub struct BedrockQuery;
 #[async_trait]
 impl QueryModeHandler for BedrockQuery {
     async fn do_query(&self, addr: &str) -> std::io::Result<StatusPayload> {
-        let mut res = resolve_addr(addr, 19132);
-        let addrs = res.get_or_insert_default();
+        let res = resolve_addr(addr, 19132);
+        let addrs = res.unwrap_or(vec![]);
         check_bedrock_server(addrs).await
     }
 }

src/mode/java.rs

@@ -1,8 +1,10 @@
 use crate::analyze::{MotdInfo, PlayerInfo, StatusPayload};
+use crate::mode::QueryMode::JAVA;
 use crate::mode::QueryModeHandler;
 use crate::network::resolve::{resolve_addr, resolve_server_srv};
 use crate::network::schema::{read_string, read_var_int_stream, write_var_int};
-use crate::network::util;
+use crate::network::util::{io_timeout, now_timestamp};
+use crate::util::make_tcp_socket;
 use async_trait::async_trait;
 use bytes::{Buf, BufMut, BytesMut};
 use clap::Args;
@@ -11,22 +13,14 @@ use std::io::ErrorKind;
 use std::net::SocketAddr;
 use std::time::Duration;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use tokio::net::TcpSocket;
-use tokio::time::timeout;
-use crate::mode::QueryMode::JAVA;
+use tokio::task::JoinSet;
 
 async fn single_ip_check(addr: &SocketAddr, protocol: i32) -> std::io::Result<StatusPayload> {
-    let timeout_time = Duration::from_secs(5);
+    let time = Duration::from_secs(5);
     let ip_str = addr.ip().to_string();
 
-    let socket = if addr.is_ipv4() {
-        log::trace!("Using IPv4 socket to {}", addr);
-        TcpSocket::new_v4()?
-    } else {
-        log::trace!("Using IPv6 socket to {}", addr);
-        TcpSocket::new_v6()?
-    };
-    let mut stream = timeout(timeout_time, socket.connect(*addr)).await??;
+    let socket = make_tcp_socket(addr)?;
+    let mut stream = io_timeout(time, socket.connect(*addr), "Connection").await??;
 
     let mut handshake = vec![0];
     write_var_int(&mut handshake, protocol); // protocol_version
@@ -44,9 +38,9 @@ async fn single_ip_check(addr: &SocketAddr, protocol: i32) -> std::io::Result<St
 
     let handshake_recv_len = read_var_int_stream(&mut stream).await?;
     let mut handshake_recv = vec![0; handshake_recv_len as usize];
-    timeout(timeout_time, stream.read_exact(&mut handshake_recv)).await??;
+    io_timeout(time, stream.read_exact(&mut handshake_recv), "Handshake").await??;
     let mut recv_buf = BytesMut::from(handshake_recv.as_slice());
-    log::trace!("Handshake received, length: {}", handshake_recv.len());
+    log::trace!("Handshake received from {}, length: {}", addr, handshake_recv.len());
 
     if recv_buf.remaining() == 0 || recv_buf.get_u8() != 0 {
         return Err(std::io::Error::new(
@@ -68,20 +62,20 @@ async fn single_ip_check(addr: &SocketAddr, protocol: i32) -> std::io::Result<St
     let mut decoded: Value = from_str(&json_str).map_err(crate::util::wrap_invalid)?;
 
     stream.write(&[9, 1]).await?; // ping_request
-    stream.write_i64(util::now_timestamp()).await?;
+    stream.write_i64(now_timestamp()).await?;
     stream.flush().await?;
     log::trace!("Ping request sent");
 
     let recv_pong = &mut [0; 10];
-    timeout(timeout_time, stream.read_exact(recv_pong)).await??;
+    io_timeout(time, stream.read_exact(recv_pong), "Ping receiving").await??;
     if recv_pong[0] != 9 || recv_pong[1] != 1 {
         return Err(std::io::Error::new(
             ErrorKind::InvalidData,
             "Invalid pong packet",
         ));
     }
     let server_clock = i64::from_be_bytes(recv_pong[2..10].try_into().expect("Recv failed"));
-    let diff = util::now_timestamp() - server_clock;
+    let diff = now_timestamp() - server_clock;
     log::trace!("Got ping time: {}", diff);
 
     let players = decoded["players"].take();
@@ -123,21 +117,34 @@ async fn single_ip_check(addr: &SocketAddr, protocol: i32) -> std::io::Result<St
     })
 }
 
+async fn safe_ip_check(addr: SocketAddr, protocol: i32) -> std::io::Result<StatusPayload> {
+    match single_ip_check(&addr, protocol).await {
+        Ok(r) => Ok(r),
+        Err(e) => {
+            log::warn!("Failed to check available server ip {}: {}", addr, e);
+            Err(e)
+        }
+    }
+}
+
 async fn check_java_server(
-    addr_vec: &Vec<SocketAddr>,
+    addr_vec: Vec<SocketAddr>,
     protocol: i32,
 ) -> std::io::Result<StatusPayload> {
+    let mut set = JoinSet::new();
+
     for addr in addr_vec {
-        match single_ip_check(&addr, protocol).await {
-            Ok(r) => {
-                return Ok(r);
-            }
-            Err(e) => {
-                log::warn!("Failed to check available server ip {}: {}", addr, e);
-                continue;
-            }
+        set.spawn(safe_ip_check(addr, protocol));
+    }
+
+    while let Some(join_res) = set.join_next().await {
+        if let Ok(res) = join_res
+            && res.is_ok()
+        {
+            return res;
         }
     }
+
     Err(std::io::Error::new(ErrorKind::NotFound, "No server found"))
 }
 
@@ -155,19 +162,23 @@ pub struct JavaQuery<'a> {
     args: &'a JavaModeArgs,
 }
 
+pub async fn add_srv(addr: &str, addresses: &mut Vec<SocketAddr>) {
+    let srv_res = resolve_server_srv(addr).await;
+    let srv = srv_res
+        .iter()
+        .filter_map(|addr| resolve_addr(addr, 25565))
+        .flatten();
+    addresses.splice(0..0, srv);
+}
+
 #[async_trait]
 impl QueryModeHandler for JavaQuery<'_> {
     async fn do_query(&self, addr: &str) -> std::io::Result<StatusPayload> {
-        let mut je_res = resolve_addr(addr, 25565);
-        let je_address = je_res.get_or_insert_default();
+        let je_res = resolve_addr(addr, 25565);
+        let mut je_address = je_res.unwrap_or(vec![]);
 
         if !self.args.no_srv {
-            let srv_res = resolve_server_srv(addr).await;
-            let srv = srv_res
-                .iter()
-                .filter_map(|addr| resolve_addr(addr, 25565))
-                .flatten();
-            je_address.splice(0..0, srv);
+            add_srv(addr, &mut je_address).await;
         }
 
         check_java_server(je_address, self.args.protocol).await

src/mode/legacy.rs

@@ -1,34 +1,29 @@
 use crate::analyze::*;
+use crate::mode::QueryMode::LEGACY;
 use crate::mode::QueryModeHandler;
-use crate::mode::java::JavaModeArgs;
-use crate::network::resolve::{resolve_addr, resolve_server_srv};
+use crate::mode::java::{JavaModeArgs, add_srv};
+use crate::network::resolve::resolve_addr;
 use crate::network::util;
+use crate::network::util::{io_timeout, now_timestamp};
+use crate::util::make_tcp_socket;
 use async_trait::async_trait;
 use bytes::BufMut;
 use serde_json::json;
 use std::io::ErrorKind;
 use std::net::SocketAddr;
 use std::time::Duration;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use tokio::net::TcpSocket;
-use tokio::time::timeout;
-use crate::mode::QueryMode::LEGACY;
+use tokio::task::JoinSet;
 
 const LEGACY_HEADER: [u8; 27] = [
     0xFE, 0x01, 0xFA, 0x00, 0x0B, 0x00, 0x4D, 0x00, 0x43, 0x00, 0x7C, 0x00, 0x50, 0x00, 0x69, 0x00,
     0x6E, 0x00, 0x67, 0x00, 0x48, 0x00, 0x6F, 0x00, 0x73, 0x00, 0x74,
 ];
 
 async fn single_ip_check(addr: &SocketAddr) -> std::io::Result<StatusPayload> {
-    let timeout_time = Duration::from_secs(5);
-    let socket = if addr.is_ipv4() {
-        log::trace!("Using IPv4 socket to {}", addr);
-        TcpSocket::new_v4()?
-    } else {
-        log::trace!("Using IPv6 socket to {}", addr);
-        TcpSocket::new_v6()?
-    };
-    let mut stream = socket.connect(*addr).await?;
+    let time = Duration::from_secs(5);
+    let socket = make_tcp_socket(addr)?;
+    let mut stream = io_timeout(time, socket.connect(*addr), "Connection").await??;
 
     let ip_str = addr.ip().to_string();
     let utf16 = ip_str.encode_utf16().collect::<Vec<_>>();
@@ -43,13 +38,13 @@ async fn single_ip_check(addr: &SocketAddr) -> std::io::Result<StatusPayload> {
     buf.put_u16(0);
     buf.put_u16(addr.port());
 
-    let send_time = util::now_timestamp();
+    let send_time = now_timestamp();
     stream.write_all(&buf).await?;
     log::trace!("Legacy query sent, packet length = {}", packet_len);
 
     let mut recv_buffer = [0u8; 3];
-    timeout(timeout_time, stream.read_exact(&mut recv_buffer)).await??;
-    let ping = util::now_timestamp() - send_time;
+    io_timeout(time, stream.read_exact(&mut recv_buffer), "Handshake").await??;
+    let ping = now_timestamp() - send_time;
     log::trace!("Legacy query received, ping = {}", ping);
     if recv_buffer[0] != 0xFF {
         return Err(std::io::Error::new(
@@ -67,7 +62,7 @@ async fn single_ip_check(addr: &SocketAddr) -> std::io::Result<StatusPayload> {
         .map(|a| u16::from_be_bytes([a[0], a[1]]))
         .collect::<Vec<_>>();
     let str = String::from_utf16(&u16buf).map_err(crate::util::wrap_invalid)?;
-    log::trace!("Legacy query received: {}", str);
+    log::trace!("Legacy query received from {}: {}", addr, str);
 
     if str.starts_with("\u{00A7}1\0") {
         log::debug!("Legacy query version 1");
@@ -122,18 +117,31 @@ async fn single_ip_check(addr: &SocketAddr) -> std::io::Result<StatusPayload> {
     }
 }
 
-async fn check_legacy_server(addr_vec: &Vec<SocketAddr>) -> std::io::Result<StatusPayload> {
+async fn safe_ip_check(addr: SocketAddr) -> std::io::Result<StatusPayload> {
+    match single_ip_check(&addr).await {
+        Ok(r) => Ok(r),
+        Err(e) => {
+            log::warn!("Failed to check available server ip {}: {}", addr, e);
+            Err(e)
+        }
+    }
+}
+
+async fn check_legacy_server(addr_vec: Vec<SocketAddr>) -> std::io::Result<StatusPayload> {
+    let mut set = JoinSet::new();
+
     for addr in addr_vec {
-        match single_ip_check(&addr).await {
-            Ok(r) => {
-                return Ok(r);
-            }
-            Err(e) => {
-                log::warn!("Failed to check available server ip {}: {}", addr, e);
-                continue;
-            }
+        set.spawn(safe_ip_check(addr));
+    }
+
+    while let Some(join_res) = set.join_next().await {
+        if let Ok(res) = join_res
+            && res.is_ok()
+        {
+            return res;
         }
     }
+
     Err(std::io::Error::new(ErrorKind::NotFound, "No server found"))
 }
 
@@ -144,16 +152,11 @@ pub struct LegacyQuery<'a> {
 #[async_trait]
 impl QueryModeHandler for LegacyQuery<'_> {
     async fn do_query(&self, addr: &str) -> std::io::Result<StatusPayload> {
-        let mut res = resolve_addr(addr, 25565);
-        let addrs = res.get_or_insert_default();
+        let res = resolve_addr(addr, 25565);
+        let mut addrs = res.unwrap_or(vec![]);
 
         if !self.args.no_srv {
-            let srv_res = resolve_server_srv(addr).await;
-            let srv = srv_res
-                .iter()
-                .filter_map(|addr| resolve_addr(addr, 25565))
-                .flatten();
-            addrs.splice(0..0, srv);
+            add_srv(addr, &mut addrs).await;
         }
 
         check_legacy_server(addrs).await