Rewrite untar to take a visitor & simplify

signature of `download_and_extract_with_filter`

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>
This commit is contained in:
Jiahao XU 2022-06-11 20:31:46 +10:00
parent 5a43ee2681
commit 90a96cabc9
No known key found for this signature in database
GPG key ID: 591C0B03040416D6
4 changed files with 64 additions and 38 deletions

View file

@ -114,7 +114,7 @@ pub async fn fetch_crate_cratesio(
Url::parse(&crate_url)?, Url::parse(&crate_url)?,
TarBasedFmt::Tgz, TarBasedFmt::Tgz,
&temp_dir, &temp_dir,
Some(move |path: &Path| path == cargo_toml || path == main || path.starts_with(&bin)), move |path: &Path| path == cargo_toml || path == main || path.starts_with(&bin),
) )
.await?; .await?;

View file

@ -92,7 +92,7 @@ pub async fn download_and_extract_with_filter<
url: Url, url: Url,
fmt: TarBasedFmt, fmt: TarBasedFmt,
path: P, path: P,
filter: Option<Filter>, filter: Filter,
) -> Result<(), BinstallError> { ) -> Result<(), BinstallError> {
debug!("Downloading from: '{url}'"); debug!("Downloading from: '{url}'");

View file

@ -1,11 +1,14 @@
use std::fmt::Debug; use std::fmt::Debug;
use std::fs; use std::fs;
use std::io::{self, Seek, Write}; use std::io::{self, Read, Seek, Write};
use std::path::Path; use std::path::{Path, PathBuf};
use std::sync::Arc;
use bytes::Bytes; use bytes::Bytes;
use futures_util::stream::{Stream, StreamExt}; use futures_util::stream::{Stream, StreamExt};
use log::debug;
use scopeguard::{guard, ScopeGuard}; use scopeguard::{guard, ScopeGuard};
use tar::Entries;
use tempfile::tempfile; use tempfile::tempfile;
use tokio::{ use tokio::{
sync::mpsc, sync::mpsc,
@ -209,17 +212,42 @@ pub async fn extract_tar_based_stream_with_filter<
stream: impl Stream<Item = Result<Bytes, E>> + Unpin, stream: impl Stream<Item = Result<Bytes, E>> + Unpin,
output: &Path, output: &Path,
fmt: TarBasedFmt, fmt: TarBasedFmt,
filter: Option<Filter>, filter: Filter,
) -> Result<(), BinstallError> ) -> Result<(), BinstallError>
where where
BinstallError: From<E>, BinstallError: From<E>,
{ {
let path = output.to_owned(); struct Visitor<F>(F, Arc<PathBuf>);
impl<F: FnMut(&Path) -> bool + Send + 'static> TarEntriesVisitor for Visitor<F> {
fn visit<R: Read>(&mut self, entries: Entries<'_, R>) -> Result<(), BinstallError> {
for res in entries {
let mut entry = res?;
let entry_path = entry.path()?;
if self.0(&entry_path) {
debug!("Extracting {entry_path:#?}");
let dst = self.1.join(entry_path);
fs::create_dir_all(dst.parent().unwrap())?;
entry.unpack(dst)?;
}
}
Ok(())
}
}
let path = Arc::new(output.to_owned());
let visitor = Visitor(filter, path.clone());
extract_impl(stream, move |mut rx| { extract_impl(stream, move |mut rx| {
fs::create_dir_all(path.parent().unwrap())?; fs::create_dir_all(path.parent().unwrap())?;
extract_compressed_from_readable(ReadableRx::new(&mut rx), fmt, &path, filter) extract_compressed_from_readable(ReadableRx::new(&mut rx), fmt, &*path, Some(visitor))
}) })
.await .await
} }
@ -232,12 +260,20 @@ pub async fn extract_tar_based_stream<E>(
where where
BinstallError: From<E>, BinstallError: From<E>,
{ {
struct DummyVisitor;
impl TarEntriesVisitor for DummyVisitor {
fn visit<R: Read>(&mut self, _entries: Entries<'_, R>) -> Result<(), BinstallError> {
unimplemented!()
}
}
let path = output.to_owned(); let path = output.to_owned();
extract_impl(stream, move |mut rx| { extract_impl(stream, move |mut rx| {
fs::create_dir_all(path.parent().unwrap())?; fs::create_dir_all(path.parent().unwrap())?;
extract_compressed_from_readable::<fn(&Path) -> bool, _>( extract_compressed_from_readable::<DummyVisitor, _>(
ReadableRx::new(&mut rx), ReadableRx::new(&mut rx),
fmt, fmt,
&path, &path,

View file

@ -1,44 +1,34 @@
use std::fs::{self, File}; use std::fs::File;
use std::io::{BufRead, Read}; use std::io::{BufRead, Read};
use std::path::Path; use std::path::Path;
use flate2::bufread::GzDecoder; use flate2::bufread::GzDecoder;
use log::debug; use log::debug;
use tar::Archive; use tar::{Archive, Entries};
use xz2::bufread::XzDecoder; use xz2::bufread::XzDecoder;
use zip::read::ZipArchive; use zip::read::ZipArchive;
use zstd::stream::Decoder as ZstdDecoder; use zstd::stream::Decoder as ZstdDecoder;
use crate::{BinstallError, TarBasedFmt}; use crate::{BinstallError, TarBasedFmt};
/// * `filter` - If Some, then it will pass the path of the file to it pub trait TarEntriesVisitor {
/// and only extract ones which filter returns `true`. fn visit<R: Read>(&mut self, entries: Entries<'_, R>) -> Result<(), BinstallError>;
/// Note that this is a best-effort and it only works when `fmt` }
/// is not `PkgFmt::Bin` or `PkgFmt::Zip`.
fn untar<Filter: FnMut(&Path) -> bool>( /// * `f` - If Some, then this function will pass
dat: impl Read, /// the entries of the `dat` to it and let it decides
/// what to do with the tar.
fn untar<R: Read, V: TarEntriesVisitor>(
dat: R,
path: &Path, path: &Path,
filter: Option<Filter>, visitor: Option<V>,
) -> Result<(), BinstallError> { ) -> Result<(), BinstallError> {
let mut tar = Archive::new(dat); let mut tar = Archive::new(dat);
if let Some(mut filter) = filter { if let Some(mut visitor) = visitor {
debug!("Untaring with filter"); debug!("Untaring with filter");
for res in tar.entries()? { visitor.visit(tar.entries()?)?;
let mut entry = res?;
let entry_path = entry.path()?;
if filter(&entry_path) {
debug!("Extracting {entry_path:#?}");
let dst = path.join(entry_path);
fs::create_dir_all(dst.parent().unwrap())?;
entry.unpack(dst)?;
}
}
} else { } else {
debug!("Untaring entire tar"); debug!("Untaring entire tar");
tar.unpack(path)?; tar.unpack(path)?;
@ -56,11 +46,11 @@ fn untar<Filter: FnMut(&Path) -> bool>(
/// and only extract ones which filter returns `true`. /// and only extract ones which filter returns `true`.
/// Note that this is a best-effort and it only works when `fmt` /// Note that this is a best-effort and it only works when `fmt`
/// is not `PkgFmt::Bin` or `PkgFmt::Zip`. /// is not `PkgFmt::Bin` or `PkgFmt::Zip`.
pub(crate) fn extract_compressed_from_readable<Filter: FnMut(&Path) -> bool, R: BufRead>( pub(crate) fn extract_compressed_from_readable<V: TarEntriesVisitor, R: BufRead>(
dat: R, dat: R,
fmt: TarBasedFmt, fmt: TarBasedFmt,
path: &Path, path: &Path,
filter: Option<Filter>, visitor: Option<V>,
) -> Result<(), BinstallError> { ) -> Result<(), BinstallError> {
use TarBasedFmt::*; use TarBasedFmt::*;
@ -69,21 +59,21 @@ pub(crate) fn extract_compressed_from_readable<Filter: FnMut(&Path) -> bool, R:
// Extract to install dir // Extract to install dir
debug!("Extracting from tar archive to `{path:?}`"); debug!("Extracting from tar archive to `{path:?}`");
untar(dat, path, filter)? untar(dat, path, visitor)?
} }
Tgz => { Tgz => {
// Extract to install dir // Extract to install dir
debug!("Decompressing from tgz archive to `{path:?}`"); debug!("Decompressing from tgz archive to `{path:?}`");
let tar = GzDecoder::new(dat); let tar = GzDecoder::new(dat);
untar(tar, path, filter)?; untar(tar, path, visitor)?;
} }
Txz => { Txz => {
// Extract to install dir // Extract to install dir
debug!("Decompressing from txz archive to `{path:?}`"); debug!("Decompressing from txz archive to `{path:?}`");
let tar = XzDecoder::new(dat); let tar = XzDecoder::new(dat);
untar(tar, path, filter)?; untar(tar, path, visitor)?;
} }
Tzstd => { Tzstd => {
// Extract to install dir // Extract to install dir
@ -94,7 +84,7 @@ pub(crate) fn extract_compressed_from_readable<Filter: FnMut(&Path) -> bool, R:
// as &[] by ZstdDecoder::new, thus ZstdDecoder::new // as &[] by ZstdDecoder::new, thus ZstdDecoder::new
// should not return any error. // should not return any error.
let tar = ZstdDecoder::with_buffer(dat)?; let tar = ZstdDecoder::with_buffer(dat)?;
untar(tar, path, filter)?; untar(tar, path, visitor)?;
} }
}; };