use iroh_base::EndpointId;
use iroh_relay::dns::DnsResolver;
pub use iroh_relay::dns::{N0_DNS_ENDPOINT_ORIGIN_PROD, N0_DNS_ENDPOINT_ORIGIN_STAGING};
use n0_future::boxed::BoxStream;
use crate::{
Endpoint,
address_lookup::{
AddressLookup, Error as AddressLookupError, IntoAddressLookup, IntoAddressLookupError,
Item as AddressLookupItem,
},
endpoint::force_staging_infra,
};
pub(crate) const DNS_STAGGERING_MS: &[u64] = &[200, 300];
#[derive(Debug)]
pub struct DnsAddressLookup {
origin_domain: String,
dns_resolver: DnsResolver,
}
#[derive(Debug)]
pub struct DnsAddressLookupBuilder {
origin_domain: String,
dns_resolver: Option<DnsResolver>,
}
impl DnsAddressLookupBuilder {
pub fn dns_resolver(mut self, dns_resolver: DnsResolver) -> Self {
self.dns_resolver = Some(dns_resolver);
self
}
pub fn build(self) -> DnsAddressLookup {
DnsAddressLookup {
dns_resolver: self.dns_resolver.unwrap_or_default(),
origin_domain: self.origin_domain,
}
}
}
impl DnsAddressLookup {
pub fn builder(origin_domain: String) -> DnsAddressLookupBuilder {
DnsAddressLookupBuilder {
origin_domain,
dns_resolver: None,
}
}
pub fn n0_dns() -> DnsAddressLookupBuilder {
if force_staging_infra() {
Self::builder(N0_DNS_ENDPOINT_ORIGIN_STAGING.to_string())
} else {
Self::builder(N0_DNS_ENDPOINT_ORIGIN_PROD.to_string())
}
}
}
impl IntoAddressLookup for DnsAddressLookupBuilder {
fn into_address_lookup(
mut self,
endpoint: &Endpoint,
) -> Result<impl AddressLookup, IntoAddressLookupError> {
if self.dns_resolver.is_none() {
self.dns_resolver = Some(endpoint.dns_resolver()?.clone());
}
Ok(self.build())
}
}
impl AddressLookup for DnsAddressLookup {
fn resolve(
&self,
endpoint_id: EndpointId,
) -> Option<BoxStream<Result<AddressLookupItem, AddressLookupError>>> {
let resolver = self.dns_resolver.clone();
let origin_domain = self.origin_domain.clone();
let fut = async move {
let endpoint_info = resolver
.lookup_endpoint_by_id_staggered(&endpoint_id, &origin_domain, DNS_STAGGERING_MS)
.await
.map_err(|e| AddressLookupError::from_err_any("dns", e))?;
Ok(AddressLookupItem::new(endpoint_info, "dns", None))
};
let stream = n0_future::stream::once_future(fut);
Some(Box::pin(stream))
}
}