1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
use std::{net::SocketAddr, process::exit, sync::OnceLock};
use hickory_resolver::{
TokioResolver,
net::NetError,
proto::rr::{IntoName, Name},
};
use rand::Rng;
use tokio::net::lookup_host;
use crate::daemon::{config::NormalizedAddress, exitcode};
// We keep the resolver globally to avoid reloading its configuration constantly.
static RESOLVER: OnceLock<TokioResolver> = OnceLock::new();
pub(crate) struct KeResolutionResult {
pub(crate) addr: SocketAddr,
pub(crate) srv_record_name: Option<String>,
}
pub(crate) async fn resolve_ke(
addr: &NormalizedAddress,
) -> Result<impl Iterator<Item = KeResolutionResult>, std::io::Error> {
// Kludge allowing us to return two types of iterator.
enum Either<A, B> {
A(A),
B(B),
}
impl<A: Iterator<Item = KeResolutionResult>, B: Iterator<Item = KeResolutionResult>> Iterator
for Either<A, B>
{
type Item = KeResolutionResult;
fn next(&mut self) -> Option<Self::Item> {
match self {
Either::A(a) => a.next(),
Either::B(b) => b.next(),
}
}
}
// First try looking up SRV records
if let Ok(srv_names) = resolve_srv(format!("_ntske._tcp.{}", addr.server_name)).await {
let mut result = vec![];
for name in srv_names.into_iter().map(|v| v.to_ascii()) {
if let Ok(lookup) = lookup_host((name.as_str(), 4460)).await {
result.extend(lookup.map(|addr| KeResolutionResult {
addr,
srv_record_name: Some(name.clone()),
}));
}
}
if !result.is_empty() {
return Ok(Either::A(result.into_iter()));
}
}
// Otherwise do a direct name lookup
Ok(Either::B(
lookup_host((addr.server_name.as_str(), addr.port))
.await?
.map(|addr| KeResolutionResult {
addr,
srv_record_name: None,
}),
))
}
async fn resolve_srv<N: IntoName>(name: N) -> Result<Vec<Name>, NetError> {
let resolver = RESOLVER.get_or_init(|| {
let mut builder = match TokioResolver::builder_tokio() {
Ok(builder) => builder,
Err(e) => {
// Abort when the resolver configuration cannot be loaded
// trying anything else is madness when the system we run
// on is this broken.
tracing::error!("Could not load resolver configuration, aborting: {e}.");
exit(exitcode::CONFIG);
}
};
builder.options_mut().validate = true;
match builder.build() {
Ok(resolver) => resolver,
Err(e) => {
tracing::error!("Could not build resolver, aborting: {e}.");
exit(exitcode::CONFIG);
}
}
});
let lookup_result = resolver.srv_lookup(name).await?;
// Unfortunately, hickory doesn't order the results for us apropriately, so we need
// to do this ourselves. See also https://github.com/hickory-dns/hickory-dns/issues/3440
//
// For this, we generate a list of all valid results, augmented by a value equal to
// T^(1/w) where w is the weight of the entry, and T a uniform random variable
// between 0 and 1. Sorting by these values in increasing order gives a random order
// respecting the weighting, since independent uniform random X and Y both between 0
// and 1, the probability X^(1/n) > Y^(1/m) is m/(n+m), which is exactly the chance
// that the item with weight m should appear before the item with weight n. (Note,
// this can be checked by calculating the area under the implicit curve
// x=t^(1/n), y=t^(1/m) in the unit square)
let mut items: Vec<_> = lookup_result
.answers()
.iter()
.filter_map(|record| {
if !record.proof.is_secure() {
return None;
}
match &record.data {
hickory_resolver::proto::rr::RData::SRV(srv) => Some(srv),
_ => None,
}
})
.map(|v| {
(
if v.weight != 0 {
rand::thread_rng()
.r#gen::<f64>()
.powf(1.0 / (f64::from(v.weight)))
} else {
// Guarantee 0 weight items end up last within their priority group
2.0 + rand::thread_rng().r#gen::<f64>()
},
v,
)
})
.collect();
// Now all that remains to be done is sorting the items by first priority and then
// the generated random value, and we get an ordering respecting RFC2782.
items.sort_by(|a, b| {
a.1.priority
.cmp(&b.1.priority)
.then(f64::total_cmp(&a.0, &b.0))
});
Ok(items.into_iter().map(|v| &v.1.target).cloned().collect())
}