monovm-whois-rust 1.0.0

//! [`ThrottleTransport`]: stay inside a registry's query budget.
//!
//! Registries rate-limit aggressively, and the punishment is silence rather than
//! an error — most port 43 servers simply stop answering an address that queries
//! too fast, and some keep refusing for hours. Pacing is therefore not politeness
//! but the only way a bulk check finishes with real answers.
//!
//! The budget is tracked per host, because that is where the limit lives. Two
//! `.com` queries share a limit; a `.com` and a `.de` query do not.

use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use crate::error::Result;
use crate::registry::Endpoint;
use crate::transport::{Query, RawResponse, Transport};

#[cfg(feature = "async")]
use crate::transport::{AsyncTransport, BoxFuture};

/// How closely together queries may be sent.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ThrottlePolicy {
    /// Minimum gap between two queries to the same host.
    per_host: Duration,
    /// Minimum gap between any two queries, whatever the host.
    global: Duration,
    /// Per-host overrides, for registries with a documented limit.
    overrides: HashMap<String, Duration>,
}

impl ThrottlePolicy {
    /// No pacing at all.
    pub fn unlimited() -> Self {
        ThrottlePolicy {
            per_host: Duration::ZERO,
            global: Duration::ZERO,
            overrides: HashMap::new(),
        }
    }

    /// A minimum gap between queries to the same host.
    pub fn per_host(gap: Duration) -> Self {
        ThrottlePolicy {
            per_host: gap,
            global: Duration::ZERO,
            overrides: HashMap::new(),
        }
    }

    /// Also apply a minimum gap between any two queries.
    pub fn with_global(mut self, gap: Duration) -> Self {
        self.global = gap;
        self
    }

    /// Override the gap for one host.
    pub fn with_host(mut self, host: impl Into<String>, gap: Duration) -> Self {
        self.overrides.insert(host.into().to_ascii_lowercase(), gap);
        self
    }

    /// The gap that applies to a host.
    pub fn gap_for(&self, host: &str) -> Duration {
        self.overrides
            .get(&host.to_ascii_lowercase())
            .copied()
            .unwrap_or(self.per_host)
    }

    /// Whether this policy would ever make a caller wait.
    pub fn is_unlimited(&self) -> bool {
        self.per_host.is_zero() && self.global.is_zero() && self.overrides.is_empty()
    }
}

impl Default for ThrottlePolicy {
    /// One second between queries to the same host, none between different hosts.
    ///
    /// Chosen because it is the slowest documented limit among the major
    /// registries that publish one, and because the cost of being wrong in this
    /// direction is a slower check, while the cost of being wrong in the other is
    /// an hour of refusals.
    fn default() -> Self {
        ThrottlePolicy::per_host(Duration::from_secs(1))
    }
}

/// The last time each host was queried.
#[derive(Debug, Default)]
struct LastSeen {
    per_host: HashMap<String, Instant>,
    any: Option<Instant>,
}

impl LastSeen {
    /// How long to wait before querying `host`, and record the intent to do so.
    ///
    /// Reserving the slot inside the lock is what makes this correct under
    /// concurrency: two threads that ask at the same moment get different waits
    /// rather than both being told zero.
    fn reserve(&mut self, host: &str, policy: &ThrottlePolicy) -> Duration {
        let now = Instant::now();

        // What is stored is when the previous query was *scheduled to be sent*, which
        // may be in the future — that is exactly what makes reservations stack. So the
        // wait is measured forward from `last + gap`. Measuring backward from `now`
        // instead, as an elapsed-time subtraction does, hands three concurrent callers
        // the same one-gap wait and defeats the whole mechanism.
        let wait_until = |last: Option<Instant>, gap: Duration| match last {
            Some(last) => (last + gap).saturating_duration_since(now),
            None => Duration::ZERO,
        };

        let wait = wait_until(self.per_host.get(host).copied(), policy.gap_for(host))
            .max(wait_until(self.any, policy.global));

        let scheduled = now + wait;
        self.per_host.insert(host.to_string(), scheduled);
        self.any = Some(scheduled);

        wait
    }
}

/// Paces queries so a registry keeps answering.
///
/// ```
/// # #[cfg(feature = "blocking")] {
/// use std::time::Duration;
/// use monovm_whois::transport::{ThrottlePolicy, ThrottleTransport, Whois43Transport};
///
/// let policy = ThrottlePolicy::per_host(Duration::from_millis(500))
///     // DENIC documents a stricter limit than most.
///     .with_host("whois.denic.de", Duration::from_secs(2));
/// let transport = ThrottleTransport::new(Whois43Transport::default(), policy);
/// # }
/// ```
#[derive(Debug)]
pub struct ThrottleTransport<T> {
    inner: T,
    policy: ThrottlePolicy,
    state: Mutex<LastSeen>,
}

impl<T> ThrottleTransport<T> {
    /// Wrap a transport.
    pub fn new(inner: T, policy: ThrottlePolicy) -> Self {
        ThrottleTransport {
            inner,
            policy,
            state: Mutex::new(LastSeen::default()),
        }
    }

    /// The policy in force.
    pub fn policy(&self) -> &ThrottlePolicy {
        &self.policy
    }

    /// The wrapped transport.
    pub fn inner(&self) -> &T {
        &self.inner
    }

    /// Unwrap, returning the transport.
    pub fn into_inner(self) -> T {
        self.inner
    }

    fn reserve(&self, endpoint: &Endpoint) -> Duration {
        if self.policy.is_unlimited() {
            return Duration::ZERO;
        }

        let host = throttle_key(endpoint);
        self.state
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .reserve(&host, &self.policy)
    }
}

impl<T: Transport> Transport for ThrottleTransport<T> {
    fn supports(&self, endpoint: &Endpoint) -> bool {
        self.inner.supports(endpoint)
    }

    fn fetch(&self, query: &Query) -> Result<RawResponse> {
        let wait = self.reserve(&query.endpoint);
        if !wait.is_zero() {
            std::thread::sleep(wait);
        }
        self.inner.fetch(query)
    }

    fn name(&self) -> String {
        format!("throttle({})", self.inner.name())
    }
}

/// The asynchronous counterpart of [`ThrottleTransport`].
#[cfg(feature = "async")]
#[derive(Debug)]
pub struct AsyncThrottleTransport<T> {
    inner: T,
    policy: ThrottlePolicy,
    state: Mutex<LastSeen>,
}

#[cfg(feature = "async")]
impl<T> AsyncThrottleTransport<T> {
    /// Wrap a transport.
    pub fn new(inner: T, policy: ThrottlePolicy) -> Self {
        AsyncThrottleTransport {
            inner,
            policy,
            state: Mutex::new(LastSeen::default()),
        }
    }

    /// The policy in force.
    pub fn policy(&self) -> &ThrottlePolicy {
        &self.policy
    }

    /// The wrapped transport.
    pub fn inner(&self) -> &T {
        &self.inner
    }

    /// Unwrap, returning the transport.
    pub fn into_inner(self) -> T {
        self.inner
    }
}

#[cfg(feature = "async")]
impl<T: AsyncTransport> AsyncTransport for AsyncThrottleTransport<T> {
    fn supports(&self, endpoint: &Endpoint) -> bool {
        self.inner.supports(endpoint)
    }

    fn fetch<'a>(&'a self, query: &'a Query) -> BoxFuture<'a, Result<RawResponse>> {
        Box::pin(async move {
            // The lock is a std `Mutex` and is released before the await: holding
            // one across an await point would serialise every task on it, which is
            // precisely what a per-host budget is meant to avoid.
            let wait = if self.policy.is_unlimited() {
                Duration::ZERO
            } else {
                let host = throttle_key(&query.endpoint);
                self.state
                    .lock()
                    .unwrap_or_else(|poisoned| poisoned.into_inner())
                    .reserve(&host, &self.policy)
            };

            if !wait.is_zero() {
                tokio::time::sleep(wait).await;
            }
            self.inner.fetch(query).await
        })
    }

    fn name(&self) -> String {
        format!("async-throttle({})", self.inner.name())
    }
}

/// The thing a budget is tracked against: a host name.
///
/// For RDAP that means the URL's authority rather than the whole URL, since two
/// paths on one server share its limit.
fn throttle_key(endpoint: &Endpoint) -> String {
    match endpoint {
        Endpoint::Whois(whois) => whois.host().to_ascii_lowercase(),
        Endpoint::Rdap(rdap) => rdap
            .base()
            .split_once("://")
            .map(|(_, rest)| rest)
            .unwrap_or_else(|| rdap.base())
            .split('/')
            .next()
            .unwrap_or_else(|| rdap.base())
            .to_ascii_lowercase(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::domain::Tld;
    use crate::transport::mock::MockTransport;

    fn query(endpoint: Endpoint) -> Query {
        Query::new(endpoint, "example.com", Tld::parse("com").unwrap())
    }

    #[test]
    fn budgets_are_tracked_per_host_not_per_url() {
        assert_eq!(
            throttle_key(&Endpoint::whois("WHOIS.Example")),
            "whois.example"
        );
        assert_eq!(
            throttle_key(&Endpoint::rdap("https://rdap.example/com/v1/domain/")),
            "rdap.example"
        );
        // Two paths on one server share its limit.
        assert_eq!(
            throttle_key(&Endpoint::rdap("https://rdap.example/a/")),
            throttle_key(&Endpoint::rdap("https://rdap.example/b/"))
        );
    }

    #[test]
    fn per_host_overrides_beat_the_default_gap() {
        let policy = ThrottlePolicy::per_host(Duration::from_millis(100))
            .with_host("Slow.Example", Duration::from_secs(5));

        assert_eq!(policy.gap_for("slow.example"), Duration::from_secs(5));
        assert_eq!(policy.gap_for("SLOW.EXAMPLE"), Duration::from_secs(5));
        assert_eq!(policy.gap_for("other.example"), Duration::from_millis(100));
    }

    #[test]
    fn the_first_query_to_a_host_never_waits() {
        let mut state = LastSeen::default();
        let policy = ThrottlePolicy::per_host(Duration::from_secs(10));
        assert_eq!(state.reserve("a.example", &policy), Duration::ZERO);
        // A different host is a different budget.
        assert_eq!(state.reserve("b.example", &policy), Duration::ZERO);
    }

    #[test]
    fn a_second_query_to_the_same_host_waits() {
        let mut state = LastSeen::default();
        let policy = ThrottlePolicy::per_host(Duration::from_secs(10));

        state.reserve("a.example", &policy);
        let wait = state.reserve("a.example", &policy);
        assert!(wait > Duration::from_secs(9), "got {wait:?}");
    }

    #[test]
    fn reservations_stack_so_concurrent_callers_get_different_slots() {
        let mut state = LastSeen::default();
        let policy = ThrottlePolicy::per_host(Duration::from_secs(1));

        assert_eq!(state.reserve("a.example", &policy), Duration::ZERO);
        let second = state.reserve("a.example", &policy);
        let third = state.reserve("a.example", &policy);

        assert!(third > second, "{third:?} should be later than {second:?}");
        assert!(third > Duration::from_millis(1900), "got {third:?}");
    }

    #[test]
    fn a_global_gap_applies_across_hosts() {
        let mut state = LastSeen::default();
        let policy = ThrottlePolicy::unlimited().with_global(Duration::from_secs(5));

        state.reserve("a.example", &policy);
        let wait = state.reserve("b.example", &policy);
        assert!(wait > Duration::from_secs(4), "got {wait:?}");
    }

    #[test]
    fn an_unlimited_policy_never_waits() {
        let policy = ThrottlePolicy::unlimited();
        assert!(policy.is_unlimited());

        let transport = ThrottleTransport::new(MockTransport::answering("ok"), policy);
        let started = Instant::now();
        for _ in 0..5 {
            transport
                .fetch(&query(Endpoint::whois("a.example")))
                .unwrap();
        }
        assert!(started.elapsed() < Duration::from_millis(200));
    }

    #[test]
    fn pacing_actually_delays_the_second_call() {
        let transport = ThrottleTransport::new(
            MockTransport::answering("ok"),
            ThrottlePolicy::per_host(Duration::from_millis(120)),
        );

        let started = Instant::now();
        transport
            .fetch(&query(Endpoint::whois("a.example")))
            .unwrap();
        transport
            .fetch(&query(Endpoint::whois("a.example")))
            .unwrap();
        assert!(
            started.elapsed() >= Duration::from_millis(100),
            "got {:?}",
            started.elapsed()
        );
    }
}