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            6 : pub async fn extract_zst_tarball(
      71            6 :     path: &Utf8Path,
      72            6 :     tarball: impl AsyncBufRead + Unpin,
      73            6 : ) -> Result<()> {
      74            6 :     let decoder = Box::pin(ZstdDecoder::new(tarball));
      75            6 :     let mut archive = Archive::new(decoder);
      76        31304 :     archive.unpack(path).await?;
      77            6 :     Ok(())
      78            6 : }