chore(download): add functions to fetch remote size

Author: motorailgun

src/download/mod.rs

@@ -240,6 +240,27 @@ async fn download_file_(
     res
 }
 
+pub(crate) async fn content_length(url: &Url, process: &Process) -> anyhow::Result<Option<u64>> {
+    if url.scheme() == "file" {
+        let path = url
+            .to_file_path()
+            .map_err(|_| anyhow::anyhow!("bogus file url: '{url}'"))?;
+        return Ok(Some(std::fs::metadata(path)?.len()));
+    }
+
+    let backend = select_backend(process)?;
+    let timeout = timeout(process)?;
+
+    match backend {
+        #[cfg(feature = "curl-backend")]
+        Backend::Curl => debug!(url = %url, "fetching content-length with curl"),
+        #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
+        Backend::Reqwest(_) => debug!(url = %url, "fetching content-length with reqwest"),
+    };
+
+    backend.content_length(url, timeout).await
+}
+
 /// User agent header value for HTTP request.
 /// See: https://github.com/rust-lang/rustup/issues/2860.
 #[cfg(feature = "curl-backend")]
@@ -402,6 +423,15 @@ impl Backend {
             Self::Reqwest(tls) => tls.download(url, resume_from, callback, timeout).await,
         }
     }
+
+    async fn content_length(self, url: &Url, timeout: Duration) -> anyhow::Result<Option<u64>> {
+        match self {
+            #[cfg(feature = "curl-backend")]
+            Self::Curl => curl::content_length(url, timeout),
+            #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
+            Self::Reqwest(tls) => tls.content_length(url, timeout).await,
+        }
+    }
 }
 
 #[cfg(any(feature = "reqwest-rustls-tls", feature = "reqwest-native-tls"))]
@@ -431,6 +461,17 @@ impl TlsBackend {
 
         reqwest_be::download(url, resume_from, callback, client).await
     }
+
+    async fn content_length(self, url: &Url, timeout: Duration) -> anyhow::Result<Option<u64>> {
+        let client = match self {
+            #[cfg(feature = "reqwest-rustls-tls")]
+            Self::Rustls => reqwest_be::rustls_client(timeout)?,
+            #[cfg(feature = "reqwest-native-tls")]
+            Self::NativeTls => reqwest_be::native_tls_client(timeout)?,
+        };
+
+        reqwest_be::content_length(url, client).await
+    }
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -448,6 +489,7 @@ type DownloadCallback<'a> = &'a dyn Fn(Event<'_>) -> anyhow::Result<()>;
 /// stack via libcurl
 #[cfg(feature = "curl-backend")]
 mod curl {
+    use std::cell::Cell;
     use std::cell::RefCell;
     use std::str;
     use std::time::Duration;
@@ -458,6 +500,46 @@ mod curl {
 
     use super::{DownloadError, Event};
 
+    pub(super) fn content_length(url: &Url, timeout: Duration) -> Result<Option<u64>> {
+        let mut handle = Easy::new();
+        handle.url(url.as_ref())?;
+        handle.follow_location(true)?;
+        handle.useragent(super::CURL_USER_AGENT)?;
+        handle.nobody(true)?;
+        handle.connect_timeout(timeout)?;
+
+        let length = Cell::new(None);
+        {
+            let mut transfer = handle.transfer();
+            transfer.header_function(|header| {
+                let Ok(data) = str::from_utf8(header) else {
+                    return true;
+                };
+                let prefix = "content-length: ";
+                let Some((dp, ds)) = data.split_at_checked(prefix.len()) else {
+                    return true;
+                };
+                if !dp.eq_ignore_ascii_case(prefix) {
+                    return true;
+                }
+                if let Ok(s) = ds.trim().parse::<u64>() {
+                    length.set(Some(s));
+                }
+                true
+            })?;
+
+            transfer.perform()?;
+        }
+
+        let code = handle.response_code()?;
+        match code {
+            0 | 200..=299 => Ok(length.get()),
+            // Some servers do not support HEAD for file assets
+            405 => Ok(None),
+            _ => Err(DownloadError::HttpStatus(code).into()),
+        }
+    }
+
     pub(super) fn download(
         url: &Url,
         resume_from: u64,
@@ -620,6 +702,22 @@ mod reqwest_be {
         Ok(())
     }
 
+    pub(super) async fn content_length(url: &Url, client: &Client) -> anyhow::Result<Option<u64>> {
+        let res = client
+            .head(url.as_str())
+            .send()
+            .await
+            .context("error fetching content length")?;
+
+        let status = res.status().into();
+        match status {
+            200..=299 => Ok(res.content_length()),
+            // Some servers do not support HEAD for file assets
+            405 => Ok(None),
+            _ => Err(DownloadError::HttpStatus(u32::from(status)).into()),
+        }
+    }
+
     fn client_generic() -> ClientBuilder {
         Client::builder()
             // HACK: set `pool_max_idle_per_host` to `0` to avoid an issue in the underlying