Skip to main content

hickory_net/client/
mod.rs

1/*
2 * Copyright (C) 2015-2016 Benjamin Fry <benjaminfry@me.com>
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     https://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! DNS Client associated classes for performing queries and other operations.
18
19use core::{
20    future::Future,
21    pin::Pin,
22    task::{Context, Poll},
23    time::Duration,
24};
25
26use futures_util::{
27    ready,
28    stream::{Stream, StreamExt},
29};
30use tracing::debug;
31
32use crate::{
33    error::NetError,
34    proto::{
35        ProtoError,
36        op::{
37            DEFAULT_MAX_PAYLOAD_LEN, DnsRequest, DnsRequestOptions, DnsResponse, Edns, Message,
38            OpCode, Query, update_message,
39        },
40        rr::{DNSClass, Name, RData, Record, RecordSet, RecordType, rdata::SOA},
41    },
42    runtime::RuntimeProvider,
43    xfer::{
44        BufDnsStreamHandle, DnsClientStream, DnsExchange, DnsExchangeBackground, DnsExchangeSend,
45        DnsHandle, DnsMultiplexer, DnsRequestSender,
46    },
47};
48
49#[cfg(all(feature = "__dnssec", feature = "tokio"))]
50pub(crate) mod dnssec_client;
51#[cfg(all(feature = "__dnssec", feature = "tokio"))]
52pub use dnssec_client::{AsyncSecureClientBuilder, DnssecClient};
53
54mod memoize_client_handle;
55pub use memoize_client_handle::MemoizeClientHandle;
56
57mod rc_stream;
58
59#[cfg(test)]
60mod tests;
61
62/// A DNS Client implemented over futures-rs.
63///
64/// This Client is generic and capable of wrapping UDP, TCP, and other underlying DNS protocol
65///  implementations.
66#[derive(Clone)]
67pub struct Client<P> {
68    exchange: DnsExchange<P>,
69    use_edns: bool,
70}
71
72impl<P: RuntimeProvider> Client<P> {
73    /// Spawns a new Client Stream. This uses a default timeout of 5 seconds for all requests.
74    ///
75    /// # Arguments
76    ///
77    /// * `stream` - A stream of bytes that can be used to send/receive DNS messages
78    ///   (see TcpClientStream or UdpClientStream)
79    /// * `stream_handle` - The handle for the `stream` on which bytes can be sent/received.
80    pub fn new<S: DnsClientStream>(
81        stream: S,
82        stream_handle: BufDnsStreamHandle,
83    ) -> (Self, DnsExchangeBackground<DnsMultiplexer<S>, P::Timer>) {
84        Self::with_timeout(stream, stream_handle, Duration::from_secs(5))
85    }
86
87    /// Spawns a new Client Stream.
88    ///
89    /// # Arguments
90    ///
91    /// * `stream` - A stream of bytes that can be used to send/receive DNS messages
92    ///   (see TcpClientStream or UdpClientStream)
93    /// * `stream_handle` - The handle for the `stream` on which bytes can be sent/received.
94    /// * `timeout_duration` - All requests may fail due to lack of response, this is the time to
95    ///   wait for a response before canceling the request.
96    pub fn with_timeout<S: DnsClientStream>(
97        stream: S,
98        stream_handle: BufDnsStreamHandle,
99        timeout_duration: Duration,
100    ) -> (Self, DnsExchangeBackground<DnsMultiplexer<S>, P::Timer>) {
101        Self::from_sender(DnsMultiplexer::new(stream, stream_handle).with_timeout(timeout_duration))
102    }
103
104    /// Creates a Client from an existing DnsRequestSender
105    pub fn from_sender<S: DnsRequestSender>(
106        sender: S,
107    ) -> (Self, DnsExchangeBackground<S, P::Timer>) {
108        let (exchange, bg) = DnsExchange::from_stream(sender);
109        (
110            Self {
111                exchange,
112                use_edns: true,
113            },
114            bg,
115        )
116    }
117
118    /// (Re-)enable usage of EDNS for outgoing messages
119    pub fn enable_edns(&mut self) {
120        self.use_edns = true;
121    }
122
123    /// Disable usage of EDNS for outgoing messages
124    pub fn disable_edns(&mut self) {
125        self.use_edns = false;
126    }
127}
128
129impl<P: RuntimeProvider> DnsHandle for Client<P> {
130    type Response = DnsExchangeSend<P>;
131    type Runtime = P;
132
133    fn send(&self, request: DnsRequest) -> Self::Response {
134        self.exchange.send(request)
135    }
136
137    fn is_using_edns(&self) -> bool {
138        self.use_edns
139    }
140}
141
142impl<T> ClientHandle for T where T: DnsHandle {}
143
144/// A trait for implementing high level functions of DNS.
145pub trait ClientHandle: 'static + Clone + DnsHandle + Send {
146    /// A *classic* DNS query
147    ///
148    /// *Note* As of now, this will not recurse on PTR or CNAME record responses, that is up to
149    ///        the caller.
150    ///
151    /// # Arguments
152    ///
153    /// * `name` - the label to lookup
154    /// * `query_class` - most likely this should always be DNSClass::IN
155    /// * `query_type` - record type to lookup
156    fn query(
157        &mut self,
158        name: Name,
159        query_class: DNSClass,
160        query_type: RecordType,
161    ) -> ClientResponse<<Self as DnsHandle>::Response> {
162        let mut query = Query::query(name, query_type);
163        query.set_query_class(query_class);
164        let mut options = DnsRequestOptions::default();
165        options.use_edns = self.is_using_edns();
166        ClientResponse(self.lookup(query, options))
167    }
168
169    /// Sends a NOTIFY message to the remote system
170    ///
171    /// [RFC 1996](https://tools.ietf.org/html/rfc1996), DNS NOTIFY, August 1996
172    ///
173    ///
174    /// ```text
175    /// 1. Rationale and Scope
176    ///
177    ///   1.1. Slow propagation of new and changed data in a DNS zone can be
178    ///   due to a zone's relatively long refresh times.  Longer refresh times
179    ///   are beneficial in that they reduce load on the Primary Zone Servers, but
180    ///   that benefit comes at the cost of long intervals of incoherence among
181    ///   authority servers whenever the zone is updated.
182    ///
183    ///   1.2. The DNS NOTIFY transaction allows Primary Zone Servers to inform Secondary
184    ///   Zone Servers when the zone has changed -- an interrupt as opposed to poll
185    ///   model -- which it is hoped will reduce propagation delay while not
186    ///   unduly increasing the masters' load.  This specification only allows
187    ///   slaves to be notified of SOA RR changes, but the architecture of
188    ///   NOTIFY is intended to be extensible to other RR types.
189    ///
190    ///   1.3. This document intentionally gives more definition to the roles
191    ///   of "Primary", "Secondary" and "Stealth" servers, their enumeration in NS
192    ///   RRs, and the SOA MNAME field.  In that sense, this document can be
193    ///   considered an addendum to [RFC1035].
194    ///
195    /// ```
196    ///
197    /// The below section describes how the Notify message should be constructed. The function
198    ///  implementation accepts a Record, but the actual data of the record should be ignored by the
199    ///  server, i.e. the server should make a request subsequent to receiving this Notification for
200    ///  the authority record, but could be used to decide to request an update or not:
201    ///
202    /// ```text
203    ///   3.7. A NOTIFY request has QDCOUNT>0, ANCOUNT>=0, AUCOUNT>=0,
204    ///   ADCOUNT>=0.  If ANCOUNT>0, then the answer section represents an
205    ///   unsecure hint at the new RRset for this <QNAME,QCLASS,QTYPE>.  A
206    ///   Secondary receiving such a hint is free to treat equivalence of this
207    ///   answer section with its local data as a "no further work needs to be
208    ///   done" indication.  If ANCOUNT=0, or ANCOUNT>0 and the answer section
209    ///   differs from the Secondary's local data, then the Secondary should query its
210    ///   known Primaries to retrieve the new data.
211    /// ```
212    ///
213    /// Client's should be ready to handle, or be aware of, a server response of NOTIMP:
214    ///
215    /// ```text
216    ///   3.12. If a NOTIFY request is received by a Secondary who does not
217    ///   implement the NOTIFY opcode, it will respond with a NOTIMP
218    ///   (unimplemented feature error) message.  A Primary Zone Server who receives
219    ///   such a NOTIMP should consider the NOTIFY transaction complete for
220    ///   that Secondary.
221    /// ```
222    ///
223    /// # Arguments
224    ///
225    /// * `name` - the label which is being notified
226    /// * `query_class` - most likely this should always be DNSClass::IN
227    /// * `query_type` - record type which has been updated
228    /// * `rrset` - the new version of the record(s) being notified
229    fn notify<R>(
230        &mut self,
231        name: Name,
232        query_class: DNSClass,
233        query_type: RecordType,
234        rrset: Option<R>,
235    ) -> ClientResponse<<Self as DnsHandle>::Response>
236    where
237        R: Into<RecordSet>,
238    {
239        debug!("notifying: {} {:?}", name, query_type);
240
241        // build the message
242        let mut message = Message::query();
243        // 3.3. NOTIFY is similar to QUERY in that it has a request message with
244        // the header QR flag "clear" and a response message with QR "set".  The
245        // response message contains no useful information, but its reception by
246        // the Primary is an indication that the Secondary has received the NOTIFY
247        // and that the Primary Zone Server can remove the Secondary from any retry queue for
248        // this NOTIFY event.
249        message.metadata.op_code = OpCode::Notify;
250
251        // Extended dns
252        if self.is_using_edns() {
253            message
254                .edns
255                .get_or_insert_with(Edns::new)
256                .set_max_payload(DEFAULT_MAX_PAYLOAD_LEN)
257                .set_version(0);
258        }
259
260        // add the query
261        let mut query: Query = Query::new();
262        query
263            .set_name(name)
264            .set_query_class(query_class)
265            .set_query_type(query_type);
266        message.add_query(query);
267
268        // add the notify message, see https://tools.ietf.org/html/rfc1996, section 3.7
269        if let Some(rrset) = rrset {
270            message.add_answers(rrset.into());
271        }
272
273        ClientResponse(self.send(DnsRequest::from(message)))
274    }
275
276    /// Sends a record to create on the server, this will fail if the record exists (atomicity
277    ///  depends on the server)
278    ///
279    /// [RFC 2136](https://tools.ietf.org/html/rfc2136), DNS Update, April 1997
280    ///
281    /// ```text
282    ///  2.4.3 - RRset Does Not Exist
283    ///
284    ///   No RRs with a specified NAME and TYPE (in the zone and class denoted
285    ///   by the Zone Section) can exist.
286    ///
287    ///   For this prerequisite, a requestor adds to the section a single RR
288    ///   whose NAME and TYPE are equal to that of the RRset whose nonexistence
289    ///   is required.  The RDLENGTH of this record is zero (0), and RDATA
290    ///   field is therefore empty.  CLASS must be specified as NONE in order
291    ///   to distinguish this condition from a valid RR whose RDLENGTH is
292    ///   naturally zero (0) (for example, the NULL RR).  TTL must be specified
293    ///   as zero (0).
294    ///
295    /// 2.5.1 - Add To An RRset
296    ///
297    ///    RRs are added to the Update Section whose NAME, TYPE, TTL, RDLENGTH
298    ///    and RDATA are those being added, and CLASS is the same as the zone
299    ///    class.  Any duplicate RRs will be silently ignored by the Primary Zone
300    ///    Server.
301    /// ```
302    ///
303    /// # Arguments
304    ///
305    /// * `rrset` - the record(s) to create
306    /// * `zone_origin` - the zone name to update, i.e. SOA name
307    ///
308    /// The update must go to a zone authority (i.e. the server used in the ClientConnection)
309    fn create<R>(
310        &mut self,
311        rrset: R,
312        zone_origin: Name,
313    ) -> ClientResponse<<Self as DnsHandle>::Response>
314    where
315        R: Into<RecordSet>,
316    {
317        let rrset = rrset.into();
318        let message = update_message::create(rrset, zone_origin, self.is_using_edns());
319        ClientResponse(self.send(DnsRequest::from(message)))
320    }
321
322    /// Appends a record to an existing rrset, optionally require the rrset to exist (atomicity
323    ///  depends on the server)
324    ///
325    /// [RFC 2136](https://tools.ietf.org/html/rfc2136), DNS Update, April 1997
326    ///
327    /// ```text
328    /// 2.4.1 - RRset Exists (Value Independent)
329    ///
330    ///   At least one RR with a specified NAME and TYPE (in the zone and class
331    ///   specified in the Zone Section) must exist.
332    ///
333    ///   For this prerequisite, a requestor adds to the section a single RR
334    ///   whose NAME and TYPE are equal to that of the zone RRset whose
335    ///   existence is required.  RDLENGTH is zero and RDATA is therefore
336    ///   empty.  CLASS must be specified as ANY to differentiate this
337    ///   condition from that of an actual RR whose RDLENGTH is naturally zero
338    ///   (0) (e.g., NULL).  TTL is specified as zero (0).
339    ///
340    /// 2.5.1 - Add To An RRset
341    ///
342    ///    RRs are added to the Update Section whose NAME, TYPE, TTL, RDLENGTH
343    ///    and RDATA are those being added, and CLASS is the same as the zone
344    ///    class.  Any duplicate RRs will be silently ignored by the Primary Zone
345    ///    Server.
346    /// ```
347    ///
348    /// # Arguments
349    ///
350    /// * `rrset` - the record(s) to append to an RRSet
351    /// * `zone_origin` - the zone name to update, i.e. SOA name
352    /// * `must_exist` - if true, the request will fail if the record does not exist
353    ///
354    /// The update must go to a zone authority (i.e. the server used in the ClientConnection). If
355    /// the rrset does not exist and must_exist is false, then the RRSet will be created.
356    fn append<R>(
357        &mut self,
358        rrset: R,
359        zone_origin: Name,
360        must_exist: bool,
361    ) -> ClientResponse<<Self as DnsHandle>::Response>
362    where
363        R: Into<RecordSet>,
364    {
365        let rrset = rrset.into();
366        let message = update_message::append(rrset, zone_origin, must_exist, self.is_using_edns());
367        ClientResponse(self.send(DnsRequest::from(message)))
368    }
369
370    /// Compares and if it matches, swaps it for the new value (atomicity depends on the server)
371    ///
372    /// ```text
373    ///  2.4.2 - RRset Exists (Value Dependent)
374    ///
375    ///   A set of RRs with a specified NAME and TYPE exists and has the same
376    ///   members with the same RDATAs as the RRset specified here in this
377    ///   section.  While RRset ordering is undefined and therefore not
378    ///   significant to this comparison, the sets be identical in their
379    ///   extent.
380    ///
381    ///   For this prerequisite, a requestor adds to the section an entire
382    ///   RRset whose preexistence is required.  NAME and TYPE are that of the
383    ///   RRset being denoted.  CLASS is that of the zone.  TTL must be
384    ///   specified as zero (0) and is ignored when comparing RRsets for
385    ///   identity.
386    ///
387    ///  2.5.4 - Delete An RR From An RRset
388    ///
389    ///   RRs to be deleted are added to the Update Section.  The NAME, TYPE,
390    ///   RDLENGTH and RDATA must match the RR being deleted.  TTL must be
391    ///   specified as zero (0) and will otherwise be ignored by the Primary
392    ///   Zone Server.  CLASS must be specified as NONE to distinguish this from an
393    ///   RR addition.  If no such RRs exist, then this Update RR will be
394    ///   silently ignored by the Primary Zone Server.
395    ///
396    ///  2.5.1 - Add To An RRset
397    ///
398    ///   RRs are added to the Update Section whose NAME, TYPE, TTL, RDLENGTH
399    ///   and RDATA are those being added, and CLASS is the same as the zone
400    ///   class.  Any duplicate RRs will be silently ignored by the Primary
401    ///   Zone Server.
402    /// ```
403    ///
404    /// # Arguments
405    ///
406    /// * `current` - the current rrset which must exist for the swap to complete
407    /// * `new` - the new rrset with which to replace the current rrset
408    /// * `zone_origin` - the zone name to update, i.e. SOA name
409    ///
410    /// The update must go to a zone authority (i.e. the server used in the ClientConnection).
411    fn compare_and_swap<C, N>(
412        &mut self,
413        current: C,
414        new: N,
415        zone_origin: Name,
416    ) -> ClientResponse<<Self as DnsHandle>::Response>
417    where
418        C: Into<RecordSet>,
419        N: Into<RecordSet>,
420    {
421        let current = current.into();
422        let new = new.into();
423
424        let message =
425            update_message::compare_and_swap(current, new, zone_origin, self.is_using_edns());
426        ClientResponse(self.send(DnsRequest::from(message)))
427    }
428
429    /// Deletes a record (by rdata) from an rrset, optionally require the rrset to exist.
430    ///
431    /// [RFC 2136](https://tools.ietf.org/html/rfc2136), DNS Update, April 1997
432    ///
433    /// ```text
434    /// 2.4.1 - RRset Exists (Value Independent)
435    ///
436    ///   At least one RR with a specified NAME and TYPE (in the zone and class
437    ///   specified in the Zone Section) must exist.
438    ///
439    ///   For this prerequisite, a requestor adds to the section a single RR
440    ///   whose NAME and TYPE are equal to that of the zone RRset whose
441    ///   existence is required.  RDLENGTH is zero and RDATA is therefore
442    ///   empty.  CLASS must be specified as ANY to differentiate this
443    ///   condition from that of an actual RR whose RDLENGTH is naturally zero
444    ///   (0) (e.g., NULL).  TTL is specified as zero (0).
445    ///
446    /// 2.5.4 - Delete An RR From An RRset
447    ///
448    ///   RRs to be deleted are added to the Update Section.  The NAME, TYPE,
449    ///   RDLENGTH and RDATA must match the RR being deleted.  TTL must be
450    ///   specified as zero (0) and will otherwise be ignored by the Primary
451    ///   Zone Server.  CLASS must be specified as NONE to distinguish this from an
452    ///   RR addition.  If no such RRs exist, then this Update RR will be
453    ///   silently ignored by the Primary Zone Server.
454    /// ```
455    ///
456    /// # Arguments
457    ///
458    /// * `rrset` - the record(s) to delete from a RRSet, the name, type and rdata must match the
459    ///   record to delete
460    /// * `zone_origin` - the zone name to update, i.e. SOA name
461    /// * `signer` - the signer, with private key, to use to sign the request
462    ///
463    /// The update must go to a zone authority (i.e. the server used in the ClientConnection). If
464    /// the rrset does not exist and must_exist is false, then the RRSet will be deleted.
465    fn delete_by_rdata<R>(
466        &mut self,
467        rrset: R,
468        zone_origin: Name,
469    ) -> ClientResponse<<Self as DnsHandle>::Response>
470    where
471        R: Into<RecordSet>,
472    {
473        let rrset = rrset.into();
474        let message = update_message::delete_by_rdata(rrset, zone_origin, self.is_using_edns());
475        ClientResponse(self.send(DnsRequest::from(message)))
476    }
477
478    /// Deletes an entire rrset, optionally require the rrset to exist.
479    ///
480    /// [RFC 2136](https://tools.ietf.org/html/rfc2136), DNS Update, April 1997
481    ///
482    /// ```text
483    /// 2.4.1 - RRset Exists (Value Independent)
484    ///
485    ///   At least one RR with a specified NAME and TYPE (in the zone and class
486    ///   specified in the Zone Section) must exist.
487    ///
488    ///   For this prerequisite, a requestor adds to the section a single RR
489    ///   whose NAME and TYPE are equal to that of the zone RRset whose
490    ///   existence is required.  RDLENGTH is zero and RDATA is therefore
491    ///   empty.  CLASS must be specified as ANY to differentiate this
492    ///   condition from that of an actual RR whose RDLENGTH is naturally zero
493    ///   (0) (e.g., NULL).  TTL is specified as zero (0).
494    ///
495    /// 2.5.2 - Delete An RRset
496    ///
497    ///   One RR is added to the Update Section whose NAME and TYPE are those
498    ///   of the RRset to be deleted.  TTL must be specified as zero (0) and is
499    ///   otherwise not used by the Primary Zone Server.  CLASS must be specified as
500    ///   ANY.  RDLENGTH must be zero (0) and RDATA must therefore be empty.
501    ///   If no such RRset exists, then this Update RR will be silently ignored
502    ///   by the Primary Zone Server.
503    /// ```
504    ///
505    /// # Arguments
506    ///
507    /// * `record` - The name, class and record_type will be used to match and delete the RecordSet
508    /// * `zone_origin` - the zone name to update, i.e. SOA name
509    ///
510    /// The update must go to a zone authority (i.e. the server used in the ClientConnection). If
511    /// the rrset does not exist and must_exist is false, then the RRSet will be deleted.
512    fn delete_rrset(
513        &mut self,
514        record: Record,
515        zone_origin: Name,
516    ) -> ClientResponse<<Self as DnsHandle>::Response> {
517        assert!(zone_origin.zone_of(&record.name));
518        let message = update_message::delete_rrset(record, zone_origin, self.is_using_edns());
519        ClientResponse(self.send(DnsRequest::from(message)))
520    }
521
522    /// Deletes all records at the specified name
523    ///
524    /// [RFC 2136](https://tools.ietf.org/html/rfc2136), DNS Update, April 1997
525    ///
526    /// ```text
527    /// 2.5.3 - Delete All RRsets From A Name
528    ///
529    ///   One RR is added to the Update Section whose NAME is that of the name
530    ///   to be cleansed of RRsets.  TYPE must be specified as ANY.  TTL must
531    ///   be specified as zero (0) and is otherwise not used by the Primary
532    ///   Zone Server.  CLASS must be specified as ANY.  RDLENGTH must be zero (0)
533    ///   and RDATA must therefore be empty.  If no such RRsets exist, then
534    ///   this Update RR will be silently ignored by the Primary Zone Server.
535    /// ```
536    ///
537    /// # Arguments
538    ///
539    /// * `name_of_records` - the name of all the record sets to delete
540    /// * `zone_origin` - the zone name to update, i.e. SOA name
541    /// * `dns_class` - the class of the SOA
542    ///
543    /// The update must go to a zone authority (i.e. the server used in the ClientConnection). This
544    /// operation attempts to delete all resource record sets the specified name regardless of
545    /// the record type.
546    fn delete_all(
547        &mut self,
548        name_of_records: Name,
549        zone_origin: Name,
550        dns_class: DNSClass,
551    ) -> ClientResponse<<Self as DnsHandle>::Response> {
552        assert!(zone_origin.zone_of(&name_of_records));
553        let message = update_message::delete_all(
554            name_of_records,
555            zone_origin,
556            dns_class,
557            self.is_using_edns(),
558        );
559        ClientResponse(self.send(DnsRequest::from(message)))
560    }
561
562    /// Download all records from a zone, or all records modified since given SOA was observed.
563    /// The request will either be a AXFR Query (ask for full zone transfer) if a SOA was not
564    /// provided, or a IXFR Query (incremental zone transfer) if a SOA was provided.
565    ///
566    /// # Arguments
567    /// * `zone_origin` - the zone name to update, i.e. SOA name
568    /// * `last_soa` - the last SOA known, if any. If provided, name must match `zone_origin`
569    fn zone_transfer(
570        &mut self,
571        zone_origin: Name,
572        last_soa: Option<SOA>,
573    ) -> ClientStreamXfr<<Self as DnsHandle>::Response> {
574        let ixfr = last_soa.is_some();
575        let message = update_message::zone_transfer(zone_origin, last_soa);
576        ClientStreamXfr::new(self.send(DnsRequest::from(message)), ixfr)
577    }
578}
579
580/// A stream result of a Client Request
581#[must_use = "stream do nothing unless polled"]
582pub struct ClientStreamingResponse<R>(pub(crate) R)
583where
584    R: Stream<Item = Result<DnsResponse, ProtoError>> + Send + Unpin + 'static;
585
586impl<R> Stream for ClientStreamingResponse<R>
587where
588    R: Stream<Item = Result<DnsResponse, ProtoError>> + Send + Unpin + 'static,
589{
590    type Item = Result<DnsResponse, ProtoError>;
591
592    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
593        Poll::Ready(ready!(self.0.poll_next_unpin(cx)))
594    }
595}
596
597/// A future result of a Client Request
598#[must_use = "futures do nothing unless polled"]
599pub struct ClientResponse<R>(pub(crate) R)
600where
601    R: Stream<Item = Result<DnsResponse, NetError>> + Send + Unpin + 'static;
602
603impl<R> Future for ClientResponse<R>
604where
605    R: Stream<Item = Result<DnsResponse, NetError>> + Send + Unpin + 'static,
606{
607    type Output = Result<DnsResponse, NetError>;
608
609    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
610        Poll::Ready(match ready!(self.0.poll_next_unpin(cx)) {
611            Some(r) => r,
612            None => Err(NetError::Timeout),
613        })
614    }
615}
616
617/// A stream result of a zone transfer Client Request
618/// Accept messages until the end of a zone transfer. For AXFR, it search for a starting and an
619/// ending SOA. For IXFR, it do so taking into account there will be other SOA inbetween
620#[must_use = "stream do nothing unless polled"]
621pub struct ClientStreamXfr<R>
622where
623    R: Stream<Item = Result<DnsResponse, NetError>> + Send + Unpin + 'static,
624{
625    state: ClientStreamXfrState<R>,
626}
627
628impl<R> ClientStreamXfr<R>
629where
630    R: Stream<Item = Result<DnsResponse, NetError>> + Send + Unpin + 'static,
631{
632    fn new(inner: R, maybe_incr: bool) -> Self {
633        Self {
634            state: ClientStreamXfrState::Start { inner, maybe_incr },
635        }
636    }
637}
638
639/// State machine for ClientStreamXfr, implementing almost all logic
640#[derive(Debug)]
641enum ClientStreamXfrState<R> {
642    Start {
643        inner: R,
644        maybe_incr: bool,
645    },
646    Second {
647        inner: R,
648        expected_serial: u32,
649        maybe_incr: bool,
650    },
651    Axfr {
652        inner: R,
653        expected_serial: u32,
654    },
655    Ixfr {
656        inner: R,
657        even: bool,
658        expected_serial: u32,
659    },
660    Ended,
661    Invalid,
662}
663
664impl<R> ClientStreamXfrState<R> {
665    /// Helper to get the stream from the enum
666    fn inner(&mut self) -> &mut R {
667        use ClientStreamXfrState::*;
668        match self {
669            Start { inner, .. } => inner,
670            Second { inner, .. } => inner,
671            Axfr { inner, .. } => inner,
672            Ixfr { inner, .. } => inner,
673            Ended | Invalid => unreachable!(),
674        }
675    }
676
677    /// Helper to ingest answer Records
678    // TODO: this is complex enough it should get its own tests
679    fn process(&mut self, answers: &[Record]) -> Result<(), NetError> {
680        use ClientStreamXfrState::*;
681        fn get_serial(r: &Record) -> Option<u32> {
682            match &r.data {
683                RData::SOA(soa) => Some(soa.serial),
684                _ => None,
685            }
686        }
687
688        if answers.is_empty() {
689            return Ok(());
690        }
691        match core::mem::replace(self, Invalid) {
692            Start { inner, maybe_incr } => {
693                if let Some(expected_serial) = get_serial(&answers[0]) {
694                    *self = Second {
695                        inner,
696                        maybe_incr,
697                        expected_serial,
698                    };
699                    self.process(&answers[1..])
700                } else {
701                    *self = Ended;
702                    Ok(())
703                }
704            }
705            Second {
706                inner,
707                maybe_incr,
708                expected_serial,
709            } => {
710                if let Some(serial) = get_serial(&answers[0]) {
711                    // maybe IXFR, or empty AXFR
712                    if serial == expected_serial {
713                        // empty AXFR
714                        *self = Ended;
715                        if answers.len() == 1 {
716                            Ok(())
717                        } else {
718                            // invalid answer : trailing records
719                            Err("invalid zone transfer, contains trailing records".into())
720                        }
721                    } else if maybe_incr {
722                        *self = Ixfr {
723                            inner,
724                            expected_serial,
725                            even: true,
726                        };
727                        self.process(&answers[1..])
728                    } else {
729                        *self = Ended;
730                        Err("invalid zone transfer, expected AXFR, got IXFR".into())
731                    }
732                } else {
733                    // standard AXFR
734                    *self = Axfr {
735                        inner,
736                        expected_serial,
737                    };
738                    self.process(&answers[1..])
739                }
740            }
741            Axfr {
742                inner,
743                expected_serial,
744            } => {
745                let soa_count = answers
746                    .iter()
747                    .filter(|a| a.record_type() == RecordType::SOA)
748                    .count();
749                match soa_count {
750                    0 => {
751                        *self = Axfr {
752                            inner,
753                            expected_serial,
754                        };
755                        Ok(())
756                    }
757                    1 => {
758                        *self = Ended;
759                        match answers.last().map(|r| r.record_type()) {
760                            Some(RecordType::SOA) => Ok(()),
761                            _ => Err("invalid zone transfer, contains trailing records".into()),
762                        }
763                    }
764                    _ => {
765                        *self = Ended;
766                        Err("invalid zone transfer, contains trailing records".into())
767                    }
768                }
769            }
770            Ixfr {
771                inner,
772                even,
773                expected_serial,
774            } => {
775                let even = answers
776                    .iter()
777                    .fold(even, |even, a| even ^ (a.record_type() == RecordType::SOA));
778                if even {
779                    if let Some(serial) = get_serial(answers.last().unwrap()) {
780                        if serial == expected_serial {
781                            *self = Ended;
782                            return Ok(());
783                        }
784                    }
785                }
786                *self = Ixfr {
787                    inner,
788                    even,
789                    expected_serial,
790                };
791                Ok(())
792            }
793            Ended | Invalid => {
794                unreachable!();
795            }
796        }
797    }
798}
799
800impl<R> Stream for ClientStreamXfr<R>
801where
802    R: Stream<Item = Result<DnsResponse, NetError>> + Send + Unpin + 'static,
803{
804    type Item = Result<DnsResponse, NetError>;
805
806    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
807        use ClientStreamXfrState::*;
808
809        if matches!(self.state, Ended) {
810            return Poll::Ready(None);
811        }
812
813        let message = ready!(self.state.inner().poll_next_unpin(cx)).map(|response| {
814            let ok = response?;
815            self.state.process(&ok.answers)?;
816            Ok(ok)
817        });
818        Poll::Ready(message)
819    }
820}