Skip to main content

max / makenotwork

7.0 KB · 228 lines History Blame Raw
1 //! DNS record verification, resolves hostnames and compares against expected values.
2
3 use std::collections::HashSet;
4
5 use hickory_resolver::proto::rr::RData;
6 use hickory_resolver::{ResolverBuilder, TokioResolver};
7 use tracing::instrument;
8
9 use crate::config::DnsRecord;
10 use crate::types::{DnsCheckResult, DnsRecordType};
11
12 /// Resolve DNS records and compare against expected values.
13 /// Returns one `DnsCheckResult` per `DnsRecord` in the input.
14 #[instrument(skip_all)]
15 pub async fn check_dns(target: &str, records: &[DnsRecord]) -> Vec<DnsCheckResult> {
16 let resolver = match TokioResolver::builder_tokio().and_then(ResolverBuilder::build) {
17 Ok(resolver) => resolver,
18 Err(e) => {
19 return records
20 .iter()
21 .map(|r| DnsCheckResult {
22 target: target.to_string(),
23 name: r.name.clone(),
24 record_type: r.record_type,
25 expected: r.expected.clone(),
26 actual: vec![],
27 matches: false,
28 checked_at: chrono::Utc::now().to_rfc3339(),
29 error: Some(format!("failed to create resolver: {e}")),
30 })
31 .collect();
32 }
33 };
34
35 let mut results = Vec::with_capacity(records.len());
36 for record in records {
37 let result = resolve_record(target, &resolver, record).await;
38 results.push(result);
39 }
40 results
41 }
42
43 async fn resolve_record(
44 target: &str,
45 resolver: &TokioResolver,
46 record: &DnsRecord,
47 ) -> DnsCheckResult {
48 let now = chrono::Utc::now().to_rfc3339();
49
50 let actual = match record.record_type {
51 DnsRecordType::A => resolve_a(resolver, &record.name).await,
52 DnsRecordType::Aaaa => resolve_aaaa(resolver, &record.name).await,
53 DnsRecordType::Cname => resolve_cname(resolver, &record.name).await,
54 DnsRecordType::Mx => resolve_mx(resolver, &record.name).await,
55 DnsRecordType::Txt => resolve_txt(resolver, &record.name).await,
56 };
57
58 match actual {
59 Ok(actual_values) => {
60 let matches = check_match(&record.expected, &actual_values);
61 DnsCheckResult {
62 target: target.to_string(),
63 name: record.name.clone(),
64 record_type: record.record_type,
65 expected: record.expected.clone(),
66 actual: actual_values,
67 matches,
68 checked_at: now,
69 error: None,
70 }
71 }
72 Err(e) => DnsCheckResult {
73 target: target.to_string(),
74 name: record.name.clone(),
75 record_type: record.record_type,
76 expected: record.expected.clone(),
77 actual: vec![],
78 matches: false,
79 checked_at: now,
80 error: Some(e),
81 },
82 }
83 }
84
85 /// Check if all expected values are found in actual (expected ⊆ actual).
86 pub fn check_match(expected: &[String], actual: &[String]) -> bool {
87 let actual_set: HashSet<&str> = actual.iter().map(std::string::String::as_str).collect();
88 expected.iter().all(|e| actual_set.contains(e.as_str()))
89 }
90
91 async fn resolve_a(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
92 let response = resolver
93 .ipv4_lookup(name)
94 .await
95 .map_err(|e| format!("A lookup failed for {name}: {e}"))?;
96 Ok(rdata(&response)
97 .filter_map(|d| match d {
98 RData::A(a) => Some(a.0.to_string()),
99 _ => None,
100 })
101 .collect())
102 }
103
104 async fn resolve_aaaa(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
105 let response = resolver
106 .ipv6_lookup(name)
107 .await
108 .map_err(|e| format!("AAAA lookup failed for {name}: {e}"))?;
109 Ok(rdata(&response)
110 .filter_map(|d| match d {
111 RData::AAAA(aaaa) => Some(aaaa.0.to_string()),
112 _ => None,
113 })
114 .collect())
115 }
116
117 async fn resolve_cname(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
118 let response = resolver
119 .lookup(name, hickory_resolver::proto::rr::RecordType::CNAME)
120 .await
121 .map_err(|e| format!("CNAME lookup failed for {name}: {e}"))?;
122 Ok(rdata(&response)
123 .filter_map(|d| match d {
124 RData::CNAME(c) => Some(trim_root(&c.0.to_string())),
125 _ => None,
126 })
127 .collect())
128 }
129
130 async fn resolve_mx(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
131 let response = resolver
132 .mx_lookup(name)
133 .await
134 .map_err(|e| format!("MX lookup failed for {name}: {e}"))?;
135 Ok(rdata(&response)
136 .filter_map(|d| match d {
137 RData::MX(mx) => Some(trim_root(&mx.exchange.to_string())),
138 _ => None,
139 })
140 .collect())
141 }
142
143 async fn resolve_txt(resolver: &TokioResolver, name: &str) -> Result<Vec<String>, String> {
144 let response = resolver
145 .txt_lookup(name)
146 .await
147 .map_err(|e| format!("TXT lookup failed for {name}: {e}"))?;
148 Ok(rdata(&response)
149 .filter_map(|d| match d {
150 RData::TXT(txt) => Some(txt.to_string()),
151 _ => None,
152 })
153 .collect())
154 }
155
156 /// Answer-section rdata for a lookup. Typed lookups return a plain `Lookup` as of
157 /// hickory 0.26, so every caller filters the answers down to the variant it asked for.
158 fn rdata(response: &hickory_resolver::lookup::Lookup) -> impl Iterator<Item = &RData> {
159 response.answers().iter().map(|record| &record.data)
160 }
161
162 fn trim_root(name: &str) -> String {
163 name.trim_end_matches('.').to_string()
164 }
165
166 #[cfg(test)]
167 mod tests {
168 use super::*;
169
170 #[test]
171 fn check_match_exact() {
172 assert!(check_match(
173 &["1.2.3.4".to_string()],
174 &["1.2.3.4".to_string()],
175 ));
176 }
177
178 #[test]
179 fn check_match_subset() {
180 assert!(check_match(
181 &["1.2.3.4".to_string()],
182 &["1.2.3.4".to_string(), "5.6.7.8".to_string()],
183 ));
184 }
185
186 #[test]
187 fn check_match_mismatch() {
188 assert!(!check_match(
189 &["1.2.3.4".to_string()],
190 &["5.6.7.8".to_string()],
191 ));
192 }
193
194 #[test]
195 fn check_match_empty_expected() {
196 assert!(check_match(&[], &["1.2.3.4".to_string()]));
197 }
198
199 #[test]
200 fn check_match_empty_actual() {
201 assert!(!check_match(&["1.2.3.4".to_string()], &[]));
202 }
203
204 #[test]
205 fn check_match_order_independent() {
206 assert!(check_match(
207 &["b".to_string(), "a".to_string()],
208 &["a".to_string(), "b".to_string(), "c".to_string()],
209 ));
210 }
211
212 #[test]
213 fn check_match_multiple_expected_all_present() {
214 assert!(check_match(
215 &["1.2.3.4".to_string(), "5.6.7.8".to_string()],
216 &["5.6.7.8".to_string(), "1.2.3.4".to_string()],
217 ));
218 }
219
220 #[test]
221 fn check_match_multiple_expected_one_missing() {
222 assert!(!check_match(
223 &["1.2.3.4".to_string(), "5.6.7.8".to_string()],
224 &["1.2.3.4".to_string()],
225 ));
226 }
227 }
228