Added qname decompressor

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<Self, &'static str> {
         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);
+    }
+}

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<u8>) {
-        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<Self, &'static str> {
         // 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!()
-    }
-}

src/question.rs

@@ -8,8 +8,8 @@ pub struct Question<'a> {
 }
 
 impl<'a> Question<'a> {
-    pub fn parse(mut data: &'a [u8]) -> Result<Self, &'static str> {
-        let qname = Qname::parse(data);
+    pub fn parse(mut data: &'a [u8], original: &'a [u8]) -> Result<Self, &'static str> {
+        let qname = Qname::parse(data, original)?;
         data = &data[qname.length()..];
 
         debug_assert!(data.len() >= 4);

src/rrecord.rs

@@ -11,8 +11,8 @@ pub struct RRecord<'a> {
 }
 
 impl<'a> RRecord<'a> {
-    pub fn parse(mut data: &'a [u8]) -> Result<Self, &'static str> {
-        let qname = Qname::parse(data);
+    pub fn parse(mut data: &'a [u8], original: &'a [u8]) -> Result<Self, &'static str> {
+        let qname = Qname::parse(data, original)?;
         data = &data[qname.length()..];
         debug_assert!(data.len() >= 10);