Line data    Source code

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