use parking_lot::RwLock;
use std::{collections::HashMap, future::Future};
use tracing::trace;
pub use trust_dns_resolver::{error::ResolveError, TokioAsyncResolver};
use trust_dns_resolver::{name_server::ConnectionProvider, AsyncResolver};
pub trait Resolver: Send + Sync + Unpin + 'static {
fn lookup_txt(&self, query: &str) -> impl Future<Output = Option<String>> + Send;
}
impl<P: ConnectionProvider> Resolver for AsyncResolver<P> {
async fn lookup_txt(&self, query: &str) -> Option<String> {
let fqn = if query.ends_with('.') { query.to_string() } else { format!("{query}.") };
match self.txt_lookup(fqn).await {
Err(err) => {
trace!(target: "disc::dns", %err, ?query, "dns lookup failed");
None
}
Ok(lookup) => {
let txt = lookup.into_iter().next()?;
let entry = txt.iter().next()?;
String::from_utf8(entry.to_vec()).ok()
}
}
}
}
#[derive(Clone, Debug)]
pub struct DnsResolver(TokioAsyncResolver);
impl DnsResolver {
pub const fn new(resolver: TokioAsyncResolver) -> Self {
Self(resolver)
}
pub fn from_system_conf() -> Result<Self, ResolveError> {
TokioAsyncResolver::tokio_from_system_conf().map(Self::new)
}
}
impl Resolver for DnsResolver {
async fn lookup_txt(&self, query: &str) -> Option<String> {
Resolver::lookup_txt(&self.0, query).await
}
}
#[derive(Debug, Default)]
pub struct MapResolver(RwLock<HashMap<String, String>>);
impl MapResolver {
pub fn insert(&self, k: String, v: String) -> Option<String> {
self.0.write().insert(k, v)
}
pub fn get(&self, k: &str) -> Option<String> {
self.0.read().get(k).cloned()
}
pub fn remove(&self, k: &str) -> Option<String> {
self.0.write().remove(k)
}
}
impl Resolver for MapResolver {
async fn lookup_txt(&self, query: &str) -> Option<String> {
self.get(query)
}
}
#[cfg(test)]
pub(crate) struct TimeoutResolver(pub(crate) std::time::Duration);
#[cfg(test)]
impl Resolver for TimeoutResolver {
async fn lookup_txt(&self, _query: &str) -> Option<String> {
tokio::time::sleep(self.0).await;
None
}
}