add checks for SRV records

Cargo.lock

@@ -118,7 +118,7 @@ checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94"
 
 [[package]]
 name = "mcstatusface"
-version = "1.1.1"
+version = "1.2.0"
 dependencies = [
  "chrono",
  "serde",

Cargo.toml

@@ -1,10 +1,10 @@
 [package]
 name = "mcstatusface"
-authors = ["ari melody <ari@arimelody.me>"]
+authors = ["ari melody <ari@arimelody.space>"]
-repository = "https://git.arimelody.me/ari/mcstatusface"
+repository = "https://codeberg.org/arimelody/mcstatusface"
 license = "MIT"
 keywords = ["minecraft", "server", "query", "web"]
-version = "1.1.1"
+version = "1.2.0"
 edition = "2024"
 
 [dependencies]

src/dns.rs

@@ -0,0 +1,128 @@
+use std::{io::{Error, ErrorKind, Result}, net::UdpSocket};
+
+const CLOUDFLARE_DNS: &str = "1.1.1.1";
+
+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>> {
+    let qname = match qname.is_ascii() {
+        true => qname,
+        false => return Err(Error::new(ErrorKind::InvalidInput, "domain is not valid ASCII")),
+    };
+
+    let txid    = 1u16;
+    let flags   = RECURSION_DESIRED;
+    let qdcount = 1u16;
+    let ancount = 0u16;
+    let nscount = 0u16;
+    let arcount = 0u16;
+
+    let header: [u16; 6] =
+        [txid, flags, qdcount, ancount, nscount, arcount]
+        .map(|el| el.to_be());
+
+    // let mut qname_bytes = bitvec![u8, Msb0;];
+    let mut query = Vec::<u8>::new();
+    for mut label in qname.split(".") {
+        label = match label.len() > u8::MAX.into() {
+            true => label.split_at(u8::MAX.into()).0,
+            false => label
+        };
+        let len: u8 = label.len().try_into().unwrap();
+        query.push(len);
+        query.append(&mut label.as_bytes().to_vec());
+    }
+    query.push(0);
+    query.append(&mut qtype.to_be_bytes().to_vec());
+    query.append(&mut qclass.to_be_bytes().to_vec());
+    
+    let mut request = Vec::<u8>::new();
+    for el in header {
+        request.push(el as u8);
+        request.push((el >> 8) as u8);
+    }
+    request.append(&mut query);
+
+    Ok(request)
+}
+
+#[derive(Debug)]
+pub struct DnsSrvResponse {
+    qref: u16,
+    rtype: u16,
+    class: u16,
+    ttl: u32,
+    priority: u16,
+    weight: u16,
+    port: u16,
+}
+
+pub fn parse_srv_response(mut ptr: usize, recv: &[u8]) -> Option<DnsSrvResponse> {
+    let mut data = DnsSrvResponse {
+        qref: 0,
+        rtype: 0,
+        class: 0,
+        ttl: 0,
+        priority: 0,
+        weight: 0,
+        port: 0,
+    };
+
+    fn read16(buf: &[u8], offset: usize) -> u16 {
+        let left = (buf[offset] as u16) << 8;
+        let right = buf[offset + 1] as u16;
+        left + right
+    }
+
+    fn read32(buf: &[u8], offset: usize) -> u32 {
+        let left = (read16(buf, offset) as u32) << 16;
+        let right = read16(buf, offset + 2) as u32;
+        left + right
+    }
+
+    data.qref = read16(recv, ptr);
+    ptr += 2;
+    data.rtype = read16(recv, ptr);
+    ptr += 2;
+    if data.rtype != RECORD_TYPE_SRV {
+        return None
+    }
+    data.class = read16(recv, ptr);
+    ptr += 2;
+    data.ttl = read32(recv, ptr);
+    ptr += 4;
+    let _rdata_len = read16(recv, ptr);
+    ptr += 2;
+    data.priority = read16(recv, ptr);
+    ptr += 2;
+    data.weight = read16(recv, ptr);
+    ptr += 2;
+    data.port = read16(recv, ptr);
+    // optional: read target?
+
+    // print!("data: {:#?}\n", data);
+
+    Some(data)
+}
+
+pub fn resolve_srv_port(domain: &String) -> Option<u16> {
+    let request = create_dns_query(
+        &("_minecraft._tcp.".to_string() + domain),
+        RECORD_TYPE_SRV,
+        RECORD_CLASS_IN).unwrap();
+
+    let socket = UdpSocket::bind("0.0.0.0:0").expect("failed to bind to port!");
+    socket.connect(CLOUDFLARE_DNS.to_string() + ":53").expect("failed to connect to dns server");
+    socket.send(&request).expect("failed to send dns request");
+
+    let mut recv_buf: [u8; 1024] = [0; 1024];
+    let _ = socket.recv_from(&mut recv_buf).unwrap();
+
+    match parse_srv_response(request.len(), &recv_buf) {
+        Some(data) => Some(data.port),
+        None => None,
+    }
+}

src/lib.rs

@@ -3,6 +3,7 @@ use std::{sync::{mpsc, Arc, Mutex}, thread};
 pub mod leb128;
 pub mod status;
 pub mod http;
+pub mod dns;
 pub use status::MinecraftStatus;
 
 pub struct ThreadPool {

src/main.rs

@@ -5,6 +5,7 @@ use std::{env};
 
 use mcstatusface::http::{HttpServer, StatusCode};
 use mcstatusface::{MinecraftStatus};
+use mcstatusface::dns::{resolve_srv_port};
 
 #[derive(serde::Serialize)]
 struct MinecraftStatusResponse<'a> {
@@ -16,6 +17,8 @@ struct MinecraftStatusResponse<'a> {
     motd: String,
 }
 
+const DEFAULT_PORT: u16 = 25565;
+
 fn main() -> Result<()> {
     let args: Vec<String> = env::args().collect();
     if args.len() < 2 {
@@ -33,7 +36,15 @@ env!("CARGO_PKG_VERSION"));
 
     if args[1] != "serve" {
         let mut address = String::from(args[1].as_str());
-        if !address.contains(":") { address.push_str(":25565"); }
+        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());
+        }
         let mut addrs_iter = address.to_socket_addrs().unwrap();
         let address = addrs_iter.next().unwrap();
 
@@ -98,7 +109,15 @@ env!("CARGO_PKG_VERSION"));
                     Some(query_address) => {
                         let mut address = query_address.to_string();
                         address = address.replace("%3A", ":");
-                        if !address.contains(":") { address.push_str(":25565"); }
+                        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);
@@ -179,7 +198,15 @@ env!("CARGO_PKG_VERSION"));
                 Some(query_address) => {
                     let mut address = query_address.to_string();
                     address = address.replace("%3A", ":");
-                    if !address.contains(":") { address.push_str(":25565"); }
+                    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);