monovm-whois-rust 1.0.0

//! The asynchronous [`AsyncWhoisClient`] and its builder.
//!
//! A mirror of the blocking client. The lookup *sequence* is not duplicated — that
//! lives in [`plan`](crate::client::plan) and is shared — but the waiting is, because
//! blocking a Tokio worker thread on a socket is exactly the thing an async client
//! exists to avoid.

use std::sync::Arc;
use std::time::Duration;

use crate::cache::{MemoryCache, NullCache, ResponseCache};
use crate::client::lookup::{Explanation, Lookup};
use crate::client::plan::{self, Plan, Preference, ReferralPolicy};
use crate::client::Parts;
use crate::detect::{DetectionEngine, Evidence, Verdict};
use crate::domain::{Availability, DomainName, Tld};
use crate::error::{Error, Result};
use crate::registry::{Endpoint, RegistryProvider, WhoisEndpoint};
use crate::transport::{
    AsyncCachingTransport, AsyncRetryTransport, AsyncRouter, AsyncThrottleTransport,
    AsyncTransport, AsyncWhois43Transport, RawResponse, RetryPolicy, ThrottlePolicy,
    TransportConfig,
};

#[cfg(feature = "parser")]
use crate::parser::{CompositeParser, WhoisRecord};

/// Looks domains up, asynchronously.
///
/// Cheap to clone and safe to share across tasks. Sharing one client is what makes
/// the rate limiter work: a hundred tasks with a hundred clients have no shared
/// budget and will get the whole process blocked by a registry.
#[derive(Debug, Clone)]
pub struct AsyncWhoisClient {
    parts: Parts,
    transport: Arc<dyn AsyncTransport>,
    concurrency: usize,
}

/// How many lookups [`AsyncWhoisClient::lookup_many`] keeps in flight by default.
///
/// Modest on purpose. The limit that matters is per-registry, and a batch of a
/// thousand `.com` names is a thousand queries to one server however many tasks
/// carry them; a high ceiling only buys sockets held open while the rate limiter
/// makes them wait.
pub const DEFAULT_CONCURRENCY: usize = 16;

impl AsyncWhoisClient {
    /// A client with the default configuration.
    pub fn new() -> Result<Self> {
        AsyncWhoisClient::builder().build()
    }

    /// Start configuring a client.
    pub fn builder() -> AsyncWhoisClientBuilder {
        AsyncWhoisClientBuilder::new()
    }

    /// Look a domain up.
    pub async fn lookup(&self, domain: &str) -> Result<Lookup> {
        let name = DomainName::parse(domain)?;
        self.lookup_name(&name).await
    }

    /// Look an already-validated name up.
    pub async fn lookup_name(&self, name: &DomainName) -> Result<Lookup> {
        let resolution = self.parts.registry.resolve(name)?;
        let plan = Plan::build(resolution, self.parts.preference)?;

        let (verdict, mut responses) = self.first_usable_answer(&plan).await?;

        if verdict.availability.is_registered() {
            self.follow_referrals(&plan, &mut responses).await;
        }

        #[cfg(feature = "parser")]
        let record = self
            .parts
            .parser
            .parse_all(&responses)
            .ok()
            .filter(|record| !record.is_empty());

        Ok(Lookup {
            queried: plan.resolution.queried.clone(),
            domain: plan.resolution.registrable.clone(),
            tld: plan.resolution.tld.clone(),
            verdict,
            responses,
            #[cfg(feature = "parser")]
            record,
        })
    }

    /// Just the verdict.
    pub async fn availability(&self, domain: &str) -> Result<Availability> {
        Ok(self.lookup(domain).await?.availability())
    }

    /// Whether the domain is free to register at the ordinary price.
    pub async fn is_available(&self, domain: &str) -> Result<bool> {
        Ok(self.lookup(domain).await?.is_available())
    }

    /// Look up and report what every detection rule made of the response.
    pub async fn explain(&self, domain: &str) -> Result<Explanation> {
        let lookup = self.lookup(domain).await?;
        let registry = self.parts.registry.get(&lookup.tld);

        let report = match lookup.primary_response() {
            Some(response) => {
                let evidence = Evidence::from_response(response, &lookup.tld, registry.as_deref());
                self.parts.engine.report(&evidence)
            }
            None => {
                let evidence = Evidence::new(
                    "",
                    crate::transport::ResponseKind::WhoisText,
                    &lookup.tld,
                    registry.as_deref(),
                );
                self.parts.engine.report(&evidence)
            }
        };

        Ok(Explanation { lookup, report })
    }

    /// The parsed registration record.
    ///
    /// # Errors
    ///
    /// As [`lookup`](AsyncWhoisClient::lookup), plus [`Error::Inconclusive`] when the
    /// domain is not registered. Decided on the verdict rather than on whether anything
    /// parsed — see [`WhoisClient::record`](crate::WhoisClient::record) for why that
    /// distinction matters.
    #[cfg(feature = "parser")]
    pub async fn record(&self, domain: &str) -> Result<WhoisRecord> {
        let lookup = self.lookup(domain).await?;
        let consulted = || {
            lookup
                .consulted()
                .iter()
                .map(Endpoint::address)
                .collect::<Vec<_>>()
                .join(", ")
        };

        if !lookup.availability().is_registered() {
            return Err(Error::Inconclusive {
                domain: lookup.domain.as_ascii().to_string(),
                consulted: consulted(),
                detail: format!(
                    "no registration record: the domain is {}",
                    lookup.availability()
                ),
            });
        }

        lookup.record.clone().ok_or_else(|| Error::Inconclusive {
            domain: lookup.domain.as_ascii().to_string(),
            consulted: consulted(),
            detail: "the domain is registered but no field of its record could be parsed"
                .to_string(),
        })
    }

    /// Look several domains up concurrently.
    ///
    /// Results come back in the order the inputs were given, not the order the
    /// servers answered, so a caller can zip them against their input list.
    ///
    /// At most [`concurrency`](AsyncWhoisClientBuilder::concurrency) lookups are in
    /// flight at once. That bound is separate from the rate limiter and both matter:
    /// the limiter paces one host, while this caps how many sockets and tasks the
    /// whole batch may hold open.
    ///
    /// ```no_run
    /// # #[cfg(all(feature = "async", feature = "rdap"))] {
    /// # async fn run() -> Result<(), monovm_whois::Error> {
    /// use monovm_whois::AsyncWhoisClient;
    ///
    /// let client = AsyncWhoisClient::new()?;
    /// for (domain, outcome) in client.lookup_many(["example.com", "example.net"]).await {
    ///     match outcome {
    ///         Ok(lookup) => println!("{domain}: {}", lookup.availability()),
    ///         Err(error) => println!("{domain}: {error}"),
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// # }
    /// ```
    pub async fn lookup_many<I, S>(&self, domains: I) -> Vec<(String, Result<Lookup>)>
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        let inputs: Vec<String> = domains.into_iter().map(Into::into).collect();
        let permits = Arc::new(tokio::sync::Semaphore::new(self.concurrency.max(1)));

        let mut tasks = tokio::task::JoinSet::new();
        for (index, input) in inputs.iter().cloned().enumerate() {
            let client = self.clone();
            let permits = Arc::clone(&permits);

            tasks.spawn(async move {
                // Dropped at the end of the task, releasing the slot. An error here
                // means the semaphore was closed, which cannot happen while this
                // function holds it.
                let _permit = permits.acquire().await;
                let outcome = client.lookup(&input).await;
                (index, outcome)
            });
        }

        let mut outcomes: Vec<Option<Result<Lookup>>> = (0..inputs.len()).map(|_| None).collect();

        while let Some(joined) = tasks.join_next().await {
            if let Ok((index, outcome)) = joined {
                outcomes[index] = Some(outcome);
            }
        }

        inputs
            .into_iter()
            .zip(outcomes)
            .map(|(input, outcome)| {
                let outcome = outcome.unwrap_or_else(|| {
                    // Only reachable if a lookup task panicked, which would be a bug
                    // in this crate rather than a condition to model.
                    Err(Error::Inconclusive {
                        domain: input.clone(),
                        consulted: "none".to_string(),
                        detail: "the lookup task did not complete".to_string(),
                    })
                });
                (input, outcome)
            })
            .collect()
    }

    /// The registry provider in use.
    pub fn registry(&self) -> &Arc<dyn RegistryProvider> {
        &self.parts.registry
    }

    /// The detection engine in use.
    pub fn engine(&self) -> &Arc<DetectionEngine> {
        &self.parts.engine
    }

    /// A description of the transport stack.
    pub fn transport_name(&self) -> String {
        self.transport.name()
    }

    /// Every suffix this client can look up.
    pub fn supported_tlds(&self) -> Vec<Tld> {
        self.parts.registry.tlds()
    }

    /// Whether a name's suffix is served.
    pub fn can_lookup(&self, domain: &str) -> bool {
        DomainName::parse(domain)
            .map(|name| self.parts.registry.can_resolve(&name))
            .unwrap_or(false)
    }

    async fn first_usable_answer(&self, plan: &Plan) -> Result<(Verdict, Vec<RawResponse>)> {
        let mut failures: Vec<(Endpoint, Error)> = Vec::new();

        for endpoint in &plan.attempts {
            let query = plan.query(endpoint);

            let response = match self.transport.fetch(&query).await {
                Ok(response) => response,
                Err(error) => {
                    if error.is_endpoint_failure() {
                        failures.push((endpoint.clone(), error));
                        continue;
                    }
                    return Err(error);
                }
            };

            match plan::interpret(&self.parts.engine, &response, &plan.resolution) {
                Ok(verdict) => return Ok((verdict, vec![response])),
                Err(error) => failures.push((endpoint.clone(), error)),
            }
        }

        Err(plan::combined_failure(
            plan.resolution.registrable.as_ascii(),
            failures,
        ))
    }

    async fn follow_referrals(&self, plan: &Plan, responses: &mut Vec<RawResponse>) {
        let mut visited: Vec<WhoisEndpoint> = responses
            .iter()
            .filter_map(|response| match response.endpoint() {
                Endpoint::Whois(endpoint) => Some(endpoint.clone()),
                Endpoint::Rdap(_) => None,
            })
            .collect();

        loop {
            let Some(last) = responses.last() else { return };
            let Some(next) = plan::next_referral(
                last,
                &plan.resolution.registry,
                self.parts.referrals,
                &visited,
            ) else {
                return;
            };

            let query = plan.query(&Endpoint::Whois(next.clone()));
            match self.transport.fetch(&query).await {
                Ok(response) => {
                    visited.push(next);
                    responses.push(response);
                }
                Err(_) => return,
            }
        }
    }
}

/// Assembles an [`AsyncWhoisClient`].
#[derive(Debug, Default)]
pub struct AsyncWhoisClientBuilder {
    registry: Option<Arc<dyn RegistryProvider>>,
    engine: Option<Arc<DetectionEngine>>,
    #[cfg(feature = "parser")]
    parser: Option<Arc<CompositeParser>>,
    transport: Option<Arc<dyn AsyncTransport>>,
    cache: Option<Arc<dyn ResponseCache>>,
    config: TransportConfig,
    retry: RetryPolicy,
    throttle: ThrottlePolicy,
    preference: Preference,
    referrals: ReferralPolicy,
    concurrency: usize,
}

impl AsyncWhoisClientBuilder {
    /// A builder holding the defaults.
    pub fn new() -> Self {
        AsyncWhoisClientBuilder {
            registry: None,
            engine: None,
            #[cfg(feature = "parser")]
            parser: None,
            transport: None,
            cache: None,
            config: TransportConfig::default(),
            retry: RetryPolicy::default(),
            throttle: ThrottlePolicy::default(),
            preference: Preference::default(),
            referrals: ReferralPolicy::default(),
            concurrency: DEFAULT_CONCURRENCY,
        }
    }

    /// Cap how many lookups [`AsyncWhoisClient::lookup_many`] runs at once.
    ///
    /// Zero is treated as one.
    pub fn concurrency(mut self, limit: usize) -> Self {
        self.concurrency = limit.max(1);
        self
    }

    /// Use a different source of registry definitions.
    pub fn registry(mut self, registry: impl RegistryProvider + 'static) -> Self {
        self.registry = Some(Arc::new(registry));
        self
    }

    /// Use an already-shared registry provider.
    pub fn shared_registry(mut self, registry: Arc<dyn RegistryProvider>) -> Self {
        self.registry = Some(registry);
        self
    }

    /// Use a different detection engine.
    pub fn engine(mut self, engine: DetectionEngine) -> Self {
        self.engine = Some(Arc::new(engine));
        self
    }

    /// Use a different record parser.
    #[cfg(feature = "parser")]
    pub fn parser(mut self, parser: CompositeParser) -> Self {
        self.parser = Some(Arc::new(parser));
        self
    }

    /// Supply the whole transport stack, overriding the timeout, retry, throttle
    /// and cache settings.
    pub fn transport(mut self, transport: impl AsyncTransport + 'static) -> Self {
        self.transport = Some(Arc::new(transport));
        self
    }

    /// Supply an already-shared transport stack.
    pub fn shared_transport(mut self, transport: Arc<dyn AsyncTransport>) -> Self {
        self.transport = Some(transport);
        self
    }

    /// Set the connect and read timeouts.
    pub fn timeouts(mut self, config: TransportConfig) -> Self {
        self.config = config;
        self
    }

    /// Set the connect timeout.
    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
        self.config.connect_timeout = timeout;
        self
    }

    /// Set the read timeout.
    pub fn read_timeout(mut self, timeout: Duration) -> Self {
        self.config.read_timeout = timeout;
        self
    }

    /// Set the retry policy.
    pub fn retry(mut self, policy: RetryPolicy) -> Self {
        self.retry = policy;
        self
    }

    /// Set the rate-limiting policy.
    pub fn throttle(mut self, policy: ThrottlePolicy) -> Self {
        self.throttle = policy;
        self
    }

    /// Set a minimum gap between queries to the same host.
    pub fn throttle_per_host(mut self, gap: Duration) -> Self {
        self.throttle = ThrottlePolicy::per_host(gap);
        self
    }

    /// Cache responses in memory for `ttl`.
    pub fn memory_cache(mut self, ttl: Duration) -> Self {
        self.cache = Some(Arc::new(MemoryCache::with_ttl(ttl)));
        self
    }

    /// Cache responses somewhere else.
    pub fn cache(mut self, cache: impl ResponseCache + 'static) -> Self {
        self.cache = Some(Arc::new(cache));
        self
    }

    /// Use an already-shared cache.
    pub fn shared_cache(mut self, cache: Arc<dyn ResponseCache>) -> Self {
        self.cache = Some(cache);
        self
    }

    /// Choose which protocol to try first.
    pub fn prefer(mut self, preference: Preference) -> Self {
        self.preference = preference;
        self
    }

    /// Set how far to chase referrals.
    pub fn referrals(mut self, policy: ReferralPolicy) -> Self {
        self.referrals = policy;
        self
    }

    /// Build the client.
    pub fn build(self) -> Result<AsyncWhoisClient> {
        let transport = match &self.transport {
            Some(transport) => Arc::clone(transport),
            None => self.default_transport()?,
        };
        let concurrency = self.concurrency.max(1);

        let parts = Parts::new(self.registry, self.engine, self.preference, self.referrals);
        #[cfg(feature = "parser")]
        let parts = parts.with_parser(self.parser);

        Ok(AsyncWhoisClient {
            parts,
            transport,
            concurrency,
        })
    }

    fn default_transport(&self) -> Result<Arc<dyn AsyncTransport>> {
        #[cfg(feature = "rdap")]
        let router = AsyncRouter::new()
            .with(AsyncWhois43Transport::with_config(self.config))
            .with(crate::transport::AsyncRdapTransport::with_config(
                self.config,
            )?);

        #[cfg(not(feature = "rdap"))]
        let router = AsyncRouter::new().with(AsyncWhois43Transport::with_config(self.config));

        let throttled = AsyncThrottleTransport::new(router, self.throttle.clone());
        let retried = AsyncRetryTransport::new(throttled, self.retry);

        Ok(match &self.cache {
            Some(cache) => Arc::new(AsyncCachingTransport::new(retried, Arc::clone(cache))),
            None => Arc::new(AsyncCachingTransport::new(retried, NullCache)),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::registry::JsonRegistry;
    use crate::transport::{MockTransport, Scripted};

    fn registry() -> JsonRegistry {
        JsonRegistry::from_json(
            r#"{"registries":[
                {"tlds":["test"],"whois":"registry.test","available":["No match for"],"thin":true},
                {"tlds":["both"],"whois":"registry.both","rdap":["https://rdap.both/"]}
            ]}"#,
            "test",
        )
        .unwrap()
    }

    fn client(transport: MockTransport) -> AsyncWhoisClient {
        AsyncWhoisClient::builder()
            .registry(registry())
            .transport(transport)
            .build()
            .unwrap()
    }

    const REGISTERED: &str = "\
Domain Name: EXAMPLE.TEST
Registrar: Example Registrar, LLC
Registrar WHOIS Server: registrar.test
Domain Status: clientTransferProhibited
";

    #[tokio::test]
    async fn an_available_domain_is_reported_available() {
        let client = client(MockTransport::answering("No match for \"NOTHERE.TEST\""));
        let lookup = client.lookup("nothere.test").await.unwrap();

        assert_eq!(lookup.availability(), Availability::Available);
        assert_eq!(lookup.verdict.rule, "registry-marker");
    }

    #[tokio::test]
    async fn a_thin_registrys_referral_is_followed() {
        let transport = MockTransport::routed([
            ("registry.test".to_string(), REGISTERED.to_string()),
            (
                "registrar.test".to_string(),
                "Registrant Name: Ada Lovelace\n".to_string(),
            ),
        ]);
        let client = client(transport.clone());

        let lookup = client.lookup("example.test").await.unwrap();
        assert!(lookup.followed_referral());
        assert_eq!(transport.contacted(), ["registry.test", "registrar.test"]);
    }

    #[tokio::test]
    async fn a_refused_endpoint_falls_through_to_the_next() {
        let transport = MockTransport::routed([
            (
                "registry.both".to_string(),
                "%% queries limit exceeded".to_string(),
            ),
            (
                "https://rdap.both/".to_string(),
                r#"{"errorCode":404}"#.to_string(),
            ),
        ]);
        let client = client(transport.clone());

        let lookup = client.lookup("nothere.both").await.unwrap();
        assert_eq!(lookup.availability(), Availability::Available);
        assert_eq!(transport.contacted().len(), 2);
    }

    #[tokio::test]
    async fn invalid_input_is_rejected_before_any_query() {
        let transport = MockTransport::answering("should not be reached");
        let client = client(transport.clone());

        assert!(matches!(
            client.lookup("localhost").await.unwrap_err(),
            Error::InvalidDomain { .. }
        ));
        assert_eq!(transport.call_count(), 0);
    }

    #[tokio::test]
    async fn every_endpoint_failing_is_reported_together() {
        let transport = MockTransport::new(vec![
            Scripted::Fail(Error::Timeout {
                server: "registry.both".into(),
                elapsed: Duration::ZERO,
            }),
            Scripted::Fail(Error::Http {
                url: "https://rdap.both/".into(),
                status: 503,
            }),
        ]);
        let client = client(transport);

        assert!(matches!(
            client.lookup("example.both").await.unwrap_err(),
            Error::Inconclusive { .. }
        ));
    }

    #[tokio::test]
    async fn many_lookups_come_back_in_input_order() {
        let client = client(MockTransport::answering("No match for \"X.TEST\""));
        let results = client.lookup_many(["c.test", "a.test", "b.test"]).await;

        let names: Vec<&str> = results.iter().map(|(name, _)| name.as_str()).collect();
        assert_eq!(names, ["c.test", "a.test", "b.test"]);
        assert!(results.iter().all(|(_, outcome)| outcome.is_ok()));
    }

    #[tokio::test]
    async fn many_lookups_report_failures_per_domain() {
        let client = client(MockTransport::answering("No match for \"X.TEST\""));
        let results = client.lookup_many(["good.test", "bad.unknown"]).await;

        assert!(results[0].1.is_ok());
        assert!(matches!(
            results[1].1.as_ref().unwrap_err(),
            Error::UnsupportedTld { .. }
        ));
    }

    #[tokio::test]
    async fn the_default_stack_is_layered_as_documented() {
        let client = AsyncWhoisClient::new().unwrap();
        let name = client.transport_name();

        assert!(name.starts_with("async-cached("), "{name}");
        assert!(name.contains("async-retry("), "{name}");
        assert!(name.contains("async-throttle("), "{name}");
    }

    #[tokio::test]
    async fn a_client_is_shareable_across_tasks() {
        let transport = MockTransport::answering("No match for \"X.TEST\"");
        let client = Arc::new(client(transport.clone()));

        let mut handles = Vec::new();
        for index in 0..4 {
            let client = Arc::clone(&client);
            handles.push(tokio::spawn(async move {
                client.lookup(&format!("d{index}.test")).await.is_ok()
            }));
        }

        for handle in handles {
            assert!(handle.await.unwrap());
        }
        assert_eq!(transport.call_count(), 4);
    }
}