feat: extract hashes, make updatecheck use almost no network

Author: Taxrosdev

src/bin/pkgsmgr-packager.rs

@@ -1,17 +1,16 @@
 use async_compression::tokio::write::ZstdEncoder;
 use clap::Parser;
-use futures_util::future::try_join_all;
+use nix::fcntl::{AT_FDCWD, RenameFlags, renameat2};
 use std::boxed::Box;
 use std::collections::HashMap;
-use std::hash::Hasher;
-use std::io::Write;
 use std::os::unix::fs::MetadataExt;
 use std::path::{Path, PathBuf};
 use tokio::fs;
 use tokio::fs::File;
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt};
 
 use pkgsmgr::types::*;
+use pkgsmgr::utils::Hasher;
 
 #[derive(Parser)]
 #[command(version, about, long_about = None)]
@@ -55,35 +54,36 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     println!("Beginning hashing and compressing...");
     let mut hashes = HashMap::new();
 
-    for file_group in files.chunks_mut(256) {
-        let futures: Vec<_> = file_group
-            .iter_mut()
-            .map(|file_path| {
-                compress_file(
-                    file_path.to_path_buf(),
-                    args.hash,
-                    args.compression,
-                    chunks_path,
-                )
-            })
-            .collect();
-
-        for (path, hash) in try_join_all(futures).await? {
-            hashes.insert(path, hash);
-        }
+    for file_path in &files {
+        let hash = hash_file(file_path, args.hash).await?;
+
+        compress(file_path, args.compression, chunks_path, &hash).await?;
+
+        if fs::hard_link(&file_path, chunks_path.join(&hash))
+            .await
+            .is_err()
+        {
+            fs::copy(&file_path, chunks_path.join(&hash)).await?;
+        };
+
+        hashes.insert(file_path, hash);
     }
 
     println!("Generating manifest...");
     let mut manifest = "".to_string();
 
     match args.compression {
-        Compression::Zstd => manifest += "Compressed: zstd\n",
+        Compression::Zstd => manifest += "Compression: zstd\n",
         Compression::None => (),
     }
+    match args.hash {
+        HashType::Blake3 => manifest += "Hasher: blake3\n",
+        HashType::Xxh3_128 => manifest += "Hasher: xxh3_128\n",
+    }
 
     manifest += "---\n";
 
-    for file in files {
+    for file in &files {
         let hash = hashes
             .get(&file)
             .expect("tried adding file to manifest that has no hash");
@@ -94,31 +94,52 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         let size = metadata.size() / 1024;
         let path = file
             .strip_prefix(&args.input_path)
-            .expect("tried adding file to manifest that is outside of input_path");
+            .expect("tried adding file to manifest that is outside of input_path")
+            .to_str()
+            .unwrap();
+
+        manifest += &format!("{mode};{size};{hash};{path}\n");
+    }
+
+    // Atomically replace on-disk manifest
+    let hash = &blake3::hash(manifest.as_bytes()).to_hex().to_string();
+    let tmp_link_path = args.output_path.join("manifest.tmp");
+    let main_link_path = args.output_path.join("manifest");
+    let manifest_path = args.output_path.join(hash);
 
-        manifest += &format!("{mode};{size};{hash};{path:?}\n");
+    fs::write(manifest_path, manifest).await?;
+    fs::write(&tmp_link_path, hash).await?;
+
+    if !&main_link_path.exists() {
+        fs::write(&main_link_path, "").await?;
     }
 
-    fs::write(args.output_path.join("manifest"), manifest).await?;
+    renameat2(
+        AT_FDCWD,
+        &tmp_link_path,
+        AT_FDCWD,
+        &main_link_path,
+        RenameFlags::RENAME_EXCHANGE,
+    )?;
+
+    fs::remove_file(&tmp_link_path).await?;
 
     Ok(())
 }
 
-async fn compress_file(
-    file_path: PathBuf,
-    hash: HashType,
-    compression: Compression,
-    chunks_path: &Path,
-) -> Result<(PathBuf, String), Box<dyn std::error::Error>> {
+async fn hash_file(
+    file_path: &Path,
+    hash_method: HashType,
+) -> Result<String, Box<dyn std::error::Error>> {
     let mut source_file = match File::open(&file_path).await {
         Ok(file) => file,
         Err(e) => {
             eprintln!("couldn't open source file: {}", file_path.display());
             panic!("{e}")
         }
     };
-    let mut xxh_hasher = xxhash_rust::xxh3::Xxh3Default::new();
-    let mut blake3_hasher = blake3::Hasher::new();
+
+    let mut hasher = Hasher::new(hash_method);
 
     let mut buf = [0; 8192];
     loop {
@@ -128,48 +149,55 @@ async fn compress_file(
         }
 
         let chunk = &buf[0..n];
-        match hash {
-            HashType::Blake3 => blake3_hasher.write_all(chunk)?,
-            HashType::Xxh3_128 => xxh_hasher.write_all(chunk)?,
-        }
+        hasher.write(chunk);
     }
 
-    let hash = match hash {
-        HashType::Xxh3_128 => hex::encode(xxh_hasher.finish().to_le_bytes()),
-        HashType::Blake3 => blake3_hasher.finalize().to_hex().to_string(),
-    };
+    let hash = hasher.digest();
 
-    let chunk_path = &chunks_path.join(&hash);
+    Ok(hash)
+}
 
-    if !chunk_path.exists() {
+async fn compress(
+    file_path: &Path,
+    compression: Compression,
+    chunks_path: &Path,
+    hash: &str,
+) -> Result<(), std::io::Error> {
+    let compressed_chunk_filename = match compression {
+        Compression::Zstd => format!("{hash}.zstd"),
+        Compression::None => panic!("Tried to compress on a non-compressable request."),
+    };
+    let compressed_chunk_path = &chunks_path.join(compressed_chunk_filename);
+
+    if !compressed_chunk_path.exists() {
         let mut source_file = File::open(&file_path).await.unwrap();
         let temp_file_path = temp_file::TempFile::new()?;
-        let temp_file = File::create(&temp_file_path).await?;
-
-        if compression == Compression::Zstd {
-            let mut zstd = ZstdEncoder::new(temp_file);
+        let mut temp_file = File::create(&temp_file_path).await?;
 
-            let mut buf = [0; 8192];
-            loop {
-                let n = source_file.read(&mut buf).await?;
-                if n == 0 {
-                    break;
-                }
+        let mut compressor: Box<dyn AsyncWrite + Sync + Unpin> = match compression {
+            Compression::Zstd => Box::new(ZstdEncoder::new(&mut temp_file)),
+            Compression::None => panic!("Tried to copmress on a non-compressable request."),
+        };
 
-                zstd.write_all(&buf).await?;
+        let mut buf = [0; 8192];
+        loop {
+            let n = source_file.read(&mut buf).await?;
+            if n == 0 {
+                break;
             }
 
-            let mut zstd_path = chunk_path.clone();
-            zstd_path.set_extension("zstd");
-            fs::copy(temp_file_path, zstd_path).await?;
+            compressor.write_all(&buf[0..n]).await?;
         }
 
-        if fs::hard_link(&file_path, chunk_path).await.is_err() {
-            fs::copy(&file_path, chunk_path).await?;
-        };
+        // Finish compressing
+        compressor.flush().await?;
+        compressor.shutdown().await?;
+
+        // Move compressed from memory and onto disk
+        fs::copy(temp_file_path, compressed_chunk_path).await?;
 
-        println!("Created chunk from path {file_path:?}");
+        println!("Compressed chunk from path {file_path:?}");
     };
 
-    Ok((file_path, hash))
+    Ok(())
 }

src/bin/pkgsmgr-updater.rs

@@ -5,8 +5,8 @@ use std::path::PathBuf;
 use std::sync::LazyLock;
 
 use pkgsmgr::chunks::{chunk_filename, clean_old_chunks, install_chunk};
-use pkgsmgr::manifest::{build_tree, parse_manifest, update_manifest};
-use pkgsmgr::types::Compression;
+use pkgsmgr::manifest::{build_tree, parse_manifest, try_update_manifest_hash, update_manifest};
+use pkgsmgr::types::{Compression, HashType};
 use pkgsmgr::utils::get;
 
 static MAJOR_VERSION: LazyLock<usize> =
@@ -20,6 +20,9 @@ struct Args {
     repo_url: String,
     #[arg(long)]
     root_path: Option<PathBuf>,
+    #[arg(long)]
+    /// Useful for installers, where the installation media may contain relevant chunks already
+    additional_cache_path: Option<PathBuf>,
 }
 
 #[tokio::main]
@@ -34,22 +37,28 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let manifests_path = &internal_path.join("manifests");
     fs::create_dir_all(manifests_path)?;
 
-    let manifest_raw = get(&format!("{}/manifest", &args.repo_url))
+    let manifest_hash = get(&format!("{}/manifest", &args.repo_url))
         .await?
+        .error_for_status()?
         .text()
-        .await
-        .expect("server responded with 200, yet not valid utf8 text.");
+        .await?;
 
-    // Quit early if nothing has changed
-    if !update_manifest(&manifest_raw, manifests_path)
-        .expect("could not update local manifest cache")
-    {
-        return Ok(());
+    if !try_update_manifest_hash(manifests_path, &manifest_hash)? {
+        println!("[INFO] Skipping, no update found.");
+        std::process::exit(0);
     }
+    println!("[INFO] Update found, downloading manifest...");
+
+    let manifest_raw = get(&format!("{}/{}", &args.repo_url, manifest_hash))
+        .await?
+        .text()
+        .await
+        .expect("server responded with 200, yet not valid utf8 text.");
 
     let (headers, chunklist) = parse_manifest(&manifest_raw);
 
     let mut compression = Compression::None;
+    let mut hasher = HashType::Blake3;
 
     for (key, value) in headers {
         match key {
@@ -78,6 +87,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                     eprintln!("Unknown compression requested: {}", value);
                 }
             },
+            "Hasher" => match value.to_lowercase().as_str() {
+                "blake3" => {
+                    hasher = HashType::Blake3;
+                }
+                "xxh3_128" => hasher = HashType::Xxh3_128,
+                _ => {
+                    eprintln!("Unknown compression requested: {}", value);
+                }
+            },
             _ => {
                 eprintln!("[WARNING] Unknown header: {key}");
             }
@@ -86,25 +104,41 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     // Install all chunks in chunklist before doing anything else.
     for chunk in &chunklist {
-        if !chunks_path.join(chunk_filename(chunk)).exists() {
-            install_chunk(chunk, &args.repo_url, chunks_path, &compression)
+        let chunk_path = chunks_path.join(chunk_filename(chunk));
+
+        if !chunk_path.exists() {
+            install_chunk(chunk, &args.repo_url, chunks_path, &compression, hasher)
                 .await
                 .expect("could not download chunk");
         }
     }
 
+    // Quit early if nothing has changed
+    if !update_manifest(&manifest_raw, manifests_path)
+        .expect("could not update local manifest cache")
+    {
+        return Ok(());
+    }
+
     build_tree(staging_path, chunks_path, &chunklist).expect("could not build staging");
 
-    let cwd = AT_FDCWD;
+    println!("[INFO] Swapping tree...");
+
+    let usr_path = root_path.join("usr");
+    if !usr_path.exists() {
+        fs::create_dir_all(&usr_path)?;
+    }
 
     renameat2(
-        cwd,
+        AT_FDCWD,
         staging_path,
-        cwd,
-        &root_path.join("usr"),
+        AT_FDCWD,
+        &usr_path,
         RenameFlags::RENAME_EXCHANGE,
     )?;
 
+    println!("[INFO] Cleaning up old chunks...");
+
     let freed_bytes =
         clean_old_chunks(manifests_path, chunks_path).expect("could not free old chunks");
     println!("Freed {}kb", freed_bytes / 1024);