v0.28.0

src/geoip/update_http.rs

@@ -0,0 +1,200 @@
+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)
+}