#![warn(missing_docs, clippy::missing_errors_doc)]
use futures_util::future::join_all;
use hickory_resolver::{TokioResolver, net::NetError, proto::rr::RData};
use std::{net::SocketAddr, num::NonZeroU16};
use tellus::cluster::SeedDiscovery;
use thiserror::Error;
#[derive(Debug)]
pub struct DnsSeeds {
resolver: TokioResolver,
query: Query,
}
impl DnsSeeds {
pub fn new(query: Query) -> Result<Self, DnsSeedsError> {
Ok(Self {
resolver: system_resolver()?,
query,
})
}
pub fn srv(name: impl Into<String>) -> Result<Self, DnsSeedsError> {
Self::new(Query::Srv(name.into()))
}
pub fn ip(name: impl Into<String>, port: NonZeroU16) -> Result<Self, DnsSeedsError> {
Self::new(Query::Ip {
name: name.into(),
port,
})
}
pub fn with_resolver(mut self, resolver: TokioResolver) -> Self {
self.resolver = resolver;
self
}
}
impl SeedDiscovery for DnsSeeds {
type Error = DnsSeedsError;
async fn resolve(&mut self) -> Result<Vec<SocketAddr>, DnsSeedsError> {
match &self.query {
Query::Srv(name) => {
let records = self.resolver.srv_lookup(name.as_str()).await?;
let lookups = records
.answers()
.iter()
.filter_map(|record| match &record.data {
RData::SRV(srv) => Some(srv),
_ => None,
})
.map(|srv| {
let resolver = self.resolver.clone();
let target = srv.target.clone();
let port = srv.port;
async move { resolver.lookup_ip(target).await.map(|ips| (port, ips)) }
});
let mut addrs = Vec::new();
let mut first_error = None;
for result in join_all(lookups).await {
match result {
Ok((port, ips)) => {
addrs.extend(ips.iter().map(|ip| SocketAddr::new(ip, port)));
}
Err(error) => {
first_error.get_or_insert(error);
}
}
}
match first_error {
Some(error) if addrs.is_empty() => Err(error.into()),
_ => Ok(addrs),
}
}
Query::Ip { name, port } => {
let ips = self.resolver.lookup_ip(name.as_str()).await?;
let port = port.get();
Ok(ips.iter().map(|ip| SocketAddr::new(ip, port)).collect())
}
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
feature = "serde",
derive(serde::Deserialize),
serde(rename_all = "snake_case", deny_unknown_fields)
)]
pub enum Query {
Srv(String),
Ip {
name: String,
port: NonZeroU16,
},
}
#[derive(Debug, Error)]
#[error(transparent)]
pub struct DnsSeedsError(#[from] NetError);
fn system_resolver() -> Result<TokioResolver, DnsSeedsError> {
let mut builder = TokioResolver::builder_tokio()?;
builder.options_mut().cache_size = 0;
Ok(builder.build()?)
}
#[cfg(test)]
mod tests {
#[cfg(feature = "serde")]
use crate::Query;
use crate::system_resolver;
#[cfg(feature = "serde")]
use std::num::NonZeroU16;
#[test]
fn the_system_resolver_keeps_no_response_cache() {
let resolver = system_resolver().expect("the system resolver is built");
assert_eq!(resolver.options().cache_size, 0);
}
#[cfg(feature = "serde")]
#[test]
fn a_query_deserializes_from_its_documented_form() {
let query = serde_json::from_str::<Query>(r#"{ "srv": "_tellus._udp.example" }"#)
.expect("the SRV form deserializes");
assert_eq!(query, Query::Srv("_tellus._udp.example".to_string()));
let query =
serde_json::from_str::<Query>(r#"{ "ip": { "name": "nodes.example", "port": 7000 } }"#)
.expect("the A/AAAA form deserializes");
assert_eq!(
query,
Query::Ip {
name: "nodes.example".to_string(),
port: NonZeroU16::new(7_000).expect("7000 is not zero"),
}
);
assert!(
serde_json::from_str::<Query>(r#"{ "ip": { "name": "nodes.example", "port": 0 } }"#)
.is_err()
);
assert!(
serde_json::from_str::<Query>(r#"{ "ip": { "nmae": "nodes.example", "port": 7000 } }"#)
.is_err()
);
assert!(serde_json::from_str::<Query>(r#"{ "svr": "_tellus._udp.example" }"#).is_err());
}
}