fix resolver

fix async-std double-close fd crash
This commit is contained in:
Christien Rioux 2024-08-10 17:28:00 -07:00
parent c7d698f211
commit c812534eb6
2 changed files with 100 additions and 55 deletions

View File

@ -17,62 +17,67 @@ cfg_if! {
if #[cfg(not(target_os = "windows"))] { if #[cfg(not(target_os = "windows"))] {
cfg_if! { cfg_if! {
if #[cfg(feature="rt-async-std")] { if #[cfg(feature="rt-async-std")] {
use async_std_resolver::{config, resolver, resolver_from_system_conf, AsyncStdResolver as AsyncResolver}; use async_std_resolver::{config, resolver, AsyncStdResolver as AsyncResolver};
use hickory_resolver::error::ResolveErrorKind; use hickory_resolver::system_conf::read_system_conf;
} else if #[cfg(feature="rt-tokio")] { } else if #[cfg(feature="rt-tokio")] {
use hickory_resolver::{config, TokioAsyncResolver as AsyncResolver, error::ResolveError, error::ResolveErrorKind}; use hickory_resolver::{config, TokioAsyncResolver as AsyncResolver, system_conf::read_system_conf};
pub async fn resolver( async fn resolver(
config: config::ResolverConfig, config: config::ResolverConfig,
options: config::ResolverOpts, options: config::ResolverOpts,
) -> AsyncResolver { ) -> AsyncResolver {
AsyncResolver::tokio(config, options) AsyncResolver::tokio(config, options)
} }
/// Constructs a new async-std based Resolver with the system configuration.
///
/// This will use `/etc/resolv.conf` on Unix OSes and the registry on Windows.
#[cfg(any(unix, target_os = "windows"))]
pub async fn resolver_from_system_conf() -> Result<AsyncResolver, ResolveError> {
AsyncResolver::tokio_from_system_conf()
}
} else { } else {
compile_error!("needs executor implementation") compile_error!("needs executor implementation")
} }
} }
struct Resolvers {
system: Option<Arc<AsyncResolver>>,
default: Arc<AsyncResolver>,
}
lazy_static::lazy_static! { lazy_static::lazy_static! {
static ref RESOLVER: Arc<AsyncMutex<Option<AsyncResolver>>> = Arc::new(AsyncMutex::new(None)); static ref RESOLVERS: AsyncMutex<Option<Arc<Resolvers>>> = AsyncMutex::new(None);
} }
} }
} }
cfg_if! { cfg_if! {
if #[cfg(not(target_os = "windows"))] { if #[cfg(not(target_os = "windows"))] {
async fn get_resolver() -> EyreResult<AsyncResolver> {
let mut resolver_lock = RESOLVER.lock().await; async fn with_resolvers<R, F: FnOnce(Arc<Resolvers>) -> SendPinBoxFuture<R>>(closure: F) -> R {
if let Some(r) = &*resolver_lock { let mut resolvers_lock = RESOLVERS.lock().await;
Ok(r.clone()) if let Some(r) = &*resolvers_lock {
return closure(r.clone()).await;
}
let (config, mut options) = (config::ResolverConfig::default(), config::ResolverOpts::default());
options.try_tcp_on_error = true;
let default = Arc::new(resolver(config, options).await);
let system = if let Ok((config, options)) = read_system_conf() {
Some(Arc::new(resolver(config, options).await))
} else { } else {
let resolver = match resolver_from_system_conf().await { None
Ok(v) => v,
Err(_) => resolver(
config::ResolverConfig::default(),
config::ResolverOpts::default(),
)
.await
}; };
let resolvers = Arc::new(Resolvers {
*resolver_lock = Some(resolver.clone()); system, default
Ok(resolver) });
} *resolvers_lock = Some(resolvers.clone());
closure(resolvers).await
} }
async fn reset_resolver() { // async fn reset_resolver(use_default: bool) {
let mut resolver_lock = RESOLVER.lock().await; // let mut resolver_lock = if use_default {
*resolver_lock = None; // DEFAULT_RESOLVER.lock().await
} // } else {
// SYSTEM_RESOLVER.lock().await
// };
// *resolver_lock = None;
// }
} }
} }
@ -112,18 +117,32 @@ pub async fn txt_lookup<S: AsRef<str>>(host: S) -> EyreResult<Vec<String>> {
Ok(out) Ok(out)
} else { } else {
let resolver = get_resolver().await?; let host = host.as_ref().to_string();
let txt_result = match resolver let txt_result = with_resolvers(|resolvers| Box::pin(async move {
.txt_lookup(host.as_ref()) // Try the default resolver config
match resolvers.default
.txt_lookup(&host)
.await { .await {
Ok(v) => v, Ok(v) => Ok(v),
Err(e) => { Err(e) => {
if !matches!(e.kind(), ResolveErrorKind::NoRecordsFound { query:_, soa:_, negative_ttl:_, response_code:_, trusted:_ }) { // Try the system resolver config if we have it
reset_resolver().await; if let Some(system_resolver) = &resolvers.system {
debug!("default resolver txt_lookup error: {}", e);
match system_resolver
.txt_lookup(&host)
.await {
Ok(v) => Ok(v),
Err(e) => {
bail!("system resolver txt_lookup error: {}", e);
} }
bail!("txt_lookup error: {}", e);
} }
}; } else {
bail!("default resolver txt_lookup error: {}", e);
}
}
}
})).await?;
let mut out = Vec::new(); let mut out = Vec::new();
for x in txt_result.iter() { for x in txt_result.iter() {
for s in x.txt_data() { for s in x.txt_data() {
@ -179,11 +198,31 @@ pub async fn ptr_lookup(ip_addr: IpAddr) -> EyreResult<String> {
} }
bail!("No records returned"); bail!("No records returned");
} else { } else {
let resolver = get_resolver().await?; let ptr_result = with_resolvers(|resolvers| Box::pin(async move {
let ptr_result = resolver // Try the default resolver config
match resolvers.default
.reverse_lookup(ip_addr) .reverse_lookup(ip_addr)
.await .await {
.wrap_err("resolver error")?; Ok(v) => Ok(v),
Err(e) => {
// Try the system resolver config if we have it
if let Some(system_resolver) = &resolvers.system {
debug!("default resolver ptr_lookup error: {}", e);
match system_resolver
.reverse_lookup(ip_addr)
.await {
Ok(v) => Ok(v),
Err(e) => {
bail!("system resolver ptr_lookup error: {}", e);
}
}
} else {
bail!("default resolver ptr_lookup error: {}", e);
}
}
}
})).await?;
if let Some(r) = ptr_result.iter().next() { if let Some(r) = ptr_result.iter().next() {
Ok(r.to_string().trim_end_matches('.').to_string()) Ok(r.to_string().trim_end_matches('.').to_string())
} else { } else {

View File

@ -138,10 +138,13 @@ impl Network {
#[cfg(all(feature = "rt-async-std", unix))] #[cfg(all(feature = "rt-async-std", unix))]
{ {
// async-std does not directly support linger on TcpStream yet // async-std does not directly support linger on TcpStream yet
use std::os::fd::{AsRawFd, FromRawFd}; use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd};
if let Err(e) = unsafe { socket2::Socket::from_raw_fd(tcp_stream.as_raw_fd()) } if let Err(e) = unsafe {
.set_linger(Some(core::time::Duration::from_secs(0))) let s = socket2::Socket::from_raw_fd(tcp_stream.as_raw_fd());
{ let res = s.set_linger(Some(core::time::Duration::from_secs(0)));
s.into_raw_fd();
res
} {
log_net!(debug "Couldn't set TCP linger: {}", e); log_net!(debug "Couldn't set TCP linger: {}", e);
return; return;
} }
@ -149,10 +152,13 @@ impl Network {
#[cfg(all(feature = "rt-async-std", windows))] #[cfg(all(feature = "rt-async-std", windows))]
{ {
// async-std does not directly support linger on TcpStream yet // async-std does not directly support linger on TcpStream yet
use std::os::windows::io::{AsRawSocket, FromRawSocket}; use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket};
if let Err(e) = unsafe { socket2::Socket::from_raw_socket(tcp_stream.as_raw_socket()) } if let Err(e) = unsafe {
.set_linger(Some(core::time::Duration::from_secs(0))) let s = socket2::Socket::from_raw_socket(tcp_stream.as_raw_socket());
{ let res = s.set_linger(Some(core::time::Duration::from_secs(0)));
s.into_raw_socket();
res
} {
log_net!(debug "Couldn't set TCP linger: {}", e); log_net!(debug "Couldn't set TCP linger: {}", e);
return; return;
} }