diff --git a/crates/tracedecay-code-index-retention/src/code_index_generations/generation_scan.rs b/crates/tracedecay-code-index-retention/src/code_index_generations/generation_scan.rs index ba1908dc2f..94e144df53 100644 --- a/crates/tracedecay-code-index-retention/src/code_index_generations/generation_scan.rs +++ b/crates/tracedecay-code-index-retention/src/code_index_generations/generation_scan.rs @@ -12,18 +12,23 @@ use super::{ }; const MAX_FORMAT_REVISION_PREFIX_BYTES: usize = 4 * 1024; +const GENERATION_METADATA_READ_CHUNK_BYTES: usize = 4 * 1024; pub(super) fn read_generation_format_revision( path: &Path, is_cancelled: &dyn Fn() -> bool, ) -> Result { - let mut file = File::open(path).map_err(deferred_if_absent)?; - let mut prefix = vec![0_u8; MAX_FORMAT_REVISION_PREFIX_BYTES]; - let bytes_read = file.read(&mut prefix).map_err(storage)?; if is_cancelled() { return Err(CodeGenerationRetentionErrorV1::Cancelled); } - prefix.truncate(bytes_read); + let file = File::open(path).map_err(deferred_if_absent)?; + let mut prefix = Vec::with_capacity(MAX_FORMAT_REVISION_PREFIX_BYTES); + file.take(MAX_FORMAT_REVISION_PREFIX_BYTES as u64) + .read_to_end(&mut prefix) + .map_err(storage)?; + if is_cancelled() { + return Err(CodeGenerationRetentionErrorV1::Cancelled); + } parse_json_u32_field(&prefix, b"format_revision").ok_or_else(|| { CodeGenerationRetentionErrorV1::UnsafeState(format!( "generation file '{}' has no readable format revision in its bounded prefix", @@ -39,28 +44,80 @@ pub(super) fn read_generation_metadata( is_cancelled: &dyn Fn() -> bool, ) -> Result<(u32, SealedGenerationManifestMetadataV1, String, u64), CodeGenerationRetentionErrorV1> { + if is_cancelled() { + return Err(CodeGenerationRetentionErrorV1::Cancelled); + } let mut file = File::open(path).map_err(deferred_if_absent)?; let size_bytes = file.metadata().map_err(storage)?.len(); let mut hasher = Sha256::new(); - let mut prefix = Vec::with_capacity(MAX_GENERATION_METADATA_PREFIX_BYTES); - let mut buffer = vec![0u8; 1024 * 1024]; - loop { - let bytes_read = file.read(&mut buffer).map_err(storage)?; + let (format_revision, manifest) = + read_generation_metadata_prefix(&mut file, path, verification, &mut hasher, is_cancelled)?; + let state_digest = match verification { + GenerationDigestVerificationV1::Full => { + // The prefix reader hashed exactly the bytes it consumed from this + // handle. Continue at that position: a small manifest must never + // turn a full integrity check into a prefix-only digest. + let mut buffer = vec![0u8; 1024 * 1024]; + loop { + if is_cancelled() { + return Err(CodeGenerationRetentionErrorV1::Cancelled); + } + let bytes_read = file.read(&mut buffer).map_err(storage)?; + if bytes_read == 0 { + break; + } + hasher.update(&buffer[..bytes_read]); + if is_cancelled() { + return Err(CodeGenerationRetentionErrorV1::Cancelled); + } + } + encode_tagged_lowercase_hex("sha256:", &hasher.finalize()) + } + GenerationDigestVerificationV1::MetadataOnly => named_state_digest(path)?, + }; + Ok((format_revision, manifest, state_digest, size_bytes)) +} + +/// Read only as far as complete metadata requires, up to the existing bound. +/// Geometrically spaced probes keep repeated prefix parsing linear in bytes +/// even when a manifest is large. Reads stay small and cancellable between +/// probes; the prefix allocation is released before full-payload hashing. +fn read_generation_metadata_prefix( + file: &mut File, + path: &Path, + verification: GenerationDigestVerificationV1, + hasher: &mut Sha256, + is_cancelled: &dyn Fn() -> bool, +) -> Result<(u32, SealedGenerationManifestMetadataV1), CodeGenerationRetentionErrorV1> { + let mut prefix = Vec::with_capacity(GENERATION_METADATA_READ_CHUNK_BYTES); + let mut buffer = [0u8; GENERATION_METADATA_READ_CHUNK_BYTES]; + let mut next_probe = GENERATION_METADATA_READ_CHUNK_BYTES; + while prefix.len() < MAX_GENERATION_METADATA_PREFIX_BYTES { + if is_cancelled() { + return Err(CodeGenerationRetentionErrorV1::Cancelled); + } + let remaining = MAX_GENERATION_METADATA_PREFIX_BYTES - prefix.len(); + let read_limit = buffer.len().min(remaining).min(next_probe - prefix.len()); + let bytes_read = file.read(&mut buffer[..read_limit]).map_err(storage)?; if bytes_read == 0 { break; } if verification == GenerationDigestVerificationV1::Full { hasher.update(&buffer[..bytes_read]); } - let remaining = MAX_GENERATION_METADATA_PREFIX_BYTES.saturating_sub(prefix.len()); - prefix.extend_from_slice(&buffer[..bytes_read.min(remaining)]); + prefix.extend_from_slice(&buffer[..bytes_read]); if is_cancelled() { return Err(CodeGenerationRetentionErrorV1::Cancelled); } - if verification == GenerationDigestVerificationV1::MetadataOnly - && prefix.len() >= MAX_GENERATION_METADATA_PREFIX_BYTES - { - break; + if prefix.len() >= next_probe { + if parse_json_u32_field(&prefix, b"format_revision").is_some() + && extract_json_object_field(&prefix, b"manifest").is_some() + { + break; + } + next_probe = next_probe + .saturating_mul(2) + .min(MAX_GENERATION_METADATA_PREFIX_BYTES); } } let format_revision = parse_json_u32_field(&prefix, b"format_revision").ok_or_else(|| { @@ -81,13 +138,7 @@ pub(super) fn read_generation_metadata( path.display() )) })?; - let state_digest = match verification { - GenerationDigestVerificationV1::Full => { - encode_tagged_lowercase_hex("sha256:", &hasher.finalize()) - } - GenerationDigestVerificationV1::MetadataOnly => named_state_digest(path)?, - }; - Ok((format_revision, manifest, state_digest, size_bytes)) + Ok((format_revision, manifest)) } fn named_state_digest(path: &Path) -> Result { @@ -112,7 +163,16 @@ fn parse_json_u32_field(prefix: &[u8], field: &[u8]) -> Option { .iter() .position(|byte| !byte.is_ascii_digit()) .map_or(prefix.len(), |offset| start + offset); - std::str::from_utf8(&prefix[start..end]).ok()?.parse().ok() + // A prefix ending in digits is not yet a complete number. In particular, + // do not accept the integer prefix of a fractional or exponent value. + let terminator = prefix[end..] + .iter() + .copied() + .find(|byte| !byte.is_ascii_whitespace())?; + if !matches!(terminator, b',' | b'}') { + return None; + } + serde_json::from_slice(&prefix[start..end]).ok() } fn extract_json_object_field<'a>(prefix: &'a [u8], field: &[u8]) -> Option<&'a [u8]> { @@ -167,3 +227,231 @@ fn json_field_value_start(prefix: &[u8], field: &[u8]) -> Option { } Some(cursor) } + +#[cfg(test)] +mod tests { + use std::cell::Cell; + use std::fs::{self, File, OpenOptions}; + use std::io::{Seek, SeekFrom, Write}; + use std::path::{Path, PathBuf}; + + use sha2::{Digest, Sha256}; + use tracedecay_code_index::production::SEALED_GENERATION_FORMAT_REVISION_V1; + use tracedecay_domain::canonical_text::encode_tagged_lowercase_hex; + + use super::{ + CodeGenerationRetentionErrorV1, GENERATION_METADATA_READ_CHUNK_BYTES, + GenerationDigestVerificationV1, MAX_GENERATION_METADATA_PREFIX_BYTES, parse_json_u32_field, + read_generation_format_revision, read_generation_metadata, read_generation_metadata_prefix, + }; + + fn generation_bytes(manifest_padding: usize, payload_bytes: usize) -> Vec { + serde_json::to_vec(&serde_json::json!({ + "format_revision": SEALED_GENERATION_FORMAT_REVISION_V1, + "manifest": { + "generation_id": "generation.v1.scan", + "padding": "x".repeat(manifest_padding), + "seal": { "sealed_at": 7 }, + }, + "payload": "x".repeat(payload_bytes), + })) + .unwrap() + } + + fn write_generation(root: &Path, bytes: &[u8]) -> (PathBuf, String) { + let digest = encode_tagged_lowercase_hex("sha256:", &Sha256::digest(bytes)); + let path = root.join(format!("generation-{}.json", &digest[7..])); + fs::write(&path, bytes).unwrap(); + (path, digest) + } + + #[test] + fn metadata_only_stops_at_the_complete_manifest() { + let root = tempfile::tempdir().unwrap(); + let bytes = generation_bytes(0, 256 * 1024); + let (path, expected_digest) = write_generation(root.path(), &bytes); + assert_eq!( + read_generation_format_revision(&path, &|| false).unwrap(), + SEALED_GENERATION_FORMAT_REVISION_V1 + ); + let mut file = File::open(&path).unwrap(); + let (revision, manifest) = read_generation_metadata_prefix( + &mut file, + &path, + GenerationDigestVerificationV1::MetadataOnly, + &mut Sha256::new(), + &|| false, + ) + .unwrap(); + assert_eq!(revision, SEALED_GENERATION_FORMAT_REVISION_V1); + assert_eq!(manifest.generation_id.as_str(), "generation.v1.scan"); + assert_eq!(manifest.seal.sealed_at.0, 7); + assert_eq!( + file.stream_position().unwrap(), + GENERATION_METADATA_READ_CHUNK_BYTES as u64 + ); + let (_, _, digest, size) = + read_generation_metadata(&path, GenerationDigestVerificationV1::MetadataOnly, &|| { + false + }) + .unwrap(); + assert_eq!(digest, expected_digest); + assert_eq!(size, bytes.len() as u64); + } + + #[test] + fn full_verification_still_hashes_the_payload_after_the_manifest() { + let root = tempfile::tempdir().unwrap(); + let bytes = generation_bytes(0, 2 * 1024 * 1024); + let (path, named_digest) = write_generation(root.path(), &bytes); + let (_, _, digest, size) = + read_generation_metadata(&path, GenerationDigestVerificationV1::Full, &|| false) + .unwrap(); + assert_eq!(digest, named_digest); + assert_eq!(size, bytes.len() as u64); + + let mut file = OpenOptions::new().write(true).open(&path).unwrap(); + file.seek(SeekFrom::End(-3)).unwrap(); + file.write_all(b"y").unwrap(); + drop(file); + let (_, _, changed_digest, _) = + read_generation_metadata(&path, GenerationDigestVerificationV1::Full, &|| false) + .unwrap(); + assert_ne!(changed_digest, named_digest); + assert_eq!( + changed_digest, + encode_tagged_lowercase_hex("sha256:", &Sha256::digest(fs::read(&path).unwrap())) + ); + } + + #[test] + fn a_manifest_spanning_multiple_chunks_is_read_completely() { + let root = tempfile::tempdir().unwrap(); + let bytes = generation_bytes(12 * 1024, 64 * 1024); + let (path, _) = write_generation(root.path(), &bytes); + let mut file = File::open(&path).unwrap(); + let (_, manifest) = read_generation_metadata_prefix( + &mut file, + &path, + GenerationDigestVerificationV1::MetadataOnly, + &mut Sha256::new(), + &|| false, + ) + .unwrap(); + assert_eq!(manifest.generation_id.as_str(), "generation.v1.scan"); + assert_eq!(manifest.seal.sealed_at.0, 7); + assert_eq!(file.stream_position().unwrap(), 16 * 1024); + } + + #[test] + fn metadata_beyond_the_existing_prefix_bound_is_refused() { + let root = tempfile::tempdir().unwrap(); + let mut bytes = vec![b' '; MAX_GENERATION_METADATA_PREFIX_BYTES]; + bytes.extend(generation_bytes(0, 0)); + let (path, _) = write_generation(root.path(), &bytes); + let mut file = File::open(&path).unwrap(); + let result = read_generation_metadata_prefix( + &mut file, + &path, + GenerationDigestVerificationV1::MetadataOnly, + &mut Sha256::new(), + &|| false, + ); + assert!(matches!( + result, + Err(CodeGenerationRetentionErrorV1::UnsafeState(_)) + )); + assert_eq!( + file.stream_position().unwrap(), + MAX_GENERATION_METADATA_PREFIX_BYTES as u64 + ); + } + + #[test] + fn malformed_manifest_metadata_is_refused_in_both_modes() { + let root = tempfile::tempdir().unwrap(); + let bytes = serde_json::to_vec(&serde_json::json!({ + "format_revision": SEALED_GENERATION_FORMAT_REVISION_V1, + "manifest": { "generation_id": "generation.v1.scan" }, + "payload": "x".repeat(64 * 1024), + })) + .unwrap(); + let (path, _) = write_generation(root.path(), &bytes); + for verification in [ + GenerationDigestVerificationV1::MetadataOnly, + GenerationDigestVerificationV1::Full, + ] { + assert!(matches!( + read_generation_metadata(&path, verification, &|| false), + Err(CodeGenerationRetentionErrorV1::UnsafeState(_)) + )); + } + } + + #[test] + fn metadata_reads_observe_cancellation_before_and_after_io() { + let root = tempfile::tempdir().unwrap(); + let (path, _) = write_generation(root.path(), &generation_bytes(0, 64 * 1024)); + let mut file = File::open(&path).unwrap(); + let result = read_generation_metadata_prefix( + &mut file, + &path, + GenerationDigestVerificationV1::MetadataOnly, + &mut Sha256::new(), + &|| true, + ); + assert!(matches!( + result, + Err(CodeGenerationRetentionErrorV1::Cancelled) + )); + assert_eq!(file.stream_position().unwrap(), 0); + + let checks = Cell::new(0); + let result = read_generation_metadata_prefix( + &mut file, + &path, + GenerationDigestVerificationV1::MetadataOnly, + &mut Sha256::new(), + &|| { + let previous = checks.get(); + checks.set(previous + 1); + previous > 0 + }, + ); + assert!(matches!( + result, + Err(CodeGenerationRetentionErrorV1::Cancelled) + )); + assert_eq!( + file.stream_position().unwrap(), + GENERATION_METADATA_READ_CHUNK_BYTES as u64 + ); + } + + #[test] + fn format_revision_requires_a_complete_unsigned_json_integer() { + for input in [ + r#"{"format_revision":21,"manifest":{}}"#, + r#"{"format_revision":21 }"#, + ] { + assert_eq!( + parse_json_u32_field(input.as_bytes(), b"format_revision"), + Some(21) + ); + } + for input in [ + r#"{"format_revision":21"#, + r#"{"format_revision":21.5}"#, + r#"{"format_revision":21e1}"#, + r#"{"format_revision":021}"#, + r#"{"format_revision":-1}"#, + r#"{"format_revision":"21"}"#, + r#"{"format_revision":4294967296}"#, + ] { + assert_eq!( + parse_json_u32_field(input.as_bytes(), b"format_revision"), + None + ); + } + } +}