feat: working on multiple request types

src/config.rs

@@ -1,8 +1,17 @@
 use anyhow::{Result, ensure};
 use serde::Deserialize;
 
+#[derive(Deserialize)]
+pub enum EntryType {
+    A,
+    AAAA,
+    MX,
+}
+
 #[derive(Deserialize)]
 pub struct Entry {
+    #[serde(rename = "type")]
+    pub entry_type: EntryType,
     pub name: String,
     pub address: String,
     #[serde(skip)]
@@ -26,7 +35,7 @@ pub struct ConfigData {
 }
 
 impl ConfigData {
-    pub fn get_address_by_domain(&self, domain: &str) -> Option<[u8; 4]> {
+    pub fn get_ipv4_address_by_domain(&self, domain: &str) -> Option<[u8; 4]> {
         self.entries
             .iter()
             .find(|e| e.name.as_str() == domain)

src/dns/answer.rs

@@ -1,6 +1,6 @@
 use std::io::Cursor;
 
-use anyhow::{Result, bail};
+use anyhow::Result;
 
 use crate::{
     dns::{domain_to_bytes, parse_domain},

src/dns/domainname.rs

@@ -43,9 +43,8 @@ fn parse_domain_inner(cursor: &mut Cursor<&[u8]>, jumps: usize) -> Result<String
             break;
         }
 
-        let domain_part_bytes = cursor.read_bytes(len)?;
-        let domain_part = String::from_utf8_lossy(&domain_part_bytes);
-        domain_parts.push(domain_part.to_string());
+        let domain_part = cursor.read_string(len)?;
+        domain_parts.push(domain_part);
     }
 
     Ok(domain_parts.join("."))

src/dns/header.rs

@@ -1,6 +1,6 @@
 use std::io::Cursor;
 
-use anyhow::{Result, bail};
+use anyhow::Result;
 
 use crate::ext::CursorExt;
 
@@ -54,6 +54,7 @@ impl Header {
 
     pub fn to_bytes(&self) -> Result<Vec<u8>> {
         let mut buffer = [0u8; 12];
+        let transaction_id = self.transaction_id.to_be_bytes();
         let number_of_questions = self.number_of_questions.to_be_bytes();
         let number_of_answers = self.number_of_answers.to_be_bytes();
         let number_of_authority_resource_records =
@@ -61,8 +62,8 @@ impl Header {
         let number_of_additional_resource_records =
             self.number_of_additional_resource_records.to_be_bytes();
 
-        buffer[0] = (self.transaction_id >> 8) as u8;
-        buffer[1] = (self.transaction_id & 0x0F) as u8;
+        buffer[0] = transaction_id[0];
+        buffer[1] = transaction_id[1];
         if self.query_response {
             buffer[2] |= 1 << 7;
         }

src/dns/question.rs

@@ -1,6 +1,6 @@
 use std::io::Cursor;
 
-use anyhow::{Result, bail};
+use anyhow::Result;
 
 use crate::{
     dns::{domain_to_bytes, parse_domain},

src/main.rs

@@ -11,7 +11,7 @@ use dotenv::dotenv;
 use tokio::net::UdpSocket;
 use tracing::{Instrument, Level};
 
-use crate::dns::*;
+use crate::{config::ConfigData, dns::*};
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -61,13 +61,6 @@ async fn handle_datagram(
 
         let mut response = message.prepare_response();
         for question in message.questions {
-            if question.question_type != RECORD_TYPE_A {
-                tracing::error!(
-                    "only A records are supported yet, requested: {}",
-                    question.question_type
-                );
-                return Ok(());
-            }
             if question.question_class != CLASS_IN {
                 tracing::error!(
                     "only IN class records are supported yet, requested: {}",
@@ -76,22 +69,13 @@ async fn handle_datagram(
                 return Ok(());
             }
 
-            let address = config.get_address_by_domain(&question.domain);
-            match address {
-                Some(address) => {
-                    let answer = Answer {
-                        domain: question.domain,
-                        record_type: RECORD_TYPE_A,
-                        class: CLASS_IN,
-                        ttl: 1,
-                        data: address.to_vec(),
-                    };
-
-                    response.answers.push(answer);
-                    response.header.number_of_answers += 1;
+            match question.question_type {
+                RECORD_TYPE_A => {
+                    handle_type_a_question(&config, &question, &mut response);
                 }
-                None => {
-                    response.header.response_code = RESPONSE_CODE_NXDOMAIN;
+                qt => {
+                    tracing::error!("unsupported question type, requested: {qt}");
+                    return Ok(());
                 }
             }
         }
@@ -106,3 +90,24 @@ async fn handle_datagram(
 
     result
 }
+
+fn handle_type_a_question(config: &ConfigData, question: &Question, response: &mut Message) {
+    let address = config.get_ipv4_address_by_domain(&question.domain);
+    match address {
+        Some(address) => {
+            let answer = Answer {
+                domain: question.domain.clone(),
+                record_type: RECORD_TYPE_A,
+                class: CLASS_IN,
+                ttl: 1,
+                data: address.to_vec(),
+            };
+
+            response.answers.push(answer);
+            response.header.number_of_answers += 1;
+        }
+        None => {
+            response.header.response_code = RESPONSE_CODE_NXDOMAIN;
+        }
+    }
+}