Files
NLnetLabs-domain/src/base/message.rs
T
Martin HoffmannandGitHub 4cdf3b62bf Drop AsRef<[u8]> and AsMut<[u8]> from OctetsBuilder trait bounds. (#130)
This makes it possible to impl OctetsBuilder for Cow<[u8]> and other
types that can only append at the price of explicitely requiring AsRef<[u8]>
and AsMut<[u8]> where needed. Unfortunately, this means lots of code
changes but hopefully, these should not cause too much havoc with code
using the crate.
2022-05-24 12:55:10 +02:00

1239 lines
41 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! Accessing existing DNS messages.
//!
//! This module defines a number of types for processing the content of a DNS
//! message in wire format. Because many components of the message are of
//! varying length, this can only be done iteratively. The type [`Message`]
//! wraps an octets sequence containing a complete message. It provides access
//! to the four sections of the message via additional types.
//!
//! For details, see the [`Message`] type.
//!
//! [`Message`]: struct.Message.html
use super::header::{Header, HeaderCounts, HeaderSection};
use super::iana::{Class, Rcode, Rtype};
use super::message_builder::{AdditionalBuilder, AnswerBuilder};
use super::name::ParsedDname;
use super::octets::{
OctetsBuilder, OctetsFrom, OctetsRef, Parse, ParseError, Parser, ShortBuf,
};
use super::opt::{Opt, OptRecord};
use super::question::Question;
use super::rdata::ParseRecordData;
use super::record::{AsRecord, ParsedRecord, Record};
use crate::rdata::rfc1035::Cname;
use core::marker::PhantomData;
use core::{fmt, mem};
//------------ Message -------------------------------------------------------
/// A DNS message.
///
/// This type wraps an octets sequence containing the complete wire-format DNS
/// message and allows access to the various components of the message.
///
/// You create a message by passing an octets sequence to the [`from_octets`]
/// associate function which does some basic sanity checks and, if they
/// succeed, returns a message for the sequence. All further parsing happens
/// lazily when you access more of the message. This means that a message is
/// not necessarily well-formatted and further parsing may fail later on.
///
/// Section 4 of [RFC 1035] defines DNS messages as being divded into five
/// sections named header, question, answer, authority, and additional.
///
/// The header section is of a fixed sized and can be accessed at any time
/// through the methods given under [Header Section]. Most likely, you will
/// be interested in the first part of the header which is
/// returned by the [`header`] method. The second part of the header
/// section contains the number of entries in the following four sections
/// and is of less interest as there are more sophisticated ways of accessing
/// these sections. If you do care, you can get access through
/// [`header_counts`].
///
/// The meaning of the next four sections depends on the type of message as
/// described by the [opcode] field of the header. Since the most common
/// type is a query, the sections are named after their function in this type
/// and the following description will focus on it.
///
/// The question section contains what was asked of the DNS by a query. It
/// contains a number of questions that consist of a domain name, a record
/// type, and class. A query asks for all records of the given record type
/// that are owned by the domain name within the class. In queries, there will
/// be exactly one question. With other opcodes, there may be multiple
/// questions.
///
/// You can acquire an iterator over the questions through the [`question`]
/// method. It returns a [`QuestionSection`] value that is an iterator over
/// questions. Since a single question is such a common case, there is a
/// convenience method [`first_question`] that returns the first question
/// only.
///
/// The following three section all contain DNS resource records. In
/// queries, they are empty in a request and may or may not contain records
/// in a response. The *answer* section contains all the records that answer
/// the given question. The *authority* section contains records declaring
/// which name server provided authoritative information for the question,
/// and the *additional* section can contain records that the name server
/// thought might be useful for processing the question. For instance, if you
/// trying to find out the mail server of a domain by asking for MX records,
/// you likely also want the IP addresses for the server, so the name server
/// may include these right away and free of charge.
///
/// There are functions to access all three sections directly: [`answer`],
/// [`authority`], and [`additional`]. Each method returns a value of type
/// [RecordSection] which acts as an iterator over the records in the
/// section. Since there are no pointers to where the later sections start,
/// accessing them directly means iterating over the previous sections. This
/// is why it is more efficitent to call [`RecordSection::next_section`] to
/// progress to a later section. Alternatively, you can use the message’s
/// [`sections`] method that gives you all four sections at once with the
/// minimal amount of iterating necessary.
///
/// When iterating over the record section, you will receive values of type
/// [`ParsedRecord`], an intermediary type that only parsed the parts common
/// to all records. In order to access the data of the record, you will want
/// to convert it into a [`Record`] which is generic over the actual record
/// type data. This can be done via [`ParsedRecord::into_record`].
///
/// Alternatively, you can trade the record section for one that only returns
/// the types you are interested in via the [`RecordSection::limit_to`]
/// method. The iterator returned by that method will quietly skip over all
/// records that aren’t of the type you are interested in.
///
/// So, if you want to iterate over the MX records in the answer section, you
/// would do something like this:
///
/// ```
/// use domain::base::Message;
/// use domain::rdata::Mx;
///
/// # let octets = vec![0; 12];
/// let msg = Message::from_octets(octets).unwrap();
/// for record in msg.answer().unwrap().limit_to::<Mx<_>>() {
/// if let Ok(record) = record {
/// // Do something with the record ...
/// }
/// }
/// ```
///
/// The `limit_to` method takes the record type as a type argument. Many
/// record types, like [`Mx`], are generic over octet sequences but the
/// compiler generally can figure out the concrete type itself, so in most
/// cases you get away with the underscore there.
///
/// Note how the iterator doesn’t actually return records but results of
/// records and parse errors. This is because only now can it check whether
/// the record is actually properly formatted. An error signals that something
/// went wrong while parsing. If only the record data is broken, the message
/// remains useful and parsing can continue with the next record. If the
/// message is fully broken, the next iteration will return `None` to signal
/// that.
///
/// [`additional`]: #method.additional
/// [`answer`]: #method.answer
/// [`authority`]: #method.authority
/// [`first_question`]: #method.first_question
/// [`from_octets`]: #method.from_octets
/// [`header`]: #method.header
/// [`header_counts`]: #method.header_counts
/// [`question`]: #method.question
/// [`sections`]: #method.sections
/// [`Mx`]: ../../rdata/rfc1035/struct.Mx.html
/// [`ParsedRecord`]: ../record/struct.ParsedRecord.html
/// [`ParsedRecord::into_record`]: ../record/struct.ParsedRecord.html#method.into_record
/// [`QuestionSection`]: struct.QuestionSection.html
/// [`Record`]: ../record/struct.Record.html
/// [`RecordSection`]: struct.RecordSection.html
/// [`RecordSection::limit_to`]: ../struct.RecordSection.html#method.limit_to
/// [`RecordSection::next_section`]: ../struct.RecordSection.html#method.next_section
/// [Header Section]: #header-section
/// [rdata]: ../../rdata/index.html
/// [opcode]: ../iana/opcode/enum.Opcode.html
/// [RFC 1035]: https://tools.ietf.org/html/rfc1035
#[derive(Clone, Copy)]
pub struct Message<Octets> {
octets: Octets,
}
/// # Creation and Conversion
///
impl<Octets> Message<Octets> {
/// Creates a message from an octets sequence.
///
/// This fails if the slice is too short to even contain a complete
/// header section. No further checks are done, though, so if this
/// function returns ok, the message may still be broken with other
/// methods returning errors later one.
pub fn from_octets(octets: Octets) -> Result<Self, ShortBuf>
where
Octets: AsRef<[u8]>,
{
if octets.as_ref().len() < mem::size_of::<HeaderSection>() {
Err(ShortBuf)
} else {
Ok(unsafe { Self::from_octets_unchecked(octets) })
}
}
/// Creates a message from a bytes value without checking.
///
/// The methods for header access rely on the octets being at least as
/// long as a header, so this is unsafe.
pub(super) unsafe fn from_octets_unchecked(octets: Octets) -> Self {
Message { octets }
}
/// Returns a reference to the underlying octets sequence.
pub fn as_octets(&self) -> &Octets {
&self.octets
}
/// Converts the message into the underlying octets sequence.
pub fn into_octets(self) -> Octets {
self.octets
}
/// Returns a slice to the underlying octets sequence.
pub fn as_slice(&self) -> &[u8]
where
Octets: AsRef<[u8]>,
{
self.octets.as_ref()
}
/// Returns a mutable slice to the underlying octets sequence.
///
/// Because it is possible to utterly break the message using this slice,
/// the method is private.
fn as_slice_mut(&mut self) -> &mut [u8]
where
Octets: AsMut<[u8]>,
{
self.octets.as_mut()
}
/// Returns a message for a slice of the octets sequence.
pub fn for_slice(&self) -> Message<&[u8]>
where
Octets: AsRef<[u8]>,
{
unsafe { Message::from_octets_unchecked(self.octets.as_ref()) }
}
}
/// # Header Section
///
impl<Octets: AsRef<[u8]>> Message<Octets> {
/// Returns the message header.
pub fn header(&self) -> Header {
*Header::for_message_slice(self.as_slice())
}
/// Returns a mutable reference to the message header.
pub fn header_mut(&mut self) -> &mut Header
where
Octets: AsMut<[u8]>,
{
Header::for_message_slice_mut(self.as_slice_mut())
}
/// Returns the header counts of the message.
pub fn header_counts(&self) -> HeaderCounts {
*HeaderCounts::for_message_slice(self.as_slice())
}
/// Returns the entire header section.
pub fn header_section(&self) -> HeaderSection {
*HeaderSection::for_message_slice(self.as_slice())
}
/// Returns whether the rcode of the header is NoError.
pub fn no_error(&self) -> bool {
self.header().rcode() == Rcode::NoError
}
/// Returns whether the rcode of the header is one of the error values.
pub fn is_error(&self) -> bool {
self.header().rcode() != Rcode::NoError
}
}
/// # Access to Sections
///
impl<Octets> Message<Octets>
where
for<'a> &'a Octets: OctetsRef,
{
/// Returns the question section.
pub fn question(&self) -> QuestionSection<&Octets> {
QuestionSection::new(&self.octets)
}
/// Returns the zone section of an UPDATE message.
///
/// This is identical to `self.question()`.
pub fn zone(&self) -> QuestionSection<&Octets> {
self.question()
}
/// Returns the answer section.
///
/// Iterates over the question section in order to access the answer
/// section. If you are accessing the question section anyway, using
/// its [`next_section`] method may be more efficient.
///
/// [`next_section`]: ../struct.QuestionSection.html#method.next_section
pub fn answer(&self) -> Result<RecordSection<&Octets>, ParseError> {
self.question().next_section()
}
/// Returns the prerequisite section of an UPDATE message.
///
/// This is identical to `self.answer()`.
pub fn prerequisite(&self) -> Result<RecordSection<&Octets>, ParseError> {
self.answer()
}
/// Returns the authority section.
///
/// Iterates over both the question and the answer sections to determine
/// the start of the authority section. If you are already accessing the
/// answer section, using [`next_section`] on it is more efficient.
///
/// [`next_section`]: ../struct.RecordSection.html#method.next_section
pub fn authority(&self) -> Result<RecordSection<&Octets>, ParseError> {
Ok(self.answer()?.next_section()?.unwrap())
}
/// Returns the update section of an UPDATE message.
///
/// This is identical to `self.authority()`.
pub fn update(&self) -> Result<RecordSection<&Octets>, ParseError> {
self.authority()
}
/// Returns the additional section.
///
/// Iterates over all three previous sections to determine the start of
/// the additional section. If you are already accessing the
/// authority section, using [`next_section`] on it is more efficient.
///
/// [`next_section`]: ../struct.RecordSection.html#method.next_section
pub fn additional(&self) -> Result<RecordSection<&Octets>, ParseError> {
Ok(self.authority()?.next_section()?.unwrap())
}
/// Returns all four sections in one fell swoop.
#[allow(clippy::type_complexity)]
pub fn sections(
&self,
) -> Result<
(
QuestionSection<&Octets>,
RecordSection<&Octets>,
RecordSection<&Octets>,
RecordSection<&Octets>,
),
ParseError,
> {
let question = self.question();
let answer = question.next_section()?;
let authority = answer.next_section()?.unwrap();
let additional = authority.next_section()?.unwrap();
Ok((question, answer, authority, additional))
}
/// Returns an iterator over the records in the message.
///
/// The iterator’s item is a pair of a [`ParsedRecord`] and the
/// [`Section`] it was found in.
///
/// As is customary, this iterator is also accessible via the
/// `IntoIterator` trait on a reference to the message.
///
/// [`ParsedRecord`]: ../record/struct.ParsedRecord.html
/// [`Section`]: enum.Section.html
pub fn iter(&self) -> MessageIter<&Octets> {
self.into_iter()
}
}
/// # Helpers for Common Tasks
///
impl<Octets> Message<Octets>
where
Octets: AsRef<[u8]>,
for<'a> &'a Octets: OctetsRef,
{
/// Returns whether this is the answer to some other message.
///
/// The method checks whether the ID fields of the headers are the same,
/// whether the QR flag is set in this message, and whether the questions
/// are the same.
pub fn is_answer<Other>(&self, query: &Message<Other>) -> bool
where
Other: AsRef<[u8]>,
for<'o> &'o Other: OctetsRef,
{
if !self.header().qr()
|| self.header().id() != query.header().id()
|| self.header_counts().qdcount()
!= query.header_counts().qdcount()
{
false
} else {
self.question() == query.question()
}
}
/// Returns the first question, if there is any.
///
/// The method will return `None` both if there are no questions or if
/// parsing fails.
pub fn first_question(&self) -> Option<Question<ParsedDname<&Octets>>> {
match self.question().next() {
None | Some(Err(..)) => None,
Some(Ok(question)) => Some(question),
}
}
/// Returns the sole question of the message.
///
/// This is like [`first_question`] but returns an error if there isn’t
/// exactly one question or there is a parse error.
///
/// [`first_question`]: #method.first_question
pub fn sole_question(
&self,
) -> Result<Question<ParsedDname<&Octets>>, ParseError> {
match self.header_counts().qdcount() {
0 => return Err(ParseError::form_error("no question")),
1 => {}
_ => return Err(ParseError::form_error("multiple questions")),
}
self.question().next().unwrap()
}
/// Returns the query type of the first question, if any.
pub fn qtype(&self) -> Option<Rtype> {
self.first_question().map(|x| x.qtype())
}
/// Returns whether the message contains answers of a given type.
pub fn contains_answer<'s, Data>(&'s self) -> bool
where
Data: ParseRecordData<&'s Octets>,
{
let answer = match self.answer() {
Ok(answer) => answer,
Err(..) => return false,
};
answer.limit_to::<Data>().next().is_some()
}
/// Resolves the canonical name of the answer.
///
/// The CNAME record allows a domain name to be an alias for a different
/// name. Aliases may be chained. The ‘canonical name’ referred to be the
/// method’s name is the last name in this chain. A recursive resolver
/// will support a stub resolver in figuring out this canonical name by
/// including all necessary CNAME records in its answer. This method can
/// be used on such an answer to determine the canonical name. As such,
/// it will only consider CNAMEs present in the message’s answer section.
///
/// It starts with the question name and follows CNAME records until there
/// is no next CNAME in the chain and then returns the last CNAME.
///
/// If the message doesn’t have a question, if there is a parse error, or
/// if there is a CNAME loop the method returns `None`.
//
// Loop detection is done by breaking off after ANCOUNT + 1 steps -- if
// there is more steps then there is records in the answer section we
// must have a loop. While the ANCOUNT could be unreasonably large, the
// iterator would break off in this case and we break out with a None
// right away.
pub fn canonical_name(&self) -> Option<ParsedDname<&Octets>> {
let question = match self.first_question() {
None => return None,
Some(question) => question,
};
let mut name = question.into_qname();
let answer = match self.answer() {
Ok(answer) => answer.limit_to::<Cname<_>>(),
Err(_) => return None,
};
for _ in 0..self.header_counts().ancount() + 1 {
let mut found = false;
for record in answer.clone() {
let record = match record {
Ok(record) => record,
Err(_) => continue,
};
if *record.owner() == name {
name = *record.data().cname();
found = true;
break;
}
}
if !found {
return Some(name);
}
}
None
}
/// Returns the OPT record from the message, if there is one.
pub fn opt(&self) -> Option<OptRecord<<&Octets as OctetsRef>::Range>> {
match self.additional() {
Ok(section) => match section.limit_to::<Opt<_>>().next() {
Some(Ok(rr)) => Some(OptRecord::from(rr)),
_ => None,
},
Err(_) => None,
}
}
/// Returns the last additional record from the message.
///
/// The method tries to parse the last record of the additional section
/// as the provided record type. If that succeeds, it returns that
/// parsed record.
///
/// If the last record is of the wrong type or parsing fails, returns
/// `None`.
pub fn get_last_additional<'s, Data: ParseRecordData<&'s Octets>>(
&'s self,
) -> Option<Record<ParsedDname<&'s Octets>, Data>> {
let mut section = match self.additional() {
Ok(section) => section,
Err(_) => return None,
};
loop {
match section.count {
Err(_) => return None,
Ok(0) => return None,
Ok(1) => break,
_ => {}
}
let _ = section.next();
}
let record = match ParsedRecord::parse(&mut section.parser) {
Ok(record) => record,
Err(_) => return None,
};
let record = match record.into_record() {
Ok(Some(record)) => record,
_ => return None,
};
Some(record)
}
/// Drops the last additional record from the message.
///
/// Does so by decreasing the ’arcount.’ Does, however, not change the
/// underlying octet sequence.
///
/// # Panics
///
/// The method panics if the additional section is empty.
pub fn remove_last_additional(&mut self)
where
Octets: AsMut<[u8]>,
{
HeaderCounts::for_message_slice_mut(self.octets.as_mut())
.dec_arcount();
}
/// Copy records from a message into the target message builder.
///
/// The method uses `op` to process records from all record sections
/// before inserting, caller can use this closure to filter or manipulate
/// records before inserting.
pub fn copy_records<'s, R, F, T, O>(
&'s self,
target: T,
mut op: F,
) -> Result<AdditionalBuilder<O>, CopyRecordsError>
where
for<'a> &'a Octets: OctetsRef,
R: AsRecord + 's,
F: FnMut(ParsedRecord<&'s Octets>) -> Option<R>,
T: Into<AnswerBuilder<O>>,
O: OctetsBuilder + AsMut<[u8]>,
{
let mut source = self.answer()?;
let mut target = target.into();
for rr in &mut source {
let rr = rr?;
if let Some(rr) = op(rr) {
target.push(rr)?;
}
}
let mut source = source.next_section()?.unwrap();
let mut target = target.authority();
for rr in &mut source {
let rr = rr?;
if let Some(rr) = op(rr) {
target.push(rr)?;
}
}
let source = source.next_section()?.unwrap();
let mut target = target.additional();
for rr in source {
let rr = rr?;
if let Some(rr) = op(rr) {
target.push(rr)?;
}
}
Ok(target)
}
}
//--- AsRef
impl<Octets> AsRef<Octets> for Message<Octets> {
fn as_ref(&self) -> &Octets {
&self.octets
}
}
impl<Octets: AsRef<[u8]>> AsRef<[u8]> for Message<Octets> {
fn as_ref(&self) -> &[u8] {
self.octets.as_ref()
}
}
//--- OctetsFrom
impl<Octets, SrcOctets> OctetsFrom<Message<SrcOctets>> for Message<Octets>
where
Octets: OctetsFrom<SrcOctets>,
{
fn octets_from(source: Message<SrcOctets>) -> Result<Self, ShortBuf> {
Octets::octets_from(source.octets)
.map(|octets| unsafe { Self::from_octets_unchecked(octets) })
}
}
//--- IntoIterator
impl<'a, Octets> IntoIterator for &'a Message<Octets>
where
for<'s> &'s Octets: OctetsRef,
{
type Item = Result<(ParsedRecord<&'a Octets>, Section), ParseError>;
type IntoIter = MessageIter<&'a Octets>;
fn into_iter(self) -> Self::IntoIter {
MessageIter {
inner: self.answer().ok(),
}
}
}
//------------ QuestionSection ----------------------------------------------
/// An iterator over the question section of a DNS message.
///
/// The iterator’s item is the result of trying to parse the question. In case
/// of a parse error, `next` will return an error once and `None` after that.
///
/// You can create a value of this type through [`Message::question`]. Use the
/// [`answer`] or [`next_section`] methods on a question section to proceed
/// to an iterator over the answer section.
///
/// [`Message::question`]: struct.Message.html#method.question
/// [`answer`]: #method.answer
/// [`next_section`]: #method.next_section
#[derive(Clone, Copy, Debug)]
pub struct QuestionSection<Ref> {
/// The parser for generating the questions.
parser: Parser<Ref>,
/// The remaining number of questions.
///
/// The `Result` is here to monitor an error during iteration.
/// It is used to fuse the iterator after an error and is also returned
/// by `answer()` should that be called after an error.
count: Result<u16, ParseError>,
}
impl<Ref: OctetsRef> QuestionSection<Ref> {
/// Creates a new question section from a reference to the message octets.
fn new(octets: Ref) -> Self {
let mut parser = Parser::from_ref(octets);
parser.advance(mem::size_of::<HeaderSection>()).unwrap();
QuestionSection {
count: Ok(
HeaderCounts::for_message_slice(parser.as_slice()).qdcount()
),
parser,
}
}
/// Returns the current position relative to the beginning of the message.
pub fn pos(&self) -> usize {
self.parser.pos()
}
/// Proceeds to the answer section.
///
/// Skips over any remaining questions and then converts itself into the
/// first [`RecordSection`].
///
/// [`RecordSection`]: struct.RecordSection.html
pub fn answer(mut self) -> Result<RecordSection<Ref>, ParseError> {
while self.next().is_some() {}
let _ = self.count?;
Ok(RecordSection::new(self.parser, Section::first()))
}
/// Proceeds to the answer section.
///
/// This is identical to [`answer`][Self::answer] and is here for
/// consistency.
pub fn next_section(self) -> Result<RecordSection<Ref>, ParseError> {
self.answer()
}
}
//--- Iterator
impl<Ref: OctetsRef> Iterator for QuestionSection<Ref> {
type Item = Result<Question<ParsedDname<Ref>>, ParseError>;
fn next(&mut self) -> Option<Self::Item> {
match self.count {
Ok(count) if count > 0 => match Question::parse(&mut self.parser)
{
Ok(question) => {
self.count = Ok(count - 1);
Some(Ok(question))
}
Err(err) => {
self.count = Err(err);
Some(Err(err))
}
},
_ => None,
}
}
}
//--- PartialEq
impl<Ref, Other> PartialEq<QuestionSection<Other>> for QuestionSection<Ref>
where
Ref: OctetsRef,
Other: OctetsRef,
{
fn eq(&self, other: &QuestionSection<Other>) -> bool {
let mut me = *self;
let mut other = *other;
loop {
match (me.next(), other.next()) {
(Some(Ok(left)), Some(Ok(right))) => {
if left != right {
return false;
}
}
(None, None) => return true,
_ => return false,
}
}
}
}
//------------ Section -------------------------------------------------------
/// A helper type enumerating the three kinds of record sections.
///
/// See the documentation of [`Message`] for what the three sections are.
///
/// [`Message`]: struct.Message.html
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd)]
pub enum Section {
Answer,
Authority,
Additional,
}
impl Section {
/// Returns the first section.
pub fn first() -> Self {
Section::Answer
}
/// Returns the correct record count for this section.
fn count(self, counts: HeaderCounts) -> u16 {
match self {
Section::Answer => counts.ancount(),
Section::Authority => counts.nscount(),
Section::Additional => counts.arcount(),
}
}
/// Returns the value for the following section or `None` if this is last.
pub(crate) fn next_section(self) -> Option<Self> {
match self {
Section::Answer => Some(Section::Authority),
Section::Authority => Some(Section::Additional),
Section::Additional => None,
}
}
}
//------------ RecordSection -----------------------------------------------
/// An iterator over the records in one of the three record sections.
///
/// The iterator’s item is the result of parsing a raw record represented by
/// [`ParsedRecord`]. This type will allow access to an unparsed record. It
/// can be converted into a concrete [`Record`] via its [`into_record`]
/// method. If parsing the raw record fails, the iterator will return an
/// error once and `None` after that.
///
/// Alternatively, you can trade in a value of this type into a
/// [`RecordIter`] that iterates over [`Record`]s of a specific type by
/// calling the [`limit_to`] method. In particular, you can use this together
/// with [`AllRecordData`] to acquire an iterator that parses all known
/// record types.
///
/// `RecordSection` values cannot be created directly. You can get one either
/// by calling the method for the section in question of a [`Message`] value
/// or by proceeding from another section via its `next_section` method.
///
/// [`limit_to`]: #method.limit_to
/// [`AllRecordData`]: ../../rdata/enum.AllRecordData.html
/// [`Message`]: struct.Message.html
/// [`ParseRecord`]: ../record/struct.ParsedRecord.html
/// [`Record`]: ../record/struct.Record.html
/// [`RecordIter`]: struct.RecordIter.html
/// [`into_record`]: ../record/struct.ParsedRecord.html#method.into_record
#[derive(Clone, Copy, Debug)]
pub struct RecordSection<Ref> {
/// The parser for generating the records.
parser: Parser<Ref>,
/// Which section are we, really?
section: Section,
/// The remaining number of records.
///
/// The `Result` is here to monitor an error during iteration.
/// It is used to fuse the iterator after an error and is also returned
/// by `answer()` should that be called after an error.
count: Result<u16, ParseError>,
}
impl<Ref: OctetsRef> RecordSection<Ref> {
/// Creates a new section from a parser.
///
/// The parser must be positioned at the beginning of this section.
fn new(parser: Parser<Ref>, section: Section) -> Self {
RecordSection {
count: Ok(section
.count(*HeaderCounts::for_message_slice(parser.as_slice()))),
section,
parser,
}
}
/// Returns the current position relative to the beginning of the message.
pub fn pos(&self) -> usize {
self.parser.pos()
}
/// Trades `self` in for an iterator limited to a concrete record type.
///
/// The record type is given through its record data type. Since the data
/// is being parsed, this type must implement [`ParseRecordData`]. For
/// record data types that are generic over domain name types, this is
/// normally achieved by giving them a [`ParsedDname`]. As a convenience,
/// type aliases for all the fundamental record data types exist in the
/// [domain::rdata::parsed] module.
///
/// The returned limited iterator will continue at the current position
/// of `self`. It will *not* start from the beginning of the section.
///
/// [`ParseRecordData`]: ../rdata/trait.ParseRecordData.html
/// [`ParsedDname`]: ../name/struct.ParsedDname.html
/// [domain::rdata::parsed]: ../../rdata/parsed/index.html
pub fn limit_to<Data: ParseRecordData<Ref>>(
self,
) -> RecordIter<Ref, Data> {
RecordIter::new(self, false)
}
/// Trades `self` in for an iterator limited to a type in IN class.
///
/// Behaves exactly like [`limit_to`] but skips over records that are not
/// of class IN.
///
/// [`limit_to`]: #method.limit_to
pub fn limit_to_in<Data: ParseRecordData<Ref>>(
self,
) -> RecordIter<Ref, Data> {
RecordIter::new(self, true)
}
/// Proceeds to the next section if there is one.
///
/// Returns an error if parsing has failed and the message is unusable
/// now.
pub fn next_section(mut self) -> Result<Option<Self>, ParseError> {
let section = match self.section.next_section() {
Some(section) => section,
None => return Ok(None),
};
while self.skip_next().is_some() {}
let _ = self.count?;
Ok(Some(RecordSection::new(self.parser, section)))
}
/// Skip the next record.
fn skip_next(&mut self) -> Option<Result<(), ParseError>> {
match self.count {
Ok(count) if count > 0 => {
match ParsedRecord::skip(&mut self.parser) {
Ok(_) => {
self.count = Ok(count - 1);
Some(Ok(()))
}
Err(err) => {
self.count = Err(err);
Some(Err(err))
}
}
}
_ => None,
}
}
}
//--- Iterator
impl<Ref: OctetsRef> Iterator for RecordSection<Ref> {
type Item = Result<ParsedRecord<Ref>, ParseError>;
fn next(&mut self) -> Option<Self::Item> {
match self.count {
Ok(count) if count > 0 => {
match ParsedRecord::parse(&mut self.parser) {
Ok(record) => {
self.count = Ok(count - 1);
Some(Ok(record))
}
Err(err) => {
self.count = Err(err);
Some(Err(err))
}
}
}
_ => None,
}
}
}
//------------ MessageIter ---------------------------------------------------
/// An iterator over the records of a message.
pub struct MessageIter<Ref> {
inner: Option<RecordSection<Ref>>,
}
impl<Ref: OctetsRef> Iterator for MessageIter<Ref> {
type Item = Result<(ParsedRecord<Ref>, Section), ParseError>;
fn next(&mut self) -> Option<Self::Item> {
// Try to get next record from current section
match self.inner {
Some(ref mut inner) => {
let item = inner.next();
if let Some(item) = item {
return Some(item.map(|item| (item, inner.section)));
}
}
None => return None,
}
// Advance to next section if possible, and retry
let inner = self.inner.take()?;
match inner.next_section() {
Ok(section) => {
self.inner = section;
self.next()
}
Err(err) => Some(Err(err)),
}
}
}
//------------ RecordIter ----------------------------------------------------
/// An iterator over specific records of a record section of a DNS message.
///
/// The iterator’s item type is the result of trying to parse a record.
/// It silently skips over all records that `Data` cannot or does not want to
/// parse. If parsing the record data fails, the iterator will return an
/// error but can continue with the next record. If parsing the entire record
/// fails the item will be an error and subsequent attempts to continue will
/// also produce errors. This case can be distinguished from an error while
/// parsing the record data by [`next_section`] returning an error, too.
///
/// You can create a value of this type through the
/// [`RecordSection::limit_to`] method.
///
/// [`next_section`]: #method.next_section
/// [`RecordSection::limit_to`]: struct.RecordSection.html#method.limit_to
#[derive(Clone, Copy, Debug)]
pub struct RecordIter<Ref, Data> {
section: RecordSection<Ref>,
in_only: bool,
marker: PhantomData<Data>,
}
impl<Ref: OctetsRef, Data: ParseRecordData<Ref>> RecordIter<Ref, Data> {
/// Creates a new record iterator.
fn new(section: RecordSection<Ref>, in_only: bool) -> Self {
RecordIter {
section,
in_only,
marker: PhantomData,
}
}
/// Trades the limited iterator for the full iterator.
///
/// The returned iterator will continue right after the last record
/// previously returned.
pub fn unwrap(self) -> RecordSection<Ref> {
self.section
}
/// Proceeds to the next section if there is one.
///
/// Returns an error if parsing the message has failed. Returns
/// `Ok(None)` if this iterator was already on the additional section.
pub fn next_section(
self,
) -> Result<Option<RecordSection<Ref>>, ParseError> {
self.section.next_section()
}
}
//--- Iterator
impl<Ref, Data> Iterator for RecordIter<Ref, Data>
where
Ref: OctetsRef,
Data: ParseRecordData<Ref>,
{
type Item = Result<Record<ParsedDname<Ref>, Data>, ParseError>;
fn next(&mut self) -> Option<Self::Item> {
loop {
let record = match self.section.next() {
Some(Ok(record)) => record,
Some(Err(err)) => return Some(Err(err)),
None => return None,
};
if self.in_only && record.class() != Class::In {
continue;
}
match record.into_record() {
Ok(Some(record)) => return Some(Ok(record)),
Err(err) => return Some(Err(err)),
Ok(None) => {}
}
}
}
}
//============ Error Types ===================================================
//------------ CopyRecordsError ----------------------------------------------
/// An error occurrd while copying records.
#[derive(Clone, Copy, Debug)]
pub enum CopyRecordsError {
/// Parsing the source message failed.
Parse(ParseError),
/// Not enough space in the target.
ShortBuf,
}
//--- From
impl From<ParseError> for CopyRecordsError {
fn from(err: ParseError) -> Self {
CopyRecordsError::Parse(err)
}
}
impl From<ShortBuf> for CopyRecordsError {
fn from(_: ShortBuf) -> Self {
CopyRecordsError::ShortBuf
}
}
//--- Display and Error
impl fmt::Display for CopyRecordsError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match *self {
CopyRecordsError::Parse(ref err) => err.fmt(f),
CopyRecordsError::ShortBuf => ShortBuf.fmt(f),
}
}
}
#[cfg(feature = "std")]
impl std::error::Error for CopyRecordsError {}
//============ Testing =======================================================
#[cfg(test)]
mod test {
use super::*;
#[cfg(feature = "std")]
use crate::base::message_builder::MessageBuilder;
#[cfg(feature = "std")]
use crate::base::name::Dname;
#[cfg(feature = "std")]
use crate::rdata::{AllRecordData, Ns};
#[cfg(feature = "std")]
use std::vec::Vec;
// Helper for test cases
#[cfg(feature = "std")]
fn get_test_message() -> Message<Vec<u8>> {
let msg = MessageBuilder::new_vec();
let mut msg = msg.answer();
msg.push((
Dname::vec_from_str("foo.example.com.").unwrap(),
86000,
Cname::new(Dname::vec_from_str("baz.example.com.").unwrap()),
))
.unwrap();
let mut msg = msg.authority();
msg.push((
Dname::vec_from_str("bar.example.com.").unwrap(),
86000,
Ns::new(Dname::vec_from_str("baz.example.com.").unwrap()),
))
.unwrap();
msg.into_message()
}
#[test]
fn short_message() {
assert!(Message::from_octets(&[0u8; 11]).is_err());
assert!(Message::from_octets(&[0u8; 12]).is_ok());
}
#[test]
#[cfg(feature = "std")]
fn canonical_name() {
use crate::rdata::A;
// Message without CNAMEs.
let mut msg = MessageBuilder::new_vec().question();
msg.push((Dname::vec_from_str("example.com.").unwrap(), Rtype::A))
.unwrap();
let msg_ref = msg.as_message();
assert_eq!(
Dname::vec_from_str("example.com.").unwrap(),
msg_ref.canonical_name().unwrap()
);
// Message with CNAMEs.
let mut msg = msg.answer();
msg.push((
Dname::vec_from_str("bar.example.com.").unwrap(),
86000,
Cname::new(Dname::vec_from_str("baz.example.com.").unwrap()),
))
.unwrap();
msg.push((
Dname::vec_from_str("example.com.").unwrap(),
86000,
Cname::new(Dname::vec_from_str("foo.example.com.").unwrap()),
))
.unwrap();
msg.push((
Dname::vec_from_str("foo.example.com.").unwrap(),
86000,
Cname::new(Dname::vec_from_str("bar.example.com.").unwrap()),
))
.unwrap();
let msg_ref = msg.as_message();
assert_eq!(
Dname::vec_from_str("baz.example.com.").unwrap(),
msg_ref.canonical_name().unwrap()
);
// CNAME loop.
msg.push((
Dname::vec_from_str("baz.example.com").unwrap(),
86000,
Cname::new(Dname::vec_from_str("foo.example.com").unwrap()),
))
.unwrap();
assert!(msg.as_message().canonical_name().is_none());
msg.push((
Dname::vec_from_str("baz.example.com").unwrap(),
86000,
A::from_octets(127, 0, 0, 1),
))
.unwrap();
assert!(msg.as_message().canonical_name().is_none());
}
#[test]
#[cfg(feature = "std")]
fn message_iterator() {
let msg = get_test_message();
let mut iter = msg.iter();
// Check that it returns a record from first section
let (_rr, section) = iter.next().unwrap().unwrap();
assert_eq!(Section::Answer, section);
// Check that it advances to next section
let (_rr, section) = iter.next().unwrap().unwrap();
assert_eq!(Section::Authority, section);
}
#[test]
#[cfg(feature = "std")]
fn copy_records() {
let msg = get_test_message();
let msg = msg.for_slice();
let target = MessageBuilder::new_vec().question();
let res = msg.copy_records(target.answer(), |rr| {
if let Ok(Some(rr)) =
rr.into_record::<AllRecordData<_, ParsedDname<_>>>()
{
if rr.rtype() == Rtype::Cname {
return Some(rr);
}
}
None
});
assert!(res.is_ok());
if let Ok(target) = res {
let msg = target.into_message();
assert_eq!(1, msg.header_counts().ancount());
assert_eq!(0, msg.header_counts().arcount());
}
}
}