From 10032ca13f74bb90597a9ad05f9cabe53ab7b3cd Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Mon, 23 Feb 2026 10:23:51 +0100
Subject: [PATCH] feat(ml): SHA-256 checksum validation for model checkpoint
 integrity

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 ml/src/checkpoint/mod.rs        |  2 +-
 ml/src/checkpoint/validation.rs | 78 +++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/ml/src/checkpoint/mod.rs b/ml/src/checkpoint/mod.rs
index 921be828d..29f42bf71 100644
--- a/ml/src/checkpoint/mod.rs
+++ b/ml/src/checkpoint/mod.rs
@@ -64,7 +64,7 @@ pub use signer::{CheckpointSigner, SignatureInfo};
 #[cfg(feature = "s3-storage")]
 pub use storage::S3CheckpointStorage;
 pub use storage::{CheckpointStorage, FileSystemStorage, MemoryStorage, StorageStats};
-pub use validation::ValidationManager;
+pub use validation::{verify_checksum, write_checksum, ValidationManager};
 pub use versioning::VersionManager;
 
 /// Checkpoint format options
diff --git a/ml/src/checkpoint/validation.rs b/ml/src/checkpoint/validation.rs
index 89387439e..eb3437c1e 100644
--- a/ml/src/checkpoint/validation.rs
+++ b/ml/src/checkpoint/validation.rs
@@ -3,6 +3,7 @@
 //! Provides checksum validation and corruption detection for checkpoints.
 
 use std::collections::HashMap;
+use std::path::Path;
 
 use sha2::{Digest, Sha256};
 use tracing::{debug, error, warn};
@@ -10,6 +11,43 @@ use tracing::{debug, error, warn};
 use super::{CheckpointMetadata, ModelType};
 use crate::MLError;
 
+/// Write SHA-256 checksum sidecar file alongside a safetensors checkpoint.
+/// Creates `{path}.sha256` containing the hex digest.
+pub fn write_checksum(safetensors_path: &Path) -> Result<(), MLError> {
+    let bytes = std::fs::read(safetensors_path).map_err(|e| {
+        MLError::CheckpointError(format!("Failed to read file for checksum: {}", e))
+    })?;
+    let hash = Sha256::digest(&bytes);
+    let hex = format!("{:x}", hash);
+    let checksum_path = safetensors_path.with_extension("sha256");
+    std::fs::write(&checksum_path, hex.as_bytes()).map_err(|e| {
+        MLError::CheckpointError(format!("Failed to write checksum: {}", e))
+    })?;
+    Ok(())
+}
+
+/// Verify SHA-256 checksum of a safetensors file against its `.sha256` sidecar.
+/// Returns `Ok(true)` if valid, `Ok(false)` if mismatch.
+/// Returns `Ok(true)` if no sidecar exists (backwards compatible).
+pub fn verify_checksum(safetensors_path: &Path) -> Result<bool, MLError> {
+    let checksum_path = safetensors_path.with_extension("sha256");
+    if !checksum_path.exists() {
+        warn!(
+            "No checksum file for {}, skipping integrity check",
+            safetensors_path.display()
+        );
+        return Ok(true);
+    }
+    let expected = std::fs::read_to_string(&checksum_path).map_err(|e| {
+        MLError::CheckpointError(format!("Failed to read checksum: {}", e))
+    })?;
+    let bytes = std::fs::read(safetensors_path).map_err(|e| {
+        MLError::CheckpointError(format!("Failed to read file for verification: {}", e))
+    })?;
+    let actual = format!("{:x}", Sha256::digest(&bytes));
+    Ok(actual.trim() == expected.trim())
+}
+
 /// Validation manager for checkpoint integrity
 #[derive(Debug)]
 pub struct ValidationManager {
@@ -522,4 +560,44 @@ mod tests {
         assert!(summary.contains("1 errors"));
         assert!(summary.contains("1 warnings"));
     }
+
+    #[test]
+    fn test_sidecar_checksum_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
+        let dir = tempfile::TempDir::new()?;
+        let path = dir.path().join("test_model.safetensors");
+        std::fs::write(&path, b"fake model data")?;
+
+        write_checksum(&path)?;
+
+        let checksum_path = path.with_extension("sha256");
+        assert!(checksum_path.exists());
+
+        assert!(verify_checksum(&path)?);
+        Ok(())
+    }
+
+    #[test]
+    fn test_sidecar_checksum_detects_corruption() -> Result<(), Box<dyn std::error::Error>> {
+        let dir = tempfile::TempDir::new()?;
+        let path = dir.path().join("test_model.safetensors");
+        std::fs::write(&path, b"original data")?;
+
+        write_checksum(&path)?;
+
+        // Corrupt the file
+        std::fs::write(&path, b"corrupted data")?;
+        assert!(!verify_checksum(&path)?);
+        Ok(())
+    }
+
+    #[test]
+    fn test_verify_without_sidecar_returns_true() -> Result<(), Box<dyn std::error::Error>> {
+        let dir = tempfile::TempDir::new()?;
+        let path = dir.path().join("test_model.safetensors");
+        std::fs::write(&path, b"no checksum file")?;
+
+        // No .sha256 file exists -- should return true (backwards compatible)
+        assert!(verify_checksum(&path)?);
+        Ok(())
+    }
 }