domain/base/
message_builder.rs

1//! Building a new DNS message.
2//!
3//! The types in this module allow building a DNS message consecutively from
4//! its parts. Since messages consist of five parts, a number of types are
5//! involved. The concept is that you start out with a [`MessageBuilder`] and
6//! work your way step by step through the sections by trading the builder in
7//! for on of another type representing the following section. The sequence
8//! is [`MessageBuilder`], [`QuestionBuilder`], [`AnswerBuilder`],
9//! [`AuthorityBuilder`], and finally [`AdditionalBuilder`].
10//!
11//! You can skip forward over unwanted sections. You can also go backwards,
12//! but then you’ll loose anything you built before. The naming of the
13//! methods that do these things is consistent across types: `builder` takes
14//! you to the message builder. The four methods `question`, `answer`,
15//! `additional`, and `authority` progress or return to the respective
16//! section. Finally, `finish` completes building.
17//!
18//! Each of the section builders offers a `push` method to add elements to
19//! the section. For the question section, the method accepts anything that
20//! resembles a [`Question`] while the three record sections except
21//! something that looks like a [`Record`]. Apart from actual values
22//! of these types, tuples of the components also work, such as a pair of a
23//! domain name and a record type for a question or a triple of the owner
24//! name, TTL, and record data for a record. If you already have a question
25//! or record, you can use the `push_ref` method to add
26//!
27//!
28//! The `push` method of the record
29//! section builders is also available via the [`RecordSectionBuilder`]
30//! trait so you can build code that works with all three record sections.
31//!
32//! The [`AdditionalBuilder`] has a special feature that helps building the
33//! OPT record for EDNS. Its [`opt`][AdditionalBuilder::opt] method allows a
34//! closure to build this record on the fly via the [`OptBuilder`] type.
35//!
36//! Building happens atop any [octets builder], so the type of buffer to use
37//! for building can be chosen. The module also provides a few helper types
38//! that provide optional features for building messages. All of these are
39//! wrappers around an octets builder and are octets builders themselves, so
40//! you can mix and match.
41//!
42//! First, the [`StreamTarget`] builds a message for use with streaming
43//! transport protocols, e.g., TCP, where the actual message is preceded by
44//! a 16 bit length counter. The stream target keeps this counter up-to-date
45//! and makes sure the message doesn’t become longer than what the counter
46//! can provide for.
47//!
48//! Two further types, [`TreeCompressor`] and [`StaticCompressor`], provide
49//! name compression. This is a mechanism to decrease the size of a DNS
50//! message by avoiding repeating domain names: Instead of including a domain
51//! name or suffix of a domain name that has been mentioned already, a pointer
52//! to the position of the original mention is provided. Since this process is
53//! somewhat expensive as you have to remember which names have already been
54//! used, it isn’t enabled by default and provided via separate octets
55//! builders instead which we call compressors.
56//!
57//! Currently, there are two different compressors. [`TreeCompressor`] stores
58//! all names it encountered in a binary tree. While it can handle any number
59//! of names, it does require an allocator and therefore cannot be used in a
60//! `no_std` environment. [`StaticCompressor`], meanwhile, has a static table
61//! for up to 24 names. It is thus becoming ineffective on large messages
62//! with lots of different names. However, 24 should be good enough for most
63//! normal messages.
64//!
65//! # Example
66//!
67//! The following example builds a message with both name compression and
68//! the stream length and simply puts two A records into it.
69//!
70#![cfg_attr(feature = "std", doc = "```")]
71#![cfg_attr(not(feature = "std"), doc = "```ignore")]
72//! use std::str::FromStr;
73//! use domain::base::{
74//!     Dname, MessageBuilder, Rtype, StaticCompressor, StreamTarget
75//! };
76//! use domain::rdata::A;
77//!
78//! // Make a domain name we can use later on.
79//! let name = Dname::<Vec<u8>>::from_str("example.com").unwrap();
80//!
81//! // Create a message builder wrapping a compressor wrapping a stream
82//! // target.
83//! let mut msg = MessageBuilder::from_target(
84//!     StaticCompressor::new(
85//!         StreamTarget::new_vec()
86//!     )
87//! ).unwrap();
88//!
89//! // Set the RD bit in the header and proceed to the question section.
90//! msg.header_mut().set_rd(true);
91//! let mut msg = msg.question();
92//!
93//! // Add a question and proceed to the answer section.
94//! msg.push((&name, Rtype::A)).unwrap();
95//! let mut msg = msg.answer();
96//!
97//! // Add two answer and proceed to the additional sections
98//! msg.push((&name, 86400, A::from_octets(192, 0, 2, 1))).unwrap();
99//! msg.push((&name, 86400, A::from_octets(192, 0, 2, 2))).unwrap();
100//! let mut msg = msg.additional();
101//!
102//! // Add an OPT record.
103//! msg.opt(|opt| {
104//!     opt.set_udp_payload_size(4096);
105//!     Ok(())
106//! }).unwrap();
107//!
108//! // Convert the builder into the actual message.
109//! let target = msg.finish().into_target();
110//!
111//! // A stream target can provide access to the data with or without the
112//! // length counter:
113//! let _ = target.as_stream_slice(); // With length
114//! let _ = target.as_dgram_slice(); // Without length
115//! ```
116//!
117//! [`MessageBuilder`]: struct.MessageBuilder.html
118//! [`QuestionBuilder`]: struct.QuestionBuilder.html
119//! [`AnswerBuilder`]: struct.AnswerBuilder.html
120//! [`AuthorityBuilder`]: struct.AuthorityBuilder.html
121//! [`AdditionalBuilder`]: struct.AdditionalBuilder.html
122//! [`AdditionalBuilder::opt`]: struct.AdditionalBuilder.html#method.opt
123//! [`OptBuilder`]: struct.OptBuilder.html
124//! [`RecordSectionBuilder`]: trait.RecordSectionBuilder.html
125//! [`StaticCompressor`]: struct.StaticCompressor.html
126//! [`StreamTarget`]: struct.StreamTarget.html
127//! [`TreeCompressor`]: struct.TreeCompressor.html
128//! [`Question`]: ../question/struct.Question.html
129//! [`Record`]: ../question/struct.Record.html
130//! [octets builder]: ../octets/trait.OctetsBuilder.html
131
132use super::header::{CountOverflow, Header, HeaderCounts, HeaderSection};
133#[cfg(feature = "rand")]
134use super::iana::Rtype;
135use super::iana::{OptRcode, OptionCode, Rcode};
136use super::message::Message;
137use super::name::{Label, ToDname};
138use super::opt::{ComposeOptData, OptHeader};
139use super::question::ComposeQuestion;
140use super::record::ComposeRecord;
141use super::wire::{Compose, Composer};
142#[cfg(feature = "bytes")]
143use bytes::BytesMut;
144#[cfg(feature = "std")]
145use core::convert::TryInto;
146use core::ops::{Deref, DerefMut};
147use core::{fmt, mem};
148#[cfg(feature = "std")]
149use octseq::array::Array;
150#[cfg(any(feature = "std", feature = "bytes"))]
151use octseq::builder::infallible;
152use octseq::builder::{FreezeBuilder, OctetsBuilder, ShortBuf, Truncate};
153use octseq::octets::Octets;
154#[cfg(feature = "std")]
155use std::collections::HashMap;
156#[cfg(feature = "std")]
157use std::vec::Vec;
158
159//------------ MessageBuilder ------------------------------------------------
160
161/// Starts building a DNS message.
162///
163/// This type wraps an [`OctetsBuilder`] and starts the process of building a
164/// message. It allows access to the header section. The message builder can
165/// be traded in for any section builder or the underlying octets builder.
166///
167/// For more details see the [module documentation].
168///
169/// [module documentation]: index.html
170/// [`OctetsBuilder`]: ../../octets/trait.OctetsBuilder.html
171#[derive(Clone, Debug)]
172pub struct MessageBuilder<Target> {
173    target: Target,
174}
175
176/// # Creating Message Builders
177///
178impl<Target: OctetsBuilder + Truncate> MessageBuilder<Target> {
179    /// Creates a new message builder using the given target.
180    ///
181    /// The target must be an [`OctetsBuilder`]. It will be truncated to zero
182    /// size before appending the header section. That is, all data that was
183    /// in the builder before will be lost.
184    ///
185    /// The function will result in an error if the builder doesn’t have
186    /// enough space for the header section.
187    pub fn from_target(
188        mut target: Target,
189    ) -> Result<Self, Target::AppendError> {
190        target.truncate(0);
191        target.append_slice(HeaderSection::new().as_slice())?;
192        Ok(MessageBuilder { target })
193    }
194}
195
196#[cfg(feature = "std")]
197impl MessageBuilder<Vec<u8>> {
198    /// Creates a new message builder atop a `Vec<u8>`.
199    #[must_use]
200    pub fn new_vec() -> Self {
201        infallible(Self::from_target(Vec::new()))
202    }
203}
204
205#[cfg(feature = "std")]
206impl MessageBuilder<StreamTarget<Vec<u8>>> {
207    /// Creates a new builder for a streamable message atop a `Vec<u8>`.
208    #[must_use]
209    pub fn new_stream_vec() -> Self {
210        Self::from_target(StreamTarget::new_vec()).unwrap()
211    }
212}
213
214#[cfg(feature = "bytes")]
215impl MessageBuilder<BytesMut> {
216    /// Creates a new message builder atop a bytes value.
217    pub fn new_bytes() -> Self {
218        infallible(Self::from_target(BytesMut::new()))
219    }
220}
221
222#[cfg(feature = "bytes")]
223impl MessageBuilder<StreamTarget<BytesMut>> {
224    /// Creates a new streamable message builder atop a bytes value.
225    pub fn new_stream_bytes() -> Self {
226        Self::from_target(StreamTarget::new_bytes()).unwrap()
227    }
228}
229
230impl<Target: Composer> MessageBuilder<Target> {
231    /// Starts creating an answer for the given message.
232    ///
233    /// Specifically, this sets the ID, QR, OPCODE, RD, and RCODE fields
234    /// in the header and attempts to push the message’s questions to the
235    /// builder. If iterating of the questions fails, it adds what it can.
236    ///
237    /// The method converts the message builder into an answer builder ready
238    /// to receive the answer for the question.
239    pub fn start_answer<Octs: Octets + ?Sized>(
240        mut self,
241        msg: &Message<Octs>,
242        rcode: Rcode,
243    ) -> Result<AnswerBuilder<Target>, PushError> {
244        {
245            let header = self.header_mut();
246            header.set_id(msg.header().id());
247            header.set_qr(true);
248            header.set_opcode(msg.header().opcode());
249            header.set_rd(msg.header().rd());
250            header.set_rcode(rcode);
251        }
252        let mut builder = self.question();
253        for item in msg.question().flatten() {
254            builder.push(item)?;
255        }
256        Ok(builder.answer())
257    }
258
259    /// Creates an AXFR request for the given domain.
260    ///
261    /// Sets a random ID, pushes the domain and the AXFR record type into
262    /// the question section, and converts the builder into an answer builder.
263    #[cfg(feature = "rand")]
264    pub fn request_axfr<N: ToDname>(
265        mut self,
266        apex: N,
267    ) -> Result<AnswerBuilder<Target>, PushError> {
268        self.header_mut().set_random_id();
269        let mut builder = self.question();
270        builder.push((apex, Rtype::Axfr))?;
271        Ok(builder.answer())
272    }
273}
274
275/// # Access to the Message Header
276///
277impl<Target: OctetsBuilder + AsRef<[u8]>> MessageBuilder<Target> {
278    /// Return the current value of the message header.
279    pub fn header(&self) -> Header {
280        *Header::for_message_slice(self.target.as_ref())
281    }
282
283    /// Return the current value of the message header counts.
284    pub fn counts(&self) -> HeaderCounts {
285        *HeaderCounts::for_message_slice(self.target.as_ref())
286    }
287}
288
289impl<Target: OctetsBuilder + AsMut<[u8]>> MessageBuilder<Target> {
290    /// Returns a mutable reference to the message header for manipulations.
291    pub fn header_mut(&mut self) -> &mut Header {
292        Header::for_message_slice_mut(self.target.as_mut())
293    }
294
295    /// Returns a mutable reference to the message header counts.
296    fn counts_mut(&mut self) -> &mut HeaderCounts {
297        HeaderCounts::for_message_slice_mut(self.target.as_mut())
298    }
299}
300
301/// # Conversions
302///
303impl<Target: Composer> MessageBuilder<Target> {
304    /// Converts the message builder into a message builder
305    ///
306    /// This is a no-op.
307    pub fn builder(self) -> MessageBuilder<Target> {
308        self
309    }
310
311    /// Converts the message builder into a question builder.
312    pub fn question(self) -> QuestionBuilder<Target> {
313        QuestionBuilder::new(self)
314    }
315
316    /// Converts the message builder into an answer builder.
317    ///
318    /// This will leave the question section empty.
319    pub fn answer(self) -> AnswerBuilder<Target> {
320        self.question().answer()
321    }
322
323    /// Converts the message builder into an authority builder.
324    ///
325    /// This will leave the question and answer sections empty.
326    pub fn authority(self) -> AuthorityBuilder<Target> {
327        self.question().answer().authority()
328    }
329
330    /// Converts the message builder into an additional builder.
331    ///
332    /// This will leave the question, answer, and authority sections empty.
333    pub fn additional(self) -> AdditionalBuilder<Target> {
334        self.question().answer().authority().additional()
335    }
336
337    /// Converts the message into the underlying octets builder.
338    ///
339    /// This will leave the all sections empty.
340    pub fn finish(self) -> Target {
341        self.target
342    }
343
344    /// Converts the builder into a message.
345    ///
346    /// The method will return a message atop whatever octets sequence the
347    /// builder’s octets builder converts into.
348    pub fn into_message(self) -> Message<<Target as FreezeBuilder>::Octets>
349    where
350        Target: FreezeBuilder,
351    {
352        unsafe { Message::from_octets_unchecked(self.target.freeze()) }
353    }
354}
355
356impl<Target> MessageBuilder<Target> {
357    /// Returns a reference to the underlying octets builder.
358    pub fn as_target(&self) -> &Target {
359        &self.target
360    }
361
362    /// Returns a mutable reference to the underlying octets builder.
363    ///
364    /// Since one could entirely mess up the message with this reference, the
365    /// method is private.
366    fn as_target_mut(&mut self) -> &mut Target {
367        &mut self.target
368    }
369
370    /// Returns an octets slice of the octets assembled so far.
371    pub fn as_slice(&self) -> &[u8]
372    where
373        Target: AsRef<[u8]>,
374    {
375        self.as_target().as_ref()
376    }
377
378    /// Returns a message atop for the octets assembled so far.
379    ///
380    /// This message is atop the octets slices derived from the builder, so
381    /// it can be created cheaply.
382    pub fn as_message(&self) -> Message<&[u8]>
383    where
384        Target: AsRef<[u8]>,
385    {
386        unsafe { Message::from_octets_unchecked(self.target.as_ref()) }
387    }
388}
389
390impl<Target: Composer> MessageBuilder<Target> {
391    fn push<Push, Inc>(
392        &mut self,
393        push: Push,
394        inc: Inc,
395    ) -> Result<(), PushError>
396    where
397        Push: FnOnce(&mut Target) -> Result<(), ShortBuf>,
398        Inc: FnOnce(&mut HeaderCounts) -> Result<(), CountOverflow>,
399    {
400        let pos = self.target.as_ref().len();
401        if let Err(err) = push(&mut self.target) {
402            self.target.truncate(pos);
403            return Err(From::from(err));
404        }
405        if inc(self.counts_mut()).is_err() {
406            self.target.truncate(pos);
407            return Err(PushError::CountOverflow);
408        }
409        Ok(())
410    }
411}
412
413//--- From
414
415impl<Target> From<QuestionBuilder<Target>> for MessageBuilder<Target>
416where
417    Target: Composer,
418{
419    fn from(src: QuestionBuilder<Target>) -> Self {
420        src.builder()
421    }
422}
423
424impl<Target> From<AnswerBuilder<Target>> for MessageBuilder<Target>
425where
426    Target: Composer,
427{
428    fn from(src: AnswerBuilder<Target>) -> Self {
429        src.builder()
430    }
431}
432
433impl<Target> From<AuthorityBuilder<Target>> for MessageBuilder<Target>
434where
435    Target: Composer,
436{
437    fn from(src: AuthorityBuilder<Target>) -> Self {
438        src.builder()
439    }
440}
441
442impl<Target> From<AdditionalBuilder<Target>> for MessageBuilder<Target>
443where
444    Target: Composer,
445{
446    fn from(src: AdditionalBuilder<Target>) -> Self {
447        src.builder()
448    }
449}
450
451//--- AsRef
452//
453// XXX Should we deref down to target?
454
455impl<Target> AsRef<Target> for MessageBuilder<Target> {
456    fn as_ref(&self) -> &Target {
457        self.as_target()
458    }
459}
460
461impl<Target: AsRef<[u8]>> AsRef<[u8]> for MessageBuilder<Target> {
462    fn as_ref(&self) -> &[u8] {
463        self.as_slice()
464    }
465}
466
467//------------ QuestionBuilder -----------------------------------------------
468
469/// Builds the question section of a DNS message.
470///
471/// A value of this type can be acquired by calling the `question` method on
472/// any other builder type. See the [module documentation] for an overview of
473/// how to build a message.
474///
475/// You can push questions to the end of the question section via the
476/// [`push`] method. It accepts various things that represent a question:
477/// question values and references; tuples of a domain name, record type, and
478/// class; and, using the regular class of IN, a pair of just a domain name
479/// and record type.
480///
481/// Once you are finished building the question section, you can progress to
482/// the answer section via the [`answer`] method or finish the message via
483/// [`finish`]. Additionally, conversions to all other builder types are
484/// available as well.
485///
486/// [`answer`]: #method.answer
487/// [`finish`]: #method.finish
488/// [`push`]: #method.push
489/// [module documentation]: index.html
490#[derive(Clone, Debug)]
491pub struct QuestionBuilder<Target> {
492    builder: MessageBuilder<Target>,
493}
494
495impl<Target: OctetsBuilder> QuestionBuilder<Target> {
496    /// Creates a new question builder from a message builder.
497    fn new(builder: MessageBuilder<Target>) -> Self {
498        Self { builder }
499    }
500}
501
502impl<Target: Composer> QuestionBuilder<Target> {
503    /// Appends a question to the question section.
504    ///
505    /// This method accepts anything that implements the [`ComposeQuestion`]
506    /// trait. Apart from an actual [`Question`][super::question::Question]
507    /// or a reference to it, this can also be a tuple of a domain name,
508    /// record type, and class or, if the class is the usual IN, a pair of
509    /// just the name and type.
510    ///
511    /// In other words, the options are:
512    ///
513    #[cfg_attr(feature = "std", doc = "```")]
514    #[cfg_attr(not(feature = "std"), doc = "```ignore")]
515    /// use domain::base::{Dname, MessageBuilder, Question, Rtype};
516    /// use domain::base::iana::Class;
517    ///
518    /// let mut msg = MessageBuilder::new_vec().question();
519    /// msg.push(Question::new_in(Dname::root_ref(), Rtype::A)).unwrap();
520    /// msg.push(&Question::new_in(Dname::root_ref(), Rtype::A)).unwrap();
521    /// msg.push((Dname::root_ref(), Rtype::A, Class::In)).unwrap();
522    /// msg.push((Dname::root_ref(), Rtype::A)).unwrap();
523    /// ```
524    pub fn push(
525        &mut self,
526        question: impl ComposeQuestion,
527    ) -> Result<(), PushError> {
528        self.builder.push(
529            |target| question.compose_question(target).map_err(Into::into),
530            |counts| counts.inc_qdcount(),
531        )
532    }
533}
534
535/// # Conversions
536///
537/// Additional conversion are available via the `Deref` implementation.
538impl<Target: Composer> QuestionBuilder<Target> {
539    /// Rewinds to an empty question section.
540    ///
541    /// All previously added questions will be lost.
542    pub fn rewind(&mut self) {
543        self.as_target_mut()
544            .truncate(mem::size_of::<HeaderSection>());
545        self.counts_mut().set_qdcount(0);
546    }
547
548    /// Converts the question builder into a message builder.
549    ///
550    /// All questions will be dropped and the question section will be empty.
551    pub fn builder(mut self) -> MessageBuilder<Target> {
552        self.rewind();
553        self.builder
554    }
555}
556
557impl<Target: Composer> QuestionBuilder<Target> {
558    /// Converts the question builder into a question builder.
559    ///
560    /// In other words, doesn’t do anything.
561    pub fn question(self) -> QuestionBuilder<Target> {
562        self
563    }
564
565    /// Converts the question builder into an answer builder.
566    pub fn answer(self) -> AnswerBuilder<Target> {
567        AnswerBuilder::new(self.builder)
568    }
569
570    /// Converts the question builder into an authority builder.
571    ///
572    /// This will leave the answer section empty.
573    pub fn authority(self) -> AuthorityBuilder<Target> {
574        self.answer().authority()
575    }
576
577    /// Converts the question builder into an additional builder.
578    ///
579    /// This will leave the answer and authority sections empty.
580    pub fn additional(self) -> AdditionalBuilder<Target> {
581        self.answer().authority().additional()
582    }
583
584    /// Converts the question builder into the underlying octets builder.
585    ///
586    /// This will leave the answer, authority, and additional sections empty.
587    pub fn finish(self) -> Target {
588        self.builder.finish()
589    }
590
591    /// Converts the question builder into the final message.
592    ///
593    /// The method will return a message atop whatever octets sequence the
594    /// builder’s octets builder converts into.
595    pub fn into_message(self) -> Message<Target::Octets>
596    where
597        Target: FreezeBuilder,
598    {
599        self.builder.into_message()
600    }
601}
602
603impl<Target> QuestionBuilder<Target> {
604    /// Returns a reference to the underlying message builder.
605    pub fn as_builder(&self) -> &MessageBuilder<Target> {
606        &self.builder
607    }
608
609    /// Returns a mutable reference to the underlying message builder.
610    pub fn as_builder_mut(&mut self) -> &mut MessageBuilder<Target> {
611        &mut self.builder
612    }
613}
614
615//--- From
616
617impl<Target> From<MessageBuilder<Target>> for QuestionBuilder<Target>
618where
619    Target: Composer,
620{
621    fn from(src: MessageBuilder<Target>) -> Self {
622        src.question()
623    }
624}
625
626impl<Target> From<AnswerBuilder<Target>> for QuestionBuilder<Target>
627where
628    Target: Composer,
629{
630    fn from(src: AnswerBuilder<Target>) -> Self {
631        src.question()
632    }
633}
634
635impl<Target> From<AuthorityBuilder<Target>> for QuestionBuilder<Target>
636where
637    Target: Composer,
638{
639    fn from(src: AuthorityBuilder<Target>) -> Self {
640        src.question()
641    }
642}
643
644impl<Target> From<AdditionalBuilder<Target>> for QuestionBuilder<Target>
645where
646    Target: Composer,
647{
648    fn from(src: AdditionalBuilder<Target>) -> Self {
649        src.question()
650    }
651}
652
653//--- Deref, DerefMut, AsRef, and AsMut
654
655impl<Target> Deref for QuestionBuilder<Target> {
656    type Target = MessageBuilder<Target>;
657
658    fn deref(&self) -> &Self::Target {
659        &self.builder
660    }
661}
662
663impl<Target> DerefMut for QuestionBuilder<Target> {
664    fn deref_mut(&mut self) -> &mut Self::Target {
665        &mut self.builder
666    }
667}
668
669impl<Target> AsRef<MessageBuilder<Target>> for QuestionBuilder<Target> {
670    fn as_ref(&self) -> &MessageBuilder<Target> {
671        self.as_builder()
672    }
673}
674
675impl<Target> AsMut<MessageBuilder<Target>> for QuestionBuilder<Target> {
676    fn as_mut(&mut self) -> &mut MessageBuilder<Target> {
677        self.as_builder_mut()
678    }
679}
680
681impl<Target> AsRef<Target> for QuestionBuilder<Target> {
682    fn as_ref(&self) -> &Target {
683        self.as_target()
684    }
685}
686
687impl<Target: AsRef<[u8]>> AsRef<[u8]> for QuestionBuilder<Target> {
688    fn as_ref(&self) -> &[u8] {
689        self.as_slice()
690    }
691}
692
693//------------ AnswerBuilder -------------------------------------------------
694
695/// Builds the answer section of a DNS message.
696///
697/// A value of this type can be acquired by calling the `answer` method on
698/// any other builder type. See the [module documentation] for an overview of
699/// how to build a message.
700///
701/// You can push records to the end of the answer section via the [`push`]
702/// method. It accepts various things that represent resource records: record
703/// values and references, tuples of an owner domain name, a class, TTL, and
704/// record data, as well as tuples of just the owner, TTL, and data, assuming
705/// the class of IN.
706///
707/// Once you are finished building the answer section, you can progress to
708/// the authority section via the [`authority`] method or finish the message
709/// via [`finish`]. Additionally, conversions to all other builder types are
710/// available as well.
711///
712/// [`authority`]: #method.authority
713/// [`finish`]: #method.finish
714/// [`push`]: #method.push
715/// [module documentation]: index.html
716#[derive(Clone, Debug)]
717pub struct AnswerBuilder<Target> {
718    /// The message builder we work on.
719    builder: MessageBuilder<Target>,
720
721    /// The index in the octets builder where the answer section starts.
722    start: usize,
723}
724
725impl<Target: Composer> AnswerBuilder<Target> {
726    /// Creates a new answer builder from an underlying message builder.
727    ///
728    /// Assumes that all three record sections are empty.
729    fn new(builder: MessageBuilder<Target>) -> Self {
730        AnswerBuilder {
731            start: builder.target.as_ref().len(),
732            builder,
733        }
734    }
735}
736
737impl<Target: Composer> AnswerBuilder<Target> {
738    /// Appends a record to the answer section.
739    ///
740    /// This methods accepts anything that implements the [`ComposeRecord`]
741    /// trait. Apart from record values and references, this are tuples of
742    /// the owner domain name, optionally the class (which is taken to be IN
743    /// if missing), the TTL, and record data.
744    ///
745    /// In other words, you can do the following things:
746    ///
747    #[cfg_attr(feature = "std", doc = "```")]
748    #[cfg_attr(not(feature = "std"), doc = "```ignore")]
749    /// use domain::base::{Dname, MessageBuilder, Record, Rtype, Ttl};
750    /// use domain::base::iana::Class;
751    /// use domain::rdata::A;
752    ///
753    /// let mut msg = MessageBuilder::new_vec().answer();
754    /// let record = Record::new(
755    ///     Dname::root_ref(), Class::In, Ttl::from_secs(86400), A::from_octets(192, 0, 2, 1)
756    /// );
757    /// msg.push(&record).unwrap();
758    /// msg.push(record).unwrap();
759    /// msg.push(
760    ///     (Dname::root_ref(), Class::In, 86400, A::from_octets(192, 0, 2, 1))
761    /// ).unwrap();
762    /// msg.push(
763    ///     (Dname::root_ref(), 86400, A::from_octets(192, 0, 2, 1))
764    /// ).unwrap();
765    /// ```
766    ///
767    pub fn push(
768        &mut self,
769        record: impl ComposeRecord,
770    ) -> Result<(), PushError> {
771        self.builder.push(
772            |target| record.compose_record(target).map_err(Into::into),
773            |counts| counts.inc_ancount(),
774        )
775    }
776}
777
778/// # Conversions
779///
780/// Additional conversion are available via the `Deref` implementation.
781impl<Target: Composer> AnswerBuilder<Target> {
782    /// Rewinds to an empty answer section.
783    ///
784    /// All previously added answers will be lost.
785    pub fn rewind(&mut self) {
786        self.builder.target.truncate(self.start);
787        self.counts_mut().set_ancount(0);
788    }
789
790    /// Converts the answer builder into a message builder.
791    ///
792    /// All questions and answers will be dropped and all sections will be
793    /// empty.
794    pub fn builder(self) -> MessageBuilder<Target> {
795        self.question().builder()
796    }
797
798    /// Converts the answer builder into a question builder.
799    ///
800    /// All answers will be dropped. All previously added questions will,
801    /// however, remain.
802    pub fn question(mut self) -> QuestionBuilder<Target> {
803        self.rewind();
804        QuestionBuilder::new(self.builder)
805    }
806}
807
808impl<Target: Composer> AnswerBuilder<Target> {
809    /// Converts the answer builder into an answer builder.
810    ///
811    /// This doesn’t do anything, really.
812    pub fn answer(self) -> AnswerBuilder<Target> {
813        self
814    }
815
816    /// Converts the answer builder into an authority builder.
817    pub fn authority(self) -> AuthorityBuilder<Target> {
818        AuthorityBuilder::new(self)
819    }
820
821    /// Converts the answer builder into an additional builder.
822    ///
823    /// This will leave the authority section empty.
824    pub fn additional(self) -> AdditionalBuilder<Target> {
825        self.authority().additional()
826    }
827
828    /// Converts the answer builder into the underlying octets builder.
829    ///
830    /// This will leave the authority and additional sections empty.
831    pub fn finish(self) -> Target {
832        self.builder.finish()
833    }
834
835    /// Converts the answer builder into the final message.
836    ///
837    /// The method will return a message atop whatever octets sequence the
838    /// builder’s octets builder converts into.
839    pub fn into_message(self) -> Message<Target::Octets>
840    where
841        Target: FreezeBuilder,
842    {
843        self.builder.into_message()
844    }
845}
846
847impl<Target> AnswerBuilder<Target> {
848    /// Returns a reference to the underlying message builder.
849    pub fn as_builder(&self) -> &MessageBuilder<Target> {
850        &self.builder
851    }
852
853    /// Returns a mutable reference to the underlying message builder.
854    pub fn as_builder_mut(&mut self) -> &mut MessageBuilder<Target> {
855        &mut self.builder
856    }
857}
858
859//--- From
860
861impl<Target> From<MessageBuilder<Target>> for AnswerBuilder<Target>
862where
863    Target: Composer,
864{
865    fn from(src: MessageBuilder<Target>) -> Self {
866        src.answer()
867    }
868}
869
870impl<Target> From<QuestionBuilder<Target>> for AnswerBuilder<Target>
871where
872    Target: Composer,
873{
874    fn from(src: QuestionBuilder<Target>) -> Self {
875        src.answer()
876    }
877}
878
879impl<Target> From<AuthorityBuilder<Target>> for AnswerBuilder<Target>
880where
881    Target: Composer,
882{
883    fn from(src: AuthorityBuilder<Target>) -> Self {
884        src.answer()
885    }
886}
887
888impl<Target> From<AdditionalBuilder<Target>> for AnswerBuilder<Target>
889where
890    Target: Composer,
891{
892    fn from(src: AdditionalBuilder<Target>) -> Self {
893        src.answer()
894    }
895}
896
897//--- Deref, DerefMut, AsRef, and AsMut
898
899impl<Target> Deref for AnswerBuilder<Target> {
900    type Target = MessageBuilder<Target>;
901
902    fn deref(&self) -> &Self::Target {
903        &self.builder
904    }
905}
906
907impl<Target> DerefMut for AnswerBuilder<Target> {
908    fn deref_mut(&mut self) -> &mut Self::Target {
909        &mut self.builder
910    }
911}
912
913impl<Target> AsRef<MessageBuilder<Target>> for AnswerBuilder<Target> {
914    fn as_ref(&self) -> &MessageBuilder<Target> {
915        self.as_builder()
916    }
917}
918
919impl<Target> AsMut<MessageBuilder<Target>> for AnswerBuilder<Target> {
920    fn as_mut(&mut self) -> &mut MessageBuilder<Target> {
921        self.as_builder_mut()
922    }
923}
924
925impl<Target> AsRef<Target> for AnswerBuilder<Target> {
926    fn as_ref(&self) -> &Target {
927        self.as_target()
928    }
929}
930
931impl<Target: AsRef<[u8]>> AsRef<[u8]> for AnswerBuilder<Target> {
932    fn as_ref(&self) -> &[u8] {
933        self.as_slice()
934    }
935}
936
937//------------ AuthorityBuilder ----------------------------------------------
938
939/// Builds the authority section of a DNS message.
940///
941/// A value of this type can be acquired by calling the `authority` method on
942/// any other builder type. See the [module documentation] for an overview of
943/// how to build a message.
944///
945/// You can push records to the end of the authority section via the [`push`]
946/// method. It accepts various things that represent resource records: record
947/// values and references, tuples of an owner domain name, a class, TTL, and
948/// record data, as well as tuples of just the owner, TTL, and data, assuming
949/// the class of IN.
950///
951/// Once you are finished building the authority section, you can progress to
952/// the additional section via the [`additional`] method or finish the message
953/// via [`finish`]. Additionally, conversions to all other builder types are
954/// available as well.
955///
956/// [`additional`]: #method.additional
957/// [`finish`]: #method.finish
958/// [`push`]: #method.push
959/// [module documentation]: index.html
960#[derive(Clone, Debug)]
961pub struct AuthorityBuilder<Target> {
962    /// The message builder we work on.
963    answer: AnswerBuilder<Target>,
964
965    /// The index in the octets builder where the authority section starts.
966    start: usize,
967}
968
969impl<Target: Composer> AuthorityBuilder<Target> {
970    /// Creates a new authority builder from an answer builder.
971    ///
972    /// Assumes that the authority and additional sections are empty.
973    fn new(answer: AnswerBuilder<Target>) -> Self {
974        AuthorityBuilder {
975            start: answer.as_target().as_ref().len(),
976            answer,
977        }
978    }
979}
980
981impl<Target: Composer> AuthorityBuilder<Target> {
982    /// Appends a record to the authority section.
983    ///
984    /// This methods accepts anything that implements the [`ComposeRecord`] trait.
985    /// Apart from record values and references, this are tuples of the owner
986    /// domain name, optionally the class (which is taken to be IN if
987    /// missing), the TTL, and record data.
988    ///
989    /// In other words, you can do the following things:
990    ///
991    #[cfg_attr(feature = "std", doc = "```")]
992    #[cfg_attr(not(feature = "std"), doc = "```ignore")]
993    /// use domain::base::{Dname, MessageBuilder, Record, Rtype, Ttl};
994    /// use domain::base::iana::Class;
995    /// use domain::rdata::A;
996    ///
997    /// let mut msg = MessageBuilder::new_vec().authority();
998    /// let record = Record::new(
999    ///     Dname::root_ref(), Class::In, Ttl::from_secs(86400), A::from_octets(192, 0, 2, 1)
1000    /// );
1001    /// msg.push(&record).unwrap();
1002    /// msg.push(record).unwrap();
1003    /// msg.push(
1004    ///     (Dname::root_ref(), Class::In, 86400, A::from_octets(192, 0, 2, 1))
1005    /// ).unwrap();
1006    /// msg.push(
1007    ///     (Dname::root_ref(), 86400, A::from_octets(192, 0, 2, 1))
1008    /// ).unwrap();
1009    /// ```
1010    pub fn push(
1011        &mut self,
1012        record: impl ComposeRecord,
1013    ) -> Result<(), PushError> {
1014        self.answer.builder.push(
1015            |target| record.compose_record(target).map_err(Into::into),
1016            |counts| counts.inc_nscount(),
1017        )
1018    }
1019}
1020
1021/// # Conversions
1022///
1023/// Additional conversion methods are available via the `Deref`
1024/// implementation.
1025impl<Target: Composer> AuthorityBuilder<Target> {
1026    /// Rewinds to an empty authority section.
1027    ///
1028    /// All previously added authority records will be lost.
1029    pub fn rewind(&mut self) {
1030        self.answer.as_target_mut().truncate(self.start);
1031        self.counts_mut().set_nscount(0);
1032    }
1033
1034    /// Converts the authority builder into a message builder.
1035    ///
1036    /// All questions, answer and authority records will be dropped and all
1037    /// sections will be empty.
1038    pub fn builder(self) -> MessageBuilder<Target> {
1039        self.question().builder()
1040    }
1041
1042    /// Converts the authority builder into a question builder.
1043    ///
1044    /// All authority and answer records will be dropped. All previously added
1045    /// questions will, however, remain.
1046    pub fn question(self) -> QuestionBuilder<Target> {
1047        self.answer().question()
1048    }
1049
1050    /// Converts the authority builder into an answer builder.
1051    ///
1052    /// All authority records will be dropped. All previously added questions
1053    /// and answer records will, however, remain.
1054    pub fn answer(mut self) -> AnswerBuilder<Target> {
1055        self.rewind();
1056        self.answer
1057    }
1058}
1059
1060impl<Target: Composer> AuthorityBuilder<Target> {
1061    /// Converts the authority builder into an authority builder.
1062    ///
1063    /// This is identical to the identity function.
1064    pub fn authority(self) -> AuthorityBuilder<Target> {
1065        self
1066    }
1067
1068    /// Converts the authority builder into an additional builder.
1069    pub fn additional(self) -> AdditionalBuilder<Target> {
1070        AdditionalBuilder::new(self)
1071    }
1072
1073    /// Converts the authority builder into the underlying octets builder.
1074    ///
1075    /// This will leave the additional section empty.
1076    pub fn finish(self) -> Target {
1077        self.answer.finish()
1078    }
1079
1080    /// Converts the authority builder into the final message.
1081    ///
1082    /// The method will return a message atop whatever octets sequence the
1083    /// builder’s octets builder converts into.
1084    pub fn into_message(self) -> Message<Target::Octets>
1085    where
1086        Target: FreezeBuilder,
1087    {
1088        self.answer.into_message()
1089    }
1090}
1091
1092impl<Target> AuthorityBuilder<Target> {
1093    /// Returns a reference to the underlying message builder.
1094    pub fn as_builder(&self) -> &MessageBuilder<Target> {
1095        self.answer.as_builder()
1096    }
1097
1098    /// Returns a mutable reference to the underlying message builder.
1099    pub fn as_builder_mut(&mut self) -> &mut MessageBuilder<Target> {
1100        self.answer.as_builder_mut()
1101    }
1102}
1103
1104//--- From
1105
1106impl<Target> From<MessageBuilder<Target>> for AuthorityBuilder<Target>
1107where
1108    Target: Composer,
1109{
1110    fn from(src: MessageBuilder<Target>) -> Self {
1111        src.authority()
1112    }
1113}
1114
1115impl<Target> From<QuestionBuilder<Target>> for AuthorityBuilder<Target>
1116where
1117    Target: Composer,
1118{
1119    fn from(src: QuestionBuilder<Target>) -> Self {
1120        src.authority()
1121    }
1122}
1123
1124impl<Target> From<AnswerBuilder<Target>> for AuthorityBuilder<Target>
1125where
1126    Target: Composer,
1127{
1128    fn from(src: AnswerBuilder<Target>) -> Self {
1129        src.authority()
1130    }
1131}
1132
1133impl<Target> From<AdditionalBuilder<Target>> for AuthorityBuilder<Target>
1134where
1135    Target: Composer,
1136{
1137    fn from(src: AdditionalBuilder<Target>) -> Self {
1138        src.authority()
1139    }
1140}
1141
1142//--- Deref, DerefMut, AsRef, and AsMut
1143
1144impl<Target> Deref for AuthorityBuilder<Target> {
1145    type Target = MessageBuilder<Target>;
1146
1147    fn deref(&self) -> &Self::Target {
1148        self.answer.deref()
1149    }
1150}
1151
1152impl<Target> DerefMut for AuthorityBuilder<Target> {
1153    fn deref_mut(&mut self) -> &mut Self::Target {
1154        self.answer.deref_mut()
1155    }
1156}
1157
1158impl<Target> AsRef<MessageBuilder<Target>> for AuthorityBuilder<Target> {
1159    fn as_ref(&self) -> &MessageBuilder<Target> {
1160        self.as_builder()
1161    }
1162}
1163
1164impl<Target> AsMut<MessageBuilder<Target>> for AuthorityBuilder<Target> {
1165    fn as_mut(&mut self) -> &mut MessageBuilder<Target> {
1166        self.as_builder_mut()
1167    }
1168}
1169
1170impl<Target> AsRef<Target> for AuthorityBuilder<Target> {
1171    fn as_ref(&self) -> &Target {
1172        self.as_target()
1173    }
1174}
1175
1176impl<Target: AsRef<[u8]>> AsRef<[u8]> for AuthorityBuilder<Target> {
1177    fn as_ref(&self) -> &[u8] {
1178        self.as_slice()
1179    }
1180}
1181
1182//------------ AdditionalBuilder ---------------------------------------------
1183
1184/// Builds the additional section of a DNS message.
1185///
1186/// A value of this type can be acquired by calling the `additional` method on
1187/// any other builder type. See the [module documentation] for an overview of
1188/// how to build a message.
1189///
1190/// You can push records to the end of the additional section via the [`push`]
1191/// method. It accepts various things that represent resource records: record
1192/// values and references, tuples of an owner domain name, a class, TTL, and
1193/// record data, as well as tuples of just the owner, TTL, and data, assuming
1194/// the class of IN.
1195///
1196/// A special method exists to make adding an OPT record to the section
1197/// easier. The [`opt`] method creates an [`OptBuilder`] and passes it to a
1198/// closure. This way, you can add and remove OPT records from additional
1199/// builders that are part of another type and cannot be traded in easily.
1200///
1201/// Once you are finished building the additional section, you can finish the
1202/// message via [`finish`]. Additionally, conversions to all other builder
1203/// types are available as well.
1204///
1205/// [`finish`]: #method.finish
1206/// [`opt`]: #method.opt
1207/// [`push`]: #method.push
1208/// [`OptBuilder`]: struct.OptBuilder.html
1209/// [module documentation]: index.html
1210#[derive(Clone, Debug)]
1211pub struct AdditionalBuilder<Target> {
1212    /// The message builder we work on.
1213    authority: AuthorityBuilder<Target>,
1214
1215    /// The index in the octets builder where the additional section starts.
1216    start: usize,
1217}
1218
1219impl<Target: Composer> AdditionalBuilder<Target> {
1220    /// Creates a new additional builder from an authority builder.
1221    ///
1222    /// Assumes that the additional section is currently empty.
1223    fn new(authority: AuthorityBuilder<Target>) -> Self {
1224        AdditionalBuilder {
1225            start: authority.as_target().as_ref().len(),
1226            authority,
1227        }
1228    }
1229}
1230
1231impl<Target: Composer> AdditionalBuilder<Target> {
1232    /// Appends a record to the additional section.
1233    ///
1234    /// This methods accepts anything that implements the
1235    /// [`ComposeRecord`] trait.
1236    /// Apart from record values and references, this are tuples of the owner
1237    /// domain name, optionally the class (which is taken to be IN if
1238    /// missing), the TTL, and record data.
1239    ///
1240    /// In other words, you can do the following things:
1241    ///
1242    #[cfg_attr(feature = "std", doc = "```")]
1243    #[cfg_attr(not(feature = "std"), doc = "```ignore")]
1244    /// use domain::base::{Dname, MessageBuilder, Record, Rtype, Ttl};
1245    /// use domain::base::iana::Class;
1246    /// use domain::rdata::A;
1247    ///
1248    /// let mut msg = MessageBuilder::new_vec().additional();
1249    /// let record = Record::new(
1250    ///     Dname::root_ref(), Class::In, Ttl::from_secs(86400), A::from_octets(192, 0, 2, 1)
1251    /// );
1252    /// msg.push(&record).unwrap();
1253    /// msg.push(record).unwrap();
1254    /// msg.push(
1255    ///     (Dname::root_ref(), Class::In, 86400, A::from_octets(192, 0, 2, 1))
1256    /// ).unwrap();
1257    /// msg.push(
1258    ///     (Dname::root_ref(), 86400, A::from_octets(192, 0, 2, 1))
1259    /// ).unwrap();
1260    /// ```
1261    pub fn push(
1262        &mut self,
1263        record: impl ComposeRecord,
1264    ) -> Result<(), PushError> {
1265        self.authority.answer.builder.push(
1266            |target| record.compose_record(target).map_err(Into::into),
1267            |counts| counts.inc_arcount(),
1268        )
1269    }
1270}
1271
1272impl<Target: Composer> AdditionalBuilder<Target> {
1273    /// Appends and builds an OPT record.
1274    ///
1275    /// The actual building of the record is handled by a closure that
1276    /// receives an [`OptBuilder`] which can both change the header of the
1277    /// record and add options.
1278    ///
1279    /// The method will return whatever the closure returns. In addition, it
1280    /// will return an error if it failed to add the header of the OPT record.
1281    ///
1282    /// [`OptBuilder`]: struct.OptBuilder.html
1283    pub fn opt<F>(&mut self, op: F) -> Result<(), PushError>
1284    where
1285        F: FnOnce(&mut OptBuilder<Target>) -> Result<(), Target::AppendError>,
1286    {
1287        self.authority.answer.builder.push(
1288            |target| OptBuilder::new(target)?.build(op).map_err(Into::into),
1289            |counts| counts.inc_arcount(),
1290        )
1291    }
1292}
1293
1294/// # Conversions
1295///
1296/// Additional conversion methods are available via the `Deref`
1297/// implementation.
1298impl<Target: Composer> AdditionalBuilder<Target> {
1299    /// Rewinds to an empty additional section.
1300    ///
1301    /// All previously added additional records will be lost.
1302    pub fn rewind(&mut self) {
1303        self.authority.as_target_mut().truncate(self.start);
1304        self.counts_mut().set_arcount(0);
1305    }
1306
1307    /// Converts the additional builder into a message builder.
1308    ///
1309    /// All questions and records will be dropped and all sections will be
1310    /// empty.
1311    pub fn builder(self) -> MessageBuilder<Target> {
1312        self.question().builder()
1313    }
1314
1315    /// Converts the additional builder into a question builder.
1316    ///
1317    /// All answer, authority, and additional records will be dropped. All
1318    /// previously added questions will, however, remain.
1319    pub fn question(self) -> QuestionBuilder<Target> {
1320        self.answer().question()
1321    }
1322
1323    /// Converts the additional builder into an answer builder.
1324    ///
1325    /// All authority and additional records will be dropped. All questions
1326    /// and answer records will remain.
1327    pub fn answer(self) -> AnswerBuilder<Target> {
1328        self.authority().answer()
1329    }
1330
1331    /// Converts the additional builder into an authority builder.
1332    ///
1333    /// All additional records will be dropped. All questions, answer, and
1334    /// authority records will remain.
1335    pub fn authority(mut self) -> AuthorityBuilder<Target> {
1336        self.rewind();
1337        self.authority
1338    }
1339}
1340
1341impl<Target: Composer> AdditionalBuilder<Target> {
1342    /// Converts the additional builder into an additional builder.
1343    ///
1344    /// In other words, does absolutely nothing.
1345    pub fn additional(self) -> AdditionalBuilder<Target> {
1346        self
1347    }
1348
1349    /// Converts the additional builder into the underlying octets builder.
1350    pub fn finish(self) -> Target {
1351        self.authority.finish()
1352    }
1353
1354    /// Converts the additional builder into the final message.
1355    ///
1356    /// The method will return a message atop whatever octets sequence the
1357    /// builder’s octets builder converts into.
1358    pub fn into_message(self) -> Message<Target::Octets>
1359    where
1360        Target: FreezeBuilder,
1361    {
1362        self.authority.into_message()
1363    }
1364}
1365
1366impl<Target> AdditionalBuilder<Target> {
1367    /// Returns a reference to the underlying message builder.
1368    pub fn as_builder(&self) -> &MessageBuilder<Target> {
1369        self.authority.as_builder()
1370    }
1371
1372    /// Returns a mutable reference to the underlying message builder.
1373    pub fn as_builder_mut(&mut self) -> &mut MessageBuilder<Target> {
1374        self.authority.as_builder_mut()
1375    }
1376}
1377
1378//--- From
1379
1380impl<Target> From<MessageBuilder<Target>> for AdditionalBuilder<Target>
1381where
1382    Target: Composer,
1383{
1384    fn from(src: MessageBuilder<Target>) -> Self {
1385        src.additional()
1386    }
1387}
1388
1389impl<Target> From<QuestionBuilder<Target>> for AdditionalBuilder<Target>
1390where
1391    Target: Composer,
1392{
1393    fn from(src: QuestionBuilder<Target>) -> Self {
1394        src.additional()
1395    }
1396}
1397
1398impl<Target> From<AnswerBuilder<Target>> for AdditionalBuilder<Target>
1399where
1400    Target: Composer,
1401{
1402    fn from(src: AnswerBuilder<Target>) -> Self {
1403        src.additional()
1404    }
1405}
1406
1407impl<Target> From<AuthorityBuilder<Target>> for AdditionalBuilder<Target>
1408where
1409    Target: Composer,
1410{
1411    fn from(src: AuthorityBuilder<Target>) -> Self {
1412        src.additional()
1413    }
1414}
1415
1416//--- Deref, DerefMut, AsRef, and AsMut
1417
1418impl<Target> Deref for AdditionalBuilder<Target> {
1419    type Target = MessageBuilder<Target>;
1420
1421    fn deref(&self) -> &Self::Target {
1422        self.as_builder()
1423    }
1424}
1425
1426impl<Target> DerefMut for AdditionalBuilder<Target> {
1427    fn deref_mut(&mut self) -> &mut Self::Target {
1428        self.as_builder_mut()
1429    }
1430}
1431
1432impl<Target> AsRef<MessageBuilder<Target>> for AdditionalBuilder<Target> {
1433    fn as_ref(&self) -> &MessageBuilder<Target> {
1434        self.as_builder()
1435    }
1436}
1437
1438impl<Target> AsMut<MessageBuilder<Target>> for AdditionalBuilder<Target> {
1439    fn as_mut(&mut self) -> &mut MessageBuilder<Target> {
1440        self.as_builder_mut()
1441    }
1442}
1443
1444impl<Target> AsRef<Target> for AdditionalBuilder<Target> {
1445    fn as_ref(&self) -> &Target {
1446        self.as_target()
1447    }
1448}
1449
1450impl<Target: AsRef<[u8]>> AsRef<[u8]> for AdditionalBuilder<Target> {
1451    fn as_ref(&self) -> &[u8] {
1452        self.as_slice()
1453    }
1454}
1455
1456//------------ RecordSectionBuilder ------------------------------------------
1457
1458/// A section that can have records pushed to it.
1459///
1460/// This trait exists to make it possible to write code that works for all
1461/// three record sections. It basically just duplicates the `push` method of
1462/// these sections.
1463///
1464/// (This method is available on the sections as a method, too, so you don’t
1465/// need to import the `RecordSectionBuilder` all the time.)
1466pub trait RecordSectionBuilder<Target: Composer> {
1467    /// Appends a record to a record section.
1468    ///
1469    /// The methods accepts anything that implements the [`ComposeRecord`] trait.
1470    /// Apart from record values and references, this are tuples of the owner
1471    /// domain name, optionally the class (which is taken to be IN if
1472    /// missing), the TTL, and record data.
1473    fn push(&mut self, record: impl ComposeRecord) -> Result<(), PushError>;
1474}
1475
1476impl<Target> RecordSectionBuilder<Target> for AnswerBuilder<Target>
1477where
1478    Target: Composer,
1479{
1480    fn push(&mut self, record: impl ComposeRecord) -> Result<(), PushError> {
1481        Self::push(self, record)
1482    }
1483}
1484
1485impl<Target: Composer> RecordSectionBuilder<Target>
1486    for AuthorityBuilder<Target>
1487{
1488    fn push(&mut self, record: impl ComposeRecord) -> Result<(), PushError> {
1489        Self::push(self, record)
1490    }
1491}
1492
1493impl<Target> RecordSectionBuilder<Target> for AdditionalBuilder<Target>
1494where
1495    Target: Composer,
1496{
1497    fn push(&mut self, record: impl ComposeRecord) -> Result<(), PushError> {
1498        Self::push(self, record)
1499    }
1500}
1501
1502//------------ OptBuilder ----------------------------------------------------
1503
1504/// Builds an OPT record.
1505///
1506/// A mutable reference of this type is passed to the closure given to
1507/// [`AdditionalBuilder::opt`] allowing this closure to manipulate both the
1508/// header values of the record and push options to the record data.
1509///
1510/// [`AdditionalBuilder::opt`]: struct.AdditonalBuilder.html#method.opt
1511pub struct OptBuilder<'a, Target: ?Sized> {
1512    start: usize,
1513    target: &'a mut Target,
1514}
1515
1516impl<'a, Target: Composer + ?Sized> OptBuilder<'a, Target> {
1517    /// Creates a new opt builder atop an additional builder.
1518    fn new(target: &'a mut Target) -> Result<Self, ShortBuf> {
1519        let start = target.as_ref().len();
1520        OptHeader::default().compose(target).map_err(Into::into)?;
1521        Ok(OptBuilder { start, target })
1522    }
1523
1524    fn build<F>(&mut self, op: F) -> Result<(), ShortBuf>
1525    where
1526        F: FnOnce(&mut Self) -> Result<(), Target::AppendError>,
1527    {
1528        self.target.append_slice(&[0; 2]).map_err(Into::into)?;
1529        let pos = self.target.as_ref().len();
1530        match op(self) {
1531            Ok(_) => match u16::try_from(self.target.as_ref().len() - pos) {
1532                Ok(len) => {
1533                    self.target.as_mut()[pos - 2..pos]
1534                        .copy_from_slice(&(len).to_be_bytes());
1535                    Ok(())
1536                }
1537                Err(_) => {
1538                    self.target.truncate(pos);
1539                    Err(ShortBuf)
1540                }
1541            },
1542            Err(_) => {
1543                self.target.truncate(pos);
1544                Err(ShortBuf)
1545            }
1546        }
1547    }
1548
1549    /// Appends an option to the OPT record.
1550    pub fn push<Opt: ComposeOptData + ?Sized>(
1551        &mut self,
1552        opt: &Opt,
1553    ) -> Result<(), Target::AppendError> {
1554        self.push_raw_option(opt.code(), opt.compose_len(), |target| {
1555            opt.compose_option(target)
1556        })
1557    }
1558
1559    /// Appends a raw option to the OPT record.
1560    ///
1561    /// The method will append an option with the given option code. The data
1562    /// of the option will be written via the closure `op`.
1563    pub fn push_raw_option<F>(
1564        &mut self,
1565        code: OptionCode,
1566        option_len: u16,
1567        op: F,
1568    ) -> Result<(), Target::AppendError>
1569    where
1570        F: FnOnce(&mut Target) -> Result<(), Target::AppendError>,
1571    {
1572        code.compose(self.target)?;
1573        option_len.compose(self.target)?;
1574        op(self.target)
1575    }
1576
1577    /// Returns the current UDP payload size field of the OPT record.
1578    ///
1579    /// This field contains the largest UDP datagram the sender can accept.
1580    /// This is not the path MTU but really what the sender can work with
1581    /// internally.
1582    #[must_use]
1583    pub fn udp_payload_size(&self) -> u16 {
1584        self.opt_header().udp_payload_size()
1585    }
1586
1587    /// Sets the UDP payload size field of the OPT record.
1588    pub fn set_udp_payload_size(&mut self, value: u16) {
1589        self.opt_header_mut().set_udp_payload_size(value)
1590    }
1591
1592    /// Returns the extended rcode of the message.
1593    ///
1594    /// The method assembles the rcode both from the message header and the
1595    /// OPT header.
1596    #[must_use]
1597    pub fn rcode(&self) -> OptRcode {
1598        self.opt_header()
1599            .rcode(*Header::for_message_slice(self.target.as_ref()))
1600    }
1601
1602    /// Sets the extended rcode of the message.
1603    //
1604    /// The method will update both the message header and the OPT header.
1605    pub fn set_rcode(&mut self, rcode: OptRcode) {
1606        Header::for_message_slice_mut(self.target.as_mut())
1607            .set_rcode(rcode.rcode());
1608        self.opt_header_mut().set_rcode(rcode)
1609    }
1610
1611    /// Returns the EDNS version of the OPT header.
1612    ///
1613    /// Only EDNS version 0 is currently defined.
1614    #[must_use]
1615    pub fn version(&self) -> u8 {
1616        self.opt_header().version()
1617    }
1618
1619    /// Sets the EDNS version of the OPT header.
1620    pub fn set_version(&mut self, version: u8) {
1621        self.opt_header_mut().set_version(version)
1622    }
1623
1624    /// Returns the value of the DNSSEC OK (DO) bit.
1625    ///
1626    /// By setting this bit, a resolver indicates that it is interested in
1627    /// also receiving the DNSSEC-related resource records necessary to
1628    /// validate an answer. The bit and the related procedures are defined in
1629    /// [RFC 3225].
1630    ///
1631    /// [RFC 3225]: https://tools.ietf.org/html/rfc3225
1632    #[must_use]
1633    pub fn dnssec_ok(&self) -> bool {
1634        self.opt_header().dnssec_ok()
1635    }
1636
1637    /// Sets the DNSSEC OK (DO) bit to the given value.
1638    pub fn set_dnssec_ok(&mut self, value: bool) {
1639        self.opt_header_mut().set_dnssec_ok(value)
1640    }
1641
1642    /// Returns a reference to the full OPT header.
1643    fn opt_header(&self) -> &OptHeader {
1644        OptHeader::for_record_slice(&self.target.as_ref()[self.start..])
1645    }
1646
1647    /// Returns a mutual reference to the full OPT header.
1648    fn opt_header_mut(&mut self) -> &mut OptHeader {
1649        let start = self.start;
1650        OptHeader::for_record_slice_mut(&mut self.target.as_mut()[start..])
1651    }
1652}
1653
1654//------------ StreamTarget --------------------------------------------------
1655
1656/// A builder target for sending messages on stream transports.
1657///
1658/// When messages are sent over stream-oriented transports such as TCP, a DNS
1659/// message is preceded by a 16 bit length value in order to determine the
1660/// end of a message. This type transparently adds this length value as the
1661/// first two octets of an octets builder and itself presents an octets
1662/// builder interface for building the actual message. Whenever data is pushed
1663/// to that builder interface, the type will update the length value.
1664///
1665/// Because the length is 16 bits long, the assembled message can be at most
1666/// 65536 octets long, independently of the maximum length the underlying
1667/// builder allows.
1668#[derive(Clone, Debug)]
1669pub struct StreamTarget<Target> {
1670    /// The underlying octets builder.
1671    target: Target,
1672}
1673
1674impl<Target: Composer> StreamTarget<Target> {
1675    /// Creates a new stream target wrapping an octets builder.
1676    ///
1677    /// The function will truncate the builder back to empty and append the
1678    /// length value. Because of the latter, this can fail if the octets
1679    /// builder doesn’t even have space for that.
1680    pub fn new(mut target: Target) -> Result<Self, Target::AppendError> {
1681        target.truncate(0);
1682        0u16.compose(&mut target)?;
1683        Ok(StreamTarget { target })
1684    }
1685}
1686
1687#[cfg(feature = "std")]
1688impl StreamTarget<Vec<u8>> {
1689    /// Creates a stream target atop an empty `Vec<u8>`.
1690    #[must_use]
1691    pub fn new_vec() -> Self {
1692        infallible(Self::new(Vec::new()))
1693    }
1694}
1695
1696#[cfg(feature = "bytes")]
1697impl StreamTarget<BytesMut> {
1698    /// Creates a stream target atop an empty `Vec<u8>`.
1699    pub fn new_bytes() -> Self {
1700        infallible(Self::new(BytesMut::new()))
1701    }
1702}
1703
1704impl<Target> StreamTarget<Target> {
1705    /// Returns a reference to the underlying octets builder.
1706    pub fn as_target(&self) -> &Target {
1707        &self.target
1708    }
1709
1710    /// Converts the stream target into the underlying octets builder.
1711    ///
1712    /// The returned builder will contain the 16 bit length value with the
1713    /// correct content and the assembled message.
1714    pub fn into_target(self) -> Target {
1715        self.target
1716    }
1717}
1718
1719impl<Target: AsRef<[u8]> + AsMut<[u8]>> StreamTarget<Target> {
1720    /// Updates the length value to the current length of the target.
1721    fn update_shim(&mut self) -> Result<(), ShortBuf> {
1722        match u16::try_from(self.target.as_ref().len() - 2) {
1723            Ok(len) => {
1724                self.target.as_mut()[..2].copy_from_slice(&len.to_be_bytes());
1725                Ok(())
1726            }
1727            Err(_) => Err(ShortBuf),
1728        }
1729    }
1730}
1731
1732impl<Target: AsRef<[u8]>> StreamTarget<Target> {
1733    /// Returns an octets slice of the message for stream transports.
1734    ///
1735    /// The slice will start with the length octets and can be send as is
1736    /// through a stream transport such as TCP.
1737    pub fn as_stream_slice(&self) -> &[u8] {
1738        self.target.as_ref()
1739    }
1740
1741    /// Returns an octets slice of the message for datagram transports.
1742    ///
1743    /// The slice will not contain the length octets but only the actual
1744    /// message itself. This slice can be used for sending via datagram
1745    /// transports such as UDP.
1746    pub fn as_dgram_slice(&self) -> &[u8] {
1747        &self.target.as_ref()[2..]
1748    }
1749}
1750
1751//--- AsRef, AsMut
1752
1753impl<Target: AsRef<[u8]>> AsRef<[u8]> for StreamTarget<Target> {
1754    fn as_ref(&self) -> &[u8] {
1755        &self.target.as_ref()[2..]
1756    }
1757}
1758
1759impl<Target: AsMut<[u8]>> AsMut<[u8]> for StreamTarget<Target> {
1760    fn as_mut(&mut self) -> &mut [u8] {
1761        &mut self.target.as_mut()[2..]
1762    }
1763}
1764
1765//--- OctetsBuilder, Truncate, Composer
1766
1767impl<Target> OctetsBuilder for StreamTarget<Target>
1768where
1769    Target: OctetsBuilder + AsRef<[u8]> + AsMut<[u8]>,
1770    Target::AppendError: Into<ShortBuf>,
1771{
1772    type AppendError = ShortBuf;
1773
1774    fn append_slice(
1775        &mut self,
1776        slice: &[u8],
1777    ) -> Result<(), Self::AppendError> {
1778        self.target.append_slice(slice).map_err(Into::into)?;
1779        self.update_shim()
1780    }
1781}
1782
1783impl<Target: Composer> Truncate for StreamTarget<Target> {
1784    fn truncate(&mut self, len: usize) {
1785        self.target
1786            .truncate(len.checked_add(2).expect("long truncate"));
1787        self.update_shim().expect("truncate grew buffer???")
1788    }
1789}
1790
1791impl<Target> Composer for StreamTarget<Target>
1792where
1793    Target: Composer,
1794    Target::AppendError: Into<ShortBuf>,
1795{
1796    fn append_compressed_dname<N: ToDname + ?Sized>(
1797        &mut self,
1798        name: &N,
1799    ) -> Result<(), Self::AppendError> {
1800        self.target
1801            .append_compressed_dname(name)
1802            .map_err(Into::into)?;
1803        self.update_shim()
1804    }
1805}
1806
1807//------------ StaticCompressor ----------------------------------------------
1808
1809/// A domain name compressor that doesn’t require an allocator.
1810///
1811/// This type wraps around an octets builder and implements domain name
1812/// compression. It does not require an allocator but because of that it
1813/// can only remember the position of up to 24 domain names. This should be
1814/// sufficient for most messages.
1815///
1816/// The position of a domain name is calculated relative to the beginning of
1817/// the underlying octets builder. This means that this builder must represent
1818/// the message only. This means that if you are using the [`StreamTarget`],
1819/// you need to place it inside this type, _not_ the other way around.
1820///
1821/// [`StreamTarget`]: struct.StreamTarget.html
1822#[derive(Clone, Debug)]
1823pub struct StaticCompressor<Target> {
1824    /// The underlying octets builder.
1825    target: Target,
1826
1827    /// The domain names we have encountered so far.
1828    ///
1829    /// The value is the position of the domain name within the message.
1830    entries: [u16; 24],
1831
1832    /// The number of entries in `entries`.
1833    len: usize,
1834}
1835
1836impl<Target> StaticCompressor<Target> {
1837    /// Creates a static compressor from an octets builder.
1838    pub fn new(target: Target) -> Self {
1839        StaticCompressor {
1840            target,
1841            entries: Default::default(),
1842            len: 0,
1843        }
1844    }
1845
1846    /// Returns a reference to the underlying octets builder.
1847    pub fn as_target(&self) -> &Target {
1848        &self.target
1849    }
1850
1851    /// Converts the static compressor into the underlying octets builder.
1852    pub fn into_target(self) -> Target {
1853        self.target
1854    }
1855
1856    /// Returns a reference to the octets slice of the content.
1857    pub fn as_slice(&self) -> &[u8]
1858    where
1859        Target: AsRef<[u8]>,
1860    {
1861        self.target.as_ref()
1862    }
1863
1864    /// Returns a reference to the octets slice of the content.
1865    pub fn as_slice_mut(&mut self) -> &mut [u8]
1866    where
1867        Target: AsMut<[u8]>,
1868    {
1869        self.target.as_mut()
1870    }
1871
1872    /// Returns a known position of a domain name if there is one.
1873    fn get<'a, N: Iterator<Item = &'a Label> + Clone>(
1874        &self,
1875        name: N,
1876    ) -> Option<u16>
1877    where
1878        Target: AsRef<[u8]>,
1879    {
1880        self.entries[..self.len].iter().find_map(|&pos| {
1881            if name
1882                .clone()
1883                .eq(Label::iter_slice(self.target.as_ref(), pos as usize))
1884            {
1885                Some(pos)
1886            } else {
1887                None
1888            }
1889        })
1890    }
1891
1892    /// Inserts the position of a new domain name if possible.
1893    fn insert(&mut self, pos: usize) -> bool {
1894        if pos < 0xc000 && self.len < self.entries.len() {
1895            self.entries[self.len] = pos as u16;
1896            self.len += 1;
1897            true
1898        } else {
1899            false
1900        }
1901    }
1902}
1903
1904//--- AsRef and AsMut
1905
1906impl<Target: AsRef<[u8]>> AsRef<[u8]> for StaticCompressor<Target> {
1907    fn as_ref(&self) -> &[u8] {
1908        self.as_slice()
1909    }
1910}
1911
1912impl<Target: AsMut<[u8]>> AsMut<[u8]> for StaticCompressor<Target> {
1913    fn as_mut(&mut self) -> &mut [u8] {
1914        self.as_slice_mut()
1915    }
1916}
1917
1918//--- OctetsBuilder
1919
1920impl<Target: OctetsBuilder> OctetsBuilder for StaticCompressor<Target> {
1921    type AppendError = Target::AppendError;
1922
1923    fn append_slice(
1924        &mut self,
1925        slice: &[u8],
1926    ) -> Result<(), Self::AppendError> {
1927        self.target.append_slice(slice)
1928    }
1929}
1930
1931impl<Target: Composer> Composer for StaticCompressor<Target> {
1932    fn append_compressed_dname<N: ToDname + ?Sized>(
1933        &mut self,
1934        name: &N,
1935    ) -> Result<(), Self::AppendError> {
1936        let mut name = name.iter_labels().peekable();
1937
1938        loop {
1939            // If the parent is root, just write that and return.
1940            // Because we do that, there will always be a label left here.
1941            if let Some(label) = name.peek() {
1942                if label.is_root() {
1943                    label.compose(self)?;
1944                    return Ok(());
1945                }
1946            }
1947
1948            // If we already know this name, append it as a compressed label.
1949            if let Some(pos) = self.get(name.clone()) {
1950                return (pos | 0xC000).compose(self);
1951            }
1952
1953            // So we don’t know the name. Try inserting it into the
1954            // compressor. If we can’t insert anymore, just write out what’s
1955            // left and return.
1956            if !self.insert(self.target.as_ref().len()) {
1957                for label in &mut name {
1958                    label.compose(self)?;
1959                }
1960                return Ok(());
1961            }
1962
1963            // Advance to the parent.
1964            let label = name.next().unwrap();
1965            label.compose(self)?;
1966        }
1967    }
1968
1969    fn can_compress(&self) -> bool {
1970        true
1971    }
1972}
1973
1974impl<Target: Truncate> Truncate for StaticCompressor<Target> {
1975    fn truncate(&mut self, len: usize) {
1976        self.target.truncate(len);
1977        if len < 0xC000 {
1978            let len = len as u16;
1979            for i in 0..self.len {
1980                if self.entries[i] >= len {
1981                    self.len = i;
1982                    break;
1983                }
1984            }
1985        }
1986    }
1987}
1988
1989//------------ TreeCompressor ------------------------------------------------
1990
1991/// A domain name compressor that uses a tree.
1992///
1993/// This type wraps around an octets builder and implements domain name
1994/// compression for it. It stores the position of any domain name it has seen
1995/// in a binary tree.
1996///
1997/// The position of a domain name is calculated relative to the beginning of
1998/// the underlying octets builder. This means that this builder must represent
1999/// the message only. This means that if you are using the [`StreamTarget`],
2000/// you need to place it inside this type, _not_ the other way around.
2001///
2002/// [`StreamTarget`]: struct.StreamTarget.html
2003#[cfg(feature = "std")]
2004#[derive(Clone, Debug)]
2005pub struct TreeCompressor<Target> {
2006    /// The underlying octetsbuilder.
2007    target: Target,
2008
2009    /// The topmost node of our tree.
2010    start: Node,
2011}
2012
2013/// A node in our tree.
2014///
2015/// The tree is a bit odd. It follows the labels of the domain names from the
2016/// root towards the left. The root node is for the root label. It contains a
2017/// map that maps all the labels encountered to the immediate left of the
2018/// name traced by this path through the tree to a node for the name resulting
2019/// by adding this label to the name constructed so far.
2020///
2021/// Each node also contains the position of that name in the message.
2022#[cfg(feature = "std")]
2023#[derive(Clone, Debug, Default)]
2024struct Node {
2025    /// The labels immediately to the left of this name and their nodes.
2026    parents: HashMap<Array<64>, Self>,
2027
2028    /// The position of this name in the message.
2029    value: Option<u16>,
2030}
2031
2032#[cfg(feature = "std")]
2033impl Node {
2034    fn drop_above(&mut self, len: u16) {
2035        self.value = match self.value {
2036            Some(value) if value < len => Some(value),
2037            _ => None,
2038        };
2039        self.parents
2040            .values_mut()
2041            .for_each(|node| node.drop_above(len))
2042    }
2043}
2044
2045#[cfg(feature = "std")]
2046impl<Target> TreeCompressor<Target> {
2047    /// Creates a new compressor from an underlying octets builder.
2048    pub fn new(target: Target) -> Self {
2049        TreeCompressor {
2050            target,
2051            start: Default::default(),
2052        }
2053    }
2054
2055    /// Returns a reference to the underlying octets builder.
2056    pub fn as_target(&self) -> &Target {
2057        &self.target
2058    }
2059
2060    /// Converts the compressor into the underlying octets builder.
2061    pub fn into_target(self) -> Target {
2062        self.target
2063    }
2064
2065    /// Returns an octets slice of the data.
2066    pub fn as_slice(&self) -> &[u8]
2067    where
2068        Target: AsRef<[u8]>,
2069    {
2070        self.target.as_ref()
2071    }
2072
2073    /// Returns an mutable octets slice of the data.
2074    pub fn as_slice_mut(&mut self) -> &mut [u8]
2075    where
2076        Target: AsMut<[u8]>,
2077    {
2078        self.target.as_mut()
2079    }
2080
2081    fn get<'a, N: Iterator<Item = &'a Label> + Clone>(
2082        &self,
2083        name: N,
2084    ) -> Option<u16> {
2085        let mut node = &self.start;
2086        for label in name {
2087            if label.is_root() {
2088                return node.value;
2089            }
2090            node = node.parents.get(label.as_ref())?;
2091        }
2092        None
2093    }
2094
2095    fn insert<'a, N: Iterator<Item = &'a Label> + Clone>(
2096        &mut self,
2097        name: N,
2098        pos: usize,
2099    ) -> bool {
2100        if pos >= 0xC000 {
2101            return false;
2102        }
2103        let pos = pos as u16;
2104        let mut node = &mut self.start;
2105        for label in name {
2106            if label.is_root() {
2107                node.value = Some(pos);
2108                break;
2109            }
2110            node = node
2111                .parents
2112                .entry(label.as_ref().try_into().unwrap())
2113                .or_default();
2114        }
2115        true
2116    }
2117}
2118
2119//--- AsRef, AsMut, and OctetsBuilder
2120
2121#[cfg(feature = "std")]
2122impl<Target: AsRef<[u8]>> AsRef<[u8]> for TreeCompressor<Target> {
2123    fn as_ref(&self) -> &[u8] {
2124        self.as_slice()
2125    }
2126}
2127
2128#[cfg(feature = "std")]
2129impl<Target: AsMut<[u8]>> AsMut<[u8]> for TreeCompressor<Target> {
2130    fn as_mut(&mut self) -> &mut [u8] {
2131        self.as_slice_mut()
2132    }
2133}
2134
2135#[cfg(feature = "std")]
2136impl<Target: OctetsBuilder> OctetsBuilder for TreeCompressor<Target> {
2137    type AppendError = Target::AppendError;
2138
2139    fn append_slice(
2140        &mut self,
2141        slice: &[u8],
2142    ) -> Result<(), Self::AppendError> {
2143        self.target.append_slice(slice)
2144    }
2145}
2146
2147#[cfg(feature = "std")]
2148impl<Target: Composer> Composer for TreeCompressor<Target> {
2149    fn append_compressed_dname<N: ToDname + ?Sized>(
2150        &mut self,
2151        name: &N,
2152    ) -> Result<(), Self::AppendError> {
2153        let mut name = name.iter_labels().peekable();
2154
2155        loop {
2156            // If the parent is root, just write that and return.
2157            // Because we do that, there will always be a label left here.
2158            if let Some(label) = name.peek() {
2159                if label.is_root() {
2160                    label.compose(self)?;
2161                    return Ok(());
2162                }
2163            }
2164
2165            // If we already know this name, append it as a compressed label.
2166            if let Some(pos) = self.get(name.clone()) {
2167                return (pos | 0xC000).compose(self);
2168            }
2169
2170            // So we don’t know the name. Try inserting it into the
2171            // compressor. If we can’t insert anymore, just write out what’s
2172            // left and return.
2173            if !self.insert(name.clone(), self.target.as_ref().len()) {
2174                for label in &mut name {
2175                    label.compose(self)?;
2176                }
2177                return Ok(());
2178            }
2179
2180            // Advance to the parent. If the parent is root, just write that
2181            // and return. Because we do that, there will always be a label
2182            // left here.
2183            let label = name.next().unwrap();
2184            label.compose(self)?;
2185        }
2186    }
2187
2188    fn can_compress(&self) -> bool {
2189        true
2190    }
2191}
2192
2193#[cfg(feature = "std")]
2194impl<Target: Composer> Truncate for TreeCompressor<Target> {
2195    fn truncate(&mut self, len: usize) {
2196        self.target.truncate(len);
2197        if len < 0xC000 {
2198            self.start.drop_above(len as u16)
2199        }
2200    }
2201}
2202
2203//============ Errors ========================================================
2204
2205#[derive(Clone, Copy, Debug)]
2206pub enum PushError {
2207    CountOverflow,
2208    ShortBuf,
2209}
2210
2211impl<T: Into<ShortBuf>> From<T> for PushError {
2212    fn from(_: T) -> Self {
2213        Self::ShortBuf
2214    }
2215}
2216
2217impl fmt::Display for PushError {
2218    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
2219        match *self {
2220            PushError::CountOverflow => f.write_str("counter overflow"),
2221            PushError::ShortBuf => ShortBuf.fmt(f),
2222        }
2223    }
2224}
2225
2226#[cfg(feature = "std")]
2227impl std::error::Error for PushError {}
2228
2229//============ Testing =======================================================
2230
2231#[cfg(test)]
2232#[cfg(feature = "std")]
2233mod test {
2234    use super::*;
2235    use crate::base::Serial;
2236    use crate::base::Ttl;
2237    use crate::base::{iana::Rtype, opt, Dname};
2238    use crate::rdata::{Ns, Soa, A};
2239    use core::str::FromStr;
2240    use std::vec::Vec;
2241
2242    #[test]
2243    fn message_builder() {
2244        // Make a domain name we can use later on.
2245        let name = Dname::<Vec<u8>>::from_str("example.com").unwrap();
2246
2247        // Create a message builder wrapping a compressor wrapping a stream
2248        // target.
2249        let mut msg = MessageBuilder::from_target(StaticCompressor::new(
2250            StreamTarget::new_vec(),
2251        ))
2252        .unwrap();
2253
2254        // Set the RD bit in the header and proceed to the question section.
2255        msg.header_mut().set_rd(true);
2256        let mut msg = msg.question();
2257
2258        // Add a question and proceed to the answer section.
2259        msg.push((&name, Rtype::A)).unwrap();
2260        let mut msg = msg.answer();
2261
2262        // Add two answer and proceed to the additional sections
2263        msg.push((&name, 86400, A::from_octets(192, 0, 2, 1)))
2264            .unwrap();
2265        msg.push((&name, 86400, A::from_octets(192, 0, 2, 2)))
2266            .unwrap();
2267
2268        // Add an authority
2269        let mut msg = msg.authority();
2270        msg.push((&name, 0, Ns::from(name.clone()))).unwrap();
2271
2272        // Add additional
2273        let mut msg = msg.additional();
2274        msg.push((&name, 86400, A::from_octets(192, 0, 2, 1)))
2275            .unwrap();
2276
2277        // Convert the builder into the actual message.
2278        let target = msg.finish().into_target();
2279
2280        // Reparse message and check contents
2281        let msg = Message::from_octets(target.as_dgram_slice()).unwrap();
2282        let q = msg.first_question().unwrap();
2283        assert_eq!(q.qname(), &name);
2284        assert_eq!(q.qtype(), Rtype::A);
2285
2286        let section = msg.answer().unwrap();
2287        let mut records = section.limit_to::<A>();
2288        assert_eq!(
2289            records.next().unwrap().unwrap().data(),
2290            &A::from_octets(192, 0, 2, 1)
2291        );
2292        assert_eq!(
2293            records.next().unwrap().unwrap().data(),
2294            &A::from_octets(192, 0, 2, 2)
2295        );
2296
2297        let section = msg.authority().unwrap();
2298        let mut records = section.limit_to::<Ns<_>>();
2299        let rr = records.next().unwrap().unwrap();
2300        assert_eq!(rr.owner(), &name);
2301        assert_eq!(rr.data().nsdname(), &name);
2302
2303        let section = msg.additional().unwrap();
2304        let mut records = section.limit_to::<A>();
2305        let rr = records.next().unwrap().unwrap();
2306        assert_eq!(rr.owner(), &name);
2307        assert_eq!(rr.data(), &A::from_octets(192, 0, 2, 1));
2308    }
2309
2310    #[test]
2311    fn opt_builder() {
2312        let mut msg = MessageBuilder::new_vec().additional();
2313
2314        // Add an OPT record.
2315        let nsid = opt::nsid::Nsid::from_octets(&b"example"[..]).unwrap();
2316        msg.opt(|o| {
2317            o.set_udp_payload_size(4096);
2318            o.push(&nsid)?;
2319            Ok(())
2320        })
2321        .unwrap();
2322
2323        let msg = msg.finish();
2324        println!("{:?}", msg);
2325        let msg = Message::from_octets(msg).unwrap();
2326        let opt = msg.opt().unwrap();
2327
2328        // Check options
2329        assert_eq!(opt.udp_payload_size(), 4096);
2330        let mut opts = opt.opt().iter::<opt::nsid::Nsid<_>>();
2331        assert_eq!(opts.next(), Some(Ok(nsid)));
2332    }
2333
2334    fn create_compressed<T: Composer>(target: T) -> T
2335    where
2336        T::AppendError: fmt::Debug,
2337    {
2338        let mut msg = MessageBuilder::from_target(target).unwrap().question();
2339        msg.header_mut().set_rcode(Rcode::NXDomain);
2340        msg.header_mut().set_rd(true);
2341        msg.header_mut().set_ra(true);
2342        msg.header_mut().set_qr(true);
2343
2344        msg.push((&"example".parse::<Dname<Vec<u8>>>().unwrap(), Rtype::Ns))
2345            .unwrap();
2346        let mut msg = msg.authority();
2347
2348        let mname: Dname<Vec<u8>> = "a.root-servers.net".parse().unwrap();
2349        let rname = "nstld.verisign-grs.com".parse().unwrap();
2350        msg.push((
2351            Dname::root_slice(),
2352            86390,
2353            Soa::new(
2354                mname,
2355                rname,
2356                Serial(2020081701),
2357                Ttl::from_secs(1800),
2358                Ttl::from_secs(900),
2359                Ttl::from_secs(604800),
2360                Ttl::from_secs(86400),
2361            ),
2362        ))
2363        .unwrap();
2364        msg.finish()
2365    }
2366
2367    #[test]
2368    fn compressor() {
2369        // An example negative response to `example. NS` with an SOA to test
2370        // various compressed name situations.
2371        let expect = &[
2372            0x00, 0x00, 0x81, 0x83, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00,
2373            0x00, 0x07, 0x65, 0x78, 0x61, 0x6d, 0x70, 0x6c, 0x65, 0x00, 0x00,
2374            0x02, 0x00, 0x01, 0x00, 0x00, 0x06, 0x00, 0x01, 0x00, 0x01, 0x51,
2375            0x76, 0x00, 0x40, 0x01, 0x61, 0x0c, 0x72, 0x6f, 0x6f, 0x74, 0x2d,
2376            0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x03, 0x6e, 0x65, 0x74,
2377            0x00, 0x05, 0x6e, 0x73, 0x74, 0x6c, 0x64, 0x0c, 0x76, 0x65, 0x72,
2378            0x69, 0x73, 0x69, 0x67, 0x6e, 0x2d, 0x67, 0x72, 0x73, 0x03, 0x63,
2379            0x6f, 0x6d, 0x00, 0x78, 0x68, 0x00, 0x25, 0x00, 0x00, 0x07, 0x08,
2380            0x00, 0x00, 0x03, 0x84, 0x00, 0x09, 0x3a, 0x80, 0x00, 0x01, 0x51,
2381            0x80,
2382        ];
2383
2384        let msg = create_compressed(StaticCompressor::new(Vec::new()));
2385        assert_eq!(&expect[..], msg.as_ref());
2386
2387        let msg = create_compressed(TreeCompressor::new(Vec::new()));
2388        assert_eq!(&expect[..], msg.as_ref());
2389    }
2390}