learning/rust-dns-lowlevel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: use cursors instead of byte arrays

Browse Source
This commit is contained in:
matthias
2026-07-04 23:20:16 +02:00
parent b0e230d4d8
commit dee178ebc2
9 changed files with 136 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/dns/answer.rs
+13 -23
View File
@@ -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> {
let domain = parse_domain(datagram, pos)?;
pub fn parse(datagram: &mut Cursor<&[u8]>) -> Result<Answer> {
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,
src/dns/domainname.rs
+24 -33
View File
@@ -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 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("."))
src/dns/header.rs
+29 -21
View File
@@ -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> {
if datagram.len() < 12 {
bail!("malformed datagram");
}
pub fn parse(cursor: &mut Cursor<&[u8]>) -> Result<Header> {
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;
}
src/dns/message.rs
+5 -4
View File
@@ -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
+7 -2
View File
@@ -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
+10 -10
View File
@@ -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> {
let domain = parse_domain(datagram, pos)?;
pub fn parse(cursor: &mut Cursor<&[u8]>) -> Result<Question> {
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,
src/ext.rs
+40
View File
@@ -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
+2 -2
View File
@@ -1,6 +1,6 @@
#[macro_export]
macro_rules! txid {
($bytes:expr) => {
format!("{:02x}{:02x}", $bytes[0], $bytes[1])
($txid:expr) => {
format!("{:04x}", $txid)
};
}
src/main.rs
+6 -4
View File
@@ -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",