toluol_proto/
lib.rs

1//! `toluol-proto` provides the definition of the DNS protocol's data types as well as the means to
2//! de-/serialize them from/to the wire format. In simpler terms, you can construct, encode, and
3//! decode DNS queries and responses with it.
4//!
5//! It is used as the backend for [`toluol`], a DNS client that aims to
6//! replace `dig`, but you can use this library on its own as well. It is possible to compile it to
7//! WASM, so you can even make DNS queries from the browser with it (using DNS over HTTPS).
8//!
9//! # Basic usage example
10//! ```rust
11//! use toluol_proto::{EdnsConfig, HeaderFlags, Message, Name, Opcode, RecordType};
12//!
13//! let flags = HeaderFlags { aa: false, tc: false, rd: true, ra: false, ad: true, cd: true };
14//! let msg = Message::new_query(
15//!     Name::from_ascii("example.com").unwrap(),
16//!     RecordType::A,
17//!     Opcode::QUERY,
18//!     flags,
19//!     Some(EdnsConfig {
20//!         do_flag: false,
21//!         bufsize: 4096,
22//!         client_cookie: None,
23//!     }),
24//! ).unwrap();
25//! let _encoded = msg.encode().unwrap();
26//! ```
27//!
28//! If you're also looking for utilities to actually send and receive DNS queries and responses,
29//! please take a look at [`toluol`].
30//!
31//! # Usage note
32//! You can construct most structs directly, without using any `new()` method. In some cases, this
33//! can lead to inconsistencies, e.g. manually creating a [`Message`] where the record counts in the
34//! header don't match the actual number of records.
35//!
36//! In these cases, you should prefer using the appropriate constructor of the struct (if there is
37//! none, please file a bug). However, this library does not force you to do so, so that you have
38//! as much freedom using it as possible. It won't stop you if you really want to create
39//! inconsistent messages, for whatever reason.
40//!
41//! [`toluol`]: https://docs.rs/toluol
42
43use std::cmp::max;
44use std::collections::HashMap;
45use std::fmt::{self, Display};
46use std::io::{Cursor, Read, Write};
47
48use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
49use owo_colors::OwoColorize;
50use rand::Rng;
51use rdata::opt::OptionCode;
52use repr_with_fallback::repr_with_fallback;
53#[cfg(feature = "serde")]
54use serde::Serialize;
55use strum_macros::EnumString;
56
57// TODO put the dnssec module behind a feature?
58pub mod dnssec;
59pub mod error;
60pub mod name;
61pub mod rdata;
62
63use error::{DnssecError, EncodeError, ParseError, ToluolError};
64use rdata::{RdataTrait, OPT};
65
66pub use name::Name;
67pub use rdata::Rdata;
68
69/// Represents a DNS OpCode.
70///
71/// See [here](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-5) for
72/// further information.
73#[cfg_attr(feature = "serde", derive(Serialize))]
74#[derive(PartialEq, Eq, Copy, Clone, Debug)]
75pub enum Opcode {
76    QUERY,
77    IQUERY,
78    STATUS,
79    NOTIFY,
80    UPDATE,
81    DSO,
82}
83
84/// Represents a DNS RCODE, including those introduced by EDNS.
85///
86/// See
87/// [here](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-6)
88/// for further information.
89#[cfg_attr(feature = "serde", derive(Serialize))]
90#[derive(PartialEq, Eq, Copy, Clone, Debug)]
91#[non_exhaustive]
92pub enum RCode {
93    NOERROR,
94    FORMERR,
95    SERVFAIL,
96    NXDOMAIN,
97    NOTIMP,
98    REFUSED,
99    YXDOMAIN,
100    YXRRSET,
101    NXRRSET,
102    NOTAUTH,
103    NOTZONE,
104    DSOTYPENI,
105    BADVERSBADSIG,
106    BADKEY,
107    BADTIME,
108    BADMODE,
109    BADNAME,
110    BADALG,
111    BADTRUNC,
112    BADCOOKIE,
113    // TODO Unknown(u16) ?
114}
115
116repr_with_fallback! {
117    /// Represents a DNS TYPE.
118    ///
119    /// See the documentation in the [`rdata`] module for explanations of the different types.
120    ///
121    /// This enum is non-exhaustive, see
122    /// [here](https://en.wikipedia.org/wiki/List_of_DNS_record_types) for a more comprehensive
123    /// overview.
124    #[cfg_attr(feature = "serde", derive(Serialize))]
125    #[derive(PartialEq, Eq, Copy, Clone, EnumString, Debug)]
126    #[non_exhaustive]
127    pub enum RecordType {
128        A = 1,
129        NS = 2,
130        CNAME = 5,
131        SOA = 6,
132        PTR = 12,
133        HINFO = 13,
134        MX = 15,
135        TXT = 16,
136        RP = 17,
137        // TODO: SIG (24) (should have the same wire format as RRSIG)
138        // TODO: KEY (25) (should have the same wire format as DNSKEY)
139        AAAA = 28,
140        LOC = 29,
141        SRV = 33,
142        NAPTR = 35,
143        CERT = 37,
144        DNAME = 39,
145        OPT = 41,
146        DS = 43,
147        SSHFP = 44,
148        // TODO: IPSECKEY (45)
149        RRSIG = 46,
150        NSEC = 47,
151        DNSKEY = 48,
152        // TODO: DHCID (49)
153        NSEC3 = 50,
154        NSEC3PARAM = 51,
155        TLSA = 52,
156        // TODO: SMIMEA (53)
157        // TODO: HIP (55)
158        // TODO: CDNSKEY (60)
159        OPENPGPKEY = 61,
160        // TODO: HTTPS (65)
161        // TODO: TKEY (249)
162        // TODO: TSIG (250)
163        CAA = 257,
164        // TODO: TA (32768)
165        // TODO: DLV (32769)
166        Unknown(u16),
167    }
168}
169
170/// Represents a DNS CLASS.
171///
172/// Other classes than `IN` and `ANY` are included only for completeness and historical reasons.
173///
174/// See [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035) for further information.
175#[cfg_attr(feature = "serde", derive(Serialize))]
176#[derive(PartialEq, Eq, Copy, Clone, Debug)]
177pub enum Class {
178    IN,
179    CH,
180    HS,
181    NONE,
182    /// See also [RFC 8482](https://www.rfc-editor.org/rfc/rfc8482).
183    ANY,
184}
185
186/// Represents the flags of a [`Header`].
187#[cfg_attr(feature = "serde", derive(Serialize))]
188#[derive(PartialEq, Eq, Copy, Clone, Debug)]
189pub struct HeaderFlags {
190    /// authoritative answer (valid in responses only)
191    /// [\[RFC 1035\]](https://www.rfc-editor.org/rfc/rfc1035)
192    pub aa: bool,
193    /// truncated (set on all truncated messages except last one)
194    /// [\[RFC 1035\]](https://www.rfc-editor.org/rfc/rfc1035)
195    pub tc: bool,
196    /// recursion desired (copied in answer if supported and accepted)
197    /// [\[RFC 1035\]](https://www.rfc-editor.org/rfc/rfc1035)
198    pub rd: bool,
199    /// valid in responses, indicating recursive query support in the name server
200    /// [\[RFC 1035\]](https://www.rfc-editor.org/rfc/rfc1035)
201    pub ra: bool,
202    /// For queries: indicates interest in the `ad` bit of the upcoming response; for responses:
203    /// indicates that the resolver side considers all resource records in the Answer section and
204    /// relevant negative response resource records in the Authority section to be authentic.
205    /// [\[RFC 4035\]](https://www.rfc-editor.org/rfc/rfc4035),
206    /// [\[RFC 6840\]](https://www.rfc-editor.org/rfc/rfc6840)
207    pub ad: bool,
208    /// disable signature validation in a security-aware name server's processing of a particular query
209    /// [\[RFC 4035\]](https://www.rfc-editor.org/rfc/rfc4035),
210    /// [\[RFC 6840\]](https://www.rfc-editor.org/rfc/rfc6840)
211    pub cd: bool,
212}
213
214/// Represents a DNS header.
215///
216/// The general format of a header is defined in [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035).
217#[cfg_attr(feature = "serde", derive(Serialize))]
218#[derive(PartialEq, Eq, Clone, Debug)]
219pub struct Header {
220    /// Supplied by questioner and reflected back unchanged by responder.
221    pub msg_id: u16,
222    /// False for queries, true for responses.
223    pub qr: bool,
224    /// The [`Opcode`] of the message.
225    pub opcode: Opcode,
226    /// The [`HeaderFlags`] of the message.
227    pub flags: HeaderFlags,
228    /// For queries: [`None`]. For responses: the return/status code of the server.
229    pub rcode: Option<RCode>,
230    /// The number of questions.
231    pub qdcount: u16,
232    /// The number of resource records.
233    pub ancount: u16,
234    /// The number of name server resource records.
235    pub nscount: u16,
236    /// The number of additional resource records.
237    pub arcount: u16,
238}
239
240/// Represents a DNS question, i.e. an entry in the question section of a DNS message.
241///
242/// See [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035) for further information.
243#[cfg_attr(feature = "serde", derive(Serialize))]
244#[derive(PartialEq, Eq, Clone, Debug)]
245pub struct Question {
246    /// The [`Name`] to query for.
247    pub qname: Name,
248    /// The [`RecordType`] to query for.
249    pub qtype: RecordType,
250    /// The query [`Class`].
251    pub qclass: Class,
252}
253
254/// Represents a DNS record, i.e. an entry in the answer, authority or additional section of a DNS
255/// message.
256///
257/// See [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035) for further information.
258#[cfg_attr(feature = "serde", derive(Serialize))]
259#[cfg_attr(feature = "serde", serde(untagged))]
260#[derive(PartialEq, Eq, Clone, Debug)]
261pub enum Record {
262    OPT(OptRecord),
263    NONOPT(NonOptRecord),
264}
265
266/// Flags for an [`OptRecord`].
267///
268/// See [RFC 6891](https://www.rfc-editor.org/rfc/rfc6891#section-6) as well as
269/// <https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-13> for
270/// further information.
271#[cfg_attr(feature = "serde", derive(Serialize))]
272#[derive(PartialEq, Eq, Copy, Clone, Debug)]
273pub enum OptFlags {
274    /// Indicates to the server that the resolver is able to accept DNSSEC security records.
275    /// [\[RFC 3225\]](https://www.rfc-editor.org/rfc/rfc3225)
276    DO,
277}
278
279/// EDNS parameters.
280pub struct EdnsConfig {
281    /// Indicates DNSSEC support, i.e. whether the server should send appropiate DNSSEC records.
282    pub do_flag: bool,
283    /// The payload size that gets sent in the `OPT` record.
284    pub bufsize: u16,
285    /// May be [`None`] to indicate no client cookie should be set.
286    ///
287    /// See [RFC 7873](https://www.rfc-editor.org/rfc/rfc7873.html) for more.
288    pub client_cookie: Option<[u8; 8]>,
289    // TODO: support padding?
290}
291
292/// The `OPT` variant of [`Record`].
293///
294/// See [RFC 6891](https://www.rfc-editor.org/rfc/rfc6891#section-6) for further information.
295#[cfg_attr(feature = "serde", derive(Serialize))]
296#[derive(PartialEq, Eq, Clone, Debug)]
297pub struct OptRecord {
298    /// Must be [`Name::root()`].
299    pub owner: Name,
300    /// The number of octets of the largest UDP payload that can be reassembled and delivered in the
301    /// requestor's network stack.
302    pub payload_size: u16,
303    /// `None` for queries. For responses, this is always the correct [`RCode`], i.e. the lower four
304    /// bits from the header are included.
305    pub rcode: Option<RCode>,
306    /// Almost always zero.
307    pub edns_version: u8,
308    /// A list of [`OptFlags`] (may be empty).
309    pub flags: Vec<OptFlags>,
310    // rdlength omitted as rdata knows its own length
311    #[cfg_attr(feature = "serde", serde(skip))]
312    encoded_rdata: Vec<u8>, // needed for encoding
313    rdata: Rdata, // this is of type Rdata and not OPT so that it nicely mirrors NonOptRecord
314}
315
316/// The `NONOPT` variant of [`Record`].
317///
318/// See [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035) for further information.
319#[cfg_attr(feature = "serde", derive(Serialize))]
320#[derive(PartialEq, Eq, Clone, Debug)]
321pub struct NonOptRecord {
322    /// The [`Name`] that this record is for.
323    pub owner: Name,
324    /// The type of this record.
325    pub rtype: RecordType,
326    /// The class of this record (will almost always be [`Class::IN`]).
327    pub class: Class,
328    /// The amount of seconds this record may be cached for.
329    pub ttl: u32,
330    // rdlength omitted as rdata knows its own length
331    #[cfg_attr(feature = "serde", serde(skip))]
332    encoded_rdata: Vec<u8>, // needed for encoding and DNSSEC
333    rdata: Rdata,
334}
335
336/// Represents a DNS message.
337///
338/// See [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035) for further information.
339#[cfg_attr(feature = "serde", derive(Serialize))]
340#[derive(PartialEq, Eq, Clone, Debug)]
341pub struct Message {
342    /// The message header.
343    pub header: Header,
344    /// The list of questions.
345    pub questions: Vec<Question>,
346    /// The list of resource records.
347    pub answers: Vec<Record>,
348    /// The list of name server resource records.
349    pub authoritative_answers: Vec<Record>,
350    /// The list of additional resource records.
351    pub additional_answers: Vec<Record>,
352}
353
354impl Opcode {
355    /// Encodes a `Opcode` as a byte.
356    pub fn encode(&self) -> u8 {
357        match self {
358            Opcode::QUERY => 0,
359            Opcode::IQUERY => 1,
360            Opcode::STATUS => 2,
361            Opcode::NOTIFY => 4,
362            Opcode::UPDATE => 5,
363            Opcode::DSO => 6,
364        }
365    }
366
367    /// Parses an encoded `Opcode` from a byte.
368    ///
369    /// Returns an error if the given byte does not represent a valid DNS OpCode.
370    pub fn parse(val: u8) -> Result<Opcode, ParseError> {
371        Ok(match val {
372            0 => Opcode::QUERY,
373            1 => Opcode::IQUERY,
374            2 => Opcode::STATUS,
375            4 => Opcode::NOTIFY,
376            5 => Opcode::UPDATE,
377            6 => Opcode::DSO,
378            x => return Err(ParseError::InvalidOpcode(x)),
379        })
380    }
381}
382
383impl Display for Opcode {
384    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
385        write!(f, "{:?}", self)
386    }
387}
388
389impl RCode {
390    /// Encodes an `RCode` as a byte (actually only the lower four bits are used).
391    ///
392    /// Note that for RCODEs `BADVERSBADSIG` and following only the lower four bits are encoded;
393    /// the upper eight bits need to be encoded in an OPT record in the additional section of the
394    /// DNS message.
395    pub fn encode(&self) -> u8 {
396        match self {
397            RCode::NOERROR => 0,
398            RCode::FORMERR => 1,
399            RCode::SERVFAIL => 2,
400            RCode::NXDOMAIN => 3,
401            RCode::NOTIMP => 4,
402            RCode::REFUSED => 5,
403            RCode::YXDOMAIN => 6,
404            RCode::YXRRSET => 7,
405            RCode::NXRRSET => 8,
406            RCode::NOTAUTH => 9,
407            RCode::NOTZONE => 10,
408            RCode::DSOTYPENI => 11,
409            RCode::BADVERSBADSIG => 16 & 0b1111,
410            RCode::BADKEY => 17 & 0b1111,
411            RCode::BADTIME => 18 & 0b1111,
412            RCode::BADMODE => 19 & 0b1111,
413            RCode::BADNAME => 20 & 0b1111,
414            RCode::BADALG => 21 & 0b1111,
415            RCode::BADTRUNC => 22 & 0b1111,
416            RCode::BADCOOKIE => 23 & 0b1111,
417        }
418    }
419
420    /// Parses an encoded `RCode` from a twelve bit value. If EDNS is used, the upper eight bits
421    /// are stored in the OPT entry of the additional section and the lower four bits are stored in
422    /// the [`Header`].
423    ///
424    /// Returns an error if the given value does not represent a valid DNS RCODE.
425    pub fn parse(val: u16) -> Result<RCode, ParseError> {
426        Ok(match val {
427            0 => RCode::NOERROR,
428            1 => RCode::FORMERR,
429            2 => RCode::SERVFAIL,
430            3 => RCode::NXDOMAIN,
431            4 => RCode::NOTIMP,
432            5 => RCode::REFUSED,
433            6 => RCode::YXDOMAIN,
434            7 => RCode::YXRRSET,
435            8 => RCode::NXRRSET,
436            9 => RCode::NOTAUTH,
437            10 => RCode::NOTZONE,
438            11 => RCode::DSOTYPENI,
439            16 => RCode::BADVERSBADSIG,
440            17 => RCode::BADKEY,
441            18 => RCode::BADTIME,
442            19 => RCode::BADMODE,
443            20 => RCode::BADNAME,
444            21 => RCode::BADALG,
445            22 => RCode::BADTRUNC,
446            23 => RCode::BADCOOKIE,
447            x => return Err(ParseError::InvalidRcode(x)),
448        })
449    }
450}
451
452impl Display for RCode {
453    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
454        write!(f, "{:?}", self)
455    }
456}
457
458impl Display for RecordType {
459    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
460        match self {
461            RecordType::Unknown(x) => write!(f, "TYPE{}", x),
462            _ => write!(f, "{:?}", self),
463        }
464    }
465}
466
467impl Class {
468    /// Encodes a `Class` as a two-byte value.
469    pub fn encode(&self) -> u16 {
470        match self {
471            Class::IN => 1,
472            Class::CH => 3,
473            Class::HS => 4,
474            Class::NONE => 254,
475            Class::ANY => 255,
476        }
477    }
478
479    /// Parses an encoded `Class` from a two-byte value.
480    ///
481    /// Returns an error if the given value does not represent a valid DNS CLASS.
482    pub fn parse(val: u16) -> Result<Class, ParseError> {
483        Ok(match val {
484            1 => Class::IN,
485            3 => Class::CH,
486            4 => Class::HS,
487            254 => Class::NONE,
488            255 => Class::ANY,
489            x => return Err(ParseError::InvalidClass(x)),
490        })
491    }
492}
493
494impl Display for Class {
495    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
496        write!(f, "{:?}", self)
497    }
498}
499
500impl HeaderFlags {
501    /// Creates a `HeaderFlags` struct from bitflags as they would appear in the second 16-octet
502    /// line of a [`Header`].
503    pub fn from_flags(flags: u16) -> Self {
504        Self {
505            aa: (flags & (1 << 10)) != 0,
506            tc: (flags & (1 << 9)) != 0,
507            rd: (flags & (1 << 8)) != 0,
508            ra: (flags & (1 << 8)) != 0,
509            ad: (flags & (1 << 5)) != 0,
510            cd: (flags & (1 << 4)) != 0,
511        }
512    }
513
514    /// Returns a u16 representing bitflags as they would appear in the second 16-octet line of a
515    /// [`Header`].
516    pub fn as_flags(&self) -> u16 {
517        let aa = if self.aa { 1 } else { 0 };
518        let tc = if self.tc { 1 } else { 0 };
519        let rd = if self.rd { 1 } else { 0 };
520        let ra = if self.ra { 1 } else { 0 };
521        let ad = if self.ad { 1 } else { 0 };
522        let cd = if self.cd { 1 } else { 0 };
523        (aa << 10) + (tc << 9) + (rd << 8) + (ra << 7) + (ad << 5) + (cd << 4)
524    }
525}
526
527impl Header {
528    /// Creates a header for a DNS response message.
529    ///
530    /// See [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035) and
531    /// [here](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-12) for
532    /// information about the parameters.
533    ///
534    /// `qdcount`, `ancount`, `nscount` and `arcount` are grouped in that order in the `counts` parameter.
535    pub fn new_response_header(
536        msg_id: u16,
537        opcode: Opcode,
538        flags: HeaderFlags,
539        rcode: RCode,
540        counts: [u16; 4],
541    ) -> Self {
542        Header {
543            msg_id,
544            qr: true,
545            opcode,
546            flags,
547            rcode: Some(rcode),
548            qdcount: counts[0],
549            ancount: counts[1],
550            nscount: counts[2],
551            arcount: counts[3],
552        }
553    }
554
555    /// Creates a header for a DNS query message.
556    ///
557    /// If the query includes an [`OPT`](rdata::opt::OPT) record, `edns` must be `true`.
558    ///
559    /// See [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035) and
560    /// [here](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-12) for
561    /// information about the other parameters.
562    pub fn new_query_header(
563        msg_id: u16,
564        opcode: Opcode,
565        flags: HeaderFlags,
566        edns: bool,
567        qdcount: u16,
568    ) -> Result<Self, EncodeError> {
569        if flags.aa || flags.ra {
570            Err(EncodeError::AaOrRaInQuery)
571        } else {
572            Ok(Header {
573                msg_id,
574                qr: false,
575                opcode,
576                flags,
577                rcode: None,
578                qdcount,
579                ancount: 0,
580                nscount: 0,
581                arcount: if edns { 1 } else { 0 },
582            })
583        }
584    }
585
586    /// Encodes a `Header` as a series of bytes.
587    ///
588    /// Returns an error if a method defined in [`byteorder::WriteBytesExt`] returns an error.
589    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
590        let mut buf = Vec::new();
591        self.encode_into(&mut buf)?;
592        Ok(buf)
593    }
594
595    /// The same as [`encode()`](Self::encode()), but encoded bytes are appended to the given writer
596    /// instead of to a newly allocated one.
597    pub fn encode_into(&self, buf: &mut impl Write) -> Result<(), EncodeError> {
598        let qr = if self.qr { 1u16 } else { 0u16 };
599        let opcode = self.opcode.encode() as u16;
600        let rcode = match &self.rcode {
601            Some(val) => val.encode() as u16,
602            None => 0u16,
603        };
604
605        let line_two = (qr << 15) + (opcode << 11) + self.flags.as_flags() + rcode;
606        buf.write_u16::<NetworkEndian>(self.msg_id)?;
607        buf.write_u16::<NetworkEndian>(line_two)?;
608        buf.write_u16::<NetworkEndian>(self.qdcount)?;
609        buf.write_u16::<NetworkEndian>(self.ancount)?;
610        buf.write_u16::<NetworkEndian>(self.nscount)?;
611        buf.write_u16::<NetworkEndian>(self.arcount)?;
612
613        Ok(())
614    }
615
616    /// Parses an encoded `Header` from a series of bytes.
617    ///
618    /// Returns an error if [`Opcode::parse()`], [`RCode::parse()`] or a method defined in
619    /// [`byteorder::ReadBytesExt`] return an error.
620    pub fn parse(header: &mut Cursor<&[u8]>) -> Result<Self, ParseError> {
621        let msg_id = header.read_u16::<NetworkEndian>()?;
622        let line_two = header.read_u16::<NetworkEndian>()?;
623        let qr = (line_two & (1 << 15)) >> 15;
624        let opcode = Opcode::parse(((line_two & (0b1111 << 11)) >> 11) as u8)?;
625        let flags = HeaderFlags::from_flags(line_two & 0b0000011110110000);
626        let rcode = RCode::parse(line_two & 0b1111)?;
627
628        Ok(Header {
629            msg_id,
630            qr: qr != 0,
631            opcode,
632            flags,
633            rcode: if qr != 0 { Some(rcode) } else { None },
634            qdcount: header.read_u16::<NetworkEndian>()?,
635            ancount: header.read_u16::<NetworkEndian>()?,
636            nscount: header.read_u16::<NetworkEndian>()?,
637            arcount: header.read_u16::<NetworkEndian>()?,
638        })
639    }
640
641    /// Creates a string containing information (id, opcode, rcode if applicable, flags) about the
642    /// header.
643    pub fn info_str(&self) -> String {
644        let mut s = String::new();
645        if let Some(rcode) = self.rcode {
646            s.push_str(
647                format!(
648                    "id: {}, opcode: {}, rcode: {}, flags: ",
649                    self.msg_id, self.opcode, rcode
650                )
651                .as_str(),
652            );
653        } else {
654            s.push_str(format!("id: {}, opcode: {}, flags: ", self.msg_id, self.opcode).as_str());
655        }
656        if self.flags.aa {
657            s.push_str("aa ")
658        }
659        if self.flags.tc {
660            s.push_str("tc ")
661        }
662        if self.flags.rd {
663            s.push_str("rd ")
664        }
665        if self.flags.ra {
666            s.push_str("ra ")
667        }
668        if self.flags.ad {
669            s.push_str("ad ")
670        }
671        if self.flags.cd {
672            s.push_str("cd ")
673        }
674        s.remove(s.len() - 1); // remove last ' '
675        s
676    }
677}
678
679impl Display for Header {
680    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
681        let mut s = String::new();
682        if self.qr {
683            s.push_str("DNS Response (");
684        } else {
685            s.push_str("DNS Query (");
686        }
687        s.push_str(&self.info_str());
688        s.push(')');
689        write!(f, "{}", s)
690    }
691}
692
693impl Question {
694    /// Creates a DNS question.
695    pub fn new(name: Name, qtype: RecordType, qclass: Class) -> Self {
696        Question {
697            qname: name,
698            qtype,
699            qclass,
700        }
701    }
702
703    /// Encodes a `Question` as a series of bytes.
704    ///
705    /// Returns an error if a method defined in [`byteorder::WriteBytesExt`] returns an error.
706    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
707        let mut buf = Vec::new();
708        self.encode_into(&mut buf)?;
709        Ok(buf)
710    }
711
712    /// The same as [`encode()`](Self::encode()), but encoded bytes are appended to the given writer
713    /// instead of to a newly allocated one.
714    pub fn encode_into(&self, buf: &mut impl Write) -> Result<(), EncodeError> {
715        self.qname.encode_into(buf)?;
716        buf.write_u16::<NetworkEndian>(self.qtype.into())?;
717        buf.write_u16::<NetworkEndian>(self.qclass.encode())?;
718        Ok(())
719    }
720
721    /// Parses an encoded `Question` from a series of bytes.
722    ///
723    /// Returns an error if [`Name::parse()`], [`Class::parse()`] or a method defined in
724    /// [`byteorder::ReadBytesExt`] return an error.
725    pub fn parse(msg: &mut Cursor<&[u8]>) -> Result<Self, ParseError> {
726        let qname = Name::parse(msg, name::Compression::Allowed)?;
727        let qtype: RecordType = msg.read_u16::<NetworkEndian>()?.into();
728        let qclass = Class::parse(msg.read_u16::<NetworkEndian>()?)?;
729
730        Ok(Question {
731            qname,
732            qtype,
733            qclass,
734        })
735    }
736
737    /// Returns a string representing the record in the canonical format, with the owner padded to
738    /// the given length.
739    ///
740    /// If `output` is [`Some`] and the specified output stream supports colours, the output will be
741    /// colourized.
742    pub fn as_padded_string(&self, owner_len: usize, output: Option<owo_colors::Stream>) -> String {
743        let mut res = String::new();
744
745        let mut owner = self.qname.to_string();
746        while owner.len() < owner_len {
747            owner.push(' ');
748        }
749
750        let mut qtype = self.qtype.to_string();
751        if let Some(stream) = output {
752            owner = owner.if_supports_color(stream, |s| s.green()).to_string();
753            qtype = qtype.if_supports_color(stream, |s| s.purple()).to_string();
754        }
755
756        res.push_str(format!("{}          {}", owner, qtype).as_str());
757
758        res
759    }
760}
761
762impl Display for Question {
763    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
764        write!(
765            f,
766            "DNS Question for '{}' (type: {}, class: {})",
767            self.qname, self.qtype, self.qclass
768        )
769    }
770}
771
772impl Record {
773    /// Encodes a `Record` as a series of bytes.
774    ///
775    /// Returns an error if a method defined in [`byteorder::WriteBytesExt`] returns an error.
776    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
777        let mut buf = Vec::new();
778        self.encode_into(&mut buf)?;
779        Ok(buf)
780    }
781
782    /// The same as [`encode()`](Self::encode()), but encoded bytes are appended to the given writer
783    /// instead of to a newly allocated one.
784    pub fn encode_into(&self, buf: &mut impl Write) -> Result<(), EncodeError> {
785        match self {
786            Record::NONOPT(nonopt) => nonopt.encode_into(buf),
787            Record::OPT(opt) => opt.encode_into(buf),
788        }
789    }
790
791    /// Parses an encoded `Record` from a series of bytes.
792    ///
793    /// Returns an error if [`Name::parse()`], [`Class::parse()`],
794    /// [`parse_rdata()`](Self::parse_rdata()) or a method defined in [`byteorder::ReadBytesExt`]
795    /// return an error, or if an `OPT` record has a name other than `"."`.
796    pub fn parse(msg: &mut Cursor<&[u8]>, rcode: Option<RCode>) -> Result<Self, ParseError> {
797        let owner = Name::parse(msg, name::Compression::Allowed)?;
798        let atype: RecordType = msg.read_u16::<NetworkEndian>()?.into();
799        if atype == RecordType::OPT {
800            return OptRecord::parse(msg, owner, rcode);
801        }
802        let class = Class::parse(msg.read_u16::<NetworkEndian>()?)?;
803        let ttl = msg.read_u32::<NetworkEndian>()?;
804        let rdlength = msg.read_u16::<NetworkEndian>()?;
805
806        let mut encoded_rdata = vec![0; rdlength as usize];
807        let pos_rdata_start = msg.position();
808        msg.read_exact(&mut encoded_rdata)?;
809        // reset position to the start of rdata for parse_rdata()
810        msg.set_position(pos_rdata_start);
811        let rdata = Record::parse_rdata(&atype, msg, rdlength)?;
812
813        Ok(Record::NONOPT(NonOptRecord {
814            owner,
815            rtype: atype,
816            class,
817            ttl,
818            encoded_rdata,
819            rdata,
820        }))
821    }
822
823    /// Parses encoded rdata into a vector of strings (canonical format).
824    ///
825    /// `atype` is the type of the record containing the rdata. `msg` is the complete response
826    /// message, which is needed for message compression. `rdlength` is the length of the RDATA in
827    /// bytes.
828    ///
829    /// Returns an error if any of the `parse_rdata()` methods in [`rdata`] or a method defined in
830    /// [`byteorder::ReadBytesExt`] return an error.
831    pub fn parse_rdata(
832        atype: &RecordType,
833        msg: &mut Cursor<&[u8]>,
834        rdlength: u16,
835    ) -> Result<Rdata, ParseError> {
836        match atype {
837            RecordType::A => rdata::A::parse_rdata(msg, rdlength),
838            RecordType::NS => rdata::NS::parse_rdata(msg, rdlength),
839            RecordType::CNAME => rdata::CNAME::parse_rdata(msg, rdlength),
840            RecordType::SOA => rdata::SOA::parse_rdata(msg, rdlength),
841            RecordType::PTR => rdata::PTR::parse_rdata(msg, rdlength),
842            RecordType::HINFO => rdata::HINFO::parse_rdata(msg, rdlength),
843            RecordType::MX => rdata::MX::parse_rdata(msg, rdlength),
844            RecordType::TXT => rdata::TXT::parse_rdata(msg, rdlength),
845            RecordType::RP => rdata::RP::parse_rdata(msg, rdlength),
846            RecordType::AAAA => rdata::AAAA::parse_rdata(msg, rdlength),
847            RecordType::LOC => rdata::LOC::parse_rdata(msg, rdlength),
848            RecordType::SRV => rdata::SRV::parse_rdata(msg, rdlength),
849            RecordType::NAPTR => rdata::NAPTR::parse_rdata(msg, rdlength),
850            RecordType::CERT => rdata::CERT::parse_rdata(msg, rdlength),
851            RecordType::DNAME => rdata::DNAME::parse_rdata(msg, rdlength),
852            RecordType::OPT => rdata::OPT::parse_rdata(msg, rdlength),
853            RecordType::DS => rdata::DS::parse_rdata(msg, rdlength),
854            RecordType::SSHFP => rdata::SSHFP::parse_rdata(msg, rdlength),
855            RecordType::RRSIG => rdata::RRSIG::parse_rdata(msg, rdlength),
856            RecordType::NSEC => rdata::NSEC::parse_rdata(msg, rdlength),
857            RecordType::DNSKEY => rdata::DNSKEY::parse_rdata(msg, rdlength),
858            RecordType::NSEC3 => rdata::NSEC3::parse_rdata(msg, rdlength),
859            RecordType::NSEC3PARAM => rdata::NSEC3PARAM::parse_rdata(msg, rdlength),
860            RecordType::TLSA => rdata::TLSA::parse_rdata(msg, rdlength),
861            RecordType::OPENPGPKEY => rdata::OPENPGPKEY::parse_rdata(msg, rdlength),
862            RecordType::CAA => rdata::CAA::parse_rdata(msg, rdlength),
863            RecordType::Unknown(_) => {
864                let mut rdata = vec![0; rdlength as usize];
865                msg.read_exact(&mut rdata)?;
866                Ok(Rdata::Unknown(rdata))
867            }
868        }
869    }
870
871    /// Returns a reference to the inner [`OptRecord`]. [`None`] for the `NONOPT` variant.
872    pub fn as_opt(&self) -> Option<&OptRecord> {
873        match self {
874            Self::OPT(opt) => Some(opt),
875            Self::NONOPT(_) => None,
876        }
877    }
878
879    /// Returns a reference to the inner [`NonOptRecord`]. [`None`] for the `OPT` variant.
880    pub fn as_nonopt(&self) -> Option<&NonOptRecord> {
881        match self {
882            Self::NONOPT(nonopt) => Some(nonopt),
883            Self::OPT(_) => None,
884        }
885    }
886
887    /// Returns the inner [`OptRecord`]. Panics if the variant is not `OPT`.
888    pub fn into_opt(self) -> OptRecord {
889        match self {
890            Self::OPT(opt) => opt,
891            Self::NONOPT(_) => panic!("Record::into_opt() called on NONOPT variant"),
892        }
893    }
894
895    /// Returns the inner [`NonOptRecord`]. Panics if the variant is not `NONOPT`.
896    pub fn into_nonopt(self) -> NonOptRecord {
897        match self {
898            Self::NONOPT(nonopt) => nonopt,
899            Self::OPT(_) => panic!("Record::into_nonopt() called on OPT variant"),
900        }
901    }
902
903    /// Returns a reference to the contained [`Rdata`].
904    pub fn rdata(&self) -> &Rdata {
905        match self {
906            Self::OPT(opt) => opt.rdata(),
907            Self::NONOPT(nonopt) => nonopt.rdata(),
908        }
909    }
910
911    /// Returns a mutable reference to the contained [`Rdata`].
912    pub fn rdata_mut(&mut self) -> &mut Rdata {
913        match self {
914            Self::OPT(opt) => opt.rdata_mut(),
915            Self::NONOPT(nonopt) => nonopt.rdata_mut(),
916        }
917    }
918}
919
920impl NonOptRecord {
921    /// Creates a new `NonOptRecord` from [`Rdata`].
922    ///
923    /// Returns an error if `rdata` is [`Rdata::OPT`] or if `rdata` could not be encoded.
924    pub fn new(owner: Name, class: Class, ttl: u32, rdata: Rdata) -> Result<Self, ToluolError> {
925        if rdata.as_opt().is_some() {
926            return Err(ToluolError::OptRdataForNonOptRecord);
927        }
928
929        let rtype = rdata.rtype();
930        let encoded_rdata = rdata.encode()?;
931
932        Ok(Self {
933            owner,
934            rtype,
935            class,
936            ttl,
937            rdata,
938            encoded_rdata,
939        })
940    }
941
942    /// Encodes a `NonOptRecord` as a series of bytes.
943    ///
944    /// Returns an error if a method defined in [`byteorder::WriteBytesExt`] returns an error.
945    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
946        let mut buf = Vec::new();
947        self.encode_into(&mut buf)?;
948        Ok(buf)
949    }
950
951    /// The same as [`encode()`](Self::encode()), but encoded bytes are appended to the given writer
952    /// instead of to a newly allocated one.
953    pub fn encode_into(&self, buf: &mut impl Write) -> Result<(), EncodeError> {
954        self.owner.encode_into(buf)?;
955        buf.write_u16::<NetworkEndian>(self.rtype.into())?;
956        buf.write_u16::<NetworkEndian>(self.class.encode())?;
957        buf.write_u32::<NetworkEndian>(self.ttl)?;
958        buf.write_u16::<NetworkEndian>(self.encoded_rdata.len() as u16)?;
959        buf.write_all(&self.encoded_rdata)?;
960        Ok(())
961    }
962
963    /// Ensures the record is in canonical format, as defined in
964    /// [RFC 4034, Section 6.2](https://www.rfc-editor.org/rfc/rfc4034#section-6.2).
965    ///
966    /// `rrsig_labels` and `original_ttl` should be the values of
967    /// [`RRSIG::labels`](rdata::RRSIG::labels) and
968    /// [`RRSIG::original_ttl`](rdata::RRSIG::original_ttl), respectively, from a
969    /// [`RRSIG`](rdata::RRSIG) record that covers this record.
970    ///
971    /// Canonical format means that
972    /// - [`Self::owner`] is in canonical format (see
973    ///   [`Name::canonicalize()`](Name::canonicalize())).
974    /// - [`Self::rdata`] is in canonical format (see [`Rdata::canonicalize()`]).
975    /// - If [`Self::owner`] was originally a wildcard name, the wildcard substitution is undone.
976    ///   For example, imagine a record for a.example.com that was generated from *.example.com.
977    ///   Canonicalizing it would (among possibly other things) set its owner to *.example.com.
978    ///   This is what the `rrsig_labels` parameter is needed for.
979    /// - [`Self::ttl`] is set to the value of `original_ttl`.
980    pub fn canonicalize(&mut self, rrsig_labels: u8, original_ttl: u32) -> Result<(), DnssecError> {
981        if self.owner.label_count() < rrsig_labels {
982            return Err(DnssecError::InvalidRrsigLabelCount(
983                self.owner.label_count(),
984                rrsig_labels,
985            ));
986        }
987
988        self.owner.canonicalize();
989        self.rdata.canonicalize();
990        self.ttl = original_ttl;
991
992        // see RFC 4035, Section 5.3.2
993        let mut popped_label = None;
994        while self.owner.label_count() > rrsig_labels {
995            popped_label = self.owner.pop_front_label();
996        }
997
998        // if we popped a label, we still need to make `self.owner` into a wildcard name
999        if popped_label.is_some() {
1000            self.owner.make_wildcard();
1001        }
1002
1003        // ensure that any changes of the canonicalization are also reflected in the encoded rdata
1004        self.encoded_rdata.clear();
1005        self.rdata.encode_into(&mut self.encoded_rdata)?;
1006
1007        Ok(())
1008    }
1009
1010    /// Returns a reference to the contained [`Rdata`].
1011    pub fn rdata(&self) -> &Rdata {
1012        &self.rdata
1013    }
1014
1015    /// Returns a mutable reference to the contained [`Rdata`].
1016    pub fn rdata_mut(&mut self) -> &mut Rdata {
1017        &mut self.rdata
1018    }
1019
1020    /// Returns a string representing the record in the format used in zone files, but without the
1021    /// redundant IN class and without trailing dots for domain names.
1022    ///
1023    /// If `separate_with_single_space` is true, the different fields of the record are always
1024    /// separated by a single space. If it is false, all fields are separated by two spaces, and the
1025    /// TTL field is always six characters long (not including separators).
1026    ///
1027    /// If `owner_len`/`atype_len` is [`Some`], the `owner`/`atype` field is padded to the specified
1028    /// length.
1029    ///
1030    /// If `output` is [`Some`] and the specified output stream supports colours, the output will
1031    /// be colourized.
1032    pub fn as_string(
1033        &self,
1034        separate_with_single_space: bool,
1035        owner_len: Option<usize>,
1036        atype_len: Option<usize>,
1037        output: Option<owo_colors::Stream>,
1038    ) -> String {
1039        let mut owner = self.owner.to_string();
1040        if let Some(len) = owner_len {
1041            while owner.len() < len {
1042                owner.push(' ');
1043            }
1044        }
1045
1046        let mut atype = self.rtype.to_string();
1047        if let Some(len) = atype_len {
1048            while atype.len() < len {
1049                atype.push(' ');
1050            }
1051        }
1052
1053        if let Some(stream) = output {
1054            owner = owner.if_supports_color(stream, |s| s.green()).to_string();
1055            atype = atype.if_supports_color(stream, |s| s.purple()).to_string();
1056        }
1057
1058        if separate_with_single_space {
1059            format!("{} {} {} {}", owner, self.ttl, atype, self.rdata,)
1060        } else {
1061            format!("{}  {:>6}  {}  {}", owner, self.ttl, atype, &self.rdata,)
1062        }
1063    }
1064}
1065
1066impl Display for NonOptRecord {
1067    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1068        write!(f, "{}", self.as_string(true, None, None, None))
1069    }
1070}
1071
1072impl Display for OptFlags {
1073    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1074        let flag = format!("{:?}", self);
1075        write!(f, "{}", flag.to_ascii_lowercase())
1076    }
1077}
1078
1079impl OptRecord {
1080    /// Creates a new `OPT` record.
1081    ///
1082    /// For the `rcode` parameter, see [`Self::rcode`].
1083    pub fn new(rcode: Option<RCode>, edns_config: EdnsConfig) -> Result<Self, EncodeError> {
1084        let mut flags = vec![];
1085        if edns_config.do_flag {
1086            flags.push(OptFlags::DO);
1087        }
1088        let mut options = HashMap::new();
1089        if let Some(cookie) = edns_config.client_cookie {
1090            options.insert(OptionCode::Cookie, cookie.to_vec());
1091        }
1092        let rdata = Rdata::OPT(OPT { options });
1093        Ok(Self {
1094            owner: Name::root(),
1095            payload_size: edns_config.bufsize,
1096            rcode,
1097            edns_version: 0,
1098            flags,
1099            encoded_rdata: rdata.encode()?,
1100            rdata,
1101        })
1102    }
1103
1104    /// Encodes a `OptRecord` as a series of bytes.
1105    ///
1106    /// Returns an error if a method defined in [`byteorder::WriteBytesExt`] returns an error.
1107    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
1108        let mut buf = Vec::new();
1109        self.encode_into(&mut buf)?;
1110        Ok(buf)
1111    }
1112
1113    /// The same as [`encode()`](Self::encode()), but encoded bytes are appended to the given writer
1114    /// instead of to a newly allocated one.
1115    pub fn encode_into(&self, buf: &mut impl Write) -> Result<(), EncodeError> {
1116        self.owner.encode_into(buf)?;
1117        buf.write_u16::<NetworkEndian>(RecordType::OPT.into())?;
1118        buf.write_u16::<NetworkEndian>(self.payload_size)?;
1119        let rcode = self.rcode.unwrap_or(RCode::NOERROR);
1120        let rcode = (((rcode.encode() as u16) & 0b111111110000) >> 4) as u8;
1121        buf.write_u8(rcode)?;
1122        buf.write_u8(self.edns_version)?;
1123        if self.flags.contains(&OptFlags::DO) {
1124            buf.write_u16::<NetworkEndian>(1 << 15)?;
1125        } else {
1126            buf.write_u16::<NetworkEndian>(0)?;
1127        }
1128        buf.write_u16::<NetworkEndian>(self.encoded_rdata.len() as u16)?;
1129        buf.write_all(&self.encoded_rdata)?;
1130        Ok(())
1131    }
1132
1133    /// Returns a string describing the `OPT` record, with the given `prefix` prepended to each
1134    /// line.
1135    ///
1136    /// If `output` is [`Some`] and the specified output stream supports colours, the output will be
1137    /// colourized.
1138    pub fn as_padded_string(&self, prefix: &str, _output: Option<owo_colors::Stream>) -> String {
1139        let mut s = prefix.to_string();
1140
1141        s.push_str(&self.to_string());
1142
1143        // TODO: don't ignore output so we get coloured output
1144
1145        if !self.opt_rdata().options.is_empty() {
1146            let options = self.rdata.to_string();
1147            let options_iter: Vec<_> = options.split(", ").collect();
1148            let options_str = options_iter.join("\n");
1149            s.push('\n');
1150            s.push_str(prefix);
1151            s.push_str(&options_str);
1152        }
1153
1154        s
1155    }
1156
1157    /// Returns a reference to the contained [`Rdata`].
1158    pub fn rdata(&self) -> &Rdata {
1159        &self.rdata
1160    }
1161
1162    /// Returns a mutable reference to the contained [`Rdata`].
1163    pub fn rdata_mut(&mut self) -> &mut Rdata {
1164        &mut self.rdata
1165    }
1166
1167    /// Returns a reference to the contained [`Rdata`].
1168    pub fn opt_rdata(&self) -> &OPT {
1169        self.rdata.as_opt().expect("OPT record had non-OPT RDATA")
1170    }
1171
1172    /// Returns a mutable reference to the contained [`Rdata`].
1173    pub fn opt_rdata_mut(&mut self) -> &mut OPT {
1174        self.rdata
1175            .as_mut_opt()
1176            .expect("OPT record had non-OPT RDATA")
1177    }
1178
1179    /// Generates and returns the string used for our `Display` impl.
1180    fn info_str(&self) -> Result<String, fmt::Error> {
1181        use fmt::Write;
1182        let mut s = String::new();
1183        write!(&mut s, "EDNS: Version {}, flags: ", self.edns_version)?;
1184        let mut wrote_flag = false;
1185        for (i, flag) in self.flags.iter().enumerate() {
1186            wrote_flag = true;
1187            write!(&mut s, "{}", flag)?;
1188            if i < self.flags.len() - 1 {
1189                write!(&mut s, " ")?;
1190            }
1191        }
1192        if !wrote_flag {
1193            write!(&mut s, "<none>, ")?;
1194        } else {
1195            write!(&mut s, ", ")?;
1196        }
1197        write!(&mut s, "payload size: {}", self.payload_size)?;
1198        Ok(s)
1199    }
1200
1201    /// Parses an encoded `OptRecord` from a series of bytes.
1202    ///
1203    /// See [`DnsRecord::parse()`] for further information.
1204    fn parse(
1205        msg: &mut Cursor<&[u8]>,
1206        owner: Name,
1207        rcode: Option<RCode>,
1208    ) -> Result<Record, ParseError> {
1209        if !owner.is_root() {
1210            return Err(ParseError::InvalidOptName(owner));
1211        }
1212
1213        let payload_size = msg.read_u16::<NetworkEndian>()?;
1214        let ext_rcode = msg.read_u8()?;
1215        let rcode = if rcode.is_some() {
1216            match ext_rcode {
1217                0 => rcode,
1218                x => Some(RCode::parse(
1219                    ((x as u16) << 4) + (rcode.unwrap().encode() as u16),
1220                )?),
1221            }
1222        } else {
1223            rcode
1224        };
1225        let edns_version = msg.read_u8()?;
1226        let mut flags = vec![];
1227        let do_flag = msg.read_u16::<NetworkEndian>()? & (1 << 15) != 0;
1228        if do_flag {
1229            flags.push(OptFlags::DO);
1230        }
1231
1232        let rdlength = msg.read_u16::<NetworkEndian>()?;
1233        let mut encoded_rdata = vec![0; rdlength as usize];
1234        let pos_rdata_start = msg.position();
1235        msg.read_exact(&mut encoded_rdata)?;
1236        // reset position to the start of rdata for parse_rdata()
1237        msg.set_position(pos_rdata_start);
1238        let rdata = Record::parse_rdata(&RecordType::OPT, msg, rdlength)?;
1239
1240        Ok(Record::OPT(OptRecord {
1241            owner,
1242            payload_size,
1243            rcode,
1244            edns_version,
1245            flags,
1246            encoded_rdata,
1247            rdata,
1248        }))
1249    }
1250}
1251
1252impl Display for OptRecord {
1253    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1254        write!(f, "{}", self.info_str()?)
1255    }
1256}
1257
1258impl Display for Record {
1259    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1260        match self {
1261            Record::NONOPT(nonopt) => write!(f, "{}", nonopt),
1262            Record::OPT(opt) => write!(f, "{}", opt),
1263        }
1264    }
1265}
1266
1267impl Message {
1268    /// Creates a DNS query.
1269    ///
1270    /// If `edns` is [`Some`], the query will contain an `OPT` record.
1271    ///
1272    /// See [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035) and the documentation of [`Header`]
1273    /// for information about the remaining parameters.
1274    ///
1275    /// Returns an error if `aa` or `ra` are set in `flags`.
1276    pub fn new_query(
1277        domain: Name,
1278        qtype: RecordType,
1279        opcode: Opcode,
1280        flags: HeaderFlags,
1281        edns: Option<EdnsConfig>,
1282    ) -> Result<Self, EncodeError> {
1283        if flags.aa || flags.ra {
1284            return Err(EncodeError::AaOrRaInQuery);
1285        }
1286
1287        let msg_id = rand::thread_rng().gen_range(0..(1u32 << 16)) as u16;
1288
1289        let header = Header::new_query_header(msg_id, opcode, flags, edns.is_some(), 1)?;
1290
1291        let mut additional_answers = Vec::new();
1292        if let Some(edns_config) = edns {
1293            additional_answers.push(Record::OPT(OptRecord::new(None, edns_config)?));
1294        }
1295
1296        Ok(Message {
1297            header,
1298            questions: vec![Question::new(domain, qtype, Class::IN)],
1299            answers: Vec::new(),
1300            authoritative_answers: Vec::new(),
1301            additional_answers,
1302        })
1303    }
1304
1305    /// Creates a DNS response.
1306    ///
1307    /// See the documentation of [`Header`] for information about the parameters.
1308    ///
1309    /// `answers`, `authoritative_answers`, `additional_answers` are grouped in that order in the `records` parameter.
1310    pub fn new_response(
1311        msg_id: u16,
1312        opcode: Opcode,
1313        flags: HeaderFlags,
1314        rcode: RCode,
1315        questions: Vec<Question>,
1316        records: [Vec<Record>; 3],
1317    ) -> Self {
1318        Message {
1319            header: Header::new_response_header(
1320                msg_id,
1321                opcode,
1322                flags,
1323                rcode,
1324                [
1325                    questions.len() as u16,
1326                    records[0].len() as u16,
1327                    records[1].len() as u16,
1328                    records[2].len() as u16,
1329                ],
1330            ),
1331            questions,
1332            answers: records[0].clone(),
1333            authoritative_answers: records[1].clone(),
1334            additional_answers: records[2].clone(),
1335        }
1336    }
1337
1338    /// Encodes a `Message` as a series of bytes.
1339    ///
1340    /// Returns an error if [`Header::encode()`], [`Question::encode()`] or [`Record::encode()`]
1341    /// return an error.
1342    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
1343        let mut buf = Vec::new();
1344        self.encode_into(&mut buf)?;
1345        Ok(buf)
1346    }
1347
1348    /// The same as [`encode()`](Self::encode()), but encoded bytes are appended to the given writer
1349    /// instead of to a newly allocated one.
1350    pub fn encode_into(&self, buf: &mut impl Write) -> Result<(), EncodeError> {
1351        self.header.encode_into(buf)?;
1352        for question in &self.questions {
1353            question.encode_into(buf)?;
1354        }
1355        for record in &self.answers {
1356            record.encode_into(buf)?;
1357        }
1358        for record in &self.authoritative_answers {
1359            record.encode_into(buf)?;
1360        }
1361        for record in &self.additional_answers {
1362            record.encode_into(buf)?;
1363        }
1364
1365        Ok(())
1366    }
1367
1368    /// Parses an encoded `Message` from a series of bytes.
1369    ///
1370    /// Returns an error if [`Header::parse()`], [`Question::parse()`] or [`Record::parse()`] return
1371    /// an error or a truncated message is received.
1372    pub fn parse(msg: &mut Cursor<&[u8]>) -> Result<Self, ParseError> {
1373        let mut header = Header::parse(msg)?;
1374
1375        if header.flags.tc {
1376            return Err(ParseError::TruncatedMessage);
1377        }
1378
1379        let qdcount = header.qdcount;
1380        let ancount = header.ancount;
1381        let nscount = header.nscount;
1382        let arcount = header.arcount;
1383        let questions = Message::parse_questions(msg, qdcount)?;
1384        let mut answers = Vec::new();
1385        let mut authoritative_answers = Vec::new();
1386        let mut additional_answers = Vec::new();
1387        if ancount > 0 {
1388            answers = Message::parse_records(msg, ancount, header.rcode)?;
1389        }
1390        if nscount > 0 {
1391            authoritative_answers = Message::parse_records(msg, nscount, header.rcode)?;
1392        }
1393        if arcount > 0 {
1394            additional_answers = Message::parse_records(msg, arcount, header.rcode)?;
1395        }
1396
1397        for answer in &additional_answers {
1398            if let Record::OPT(OptRecord { rcode, .. }) = answer {
1399                header.rcode = *rcode;
1400            }
1401        }
1402
1403        Ok(Message {
1404            header,
1405            questions,
1406            answers,
1407            authoritative_answers,
1408            additional_answers,
1409        })
1410    }
1411
1412    /// Returns a string verbosely describing the message (i.e. header and all the other sections).
1413    ///
1414    /// If `output` is [`Some`] and the specified output stream supports colours, the output will
1415    /// be colourized.
1416    pub fn as_string(&self, output: Option<owo_colors::Stream>) -> String {
1417        let section_name = |s: &str, o: Option<owo_colors::Stream>| {
1418            let mut s = s.to_string();
1419            if let Some(stream) = o {
1420                s = s.if_supports_color(stream, |s| s.yellow()).to_string();
1421            }
1422            s
1423        };
1424
1425        let mut res = String::new();
1426
1427        let mut additional_answers = self.additional_answers.clone();
1428        let mut opt_index = None;
1429
1430        let mut max_owner_len = 0;
1431        let mut max_type_len = 0;
1432
1433        for q in &self.questions {
1434            max_owner_len = max(max_owner_len, q.qname.string_len());
1435            max_type_len = max(max_type_len, q.qtype.to_string().len());
1436        }
1437
1438        let answers = [
1439            &self.answers,
1440            &self.authoritative_answers,
1441            &self.additional_answers,
1442        ];
1443        let answers_iter = answers.iter().flat_map(|a| a.iter());
1444        for (i, answer) in answers_iter.enumerate() {
1445            match answer {
1446                Record::OPT(_) => {
1447                    // the iterator runs over self.answers, self.authoritative_answers and self.additional_answers,
1448                    // but we want the index of answer with respect to self.additional_answers
1449                    opt_index = Some(i - self.answers.len() - self.authoritative_answers.len());
1450                }
1451                Record::NONOPT(NonOptRecord {
1452                    owner: name,
1453                    rtype: atype,
1454                    ..
1455                }) => {
1456                    max_owner_len = max(max_owner_len, name.string_len());
1457                    max_type_len = max(max_type_len, atype.to_string().len());
1458                }
1459            }
1460        }
1461
1462        // Header
1463        // TODO: coloured header output?
1464        res.push_str(section_name("Header:\n\t", output).as_str());
1465        res.push_str(format!("{}\n\n", self.header.info_str()).as_str());
1466
1467        // OPT Pseudosection (if present)
1468        if let Some(idx) = opt_index {
1469            let opt = additional_answers.remove(idx);
1470            let opt = opt
1471                .as_opt()
1472                .expect("Calculated incorrect index for OPT record");
1473            res.push_str(section_name("OPT Pseudosection:\n", output).as_str());
1474            res.push_str(&opt.as_padded_string("\t", output));
1475            res.push_str("\n\n");
1476        }
1477
1478        res.push_str(section_name("Question Section:\n", output).as_str());
1479        for question in &self.questions {
1480            res.push('\t');
1481            // question doesn't need max_type_len because nothing gets printed after its qtype
1482            res.push_str(question.as_padded_string(max_owner_len, output).as_str());
1483            res.push('\n');
1484        }
1485        res.push('\n');
1486
1487        if !self.answers.is_empty() {
1488            res.push_str(section_name("Answer Section:\n", output).as_str());
1489            for answer in &self.answers {
1490                res.push('\t');
1491                res.push_str(
1492                    answer
1493                        .as_nonopt()
1494                        .expect("Unexpected OPT record")
1495                        .as_string(false, Some(max_owner_len), Some(max_type_len), output)
1496                        .as_str(),
1497                );
1498                res.push('\n');
1499            }
1500            res.push('\n');
1501        }
1502
1503        if !self.authoritative_answers.is_empty() {
1504            res.push_str(section_name("Authoritative Section:\n", output).as_str());
1505            for answer in &self.authoritative_answers {
1506                res.push('\t');
1507                res.push_str(
1508                    answer
1509                        .as_nonopt()
1510                        .expect("Unexpected OPT record")
1511                        .as_string(false, Some(max_owner_len), Some(max_type_len), output)
1512                        .as_str(),
1513                );
1514                res.push('\n');
1515            }
1516            res.push('\n');
1517        }
1518
1519        if !additional_answers.is_empty() {
1520            res.push_str(section_name("Additional Section:\n", output).as_str());
1521            for answer in &additional_answers {
1522                res.push('\t');
1523                res.push_str(
1524                    answer
1525                        .as_nonopt()
1526                        .expect("Unexpected OPT record")
1527                        .as_string(false, Some(max_owner_len), Some(max_type_len), output)
1528                        .as_str(),
1529                );
1530                res.push('\n');
1531            }
1532        }
1533
1534        // remove trailing '\n's
1535        while res.chars().nth(res.len() - 1).unwrap() == '\n' {
1536            res.remove(res.len() - 1);
1537        }
1538
1539        res
1540    }
1541
1542    /// Parses the question section of a DNS message.
1543    fn parse_questions(msg: &mut Cursor<&[u8]>, qdcount: u16) -> Result<Vec<Question>, ParseError> {
1544        let mut questions = Vec::with_capacity(qdcount as usize);
1545        for _i in 0..qdcount {
1546            questions.push(Question::parse(msg)?);
1547        }
1548
1549        Ok(questions)
1550    }
1551
1552    /// Parses an answer section (i. e. answer, authoritative or additional) of a DNS message.
1553    fn parse_records(
1554        msg: &mut Cursor<&[u8]>,
1555        ancount: u16,
1556        rcode: Option<RCode>,
1557    ) -> Result<Vec<Record>, ParseError> {
1558        let mut answers = Vec::with_capacity(ancount as usize);
1559        for _i in 0..ancount {
1560            answers.push(Record::parse(msg, rcode)?);
1561        }
1562
1563        Ok(answers)
1564    }
1565}