1 files changed, 70 insertions, 26 deletions

diff --git a/packet/src/buffer.rs b/src/packet/buffer.rs
index 4394705..4ecc605 100644
--- a/packet/src/buffer.rs
+++ b/src/packet/buffer.rs
@@ -1,15 +1,27 @@
 use super::Result;
 
 pub struct PacketBuffer {
-    pub buf: [u8; 512],
+    pub buf: Vec<u8>,
     pub pos: usize,
+    pub size: usize,
 }
 
 impl PacketBuffer {
-    pub fn new() -> Self {
+    pub fn new(buf: Vec<u8>) -> Self {
         Self {
-            buf: [0; 512],
+            buf,
             pos: 0,
+            size: 0,
+        }
+    }
+
+    fn check(&mut self, pos: usize) {
+        if self.size < pos {
+            self.size = pos;
+        }
+
+        if self.buf.len() <= self.size {
+            self.buf.resize(self.size + 1, 0x00);
         }
     }
 
@@ -30,9 +42,11 @@ impl PacketBuffer {
     }
 
     pub fn read(&mut self) -> Result<u8> {
-        if self.pos >= 512 {
-            return Err("End of buffer".into());
-        }
+        // if self.pos >= 512 {
+        //     error!("Tried to read past end of buffer");
+        //     return Err("End of buffer".into());
+        // }
+        self.check(self.pos);
         let res = self.buf[self.pos];
         self.pos += 1;
 
@@ -40,16 +54,20 @@ impl PacketBuffer {
     }
 
     pub fn get(&mut self, pos: usize) -> Result<u8> {
-        if pos >= 512 {
-            return Err("End of buffer".into());
-        }
+        // if pos >= 512 {
+        //     error!("Tried to read past end of buffer");
+        //     return Err("End of buffer".into());
+        // }
+        self.check(pos);
         Ok(self.buf[pos])
     }
 
     pub fn get_range(&mut self, start: usize, len: usize) -> Result<&[u8]> {
-        if start + len >= 512 {
-            return Err("End of buffer".into());
-        }
+        // if start + len >= 512 {
+        //     error!("Tried to read past end of buffer");
+        //     return Err("End of buffer".into());
+        // }
+        self.check(start + len);
         Ok(&self.buf[start..start + len])
     }
 
@@ -85,13 +103,7 @@ impl PacketBuffer {
 
             let len = self.get(pos)?;
 
-            // A two byte sequence, where the two highest bits of the first byte is
-            // set, represents a offset relative to the start of the buffer. We
-            // handle this by jumping to the offset, setting a flag to indicate
-            // that we shouldn't update the shared buffer position once done.
             if (len & 0xC0) == 0xC0 {
-                // When a jump is performed, we only modify the shared buffer
-                // position once, and avoid making the change later on.
                 if !jumped {
                     self.seek(pos + 2)?;
                 }
@@ -106,7 +118,6 @@ impl PacketBuffer {
 
             pos += 1;
 
-            // Names are terminated by an empty label of length 0
             if len == 0 {
                 break;
             }
@@ -128,10 +139,38 @@ impl PacketBuffer {
         Ok(())
     }
 
-    pub fn write(&mut self, val: u8) -> Result<()> {
-        if self.pos >= 512 {
-            return Err("End of buffer".into());
+    pub fn read_string(&mut self, outstr: &mut String) -> Result<()> {
+        let len = self.read()?;
+
+        self.read_string_n(outstr, len)?;
+
+        Ok(())
+    }
+
+    pub fn read_string_n(&mut self, outstr: &mut String, len: u8) -> Result<()> {
+        let mut pos = self.pos;
+
+        let str_buffer = self.get_range(pos, len as usize)?;
+
+        let mut i = 0;
+        for b in str_buffer {
+            let c = *b as char;
+            if c == '\0' {
+                break;
+            }
+            outstr.push(c);
+            i += 1;
         }
+
+        pos += i;
+        self.seek(pos)?;
+
+        Ok(())
+    }
+
+    pub fn write(&mut self, val: u8) -> Result<()> {
+        self.check(self.pos);
+
         self.buf[self.pos] = val;
         self.pos += 1;
         Ok(())
@@ -162,9 +201,6 @@ impl PacketBuffer {
     pub fn write_qname(&mut self, qname: &str) -> Result<()> {
         for label in qname.split('.') {
             let len = label.len();
-            if len > 0x34 {
-                return Err("Single label exceeds 63 characters of length".into());
-            }
 
             self.write_u8(len as u8)?;
             for b in label.as_bytes() {
@@ -177,6 +213,14 @@ impl PacketBuffer {
         Ok(())
     }
 
+    pub fn write_string(&mut self, text: &str) -> Result<()> {
+        for b in text.as_bytes() {
+            self.write_u8(*b)?;
+        }
+
+        Ok(())
+    }
+
     pub fn set(&mut self, pos: usize, val: u8) -> Result<()> {
         self.buf[pos] = val;
 
@@ -189,4 +233,4 @@ impl PacketBuffer {
 
         Ok(())
     }
-}
\ No newline at end of file
+}