Line data Source code
1 : use std::io::SeekFrom;
2 :
3 : use anyhow::{Context, Result};
4 : use async_compression::{
5 : tokio::{bufread::ZstdDecoder, write::ZstdEncoder},
6 : zstd::CParameter,
7 : Level,
8 : };
9 : use camino::Utf8Path;
10 : use nix::NixPath;
11 : use tokio::{
12 : fs::{File, OpenOptions},
13 : io::AsyncBufRead,
14 : io::AsyncSeekExt,
15 : io::AsyncWriteExt,
16 : };
17 : use tokio_tar::{Archive, Builder, HeaderMode};
18 : use walkdir::WalkDir;
19 :
20 : /// Creates a Zstandard tarball.
21 0 : pub async fn create_zst_tarball(path: &Utf8Path, tarball: &Utf8Path) -> Result<(File, u64)> {
22 0 : let file = OpenOptions::new()
23 0 : .create(true)
24 0 : .truncate(true)
25 0 : .read(true)
26 0 : .write(true)
27 0 : .open(&tarball)
28 0 : .await
29 0 : .with_context(|| format!("tempfile creation {tarball}"))?;
30 :
31 0 : let mut paths = Vec::new();
32 0 : for entry in WalkDir::new(path) {
33 0 : let entry = entry?;
34 0 : let metadata = entry.metadata().expect("error getting dir entry metadata");
35 0 : // Also allow directories so that we also get empty directories
36 0 : if !(metadata.is_file() || metadata.is_dir()) {
37 0 : continue;
38 0 : }
39 0 : let path = entry.into_path();
40 0 : paths.push(path);
41 : }
42 : // Do a sort to get a more consistent listing
43 0 : paths.sort_unstable();
44 0 : let zstd = ZstdEncoder::with_quality_and_params(
45 0 : file,
46 0 : Level::Default,
47 0 : &[CParameter::enable_long_distance_matching(true)],
48 0 : );
49 0 : let mut builder = Builder::new(zstd);
50 0 : // Use reproducible header mode
51 0 : builder.mode(HeaderMode::Deterministic);
52 0 : for p in paths {
53 0 : let rel_path = p.strip_prefix(path)?;
54 0 : if rel_path.is_empty() {
55 : // The top directory should not be compressed,
56 : // the tar crate doesn't like that
57 0 : continue;
58 0 : }
59 0 : builder.append_path_with_name(&p, rel_path).await?;
60 : }
61 0 : let mut zstd = builder.into_inner().await?;
62 0 : zstd.shutdown().await?;
63 0 : let mut compressed = zstd.into_inner();
64 0 : let compressed_len = compressed.metadata().await?.len();
65 0 : compressed.seek(SeekFrom::Start(0)).await?;
66 0 : Ok((compressed, compressed_len))
67 0 : }
68 :
69 : /// Creates a Zstandard tarball.
70 2 : pub async fn extract_zst_tarball(
71 2 : path: &Utf8Path,
72 2 : tarball: impl AsyncBufRead + Unpin,
73 2 : ) -> Result<()> {
74 2 : let decoder = Box::pin(ZstdDecoder::new(tarball));
75 2 : let mut archive = Archive::new(decoder);
76 10670 : archive.unpack(path).await?;
77 2 : Ok(())
78 2 : }
|