Skip to content

Commit

Permalink
feat(ext/fetch): support custom DNS resolver
Browse files Browse the repository at this point in the history
  • Loading branch information
zino23 committed Jan 20, 2025
1 parent e4a16e9 commit 8efc9e8
Showing 1 changed file with 51 additions and 0 deletions.
51 changes: 51 additions & 0 deletions ext/fetch/dns.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use std::future::Future;
use std::io;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::task::Poll;
use std::task::{self};
use std::vec;
Expand All @@ -19,6 +20,16 @@ pub enum Resolver {
Gai(GaiResolver),
/// hickory-resolver's userspace resolver.
Hickory(hickory_resolver::Resolver<TokioConnectionProvider>),
/// A custom resolver that implements `Resolve`.
Custom(Arc<dyn Resolve>),
}

/// Alias for `Future` type returned by a custom DNS resolver.
pub type Resolving =
Pin<Box<dyn Future<Output = Result<SocketAddrs, io::Error>> + Send>>;

pub trait Resolve: Send + Sync + std::fmt::Debug {
fn resolve(&self, name: Name) -> Resolving;
}

impl Default for Resolver {
Expand Down Expand Up @@ -107,7 +118,47 @@ impl Service<Name> for Resolver {
Ok(iter)
})
}
Resolver::Custom(resolver) => {
let resolver = resolver.clone();
tokio::spawn(async move {
let result = resolver.resolve(name).await;

result
})
}
};
ResolveFut { inner: task }
}
}

#[cfg(test)]
mod tests {
use std::str::FromStr;

use super::*;

// A resolver that resolves any name into the same address.
#[derive(Debug)]
struct DebugResolver(SocketAddr);

impl Resolve for DebugResolver {
fn resolve(&self, _name: Name) -> Resolving {
let addr = self.0;
Box::pin(async move { Ok(vec![addr].into_iter()) })
}
}

#[tokio::test]
async fn custom_dns_resolver() {
let mut resolver = Resolver::Custom(Arc::new(DebugResolver(
"127.0.0.1:8080".parse().unwrap(),
)));
let mut addr = resolver
.call(Name::from_str("foo.com").unwrap())
.await
.unwrap();

let addr = addr.next().unwrap();
assert_eq!(addr, "127.0.0.1:8080".parse().unwrap());
}
}

0 comments on commit 8efc9e8

Please sign in to comment.