From 62238e5245ee3f0dd4ae26afcb383d1243282c8d Mon Sep 17 00:00:00 2001 From: Ishan Jain Date: Tue, 12 Dec 2023 03:47:19 +0530 Subject: [PATCH] Added qname decompressor --- src/main.rs | 55 ++++++++++++++++++----- src/qname.rs | 115 ++++++++++++++++++++++++++++++++++++++++-------- src/question.rs | 4 +- src/rrecord.rs | 4 +- 4 files changed, 143 insertions(+), 35 deletions(-) diff --git a/src/main.rs b/src/main.rs index d450916..ff5ab21 100644 --- a/src/main.rs +++ b/src/main.rs @@ -29,19 +29,20 @@ fn main() { recv_packet.header.recursion_avail = false; recv_packet.header.reserved = 0; recv_packet.header.rcode = if recv_packet.header.opcode == 0 { 0 } else { 4 }; - recv_packet.header.qd_count = 1; - recv_packet.header.an_count = 1; + recv_packet.header.an_count = recv_packet.header.qd_count; recv_packet.header.authority_records = 0; recv_packet.header.additional_records = 0; - recv_packet.answers = vec![RRecord { - name: recv_packet.questions[0].name.clone(), - class: 1, - r_type: 1, - ttl: 1337, - rdlength: 4, - rdata: RData::A([0x8, 0x8, 0x8, 0x8]), - }]; + for question in recv_packet.questions.iter() { + recv_packet.answers.push(RRecord { + name: question.name.clone(), + r_type: 1, + class: 1, + ttl: 1337, + rdlength: 4, + rdata: RData::A([0x8, 0x8, 0x8, 0x8]), + }) + } recv_packet.write_to(&mut response); @@ -67,18 +68,23 @@ struct Packet<'a> { impl<'a> Packet<'a> { pub fn parse(mut data: &'a [u8]) -> Result { let header = Header::parse(&data[..12])?; + let original = data; data = &data[12..]; let mut questions = Vec::with_capacity(header.qd_count as usize); for _ in 0..header.qd_count { - let rec = Question::parse(data)?; + let rec = Question::parse(data, original)?; + + println!("label = {}", rec.name); + data = &data[rec.length()..]; questions.push(rec); } let mut answers = Vec::with_capacity(header.an_count as usize); for _ in 0..header.an_count { - let rec = RRecord::parse(data)?; + let rec = RRecord::parse(data, original)?; + data = &data[rec.length()..]; answers.push(rec); } @@ -102,3 +108,28 @@ impl<'a> Packet<'a> { } } } + +#[cfg(test)] +mod test { + use crate::Packet; + + #[test] + fn parse_compressed_packet() { + let data = [ + // Header + 0xAA, 0xAA, 0x01, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + // codecrafters.io, A, Class 1 + 0x0C, 0x63, 0x6F, 0x64, 0x65, 0x63, 0x72, 0x61, 0x66, 0x74, 0x65, 0x72, 0x73, 0x02, + 0x69, 0x6F, 0x00, 0x00, 0x01, 0x00, 0x01, + // + // testing.codecrafters.io, A, class 1 + 0x7, 0x74, 0x65, 0x73, 0x74, 0x69, 0x6e, 0x67, 0xc0, 0xc, 0x00, 0x01, 0x00, 0x01, + ]; + + let packet = Packet::parse(&data); + + assert!(packet.is_ok()); + + println!("{:?}", packet); + } +} diff --git a/src/qname.rs b/src/qname.rs index caf9355..ab7d2fe 100644 --- a/src/qname.rs +++ b/src/qname.rs @@ -1,61 +1,138 @@ use std::{ borrow::Cow, - fmt::{Display, Formatter, Result as FmtResult}, + fmt::{Display, Formatter, Result as FmtResult, Write}, }; #[derive(Debug, Clone)] pub struct Qname<'a> { - labels: Cow<'a, [u8]>, + name: Cow<'a, [u8]>, + original: Cow<'a, [u8]>, + length: usize, } impl<'a> Qname<'a> { pub fn write_to(&self, buf: &mut Vec) { - buf.extend(self.labels.iter()); + buf.extend(self.name.iter()); } // Assume data _starts_ at the question section - pub fn parse(data: &'a [u8]) -> Self { + pub fn parse(data: &'a [u8], original: &'a [u8]) -> Result { // Do a quick check to make sure every thing is valid // and save the slice let mut i = 0; - while data[i] != 0 { - let length = data[i] as usize; - i += length + 1; + while i < data.len() - 1 && data[i + 1] != 0 { + let is_pointer = (data[i] & 0xc0) == 0xc0; + + if is_pointer { + let offset: u16 = ((data[i] & !0xc0) as u16) << 8 | data[i + 1] as u16; + + // Make sure offset points to a location we have already read + // 12 bytes for the header + if offset > i as u16 + 12 { + return Err("offset points to a location ahead of itself in the packet"); + } + + i += 1; + } else { + let length = data[i] as usize; + i += length + 1; + } } - // NULL byte + i += 1; - Self { - labels: Cow::Borrowed(&data[0..i]), - } + let start = original.len() - data.len(); + + Ok(Self { + name: Cow::Borrowed(&data[0..i]), + original: Cow::Borrowed(&original[0..start + i]), + length: i, + }) } pub fn length(&self) -> usize { - self.labels.len() + self.length + } + + pub fn labels(&'a self) -> Vec<&'a [u8]> { + let mut lookup: &[u8] = &self.name; + + let mut out = vec![]; + + let mut i = 0; + while i < lookup.len() - 1 && lookup[i + 1] != 0 { + let is_pointer = (lookup[i] & 0xc0) == 0xc0; + + if is_pointer { + let offset: u16 = (((lookup[i] & !0xc0) as u16) << 8) | lookup[i + 1] as u16; + + i = offset as usize; + lookup = &self.original; + } else { + let length = lookup[i] as usize; + + out.push(&lookup[i + 1..i + length + 1]); + i += length + 1; + } + } + + out + } +} + +impl<'a> Display for Qname<'a> { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + let mut i = 0; + let mut lookup: &[u8] = &self.name; + + while i < lookup.len() - 1 && lookup[i + 1] != 0 { + let is_pointer = (lookup[i] & 0xc0) == 0xc0; + + if is_pointer { + let offset: u16 = (((lookup[i] & !0xc0) as u16) << 8) | lookup[i + 1] as u16; + i = offset as usize; + lookup = &self.original; + } else { + let length = lookup[i] as usize; + + for c in &lookup[i + 1..i + 1 + length] { + f.write_char(*c as char).unwrap(); + } + + f.write_char('.').unwrap(); + + i += length + 1; + } + } + + Ok(()) } } impl<'a> From<&'a str> for Qname<'a> { fn from(value: &'a str) -> Self { let mut labels = vec![]; + let mut length = 0; for label in value.split('.') { labels.push(label.len() as u8); labels.extend(label.bytes()); + length += 1 + label.len(); } // for null byte labels.push(0); + length += 1; Qname { - labels: Cow::Owned(labels), + name: Cow::Owned(labels), + // This should not cause any problems since + // we are not yet compressing any thing + // it might be better to get rid of this completely and provide + // a custom implementation which also allows compression ? + original: Cow::Owned(vec![]), + length, } } } - -impl<'a> Display for Qname<'a> { - fn fmt(&self, f: &mut Formatter) -> FmtResult { - todo!() - } -} diff --git a/src/question.rs b/src/question.rs index 31db36b..f30f039 100644 --- a/src/question.rs +++ b/src/question.rs @@ -8,8 +8,8 @@ pub struct Question<'a> { } impl<'a> Question<'a> { - pub fn parse(mut data: &'a [u8]) -> Result { - let qname = Qname::parse(data); + pub fn parse(mut data: &'a [u8], original: &'a [u8]) -> Result { + let qname = Qname::parse(data, original)?; data = &data[qname.length()..]; debug_assert!(data.len() >= 4); diff --git a/src/rrecord.rs b/src/rrecord.rs index 56fcb8b..4ec82da 100644 --- a/src/rrecord.rs +++ b/src/rrecord.rs @@ -11,8 +11,8 @@ pub struct RRecord<'a> { } impl<'a> RRecord<'a> { - pub fn parse(mut data: &'a [u8]) -> Result { - let qname = Qname::parse(data); + pub fn parse(mut data: &'a [u8], original: &'a [u8]) -> Result { + let qname = Qname::parse(data, original)?; data = &data[qname.length()..]; debug_assert!(data.len() >= 10);