reserve-core 0.2.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 =
            crate::lookup::read_capped(path, crate::lookup::MAX_TABLE_BYTES).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> {
    // @docgen One deadline covers the whole lookup, because a per-step timeout let a slow server spend the budget three times over.
    let deadline = deadline_from(timeout);

    let mut stream = connect(resolver, host, deadline).await?;
    let request = format_request(host, domain);

    tokio::time::timeout_at(deadline, 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_at(
        deadline,
        (&mut stream)
            .take(MAX_ANSWER_BYTES as u64)
            .read_to_end(&mut buffer),
    )
    .await
    .map_err(|_| Reason::TimedOut)?;

    answer_from(buffer, read.is_err())
}

/// @docgen A registry that says its piece and then hangs up has answered. Throwing those bytes away turned "quota exceeded" into "could not connect", which then read as the whole machine being offline.
fn answer_from(buffer: Vec<u8>, cut_short: bool) -> std::result::Result<String, Reason> {
    if cut_short && buffer.is_empty() {
        return Err(Reason::Unreachable);
    }
    Ok(String::from_utf8_lossy(&buffer).replace("\r\n", "\n"))
}

/// @docgen Adding an unbounded caller-supplied timeout to an Instant panics on overflow.
fn deadline_from(timeout: Duration) -> tokio::time::Instant {
    tokio::time::Instant::now()
        .checked_add(timeout)
        .unwrap_or_else(|| tokio::time::Instant::now() + Duration::from_secs(3600))
}

/// @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,
    deadline: tokio::time::Instant,
) -> std::result::Result<TcpStream, Reason> {
    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")
    }

    /// @docgen Austria refuses a fast caller by printing "% Quota exceeded" and resetting the connection, which is an answer and must not read as a dead network.
    #[test]
    fn a_reply_cut_short_after_the_registry_spoke_is_still_the_reply() {
        let spoken = b"% Copyright NIC.AT\r\n%\r\n% Quota exceeded\r\n".to_vec();
        let answer = answer_from(spoken, true).expect("the registry answered before it hung up");
        assert!(
            answer.contains("Quota exceeded"),
            "the refusal survives the reset: {answer}"
        );
        assert!(!answer.contains('\r'), "line endings are still normalised");
    }

    #[test]
    fn a_connection_that_carried_nothing_is_still_a_failure_to_reach() {
        assert_eq!(
            answer_from(Vec::new(), true).unwrap_err(),
            Reason::Unreachable,
            "no bytes means no answer, whatever the socket did"
        );
    }

    #[test]
    fn a_clean_read_is_unaffected() {
        let body = b"Domain not found.\r\n".to_vec();
        assert_eq!(
            answer_from(body, false).expect("a clean read"),
            "Domain not found.\n"
        );
    }

    #[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());
    }
}