use crate::{
Error, Result,
bytes::{Cursor, Reader},
constants::HEADER_LENGTH,
message::Flags,
};
#[derive(Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct Header {
pub id: u16,
pub flags: Flags,
pub qd_count: u16,
pub an_count: u16,
pub ns_count: u16,
pub ar_count: u16,
}
cfg_any_client! {
impl crate::bytes::Writer<Header> for crate::bytes::WCursor<'_> {
fn write(&mut self, h: &Header) -> Result<usize> {
if self.len() >= HEADER_LENGTH {
unsafe {
self.u16_be_unchecked(h.id);
self.u16_be_unchecked(h.flags.into());
self.u16_be_unchecked(h.qd_count);
self.u16_be_unchecked(h.an_count);
self.u16_be_unchecked(h.ns_count);
self.u16_be_unchecked(h.ar_count);
}
Ok(HEADER_LENGTH)
} else {
Err(Error::EndOfBuffer)
}
}
}
}
impl Reader<Header> for Cursor<'_> {
fn read(&mut self) -> Result<Header> {
if self.len() >= HEADER_LENGTH {
unsafe {
Ok(Header {
id: self.u16_be_unchecked(),
flags: Flags::from(self.u16_be_unchecked()),
qd_count: self.u16_be_unchecked(),
an_count: self.u16_be_unchecked(),
ns_count: self.u16_be_unchecked(),
ar_count: self.u16_be_unchecked(),
})
}
} else {
Err(Error::EndOfBuffer)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
bytes::{WCursor, Writer},
message::{OpCode, RCode},
};
use rand::seq::IteratorRandom;
#[test]
fn test_serialization() {
let mut rng = rand::rng();
let flags = *Flags::new()
.set_message_type(rand::random::<bool>().into())
.set_authoritative_answer(rand::random())
.set_recursion_available(rand::random())
.set_recursion_desired(rand::random())
.set_truncated(rand::random())
.set_opcode(*OpCode::VALUES.iter().choose(&mut rng).unwrap())
.set_response_code(*RCode::VALUES.iter().choose(&mut rng).unwrap());
let header = Header {
id: rand::random::<u16>(),
flags,
qd_count: rand::random::<u16>(),
an_count: rand::random::<u16>(),
ns_count: rand::random::<u16>(),
ar_count: rand::random::<u16>(),
};
let mut buf = [0u8; HEADER_LENGTH];
{
let mut wcursor = WCursor::new(&mut buf[..]);
wcursor.write(&header).unwrap();
}
let mut cursor = Cursor::new(&buf[..]);
let another = cursor.read().unwrap();
assert_eq!(header, another);
}
#[test]
fn test_serializaton_end_of_buffer() {
let mut empty_arr = [0u8; 0];
let mut small_arr = [0u8; HEADER_LENGTH - 1];
let res: Result<Header> = Cursor::new(&empty_arr[..]).read();
assert!(matches!(res, Err(Error::EndOfBuffer)));
let res: Result<Header> = Cursor::new(&small_arr[..]).read();
assert!(matches!(res, Err(Error::EndOfBuffer)));
let header = Header::default();
assert!(matches!(
WCursor::new(&mut empty_arr[..]).write(&header),
Err(Error::EndOfBuffer)
));
assert!(matches!(
WCursor::new(&mut small_arr[..]).write(&header),
Err(Error::EndOfBuffer)
));
}
}