kibble 0.1.0

chew through any source into clean datasets — a fast ingestion, RAG & fine-tuning toolkit
Documentation
use std::net::{IpAddr, ToSocketAddrs};

/// SSRF guard shared by `serve`, `crawl`, and `websearch` — pure and dep-light,
/// so it lives in `net` (lean) rather than `serve` (full-only, owns axum).
pub fn url_host(url: &str) -> Option<String> {
    let after = url.split_once("://").map(|(_, r)| r).unwrap_or(url);
    let host = after.split(['/', '?', '#']).next().unwrap_or("");
    let host = host.rsplit('@').next().unwrap_or(host);
    // bracketed IPv6: [::1]:port → ::1
    if let Some(end) = host.strip_prefix('[').and_then(|h| h.split_once(']')).map(|(h, _)| h) {
        return Some(end.to_lowercase());
    }
    // non-bracketed host[:port] → drop the port
    let host = host.split(':').next().unwrap_or(host);
    if host.is_empty() { None } else { Some(host.to_lowercase()) }
}

fn ip_is_blocked(ip: &IpAddr) -> bool {
    match ip {
        IpAddr::V4(v4) => {
            v4.is_loopback() || v4.is_private() || v4.is_link_local()
                || v4.is_unspecified() || v4.is_multicast() || v4.is_broadcast()
        }
        IpAddr::V6(v6) => {
            // IPv4-mapped (::ffff:a.b.c.d) targets a v4 address — classify it as v4.
            if let Some(v4) = v6.to_ipv4_mapped() {
                return ip_is_blocked(&IpAddr::V4(v4));
            }
            v6.is_loopback() || v6.is_unspecified() || v6.is_multicast()
                // unique-local fc00::/7 and link-local fe80::/10
                || (v6.segments()[0] & 0xfe00) == 0xfc00
                || (v6.segments()[0] & 0xffc0) == 0xfe80
        }
    }
}

pub fn url_is_allowed(url: &str, allow_hosts: &[String]) -> Result<(), String> {
    let host = url_host(url).ok_or_else(|| format!("no host in url: {url}"))?;
    if allow_hosts.iter().any(|h| h.to_lowercase() == host) {
        return Ok(());
    }
    // literal IP?
    if let Ok(ip) = host.parse::<IpAddr>() {
        return if ip_is_blocked(&ip) { Err(format!("blocked ip: {host}")) } else { Ok(()) };
    }
    // resolve hostname; reject if any resolved ip is blocked
    let addrs = (host.as_str(), 80u16)
        .to_socket_addrs()
        .map_err(|e| format!("cannot resolve {host}: {e}"))?;
    let mut any = false;
    for a in addrs {
        any = true;
        if ip_is_blocked(&a.ip()) {
            return Err(format!("blocked resolved ip {} for {host}", a.ip()));
        }
    }
    if !any {
        return Err(format!("no addresses for {host}"));
    }
    Ok(())
}

/// Resolve the effective proxy: config wins, else env in priority ALL_ > HTTPS_ > HTTP_.
pub fn resolve_proxy(
    config_proxy: Option<&str>,
    lookup: impl Fn(&str) -> Option<String>,
) -> Option<String> {
    if let Some(p) = config_proxy {
        if !p.is_empty() {
            return Some(p.to_string());
        }
    }
    for key in ["ALL_PROXY", "HTTPS_PROXY", "HTTP_PROXY"] {
        if let Some(v) = lookup(key) {
            if !v.is_empty() {
                return Some(v);
            }
        }
    }
    None
}

/// Default per-request timeout. Without it a stuck/slow backend (e.g. a reasoning LLM that
/// stalls mid-generation) hangs the whole run forever; a bounded timeout lets callers fail-soft.
pub const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);

/// Build a reqwest client, routing through `proxy` (http/https/socks5 URL) when set.
pub fn build_client(proxy: Option<&str>) -> reqwest::Result<reqwest::Client> {
    let mut builder = reqwest::Client::builder().timeout(REQUEST_TIMEOUT);
    if let Some(p) = proxy {
        builder = builder.proxy(reqwest::Proxy::all(p)?);
    }
    builder.build()
}

/// Build a reqwest client like [`build_client`], but with a custom `User-Agent` header
/// so both `robots.txt` and page fetches identify themselves correctly.
pub fn build_client_ua(proxy: Option<&str>, user_agent: &str) -> reqwest::Result<reqwest::Client> {
    let mut builder = reqwest::Client::builder().user_agent(user_agent).timeout(REQUEST_TIMEOUT);
    if let Some(p) = proxy {
        builder = builder.proxy(reqwest::Proxy::all(p)?);
    }
    builder.build()
}

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

    fn env<'a>(map: &'a HashMap<&'a str, &'a str>) -> impl Fn(&str) -> Option<String> + 'a {
        move |k| map.get(k).map(|v| v.to_string())
    }

    #[test]
    fn config_proxy_wins_over_env() {
        let mut m = HashMap::new();
        m.insert("ALL_PROXY", "socks5://env:1");
        let got = resolve_proxy(Some("http://cfg:2"), env(&m));
        assert_eq!(got.as_deref(), Some("http://cfg:2"));
    }

    #[test]
    fn falls_back_to_env_in_priority_order() {
        let mut m = HashMap::new();
        m.insert("HTTP_PROXY", "http://low:1");
        m.insert("HTTPS_PROXY", "http://mid:1");
        assert_eq!(resolve_proxy(None, env(&m)).as_deref(), Some("http://mid:1"));
        m.insert("ALL_PROXY", "socks5://top:1");
        assert_eq!(resolve_proxy(None, env(&m)).as_deref(), Some("socks5://top:1"));
    }

    #[test]
    fn none_when_unset() {
        let m: HashMap<&str, &str> = HashMap::new();
        assert_eq!(resolve_proxy(None, env(&m)), None);
    }

    #[test]
    fn builds_client_with_and_without_proxy() {
        assert!(build_client(None).is_ok());
        assert!(build_client(Some("socks5://127.0.0.1:9050")).is_ok());
    }

    #[test]
    fn builds_client_ua_with_and_without_proxy() {
        assert!(build_client_ua(None, "kibble-crawler/1.0").is_ok());
        assert!(build_client_ua(Some("socks5://127.0.0.1:9050"), "kibble-crawler/1.0").is_ok());
    }

    #[test]
    fn rejects_loopback_and_private_literals() {
        assert!(url_is_allowed("http://127.0.0.1/x", &[]).is_err());
        assert!(url_is_allowed("http://10.0.0.5/x", &[]).is_err());
        assert!(url_is_allowed("http://169.254.169.254/latest/meta-data/", &[]).is_err());
        assert!(url_is_allowed("http://[::1]/x", &[]).is_err());
    }

    #[test]
    fn allow_hosts_bypasses_guard() {
        assert!(url_is_allowed("http://127.0.0.1:8080/x", &["127.0.0.1".to_string()]).is_ok());
    }

    #[test]
    fn rejects_unresolvable_host() {
        assert!(url_is_allowed("http://no.such.host.invalid./x", &[]).is_err());
    }

    #[test]
    fn rejects_ipv4_mapped_ipv6_loopback() {
        // ::ffff:127.0.0.1 must be blocked (maps to 127.0.0.1)
        assert!(url_is_allowed("http://[::ffff:127.0.0.1]/x", &[]).is_err());
        // ::ffff:10.0.0.1 (mapped private) blocked too
        assert!(url_is_allowed("http://[::ffff:10.0.0.1]/x", &[]).is_err());
    }

    #[test]
    fn userinfo_at_host_uses_connect_target() {
        // http://evil.com@127.0.0.1/ connects to 127.0.0.1 (after @), which must be blocked
        assert!(url_is_allowed("http://evil.com@127.0.0.1/x", &[]).is_err());
    }
}