diff --git a/Cargo.lock b/Cargo.lock index 637202fd5..c0a5e1b69 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9587,6 +9587,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "sha2 0.11.0", "temp-env", "tempfile", "thiserror 2.0.20", diff --git a/crates/heal/Cargo.toml b/crates/heal/Cargo.toml index 4d4fb355b..7ea3b1854 100644 --- a/crates/heal/Cargo.toml +++ b/crates/heal/Cargo.toml @@ -91,6 +91,7 @@ metrics = { workspace = true } base64 = { workspace = true } bytes = { workspace = true } crc-fast = { workspace = true } +sha2 = { workspace = true } [dev-dependencies] serde_json = { workspace = true, features = ["raw_value"] } diff --git a/crates/heal/src/heal/resume/checkpoint.rs b/crates/heal/src/heal/resume/checkpoint.rs index 584b13746..261a434ad 100644 --- a/crates/heal/src/heal/resume/checkpoint.rs +++ b/crates/heal/src/heal/resume/checkpoint.rs @@ -13,7 +13,9 @@ // limitations under the License. use crate::{Error, Result}; +use base64::Engine as _; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; use std::collections::HashSet; use std::path::Path; use std::sync::{Arc, Mutex}; @@ -29,6 +31,7 @@ use super::{ }; const EVENT_HEAL_CHECKPOINT_STATE: &str = "heal_checkpoint_state"; +const RESUME_CHECKPOINT_DIGEST_FILE: &str = "ahm_checkpoint.sha256"; /// Current on-disk schema version for `ResumeCheckpoint`. Same rationale as /// `CURRENT_RESUME_SCHEMA`: pre-per-version dedup identities are not comparable @@ -145,19 +148,22 @@ impl CheckpointManager { /// Validate the checkpoint while enumerating resumable state. This reads /// the checkpoint once and also isolates malformed or unsupported data. - pub(crate) async fn is_resumable(disk: &DiskStore, task_id: &str) -> bool { - if validate_resume_task_id(task_id).is_err() || Self::is_blocked(disk, task_id).await { - return false; + pub(crate) async fn is_resumable(disk: &DiskStore, task_id: &str) -> Result { + validate_resume_task_id(task_id)?; + if Self::is_blocked(disk, task_id).await { + return Err(Error::InvalidCheckpoint(format!("Resume task {task_id} has a blocked checkpoint"))); } let file_path = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_CHECKPOINT_FILE}")); let Ok(path) = path_to_str(&file_path) else { - return false; + return Err(Error::InvalidCheckpoint("Resume checkpoint path is not valid UTF-8".to_string())); }; match HealDiskExt::read_all(disk.as_ref(), RUSTFS_META_BUCKET, path).await { - Ok(bytes) if bytes.is_empty() => true, - Ok(bytes) => Self::load_from_data(disk.clone(), task_id, bytes.to_vec()).await.is_ok(), - Err(crate::heal::DiskError::FileNotFound) => true, - Err(_) => false, + Ok(bytes) if bytes.is_empty() => Ok(true), + Ok(bytes) => Self::load_from_data(disk.clone(), task_id, bytes.to_vec()) + .await + .map(|_| true), + Err(crate::heal::DiskError::FileNotFound) => Ok(true), + Err(error) => Err(error.into()), } } @@ -399,6 +405,7 @@ impl CheckpointManager { let checkpoint_file = Path::new(BUCKET_META_PREFIX).join(format!("{task_id}_{RESUME_CHECKPOINT_FILE}")); delete_resume_file(&self.disk, &checkpoint_file).await?; + delete_resume_file(&self.disk, &Self::digest_path(&task_id)).await?; delete_resume_file(&self.disk, &Self::blocked_path(&task_id)).await?; debug!( @@ -531,6 +538,19 @@ 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(), })?; @@ -554,11 +574,35 @@ 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)?; - HealDiskExt::read_all(disk.as_ref(), RUSTFS_META_BUCKET, path_str) + let checkpoint = 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 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}" + ))); + } + } + Err(crate::heal::DiskError::FileNotFound) => {} + Err(error) => { + return Err(Error::TaskExecutionFailed { + message: format!("Failed to read checkpoint digest: {error}"), + }); + } + } + Ok(checkpoint) + } + + 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 d8331b93e..a95d20575 100644 --- a/crates/heal/src/heal/resume/tests.rs +++ b/crates/heal/src/heal/resume/tests.rs @@ -1733,6 +1733,41 @@ async fn checkpoint_save_does_not_replace_a_future_schema_snapshot() { temp_dir.close().expect("remove future schema test directory"); } +#[tokio::test] +async fn checkpoint_digest_rejects_same_length_progress_tampering() { + 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 + .expect("create checkpoint manager"); + manager + .add_processed_object("victim-a".to_string()) + .await + .expect("persist checkpoint progress"); + manager.update_position(1, 1).await.expect("flush checkpoint progress"); + let checkpoint_path = format!("{BUCKET_META_PREFIX}/{task_id}_{RESUME_CHECKPOINT_FILE}"); + let original = disk + .read_all(RUSTFS_META_BUCKET, &checkpoint_path) + .await + .expect("read checkpoint fixture"); + let tampered = original + .windows(b"victim-a".len()) + .position(|window| window == b"victim-a") + .map(|index| { + let mut bytes = original.to_vec(); + bytes[index..index + b"victim-a".len()].copy_from_slice(b"victim-b"); + bytes + }) + .expect("checkpoint should contain the processed object"); + disk.write_all(RUSTFS_META_BUCKET, &checkpoint_path, tampered.into()) + .await + .expect("write tampered checkpoint fixture"); + + assert!(CheckpointManager::load_from_disk(disk.clone(), &task_id).await.is_err()); + assert!(CheckpointManager::is_blocked(&disk, &task_id).await); + temp_dir.close().expect("remove digest test directory"); +} + #[tokio::test] async fn new_checkpoint_manager_rebuilds_an_empty_snapshot() { let (temp_dir, disk) = schema_test_disk().await; @@ -1784,7 +1819,7 @@ async fn an_empty_blocked_marker_still_blocks_resume_selection() { .expect("write empty blocked marker fixture"); assert!(CheckpointManager::is_blocked(&disk, &task_id).await); - assert!(!CheckpointManager::is_resumable(&disk, &task_id).await); + assert!(CheckpointManager::is_resumable(&disk, &task_id).await.is_err()); // Recovery requires replacing/cleaning the snapshot, then removing the // marker; ordinary selector retries are intentionally not an unlock path. manager.cleanup().await.expect("clean blocked checkpoint"); @@ -1822,12 +1857,7 @@ async fn resumable_selector_skips_healthy_tasks_with_blocked_markers() { .await .expect("write blocked marker"); - assert!( - ResumeUtils::get_resumable_tasks(&disk) - .await - .expect("filter blocked healthy task") - .is_empty() - ); + assert!(ResumeUtils::get_resumable_tasks(&disk).await.is_err()); assert_eq!( disk.read_all(RUSTFS_META_BUCKET, &checkpoint_path) .await @@ -1899,18 +1929,9 @@ async fn resumable_selector_isolates_future_and_corrupt_checkpoints() { .await .expect("write corrupt checkpoint fixture"); - assert!( - ResumeUtils::get_resumable_tasks(&disk) - .await - .expect("filter malformed resumable tasks") - .is_empty() - ); - assert!( - ResumeUtils::get_resumable_tasks(&disk) - .await - .expect("filter blocked resumable tasks") - .is_empty() - ); + assert!(CheckpointManager::is_resumable(&disk, &future_task).await.is_err()); + assert!(CheckpointManager::is_resumable(&disk, &corrupt_task).await.is_err()); + assert!(ResumeUtils::get_resumable_tasks(&disk).await.is_err()); for (task_id, path, bytes) in [ (&future_task, future_path, future_bytes), (&corrupt_task, corrupt_path, corrupt_bytes.to_vec()), diff --git a/crates/heal/src/heal/resume/utils.rs b/crates/heal/src/heal/resume/utils.rs index 2ef0c4fae..557269d6e 100644 --- a/crates/heal/src/heal/resume/utils.rs +++ b/crates/heal/src/heal/resume/utils.rs @@ -67,7 +67,7 @@ impl ResumeUtils { // Extract task ID from filename: {task_id}_ahm_resume_state.json if let Some(task_id) = entry.strip_suffix(&format!("_{RESUME_STATE_FILE}")) && validate_resume_task_id(task_id).is_ok() - && CheckpointManager::is_resumable(disk, task_id).await + && CheckpointManager::is_resumable(disk, task_id).await? { task_ids.push(task_id.to_string()); }