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,90 +117,71 @@ 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 {
-        match current_length {
+        if questions.len() != total_questions as usize {
-            None => { // next question, init lengths
+            match current_length {
-                current_length = Some(byte);
+                None => { // next question, init lengths
-                remaining_length = Box::from(byte);
-                current_query = Some(Vec::with_capacity(10));
-            }
-            Some(_) => {
-                if byte == 0 {
-                    trailers_reached = true;
-                    continue
-                }
-
-                if *remaining_length == 0 && !trailers_reached {
-                    current_query.as_mut().unwrap().push('.' as u8);
                     current_length = Some(byte);
                     remaining_length = Box::from(byte);
+                    current_query = Some(Vec::with_capacity(10));
                 }
-                else if trailers_reached { // trailer fields
+                Some(_) => {
-                    match current_qtype {
+                    if byte == 0 {
-                        None => {
+                        trailers_reached = true;
-                            current_qtype = Some(byte);
+                        continue
-                        }
+                    }
-                        Some(qtype_b) => {
-                            match current_qclass {
-                                None => {
-                                    current_qclass = Some(byte);
-                                }
-                                Some(qclass_b) => {
 
-                                    match (qtype_b.try_into(), qclass_b.try_into()) {
+                    if *remaining_length == 0 && !trailers_reached {
-                                        (Ok(qtype), Ok(qclass)) => {
+                        current_query.as_mut().unwrap().push('.' as u8);
-                                            questions.push(DNSQuestion {
+                        current_length = Some(byte);
-                                                qname: String::from_utf8(current_query.unwrap()).unwrap(),
+                        remaining_length = Box::from(byte);
-                                                qtype,
+                    }
-                                                qclass
+                    else if trailers_reached { // trailer fields
-                                            });
+                        match current_qtype {
+                            None => {
+                                current_qtype = Some(byte);
+                            }
+                            Some(qtype_b) => {
+                                match current_qclass {
+                                    None => {
+                                        // current_qclass = Some(byte);
+                                        match (qtype_b.try_into(), byte.try_into()) {
+                                            (Ok(qtype), Ok(qclass)) => {
+                                                questions.push(DNSQuestion {
+                                                    qname: String::from_utf8(current_query.unwrap()).unwrap(),
+                                                    qtype,
+                                                    qclass
+                                                });
 
-                                            current_length = Some(byte);
+                                                current_length = None;
-                                            remaining_length = Box::from(byte);
+                                                remaining_length = Box::from(byte);
-                                            current_query = Some(Vec::with_capacity(10));
+                                                current_query = None;
-                                            current_qtype = None;
+                                                current_qtype = None;
-                                            current_qclass = None;
+                                                current_qclass = None;
-                                            trailers_reached = false;
+                                                trailers_reached = false;
-                                        }
+                                            }
-                                        _ => {
+                                            _ => {
-                                            return Err(());
+                                                return Err(());
+                                            }
                                         }
                                     }
+                                    Some(_) => {
+
+                                    }
                                 }
                             }
                         }
+                    } else {
+                        current_query.as_mut().unwrap().push(byte);
+                        *remaining_length = remaining_length.sub(1);
                     }
                 }
-                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")