mcstatusface/src/dns.rs

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