finish dns and start webserver

1 files changed, 21 insertions, 30 deletions

diff --git a/src/packet/buffer.rs b/src/dns/packet/buffer.rs
index 4ecc605..058156e 100644
--- a/src/packet/buffer.rs
+++ b/src/dns/packet/buffer.rs
@@ -1,4 +1,4 @@
-use super::Result;
+use crate::Result;
 
 pub struct PacketBuffer {
     pub buf: Vec<u8>,
@@ -9,19 +9,9 @@ pub struct PacketBuffer {
 impl PacketBuffer {
     pub fn new(buf: Vec<u8>) -> Self {
         Self {
+            size: buf.len(),
             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);
         }
     }
 
@@ -42,32 +32,25 @@ impl PacketBuffer {
     }
 
     pub fn read(&mut self) -> Result<u8> {
-        // if self.pos >= 512 {
-        //     error!("Tried to read past end of buffer");
-        //     return Err("End of buffer".into());
-        // }
-        self.check(self.pos);
+        if self.pos >= self.size {
+            return Err("Tried to read past end of buffer".into());
+        }
         let res = self.buf[self.pos];
         self.pos += 1;
-
         Ok(res)
     }
 
     pub fn get(&mut self, pos: usize) -> Result<u8> {
-        // if pos >= 512 {
-        //     error!("Tried to read past end of buffer");
-        //     return Err("End of buffer".into());
-        // }
-        self.check(pos);
+        if pos >= self.size {
+            return Err("Tried to read past end of buffer".into());
+        }
         Ok(self.buf[pos])
     }
 
     pub fn get_range(&mut self, start: usize, len: usize) -> Result<&[u8]> {
-        // if start + len >= 512 {
-        //     error!("Tried to read past end of buffer");
-        //     return Err("End of buffer".into());
-        // }
-        self.check(start + len);
+        if start + len >= self.size {
+            return Err("Tried to read past end of buffer".into());
+        }
         Ok(&self.buf[start..start + len])
     }
 
@@ -169,7 +152,13 @@ impl PacketBuffer {
     }
 
     pub fn write(&mut self, val: u8) -> Result<()> {
-        self.check(self.pos);
+        if self.size < self.pos {
+            self.size = self.pos;
+        }
+
+        if self.buf.len() <= self.size {
+            self.buf.resize(self.size + 1, 0x00);
+        }
 
         self.buf[self.pos] = val;
         self.pos += 1;
@@ -208,7 +197,9 @@ impl PacketBuffer {
             }
         }
 
-        self.write_u8(0)?;
+        if !qname.is_empty() {
+            self.write_u8(0)?;
+        }
 
         Ok(())
     }