diff --git a/crates/vite_install/src/request.rs b/crates/vite_install/src/request.rs index 900f9f4133..db6a749f9f 100644 --- a/crates/vite_install/src/request.rs +++ b/crates/vite_install/src/request.rs @@ -54,23 +54,14 @@ impl HttpClient { /// * `Err(e)` - If the request fails pub async fn get_bytes(&self, url: &str) -> Result, Error> { tracing::debug!("Fetching bytes from: {}", url); - let response = self.get(url).await?; - Ok(response.bytes().await?.to_vec()) - } - - async fn get(&self, url: &str) -> Result { - self.get_with_accept(url, None).await - } - async fn get_with_accept(&self, url: &str, accept: Option<&str>) -> Result { let client = vite_shared::shared_http_client(); - let response = (|| async { - let mut request = client.get(url); - if let Some(accept) = accept { - request = request.header(reqwest::header::ACCEPT, accept); - } - request.send().await?.error_for_status() + // Read the body inside the retry so a mid-body connection drop gets + // retried instead of failing outright, like `download_file`. + let bytes = (|| async { + let response = client.get(url).send().await?.error_for_status()?; + Ok::<_, Error>(response.bytes().await?) }) .retry( ExponentialBuilder::default() @@ -80,7 +71,7 @@ impl HttpClient { ) .await?; - Ok(response) + Ok(bytes.to_vec()) } /// Get JSON data from a URL @@ -739,6 +730,69 @@ mod tests { ); } + /// `get_bytes` used to read the body outside the retry, so a connection + /// dropped mid-body never got retried. The server truncates the first + /// response, then sends the whole body — `get_bytes` should retry and succeed. + #[tokio::test] + async fn test_get_bytes_retries_on_truncated_body() { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; + + let body = b"the complete body payload that must arrive intact"; + let len = body.len(); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let connections = Arc::new(AtomicUsize::new(0)); + + let server_connections = Arc::clone(&connections); + let server = tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + let attempt = server_connections.fetch_add(1, Ordering::SeqCst); + + // Drain the request before replying. + let mut scratch = [0u8; 1024]; + let _ = socket.read(&mut scratch).await; + + let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {len}\r\n\r\n"); + socket.write_all(head.as_bytes()).await.unwrap(); + if attempt == 0 { + // First attempt: send half the body, then drop the connection. + socket.write_all(&body[..len / 2]).await.unwrap(); + } else { + // Retry: send the whole body. + socket.write_all(body).await.unwrap(); + } + socket.flush().await.unwrap(); + } + }); + + let client = HttpClient::with_config(3, 10); + let url = format!("http://{addr}/"); + let result = client.get_bytes(&url).await; + + server.abort(); + + let attempts = connections.load(Ordering::SeqCst); + assert!( + result.is_ok(), + "get_bytes must retry a truncated body and eventually succeed, but got {result:?} after {attempts} attempt(s)" + ); + assert_eq!(result.unwrap(), body); + assert!( + attempts >= 2, + "a body-level failure must be retried, but get_bytes only made {attempts} connection(s)" + ); + } + #[tokio::test] #[ignore] // Flaky on musl/Alpine — temp file race condition async fn test_verify_file_hash_sha1() {