use std::sync::Arc; use async_stream::try_stream; use async_zip::StoredZipEntry; use async_zip::base::read::WithoutEntry; use futures_util::{Stream, TryStreamExt, pin_mut}; use tokio_util::compat::TokioAsyncWriteCompatExt; use crate::{IOError, PolyIOResult}; #[tracing::instrument(skip(data, f), level = "debug")] pub async fn read_zip_entries_bytes(data: Vec, mut f: F) -> PolyIOResult<()> where F: AsyncFnMut( usize, &StoredZipEntry, &mut async_zip::base::read::ZipEntryReader< '_, futures_lite::io::Cursor<&[u8]>, WithoutEntry, >, ) -> PolyIOResult<()>, { let reader = async_zip::base::read::mem::ZipFileReader::new(data).await?; let entries = reader.file().entries(); for index in 0..entries.len() { let entry = entries.get(index).expect("expected more zip entries"); let mut entry_reader = reader .reader_without_entry(index) .await .unwrap_or_else(|_| panic!("expected zip entry at index '{index}'")); f(index, entry, &mut entry_reader).await?; } Ok(()) } #[tracing::instrument(skip(data), level = "debug")] pub fn stream_zip_entries_bytes( data: Vec, ) -> impl Stream< Item = Result< ( usize, StoredZipEntry, Arc, ), IOError, >, > { try_stream! { let reader = Arc::new(async_zip::base::read::mem::ZipFileReader::new(data).await?); let entries = reader.file().entries(); for index in 0..entries.len() { let entry = entries.get(index).expect("expected more zip entries"); yield (index, entry.clone(), reader.clone()); } } } #[tracing::instrument( level = "debug", skip(data, dest_path), fields( dest_path = %dest_path.as_ref().display() ) )] pub async fn unzip_bytes( data: Vec, dest_path: impl AsRef, ) -> PolyIOResult<()> { unzip_bytes_filtered(data, None:: bool>, dest_path).await } #[tracing::instrument( level = "debug", skip(data, filter_entries, dest_path), fields( dest_path = %dest_path.as_ref().display() ) )] pub async fn unzip_bytes_filtered( data: Vec, filter_entries: Option bool + Send + Sync>, dest_path: impl AsRef, ) -> PolyIOResult<()> { let stream = stream_zip_entries_bytes(data); pin_mut!(stream); while let Some((index, entry, reader)) = stream.try_next().await? { let file_name = entry.filename().as_str()?; if let Some(filter) = &filter_entries && !filter(file_name) { continue; } let path = dest_path.as_ref().join(crate::sanitize_path(file_name)); let is_dir = entry.dir()?; if is_dir { crate::create_dir_all(path).await?; } else { if let Some(parent) = path.parent() { crate::create_dir_all(parent).await?; } let entry_reader = reader.reader_without_entry(index).await?; let file = tokio::fs::File::create(&path).await?; let writer = tokio::io::BufWriter::new(file); futures_lite::io::copy(entry_reader, &mut writer.compat_write()).await?; } } Ok(()) } #[tracing::instrument( level = "debug", skip(zip_path, dest_path), fields( zip_path = %zip_path.as_ref().display(), dest_path = %dest_path.as_ref().display() ) )] pub async fn extract_zip( zip_path: impl AsRef, dest_path: impl AsRef, ) -> PolyIOResult<()> { extract_zip_filtered( zip_path, dest_path, None:: bool>, None:: String>, ) .await } #[tracing::instrument( level = "debug", skip(zip_path, dest_path, filter_entries, modify_entry_name), fields( zip_path = %zip_path.as_ref().display(), dest_path = %dest_path.as_ref().display() ) )] pub async fn extract_zip_filtered( zip_path: impl AsRef, dest_path: impl AsRef, filter_entries: Option bool + Send + Sync>, modify_entry_name: Option String>, ) -> PolyIOResult<()> { let zip_path = zip_path.as_ref(); let dest_path = dest_path.as_ref(); let reader = async_zip::tokio::read::fs::ZipFileReader::new(zip_path).await?; let entries = reader.file().entries(); for index in 0..entries.len() { let entry = entries.get(index).expect("expected more zip entries"); if let Some(filter) = &filter_entries && !filter(entry) { continue; } let old_name = entry.filename().as_str()?; let name: String = modify_entry_name .as_ref() .map_or_else(|| old_name.to_string(), |modify| modify(old_name)); let name = crate::sanitize_path(name); let path = dest_path.join(name); let is_dir = entry.dir()?; if is_dir { crate::create_dir_all(&path).await?; } else { if let Some(parent) = path.parent() { crate::create_dir_all(parent).await?; } let file = tokio::fs::File::create(&path).await?; let writer = tokio::io::BufWriter::new(file); let entry_reader = reader.reader_without_entry(index).await?; futures_lite::io::copy(entry_reader, &mut writer.compat_write()).await?; } } Ok(()) } /// Each matching entry is read fully into memory so this is for small files /// not large archives #[tracing::instrument( level = "debug", skip(zip_path, filter), fields(zip_path = %zip_path.as_ref().display()) )] pub async fn read_zip_file_entries( zip_path: impl AsRef, filter: impl Fn(&str) -> bool, ) -> PolyIOResult)>> { let zip_path = zip_path.as_ref(); let reader = async_zip::tokio::read::fs::ZipFileReader::new(zip_path).await?; let entries = reader.file().entries(); let mut out = Vec::new(); for index in 0..entries.len() { let entry = entries.get(index).expect("expected more zip entries"); let Ok(name) = entry.filename().as_str() else { continue; }; let name = name.to_string(); if entry.dir().unwrap_or(false) || !filter(&name) { continue; } let mut entry_reader = reader.reader_without_entry(index).await?; let mut data = Vec::new(); futures_lite::AsyncReadExt::read_to_end(&mut entry_reader, &mut data).await?; out.push((name, data)); } Ok(out) } /// Returns a zip file entry's bytes without reading the entire file into memory #[tracing::instrument(level = "debug", skip(reader))] pub async fn try_read_zip_entry_bytes(reader: R, file_name: &str) -> PolyIOResult> where R: tokio::io::AsyncRead + tokio::io::AsyncBufRead + tokio::io::AsyncSeek + Unpin, { use tokio_util::compat::TokioAsyncReadCompatExt; let compat = reader.compat(); let mut zip_reader = async_zip::base::read::seek::ZipFileReader::new(compat).await?; let index = zip_reader .file() .entries() .iter() .position(|entry| entry.filename().as_str().is_ok_and(|n| n == file_name)) .ok_or_else(|| IOError::FileNotFoundInZip { file_name: file_name.to_string(), })?; let mut entry_reader = zip_reader.reader_without_entry(index).await?; let mut data: Vec = Vec::new(); futures_lite::AsyncReadExt::read_to_end(&mut entry_reader, &mut data).await?; Ok(data) } #[tracing::instrument( level = "debug", skip(archive, dest), fields( archive = %archive.as_ref().display(), dest = %dest.as_ref().display(), ) )] pub async fn extract_tar_gz( archive: impl AsRef, dest: impl AsRef, ) -> PolyIOResult<()> { crate::create_dir_all(&dest).await?; let file = tokio::fs::File::open(archive).await?; let buf_reader = tokio::io::BufReader::new(file); let gzip_decoder = async_compression::tokio::bufread::GzipDecoder::new(buf_reader); let mut tar_archive = tokio_tar::Archive::new(gzip_decoder); tar_archive.unpack(dest).await?; Ok(()) }