fix: retry truncated response bodies

packages/zpm/src/builtins/node.rs

@@ -30,7 +30,7 @@ pub async fn resolve_nodejs_version(context: &InstallContext<'_>, range: &zpm_se
         = format!("{}/index.json", project.config.settings.node_dist_url.value);
 
     let text
-        = project.http_client.get(&release_url)?.send().await?.text().await?;
+        = project.http_client.get(&release_url)?.send_text().await?;
 
     #[derive(Deserialize)]
     struct NodejsManifest {
@@ -173,11 +173,9 @@ pub async fn fetch_nodejs_locator<'a>(context: &InstallContext<'a>, locator: &Lo
         = system.arch.clone();
 
     let cached_blob = package_cache.ensure_blob(locator.clone(), ".zip", || async move {
-        let bytes
+        let (_, bytes)
             = project.http_client.get(&url)?
-                .send().await?
-                .error_for_status()?
-                .bytes().await?;
+                .send_bytes().await?;
 
         let archive = tokio::task::spawn_blocking(move || -> Result<Vec<u8>, Error> {
             let tar_data

packages/zpm/src/commands/debug/http.rs

@@ -29,9 +29,7 @@ impl Http {
         with_report_result(report, async {
             project.http_client
                 .get(&self.url)?
-                .send()
-                .await?
-                .text()
+                .send_text()
                 .await?;
 
             Ok(())

packages/zpm/src/fetchers/pypi.rs

@@ -97,13 +97,11 @@ pub async fn fetch_locator<'a>(context: &InstallContext<'a>, locator: &Locator,
 
     let cached_blob
         = package_cache.ensure_blob(locator.clone(), ".zip", || async {
-            let response
+            let (_, bytes)
                 = project.http_client.get(&artifact_url)?
-                    .send()
+                    .send_bytes()
                     .await?;
 
-            let bytes
-                = response.bytes().await?;
             Ok(bytes.to_vec())
         }).await?.into_info();
 

packages/zpm/src/fetchers/url.rs

@@ -64,12 +64,16 @@ pub async fn fetch_locator<'a>(context: &InstallContext<'a>, locator: &Locator,
     };
 
     let cached_blob = package_cache.upsert_blob(locator.clone(), ".zip", || async {
-        let response = project.http_client.get(&params.url)?
+        let (_, tgz_data) = project.http_client.get(&params.url)?
             .header("authorization", authorization.as_deref())
-            .send().await?;
-
-        let tgz_data = response.bytes().await
-            .map_err(|err| Error::RemoteRegistryError(Arc::new(err)))?;
+            .send_bytes().await
+            .map_err(|err| {
+                if err.is_body() || err.is_decode() {
+                    Error::RemoteRegistryError(Arc::new(err))
+                } else {
+                    err.into()
+                }
+            })?;
         let archive = tokio::task::spawn_blocking(move || -> Result<Vec<u8>, Error> {
             let tar_data
                 = zpm_formats::tar::unpack_tgz(&tgz_data)?;

packages/zpm/src/github.rs

@@ -20,12 +20,10 @@ pub async fn download_into(source: &GitSource, commit: &str, download_dir: &Path
     };
 
     let response
-        = http_client.get(public_tarball_url(owner, &repository, commit))?.send().await;
+        = http_client.get(public_tarball_url(owner, &repository, commit))?.send_bytes().await;
 
     let tgz_data = match response {
-        Ok(response) => {
-            response.bytes().await.map_err(|_| Error::ReplaceMe)?
-        },
+        Ok((_, tgz_data)) => tgz_data,
 
         Err(err) if err.status() == Some(StatusCode::NOT_FOUND) => {
             return Ok(None);

packages/zpm/src/http.rs

@@ -1,6 +1,6 @@
-use std::{collections::HashSet, net::SocketAddr, sync::{Arc, LazyLock, OnceLock}, time::Duration};
+use std::{collections::HashSet, future::Future, net::SocketAddr, sync::{Arc, LazyLock, OnceLock}, time::Duration};
 
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use dashmap::DashMap;
 use hickory_resolver::{config::LookupIpStrategy, TokioResolver};
 use http::HeaderMap;
@@ -181,7 +181,11 @@ impl<'a> HttpRequest<'a> {
         self
     }
 
-    pub async fn send(self) -> Result<Response, reqwest::Error> {
+    async fn send_with<T, F, Fut>(self, consume: F) -> Result<T, reqwest::Error>
+    where
+        F: Fn(Response) -> Fut,
+        Fut: Future<Output = Result<T, reqwest::Error>>,
+    {
         let mut retry_count
             = 0;
 
@@ -222,33 +226,91 @@ impl<'a> HttpRequest<'a> {
                 }
             };
 
-            if self.enable_retry && retry_count < self.client.config.http_retry {
-                let is_failure = match &response {
-                    Ok(response) => response.status().is_server_error() || matches!(response.status().as_u16(), 408 | 413 | 429),
-                    Err(_) => true,
-                };
+            let is_failure = match &response {
+                Ok(response) => response.status().is_server_error() || matches!(response.status().as_u16(), 408 | 413 | 429),
+                Err(_) => true,
+            };
 
-                if is_failure {
-                    retry_count += 1;
+            if self.enable_retry && retry_count < self.client.config.http_retry && is_failure {
+                retry_count += 1;
 
-                    let sleep_duration
-                        = 2_u64.saturating_pow(retry_count as u32);
-                    let bounded_sleep_duration
-                        = std::cmp::min(sleep_duration, 10);
+                let sleep_duration
+                    = 2_u64.saturating_pow(retry_count as u32);
+                let bounded_sleep_duration
+                    = std::cmp::min(sleep_duration, 10);
 
-                    tokio::time::sleep(Duration::from_secs(bounded_sleep_duration)).await;
-                    continue;
-                }
+                tokio::time::sleep(Duration::from_secs(bounded_sleep_duration)).await;
+                continue;
             }
 
-            return if self.enable_status_check {
-                response?.error_for_status()
+            let response
+                = response?;
+
+            let response = if self.enable_status_check {
+                response.error_for_status()?
             } else {
                 response
             };
+
+            let result
+                = consume(response).await;
+
+            if self.enable_retry && retry_count < self.client.config.http_retry && result.is_err() {
+                retry_count += 1;
+
+                let sleep_duration
+                    = 2_u64.saturating_pow(retry_count as u32);
+                let bounded_sleep_duration
+                    = std::cmp::min(sleep_duration, 10);
+
+                tokio::time::sleep(Duration::from_secs(bounded_sleep_duration)).await;
+                continue;
+            }
+
+            return result;
         }
     }
 
+    pub async fn send(self) -> Result<Response, reqwest::Error> {
+        self.send_with(|response| async move {
+            Ok(response)
+        }).await
+    }
+
+    pub async fn send_text(self) -> Result<String, reqwest::Error> {
+        self.send_with(|response| response.text()).await
+    }
+
+    /// Buffers the response body inside the retry loop while retaining the
+    /// drained response so callers can inspect its status and headers.
+    pub async fn send_bytes(self) -> Result<(Response, Bytes), reqwest::Error> {
+        let enable_status_check
+            = self.enable_status_check;
+
+        self.send_with(move |mut response| async move {
+            if !enable_status_check
+                && (response.status().is_client_error()
+                    || response.status().is_server_error()
+                    || response.status().as_u16() == 304)
+            {
+                return Ok((response, Bytes::new()));
+            }
+
+            let capacity
+                = response.content_length()
+                    .and_then(|length| usize::try_from(length).ok())
+                    .unwrap_or_default();
+            let mut body
+                = BytesMut::with_capacity(capacity);
+
+            while let Some(chunk) = response.chunk().await? {
+                body.extend_from_slice(&chunk);
+            }
+
+            Ok((response, body.freeze()))
+        }).await
+    }
+
     pub fn headers(&self) -> HeaderMap {
         // TODO: This is filthy
         self.builder.try_clone().unwrap().build().unwrap().headers().clone()
@@ -523,11 +585,8 @@ impl HttpClient {
             let request
                 = self.get(&url_str)?;
 
-            let result
-                = request.send().await?;
-
-            let bytes
-                = result.bytes().await?;
+            let (_, bytes)
+                = request.send_bytes().await?;
 
             Ok(bytes)
         }).await;

packages/zpm/src/http_npm.rs

@@ -340,15 +340,12 @@ pub async fn get_id_token(options: &GetIdTokenOptions<'_>) -> Result<Option<Stri
     actions_id_token_request_url.query_pairs_mut()
         .append_pair("audience", options.audience);
 
-    let response
+    let body
         = options.http_client.get(actions_id_token_request_url)?
             .header("authorization", Some(format!("Bearer {}", actions_id_token_request_token)))
-            .send()
+            .send_text()
             .await?;
 
-    let body
-        = response.text().await?;
-
     #[derive(Deserialize)]
     struct ActionsIdTokenResponse {
         value: String,
@@ -485,14 +482,15 @@ pub async fn get(params: &NpmHttpParams<'_>) -> Result<Bytes, Error> {
 
     let bytes = match params.authorization {
         Some(authorization) => {
-            let response = params.http_client.get(&url)?
+            let (response, bytes) = params.http_client.get(&url)?
                 .header("authorization", Some(authorization))
                 .enable_status_check(false)
-                .send().await?;
+                .send_bytes().await?;
 
             handle_invalid_authentication_error(params, &response).await?;
 
-            response.error_for_status()?.bytes().await?
+            response.error_for_status()?;
+            bytes
         },
 
         None => {
@@ -510,16 +508,17 @@ pub async fn get_uncached(params: &NpmHttpParams<'_>) -> Result<Bytes, Error> {
     let url
         = format!("{}{}", params.registry, params.path);
 
-    let response = params.http_client.get(&url)?
+    let (response, bytes) = params.http_client.get(&url)?
         .header("authorization", params.authorization)
         .enable_status_check(false)
-        .send().await?;
+        .send_bytes().await?;
 
     if params.authorization.is_some() {
         handle_invalid_authentication_error(params, &response).await?;
     }
 
-    Ok(response.error_for_status()?.bytes().await?)
+    response.error_for_status()?;
+    Ok(bytes)
 }
 
 const CACHED_VERSION_FIELDS: &[&str] = &[
@@ -740,8 +739,8 @@ async fn fetch_metadata_with_disk_cache(params: &GetPackageMetadataParams<'_>) -
         }
     }
 
-    let response
-        = request.send().await?;
+    let (response, fresh_body)
+        = request.send_bytes().await?;
 
     if params.authorization.is_some() {
         let npm_params = NpmHttpParams {
@@ -769,8 +768,7 @@ async fn fetch_metadata_with_disk_cache(params: &GetPackageMetadataParams<'_>) -
             .and_then(|v| v.to_str().ok())
             .map(|s| s.to_string());
 
-    let fresh_body
-        = response.error_for_status()?.bytes().await?;
+    response.error_for_status()?;
 
     // Keep stale version entries the fresh response omits so
     // resolution still works when a published version is later