c-ares-resolver 12.0.0

use std::fmt;
use std::future::Future;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use crate::error::Error;
use crate::host::HostResults;
use crate::nameinfo::NameInfoResult;
use crate::resolver::{Options, Resolver};

#[cfg(cares1_29)]
use c_ares::ServerStateFlags;

/// The type of future returned by methods on the `FutureResolver`.
#[must_use]
pub struct CAresFuture<T> {
    inner: futures_channel::oneshot::Receiver<c_ares::Result<T>>,
    _resolver: Arc<Resolver>,
}

impl<T> fmt::Debug for CAresFuture<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("CAresFuture").finish_non_exhaustive()
    }
}

impl<T> CAresFuture<T> {
    fn new(
        promise: futures_channel::oneshot::Receiver<c_ares::Result<T>>,
        resolver: Arc<Resolver>,
    ) -> Self {
        Self {
            inner: promise,
            _resolver: resolver,
        }
    }

    fn pin_get_inner(
        self: Pin<&mut Self>,
    ) -> Pin<&mut futures_channel::oneshot::Receiver<c_ares::Result<T>>> {
        unsafe { self.map_unchecked_mut(|s| &mut s.inner) }
    }
}

impl<T> Future for CAresFuture<T> {
    type Output = c_ares::Result<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        self.pin_get_inner()
            .poll(cx)
            .map(|result| result.unwrap_or(Err(c_ares::Error::ECANCELLED)))
    }
}

/// An asynchronous DNS resolver, which returns results as `futures::Future`s.
///
/// Note that dropping the `FutureResolver` does *not* cause outstanding queries to fail - contrast
/// the `Resolver` - because the returned futures hold a reference to the underlying resolver.
///
/// The `c-ares` library returns results via callbacks, and some of those callbacks receive
/// borrowed data.  To return owned results, this resolver must clone or copy where the
/// callback-based `Resolver` need not.
///
/// Therefore: if you are trying very hard to avoid unnecessary allocations - prefer the
/// `Resolver`.
pub struct FutureResolver {
    inner: Arc<Resolver>,
}

impl fmt::Debug for FutureResolver {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("FutureResolver").finish_non_exhaustive()
    }
}

// Most query implementations follow the same pattern: call through to the `Resolver`, arranging
// that the callback completes a future.
macro_rules! futurize {
    ($resolver:expr, $query:ident, $($arg:expr),+ $(,)?) => {{
        let (sender, receiver) = futures_channel::oneshot::channel();
        $resolver.$query($($arg,)+ |result| {
            let _ = sender.send(result);
        });
        let resolver = Arc::clone(&$resolver);
        CAresFuture::new(receiver, resolver)
    }};
}

impl FutureResolver {
    /// Create a new `FutureResolver`, using default `Options`.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let resolver = c_ares_resolver::FutureResolver::new().unwrap();
    /// let future = resolver.query_a("example.com");
    /// // await or spawn the future...
    /// ```
    pub fn new() -> Result<Self, Error> {
        let options = Options::default();
        Self::with_options(options)
    }

    /// Create a new `FutureResolver`, with the given `Options`.
    pub fn with_options(options: Options) -> Result<Self, Error> {
        let inner = Resolver::with_options(options)?;
        let resolver = Self {
            inner: Arc::new(inner),
        };
        Ok(resolver)
    }

    /// Reinitialize a channel from system configuration.
    #[cfg(cares1_22)]
    pub fn reinit(&self) -> c_ares::Result<&Self> {
        self.inner.reinit()?;
        Ok(self)
    }

    /// Set the list of servers to contact, instead of the servers specified in resolv.conf or the
    /// local named.
    ///
    /// String format is `host[:port]`.  IPv6 addresses with ports require square brackets eg
    /// `[2001:4860:4860::8888]:53`.
    pub fn set_servers<I, S>(&self, servers: I) -> c_ares::Result<&Self>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        self.inner.set_servers(servers)?;
        Ok(self)
    }

    /// Retrieves the list of configured servers.
    ///
    /// Each entry is in `host[:port]` format, matching what [`set_servers`](Self::set_servers)
    /// accepts.
    #[cfg(cares1_24)]
    pub fn servers(&self) -> Vec<String> {
        self.inner.servers()
    }

    /// Set the local IPv4 address from which to make queries.
    pub fn set_local_ipv4(&self, ipv4: Ipv4Addr) -> &Self {
        self.inner.set_local_ipv4(ipv4);
        self
    }

    /// Set the local IPv6 address from which to make queries.
    pub fn set_local_ipv6(&self, ipv6: Ipv6Addr) -> &Self {
        self.inner.set_local_ipv6(ipv6);
        self
    }

    /// Set the local device from which to make queries.
    pub fn set_local_device(&self, device: &str) -> c_ares::Result<&Self> {
        self.inner.set_local_device(device)?;
        Ok(self)
    }

    /// Initializes an address sortlist configuration, so that addresses returned by
    /// `get_host_by_name()` are sorted according to the sortlist.
    ///
    /// Each element of the sortlist holds an IP-address/netmask pair. The netmask is optional but
    /// follows the address after a slash if present. For example: "130.155.160.0/255.255.240.0",
    /// or "130.155.0.0".
    pub fn set_sortlist<I, S>(&self, sortlist: I) -> c_ares::Result<&Self>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        self.inner.set_sortlist(sortlist)?;
        Ok(self)
    }

    /// Set a callback function to be invoked whenever a query on the channel completes.
    ///
    /// `callback(server, success, flags)` will be called when a query completes.
    ///
    /// - `server` indicates the DNS server that was used for the query.
    /// - `success` indicates whether the query succeeded or not.
    /// - `flags` is a bitmask of flags describing various aspects of the query.
    #[cfg(cares1_29)]
    pub fn set_server_state_callback<F>(&self, callback: F) -> &Self
    where
        F: Fn(&str, bool, ServerStateFlags) + Send + Sync + 'static,
    {
        self.inner.set_server_state_callback(callback);
        self
    }

    /// Look up the A records associated with `name`.
    pub fn query_a(&self, name: &str) -> CAresFuture<c_ares::AResults> {
        futurize!(self.inner, query_a, name)
    }

    /// Search for the A records associated with `name`.
    pub fn search_a(&self, name: &str) -> CAresFuture<c_ares::AResults> {
        futurize!(self.inner, search_a, name)
    }

    /// Look up the AAAA records associated with `name`.
    pub fn query_aaaa(&self, name: &str) -> CAresFuture<c_ares::AAAAResults> {
        futurize!(self.inner, query_aaaa, name)
    }

    /// Search for the AAAA records associated with `name`.
    pub fn search_aaaa(&self, name: &str) -> CAresFuture<c_ares::AAAAResults> {
        futurize!(self.inner, search_aaaa, name)
    }

    /// Look up the CAA records associated with `name`.
    pub fn query_caa(&self, name: &str) -> CAresFuture<c_ares::CAAResults> {
        futurize!(self.inner, query_caa, name)
    }

    /// Search for the CAA records associated with `name`.
    pub fn search_caa(&self, name: &str) -> CAresFuture<c_ares::CAAResults> {
        futurize!(self.inner, search_caa, name)
    }

    /// Look up the CNAME records associated with `name`.
    pub fn query_cname(&self, name: &str) -> CAresFuture<c_ares::CNameResults> {
        futurize!(self.inner, query_cname, name)
    }

    /// Search for the CNAME records associated with `name`.
    pub fn search_cname(&self, name: &str) -> CAresFuture<c_ares::CNameResults> {
        futurize!(self.inner, search_cname, name)
    }

    /// Look up the MX records associated with `name`.
    pub fn query_mx(&self, name: &str) -> CAresFuture<c_ares::MXResults> {
        futurize!(self.inner, query_mx, name)
    }

    /// Search for the MX records associated with `name`.
    pub fn search_mx(&self, name: &str) -> CAresFuture<c_ares::MXResults> {
        futurize!(self.inner, search_mx, name)
    }

    /// Look up the NAPTR records associated with `name`.
    pub fn query_naptr(&self, name: &str) -> CAresFuture<c_ares::NAPTRResults> {
        futurize!(self.inner, query_naptr, name)
    }

    /// Search for the NAPTR records associated with `name`.
    pub fn search_naptr(&self, name: &str) -> CAresFuture<c_ares::NAPTRResults> {
        futurize!(self.inner, search_naptr, name)
    }

    /// Look up the NS records associated with `name`.
    pub fn query_ns(&self, name: &str) -> CAresFuture<c_ares::NSResults> {
        futurize!(self.inner, query_ns, name)
    }

    /// Search for the NS records associated with `name`.
    pub fn search_ns(&self, name: &str) -> CAresFuture<c_ares::NSResults> {
        futurize!(self.inner, search_ns, name)
    }

    /// Look up the PTR records associated with `name`.
    pub fn query_ptr(&self, name: &str) -> CAresFuture<c_ares::PTRResults> {
        futurize!(self.inner, query_ptr, name)
    }

    /// Search for the PTR records associated with `name`.
    pub fn search_ptr(&self, name: &str) -> CAresFuture<c_ares::PTRResults> {
        futurize!(self.inner, search_ptr, name)
    }

    /// Look up the SOA records associated with `name`.
    pub fn query_soa(&self, name: &str) -> CAresFuture<c_ares::SOAResult> {
        futurize!(self.inner, query_soa, name)
    }

    /// Search for the SOA records associated with `name`.
    pub fn search_soa(&self, name: &str) -> CAresFuture<c_ares::SOAResult> {
        futurize!(self.inner, search_soa, name)
    }

    /// Look up the SRV records associated with `name`.
    pub fn query_srv(&self, name: &str) -> CAresFuture<c_ares::SRVResults> {
        futurize!(self.inner, query_srv, name)
    }

    /// Search for the SRV records associated with `name`.
    pub fn search_srv(&self, name: &str) -> CAresFuture<c_ares::SRVResults> {
        futurize!(self.inner, search_srv, name)
    }

    /// Look up the TXT records associated with `name`.
    pub fn query_txt(&self, name: &str) -> CAresFuture<c_ares::TXTResults> {
        futurize!(self.inner, query_txt, name)
    }

    /// Search for the TXT records associated with `name`.
    pub fn search_txt(&self, name: &str) -> CAresFuture<c_ares::TXTResults> {
        futurize!(self.inner, search_txt, name)
    }

    /// Look up the URI records associated with `name`.
    pub fn query_uri(&self, name: &str) -> CAresFuture<c_ares::URIResults> {
        futurize!(self.inner, query_uri, name)
    }

    /// Search for the URI records associated with `name`.
    pub fn search_uri(&self, name: &str) -> CAresFuture<c_ares::URIResults> {
        futurize!(self.inner, search_uri, name)
    }

    /// Perform a host query by address.
    pub fn get_host_by_address(&self, address: &IpAddr) -> CAresFuture<HostResults> {
        let (sender, receiver) = futures_channel::oneshot::channel();
        self.inner.get_host_by_address(address, |result| {
            let _ = sender.send(result.map(Into::into));
        });
        let resolver = Arc::clone(&self.inner);
        CAresFuture::new(receiver, resolver)
    }

    /// Perform a host query by name.
    pub fn get_host_by_name(
        &self,
        name: &str,
        family: c_ares::AddressFamily,
    ) -> CAresFuture<HostResults> {
        let (sender, receiver) = futures_channel::oneshot::channel();
        self.inner.get_host_by_name(name, family, |result| {
            let _ = sender.send(result.map(Into::into));
        });
        let resolver = Arc::clone(&self.inner);
        CAresFuture::new(receiver, resolver)
    }

    /// Address-to-nodename translation in protocol-independent manner.
    pub fn get_name_info(
        &self,
        address: &SocketAddr,
        flags: c_ares::NIFlags,
    ) -> CAresFuture<NameInfoResult> {
        let (sender, receiver) = futures_channel::oneshot::channel();
        self.inner.get_name_info(address, flags, |result| {
            let _ = sender.send(result.map(Into::into));
        });
        let resolver = Arc::clone(&self.inner);
        CAresFuture::new(receiver, resolver)
    }

    /// Initiate a host query by name and service.
    pub fn get_addrinfo(
        &self,
        name: &str,
        service: Option<&str>,
        hints: &c_ares::AddrInfoHints,
    ) -> CAresFuture<c_ares::AddrInfoResults> {
        futurize!(self.inner, get_addrinfo, name, service, hints)
    }

    /// Initiate a single-question DNS query for `name`.  The class and type of the query are per
    /// the provided parameters, taking values as defined in `arpa/nameser.h`.
    ///
    /// This method is provided so that users can query DNS types for which `c-ares` does not
    /// provide a parser; or in case a third-party parser is preferred.  Usually, if a suitable
    /// `query_xxx()` is available, that should be used.
    pub fn query(&self, name: &str, dns_class: u16, query_type: u16) -> CAresFuture<Vec<u8>> {
        let (sender, receiver) = futures_channel::oneshot::channel();
        self.inner.query(name, dns_class, query_type, |result| {
            let _ = sender.send(result.map(std::borrow::ToOwned::to_owned));
        });
        let resolver = Arc::clone(&self.inner);
        CAresFuture::new(receiver, resolver)
    }

    /// Initiate a series of single-question DNS queries for `name`.  The class and type of the
    /// query are per the provided parameters, taking values as defined in `arpa/nameser.h`.
    ///
    /// This method is provided so that users can search DNS types for which `c-ares` does not
    /// provide a parser; or in case a third-party parser is preferred.  Usually, if a suitable
    /// `search_xxx()` is available, that should be used.
    pub fn search(&self, name: &str, dns_class: u16, query_type: u16) -> CAresFuture<Vec<u8>> {
        let (sender, receiver) = futures_channel::oneshot::channel();
        self.inner.search(name, dns_class, query_type, |result| {
            let _ = sender.send(result.map(std::borrow::ToOwned::to_owned));
        });
        let resolver = Arc::clone(&self.inner);
        CAresFuture::new(receiver, resolver)
    }

    /// Send a DNS query using a pre-built [`c_ares::DnsRecord`].
    ///
    /// Returns a tuple of `(query_id, future)`.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use c_ares::*;
    /// use futures_executor::block_on;
    ///
    /// let resolver = c_ares_resolver::FutureResolver::new().unwrap();
    /// let mut query = DnsRecord::new(0, DnsFlags::RD, DnsOpcode::Query, DnsRcode::NoError).unwrap();
    /// query.query_add("example.com", DnsRecordType::A, DnsCls::IN).unwrap();
    /// let response = block_on(resolver.send_dnsrec(&query).unwrap()).unwrap();
    /// for rr in response.rrs(DnsSection::Answer) {
    ///     if let Some(addr) = rr.get_addr(DnsRrKey::A_ADDR) {
    ///         println!("address: {addr}");
    ///     }
    /// }
    /// ```
    #[cfg(cares1_28)]
    pub fn send_dnsrec(
        &self,
        dnsrec: &c_ares::DnsRecord,
    ) -> c_ares::Result<CAresFuture<c_ares::DnsRecord>> {
        let (sender, receiver) = futures_channel::oneshot::channel();
        self.inner.send_dnsrec(dnsrec, |result| {
            let _ = sender.send(result.and_then(c_ares::DnsRecord::try_clone));
        })?;
        let resolver = Arc::clone(&self.inner);
        Ok(CAresFuture::new(receiver, resolver))
    }

    /// Initiate a DNS query for `name` with the given class and type, receiving a parsed
    /// [`c_ares::DnsRecord`].
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use c_ares::{DnsCls, DnsRecordType, DnsRrKey, DnsSection};
    /// use futures_executor::block_on;
    ///
    /// let resolver = c_ares_resolver::FutureResolver::new().unwrap();
    /// let future = resolver
    ///     .query_dnsrec("example.com", DnsCls::IN, DnsRecordType::A)
    ///     .unwrap();
    /// let record = block_on(future).unwrap();
    /// for rr in record.rrs(DnsSection::Answer) {
    ///     if let Some(addr) = rr.get_addr(DnsRrKey::A_ADDR) {
    ///         println!("address: {addr}");
    ///     }
    /// }
    /// ```
    #[cfg(cares1_28)]
    pub fn query_dnsrec(
        &self,
        name: &str,
        dns_class: c_ares::DnsCls,
        query_type: c_ares::DnsRecordType,
    ) -> c_ares::Result<CAresFuture<c_ares::DnsRecord>> {
        let (sender, receiver) = futures_channel::oneshot::channel();
        self.inner
            .query_dnsrec(name, dns_class, query_type, |result| {
                let _ = sender.send(result.and_then(c_ares::DnsRecord::try_clone));
            })?;
        let resolver = Arc::clone(&self.inner);
        Ok(CAresFuture::new(receiver, resolver))
    }

    /// Initiate a series of DNS queries using a pre-built [`c_ares::DnsRecord`], receiving a
    /// parsed [`c_ares::DnsRecord`].
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use c_ares::*;
    /// use futures_executor::block_on;
    ///
    /// let resolver = c_ares_resolver::FutureResolver::new().unwrap();
    /// let mut query = DnsRecord::new(0, DnsFlags::RD, DnsOpcode::Query, DnsRcode::NoError).unwrap();
    /// query.query_add("example.com", DnsRecordType::A, DnsCls::IN).unwrap();
    /// let response = block_on(resolver.search_dnsrec(&query).unwrap()).unwrap();
    /// for rr in response.rrs(DnsSection::Answer) {
    ///     if let Some(addr) = rr.get_addr(DnsRrKey::A_ADDR) {
    ///         println!("address: {addr}");
    ///     }
    /// }
    /// ```
    #[cfg(cares1_28)]
    pub fn search_dnsrec(
        &self,
        dnsrec: &c_ares::DnsRecord,
    ) -> c_ares::Result<CAresFuture<c_ares::DnsRecord>> {
        let (sender, receiver) = futures_channel::oneshot::channel();
        self.inner.search_dnsrec(dnsrec, |result| {
            let _ = sender.send(result.and_then(c_ares::DnsRecord::try_clone));
        })?;
        let resolver = Arc::clone(&self.inner);
        Ok(CAresFuture::new(receiver, resolver))
    }

    /// Block until notified that there are no longer any queries in queue, or the specified
    /// timeout has expired.
    ///
    /// Pass `None` to wait indefinitely.
    #[cfg(cares1_27)]
    pub fn queue_wait_empty(&self, timeout: Option<std::time::Duration>) -> c_ares::Result<()> {
        self.inner.queue_wait_empty(timeout)
    }

    /// Retrieve the total number of active queries pending answers from servers.
    #[cfg(cares1_27)]
    pub fn queue_active_queries(&self) -> usize {
        self.inner.queue_active_queries()
    }

    /// Cancel all requests made on this `FutureResolver`.
    pub fn cancel(&self) {
        self.inner.cancel();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    fn assert_send<T: Send>() {}
    fn assert_sync<T: Sync>() {}

    #[test]
    fn future_resolver_is_send() {
        assert_send::<FutureResolver>();
    }

    #[test]
    fn future_resolver_is_sync() {
        assert_sync::<FutureResolver>();
    }

    #[test]
    fn c_ares_future_is_send() {
        assert_send::<CAresFuture<c_ares::AResults>>();
    }

    #[test]
    fn c_ares_future_is_sync() {
        assert_sync::<CAresFuture<c_ares::AResults>>();
    }

    #[test]
    fn future_resolver_new() {
        let resolver = FutureResolver::new();
        assert!(resolver.is_ok());
    }

    #[test]
    fn future_resolver_with_options() {
        let options = Options::new();
        let resolver = FutureResolver::with_options(options);
        assert!(resolver.is_ok());
    }

    #[test]
    fn future_resolver_with_custom_options() {
        let mut options = Options::new();
        options.set_timeout(Duration::from_secs(2)).set_tries(2);
        let resolver = FutureResolver::with_options(options);
        assert!(resolver.is_ok());
    }

    #[test]
    fn future_resolver_set_local_ipv4() {
        let resolver = FutureResolver::new().unwrap();
        let result = resolver.set_local_ipv4(Ipv4Addr::LOCALHOST);
        assert!(std::ptr::eq(result, &raw const resolver));
    }

    #[test]
    fn future_resolver_set_local_ipv6() {
        let resolver = FutureResolver::new().unwrap();
        let ipv6 = Ipv6Addr::LOCALHOST;
        let result = resolver.set_local_ipv6(ipv6);
        assert!(std::ptr::eq(result, &raw const resolver));
    }

    #[test]
    fn future_resolver_set_local_device() {
        let resolver = FutureResolver::new().unwrap();
        let result = resolver.set_local_device("lo").unwrap();
        assert!(std::ptr::eq(result, &raw const resolver));
    }

    #[test]
    fn future_resolver_set_servers_valid() {
        let resolver = FutureResolver::new().unwrap();
        let result = resolver.set_servers(["8.8.8.8", "8.8.4.4"]);
        assert!(result.is_ok());
    }

    #[test]
    fn future_resolver_set_sortlist_valid() {
        let resolver = FutureResolver::new().unwrap();
        let result = resolver.set_sortlist(["130.155.160.0/255.255.240.0"]);
        assert!(result.is_ok());
    }

    #[test]
    fn future_resolver_cancel() {
        let resolver = FutureResolver::new().unwrap();
        resolver.cancel(); // Should not panic
    }

    #[test]
    #[cfg(cares1_22)]
    fn future_resolver_reinit() {
        let resolver = FutureResolver::new().unwrap();
        let result = resolver.reinit();
        assert!(result.is_ok());
    }

    #[test]
    #[cfg(cares1_24)]
    fn future_resolver_servers() {
        let resolver = FutureResolver::new().unwrap();
        let _ = resolver.set_servers(["8.8.8.8"]);
        let servers = resolver.servers();
        assert!(!servers.is_empty());
    }

    #[test]
    #[cfg(cares1_27)]
    fn future_resolver_queue_active_queries() {
        let resolver = FutureResolver::new().unwrap();
        assert_eq!(resolver.queue_active_queries(), 0);
    }

    #[test]
    #[cfg(cares1_27)]
    fn future_resolver_queue_wait_empty() {
        let resolver = FutureResolver::new().unwrap();
        let result = resolver.queue_wait_empty(Some(std::time::Duration::ZERO));
        assert!(result.is_ok() || result == Err(c_ares::Error::ENOTIMP));
    }

    #[test]
    #[cfg(cares1_29)]
    fn future_resolver_set_server_state_callback() {
        let resolver = FutureResolver::new().unwrap();
        let result = resolver.set_server_state_callback(|_server, _success, _flags| {});
        assert!(std::ptr::eq(result, &raw const resolver));
    }

    #[test]
    fn debug_future_resolver() {
        let resolver = FutureResolver::new().unwrap();
        let debug = format!("{resolver:?}");
        assert!(debug.contains("FutureResolver"));
    }

    #[test]
    fn debug_cares_future() {
        let resolver = FutureResolver::new().unwrap();
        let future = resolver.query_a("example.com");
        let debug = format!("{future:?}");
        assert!(debug.contains("CAresFuture"));
    }
}