Support UDP Proxying

Author: Nickid2018

Cargo.lock

@@ -11,9 +11,10 @@ ping-legacy = []
 anyhow = "1.0.100"
 async-trait = "0.1.89"
 bytes = "1.11.0"
-clap = { version = "4.5.52", features = ["derive"] }
+clap = { version = "4.5.53", features = ["derive"] }
 colored = "3.0.0"
 data-url = "0.3.2"
+fast-socks5 = "0.10.0"
 log = { version = "0.4.28", features = ["std"] }
 proxied = "1.1.0"
 regex = "1.12.2"

Cargo.toml

@@ -1,9 +1,9 @@
 use crate::analyze::{Analyzer, StatusPayload};
 use crate::network::schema::{read_string, read_var_int_buf};
+use anyhow::{Result, anyhow};
 use async_trait::async_trait;
 use bytes::{Buf, BufMut, BytesMut};
 use clap::Args;
-use std::io::ErrorKind;
 
 #[derive(Args, Debug)]
 pub struct ForgeInfoArgs {
@@ -16,13 +16,10 @@ pub struct ForgeInfo<'a> {
     args: &'a ForgeInfoArgs,
 }
 
-async fn try_analyze_encoded(data: &str, display_channels: bool) -> std::io::Result<()> {
+async fn try_analyze_encoded(data: &str, display_channels: bool) -> Result<()> {
     let chars = data.encode_utf16().collect::<Vec<u16>>();
     if chars.len() < 2 {
-        return Err(std::io::Error::new(
-            ErrorKind::InvalidData,
-            "ForgeData too short",
-        ));
+        return Err(anyhow!("ForgeData too short"));
     }
 
     let buffer_len = chars[0] as u32 | (chars[1] as u32) << 15;
@@ -132,7 +129,7 @@ impl Analyzer for ForgeInfo<'_> {
             if forge_data["truncated"].as_bool().unwrap_or(false) {
                 log::info!("Server truncated mod information");
             }
-            
+
             let default_vec = vec![];
             let mod_list = forge_data["mods"].as_array().unwrap_or(&default_vec);
             let ch_list = forge_data["channels"].as_array().unwrap_or(&default_vec);

src/analyze/forge_info.rs

@@ -1,28 +1,23 @@
 use crate::analyze::{MotdInfo, StatusPayload};
 use crate::mode::QueryMode::BEDROCK;
 use crate::mode::QueryModeHandler;
-use crate::network::resolve::{resolve_addr, sanitize_addr};
-use crate::network::util;
-use crate::network::util::io_timeout;
+use crate::network::connection::{ProxyableUdpSocket, UdpTarget, udp_socket};
+use crate::network::resolve::sanitize_addr;
+use crate::network::util::{generic_timeout, now_timestamp};
 use anyhow::{Result, anyhow};
 use async_trait::async_trait;
 use bytes::{Buf, BufMut, BytesMut};
 use serde_json::json;
-use std::net::SocketAddr;
 use std::time::Duration;
-use tokio::net::UdpSocket;
 use tokio::task::JoinSet;
 
 const MAGIC_HIGH: u64 = 0x00ffff00fefefefeu64;
 const MAGIC_LOW: u64 = 0xfdfdfdfd12345678u64;
 
-async fn single_ip_check(addr: SocketAddr) -> Result<StatusPayload> {
+async fn single_ip_check(addr: &UdpTarget, socket: ProxyableUdpSocket) -> Result<StatusPayload> {
     let timeout_time = Duration::from_secs(5);
-    let socket = UdpSocket::bind("0.0.0.0:0").await?;
-    socket.connect(addr).await?;
-    log::trace!("Connected to {}", addr);
 
-    let timestamp = util::now_timestamp();
+    let timestamp = now_timestamp();
     let mut packet = Vec::from([1u8]);
     packet.put_i64(timestamp);
     packet.put_u64(MAGIC_HIGH);
@@ -32,10 +27,10 @@ async fn single_ip_check(addr: SocketAddr) -> Result<StatusPayload> {
     log::trace!("Sent Unconnected Ping packet");
 
     let mut recv_buf = [0u8; 1024];
-    let recv = io_timeout(timeout_time, socket.recv_from(&mut recv_buf), "Recv").await?;
+    let recv = generic_timeout(timeout_time, socket.recv_from(&mut recv_buf), "Recv").await?;
     log::trace!("Received response from {}", addr);
 
-    let mut bytes = BytesMut::from(&recv_buf[..recv.0]);
+    let mut bytes = BytesMut::from(&recv_buf[..recv]);
     if bytes.get_u8() != 0x1C {
         return Err(anyhow!("Unexpected response byte"));
     }
@@ -49,11 +44,11 @@ async fn single_ip_check(addr: SocketAddr) -> Result<StatusPayload> {
     if magic_high != MAGIC_HIGH || magic_low != MAGIC_LOW {
         return Err(anyhow!("Invalid magic number"));
     }
-    if str_len as usize != recv.0 - 35 {
-        return Err(anyhow!("Invalid string length",));
+    if str_len as usize != recv - 35 {
+        return Err(anyhow!("Invalid string length"));
     }
 
-    let ping = util::now_timestamp() - server_clock;
+    let ping = now_timestamp() - server_clock;
     let resp = String::from_utf8(bytes.to_vec())?;
     log::trace!("Ping response: {}", resp);
 
@@ -65,55 +60,47 @@ async fn single_ip_check(addr: SocketAddr) -> Result<StatusPayload> {
     Ok(StatusPayload {
         mode: BEDROCK,
         ping,
-        max_players: Some(util::parse_i64(parts[5])?),
-        player_count: Some(util::parse_i64(parts[4])?),
+        max_players: Some(parts[5].parse()?),
+        player_count: Some(parts[4].parse()?),
         players: None,
         motd: Some(MotdInfo::String(format!("{}\n{}", parts[1], parts[7]))),
-        protocol: Some(util::parse_i64(parts[2])?),
+        protocol: Some(parts[2].parse()?),
         version_name: Some(parts[3].to_string()),
         favicon: None,
         full_extra: Some(json!({"server_guid": server_guid, "game_mode": parts[8].to_string()})),
     })
 }
 
-async fn safe_ip_check(addr: SocketAddr) -> Result<StatusPayload> {
-    match single_ip_check(addr).await {
+async fn safe_ip_check(addr: UdpTarget, socket: ProxyableUdpSocket) -> Result<StatusPayload> {
+    match single_ip_check(&addr, socket).await {
         Ok(status) => Ok(status),
-        Err(e) => Err(anyhow!(
-            "Protocol error in <{}:{}>: {}",
-            addr.ip(),
-            addr.port(),
-            e
-        )),
+        Err(e) => Err(anyhow!("Protocol error in <{}>: {}", addr, e)),
     }
 }
 
-async fn check_bedrock_server(addr_vec: Vec<SocketAddr>) -> Result<StatusPayload> {
-    let mut set = JoinSet::new();
-
-    for addr in addr_vec {
-        set.spawn(safe_ip_check(addr));
-    }
-
-    while let Some(join_res) = set.join_next().await {
-        if let Ok(res) = join_res {
-            match res {
-                Ok(res) => return Ok(res),
-                Err(e) => log::warn!("{}", e),
-            }
-        }
-    }
-
-    Err(anyhow!("No server found"))
-}
-
 pub struct BedrockQuery;
 
 #[async_trait]
 impl QueryModeHandler for BedrockQuery {
     async fn do_query(&self, addr: &str) -> Result<StatusPayload> {
         let (host, port) = sanitize_addr(addr, 19132)?;
-        check_bedrock_server(resolve_addr(&host, port)).await
+        let socks = udp_socket(&host, port, Duration::new(5, 0)).await?;
+        let mut set = JoinSet::new();
+
+        for sock in socks {
+            set.spawn(safe_ip_check(sock.0, sock.1));
+        }
+
+        while let Some(join_res) = set.join_next().await {
+            if let Ok(res) = join_res {
+                match res {
+                    Ok(res) => return Ok(res),
+                    Err(e) => log::warn!("{}", e),
+                }
+            }
+        }
+
+        Err(anyhow!("No server found"))
     }
 }
 

src/mode/bedrock.rs

@@ -1,10 +1,9 @@
-use crate::analyze::*;
+use crate::analyze::{MotdInfo, StatusPayload};
 use crate::mode::QueryMode::LEGACY;
 use crate::mode::QueryModeHandler;
 use crate::mode::java::JavaModeArgs;
 use crate::network::connection::connect_tcp;
 use crate::network::resolve::{resolve_server_srv, sanitize_addr};
-use crate::network::util;
 use crate::network::util::{io_timeout, now_timestamp};
 use anyhow::{Result, anyhow};
 use async_trait::async_trait;
@@ -71,11 +70,11 @@ async fn single_ip_check(addr: &str, port: u16, stream: &mut TcpStream) -> Resul
         Ok(StatusPayload {
             mode: LEGACY,
             ping,
-            max_players: Some(util::parse_i64(parts[5])?),
-            player_count: Some(util::parse_i64(parts[4])?),
+            max_players: Some(parts[5].parse()?),
+            player_count: Some(parts[4].parse()?),
             players: None,
             motd: Some(MotdInfo::String(parts[3].to_string())),
-            protocol: Some(util::parse_i64(parts[1])?),
+            protocol: Some(parts[1].parse()?),
             version_name: Some(parts[2].to_string()),
             favicon: None,
             full_extra: Some(json!({"legacy_version": 1})),
@@ -93,8 +92,8 @@ async fn single_ip_check(addr: &str, port: u16, stream: &mut TcpStream) -> Resul
         Ok(StatusPayload {
             mode: LEGACY,
             ping,
-            max_players: Some(util::parse_i64(parts[2])?),
-            player_count: Some(util::parse_i64(parts[1])?),
+            max_players: Some(parts[2].parse()?),
+            player_count: Some(parts[1].parse()?),
             players: None,
             motd: Some(MotdInfo::String(parts[0].to_string())),
             protocol: None,