diff --git a/crates/heal/src/heal/resume/checkpoint.rs b/crates/heal/src/heal/resume/checkpoint.rs index 936b41588..119c7cc6d 100644 --- a/crates/heal/src/heal/resume/checkpoint.rs +++ b/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> { diff --git a/crates/heal/src/heal/resume/tests.rs b/crates/heal/src/heal/resume/tests.rs index fa70a6c75..c938969bc 100644 --- a/crates/heal/src/heal/resume/tests.rs +++ b/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;