RustProxy/src/geoip/update_http.rs

use anyhow::{ensure, Context, Result};
use axum::body::Body;
use http::{header, Request, Uri};
use http_body_util::{BodyExt, Limited};
use hyper_rustls::HttpsConnectorBuilder;
use hyper_util::{
    client::legacy::{connect::HttpConnector, Client},
    rt::TokioExecutor,
};
use rustls::ClientConfig;
use sha2::{Digest, Sha256};
use std::{fs::OpenOptions, path::Path, sync::Arc, time::Duration};
use tokio::io::AsyncWriteExt;
use url::Url;

const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_REDIRECTS: usize = 5;

pub(super) async fn fetch_bytes(url: &str, max_bytes: usize) -> Result<Vec<u8>> {
    let client = client()?;
    let mut current = initial_url(url)?;
    for redirects in 0..=MAX_REDIRECTS {
        let response = request(&client, &current).await?;
        if response.status().is_redirection() {
            current = redirect(&current, &response, redirects)?;
            continue;
        }
        ensure_success_and_length(&response, max_bytes as u64)?;
        let body = tokio::time::timeout(
            REQUEST_TIMEOUT,
            Limited::new(response.into_body(), max_bytes).collect(),
        )
        .await
        .context("GeoIP response body timed out")?
        .map_err(|error| anyhow::anyhow!("GeoIP download is too large or incomplete: {error}"))?
        .to_bytes();
        return Ok(body.to_vec());
    }
    unreachable!("GeoIP redirect loop always returns or fails")
}

pub(super) async fn download(
    url: &str,
    target: &Path,
    max_bytes: u64,
    expected_digest: Option<&str>,
) -> Result<String> {
    let client = client()?;
    let mut current = initial_url(url)?;
    for redirects in 0..=MAX_REDIRECTS {
        let response = request(&client, &current).await?;
        if response.status().is_redirection() {
            current = redirect(&current, &response, redirects)?;
            continue;
        }
        ensure_success_and_length(&response, max_bytes)?;
        return stream_to_file(response.into_body(), target, max_bytes, expected_digest).await;
    }
    unreachable!("GeoIP redirect loop always returns or fails")
}

fn client() -> Result<Client<hyper_rustls::HttpsConnector<HttpConnector>, Body>> {
    let roots = crate::update_check::load_root_store();
    let tls =
        ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
            .with_safe_default_protocol_versions()?
            .with_root_certificates(roots)
            .with_no_client_auth();
    let mut http = HttpConnector::new();
    http.enforce_http(false);
    http.set_connect_timeout(Some(REQUEST_TIMEOUT));
    http.set_nodelay(true);
    let connector = HttpsConnectorBuilder::new()
        .with_tls_config(tls)
        .https_only()
        .enable_http1()
        .wrap_connector(http);
    Ok(Client::builder(TokioExecutor::new()).build(connector))
}

fn initial_url(url: &str) -> Result<Url> {
    let parsed = Url::parse(url).context("invalid GeoIP URL")?;
    ensure!(parsed.scheme() == "https", "GeoIP downloads require HTTPS");
    Ok(parsed)
}

async fn request(
    client: &Client<hyper_rustls::HttpsConnector<HttpConnector>, Body>,
    current: &Url,
) -> Result<hyper::Response<hyper::body::Incoming>> {
    let uri: Uri = current.as_str().parse().context("invalid GeoIP URI")?;
    let request = Request::builder()
        .method(http::Method::GET)
        .uri(uri)
        .header(header::ACCEPT, "application/octet-stream, text/plain;q=0.9")
        .header(
            header::USER_AGENT,
            concat!("RustProxy/", env!("CARGO_PKG_VERSION"), " GeoIP-update"),
        )
        .body(Body::empty())?;
    tokio::time::timeout(REQUEST_TIMEOUT, client.request(request))
        .await
        .context("GeoIP request timed out")?
        .with_context(|| format!("GeoIP request failed for {current}"))
}

fn redirect(
    current: &Url,
    response: &hyper::Response<hyper::body::Incoming>,
    redirects: usize,
) -> Result<Url> {
    ensure!(redirects < MAX_REDIRECTS, "GeoIP redirect limit exceeded");
    let location = response
        .headers()
        .get(header::LOCATION)
        .context("GeoIP redirect is missing Location")?
        .to_str()
        .context("GeoIP redirect Location is not UTF-8")?;
    let target = crate::update_check::redirect_target(current, location)?;
    ensure!(
        target.scheme() == "https",
        "GeoIP redirect must remain HTTPS"
    );
    Ok(target)
}

fn ensure_success_and_length(
    response: &hyper::Response<hyper::body::Incoming>,
    max_bytes: u64,
) -> Result<()> {
    ensure!(
        response.status().is_success(),
        "GeoIP request returned HTTP {}",
        response.status()
    );
    if let Some(length) = response.headers().get(header::CONTENT_LENGTH) {
        let length = length
            .to_str()?
            .parse::<u64>()
            .context("invalid GeoIP Content-Length")?;
        ensure!(
            length <= max_bytes,
            "GeoIP download is larger than the configured safety limit"
        );
    }
    Ok(())
}

async fn stream_to_file(
    mut body: hyper::body::Incoming,
    target: &Path,
    max_bytes: u64,
    expected_digest: Option<&str>,
) -> Result<String> {
    let mut options = OpenOptions::new();
    options.write(true).create_new(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt;
        options.mode(0o600).custom_flags(libc::O_NOFOLLOW);
    }
    let file = options
        .open(target)
        .with_context(|| format!("create GeoIP staging file {}", target.display()))?;
    let mut output = tokio::fs::File::from_std(file);
    let mut digest = Sha256::new();
    let mut written = 0u64;

    while let Some(frame) = tokio::time::timeout(REQUEST_TIMEOUT, body.frame())
        .await
        .context("GeoIP response body timed out")?
    {
        let frame = frame.context("read GeoIP response body")?;
        if let Some(data) = frame.data_ref() {
            written = written
                .checked_add(data.len() as u64)
                .context("GeoIP download size overflow")?;
            ensure!(
                written <= max_bytes,
                "GeoIP download is larger than the configured safety limit"
            );
            digest.update(data);
            output.write_all(data).await?;
        }
    }
    output.flush().await?;
    output.sync_all().await?;
    let actual = format!("{:x}", digest.finalize());
    if let Some(expected) = expected_digest {
        ensure!(
            actual == expected,
            "GeoIP SHA-256 mismatch for {}",
            target.file_name().unwrap_or_default().to_string_lossy()
        );
    }
    Ok(actual)
}