1#![cfg(any(feature = "alloc", feature = "std", feature = "no-atomic"))]
13
14use super::{HEADER_SIZE, Header, ResourceClass, ResourceType};
15use crate::{
16 Name,
17 constants::{MAX_LABEL_BYTES, MAX_NAME_BYTES},
18 error::{BufferTooSmallDetail, EncodeError},
19};
20
21pub const DEFAULT_COMPRESSION_TABLE: usize = 32;
24
25#[derive(Debug, Copy, Clone)]
28pub struct CompressionTable<const COMP_N: usize> {
29 entries: [(u64, u16); COMP_N],
30 used: usize,
31}
32
33impl<const COMP_N: usize> CompressionTable<COMP_N> {
34 pub const fn new() -> Self {
36 Self {
37 entries: [(0, 0); COMP_N],
38 used: 0,
39 }
40 }
41
42 pub fn lookup(&self, hash: u64) -> Option<u16> {
44 let mut i = 0usize;
45 while i < self.used {
46 let entry = self.entries.get(i)?;
47 if entry.0 == hash {
48 return Some(entry.1);
49 }
50 i = i.saturating_add(1);
51 }
52 None
53 }
54
55 pub fn insert(&mut self, hash: u64, offset: u16) {
58 if self.used < COMP_N
59 && let Some(slot) = self.entries.get_mut(self.used)
60 {
61 *slot = (hash, offset);
62 self.used = self.used.saturating_add(1);
63 }
64 }
65}
66
67impl<const COMP_N: usize> Default for CompressionTable<COMP_N> {
68 fn default() -> Self {
69 Self::new()
70 }
71}
72
73#[allow(clippy::arithmetic_side_effects)]
76fn hash_suffix(labels: &[&[u8]]) -> u64 {
77 const FNV_BASIS: u64 = 0xcbf29ce484222325;
78 const FNV_PRIME: u64 = 0x100000001b3;
79 let mut h: u64 = FNV_BASIS;
80 for label in labels {
81 for &b in *label {
82 h ^= b.to_ascii_lowercase() as u64;
83 h = h.wrapping_mul(FNV_PRIME);
84 }
85 h ^= b'.' as u64;
86 h = h.wrapping_mul(FNV_PRIME);
87 }
88 h
89}
90
91#[derive(Copy, Clone)]
96pub struct BuilderCheckpoint {
97 cursor: usize,
98 header: Header,
99 compression_used: usize,
100}
101
102pub struct MessageBuilder<'a, const COMP_N: usize = DEFAULT_COMPRESSION_TABLE> {
104 out: &'a mut [u8],
105 cursor: usize,
106 header: Header,
107 compression: CompressionTable<COMP_N>,
108}
109
110impl<'a, const COMP_N: usize> MessageBuilder<'a, COMP_N> {
111 pub fn try_new(out: &'a mut [u8], header: Header) -> Result<Self, EncodeError> {
114 if out.len() < HEADER_SIZE {
115 return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
116 HEADER_SIZE,
117 out.len(),
118 )));
119 }
120 Ok(Self {
121 out,
122 cursor: HEADER_SIZE,
123 header,
124 compression: CompressionTable::new(),
125 })
126 }
127
128 pub fn checkpoint(&self) -> BuilderCheckpoint {
131 BuilderCheckpoint {
132 cursor: self.cursor,
133 header: self.header,
134 compression_used: self.compression.used,
135 }
136 }
137
138 pub fn restore(&mut self, cp: BuilderCheckpoint) {
146 self.cursor = cp.cursor;
147 self.header = cp.header;
148 self.compression.used = cp.compression_used;
149 }
150
151 pub fn push_question(
153 &mut self,
154 name: &Name,
155 qtype: ResourceType,
156 qclass: ResourceClass,
157 unicast_response: bool,
158 ) -> Result<(), EncodeError> {
159 self.write_name(name)?;
160 self.write_u16(qtype.to_u16())?;
161 let class_raw = qclass.to_u16()
162 | if unicast_response {
163 super::resource_class::UNICAST_RESPONSE_BIT
164 } else {
165 0
166 };
167 self.write_u16(class_raw)?;
168 self
169 .header
170 .set_question_count(self.header.question_count().saturating_add(1));
171 Ok(())
172 }
173
174 pub fn push_a_answer(
180 &mut self,
181 name: &Name,
182 ttl: u32,
183 addr: core::net::Ipv4Addr,
184 cache_flush: bool,
185 ) -> Result<(), EncodeError> {
186 self.write_name(name)?;
187 self.write_u16(ResourceType::A.to_u16())?;
188 let class_raw = ResourceClass::In.to_u16()
189 | if cache_flush {
190 super::resource_class::CACHE_FLUSH_BIT
191 } else {
192 0
193 };
194 self.write_u16(class_raw)?;
195 self.write_u32(ttl)?;
196 self.write_u16(4)?;
197 for &b in &addr.octets() {
198 self.write_byte(b)?;
199 }
200 self
201 .header
202 .set_answer_count(self.header.answer_count().saturating_add(1));
203 Ok(())
204 }
205
206 pub fn push_aaaa_answer(
211 &mut self,
212 name: &Name,
213 ttl: u32,
214 addr: core::net::Ipv6Addr,
215 cache_flush: bool,
216 ) -> Result<(), EncodeError> {
217 self.write_name(name)?;
218 self.write_u16(ResourceType::AAAA.to_u16())?;
219 let class_raw = ResourceClass::In.to_u16()
220 | if cache_flush {
221 super::resource_class::CACHE_FLUSH_BIT
222 } else {
223 0
224 };
225 self.write_u16(class_raw)?;
226 self.write_u32(ttl)?;
227 self.write_u16(16)?;
228 for &b in &addr.octets() {
229 self.write_byte(b)?;
230 }
231 self
232 .header
233 .set_answer_count(self.header.answer_count().saturating_add(1));
234 Ok(())
235 }
236
237 #[allow(clippy::too_many_arguments)]
242 pub fn push_srv_answer(
243 &mut self,
244 name: &Name,
245 ttl: u32,
246 priority: u16,
247 weight: u16,
248 port: u16,
249 target: &Name,
250 cache_flush: bool,
251 ) -> Result<(), EncodeError> {
252 self.write_name(name)?;
253 self.write_u16(ResourceType::Srv.to_u16())?;
254 let class_raw = ResourceClass::In.to_u16()
255 | if cache_flush {
256 super::resource_class::CACHE_FLUSH_BIT
257 } else {
258 0
259 };
260 self.write_u16(class_raw)?;
261 self.write_u32(ttl)?;
262 let len_offset = self.cursor;
263 self.write_u16(0)?;
264 let rdata_start = self.cursor;
265 self.write_u16(priority)?;
266 self.write_u16(weight)?;
267 self.write_u16(port)?;
268 self.write_name(target)?;
269 let rdata_end = self.cursor;
270 let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
271 let slot = self
272 .out
273 .get_mut(len_offset..len_offset.saturating_add(2))
274 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
275 slot.copy_from_slice(&rdlen.to_be_bytes());
276 self
277 .header
278 .set_answer_count(self.header.answer_count().saturating_add(1));
279 Ok(())
280 }
281
282 pub fn push_txt_answer<I, S>(
287 &mut self,
288 name: &Name,
289 ttl: u32,
290 segments: I,
291 cache_flush: bool,
292 ) -> Result<(), EncodeError>
293 where
294 I: IntoIterator<Item = S>,
295 S: AsRef<[u8]>,
296 {
297 self.write_name(name)?;
298 self.write_u16(ResourceType::Txt.to_u16())?;
299 let class_raw = ResourceClass::In.to_u16()
300 | if cache_flush {
301 super::resource_class::CACHE_FLUSH_BIT
302 } else {
303 0
304 };
305 self.write_u16(class_raw)?;
306 self.write_u32(ttl)?;
307 let len_offset = self.cursor;
308 self.write_u16(0)?;
309 let rdata_start = self.cursor;
310 let mut wrote_any = false;
311 for seg in segments {
312 let s = seg.as_ref();
313 if s.len() > 255 {
314 return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
315 s.len(),
316 255,
317 )));
318 }
319 #[allow(clippy::cast_possible_truncation)]
320 self.write_byte(s.len() as u8)?;
321 for &b in s {
322 self.write_byte(b)?;
323 }
324 wrote_any = true;
325 }
326 if !wrote_any {
330 self.write_byte(0)?;
331 }
332 let rdata_end = self.cursor;
333 let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
334 let slot = self
335 .out
336 .get_mut(len_offset..len_offset.saturating_add(2))
337 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
338 slot.copy_from_slice(&rdlen.to_be_bytes());
339 self
340 .header
341 .set_answer_count(self.header.answer_count().saturating_add(1));
342 Ok(())
343 }
344
345 pub fn push_ptr_answer(
347 &mut self,
348 name: &Name,
349 ttl: u32,
350 target: &Name,
351 ) -> Result<(), EncodeError> {
352 self.write_name(name)?;
353 self.write_u16(ResourceType::Ptr.to_u16())?;
354 self.write_u16(ResourceClass::In.to_u16())?;
355 self.write_u32(ttl)?;
356 let len_offset = self.cursor;
357 self.write_u16(0)?;
358 let rdata_start = self.cursor;
359 self.write_name(target)?;
360 let rdata_end = self.cursor;
361 let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
362 let slot = self
363 .out
364 .get_mut(len_offset..len_offset.saturating_add(2))
365 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
366 slot.copy_from_slice(&rdlen.to_be_bytes());
367 self
368 .header
369 .set_answer_count(self.header.answer_count().saturating_add(1));
370 Ok(())
371 }
372
373 pub fn push_a_authority(
379 &mut self,
380 name: &Name,
381 ttl: u32,
382 addr: core::net::Ipv4Addr,
383 ) -> Result<(), EncodeError> {
384 self.write_name(name)?;
385 self.write_u16(ResourceType::A.to_u16())?;
386 self.write_u16(ResourceClass::In.to_u16())?;
387 self.write_u32(ttl)?;
388 self.write_u16(4)?;
389 for &b in &addr.octets() {
390 self.write_byte(b)?;
391 }
392 self
393 .header
394 .set_authority_count(self.header.authority_count().saturating_add(1));
395 Ok(())
396 }
397
398 pub fn push_aaaa_authority(
400 &mut self,
401 name: &Name,
402 ttl: u32,
403 addr: core::net::Ipv6Addr,
404 ) -> Result<(), EncodeError> {
405 self.write_name(name)?;
406 self.write_u16(ResourceType::AAAA.to_u16())?;
407 self.write_u16(ResourceClass::In.to_u16())?;
408 self.write_u32(ttl)?;
409 self.write_u16(16)?;
410 for &b in &addr.octets() {
411 self.write_byte(b)?;
412 }
413 self
414 .header
415 .set_authority_count(self.header.authority_count().saturating_add(1));
416 Ok(())
417 }
418
419 pub fn push_srv_authority(
421 &mut self,
422 name: &Name,
423 ttl: u32,
424 priority: u16,
425 weight: u16,
426 port: u16,
427 target: &Name,
428 ) -> Result<(), EncodeError> {
429 self.write_name(name)?;
430 self.write_u16(ResourceType::Srv.to_u16())?;
431 self.write_u16(ResourceClass::In.to_u16())?;
432 self.write_u32(ttl)?;
433 let len_offset = self.cursor;
434 self.write_u16(0)?;
435 let rdata_start = self.cursor;
436 self.write_u16(priority)?;
437 self.write_u16(weight)?;
438 self.write_u16(port)?;
439 self.write_name(target)?;
440 let rdata_end = self.cursor;
441 let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
442 let slot = self
443 .out
444 .get_mut(len_offset..len_offset.saturating_add(2))
445 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
446 slot.copy_from_slice(&rdlen.to_be_bytes());
447 self
448 .header
449 .set_authority_count(self.header.authority_count().saturating_add(1));
450 Ok(())
451 }
452
453 pub fn push_txt_authority<I, S>(
455 &mut self,
456 name: &Name,
457 ttl: u32,
458 segments: I,
459 ) -> Result<(), EncodeError>
460 where
461 I: IntoIterator<Item = S>,
462 S: AsRef<[u8]>,
463 {
464 self.write_name(name)?;
465 self.write_u16(ResourceType::Txt.to_u16())?;
466 self.write_u16(ResourceClass::In.to_u16())?;
467 self.write_u32(ttl)?;
468 let len_offset = self.cursor;
469 self.write_u16(0)?;
470 let rdata_start = self.cursor;
471 let mut wrote_any = false;
472 for seg in segments {
473 let s = seg.as_ref();
474 if s.len() > 255 {
475 return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
476 s.len(),
477 255,
478 )));
479 }
480 #[allow(clippy::cast_possible_truncation)]
481 self.write_byte(s.len() as u8)?;
482 for &b in s {
483 self.write_byte(b)?;
484 }
485 wrote_any = true;
486 }
487 if !wrote_any {
491 self.write_byte(0)?;
492 }
493 let rdata_end = self.cursor;
494 let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
495 let slot = self
496 .out
497 .get_mut(len_offset..len_offset.saturating_add(2))
498 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
499 slot.copy_from_slice(&rdlen.to_be_bytes());
500 self
501 .header
502 .set_authority_count(self.header.authority_count().saturating_add(1));
503 Ok(())
504 }
505
506 pub fn push_ptr_authority(
508 &mut self,
509 name: &Name,
510 ttl: u32,
511 target: &Name,
512 ) -> Result<(), EncodeError> {
513 self.write_name(name)?;
514 self.write_u16(ResourceType::Ptr.to_u16())?;
515 self.write_u16(ResourceClass::In.to_u16())?;
516 self.write_u32(ttl)?;
517 let len_offset = self.cursor;
518 self.write_u16(0)?;
519 let rdata_start = self.cursor;
520 self.write_name(target)?;
521 let rdata_end = self.cursor;
522 let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
523 let slot = self
524 .out
525 .get_mut(len_offset..len_offset.saturating_add(2))
526 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
527 slot.copy_from_slice(&rdlen.to_be_bytes());
528 self
529 .header
530 .set_authority_count(self.header.authority_count().saturating_add(1));
531 Ok(())
532 }
533
534 pub fn push_nsec_additional(
548 &mut self,
549 name: &Name,
550 ttl: u32,
551 present_types: &[u16],
552 cache_flush: bool,
553 ) -> Result<(), EncodeError> {
554 self.write_name(name)?;
555 self.write_u16(ResourceType::Nsec.to_u16())?;
556 let class_raw = ResourceClass::In.to_u16()
557 | if cache_flush {
558 super::resource_class::CACHE_FLUSH_BIT
559 } else {
560 0
561 };
562 self.write_u16(class_raw)?;
563 self.write_u32(ttl)?;
564 let len_offset = self.cursor;
565 self.write_u16(0)?;
566 let rdata_start = self.cursor;
567 self.write_name(name)?;
569 let mut bitmap = [0u8; 32];
573 let mut max_byte: Option<usize> = None;
574 #[allow(clippy::arithmetic_side_effects, clippy::cast_possible_truncation)]
575 for &t in present_types {
576 if t >= 256 {
577 continue;
578 }
579 let byte_idx = (t >> 3) as usize; let mask = 0x80u8 >> (t & 0x07); if let Some(slot) = bitmap.get_mut(byte_idx) {
582 *slot |= mask;
583 max_byte = Some(max_byte.map_or(byte_idx, |m| m.max(byte_idx)));
584 }
585 }
586 if let Some(max_byte) = max_byte {
587 let blen = max_byte.saturating_add(1);
588 self.write_byte(0)?; #[allow(clippy::cast_possible_truncation)]
590 self.write_byte(blen as u8)?; let mut i = 0usize;
592 while i < blen {
593 self.write_byte(bitmap.get(i).copied().unwrap_or(0))?;
594 i = i.saturating_add(1);
595 }
596 }
597 let rdata_end = self.cursor;
598 let rdlen = u16::try_from(rdata_end.saturating_sub(rdata_start)).map_err(EncodeError::from)?;
599 let slot = self
600 .out
601 .get_mut(len_offset..len_offset.saturating_add(2))
602 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
603 slot.copy_from_slice(&rdlen.to_be_bytes());
604 self
605 .header
606 .set_additional_count(self.header.additional_count().saturating_add(1));
607 Ok(())
608 }
609
610 fn write_name(&mut self, name: &Name) -> Result<(), EncodeError> {
613 let s = name.as_str();
614 if s.is_empty() {
615 return self.write_byte(0);
616 }
617 let trimmed = match s.strip_suffix('.') {
618 Some(r) => r,
619 None => s,
620 };
621 let mut labels: [&[u8]; 128] = [&[]; 128];
623 let mut nlabels = 0usize;
624 for label in trimmed.split('.') {
625 if nlabels >= labels.len() {
626 return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
627 MAX_NAME_BYTES,
628 self.out.len(),
629 )));
630 }
631 if let Some(slot) = labels.get_mut(nlabels) {
632 *slot = label.as_bytes();
633 nlabels = nlabels.saturating_add(1);
634 }
635 }
636
637 let mut suffix_start = nlabels;
639 let mut pointer_target: Option<u16> = None;
640 let mut i = 0usize;
641 while i < nlabels {
642 let suffix = match labels.get(i..nlabels) {
643 Some(s) => s,
644 None => break,
645 };
646 let h = hash_suffix(suffix);
647 if let Some(off) = self.compression.lookup(h) {
648 suffix_start = i;
649 pointer_target = Some(off);
650 break;
651 }
652 i = i.saturating_add(1);
653 }
654
655 let mut j = 0usize;
657 while j < suffix_start {
658 let label = match labels.get(j) {
659 Some(l) => *l,
660 None => break,
661 };
662 if label.len() > MAX_LABEL_BYTES as usize {
663 return Err(EncodeError::BufferTooSmall(BufferTooSmallDetail::new(
664 label.len(),
665 self.out.len(),
666 )));
667 }
668 let suffix_slice = match labels.get(j..nlabels) {
670 Some(s) => s,
671 None => break,
672 };
673 let h = hash_suffix(suffix_slice);
674 let offset = u16::try_from(self.cursor).map_err(EncodeError::from)?;
675 self.compression.insert(h, offset);
676 #[allow(clippy::cast_possible_truncation)]
677 self.write_byte(label.len() as u8)?;
678 for &b in label {
679 self.write_byte(b.to_ascii_lowercase())?;
680 }
681 j = j.saturating_add(1);
682 }
683 match pointer_target {
684 Some(off) => {
685 #[allow(clippy::arithmetic_side_effects)]
686 let hi = ((off >> 8) as u8) | 0b1100_0000;
687 #[allow(clippy::cast_possible_truncation)]
688 let lo = (off & 0x00ff) as u8;
689 self.write_byte(hi)?;
690 self.write_byte(lo)?;
691 }
692 None => {
693 self.write_byte(0)?;
694 }
695 }
696 Ok(())
697 }
698
699 fn write_byte(&mut self, b: u8) -> Result<(), EncodeError> {
700 let slot = self
701 .out
702 .get_mut(self.cursor)
703 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(1, 0)))?;
704 *slot = b;
705 self.cursor = self.cursor.saturating_add(1);
706 Ok(())
707 }
708
709 fn write_u16(&mut self, v: u16) -> Result<(), EncodeError> {
710 let bytes = v.to_be_bytes();
711 let dest = self
712 .out
713 .get_mut(self.cursor..self.cursor.saturating_add(2))
714 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(2, 0)))?;
715 dest.copy_from_slice(&bytes);
716 self.cursor = self.cursor.saturating_add(2);
717 Ok(())
718 }
719
720 fn write_u32(&mut self, v: u32) -> Result<(), EncodeError> {
721 let bytes = v.to_be_bytes();
722 let dest = self
723 .out
724 .get_mut(self.cursor..self.cursor.saturating_add(4))
725 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(4, 0)))?;
726 dest.copy_from_slice(&bytes);
727 self.cursor = self.cursor.saturating_add(4);
728 Ok(())
729 }
730
731 pub fn finish(self) -> Result<usize, EncodeError> {
734 let head_slot = self
735 .out
736 .get_mut(..HEADER_SIZE)
737 .ok_or_else(|| EncodeError::BufferTooSmall(BufferTooSmallDetail::new(HEADER_SIZE, 0)))?;
738 self.header.write(head_slot)?;
739 Ok(self.cursor)
740 }
741}
742
743#[cfg(test)]
744#[cfg(any(feature = "alloc", feature = "std"))]
745#[allow(clippy::unwrap_used)]
746mod tests;