diff --git a/src/dns/answer.rs b/src/dns/answer.rs index 2b67d31..30691ae 100644 --- a/src/dns/answer.rs +++ b/src/dns/answer.rs @@ -1,6 +1,11 @@ +use std::io::Cursor; + use anyhow::{Result, bail}; -use crate::dns::{domain_to_bytes, parse_domain}; +use crate::{ + dns::{domain_to_bytes, parse_domain}, + ext::CursorExt, +}; #[derive(Debug, Clone)] pub struct Answer { @@ -12,30 +17,15 @@ pub struct Answer { } impl Answer { - pub fn parse(datagram: &[u8], pos: &mut usize) -> Result { - let domain = parse_domain(datagram, pos)?; + pub fn parse(datagram: &mut Cursor<&[u8]>) -> Result { + let domain = parse_domain(datagram)?; - if *pos + 10 > datagram.len() { - bail!("malformed datagram"); - } + let record_type = datagram.read_u16()?; + let class = datagram.read_u16()?; + let ttl = datagram.read_u32()?; + let data_length = datagram.read_u16()? as usize; - let record_type = u16::from_be_bytes([datagram[*pos], datagram[*pos + 1]]); - let class = u16::from_be_bytes([datagram[*pos + 2], datagram[*pos + 3]]); - let ttl = u32::from_be_bytes([ - datagram[*pos + 4], - datagram[*pos + 5], - datagram[*pos + 6], - datagram[*pos + 7], - ]); - let data_length = u16::from_be_bytes([datagram[*pos + 8], datagram[*pos + 9]]) as usize; - *pos += 10; - - if *pos + data_length > datagram.len() { - bail!("malformed datagram"); - } - - let data = datagram[*pos..*pos + data_length].to_vec(); - *pos += data_length; + let data = datagram.read_bytes(data_length)?; Ok(Answer { domain, diff --git a/src/dns/domainname.rs b/src/dns/domainname.rs index 652fd10..d89c633 100644 --- a/src/dns/domainname.rs +++ b/src/dns/domainname.rs @@ -1,60 +1,51 @@ +use std::io::Cursor; + use anyhow::{Result, bail}; -pub fn parse_domain(datagram: &[u8], pos: &mut usize) -> Result { +use crate::ext::CursorExt; + +pub fn parse_domain(cursor: &mut Cursor<&[u8]>) -> Result { + parse_domain_inner(cursor, 0) +} + +fn parse_domain_inner(cursor: &mut Cursor<&[u8]>, jumps: usize) -> Result { let mut domain_parts: Vec = vec![]; - let mut cursor = *pos; - let mut jumped = false; - let mut jumps = 0usize; + let datagram_len = cursor.get_ref().len(); loop { - if cursor >= datagram.len() { - bail!("malformed datagram"); - } - - let len = datagram[cursor] as usize; + let len = cursor.read_u8()? as usize; if (len & 0xC0) == 0xC0 { - if cursor + 1 >= datagram.len() { + let pointer = ((len & 0x3F) << 8) | cursor.read_u8()? as usize; + if pointer >= datagram_len { + bail!("malformed datagram"); + } + if jumps > datagram_len { bail!("malformed datagram"); } - let pointer = ((len & 0x3F) << 8) | datagram[cursor + 1] as usize; - if pointer >= datagram.len() { - bail!("malformed datagram"); - } - if !jumped { - *pos = cursor + 2; - jumped = true; - } + let mut pointed_cursor = cursor.clone(); + pointed_cursor.set_position(pointer as u64); - jumps += 1; - if jumps > datagram.len() { - bail!("malformed datagram"); + let pointed_domain = parse_domain_inner(&mut pointed_cursor, jumps + 1)?; + if !pointed_domain.is_empty() { + domain_parts.push(pointed_domain); } - - cursor = pointer; - continue; + break; } if (len & 0xC0) != 0 { + // True when 0b01.. and 0b10... bail!("malformed datagram"); } if len > 63 { bail!("malformed datagram"); } if len == 0 { - if !jumped { - *pos = cursor + 1; - } break; } - if (cursor + 1 + len) > datagram.len() { - bail!("malformed datagram"); - } - - let domain_part_bytes = &datagram[(cursor + 1)..(cursor + 1 + len)]; + let domain_part_bytes = cursor.read_bytes(len)?; let domain_part = String::from_utf8_lossy(&domain_part_bytes); domain_parts.push(domain_part.to_string()); - cursor += 1 + len; } Ok(domain_parts.join(".")) diff --git a/src/dns/header.rs b/src/dns/header.rs index a693be1..bf0c0ca 100644 --- a/src/dns/header.rs +++ b/src/dns/header.rs @@ -1,8 +1,12 @@ +use std::io::Cursor; + use anyhow::{Result, bail}; +use crate::ext::CursorExt; + #[derive(Debug, Clone)] pub struct Header { - pub transaction_id: [u8; 2], + pub transaction_id: u16, pub query_response: bool, pub op_code: u8, // 4 bits pub authoritative_answer: bool, @@ -19,26 +23,30 @@ pub struct Header { } impl Header { - pub fn parse(datagram: &[u8]) -> Result
{ - if datagram.len() < 12 { - bail!("malformed datagram"); - } + pub fn parse(cursor: &mut Cursor<&[u8]>) -> Result
{ + let transaction_id = cursor.read_u16()?; + let flags1 = cursor.read_u8()?; + let flags2 = cursor.read_u8()?; + let number_of_questions = cursor.read_u16()?; + let number_of_answers = cursor.read_u16()?; + let number_of_authority_resource_records = cursor.read_u16()?; + let number_of_additional_resource_records = cursor.read_u16()?; let datagram = Header { - transaction_id: [datagram[0], datagram[1]], - query_response: datagram[2] & (1 << 7) > 0, - op_code: (datagram[2] >> 3) & 0x0F, - authoritative_answer: datagram[2] & (1 << 2) > 0, - truncation: datagram[2] & (1 << 1) > 0, - recursion_desired: datagram[2] & (1 << 0) > 0, - recursion_available: datagram[3] & (1 << 7) > 0, - authentic_data: datagram[3] & (1 << 5) > 0, - checking_disabled: datagram[3] & (1 << 4) > 0, - response_code: datagram[3] & 0xF, - number_of_questions: u16::from_be_bytes([datagram[4], datagram[5]]), - number_of_answers: u16::from_be_bytes([datagram[6], datagram[7]]), - number_of_authority_resource_records: u16::from_be_bytes([datagram[8], datagram[9]]), - number_of_additional_resource_records: u16::from_be_bytes([datagram[10], datagram[11]]), + transaction_id, + query_response: flags1 & (1 << 7) > 0, + op_code: (flags1 >> 3) & 0x0F, + authoritative_answer: flags1 & (1 << 2) > 0, + truncation: flags1 & (1 << 1) > 0, + recursion_desired: flags1 & (1 << 0) > 0, + recursion_available: flags2 & (1 << 7) > 0, + authentic_data: flags2 & (1 << 5) > 0, + checking_disabled: flags2 & (1 << 4) > 0, + response_code: flags2 & 0xF, + number_of_questions, + number_of_answers, + number_of_authority_resource_records, + number_of_additional_resource_records, }; Ok(datagram) @@ -53,8 +61,8 @@ impl Header { let number_of_additional_resource_records = self.number_of_additional_resource_records.to_be_bytes(); - buffer[0] = self.transaction_id[0]; - buffer[1] = self.transaction_id[1]; + buffer[0] = (self.transaction_id >> 8) as u8; + buffer[1] = (self.transaction_id & 0x0F) as u8; if self.query_response { buffer[2] |= 1 << 7; } diff --git a/src/dns/message.rs b/src/dns/message.rs index 925dbe2..9de8c35 100644 --- a/src/dns/message.rs +++ b/src/dns/message.rs @@ -1,3 +1,5 @@ +use std::io::Cursor; + use anyhow::Result; use crate::dns::{Answer, Header, Question, RESPONSE_CODE_NO_ERROR}; @@ -10,7 +12,7 @@ pub struct Message { } impl Message { - pub fn parse(datagram: &[u8]) -> Result { + pub fn parse(datagram: &mut Cursor<&[u8]>) -> Result { let header = Header::parse(datagram)?; let mut message = Message { @@ -19,13 +21,12 @@ impl Message { answers: vec![], }; - let mut pos = 12usize; for _ in 0..(message.header.number_of_questions) { - let question = Question::parse(datagram, &mut pos)?; + let question = Question::parse(datagram)?; message.questions.push(question); } for _ in 0..(message.header.number_of_answers) { - let answer = Answer::parse(datagram, &mut pos)?; + let answer = Answer::parse(datagram)?; message.answers.push(answer); } diff --git a/src/dns/mod.rs b/src/dns/mod.rs index 717d9b4..01fb810 100644 --- a/src/dns/mod.rs +++ b/src/dns/mod.rs @@ -14,6 +14,8 @@ pub use question::*; #[cfg(test)] mod tests { + use std::io::Cursor; + use super::*; fn a_query_datagram() -> Vec { @@ -28,7 +30,9 @@ mod tests { #[test] fn response_keeps_original_question_section() { - let query = Message::parse(&a_query_datagram()).unwrap(); + let datagram: &[u8] = &a_query_datagram(); + let mut cursor = Cursor::new(datagram); + let query = Message::parse(&mut cursor).unwrap(); let mut response = query.prepare_response(); response.answers.push(Answer { domain: "example.com".to_string(), @@ -56,7 +60,8 @@ mod tests { 0xC0, 0x0C, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x04, 127, 0, 0, 1, ]); - let message = Message::parse(&datagram).unwrap(); + let mut cursor = Cursor::new(datagram.as_slice()); + let message = Message::parse(&mut cursor).unwrap(); assert_eq!(message.answers.len(), 1); assert_eq!(message.answers[0].domain, "example.com"); assert_eq!(message.answers[0].data, vec![127, 0, 0, 1]); diff --git a/src/dns/question.rs b/src/dns/question.rs index 37eb353..6b32570 100644 --- a/src/dns/question.rs +++ b/src/dns/question.rs @@ -1,6 +1,11 @@ +use std::io::Cursor; + use anyhow::{Result, bail}; -use crate::dns::{domain_to_bytes, parse_domain}; +use crate::{ + dns::{domain_to_bytes, parse_domain}, + ext::CursorExt, +}; #[derive(Debug, Clone)] pub struct Question { @@ -10,16 +15,11 @@ pub struct Question { } impl Question { - pub fn parse(datagram: &[u8], pos: &mut usize) -> Result { - let domain = parse_domain(datagram, pos)?; + pub fn parse(cursor: &mut Cursor<&[u8]>) -> Result { + let domain = parse_domain(cursor)?; - if *pos + 4 > datagram.len() { - bail!("malformed datagram"); - } - - let question_type = ((datagram[*pos] as u16) << 8) | (datagram[*pos + 1] as u16); - let question_class = ((datagram[*pos + 2] as u16) << 8) | (datagram[*pos + 3] as u16); - *pos += 4; + let question_type = cursor.read_u16()?; + let question_class = cursor.read_u16()?; let question = Question { domain, diff --git a/src/ext.rs b/src/ext.rs new file mode 100644 index 0000000..c69387b --- /dev/null +++ b/src/ext.rs @@ -0,0 +1,40 @@ +use std::io::{Cursor, Read}; + +pub trait CursorExt { + fn read_u8(&mut self) -> anyhow::Result; + fn read_u16(&mut self) -> anyhow::Result; + fn read_u32(&mut self) -> anyhow::Result; + fn read_bytes(&mut self, len: usize) -> anyhow::Result>; + fn read_string(&mut self, len: usize) -> anyhow::Result; +} + +impl CursorExt for Cursor<&[u8]> { + fn read_u8(&mut self) -> anyhow::Result { + let mut result = [0u8; 1]; + self.read_exact(&mut result)?; + Ok(result[0]) + } + + fn read_u16(&mut self) -> anyhow::Result { + let mut result = [0u8; 2]; + self.read_exact(&mut result)?; + Ok(u16::from_be_bytes(result)) + } + + fn read_u32(&mut self) -> anyhow::Result { + let mut result = [0u8; 4]; + self.read_exact(&mut result)?; + Ok(u32::from_be_bytes(result)) + } + + fn read_bytes(&mut self, len: usize) -> anyhow::Result> { + let mut result = vec![0u8; len]; + self.read_exact(&mut result)?; + Ok(result) + } + + fn read_string(&mut self, len: usize) -> anyhow::Result { + let result = self.read_bytes(len)?; + Ok(String::from_utf8_lossy(&result).to_string()) + } +} diff --git a/src/macros.rs b/src/macros.rs index 96a8bb7..dfae1f7 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -1,6 +1,6 @@ #[macro_export] macro_rules! txid { - ($bytes:expr) => { - format!("{:02x}{:02x}", $bytes[0], $bytes[1]) + ($txid:expr) => { + format!("{:04x}", $txid) }; } diff --git a/src/main.rs b/src/main.rs index 33f7ee6..846cdc0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,9 +1,10 @@ mod config; mod dns; +mod ext; mod logging; mod macros; -use std::net::SocketAddr; +use std::{io::Cursor, net::SocketAddr}; use anyhow::Result; use dotenv::dotenv; @@ -24,7 +25,8 @@ async fn main() -> Result<()> { loop { let (size, addr) = sock.recv_from(&mut buffer).await?; - let result = handle_datagram(&config, &sock, &addr, &buffer[..size]).await; + let mut cursor = Cursor::new(&buffer[..size]); + let result = handle_datagram(&config, &sock, &addr, &mut cursor).await; if let Err(err) = result { tracing::error!("failed to handle datagram: {err}"); } @@ -35,9 +37,9 @@ async fn handle_datagram( config: &config::ConfigData, sock: &UdpSocket, addr: &SocketAddr, - datagram: &[u8], + cursor: &mut Cursor<&[u8]>, ) -> Result<()> { - let message = Message::parse(datagram)?; + let message = Message::parse(cursor)?; let span = tracing::span!( Level::INFO, "handle_datagram",