merging request and response

dnstp-client/src/main.rs

@@ -7,7 +7,7 @@ use log::{info, LevelFilter};
 use rand::RngCore;
 use simplelog::*;
 
-use dnstplib::message::DNSRequest;
+use dnstplib::message::DNSMessage;
 use dnstplib::net::{DNSSocket, NetworkMessage};
 use dnstplib::processor::ResponseProcesor;
 
@@ -60,7 +60,7 @@ fn main() {
 
         info!("sending...");
 
-        let message = DNSRequest::from_hostname(address, rng.next_u32() as u16, domain.clone());
+        let message = DNSMessage::from_hostname(address, rng.next_u32() as u16, domain.clone());
 
         let bytes = message.to_bytes();
 

dnstp/src/lib.rs

@@ -1,4 +1,4 @@
-pub mod request_parser;
+pub mod message_parser;
 
 mod byte;
 pub mod processor;

dnstp/src/message/header.rs

@@ -102,29 +102,29 @@ impl DNSHeader {
     {
         let mut header_bytes: [u8; 12] = [0; 12];
 
-        apply_split_bytes(&mut header_bytes, self.id, crate::request_parser::ID_START);
+        apply_split_bytes(&mut header_bytes, self.id, crate::message_parser::ID_START);
 
         let mut flags: u16 = 0;
 
         if self.direction == Response {
-            flags |= 0b1 << crate::request_parser::DIRECTION_SHIFT;
+            flags |= 0b1 << crate::message_parser::DIRECTION_SHIFT;
         }
 
-        flags |= (self.opcode as u16) << crate::request_parser::OPCODE_SHIFT;
+        flags |= (self.opcode as u16) << crate::message_parser::OPCODE_SHIFT;
 
-        flags |= (self.authoritative as u16) << crate::request_parser::AUTHORITATIVE_SHIFT;
-        flags |= (self.truncation as u16) << crate::request_parser::TRUNCATION_SHIFT;
-        flags |= (self.recursion_desired as u16) << crate::request_parser::RECURSION_DESIRED_SHIFT;
-        flags |= (self.recursion_available as u16) << crate::request_parser::RECURSION_AVAILABLE_SHIFT;
+        flags |= (self.authoritative as u16) << crate::message_parser::AUTHORITATIVE_SHIFT;
+        flags |= (self.truncation as u16) << crate::message_parser::TRUNCATION_SHIFT;
+        flags |= (self.recursion_desired as u16) << crate::message_parser::RECURSION_DESIRED_SHIFT;
+        flags |= (self.recursion_available as u16) << crate::message_parser::RECURSION_AVAILABLE_SHIFT;
 
         flags |= self.response as u16;
 
-        apply_split_bytes(&mut header_bytes, flags, crate::request_parser::FLAGS_START);
+        apply_split_bytes(&mut header_bytes, flags, crate::message_parser::FLAGS_START);
 
-        apply_split_bytes(&mut header_bytes, self.question_count, crate::request_parser::QUESTION_COUNT_START);
-        apply_split_bytes(&mut header_bytes, self.answer_record_count, crate::request_parser::ANSWER_RECORD_COUNT_START);
-        apply_split_bytes(&mut header_bytes, self.authority_record_count, crate::request_parser::AUTHORITY_RECORD_COUNT_START);
-        apply_split_bytes(&mut header_bytes, self.additional_record_count, crate::request_parser::ADDITIONAL_RECORD_COUNT_START);
+        apply_split_bytes(&mut header_bytes, self.question_count, crate::message_parser::QUESTION_COUNT_START);
+        apply_split_bytes(&mut header_bytes, self.answer_record_count, crate::message_parser::ANSWER_RECORD_COUNT_START);
+        apply_split_bytes(&mut header_bytes, self.authority_record_count, crate::message_parser::AUTHORITY_RECORD_COUNT_START);
+        apply_split_bytes(&mut header_bytes, self.additional_record_count, crate::message_parser::ADDITIONAL_RECORD_COUNT_START);
 
         header_bytes
     }

dnstp/src/message/message.rs

@@ -0,0 +1,102 @@
+use std::net::{Ipv4Addr, SocketAddr};
+use crate::message::{DNSQuestion, DNSHeader, questions_to_bytes, Direction, Opcode, ResponseCode, QType, QClass, ResourceRecord, records_to_bytes, ARdata};
+
+#[derive(Debug)]
+pub struct DNSMessage {
+    pub header: DNSHeader,
+    pub questions: Vec<DNSQuestion>,
+    pub answer_records: Vec<ResourceRecord>,
+    pub authority_records: Vec<ResourceRecord>,
+    pub additional_records: Vec<ResourceRecord>,
+    pub peer: SocketAddr
+}
+
+impl DNSMessage {
+
+    pub fn to_bytes(& self) -> Vec<u8>
+    {
+        let mut header_bytes = self.header.to_bytes().to_vec();
+        let mut body_bytes = questions_to_bytes(&self.questions);
+        let mut answer_bytes = records_to_bytes(&self.answer_records);
+        let mut authority_bytes = records_to_bytes(&self.authority_records);
+        let mut additional_bytes = records_to_bytes(&self.additional_records);
+
+        header_bytes.append(&mut body_bytes);
+        header_bytes.append(&mut answer_bytes);
+        header_bytes.append(&mut authority_bytes);
+        header_bytes.append(&mut additional_bytes);
+
+        return header_bytes
+    }
+
+    pub fn from_hostname(peer: SocketAddr, id: u16, hostname: String) -> DNSMessage
+    {
+        DNSMessage {
+            header: DNSHeader::new_request(id, None),
+            questions: vec![
+                DNSQuestion {
+                    qname: hostname,
+                    qtype: QType::A,
+                    qclass: QClass::Internet
+                }
+            ],
+            answer_records: vec![],
+            authority_records: vec![],
+            additional_records: vec![],
+            peer
+        }
+    }
+
+    pub fn from_hostnames(peer: SocketAddr, id: u16, hostnames: Vec<String>) -> DNSMessage
+    {
+        DNSMessage {
+            header: DNSHeader::new_request(id, Some(hostnames.len() as u16)),
+            questions: hostnames
+                .into_iter()
+                .map(|n|
+                    DNSQuestion {
+                        qname: n,
+                        qclass: QClass::Internet,
+                        qtype: QType::A
+                    })
+                .collect(),
+            answer_records: vec![],
+            authority_records: vec![],
+            additional_records: vec![],
+            peer
+        }
+    }
+
+    pub fn a_from_request(request: &DNSMessage, ip: impl Fn(&DNSQuestion) -> Ipv4Addr) -> DNSMessage
+    {
+        let mut response = DNSMessage{
+            header: request.header.clone(),
+            questions: request.questions.clone(),
+            answer_records: vec![],
+            authority_records: vec![],
+            additional_records: vec![],
+            peer: request.peer
+        };
+
+        response.answer_records = request.questions
+            .iter()
+            .map(|x|
+                ResourceRecord::from_query(x,
+                                           12,
+                                           Box::from(ARdata::from(ip(x))),
+                                           None))
+            .collect();
+
+        response.header.direction = Direction::Response;
+        response.header.response = ResponseCode::NoError;
+        response.header.answer_record_count = response.answer_records.len() as u16;
+        response.header.authority_record_count = 0;
+        response.header.additional_record_count = 0;
+
+        if response.header.recursion_desired {
+            response.header.recursion_available = true;
+        }
+
+        response
+    }
+}

dnstp/src/message/mod.rs

@@ -1,12 +1,10 @@
 //! Structures making up the DNS workflow including requests, responses and headers
 pub mod header;
 pub mod question;
-pub mod request;
+pub mod message;
 pub mod answer;
-pub mod response;
 
 pub use question::{DNSQuestion, QClass, QType, QuestionParseError, questions_to_bytes, questions_from_bytes};
 pub use answer::{ResourceRecord, RawRData, RData, ARdata, AAAARdata, TXTRdata, RecordParseError, records_to_bytes, records_from_bytes};
 pub use header::{DNSHeader, Direction, Opcode, ResponseCode, HEADER_SIZE};
-pub use request::DNSRequest;
-pub use response::DNSResponse;
+pub use message::DNSMessage;

dnstp/src/request_parser.rs → dnstp/src/message_parser.rs

@@ -1,7 +1,7 @@
 use crate::byte;
-use crate::message::{DNSRequest, Direction, DNSHeader, Opcode, ResponseCode, QuestionParseError, questions_from_bytes, records_from_bytes, RecordParseError, ResourceRecord};
+use crate::message::{DNSMessage, Direction, DNSHeader, Opcode, ResponseCode, QuestionParseError, questions_from_bytes, records_from_bytes, RecordParseError, ResourceRecord};
 use crate::net::NetworkMessage;
-use crate::request_parser::RequestParseError::{HeaderParse, QuesionsParse};
+use crate::message_parser::RequestParseError::{HeaderParse, QuesionsParse};
 
 pub const ID_START: usize = 0;
 pub const FLAGS_START: usize = 2;
@@ -82,7 +82,7 @@ pub enum RequestParseError {
     RecordCount(u16, usize),
 }
 
-pub fn parse_request(msg: NetworkMessage) -> Result<DNSRequest, RequestParseError>
+pub fn parse_message(msg: NetworkMessage) -> Result<DNSMessage, RequestParseError>
 {
     let header = parse_header(msg.buffer[0..12].try_into().unwrap());
 
@@ -108,7 +108,7 @@ pub fn parse_request(msg: NetworkMessage) -> Result<DNSRequest, RequestParseErro
                                     let answer_records = answers.drain(0 .. (header.answer_record_count as usize)).collect();
                                     let authority_records = answers.drain(0 .. (header.authority_record_count as usize)).collect();
 
-                                    return Ok(DNSRequest {
+                                    return Ok(DNSMessage {
                                         header,
                                         questions,
                                         peer: msg.peer,
@@ -122,7 +122,7 @@ pub fn parse_request(msg: NetworkMessage) -> Result<DNSRequest, RequestParseErro
                         }
                     }
                     else {
-                        return Ok(DNSRequest {
+                        return Ok(DNSMessage {
                             header,
                             questions,
                             peer: msg.peer,

dnstp/src/processor/request.rs

@@ -5,9 +5,9 @@ use std::thread;
 use log::{error, info};
 use crate::config::DomainConfig;
 
-use crate::message::{QuestionParseError, DNSResponse, RecordParseError};
+use crate::message::{DNSMessage, QuestionParseError, RecordParseError};
 use crate::net::{NetworkMessage, NetworkMessagePtr};
-use crate::request_parser::{HeaderParseError, parse_request, RequestParseError};
+use crate::message_parser::{HeaderParseError, parse_message, RequestParseError};
 
 pub struct RequestProcesor {
     message_channel: Option<Sender<NetworkMessagePtr>>,
@@ -37,13 +37,13 @@ impl RequestProcesor {
             {
                 let peer = m.peer.clone();
 
-                match parse_request(*m) {
+                match parse_message(*m) {
                     Ok(r) => {
                         info!("received dns message: {:?}", r);
 
                         if r.questions.iter().any(|q| q.qname.ends_with(&base_domain_equality))
                         {
-                            let mut response = DNSResponse::a_from_request(&r, |q| Ipv4Addr::from([127, 0, 0, 1]));
+                            let mut response = DNSMessage::a_from_request(&r, |q| Ipv4Addr::from([127, 0, 0, 1]));
 
                             sending_channel.send(Box::from(
                                 NetworkMessage {
@@ -53,7 +53,7 @@ impl RequestProcesor {
                             ));
                         }
                         else {
-                            let mut response = DNSResponse::a_from_request(&r, |q| Ipv4Addr::from([127, 0, 0, 1]));
+                            let mut response = DNSMessage::a_from_request(&r, |q| Ipv4Addr::from([127, 0, 0, 1]));
 
                             sending_channel.send(Box::from(
                                 NetworkMessage {