From 1f6593263a76e55d46c6ae6bb890d08768310583 Mon Sep 17 00:00:00 2001 From: Andy Pack Date: Mon, 29 Jan 2024 23:19:09 +0000 Subject: [PATCH] tweaking dns question parsing --- dnstp/src/dns_question.rs | 128 ++++++++++++++------------------- dnstp/src/dns_request.rs | 10 +++ dnstp/src/dns_socket.rs | 5 +- dnstp/src/lib.rs | 5 +- dnstp/src/request_parser.rs | 27 +++++++ dnstp/src/request_processor.rs | 23 ++++-- 6 files changed, 114 insertions(+), 84 deletions(-) create mode 100644 dnstp/src/dns_request.rs diff --git a/dnstp/src/dns_question.rs b/dnstp/src/dns_question.rs index def4b54..c8a3aa6 100644 --- a/dnstp/src/dns_question.rs +++ b/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 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 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, total_questions: u8) -> Result, ()> +pub fn questions_from_bytes(bytes: Vec, total_questions: u16) -> Result, ()> { - if (bytes.len() < 4) + if bytes.len() < 4 { return Err(()); } @@ -117,90 +117,71 @@ pub fn questions_from_bytes(bytes: Vec, total_questions: u8) -> Result = None; let mut current_qclass: Option = None; let mut trailers_reached = false; + // let mut finished = false; for byte in bytes { - match current_length { - None => { // next question, init lengths - current_length = Some(byte); - 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); + if questions.len() != total_questions as usize { + match current_length { + None => { // next question, init lengths current_length = Some(byte); remaining_length = Box::from(byte); + current_query = Some(Vec::with_capacity(10)); } - 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); - } - Some(qclass_b) => { + Some(_) => { + if byte == 0 { + trailers_reached = true; + continue + } - match (qtype_b.try_into(), qclass_b.try_into()) { - (Ok(qtype), Ok(qclass)) => { - questions.push(DNSQuestion { - qname: String::from_utf8(current_query.unwrap()).unwrap(), - qtype, - qclass - }); + if *remaining_length == 0 && !trailers_reached { + current_query.as_mut().unwrap().push('.' as u8); + current_length = Some(byte); + remaining_length = Box::from(byte); + } + 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); - remaining_length = Box::from(byte); - current_query = Some(Vec::with_capacity(10)); - current_qtype = None; - current_qclass = None; - trailers_reached = false; - } - _ => { - return Err(()); + current_length = None; + remaining_length = Box::from(byte); + current_query = None; + current_qtype = None; + current_qclass = None; + trailers_reached = false; + } + _ => { + 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); diff --git a/dnstp/src/dns_request.rs b/dnstp/src/dns_request.rs new file mode 100644 index 0000000..8d5fda8 --- /dev/null +++ b/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, + pub peer: SocketAddr +} \ No newline at end of file diff --git a/dnstp/src/dns_socket.rs b/dnstp/src/dns_socket.rs index b8cd9b4..b7ebc54 100644 --- a/dnstp/src/dns_socket.rs +++ b/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(); - info!("received [{}] from [{}]", res_str, peer); + // let res_str = str::from_utf8(&(*buf)).unwrap(); + // info!("received [{}] from [{}]", res_str, peer); if read_count > HEADER_SIZE { message_sender.send(Box::new(NetworkMessage { diff --git a/dnstp/src/lib.rs b/dnstp/src/lib.rs index 4b2fce8..9e0fd87 100644 --- a/dnstp/src/lib.rs +++ b/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; \ No newline at end of file +pub mod dns_question; +pub mod dns_request; \ No newline at end of file diff --git a/dnstp/src/request_parser.rs b/dnstp/src/request_parser.rs index cb744cd..b2fcd19 100644 --- a/dnstp/src/request_parser.rs +++ b/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 +{ + 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::*; diff --git a/dnstp/src/request_processor.rs b/dnstp/src/request_processor.rs index 470c3f9..7c05235 100644 --- a/dnstp/src/request_processor.rs +++ b/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> @@ -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) { - Ok(_) => {} - Err(_) => {} + match request { + Ok(r) => { + info!("received dns message: {:?}", r); + } + Err(_) => { + error!("failed to parse message"); + } } + + // match sending_channel.send(m) { + // Ok(_) => {} + // Err(_) => {} + // } } info!("message processing thread finishing")