edge_mdns/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(clippy::large_futures)]
3#![allow(async_fn_in_trait)]
4#![allow(clippy::uninlined_format_args)]
5#![allow(unknown_lints)]
6
7use core::cmp::Ordering;
8use core::fmt::Display;
9use core::ops::RangeBounds;
10
11use domain::base::header::Flags;
12use domain::base::iana::{Opcode, Rcode};
13use domain::base::message::ShortMessage;
14use domain::base::message_builder::PushError;
15use domain::base::name::{FromStrError, Label, ToLabelIter};
16use domain::base::rdata::ComposeRecordData;
17use domain::base::wire::{Composer, ParseError};
18use domain::base::{
19    Message, MessageBuilder, ParsedName, Question, Record, RecordData, Rtype, ToName,
20};
21use domain::dep::octseq::{FreezeBuilder, FromBuilder, Octets, OctetsBuilder, ShortBuf, Truncate};
22use domain::rdata::AllRecordData;
23
24// This mod MUST go first, so that the others see its macros.
25pub(crate) mod fmt;
26
27#[cfg(feature = "io")]
28pub mod buf; // TODO: Maybe move to a generic `edge-buf` crate in future
29/// Re-export the domain lib if the user would like to directly
30/// assemble / parse mDNS messages.
31pub mod domain {
32    pub use domain::*;
33}
34pub mod host;
35#[cfg(feature = "io")]
36pub mod io;
37
38/// The DNS-SD owner name.
39pub const DNS_SD_OWNER: NameSlice = NameSlice::new(&["_services", "_dns-sd", "_udp", "local"]);
40
41/// A wrapper type for the errors returned by the `domain` library during parsing and
42/// constructing mDNS messages.
43#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
44pub enum MdnsError {
45    ShortBuf,
46    InvalidMessage,
47}
48
49impl Display for MdnsError {
50    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
51        match self {
52            Self::ShortBuf => write!(f, "ShortBuf"),
53            Self::InvalidMessage => write!(f, "InvalidMessage"),
54        }
55    }
56}
57
58#[cfg(feature = "defmt")]
59impl defmt::Format for MdnsError {
60    fn format(&self, f: defmt::Formatter<'_>) {
61        match self {
62            Self::ShortBuf => defmt::write!(f, "ShortBuf"),
63            Self::InvalidMessage => defmt::write!(f, "InvalidMessage"),
64        }
65    }
66}
67
68impl core::error::Error for MdnsError {}
69
70impl From<ShortBuf> for MdnsError {
71    fn from(_: ShortBuf) -> Self {
72        Self::ShortBuf
73    }
74}
75
76impl From<PushError> for MdnsError {
77    fn from(_: PushError) -> Self {
78        Self::ShortBuf
79    }
80}
81
82impl From<FromStrError> for MdnsError {
83    fn from(_: FromStrError) -> Self {
84        Self::InvalidMessage
85    }
86}
87
88impl From<ShortMessage> for MdnsError {
89    fn from(_: ShortMessage) -> Self {
90        Self::InvalidMessage
91    }
92}
93
94impl From<ParseError> for MdnsError {
95    fn from(_: ParseError) -> Self {
96        Self::InvalidMessage
97    }
98}
99
100/// This newtype struct allows the construction of a `domain` lib Name from
101/// a bunch of `&str` labels represented as a slice.
102///
103/// Implements the `domain` lib `ToName` trait.
104#[derive(Debug, Clone)]
105pub struct NameSlice<'a>(&'a [&'a str]);
106
107impl<'a> NameSlice<'a> {
108    /// Create a new `NameSlice` instance from a slice of `&str` labels.
109    pub const fn new(labels: &'a [&'a str]) -> Self {
110        Self(labels)
111    }
112}
113
114impl core::fmt::Display for NameSlice<'_> {
115    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
116        for label in self.0 {
117            write!(f, "{}.", label)?;
118        }
119
120        Ok(())
121    }
122}
123
124#[cfg(feature = "defmt")]
125impl defmt::Format for NameSlice<'_> {
126    fn format(&self, f: defmt::Formatter<'_>) {
127        for label in self.0 {
128            defmt::write!(f, "{}.", label);
129        }
130    }
131}
132
133impl ToName for NameSlice<'_> {}
134
135/// An iterator over the labels in a `NameSlice` instance.
136#[derive(Clone)]
137pub struct NameSliceIter<'a> {
138    name: &'a NameSlice<'a>,
139    index: usize,
140}
141
142impl<'a> Iterator for NameSliceIter<'a> {
143    type Item = &'a Label;
144
145    fn next(&mut self) -> Option<Self::Item> {
146        match self.index.cmp(&self.name.0.len()) {
147            Ordering::Less => {
148                let label = unwrap!(
149                    Label::from_slice(self.name.0[self.index].as_bytes()),
150                    "Unreachable"
151                );
152                self.index += 1;
153                Some(label)
154            }
155            Ordering::Equal => {
156                let label = Label::root();
157                self.index += 1;
158                Some(label)
159            }
160            Ordering::Greater => None,
161        }
162    }
163}
164
165impl DoubleEndedIterator for NameSliceIter<'_> {
166    fn next_back(&mut self) -> Option<Self::Item> {
167        if self.index > 0 {
168            self.index -= 1;
169            if self.index == self.name.0.len() {
170                let label = Label::root();
171                Some(label)
172            } else {
173                let label = unwrap!(
174                    Label::from_slice(self.name.0[self.index].as_bytes()),
175                    "Unreachable"
176                );
177                Some(label)
178            }
179        } else {
180            None
181        }
182    }
183}
184
185impl ToLabelIter for NameSlice<'_> {
186    type LabelIter<'t>
187        = NameSliceIter<'t>
188    where
189        Self: 't;
190
191    fn iter_labels(&self) -> Self::LabelIter<'_> {
192        NameSliceIter {
193            name: self,
194            index: 0,
195        }
196    }
197}
198
199/// A custom struct for representing a TXT data record off from a slice of
200/// key-value `&str` pairs.
201#[derive(Debug, Clone)]
202pub struct Txt<'a>(&'a [(&'a str, &'a str)]);
203
204impl<'a> Txt<'a> {
205    pub const fn new(txt: &'a [(&'a str, &'a str)]) -> Self {
206        Self(txt)
207    }
208}
209
210impl core::fmt::Display for Txt<'_> {
211    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
212        write!(f, "Txt [")?;
213
214        for (i, (k, v)) in self.0.iter().enumerate() {
215            if i > 0 {
216                write!(f, ", {}={}", k, v)?;
217            } else {
218                write!(f, "{}={}", k, v)?;
219            }
220        }
221
222        write!(f, "]")?;
223
224        Ok(())
225    }
226}
227
228#[cfg(feature = "defmt")]
229impl defmt::Format for Txt<'_> {
230    fn format(&self, f: defmt::Formatter<'_>) {
231        defmt::write!(f, "Txt [");
232
233        for (i, (k, v)) in self.0.iter().enumerate() {
234            if i > 0 {
235                defmt::write!(f, ", {}={}", k, v);
236            } else {
237                defmt::write!(f, "{}={}", k, v);
238            }
239        }
240
241        defmt::write!(f, "]");
242    }
243}
244
245impl RecordData for Txt<'_> {
246    fn rtype(&self) -> Rtype {
247        Rtype::TXT
248    }
249}
250
251impl ComposeRecordData for Txt<'_> {
252    fn rdlen(&self, _compress: bool) -> Option<u16> {
253        None
254    }
255
256    fn compose_rdata<Target: Composer + ?Sized>(
257        &self,
258        target: &mut Target,
259    ) -> Result<(), Target::AppendError> {
260        if self.0.is_empty() {
261            target.append_slice(&[0])?;
262        } else {
263            // TODO: Will not work for (k, v) pairs larger than 254 bytes in length
264            for (k, v) in self.0 {
265                target.append_slice(&[(k.len() + v.len() + 1) as u8])?;
266                target.append_slice(k.as_bytes())?;
267                target.append_slice(b"=")?;
268                target.append_slice(v.as_bytes())?;
269            }
270        }
271
272        Ok(())
273    }
274
275    fn compose_canonical_rdata<Target: Composer + ?Sized>(
276        &self,
277        target: &mut Target,
278    ) -> Result<(), Target::AppendError> {
279        self.compose_rdata(target)
280    }
281}
282
283/// A custom struct allowing to chain together multiple custom record data types.
284/// Allows e.g. using the custom `Txt` struct from above and chain it with `domain`'s `AllRecordData`,
285#[derive(Debug, Clone)]
286pub enum RecordDataChain<T, U> {
287    This(T),
288    Next(U),
289}
290
291impl<T, U> core::fmt::Display for RecordDataChain<T, U>
292where
293    T: core::fmt::Display,
294    U: core::fmt::Display,
295{
296    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
297        match self {
298            Self::This(data) => write!(f, "{}", data),
299            Self::Next(data) => write!(f, "{}", data),
300        }
301    }
302}
303
304#[cfg(feature = "defmt")]
305impl<T, U> defmt::Format for RecordDataChain<T, U>
306where
307    T: defmt::Format,
308    U: defmt::Format,
309{
310    fn format(&self, f: defmt::Formatter<'_>) {
311        match self {
312            Self::This(data) => defmt::write!(f, "{}", data),
313            Self::Next(data) => defmt::write!(f, "{}", data),
314        }
315    }
316}
317
318impl<T, U> RecordData for RecordDataChain<T, U>
319where
320    T: RecordData,
321    U: RecordData,
322{
323    fn rtype(&self) -> Rtype {
324        match self {
325            Self::This(data) => data.rtype(),
326            Self::Next(data) => data.rtype(),
327        }
328    }
329}
330
331impl<T, U> ComposeRecordData for RecordDataChain<T, U>
332where
333    T: ComposeRecordData,
334    U: ComposeRecordData,
335{
336    fn rdlen(&self, compress: bool) -> Option<u16> {
337        match self {
338            Self::This(data) => data.rdlen(compress),
339            Self::Next(data) => data.rdlen(compress),
340        }
341    }
342
343    fn compose_rdata<Target: Composer + ?Sized>(
344        &self,
345        target: &mut Target,
346    ) -> Result<(), Target::AppendError> {
347        match self {
348            Self::This(data) => data.compose_rdata(target),
349            Self::Next(data) => data.compose_rdata(target),
350        }
351    }
352
353    fn compose_canonical_rdata<Target: Composer + ?Sized>(
354        &self,
355        target: &mut Target,
356    ) -> Result<(), Target::AppendError> {
357        match self {
358            Self::This(data) => data.compose_canonical_rdata(target),
359            Self::Next(data) => data.compose_canonical_rdata(target),
360        }
361    }
362}
363
364/// This struct allows one to use a regular `&mut [u8]` slice as an octet buffer
365/// with the `domain` library.
366///
367/// Useful when a `domain` message needs to be constructed in a `&mut [u8]` slice.
368pub struct Buf<'a>(pub &'a mut [u8], pub usize);
369
370impl<'a> Buf<'a> {
371    /// Create a new `Buf` instance from a mutable slice.
372    pub fn new(buf: &'a mut [u8]) -> Self {
373        Self(buf, 0)
374    }
375}
376
377impl FreezeBuilder for Buf<'_> {
378    type Octets = Self;
379
380    fn freeze(self) -> Self {
381        self
382    }
383}
384
385impl Octets for Buf<'_> {
386    type Range<'r>
387        = &'r [u8]
388    where
389        Self: 'r;
390
391    fn range(&self, range: impl RangeBounds<usize>) -> Self::Range<'_> {
392        self.0[..self.1].range(range)
393    }
394}
395
396impl<'a> FromBuilder for Buf<'a> {
397    type Builder = Buf<'a>;
398
399    fn from_builder(builder: Self::Builder) -> Self {
400        Buf(&mut builder.0[builder.1..], 0)
401    }
402}
403
404impl Composer for Buf<'_> {}
405
406impl OctetsBuilder for Buf<'_> {
407    type AppendError = ShortBuf;
408
409    fn append_slice(&mut self, slice: &[u8]) -> Result<(), Self::AppendError> {
410        if self.1 + slice.len() <= self.0.len() {
411            let end = self.1 + slice.len();
412            self.0[self.1..end].copy_from_slice(slice);
413            self.1 = end;
414
415            Ok(())
416        } else {
417            Err(ShortBuf)
418        }
419    }
420}
421
422impl Truncate for Buf<'_> {
423    fn truncate(&mut self, len: usize) {
424        self.1 = len;
425    }
426}
427
428impl AsMut<[u8]> for Buf<'_> {
429    fn as_mut(&mut self) -> &mut [u8] {
430        &mut self.0[..self.1]
431    }
432}
433
434impl AsRef<[u8]> for Buf<'_> {
435    fn as_ref(&self) -> &[u8] {
436        &self.0[..self.1]
437    }
438}
439
440/// Type of request for `MdnsHandler::handle`.
441#[derive(Debug, Clone, Eq, PartialEq)]
442#[cfg_attr(feature = "defmt", derive(defmt::Format))]
443pub enum MdnsRequest<'a> {
444    /// No incoming mDNS request. Send a broadcast message
445    None,
446    /// Incoming mDNS request
447    Request {
448        /// Whether it is a legacy request (i.e. UDP packet source port is not 5353, as per spec)
449        legacy: bool,
450        /// Whether the request arrived on the multicast address
451        multicast: bool,
452        /// The data of the request
453        data: &'a [u8],
454    },
455}
456
457/// Return type for `MdnsHandler::handle`.
458#[derive(Debug, Clone, Eq, PartialEq)]
459#[cfg_attr(feature = "defmt", derive(defmt::Format))]
460pub enum MdnsResponse<'a> {
461    None,
462    Reply { data: &'a [u8], delay: bool },
463}
464
465/// A trait that abstracts the processing logic for an incoming mDNS message.
466///
467/// Handles an incoming mDNS message by parsing it and potentially preparing a response.
468///
469/// If request is `None`, the handler should prepare a broadcast message with
470/// all its data (i.e. mDNS responder brodcasts on internal state changes).
471///
472/// Returns an `MdnsResponse` instance that instructs the caller
473/// what data to send as a response (if any) and whether to generate a random delay
474/// before sending (as per spec).
475pub trait MdnsHandler {
476    fn handle<'a>(
477        &mut self,
478        request: MdnsRequest<'_>,
479        response_buf: &'a mut [u8],
480    ) -> Result<MdnsResponse<'a>, MdnsError>;
481}
482
483impl<T> MdnsHandler for &mut T
484where
485    T: MdnsHandler,
486{
487    fn handle<'a>(
488        &mut self,
489        request: MdnsRequest<'_>,
490        response_buf: &'a mut [u8],
491    ) -> Result<MdnsResponse<'a>, MdnsError> {
492        (**self).handle(request, response_buf)
493    }
494}
495
496/// A structure representing a handler that does not do any processing.
497///
498/// Useful only when chaining multiple `MdnsHandler` instances.
499pub struct NoHandler;
500
501impl NoHandler {
502    /// Chains a `NoHandler` with another handler.
503    pub fn chain<T>(self, handler: T) -> ChainedHandler<T, Self> {
504        ChainedHandler::new(handler, self)
505    }
506}
507
508impl MdnsHandler for NoHandler {
509    fn handle<'a>(
510        &mut self,
511        _request: MdnsRequest<'_>,
512        _response_buf: &'a mut [u8],
513    ) -> Result<MdnsResponse<'a>, MdnsError> {
514        Ok(MdnsResponse::None)
515    }
516}
517
518/// A composite handler that chains two handlers together.
519pub struct ChainedHandler<T, U> {
520    first: T,
521    second: U,
522}
523
524impl<T, U> ChainedHandler<T, U> {
525    /// Create a new `ChainedHandler` instance from two handlers.
526    pub const fn new(first: T, second: U) -> Self {
527        Self { first, second }
528    }
529
530    /// Chains a `ChainedHandler` with another handler,
531    /// where our instance would be the first one to be called.
532    ///
533    /// Chaining works by calling each chained handler from the first to the last,
534    /// until a handler in the chain returns a non-zero `usize` result.
535    ///
536    /// Once that happens, traversing the handlers down the chain stops.
537    pub fn chain<V>(self, handler: V) -> ChainedHandler<V, Self> {
538        ChainedHandler::new(handler, self)
539    }
540}
541
542impl<T, U> MdnsHandler for ChainedHandler<T, U>
543where
544    T: MdnsHandler,
545    U: MdnsHandler,
546{
547    fn handle<'a>(
548        &mut self,
549        request: MdnsRequest<'_>,
550        response_buf: &'a mut [u8],
551    ) -> Result<MdnsResponse<'a>, MdnsError> {
552        match self.first.handle(request.clone(), response_buf)? {
553            MdnsResponse::None => self.second.handle(request, response_buf),
554            MdnsResponse::Reply { data, delay } => {
555                let len = data.len();
556
557                Ok(MdnsResponse::Reply {
558                    data: &response_buf[..len],
559                    delay,
560                })
561            }
562        }
563    }
564}
565
566/// A type alias for the answer which is expected to be returned by instances
567/// implementing the `HostAnswers` trait.
568pub type HostAnswer<'a> =
569    Record<NameSlice<'a>, RecordDataChain<Txt<'a>, AllRecordData<&'a [u8], NameSlice<'a>>>>;
570
571/// A trait that abstracts the logic for providing answers to incoming mDNS queries.
572///
573/// The visitor-pattern-with-a-callback is chosen on purpose, as that allows `domain`
574/// Names to be constructed on-the-fly, possibly without interim buffer allocations.
575///
576/// Look at the implementation of `HostAnswers` for `host::Host` and `host::Service`
577/// for examples of this technique.
578pub trait HostAnswers {
579    /// Visits an entity that does have answers to mDNS queries.
580    ///
581    /// The answers will be provided to the supplied `f` callback.
582    ///
583    /// Note that the entity should provide ALL of its answers, regardless of the
584    /// concrete questions.
585    ///
586    /// The filtering of the answers relevant for the asked questions is done by the caller,
587    /// and only if necessary (i.e. only if these answers are used to reply to a concrete mDNS query,
588    /// rather than just broadcasting all answers that the entity has, which is also a valid mDNS
589    /// operation, that should be done when the entity providing answers has changes ot its internal state).
590    fn visit<F, E>(&self, f: F) -> Result<(), E>
591    where
592        F: FnMut(HostAnswer) -> Result<(), E>,
593        E: From<MdnsError>;
594}
595
596impl<T> HostAnswers for &T
597where
598    T: HostAnswers,
599{
600    fn visit<F, E>(&self, f: F) -> Result<(), E>
601    where
602        F: FnMut(HostAnswer) -> Result<(), E>,
603        E: From<MdnsError>,
604    {
605        (*self).visit(f)
606    }
607}
608
609impl<T> HostAnswers for &mut T
610where
611    T: HostAnswers,
612{
613    fn visit<F, E>(&self, f: F) -> Result<(), E>
614    where
615        F: FnMut(HostAnswer) -> Result<(), E>,
616        E: From<MdnsError>,
617    {
618        (**self).visit(f)
619    }
620}
621
622/// A type alias for the question which is expected to be returned by instances
623/// implementing the `HostQuestions` trait.
624pub type HostQuestion<'a> = Question<NameSlice<'a>>;
625
626/// A trait that abstracts the logic for providing questions to outgoing mDNS queries.
627///
628/// The visitor-pattern-with-a-callback is chosen on purpose, as that allows `domain`
629/// Names to be constructed on-the-fly, possibly without interim buffer allocations.
630pub trait HostQuestions {
631    /// Visits an entity that does have questions.
632    ///
633    /// The questions will be provided to the supplied `f` callback.
634    fn visit<F, E>(&self, f: F) -> Result<(), E>
635    where
636        F: FnMut(HostQuestion) -> Result<(), E>,
637        E: From<MdnsError>;
638
639    /// A function that constructs an mDNS query message in a `&mut [u8]` buffer
640    /// using questions generated by this trait.
641    fn query(&self, id: u16, buf: &mut [u8]) -> Result<usize, MdnsError> {
642        let buf = Buf(buf, 0);
643
644        let mut mb = MessageBuilder::from_target(buf)?;
645
646        set_header(&mut mb, id, false);
647
648        let mut qb = mb.question();
649
650        let mut pushed = false;
651
652        self.visit(|question| {
653            qb.push(question)?;
654
655            pushed = true;
656
657            Ok::<_, MdnsError>(())
658        })?;
659
660        let buf = qb.finish();
661
662        if pushed {
663            Ok(buf.1)
664        } else {
665            Ok(0)
666        }
667    }
668}
669
670impl<T> HostQuestions for &T
671where
672    T: HostQuestions,
673{
674    fn visit<F, E>(&self, f: F) -> Result<(), E>
675    where
676        F: FnMut(HostQuestion) -> Result<(), E>,
677        E: From<MdnsError>,
678    {
679        (*self).visit(f)
680    }
681}
682
683impl<T> HostQuestions for &mut T
684where
685    T: HostQuestions,
686{
687    fn visit<F, E>(&self, f: F) -> Result<(), E>
688    where
689        F: FnMut(HostQuestion) -> Result<(), E>,
690        E: From<MdnsError>,
691    {
692        (**self).visit(f)
693    }
694}
695
696/// A structure modeling an entity that does not generate any questions.
697///
698/// Useful only when chaining multiple `HostQuestions` instances.
699pub struct NoHostQuestions;
700
701impl NoHostQuestions {
702    /// Chains a `HostQuestions` with another `HostAnswers` instance.
703    pub fn chain<T>(self, questions: T) -> ChainedHostQuestions<T, Self> {
704        ChainedHostQuestions::new(questions, self)
705    }
706}
707
708impl HostQuestions for NoHostQuestions {
709    fn visit<F, E>(&self, _f: F) -> Result<(), E>
710    where
711        F: FnMut(HostQuestion) -> Result<(), E>,
712    {
713        Ok(())
714    }
715}
716
717/// A composite `HostQuestions` that chains two `HostQuestions` instances together.
718pub struct ChainedHostQuestions<T, U> {
719    first: T,
720    second: U,
721}
722
723impl<T, U> ChainedHostQuestions<T, U> {
724    /// Create a new `ChainedHostQuestions` instance from two `HostQuestions` instances.
725    pub const fn new(first: T, second: U) -> Self {
726        Self { first, second }
727    }
728
729    /// Chains this instance with another `HostQuestions` instance,
730    pub fn chain<V>(self, answers: V) -> ChainedHostQuestions<V, Self> {
731        ChainedHostQuestions::new(answers, self)
732    }
733}
734
735impl<T, U> HostQuestions for ChainedHostQuestions<T, U>
736where
737    T: HostQuestions,
738    U: HostQuestions,
739{
740    fn visit<F, E>(&self, mut f: F) -> Result<(), E>
741    where
742        F: FnMut(HostQuestion) -> Result<(), E>,
743        E: From<MdnsError>,
744    {
745        self.first.visit(&mut f)?;
746        self.second.visit(f)
747    }
748}
749
750/// A structure modeling an entity that does not generate any answers.
751///
752/// Useful only when chaining multiple `HostAnswers` instances.
753pub struct NoHostAnswers;
754
755impl NoHostAnswers {
756    /// Chains a `NoHostAnswers` with another `HostAnswers` instance.
757    pub fn chain<T>(self, answers: T) -> ChainedHostAnswers<T, Self> {
758        ChainedHostAnswers::new(answers, self)
759    }
760}
761
762impl HostAnswers for NoHostAnswers {
763    fn visit<F, E>(&self, _f: F) -> Result<(), E>
764    where
765        F: FnMut(HostAnswer) -> Result<(), E>,
766    {
767        Ok(())
768    }
769}
770
771/// A composite `HostAnswers` that chains two `HostAnswers` instances together.
772pub struct ChainedHostAnswers<T, U> {
773    first: T,
774    second: U,
775}
776
777impl<T, U> ChainedHostAnswers<T, U> {
778    /// Create a new `ChainedHostAnswers` instance from two `HostAnswers` instances.
779    pub const fn new(first: T, second: U) -> Self {
780        Self { first, second }
781    }
782
783    /// Chains this instance with another `HostAnswers` instance,
784    pub fn chain<V>(self, answers: V) -> ChainedHostAnswers<V, Self> {
785        ChainedHostAnswers::new(answers, self)
786    }
787}
788
789impl<T, U> HostAnswers for ChainedHostAnswers<T, U>
790where
791    T: HostAnswers,
792    U: HostAnswers,
793{
794    fn visit<F, E>(&self, mut f: F) -> Result<(), E>
795    where
796        F: FnMut(HostAnswer) -> Result<(), E>,
797        E: From<MdnsError>,
798    {
799        self.first.visit(&mut f)?;
800        self.second.visit(f)
801    }
802}
803
804/// An `MdnsHandler` implementation that answers mDNS queries with the answers
805/// provided by an entity implementing the `HostAnswers` trait.
806///
807/// Typically, this structure will be used to provide answers to other peers that broadcast
808/// mDNS queries - i.e. this is the "responder" aspect of the mDNS protocol.
809pub struct HostAnswersMdnsHandler<T> {
810    answers: T,
811}
812
813impl<T> HostAnswersMdnsHandler<T> {
814    /// Create a new `HostAnswersMdnsHandler` instance from an entity that provides answers.
815    pub const fn new(answers: T) -> Self {
816        Self { answers }
817    }
818}
819
820impl<T> MdnsHandler for HostAnswersMdnsHandler<T>
821where
822    T: HostAnswers,
823{
824    fn handle<'a>(
825        &mut self,
826        request: MdnsRequest<'_>,
827        response_buf: &'a mut [u8],
828    ) -> Result<MdnsResponse<'a>, MdnsError> {
829        let buf = Buf(response_buf, 0);
830
831        let mut mb = MessageBuilder::from_target(buf)?;
832
833        let mut pushed = false;
834
835        let buf = if let MdnsRequest::Request { legacy, data, .. } = request {
836            let message = Message::from_octets(data)?;
837
838            if !matches!(message.header().opcode(), Opcode::QUERY)
839                || !matches!(message.header().rcode(), Rcode::NOERROR)
840                || message.header().qr()
841            // Not a query but a response
842            {
843                return Ok(MdnsResponse::None);
844            }
845
846            let mut ab = if legacy {
847                set_header(&mut mb, message.header().id(), true);
848
849                let mut qb = mb.question();
850
851                // As per spec, for legacy requests we need to fill-in the questions section
852                for question in message.question() {
853                    qb.push(question?)?;
854                }
855
856                qb.answer()
857            } else {
858                set_header(&mut mb, 0, true);
859
860                mb.answer()
861            };
862
863            let mut additional_a = false;
864            let mut additional_srv_txt = false;
865
866            for question in message.question() {
867                let question = question?;
868
869                self.answers.visit(|answer| {
870                    if matches!(answer.data(), RecordDataChain::Next(AllRecordData::Srv(_))) {
871                        additional_a = true;
872                    }
873
874                    if !answer.owner().name_eq(&DNS_SD_OWNER)
875                        && matches!(answer.data(), RecordDataChain::Next(AllRecordData::Ptr(_)))
876                    {
877                        additional_a = true;
878
879                        // Over-simplifying here in that we'll send all our SRV and TXT records, however
880                        // sending only some SRV and PTR records is too complex to implement.
881                        additional_srv_txt = true;
882                    }
883
884                    if question.qname().name_eq(&answer.owner()) {
885                        debug!(
886                            "Answering question [{}] with: [{}]",
887                            debug2format!(question),
888                            debug2format!(answer)
889                        );
890
891                        ab.push(answer)?;
892
893                        pushed = true;
894                    }
895
896                    Ok::<_, MdnsError>(())
897                })?;
898            }
899
900            if additional_a || additional_srv_txt {
901                // Fill-in the additional section as well
902
903                let mut aa = ab.additional();
904
905                self.answers.visit(|answer| {
906                    if matches!(
907                        answer.data(),
908                        RecordDataChain::Next(AllRecordData::A(_))
909                            | RecordDataChain::Next(AllRecordData::Aaaa(_))
910                            | RecordDataChain::Next(AllRecordData::Srv(_))
911                            | RecordDataChain::Next(AllRecordData::Txt(_))
912                            | RecordDataChain::This(Txt(_))
913                    ) {
914                        debug!("Additional answer: [{}]", debug2format!(answer));
915
916                        aa.push(answer)?;
917                    }
918
919                    Ok::<_, MdnsError>(())
920                })?;
921
922                aa.finish()
923            } else {
924                ab.finish()
925            }
926        } else {
927            set_header(&mut mb, 0, true);
928
929            let mut ab = mb.answer();
930
931            self.answers.visit(|answer| {
932                ab.push(answer)?;
933
934                pushed = true;
935
936                Ok::<_, MdnsError>(())
937            })?;
938
939            ab.finish()
940        };
941
942        if pushed {
943            Ok(MdnsResponse::Reply {
944                data: &buf.0[..buf.1],
945                delay: false,
946            })
947        } else {
948            Ok(MdnsResponse::None)
949        }
950    }
951}
952
953/// A type alias for the answer which is expected to be returned by instances
954/// implementing the `PeerAnswers` trait.
955pub type PeerAnswer<'a> =
956    Record<ParsedName<&'a [u8]>, AllRecordData<&'a [u8], ParsedName<&'a [u8]>>>;
957
958/// A trait that abstracts the logic for processing answers from incoming mDNS queries.
959///
960/// Rather than dealing with the whole mDNS message, this trait is focused on processing
961/// the answers from the message (in the `answer` and `additional` mDNS message sections).
962pub trait PeerAnswers {
963    /// Processes the answers from an incoming mDNS message.
964    fn answers<'a, T, A>(&self, answers: T, additional: A) -> Result<(), MdnsError>
965    where
966        T: IntoIterator<Item = Result<PeerAnswer<'a>, MdnsError>> + Clone + 'a,
967        A: IntoIterator<Item = Result<PeerAnswer<'a>, MdnsError>> + Clone + 'a;
968}
969
970impl<T> PeerAnswers for &mut T
971where
972    T: PeerAnswers,
973{
974    fn answers<'a, U, V>(&self, answers: U, additional: V) -> Result<(), MdnsError>
975    where
976        U: IntoIterator<Item = Result<PeerAnswer<'a>, MdnsError>> + Clone + 'a,
977        V: IntoIterator<Item = Result<PeerAnswer<'a>, MdnsError>> + Clone + 'a,
978    {
979        (**self).answers(answers, additional)
980    }
981}
982
983impl<T> PeerAnswers for &T
984where
985    T: PeerAnswers,
986{
987    fn answers<'a, U, V>(&self, answers: U, additional: V) -> Result<(), MdnsError>
988    where
989        U: IntoIterator<Item = Result<PeerAnswer<'a>, MdnsError>> + Clone + 'a,
990        V: IntoIterator<Item = Result<PeerAnswer<'a>, MdnsError>> + Clone + 'a,
991    {
992        (*self).answers(answers, additional)
993    }
994}
995
996/// A structure implementing the `MdnsHandler` trait by processing all answers from an
997/// incoming mDNS message via delegating to an entity implementing the `PeerAnswers` trait.
998///
999/// Typically, this structure will be used to process answers which are replies to mDNS
1000/// queries that we have sent using the `HostQuestions::query` method, i.e., this is the
1001/// "querying" part of the mDNS protocol.
1002///
1003/// Since the "querying" aspect of the mDNS protocol is modeled here, this handler never
1004/// answers anything, i.e. it always returns a 0 `usize`, because - unlike the
1005/// `HostAnswersMdnsHandler` - it does not have any answers to provide, as it - itself -
1006/// processes answers provided by peers, which were themselves sent because we issued a query
1007/// using e.g. the `HostQuestions::query` method at an earlier point in time.
1008pub struct PeerAnswersMdnsHandler<T> {
1009    answers: T,
1010}
1011
1012impl<T> PeerAnswersMdnsHandler<T> {
1013    /// Create a new `PeerAnswersMdnsHandler` instance from an entity that processes answers.
1014    pub const fn new(answers: T) -> Self {
1015        Self { answers }
1016    }
1017}
1018
1019impl<T> MdnsHandler for PeerAnswersMdnsHandler<T>
1020where
1021    T: PeerAnswers,
1022{
1023    fn handle<'a>(
1024        &mut self,
1025        request: MdnsRequest<'_>,
1026        _response_buf: &'a mut [u8],
1027    ) -> Result<MdnsResponse<'a>, MdnsError> {
1028        let MdnsRequest::Request { data, legacy, .. } = request else {
1029            return Ok(MdnsResponse::None);
1030        };
1031
1032        if legacy {
1033            // Legacy packets should not contain mDNS answers anyway, per spec
1034            return Ok(MdnsResponse::None);
1035        }
1036
1037        let message = Message::from_octets(data)?;
1038
1039        if !matches!(message.header().opcode(), Opcode::QUERY)
1040            || !matches!(message.header().rcode(), Rcode::NOERROR)
1041            || !message.header().qr()
1042        // Not a response but a query
1043        {
1044            return Ok(MdnsResponse::None);
1045        }
1046
1047        let answers = message.answer()?;
1048        let additional = message.additional()?;
1049
1050        let answers = answers.filter_map(|answer| {
1051            match answer {
1052                Ok(answer) => answer.into_record::<AllRecordData<_, _>>(),
1053                Err(e) => Err(e),
1054            }
1055            .map_err(|_| MdnsError::InvalidMessage)
1056            .transpose()
1057        });
1058
1059        let additional = additional.filter_map(|answer| {
1060            match answer {
1061                Ok(answer) => answer.into_record::<AllRecordData<_, _>>(),
1062                Err(e) => Err(e),
1063            }
1064            .map_err(|_| MdnsError::InvalidMessage)
1065            .transpose()
1066        });
1067
1068        self.answers.answers(answers, additional)?;
1069
1070        Ok(MdnsResponse::None)
1071    }
1072}
1073
1074/// Utility function that sets the header of an mDNS `domain` message builder
1075/// to be a response or a query.
1076pub fn set_header<T: Composer>(answer: &mut MessageBuilder<T>, id: u16, response: bool) {
1077    let header = answer.header_mut();
1078    header.set_id(id);
1079    header.set_opcode(Opcode::QUERY);
1080    header.set_rcode(Rcode::NOERROR);
1081
1082    let mut flags = Flags::new();
1083    flags.qr = response;
1084    flags.aa = response;
1085    header.set_flags(flags);
1086}