Optimize extract_zip: Use async_zip::read::stream::ZipFileReader to avoid temporary file (#590)

* Add new dep async_zip v0.0.9 to binstalk-downloader
   with features "gzip", "zstd", "xz", "bzip2", "tokio".
* Refactor: Simplify `async_extracter::extract_*` API
* Refactor: Create newtype wrapper of `ZipError`
   so that the zip can be upgraded without affecting API of this crate.
* Enable feature fs of dep tokio in binstalk-downloader
* Rewrite `extract_zip` to use `async_zip::read::stream::ZipFileReader`
   which avoids writing the zip file to a temporary file and then read it
   back into memory.
* Refactor: Impl new fn `await_on_option` and use it
* Optimize `tokio::select!`: Make them biased and check for cancellation first
  to make cancellation takes effect ASAP.
* Rm unused dep zip from binstalk-downloader

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>

Cargo.lock

@@ -52,6 +52,15 @@ dependencies = [
  "alloc-no-stdlib",
 ]
 
+[[package]]
+name = "android_system_properties"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "async-compression"
 version = "0.3.15"
@@ -81,6 +90,29 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "async_io_utilities"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b20cffc5590f4bf33f05f97a3ea587feba9c50d20325b401daa096b92ff7da0"
+dependencies = [
+ "tokio",
+]
+
+[[package]]
+name = "async_zip"
+version = "0.0.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a36d43bdefc7215b2b3a97edd03b1553b7969ad76551025eedd3b913c645f6e"
+dependencies = [
+ "async-compression",
+ "async_io_utilities",
+ "chrono",
+ "crc32fast",
+ "thiserror",
+ "tokio",
+]
+
 [[package]]
 name = "atty"
 version = "0.2.14"
@@ -156,6 +188,7 @@ version = "0.1.0"
 dependencies = [
  "async-compression",
  "async-trait",
+ "async_zip",
  "binstalk-types",
  "binstall-tar",
  "bytes",
@@ -177,7 +210,6 @@ dependencies = [
  "trust-dns-resolver",
  "url",
  "xz2",
- "zip",
  "zstd 0.12.1+zstd.1.5.2",
 ]
 
@@ -366,9 +398,14 @@ version = "0.4.23"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "16b0a3d9ed01224b22057780a37bb8c5dbfe1be8ba48678e7bf57ec4b385411f"
 dependencies = [
+ "iana-time-zone",
+ "js-sys",
  "num-integer",
  "num-traits",
  "serde",
+ "time",
+ "wasm-bindgen",
+ "winapi",
 ]
 
 [[package]]
@@ -417,6 +454,16 @@ dependencies = [
  "cc",
 ]
 
+[[package]]
+name = "codespan-reporting"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3538270d33cc669650c4b093848450d380def10c331d38c768e34cac80576e6e"
+dependencies = [
+ "termcolor",
+ "unicode-width",
+]
+
 [[package]]
 name = "combine"
 version = "4.6.6"
@@ -482,15 +529,6 @@ dependencies = [
  "cfg-if",
 ]
 
-[[package]]
-name = "crossbeam-utils"
-version = "0.8.13"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "422f23e724af1240ec469ea1e834d87a4b59ce2efe2c6a96256b0c47e2fd86aa"
-dependencies = [
- "cfg-if",
-]
-
 [[package]]
 name = "crypto-common"
 version = "0.1.6"
@@ -501,6 +539,50 @@ dependencies = [
  "typenum",
 ]
 
+[[package]]
+name = "cxx"
+version = "1.0.83"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bdf07d07d6531bfcdbe9b8b739b104610c6508dcc4d63b410585faf338241daf"
+dependencies = [
+ "cc",
+ "cxxbridge-flags",
+ "cxxbridge-macro",
+ "link-cplusplus",
+]
+
+[[package]]
+name = "cxx-build"
+version = "1.0.83"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d2eb5b96ecdc99f72657332953d4d9c50135af1bac34277801cc3937906ebd39"
+dependencies = [
+ "cc",
+ "codespan-reporting",
+ "once_cell",
+ "proc-macro2",
+ "quote",
+ "scratch",
+ "syn",
+]
+
+[[package]]
+name = "cxxbridge-flags"
+version = "1.0.83"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac040a39517fd1674e0f32177648334b0f4074625b5588a64519804ba0553b12"
+
+[[package]]
+name = "cxxbridge-macro"
+version = "1.0.83"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1362b0ddcfc4eb0a1f57b68bd77dd99f0e826958a96abd0ae9bd092e114ffed6"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "data-encoding"
 version = "2.3.2"
@@ -810,7 +892,7 @@ checksum = "c05aeb6a22b8f62540c194aac980f2115af067bfe15a0734d7277a768d396b31"
 dependencies = [
  "cfg-if",
  "libc",
- "wasi",
+ "wasi 0.11.0+wasi-snapshot-preview1",
 ]
 
 [[package]]
@@ -987,6 +1069,30 @@ dependencies = [
  "tokio-native-tls",
 ]
 
+[[package]]
+name = "iana-time-zone"
+version = "0.1.53"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "64c122667b287044802d6ce17ee2ddf13207ed924c712de9a66a5814d5b64765"
+dependencies = [
+ "android_system_properties",
+ "core-foundation-sys",
+ "iana-time-zone-haiku",
+ "js-sys",
+ "wasm-bindgen",
+ "winapi",
+]
+
+[[package]]
+name = "iana-time-zone-haiku"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0703ae284fc167426161c2e3f1da3ea71d94b21bedbcc9494e92b28e334e3dca"
+dependencies = [
+ "cxx",
+ "cxx-build",
+]
+
 [[package]]
 name = "idna"
 version = "0.2.3"
@@ -1158,6 +1264,15 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "link-cplusplus"
+version = "1.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9272ab7b96c9046fbc5bc56c06c117cb639fe2d509df0c421cad82d2915cf369"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "linked-hash-map"
 version = "0.5.6"
@@ -1305,7 +1420,7 @@ checksum = "e5d732bc30207a6423068df043e3d02e0735b155ad7ce1a6f76fe2baa5b158de"
 dependencies = [
  "libc",
  "log",
- "wasi",
+ "wasi 0.11.0+wasi-snapshot-preview1",
  "windows-sys 0.42.0",
 ]
 
@@ -1881,6 +1996,12 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
 
+[[package]]
+name = "scratch"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c8132065adcfd6e02db789d9285a0deb2f3fcb04002865ab67d5fb103533898"
+
 [[package]]
 name = "sct"
 version = "0.7.0"
@@ -2176,6 +2297,17 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "time"
+version = "0.1.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1b797afad3f312d1c66a56d11d0316f916356d11bd158fbc6ca6389ff6bf805a"
+dependencies = [
+ "libc",
+ "wasi 0.10.0+wasi-snapshot-preview1",
+ "winapi",
+]
+
 [[package]]
 name = "tinytemplate"
 version = "1.2.1"
@@ -2599,6 +2731,12 @@ dependencies = [
  "try-lock",
 ]
 
+[[package]]
+name = "wasi"
+version = "0.10.0+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1a143597ca7c7793eff794def352d41792a93c481eb1042423ff7ff72ba2c31f"
+
 [[package]]
 name = "wasi"
 version = "0.11.0+wasi-snapshot-preview1"
@@ -2864,20 +3002,6 @@ dependencies = [
  "lzma-sys",
 ]
 
-[[package]]
-name = "zip"
-version = "0.6.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "537ce7411d25e54e8ae21a7ce0b15840e7bfcff15b51d697ec3266cc76bdf080"
-dependencies = [
- "byteorder",
- "bzip2",
- "crc32fast",
- "crossbeam-utils",
- "flate2",
- "zstd 0.11.2+zstd.1.5.2",
-]
-
 [[package]]
 name = "zstd"
 version = "0.11.2+zstd.1.5.2"

crates/binstalk-downloader/Cargo.toml

@@ -12,6 +12,7 @@ license = "GPL-3.0"
 [dependencies]
 async-trait = "0.1.59"
 async-compression = { version = "0.3.15", features = ["gzip", "zstd", "xz", "bzip2", "tokio"] }
+async_zip = { version = "0.0.9", features = ["deflate", "bzip2", "lzma", "zstd", "xz"] }
 binstalk-types = { version = "0.1.0", path = "../binstalk-types" }
 bytes = "1.3.0"
 bzip2 = "0.4.3"
@@ -29,7 +30,7 @@ scopeguard = "1.1.0"
 tar = { package = "binstall-tar", version = "0.4.39" }
 tempfile = "3.3.0"
 thiserror = "1.0.37"
-tokio = { version = "1.22.0", features = ["macros", "rt-multi-thread", "sync", "time"], default-features = false }
+tokio = { version = "1.22.0", features = ["macros", "rt-multi-thread", "sync", "time", "fs"], default-features = false }
 tokio-tar = "0.3.0"
 tokio-util = { version = "0.7.4", features = ["io"] }
 tower = { version = "0.4.13", features = ["limit", "util"] }
@@ -39,14 +40,6 @@ url = "2.3.1"
 
 xz2 = "0.1.7"
 
-# Disable all features of zip except for features of compression algorithms:
-# Disabled features include:
-#  - aes-crypto: Enables decryption of files which were encrypted with AES, absolutely zero use for
-#    this crate.
-#  - time: Enables features using the [time](https://github.com/time-rs/time) crate,
-#    which is not used by this crate.
-zip = { version = "0.6.3", default-features = false, features = ["deflate", "bzip2", "zstd"] }
-
 # zstd is also depended by zip.
 # Since zip 0.6.3 depends on zstd 0.11, we can use 0.12.0 here
 # because it uses the same zstd-sys version.

crates/binstalk-downloader/src/download.rs

@@ -7,7 +7,6 @@ use thiserror::Error as ThisError;
 use tracing::{debug, instrument};
 
 pub use binstalk_types::cargo_toml_binstall::PkgFmt;
-pub use zip::result::ZipError;
 
 use crate::remote::{Client, Error as RemoteError, Url};
 
@@ -20,6 +19,12 @@ pub use async_tar_visitor::*;
 mod extracter;
 mod stream_readable;
 
+mod zip_extraction;
+pub use zip_extraction::ZipError;
+
+mod utils;
+use utils::await_on_option;
+
 pub type CancellationFuture = Option<Pin<Box<dyn Future<Output = Result<(), io::Error>> + Send>>>;
 
 #[derive(Debug, ThisError)]
@@ -112,15 +117,13 @@ impl Download {
 
         debug!("Downloading and extracting then in-memory processing");
 
-        let ret = if let Some(cancellation_future) = cancellation_future {
+        let ret = tokio::select! {
-            tokio::select! {
+            biased;
-                res = extract_tar_based_stream_and_visit(stream, fmt, visitor) => res?,
+
-                res = cancellation_future => {
+            res = await_on_option(cancellation_future) => {
                 Err(res.err().unwrap_or_else(|| io::Error::from(DownloadError::UserAbort)))?
             }
-            }
+            res = extract_tar_based_stream_and_visit(stream, fmt, visitor) => res?,
-        } else {
-            extract_tar_based_stream_and_visit(stream, fmt, visitor).await?
         };
 
         debug!("Download, extraction and in-memory procession OK");
@@ -145,7 +148,11 @@ impl Download {
             path: &Path,
             cancellation_future: CancellationFuture,
         ) -> Result<(), DownloadError> {
-            let stream = this.client.get_stream(this.url).await?;
+            let stream = this
+                .client
+                .get_stream(this.url)
+                .await?
+                .map(|res| res.map_err(DownloadError::from));
 
             debug!("Downloading and extracting to: '{}'", path.display());
 

crates/binstalk-downloader/src/download/async_extracter.rs

@@ -1,24 +1,25 @@
-use std::{fs, io::Seek, path::Path};
+use std::{fs, path::Path};
 
+use async_zip::read::stream::ZipFileReader;
 use bytes::Bytes;
 use futures_util::stream::Stream;
 use scopeguard::{guard, ScopeGuard};
-use tempfile::tempfile;
 use tokio::task::block_in_place;
+use tokio_util::io::StreamReader;
 use tracing::debug;
 
 use super::{
-    extracter::*, stream_readable::StreamReadable, CancellationFuture, DownloadError, TarBasedFmt,
+    await_on_option, extracter::*, stream_readable::StreamReadable,
+    zip_extraction::extract_zip_entry, CancellationFuture, DownloadError, TarBasedFmt, ZipError,
 };
 
-pub async fn extract_bin<S, E>(
+pub async fn extract_bin<S>(
     stream: S,
     path: &Path,
     cancellation_future: CancellationFuture,
 ) -> Result<(), DownloadError>
 where
-    S: Stream<Item = Result<Bytes, E>> + Unpin + 'static,
+    S: Stream<Item = Result<Bytes, DownloadError>> + Unpin + 'static,
-    DownloadError: From<E>,
 {
     let mut reader = StreamReadable::new(stream, cancellation_future).await;
     block_in_place(move || {
@@ -42,39 +43,45 @@ where
     })
 }
 
-pub async fn extract_zip<S, E>(
+pub async fn extract_zip<S>(
     stream: S,
     path: &Path,
     cancellation_future: CancellationFuture,
 ) -> Result<(), DownloadError>
 where
-    S: Stream<Item = Result<Bytes, E>> + Unpin + 'static,
+    S: Stream<Item = Result<Bytes, DownloadError>> + Unpin + Send + Sync + 'static,
-    DownloadError: From<E>,
 {
-    let mut reader = StreamReadable::new(stream, cancellation_future).await;
+    debug!("Decompressing from zip archive to `{}`", path.display());
-    block_in_place(move || {
-        fs::create_dir_all(path.parent().unwrap())?;
 
-        let mut file = tempfile()?;
+    let extract_future = Box::pin(async move {
+        let reader = StreamReader::new(stream);
+        let mut zip = ZipFileReader::new(reader);
 
-        reader.copy(&mut file)?;
+        while let Some(entry) = zip.entry_reader().await.map_err(ZipError::from_inner)? {
+            extract_zip_entry(entry, path).await?;
+        }
 
-        // rewind it so that we can pass it to unzip
+        Ok(())
-        file.rewind()?;
+    });
 
-        unzip(file, path)
+    tokio::select! {
-    })
+        biased;
+
+        res = await_on_option(cancellation_future) => {
+            Err(res.err().map(DownloadError::from).unwrap_or(DownloadError::UserAbort))
+        }
+        res = extract_future => res,
+    }
 }
 
-pub async fn extract_tar_based_stream<S, E>(
+pub async fn extract_tar_based_stream<S>(
     stream: S,
     path: &Path,
     fmt: TarBasedFmt,
     cancellation_future: CancellationFuture,
 ) -> Result<(), DownloadError>
 where
-    S: Stream<Item = Result<Bytes, E>> + Unpin + 'static,
+    S: Stream<Item = Result<Bytes, DownloadError>> + Unpin + 'static,
-    DownloadError: From<E>,
 {
     let reader = StreamReadable::new(stream, cancellation_future).await;
     block_in_place(move || {

crates/binstalk-downloader/src/download/extracter.rs

@@ -1,18 +1,12 @@
-use std::{
+use std::io::{self, BufRead, Read};
-    fs::File,
-    io::{self, BufRead, Read},
-    path::Path,
-};
 
 use bzip2::bufread::BzDecoder;
 use flate2::bufread::GzDecoder;
 use tar::Archive;
-use tracing::debug;
 use xz2::bufread::XzDecoder;
-use zip::read::ZipArchive;
 use zstd::stream::Decoder as ZstdDecoder;
 
-use super::{DownloadError, TarBasedFmt};
+use super::TarBasedFmt;
 
 pub fn create_tar_decoder(
     dat: impl BufRead + 'static,
@@ -35,12 +29,3 @@ pub fn create_tar_decoder(
 
     Ok(Archive::new(r))
 }
-
-pub fn unzip(dat: File, dst: &Path) -> Result<(), DownloadError> {
-    debug!("Decompressing from zip archive to `{dst:?}`");
-
-    let mut zip = ZipArchive::new(dat)?;
-    zip.extract(dst)?;
-
-    Ok(())
-}

crates/binstalk-downloader/src/download/stream_readable.rs

@@ -7,7 +7,7 @@ use bytes::{Buf, Bytes};
 use futures_util::stream::{Stream, StreamExt};
 use tokio::runtime::Handle;
 
-use super::{CancellationFuture, DownloadError};
+use super::{await_on_option, CancellationFuture, DownloadError};
 
 /// This wraps an AsyncIterator as a `Read`able.
 /// It must be used in non-async context only,
@@ -120,18 +120,15 @@ where
 
         if !bytes.has_remaining() {
             let option = self.handle.block_on(async {
-                if let Some(cancellation_future) = self.cancellation_future.as_mut() {
+                let cancellation_future = self.cancellation_future.as_mut();
                 tokio::select! {
                     biased;
 
-                        res = cancellation_future => {
+                    res = await_on_option(cancellation_future) => {
                         Err(res.err().unwrap_or_else(|| io::Error::from(DownloadError::UserAbort)))
                     },
                     res = next_stream(&mut self.stream) => res,
                 }
-                } else {
-                    next_stream(&mut self.stream).await
-                }
             })?;
 
             if let Some(new_bytes) = option {

crates/binstalk-downloader/src/download/utils.rs

@@ -0,0 +1,16 @@
+use std::future::{pending, Future};
+
+/// Await on `future` if it is not `None`, or call [`pending`]
+/// so that this branch would never get selected again.
+///
+/// Designed to use with [`tokio::select`].
+pub(super) async fn await_on_option<Fut, R>(future: Option<Fut>) -> R
+where
+    Fut: Future<Output = R>,
+{
+    if let Some(future) = future {
+        future.await
+    } else {
+        pending().await
+    }
+}

crates/binstalk-downloader/src/download/zip_extraction.rs

@@ -0,0 +1,149 @@
+use std::{
+    io,
+    path::{Component, Path, PathBuf},
+};
+
+use async_zip::{read::ZipEntryReader, ZipEntryExt};
+use thiserror::Error as ThisError;
+use tokio::{fs, io::AsyncRead, task::spawn_blocking};
+
+use super::DownloadError;
+
+#[derive(Debug, ThisError)]
+enum ZipErrorInner {
+    #[error(transparent)]
+    Inner(#[from] async_zip::error::ZipError),
+
+    #[error("Invalid file path: {0}")]
+    InvalidFilePath(Box<str>),
+}
+
+#[derive(Debug, ThisError)]
+#[error(transparent)]
+pub struct ZipError(#[from] ZipErrorInner);
+
+impl ZipError {
+    pub(super) fn from_inner(err: async_zip::error::ZipError) -> Self {
+        Self(ZipErrorInner::Inner(err))
+    }
+}
+
+pub(super) async fn extract_zip_entry<R>(
+    entry: ZipEntryReader<'_, R>,
+    path: &Path,
+) -> Result<(), DownloadError>
+where
+    R: AsyncRead + Unpin + Send + Sync,
+{
+    // Sanitize filename
+    let raw_filename = entry.entry().filename();
+    let filename = check_filename_and_normalize(raw_filename)
+        .ok_or_else(|| ZipError(ZipErrorInner::InvalidFilePath(raw_filename.into())))?;
+
+    // Calculates the outpath
+    let outpath = path.join(filename);
+
+    // Get permissions
+    let mut perms = None;
+
+    #[cfg(unix)]
+    {
+        use std::{fs::Permissions, os::unix::fs::PermissionsExt};
+
+        if let Some(mode) = entry.entry().unix_permissions() {
+            let mode: u16 = mode;
+            perms = Some(Permissions::from_mode(mode as u32));
+        }
+    }
+
+    if raw_filename.ends_with('/') {
+        // This entry is a dir.
+        asyncify(move || {
+            std::fs::create_dir_all(&outpath)?;
+            if let Some(perms) = perms {
+                std::fs::set_permissions(&outpath, perms)?;
+            }
+
+            Ok(())
+        })
+        .await?;
+    } else {
+        // This entry is a file.
+        let mut outfile = asyncify(move || {
+            if let Some(p) = outpath.parent() {
+                std::fs::create_dir_all(p)?;
+            }
+            let outfile = std::fs::File::create(&outpath)?;
+
+            if let Some(perms) = perms {
+                outfile.set_permissions(perms)?;
+            }
+
+            Ok(outfile)
+        })
+        .await
+        .map(fs::File::from_std)?;
+
+        entry
+            .copy_to_end_crc(&mut outfile, 64 * 1024)
+            .await
+            .map_err(ZipError::from_inner)?;
+    }
+
+    Ok(())
+}
+
+/// Ensure the file path is safe to use as a [`Path`].
+///
+/// - It can't contain NULL bytes
+/// - It can't resolve to a path outside the current directory
+///   > `foo/../bar` is fine, `foo/../../bar` is not.
+/// - It can't be an absolute path
+///
+/// It will then return a normalized path.
+///
+/// This will read well-formed ZIP files correctly, and is resistant
+/// to path-based exploits.
+///
+/// This function is adapted from `zip::ZipFile::enclosed_name`.
+fn check_filename_and_normalize(filename: &str) -> Option<PathBuf> {
+    if filename.contains('\0') {
+        return None;
+    }
+
+    let mut path = PathBuf::new();
+
+    // The following loop is adapted from
+    // `normalize_path::NormalizePath::normalize`.
+    for component in Path::new(filename).components() {
+        match component {
+            Component::Prefix(_) | Component::RootDir => return None,
+            Component::CurDir => (),
+            Component::ParentDir => {
+                if !path.pop() {
+                    // `PathBuf::pop` returns false if there is no parent.
+                    // which means the path is invalid.
+                    return None;
+                }
+            }
+            Component::Normal(c) => path.push(c),
+        }
+    }
+
+    Some(path)
+}
+
+/// Copied from tokio https://docs.rs/tokio/latest/src/tokio/fs/mod.rs.html#132
+async fn asyncify<F, T>(f: F) -> io::Result<T>
+where
+    F: FnOnce() -> io::Result<T> + Send + 'static,
+    T: Send + 'static,
+{
+    match spawn_blocking(f).await {
+        Ok(res) => res,
+        Err(_) => Err(io::Error::new(
+            io::ErrorKind::Other,
+            "background task failed",
+        )),
+    }
+}