dns.rs

use std::io;
use std::net::UdpSocket;

const QUESTION_OFFSET: usize = 12;

fn encode_label_format(name: &str) -> Vec<u8> {
    let mut encoded = Vec::new();
    for label in name.split('.') {
        encoded.push(label.len() as u8);
        encoded.extend_from_slice(label.as_bytes());
    }
    encoded.push(0);
    encoded
}
struct DNSHeader {
    id: u16,
    flags: u16,
    qdcount: u16,
    ancount: u16,
    nscount: u16,
    arcount: u16,
}

impl DNSHeader {
    fn from_bytes(buf: &[u8]) -> Self {
        if buf.len() < 12 {
            panic!("Buffer too short for DNS header");
        }
        DNSHeader {
            id: u16::from_be_bytes([buf[0], buf[1]]),
            flags: u16::from_be_bytes([buf[2], buf[3]]),
            qdcount: u16::from_be_bytes([buf[4], buf[5]]),
            ancount: u16::from_be_bytes([buf[6], buf[7]]),
            nscount: u16::from_be_bytes([buf[8], buf[9]]),
            arcount: u16::from_be_bytes([buf[10], buf[11]]),
        }
    }

    fn print(&self) {
        println!(
            "ID: {}, Flags: {}, QDCount: {}, ANCount: {}, NSCount: {}, ARCount: {}",
            self.id, self.flags, self.qdcount, self.ancount, self.nscount, self.arcount
        );
    }

    fn serialize(&self) -> Vec<u8> {
        let response = [
            self.id.to_be_bytes(),
            self.flags.to_be_bytes(),
            self.qdcount.to_be_bytes(),
            self.ancount.to_be_bytes(),
            self.nscount.to_be_bytes(),
            self.arcount.to_be_bytes(),
        ]
        .concat();
        response
    }
}

struct DNSQuestion {
    qname: String,
    qtype: u16,
    qclass: u16,
}

impl DNSQuestion {
    fn from_bytes(buf: &[u8]) -> Self {
        let mut i = 0;
        let mut labels = Vec::new();

        while i < buf.len() {
            let len = buf[i] as usize;
            if len == 0 {
                i += 1; // null terminator
                break;
            }
            i += 1;
            if i + len > buf.len() {
                panic!("QNAME label too long");
            }
            labels.push(String::from_utf8_lossy(&buf[i..i + len]).to_string());
            i += len;
        }

        let qname = labels.join(".");
        if i + 4 > buf.len() {
            panic!("Missing QTYPE/QCLASS");
        }

        let qtype = u16::from_be_bytes([buf[i], buf[i + 1]]);
        let qclass = u16::from_be_bytes([buf[i + 2], buf[i + 3]]);
        DNSQuestion {
            qname,
            qtype,
            qclass,
        }
    }

    fn serialize(&self) -> Vec<u8> {
        let response = [
            encode_label_format(&self.qname).as_slice(),
            &self.qtype.to_be_bytes(),
            &self.qclass.to_be_bytes(),
        ]
        .concat();
        response
    }

    fn print(&self) {
        println!(
            "Question: QNAME: {}, QTYPE: {}, QCLASS: {}",
            self.qname, self.qtype, self.qclass
        )
    }
}

struct DNSAnswer {
    name: String,
    rtype: u16,
    rclass: u16,
    ttl: u32,
    rdlength: u16,
    rdata: Vec<u8>,
}

impl DNSAnswer {
    fn from_string(data: &str) -> Self {
        let rdata = vec![127, 0, 0, 1];
        DNSAnswer {
            name: data.to_string(),
            rtype: 1,
            rclass: 1,
            ttl: 300,
            rdlength: rdata.len() as u16,
            rdata: rdata,
        }
    }

    fn print(&self) {
        println!(
            "Answer: NAME: {}, RTYPE: {}, RCLASS: {}, TTL: {}, RDLENGTH: {}, RDATA: {:?}",
            self.name, self.rtype, self.rclass, self.ttl, self.rdlength, self.rdata
        );
    }

    fn serialize(&self) -> Vec<u8> {
        let response: Vec<u8> = [
            encode_label_format(&self.name).as_slice(),
            &self.rtype.to_be_bytes(),
            &self.rclass.to_be_bytes(),
            &self.ttl.to_be_bytes(),
            &self.rdlength.to_be_bytes(),
            &self.rdata,
        ]
        .concat();
        response
    }
}

struct DNSResponse {
    header: DNSHeader,
    question: DNSQuestion,
    answers: Vec<DNSAnswer>,
}

impl DNSResponse {
    fn new(header: DNSHeader, question: DNSQuestion, answers: Vec<DNSAnswer>) -> Self {
        DNSResponse {
            header,
            question,
            answers,
        }
    }

    fn get(&mut self) -> Vec<u8> {
        self.header.flags = 0x8180;
        self.header.ancount = self.answers.len() as u16;

        let mut response = self.header.serialize();
        response.extend_from_slice(&self.question.serialize());

        for answer in self.answers.iter() {
            response.extend_from_slice(&answer.serialize());
        }
        return response;
    }
}

fn main() -> io::Result<()> {
    let socket = UdpSocket::bind("0.0.0.0:2053")?;
    println!("Listening on UDP port 2053...");
    let mut buf = [0u8; 512];

    loop {
        let (amt, src) = socket.recv_from(&mut buf)?;
        println!("Received {} bytes from {}", amt, src);

        let header = DNSHeader::from_bytes(&buf[..amt]);
        header.print();

        let question = DNSQuestion::from_bytes(&buf[QUESTION_OFFSET..amt]);
        question.print();

        let answer = DNSAnswer::from_string(&question.qname);
        answer.print();
        let answer2 = DNSAnswer::from_string("jamesnewport.isgay");
        answer2.print();

        let mut response = DNSResponse::new(header, question, vec![answer, answer2]);
        socket.send_to(&response.get(), src)?;
    }
}