Skip to main content

hickory_server/zone_handler/
auth_lookup.rs

1// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// https://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// https://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8use std::mem;
9use std::slice::Iter;
10use std::sync::Arc;
11
12use crate::proto::{
13    op::Message,
14    rr::{Record, RecordSet, RecordType, RrsetRecords},
15};
16#[cfg(feature = "resolver")]
17use crate::resolver::lookup::Lookup;
18use crate::zone_handler::LookupOptions;
19
20/// The result of a lookup on a ZoneHandler
21#[allow(clippy::large_enum_variant)]
22#[derive(Debug, Default)]
23#[non_exhaustive]
24pub enum AuthLookup {
25    /// No records
26    #[default]
27    Empty,
28    // TODO: change the result of a lookup to a set of chained iterators...
29    /// Records
30    Records {
31        /// Authoritative answers
32        answers: LookupRecords,
33        /// Optional set of LookupRecords
34        additionals: Option<LookupRecords>,
35    },
36    /// Records resulting from a resolver lookup
37    #[cfg(feature = "resolver")]
38    Resolved(Lookup),
39    /// A response message
40    Response(Message),
41}
42
43impl AuthLookup {
44    /// Construct an answer with additional section
45    pub fn answers(answers: LookupRecords, additionals: Option<LookupRecords>) -> Self {
46        Self::Records {
47            answers,
48            additionals,
49        }
50    }
51
52    /// Returns true if either the associated Records are empty, or this is a NameExists or NxDomain
53    pub fn is_empty(&self) -> bool {
54        // TODO: this needs to be cheap
55        self.was_empty()
56    }
57
58    /// This is an NxDomain or NameExists, and has no associated records
59    ///
60    /// this consumes the iterator, and verifies it is empty
61    pub fn was_empty(&self) -> bool {
62        self.iter().count() == 0
63    }
64
65    /// Conversion to an iterator
66    pub fn iter(&self) -> AuthLookupIter<'_> {
67        self.into_iter()
68    }
69
70    /// Does not panic, but will return no records if it is not of that type
71    pub fn unwrap_records(self) -> LookupRecords {
72        match self {
73            // TODO: this is ugly, what about the additionals?
74            Self::Records { answers, .. } => answers,
75            _ => LookupRecords::default(),
76        }
77    }
78
79    /// Iterates over the records from the Authority section, if present.
80    pub fn authorities(&self) -> Option<LookupRecordsIter<'_>> {
81        match self {
82            Self::Response(message) => {
83                Some(LookupRecordsIter::SliceIter(message.authorities.iter()))
84            }
85            _ => None,
86        }
87    }
88
89    /// Iterates over the records from the Additional section, if present.
90    pub fn additionals(&self) -> Option<LookupRecordsIter<'_>> {
91        match self {
92            Self::Records { additionals, .. } => additionals.as_ref().map(|l| l.iter()),
93            Self::Response(message) => {
94                Some(LookupRecordsIter::SliceIter(message.additionals.iter()))
95            }
96            _ => None,
97        }
98    }
99
100    /// Takes the additional records, leaving behind None
101    pub fn take_additionals(&mut self) -> Option<LookupRecords> {
102        match self {
103            Self::Records { additionals, .. } => additionals.take(),
104            Self::Response(message) => {
105                Some(LookupRecords::Section(mem::take(&mut message.additionals)))
106            }
107            _ => None,
108        }
109    }
110}
111
112#[cfg(feature = "resolver")]
113impl From<Lookup> for AuthLookup {
114    fn from(lookup: Lookup) -> Self {
115        Self::Resolved(lookup)
116    }
117}
118
119impl<'a> IntoIterator for &'a AuthLookup {
120    type Item = &'a Record;
121    type IntoIter = AuthLookupIter<'a>;
122
123    fn into_iter(self) -> Self::IntoIter {
124        match self {
125            AuthLookup::Empty => AuthLookupIter::Empty,
126            // TODO: what about the additionals? is IntoIterator a bad idea?
127            AuthLookup::Records { answers: r, .. } => AuthLookupIter::Records(r.into_iter()),
128            #[cfg(feature = "resolver")]
129            AuthLookup::Resolved(lookup) => AuthLookupIter::Response(lookup.answers().iter()),
130            AuthLookup::Response(message) => AuthLookupIter::Response(message.answers.iter()),
131        }
132    }
133}
134
135/// An iterator over a ZoneHandler lookup
136#[derive(Default)]
137pub enum AuthLookupIter<'r> {
138    /// The empty set
139    #[default]
140    Empty,
141    /// An iteration over a set of Records
142    Records(LookupRecordsIter<'r>),
143    /// An iterator over the answer section of a response message
144    Response(Iter<'r, Record>),
145}
146
147impl<'r> Iterator for AuthLookupIter<'r> {
148    type Item = &'r Record;
149
150    fn next(&mut self) -> Option<Self::Item> {
151        match self {
152            AuthLookupIter::Empty => None,
153            AuthLookupIter::Records(i) => i.next(),
154            AuthLookupIter::Response(i) => i.next(),
155        }
156    }
157}
158
159impl From<LookupRecords> for AuthLookup {
160    fn from(lookup: LookupRecords) -> Self {
161        Self::Records {
162            answers: lookup,
163            additionals: None,
164        }
165    }
166}
167
168/// A collection of [`RecordSet`]s for an AXFR.
169///
170/// This omits the SOA record during iteration.
171#[derive(Debug)]
172pub struct AxfrRecords {
173    dnssec_ok: bool,
174    rrsets: Vec<Arc<RecordSet>>,
175}
176
177impl AxfrRecords {
178    /// Construct this wrapper around the contents of a zone.
179    pub fn new(dnssec_ok: bool, rrsets: Vec<Arc<RecordSet>>) -> Self {
180        Self { dnssec_ok, rrsets }
181    }
182
183    fn iter(&self) -> AxfrRecordsIter<'_> {
184        self.into_iter()
185    }
186}
187
188impl<'r> IntoIterator for &'r AxfrRecords {
189    type Item = &'r Record;
190    type IntoIter = AxfrRecordsIter<'r>;
191
192    fn into_iter(self) -> Self::IntoIter {
193        AxfrRecordsIter {
194            dnssec_ok: self.dnssec_ok,
195            rrsets: self.rrsets.iter(),
196            records: None,
197        }
198    }
199}
200
201/// An iterator over all records in a zone, except the SOA record.
202pub struct AxfrRecordsIter<'r> {
203    #[cfg_attr(not(feature = "__dnssec"), allow(dead_code))]
204    dnssec_ok: bool,
205    rrsets: Iter<'r, Arc<RecordSet>>,
206    records: Option<RrsetRecords<'r>>,
207}
208
209impl<'r> Iterator for AxfrRecordsIter<'r> {
210    type Item = &'r Record;
211
212    fn next(&mut self) -> Option<Self::Item> {
213        loop {
214            if let Some(records) = &mut self.records {
215                if let Some(record) = records
216                    .by_ref()
217                    .find(|record| record.record_type() != RecordType::SOA)
218                {
219                    return Some(record);
220                }
221            }
222
223            // Return if there are no more RRsets.
224            let rrset = self.rrsets.next()?;
225
226            #[cfg(feature = "__dnssec")]
227            let records = rrset.records(self.dnssec_ok);
228
229            #[cfg(not(feature = "__dnssec"))]
230            let records = rrset.records_without_rrsigs();
231
232            self.records = Some(records);
233        }
234    }
235}
236
237/// The result of a lookup
238#[derive(Debug, Default)]
239pub enum LookupRecords {
240    /// The empty set of records
241    #[default]
242    Empty,
243    /// The associate records
244    Records {
245        /// LookupOptions for the request, e.g. dnssec
246        lookup_options: LookupOptions,
247        /// the records found based on the query
248        records: Arc<RecordSet>,
249    },
250    /// Vec of disjoint record sets
251    ManyRecords(LookupOptions, Vec<Arc<RecordSet>>),
252    /// A section from a response message
253    Section(Vec<Record>),
254}
255
256impl LookupRecords {
257    /// Construct a new LookupRecords
258    pub fn new(lookup_options: LookupOptions, records: Arc<RecordSet>) -> Self {
259        Self::Records {
260            lookup_options,
261            records,
262        }
263    }
264
265    /// Construct a new LookupRecords over a set of RecordSets
266    pub fn many(lookup_options: LookupOptions, mut records: Vec<Arc<RecordSet>>) -> Self {
267        // we're reversing the records because they are output in reverse order, via pop()
268        records.reverse();
269        Self::ManyRecords(lookup_options, records)
270    }
271
272    /// This is an NxDomain or NameExists, and has no associated records
273    ///
274    /// this consumes the iterator, and verifies it is empty
275    pub fn was_empty(&self) -> bool {
276        self.iter().count() == 0
277    }
278
279    /// Conversion to an iterator
280    pub fn iter(&self) -> LookupRecordsIter<'_> {
281        self.into_iter()
282    }
283}
284
285impl<'a> IntoIterator for &'a LookupRecords {
286    type Item = &'a Record;
287    type IntoIter = LookupRecordsIter<'a>;
288
289    fn into_iter(self) -> Self::IntoIter {
290        match self {
291            LookupRecords::Empty => LookupRecordsIter::Empty,
292            LookupRecords::Records {
293                lookup_options,
294                records,
295            } => LookupRecordsIter::RecordsIter(lookup_options.rrset_with_rrigs(records)),
296            LookupRecords::ManyRecords(lookup_options, r) => LookupRecordsIter::ManyRecordsIter(
297                r.iter()
298                    .map(|r| lookup_options.rrset_with_rrigs(r))
299                    .collect(),
300                None,
301            ),
302            LookupRecords::Section(vec) => LookupRecordsIter::SliceIter(vec.iter()),
303        }
304    }
305}
306
307/// Iterator over lookup records
308#[derive(Default)]
309pub enum LookupRecordsIter<'r> {
310    /// An iteration over a single RecordSet
311    RecordsIter(RrsetRecords<'r>),
312    /// An iteration over many rrsets
313    ManyRecordsIter(Vec<RrsetRecords<'r>>, Option<RrsetRecords<'r>>),
314    /// An iteration over a slice of records
315    SliceIter(Iter<'r, Record>),
316    /// An empty set
317    #[default]
318    Empty,
319}
320
321impl<'r> Iterator for LookupRecordsIter<'r> {
322    type Item = &'r Record;
323
324    fn next(&mut self) -> Option<Self::Item> {
325        match self {
326            LookupRecordsIter::Empty => None,
327            LookupRecordsIter::RecordsIter(current) => current.next(),
328            LookupRecordsIter::SliceIter(current) => current.next(),
329            LookupRecordsIter::ManyRecordsIter(set, current) => loop {
330                if let Some(o) = current.as_mut().and_then(Iterator::next) {
331                    return Some(o);
332                }
333
334                *current = set.pop();
335                if current.is_none() {
336                    return None;
337                }
338            },
339        }
340    }
341}
342
343/// A copy of all data in a zone.
344///
345/// This is used in the AXFR sub-protocol.
346#[derive(Debug)]
347pub struct ZoneTransfer {
348    /// The SOA record, plus its RRSIG.
349    ///
350    /// This is sent at the start of the first message of the response.
351    pub start_soa: LookupRecords,
352    /// All the records in the zone.
353    pub records: AxfrRecords,
354    /// The SOA record again.
355    ///
356    /// This is sent at the end of the last message of the response.
357    pub end_soa: LookupRecords,
358}
359
360impl ZoneTransfer {
361    /// Iterate over all the records, starting and ending with the SOA record.
362    pub fn iter(&self) -> impl Iterator<Item = &Record> {
363        self.start_soa
364            .iter()
365            .chain(self.records.iter())
366            .chain(self.end_soa.iter())
367    }
368}