Simplify download API abstraction

djc · rami3l · commit 0788cd6d2023 · 2025-04-07T15:52:23.000Z
diff --git a/src/cli/self_update.rs b/src/cli/self_update.rs
@@ -48,6 +48,7 @@ use same_file::Handle;
 use serde::{Deserialize, Serialize};
 use tracing::{error, info, trace, warn};
 
+use crate::download::download_file;
 use crate::{
     DUP_TOOLS, TOOLS,
     cli::{
@@ -1134,7 +1135,7 @@ pub(crate) async fn prepare_update(process: &Process) -> Result<Option<PathBuf>>
 
     // Download new version
     info!("downloading self-update");
-    utils::download_file(&download_url, &setup_path, None, &|_| (), process).await?;
+    download_file(&download_url, &setup_path, None, &|_| (), process).await?;
 
     // Mark as executable
     utils::make_executable(&setup_path)?;
@@ -1153,7 +1154,7 @@ async fn get_available_rustup_version(process: &Process) -> Result<String> {
     let release_file_url = format!("{update_root}/release-stable.toml");
     let release_file_url = utils::parse_url(&release_file_url)?;
     let release_file = tempdir.path().join("release-stable.toml");
-    utils::download_file(&release_file_url, &release_file, None, &|_| (), process).await?;
+    download_file(&release_file_url, &release_file, None, &|_| (), process).await?;
     let release_toml_str = utils::read_file("rustup release", &release_file)?;
     let release_toml = toml::from_str::<RustupManifest>(&release_toml_str)
         .context("unable to parse rustup release file")?;
diff --git a/src/cli/self_update/windows.rs b/src/cli/self_update/windows.rs
@@ -22,6 +22,7 @@ use super::common;
 use super::{InstallOpts, install_bins, report_error};
 use crate::cli::{download_tracker::DownloadTracker, markdown::md};
 use crate::dist::TargetTriple;
+use crate::download::download_file;
 use crate::process::{Process, terminalsource::ColorableTerminal};
 use crate::utils::{self, Notification};
 
@@ -276,7 +277,7 @@ pub(crate) async fn try_install_msvc(
     download_tracker.lock().unwrap().download_finished();
 
     info!("downloading Visual Studio installer");
-    utils::download_file(
+    download_file(
         &visual_studio_url,
         &visual_studio,
         None,
diff --git a/src/dist/download.rs b/src/dist/download.rs
@@ -8,6 +8,8 @@ use url::Url;
 
 use crate::dist::notifications::*;
 use crate::dist::temp;
+use crate::download::download_file;
+use crate::download::download_file_with_resume;
 use crate::errors::*;
 use crate::process::Process;
 use crate::utils;
@@ -73,7 +75,7 @@ impl<'a> DownloadCfg<'a> {
 
         let mut hasher = Sha256::new();
 
-        if let Err(e) = utils::download_file_with_resume(
+        if let Err(e) = download_file_with_resume(
             url,
             &partial_file_path,
             Some(&mut hasher),
@@ -134,7 +136,7 @@ impl<'a> DownloadCfg<'a> {
         let hash_url = utils::parse_url(&(url.to_owned() + ".sha256"))?;
         let hash_file = self.tmp_cx.new_file()?;
 
-        utils::download_file(
+        download_file(
             &hash_url,
             &hash_file,
             None,
@@ -179,7 +181,7 @@ impl<'a> DownloadCfg<'a> {
         let file = self.tmp_cx.new_file_with_ext("", ext)?;
 
         let mut hasher = Sha256::new();
-        utils::download_file(
+        download_file(
             &url,
             &file,
             Some(&mut hasher),
diff --git a/src/dist/manifestation/tests.rs b/src/dist/manifestation/tests.rs
@@ -23,6 +23,7 @@ use crate::{
         prefix::InstallPrefix,
         temp,
     },
+    download::download_file,
     errors::RustupError,
     process::TestProcess,
     test::{
@@ -495,7 +496,7 @@ impl TestContext {
         // Download the dist manifest and place it into the installation prefix
         let manifest_url = make_manifest_url(&self.url, &self.toolchain)?;
         let manifest_file = self.tmp_cx.new_file()?;
-        utils::download_file(&manifest_url, &manifest_file, None, &|_| {}, dl_cfg.process).await?;
+        download_file(&manifest_url, &manifest_file, None, &|_| {}, dl_cfg.process).await?;
         let manifest_str = utils::read_file("manifest", &manifest_file)?;
         let manifest = Manifest::parse(&manifest_str)?;
 
diff --git a/src/download/mod.rs b/src/download/mod.rs
@@ -10,12 +10,192 @@ use std::path::Path;
 ))]
 use anyhow::anyhow;
 use anyhow::{Context, Result};
+use sha2::Sha256;
 use thiserror::Error;
+#[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
+use tracing::info;
 use url::Url;
 
+use crate::{errors::RustupError, process::Process, utils::Notification};
+
 #[cfg(test)]
 mod tests;
 
+pub(crate) async fn download_file(
+    url: &Url,
+    path: &Path,
+    hasher: Option<&mut Sha256>,
+    notify_handler: &dyn Fn(Notification<'_>),
+    process: &Process,
+) -> Result<()> {
+    download_file_with_resume(url, path, hasher, false, &notify_handler, process).await
+}
+
+pub(crate) async fn download_file_with_resume(
+    url: &Url,
+    path: &Path,
+    hasher: Option<&mut Sha256>,
+    resume_from_partial: bool,
+    notify_handler: &dyn Fn(Notification<'_>),
+    process: &Process,
+) -> Result<()> {
+    use crate::download::DownloadError as DEK;
+    match download_file_(
+        url,
+        path,
+        hasher,
+        resume_from_partial,
+        notify_handler,
+        process,
+    )
+    .await
+    {
+        Ok(_) => Ok(()),
+        Err(e) => {
+            if e.downcast_ref::<std::io::Error>().is_some() {
+                return Err(e);
+            }
+            let is_client_error = match e.downcast_ref::<DEK>() {
+                // Specifically treat the bad partial range error as not our
+                // fault in case it was something odd which happened.
+                Some(DEK::HttpStatus(416)) => false,
+                Some(DEK::HttpStatus(400..=499)) | Some(DEK::FileNotFound) => true,
+                _ => false,
+            };
+            Err(e).with_context(|| {
+                if is_client_error {
+                    RustupError::DownloadNotExists {
+                        url: url.clone(),
+                        path: path.to_path_buf(),
+                    }
+                } else {
+                    RustupError::DownloadingFile {
+                        url: url.clone(),
+                        path: path.to_path_buf(),
+                    }
+                }
+            })
+        }
+    }
+}
+
+async fn download_file_(
+    url: &Url,
+    path: &Path,
+    hasher: Option<&mut Sha256>,
+    resume_from_partial: bool,
+    notify_handler: &dyn Fn(Notification<'_>),
+    process: &Process,
+) -> Result<()> {
+    #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
+    use crate::download::{Backend, Event, TlsBackend};
+    use sha2::Digest;
+    use std::cell::RefCell;
+
+    notify_handler(Notification::DownloadingFile(url, path));
+
+    let hasher = RefCell::new(hasher);
+
+    // This callback will write the download to disk and optionally
+    // hash the contents, then forward the notification up the stack
+    let callback: &dyn Fn(Event<'_>) -> anyhow::Result<()> = &|msg| {
+        if let Event::DownloadDataReceived(data) = msg {
+            if let Some(h) = hasher.borrow_mut().as_mut() {
+                h.update(data);
+            }
+        }
+
+        match msg {
+            Event::DownloadContentLengthReceived(len) => {
+                notify_handler(Notification::DownloadContentLengthReceived(len));
+            }
+            Event::DownloadDataReceived(data) => {
+                notify_handler(Notification::DownloadDataReceived(data));
+            }
+            Event::ResumingPartialDownload => {
+                notify_handler(Notification::ResumingPartialDownload);
+            }
+        }
+
+        Ok(())
+    };
+
+    // Download the file
+
+    // Keep the curl env var around for a bit
+    let use_curl_backend = process.var_os("RUSTUP_USE_CURL").map(|it| it != "0");
+    let use_rustls = process.var_os("RUSTUP_USE_RUSTLS").map(|it| it != "0");
+
+    let backend = match (use_curl_backend, use_rustls) {
+        // If environment specifies a backend that's unavailable, error out
+        #[cfg(not(feature = "reqwest-rustls-tls"))]
+        (_, Some(true)) => {
+            return Err(anyhow!(
+                "RUSTUP_USE_RUSTLS is set, but this rustup distribution was not built with the reqwest-rustls-tls feature"
+            ));
+        }
+        #[cfg(not(feature = "reqwest-native-tls"))]
+        (_, Some(false)) => {
+            return Err(anyhow!(
+                "RUSTUP_USE_RUSTLS is set to false, but this rustup distribution was not built with the reqwest-native-tls feature"
+            ));
+        }
+        #[cfg(not(feature = "curl-backend"))]
+        (Some(true), _) => {
+            return Err(anyhow!(
+                "RUSTUP_USE_CURL is set, but this rustup distribution was not built with the curl-backend feature"
+            ));
+        }
+
+        // Positive selections, from least preferred to most preferred
+        #[cfg(feature = "curl-backend")]
+        (Some(true), None) => Backend::Curl,
+        #[cfg(feature = "reqwest-native-tls")]
+        (_, Some(false)) => {
+            if use_curl_backend == Some(true) {
+                info!(
+                    "RUSTUP_USE_CURL is set and RUSTUP_USE_RUSTLS is set to off, using reqwest with native-tls"
+                );
+            }
+            Backend::Reqwest(TlsBackend::NativeTls)
+        }
+        #[cfg(feature = "reqwest-rustls-tls")]
+        _ => {
+            if use_curl_backend == Some(true) {
+                info!(
+                    "both RUSTUP_USE_CURL and RUSTUP_USE_RUSTLS are set, using reqwest with rustls"
+                );
+            }
+            Backend::Reqwest(TlsBackend::Rustls)
+        }
+
+        // Falling back if only one backend is available
+        #[cfg(all(not(feature = "reqwest-rustls-tls"), feature = "reqwest-native-tls"))]
+        _ => Backend::Reqwest(TlsBackend::NativeTls),
+        #[cfg(all(
+            not(feature = "reqwest-rustls-tls"),
+            not(feature = "reqwest-native-tls"),
+            feature = "curl-backend"
+        ))]
+        _ => Backend::Curl,
+    };
+
+    notify_handler(match backend {
+        #[cfg(feature = "curl-backend")]
+        Backend::Curl => Notification::UsingCurl,
+        #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
+        Backend::Reqwest(_) => Notification::UsingReqwest,
+    });
+
+    let res = backend
+        .download_to_path(url, path, resume_from_partial, Some(callback))
+        .await;
+
+    notify_handler(Notification::DownloadFinished);
+
+    res
+}
+
 /// User agent header value for HTTP request.
 /// See: https://github.com/rust-lang/rustup/issues/2860.
 #[cfg(feature = "curl-backend")]
@@ -33,15 +213,15 @@ const REQWEST_RUSTLS_TLS_USER_AGENT: &str =
     concat!("rustup/", env!("CARGO_PKG_VERSION"), " (reqwest; rustls)");
 
 #[derive(Debug, Copy, Clone)]
-pub enum Backend {
+enum Backend {
     #[cfg(feature = "curl-backend")]
     Curl,
     #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
     Reqwest(TlsBackend),
 }
 
 impl Backend {
-    pub async fn download_to_path(
+    async fn download_to_path(
         self,
         url: &Url,
         path: &Path,
@@ -175,7 +355,7 @@ impl Backend {
 
 #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
 #[derive(Debug, Copy, Clone)]
-pub enum TlsBackend {
+enum TlsBackend {
     #[cfg(feature = "reqwest-rustls-tls")]
     Rustls,
     #[cfg(feature = "reqwest-native-tls")]
@@ -202,7 +382,7 @@ impl TlsBackend {
 }
 
 #[derive(Debug, Copy, Clone)]
-pub enum Event<'a> {
+enum Event<'a> {
     ResumingPartialDownload,
     /// Received the Content-Length of the to-be downloaded data.
     DownloadContentLengthReceived(u64),
@@ -215,7 +395,7 @@ type DownloadCallback<'a> = &'a dyn Fn(Event<'_>) -> Result<()>;
 /// Download via libcurl; encrypt with the native (or OpenSSl) TLS
 /// stack via libcurl
 #[cfg(feature = "curl-backend")]
-pub mod curl {
+mod curl {
     use std::cell::RefCell;
     use std::str;
     use std::time::Duration;
@@ -226,7 +406,7 @@ pub mod curl {
 
     use super::{DownloadError, Event};
 
-    pub fn download(
+    pub(super) fn download(
         url: &Url,
         resume_from: u64,
         callback: &dyn Fn(Event<'_>) -> Result<()>,
@@ -327,7 +507,7 @@ pub mod curl {
 }
 
 #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
-pub mod reqwest_be {
+mod reqwest_be {
     use std::io;
     #[cfg(feature = "reqwest-rustls-tls")]
     use std::sync::Arc;
@@ -346,7 +526,7 @@ pub mod reqwest_be {
 
     use super::{DownloadError, Event};
 
-    pub async fn download(
+    pub(super) async fn download(
         url: &Url,
         resume_from: u64,
         callback: &dyn Fn(Event<'_>) -> Result<()>,
@@ -491,7 +671,7 @@ pub mod reqwest_be {
 }
 
 #[derive(Debug, Error)]
-pub enum DownloadError {
+enum DownloadError {
     #[error("http request returned an unsuccessful status code: {0}")]
     HttpStatus(u32),
     #[error("file not found")]
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
