reserve-core 0.5.1

//! 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 crate::lookup::resolve::Resolvers;
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>,
    /// @docgen Only the entries a person supplied are their own choice; the bundled table is ours and stays guarded.
    supplied: std::collections::HashSet<String>,
}

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(),
            supplied: std::collections::HashSet::new(),
        })
    }

    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.supplied.extend(other.by_suffix.keys().cloned());
        self.by_suffix.extend(other.by_suffix);
    }

    /// @docgen Lifting the address guard for the whole table would hand a hostile bundled entry the same trust the user meant for their own.
    #[must_use]
    pub fn was_supplied(&self, suffix: &str) -> bool {
        let suffix = suffix.trim_matches('.').to_lowercase();
        let mut rest = suffix.as_str();
        loop {
            if self.by_suffix.contains_key(rest) {
                return self.supplied.contains(rest);
            }
            match rest.split_once('.') {
                Some((_, tail)) if !tail.is_empty() => rest = tail,
                _ => return false,
            }
        }
    }

    /// @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 { raw: String },
    Taken { raw: String },
    Unknown(Reason),
}

/// @docgen A host IANA or the bundled table named must resolve to a public address; one the user listed is their own choice.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum HostGuard {
    Enforce,
    Trusted,
}

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

    let outcome = ask(resolver, &server.host, domain, timeout, guard).await;
    drop(permit);

    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 { raw }
                }
                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;
                    }
                    // @docgen An answer nobody could read is not a success either, and crediting one ramped the rate up against a server that was struggling.
                    Verdict::Unknown(reason)
                }
            }
        }
    }
}

async fn ask(
    resolver: &Resolvers,
    host: &str,
    domain: &str,
    timeout: Duration,
    guard: HostGuard,
) -> 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, guard).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;

    // @docgen A reply that reached the ceiling is missing its end, and its end is where a refusal is usually written.
    if buffer.len() >= MAX_ANSWER_BYTES {
        return Err(Reason::Malformed {
            detail: "the reply was longer than this tool will read".to_owned(),
        });
    }

    match read {
        Ok(outcome) => answer_from(buffer, outcome.is_err()),
        // @docgen A server that says its piece then holds the socket open has answered, so the bytes it sent must not be dropped.
        // @docgen One cut off mid-line has not, and its end is where a refusal is usually written.
        Err(_) if buffer.last() == Some(&b'\n') => answer_from(buffer, false),
        Err(_) if !buffer.is_empty() => Err(Reason::Malformed {
            detail: "the reply stopped in the middle of a line".to_owned(),
        }),
        Err(_) => Err(Reason::TimedOut),
    }
}

/// @docgen A registry that speaks then hangs up has answered; dropping those bytes turned "quota exceeded" into a dead machine.
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: &Resolvers,
    host: &str,
    deadline: tokio::time::Instant,
    guard: HostGuard,
) -> 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() {
        // @docgen A public name can answer with an internal address, so the resolved address decides, not the name.
        if guard == HostGuard::Enforce && !crate::lookup::registry::is_public_ip(ip) {
            last_reason = Reason::Blocked;
            continue;
        }
        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, which is an answer, not 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(|server| server.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(|server| server.host.clone());
        let child = servers
            .for_suffix("nothing-here.com")
            .map(|server| server.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","available_phrase":"nothing here"}}}"#,
        )
        .unwrap();
        servers.merge(custom);
        assert_eq!(
            servers.for_suffix("com").map(|server| server.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());
    }

    #[test]
    fn only_the_entry_a_person_supplied_is_treated_as_their_own_choice() {
        let mut servers = Servers::bundled().expect("the bundled table parses");
        let bundled_suffix = servers
            .by_suffix
            .keys()
            .next()
            .cloned()
            .expect("the bundled table holds entries");
        assert!(
            !servers.was_supplied(&bundled_suffix),
            "a bundled host is ours and stays guarded"
        );

        let mine =
            Servers::parse(r#"{"servers": {"example-zone": {"host": "whois.example.test"}}}"#)
                .expect("a supplied table parses");
        servers.merge(mine);

        assert!(
            servers.was_supplied("example-zone"),
            "the entry the user wrote is their own choice"
        );
        assert!(
            !servers.was_supplied(&bundled_suffix),
            "supplying one entry must not lift the guard for the whole table"
        );
    }
}