Skip to content

Commit

Permalink
refactor(iroh-net): small improvements to dns code (#2301)
Browse files Browse the repository at this point in the history 
## Description

- Move helper methods to a trait extension to make them more ergonomic
to use.
- Avoid temporary allocations via `collect` in dns lookups.
- Avoid double iterations on lookup results in the relay.
- Remove unnecessary `TryParseIp` bound on lookup functions.
- Modify comments about ipv4/ipv6 lookup to clarify why we avoid the
lookup
  strategy that would naturally fit this use.

## Breaking Changes

- `iroh_net::dns::lookup_ipv4_ipv6` is removed in favor of
`iroh_net::dns::ResolverExt::lookup_ipv4_ipv6`

## Notes & open questions

Was `iroh_net::dns::lookup_ipv4_ipv6` meant to be part of the public
api? I don't really see why.

## Change checklist

- [x] Self-review.
- [x] Documentation updates if relevant.
- [ ] ~Tests if relevant.~
- [x] All breaking changes documented.
  • Loading branch information
@divagant-martian
divagant-martian authored May 16, 2024
1 parent 491012c commit b93dd34
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 91 deletions.
157 changes: 87 additions & 70 deletions iroh-net/src/dns.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@ use std::net::{IpAddr, Ipv6Addr};
use std::time::Duration;

use anyhow::Result;
use hickory_resolver::{AsyncResolver, IntoName, TokioAsyncResolver, TryParseIp};
use futures_lite::Future;
use hickory_resolver::{AsyncResolver, IntoName, TokioAsyncResolver};
use once_cell::sync::Lazy;

pub mod node_info;
Expand Down Expand Up @@ -67,83 +68,97 @@ fn create_default_resolver() -> Result<TokioAsyncResolver> {
}
}

// lookup IPv4 and IPv6 in parallel
// see [`ResolverExt::lookup_ipv4_ipv6`] for info on why we avoid `LookupIpStrategy::Ipv4AndIpv6`
options.ip_strategy = hickory_resolver::config::LookupIpStrategy::Ipv4thenIpv6;

let resolver = AsyncResolver::tokio(config, options);
Ok(resolver)
}

pub(crate) async fn lookup_ipv4<N: IntoName + TryParseIp + Clone>(
resolver: &DnsResolver,
host: N,
timeout: Duration,
) -> Result<Vec<IpAddr>> {
let addrs = tokio::time::timeout(timeout, resolver.ipv4_lookup(host)).await??;
Ok(addrs.into_iter().map(|ip| IpAddr::V4(ip.0)).collect())
/// Extension trait to [`DnsResolver`].
pub trait ResolverExt {
/// Perform an ipv4 lookup with a timeout.
fn lookup_ipv4<N: IntoName + Clone>(
&self,
host: N,
timeout: Duration,
) -> impl Future<Output = Result<impl Iterator<Item = IpAddr>>>;

/// Perform an ipv6 lookup with a timeout.
fn lookup_ipv6<N: IntoName + Clone>(
&self,
host: N,
timeout: Duration,
) -> impl Future<Output = Result<impl Iterator<Item = IpAddr>>>;

/// Race an ipv4 and ipv6 lookup with a timeout.
fn lookup_ipv4_ipv6<N: IntoName + Clone>(
&self,
host: N,
timeout: Duration,
) -> impl Future<Output = Result<impl Iterator<Item = IpAddr>>>;
}

pub(crate) async fn lookup_ipv6<N: IntoName + TryParseIp + Clone>(
resolver: &DnsResolver,
host: N,
timeout: Duration,
) -> Result<Vec<IpAddr>> {
let addrs = tokio::time::timeout(timeout, resolver.ipv6_lookup(host)).await??;
Ok(addrs.into_iter().map(|ip| IpAddr::V6(ip.0)).collect())
}
impl ResolverExt for DnsResolver {
async fn lookup_ipv4<N: IntoName + Clone>(
&self,
host: N,
timeout: Duration,
) -> Result<impl Iterator<Item = IpAddr>> {
let addrs = tokio::time::timeout(timeout, self.ipv4_lookup(host)).await??;
Ok(addrs.into_iter().map(|ip| IpAddr::V4(ip.0)))
}

/// Resolve IPv4 and IPv6 in parallel.
///
/// `LookupIpStrategy::Ipv4AndIpv6` will wait for ipv6 resolution timeout, even if it is
/// not usable on the stack, so we manually query both lookups concurrently and time them out
/// individually.
pub async fn lookup_ipv4_ipv6<N: IntoName + TryParseIp + Clone>(
resolver: &DnsResolver,
host: N,
timeout: Duration,
) -> Result<Vec<IpAddr>> {
let ipv4 = resolver.ipv4_lookup(host.clone());
let ipv6 = resolver.ipv6_lookup(host);
let ipv4 = tokio::time::timeout(timeout, ipv4);
let ipv6 = tokio::time::timeout(timeout, ipv6);

let res = tokio::join!(ipv4, ipv6);
match res {
(Ok(Ok(ipv4)), Ok(Ok(ipv6))) => {
let res = ipv4
.into_iter()
.map(|ip| IpAddr::V4(ip.0))
.chain(ipv6.into_iter().map(|ip| IpAddr::V6(ip.0)))
.collect();
Ok(res)
}
(Ok(Ok(ipv4)), Err(_timeout)) => {
let res = ipv4.into_iter().map(|ip| IpAddr::V4(ip.0)).collect();
Ok(res)
}
(Ok(Ok(ipv4)), Ok(Err(_err))) => {
let res = ipv4.into_iter().map(|ip| IpAddr::V4(ip.0)).collect();
Ok(res)
}
(Ok(Err(_err)), Ok(Ok(ipv6))) => {
let res = ipv6.into_iter().map(|ip| IpAddr::V6(ip.0)).collect();
Ok(res)
}
(Ok(Err(err1)), Ok(Err(err2))) => {
anyhow::bail!("Ipv4: {:?}, Ipv6: {:?}", err1, err2);
}
(Ok(Err(err1)), Err(err2)) => {
anyhow::bail!("Ipv4: {:?}, Ipv6: {:?}", err1, err2);
}
(Err(_timeout), Ok(Ok(ipv6))) => {
let res = ipv6.into_iter().map(|ip| IpAddr::V6(ip.0)).collect();
Ok(res)
}
(Err(err1), Ok(Err(err2))) => {
anyhow::bail!("Ipv4: {:?}, Ipv6: {:?}", err1, err2);
async fn lookup_ipv6<N: IntoName + Clone>(
&self,
host: N,
timeout: Duration,
) -> Result<impl Iterator<Item = IpAddr>> {
let addrs = tokio::time::timeout(timeout, self.ipv6_lookup(host)).await??;
Ok(addrs.into_iter().map(|ip| IpAddr::V6(ip.0)))
}

/// Resolve IPv4 and IPv6 in parallel.
///
/// `LookupIpStrategy::Ipv4AndIpv6` will wait for ipv6 resolution timeout, even if it is
/// not usable on the stack, so we manually query both lookups concurrently and time them out
/// individually.
async fn lookup_ipv4_ipv6<N: IntoName + Clone>(
&self,
host: N,
timeout: Duration,
) -> Result<impl Iterator<Item = IpAddr>> {
let res = tokio::join!(
self.lookup_ipv4(host.clone(), timeout),
self.lookup_ipv6(host, timeout)
);

match res {
(Ok(ipv4), Ok(ipv6)) => Ok(LookupIter::Both(ipv4.chain(ipv6))),
(Ok(ipv4), Err(_)) => Ok(LookupIter::Ipv4(ipv4)),
(Err(_), Ok(ipv6)) => Ok(LookupIter::Ipv6(ipv6)),
(Err(ipv4_err), Err(ipv6_err)) => {
anyhow::bail!("Ipv4: {:?}, Ipv6: {:?}", ipv4_err, ipv6_err)
}
}
(Err(timeout1), Err(timeout2)) => {
anyhow::bail!("Ipv4: {:?}, Ipv6: {:?}", timeout1, timeout2);
}
}

/// Helper enum to give a unified type to the iterators of [`ResolverExt::lookup_ipv4_ipv6`]
enum LookupIter<A, B> {
Ipv4(A),
Ipv6(B),
Both(std::iter::Chain<A, B>),
}

impl<A: Iterator<Item = IpAddr>, B: Iterator<Item = IpAddr>> Iterator for LookupIter<A, B> {
type Item = IpAddr;

fn next(&mut self) -> Option<Self::Item> {
match self {
LookupIter::Ipv4(iter) => iter.next(),
LookupIter::Ipv6(iter) => iter.next(),
LookupIter::Both(iter) => iter.next(),
}
}
}
Expand All @@ -170,9 +185,11 @@ pub(crate) mod tests {
async fn test_dns_lookup_ipv4_ipv6() {
let _logging = iroh_test::logging::setup();
let resolver = default_resolver();
let res = lookup_ipv4_ipv6(resolver, NA_RELAY_HOSTNAME, Duration::from_secs(5))
let res: Vec<_> = resolver
.lookup_ipv4_ipv6(NA_RELAY_HOSTNAME, Duration::from_secs(5))
.await
.unwrap();
.unwrap()
.collect();
assert!(!res.is_empty());
dbg!(res);
}
Expand Down
18 changes: 9 additions & 9 deletions iroh-net/src/netcheck/reportgen.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ use tracing::{debug, debug_span, error, info_span, trace, warn, Instrument, Span

use super::NetcheckMetrics;
use crate::defaults::DEFAULT_RELAY_STUN_PORT;
use crate::dns::{lookup_ipv4, lookup_ipv6, DnsResolver};
use crate::dns::{DnsResolver, ResolverExt};
use crate::net::interfaces;
use crate::net::ip;
use crate::net::UdpSocket;
Expand Down Expand Up @@ -945,10 +945,10 @@ async fn get_relay_addr(
ProbeProto::StunIpv4 | ProbeProto::IcmpV4 => match relay_node.url.host() {
Some(url::Host::Domain(hostname)) => {
debug!(?proto, %hostname, "Performing DNS A lookup for relay addr");
match lookup_ipv4(dns_resolver, hostname, DNS_TIMEOUT).await {
Ok(addrs) => addrs
.first()
.map(|addr| ip::to_canonical(*addr))
match dns_resolver.lookup_ipv4(hostname, DNS_TIMEOUT).await {
Ok(mut addrs) => addrs
.next()
.map(ip::to_canonical)
.map(|addr| SocketAddr::new(addr, port))
.ok_or(anyhow!("No suitable relay addr found")),
Err(err) => Err(err.context("No suitable relay addr found")),
Expand All @@ -962,10 +962,10 @@ async fn get_relay_addr(
ProbeProto::StunIpv6 | ProbeProto::IcmpV6 => match relay_node.url.host() {
Some(url::Host::Domain(hostname)) => {
debug!(?proto, %hostname, "Performing DNS AAAA lookup for relay addr");
match lookup_ipv6(dns_resolver, hostname, DNS_TIMEOUT).await {
Ok(addrs) => addrs
.first()
.map(|addr| ip::to_canonical(*addr))
match dns_resolver.lookup_ipv6(hostname, DNS_TIMEOUT).await {
Ok(mut addrs) => addrs
.next()
.map(ip::to_canonical)
.map(|addr| SocketAddr::new(addr, port))
.ok_or(anyhow!("No suitable relay addr found")),
Err(err) => Err(err.context("No suitable relay addr found")),
Expand Down
25 changes: 13 additions & 12 deletions iroh-net/src/relay/http/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use tokio::time::Instant;
use tracing::{debug, error, info_span, trace, warn, Instrument};
use url::Url;

use crate::dns::{lookup_ipv4_ipv6, DnsResolver};
use crate::dns::{DnsResolver, ResolverExt};
use crate::key::{PublicKey, SecretKey};
use crate::relay::RelayUrl;
use crate::relay::{
Expand Down Expand Up @@ -862,19 +862,20 @@ async fn resolve_host(
match host {
url::Host::Domain(domain) => {
// Need to do a DNS lookup
let addrs = lookup_ipv4_ipv6(resolver, domain, DNS_TIMEOUT)
let mut addrs = resolver
.lookup_ipv4_ipv6(domain, DNS_TIMEOUT)
.await
.map_err(|e| ClientError::Dns(Some(e)))?;
.map_err(|e| ClientError::Dns(Some(e)))?
.peekable();

if prefer_ipv6 {
if let Some(addr) = addrs.iter().find(|addr| addr.is_ipv6()) {
return Ok(*addr);
}
}
addrs
.into_iter()
.next()
.ok_or_else(|| ClientError::Dns(None))
let found = if prefer_ipv6 {
let first = addrs.peek().copied();
addrs.find(IpAddr::is_ipv6).or(first)
} else {
addrs.next()
};

found.ok_or_else(|| ClientError::Dns(None))
}
url::Host::Ipv4(ip) => Ok(IpAddr::V4(ip)),
url::Host::Ipv6(ip) => Ok(IpAddr::V6(ip)),
Expand Down

0 comments on commit b93dd34

Please sign in to comment.