hoardom/src/tlds.rs

use std::collections::HashMap;
use std::sync::OnceLock;

/// whois server overrides from Lists.toml ("io:whois.nic.io" style entries)
#[derive(Debug, Clone, Default)]
pub struct WhoisOverrides {
    map: HashMap<String, String>,
}

impl WhoisOverrides {
    pub fn get_server(&self, tld: &str) -> Option<&str> {
        self.map.get(&tld.to_lowercase()).map(|s| s.as_str())
    }
}

/// a named TLD list from Lists.toml
struct NamedList {
    name: String,
    tlds: Vec<String>,
}

struct ParsedLists {
    lists: Vec<NamedList>,
    whois_overrides: WhoisOverrides,
}

/// parse a single entry: "tld" or "tld:whois_server"
fn parse_entry(entry: &str) -> (String, Option<String>) {
    if let Some(pos) = entry.find(':') {
        (entry[..pos].to_string(), Some(entry[pos + 1..].to_string()))
    } else {
        (entry.to_string(), None)
    }
}

/// parse entries, pull out TLD names and whois overrides
fn parse_list(entries: &[toml::Value], overrides: &mut HashMap<String, String>) -> Vec<String> {
    entries
        .iter()
        .filter_map(|v| v.as_str())
        .map(|entry| {
            let (tld, server) = parse_entry(entry);
            if let Some(s) = server {
                overrides.insert(tld.to_lowercase(), s);
            }
            tld
        })
        .collect()
}

static PARSED_LISTS: OnceLock<ParsedLists> = OnceLock::new();

fn parsed_lists() -> &'static ParsedLists {
    PARSED_LISTS.get_or_init(|| {
        let raw: toml::Value = toml::from_str(include_str!("../Lists.toml"))
            .expect("Lists.toml must be valid TOML");

        let table = raw.as_table().expect("Lists.toml must be a TOML table");

        // Build list names in the order build.rs discovered them
        let ordered_names: Vec<&str> = env!("HOARDOM_LIST_NAMES").split(',').collect();

        let mut overrides = HashMap::new();
        let mut lists = Vec::new();

        for name in &ordered_names {
            if let Some(toml::Value::Array(arr)) = table.get(*name) {
                let tlds = parse_list(arr, &mut overrides);
                lists.push(NamedList {
                    name: name.to_string(),
                    tlds,
                });
            }
        }

        ParsedLists {
            lists,
            whois_overrides: WhoisOverrides { map: overrides },
        }
    })
}

/// list names from Lists.toml, in order
pub fn list_names() -> Vec<&'static str> {
    parsed_lists()
        .lists
        .iter()
        .map(|l| l.name.as_str())
        .collect()
}

/// first list name (the default)
pub fn default_list_name() -> &'static str {
    list_names().first().copied().unwrap_or("standard")
}

/// get TLDs for a list name (case insensitive), None if not found
pub fn get_tlds(name: &str) -> Option<Vec<&'static str>> {
    let lower = name.to_lowercase();
    parsed_lists()
        .lists
        .iter()
        .find(|l| l.name == lower)
        .map(|l| l.tlds.iter().map(String::as_str).collect())
}

/// get TLDs for a list name, falls back to default if not found
pub fn get_tlds_or_default(name: &str) -> Vec<&'static str> {
    get_tlds(name).unwrap_or_else(|| get_tlds(default_list_name()).unwrap_or_default())
}

/// the builtin whois overrides from Lists.toml
pub fn whois_overrides() -> &'static WhoisOverrides {
    &parsed_lists().whois_overrides
}

pub fn apply_top_tlds(tlds: Vec<&'static str>, top: &[String]) -> Vec<&'static str> {
    let mut result: Vec<&'static str> = Vec::with_capacity(tlds.len());
    // first add the top ones in the order specified
    for t in top {
        let lower = t.to_lowercase();
        if let Some(&found) = tlds.iter().find(|&&tld| tld.to_lowercase() == lower) {
            if !result.contains(&found) {
                result.push(found);
            }
        }
    }
    // then add the rest
    for tld in &tlds {
        if !result.contains(tld) {
            result.push(tld);
        }
    }
    result
}

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

    #[test]
    fn test_parse_entry_bare() {
        let (tld, server) = parse_entry("com");
        assert_eq!(tld, "com");
        assert_eq!(server, None);
    }

    #[test]
    fn test_parse_entry_with_override() {
        let (tld, server) = parse_entry("io:whois.nic.io");
        assert_eq!(tld, "io");
        assert_eq!(server, Some("whois.nic.io".to_string()));
    }

    #[test]
    fn test_whois_overrides_populated() {
        let overrides = whois_overrides();
        // io should have an override since our Lists.toml has "io:whois.nic.io"
        assert!(overrides.get_server("io").is_some());
        // com should not (it has RDAP)
        assert!(overrides.get_server("com").is_none());
    }

    #[test]
    fn test_top_tlds_reorder() {
        let tlds = vec!["com", "net", "org", "ch", "de"];
        let top = vec!["ch".to_string(), "de".to_string()];
        let result = apply_top_tlds(tlds, &top);
        assert_eq!(result, vec!["ch", "de", "com", "net", "org"]);
    }

    #[test]
    fn test_top_tlds_missing_ignored() {
        let tlds = vec!["com", "net"];
        let top = vec!["swiss".to_string()];
        let result = apply_top_tlds(tlds, &top);
        assert_eq!(result, vec!["com", "net"]);
    }
}