reserve-core 0.1.0

//! The older text protocol on port 43, the only service many country registries answer on.

use std::collections::HashMap;
use std::net::SocketAddr;
use std::path::Path;
use std::time::Duration;

use hickory_resolver::TokioResolver;
use serde::Deserialize;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

use crate::error::{Error, Result};
use crate::limit::{Pacer, Refusal};
use crate::lookup::outcome::Reason;
use crate::lookup::verdict::{self, TextVerdict};

const BUNDLED: &str = include_str!("../../data/whois-servers.json");

/// @docgen Caps the reply so a hostile or looping server cannot exhaust memory.
const MAX_ANSWER_BYTES: usize = 256 * 1024;

#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct Server {
    pub host: String,
    #[serde(default)]
    pub available_phrase: String,
}

#[derive(Debug, Deserialize)]
struct ServerFile {
    servers: HashMap<String, Server>,
}

#[derive(Debug, Clone, Default)]
pub struct Servers {
    by_suffix: HashMap<String, Server>,
}

impl Servers {
    pub fn bundled() -> Result<Self> {
        Self::parse(BUNDLED)
    }

    pub fn parse(text: &str) -> Result<Self> {
        let file: ServerFile =
            serde_json::from_str(text).map_err(|source| Error::CatalogMalformed {
                source: Box::new(source),
            })?;
        Ok(Self {
            by_suffix: file
                .servers
                .into_iter()
                .map(|(suffix, server)| (suffix.to_lowercase(), server))
                .collect(),
        })
    }

    pub fn from_file(path: &Path) -> Result<Self> {
        let text = std::fs::read_to_string(path).map_err(|source| Error::FileUnreadable {
            path: path.to_path_buf(),
            source,
        })?;
        let parsed = Self::parse(&text)?;
        if parsed.by_suffix.is_empty() {
            return Err(Error::CatalogEmptySelection);
        }
        Ok(parsed)
    }

    pub fn merge(&mut self, other: Self) {
        self.by_suffix.extend(other.by_suffix);
    }

    /// @docgen Matched longest suffix first so a multi-label extension uses its own registry before falling back to the parent.
    #[must_use]
    pub fn for_suffix(&self, suffix: &str) -> Option<&Server> {
        let suffix = suffix.trim_matches('.').to_lowercase();
        let mut rest = suffix.as_str();
        loop {
            if let Some(server) = self.by_suffix.get(rest) {
                return Some(server);
            }
            match rest.split_once('.') {
                Some((_, tail)) if !tail.is_empty() => rest = tail,
                _ => return None,
            }
        }
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.by_suffix.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.by_suffix.is_empty()
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum Verdict {
    Available,
    Taken { raw: String },
    Unknown(Reason),
}

pub(crate) async fn query(
    resolver: &TokioResolver,
    pacer: &Pacer,
    server: &Server,
    domain: &str,
    timeout: Duration,
) -> Verdict {
    let Ok(lease) = pacer.acquire_patiently(&server.host, timeout).await else {
        return Verdict::Unknown(Reason::RateLimited);
    };

    let outcome = fetch_reply(resolver, &server.host, domain, timeout).await;
    drop(lease);

    match outcome {
        Err(reason) => {
            pacer.record_refusal(&server.host, &Refusal::Dropped).await;
            Verdict::Unknown(reason)
        }
        Ok(raw) => {
            match verdict::classify(&raw, &server.available_phrase, domain) {
                TextVerdict::Available => {
                    pacer.record_success(&server.host).await;
                    Verdict::Available
                }
                TextVerdict::Taken => {
                    pacer.record_success(&server.host).await;
                    Verdict::Taken { raw }
                }
                TextVerdict::Unknown(reason) => {
                    // @docgen Being told to slow down is pushback; any other unreadable answer must not pause the registry.
                    if matches!(reason, Reason::RateLimited | Reason::Blocked) {
                        pacer
                            .record_refusal(&server.host, &Refusal::Throttled { retry_after: None })
                            .await;
                    } else {
                        pacer.record_success(&server.host).await;
                    }
                    Verdict::Unknown(reason)
                }
            }
        }
    }
}

async fn fetch_reply(
    resolver: &TokioResolver,
    host: &str,
    domain: &str,
    timeout: Duration,
) -> std::result::Result<String, Reason> {
    let mut stream = connect(resolver, host, timeout).await?;
    let request = format_request(host, domain);

    tokio::time::timeout(timeout, stream.write_all(request.as_bytes()))
        .await
        .map_err(|_| Reason::TimedOut)?
        .map_err(|_| Reason::Unreachable)?;
    let _ = stream.flush().await;

    let mut buffer = Vec::new();
    let read = tokio::time::timeout(
        timeout,
        (&mut stream)
            .take(MAX_ANSWER_BYTES as u64)
            .read_to_end(&mut buffer),
    )
    .await
    .map_err(|_| Reason::TimedOut)?;
    read.map_err(|_| Reason::Unreachable)?;

    Ok(String::from_utf8_lossy(&buffer).replace("\r\n", "\n"))
}

/// @docgen Resolving and connecting by hand because handing a host and port to the platform resolver is unusable on some targets.
async fn connect(
    resolver: &TokioResolver,
    host: &str,
    timeout: Duration,
) -> std::result::Result<TcpStream, Reason> {
    // @docgen Adding an unbounded caller-supplied timeout to an Instant panics on overflow.
    let deadline = tokio::time::Instant::now()
        .checked_add(timeout)
        .unwrap_or_else(|| tokio::time::Instant::now() + Duration::from_secs(3600));

    let addresses = tokio::time::timeout_at(deadline, resolver.lookup_ip(host))
        .await
        .map_err(|_| Reason::TimedOut)?
        .map_err(|_| Reason::Unreachable)?;

    let mut last_reason = Reason::Unreachable;
    for ip in addresses.iter() {
        let address = SocketAddr::new(ip, 43);
        match tokio::time::timeout_at(deadline, TcpStream::connect(address)).await {
            Ok(Ok(stream)) => return Ok(stream),
            Ok(Err(_)) => last_reason = Reason::Unreachable,
            Err(_) => last_reason = Reason::TimedOut,
        }
    }
    Err(last_reason)
}

/// @docgen Some registries need their own query format on the wire: `domain <name>`, `-T dn <name>`, `<name>/e`.
fn format_request(host: &str, domain: &str) -> String {
    let host = host.to_lowercase();
    if host.contains("verisign-grs") || host.contains("crsnic") || host.contains("internic") {
        // @docgen A bare name makes these servers do a fuzzy match, so the exact registry record must be asked for.
        format!("domain {domain}\r\n")
    } else if host.contains("denic") {
        format!("-T dn {domain}\r\n")
    } else if host.contains("jprs") {
        format!("{domain}/e\r\n")
    } else if host.contains("dk-hostmaster") || host.contains("arnes.si") {
        format!("--show-handles {domain}\r\n")
    } else {
        format!("{domain}\r\n")
    }
}

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

    fn servers() -> Servers {
        Servers::bundled().expect("the bundled table must parse")
    }

    #[test]
    fn the_bundled_table_loads_and_is_substantial() {
        let servers = servers();
        assert!(!servers.is_empty());
        assert!(
            servers.len() > 500,
            "only {} servers, the table looks truncated",
            servers.len()
        );
    }

    #[test]
    fn bangladesh_is_covered_at_every_level_it_registers() {
        let servers = servers();
        for suffix in ["bd", "com.bd", "net.bd", "org.bd", "co.bd", "ai.bd"] {
            let server = servers
                .for_suffix(suffix)
                .unwrap_or_else(|| panic!(".{suffix} has no server"));
            assert_eq!(server.host, "whois.get.bd");
            assert!(
                !server.available_phrase.is_empty(),
                ".{suffix} has no available-name phrase"
            );
        }
    }

    #[test]
    fn the_big_extensions_are_covered() {
        let servers = servers();
        for suffix in ["com", "net", "org", "de", "in", "nl", "br"] {
            assert!(servers.for_suffix(suffix).is_some(), ".{suffix} missing");
        }
    }

    #[test]
    fn an_extension_that_retired_this_protocol_is_absent_by_design() {
        // @docgen Some registries withdrew their port-43 service, so absence from the table is correct rather than a gap.
        let servers = servers();
        assert!(
            servers.for_suffix("uk").is_none(),
            "the table should follow the published record rather than keep a dead host"
        );
    }

    #[test]
    fn a_multi_label_suffix_prefers_its_own_registry() {
        let servers = servers();
        // @docgen .bd registers at the third level and runs its own server there.
        let direct = servers.for_suffix("com.bd").map(|s| s.host.as_str());
        assert_eq!(direct, Some("whois.get.bd"));
    }

    #[test]
    fn an_unknown_suffix_falls_back_to_its_parent() {
        let servers = servers();
        let parent = servers.for_suffix("com").map(|s| s.host.clone());
        let child = servers
            .for_suffix("nothing-here.com")
            .map(|s| s.host.clone());
        assert_eq!(parent, child);
    }

    #[test]
    fn a_wholly_unknown_extension_has_no_server() {
        assert!(servers().for_suffix("zzzz-not-a-real-extension").is_none());
    }

    #[test]
    fn registries_that_need_a_special_request_get_one() {
        assert_eq!(
            format_request("whois.verisign-grs.com", "x.com"),
            "domain x.com\r\n"
        );
        assert_eq!(format_request("whois.denic.de", "x.de"), "-T dn x.de\r\n");
        assert_eq!(format_request("whois.jprs.jp", "x.jp"), "x.jp/e\r\n");
        assert_eq!(format_request("whois.get.bd", "x.bd"), "x.bd\r\n");
    }

    #[test]
    fn a_custom_table_overlays_the_bundled_one() {
        let mut servers = servers();
        let custom = Servers::parse(
            r#"{"servers":{"com":{"host":"whois.mine.example","free":"nothing here"}}}"#,
        )
        .unwrap();
        servers.merge(custom);
        assert_eq!(
            servers.for_suffix("com").map(|s| s.host.as_str()),
            Some("whois.mine.example")
        );
    }

    #[test]
    fn rubbish_is_refused() {
        assert!(Servers::parse("not json").is_err());
        assert!(Servers::parse(r#"{"servers":[]}"#).is_err());
    }
}