Add hashes to metadata

Author: danieldk

Cargo.lock

@@ -25,6 +25,9 @@ use init::{run_init, InitArgs};
 mod list_variants;
 use list_variants::list_variants;
 
+mod sign;
+use sign::sign;
+
 mod upload;
 use upload::{run_upload, RepoTypeArg, UploadArgs};
 
@@ -201,6 +204,12 @@ enum Commands {
         arch: bool,
     },
 
+    /// Sign the builds of a kernel.
+    Sign {
+        #[arg(name = "KERNEL_DIR")]
+        kernel_dir: Option<PathBuf>,
+    },
+
     /// Spawn a kernel test shell.
     Testshell {
         #[arg(name = "KERNEL_DIR")]
@@ -362,6 +371,7 @@ fn main() -> Result<()> {
             variant,
         ),
         Commands::ListVariants { kernel_dir, arch } => list_variants(kernel_dir, arch),
+        Commands::Sign { kernel_dir } => sign(kernel_dir),
         Commands::Testshell {
             kernel_dir,
             variant,

kernel-builder/src/main.rs

@@ -49,6 +49,7 @@ pub fn write_metadata(
                 archs: None,
                 backend_type: *backend,
             },
+            source_digest: None,
         };
 
         serde_json::to_writer_pretty(writer, &metadata)?;

kernel-builder/src/pyproject/common.rs

@@ -0,0 +1,45 @@
+use std::{
+    fs::File,
+    io::{BufReader, BufWriter},
+    path::PathBuf,
+};
+
+use eyre::{Context, Result};
+use kernels_data::metadata::{DigestAlgorithm, Metadata, SourceDigest};
+
+use crate::util::{check_or_infer_kernel_dir, discover_variants};
+
+pub fn sign(kernel_dir: Option<PathBuf>) -> Result<()> {
+    let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?;
+    let (build_dir, variants) = discover_variants(&kernel_dir)?;
+
+    for variant in variants {
+        let variant_path = build_dir.join(&variant);
+        let metadata_path = variant_path.join("metadata.json");
+
+        eprintln!(
+            "Signing variant `{}`...",
+            variant.file_name().unwrap().to_string_lossy()
+        );
+
+        let f = File::open(&metadata_path).context(format!(
+            "Cannot open `{}` for reading",
+            metadata_path.to_string_lossy()
+        ))?;
+        let mut metadata: Metadata = serde_json::from_reader(BufReader::new(f))?;
+
+        let source_digest = SourceDigest::update_hashes(DigestAlgorithm::SHA256, &variant_path)?;
+        metadata.source_digest = Some(source_digest);
+
+        let f = File::create(&metadata_path).context(format!(
+            "Cannot open `{}` for writing file hashes",
+            metadata_path.to_string_lossy()
+        ))?;
+        serde_json::to_writer_pretty(BufWriter::new(f), &metadata).context(format!(
+            "Cannot write updated metadata to `{}`",
+            metadata_path.to_string_lossy()
+        ))?;
+    }
+
+    Ok(())
+}

kernel-builder/src/sign.rs

@@ -49,6 +49,41 @@ pub(crate) fn check_or_infer_target_dir(
     }
 }
 
+/// Discover build variant directories (contain `metadata.json`).
+/// Checks `result` symlink (Nix store output) first, then falls back to `build/`.
+pub(crate) fn discover_variants(kernel_dir: &Path) -> Result<(PathBuf, Vec<PathBuf>)> {
+    let candidates = [
+        kernel_dir.join("result"),
+        kernel_dir.join("build"),
+        kernel_dir.to_path_buf(),
+    ];
+
+    for candidate in &candidates {
+        if !candidate.is_dir() {
+            continue;
+        }
+
+        let mut variants: Vec<PathBuf> = fs::read_dir(candidate)
+            .wrap_err_with(|| format!("Cannot read `{}`", candidate.display()))?
+            .filter_map(|e| e.ok())
+            .map(|e| e.path())
+            .filter(|p| p.is_dir() && p.join("metadata.json").is_file())
+            .collect();
+
+        if !variants.is_empty() {
+            variants.sort();
+            return Ok((candidate.clone(), variants));
+        }
+    }
+
+    bail!(
+        "No build variants found in `{}`, `{}`, or `{}`",
+        candidates[0].display(),
+        candidates[1].display(),
+        candidates[2].display(),
+    );
+}
+
 pub(crate) fn parse_and_validate(kernel_dir: impl AsRef<Path>) -> Result<BuildCompat> {
     let build_toml = kernel_dir.as_ref().join("build.toml");
     let mut toml_data = String::new();

kernel-builder/src/util.rs

@@ -8,11 +8,15 @@ license = "Apache-2.0"
 repository = "https://github.com/huggingface/kernels"
 
 [dependencies]
+base64 = "0.22"
+digest = "0.11"
 eyre = "0.6.12"
 itertools = "0.13"
 regex = "1"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 serde-value = "0.7"
+sha2 = "0.11"
+walkdir = "2"
 thiserror = "1"
 url = { version = "2", features = ["serde"] }

kernels-data/Cargo.toml

@@ -1,7 +1,12 @@
 use std::str::FromStr;
+use std::{collections::BTreeMap, fs, path::Path};
 
-use eyre::Result;
+use base64::prelude::{BASE64_STANDARD, Engine as _};
+use digest::{Digest, DynDigest};
+use eyre::{Context, Result};
 use serde::{Deserialize, Serialize};
+use sha2::{Sha256, Sha512};
+use walkdir::WalkDir;
 
 use crate::config::{Backend, KernelName};
 
@@ -26,6 +31,75 @@ pub struct Metadata {
     pub upstream: Option<url::Url>,
     pub python_depends: Vec<String>,
     pub backend: BackendInfo,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub source_digest: Option<SourceDigest>,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct SourceDigest {
+    algorithm: DigestAlgorithm,
+    files: BTreeMap<String, String>,
+}
+
+impl SourceDigest {
+    pub fn update_hashes(
+        digest_algorithm: DigestAlgorithm,
+        variant_path: impl AsRef<Path>,
+    ) -> Result<Self> {
+        let variant_path = variant_path.as_ref();
+
+        let mut files = BTreeMap::new();
+        for entry in WalkDir::new(variant_path) {
+            let entry = entry.wrap_err("Failed to read directory entry for hashing")?;
+            if !entry.file_type().is_file()
+                || entry.path().extension().and_then(|e| e.to_str()) == Some("pyc")
+            {
+                continue;
+            }
+
+            let path = entry.path();
+
+            // Read and hash contents.
+            let contents = fs::read(path)
+                .wrap_err_with(|| format!("Cannot read `{}` for hashing", path.display()))?;
+            let mut hasher: Box<dyn DynDigest> = digest_algorithm.into();
+            hasher.update(&contents);
+
+            let relative_path = path.strip_prefix(variant_path).wrap_err_with(|| {
+                format!("Cannot strip prefix from `{}` for hashing", path.display())
+            })?;
+
+            // Normalize Windows directory separators.
+            let relative_path_str = relative_path.to_string_lossy().replace('\\', "/");
+
+            let hash_base64 = BASE64_STANDARD.encode(hasher.finalize_reset());
+
+            files.insert(relative_path_str, hash_base64);
+        }
+
+        Ok(SourceDigest {
+            files,
+            algorithm: digest_algorithm,
+        })
+    }
+}
+
+#[derive(Copy, Clone, Debug, Deserialize, Serialize)]
+pub enum DigestAlgorithm {
+    #[serde(rename = "sha256")]
+    SHA256,
+
+    #[serde(rename = "sha512")]
+    SHA512,
+}
+
+impl From<DigestAlgorithm> for Box<dyn DynDigest> {
+    fn from(digest_algorithm: DigestAlgorithm) -> Box<dyn DynDigest> {
+        match digest_algorithm {
+            DigestAlgorithm::SHA256 => Box::new(Sha256::new()),
+            DigestAlgorithm::SHA512 => Box::new(Sha512::new()),
+        }
+    }
 }
 
 impl Metadata {