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> { let client = client()?; let mut current = initial_url(url)?; for redirects in 0..=MAX_REDIRECTS { let response = request(&client, ¤t).await?; if response.status().is_redirection() { current = redirect(¤t, &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 { let client = client()?; let mut current = initial_url(url)?; for redirects in 0..=MAX_REDIRECTS { let response = request(&client, ¤t).await?; if response.status().is_redirection() { current = redirect(¤t, &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, 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 { let parsed = Url::parse(url).context("invalid GeoIP URL")?; ensure!(parsed.scheme() == "https", "GeoIP downloads require HTTPS"); Ok(parsed) } async fn request( client: &Client, Body>, current: &Url, ) -> Result> { 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, redirects: usize, ) -> Result { 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, 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::() .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 { 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) }