monovm-whois-rust 1.0.0

//! [`MockTransport`]: a transport that answers from a script instead of a socket.
//!
//! Available to dependents behind the `mock` feature. Detection logic, referral
//! chasing and retry policy are all reachable without a network, and a test that
//! needs the network is a test that fails when a registry is having a bad day.

use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use crate::error::{Error, Result};
use crate::registry::Endpoint;
use crate::transport::{AsyncTransport, BoxFuture, Query, RawResponse, ResponseKind, Transport};

/// One scripted outcome.
#[derive(Debug)]
pub enum Scripted {
    /// Answer with this WHOIS text.
    Answer(String),
    /// Answer with this RDAP JSON.
    Rdap(String),
    /// Fail with this error.
    Fail(Error),
}

impl Scripted {
    fn into_result(self, query: &Query) -> Result<RawResponse> {
        match self {
            Scripted::Answer(text) => Ok(RawResponse::new(
                query.endpoint.clone(),
                ResponseKind::WhoisText,
                text,
                Duration::from_millis(1),
            )),
            Scripted::Rdap(json) => Ok(RawResponse::new(
                query.endpoint.clone(),
                ResponseKind::RdapJson,
                json,
                Duration::from_millis(1),
            )),
            Scripted::Fail(error) => Err(error),
        }
    }
}

#[derive(Debug)]
enum Behaviour {
    /// Outcomes consumed in order, one per call.
    Script(Mutex<VecDeque<Scripted>>),
    /// A reusable answer per endpoint address; anything else fails.
    Routed(HashMap<String, String>),
}

/// A transport whose answers are decided up front.
///
/// Cloning shares the script and the call log, so a test can keep a handle for
/// assertions after moving one into a decorator:
///
/// ```
/// # #[cfg(feature = "mock")] {
/// use monovm_whois::transport::{MockTransport, Scripted, Transport};
///
/// let transport = MockTransport::new(vec![Scripted::Answer("No match".into())]);
/// let handle = transport.clone();
/// // ... move `transport` into a client ...
/// assert_eq!(handle.call_count(), 0);
/// # }
/// ```
#[derive(Debug, Clone)]
pub struct MockTransport {
    behaviour: Arc<Behaviour>,
    calls: Arc<AtomicUsize>,
    seen: Arc<Mutex<Vec<Query>>>,
    supports_rdap: bool,
}

impl MockTransport {
    /// A transport that plays the given outcomes in order.
    pub fn new(script: impl IntoIterator<Item = Scripted>) -> Self {
        MockTransport {
            behaviour: Arc::new(Behaviour::Script(Mutex::new(script.into_iter().collect()))),
            calls: Arc::new(AtomicUsize::new(0)),
            seen: Arc::new(Mutex::new(Vec::new())),
            supports_rdap: true,
        }
    }

    /// A transport that answers every call with the same WHOIS text.
    pub fn answering(text: impl Into<String>) -> Self {
        let text = text.into();
        MockTransport {
            behaviour: Arc::new(Behaviour::Routed(HashMap::from([("*".to_string(), text)]))),
            calls: Arc::new(AtomicUsize::new(0)),
            seen: Arc::new(Mutex::new(Vec::new())),
            supports_rdap: true,
        }
    }

    /// A transport that answers based on which endpoint was asked.
    ///
    /// Keys are endpoint addresses — a host for WHOIS, a base URL for RDAP — and
    /// `"*"` is the fallback. This is what makes referral chasing testable: the
    /// registry and the registrar can be given different records.
    pub fn routed(routes: impl IntoIterator<Item = (String, String)>) -> Self {
        MockTransport {
            behaviour: Arc::new(Behaviour::Routed(routes.into_iter().collect())),
            calls: Arc::new(AtomicUsize::new(0)),
            seen: Arc::new(Mutex::new(Vec::new())),
            supports_rdap: true,
        }
    }

    /// Restrict this mock to WHOIS endpoints, so RDAP falls through to another
    /// transport in a router.
    pub fn whois_only(mut self) -> Self {
        self.supports_rdap = false;
        self
    }

    /// How many times the transport has been called.
    pub fn call_count(&self) -> usize {
        self.calls.load(Ordering::SeqCst)
    }

    /// Every query received, in order.
    pub fn queries(&self) -> Vec<Query> {
        self.seen.lock().expect("mock lock poisoned").clone()
    }

    /// The wire names received, in order. Convenient for asserting that punycode
    /// or a registry's query flags came out right.
    pub fn wire_names(&self) -> Vec<String> {
        self.queries()
            .into_iter()
            .map(|query| query.wire_name)
            .collect()
    }

    /// The endpoint addresses contacted, in order.
    pub fn contacted(&self) -> Vec<String> {
        self.queries()
            .iter()
            .map(|query| query.endpoint.address())
            .collect()
    }

    fn answer(&self, query: &Query) -> Result<RawResponse> {
        self.calls.fetch_add(1, Ordering::SeqCst);
        self.seen
            .lock()
            .expect("mock lock poisoned")
            .push(query.clone());

        match &*self.behaviour {
            Behaviour::Script(script) => {
                let next = script.lock().expect("mock lock poisoned").pop_front();
                match next {
                    Some(scripted) => scripted.into_result(query),
                    None => Err(Error::Io {
                        server: query.endpoint.address(),
                        source: std::io::Error::other("mock script exhausted"),
                    }),
                }
            }
            Behaviour::Routed(routes) => {
                let address = query.endpoint.address();
                let body = routes.get(&address).or_else(|| routes.get("*"));

                match body {
                    Some(text) => {
                        let kind = if query.endpoint.is_rdap() {
                            ResponseKind::RdapJson
                        } else {
                            ResponseKind::WhoisText
                        };
                        Ok(RawResponse::new(
                            query.endpoint.clone(),
                            kind,
                            text.clone(),
                            Duration::from_millis(1),
                        ))
                    }
                    None => Err(Error::Connect {
                        server: address,
                        source: std::io::Error::other("no mock route"),
                    }),
                }
            }
        }
    }
}

impl Transport for MockTransport {
    fn supports(&self, endpoint: &Endpoint) -> bool {
        self.supports_rdap || endpoint.is_whois()
    }

    fn fetch(&self, query: &Query) -> Result<RawResponse> {
        self.answer(query)
    }

    fn name(&self) -> String {
        "mock".to_string()
    }
}

impl AsyncTransport for MockTransport {
    fn supports(&self, endpoint: &Endpoint) -> bool {
        self.supports_rdap || endpoint.is_whois()
    }

    fn fetch<'a>(&'a self, query: &'a Query) -> BoxFuture<'a, Result<RawResponse>> {
        Box::pin(async move { self.answer(query) })
    }

    fn name(&self) -> String {
        "async-mock".to_string()
    }
}

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

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

    #[test]
    fn a_script_is_played_in_order_then_runs_out() {
        let transport = MockTransport::new(vec![
            Scripted::Answer("first".into()),
            Scripted::Answer("second".into()),
        ]);

        assert_eq!(
            Transport::fetch(&transport, &query("a")).unwrap().text(),
            "first"
        );
        assert_eq!(
            Transport::fetch(&transport, &query("a")).unwrap().text(),
            "second"
        );
        assert!(Transport::fetch(&transport, &query("a")).is_err());
        assert_eq!(transport.call_count(), 3);
    }

    #[test]
    fn routes_pick_the_answer_by_endpoint() {
        let transport = MockTransport::routed([
            ("registry.example".to_string(), "thin record".to_string()),
            ("registrar.example".to_string(), "thick record".to_string()),
        ]);

        assert_eq!(
            Transport::fetch(&transport, &query("registry.example"))
                .unwrap()
                .text(),
            "thin record"
        );
        assert_eq!(
            Transport::fetch(&transport, &query("registrar.example"))
                .unwrap()
                .text(),
            "thick record"
        );
        assert!(Transport::fetch(&transport, &query("unknown.example")).is_err());
        assert_eq!(
            transport.contacted(),
            ["registry.example", "registrar.example", "unknown.example"]
        );
    }

    #[test]
    fn a_wildcard_route_catches_everything() {
        let transport = MockTransport::answering("same everywhere");
        assert_eq!(
            Transport::fetch(&transport, &query("a")).unwrap().text(),
            "same everywhere"
        );
        assert_eq!(
            Transport::fetch(&transport, &query("b")).unwrap().text(),
            "same everywhere"
        );
    }

    #[test]
    fn rdap_endpoints_produce_json_responses() {
        let transport = MockTransport::answering("{}");
        let query = Query::new(
            Endpoint::rdap("https://rdap.example/"),
            "example.com",
            Tld::parse("com").unwrap(),
        );
        assert_eq!(
            Transport::fetch(&transport, &query).unwrap().kind(),
            ResponseKind::RdapJson
        );
    }

    #[test]
    fn whois_only_declines_rdap() {
        let transport = MockTransport::answering("x").whois_only();
        assert!(Transport::supports(
            &transport,
            &Endpoint::whois("w.example")
        ));
        assert!(!Transport::supports(
            &transport,
            &Endpoint::rdap("https://rdap.example/")
        ));
    }

    #[test]
    fn clones_share_the_script_and_the_log() {
        let transport = MockTransport::new(vec![Scripted::Answer("only".into())]);
        let handle = transport.clone();

        assert_eq!(
            Transport::fetch(&transport, &query("a")).unwrap().text(),
            "only"
        );
        assert_eq!(handle.call_count(), 1);
        assert!(
            Transport::fetch(&handle, &query("a")).is_err(),
            "the script is shared"
        );
    }

    #[tokio::test]
    async fn works_as_an_async_transport_too() {
        let transport = MockTransport::answering("async answer");
        let response = AsyncTransport::fetch(&transport, &query("a"))
            .await
            .unwrap();
        assert_eq!(response.text(), "async answer");
    }
}