use crate::Result; const BUF_SIZE: usize = 4096; pub struct BytePacketBuffer { pub buf: [u8; BUF_SIZE], pub pos: usize, } impl Default for BytePacketBuffer { fn default() -> Self { Self::new() } } impl BytePacketBuffer { pub fn new() -> BytePacketBuffer { BytePacketBuffer { buf: [0; BUF_SIZE], pos: 0, } } pub fn from_bytes(data: &[u8]) -> Self { let mut buf = Self::new(); let len = data.len().min(BUF_SIZE); buf.buf[..len].copy_from_slice(&data[..len]); buf } pub fn pos(&self) -> usize { self.pos } pub fn filled(&self) -> &[u8] { &self.buf[..self.pos] } pub fn step(&mut self, steps: usize) -> Result<()> { self.pos += steps; Ok(()) } pub fn seek(&mut self, pos: usize) -> Result<()> { self.pos = pos; Ok(()) } pub fn read(&mut self) -> Result { if self.pos >= BUF_SIZE { return Err("End of buffer".into()); } let res = self.buf[self.pos]; self.pos += 1; Ok(res) } pub fn get(&self, pos: usize) -> Result { if pos >= BUF_SIZE { return Err("End of buffer".into()); } Ok(self.buf[pos]) } pub fn get_range(&self, start: usize, len: usize) -> Result<&[u8]> { if start + len > BUF_SIZE { return Err("End of buffer".into()); } Ok(&self.buf[start..start + len]) } pub fn read_u16(&mut self) -> Result { let res = ((self.read()? as u16) << 8) | (self.read()? as u16); Ok(res) } pub fn read_u32(&mut self) -> Result { let res = ((self.read()? as u32) << 24) | ((self.read()? as u32) << 16) | ((self.read()? as u32) << 8) | (self.read()? as u32); Ok(res) } /// Read a qname, handling label compression (pointer jumps). /// Converts wire format like [3]www[6]google[3]com[0] into "www.google.com". pub fn read_qname(&mut self, outstr: &mut String) -> Result<()> { let mut pos = self.pos(); let mut jumped = false; let max_jumps = 5; let mut jumps_performed = 0; let mut delim = ""; loop { if jumps_performed > max_jumps { return Err(format!("Limit of {} jumps exceeded", max_jumps).into()); } let len = self.get(pos)?; if (len & 0xC0) == 0xC0 { if !jumped { self.seek(pos + 2)?; } let b2 = self.get(pos + 1)? as u16; let offset = (((len as u16) ^ 0xC0) << 8) | b2; pos = offset as usize; jumped = true; jumps_performed += 1; continue; } else { pos += 1; if len == 0 { break; } outstr.push_str(delim); let str_buffer = self.get_range(pos, len as usize)?; for &b in str_buffer { outstr.push(b.to_ascii_lowercase() as char); } delim = "."; pos += len as usize; } } if !jumped { self.seek(pos)?; } Ok(()) } pub fn write(&mut self, val: u8) -> Result<()> { if self.pos >= BUF_SIZE { return Err("End of buffer".into()); } self.buf[self.pos] = val; self.pos += 1; Ok(()) } pub fn write_u8(&mut self, val: u8) -> Result<()> { self.write(val) } pub fn write_u16(&mut self, val: u16) -> Result<()> { self.write((val >> 8) as u8)?; self.write((val & 0xFF) as u8)?; Ok(()) } pub fn write_u32(&mut self, val: u32) -> Result<()> { self.write(((val >> 24) & 0xFF) as u8)?; self.write(((val >> 16) & 0xFF) as u8)?; self.write(((val >> 8) & 0xFF) as u8)?; self.write((val & 0xFF) as u8)?; Ok(()) } pub fn write_qname(&mut self, qname: &str) -> Result<()> { if qname.is_empty() || qname == "." { self.write_u8(0)?; return Ok(()); } for label in qname.split('.') { let len = label.len(); if len == 0 { continue; // skip empty labels from trailing dot } if len > 0x3f { return Err("Single label exceeds 63 characters of length".into()); } self.write_u8(len as u8)?; for b in label.as_bytes() { self.write_u8(*b)?; } } self.write_u8(0)?; Ok(()) } pub fn write_bytes(&mut self, data: &[u8]) -> Result<()> { let end = self.pos + data.len(); if end > BUF_SIZE { return Err("End of buffer".into()); } self.buf[self.pos..end].copy_from_slice(data); self.pos = end; Ok(()) } pub fn set(&mut self, pos: usize, val: u8) -> Result<()> { if pos >= BUF_SIZE { return Err("End of buffer".into()); } self.buf[pos] = val; Ok(()) } pub fn set_u16(&mut self, pos: usize, val: u16) -> Result<()> { self.set(pos, (val >> 8) as u8)?; self.set(pos + 1, (val & 0xFF) as u8)?; Ok(()) } }