ipv6 support, fix some awful string handling

src/dns.rs

@@ -108,7 +108,7 @@ pub fn parse_srv_response(mut ptr: usize, recv: &[u8]) -> Option<DnsSrvResponse>
     Some(data)
 }
 
-pub fn resolve_srv_port(domain: &String) -> Option<u16> {
+pub fn resolve_srv_port(domain: &str) -> Option<u16> {
     let request = create_dns_query(
         &("_minecraft._tcp.".to_string() + domain),
         RECORD_TYPE_SRV,

src/http.rs

@@ -48,8 +48,8 @@ pub struct Request<'a> {
     path: &'a str,
     method: &'a str,
     version: &'a str,
-    headers: HashMap<String, &'a str>,
+    headers: HashMap<String, String>,
-    query: HashMap<&'a str, &'a str>,
+    query: HashMap<String, String>,
     body: Option<String>,
     real_address: IpAddr,
 }
@@ -65,15 +65,15 @@ impl<'a> Request<'a> {
         let mut path = request_line_split[1];
         let version = request_line_split[2];
 
-        let mut query: HashMap<&'a str, &'a str> = HashMap::new();
+        let mut query: HashMap<String, String> = HashMap::new();
         match request_line_split[1].split_once("?") {
             Some((path_without_query, query_string)) => {
                 path = path_without_query;
-                let query_splits: Vec<&'a str> = query_string.split("&").collect();
+                let query_splits: Vec<&str> = query_string.split("&").collect();
                 for pair in query_splits {
                     match pair.split_once("=") {
                         Some((name, value)) => {
-                            query.insert(name, value);
+                            query.insert(name.to_owned(), value.to_owned());
                         }
                         None => { continue; }
                     }
@@ -82,7 +82,7 @@ impl<'a> Request<'a> {
             None => {},
         }
 
-        let mut headers: HashMap<String, &'a str> = HashMap::new();
+        let mut headers: HashMap<String, String> = HashMap::new();
         if lines.len() > 1 {
             let mut i: usize = 1;
             loop {
@@ -90,7 +90,9 @@ impl<'a> Request<'a> {
                 let line = &lines[i];
                 if line.len() == 0 || !line.contains(":") { break; }
                 let (name, value) = line.split_once(":").unwrap();
-                headers.insert(name.to_lowercase(), value.trim());
+                headers.insert(
+                    name.to_lowercase().to_owned(),
+                    value.trim().to_owned());
                 i += 1;
             }
         }
@@ -149,10 +151,10 @@ impl<'a> Request<'a> {
     pub fn body(&self) -> &Option<String> {
         &self.body
     }
-    pub fn headers(&self) -> &HashMap<String, &'a str> {
+    pub fn headers(&self) -> &HashMap<String, String> {
         &self.headers
     }
-    pub fn query(&self) -> &HashMap<&'a str, &'a str> {
+    pub fn query(&self) -> &HashMap<String, String> {
         &self.query
     }
 }

src/main.rs

@@ -1,24 +1,199 @@
-use std::io::Result;
+use std::io::{Error, ErrorKind, Result};
-use std::net::{IpAddr, ToSocketAddrs};
+use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
 use std::str::FromStr;
 use std::env;
 
-use mcstatusface::http::{HttpServer, StatusCode};
+use mcstatusface::http::{HttpServer, Request, Response, StatusCode};
 use mcstatusface::status::MinecraftStatus;
 use mcstatusface::dns::resolve_srv_port;
 
 #[derive(serde::Serialize)]
 struct MinecraftStatusResponse<'a> {
-    version: &'a String,
+    version: &'a str,
     players: u32,
     max_players: u32,
     enforces_secure_chat: bool,
-    favicon: Option<&'a String>,
+    favicon: Option<&'a str>,
-    motd: String,
+    motd: &'a str,
 }
 
 const DEFAULT_PORT: u16 = 25565;
 
+fn parse_address(address: &str) -> Result<SocketAddr> {
+    // attempt to parse as IP + port
+    if let Ok(socket_addr) = SocketAddr::from_str(address) {
+        return Ok(socket_addr);
+    }
+
+    // handle IPv6 address between brackets
+    let mut address = address;
+    if address.starts_with("[") && address.ends_with("]") {
+        address = &address[1..address.len() - 1];
+    }
+
+    let port: u16 = match resolve_srv_port(address) {
+        Some(port) => port,
+        None => DEFAULT_PORT,
+    };
+
+    // handle IP address with no port
+    if let Ok(ip) = IpAddr::from_str(address) {
+        return Ok(SocketAddr::new(ip, port))
+    }
+
+    match format!("{}:{}", address, port).to_socket_addrs() {
+        Ok(mut addr) => {
+            return Ok(addr.next().unwrap())
+        }
+        Err(_) => {
+            return Err(Error::new(
+                    ErrorKind::InvalidInput,
+                    format!("Invalid server address: {}", address)))
+        }
+    }
+}
+
+fn handle_html_request(request: &Request, response: &mut Response) -> Result<StatusCode> {
+    let content = include_str!("views/index.html");
+    response.set_header("Content-Type", "text/html".to_string());
+    response.status(StatusCode::OK);
+    let mut query_response = "".to_string();
+
+    let query_address = request.query().get("s");
+    if let Some(query_address) = query_address {
+        match parse_address(&query_address.replace("%3A", ":")) {
+            Err(_) => {
+                response.status(StatusCode::BadRequest);
+                query_response = format!(
+                    "<hr/>
+                            <h2>Server Details</h2>
+                            <pre><code>Invalid server address: {}.</pre></code>",
+                            sanitize_html(query_address),
+                );
+            },
+            Ok(address) => {
+                match MinecraftStatus::fetch(&address) {
+                    Err(err) => {
+                        println!(
+                            "Failed to connect to {} ({}): {}",
+                            query_address, address, err);
+
+                        response.body(format!(
+                                "Failed to connect to {} ({}).\n",
+                                query_address, address));
+                        response.status(StatusCode::InternalServerError);
+                        query_response = format!(
+                            "<hr/>
+                            <h2>Server Details</h2>
+                            <pre><code>Failed to connect to {} ({}).</pre></code>",
+                            sanitize_html(query_address),
+                            sanitize_html(&address.to_string()),
+                        );
+                    }
+                    Ok(status) => {
+                        let minecraft_status = MinecraftStatusResponse{
+                            version: &status.version.name,
+                            players: status.players.online,
+                            max_players: status.players.max,
+                            enforces_secure_chat: status.enforces_secure_chat(),
+                            favicon: status.favicon.as_deref(),
+                            motd: &status.parse_description(),
+                        };
+
+                        query_response = format!(
+                            "<hr/>
+                            <h2>Server Details</h2>
+                            <p>
+                            <strong>Favicon:</strong></br>
+                            <img width=\"64\" height=\"64\" src=\"{}\"></br>
+                            <strong>Version:</strong> <code>{}</code><br/>
+                            <strong>Players:</strong> <code>{}/{}</code><br/>
+                            <strong>Enforces Secure Chat:</strong> <code>{}</code></br>
+                            <strong>MOTD:</strong>
+                            </p>
+                            <pre id=\"motd\"><code>{}</code></pre>",
+                            minecraft_status.favicon.map_or("", |fav| fav),
+                            sanitize_html(minecraft_status.version).to_string(),
+                            minecraft_status.players,
+                            minecraft_status.max_players,
+                            if minecraft_status.enforces_secure_chat { "true" } else { "false" },
+                            sanitize_html(&minecraft_status.motd).to_string(),
+                        );
+                    }
+                }
+            },
+        };
+    }
+
+    let response_content = content
+        .replace("{{response}}", &query_response)
+        .replace("{{host}}", match request.headers().get("host") {
+            None => { "mcq.bliss.town" }
+            Some(host) => { host }
+        });
+    response.set_header("Content-Type", "text/html".to_string());
+    response.body(response_content.to_string());
+    return response.send();
+}
+
+fn handle_json_request(request: &Request, response: &mut Response) -> Result<StatusCode> {
+    response.set_header("Content-Type", "text/plain".to_string());
+    response.status(StatusCode::OK);
+
+    let query_address = request.query().get("s");
+
+    if let Some(query_address) = query_address {
+        match parse_address(&query_address.replace("%3A", ":")) {
+            Err(_) => {
+                response.status(StatusCode::BadRequest);
+                response.body("Invalid server address.\n".to_string());
+            }
+            Ok(address) => {
+                match MinecraftStatus::fetch(&address) {
+                    Err(err) => {
+                        println!(
+                            "Failed to connect to {} ({}): {}",
+                            query_address, address, err);
+
+                        response.status(StatusCode::InternalServerError);
+                        response.body(format!(
+                                "Failed to connect to {} ({}).\n",
+                                query_address, address));
+                    }
+                    Ok(status) => {
+                        let minecraft_status = MinecraftStatusResponse{
+                            version: &status.version.name,
+                            players: status.players.online,
+                            max_players: status.players.max,
+                            enforces_secure_chat: status.enforces_secure_chat(),
+                            favicon: status.favicon.as_deref(),
+                            motd: &status.parse_description(),
+                        };
+
+                        match serde_json::to_string(&minecraft_status) {
+                            Err(e) => {
+                                eprintln!("Failed to format response from {} to JSON: {}", address, e);
+                                response.status(StatusCode::InternalServerError);
+                                response.body("Internal Server Error\n".to_string());
+                            }
+                            Ok(json) => {
+                                response.status(StatusCode::OK);
+                                response.set_header("Content-Type", "application/json".to_string());
+                                response.body(json + "\n");
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    } else {
+        response.status(StatusCode::OK);
+        response.body("?s=<server address>\n".to_string());
+    }
+
+    return response.send();
+}
+
 fn main() -> Result<()> {
     let args: Vec<String> = env::args().collect();
     if args.len() < 2 {
@@ -35,20 +210,8 @@ env!("CARGO_PKG_VERSION"));
     }
 
     if args[1] != "serve" {
-        let mut address = String::from(args[1].as_str());
+        let address = parse_address(&args[1]).expect("Failed to parse address");
-        if !address.contains(":") {
+        let status = MinecraftStatus::fetch(&address).unwrap();
-            let port: u16 = match resolve_srv_port(&address) {
-                Some(port) => port,
-                None => DEFAULT_PORT,
-            };
-
-            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 status = MinecraftStatus::fetch(address).unwrap();
 
         println!("Version: {} ({})", status.version.name, status.version.protocol);
         println!("Players: {}/{}", status.players.online, status.players.max);
@@ -101,173 +264,22 @@ env!("CARGO_PKG_VERSION"));
             if request.headers().get("accept").is_some_and(
                 |accept| accept.contains("text/html")
             ) {
-                // HTML response
+                return handle_html_request(request, &mut response);
-                let content = include_str!("views/index.html");
-                response.set_header("Content-Type", "text/html".to_string());
-                response.status(StatusCode::OK);
-                let query_response: String;
-                match request.query().get("s") {
-                    None => {
-                        query_response = String::from("");
-                    }
-                    Some(query_address) => {
-                        let mut address = query_address.to_string();
-                        address = address.replace("%3A", ":");
-                        if !address.contains(":") {
-                            let port: u16 = match resolve_srv_port(&address) {
-                                Some(port) => port,
-                                None => DEFAULT_PORT,
-                            };
-
-                            address.push_str(":");
-                            address.push_str(port.to_string().as_str());
-                        }
-                        match address.to_socket_addrs() {
-                            Err(_) => {
-                                response.status(StatusCode::BadRequest);
-                                query_response = format!(
-                                    "<hr/>
-                            <h2>Server Details</h2>
-                            <pre><code>Invalid server address: {}.</pre></code>",
-                            sanitize_html(&address.to_string()),
-                                );
-                            }
-                            Ok(mut addrs_iter) => {
-                                let address = addrs_iter.next().unwrap();
-
-                                match MinecraftStatus::fetch(address) {
-                                    Err(e) => {
-                                        println!(
-                                            "Failed to connect to {} ({}): {}",
-                                            query_address, address, e);
-                                        response.status(StatusCode::InternalServerError);
-                                        query_response = format!(
-                                            "<hr/>
-                            <h2>Server Details</h2>
-                            <pre><code>Failed to connect to {}.</pre></code>",
-                                            sanitize_html(&address.to_string()),
-                                            );
-                                    }
-                                    Ok(status) => {
-                                        let minecraft_status = MinecraftStatusResponse{
-                                            version: &status.version.name,
-                                            players: status.players.online,
-                                            max_players: status.players.max,
-                                            enforces_secure_chat: status.enforces_secure_chat(),
-                                            favicon: status.favicon.as_ref(),
-                                            motd: status.parse_description(),
-                                        };
-
-                                        query_response = format!(
-                                            "<hr/>
-                            <h2>Server Details</h2>
-                            <p>
-                            <strong>Favicon:</strong></br>
-                            <img width=\"64\" height=\"64\" src=\"{}\"></br>
-                            <strong>Version:</strong> <code>{}</code><br/>
-                            <strong>Players:</strong> <code>{}/{}</code><br/>
-                            <strong>Enforces Secure Chat:</strong> <code>{}</code></br>
-                            <strong>MOTD:</strong>
-                            </p>
-                            <pre id=\"motd\"><code>{}</code></pre>",
-                            sanitize_html(&minecraft_status.favicon.map_or("", |s| s).to_string()),
-                            sanitize_html(minecraft_status.version).to_string(),
-                            minecraft_status.players,
-                            minecraft_status.max_players,
-                            if minecraft_status.enforces_secure_chat { "true" } else { "false" },
-                            sanitize_html(&minecraft_status.motd).to_string(),
-                                        );
-                                    }
-                                }
-                            }
-                        }
-                    }
-                }
-
-                let response_content = content
-                    .replace("{{response}}", &query_response)
-                    .replace("{{host}}", match request.headers().get("host") {
-                        None => { "mcq.bliss.town" }
-                        Some(host) => { host }
-                    });
-                response.set_header("Content-Type", "text/html".to_string());
-                response.body(response_content.to_string());
-                return response.send();
             }
 
             // JSON response
-            match request.query().get("s") {
+            return handle_json_request(request, &mut response);
-                None => {
-                    response.status(StatusCode::OK);
-                    response.body("?s=<server address>\n".to_string());
-                    return response.send();
-                }
-                Some(query_address) => {
-                    let mut address = query_address.to_string();
-                    address = address.replace("%3A", ":");
-                    if !address.contains(":") {
-                        let port: u16 = match resolve_srv_port(&address) {
-                            Some(port) => port,
-                            None => DEFAULT_PORT,
-                        };
-
-                        address.push_str(":");
-                        address.push_str(port.to_string().as_str());
-                    }
-                    match address.to_socket_addrs() {
-                        Err(_) => {
-                            response.status(StatusCode::BadRequest);
-                            response.body("Invalid server address.\n".to_string());
-                        }
-                        Ok(mut addrs_iter) => {
-                            let address = addrs_iter.next().unwrap();
-
-                            match MinecraftStatus::fetch(address) {
-                                Err(_) => {
-                                    response.status(StatusCode::InternalServerError);
-                                    response.body(format!("Failed to connect to {address}.\n"));
-                                }
-                                Ok(status) => {
-                                    let minecraft_status = MinecraftStatusResponse{
-                                        version: &status.version.name,
-                                        players: status.players.online,
-                                        max_players: status.players.max,
-                                        enforces_secure_chat: status.enforces_secure_chat(),
-                                        favicon: status.favicon.as_ref(),
-                                        motd: status.parse_description(),
-                                    };
-
-                                    match serde_json::to_string(&minecraft_status) {
-                                        Err(e) => {
-                                            eprintln!("Failed to format response from {} to JSON: {}", address, e);
-                                            response.status(StatusCode::InternalServerError);
-                                            response.body("Internal Server Error\n".to_string());
-                                        }
-                                        Ok(json) => {
-                                            response.status(StatusCode::OK);
-                                            response.set_header("Content-Type", "application/json".to_string());
-                                            response.body(json + "\n");
-                                        }
-                                    }
-                                }
-                            }
-
-                        }
-                    }
-                    return response.send()
-                }
-            }
         }
 
         response.status(StatusCode::NotFound);
         response.body("Not Found".to_string());
         return response.send();
-    }).unwrap();
+    }).expect("Failed to start HTTP server");
 
     Ok(())
 }
 
-fn sanitize_html(input: &String) -> String {
+fn sanitize_html(input: &str) -> String {
     input
         .replace("&", "&amp;")
         .replace("<", "&lt;")

src/status.rs

@@ -56,10 +56,10 @@ pub struct MinecraftStatus {
 }
 
 impl MinecraftStatus {
-    pub fn fetch(address: SocketAddr) -> Result<MinecraftStatus> {
+    pub fn fetch(address: &SocketAddr) -> Result<MinecraftStatus> {
         //println!("Connecting to {address}...");
 
-        let stream = TcpStream::connect_timeout(&address, Duration::new(5, 0));
+        let stream = TcpStream::connect_timeout(address, Duration::new(5, 0));
         if stream.is_err() { return Err(stream.unwrap_err()); }
 
         //println!("Connected!");
@@ -124,9 +124,11 @@ impl MinecraftStatus {
                 break;
             }
         }
-        let msg = std::str::from_utf8(&data[offset..]).unwrap().trim();
+
-        let sanitised: String = msg.chars().filter(|&c| c >= '\u{20}' || c == '\n' || c == '\r' || c == '\t').collect();
+        let data = std::str::from_utf8(&data[offset..]).expect("Failed to parse UTF-8 data");
-        let status: MinecraftStatus = serde_json::from_slice(sanitised.as_bytes()).unwrap();
+        let sanitised: String = data.chars().filter(
+            |&c| c >= '\u{20}' || c == '\n' || c == '\r' || c == '\t').collect();
+        let status: MinecraftStatus = serde_json::from_str(&sanitised).expect("Failed to parse JSON");
 
         Ok(status)
     }
@@ -155,12 +157,12 @@ fn _description(description: &MinecraftDescription) -> String {
                         }
                     }
                 }
-                return description.text.clone() + &extras;
+                return description.text.to_owned() + &extras;
             }
-            description.text.clone()
+            description.text.to_owned()
         }
         MinecraftDescription::Plain(description) => {
-            description.clone()
+            description.to_owned()
         }
     }
 }