add DNS SRV lookup functionality

Cargo.toml

@@ -17,6 +17,7 @@ smart-default = "0.6.0"
 structopt = "0.3.23"
 term_size = "0.3.2"
 termcolor = "1.1.2"
+trust-dns-resolver = { version = "0.20.3", features = ["tokio-runtime"] }
 unicode-width = "0.1.9"
 viuer = "0.5.2"
 

src/lib.rs

@@ -9,7 +9,11 @@ use crossterm::{
 use image::{DynamicImage, ImageFormat};
 use itertools::Itertools;
 use miette::{bail, miette, IntoDiagnostic, WrapErr};
-use std::io::{self, Cursor, Write};
+use std::{
+    io::{self, Cursor, Write},
+    net::IpAddr,
+};
+use trust_dns_resolver::TokioAsyncResolver;
 
 pub mod output;
 
@@ -30,6 +34,54 @@ macro_rules! none_if_empty {
     }};
 }
 
+pub async fn resolve_address(addr_and_port: &str) -> miette::Result<(String, u16)> {
+    let addr;
+    let port;
+    if let Some((addr_, port_)) = addr_and_port.split_once(':') {
+        addr = addr_;
+        port = Some(
+            port_
+                .parse()
+                .into_diagnostic()
+                .wrap_err("User provided port is invalid")?,
+        );
+    } else {
+        addr = addr_and_port;
+        port = None;
+    }
+
+    if let Some(port) = port {
+        Ok((addr.to_string(), port))
+    } else if addr.parse::<IpAddr>().is_ok() {
+        // if we only have an IP and no port, there is no domain to lookup so we can
+        // only default to port 25565.
+        Ok((addr.to_string(), 25565))
+    } else {
+        let dns = TokioAsyncResolver::tokio_from_system_conf()
+            .into_diagnostic()
+            .wrap_err("Failed to create DNS resolver")?;
+
+        let lookup = dns.srv_lookup(format!("_minecraft._tcp.{}.", addr)).await;
+
+        if let Ok(lookup) = lookup {
+            let srv = lookup
+                .iter()
+                .next()
+                .ok_or_else(|| miette!("No SRV record found"))?;
+
+            let addr = srv.target().to_string();
+            let addr = addr.trim_end_matches('.');
+
+            let port = srv.port();
+
+            Ok((addr.to_string(), port))
+        } else {
+            // if there is no SRV record, we have to default to port 25565
+            Ok((addr.to_string(), 25565))
+        }
+    }
+}
+
 /// Print mincraft-formatted text to `out` using crossterm
 pub fn print_mc_formatted(s: &str, mut out: impl Write) -> io::Result<()> {
     macro_rules! exec {
@@ -82,7 +134,7 @@ pub fn print_mc_formatted(s: &str, mut out: impl Write) -> io::Result<()> {
                 'n' => exec!(at, Underlined),
                 'o' => exec!(at, Italic),
                 'r' => exec!(ResetColor),
-                _ => {}
+                _ => {},
             }
             exec!(Print(&split[1..]));
         }

src/main.rs

@@ -1,12 +1,19 @@
 use async_minecraft_ping::{ConnectionConfig, ServerDescription, StatusResponse};
 
 use itertools::Itertools;
-use miette::{miette, IntoDiagnostic, WrapErr};
+use miette::{IntoDiagnostic, WrapErr};
 use structopt::StructOpt;
 use time::{Duration, Instant};
 use tokio::time;
 
-use mcstat::{get_table, mc_formatted_to_ansi, none_if_empty, output::Table, parse_base64_image};
+use mcstat::{
+    get_table,
+    mc_formatted_to_ansi,
+    none_if_empty,
+    output::Table,
+    parse_base64_image,
+    resolve_address,
+};
 
 #[derive(Debug, StructOpt)]
 #[structopt(
@@ -16,8 +23,8 @@ use mcstat::{get_table, mc_formatted_to_ansi, none_if_empty, output::Table, pars
 struct Opt {
     #[structopt(
         index = 1,
-        help = "the ip of the server to ping. you may also specify the port, if it is not \
-                specified or invalid it will default to 25565"
+        help = "The Address to ping. By default, a SRV lookup will be made to resolve this, \
+                unless the port is specified."
     )]
     ip: String,
 
@@ -77,16 +84,13 @@ impl Opt {
 async fn main() -> miette::Result<()> {
     let opt = Opt::from_args();
 
-    let mut ip = opt.ip.splitn(2, ':');
-    let config =
-        ConnectionConfig::build(ip.next().ok_or_else(|| miette!("invalid ip"))?.to_owned())
-            .with_port(
-                ip.next()
-                    .map_or(Err(()), |p| p.parse::<u16>().map_err(|_| ()))
-                    .and_then(|p| if p > 0 { Ok(p) } else { Err(()) })
-                    .unwrap_or(25565),
-            )
-            .with_protocol_version(opt.protocol_version);
+    let (addr, port) = resolve_address(&opt.ip)
+        .await
+        .wrap_err("Error resolving address")?;
+
+    let config = ConnectionConfig::build(addr)
+        .with_port(port)
+        .with_protocol_version(opt.protocol_version);
 
     // create timeout for server connection
     let (raw_response, ping) = time::timeout(Duration::from_millis(opt.timeout), async {