#![cfg(any(
feature = "alloc",
feature = "std",
feature = "heapless",
feature = "no-atomic"
))]
use super::{HEADER_SIZE, Header, ResourceClass, ResourceType};
use crate::{
Name,
constants::{MAX_LABEL_BYTES, MAX_NAME_BYTES},
error::{BufferTooSmallDetail, EncodeError},
};
pub const DEFAULT_COMPRESSION_TABLE: usize = 32;
#[derive(Debug, Copy, Clone)]
pub struct CompressionTable<const COMP_N: usize> {
entries: [(u64, u16); COMP_N],
used: usize,
}
impl<const COMP_N: usize> CompressionTable<COMP_N> {
pub const fn new() -> Self {
Self {
entries: [(0, 0); COMP_N],
used: 0,
}
}
pub fn lookup(&self, hash: u64) -> Option<u16> {
let mut i = 0usize;
while i < self.used {
let entry = self.entries.get(i)?;
if entry.0 == hash {
return Some(entry.1);
}
i = i.saturating_add(1);
}
None
}
pub fn insert(&mut self, hash: u64, offset: u16) {
if self.used < COMP_N
&& let Some(slot) = self.entries.get_mut(self.used)
{
*slot = (hash, offset);
self.used = self.used.saturating_add(1);
}
}
}
impl<const COMP_N: usize> Default for CompressionTable<COMP_N> {
fn default() -> Self {
Self::new()
}
}
#[allow(clippy::arithmetic_side_effects)]
fn hash_suffix(labels: &[&[u8]]) -> u64 {
const FNV_BASIS: u64 = 0xcbf29ce484222325;
const FNV_PRIME: u64 = 0x100000001b3;
let mut h: u64 = FNV_BASIS;
for label in labels {
for &b in *label {
h ^= b.to_ascii_lowercase() as u64;
h = h.wrapping_mul(FNV_PRIME);
}
h ^= b'.' as u64;
h = h.wrapping_mul(FNV_PRIME);
}
h
}
#[derive(Copy, Clone)]
pub struct BuilderCheckpoint {
cursor: usize,
header: Header,
compression_used: usize,
}
pub struct MessageBuilder<'a, const COMP_N: usize = DEFAULT_COMPRESSION_TABLE> {
out: &'a mut [u8],
cursor: usize,
header: Header,
compression: CompressionTable<COMP_N>,
}
impl<'a, const COMP_N: usize> MessageBuilder<'a, COMP_N> {
pub fn try_new(out: &'a mut [u8], header: Header) -> Result<Self, EncodeError> {
if out.len() < HEADER_SIZE {
return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
HEADER_SIZE,
out.len(),
)));
}
Ok(Self {
out,
cursor: HEADER_SIZE,
header,
compression: CompressionTable::new(),
})
}
pub fn checkpoint(&self) -> BuilderCheckpoint {
BuilderCheckpoint {
cursor: self.cursor,
header: self.header,
compression_used: self.compression.used,
}
}
pub fn restore(&mut self, cp: BuilderCheckpoint) {
self.cursor = cp.cursor;
self.header = cp.header;
self.compression.used = cp.compression_used;
}
pub fn push_question(
&mut self,
name: &Name,
qtype: ResourceType,
qclass: ResourceClass,
unicast_response: bool,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(qtype.to_u16())?;
let class_raw = qclass.to_u16()
| if unicast_response {
super::resource_class::UNICAST_RESPONSE_BIT
} else {
0
};
self.write_u16(class_raw)?;
self
.header
.set_question_count(self.header.question_count().saturating_add(1));
Ok(())
}
pub fn push_a_answer(
&mut self,
name: &Name,
ttl: u32,
addr: core::net::Ipv4Addr,
cache_flush: bool,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(ResourceType::A.to_u16())?;
let class_raw = ResourceClass::In.to_u16()
| if cache_flush {
super::resource_class::CACHE_FLUSH_BIT
} else {
0
};
self.write_u16(class_raw)?;
self.write_u32(ttl)?;
self.write_u16(4)?;
for &b in &addr.octets() {
self.write_byte(b)?;
}
self
.header
.set_answer_count(self.header.answer_count().saturating_add(1));
Ok(())
}
pub fn push_aaaa_answer(
&mut self,
name: &Name,
ttl: u32,
addr: core::net::Ipv6Addr,
cache_flush: bool,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(ResourceType::AAAA.to_u16())?;
let class_raw = ResourceClass::In.to_u16()
| if cache_flush {
super::resource_class::CACHE_FLUSH_BIT
} else {
0
};
self.write_u16(class_raw)?;
self.write_u32(ttl)?;
self.write_u16(16)?;
for &b in &addr.octets() {
self.write_byte(b)?;
}
self
.header
.set_answer_count(self.header.answer_count().saturating_add(1));
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub fn push_srv_answer(
&mut self,
name: &Name,
ttl: u32,
priority: u16,
weight: u16,
port: u16,
target: &Name,
cache_flush: bool,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(ResourceType::Srv.to_u16())?;
let class_raw = ResourceClass::In.to_u16()
| if cache_flush {
super::resource_class::CACHE_FLUSH_BIT
} else {
0
};
self.write_u16(class_raw)?;
self.write_u32(ttl)?;
let len_offset = self.cursor;
self.write_u16(0)?;
let rdata_start = self.cursor;
self.write_u16(priority)?;
self.write_u16(weight)?;
self.write_u16(port)?;
self.write_name(target)?;
let rdata_end = self.cursor;
let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
let slot = self
.out
.get_mut(len_offset..len_offset.saturating_add(2))
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
slot.copy_from_slice(&rdlen.to_be_bytes());
self
.header
.set_answer_count(self.header.answer_count().saturating_add(1));
Ok(())
}
pub fn push_txt_answer<I, S>(
&mut self,
name: &Name,
ttl: u32,
segments: I,
cache_flush: bool,
) -> Result<(), EncodeError>
where
I: IntoIterator<Item = S>,
S: AsRef<[u8]>,
{
self.write_name(name)?;
self.write_u16(ResourceType::Txt.to_u16())?;
let class_raw = ResourceClass::In.to_u16()
| if cache_flush {
super::resource_class::CACHE_FLUSH_BIT
} else {
0
};
self.write_u16(class_raw)?;
self.write_u32(ttl)?;
let len_offset = self.cursor;
self.write_u16(0)?;
let rdata_start = self.cursor;
let mut wrote_any = false;
for seg in segments {
let s = seg.as_ref();
if s.len() > 255 {
return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
s.len(),
255,
)));
}
#[allow(clippy::cast_possible_truncation)]
self.write_byte(s.len() as u8)?;
for &b in s {
self.write_byte(b)?;
}
wrote_any = true;
}
if !wrote_any {
self.write_byte(0)?;
}
let rdata_end = self.cursor;
let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
let slot = self
.out
.get_mut(len_offset..len_offset.saturating_add(2))
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
slot.copy_from_slice(&rdlen.to_be_bytes());
self
.header
.set_answer_count(self.header.answer_count().saturating_add(1));
Ok(())
}
pub fn push_ptr_answer(
&mut self,
name: &Name,
ttl: u32,
target: &Name,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(ResourceType::Ptr.to_u16())?;
self.write_u16(ResourceClass::In.to_u16())?;
self.write_u32(ttl)?;
let len_offset = self.cursor;
self.write_u16(0)?;
let rdata_start = self.cursor;
self.write_name(target)?;
let rdata_end = self.cursor;
let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
let slot = self
.out
.get_mut(len_offset..len_offset.saturating_add(2))
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
slot.copy_from_slice(&rdlen.to_be_bytes());
self
.header
.set_answer_count(self.header.answer_count().saturating_add(1));
Ok(())
}
pub fn push_a_authority(
&mut self,
name: &Name,
ttl: u32,
addr: core::net::Ipv4Addr,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(ResourceType::A.to_u16())?;
self.write_u16(ResourceClass::In.to_u16())?;
self.write_u32(ttl)?;
self.write_u16(4)?;
for &b in &addr.octets() {
self.write_byte(b)?;
}
self
.header
.set_authority_count(self.header.authority_count().saturating_add(1));
Ok(())
}
pub fn push_aaaa_authority(
&mut self,
name: &Name,
ttl: u32,
addr: core::net::Ipv6Addr,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(ResourceType::AAAA.to_u16())?;
self.write_u16(ResourceClass::In.to_u16())?;
self.write_u32(ttl)?;
self.write_u16(16)?;
for &b in &addr.octets() {
self.write_byte(b)?;
}
self
.header
.set_authority_count(self.header.authority_count().saturating_add(1));
Ok(())
}
pub fn push_srv_authority(
&mut self,
name: &Name,
ttl: u32,
priority: u16,
weight: u16,
port: u16,
target: &Name,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(ResourceType::Srv.to_u16())?;
self.write_u16(ResourceClass::In.to_u16())?;
self.write_u32(ttl)?;
let len_offset = self.cursor;
self.write_u16(0)?;
let rdata_start = self.cursor;
self.write_u16(priority)?;
self.write_u16(weight)?;
self.write_u16(port)?;
self.write_name(target)?;
let rdata_end = self.cursor;
let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
let slot = self
.out
.get_mut(len_offset..len_offset.saturating_add(2))
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
slot.copy_from_slice(&rdlen.to_be_bytes());
self
.header
.set_authority_count(self.header.authority_count().saturating_add(1));
Ok(())
}
pub fn push_txt_authority<I, S>(
&mut self,
name: &Name,
ttl: u32,
segments: I,
) -> Result<(), EncodeError>
where
I: IntoIterator<Item = S>,
S: AsRef<[u8]>,
{
self.write_name(name)?;
self.write_u16(ResourceType::Txt.to_u16())?;
self.write_u16(ResourceClass::In.to_u16())?;
self.write_u32(ttl)?;
let len_offset = self.cursor;
self.write_u16(0)?;
let rdata_start = self.cursor;
let mut wrote_any = false;
for seg in segments {
let s = seg.as_ref();
if s.len() > 255 {
return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
s.len(),
255,
)));
}
#[allow(clippy::cast_possible_truncation)]
self.write_byte(s.len() as u8)?;
for &b in s {
self.write_byte(b)?;
}
wrote_any = true;
}
if !wrote_any {
self.write_byte(0)?;
}
let rdata_end = self.cursor;
let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
let slot = self
.out
.get_mut(len_offset..len_offset.saturating_add(2))
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
slot.copy_from_slice(&rdlen.to_be_bytes());
self
.header
.set_authority_count(self.header.authority_count().saturating_add(1));
Ok(())
}
pub fn push_ptr_authority(
&mut self,
name: &Name,
ttl: u32,
target: &Name,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(ResourceType::Ptr.to_u16())?;
self.write_u16(ResourceClass::In.to_u16())?;
self.write_u32(ttl)?;
let len_offset = self.cursor;
self.write_u16(0)?;
let rdata_start = self.cursor;
self.write_name(target)?;
let rdata_end = self.cursor;
let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
let slot = self
.out
.get_mut(len_offset..len_offset.saturating_add(2))
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
slot.copy_from_slice(&rdlen.to_be_bytes());
self
.header
.set_authority_count(self.header.authority_count().saturating_add(1));
Ok(())
}
pub fn push_nsec_additional(
&mut self,
name: &Name,
ttl: u32,
present_types: &[u16],
cache_flush: bool,
) -> Result<(), EncodeError> {
self.write_name(name)?;
self.write_u16(ResourceType::Nsec.to_u16())?;
let class_raw = ResourceClass::In.to_u16()
| if cache_flush {
super::resource_class::CACHE_FLUSH_BIT
} else {
0
};
self.write_u16(class_raw)?;
self.write_u32(ttl)?;
let len_offset = self.cursor;
self.write_u16(0)?;
let rdata_start = self.cursor;
self.write_name(name)?;
let mut bitmap = [0u8; 32];
let mut max_byte: Option<usize> = None;
#[allow(clippy::arithmetic_side_effects, clippy::cast_possible_truncation)]
for &t in present_types {
if t >= 256 {
continue;
}
let byte_idx = (t >> 3) as usize; let mask = 0x80u8 >> (t & 0x07); if let Some(slot) = bitmap.get_mut(byte_idx) {
*slot |= mask;
max_byte = Some(max_byte.map_or(byte_idx, |m| m.max(byte_idx)));
}
}
if let Some(max_byte) = max_byte {
let blen = max_byte.saturating_add(1);
self.write_byte(0)?; #[allow(clippy::cast_possible_truncation)]
self.write_byte(blen as u8)?; let mut i = 0usize;
while i < blen {
self.write_byte(bitmap.get(i).copied().unwrap_or(0))?;
i = i.saturating_add(1);
}
}
let rdata_end = self.cursor;
let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
let slot = self
.out
.get_mut(len_offset..len_offset.saturating_add(2))
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
slot.copy_from_slice(&rdlen.to_be_bytes());
self
.header
.set_additional_count(self.header.additional_count().saturating_add(1));
Ok(())
}
fn write_name(&mut self, name: &Name) -> Result<(), EncodeError> {
let s = name.as_str();
if s.is_empty() {
return self.write_byte(0);
}
let trimmed = match s.strip_suffix('.') {
Some(r) => r,
None => s,
};
let mut labels: [&[u8]; 128] = [&[]; 128];
let mut nlabels = 0usize;
for label in trimmed.split('.') {
if nlabels >= labels.len() {
return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
MAX_NAME_BYTES,
self.out.len(),
)));
}
if let Some(slot) = labels.get_mut(nlabels) {
*slot = label.as_bytes();
nlabels = nlabels.saturating_add(1);
}
}
let mut suffix_start = nlabels;
let mut pointer_target: Option<u16> = None;
let mut i = 0usize;
while i < nlabels {
let suffix = match labels.get(i..nlabels) {
Some(s) => s,
None => break,
};
let h = hash_suffix(suffix);
if let Some(off) = self.compression.lookup(h) {
suffix_start = i;
pointer_target = Some(off);
break;
}
i = i.saturating_add(1);
}
let mut j = 0usize;
while j < suffix_start {
let label = match labels.get(j) {
Some(l) => *l,
None => break,
};
if label.len() > MAX_LABEL_BYTES as usize {
return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
label.len(),
self.out.len(),
)));
}
let suffix_slice = match labels.get(j..nlabels) {
Some(s) => s,
None => break,
};
let h = hash_suffix(suffix_slice);
let offset = u16::try_from(self.cursor).map_err(EncodeError::from)?;
self.compression.insert(h, offset);
#[allow(clippy::cast_possible_truncation)]
self.write_byte(label.len() as u8)?;
for &b in label {
self.write_byte(b.to_ascii_lowercase())?;
}
j = j.saturating_add(1);
}
match pointer_target {
Some(off) => {
#[allow(clippy::arithmetic_side_effects)]
let hi = ((off >> 8) as u8) | 0b1100_0000;
#[allow(clippy::cast_possible_truncation)]
let lo = (off & 0x00ff) as u8;
self.write_byte(hi)?;
self.write_byte(lo)?;
}
None => {
self.write_byte(0)?;
}
}
Ok(())
}
fn write_byte(&mut self, b: u8) -> Result<(), EncodeError> {
let slot = self
.out
.get_mut(self.cursor)
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(1, 0)))?;
*slot = b;
self.cursor = self.cursor.saturating_add(1);
Ok(())
}
fn write_u16(&mut self, v: u16) -> Result<(), EncodeError> {
let bytes = v.to_be_bytes();
let dest = self
.out
.get_mut(self.cursor..self.cursor.saturating_add(2))
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
dest.copy_from_slice(&bytes);
self.cursor = self.cursor.saturating_add(2);
Ok(())
}
fn write_u32(&mut self, v: u32) -> Result<(), EncodeError> {
let bytes = v.to_be_bytes();
let dest = self
.out
.get_mut(self.cursor..self.cursor.saturating_add(4))
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(4, 0)))?;
dest.copy_from_slice(&bytes);
self.cursor = self.cursor.saturating_add(4);
Ok(())
}
pub fn finish(self) -> Result<usize, EncodeError> {
let head_slot = self
.out
.get_mut(..HEADER_SIZE)
.ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(HEADER_SIZE, 0)))?;
self.header.write(head_slot)?;
Ok(self.cursor)
}
}
#[cfg(test)]
#[cfg(any(feature = "alloc", feature = "std"))]
#[allow(clippy::unwrap_used)]
mod tests;