kcode-web-fetch 0.1.2

//! Bounded public-web fetching with SSRF protection and text extraction.

#![deny(missing_docs)]
#![forbid(unsafe_code)]

use std::{
    error::Error as StdError,
    fmt,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    time::{Duration, SystemTime},
};

use reqwest::{Client, StatusCode, Url, header};

/// Stable web-fetch failure classification.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ErrorKind {
    /// The URL or limits were invalid.
    InvalidInput,
    /// The URL did not resolve exclusively to public web addresses.
    UnsafeDestination,
    /// The request exceeded its configured timeout.
    Timeout,
    /// The destination could not be reached or returned malformed transport data.
    Transport,
    /// The destination returned a non-success HTTP status.
    HttpStatus,
    /// The destination returned a non-text content type.
    UnsupportedContent,
    /// The page contained no readable text.
    EmptyContent,
}

/// A sanitized URL, transport, HTTP, or content failure.
#[derive(Debug)]
pub struct Error {
    kind: ErrorKind,
    message: String,
    status: Option<u16>,
}

impl Error {
    fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
        Self {
            kind,
            message: message.into(),
            status: None,
        }
    }

    fn with_status(mut self, status: StatusCode) -> Self {
        self.status = Some(status.as_u16());
        self
    }

    /// Returns the stable failure classification.
    pub const fn kind(&self) -> ErrorKind {
        self.kind
    }

    /// Returns the destination HTTP status, when one caused the failure.
    pub const fn status(&self) -> Option<u16> {
        self.status
    }

    /// Returns the sanitized failure detail.
    pub fn message(&self) -> &str {
        &self.message
    }
}

impl fmt::Display for Error {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.message)
    }
}

impl StdError for Error {}

/// Result type returned by web-fetch operations.
pub type Result<T> = std::result::Result<T, Error>;

/// Resource and redirect limits applied to every fetch.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FetchLimits {
    /// Timeout applied to each HTTP request.
    pub request_timeout: Duration,
    /// Maximum downloaded response bytes.
    pub max_bytes: usize,
    /// Maximum returned Unicode characters.
    pub max_characters: usize,
    /// Maximum manually validated redirects.
    pub max_redirects: usize,
}

impl Default for FetchLimits {
    fn default() -> Self {
        Self {
            request_timeout: Duration::from_secs(90),
            max_bytes: 2_000_000,
            max_characters: 50_000,
            max_redirects: 5,
        }
    }
}

/// A fetched page converted into readable text.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FetchedPage {
    /// Final URL after validated redirects.
    pub url: String,
    /// HTML title, when one was present.
    pub title: Option<String>,
    /// Normalized response media type without parameters.
    pub content_type: String,
    /// Readable text content.
    pub content: String,
    /// Whether byte or character limits truncated the page.
    pub truncated: bool,
    /// Local time at which retrieval completed.
    pub retrieved_at: SystemTime,
}

/// Cloneable public-web fetcher.
#[derive(Clone, Debug)]
pub struct WebFetcher {
    limits: FetchLimits,
}

impl WebFetcher {
    /// Constructs a fetcher after validating nonzero limits.
    pub fn new(limits: FetchLimits) -> Result<Self> {
        if limits.request_timeout.is_zero() || limits.max_bytes == 0 || limits.max_characters == 0 {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                "web-fetch timeout and content limits must be greater than zero",
            ));
        }
        Ok(Self { limits })
    }

    /// Fetches one public HTTP(S) URL and returns bounded readable text.
    pub async fn fetch(&self, value: &str) -> Result<FetchedPage> {
        let requested = parse_public_web_url(value.trim())?;
        let fetched = fetch_page(&requested, &self.limits).await?;
        let raw = String::from_utf8_lossy(&fetched.body);
        let title = is_html_content(&fetched.content_type)
            .then(|| extract_html_title(&raw))
            .flatten();
        let readable = if is_html_content(&fetched.content_type) {
            html2text::from_read(raw.as_bytes(), 100).map_err(|_| {
                Error::new(
                    ErrorKind::Transport,
                    "the page could not be converted to readable text",
                )
            })?
        } else {
            raw.into_owned()
        };
        let (content, character_truncated) =
            truncate_characters(readable.trim(), self.limits.max_characters);
        if content.is_empty() {
            return Err(Error::new(
                ErrorKind::EmptyContent,
                "the page contained no readable text",
            ));
        }
        Ok(FetchedPage {
            url: fetched.url.to_string(),
            title,
            content_type: fetched.content_type,
            content,
            truncated: fetched.truncated || character_truncated,
            retrieved_at: SystemTime::now(),
        })
    }
}

impl Default for WebFetcher {
    fn default() -> Self {
        Self::new(FetchLimits::default()).expect("default web-fetch limits are valid")
    }
}

struct RawPage {
    url: Url,
    content_type: String,
    body: Vec<u8>,
    truncated: bool,
}

async fn fetch_page(url: &Url, limits: &FetchLimits) -> Result<RawPage> {
    let mut current = url.clone();
    for redirect_count in 0..=limits.max_redirects {
        let client = safe_client(&current, limits.request_timeout).await?;
        let mut response = client
            .get(current.clone())
            .header(
                header::ACCEPT,
                "text/html,application/xhtml+xml,text/plain,application/json;q=0.8",
            )
            .send()
            .await
            .map_err(transport_error)?;
        if response.status().is_redirection() {
            if redirect_count == limits.max_redirects {
                return Err(Error::new(
                    ErrorKind::Transport,
                    "the page exceeded the redirect limit",
                ));
            }
            let location = response
                .headers()
                .get(header::LOCATION)
                .and_then(|value| value.to_str().ok())
                .ok_or_else(|| {
                    Error::new(
                        ErrorKind::Transport,
                        "the page returned an invalid redirect",
                    )
                })?;
            current = parse_public_web_url(
                current
                    .join(location)
                    .map_err(|_| {
                        Error::new(
                            ErrorKind::Transport,
                            "the page returned an invalid redirect URL",
                        )
                    })?
                    .as_str(),
            )?;
            continue;
        }
        if !response.status().is_success() {
            let status = response.status();
            return Err(Error::new(
                ErrorKind::HttpStatus,
                format!("the page returned HTTP {status}"),
            )
            .with_status(status));
        }
        let content_type = response
            .headers()
            .get(header::CONTENT_TYPE)
            .and_then(|value| value.to_str().ok())
            .unwrap_or("application/octet-stream")
            .split(';')
            .next()
            .unwrap_or("application/octet-stream")
            .trim()
            .to_ascii_lowercase();
        if !is_supported_text_content(&content_type) {
            return Err(Error::new(
                ErrorKind::UnsupportedContent,
                format!("the page returned unsupported content type {content_type}"),
            ));
        }
        let mut body = Vec::new();
        let mut truncated = false;
        while let Some(chunk) = response.chunk().await.map_err(transport_error)? {
            let remaining = limits.max_bytes.saturating_sub(body.len());
            if chunk.len() > remaining {
                body.extend_from_slice(&chunk[..remaining]);
                truncated = true;
                break;
            }
            body.extend_from_slice(&chunk);
            if body.len() == limits.max_bytes {
                truncated = true;
                break;
            }
        }
        return Ok(RawPage {
            url: current,
            content_type,
            body,
            truncated,
        });
    }
    unreachable!("redirect loop returns or continues within its bound")
}

fn transport_error(error: reqwest::Error) -> Error {
    if error.is_timeout() {
        Error::new(ErrorKind::Timeout, "the page fetch timed out")
    } else {
        Error::new(ErrorKind::Transport, "the page could not be fetched")
    }
}

fn parse_public_web_url(value: &str) -> Result<Url> {
    if value.is_empty() || value.len() > 4_096 {
        return Err(Error::new(
            ErrorKind::InvalidInput,
            "URL must contain between 1 and 4096 bytes",
        ));
    }
    let url = Url::parse(value)
        .map_err(|_| Error::new(ErrorKind::InvalidInput, "URL must be absolute"))?;
    if !matches!(url.scheme(), "http" | "https") {
        return Err(Error::new(
            ErrorKind::InvalidInput,
            "URL must use HTTP or HTTPS",
        ));
    }
    if !url.username().is_empty() || url.password().is_some() {
        return Err(Error::new(
            ErrorKind::InvalidInput,
            "URL must not contain credentials",
        ));
    }
    let host = url
        .host_str()
        .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "URL must contain a host"))?;
    let lookup_host = host.trim_start_matches('[').trim_end_matches(']');
    let normalized = lookup_host.trim_end_matches('.').to_ascii_lowercase();
    if normalized == "localhost" || normalized.ends_with(".localhost") {
        return Err(unsafe_destination());
    }
    if !matches!(url.port_or_known_default(), Some(80 | 443)) {
        return Err(Error::new(
            ErrorKind::InvalidInput,
            "URL must use standard HTTP or HTTPS ports",
        ));
    }
    if lookup_host
        .parse::<IpAddr>()
        .is_ok_and(|address| !is_public_ip(address))
    {
        return Err(unsafe_destination());
    }
    Ok(url)
}

fn unsafe_destination() -> Error {
    Error::new(
        ErrorKind::UnsafeDestination,
        "URL does not refer to a public web destination",
    )
}

async fn safe_client(url: &Url, timeout: Duration) -> Result<Client> {
    let host = url
        .host_str()
        .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "URL must contain a host"))?;
    let lookup_name = host.trim_start_matches('[').trim_end_matches(']');
    let port = url
        .port_or_known_default()
        .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "URL must contain a valid port"))?;
    let addresses = tokio::net::lookup_host((lookup_name, port))
        .await
        .map_err(|_| Error::new(ErrorKind::Transport, "the page host could not be resolved"))?
        .collect::<Vec<SocketAddr>>();
    if addresses.is_empty() || addresses.iter().any(|address| !is_public_ip(address.ip())) {
        return Err(unsafe_destination());
    }
    Client::builder()
        .timeout(timeout)
        .redirect(reqwest::redirect::Policy::none())
        .retry(reqwest::retry::never())
        .referer(false)
        .no_proxy()
        .user_agent(concat!("kcode-web-fetch/", env!("CARGO_PKG_VERSION")))
        .resolve_to_addrs(lookup_name, &addresses)
        .build()
        .map_err(|_| {
            Error::new(
                ErrorKind::Transport,
                "the safe page-fetch client could not be created",
            )
        })
}

fn is_public_ip(address: IpAddr) -> bool {
    match address {
        IpAddr::V4(address) => is_public_ipv4(address),
        IpAddr::V6(address) => is_public_ipv6(address),
    }
}

fn is_public_ipv4(address: Ipv4Addr) -> bool {
    let [a, b, c, _] = address.octets();
    !(a == 0
        || a == 10
        || a == 127
        || (a == 100 && (64..=127).contains(&b))
        || (a == 169 && b == 254)
        || (a == 172 && (16..=31).contains(&b))
        || (a == 192 && b == 0 && c == 0)
        || (a == 192 && b == 0 && c == 2)
        || (a == 192 && b == 168)
        || (a == 198 && (b == 18 || b == 19))
        || (a == 198 && b == 51 && c == 100)
        || (a == 203 && b == 0 && c == 113)
        || a >= 224)
}

fn is_public_ipv6(address: Ipv6Addr) -> bool {
    if let Some(mapped) = address.to_ipv4_mapped() {
        return is_public_ipv4(mapped);
    }
    let segments = address.segments();
    !(address.is_unspecified()
        || address.is_loopback()
        || address.is_multicast()
        || segments[0] & 0xfe00 == 0xfc00
        || segments[0] & 0xffc0 == 0xfe80
        || (segments[0] == 0x2001 && segments[1] == 0x0db8))
}

fn is_html_content(content_type: &str) -> bool {
    matches!(content_type, "text/html" | "application/xhtml+xml")
}

fn is_supported_text_content(content_type: &str) -> bool {
    is_html_content(content_type)
        || content_type.starts_with("text/")
        || content_type == "application/json"
}

fn extract_html_title(html: &str) -> Option<String> {
    let lowercase = html.to_ascii_lowercase();
    let start = lowercase.find("<title")?;
    let content_start = lowercase[start..].find('>')? + start + 1;
    let end = lowercase[content_start..].find("</title>")? + content_start;
    let title = html[content_start..end]
        .split_whitespace()
        .collect::<Vec<_>>()
        .join(" ");
    (!title.is_empty()).then_some(title)
}

fn truncate_characters(value: &str, limit: usize) -> (String, bool) {
    let mut iter = value.char_indices();
    let Some((boundary, _)) = iter.nth(limit) else {
        return (value.to_owned(), false);
    };
    (value[..boundary].trim_end().to_owned(), true)
}

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

    #[test]
    fn private_and_credentialed_urls_are_rejected() {
        for value in [
            "http://127.0.0.1/",
            "http://[::1]/",
            "http://localhost/",
            "https://user:secret@example.com/",
            "file:///etc/passwd",
            "https://example.com:8443/",
        ] {
            assert!(parse_public_web_url(value).is_err(), "accepted {value}");
        }
        assert!(parse_public_web_url("https://example.com/path").is_ok());
    }

    #[test]
    fn readable_helpers_are_bounded() {
        assert_eq!(
            extract_html_title("<TITLE>  Example page </TITLE>"),
            Some("Example page".into())
        );
        assert_eq!(truncate_characters("éclair", 2), ("éc".into(), true));
        assert_eq!(truncate_characters("short", 20), ("short".into(), false));
    }
}