fix(heal): require current checkpoint digest

crates/heal/src/heal/resume/checkpoint.rs

@@ -32,11 +32,12 @@ use super::{
 
 const EVENT_HEAL_CHECKPOINT_STATE: &str = "heal_checkpoint_state";
 const RESUME_CHECKPOINT_DIGEST_FILE: &str = "ahm_checkpoint.sha256";
+const CHECKPOINT_PER_VERSION_SCHEMA: u32 = 5;
 
 /// Current on-disk schema version for `ResumeCheckpoint`. Same rationale as
 /// `CURRENT_RESUME_SCHEMA`: pre-per-version dedup identities are not comparable
 /// to the new `compose_key` identities, so a stale checkpoint is discarded.
-pub(super) const CURRENT_CHECKPOINT_SCHEMA: u32 = 5;
+pub(super) const CURRENT_CHECKPOINT_SCHEMA: u32 = 6;
 
 /// resume checkpoint
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -287,7 +288,43 @@ impl CheckpointManager {
                 ),
             });
         }
-        if checkpoint.schema_version < CURRENT_CHECKPOINT_SCHEMA {
+
+        if let Some(expected) = checkpoint.integrity_digest.as_deref() {
+            let actual = Self::checkpoint_digest(&Self::serialize_without_digest(&checkpoint)?);
+            if expected != actual {
+                Self::block_invalid_snapshot(&disk, task_id).await;
+                return Err(Error::InvalidCheckpoint(format!(
+                    "Resume checkpoint digest does not match task {task_id}"
+                )));
+            }
+        } else if checkpoint.schema_version >= CURRENT_CHECKPOINT_SCHEMA {
+            Self::block_invalid_snapshot(&disk, task_id).await;
+            return Err(Error::InvalidCheckpoint(format!(
+                "Resume checkpoint digest is missing for task {task_id}"
+            )));
+        } else {
+            let digest_path = Self::digest_path(task_id);
+            let digest_path = path_to_str(&digest_path)?;
+            match HealDiskExt::read_all(disk.as_ref(), RUSTFS_META_BUCKET, digest_path).await {
+                Ok(expected) => {
+                    let actual = Self::checkpoint_digest(&checkpoint_data);
+                    if expected.as_ref() != actual.as_bytes() {
+                        Self::block_invalid_snapshot(&disk, task_id).await;
+                        return Err(Error::InvalidCheckpoint(format!(
+                            "Resume checkpoint digest does not match task {task_id}"
+                        )));
+                    }
+                }
+                Err(crate::heal::DiskError::FileNotFound) => {}
+                Err(error) => {
+                    return Err(Error::TaskExecutionFailed {
+                        message: format!("Failed to read checkpoint digest: {error}"),
+                    });
+                }
+            }
+        }
+
+        if checkpoint.schema_version < CHECKPOINT_PER_VERSION_SCHEMA {
             warn!(
                 target: "rustfs::heal::resume",
                 event = EVENT_HEAL_CHECKPOINT_STATE,
@@ -304,8 +341,8 @@ impl CheckpointManager {
             checkpoint.skipped_objects.clear();
             checkpoint.current_bucket_index = 0;
             checkpoint.current_object_index = 0;
-            checkpoint.schema_version = CURRENT_CHECKPOINT_SCHEMA;
         }
+        checkpoint.schema_version = CURRENT_CHECKPOINT_SCHEMA;
 
         Ok(Self {
             disk,
@@ -571,45 +608,12 @@ impl CheckpointManager {
         let file_path = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_CHECKPOINT_FILE}"));
 
         let path_str = path_to_str(&file_path)?;
-        let checkpoint = HealDiskExt::read_all(disk.as_ref(), RUSTFS_META_BUCKET, path_str)
+        HealDiskExt::read_all(disk.as_ref(), RUSTFS_META_BUCKET, path_str)
             .await
             .map(|bytes| bytes.to_vec())
             .map_err(|e| Error::TaskExecutionFailed {
                 message: format!("Failed to read checkpoint file: {e}"),
-            })?;
+            })
-        let parsed: ResumeCheckpoint = serde_json::from_slice(&checkpoint).map_err(|error| Error::TaskExecutionFailed {
-            message: format!("Failed to deserialize checkpoint for integrity validation: {error}"),
-        })?;
-        if let Some(expected) = parsed.integrity_digest.as_deref() {
-            let actual = Self::checkpoint_digest(&Self::serialize_without_digest(&parsed)?);
-            if expected != actual {
-                Self::block_invalid_snapshot(disk, task_id).await;
-                return Err(Error::InvalidCheckpoint(format!(
-                    "Resume checkpoint digest does not match task {task_id}"
-                )));
-            }
-        } else {
-            let digest_path = Self::digest_path(task_id);
-            let digest_path = path_to_str(&digest_path)?;
-            match HealDiskExt::read_all(disk.as_ref(), RUSTFS_META_BUCKET, digest_path).await {
-                Ok(expected) => {
-                    let actual = Self::checkpoint_digest(&checkpoint);
-                    if expected.as_ref() != actual.as_bytes() {
-                        Self::block_invalid_snapshot(disk, task_id).await;
-                        return Err(Error::InvalidCheckpoint(format!(
-                            "Resume checkpoint digest does not match task {task_id}"
-                        )));
-                    }
-                }
-                Err(crate::heal::DiskError::FileNotFound) => {}
-                Err(error) => {
-                    return Err(Error::TaskExecutionFailed {
-                        message: format!("Failed to read checkpoint digest: {error}"),
-                    });
-                }
-            }
-        }
-        Ok(checkpoint)
     }
 
     fn serialize_without_digest(checkpoint: &ResumeCheckpoint) -> Result<Vec<u8>> {

crates/heal/src/heal/resume/tests.rs

@@ -1600,6 +1600,29 @@ async fn test_checkpoint_schema_v4_discarded_on_load() {
     temp_dir.close().expect("remove schema test directory");
 }
 
+#[tokio::test]
+async fn unsigned_previous_checkpoint_schema_preserves_progress() {
+    let (temp_dir, disk) = schema_test_disk().await;
+    let task_id = ResumeUtils::generate_task_id();
+    let mut legacy = ResumeCheckpoint::new(task_id.clone());
+    legacy.schema_version = CURRENT_CHECKPOINT_SCHEMA - 1;
+    legacy.update_position(2, 500);
+    legacy.add_processed_object("object".to_string());
+    legacy.integrity_digest = None;
+    let checkpoint_path = format!("{BUCKET_META_PREFIX}/{task_id}_{RESUME_CHECKPOINT_FILE}");
+    disk.write_all(RUSTFS_META_BUCKET, &checkpoint_path, serde_json::to_vec(&legacy).unwrap().into())
+        .await
+        .expect("write previous-schema checkpoint");
+
+    let manager = CheckpointManager::load_from_disk(disk, &task_id).await.unwrap();
+    let checkpoint = manager.get_checkpoint().await;
+    assert_eq!(checkpoint.schema_version, CURRENT_CHECKPOINT_SCHEMA);
+    assert_eq!(checkpoint.current_bucket_index, 2);
+    assert_eq!(checkpoint.current_object_index, 500);
+    assert!(checkpoint.processed_objects.contains("object"));
+    temp_dir.close().unwrap();
+}
+
 #[tokio::test]
 async fn current_normal_resume_schema_preserves_progress() {
     let (temp_dir, disk) = schema_test_disk().await;
@@ -1804,6 +1827,33 @@ async fn checkpoint_integrity_survives_multi_object_reload() {
     temp_dir.close().unwrap();
 }
 
+#[tokio::test]
+async fn checkpoint_integrity_rejects_a_removed_embedded_digest() {
+    let (temp_dir, disk) = schema_test_disk().await;
+    let task_id = ResumeUtils::generate_task_id();
+    let manager = CheckpointManager::new(disk.clone(), task_id.clone()).await.unwrap();
+    manager.update_position(2, 9).await.unwrap();
+
+    let checkpoint_path = format!("{BUCKET_META_PREFIX}/{task_id}_{RESUME_CHECKPOINT_FILE}");
+    let bytes = disk
+        .read_all(RUSTFS_META_BUCKET, &checkpoint_path)
+        .await
+        .expect("read checkpoint fixture");
+    let mut value: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
+    value["current_object_index"] = serde_json::json!(10);
+    value.as_object_mut().unwrap().remove("integrity_digest");
+    disk.write_all(RUSTFS_META_BUCKET, &checkpoint_path, serde_json::to_vec(&value).unwrap().into())
+        .await
+        .expect("write tampered checkpoint fixture");
+
+    assert!(
+        CheckpointManager::load_from_disk(disk.clone(), &task_id).await.is_err(),
+        "a current checkpoint without its embedded digest must fail closed"
+    );
+    assert!(CheckpointManager::is_blocked(&disk, &task_id).await);
+    temp_dir.close().unwrap();
+}
+
 #[tokio::test]
 async fn new_checkpoint_manager_rebuilds_an_empty_snapshot() {
     let (temp_dir, disk) = schema_test_disk().await;