hick_compio/discovery/
mod.rs

1//! DNS-SD service discovery — data types and parsing helpers.
2//!
3//! This module hosts the public types ([`ServiceEntry`], [`QueryParam`]) and the
4//! crate-internal helpers (`push_capped`, `fold`, `parse_name`,
5//! `parse_srv`, `parse_txt`, `name_from_ref`) that the resolver state
6//! machine and [`Lookup`] handle build on. The
7//! [`Endpoint::browse`](crate::Endpoint::browse) /
8//! [`lookup`](crate::Endpoint::lookup) /
9//! [`resolve_host`](crate::Endpoint::resolve_host) /
10//! [`resolve_instance`](crate::Endpoint::resolve_instance) APIs wire these
11//! together to drive the RFC 6763 browse → resolve chain.
12
13use std::{
14  collections::{HashMap, HashSet, VecDeque},
15  sync::Arc,
16};
17
18use core::{
19  net::{IpAddr, Ipv4Addr, Ipv6Addr},
20  time::Duration,
21};
22
23use bytes::Bytes;
24use futures::future::select_all;
25use mdns_proto::{
26  Name, QuerySpec,
27  wire::{A, AAAA, NameRef, Ptr, ResourceType, Srv, Txt},
28};
29use smol_str::SmolStr;
30
31use crate::{
32  endpoint::Endpoint,
33  query::{Query, QueryEvent},
34};
35
36#[cfg(test)]
37mod tests;
38
39/// Default cap on the number of distinct instances a lookup tracks, so a chatty
40/// or hostile responder flooding PTR answers cannot grow the in-progress state
41/// without bound. Used by [`QueryParam::new`].
42pub const DEFAULT_MAX_ENTRIES: usize = 64;
43
44/// Cap on the addresses kept per host per family (and per instance), so a
45/// responder flooding distinct A/AAAA records for one host cannot grow the
46/// address vectors without limit.
47pub(crate) const MAX_ADDRS_PER_HOST: usize = 16;
48
49/// A resolved DNS-SD service instance.
50///
51/// Produced by the resolver once the instance has SRV (host + port), TXT, and
52/// at least one address.
53#[derive(Debug, Clone)]
54pub struct ServiceEntry {
55  pub(crate) instance: Name,
56  pub(crate) host: Name,
57  pub(crate) port: u16,
58  pub(crate) ipv4: Arc<[Ipv4Addr]>,
59  pub(crate) ipv6: Arc<[Ipv6Addr]>,
60  pub(crate) txt: Arc<[Bytes]>,
61}
62
63impl ServiceEntry {
64  /// The fully-qualified service instance name (the PTR target), e.g.
65  /// `myprinter._ipp._tcp.local.`.
66  #[inline]
67  pub const fn instance_name(&self) -> &Name {
68    &self.instance
69  }
70
71  /// The target host the SRV record points at, e.g. `printer.local.`.
72  #[inline]
73  pub const fn host(&self) -> &Name {
74    &self.host
75  }
76
77  /// The service port from the SRV record.
78  #[inline]
79  pub const fn port(&self) -> u16 {
80    self.port
81  }
82
83  /// The host's IPv4 addresses (may be empty if only IPv6 resolved).
84  #[inline]
85  pub fn ipv4_addresses(&self) -> &[Ipv4Addr] {
86    &self.ipv4
87  }
88
89  /// The host's IPv6 addresses (may be empty if only IPv4 resolved).
90  #[inline]
91  pub fn ipv6_addresses(&self) -> &[Ipv6Addr] {
92    &self.ipv6
93  }
94
95  /// All resolved addresses (IPv4 first, then IPv6).
96  pub fn addresses(&self) -> impl Iterator<Item = IpAddr> + '_ {
97    self
98      .ipv4
99      .iter()
100      .copied()
101      .map(IpAddr::V4)
102      .chain(self.ipv6.iter().copied().map(IpAddr::V6))
103  }
104
105  /// The raw TXT record segments (each a length-prefixed `key=value` string in
106  /// DNS-SD usage, but treated as opaque bytes since TXT data may be binary). A
107  /// service with no metadata has a single empty segment (RFC 6763 §6.1).
108  #[inline]
109  pub fn txt(&self) -> &[Bytes] {
110    &self.txt
111  }
112}
113
114/// Parameters for a DNS-SD browse / lookup.
115#[derive(Debug, Clone)]
116pub struct QueryParam {
117  pub(crate) service: Name,
118  pub(crate) timeout: Duration,
119  pub(crate) resolve_timeout: Option<Duration>,
120  pub(crate) unicast_response: bool,
121  pub(crate) max_entries: usize,
122}
123
124impl QueryParam {
125  /// Browse the given fully-qualified service type, e.g.
126  /// `Name::try_from_str("_ipp._tcp.local.")`.
127  #[inline(always)]
128  pub const fn new(service: Name) -> Self {
129    Self {
130      service,
131      timeout: Duration::from_secs(1),
132      resolve_timeout: None,
133      unicast_response: false,
134      max_entries: DEFAULT_MAX_ENTRIES,
135    }
136  }
137
138  /// How long the browse (the PTR query) runs before terminating.
139  #[must_use]
140  #[inline(always)]
141  pub const fn with_timeout(mut self, timeout: Duration) -> Self {
142    self.timeout = timeout;
143    self
144  }
145
146  /// How long each per-instance resolve query (SRV / TXT / A / AAAA) runs.
147  /// Defaults to the browse timeout when unset.
148  #[must_use]
149  #[inline(always)]
150  pub const fn with_resolve_timeout(mut self, timeout: Duration) -> Self {
151    self.resolve_timeout = Some(timeout);
152    self
153  }
154
155  /// Request unicast responses on the issued queries (RFC 6762 §5.4). Defaults
156  /// to `false` (standard multicast browse).
157  #[must_use]
158  #[inline(always)]
159  pub const fn with_unicast_response(mut self, unicast: bool) -> Self {
160    self.unicast_response = unicast;
161    self
162  }
163
164  /// Cap on the number of distinct instances tracked; instances discovered
165  /// beyond it are dropped. Defaults to [`DEFAULT_MAX_ENTRIES`]. A value of
166  /// `0` is treated as `1`.
167  #[must_use]
168  #[inline(always)]
169  pub const fn with_max_entries(mut self, max: usize) -> Self {
170    self.max_entries = if max == 0 { 1 } else { max };
171    self
172  }
173}
174
175/// Push `item` into `v` if it is new (by equality) and `v` is still under
176/// [`MAX_ADDRS_PER_HOST`]. Returns `true` if it was added (so the caller can
177/// decide to (re)emit), `false` if it was a duplicate or the cap was reached.
178pub(crate) fn push_capped<T: PartialEq>(v: &mut Vec<T>, item: T) -> bool {
179  if v.len() >= MAX_ADDRS_PER_HOST || v.contains(&item) {
180    return false;
181  }
182  v.push(item);
183  true
184}
185
186/// Case-fold a [`Name`] to its lookup key. DNS names are case-insensitive
187/// (RFC 6762 §16), so the lower-cased string form is the canonical key for
188/// the resolver's per-instance / per-host maps.
189pub(crate) fn fold(name: &Name) -> SmolStr {
190  // `Name` is already stored canonical-lowercase (RFC 6762 §16), so a plain
191  // `SmolStr::new` suffices — and inlines names ≤23 bytes, whereas collecting
192  // from a `char` iterator always takes SmolStr's heap path.
193  SmolStr::new(name.as_str())
194}
195
196/// Decode an owner-less wire-form domain name (a decompressed PTR/SRV target as
197/// stored in a [`mdns_proto::CollectedAnswer`]) into an owned [`Name`].
198///
199/// Returns `None` for any name the [`Name`] type cannot represent faithfully:
200/// a label containing a `.` (which is always treated as a separator by
201/// [`Name`]) or a non-ASCII byte. Such inputs are skipped rather than silently
202/// corrupted.
203fn name_from_ref(nr: &NameRef<'_>) -> Option<Name> {
204  let mut buf = String::new();
205  for label in nr.labels() {
206    let label = label.ok()?;
207    if label.is_empty() {
208      break; // root terminator
209    }
210    if label.iter().any(|&b| b >= 0x80 || b == b'.') {
211      return None; // not representable by `Name` without corruption
212    }
213    // Verified ASCII above, so this slice is always valid UTF-8.
214    buf.push_str(core::str::from_utf8(label).ok()?);
215    buf.push('.');
216  }
217  if buf.is_empty() {
218    return None;
219  }
220  Name::try_from_str(&buf).ok()
221}
222
223/// Parse a PTR rdata slice (a decompressed wire-form name) into a [`Name`].
224/// Returns `None` for a malformed or unrepresentable target.
225pub(crate) fn parse_name(rdata: &[u8]) -> Option<Name> {
226  let ptr = Ptr::try_from_message(rdata, 0, rdata.len()).ok()?;
227  name_from_ref(ptr.target())
228}
229
230/// Parse an SRV rdata slice into `(target host, port)`. Returns `None` for a
231/// malformed record or an unrepresentable target host name.
232pub(crate) fn parse_srv(rdata: &[u8]) -> Option<(Name, u16)> {
233  let srv = Srv::try_from_message(rdata, 0, rdata.len()).ok()?;
234  let host = name_from_ref(srv.target())?;
235  Some((host, srv.port()))
236}
237
238/// Parse a TXT rdata slice into its length-prefixed segments
239/// (RFC 1035 §3.3.14 + RFC 6763 §6). Drops a malformed tail rather than
240/// failing the whole record.
241pub(crate) fn parse_txt(rdata: &[u8]) -> Vec<Bytes> {
242  Txt::from_rdata(rdata)
243    .segments()
244    .map_while(Result::ok)
245    .map(Bytes::copy_from_slice)
246    .collect()
247}
248
249/// Which resolve step an answer belongs to. The string is the case-folded key
250/// of the owning instance (SRV/TXT) or host (A/AAAA).
251#[derive(Clone)]
252#[allow(clippy::upper_case_acronyms)] // `AAAA` mirrors the DNS record-type name
253pub(crate) enum Step {
254  Ptr,
255  Srv(SmolStr),
256  Txt(SmolStr),
257  A(SmolStr),
258  AAAA(SmolStr),
259}
260
261/// A follow-up query the driver should launch as a result of feeding the
262/// [`Resolver`] an answer.
263#[derive(Clone)]
264pub(crate) struct Start {
265  pub(crate) name: Name,
266  pub(crate) step: Step,
267}
268
269impl Start {
270  pub(crate) const fn qtype(&self) -> ResourceType {
271    match self.step {
272      Step::Srv(_) => ResourceType::Srv,
273      Step::Txt(_) => ResourceType::Txt,
274      Step::A(_) => ResourceType::A,
275      Step::AAAA(_) => ResourceType::AAAA,
276      Step::Ptr => ResourceType::Ptr,
277    }
278  }
279}
280
281/// Addresses already learned for a host, so a later instance whose SRV resolves
282/// to the same host picks them up even though the A/AAAA answer has come and
283/// gone (the shared-host, SRV-after-A ordering).
284#[derive(Default)]
285pub(crate) struct HostAddrs {
286  pub(crate) ipv4: Vec<Ipv4Addr>,
287  pub(crate) ipv6: Vec<Ipv6Addr>,
288}
289
290/// In-progress aggregation of one service instance.
291pub(crate) struct Builder {
292  pub(crate) instance: Name,
293  pub(crate) host: Option<Name>,
294  pub(crate) host_key: Option<SmolStr>,
295  /// Whether an SRV record has been seen. Tracked separately from `port`
296  /// because `0` is a valid SRV port (the full `u16` range is parsed and the
297  /// registration API does not reject it), so it cannot double as a sentinel.
298  pub(crate) has_srv: bool,
299  pub(crate) port: u16,
300  pub(crate) ipv4: Vec<Ipv4Addr>,
301  pub(crate) ipv6: Vec<Ipv6Addr>,
302  pub(crate) txt: Option<Vec<Bytes>>,
303  pub(crate) emitted: bool,
304}
305
306impl Builder {
307  fn new(instance: Name) -> Self {
308    Self {
309      instance,
310      host: None,
311      host_key: None,
312      has_srv: false,
313      port: 0,
314      ipv4: Vec::new(),
315      ipv6: Vec::new(),
316      txt: None,
317      emitted: false,
318    }
319  }
320
321  /// Complete once it has an SRV (host + port), TXT, and at least one address.
322  fn complete(&self) -> bool {
323    self.has_srv && self.txt.is_some() && !(self.ipv4.is_empty() && self.ipv6.is_empty())
324  }
325
326  fn finalize(&self) -> Option<ServiceEntry> {
327    Some(ServiceEntry {
328      instance: self.instance.clone(),
329      host: self.host.clone()?,
330      port: self.port,
331      ipv4: self.ipv4.as_slice().into(),
332      ipv6: self.ipv6.as_slice().into(),
333      txt: self.txt.as_deref()?.into(),
334    })
335  }
336}
337
338/// Pure browse/resolve aggregation state machine — no I/O. The lookup driver
339/// feeds it parsed answers and launches the follow-up queries it requests.
340pub(crate) struct Resolver {
341  pub(crate) builders: HashMap<SmolStr, Builder>,
342  pub(crate) host_addrs: HashMap<SmolStr, HostAddrs>,
343  pub(crate) hosts_queried: HashSet<SmolStr>,
344  pub(crate) ready: VecDeque<ServiceEntry>,
345  /// Cap on distinct instances tracked.
346  pub(crate) max_entries: usize,
347  /// Cap on distinct hosts A/AAAA-queried. Set equal to `max_entries`: honest
348  /// browsing has at most one host per instance, so this never bites a real
349  /// responder, but it bounds an instance that floods distinct SRV targets
350  /// (which would otherwise grow `hosts_queried`/`host_addrs` and the in-flight
351  /// A/AAAA sub-query set without limit).
352  pub(crate) max_hosts: usize,
353  pub(crate) dropped: u64,
354}
355
356impl Resolver {
357  pub(crate) fn new(max_entries: usize) -> Self {
358    Self {
359      builders: HashMap::new(),
360      host_addrs: HashMap::new(),
361      hosts_queried: HashSet::new(),
362      ready: VecDeque::new(),
363      max_entries,
364      max_hosts: max_entries,
365      dropped: 0,
366    }
367  }
368
369  /// A newly discovered instance: register it (subject to the cap) and request
370  /// its SRV + TXT resolves.
371  pub(crate) fn on_ptr(&mut self, instance: Name) -> Vec<Start> {
372    let key = fold(&instance);
373    if self.builders.contains_key(&key) {
374      return Vec::new(); // already discovered
375    }
376    if self.builders.len() >= self.max_entries {
377      self.dropped = self.dropped.saturating_add(1);
378      return Vec::new();
379    }
380    self
381      .builders
382      .insert(key.clone(), Builder::new(instance.clone()));
383    vec![
384      Start {
385        name: instance.clone(),
386        step: Step::Srv(key.clone()),
387      },
388      Start {
389        name: instance,
390        step: Step::Txt(key),
391      },
392    ]
393  }
394
395  /// An instance's SRV: record host + port, adopt any addresses already learned
396  /// for that host, and request A/AAAA the first time we see the host — subject
397  /// to the distinct-host cap, which bounds an SRV-target flood.
398  pub(crate) fn on_srv(&mut self, inst_key: &str, host: Name, port: u16) -> Vec<Start> {
399    let host_key = fold(&host);
400    let cached = self
401      .host_addrs
402      .get(&host_key)
403      .map(|h| (h.ipv4.clone(), h.ipv6.clone()));
404    let mut changed = false;
405    if let Some(b) = self.builders.get_mut(inst_key) {
406      let host_changed = b.host_key.as_deref() != Some(host_key.as_str());
407      // SRV retargeting the instance to a DIFFERENT host invalidates the old
408      // host's addresses — drop them before adopting the new host's, so a
409      // re-emit never yields an entry whose host is the new target but whose
410      // addresses still belong to the old one.
411      if host_changed {
412        b.ipv4.clear();
413        b.ipv6.clear();
414      }
415      // A host or port change to an already-surfaced instance must re-emit
416      // even when the new host's addresses are already cached — no fresh
417      // A/AAAA event will arrive to trigger the re-emit, so without this the
418      // consumer keeps a stale host/port.
419      changed = host_changed || b.port != port;
420      b.has_srv = true;
421      b.host = Some(host.clone());
422      b.host_key = Some(host_key.clone());
423      b.port = port;
424      if let Some((v4, v6)) = cached {
425        for a in v4 {
426          push_capped(&mut b.ipv4, a);
427        }
428        for a in v6 {
429          push_capped(&mut b.ipv6, a);
430        }
431      }
432    }
433    self.try_emit(inst_key, changed);
434    if self.hosts_queried.contains(&host_key) {
435      return Vec::new(); // already querying this host
436    }
437    if self.hosts_queried.len() >= self.max_hosts {
438      // SRV-target flood guard: refuse to query an unbounded set of hosts.
439      // Counted so the resulting partial view is observable via `dropped`.
440      self.dropped = self.dropped.saturating_add(1);
441      return Vec::new();
442    }
443    self.hosts_queried.insert(host_key.clone());
444    vec![
445      Start {
446        name: host.clone(),
447        step: Step::A(host_key.clone()),
448      },
449      Start {
450        name: host,
451        step: Step::AAAA(host_key),
452      },
453    ]
454  }
455
456  pub(crate) fn on_txt(&mut self, inst_key: &str, segs: Vec<Bytes>) {
457    let mut changed = false;
458    if let Some(b) = self.builders.get_mut(inst_key) {
459      // A TXT change to an already-surfaced instance must re-emit so the
460      // consumer sees the new metadata; a duplicate TXT must not (no spurious
461      // re-emit on a refresh that carries the same payload).
462      changed = b.txt.as_deref() != Some(segs.as_slice());
463      b.txt = Some(segs);
464    }
465    self.try_emit(inst_key, changed);
466  }
467
468  pub(crate) fn on_addr(&mut self, host_key: &str, addr: IpAddr) {
469    let cache = self.host_addrs.entry(SmolStr::from(host_key)).or_default();
470    match addr {
471      IpAddr::V4(a) => {
472        push_capped(&mut cache.ipv4, a);
473      }
474      IpAddr::V6(a) => {
475        push_capped(&mut cache.ipv6, a);
476      }
477    }
478    let keys: Vec<SmolStr> = self
479      .builders
480      .iter()
481      .filter(|(_, b)| b.host_key.as_deref() == Some(host_key))
482      .map(|(k, _)| k.clone())
483      .collect();
484    for k in keys {
485      let added = match self.builders.get_mut(&k) {
486        Some(b) => match addr {
487          IpAddr::V4(a) => push_capped(&mut b.ipv4, a),
488          IpAddr::V6(a) => push_capped(&mut b.ipv6, a),
489        },
490        None => false,
491      };
492      if added {
493        self.try_emit(&k, true);
494      }
495    }
496  }
497
498  pub(crate) fn try_emit(&mut self, inst_key: &str, allow_reemit: bool) {
499    if let Some(b) = self.builders.get_mut(inst_key) {
500      if !b.complete() {
501        return;
502      }
503      if !b.emitted {
504        if let Some(e) = b.finalize() {
505          b.emitted = true;
506          self.ready.push_back(e);
507        }
508      } else if allow_reemit && let Some(e) = b.finalize() {
509        self.ready.push_back(e);
510      }
511    }
512  }
513
514  pub(crate) fn take_ready(&mut self) -> Option<ServiceEntry> {
515    self.ready.pop_front()
516  }
517}
518
519/// A running DNS-SD lookup.
520///
521/// Call [`Self::next`] to receive resolved [`ServiceEntry`] values as they
522/// complete; it returns `None` once every query (browse + resolves) has timed
523/// out and the resolver has drained.  Resolution progresses while
524/// [`Self::next`] is awaited — the lookup races the underlying sub-queries
525/// concurrently via [`futures::future::select_all`], so the first query with
526/// a buffered answer is processed without waiting on the others.
527///
528/// An instance may be yielded more than once as additional addresses resolve
529/// (e.g. a late AAAA after the entry was first surfaced on its A address); a
530/// later yield for the same [`ServiceEntry::instance_name`] supersedes the
531/// earlier one.  Dropping the `Lookup` cancels every in-flight sub-query.
532pub struct Lookup {
533  pub(crate) endpoint: Endpoint,
534  pub(crate) queries: Vec<(Query, Step)>,
535  pub(crate) resolver: Resolver,
536  pub(crate) resolve_timeout: Duration,
537  pub(crate) unicast: bool,
538  pub(crate) finished: bool,
539}
540
541impl Lookup {
542  /// Wait for the next resolved service instance, or `None` when the lookup
543  /// is finished (every sub-query has terminated and the resolver has
544  /// drained).
545  ///
546  /// The borrow on `&mut self` matches the single-consumer mailbox contract
547  /// inherited from [`Query::next`]: at most one [`Self::next`] is in flight
548  /// per [`Lookup`].
549  pub async fn next(&mut self) -> Option<ServiceEntry> {
550    loop {
551      if let Some(e) = self.resolver.take_ready() {
552        return Some(e);
553      }
554      if self.finished || self.queries.is_empty() {
555        self.finished = true;
556        return None;
557      }
558      // Race every live sub-query so a buffered answer on any of them is
559      // processed without parking on the rest.  Awaiting `queries[0].next()`
560      // sequentially would starve queries 1..N (the plan's known bug).
561      let (result, idx) = {
562        let futs: Vec<_> = self
563          .queries
564          .iter()
565          .map(|(q, _)| Box::pin(q.next()))
566          .collect();
567        let (result, idx, _remaining) = select_all(futs).await;
568        // `_remaining` drops here, releasing the immutable borrow on
569        // `self.queries`, so the swap_remove below is sound.
570        (result, idx)
571      };
572      match result {
573        Some(QueryEvent::Answer(answer)) => {
574          let step = self.queries[idx].1.clone();
575          let starts = feed(&mut self.resolver, step, QueryEvent::Answer(answer));
576          self.launch_starts(starts).await;
577        }
578        Some(QueryEvent::Terminal(_)) | None => {
579          self.queries.swap_remove(idx);
580          if self.queries.is_empty() {
581            self.finished = true;
582          }
583        }
584      }
585    }
586  }
587
588  async fn launch_starts(&mut self, starts: Vec<Start>) {
589    for start in starts {
590      let qtype = start.qtype();
591      let step = start.step;
592      let spec = QuerySpec::new(start.name, qtype)
593        .with_timeout(self.resolve_timeout)
594        .with_unicast_response(self.unicast);
595      // A `StartQueryError::StorageFull` here just means we can't make
596      // forward progress on this sub-step; the other live queries continue,
597      // and the resolver still terminates when they drain.
598      if let Ok(q) = self.endpoint.start_query(spec).await {
599        self.queries.push((q, step));
600      }
601    }
602  }
603}
604
605/// Decode one tagged answer, fold it into `resolver`, and return any
606/// follow-up queries it requests.  Shared by [`Lookup::next`] and the
607/// one-shot [`Endpoint::resolve_instance`] convenience.
608pub(crate) fn feed(resolver: &mut Resolver, step: Step, event: QueryEvent) -> Vec<Start> {
609  let answer = match event {
610    QueryEvent::Answer(a) => a,
611    QueryEvent::Terminal(_) => return Vec::new(),
612  };
613  match step {
614    Step::Ptr => {
615      if answer.rtype() != ResourceType::Ptr {
616        return Vec::new();
617      }
618      match parse_name(answer.rdata_slice()) {
619        Some(instance) => resolver.on_ptr(instance),
620        None => Vec::new(),
621      }
622    }
623    Step::Srv(inst_key) => {
624      if answer.rtype() != ResourceType::Srv {
625        return Vec::new();
626      }
627      match parse_srv(answer.rdata_slice()) {
628        Some((host, port)) => resolver.on_srv(&inst_key, host, port),
629        None => Vec::new(),
630      }
631    }
632    Step::Txt(inst_key) => {
633      if answer.rtype() != ResourceType::Txt {
634        return Vec::new();
635      }
636      resolver.on_txt(&inst_key, parse_txt(answer.rdata_slice()));
637      Vec::new()
638    }
639    Step::A(host_key) => {
640      if answer.rtype() == ResourceType::A
641        && let Ok(r) = A::try_from_rdata(answer.rdata_slice())
642      {
643        resolver.on_addr(&host_key, IpAddr::V4(r.addr()));
644      }
645      Vec::new()
646    }
647    Step::AAAA(host_key) => {
648      if answer.rtype() == ResourceType::AAAA
649        && let Ok(r) = AAAA::try_from_rdata(answer.rdata_slice())
650      {
651        resolver.on_addr(&host_key, IpAddr::V6(r.addr()));
652      }
653      Vec::new()
654    }
655  }
656}
hick_compio/discovery/mod.rs

hick_compio/discovery/
mod.rs
