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}