feat: use cursors instead of byte arrays

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<Answer> {
+    pub fn parse(datagram: &mut Cursor<&[u8]>) -> Result<Answer> {
-        let domain = parse_domain(datagram, pos)?;
+        let domain = parse_domain(datagram)?;
 
-        if *pos + 10 > datagram.len() {
+        let record_type = datagram.read_u16()?;
-            bail!("malformed datagram");
+        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 data = datagram.read_bytes(data_length)?;
-        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;
 
         Ok(Answer {
             domain,

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<String> {
+use crate::ext::CursorExt;
+
+pub fn parse_domain(cursor: &mut Cursor<&[u8]>) -> Result<String> {
+    parse_domain_inner(cursor, 0)
+}
+
+fn parse_domain_inner(cursor: &mut Cursor<&[u8]>, jumps: usize) -> Result<String> {
     let mut domain_parts: Vec<String> = vec![];
-    let mut cursor = *pos;
+    let datagram_len = cursor.get_ref().len();
-    let mut jumped = false;
-    let mut jumps = 0usize;
 
     loop {
-        if cursor >= datagram.len() {
+        let len = cursor.read_u8()? as usize;
-            bail!("malformed datagram");
-        }
-
-        let len = datagram[cursor] 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;
+            let mut pointed_cursor = cursor.clone();
-            if pointer >= datagram.len() {
+            pointed_cursor.set_position(pointer as u64);
-                bail!("malformed datagram");
-            }
-            if !jumped {
-                *pos = cursor + 2;
-                jumped = true;
-            }
 
-            jumps += 1;
+            let pointed_domain = parse_domain_inner(&mut pointed_cursor, jumps + 1)?;
-            if jumps > datagram.len() {
+            if !pointed_domain.is_empty() {
-                bail!("malformed datagram");
+                domain_parts.push(pointed_domain);
             }
-
+            break;
-            cursor = pointer;
-            continue;
         }
         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() {
+        let domain_part_bytes = cursor.read_bytes(len)?;
-            bail!("malformed datagram");
-        }
-
-        let domain_part_bytes = &datagram[(cursor + 1)..(cursor + 1 + 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("."))

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<Header> {
+    pub fn parse(cursor: &mut Cursor<&[u8]>) -> Result<Header> {
-        if datagram.len() < 12 {
+        let transaction_id = cursor.read_u16()?;
-            bail!("malformed datagram");
+        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]],
+            transaction_id,
-            query_response: datagram[2] & (1 << 7) > 0,
+            query_response: flags1 & (1 << 7) > 0,
-            op_code: (datagram[2] >> 3) & 0x0F,
+            op_code: (flags1 >> 3) & 0x0F,
-            authoritative_answer: datagram[2] & (1 << 2) > 0,
+            authoritative_answer: flags1 & (1 << 2) > 0,
-            truncation: datagram[2] & (1 << 1) > 0,
+            truncation: flags1 & (1 << 1) > 0,
-            recursion_desired: datagram[2] & (1 << 0) > 0,
+            recursion_desired: flags1 & (1 << 0) > 0,
-            recursion_available: datagram[3] & (1 << 7) > 0,
+            recursion_available: flags2 & (1 << 7) > 0,
-            authentic_data: datagram[3] & (1 << 5) > 0,
+            authentic_data: flags2 & (1 << 5) > 0,
-            checking_disabled: datagram[3] & (1 << 4) > 0,
+            checking_disabled: flags2 & (1 << 4) > 0,
-            response_code: datagram[3] & 0xF,
+            response_code: flags2 & 0xF,
-            number_of_questions: u16::from_be_bytes([datagram[4], datagram[5]]),
+            number_of_questions,
-            number_of_answers: u16::from_be_bytes([datagram[6], datagram[7]]),
+            number_of_answers,
-            number_of_authority_resource_records: u16::from_be_bytes([datagram[8], datagram[9]]),
+            number_of_authority_resource_records,
-            number_of_additional_resource_records: u16::from_be_bytes([datagram[10], datagram[11]]),
+            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[0] = (self.transaction_id >> 8) as u8;
-        buffer[1] = self.transaction_id[1];
+        buffer[1] = (self.transaction_id & 0x0F) as u8;
         if self.query_response {
             buffer[2] |= 1 << 7;
         }

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<Message> {
+    pub fn parse(datagram: &mut Cursor<&[u8]>) -> Result<Message> {
         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);
         }
 

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<u8> {
@@ -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]);

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<Question> {
+    pub fn parse(cursor: &mut Cursor<&[u8]>) -> Result<Question> {
-        let domain = parse_domain(datagram, pos)?;
+        let domain = parse_domain(cursor)?;
 
-        if *pos + 4 > datagram.len() {
+        let question_type = cursor.read_u16()?;
-            bail!("malformed datagram");
+        let question_class = cursor.read_u16()?;
-        }
-
-        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 = Question {
             domain,

src/ext.rs

@@ -0,0 +1,40 @@
+use std::io::{Cursor, Read};
+
+pub trait CursorExt {
+    fn read_u8(&mut self) -> anyhow::Result<u8>;
+    fn read_u16(&mut self) -> anyhow::Result<u16>;
+    fn read_u32(&mut self) -> anyhow::Result<u32>;
+    fn read_bytes(&mut self, len: usize) -> anyhow::Result<Vec<u8>>;
+    fn read_string(&mut self, len: usize) -> anyhow::Result<String>;
+}
+
+impl CursorExt for Cursor<&[u8]> {
+    fn read_u8(&mut self) -> anyhow::Result<u8> {
+        let mut result = [0u8; 1];
+        self.read_exact(&mut result)?;
+        Ok(result[0])
+    }
+
+    fn read_u16(&mut self) -> anyhow::Result<u16> {
+        let mut result = [0u8; 2];
+        self.read_exact(&mut result)?;
+        Ok(u16::from_be_bytes(result))
+    }
+
+    fn read_u32(&mut self) -> anyhow::Result<u32> {
+        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<Vec<u8>> {
+        let mut result = vec![0u8; len];
+        self.read_exact(&mut result)?;
+        Ok(result)
+    }
+
+    fn read_string(&mut self, len: usize) -> anyhow::Result<String> {
+        let result = self.read_bytes(len)?;
+        Ok(String::from_utf8_lossy(&result).to_string())
+    }
+}

src/macros.rs

@@ -1,6 +1,6 @@
 #[macro_export]
 macro_rules! txid {
-    ($bytes:expr) => {
+    ($txid:expr) => {
-        format!("{:02x}{:02x}", $bytes[0], $bytes[1])
+        format!("{:04x}", $txid)
     };
 }

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",