feat: refactoring

src/dns.rs

@@ -1,11 +1,13 @@
 #![allow(unused)]
 use std::fmt::Display;
 
+use anyhow::{Result, bail};
+
 const QUESTION_TYPE_A: u16 = 0x01;
 const QUESTION_TYPE_NS: u16 = 0x02;
 const QUESTION_TYPE_CNAME: u16 = 0x05;
+const QUESTION_TYPE_AAAA: u16 = 0x1C;
 const QUESTION_TYPE_MX: u16 = 0x0F;
-const QUESTION_TYPE_AAAA: u16 = 0x40;
 
 const QUESTION_CLASS_IN: u16 = 0x01;
 const QUESTION_CLASS_CHAOS: u16 = 0x03;
@@ -29,6 +31,33 @@ pub struct DnsHeader {
     pub number_of_additional_resource_records: u16,
 }
 
+impl DnsHeader {
+    pub fn parse(packet: &[u8]) -> Result<DnsHeader> {
+        if packet.len() < 12 {
+            bail!("malformed datagram");
+        }
+
+        let datagram = DnsHeader {
+            transaction_id: [packet[0], packet[1]],
+            query_response: packet[2] & (1 << 7) > 0,
+            op_code: (packet[2] >> 3) & 0x0F,
+            authoritative_answer: packet[2] & (1 << 2) > 0,
+            truncation: packet[2] & (1 << 1) > 0,
+            recursion_desired: packet[2] & (1 << 0) > 0,
+            recursion_available: packet[3] & (1 << 7) > 0,
+            authentic_data: packet[3] & (1 << 5) > 0,
+            checking_disabled: packet[3] & (1 << 4) > 0,
+            response_code: packet[3] & 0xF,
+            number_of_questions: u16::from_be_bytes([packet[4], packet[5]]),
+            number_of_answers: u16::from_be_bytes([packet[6], packet[7]]),
+            number_of_authority_resource_records: u16::from_be_bytes([packet[8], packet[9]]),
+            number_of_additional_resource_records: u16::from_be_bytes([packet[10], packet[11]]),
+        };
+
+        Ok(datagram)
+    }
+}
+
 impl Display for DnsHeader {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         writeln!(
@@ -67,6 +96,28 @@ pub struct DnsQuestion {
     pub question_class: u16,
 }
 
+impl DnsQuestion {
+    pub fn parse(packet: &[u8], pos: &mut usize) -> Result<DnsQuestion> {
+        let domain = parse_domain(packet, pos)?;
+
+        if *pos + 4 > packet.len() {
+            bail!("malformed datagram");
+        }
+
+        let question_type = ((packet[*pos] as u16) << 8) | (packet[*pos + 1] as u16);
+        let question_class = ((packet[*pos + 2] as u16) << 8) | (packet[*pos + 3] as u16);
+        *pos += 4;
+
+        let question = DnsQuestion {
+            domain,
+            question_type,
+            question_class,
+        };
+
+        Ok(question)
+    }
+}
+
 impl Display for DnsQuestion {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(
@@ -79,6 +130,7 @@ impl Display for DnsQuestion {
 
 pub struct DnsAnswer {
     pub domain: String,
+    // TODO missing fields
 }
 
 impl DnsAnswer {
@@ -90,6 +142,38 @@ impl DnsAnswer {
             part.bytes().for_each(|c| result.push(c));
         }
 
+        result.push(0);
         result
     }
 }
+
+fn parse_domain(packet: &[u8], pos: &mut usize) -> Result<String> {
+    let packet_size = packet.len();
+    let mut domain_parts: Vec<String> = vec![];
+
+    loop {
+        if *pos >= packet_size {
+            bail!("malformed datagram");
+        }
+
+        let len = packet[*pos] as usize;
+        if len > 63 {
+            bail!("malformed datagram");
+        }
+        if len == 0 {
+            *pos += 1;
+            break;
+        }
+
+        if (*pos + 1 + len) > packet_size {
+            bail!("malformed datagram");
+        }
+
+        let domain_part_bytes = &packet[(*pos + 1)..(*pos + 1 + len)];
+        let domain_part = String::from_utf8_lossy(&domain_part_bytes);
+        domain_parts.push(domain_part.to_string());
+        *pos += 1 + len;
+    }
+
+    Ok(domain_parts.join("."))
+}

src/main.rs

@@ -20,58 +20,23 @@ async fn main() -> Result<()> {
     let mut buffer = [0u8; 512];
 
     loop {
-        let size = sock.recv(&mut buffer).await?;
-        handle_datagram(buffer.clone(), size);
+        let (size, addr) = sock.recv_from(&mut buffer).await?;
+        tracing::info!("Received datagram from {}", addr.to_string());
+        if let Err(err) = handle_datagram(&buffer[..size]) {
+            tracing::error!("failed to handle datagram: {err}");
+        }
     }
 }
 
-fn handle_datagram(buffer: [u8; 512], _buffer_size: usize) {
-    let datagram = DnsHeader {
-        transaction_id: [buffer[0], buffer[1]],
-        query_response: buffer[2] & (1 << 7) > 0,
-        op_code: (buffer[2] >> 3) & 0x0F,
-        authoritative_answer: buffer[2] & (1 << 2) > 0,
-        truncation: buffer[2] & (1 << 1) > 0,
-        recursion_desired: buffer[2] & (1 << 0) > 0,
-        recursion_available: buffer[3] & (1 << 7) > 0,
-        authentic_data: buffer[3] & (1 << 5) > 0,
-        checking_disabled: buffer[3] & (1 << 4) > 0,
-        response_code: buffer[3] & 0xF,
-        number_of_questions: ((buffer[4] as u16) << 8) | (buffer[5] as u16),
-        number_of_answers: ((buffer[6] as u16) << 8) | (buffer[7] as u16),
-        number_of_authority_resource_records: ((buffer[8] as u16) << 8) | (buffer[9] as u16),
-        number_of_additional_resource_records: ((buffer[10] as u16) << 8) | (buffer[11] as u16),
-    };
-
-    println!("Received:\n{}", datagram);
+fn handle_datagram(packet: &[u8]) -> Result<()> {
+    let header = DnsHeader::parse(packet)?;
+    println!("Received:\n{}", header);
 
     let mut pos = 12usize;
-    for _ in 0..(datagram.number_of_questions) {
-        let mut domain_parts: Vec<String> = vec![];
-
-        loop {
-            let len = buffer[pos] as usize;
-            if len == 0 {
-                pos += 1;
-                break;
-            }
-
-            let domain_part_bytes = &buffer[(pos + 1)..(pos + 1 + len)];
-            let domain_part = String::from_utf8_lossy(&domain_part_bytes);
-            domain_parts.push(domain_part.to_string());
-            pos += 1 + len;
-        }
-
-        let question_type = ((buffer[pos] as u16) << 8) | (buffer[pos + 1] as u16);
-        let question_class = ((buffer[pos + 2] as u16) << 8) | (buffer[pos + 3] as u16);
-
-        let question = DnsQuestion {
-            domain: domain_parts.join("."),
-            question_type,
-            question_class,
-        };
+    for _ in 0..(header.number_of_questions) {
+        let question = DnsQuestion::parse(packet, &mut pos)?;
         println!("Question: {}", question);
-
-        pos += 4;
     }
+
+    Ok(())
 }