tweaking dns question parsing

dnstp/src/dns_question.rs

@@ -2,7 +2,7 @@ use std::ops::Sub;
 use urlencoding::{encode, decode};
 
 #[derive(Ord, PartialOrd, Eq, PartialEq, Debug, Copy, Clone)]
-enum QType {
+pub enum QType {
     A = 1,
     NS = 2,
     CNAME = 5,
@@ -42,7 +42,7 @@ impl TryFrom<u8> for QType {
 }
 
 #[derive(Ord, PartialOrd, Eq, PartialEq, Debug, Copy, Clone)]
-enum QClass {
+pub enum QClass {
     Internet = 1,
     Chaos = 3,
     Hesiod = 4,
@@ -62,7 +62,7 @@ impl TryFrom<u8> for QClass {
 }
 
 #[derive(Ord, PartialOrd, Eq, PartialEq, Debug)]
-struct DNSQuestion {
+pub struct DNSQuestion {
     qname: String,
     qtype: QType,
     qclass: QClass
@@ -102,9 +102,9 @@ impl DNSQuestion {
     }
 }
 
-pub fn questions_from_bytes(bytes: Vec<u8>, total_questions: u8) -> Result<Vec<DNSQuestion>, ()>
+pub fn questions_from_bytes(bytes: Vec<u8>, total_questions: u16) -> Result<Vec<DNSQuestion>, ()>
 {
-    if (bytes.len() < 4)
+    if bytes.len() < 4
     {
         return Err(());
     }
@@ -117,8 +117,10 @@ pub fn questions_from_bytes(bytes: Vec<u8>, total_questions: u8) -> Result<Vec<D
     let mut current_qtype: Option<u8> = None;
     let mut current_qclass: Option<u8> = None;
     let mut trailers_reached = false;
+    // let mut finished = false;
 
     for byte in bytes {
+        if questions.len() != total_questions as usize {
             match current_length {
                 None => { // next question, init lengths
                     current_length = Some(byte);
@@ -144,11 +146,8 @@ pub fn questions_from_bytes(bytes: Vec<u8>, total_questions: u8) -> Result<Vec<D
                             Some(qtype_b) => {
                                 match current_qclass {
                                     None => {
-                                    current_qclass = Some(byte);
+                                        // current_qclass = Some(byte);
-                                }
+                                        match (qtype_b.try_into(), byte.try_into()) {
-                                Some(qclass_b) => {
-
-                                    match (qtype_b.try_into(), qclass_b.try_into()) {
                                             (Ok(qtype), Ok(qclass)) => {
                                                 questions.push(DNSQuestion {
                                                     qname: String::from_utf8(current_query.unwrap()).unwrap(),
@@ -156,9 +155,9 @@ pub fn questions_from_bytes(bytes: Vec<u8>, total_questions: u8) -> Result<Vec<D
                                                     qclass
                                                 });
 
-                                            current_length = Some(byte);
+                                                current_length = None;
                                                 remaining_length = Box::from(byte);
-                                            current_query = Some(Vec::with_capacity(10));
+                                                current_query = None;
                                                 current_qtype = None;
                                                 current_qclass = None;
                                                 trailers_reached = false;
@@ -168,37 +167,19 @@ pub fn questions_from_bytes(bytes: Vec<u8>, total_questions: u8) -> Result<Vec<D
                                             }
                                         }
                                     }
+                                    Some(_) => {
+
                                     }
                                 }
                             }
                         }
-                else
+                    } else {
-                {
                         current_query.as_mut().unwrap().push(byte);
                         *remaining_length = remaining_length.sub(1);
                     }
                 }
             }
         }
-
-    match (current_qtype, current_qclass) {
-        (Some(qtype), Some(qclass)) => {
-            match (qtype.try_into(), qclass.try_into()) {
-                (Ok(qtype), Ok(qclass)) => {
-                    questions.push(DNSQuestion {
-                        qname: String::from_utf8(current_query.unwrap()).unwrap(),
-                        qtype,
-                        qclass
-                    });
-                }
-                _ => {
-                    return Err(());
-                }
-            }
-        }
-        _ => {
-            return Err(());
-        }
     }
 
     Ok(questions)
@@ -216,7 +197,8 @@ mod tests {
             qtype: QType::A
         };
 
-        let q_bytes = q.to_bytes();
+        let mut q_bytes = q.to_bytes();
+        q_bytes.append(&mut vec![0, 0, 0, 0, 0, 0]);
 
         let q_reconstructed = questions_from_bytes(q_bytes, 1).unwrap();
 
@@ -282,7 +264,7 @@ mod tests {
         q_bytes.append(&mut q2_bytes);
         q_bytes.append(&mut q3_bytes);
 
-        let q_reconstructed = questions_from_bytes(q_bytes, 2).unwrap();
+        let q_reconstructed = questions_from_bytes(q_bytes, 3).unwrap();
 
         assert_eq!(q.qname, q_reconstructed[0].qname);
         assert_eq!(q.qclass, q_reconstructed[0].qclass);

dnstp/src/dns_request.rs

@@ -0,0 +1,10 @@
+use std::net::SocketAddr;
+use crate::dns_header::DNSHeader;
+use crate::dns_question::DNSQuestion;
+
+#[derive(Ord, PartialOrd, Eq, PartialEq, Debug)]
+pub struct DNSRequest {
+    pub header: DNSHeader,
+    pub questions: Vec<DNSQuestion>,
+    pub peer: SocketAddr
+}

dnstp/src/dns_socket.rs

@@ -1,5 +1,4 @@
 use std::net::{SocketAddr, UdpSocket};
-use std::ptr::read;
 use std::thread;
 use std::thread::{JoinHandle};
 use log::{debug, error, info};
@@ -83,8 +82,8 @@ impl DNSSocket {
 
                         match res {
                             Ok((read_count, peer)) => {
-                                let res_str = str::from_utf8(&(*buf)).unwrap();
+                                // let res_str = str::from_utf8(&(*buf)).unwrap();
-                                info!("received [{}] from [{}]", res_str, peer);
+                                // info!("received [{}] from [{}]", res_str, peer);
 
                                 if read_count > HEADER_SIZE {
                                     message_sender.send(Box::new(NetworkMessage {

dnstp/src/lib.rs

@@ -1,7 +1,8 @@
 pub mod dns_socket;
 pub mod request_parser;
-mod dns_header;
+pub mod dns_header;
 pub mod request_processor;
 pub mod response_processor;
 pub mod raw_request;
-mod dns_question;
+pub mod dns_question;
+pub mod dns_request;

dnstp/src/request_parser.rs

@@ -1,5 +1,8 @@
 use crate::dns_header::{Direction, DNSHeader, Opcode, ResponseCode};
 use crate::dns_header::Direction::Response;
+use crate::dns_question::{DNSQuestion, questions_from_bytes};
+use crate::dns_request::DNSRequest;
+use crate::raw_request::NetworkMessage;
 
 fn two_byte_extraction(buffer: &[u8], idx: usize) -> u16
 {
@@ -114,6 +117,30 @@ pub fn parse_header_to_bytes(header: &DNSHeader) -> [u8; 12]
     header_bytes
 }
 
+pub fn parse_request(msg: NetworkMessage) -> Result<DNSRequest, ()>
+{
+    let header = parse_header(msg.buffer[0..12].try_into().unwrap());
+
+    match header {
+        Ok(header) => {
+            let mut trimmed = msg.buffer.to_vec();
+            trimmed.drain(0 .. 12);
+            match questions_from_bytes(trimmed, header.question_count)
+            {
+                Ok(questions) => {
+                    Ok(DNSRequest {
+                        header,
+                        questions,
+                        peer: msg.peer
+                    })
+                }
+                Err(_) => Err(())
+            }
+        },
+        Err(_) => Err(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

dnstp/src/request_processor.rs

@@ -1,9 +1,11 @@
 use std::sync::mpsc;
 use std::sync::mpsc::{Receiver, Sender};
 use std::thread;
-use log::info;
+use log::{error, info};
 use std::str;
+use crate::dns_request::DNSRequest;
 use crate::raw_request::NetworkMessagePtr;
+use crate::request_parser::parse_request;
 
 pub struct RequestProcesor {
     message_channel: Option<Sender<NetworkMessagePtr>>
@@ -25,14 +27,23 @@ impl RequestProcesor {
 
             for mut m in rx
             {
-                info!("processing: {}", str::from_utf8(&(*(*m).buffer)).unwrap());
+                // info!("processing: {}", str::from_utf8(&(*(*m).buffer)).unwrap());
 
-                (*(*m).buffer).reverse();
+                let request = parse_request(*m);
 
-                match sending_channel.send(m) {
+                match request {
-                    Ok(_) => {}
+                    Ok(r) => {
-                    Err(_) => {}
+                        info!("received dns message: {:?}", r);
                     }
+                    Err(_) => {
+                        error!("failed to parse message");
+                    }
+                }
+
+                // match sending_channel.send(m) {
+                //     Ok(_) => {}
+                //     Err(_) => {}
+                // }
             }
 
             info!("message processing thread finishing")