add checks for SRV records

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,
+    }
+}