From cfd22c8395055818fbb65b2dc34bc89e22772c4a Mon Sep 17 00:00:00 2001 From: Bartek Plotka Date: Tue, 10 Nov 2026 16:49:11 +0000 Subject: [PATCH 1/2] discovery/dns: add DnsSdConfig, ResolvConf, and target group mapping Signed-off-by: Bartek Plotka --- src/discovery.rs | 41 +++- src/discovery/dns.rs | 431 +++++++++++++++++++++++++++++++++++++++++ tests/discovery_dns.rs | 308 +++++++++++++++++++++++++++++ 3 files changed, 779 insertions(+), 1 deletion(-) create mode 100644 src/discovery/dns.rs create mode 100644 tests/discovery_dns.rs diff --git a/src/discovery.rs b/src/discovery.rs index cd23405..aacef43 100644 --- a/src/discovery.rs +++ b/src/discovery.rs @@ -11,13 +11,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Prometheus-compatible target discovery configuration and target group types. +//! Prometheus-compatible target discovery (`static_configs` and `dns_sd_configs`). + +pub mod dns; use std::collections::BTreeMap; use std::time::Duration; use serde::{Deserialize, Deserializer, Serialize, Serializer}; +pub use dns::{DnsDiscovery, DnsRecord, DnsRecordType, DnsSdConfig, ResolvConf}; + /// Reserved Prometheus label holding the `:` scrape target address. pub const ADDRESS_LABEL: &str = "__address__"; @@ -33,6 +37,9 @@ pub enum DiscoveryError { /// I/O or filesystem error. #[error("I/O error: {0}")] Io(#[from] std::io::Error), + /// DNS protocol or lookup error. + #[error("DNS error: {0}")] + Dns(String), /// JSON payload decoding error. #[error("JSON decode error: {0}")] Json(#[from] serde_json::Error), @@ -129,11 +136,17 @@ pub struct DiscoveryConfigs { /// Static target group configurations (`static_configs`). #[serde(default, skip_serializing_if = "Vec::is_empty")] pub static_configs: Vec, + /// DNS service discovery configurations (`dns_sd_configs`). + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub dns_sd_configs: Vec, } impl DiscoveryConfigs { /// Validates all configured service discovery sections according to Prometheus rules. pub fn validate(&mut self) -> Result<(), DiscoveryError> { + for dns_cfg in &mut self.dns_sd_configs { + dns_cfg.validate()?; + } Ok(()) } @@ -143,6 +156,15 @@ impl DiscoveryConfigs { } } +/// Joins a host and port into `"host:port"` or `"[ipv6]:port"`, matching Go `net.JoinHostPort`. +pub(crate) fn join_host_port(host: &str, port: &str) -> String { + if host.contains(':') { + format!("[{host}]:{port}") + } else { + format!("{host}:{port}") + } +} + /// Parses a Prometheus duration string (such as `"30s"`, `"1m"`, `"500ms"`, `"1h30m"`, `"7d"`, /// `"1w"`, `"1y"`, or `"0"`) into a [`Duration`], matching `prometheus/common/model.ParseDuration`. pub fn parse_duration(s: &str) -> Result { @@ -248,3 +270,20 @@ pub fn format_duration(d: Duration) -> String { out } + +pub(crate) mod serde_duration { + use std::time::Duration; + + use serde::{Deserialize, Deserializer, Serializer}; + + use super::{format_duration, parse_duration}; + + pub fn serialize(duration: &Duration, serializer: S) -> Result { + serializer.serialize_str(&format_duration(*duration)) + } + + pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result { + let s = String::deserialize(deserializer)?; + parse_duration(&s).map_err(serde::de::Error::custom) + } +} diff --git a/src/discovery/dns.rs b/src/discovery/dns.rs new file mode 100644 index 0000000..88aeb58 --- /dev/null +++ b/src/discovery/dns.rs @@ -0,0 +1,431 @@ +// Copyright 2026 The Prometheus Authors +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Prometheus-compatible DNS service discovery (`dns_sd_configs`). + +use std::collections::BTreeMap; +use std::fs; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::path::Path; +use std::time::Duration; + +use serde::{Deserialize, Serialize}; + +use super::{ADDRESS_LABEL, DiscoveryError, LabelSet, TargetGroup, join_host_port, serde_duration}; + +/// Meta label storing the queried record name. +pub const DNS_NAME_LABEL: &str = "__meta_dns_name"; +/// Meta label storing the target field of an SRV record. +pub const DNS_SRV_RECORD_TARGET_LABEL: &str = "__meta_dns_srv_record_target"; +/// Meta label storing the port field of an SRV record. +pub const DNS_SRV_RECORD_PORT_LABEL: &str = "__meta_dns_srv_record_port"; +/// Meta label storing the target field of an MX record. +pub const DNS_MX_RECORD_TARGET_LABEL: &str = "__meta_dns_mx_record_target"; +/// Meta label storing the target field of an NS record. +pub const DNS_NS_RECORD_TARGET_LABEL: &str = "__meta_dns_ns_record_target"; + +const fn default_refresh_interval() -> Duration { + Duration::from_secs(30) +} + +fn default_record_type() -> String { + "SRV".to_owned() +} + +/// Supported DNS record query types in Prometheus `dns_sd_configs`. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum DnsRecordType { + /// RFC 2782 `SRV` service record (default). + #[default] + Srv, + /// IPv4 `A` address record. + A, + /// IPv6 `AAAA` address record. + Aaaa, + /// Mail exchange `MX` record. + Mx, + /// Name server `NS` record. + Ns, +} + +impl DnsRecordType { + /// Parses a case-insensitive DNS record type string (`"SRV"`, `"A"`, `"AAAA"`, `"MX"`, `"NS"`). + pub fn parse(s: &str) -> Result { + match s.to_ascii_uppercase().as_str() { + "SRV" => Ok(Self::Srv), + "A" => Ok(Self::A), + "AAAA" => Ok(Self::Aaaa), + "MX" => Ok(Self::Mx), + "NS" => Ok(Self::Ns), + _ => Err(DiscoveryError::InvalidConfig(format!( + "invalid DNS-SD records type {s}" + ))), + } + } +} + +/// Parsed DNS answer record used to populate discovered targets. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DnsRecord { + /// `SRV` answer record. + Srv { + /// Priority of the target host. + priority: u16, + /// Relative weight for records with the same priority. + weight: u16, + /// TCP/UDP port on the target host. + port: u16, + /// Target domain name (typically with trailing dot). + target: String, + }, + /// IPv4 `A` answer record. + A(Ipv4Addr), + /// IPv6 `AAAA` answer record. + Aaaa(Ipv6Addr), + /// `MX` answer record. + Mx { + /// Preference value given to this RR. + preference: u16, + /// Mail exchange host domain name. + target: String, + }, + /// `NS` answer record. + Ns { + /// Authoritative name server host domain name. + target: String, + }, + /// `CNAME` alias record (skipped when constructing targets, matching Prometheus). + Cname(String), +} + +/// DNS service discovery configuration matching Prometheus `discovery/dns.SDConfig`. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct DnsSdConfig { + /// DNS domain names to query. + #[serde(default)] + pub names: Vec, + /// Refresh interval between DNS lookups (defaults to `30s`). + #[serde(default = "default_refresh_interval", with = "serde_duration")] + pub refresh_interval: Duration, + /// DNS record type to query: `"SRV"` (default), `"A"`, `"AAAA"`, `"MX"`, or `"NS"`. + #[serde(rename = "type", default = "default_record_type")] + pub record_type: String, + /// Port number used for record types other than `"SRV"` (ignored for `"SRV"`). + #[serde(default, skip_serializing_if = "is_zero_port")] + pub port: u16, +} + +const fn is_zero_port(port: &u16) -> bool { + *port == 0 +} + +impl Default for DnsSdConfig { + fn default() -> Self { + Self { + names: Vec::new(), + refresh_interval: default_refresh_interval(), + record_type: default_record_type(), + port: 0, + } + } +} + +impl DnsSdConfig { + /// Validates this configuration according to Prometheus `dns.SDConfig.UnmarshalYAML`. + pub fn validate(&mut self) -> Result<(), DiscoveryError> { + if self.names.is_empty() { + return Err(DiscoveryError::InvalidConfig( + "DNS-SD config must contain at least one SRV record name".to_owned(), + )); + } + if self.record_type.is_empty() { + self.record_type = default_record_type(); + } + let parsed = DnsRecordType::parse(&self.record_type)?; + if parsed != DnsRecordType::Srv && self.port == 0 { + return Err(DiscoveryError::InvalidConfig( + "a port is required in DNS-SD configs for all record types except SRV".to_owned(), + )); + } + Ok(()) + } + + /// Returns the parsed [`DnsRecordType`] after validation. + pub fn parsed_record_type(&self) -> Result { + DnsRecordType::parse(&self.record_type) + } +} + +/// Parsed `/etc/resolv.conf` settings used for search-path DNS resolution. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ResolvConf { + /// Nameserver socket addresses (default `127.0.0.1:53`). + pub servers: Vec, + /// Search domain suffixes. + pub search: Vec, + /// Threshold of dots in a query name before initial absolute lookup (`ndots`, default `1`). + pub ndots: usize, + /// Per-query timeout (default `5s`). + pub timeout: Duration, +} + +impl Default for ResolvConf { + fn default() -> Self { + Self { + servers: vec![SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 53)], + search: Vec::new(), + ndots: 1, + timeout: Duration::from_secs(5), + } + } +} + +impl ResolvConf { + /// Parses a `resolv.conf` file from disk, returning an I/O error if unreadable. + pub fn from_file(path: &Path) -> Result { + let content = fs::read_to_string(path)?; + Ok(Self::parse(&content)) + } + + /// Parses the textual contents of an `/etc/resolv.conf` file. + pub fn parse(content: &str) -> Self { + let mut servers = Vec::new(); + let mut search = Vec::new(); + let mut ndots = 1usize; + let mut timeout = Duration::from_secs(5); + + for raw_line in content.lines() { + let line = raw_line.split(['#', ';']).next().unwrap_or("").trim(); + if line.is_empty() { + continue; + } + let mut parts = line.split_whitespace(); + let Some(directive) = parts.next() else { + continue; + }; + match directive { + "nameserver" => { + if let Some(addr_str) = parts.next() { + if let Ok(ip) = addr_str.parse::() { + servers.push(SocketAddr::new(ip, 53)); + } else if let Ok(sa) = addr_str.parse::() { + servers.push(sa); + } + } + } + "domain" => { + search.clear(); + if let Some(dom) = parts.next() { + let clean = dom.trim_end_matches('.'); + if !clean.is_empty() { + search.push(clean.to_owned()); + } + } + } + "search" => { + search.clear(); + for dom in parts { + let clean = dom.trim_end_matches('.'); + if !clean.is_empty() { + search.push(clean.to_owned()); + } + } + } + "options" => { + for opt in parts { + if let Some(val) = opt.strip_prefix("ndots:") { + if let Ok(n) = val.parse::() { + ndots = n.min(15); + } + } else if let Some(val) = opt.strip_prefix("timeout:") + && let Ok(secs) = val.parse::() + { + timeout = Duration::from_secs(secs.clamp(1, 30)); + } + } + } + _ => {} + } + } + + if servers.is_empty() { + servers.push(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 53)); + } + + Self { + servers, + search, + ndots, + timeout, + } + } +} + +/// DNS service discovery runner matching Prometheus `discovery/dns.Discovery`. +#[derive(Clone, Debug)] +pub struct DnsDiscovery { + config: DnsSdConfig, + record_type: DnsRecordType, + resolv_conf: Option, +} + +impl DnsDiscovery { + /// Creates a new [`DnsDiscovery`] after validating `config`, reading `/etc/resolv.conf` on each refresh. + pub fn new(mut config: DnsSdConfig) -> Result { + config.validate()?; + let record_type = config.parsed_record_type()?; + Ok(Self { + config, + record_type, + resolv_conf: None, + }) + } + + /// Creates a new [`DnsDiscovery`] with an explicit [`ResolvConf`]. + pub fn with_resolv_conf( + mut config: DnsSdConfig, + resolv_conf: ResolvConf, + ) -> Result { + config.validate()?; + let record_type = config.parsed_record_type()?; + Ok(Self { + config, + record_type, + resolv_conf: Some(resolv_conf), + }) + } + + /// Returns the validated [`DnsSdConfig`]. + pub fn config(&self) -> &DnsSdConfig { + &self.config + } + + /// Returns the explicit [`ResolvConf`] if configured via [`DnsDiscovery::with_resolv_conf`]. + pub fn resolv_conf(&self) -> Option<&ResolvConf> { + self.resolv_conf.as_ref() + } + + /// Performs a single discovery refresh using a caller-supplied synchronous lookup function. + /// + /// If the lookup function returns an empty record list (for example on `NXDOMAIN`), an empty + /// [`TargetGroup`] with `source = name` is emitted so previously discovered targets for that + /// source are cleared; if it returns `Err`, that name is skipped so previous targets are + /// preserved while remaining names continue, matching Prometheus `dns.Discovery.refresh`. + pub fn refresh_with(&self, mut lookup_fn: F) -> Result, DiscoveryError> + where + F: FnMut(&str, DnsRecordType) -> Result, DiscoveryError>, + { + let mut groups = Vec::with_capacity(self.config.names.len()); + for name in &self.config.names { + if let Ok(records) = lookup_fn(name, self.record_type) { + groups.push(records_to_target_group( + name, + self.record_type, + self.config.port, + &records, + )); + } + } + Ok(groups) + } +} + +/// Converts a slice of resolved [`DnsRecord`]s for `name` into a Prometheus [`TargetGroup`], +/// matching `discovery/dns.Discovery.refreshOne`. +pub(crate) fn records_to_target_group( + name: &str, + record_type: DnsRecordType, + default_port: u16, + records: &[DnsRecord], +) -> TargetGroup { + let mut targets = Vec::new(); + for record in records { + let mut target_host_raw = String::new(); + let mut srv_target = String::new(); + let mut srv_port = String::new(); + let mut mx_target = String::new(); + let mut ns_target = String::new(); + let addr_port: u16; + + match (record_type, record) { + (_, DnsRecord::Cname(_)) => continue, + (DnsRecordType::Srv, DnsRecord::Srv { port, target, .. }) => { + if !is_safe_domain_target(target) { + continue; + } + target_host_raw.clone_from(target); + srv_target.clone_from(target); + srv_port = port.to_string(); + addr_port = *port; + } + (DnsRecordType::Mx, DnsRecord::Mx { target, .. }) => { + if !is_safe_domain_target(target) { + continue; + } + target_host_raw.clone_from(target); + mx_target.clone_from(target); + addr_port = default_port; + } + (DnsRecordType::Ns, DnsRecord::Ns { target }) => { + if !is_safe_domain_target(target) { + continue; + } + target_host_raw.clone_from(target); + ns_target.clone_from(target); + addr_port = default_port; + } + (DnsRecordType::A, DnsRecord::A(ip)) => { + target_host_raw = ip.to_string(); + addr_port = default_port; + } + (DnsRecordType::Aaaa, DnsRecord::Aaaa(ip)) => { + target_host_raw = ip.to_string(); + addr_port = default_port; + } + _ => continue, + } + + // Remove the trailing dot from rooted DNS names to make them look more usual. + let host = target_host_raw.trim_end_matches('.'); + let address = join_host_port(host, &addr_port.to_string()); + + let mut labels = LabelSet::new(); + labels.insert(ADDRESS_LABEL.to_owned(), address); + labels.insert(DNS_NAME_LABEL.to_owned(), name.to_owned()); + labels.insert(DNS_SRV_RECORD_TARGET_LABEL.to_owned(), srv_target); + labels.insert(DNS_SRV_RECORD_PORT_LABEL.to_owned(), srv_port); + labels.insert(DNS_MX_RECORD_TARGET_LABEL.to_owned(), mx_target); + labels.insert(DNS_NS_RECORD_TARGET_LABEL.to_owned(), ns_target); + targets.push(labels); + } + + TargetGroup { + targets, + labels: BTreeMap::new(), + source: name.to_owned(), + } +} + +fn is_safe_dns_label(label: &str) -> bool { + !label.bytes().any(|b| { + b.is_ascii_control() + || b.is_ascii_whitespace() + || matches!(b, b'@' | b'/' | b'\\' | b'?' | b'#' | b':' | b'[' | b']') + }) +} + +fn is_safe_domain_target(target: &str) -> bool { + let trimmed = target.trim_end_matches('.'); + trimmed.is_empty() || trimmed.split('.').all(is_safe_dns_label) +} diff --git a/tests/discovery_dns.rs b/tests/discovery_dns.rs new file mode 100644 index 0000000..420f6c9 --- /dev/null +++ b/tests/discovery_dns.rs @@ -0,0 +1,308 @@ +// Copyright 2026 The Prometheus Authors +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Integration tests for Prometheus DNS service discovery (`dns_sd_configs`). + +#![cfg(feature = "discovery")] + +use std::collections::BTreeMap; +use std::net::{Ipv4Addr, Ipv6Addr}; +use std::time::Duration; + +use expfmt::discovery::{ + DiscoveryConfigs, DiscoveryError, DnsDiscovery, DnsRecord, DnsRecordType, DnsSdConfig, + LabelSet, ResolvConf, TargetGroup, +}; + +fn label_set(pairs: &[(&str, &str)]) -> LabelSet { + pairs + .iter() + .map(|(k, v)| ((*k).to_owned(), (*v).to_owned())) + .collect() +} + +#[test] +fn dns_sd_config_validation_matches_prometheus() { + // Ported from TestDNSSDConfigValidation in discovery/dns/dns_test.go. + let valid_cases = [ + r#"{"names": ["example.com"]}"#, + r#"{"names": ["example.com"], "type": "SRV"}"#, + r#"{"names": ["example.com"], "type": "A", "port": 80}"#, + r#"{"names": ["example.com"], "type": "AAAA", "port": 80}"#, + r#"{"names": ["example.com"], "type": "MX", "port": 25}"#, + r#"{"names": ["example.com"], "type": "NS", "port": 53}"#, + r#"{"names": ["example.com"], "refresh_interval": "15s", "type": "srv"}"#, + ]; + for case in valid_cases { + let mut cfg: DnsSdConfig = serde_json::from_str(case).expect(case); + assert!(cfg.validate().is_ok(), "expected valid config: {case}"); + } + + let invalid_cases = [ + ( + r#"{}"#, + "DNS-SD config must contain at least one SRV record name", + ), + ( + r#"{"names": []}"#, + "DNS-SD config must contain at least one SRV record name", + ), + ( + r#"{"names": ["example.com"], "type": "A"}"#, + "a port is required in DNS-SD configs for all record types except SRV", + ), + ( + r#"{"names": ["example.com"], "type": "AAAA"}"#, + "a port is required in DNS-SD configs for all record types except SRV", + ), + ( + r#"{"names": ["example.com"], "type": "MX"}"#, + "a port is required in DNS-SD configs for all record types except SRV", + ), + ( + r#"{"names": ["example.com"], "type": "NS"}"#, + "a port is required in DNS-SD configs for all record types except SRV", + ), + ( + r#"{"names": ["example.com"], "type": "PTR", "port": 80}"#, + "invalid DNS-SD records type PTR", + ), + ]; + for (case, expected_err) in invalid_cases { + let mut cfg: DnsSdConfig = serde_json::from_str(case).expect(case); + let err = cfg.validate().expect_err(case).to_string(); + assert_eq!(err, expected_err, "case: {case}"); + } + + // Unitless integers, overflowing duration strings, and unknown fields are rejected. + assert!( + serde_json::from_str::( + r#"{"names": ["example.com"], "refresh_interval": 30}"# + ) + .is_err() + ); + assert!( + serde_json::from_str::( + r#"{"names": ["example.com"], "refresh_interval": "999999999999999y"}"# + ) + .is_err() + ); + assert!(serde_json::from_str::(r#"{"dns_sd_config": []}"#).is_err()); +} + +#[test] +fn dns_discovery_records_match_prometheus_dns_test() { + // Ported from TestDNS in discovery/dns/dns_test.go. + // 1. A record discovery. + let sd = DnsDiscovery::new(DnsSdConfig { + names: vec!["web.example.com.".to_owned()], + refresh_interval: Duration::from_secs(30), + record_type: "A".to_owned(), + port: 80, + }) + .unwrap(); + let groups = sd + .refresh_with(|name, qtype| { + assert_eq!(name, "web.example.com."); + assert_eq!(qtype, DnsRecordType::A); + Ok(vec![DnsRecord::A(Ipv4Addr::new(192, 0, 2, 2))]) + }) + .unwrap(); + assert_eq!( + groups, + vec![TargetGroup { + source: "web.example.com.".to_owned(), + labels: BTreeMap::new(), + targets: vec![label_set(&[ + ("__address__", "192.0.2.2:80"), + ("__meta_dns_name", "web.example.com."), + ("__meta_dns_srv_record_target", ""), + ("__meta_dns_srv_record_port", ""), + ("__meta_dns_mx_record_target", ""), + ("__meta_dns_ns_record_target", ""), + ])], + }] + ); + + // 2. AAAA record discovery. + let sd = DnsDiscovery::new(DnsSdConfig { + names: vec!["web.example.com.".to_owned()], + refresh_interval: Duration::from_secs(30), + record_type: "AAAA".to_owned(), + port: 80, + }) + .unwrap(); + let groups = sd + .refresh_with(|_, qtype| { + assert_eq!(qtype, DnsRecordType::Aaaa); + Ok(vec![DnsRecord::Aaaa( + "2001:db8::2".parse::().unwrap(), + )]) + }) + .unwrap(); + assert_eq!( + groups, + vec![TargetGroup { + source: "web.example.com.".to_owned(), + labels: BTreeMap::new(), + targets: vec![label_set(&[ + ("__address__", "[2001:db8::2]:80"), + ("__meta_dns_name", "web.example.com."), + ("__meta_dns_srv_record_target", ""), + ("__meta_dns_srv_record_port", ""), + ("__meta_dns_mx_record_target", ""), + ("__meta_dns_ns_record_target", ""), + ])], + }] + ); + + // 3. SRV record discovery with CNAME in answer section. + let sd = DnsDiscovery::new(DnsSdConfig { + names: vec!["_mysql._tcp.db.example.com.".to_owned()], + refresh_interval: Duration::from_secs(30), + record_type: "SRV".to_owned(), + port: 0, + }) + .unwrap(); + let groups = sd + .refresh_with(|_, qtype| { + assert_eq!(qtype, DnsRecordType::Srv); + Ok(vec![ + DnsRecord::Cname("_mysql._tcp.alias.example.com.".to_owned()), + DnsRecord::Srv { + priority: 0, + weight: 0, + port: 3306, + target: "db1.example.com.".to_owned(), + }, + DnsRecord::Srv { + priority: 0, + weight: 0, + port: 3306, + target: "db2.example.com.".to_owned(), + }, + ]) + }) + .unwrap(); + assert_eq!( + groups, + vec![TargetGroup { + source: "_mysql._tcp.db.example.com.".to_owned(), + labels: BTreeMap::new(), + targets: vec![ + label_set(&[ + ("__address__", "db1.example.com:3306"), + ("__meta_dns_name", "_mysql._tcp.db.example.com."), + ("__meta_dns_srv_record_target", "db1.example.com."), + ("__meta_dns_srv_record_port", "3306"), + ("__meta_dns_mx_record_target", ""), + ("__meta_dns_ns_record_target", ""), + ]), + label_set(&[ + ("__address__", "db2.example.com:3306"), + ("__meta_dns_name", "_mysql._tcp.db.example.com."), + ("__meta_dns_srv_record_target", "db2.example.com."), + ("__meta_dns_srv_record_port", "3306"), + ("__meta_dns_mx_record_target", ""), + ("__meta_dns_ns_record_target", ""), + ]), + ], + }] + ); + + // 4. MX record discovery. + let sd = DnsDiscovery::new(DnsSdConfig { + names: vec!["example.com.".to_owned()], + refresh_interval: Duration::from_secs(30), + record_type: "MX".to_owned(), + port: 25, + }) + .unwrap(); + let groups = sd + .refresh_with(|_, _| { + Ok(vec![ + DnsRecord::Mx { + preference: 10, + target: "smtp1.example.com.".to_owned(), + }, + DnsRecord::Mx { + preference: 20, + target: "smtp2.example.com.".to_owned(), + }, + ]) + }) + .unwrap(); + assert_eq!( + groups[0].targets, + vec![ + label_set(&[ + ("__address__", "smtp1.example.com:25"), + ("__meta_dns_name", "example.com."), + ("__meta_dns_srv_record_target", ""), + ("__meta_dns_srv_record_port", ""), + ("__meta_dns_mx_record_target", "smtp1.example.com."), + ("__meta_dns_ns_record_target", ""), + ]), + label_set(&[ + ("__address__", "smtp2.example.com:25"), + ("__meta_dns_name", "example.com."), + ("__meta_dns_srv_record_target", ""), + ("__meta_dns_srv_record_port", ""), + ("__meta_dns_mx_record_target", "smtp2.example.com."), + ("__meta_dns_ns_record_target", ""), + ]), + ] + ); + + // 5. NS record discovery + NXDOMAIN empty answer + per-name lookup failure isolation. + let sd = DnsDiscovery::new(DnsSdConfig { + names: vec![ + "servfail.example.com.".to_owned(), + "example.com.".to_owned(), + "nxdomain.example.com.".to_owned(), + ], + refresh_interval: Duration::from_secs(30), + record_type: "NS".to_owned(), + port: 53, + }) + .unwrap(); + let groups = sd + .refresh_with(|name, _| { + if name == "servfail.example.com." { + Err(DiscoveryError::Dns("SERVFAIL".to_owned())) + } else if name == "nxdomain.example.com." { + Ok(vec![]) + } else { + Ok(vec![DnsRecord::Ns { + target: "ns1.example.com.".to_owned(), + }]) + } + }) + .unwrap(); + assert_eq!(groups.len(), 2); + assert_eq!( + groups[0].targets, + vec![label_set(&[ + ("__address__", "ns1.example.com:53"), + ("__meta_dns_name", "example.com."), + ("__meta_dns_srv_record_target", ""), + ("__meta_dns_srv_record_port", ""), + ("__meta_dns_mx_record_target", ""), + ("__meta_dns_ns_record_target", "ns1.example.com."), + ])] + ); + assert_eq!(groups[1].source, "nxdomain.example.com."); + assert!(groups[1].targets.is_empty()); + + assert!(ResolvConf::from_file(std::path::Path::new("/nonexistent/resolv.conf")).is_err()); +} From 809d4d952d4c978faf91d1c80cf23ccc52b9e13d Mon Sep 17 00:00:00 2001 From: Bartek Plotka Date: Tue, 10 Nov 2026 16:51:07 +0000 Subject: [PATCH 2/2] discovery/dns: add async UDP/EDNS0 and TCP wire resolver Signed-off-by: Bartek Plotka --- Cargo.lock | 48 +++ Cargo.toml | 4 +- src/discovery/dns.rs | 709 +++++++++++++++++++++++++++++++++++++++++ tests/discovery_dns.rs | 281 +++++++++++++++- 4 files changed, 1038 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e7d8c44..5d5e465 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -214,6 +214,7 @@ dependencies = [ "serde", "serde_json", "thiserror", + "tokio", ] [[package]] @@ -470,6 +471,17 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "mio" +version = "1.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1788edb87fdc09c7e26304471e2f5be8cdefb1b6930d6e3985fc02ff53bf86ee" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + [[package]] name = "multimap" version = "0.10.1" @@ -764,6 +776,16 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys", +] + [[package]] name = "syn" version = "2.0.119" @@ -827,6 +849,32 @@ dependencies = [ "crunchy", ] +[[package]] +name = "tokio" +version = "1.54.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b52fd3255258d26fac7264ab5a76b53c13367136dc011d3d6ad4f8ed070396a3" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + [[package]] name = "unicode-ident" version = "1.0.27" diff --git a/Cargo.toml b/Cargo.toml index a37a80e..a577dfa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,7 @@ exclude = ["tests/otap-compat"] [features] default = ["arrow", "dto"] arrow = ["dep:arrow-array", "dep:arrow-buffer", "dep:arrow-schema"] -discovery = ["dep:serde", "dep:serde_json"] +discovery = ["dep:serde", "dep:serde_json", "dep:tokio"] dto = ["dep:prost", "dep:prost-types"] [dependencies] @@ -27,11 +27,13 @@ prost-types = { version = "0.14", optional = true } serde = { version = "1", features = ["derive"], optional = true } serde_json = { version = "1", optional = true } thiserror = "2" +tokio = { version = "1", default-features = false, features = ["io-util", "net", "time"], optional = true } [dev-dependencies] benchfmt = { git = "https://github.com/bwplotka/benchfmt.rs" } prost-build = "=0.14.4" protox = "0.10" +tokio = { version = "1", default-features = false, features = ["macros", "rt"] } [[bench]] name = "decode" diff --git a/src/discovery/dns.rs b/src/discovery/dns.rs index 88aeb58..af843d6 100644 --- a/src/discovery/dns.rs +++ b/src/discovery/dns.rs @@ -14,12 +14,17 @@ //! Prometheus-compatible DNS service discovery (`dns_sd_configs`). use std::collections::BTreeMap; +use std::collections::hash_map::RandomState; use std::fs; +use std::hash::BuildHasher; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::path::Path; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use serde::{Deserialize, Serialize}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpStream, UdpSocket}; use super::{ADDRESS_LABEL, DiscoveryError, LabelSet, TargetGroup, join_host_port, serde_duration}; @@ -34,6 +39,18 @@ pub const DNS_MX_RECORD_TARGET_LABEL: &str = "__meta_dns_mx_record_target"; /// Meta label storing the target field of an NS record. pub const DNS_NS_RECORD_TARGET_LABEL: &str = "__meta_dns_ns_record_target"; +const QTYPE_A: u16 = 1; +const QTYPE_NS: u16 = 2; +const QTYPE_CNAME: u16 = 5; +const QTYPE_MX: u16 = 15; +const QTYPE_AAAA: u16 = 28; +const QTYPE_SRV: u16 = 33; +const QTYPE_OPT: u16 = 41; +const QCLASS_IN: u16 = 1; + +const RCODE_SUCCESS: u8 = 0; +const RCODE_NXDOMAIN: u8 = 3; + const fn default_refresh_interval() -> Duration { Duration::from_secs(30) } @@ -72,6 +89,17 @@ impl DnsRecordType { ))), } } + + /// Returns the DNS wire protocol QTYPE code for this record type. + pub(crate) const fn qtype(self) -> u16 { + match self { + Self::A => QTYPE_A, + Self::Ns => QTYPE_NS, + Self::Mx => QTYPE_MX, + Self::Aaaa => QTYPE_AAAA, + Self::Srv => QTYPE_SRV, + } + } } /// Parsed DNS answer record used to populate discovered targets. @@ -316,6 +344,35 @@ impl DnsDiscovery { self.resolv_conf.as_ref() } + /// Performs a single discovery refresh across all configured DNS names using the configured + /// nameservers and search paths. + /// + /// Lookup failures for individual names are skipped so previously discovered targets for + /// those sources are preserved while healthy names still emit [`TargetGroup`]s, matching + /// Prometheus `dns.Discovery.refresh`. + pub async fn refresh(&self) -> Result, DiscoveryError> { + let owned_conf; + let conf = match &self.resolv_conf { + Some(c) => c, + None => { + owned_conf = ResolvConf::from_file(Path::new("/etc/resolv.conf"))?; + &owned_conf + } + }; + let mut groups = Vec::with_capacity(self.config.names.len()); + for name in &self.config.names { + if let Ok(records) = self.lookup_with_search_path(name, conf).await { + groups.push(records_to_target_group( + name, + self.record_type, + self.config.port, + &records, + )); + } + } + Ok(groups) + } + /// Performs a single discovery refresh using a caller-supplied synchronous lookup function. /// /// If the lookup function returns an empty record list (for example on `NXDOMAIN`), an empty @@ -339,6 +396,54 @@ impl DnsDiscovery { } Ok(groups) } + + async fn lookup_with_search_path( + &self, + name: &str, + resolv_conf: &ResolvConf, + ) -> Result, DiscoveryError> { + let candidates = candidate_query_names(name, resolv_conf); + let mut last_err: Option = None; + + for candidate in &candidates { + match self.lookup_all_servers(candidate, resolv_conf).await { + Ok(Some(records)) => return Ok(records), + Ok(None) => { + // NXDOMAIN: continue trying remaining search path candidates, and if none + // match, return an empty target group (matching Prometheus). + } + Err(e) => { + last_err = Some(e); + } + } + } + + if let Some(err) = last_err { + Err(err) + } else { + Ok(Vec::new()) + } + } + + async fn lookup_all_servers( + &self, + fqdn: &str, + resolv_conf: &ResolvConf, + ) -> Result>, DiscoveryError> { + let qtype = self.record_type.qtype(); + let mut last_err = None; + + for &server in &resolv_conf.servers { + match query_dns_server(server, fqdn, qtype, resolv_conf.timeout).await { + Ok(DnsLookupOutcome::Records(records)) => return Ok(Some(records)), + Ok(DnsLookupOutcome::NxDomain) => return Ok(None), + Err(e) => last_err = Some(e), + } + } + + Err(last_err + .unwrap_or_else(|| DiscoveryError::Dns("no DNS nameservers configured".to_owned()))) + } } /// Converts a slice of resolved [`DnsRecord`]s for `name` into a Prometheus [`TargetGroup`], @@ -417,6 +522,45 @@ pub(crate) fn records_to_target_group( } } +fn candidate_query_names(name: &str, conf: &ResolvConf) -> Vec { + if name.ends_with('.') { + return vec![name.to_owned()]; + } + let dots = name.bytes().filter(|&b| b == b'.').count(); + let mut search_names: Vec = conf + .search + .iter() + .map(|domain| format!("{name}.{domain}.")) + .filter(|cand| wire_qname_len(cand).is_some_and(|n| n <= 255)) + .collect(); + let rooted = format!("{name}."); + + if dots >= conf.ndots { + let mut out = Vec::with_capacity(1 + search_names.len()); + out.push(rooted); + out.extend(search_names); + out + } else { + search_names.push(rooted); + search_names + } +} + +fn wire_qname_len(fqdn: &str) -> Option { + let trimmed = fqdn.trim_end_matches('.'); + if trimmed.is_empty() { + return Some(1); + } + let mut len = 1usize; + for label in trimmed.split('.') { + if label.is_empty() || label.len() > 63 { + return None; + } + len += 1 + label.len(); + } + Some(len) +} + fn is_safe_dns_label(label: &str) -> bool { !label.bytes().any(|b| { b.is_ascii_control() @@ -429,3 +573,568 @@ fn is_safe_domain_target(target: &str) -> bool { let trimmed = target.trim_end_matches('.'); trimmed.is_empty() || trimmed.split('.').all(is_safe_dns_label) } + +enum DnsLookupOutcome { + Records(Vec), + NxDomain, +} + +fn next_dns_txid() -> u16 { + static NONCE: AtomicU64 = AtomicU64::new(0); + let n = NONCE.fetch_add(1, Ordering::Relaxed); + RandomState::new().hash_one(n) as u16 +} + +async fn query_dns_server( + server: SocketAddr, + fqdn: &str, + qtype: u16, + timeout: Duration, +) -> Result { + let udp_id = next_dns_txid(); + let query_pkt = encode_dns_query(udp_id, fqdn, qtype, true)?; + let bind_addr = if server.is_ipv4() { + SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0) + } else { + SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0) + }; + let udp = UdpSocket::bind(bind_addr).await?; + udp.connect(server).await?; + udp.send(&query_pkt).await?; + + let mut buf = vec![0u8; 4096]; + let len = tokio::time::timeout(timeout, udp.recv(&mut buf)) + .await + .map_err(|_| DiscoveryError::Dns(format!("DNS UDP read timeout from {server}")))??; + match decode_dns_response(udp_id, fqdn, qtype, &buf[..len])? { + DecodedDnsMessage::Truncated => {} + DecodedDnsMessage::NxDomain => return Ok(DnsLookupOutcome::NxDomain), + DecodedDnsMessage::Records(records) => return Ok(DnsLookupOutcome::Records(records)), + } + + // Fallback to TCP on TC=1 truncation bit without EDNS0 OPT, matching Prometheus askServerForName. + let tcp_id = next_dns_txid(); + let tcp_query_pkt = encode_dns_query(tcp_id, fqdn, qtype, false)?; + let tcp_query_len = u16::try_from(tcp_query_pkt.len()) + .map_err(|_| DiscoveryError::Dns("DNS TCP query packet exceeds u16 length".to_owned()))?; + let mut stream = tokio::time::timeout(timeout, TcpStream::connect(server)) + .await + .map_err(|_| DiscoveryError::Dns(format!("DNS TCP connect timeout to {server}")))??; + stream.write_all(&tcp_query_len.to_be_bytes()).await?; + stream.write_all(&tcp_query_pkt).await?; + + let mut resp_len_buf = [0u8; 2]; + tokio::time::timeout(timeout, stream.read_exact(&mut resp_len_buf)) + .await + .map_err(|_| DiscoveryError::Dns(format!("DNS TCP read timeout from {server}")))??; + let resp_len = u16::from_be_bytes(resp_len_buf) as usize; + let mut resp_buf = vec![0u8; resp_len]; + tokio::time::timeout(timeout, stream.read_exact(&mut resp_buf)) + .await + .map_err(|_| DiscoveryError::Dns(format!("DNS TCP body timeout from {server}")))??; + + match decode_dns_response(tcp_id, fqdn, qtype, &resp_buf)? { + DecodedDnsMessage::Truncated => Err(DiscoveryError::Dns( + "truncated DNS response over TCP".to_owned(), + )), + DecodedDnsMessage::NxDomain => Ok(DnsLookupOutcome::NxDomain), + DecodedDnsMessage::Records(records) => Ok(DnsLookupOutcome::Records(records)), + } +} + +/// Encodes an RFC 1035 DNS wire query message, optionally including an EDNS0 OPT pseudo-record +/// (4096-byte UDP payload). +pub(crate) fn encode_dns_query( + id: u16, + fqdn: &str, + qtype: u16, + edns: bool, +) -> Result, DiscoveryError> { + let mut pkt = Vec::with_capacity(64); + pkt.extend_from_slice(&id.to_be_bytes()); + // Flags: RD=1 (recursion desired, 0x0100). + pkt.extend_from_slice(&0x0100u16.to_be_bytes()); + // QDCOUNT = 1, ANCOUNT = 0, NSCOUNT = 0, ARCOUNT = 1 when EDNS0 is enabled. + let arcount: u16 = u16::from(edns); + pkt.extend_from_slice(&1u16.to_be_bytes()); + pkt.extend_from_slice(&0u16.to_be_bytes()); + pkt.extend_from_slice(&0u16.to_be_bytes()); + pkt.extend_from_slice(&arcount.to_be_bytes()); + + let trimmed = fqdn.trim_end_matches('.'); + let mut wire_qname_len = 1usize; + if !trimmed.is_empty() { + for label in trimmed.split('.') { + if label.is_empty() || label.len() > 63 || !is_safe_dns_label(label) { + return Err(DiscoveryError::Dns(format!( + "invalid DNS label in query name {fqdn:?}" + ))); + } + wire_qname_len += 1 + label.len(); + if wire_qname_len > 255 { + return Err(DiscoveryError::Dns(format!( + "DNS query name {fqdn:?} exceeds 255 octets" + ))); + } + pkt.push(label.len() as u8); + pkt.extend_from_slice(label.as_bytes()); + } + } + pkt.push(0); + pkt.extend_from_slice(&qtype.to_be_bytes()); + pkt.extend_from_slice(&QCLASS_IN.to_be_bytes()); + + if edns { + // EDNS0 OPT record: root name (0), TYPE = 41 (OPT), CLASS = 4096 (UDP payload size), + // TTL = 0, RDLENGTH = 0. + pkt.push(0); + pkt.extend_from_slice(&QTYPE_OPT.to_be_bytes()); + pkt.extend_from_slice(&4096u16.to_be_bytes()); + pkt.extend_from_slice(&0u32.to_be_bytes()); + pkt.extend_from_slice(&0u16.to_be_bytes()); + } + + Ok(pkt) +} + +enum DecodedDnsMessage { + Records(Vec), + NxDomain, + Truncated, +} + +fn canonical_fqdn(fqdn: &str) -> String { + let trimmed = fqdn.trim_end_matches('.'); + if trimmed.is_empty() { + ".".to_owned() + } else { + format!("{trimmed}.") + } +} + +fn decode_dns_response( + expected_id: u16, + expected_qname: &str, + expected_qtype: u16, + msg: &[u8], +) -> Result { + if msg.len() < 12 { + return Err(DiscoveryError::Dns("DNS response too short".to_owned())); + } + let resp_id = u16::from_be_bytes([msg[0], msg[1]]); + if resp_id != expected_id { + return Err(DiscoveryError::Dns(format!( + "DNS response transaction ID mismatch: expected {expected_id:#06x}, got {resp_id:#06x}" + ))); + } + let flags = u16::from_be_bytes([msg[2], msg[3]]); + if (flags & 0x8000) == 0 { + return Err(DiscoveryError::Dns( + "DNS packet is a query (QR=0), expected response (QR=1)".to_owned(), + )); + } + let tc = (flags & 0x0200) != 0; + if tc { + return Ok(DecodedDnsMessage::Truncated); + } + + let qdcount = u16::from_be_bytes([msg[4], msg[5]]) as usize; + let ancount = u16::from_be_bytes([msg[6], msg[7]]) as usize; + if qdcount != 1 { + return Err(DiscoveryError::Dns(format!( + "unexpected DNS question count {qdcount}, expected 1" + ))); + } + + let (qname, next) = decode_dns_name(msg, 12)?; + if next + 4 > msg.len() { + return Err(DiscoveryError::Dns("truncated DNS question".to_owned())); + } + let qtype = u16::from_be_bytes([msg[next], msg[next + 1]]); + let qclass = u16::from_be_bytes([msg[next + 2], msg[next + 3]]); + let want_qname = canonical_fqdn(expected_qname); + if !qname.eq_ignore_ascii_case(&want_qname) || qtype != expected_qtype || qclass != QCLASS_IN { + return Err(DiscoveryError::Dns(format!( + "DNS response question mismatch: expected ({want_qname}, {expected_qtype}, {QCLASS_IN}), got ({qname}, {qtype}, {qclass})" + ))); + } + let mut offset = next + 4; + + let rcode = (flags & 0x000F) as u8; + if rcode == RCODE_NXDOMAIN { + return Ok(DecodedDnsMessage::NxDomain); + } + if rcode != RCODE_SUCCESS { + return Err(DiscoveryError::Dns(format!( + "DNS response returned non-zero RCODE {rcode}" + ))); + } + + // Each resource record requires at least 1 byte for a root/pointer name plus 10 header bytes. + let max_possible_rrs = msg.len().saturating_sub(offset) / 11; + let mut records = Vec::with_capacity(ancount.min(max_possible_rrs)); + for _ in 0..ancount { + let (_, next) = decode_dns_name(msg, offset)?; + if next + 10 > msg.len() { + return Err(DiscoveryError::Dns( + "truncated DNS resource record header".to_owned(), + )); + } + let rtype = u16::from_be_bytes([msg[next], msg[next + 1]]); + let rclass = u16::from_be_bytes([msg[next + 2], msg[next + 3]]); + let rdlength = u16::from_be_bytes([msg[next + 8], msg[next + 9]]) as usize; + let rdata_start = next + 10; + let rdata_end = rdata_start + rdlength; + if rdata_end > msg.len() { + return Err(DiscoveryError::Dns( + "truncated DNS resource record rdata".to_owned(), + )); + } + if rclass != QCLASS_IN { + offset = rdata_end; + continue; + } + let rdata = &msg[rdata_start..rdata_end]; + + match rtype { + QTYPE_A => { + if rdata.len() != 4 { + return Err(DiscoveryError::Dns("invalid A record length".to_owned())); + } + records.push(DnsRecord::A(Ipv4Addr::new( + rdata[0], rdata[1], rdata[2], rdata[3], + ))); + } + QTYPE_AAAA => { + if rdata.len() != 16 { + return Err(DiscoveryError::Dns("invalid AAAA record length".to_owned())); + } + let mut octets = [0u8; 16]; + octets.copy_from_slice(rdata); + records.push(DnsRecord::Aaaa(Ipv6Addr::from(octets))); + } + QTYPE_SRV => { + if rdata.len() < 7 { + return Err(DiscoveryError::Dns("invalid SRV record length".to_owned())); + } + let priority = u16::from_be_bytes([rdata[0], rdata[1]]); + let weight = u16::from_be_bytes([rdata[2], rdata[3]]); + let port = u16::from_be_bytes([rdata[4], rdata[5]]); + let (target, next_rdata) = + decode_dns_name_bounded(msg, rdata_start + 6, rdata_end)?; + if next_rdata != rdata_end { + return Err(DiscoveryError::Dns("invalid SRV record length".to_owned())); + } + records.push(DnsRecord::Srv { + priority, + weight, + port, + target, + }); + } + QTYPE_MX => { + if rdata.len() < 3 { + return Err(DiscoveryError::Dns("invalid MX record length".to_owned())); + } + let preference = u16::from_be_bytes([rdata[0], rdata[1]]); + let (target, next_rdata) = + decode_dns_name_bounded(msg, rdata_start + 2, rdata_end)?; + if next_rdata != rdata_end { + return Err(DiscoveryError::Dns("invalid MX record length".to_owned())); + } + records.push(DnsRecord::Mx { preference, target }); + } + QTYPE_NS => { + let (target, next_rdata) = decode_dns_name_bounded(msg, rdata_start, rdata_end)?; + if next_rdata != rdata_end { + return Err(DiscoveryError::Dns("invalid NS record length".to_owned())); + } + records.push(DnsRecord::Ns { target }); + } + QTYPE_CNAME => { + let (target, next_rdata) = decode_dns_name_bounded(msg, rdata_start, rdata_end)?; + if next_rdata != rdata_end { + return Err(DiscoveryError::Dns( + "invalid CNAME record length".to_owned(), + )); + } + records.push(DnsRecord::Cname(target)); + } + _ => {} + } + + offset = rdata_end; + } + + Ok(DecodedDnsMessage::Records(records)) +} + +fn decode_dns_name(msg: &[u8], offset: usize) -> Result<(String, usize), DiscoveryError> { + decode_dns_name_bounded(msg, offset, msg.len()) +} + +fn decode_dns_name_bounded( + msg: &[u8], + mut offset: usize, + inline_bound: usize, +) -> Result<(String, usize), DiscoveryError> { + let mut labels = Vec::new(); + let mut end_offset = None; + let mut jumps = 0usize; + let mut total_len = 0usize; + + loop { + let active_bound = if end_offset.is_none() { + inline_bound + } else { + msg.len() + }; + if offset >= active_bound { + return Err(DiscoveryError::Dns( + "unexpected EOF while decoding DNS domain name".to_owned(), + )); + } + let len = msg[offset]; + if len == 0 { + let final_offset = end_offset.unwrap_or(offset + 1); + let fqdn = if labels.is_empty() { + ".".to_owned() + } else { + let mut s = labels.join("."); + s.push('.'); + s + }; + return Ok((fqdn, final_offset)); + } + if (len & 0xC0) == 0xC0 { + if offset + 1 >= active_bound { + return Err(DiscoveryError::Dns( + "truncated DNS name compression pointer".to_owned(), + )); + } + if end_offset.is_none() { + end_offset = Some(offset + 2); + } + let ptr = (((len & 0x3F) as usize) << 8) | (msg[offset + 1] as usize); + if ptr >= msg.len() { + return Err(DiscoveryError::Dns( + "DNS name compression pointer out of bounds".to_owned(), + )); + } + offset = ptr; + jumps += 1; + if jumps > 64 { + return Err(DiscoveryError::Dns( + "DNS name compression pointer loop detected".to_owned(), + )); + } + continue; + } + if (len & 0xC0) != 0 { + return Err(DiscoveryError::Dns( + "invalid DNS label length byte".to_owned(), + )); + } + let label_len = len as usize; + let start = offset + 1; + let end = start + label_len; + if end > active_bound { + return Err(DiscoveryError::Dns("truncated DNS label".to_owned())); + } + total_len += label_len + 1; + if total_len > 255 { + return Err(DiscoveryError::Dns( + "DNS domain name exceeds 255 bytes".to_owned(), + )); + } + let label = std::str::from_utf8(&msg[start..end]) + .map_err(|e| DiscoveryError::Dns(format!("non-UTF-8 DNS label: {e}")))?; + if !is_safe_dns_label(label) { + return Err(DiscoveryError::Dns(format!( + "unsafe character in DNS label {label:?}" + ))); + } + labels.push(label.to_owned()); + offset = end; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn decode_dns_response_compression_pointers_and_rdata_bounds() { + // Build a response with QDCOUNT=1 (`_http._tcp.example.com.`) at offset 12, + // where `example.com.` starts at offset 23 (12 + 6 ("_http") + 5 ("_tcp") = 23). + let mut pkt = Vec::new(); + pkt.extend_from_slice(&0x5052u16.to_be_bytes()); // ID + pkt.extend_from_slice(&0x8180u16.to_be_bytes()); // QR=1, RD=1, RA=1 + pkt.extend_from_slice(&1u16.to_be_bytes()); // QDCOUNT=1 + pkt.extend_from_slice(&2u16.to_be_bytes()); // ANCOUNT=2 (CNAME + SRV) + pkt.extend_from_slice(&0u16.to_be_bytes()); + pkt.extend_from_slice(&0u16.to_be_bytes()); + + // Question: _http._tcp.example.com. + for label in ["_http", "_tcp", "example", "com"] { + pkt.push(label.len() as u8); + pkt.extend_from_slice(label.as_bytes()); + } + pkt.push(0); + pkt.extend_from_slice(&QTYPE_SRV.to_be_bytes()); + pkt.extend_from_slice(&QCLASS_IN.to_be_bytes()); + + // Answer 1: CNAME with owner pointer 0xC00C -> `alias` + pointer 0xC017 (`example.com.`) + pkt.extend_from_slice(&[0xC0, 0x0C]); // Name pointer to offset 12 + pkt.extend_from_slice(&QTYPE_CNAME.to_be_bytes()); + pkt.extend_from_slice(&QCLASS_IN.to_be_bytes()); + pkt.extend_from_slice(&60u32.to_be_bytes()); + let cname_rdata = [5, b'a', b'l', b'i', b'a', b's', 0xC0, 0x17]; + pkt.extend_from_slice(&(cname_rdata.len() as u16).to_be_bytes()); + pkt.extend_from_slice(&cname_rdata); + + // Answer 2: SRV with owner pointer 0xC00C -> target `node1` + pointer 0xC017 (`example.com.`) + pkt.extend_from_slice(&[0xC0, 0x0C]); + pkt.extend_from_slice(&QTYPE_SRV.to_be_bytes()); + pkt.extend_from_slice(&QCLASS_IN.to_be_bytes()); + pkt.extend_from_slice(&60u32.to_be_bytes()); + let mut srv_rdata = Vec::new(); + srv_rdata.extend_from_slice(&10u16.to_be_bytes()); + srv_rdata.extend_from_slice(&20u16.to_be_bytes()); + srv_rdata.extend_from_slice(&8080u16.to_be_bytes()); + srv_rdata.extend_from_slice(&[5, b'n', b'o', b'd', b'e', b'1', 0xC0, 0x17]); + pkt.extend_from_slice(&(srv_rdata.len() as u16).to_be_bytes()); + pkt.extend_from_slice(&srv_rdata); + + let DecodedDnsMessage::Records(records) = + decode_dns_response(0x5052, "_http._tcp.example.com.", QTYPE_SRV, &pkt).unwrap() + else { + panic!("expected decoded DNS records"); + }; + assert_eq!( + records, + vec![ + DnsRecord::Cname("alias.example.com.".to_owned()), + DnsRecord::Srv { + priority: 10, + weight: 20, + port: 8080, + target: "node1.example.com.".to_owned(), + }, + ] + ); + + // Wrong transaction ID or mismatched question is rejected. + assert!(decode_dns_response(0x1234, "_http._tcp.example.com.", QTYPE_SRV, &pkt).is_err()); + assert!(decode_dns_response(0x5052, "_other._tcp.example.com.", QTYPE_SRV, &pkt).is_err()); + assert!(decode_dns_response(0x5052, "_http._tcp.example.com.", QTYPE_A, &pkt).is_err()); + + // Malformed NS record with RDLENGTH=0 followed by trailing root byte must fail instead of reading past RDATA. + let mut bad_ns = Vec::new(); + bad_ns.extend_from_slice(&0x5052u16.to_be_bytes()); + bad_ns.extend_from_slice(&0x8180u16.to_be_bytes()); + bad_ns.extend_from_slice(&1u16.to_be_bytes()); // QDCOUNT = 1 + bad_ns.extend_from_slice(&1u16.to_be_bytes()); + bad_ns.extend_from_slice(&0u16.to_be_bytes()); + bad_ns.extend_from_slice(&1u16.to_be_bytes()); + bad_ns.push(0); // question: root + bad_ns.extend_from_slice(&QTYPE_NS.to_be_bytes()); + bad_ns.extend_from_slice(&QCLASS_IN.to_be_bytes()); + bad_ns.push(0); // owner: root + bad_ns.extend_from_slice(&QTYPE_NS.to_be_bytes()); + bad_ns.extend_from_slice(&QCLASS_IN.to_be_bytes()); + bad_ns.extend_from_slice(&60u32.to_be_bytes()); + bad_ns.extend_from_slice(&0u16.to_be_bytes()); // RDLENGTH = 0 (invalid for NS) + bad_ns.push(0); // trailing OPT root name byte + assert!(decode_dns_response(0x5052, ".", QTYPE_NS, &bad_ns).is_err()); + + // QR=0 query packet is rejected. + let query_pkt = encode_dns_query(0x5052, "example.com.", QTYPE_A, true).unwrap(); + assert!(decode_dns_response(0x5052, "example.com.", QTYPE_A, &query_pkt).is_err()); + + // EDNS0 controls ARCOUNT (1 vs 0), and >255-octet or unsafe QNAMEs are rejected. + let no_edns_pkt = encode_dns_query(0x5052, "example.com.", QTYPE_A, false).unwrap(); + assert_eq!(u16::from_be_bytes([query_pkt[10], query_pkt[11]]), 1); + assert_eq!(u16::from_be_bytes([no_edns_pkt[10], no_edns_pkt[11]]), 0); + let overlong_qname = format!("{}.", vec!["a".repeat(60); 5].join(".")); + assert!(encode_dns_query(1, &overlong_qname, QTYPE_A, true).is_err()); + assert!(encode_dns_query(1, &"a".repeat(64), QTYPE_A, true).is_err()); + assert!(encode_dns_query(1, "bad@host.example.com.", QTYPE_A, true).is_err()); + + // Compression pointer loop (self-pointer at offset 12) is rejected. + let mut loop_pkt = Vec::new(); + loop_pkt.extend_from_slice(&0x5052u16.to_be_bytes()); + loop_pkt.extend_from_slice(&0x8180u16.to_be_bytes()); + loop_pkt.extend_from_slice(&1u16.to_be_bytes()); + loop_pkt.extend_from_slice(&0u16.to_be_bytes()); + loop_pkt.extend_from_slice(&0u16.to_be_bytes()); + loop_pkt.extend_from_slice(&0u16.to_be_bytes()); + loop_pkt.extend_from_slice(&[0xC0, 0x0C, 0, 1, 0, 1]); + assert!(decode_dns_response(0x5052, "example.com.", QTYPE_A, &loop_pkt).is_err()); + + // Out-of-bounds compression pointer is rejected. + let mut oob_pkt = loop_pkt.clone(); + oob_pkt[12] = 0xC0; + oob_pkt[13] = 0xFF; + assert!(decode_dns_response(0x5052, "example.com.", QTYPE_A, &oob_pkt).is_err()); + + // Unsafe characters in decoded SRV/NS target hostnames are rejected. + let mut unsafe_ns = Vec::new(); + unsafe_ns.extend_from_slice(&0x5052u16.to_be_bytes()); + unsafe_ns.extend_from_slice(&0x8180u16.to_be_bytes()); + unsafe_ns.extend_from_slice(&1u16.to_be_bytes()); + unsafe_ns.extend_from_slice(&1u16.to_be_bytes()); + unsafe_ns.extend_from_slice(&0u16.to_be_bytes()); + unsafe_ns.extend_from_slice(&0u16.to_be_bytes()); + unsafe_ns.push(0); + unsafe_ns.extend_from_slice(&QTYPE_NS.to_be_bytes()); + unsafe_ns.extend_from_slice(&QCLASS_IN.to_be_bytes()); + unsafe_ns.push(0); + unsafe_ns.extend_from_slice(&QTYPE_NS.to_be_bytes()); + unsafe_ns.extend_from_slice(&QCLASS_IN.to_be_bytes()); + unsafe_ns.extend_from_slice(&60u32.to_be_bytes()); + let bad_target = b"\x0bad@evil.com\x00"; + unsafe_ns.extend_from_slice(&(bad_target.len() as u16).to_be_bytes()); + unsafe_ns.extend_from_slice(bad_target); + assert!(decode_dns_response(0x5052, ".", QTYPE_NS, &unsafe_ns).is_err()); + + // Trailing extra byte inside NS RDATA is rejected. + let mut trailing_ns = Vec::new(); + trailing_ns.extend_from_slice(&0x5052u16.to_be_bytes()); + trailing_ns.extend_from_slice(&0x8180u16.to_be_bytes()); + trailing_ns.extend_from_slice(&1u16.to_be_bytes()); + trailing_ns.extend_from_slice(&1u16.to_be_bytes()); + trailing_ns.extend_from_slice(&0u16.to_be_bytes()); + trailing_ns.extend_from_slice(&0u16.to_be_bytes()); + trailing_ns.push(0); + trailing_ns.extend_from_slice(&QTYPE_NS.to_be_bytes()); + trailing_ns.extend_from_slice(&QCLASS_IN.to_be_bytes()); + trailing_ns.push(0); + trailing_ns.extend_from_slice(&QTYPE_NS.to_be_bytes()); + trailing_ns.extend_from_slice(&QCLASS_IN.to_be_bytes()); + trailing_ns.extend_from_slice(&60u32.to_be_bytes()); + trailing_ns.extend_from_slice(&5u16.to_be_bytes()); + trailing_ns.extend_from_slice(b"\x02ns\x00\x00"); + assert!(decode_dns_response(0x5052, ".", QTYPE_NS, &trailing_ns).is_err()); + + // candidate_query_names honors ndots ordering, trailing dot, and 255-octet search filtering. + let conf = ResolvConf { + servers: ResolvConf::default().servers, + search: vec![ + "svc.cluster.local".to_owned(), + vec!["x".repeat(60); 5].join("."), + ], + ndots: 2, + timeout: Duration::from_secs(5), + }; + assert_eq!( + candidate_query_names("api.ns", &conf), + vec!["api.ns.svc.cluster.local.", "api.ns."] + ); + assert_eq!( + candidate_query_names("a.b.c", &conf), + vec!["a.b.c.", "a.b.c.svc.cluster.local."] + ); + assert_eq!(candidate_query_names("rooted.", &conf), vec!["rooted."]); + } +} diff --git a/tests/discovery_dns.rs b/tests/discovery_dns.rs index 420f6c9..0d57906 100644 --- a/tests/discovery_dns.rs +++ b/tests/discovery_dns.rs @@ -16,13 +16,15 @@ #![cfg(feature = "discovery")] use std::collections::BTreeMap; -use std::net::{Ipv4Addr, Ipv6Addr}; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::time::Duration; use expfmt::discovery::{ - DiscoveryConfigs, DiscoveryError, DnsDiscovery, DnsRecord, DnsRecordType, DnsSdConfig, - LabelSet, ResolvConf, TargetGroup, + ADDRESS_LABEL, DiscoveryConfigs, DiscoveryError, DnsDiscovery, DnsRecord, DnsRecordType, + DnsSdConfig, LabelSet, ResolvConf, TargetGroup, }; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, UdpSocket}; fn label_set(pairs: &[(&str, &str)]) -> LabelSet { pairs @@ -306,3 +308,276 @@ fn dns_discovery_records_match_prometheus_dns_test() { assert!(ResolvConf::from_file(std::path::Path::new("/nonexistent/resolv.conf")).is_err()); } + +fn decode_qname_from_query(qbuf: &[u8]) -> (String, usize) { + let mut offset = 12; + let mut labels = Vec::new(); + while qbuf[offset] != 0 { + let len = qbuf[offset] as usize; + let label = std::str::from_utf8(&qbuf[offset + 1..offset + 1 + len]).unwrap(); + labels.push(label.to_owned()); + offset += 1 + len; + } + let qend = offset + 1 + 4; + let fqdn = if labels.is_empty() { + ".".to_owned() + } else { + format!("{}.", labels.join(".")) + }; + (fqdn, qend) +} + +#[tokio::test] +async fn dns_discovery_live_udp_and_tcp_fallback_with_search_path() { + let parsed_resolv = ResolvConf::parse( + r#" + # Comment line + ; Semicolon comment + nameserver 10.96.0.10 + nameserver [::1]:5353 + domain old.example.com. + search svc.cluster.local. cluster.local. + options ndots:3 timeout:4 + "#, + ); + assert_eq!( + parsed_resolv.servers, + vec![ + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 96, 0, 10)), 53), + SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 5353), + ] + ); + assert_eq!( + parsed_resolv.search, + vec!["svc.cluster.local".to_owned(), "cluster.local".to_owned()] + ); + assert_eq!(parsed_resolv.ndots, 3); + assert_eq!(parsed_resolv.timeout, Duration::from_secs(4)); + + let udp_listener = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let server_addr = udp_listener.local_addr().unwrap(); + let tcp_listener = TcpListener::bind(server_addr).await.unwrap(); + + // Spawn a local DNS server that returns TC=1 (truncated) over UDP so the client falls back + // to TCP, and returns an SRV record over TCP for `_metrics._tcp.svc.cluster.local.`. + tokio::spawn(async move { + let mut buf = [0u8; 512]; + if let Ok((len, peer)) = udp_listener.recv_from(&mut buf).await { + assert!(len >= 12); + let mut resp = Vec::with_capacity(12); + resp.extend_from_slice(&buf[0..2]); + resp.extend_from_slice(&0x8200u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + let _ = udp_listener.send_to(&resp, peer).await; + } + }); + + tokio::spawn(async move { + if let Ok((mut stream, _)) = tcp_listener.accept().await { + let mut len_buf = [0u8; 2]; + stream.read_exact(&mut len_buf).await.unwrap(); + let qlen = u16::from_be_bytes(len_buf) as usize; + let mut qbuf = vec![0u8; qlen]; + stream.read_exact(&mut qbuf).await.unwrap(); + + let (_, qname_end) = decode_qname_from_query(&qbuf); + let mut resp = Vec::new(); + resp.extend_from_slice(&qbuf[0..2]); + resp.extend_from_slice(&0x8180u16.to_be_bytes()); + resp.extend_from_slice(&1u16.to_be_bytes()); + resp.extend_from_slice(&1u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&qbuf[12..qname_end]); + + resp.push(0); + resp.extend_from_slice(&33u16.to_be_bytes()); // SRV + resp.extend_from_slice(&1u16.to_be_bytes()); // IN + resp.extend_from_slice(&60u32.to_be_bytes()); + + let mut rdata = Vec::new(); + rdata.extend_from_slice(&10u16.to_be_bytes()); + rdata.extend_from_slice(&5u16.to_be_bytes()); + rdata.extend_from_slice(&9100u16.to_be_bytes()); + for label in ["pod1", "cluster", "local"] { + rdata.push(label.len() as u8); + rdata.extend_from_slice(label.as_bytes()); + } + rdata.push(0); + + resp.extend_from_slice(&(rdata.len() as u16).to_be_bytes()); + resp.extend_from_slice(&rdata); + + stream + .write_all(&(resp.len() as u16).to_be_bytes()) + .await + .unwrap(); + stream.write_all(&resp).await.unwrap(); + } + }); + + let resolv = ResolvConf { + servers: vec![server_addr], + search: vec!["svc.cluster.local".to_owned()], + ndots: 5, + timeout: Duration::from_secs(2), + }; + let sd = DnsDiscovery::with_resolv_conf( + DnsSdConfig { + names: vec!["_metrics._tcp".to_owned()], + refresh_interval: Duration::from_secs(15), + record_type: "SRV".to_owned(), + port: 0, + }, + resolv, + ) + .unwrap(); + + assert!(sd.resolv_conf().is_some()); + let groups = sd.refresh().await.unwrap(); + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].source, "_metrics._tcp"); + assert_eq!( + groups[0].targets, + vec![label_set(&[ + ("__address__", "pod1.cluster.local:9100"), + ("__meta_dns_name", "_metrics._tcp"), + ("__meta_dns_srv_record_target", "pod1.cluster.local."), + ("__meta_dns_srv_record_port", "9100"), + ("__meta_dns_mx_record_target", ""), + ("__meta_dns_ns_record_target", ""), + ])] + ); +} + +#[test] +fn dns_multi_server_fallback_search_path_and_localset() { + let clamped = ResolvConf::parse( + "domain first.example.com. ignored.example.com.\noptions ndots:99 timeout:999\n", + ); + assert_eq!(clamped.search, vec!["first.example.com".to_owned()]); + assert_eq!(clamped.ndots, 15); + assert_eq!(clamped.timeout, Duration::from_secs(30)); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let local = tokio::task::LocalSet::new(); + + local.block_on(&rt, async { + let udp1 = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let srv1_addr = udp1.local_addr().unwrap(); + let tcp1 = TcpListener::bind(srv1_addr).await.unwrap(); + + let udp2 = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let srv2_addr = udp2.local_addr().unwrap(); + + tokio::task::spawn_local(async move { + if tcp1.accept().await.is_ok() { + panic!("non-truncated UDP response must not trigger TCP fallback"); + } + }); + + let srv2_queries = std::rc::Rc::new(std::cell::RefCell::new(Vec::::new())); + let srv2_queries_clone = std::rc::Rc::clone(&srv2_queries); + + tokio::task::spawn_local(async move { + let mut buf = [0u8; 512]; + while let Ok((len, peer)) = udp1.recv_from(&mut buf).await { + let (qname, qend) = decode_qname_from_query(&buf[..len]); + let rcode: u16 = match qname.as_str() { + "hit.two.local." | "flaky.one.local." => 2, // SERVFAIL + "nodata.one.local." => 0, // NODATA (ANCOUNT=0) + _ => 3, // NXDOMAIN + }; + let mut resp = Vec::new(); + resp.extend_from_slice(&buf[0..2]); + resp.extend_from_slice(&(0x8180u16 | rcode).to_be_bytes()); + resp.extend_from_slice(&1u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&buf[12..qend]); + let _ = udp1.send_to(&resp, peer).await; + } + }); + + tokio::task::spawn_local(async move { + let mut buf = [0u8; 512]; + while let Ok((len, peer)) = udp2.recv_from(&mut buf).await { + let (qname, qend) = decode_qname_from_query(&buf[..len]); + srv2_queries_clone.borrow_mut().push(qname.clone()); + let mut resp = Vec::new(); + resp.extend_from_slice(&buf[0..2]); + if qname == "hit.two.local." { + resp.extend_from_slice(&0x8180u16.to_be_bytes()); + resp.extend_from_slice(&1u16.to_be_bytes()); + resp.extend_from_slice(&1u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&buf[12..qend]); + resp.push(0); + resp.extend_from_slice(&1u16.to_be_bytes()); // A + resp.extend_from_slice(&1u16.to_be_bytes()); // IN + resp.extend_from_slice(&60u32.to_be_bytes()); + resp.extend_from_slice(&4u16.to_be_bytes()); + resp.extend_from_slice(&[10, 1, 2, 3]); + } else { + resp.extend_from_slice(&0x8182u16.to_be_bytes()); // SERVFAIL + resp.extend_from_slice(&1u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&0u16.to_be_bytes()); + resp.extend_from_slice(&buf[12..qend]); + } + let _ = udp2.send_to(&resp, peer).await; + } + }); + + let resolv = ResolvConf { + servers: vec![srv1_addr, srv2_addr], + search: vec!["one.local".to_owned(), "two.local".to_owned()], + ndots: 2, + timeout: Duration::from_secs(2), + }; + let sd = DnsDiscovery::with_resolv_conf( + DnsSdConfig { + names: vec![ + "hit".to_owned(), + "gone".to_owned(), + "flaky".to_owned(), + "nodata".to_owned(), + ], + refresh_interval: Duration::from_secs(15), + record_type: "A".to_owned(), + port: 9100, + }, + resolv, + ) + .unwrap(); + + let groups = tokio::task::spawn_local(async move { sd.refresh().await }) + .await + .unwrap() + .unwrap(); + assert_eq!(groups.len(), 3); + assert_eq!(groups[0].source, "hit"); + assert_eq!( + groups[0].targets[0].get(ADDRESS_LABEL), + Some(&"10.1.2.3:9100".to_owned()) + ); + assert_eq!(groups[1].source, "gone"); + assert!(groups[1].targets.is_empty()); + assert_eq!(groups[2].source, "nodata"); + assert!(groups[2].targets.is_empty()); + + assert_eq!( + *srv2_queries.borrow(), + vec!["hit.two.local.".to_owned(), "flaky.one.local.".to_owned()] + ); + }); +}