handle all resolved addresses, add verbose logging

src/dns.rs

@@ -7,7 +7,7 @@ const RECURSION_DESIRED: u16 = 1 << 8;
 const RECORD_TYPE_SRV: u16 = 33;
 const RECORD_CLASS_IN: u16 = 1;
 
-pub fn create_dns_query(qname: &String, qtype: u16, qclass: u16) -> Result<Vec<u8>> {
+pub fn create_dns_query(qname: &str, qtype: u16, qclass: u16) -> Result<Vec<u8>> {
     let qname = match qname.is_ascii() {
         true => qname,
         false => return Err(Error::new(ErrorKind::InvalidInput, "domain is not valid ASCII")),
@@ -110,7 +110,7 @@ pub fn parse_srv_response(mut ptr: usize, recv: &[u8]) -> Option<DnsSrvResponse>
 
 pub fn resolve_srv_port(domain: &String) -> Option<u16> {
     let request = create_dns_query(
-        &("_minecraft._tcp.".to_string() + domain),
+        &format!("_minecraft._tcp.{}", domain),
         RECORD_TYPE_SRV,
         RECORD_CLASS_IN).unwrap();
 

src/main.rs

@@ -1,4 +1,4 @@
-use std::io::Result;
+use std::io::{Error, ErrorKind, Result};
 use std::net::{IpAddr, ToSocketAddrs};
 use std::str::FromStr;
 use std::env;
@@ -34,8 +34,10 @@ env!("CARGO_PKG_VERSION"));
         std::process::exit(0);
     }
 
+    let verbose = args.contains(&"-v".to_string());
+
     if args[1] != "serve" {
-        let mut address = String::from(args[1].as_str());
+        let mut address = String::from(args[args.len() - 1].as_str());
         if !address.contains(":") {
             let port: u16 = match resolve_srv_port(&address) {
                 Some(port) => port,
@@ -45,10 +47,34 @@ env!("CARGO_PKG_VERSION"));
             address.push_str(":");
             address.push_str(port.to_string().as_str());
         }
-        let mut addrs_iter = address.to_socket_addrs().unwrap();
-        let address = addrs_iter.next().unwrap();
+        let addrs_iter = match address.to_socket_addrs() {
+            Ok(addr) => addr,
+            Err(e) => {
+                panic!("Failed to parse address {}: {}", address, e)
+            }
+        };
 
-        let status = MinecraftStatus::fetch(address).unwrap();
+        let mut error: Option<Error> = None;
+
+        if verbose {
+            println!(
+                "{} address{} available.",
+                addrs_iter.len(),
+                // horrible no good code but it's really funny to look at
+                // i'm sure this is awesome if you're into functional programming
+                (addrs_iter.len() == 1).then(|| "es").unwrap_or(""))
+        }
+
+        for address in addrs_iter {
+            let status = match MinecraftStatus::fetch(address, verbose) {
+                Ok(status) => status,
+                Err (e) => {
+                    error = Some(Error::new(
+                        ErrorKind::Other,
+                        format!("Failed to fetch status: {}", e)));
+                    continue
+                }
+            };
 
             println!("Version: {} ({})", status.version.name, status.version.protocol);
             println!("Players: {}/{}", status.players.online, status.players.max);
@@ -62,6 +88,11 @@ env!("CARGO_PKG_VERSION"));
             return Ok(());
         }
 
+        if let Some(error) = error {
+            return Err(error);
+        }
+    }
+
     let mut address = "0.0.0.0:8080".to_string();
     if args.len() > 2 { address = args[2].to_string() }
     let trusted_proxies: Vec<IpAddr> =
@@ -135,7 +166,7 @@ env!("CARGO_PKG_VERSION"));
                             Ok(mut addrs_iter) => {
                                 let address = addrs_iter.next().unwrap();
 
-                                match MinecraftStatus::fetch(address) {
+                                match MinecraftStatus::fetch(address, false) {
                                     Err(_) => {
                                         response.status(StatusCode::InternalServerError);
                                         query_response = format!(
@@ -219,7 +250,7 @@ env!("CARGO_PKG_VERSION"));
                         Ok(mut addrs_iter) => {
                             let address = addrs_iter.next().unwrap();
 
-                            match MinecraftStatus::fetch(address) {
+                            match MinecraftStatus::fetch(address, false) {
                                 Err(_) => {
                                     response.status(StatusCode::InternalServerError);
                                     response.body(format!("Failed to connect to {address}.\n"));

src/status.rs

@@ -1,8 +1,11 @@
 use std::io::{Error, ErrorKind, Read, Write, Result};
 use std::net::{SocketAddr, TcpStream};
+use std::time::Duration;
 
 use crate::leb128::{read_leb128, write_leb128};
 
+const TIMEOUT_SECS: u64 = 10;
+
 #[derive(serde::Serialize, serde::Deserialize)]
 pub struct MinecraftVersion {
     pub name: String,
@@ -55,17 +58,18 @@ pub struct MinecraftStatus {
 }
 
 impl MinecraftStatus {
-    pub fn fetch(address: SocketAddr) -> Result<MinecraftStatus> {
-        // println!("Connecting to {address}...");
+    pub fn fetch(address: SocketAddr, verbose: bool) -> Result<MinecraftStatus> {
+        if verbose { println!("Connecting to {address}..."); }
 
-        let stream = TcpStream::connect(address.to_string());
+        let stream = TcpStream::connect_timeout(
+            &address, Duration::from_secs(TIMEOUT_SECS));
         if stream.is_err() { return Err(stream.unwrap_err()); }
-        // println!("Connected!");
+        if verbose { println!("Connected!"); }
         let mut stream = stream.unwrap();
 
         let mut send_buffer: Vec<u8> = Vec::new();
 
-        // println!("Sending payload...");
+        if verbose { println!("Sending payload..."); }
         send_buffer.push(0x00);
         write_leb128(&mut send_buffer, 769); // 1.21.4
         write_leb128(&mut send_buffer, address.ip().to_string().len().try_into().unwrap());
@@ -76,6 +80,8 @@ impl MinecraftStatus {
 
         send_packet(&mut stream, &[0x00]).unwrap();
 
+        if verbose { println!("Awaiting response..."); }
+
         let mut data: Vec<u8> = Vec::new();
         let mut len: usize = 0;
         let mut msg_len: usize = 0;