From e3a362989f11515dda85699dc3145902ae7ceb4d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=A9=AC=E7=99=BB=E5=B1=B1?= Date: Sat, 22 Aug 2026 22:05:23 +0800 Subject: [PATCH] fix(heal): atomically authenticate checkpoints --- crates/heal/src/heal/resume/checkpoint.rs | 82 ++++++++++++++--------- crates/heal/src/heal/resume/tests.rs | 18 +++++ 2 files changed, 70 insertions(+), 30 deletions(-) diff --git a/crates/heal/src/heal/resume/checkpoint.rs b/crates/heal/src/heal/resume/checkpoint.rs index 261a434ad..43b5164f7 100644 --- a/crates/heal/src/heal/resume/checkpoint.rs +++ b/crates/heal/src/heal/resume/checkpoint.rs @@ -61,6 +61,11 @@ pub struct ResumeCheckpoint { pub failed_objects: HashSet, /// skipped objects pub skipped_objects: HashSet, + /// Integrity digest over the checkpoint with this field set to `None`. + /// Keeping it in the checkpoint makes the payload and its authentication + /// record one CAS generation instead of two independently-written files. + #[serde(default)] + pub integrity_digest: Option, } impl ResumeCheckpoint { @@ -74,6 +79,7 @@ impl ResumeCheckpoint { processed_objects: HashSet::new(), failed_objects: HashSet::new(), skipped_objects: HashSet::new(), + integrity_digest: None, } } @@ -427,8 +433,12 @@ impl CheckpointManager { let _save_guard = self.save_lock.lock().await; let checkpoint = self.checkpoint.read().await.clone(); validate_resume_task_id(&checkpoint.task_id)?; + let unsigned_checkpoint_data = Self::serialize_without_digest(&checkpoint)?; + let digest = Self::checkpoint_digest(&unsigned_checkpoint_data); + let mut persisted_checkpoint = checkpoint.clone(); + persisted_checkpoint.integrity_digest = Some(digest); let checkpoint_data = - EcstoreDiskBytes::from(serde_json::to_vec(&checkpoint).map_err(|e| Error::TaskExecutionFailed { + EcstoreDiskBytes::from(serde_json::to_vec(&persisted_checkpoint).map_err(|e| Error::TaskExecutionFailed { message: format!("Failed to serialize checkpoint: {e}"), })?); @@ -538,19 +548,6 @@ impl CheckpointManager { } } - let digest_path = Self::digest_path(&checkpoint.task_id); - let digest = base64::engine::general_purpose::STANDARD.encode(Sha256::digest(checkpoint_data.as_ref())); - HealDiskExt::write_all( - self.disk.as_ref(), - RUSTFS_META_BUCKET, - path_to_str(&digest_path)?, - EcstoreDiskBytes::from(digest.into_bytes()), - ) - .await - .map_err(|e| Error::TaskExecutionFailed { - message: format!("Failed to save checkpoint digest: {e}"), - })?; - let mut last_saved = self.last_saved.lock().map_err(|_| Error::TaskExecutionFailed { message: "Checkpoint save state lock is poisoned after save".to_string(), })?; @@ -580,28 +577,53 @@ impl CheckpointManager { .map_err(|e| Error::TaskExecutionFailed { message: format!("Failed to read checkpoint file: {e}"), })?; - 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 = base64::engine::general_purpose::STANDARD.encode(Sha256::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}" - ))); - } + 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}" + ))); } - Err(crate::heal::DiskError::FileNotFound) => {} - Err(error) => { - return Err(Error::TaskExecutionFailed { - message: format!("Failed to read checkpoint digest: {error}"), - }); + } 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> { + let mut unsigned = checkpoint.clone(); + unsigned.integrity_digest = None; + serde_json::to_vec(&unsigned).map_err(|e| Error::TaskExecutionFailed { + message: format!("Failed to serialize checkpoint: {e}"), + }) + } + + fn checkpoint_digest(checkpoint_data: &[u8]) -> String { + base64::engine::general_purpose::STANDARD.encode(Sha256::digest(checkpoint_data)) + } + fn digest_path(task_id: &str) -> std::path::PathBuf { Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_CHECKPOINT_DIGEST_FILE}")) } diff --git a/crates/heal/src/heal/resume/tests.rs b/crates/heal/src/heal/resume/tests.rs index a95d20575..593dbb235 100644 --- a/crates/heal/src/heal/resume/tests.rs +++ b/crates/heal/src/heal/resume/tests.rs @@ -1768,6 +1768,24 @@ async fn checkpoint_digest_rejects_same_length_progress_tampering() { temp_dir.close().expect("remove digest test directory"); } +#[tokio::test] +async fn checkpoint_integrity_survives_missing_legacy_sidecar() { + 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 digest_path = format!("{BUCKET_META_PREFIX}/{task_id}_ahm_checkpoint.sha256"); + delete_resume_file(&disk, Path::new(&digest_path)).await.unwrap(); + + let restored = CheckpointManager::load_from_disk(disk, &task_id).await.unwrap(); + let checkpoint = restored.get_checkpoint().await; + assert_eq!(checkpoint.current_bucket_index, 2); + assert_eq!(checkpoint.current_object_index, 9); + assert!(checkpoint.integrity_digest.is_some()); + temp_dir.close().unwrap(); +} + #[tokio::test] async fn new_checkpoint_manager_rebuilds_an_empty_snapshot() { let (temp_dir, disk) = schema_test_disk().await;