From c64691de924beaaa5e4d8a914c8da1d42f2338e5 Mon Sep 17 00:00:00 2001 From: Andy Pack Date: Thu, 8 Feb 2024 22:51:03 +0000 Subject: [PATCH] merging request and response --- dnstp-client/src/main.rs | 4 +- dnstp/src/lib.rs | 2 +- dnstp/src/message/header.rs | 24 ++--- dnstp/src/message/message.rs | 102 ++++++++++++++++++ dnstp/src/message/mod.rs | 6 +- dnstp/src/message/request.rs | 63 ----------- dnstp/src/message/response.rs | 64 ----------- .../{request_parser.rs => message_parser.rs} | 10 +- dnstp/src/processor/request.rs | 10 +- dnstp/src/processor/response.rs | 6 +- 10 files changed, 132 insertions(+), 159 deletions(-) create mode 100644 dnstp/src/message/message.rs delete mode 100644 dnstp/src/message/request.rs delete mode 100644 dnstp/src/message/response.rs rename dnstp/src/{request_parser.rs => message_parser.rs} (96%) diff --git a/dnstp-client/src/main.rs b/dnstp-client/src/main.rs index 635ed91..d4bd9a1 100644 --- a/dnstp-client/src/main.rs +++ b/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(); diff --git a/dnstp/src/lib.rs b/dnstp/src/lib.rs index 95908aa..c1e3eb9 100644 --- a/dnstp/src/lib.rs +++ b/dnstp/src/lib.rs @@ -1,4 +1,4 @@ -pub mod request_parser; +pub mod message_parser; mod byte; pub mod processor; diff --git a/dnstp/src/message/header.rs b/dnstp/src/message/header.rs index 43babb8..5cc436a 100644 --- a/dnstp/src/message/header.rs +++ b/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 } diff --git a/dnstp/src/message/message.rs b/dnstp/src/message/message.rs new file mode 100644 index 0000000..663b768 --- /dev/null +++ b/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, + pub answer_records: Vec, + pub authority_records: Vec, + pub additional_records: Vec, + pub peer: SocketAddr +} + +impl DNSMessage { + + pub fn to_bytes(& self) -> Vec + { + 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) -> 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 + } +} \ No newline at end of file diff --git a/dnstp/src/message/mod.rs b/dnstp/src/message/mod.rs index 20b5af7..c3cb4cf 100644 --- a/dnstp/src/message/mod.rs +++ b/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; \ No newline at end of file +pub use message::DNSMessage; \ No newline at end of file diff --git a/dnstp/src/message/request.rs b/dnstp/src/message/request.rs deleted file mode 100644 index 259d4c2..0000000 --- a/dnstp/src/message/request.rs +++ /dev/null @@ -1,63 +0,0 @@ -use std::net::SocketAddr; -use crate::message::{DNSQuestion, DNSHeader, questions_to_bytes, Direction, Opcode, ResponseCode, QType, QClass, ResourceRecord}; - -#[derive(Debug)] -pub struct DNSRequest { - pub header: DNSHeader, - pub questions: Vec, - pub answer_records: Vec, - pub authority_records: Vec, - pub additional_records: Vec, - pub peer: SocketAddr -} - -impl DNSRequest { - - pub fn to_bytes(& self) -> Vec - { - let mut header_bytes = self.header.to_bytes().to_vec(); - let mut body_bytes = questions_to_bytes(&self.questions); - - header_bytes.append(&mut body_bytes); - - return header_bytes - } - - pub fn from_hostname(peer: SocketAddr, id: u16, hostname: String) -> DNSRequest - { - DNSRequest { - 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) -> DNSRequest - { - DNSRequest { - 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 - } - } -} \ No newline at end of file diff --git a/dnstp/src/message/response.rs b/dnstp/src/message/response.rs deleted file mode 100644 index c9e47b7..0000000 --- a/dnstp/src/message/response.rs +++ /dev/null @@ -1,64 +0,0 @@ -use std::net::{Ipv4Addr, SocketAddr}; -use crate::message::{Direction, DNSHeader, DNSRequest, ResponseCode, records_to_bytes, ARdata, ResourceRecord, DNSQuestion, questions_to_bytes}; - -#[derive(Debug)] -pub struct DNSResponse { - pub header: DNSHeader, - pub questions: Vec, - pub answers: Vec, - pub authorities: Vec, - pub additionals: Vec, - pub peer: SocketAddr -} - -impl DNSResponse { - - pub fn to_bytes(& self) -> Vec - { - 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.answers); - let mut authority_bytes = records_to_bytes(&self.authorities); - let mut additional_bytes = records_to_bytes(&self.additionals); - - 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 a_from_request(request: &DNSRequest, ip: impl Fn(&DNSQuestion) -> Ipv4Addr) -> DNSResponse - { - let mut response = DNSResponse{ - header: request.header.clone(), - questions: request.questions.clone(), - answers: vec![], - authorities: vec![], - additionals: vec![], - peer: request.peer - }; - - response.answers = 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.answers.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 - } -} \ No newline at end of file diff --git a/dnstp/src/request_parser.rs b/dnstp/src/message_parser.rs similarity index 96% rename from dnstp/src/request_parser.rs rename to dnstp/src/message_parser.rs index f2126c9..956210f 100644 --- a/dnstp/src/request_parser.rs +++ b/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 +pub fn parse_message(msg: NetworkMessage) -> Result { let header = parse_header(msg.buffer[0..12].try_into().unwrap()); @@ -108,7 +108,7 @@ pub fn parse_request(msg: NetworkMessage) -> Result Result>, @@ -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 { diff --git a/dnstp/src/processor/response.rs b/dnstp/src/processor/response.rs index ee6ad2c..56c29c6 100644 --- a/dnstp/src/processor/response.rs +++ b/dnstp/src/processor/response.rs @@ -4,10 +4,10 @@ use std::sync::mpsc::{Receiver, Sender}; use std::thread; use log::{error, info}; use std::str; -use crate::message::{DNSResponse, QuestionParseError, RecordParseError}; +use crate::message::{QuestionParseError, RecordParseError}; use crate::net::NetworkMessage; use crate::net::raw_request::NetworkMessagePtr; -use crate::request_parser::{HeaderParseError, parse_request, RequestParseError}; +use crate::message_parser::{HeaderParseError, parse_message, RequestParseError}; pub struct ResponseProcesor { message_channel: Option> @@ -31,7 +31,7 @@ impl ResponseProcesor { { let peer = m.peer.clone(); - match parse_request(*m) { + match parse_message(*m) { Ok(r) => { info!("received dns message: {:?}", r); }