Skip to content

Commit

Permalink
allow customize ResolverOpts
Browse files Browse the repository at this point in the history
  • Loading branch information
zonyitoo committed Aug 7, 2023
1 parent eab914a commit 0ed488f
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 17 deletions.
20 changes: 18 additions & 2 deletions crates/shadowsocks-service/src/dns/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

use log::trace;
use shadowsocks::{dns_resolver::DnsResolver, net::ConnectOpts};
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::config::ResolverOpts;

use crate::config::DnsConfig;

Expand All @@ -28,7 +30,14 @@ pub async fn build_dns_resolver(
};

if !force_system_builtin {
return match DnsResolver::trust_dns_system_resolver(dns_cache_size, connect_opts.clone()).await {
let mut opts_opt = None;
if let Some(dns_cache_size) = dns_cache_size {
let mut opts = ResolverOpts::default();
opts.cache_size = dns_cache_size;
opts_opt = Some(opts);
}

return match DnsResolver::trust_dns_system_resolver(opts_opt, connect_opts.clone()).await {
Ok(r) => Some(r),
Err(err) => {
warn!(
Expand All @@ -47,7 +56,14 @@ pub async fn build_dns_resolver(
}
#[cfg(feature = "trust-dns")]
DnsConfig::TrustDns(dns) => {
match DnsResolver::trust_dns_resolver(dns, dns_cache_size, connect_opts.clone()).await {
let mut opts_opt = None;
if let Some(dns_cache_size) = dns_cache_size {
let mut opts = ResolverOpts::default();
opts.cache_size = dns_cache_size;
opts_opt = Some(opts);
}

match DnsResolver::trust_dns_resolver(dns, opts_opt, connect_opts.clone()).await {
Ok(r) => Some(r),
Err(err) => {
use log::warn;
Expand Down
16 changes: 9 additions & 7 deletions crates/shadowsocks/src/dns_resolver/resolver.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@ use tokio::net::lookup_host;
use tokio::task::JoinHandle;
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::config::ResolverConfig;
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::config::ResolverOpts;

#[cfg(feature = "trust-dns")]
use crate::net::ConnectOpts;
Expand All @@ -39,7 +41,7 @@ pub trait DnsResolve {
pub struct TrustDnsSystemResolver {
resolver: ArcSwap<TrustDnsResolver>,
connect_opts: ConnectOpts,
dns_cache_size: Option<usize>,
opts: Option<ResolverOpts>,
}

/// Collections of DNS resolver
Expand Down Expand Up @@ -196,7 +198,7 @@ async fn trust_dns_notify_update_dns(resolver: Arc<TrustDnsSystemResolver>) -> n
// Update once for all those Modify events
time::sleep(Duration::from_secs(1)).await;

match create_resolver(None, resolver.dns_cache_size, resolver.connect_opts.clone()).await {
match create_resolver(None, resolver.opts, resolver.connect_opts.clone()).await {
Ok(r) => {
debug!("auto-reload {DNS_RESOLV_FILE_PATH}");

Expand Down Expand Up @@ -228,17 +230,17 @@ impl DnsResolver {
/// On *nix system, it will try to read configurations from `/etc/resolv.conf`.
#[cfg(feature = "trust-dns")]
pub async fn trust_dns_system_resolver(
dns_cache_size: Option<usize>,
opts: Option<ResolverOpts>,
connect_opts: ConnectOpts,
) -> io::Result<DnsResolver> {
use super::trust_dns_resolver::create_resolver;

let resolver = create_resolver(None, dns_cache_size, connect_opts.clone()).await?;
let resolver = create_resolver(None, opts, connect_opts.clone()).await?;

let inner = Arc::new(TrustDnsSystemResolver {
resolver: ArcSwap::from(Arc::new(resolver)),
connect_opts,
dns_cache_size,
opts,
});

cfg_if! {
Expand All @@ -263,12 +265,12 @@ impl DnsResolver {
#[cfg(feature = "trust-dns")]
pub async fn trust_dns_resolver(
dns: ResolverConfig,
dns_cache_size: Option<usize>,
opts: Option<ResolverOpts>,
connect_opts: ConnectOpts,
) -> io::Result<DnsResolver> {
use super::trust_dns_resolver::create_resolver;
Ok(DnsResolver::TrustDns(
create_resolver(Some(dns), dns_cache_size, connect_opts).await?,
create_resolver(Some(dns), opts, connect_opts).await?,
))
}

Expand Down
11 changes: 3 additions & 8 deletions crates/shadowsocks/src/dns_resolver/trust_dns_resolver.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,7 @@ use trust_dns_resolver::{
error::ResolveResult,
name_server::{GenericConnector, RuntimeProvider},
proto::{iocompat::AsyncIoTokioAsStd, udp::DnsUdpSocket, TokioTime},
AsyncResolver,
TokioHandle,
AsyncResolver, TokioHandle,
};

use crate::net::{tcp::TcpStream as ShadowTcpStream, udp::UdpSocket as ShadowUdpSocket, ConnectOpts};
Expand Down Expand Up @@ -101,21 +100,17 @@ pub type DnsResolver = AsyncResolver<ShadowDnsConnectionProvider>;
/// Create a `trust-dns` asynchronous DNS resolver
pub async fn create_resolver(
dns: Option<ResolverConfig>,
dns_cache_size: Option<usize>,
opts: Option<ResolverOpts>,
connect_opts: ConnectOpts,
) -> ResolveResult<DnsResolver> {
// Customized dns resolution
match dns {
Some(conf) => {
let mut resolver_opts = ResolverOpts::default();
let mut resolver_opts = opts.unwrap_or_default();
// Use Ipv4AndIpv6 strategy. Because Ipv4ThenIpv6 or Ipv6ThenIpv4 will return if the first query returned.
// Since we want to use Happy Eyeballs to connect to both IPv4 and IPv6 addresses, we need both A and AAAA records.
resolver_opts.ip_strategy = LookupIpStrategy::Ipv4AndIpv6;

if let Some(size) = dns_cache_size {
resolver_opts.cache_size = size
}

trace!(
"initializing DNS resolver with config {:?} opts {:?}",
conf,
Expand Down

0 comments on commit 0ed488f

Please sign in to comment.