fix: flaky package extract error (#535)

baszalmstra · web-flow · commit c1b7ff235b55 · 2024-02-23T13:59:45.000+01:00
This adds a test to check that the flaky download behavior has been
fixed. This was caused by a Read implementation to consistently return
an error at some point which caused the zip crate to panic.
diff --git a/crates/rattler/src/package_cache.rs b/crates/rattler/src/package_cache.rs
@@ -311,11 +311,16 @@ mod test {
         routing::get_service,
         Router,
     };
+    use bytes::Bytes;
+    use futures::stream;
     use rattler_conda_types::package::{ArchiveIdentifier, PackageFile, PathsJson};
     use rattler_networking::retry_policies::{DoNotRetryPolicy, ExponentialBackoffBuilder};
-    use std::{fs::File, future::IntoFuture, net::SocketAddr, path::Path, sync::Arc};
+    use std::{
+        convert::Infallible, fs::File, future::IntoFuture, net::SocketAddr, path::Path, sync::Arc,
+    };
     use tempfile::tempdir;
     use tokio::sync::Mutex;
+    use tokio_stream::StreamExt;
     use tower_http::services::ServeDir;
     use url::Url;
 
@@ -382,22 +387,69 @@ mod test {
         Ok(next.run(req).await)
     }
 
-    #[tokio::test]
-    pub async fn test_flaky_package_cache() {
+    /// A helper middleware function that fails the first two requests.
+    async fn fail_with_half_package(
+        State((count, bytes)): State<(Arc<Mutex<i32>>, Arc<Mutex<usize>>)>,
+        req: Request<Body>,
+        next: Next,
+    ) -> Result<Response, StatusCode> {
+        let count = {
+            let mut count = count.lock().await;
+            *count += 1;
+            *count
+        };
+
+        println!("Running middleware for request #{count} for {}", req.uri());
+        let response = next.run(req).await;
+
+        if count <= 2 {
+            // println!("Cutting response body in half");
+            let body = response.into_body();
+            let mut body = body.into_data_stream();
+            let mut buffer = Vec::new();
+            while let Some(Ok(chunk)) = body.next().await {
+                buffer.extend(chunk);
+            }
+
+            let byte_count = bytes.lock().await.clone();
+            let bytes = buffer.into_iter().take(byte_count).collect::<Vec<u8>>();
+            // Create a stream that ends prematurely
+            let stream = stream::iter(vec![
+                Ok::<_, Infallible>(Bytes::from_iter(bytes.into_iter())),
+                // The stream ends after sending partial data, simulating a premature close
+            ]);
+            let body = Body::from_stream(stream);
+            return Ok(Response::new(body));
+        }
+
+        Ok(response)
+    }
+
+    enum Middleware {
+        FailTheFirstTwoRequests,
+        FailAfterBytes(usize),
+    }
+
+    async fn test_flaky_package_cache(archive_name: &str, middleware: Middleware) {
         let static_dir = get_test_data_dir();
         println!("Serving files from {}", static_dir.display());
         // Construct a service that serves raw files from the test directory
         let service = get_service(ServeDir::new(static_dir));
 
         // Construct a router that returns data from the static dir but fails the first try.
         let request_count = Arc::new(Mutex::new(0));
-        let router =
-            Router::new()
-                .route_service("/*key", service)
-                .layer(middleware::from_fn_with_state(
-                    request_count.clone(),
-                    fail_the_first_two_requests,
-                ));
+        let router = Router::new().route_service("/*key", service);
+
+        let router = match middleware {
+            Middleware::FailTheFirstTwoRequests => router.layer(middleware::from_fn_with_state(
+                request_count.clone(),
+                fail_the_first_two_requests,
+            )),
+            Middleware::FailAfterBytes(size) => router.layer(middleware::from_fn_with_state(
+                (request_count.clone(), Arc::new(Mutex::new(size))),
+                fail_with_half_package,
+            )),
+        };
 
         // Construct the server that will listen on localhost but with a *random port*. The random
         // port is very important because it enables creating multiple instances at the same time.
@@ -412,7 +464,6 @@ mod test {
         let packages_dir = tempdir().unwrap();
         let cache = PackageCache::new(packages_dir.path());
 
-        let archive_name = "ros-noetic-rosbridge-suite-0.11.14-py39h6fdeb60_14.tar.bz2";
         let server_url = Url::parse(&format!("http://localhost:{}", addr.port())).unwrap();
 
         // Do the first request without
@@ -448,4 +499,17 @@ mod test {
             assert_eq!(*request_count_lock, 3, "Expected there to be 3 requests");
         }
     }
+
+    #[tokio::test]
+    async fn test_flaky() {
+        let tar_bz2 = "ros-noetic-rosbridge-suite-0.11.14-py39h6fdeb60_14.tar.bz2";
+        let conda = "conda-22.11.1-py38haa244fe_1.conda";
+
+        test_flaky_package_cache(tar_bz2, Middleware::FailTheFirstTwoRequests).await;
+        test_flaky_package_cache(conda, Middleware::FailTheFirstTwoRequests).await;
+
+        test_flaky_package_cache(tar_bz2, Middleware::FailAfterBytes(1000)).await;
+        test_flaky_package_cache(conda, Middleware::FailAfterBytes(1000)).await;
+        test_flaky_package_cache(conda, Middleware::FailAfterBytes(50)).await;
+    }
 }
diff --git a/crates/rattler_package_streaming/Cargo.toml b/crates/rattler_package_streaming/Cargo.toml
@@ -43,3 +43,4 @@ tokio = { workspace = true, features = ["rt", "macros"] }
 walkdir = { workspace = true }
 rstest = { workspace = true }
 rstest_reuse = { workspace = true }
+assert_matches = { workspace = true }
diff --git a/crates/rattler_package_streaming/src/lib.rs b/crates/rattler_package_streaming/src/lib.rs
@@ -3,6 +3,7 @@
 //! This crate provides the ability to extract a Conda package archive or specific parts of it.
 
 use std::path::PathBuf;
+use zip::result::ZipError;
 
 use rattler_digest::{Md5Hash, Sha256Hash};
 
@@ -27,7 +28,7 @@ pub enum ExtractError {
     CouldNotCreateDestination(#[source] std::io::Error),
 
     #[error("invalid zip archive")]
-    ZipError(#[from] zip::result::ZipError),
+    ZipError(#[source] zip::result::ZipError),
 
     #[error("a component is missing from the Conda archive")]
     MissingComponent,
@@ -49,6 +50,15 @@ pub enum ExtractError {
     ArchiveMemberParseError(PathBuf, #[source] std::io::Error),
 }
 
+impl From<ZipError> for ExtractError {
+    fn from(value: ZipError) -> Self {
+        match value {
+            ZipError::Io(io) => Self::IoError(io),
+            e => Self::ZipError(e),
+        }
+    }
+}
+
 #[cfg(feature = "reqwest")]
 impl From<::reqwest::Error> for ExtractError {
     fn from(err: ::reqwest::Error) -> Self {
diff --git a/crates/rattler_package_streaming/src/read.rs b/crates/rattler_package_streaming/src/read.rs
@@ -2,6 +2,7 @@
 //! [`std::io::Read`] trait.
 
 use super::{ExtractError, ExtractResult};
+use std::mem::ManuallyDrop;
 use std::{ffi::OsStr, io::Read, path::Path};
 use zip::read::read_zipfile_from_stream;
 
@@ -55,24 +56,29 @@ pub fn extract_conda(reader: impl Read, destination: &Path) -> Result<ExtractRes
 
     // Iterate over all entries in the zip-file and extract them one-by-one
     while let Some(file) = read_zipfile_from_stream(&mut md5_reader)? {
+        // If an error occurs while we are reading the contents of the zip we don't want to
+        // seek to the end of the file. Using [`ManuallyDrop`] we prevent `drop` to be called on
+        // the `file` in case the stack unwinds.
+        let mut file = ManuallyDrop::new(file);
+
         if file
             .mangled_name()
             .file_name()
             .map(OsStr::to_string_lossy)
             .map_or(false, |file_name| file_name.ends_with(".tar.zst"))
         {
-            stream_tar_zst(file)?.unpack(destination)?;
+            stream_tar_zst(&mut *file)?.unpack(destination)?;
+        } else {
+            // Manually read to the end of the stream if that didn't happen.
+            std::io::copy(&mut *file, &mut std::io::sink())?;
         }
+
+        // Take the file out of the [`ManuallyDrop`] to properly drop it.
+        let _ = ManuallyDrop::into_inner(file);
     }
 
     // Read the file to the end to make sure the hash is properly computed.
-    let mut buf = [0; 1 << 14];
-    loop {
-        let bytes_read = md5_reader.read(&mut buf)?;
-        if bytes_read == 0 {
-            break;
-        }
-    }
+    std::io::copy(&mut md5_reader, &mut std::io::sink())?;
 
     // Get the hashes
     let (sha256_reader, md5) = md5_reader.finalize();
diff --git a/crates/rattler_package_streaming/tests/extract.rs b/crates/rattler_package_streaming/tests/extract.rs
@@ -1,8 +1,10 @@
 use rattler_conda_types::package::IndexJson;
 use rattler_package_streaming::read::{extract_conda, extract_tar_bz2};
+use rattler_package_streaming::ExtractError;
 use rstest::rstest;
 use rstest_reuse::{self, apply, template};
 use std::fs::File;
+use std::io::Read;
 use std::path::{Path, PathBuf};
 
 fn test_data_dir() -> PathBuf {
@@ -235,3 +237,42 @@ async fn test_extract_url_async(#[case] url: &str, #[case] sha256: &str, #[case]
     assert_eq!(&format!("{:x}", result.sha256), sha256);
     assert_eq!(&format!("{:x}", result.md5), md5);
 }
+
+#[rstest]
+fn test_extract_flaky_conda(#[values(0, 1, 13, 50, 74, 150, 8096, 16384, 20000)] cutoff: usize) {
+    let input = "conda-22.11.1-py38haa244fe_1.conda";
+    let temp_dir = Path::new(env!("CARGO_TARGET_TMPDIR"));
+    println!("Target dir: {}", temp_dir.display());
+    let file_path = Path::new(input);
+    let target_dir = temp_dir.join(file_path.file_stem().unwrap());
+    let result = extract_conda(
+        FlakyReader {
+            reader: File::open(test_data_dir().join(file_path)).unwrap(),
+            total_read: 0,
+            cutoff,
+        },
+        &target_dir,
+    )
+    .expect_err("this should error out and not panic");
+
+    assert_matches::assert_matches!(result, ExtractError::IoError(_));
+}
+
+struct FlakyReader<R: Read> {
+    reader: R,
+    cutoff: usize,
+    total_read: usize,
+}
+
+impl<R: Read> Read for FlakyReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let remaining = self.cutoff.saturating_sub(self.total_read);
+        if remaining == 0 {
+            return Err(std::io::Error::new(std::io::ErrorKind::Other, "flaky"));
+        }
+        let max_read = buf.len().min(remaining);
+        let bytes_read = self.reader.read(&mut buf[..max_read])?;
+        self.total_read += bytes_read;
+        Ok(bytes_read)
+    }
+}
