feat: refactoring

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(())
 }