1use std::{
4 collections::BTreeMap,
5 fs::{File, OpenOptions},
6 io::{self, Read, Seek, SeekFrom, Write},
7 path::{Path, PathBuf},
8 time::Duration,
9};
10
11use hyphae_core::{
12 DISK_FORMAT_VERSION, MIN_DISK_FORMAT_VERSION, Q15Vector, VectorMetric, VectorSpaceDefinition,
13 VectorSpaceName,
14};
15use hyphae_query::FieldPath;
16use hyphae_retrieval::{
17 LexicalError, LexicalField, LexicalIndexDefinition, MAX_LEXICAL_FIELDS,
18 MAX_LEXICAL_PATH_SEGMENT_BYTES, MAX_LEXICAL_PATH_SEGMENTS,
19};
20use thiserror::Error;
21
22use crate::{
23 CommitReceipt, MAX_KEY_BYTES, MaterializedIndexError, StorageLimitError,
24 index::MaterializedIndex,
25 limits::{OperationDeadline, limit_io_error, storage_limit_from_io},
26 log::MAX_OPERATION_BYTES,
27};
28
29const MAGIC: [u8; 8] = *b"HYSNAP01";
30const HEADER_LENGTH: usize = 112;
31const HEADER_LENGTH_U64: u64 = 112;
32const CHECKSUM_PREFIX_LENGTH: usize = 76;
33const DIGEST_PREFIX_LENGTH: usize = 80;
34const ENTRY_HEADER_LENGTH: usize = 12;
35const ENTRY_HEADER_LENGTH_U64: u64 = 12;
36const RECEIPT_LENGTH: usize = 88;
37const RECEIPT_LENGTH_U64: u64 = 88;
38const V2_COUNTS_LENGTH: usize = 24;
39const V2_COUNTS_LENGTH_U64: u64 = 24;
40const VECTOR_SPACE_FIXED_LENGTH_U64: u64 = 5;
41const VECTOR_FIXED_LENGTH_U64: u64 = 7;
42const COPY_BUFFER_LENGTH: usize = 64 * 1024;
43const COPY_BUFFER_LENGTH_U64: u64 = 64 * 1024;
44
45#[derive(Clone, Debug, Eq, PartialEq)]
47pub struct SnapshotInfo {
48 pub path: PathBuf,
50 pub disk_format_version: u16,
52 pub checkpoint_sequence: u64,
54 pub checkpoint_digest: Option<[u8; 32]>,
56 pub entry_count: u64,
58 pub vector_space_count: u64,
60 pub vector_count: u64,
62 pub lexical_index_count: u64,
64 pub receipt_count: u64,
66 pub snapshot_digest: [u8; 32],
68 pub file_bytes: u64,
70}
71
72#[derive(Clone, Debug, Eq, PartialEq)]
75pub struct SnapshotReadLimits {
76 pub file_bytes: u64,
78 pub entries: u64,
81 pub decoded_bytes: u64,
83}
84
85impl Default for SnapshotReadLimits {
86 fn default() -> Self {
87 Self {
88 file_bytes: 2 * 1024 * 1024 * 1024,
89 entries: 1_000_000,
90 decoded_bytes: 1024 * 1024 * 1024,
91 }
92 }
93}
94
95#[derive(Clone, Debug, Eq, PartialEq)]
97pub struct SnapshotEntry {
98 pub key: Vec<u8>,
100 pub value: Vec<u8>,
102}
103
104#[derive(Clone, Debug, Eq, PartialEq)]
106pub struct SnapshotContents {
107 pub info: SnapshotInfo,
109 pub entries: Vec<SnapshotEntry>,
111 pub vector_spaces: Vec<VectorSpaceDefinition>,
113 pub vectors: Vec<SnapshotVectorEntry>,
115 pub lexical_indexes: Vec<LexicalIndexDefinition>,
117}
118
119#[derive(Clone, Debug, Eq, PartialEq)]
121pub struct SnapshotVectorEntry {
122 pub space: VectorSpaceName,
124 pub key: Vec<u8>,
126 pub vector: Q15Vector,
128}
129
130#[derive(Debug, Error)]
132pub enum SnapshotError {
133 #[error(transparent)]
135 Io(#[from] io::Error),
136
137 #[error("materialized index failure during snapshot: {source}")]
139 Index {
140 #[source]
142 source: Box<MaterializedIndexError>,
143 },
144
145 #[error("invalid snapshot: {reason}")]
147 Invalid {
148 reason: &'static str,
150 },
151
152 #[error("unsupported snapshot format {found}; supported format is {supported}")]
154 UnsupportedVersion {
155 found: u16,
157 supported: u16,
159 },
160
161 #[error("snapshot sequence {sequence} already exists for a different commit")]
163 CheckpointConflict {
164 sequence: u64,
166 },
167
168 #[error("snapshot file length {actual} exceeds verification limit {maximum}")]
170 FileLimitExceeded {
171 actual: u64,
173 maximum: u64,
175 },
176
177 #[error("snapshot entry count {actual} exceeds verification limit {maximum}")]
179 EntryLimitExceeded {
180 actual: u64,
182 maximum: u64,
184 },
185
186 #[error("snapshot decoded bytes exceed verification limit {maximum}")]
188 DecodedBytesLimitExceeded {
189 maximum: u64,
191 },
192}
193
194impl From<StorageLimitError> for SnapshotError {
195 fn from(source: StorageLimitError) -> Self {
196 Self::Io(limit_io_error(source))
197 }
198}
199
200impl SnapshotError {
201 pub fn storage_limit(&self) -> Option<&StorageLimitError> {
204 if let Self::Io(source) = self
205 && let Some(source) = storage_limit_from_io(source)
206 {
207 return Some(source);
208 }
209 let mut current: &(dyn std::error::Error + 'static) = self;
210 loop {
211 if let Some(source) = current.downcast_ref::<StorageLimitError>() {
212 return Some(source);
213 }
214 let source = current.source()?;
215 current = source;
216 }
217 }
218
219 pub fn is_timeout(&self) -> bool {
222 if matches!(self, Self::Io(source) if source.kind() == io::ErrorKind::TimedOut) {
223 return true;
224 }
225 if matches!(self.storage_limit(), Some(StorageLimitError::TimedOut)) {
226 return true;
227 }
228 let mut current: &(dyn std::error::Error + 'static) = self;
229 loop {
230 if matches!(
231 current.downcast_ref::<LexicalError>(),
232 Some(LexicalError::TimedOut)
233 ) {
234 return true;
235 }
236 let Some(source) = current.source() else {
237 return false;
238 };
239 current = source;
240 }
241 }
242}
243
244impl From<MaterializedIndexError> for SnapshotError {
245 fn from(source: MaterializedIndexError) -> Self {
246 Self::Index {
247 source: Box::new(source),
248 }
249 }
250}
251
252#[allow(clippy::too_many_lines)]
253pub(crate) fn create_snapshot(
254 index: &MaterializedIndex,
255 snapshots_directory: &Path,
256 temporary_directory: &Path,
257 disk_format_version: u16,
258 limits: &SnapshotReadLimits,
259 deadline: &OperationDeadline,
260) -> Result<SnapshotInfo, SnapshotError> {
261 check_snapshot_deadline(Some(deadline))?;
262 let checkpoint = index.checkpoint()?;
263 if checkpoint.sequence == 0 && checkpoint.digest.is_some() {
264 return Err(SnapshotError::Invalid {
265 reason: "empty checkpoint has a digest",
266 });
267 }
268
269 let measurements = measure_payload(
270 index,
271 checkpoint.sequence,
272 disk_format_version,
273 limits,
274 deadline,
275 )?;
276 validate_measurement_limits(&measurements, limits)?;
277 let final_path =
278 snapshots_directory.join(format!("snapshot-{:020}.hysnap", checkpoint.sequence));
279 if final_path.exists() {
280 let mut existing_file = open_snapshot_file(&final_path)?;
281 let existing = verify_snapshot_file(
282 &mut existing_file,
283 &final_path,
284 Some(limits),
285 Some(deadline),
286 )?;
287 if existing.checkpoint_digest != checkpoint.digest {
288 return Err(SnapshotError::CheckpointConflict {
289 sequence: checkpoint.sequence,
290 });
291 }
292 return Ok(existing);
293 }
294
295 let mut header = [0_u8; HEADER_LENGTH];
296 header[0..8].copy_from_slice(&MAGIC);
297 header[8..10].copy_from_slice(&disk_format_version.to_le_bytes());
298 header[10..12].copy_from_slice(&0_u16.to_le_bytes());
299 header[12..20].copy_from_slice(&checkpoint.sequence.to_le_bytes());
300 header[20..52].copy_from_slice(&checkpoint.digest.unwrap_or([0; 32]));
301 header[52..60].copy_from_slice(&measurements.entry_count.to_le_bytes());
302 header[60..68].copy_from_slice(&measurements.receipt_count.to_le_bytes());
303 header[68..76].copy_from_slice(&measurements.payload_length.to_le_bytes());
304
305 let mut checksum = crc32c::crc32c(&header[..CHECKSUM_PREFIX_LENGTH]);
306 if disk_format_version >= 2 {
307 checksum = crc32c::crc32c_append(checksum, &measurements.v2_counts());
308 }
309 let mut checksum_error = None;
310 index.for_each_entry(|key, value| {
311 if checksum_error.is_some() {
312 return;
313 }
314 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
315 checksum_error = Some(source);
316 return;
317 }
318 match encode_entry_header(key, value) {
319 Ok(entry_header) => {
320 checksum = crc32c::crc32c_append(checksum, &entry_header);
321 checksum = crc32c::crc32c_append(checksum, key);
322 checksum = crc32c::crc32c_append(checksum, value);
323 }
324 Err(source) => checksum_error = Some(source),
325 }
326 })?;
327 if let Some(source) = checksum_error {
328 return Err(source);
329 }
330 let mut vector_checksum_error = None;
331 if disk_format_version >= 2 {
332 index.for_each_vector_space(|definition| {
333 if vector_checksum_error.is_none() {
334 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
335 vector_checksum_error = Some(source);
336 return;
337 }
338 match encode_vector_space(definition) {
339 Ok(encoded) => checksum = crc32c::crc32c_append(checksum, &encoded),
340 Err(source) => vector_checksum_error = Some(source),
341 }
342 }
343 })?;
344 index.for_each_vector(|space, key, vector| {
345 if vector_checksum_error.is_none() {
346 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
347 vector_checksum_error = Some(source);
348 return;
349 }
350 match encode_vector(space, key, vector) {
351 Ok(encoded) => checksum = crc32c::crc32c_append(checksum, &encoded),
352 Err(source) => vector_checksum_error = Some(source),
353 }
354 }
355 })?;
356 index.for_each_lexical_index(|definition| {
357 if vector_checksum_error.is_none() {
358 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
359 vector_checksum_error = Some(source);
360 return;
361 }
362 match encode_lexical_index(definition) {
363 Ok(encoded) => checksum = crc32c::crc32c_append(checksum, &encoded),
364 Err(source) => vector_checksum_error = Some(source),
365 }
366 }
367 })?;
368 }
369 index.for_each_receipt(|receipt| {
370 if vector_checksum_error.is_none() {
371 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
372 vector_checksum_error = Some(source);
373 } else {
374 checksum = crc32c::crc32c_append(checksum, &encode_receipt(receipt));
375 }
376 }
377 })?;
378 if let Some(source) = vector_checksum_error {
379 return Err(source);
380 }
381 header[76..80].copy_from_slice(&checksum.to_le_bytes());
382
383 let temporary_path = temporary_directory.join(format!(
384 "snapshot-{:020}-{}.tmp",
385 checkpoint.sequence,
386 uuid::Uuid::now_v7()
387 ));
388 let mut temporary_guard = TemporaryFileGuard::new(temporary_path.clone());
389 let mut file = OpenOptions::new()
390 .create_new(true)
391 .read(true)
392 .write(true)
393 .open(&temporary_path)?;
394 file.write_all(&header)?;
395 let mut hasher = blake3::Hasher::new();
396 hasher.update(&header[..DIGEST_PREFIX_LENGTH]);
397 if disk_format_version >= 2 {
398 let counts = measurements.v2_counts();
399 file.write_all(&counts)?;
400 hasher.update(&counts);
401 }
402 let mut write_error = None;
403 index.for_each_entry(|key, value| {
404 if write_error.is_none()
405 && let Err(source) = write_entry(&mut file, &mut hasher, key, value, Some(deadline))
406 {
407 write_error = Some(source);
408 }
409 })?;
410 if let Some(source) = write_error {
411 return Err(source);
412 }
413 let mut vector_write_error = None;
414 if disk_format_version >= 2 {
415 index.for_each_vector_space(|definition| {
416 if vector_write_error.is_none()
417 && let Err(source) = write_encoded_with_deadline(
418 &mut file,
419 &mut hasher,
420 encode_vector_space(definition),
421 Some(deadline),
422 )
423 {
424 vector_write_error = Some(source);
425 }
426 })?;
427 index.for_each_vector(|space, key, vector| {
428 if vector_write_error.is_none()
429 && let Err(source) = write_encoded_with_deadline(
430 &mut file,
431 &mut hasher,
432 encode_vector(space, key, vector),
433 Some(deadline),
434 )
435 {
436 vector_write_error = Some(source);
437 }
438 })?;
439 index.for_each_lexical_index(|definition| {
440 if vector_write_error.is_none()
441 && let Err(source) = write_encoded_with_deadline(
442 &mut file,
443 &mut hasher,
444 encode_lexical_index(definition),
445 Some(deadline),
446 )
447 {
448 vector_write_error = Some(source);
449 }
450 })?;
451 }
452 if let Some(source) = vector_write_error {
453 return Err(source);
454 }
455 let mut receipt_write_error = None;
456 index.for_each_receipt(|receipt| {
457 if receipt_write_error.is_none()
458 && let Err(source) = check_snapshot_deadline(Some(deadline))
459 {
460 receipt_write_error = Some(source);
461 } else if receipt_write_error.is_none()
462 && let Err(source) = write_receipt(&mut file, &mut hasher, receipt)
463 {
464 receipt_write_error = Some(source);
465 }
466 })?;
467 if let Some(source) = receipt_write_error {
468 return Err(source);
469 }
470 let snapshot_digest = *hasher.finalize().as_bytes();
471 file.seek(SeekFrom::Start(80))?;
472 file.write_all(&snapshot_digest)?;
473 file.sync_all()?;
474 drop(file);
475
476 let mut temporary_file = open_snapshot_file(&temporary_path)?;
477 let temporary_info = verify_snapshot_file(
478 &mut temporary_file,
479 &temporary_path,
480 Some(limits),
481 Some(deadline),
482 )?;
483 check_snapshot_deadline(Some(deadline))?;
484 std::fs::rename(&temporary_path, &final_path)?;
485 temporary_guard.disarm();
486 #[cfg(unix)]
487 sync_directory(snapshots_directory)?;
488 Ok(SnapshotInfo {
489 path: final_path,
490 ..temporary_info
491 })
492}
493
494pub fn verify_snapshot(path: impl AsRef<Path>) -> Result<SnapshotInfo, SnapshotError> {
501 let path = path.as_ref();
502 let mut file = open_snapshot_file(path)?;
503 verify_snapshot_file(&mut file, path, None, None)
504}
505
506pub fn verify_snapshot_with_limits(
515 path: impl AsRef<Path>,
516 limits: &SnapshotReadLimits,
517 timeout: Duration,
518) -> Result<SnapshotInfo, SnapshotError> {
519 let deadline = OperationDeadline::new(timeout);
520 verify_snapshot_with_policy(path.as_ref(), limits, &deadline)
521}
522
523pub fn open_verified_snapshot_with_limits(
534 path: impl AsRef<Path>,
535 limits: &SnapshotReadLimits,
536 timeout: Duration,
537) -> Result<(File, SnapshotInfo), SnapshotError> {
538 let path = path.as_ref();
539 let deadline = OperationDeadline::new(timeout);
540 let mut file = open_snapshot_file(path)?;
541 let info = verify_snapshot_file(&mut file, path, Some(limits), Some(&deadline))?;
542 check_snapshot_deadline(Some(&deadline))?;
543 file.seek(SeekFrom::Start(0))?;
544 ensure_snapshot_file_length_unchanged(&file, info.file_bytes)?;
545 check_snapshot_deadline(Some(&deadline))?;
546 Ok((file, info))
547}
548
549pub(crate) fn verify_snapshot_with_policy(
550 path: &Path,
551 limits: &SnapshotReadLimits,
552 deadline: &OperationDeadline,
553) -> Result<SnapshotInfo, SnapshotError> {
554 let mut file = open_snapshot_file(path)?;
555 verify_snapshot_file(&mut file, path, Some(limits), Some(deadline))
556}
557
558fn verify_snapshot_file(
559 file: &mut File,
560 path: &Path,
561 limits: Option<&SnapshotReadLimits>,
562 deadline: Option<&OperationDeadline>,
563) -> Result<SnapshotInfo, SnapshotError> {
564 check_snapshot_deadline(deadline)?;
565 file.seek(SeekFrom::Start(0))?;
566 let file_bytes = snapshot_file_length(file)?;
567 if let Some(limits) = limits {
568 validate_file_limit(file_bytes, limits)?;
569 }
570 let mut header = [0_u8; HEADER_LENGTH];
571 read_exact_or_invalid(file, &mut header, "truncated header")?;
572 let decoded = decode_header(&header, file_bytes)?;
573 let verified_counts = verify_payload(file, &header, &decoded, limits, deadline, None)?;
574 let verified = snapshot_info(
575 path,
576 file_bytes,
577 &decoded,
578 verified_counts.0,
579 verified_counts.1,
580 verified_counts.2,
581 );
582 if let Some(limits) = limits {
583 validate_read_limits(&verified, limits)?;
584 }
585 check_snapshot_deadline(deadline)?;
586 ensure_snapshot_file_length_unchanged(file, file_bytes)?;
587 check_snapshot_deadline(deadline)?;
588 Ok(verified)
589}
590
591fn snapshot_info(
592 path: &Path,
593 file_bytes: u64,
594 decoded: &DecodedHeader,
595 vector_space_count: u64,
596 vector_count: u64,
597 lexical_index_count: u64,
598) -> SnapshotInfo {
599 SnapshotInfo {
600 path: path.to_path_buf(),
601 disk_format_version: decoded.disk_format_version,
602 checkpoint_sequence: decoded.checkpoint_sequence,
603 checkpoint_digest: decoded.checkpoint_digest,
604 entry_count: decoded.entry_count,
605 vector_space_count,
606 vector_count,
607 lexical_index_count,
608 receipt_count: decoded.receipt_count,
609 snapshot_digest: decoded.expected_digest,
610 file_bytes,
611 }
612}
613
614pub fn load_snapshot(
626 path: impl AsRef<Path>,
627 limits: &SnapshotReadLimits,
628) -> Result<SnapshotContents, SnapshotError> {
629 load_snapshot_inner(path.as_ref(), limits, None)
630}
631
632pub fn load_snapshot_with_timeout(
641 path: impl AsRef<Path>,
642 limits: &SnapshotReadLimits,
643 timeout: Duration,
644) -> Result<SnapshotContents, SnapshotError> {
645 let deadline = OperationDeadline::new(timeout);
646 load_snapshot_inner(path.as_ref(), limits, Some(&deadline))
647}
648
649fn load_snapshot_inner(
650 path: &Path,
651 limits: &SnapshotReadLimits,
652 deadline: Option<&OperationDeadline>,
653) -> Result<SnapshotContents, SnapshotError> {
654 check_snapshot_deadline(deadline)?;
655 let mut collector = SnapshotCollector {
656 entries: Vec::new(),
657 vector_spaces: Vec::new(),
658 vectors: Vec::new(),
659 lexical_indexes: Vec::new(),
660 decoded_bytes: 0,
661 limits,
662 };
663 let info = read_snapshot_records_with_limits(path, &mut collector, Some(limits), deadline)?;
664 Ok(SnapshotContents {
665 info,
666 entries: collector.entries,
667 vector_spaces: collector.vector_spaces,
668 vectors: collector.vectors,
669 lexical_indexes: collector.lexical_indexes,
670 })
671}
672
673pub(crate) trait SnapshotRecordVisitor {
679 fn put(&mut self, key: &[u8], value: &[u8]) -> Result<(), SnapshotError>;
680 fn vector_space(&mut self, _definition: &VectorSpaceDefinition) -> Result<(), SnapshotError> {
681 Ok(())
682 }
683 fn vector(
684 &mut self,
685 _space: &VectorSpaceName,
686 _key: &[u8],
687 _vector: &Q15Vector,
688 ) -> Result<(), SnapshotError> {
689 Ok(())
690 }
691 fn lexical_index(&mut self, _definition: &LexicalIndexDefinition) -> Result<(), SnapshotError> {
692 Ok(())
693 }
694 fn receipt(&mut self, receipt: &CommitReceipt) -> Result<(), SnapshotError>;
695}
696
697pub(crate) fn read_snapshot_records_with_policy(
698 path: &Path,
699 visitor: &mut impl SnapshotRecordVisitor,
700 limits: &SnapshotReadLimits,
701 deadline: &OperationDeadline,
702) -> Result<SnapshotInfo, SnapshotError> {
703 read_snapshot_records_with_limits(path, visitor, Some(limits), Some(deadline))
704}
705
706fn read_snapshot_records_with_limits(
707 path: &Path,
708 visitor: &mut impl SnapshotRecordVisitor,
709 limits: Option<&SnapshotReadLimits>,
710 deadline: Option<&OperationDeadline>,
711) -> Result<SnapshotInfo, SnapshotError> {
712 check_snapshot_deadline(deadline)?;
713 let mut file = open_snapshot_file(path)?;
714 let file_bytes = snapshot_file_length(&file)?;
715 if let Some(limits) = limits {
716 validate_file_limit(file_bytes, limits)?;
717 }
718 let mut header = [0_u8; HEADER_LENGTH];
719 read_exact_or_invalid(&mut file, &mut header, "truncated header")?;
720 let decoded = decode_header(&header, file_bytes)?;
721 let (vector_space_count, vector_count, lexical_index_count) = verify_payload(
722 &mut file,
723 &header,
724 &decoded,
725 limits,
726 deadline,
727 Some(visitor),
728 )?;
729 let verified = snapshot_info(
730 path,
731 file_bytes,
732 &decoded,
733 vector_space_count,
734 vector_count,
735 lexical_index_count,
736 );
737 if let Some(limits) = limits {
738 validate_read_limits(&verified, limits)?;
739 }
740 check_snapshot_deadline(deadline)?;
741 ensure_snapshot_file_length_unchanged(&file, file_bytes)?;
742 check_snapshot_deadline(deadline)?;
743 Ok(verified)
744}
745
746fn open_snapshot_file(path: &Path) -> Result<File, SnapshotError> {
747 if !std::fs::metadata(path)?.is_file() {
748 return Err(SnapshotError::Invalid {
749 reason: "snapshot is not a regular file",
750 });
751 }
752 let file = File::open(path)?;
753 snapshot_file_length(&file)?;
754 Ok(file)
755}
756
757fn snapshot_file_length(file: &File) -> Result<u64, SnapshotError> {
758 let metadata = file.metadata()?;
759 if !metadata.is_file() {
760 return Err(SnapshotError::Invalid {
761 reason: "snapshot is not a regular file",
762 });
763 }
764 Ok(metadata.len())
765}
766
767fn ensure_snapshot_file_length_unchanged(file: &File, expected: u64) -> Result<(), SnapshotError> {
768 if snapshot_file_length(file)? != expected {
769 return Err(SnapshotError::Invalid {
770 reason: "snapshot changed while being read",
771 });
772 }
773 Ok(())
774}
775
776fn validate_file_limit(file_bytes: u64, limits: &SnapshotReadLimits) -> Result<(), SnapshotError> {
777 if file_bytes > limits.file_bytes {
778 return Err(SnapshotError::FileLimitExceeded {
779 actual: file_bytes,
780 maximum: limits.file_bytes,
781 });
782 }
783 Ok(())
784}
785
786fn check_snapshot_deadline(deadline: Option<&OperationDeadline>) -> Result<(), SnapshotError> {
787 deadline
788 .map(OperationDeadline::check)
789 .transpose()
790 .map(|_| ())
791 .map_err(SnapshotError::from)
792}
793
794fn validate_read_limits(
795 info: &SnapshotInfo,
796 limits: &SnapshotReadLimits,
797) -> Result<(), SnapshotError> {
798 validate_file_limit(info.file_bytes, limits)?;
799 validate_logical_record_limit(
800 info.entry_count,
801 info.vector_space_count,
802 info.vector_count,
803 info.lexical_index_count,
804 limits,
805 )
806}
807
808fn validate_logical_record_limit(
809 entry_count: u64,
810 vector_space_count: u64,
811 vector_count: u64,
812 lexical_index_count: u64,
813 limits: &SnapshotReadLimits,
814) -> Result<(), SnapshotError> {
815 let logical_records = entry_count
816 .checked_add(vector_space_count)
817 .and_then(|count| count.checked_add(vector_count))
818 .and_then(|count| count.checked_add(lexical_index_count))
819 .ok_or(SnapshotError::EntryLimitExceeded {
820 actual: u64::MAX,
821 maximum: limits.entries,
822 })?;
823 if logical_records > limits.entries {
824 return Err(SnapshotError::EntryLimitExceeded {
825 actual: logical_records,
826 maximum: limits.entries,
827 });
828 }
829 Ok(())
830}
831
832struct SnapshotCollector<'limits> {
833 entries: Vec<SnapshotEntry>,
834 vector_spaces: Vec<VectorSpaceDefinition>,
835 vectors: Vec<SnapshotVectorEntry>,
836 lexical_indexes: Vec<LexicalIndexDefinition>,
837 decoded_bytes: u64,
838 limits: &'limits SnapshotReadLimits,
839}
840
841impl SnapshotRecordVisitor for SnapshotCollector<'_> {
842 fn put(&mut self, key: &[u8], value: &[u8]) -> Result<(), SnapshotError> {
843 let next_entry_count = u64::try_from(self.entries.len())
844 .ok()
845 .and_then(|count| count.checked_add(1))
846 .ok_or(SnapshotError::EntryLimitExceeded {
847 actual: u64::MAX,
848 maximum: self.limits.entries,
849 })?;
850 if next_entry_count > self.limits.entries {
851 return Err(SnapshotError::EntryLimitExceeded {
852 actual: next_entry_count,
853 maximum: self.limits.entries,
854 });
855 }
856 let entry_bytes = u64::try_from(key.len())
857 .ok()
858 .and_then(|key_bytes| {
859 u64::try_from(value.len())
860 .ok()
861 .and_then(|value_bytes| key_bytes.checked_add(value_bytes))
862 })
863 .ok_or(SnapshotError::DecodedBytesLimitExceeded {
864 maximum: self.limits.decoded_bytes,
865 })?;
866 self.decoded_bytes = self.decoded_bytes.checked_add(entry_bytes).ok_or(
867 SnapshotError::DecodedBytesLimitExceeded {
868 maximum: self.limits.decoded_bytes,
869 },
870 )?;
871 if self.decoded_bytes > self.limits.decoded_bytes {
872 return Err(SnapshotError::DecodedBytesLimitExceeded {
873 maximum: self.limits.decoded_bytes,
874 });
875 }
876 self.entries.push(SnapshotEntry {
877 key: key.to_vec(),
878 value: value.to_vec(),
879 });
880 Ok(())
881 }
882
883 fn receipt(&mut self, _receipt: &CommitReceipt) -> Result<(), SnapshotError> {
884 Ok(())
885 }
886
887 fn vector_space(&mut self, definition: &VectorSpaceDefinition) -> Result<(), SnapshotError> {
888 self.add_decoded_bytes(definition.name.as_str().len())?;
889 self.vector_spaces.push(definition.clone());
890 Ok(())
891 }
892
893 fn vector(
894 &mut self,
895 space: &VectorSpaceName,
896 key: &[u8],
897 vector: &Q15Vector,
898 ) -> Result<(), SnapshotError> {
899 let vector_bytes = vector
900 .as_slice()
901 .len()
902 .checked_mul(2)
903 .and_then(|length| length.checked_add(space.as_str().len()))
904 .and_then(|length| length.checked_add(key.len()))
905 .ok_or(SnapshotError::DecodedBytesLimitExceeded {
906 maximum: self.limits.decoded_bytes,
907 })?;
908 self.add_decoded_bytes(vector_bytes)?;
909 self.vectors.push(SnapshotVectorEntry {
910 space: space.clone(),
911 key: key.to_vec(),
912 vector: vector.clone(),
913 });
914 Ok(())
915 }
916
917 fn lexical_index(&mut self, definition: &LexicalIndexDefinition) -> Result<(), SnapshotError> {
918 let encoded_length = encode_lexical_index(definition)?.len();
919 self.add_decoded_bytes(encoded_length)?;
920 self.lexical_indexes.push(definition.clone());
921 Ok(())
922 }
923}
924
925impl SnapshotCollector<'_> {
926 fn add_decoded_bytes(&mut self, bytes: usize) -> Result<(), SnapshotError> {
927 let bytes = u64::try_from(bytes).map_err(|_| SnapshotError::DecodedBytesLimitExceeded {
928 maximum: self.limits.decoded_bytes,
929 })?;
930 self.decoded_bytes = self.decoded_bytes.checked_add(bytes).ok_or(
931 SnapshotError::DecodedBytesLimitExceeded {
932 maximum: self.limits.decoded_bytes,
933 },
934 )?;
935 if self.decoded_bytes > self.limits.decoded_bytes {
936 return Err(SnapshotError::DecodedBytesLimitExceeded {
937 maximum: self.limits.decoded_bytes,
938 });
939 }
940 Ok(())
941 }
942}
943
944#[derive(Clone, Copy, Debug)]
945struct DecodedHeader {
946 disk_format_version: u16,
947 checkpoint_sequence: u64,
948 checkpoint_digest: Option<[u8; 32]>,
949 entry_count: u64,
950 receipt_count: u64,
951 payload_length: u64,
952 expected_checksum: u32,
953 expected_digest: [u8; 32],
954}
955
956fn decode_header(
957 header: &[u8; HEADER_LENGTH],
958 file_bytes: u64,
959) -> Result<DecodedHeader, SnapshotError> {
960 if header[0..8] != MAGIC {
961 return Err(SnapshotError::Invalid {
962 reason: "bad magic",
963 });
964 }
965 let version = u16::from_le_bytes(copy_array(&header[8..10]));
966 if !(MIN_DISK_FORMAT_VERSION..=DISK_FORMAT_VERSION).contains(&version) {
967 return Err(SnapshotError::UnsupportedVersion {
968 found: version,
969 supported: DISK_FORMAT_VERSION,
970 });
971 }
972 if u16::from_le_bytes(copy_array(&header[10..12])) != 0 {
973 return Err(SnapshotError::Invalid {
974 reason: "unsupported flags",
975 });
976 }
977
978 let checkpoint_sequence = u64::from_le_bytes(copy_array(&header[12..20]));
979 let raw_checkpoint_digest: [u8; 32] = copy_array(&header[20..52]);
980 let checkpoint_digest = if checkpoint_sequence == 0 {
981 if raw_checkpoint_digest != [0; 32] {
982 return Err(SnapshotError::Invalid {
983 reason: "empty checkpoint has a digest",
984 });
985 }
986 None
987 } else {
988 Some(raw_checkpoint_digest)
989 };
990 let entry_count = u64::from_le_bytes(copy_array(&header[52..60]));
991 let receipt_count = u64::from_le_bytes(copy_array(&header[60..68]));
992 if checkpoint_sequence == 0 && receipt_count != 0 {
993 return Err(SnapshotError::Invalid {
994 reason: "empty checkpoint has idempotency receipts",
995 });
996 }
997 let payload_length = u64::from_le_bytes(copy_array(&header[68..76]));
998 let expected_file_bytes =
999 HEADER_LENGTH_U64
1000 .checked_add(payload_length)
1001 .ok_or(SnapshotError::Invalid {
1002 reason: "file length overflow",
1003 })?;
1004 if file_bytes != expected_file_bytes {
1005 return Err(SnapshotError::Invalid {
1006 reason: "file length mismatch",
1007 });
1008 }
1009
1010 Ok(DecodedHeader {
1011 disk_format_version: version,
1012 checkpoint_sequence,
1013 checkpoint_digest,
1014 entry_count,
1015 receipt_count,
1016 payload_length,
1017 expected_checksum: u32::from_le_bytes(copy_array(&header[76..80])),
1018 expected_digest: copy_array(&header[80..112]),
1019 })
1020}
1021
1022#[allow(clippy::too_many_lines)]
1023fn verify_payload(
1024 file: &mut File,
1025 header: &[u8; HEADER_LENGTH],
1026 decoded: &DecodedHeader,
1027 limits: Option<&SnapshotReadLimits>,
1028 deadline: Option<&OperationDeadline>,
1029 mut visitor: Option<&mut dyn SnapshotRecordVisitor>,
1030) -> Result<(u64, u64, u64), SnapshotError> {
1031 check_snapshot_deadline(deadline)?;
1032 let mut checksum = crc32c::crc32c(&header[..CHECKSUM_PREFIX_LENGTH]);
1033 let mut hasher = blake3::Hasher::new();
1034 hasher.update(&header[..DIGEST_PREFIX_LENGTH]);
1035 let mut consumed = 0_u64;
1036 let mut decoded_bytes = 0_u64;
1037 let mut counts_bytes = [0_u8; V2_COUNTS_LENGTH];
1038 let (vector_space_count, vector_count, lexical_index_count) =
1039 if decoded.disk_format_version >= 2 {
1040 read_payload_exact(
1041 file,
1042 &mut counts_bytes,
1043 &mut consumed,
1044 decoded.payload_length,
1045 )?;
1046 checksum = crc32c::crc32c_append(checksum, &counts_bytes);
1047 hasher.update(&counts_bytes);
1048 (
1049 u64::from_le_bytes(copy_array(&counts_bytes[..8])),
1050 u64::from_le_bytes(copy_array(&counts_bytes[8..16])),
1051 u64::from_le_bytes(copy_array(&counts_bytes[16..24])),
1052 )
1053 } else {
1054 (0, 0, 0)
1055 };
1056 if let Some(limits) = limits {
1057 validate_logical_record_limit(
1058 decoded.entry_count,
1059 vector_space_count,
1060 vector_count,
1061 lexical_index_count,
1062 limits,
1063 )?;
1064 }
1065 let mut previous_key: Option<Vec<u8>> = None;
1066 let mut buffer = vec![0_u8; COPY_BUFFER_LENGTH].into_boxed_slice();
1067 for _ in 0..decoded.entry_count {
1068 check_snapshot_deadline(deadline)?;
1069 let mut entry_header = [0_u8; ENTRY_HEADER_LENGTH];
1070 read_payload_exact(
1071 file,
1072 &mut entry_header,
1073 &mut consumed,
1074 decoded.payload_length,
1075 )?;
1076 checksum = crc32c::crc32c_append(checksum, &entry_header);
1077 hasher.update(&entry_header);
1078 let key_length = usize::try_from(u32::from_le_bytes(copy_array(&entry_header[..4])))
1079 .map_err(|_| SnapshotError::Invalid {
1080 reason: "key length overflow",
1081 })?;
1082 let value_length = u64::from_le_bytes(copy_array(&entry_header[4..12]));
1083 if key_length == 0 || key_length > MAX_KEY_BYTES {
1084 return Err(SnapshotError::Invalid {
1085 reason: "invalid key length",
1086 });
1087 }
1088 let key_bytes =
1089 u64::try_from(key_length).map_err(|_| SnapshotError::DecodedBytesLimitExceeded {
1090 maximum: limits.map_or(u64::MAX, |limits| limits.decoded_bytes),
1091 })?;
1092 account_decoded_bytes(
1093 &mut decoded_bytes,
1094 key_bytes.checked_add(value_length),
1095 limits,
1096 )?;
1097 if visitor.is_some() && value_length > MAX_OPERATION_BYTES as u64 {
1098 return Err(SnapshotError::Invalid {
1099 reason: "record exceeds restore bounds",
1100 });
1101 }
1102
1103 let mut key = vec![0_u8; key_length];
1104 read_payload_exact_with_deadline(
1105 file,
1106 &mut key,
1107 &mut consumed,
1108 decoded.payload_length,
1109 deadline,
1110 )?;
1111 checksum = crc32c::crc32c_append(checksum, &key);
1112 hasher.update(&key);
1113 if previous_key
1114 .as_ref()
1115 .is_some_and(|previous| previous >= &key)
1116 {
1117 return Err(SnapshotError::Invalid {
1118 reason: "keys are not strictly sorted",
1119 });
1120 }
1121 if let Some(visitor) = visitor.as_deref_mut() {
1122 let value_length =
1123 usize::try_from(value_length).map_err(|_| SnapshotError::Invalid {
1124 reason: "value length overflow",
1125 })?;
1126 let mut value = vec![0_u8; value_length];
1127 read_payload_exact_with_deadline(
1128 file,
1129 &mut value,
1130 &mut consumed,
1131 decoded.payload_length,
1132 deadline,
1133 )?;
1134 checksum = crc32c::crc32c_append(checksum, &value);
1135 hasher.update(&value);
1136 visitor.put(&key, &value)?;
1137 } else {
1138 let mut remaining = value_length;
1139 while remaining > 0 {
1140 check_snapshot_deadline(deadline)?;
1141 let chunk_length =
1142 usize::try_from(remaining.min(COPY_BUFFER_LENGTH_U64)).map_err(|_| {
1143 SnapshotError::Invalid {
1144 reason: "value length overflow",
1145 }
1146 })?;
1147 let chunk = &mut buffer[..chunk_length];
1148 read_payload_exact(file, chunk, &mut consumed, decoded.payload_length)?;
1149 checksum = crc32c::crc32c_append(checksum, chunk);
1150 hasher.update(chunk);
1151 remaining -= u64::try_from(chunk_length).map_err(|_| SnapshotError::Invalid {
1152 reason: "value length overflow",
1153 })?;
1154 }
1155 }
1156 previous_key = Some(key);
1157 }
1158 let mut definitions = BTreeMap::new();
1159 let mut previous_space: Option<VectorSpaceName> = None;
1160 for _ in 0..vector_space_count {
1161 check_snapshot_deadline(deadline)?;
1162 let encoded = read_encoded_vector_space(file, decoded, &mut consumed, deadline)?;
1163 checksum = crc32c::crc32c_append(checksum, &encoded);
1164 hasher.update(&encoded);
1165 let definition = decode_vector_space(&encoded)?;
1166 if previous_space
1167 .as_ref()
1168 .is_some_and(|previous| previous >= &definition.name)
1169 {
1170 return Err(SnapshotError::Invalid {
1171 reason: "vector spaces are not strictly sorted",
1172 });
1173 }
1174 account_decoded_bytes(
1175 &mut decoded_bytes,
1176 u64::try_from(definition.name.as_str().len()).ok(),
1177 limits,
1178 )?;
1179 if let Some(visitor) = visitor.as_deref_mut() {
1180 visitor.vector_space(&definition)?;
1181 }
1182 previous_space = Some(definition.name.clone());
1183 definitions.insert(definition.name.clone(), definition);
1184 }
1185 let mut previous_vector_identity: Option<(VectorSpaceName, Vec<u8>)> = None;
1186 for _ in 0..vector_count {
1187 check_snapshot_deadline(deadline)?;
1188 let encoded = read_encoded_vector(file, decoded, &mut consumed, deadline)?;
1189 checksum = crc32c::crc32c_append(checksum, &encoded);
1190 hasher.update(&encoded);
1191 let (space, key, vector) = decode_vector(&encoded)?;
1192 if previous_vector_identity
1193 .as_ref()
1194 .is_some_and(|previous| previous >= &(space.clone(), key.clone()))
1195 {
1196 return Err(SnapshotError::Invalid {
1197 reason: "vectors are not strictly sorted",
1198 });
1199 }
1200 let definition = definitions.get(&space).ok_or(SnapshotError::Invalid {
1201 reason: "vector references an undefined space",
1202 })?;
1203 definition
1204 .validate_vector(&vector)
1205 .map_err(|_| SnapshotError::Invalid {
1206 reason: "vector dimension does not match its space",
1207 })?;
1208 let vector_bytes = vector
1209 .as_slice()
1210 .len()
1211 .checked_mul(2)
1212 .and_then(|length| length.checked_add(space.as_str().len()))
1213 .and_then(|length| length.checked_add(key.len()))
1214 .and_then(|length| u64::try_from(length).ok());
1215 account_decoded_bytes(&mut decoded_bytes, vector_bytes, limits)?;
1216 if let Some(visitor) = visitor.as_deref_mut() {
1217 visitor.vector(&space, &key, &vector)?;
1218 }
1219 previous_vector_identity = Some((space, key));
1220 }
1221 let mut previous_lexical_name: Option<VectorSpaceName> = None;
1222 for _ in 0..lexical_index_count {
1223 check_snapshot_deadline(deadline)?;
1224 let encoded = read_encoded_lexical_index(file, decoded, &mut consumed, deadline)?;
1225 checksum = crc32c::crc32c_append(checksum, &encoded);
1226 hasher.update(&encoded);
1227 let definition = decode_lexical_index(&encoded)?;
1228 if previous_lexical_name
1229 .as_ref()
1230 .is_some_and(|previous| previous >= &definition.name)
1231 {
1232 return Err(SnapshotError::Invalid {
1233 reason: "lexical indexes are not strictly sorted",
1234 });
1235 }
1236 account_decoded_bytes(
1237 &mut decoded_bytes,
1238 u64::try_from(encoded.len()).ok(),
1239 limits,
1240 )?;
1241 if let Some(visitor) = visitor.as_deref_mut() {
1242 visitor.lexical_index(&definition)?;
1243 }
1244 previous_lexical_name = Some(definition.name);
1245 }
1246 let mut previous_transaction_id = None;
1247 for _ in 0..decoded.receipt_count {
1248 check_snapshot_deadline(deadline)?;
1249 let mut encoded = [0_u8; RECEIPT_LENGTH];
1250 read_payload_exact(file, &mut encoded, &mut consumed, decoded.payload_length)?;
1251 checksum = crc32c::crc32c_append(checksum, &encoded);
1252 hasher.update(&encoded);
1253
1254 let transaction_id: [u8; 16] = copy_array(&encoded[..16]);
1255 if previous_transaction_id
1256 .as_ref()
1257 .is_some_and(|previous| previous >= &transaction_id)
1258 {
1259 return Err(SnapshotError::Invalid {
1260 reason: "transaction identifiers are not strictly sorted",
1261 });
1262 }
1263 previous_transaction_id = Some(transaction_id);
1264 let commit_sequence = u64::from_le_bytes(copy_array(&encoded[16..24]));
1265 if commit_sequence == 0 || commit_sequence > decoded.checkpoint_sequence {
1266 return Err(SnapshotError::Invalid {
1267 reason: "idempotency receipt exceeds snapshot checkpoint",
1268 });
1269 }
1270 if let Some(visitor) = visitor.as_deref_mut() {
1271 visitor.receipt(&decode_snapshot_receipt(&encoded))?;
1272 }
1273 }
1274 check_snapshot_deadline(deadline)?;
1275 if consumed != decoded.payload_length {
1276 return Err(SnapshotError::Invalid {
1277 reason: "record counts do not consume payload",
1278 });
1279 }
1280 if checksum != decoded.expected_checksum {
1281 return Err(SnapshotError::Invalid {
1282 reason: "CRC32C mismatch",
1283 });
1284 }
1285 let actual_digest = *hasher.finalize().as_bytes();
1286 if actual_digest != decoded.expected_digest {
1287 return Err(SnapshotError::Invalid {
1288 reason: "BLAKE3 digest mismatch",
1289 });
1290 }
1291 Ok((vector_space_count, vector_count, lexical_index_count))
1292}
1293
1294fn account_decoded_bytes(
1295 total: &mut u64,
1296 bytes: Option<u64>,
1297 limits: Option<&SnapshotReadLimits>,
1298) -> Result<(), SnapshotError> {
1299 let Some(limits) = limits else {
1300 return Ok(());
1301 };
1302 let bytes = bytes.ok_or(SnapshotError::DecodedBytesLimitExceeded {
1303 maximum: limits.decoded_bytes,
1304 })?;
1305 *total = total
1306 .checked_add(bytes)
1307 .ok_or(SnapshotError::DecodedBytesLimitExceeded {
1308 maximum: limits.decoded_bytes,
1309 })?;
1310 if *total > limits.decoded_bytes {
1311 return Err(SnapshotError::DecodedBytesLimitExceeded {
1312 maximum: limits.decoded_bytes,
1313 });
1314 }
1315 Ok(())
1316}
1317
1318#[derive(Clone, Copy, Debug)]
1319struct Measurements {
1320 entry_count: u64,
1321 vector_space_count: u64,
1322 vector_count: u64,
1323 lexical_index_count: u64,
1324 receipt_count: u64,
1325 payload_length: u64,
1326}
1327
1328impl Measurements {
1329 fn v2_counts(self) -> [u8; V2_COUNTS_LENGTH] {
1330 let mut encoded = [0_u8; V2_COUNTS_LENGTH];
1331 encoded[..8].copy_from_slice(&self.vector_space_count.to_le_bytes());
1332 encoded[8..16].copy_from_slice(&self.vector_count.to_le_bytes());
1333 encoded[16..24].copy_from_slice(&self.lexical_index_count.to_le_bytes());
1334 encoded
1335 }
1336}
1337
1338#[allow(clippy::too_many_lines)]
1339fn measure_payload(
1340 index: &MaterializedIndex,
1341 checkpoint_sequence: u64,
1342 disk_format_version: u16,
1343 limits: &SnapshotReadLimits,
1344 deadline: &OperationDeadline,
1345) -> Result<Measurements, SnapshotError> {
1346 check_snapshot_deadline(Some(deadline))?;
1347 if !(MIN_DISK_FORMAT_VERSION..=DISK_FORMAT_VERSION).contains(&disk_format_version) {
1348 return Err(SnapshotError::UnsupportedVersion {
1349 found: disk_format_version,
1350 supported: DISK_FORMAT_VERSION,
1351 });
1352 }
1353 let mut entry_count = Some(0_u64);
1354 let mut payload_length = Some(0_u64);
1355 let mut logical_records = 0_u64;
1356 let mut decoded_bytes = 0_u64;
1357 let mut valid = true;
1358 let mut limit_error = None;
1359 index.for_each_entry(|key, value| {
1360 if limit_error.is_some() || !valid {
1361 return;
1362 }
1363 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
1364 limit_error = Some(source);
1365 return;
1366 }
1367 if key.is_empty() || key.len() > MAX_KEY_BYTES {
1368 valid = false;
1369 return;
1370 }
1371 let Ok(key_length) = u64::try_from(key.len()) else {
1372 valid = false;
1373 return;
1374 };
1375 let Ok(value_length) = u64::try_from(value.len()) else {
1376 valid = false;
1377 return;
1378 };
1379 if let Err(source) = account_measurement_record(&mut logical_records, limits) {
1380 limit_error = Some(source);
1381 return;
1382 }
1383 if let Err(source) = account_decoded_bytes(
1384 &mut decoded_bytes,
1385 key_length.checked_add(value_length),
1386 Some(limits),
1387 ) {
1388 limit_error = Some(source);
1389 return;
1390 }
1391 entry_count = entry_count.and_then(|count| count.checked_add(1));
1392 payload_length = payload_length.and_then(|length| {
1393 length
1394 .checked_add(ENTRY_HEADER_LENGTH_U64)
1395 .and_then(|length| length.checked_add(key_length))
1396 .and_then(|length| length.checked_add(value_length))
1397 });
1398 if let Err(source) = validate_measured_file_bytes(payload_length, limits) {
1399 limit_error = Some(source);
1400 }
1401 })?;
1402 if let Some(source) = limit_error.take() {
1403 return Err(source);
1404 }
1405 let mut vector_space_count = Some(0_u64);
1406 let mut vector_count = Some(0_u64);
1407 let mut lexical_index_count = Some(0_u64);
1408 if disk_format_version >= 2 {
1409 payload_length = payload_length.and_then(|length| length.checked_add(V2_COUNTS_LENGTH_U64));
1410 index.for_each_vector_space(|definition| {
1411 if limit_error.is_some() || !valid {
1412 return;
1413 }
1414 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
1415 limit_error = Some(source);
1416 return;
1417 }
1418 let Ok(name_length) = u64::try_from(definition.name.as_str().len()) else {
1419 valid = false;
1420 return;
1421 };
1422 if let Err(source) = account_measurement_record(&mut logical_records, limits) {
1423 limit_error = Some(source);
1424 return;
1425 }
1426 if let Err(source) =
1427 account_decoded_bytes(&mut decoded_bytes, Some(name_length), Some(limits))
1428 {
1429 limit_error = Some(source);
1430 return;
1431 }
1432 vector_space_count = vector_space_count.and_then(|count| count.checked_add(1));
1433 payload_length = payload_length.and_then(|length| {
1434 length
1435 .checked_add(VECTOR_SPACE_FIXED_LENGTH_U64)
1436 .and_then(|length| length.checked_add(name_length))
1437 });
1438 if let Err(source) = validate_measured_file_bytes(payload_length, limits) {
1439 limit_error = Some(source);
1440 }
1441 })?;
1442 if let Some(source) = limit_error.take() {
1443 return Err(source);
1444 }
1445 index.for_each_vector(|space, key, vector| {
1446 if limit_error.is_some() || !valid {
1447 return;
1448 }
1449 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
1450 limit_error = Some(source);
1451 return;
1452 }
1453 let Ok(name_length) = u64::try_from(space.as_str().len()) else {
1454 valid = false;
1455 return;
1456 };
1457 let Ok(key_length) = u64::try_from(key.len()) else {
1458 valid = false;
1459 return;
1460 };
1461 let Ok(vector_bytes) = u64::try_from(vector.as_slice().len().saturating_mul(2)) else {
1462 valid = false;
1463 return;
1464 };
1465 if let Err(source) = account_measurement_record(&mut logical_records, limits) {
1466 limit_error = Some(source);
1467 return;
1468 }
1469 if let Err(source) = account_decoded_bytes(
1470 &mut decoded_bytes,
1471 name_length
1472 .checked_add(key_length)
1473 .and_then(|length| length.checked_add(vector_bytes)),
1474 Some(limits),
1475 ) {
1476 limit_error = Some(source);
1477 return;
1478 }
1479 vector_count = vector_count.and_then(|count| count.checked_add(1));
1480 payload_length = payload_length.and_then(|length| {
1481 length
1482 .checked_add(VECTOR_FIXED_LENGTH_U64)
1483 .and_then(|length| length.checked_add(name_length))
1484 .and_then(|length| length.checked_add(key_length))
1485 .and_then(|length| length.checked_add(vector_bytes))
1486 });
1487 if let Err(source) = validate_measured_file_bytes(payload_length, limits) {
1488 limit_error = Some(source);
1489 }
1490 })?;
1491 if let Some(source) = limit_error.take() {
1492 return Err(source);
1493 }
1494 index.for_each_lexical_index(|definition| {
1495 if limit_error.is_some() || !valid {
1496 return;
1497 }
1498 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
1499 limit_error = Some(source);
1500 return;
1501 }
1502 let Ok(encoded) = encode_lexical_index(definition) else {
1503 valid = false;
1504 return;
1505 };
1506 let Ok(record_bytes) = u64::try_from(encoded.len()) else {
1507 valid = false;
1508 return;
1509 };
1510 if let Err(source) = account_measurement_record(&mut logical_records, limits) {
1511 limit_error = Some(source);
1512 return;
1513 }
1514 if let Err(source) =
1515 account_decoded_bytes(&mut decoded_bytes, Some(record_bytes), Some(limits))
1516 {
1517 limit_error = Some(source);
1518 return;
1519 }
1520 lexical_index_count = lexical_index_count.and_then(|count| count.checked_add(1));
1521 payload_length = payload_length.and_then(|length| length.checked_add(record_bytes));
1522 if let Err(source) = validate_measured_file_bytes(payload_length, limits) {
1523 limit_error = Some(source);
1524 }
1525 })?;
1526 if let Some(source) = limit_error.take() {
1527 return Err(source);
1528 }
1529 }
1530 let mut receipt_count = Some(0_u64);
1531 index.for_each_receipt(|receipt| {
1532 if limit_error.is_some() || !valid {
1533 return;
1534 }
1535 if let Err(source) = check_snapshot_deadline(Some(deadline)) {
1536 limit_error = Some(source);
1537 return;
1538 }
1539 if receipt.commit_sequence == 0 || receipt.commit_sequence > checkpoint_sequence {
1540 valid = false;
1541 return;
1542 }
1543 receipt_count = receipt_count.and_then(|count| count.checked_add(1));
1544 payload_length = payload_length.and_then(|length| length.checked_add(RECEIPT_LENGTH_U64));
1545 if let Err(source) = validate_measured_file_bytes(payload_length, limits) {
1546 limit_error = Some(source);
1547 }
1548 })?;
1549 if let Some(source) = limit_error {
1550 return Err(source);
1551 }
1552 if !valid {
1553 return Err(SnapshotError::Invalid {
1554 reason: "index contains an invalid key or idempotency receipt",
1555 });
1556 }
1557 let Some(entry_count) = entry_count else {
1558 return Err(SnapshotError::Invalid {
1559 reason: "entry count overflow",
1560 });
1561 };
1562 let Some(payload_length) = payload_length else {
1563 return Err(SnapshotError::Invalid {
1564 reason: "payload length overflow",
1565 });
1566 };
1567 let Some(receipt_count) = receipt_count else {
1568 return Err(SnapshotError::Invalid {
1569 reason: "receipt count overflow",
1570 });
1571 };
1572 let Some(vector_space_count) = vector_space_count else {
1573 return Err(SnapshotError::Invalid {
1574 reason: "vector-space count overflow",
1575 });
1576 };
1577 let Some(vector_count) = vector_count else {
1578 return Err(SnapshotError::Invalid {
1579 reason: "vector count overflow",
1580 });
1581 };
1582 let Some(lexical_index_count) = lexical_index_count else {
1583 return Err(SnapshotError::Invalid {
1584 reason: "lexical-index count overflow",
1585 });
1586 };
1587 Ok(Measurements {
1588 entry_count,
1589 vector_space_count,
1590 vector_count,
1591 lexical_index_count,
1592 receipt_count,
1593 payload_length,
1594 })
1595}
1596
1597fn account_measurement_record(
1598 logical_records: &mut u64,
1599 limits: &SnapshotReadLimits,
1600) -> Result<(), SnapshotError> {
1601 *logical_records = logical_records
1602 .checked_add(1)
1603 .ok_or(SnapshotError::EntryLimitExceeded {
1604 actual: u64::MAX,
1605 maximum: limits.entries,
1606 })?;
1607 if *logical_records > limits.entries {
1608 return Err(SnapshotError::EntryLimitExceeded {
1609 actual: *logical_records,
1610 maximum: limits.entries,
1611 });
1612 }
1613 Ok(())
1614}
1615
1616fn validate_measured_file_bytes(
1617 payload_length: Option<u64>,
1618 limits: &SnapshotReadLimits,
1619) -> Result<(), SnapshotError> {
1620 let file_bytes = payload_length
1621 .and_then(|length| HEADER_LENGTH_U64.checked_add(length))
1622 .ok_or(SnapshotError::FileLimitExceeded {
1623 actual: u64::MAX,
1624 maximum: limits.file_bytes,
1625 })?;
1626 validate_file_limit(file_bytes, limits)
1627}
1628
1629fn validate_measurement_limits(
1630 measurements: &Measurements,
1631 limits: &SnapshotReadLimits,
1632) -> Result<(), SnapshotError> {
1633 let info = SnapshotInfo {
1634 path: PathBuf::new(),
1635 disk_format_version: DISK_FORMAT_VERSION,
1636 checkpoint_sequence: 0,
1637 checkpoint_digest: None,
1638 entry_count: measurements.entry_count,
1639 vector_space_count: measurements.vector_space_count,
1640 vector_count: measurements.vector_count,
1641 lexical_index_count: measurements.lexical_index_count,
1642 receipt_count: measurements.receipt_count,
1643 snapshot_digest: [0; 32],
1644 file_bytes: HEADER_LENGTH_U64
1645 .checked_add(measurements.payload_length)
1646 .ok_or(SnapshotError::FileLimitExceeded {
1647 actual: u64::MAX,
1648 maximum: limits.file_bytes,
1649 })?,
1650 };
1651 validate_read_limits(&info, limits)
1652}
1653
1654fn encode_entry_header(
1655 key: &[u8],
1656 value: &[u8],
1657) -> Result<[u8; ENTRY_HEADER_LENGTH], SnapshotError> {
1658 let key_length = u32::try_from(key.len()).map_err(|_| SnapshotError::Invalid {
1659 reason: "key length overflow",
1660 })?;
1661 let value_length = u64::try_from(value.len()).map_err(|_| SnapshotError::Invalid {
1662 reason: "value length overflow",
1663 })?;
1664 let mut entry_header = [0_u8; ENTRY_HEADER_LENGTH];
1665 entry_header[..4].copy_from_slice(&key_length.to_le_bytes());
1666 entry_header[4..].copy_from_slice(&value_length.to_le_bytes());
1667 Ok(entry_header)
1668}
1669
1670fn write_entry(
1671 writer: &mut impl Write,
1672 hasher: &mut blake3::Hasher,
1673 key: &[u8],
1674 value: &[u8],
1675 deadline: Option<&OperationDeadline>,
1676) -> Result<(), SnapshotError> {
1677 let entry_header = encode_entry_header(key, value)?;
1678 for bytes in [&entry_header[..], key, value] {
1679 for chunk in bytes.chunks(COPY_BUFFER_LENGTH) {
1680 check_snapshot_deadline(deadline)?;
1681 writer.write_all(chunk)?;
1682 hasher.update(chunk);
1683 }
1684 }
1685 Ok(())
1686}
1687
1688fn encode_vector_space(definition: &VectorSpaceDefinition) -> Result<Vec<u8>, SnapshotError> {
1689 let name = definition.name.as_str().as_bytes();
1690 let name_length = u8::try_from(name.len()).map_err(|_| SnapshotError::Invalid {
1691 reason: "vector-space name length overflow",
1692 })?;
1693 let mut encoded = Vec::with_capacity(name.len() + 5);
1694 encoded.push(name_length);
1695 encoded.extend_from_slice(name);
1696 encoded.extend_from_slice(&definition.dimension.to_le_bytes());
1697 encoded.push(definition.metric as u8);
1698 encoded.push(1);
1699 Ok(encoded)
1700}
1701
1702fn decode_vector_space(encoded: &[u8]) -> Result<VectorSpaceDefinition, SnapshotError> {
1703 let name_length = usize::from(*encoded.first().ok_or(SnapshotError::Invalid {
1704 reason: "truncated vector-space record",
1705 })?);
1706 let expected_length = name_length.checked_add(5).ok_or(SnapshotError::Invalid {
1707 reason: "vector-space record length overflow",
1708 })?;
1709 if encoded.len() != expected_length || name_length == 0 {
1710 return Err(SnapshotError::Invalid {
1711 reason: "invalid vector-space record length",
1712 });
1713 }
1714 let name =
1715 std::str::from_utf8(&encoded[1..=name_length]).map_err(|_| SnapshotError::Invalid {
1716 reason: "invalid vector-space name",
1717 })?;
1718 let name = VectorSpaceName::new(name.to_owned()).map_err(|_| SnapshotError::Invalid {
1719 reason: "invalid vector-space name",
1720 })?;
1721 let dimension = u16::from_le_bytes(copy_array(&encoded[1 + name_length..3 + name_length]));
1722 if encoded[3 + name_length] != VectorMetric::Cosine as u8 || encoded[4 + name_length] != 1 {
1723 return Err(SnapshotError::Invalid {
1724 reason: "unsupported vector-space tags",
1725 });
1726 }
1727 VectorSpaceDefinition::cosine(name, dimension).map_err(|_| SnapshotError::Invalid {
1728 reason: "invalid vector-space dimension",
1729 })
1730}
1731
1732fn encode_vector(
1733 space: &VectorSpaceName,
1734 key: &[u8],
1735 vector: &Q15Vector,
1736) -> Result<Vec<u8>, SnapshotError> {
1737 if key.is_empty() || key.len() > MAX_KEY_BYTES {
1738 return Err(SnapshotError::Invalid {
1739 reason: "invalid vector object key",
1740 });
1741 }
1742 let space_name = space.as_str().as_bytes();
1743 let space_length = u8::try_from(space_name.len()).map_err(|_| SnapshotError::Invalid {
1744 reason: "vector-space name length overflow",
1745 })?;
1746 let key_length = u32::try_from(key.len()).map_err(|_| SnapshotError::Invalid {
1747 reason: "vector key length overflow",
1748 })?;
1749 let mut encoded =
1750 Vec::with_capacity(space_name.len() + key.len() + vector.as_slice().len() * 2 + 7);
1751 encoded.push(space_length);
1752 encoded.extend_from_slice(space_name);
1753 encoded.extend_from_slice(&key_length.to_le_bytes());
1754 encoded.extend_from_slice(key);
1755 encoded.extend_from_slice(&vector.dimension().to_le_bytes());
1756 for value in vector.as_slice() {
1757 encoded.extend_from_slice(&value.to_le_bytes());
1758 }
1759 Ok(encoded)
1760}
1761
1762fn decode_vector(encoded: &[u8]) -> Result<(VectorSpaceName, Vec<u8>, Q15Vector), SnapshotError> {
1763 let space_length = usize::from(*encoded.first().ok_or(SnapshotError::Invalid {
1764 reason: "truncated vector record",
1765 })?);
1766 let key_length_offset = 1_usize
1767 .checked_add(space_length)
1768 .ok_or(SnapshotError::Invalid {
1769 reason: "vector record length overflow",
1770 })?;
1771 let key_length_end = key_length_offset
1772 .checked_add(4)
1773 .ok_or(SnapshotError::Invalid {
1774 reason: "vector record length overflow",
1775 })?;
1776 let key_length = usize::try_from(u32::from_le_bytes(copy_array(
1777 encoded
1778 .get(key_length_offset..key_length_end)
1779 .ok_or(SnapshotError::Invalid {
1780 reason: "truncated vector record",
1781 })?,
1782 )))
1783 .map_err(|_| SnapshotError::Invalid {
1784 reason: "vector key length overflow",
1785 })?;
1786 if space_length == 0 || key_length == 0 || key_length > MAX_KEY_BYTES {
1787 return Err(SnapshotError::Invalid {
1788 reason: "invalid vector identity",
1789 });
1790 }
1791 let key_end = key_length_end
1792 .checked_add(key_length)
1793 .ok_or(SnapshotError::Invalid {
1794 reason: "vector record length overflow",
1795 })?;
1796 let dimension_end = key_end.checked_add(2).ok_or(SnapshotError::Invalid {
1797 reason: "vector record length overflow",
1798 })?;
1799 let dimension = usize::from(u16::from_le_bytes(copy_array(
1800 encoded
1801 .get(key_end..dimension_end)
1802 .ok_or(SnapshotError::Invalid {
1803 reason: "truncated vector record",
1804 })?,
1805 )));
1806 let expected_length = dimension
1807 .checked_mul(2)
1808 .and_then(|length| length.checked_add(dimension_end))
1809 .ok_or(SnapshotError::Invalid {
1810 reason: "vector record length overflow",
1811 })?;
1812 if encoded.len() != expected_length {
1813 return Err(SnapshotError::Invalid {
1814 reason: "invalid vector record length",
1815 });
1816 }
1817 let space = std::str::from_utf8(&encoded[1..key_length_offset]).map_err(|_| {
1818 SnapshotError::Invalid {
1819 reason: "invalid vector-space name",
1820 }
1821 })?;
1822 let space = VectorSpaceName::new(space.to_owned()).map_err(|_| SnapshotError::Invalid {
1823 reason: "invalid vector-space name",
1824 })?;
1825 let key = encoded[key_length_end..key_end].to_vec();
1826 let values = encoded[dimension_end..]
1827 .chunks_exact(2)
1828 .map(|chunk| i16::from_le_bytes(copy_array(chunk)))
1829 .collect::<Vec<_>>();
1830 let vector = Q15Vector::new(values).map_err(|_| SnapshotError::Invalid {
1831 reason: "invalid Q15 vector",
1832 })?;
1833 Ok((space, key, vector))
1834}
1835
1836fn encode_lexical_index(definition: &LexicalIndexDefinition) -> Result<Vec<u8>, SnapshotError> {
1837 let name = definition.name.as_str().as_bytes();
1838 let name_length = u8::try_from(name.len()).map_err(|_| SnapshotError::Invalid {
1839 reason: "lexical-index name length overflow",
1840 })?;
1841 let field_count =
1842 u8::try_from(definition.fields.len()).map_err(|_| SnapshotError::Invalid {
1843 reason: "lexical-index field count overflow",
1844 })?;
1845 let mut encoded = Vec::new();
1846 encoded.push(name_length);
1847 encoded.extend_from_slice(name);
1848 encoded.push(1);
1849 encoded.push(field_count);
1850 for field in &definition.fields {
1851 let segment_count =
1852 u8::try_from(field.path.segments().len()).map_err(|_| SnapshotError::Invalid {
1853 reason: "lexical-index segment count overflow",
1854 })?;
1855 encoded.push(segment_count);
1856 for segment in field.path.segments() {
1857 let segment_length =
1858 u16::try_from(segment.len()).map_err(|_| SnapshotError::Invalid {
1859 reason: "lexical-index segment length overflow",
1860 })?;
1861 encoded.extend_from_slice(&segment_length.to_le_bytes());
1862 encoded.extend_from_slice(segment.as_bytes());
1863 }
1864 encoded.extend_from_slice(&field.weight_micros.to_le_bytes());
1865 }
1866 Ok(encoded)
1867}
1868
1869#[allow(clippy::too_many_lines)]
1870fn decode_lexical_index(encoded: &[u8]) -> Result<LexicalIndexDefinition, SnapshotError> {
1871 let name_length = usize::from(*encoded.first().ok_or(SnapshotError::Invalid {
1872 reason: "truncated lexical-index record",
1873 })?);
1874 let name_end = 1_usize
1875 .checked_add(name_length)
1876 .ok_or(SnapshotError::Invalid {
1877 reason: "lexical-index record length overflow",
1878 })?;
1879 if name_length == 0 || encoded.get(name_end) != Some(&1) {
1880 return Err(SnapshotError::Invalid {
1881 reason: "invalid lexical-index record prefix",
1882 });
1883 }
1884 let name = std::str::from_utf8(encoded.get(1..name_end).ok_or(SnapshotError::Invalid {
1885 reason: "truncated lexical-index name",
1886 })?)
1887 .map_err(|_| SnapshotError::Invalid {
1888 reason: "invalid lexical-index name",
1889 })?;
1890 let name = VectorSpaceName::new(name.to_owned()).map_err(|_| SnapshotError::Invalid {
1891 reason: "invalid lexical-index name",
1892 })?;
1893 let mut cursor = name_end.checked_add(1).ok_or(SnapshotError::Invalid {
1894 reason: "lexical-index record length overflow",
1895 })?;
1896 let field_count = usize::from(*encoded.get(cursor).ok_or(SnapshotError::Invalid {
1897 reason: "truncated lexical-index field count",
1898 })?);
1899 cursor = cursor.checked_add(1).ok_or(SnapshotError::Invalid {
1900 reason: "lexical-index record length overflow",
1901 })?;
1902 if field_count == 0 || field_count > MAX_LEXICAL_FIELDS {
1903 return Err(SnapshotError::Invalid {
1904 reason: "invalid lexical-index field count",
1905 });
1906 }
1907 let mut fields = Vec::with_capacity(field_count);
1908 for _ in 0..field_count {
1909 let segment_count = usize::from(*encoded.get(cursor).ok_or(SnapshotError::Invalid {
1910 reason: "truncated lexical-index path",
1911 })?);
1912 cursor = cursor.checked_add(1).ok_or(SnapshotError::Invalid {
1913 reason: "lexical-index record length overflow",
1914 })?;
1915 if segment_count == 0 || segment_count > MAX_LEXICAL_PATH_SEGMENTS {
1916 return Err(SnapshotError::Invalid {
1917 reason: "invalid lexical-index path",
1918 });
1919 }
1920 let mut segments = Vec::with_capacity(segment_count);
1921 for _ in 0..segment_count {
1922 let length_end = cursor.checked_add(2).ok_or(SnapshotError::Invalid {
1923 reason: "lexical-index record length overflow",
1924 })?;
1925 let length = usize::from(u16::from_le_bytes(copy_array(
1926 encoded
1927 .get(cursor..length_end)
1928 .ok_or(SnapshotError::Invalid {
1929 reason: "truncated lexical-index segment length",
1930 })?,
1931 )));
1932 cursor = length_end;
1933 if length == 0 || length > MAX_LEXICAL_PATH_SEGMENT_BYTES {
1934 return Err(SnapshotError::Invalid {
1935 reason: "invalid lexical-index segment length",
1936 });
1937 }
1938 let segment_end = cursor.checked_add(length).ok_or(SnapshotError::Invalid {
1939 reason: "lexical-index record length overflow",
1940 })?;
1941 let segment = std::str::from_utf8(encoded.get(cursor..segment_end).ok_or(
1942 SnapshotError::Invalid {
1943 reason: "truncated lexical-index segment",
1944 },
1945 )?)
1946 .map_err(|_| SnapshotError::Invalid {
1947 reason: "invalid lexical-index segment",
1948 })?
1949 .to_owned();
1950 cursor = segment_end;
1951 segments.push(segment);
1952 }
1953 let weight_end = cursor.checked_add(4).ok_or(SnapshotError::Invalid {
1954 reason: "lexical-index record length overflow",
1955 })?;
1956 let weight_micros = u32::from_le_bytes(copy_array(encoded.get(cursor..weight_end).ok_or(
1957 SnapshotError::Invalid {
1958 reason: "truncated lexical-index field weight",
1959 },
1960 )?));
1961 cursor = weight_end;
1962 fields.push(LexicalField {
1963 path: FieldPath::new(segments),
1964 weight_micros,
1965 });
1966 }
1967 if cursor != encoded.len() {
1968 return Err(SnapshotError::Invalid {
1969 reason: "invalid lexical-index record length",
1970 });
1971 }
1972 LexicalIndexDefinition::new(name, fields).map_err(|_| SnapshotError::Invalid {
1973 reason: "invalid lexical-index definition",
1974 })
1975}
1976
1977fn write_encoded_with_deadline(
1978 writer: &mut impl Write,
1979 hasher: &mut blake3::Hasher,
1980 encoded: Result<Vec<u8>, SnapshotError>,
1981 deadline: Option<&OperationDeadline>,
1982) -> Result<(), SnapshotError> {
1983 let encoded = encoded?;
1984 for chunk in encoded.chunks(COPY_BUFFER_LENGTH) {
1985 check_snapshot_deadline(deadline)?;
1986 writer.write_all(chunk)?;
1987 hasher.update(chunk);
1988 }
1989 Ok(())
1990}
1991
1992fn read_encoded_vector_space(
1993 reader: &mut impl Read,
1994 decoded: &DecodedHeader,
1995 consumed: &mut u64,
1996 deadline: Option<&OperationDeadline>,
1997) -> Result<Vec<u8>, SnapshotError> {
1998 let mut name_length = [0_u8; 1];
1999 read_payload_exact_with_deadline(
2000 reader,
2001 &mut name_length,
2002 consumed,
2003 decoded.payload_length,
2004 deadline,
2005 )?;
2006 let remaining = usize::from(name_length[0])
2007 .checked_add(4)
2008 .ok_or(SnapshotError::Invalid {
2009 reason: "vector-space record length overflow",
2010 })?;
2011 let mut encoded = vec![name_length[0]];
2012 let mut tail = vec![0_u8; remaining];
2013 read_payload_exact_with_deadline(
2014 reader,
2015 &mut tail,
2016 consumed,
2017 decoded.payload_length,
2018 deadline,
2019 )?;
2020 encoded.extend_from_slice(&tail);
2021 Ok(encoded)
2022}
2023
2024fn read_encoded_vector(
2025 reader: &mut impl Read,
2026 decoded: &DecodedHeader,
2027 consumed: &mut u64,
2028 deadline: Option<&OperationDeadline>,
2029) -> Result<Vec<u8>, SnapshotError> {
2030 let mut space_length = [0_u8; 1];
2031 read_payload_exact_with_deadline(
2032 reader,
2033 &mut space_length,
2034 consumed,
2035 decoded.payload_length,
2036 deadline,
2037 )?;
2038 let space_length = usize::from(space_length[0]);
2039 let mut prefix_tail = vec![0_u8; space_length + 4];
2040 read_payload_exact_with_deadline(
2041 reader,
2042 &mut prefix_tail,
2043 consumed,
2044 decoded.payload_length,
2045 deadline,
2046 )?;
2047 let key_length = usize::try_from(u32::from_le_bytes(copy_array(&prefix_tail[space_length..])))
2048 .map_err(|_| SnapshotError::Invalid {
2049 reason: "vector key length overflow",
2050 })?;
2051 if space_length == 0 || key_length == 0 || key_length > MAX_KEY_BYTES {
2052 return Err(SnapshotError::Invalid {
2053 reason: "invalid vector identity",
2054 });
2055 }
2056 let mut key_and_dimension = vec![0_u8; key_length + 2];
2057 read_payload_exact_with_deadline(
2058 reader,
2059 &mut key_and_dimension,
2060 consumed,
2061 decoded.payload_length,
2062 deadline,
2063 )?;
2064 let dimension = usize::from(u16::from_le_bytes(copy_array(
2065 &key_and_dimension[key_length..],
2066 )));
2067 let vector_bytes = dimension.checked_mul(2).ok_or(SnapshotError::Invalid {
2068 reason: "vector record length overflow",
2069 })?;
2070 let mut values = vec![0_u8; vector_bytes];
2071 read_payload_exact_with_deadline(
2072 reader,
2073 &mut values,
2074 consumed,
2075 decoded.payload_length,
2076 deadline,
2077 )?;
2078 let mut encoded =
2079 Vec::with_capacity(1 + prefix_tail.len() + key_and_dimension.len() + values.len());
2080 encoded.push(
2081 u8::try_from(space_length).map_err(|_| SnapshotError::Invalid {
2082 reason: "vector-space name length overflow",
2083 })?,
2084 );
2085 encoded.extend_from_slice(&prefix_tail);
2086 encoded.extend_from_slice(&key_and_dimension);
2087 encoded.extend_from_slice(&values);
2088 Ok(encoded)
2089}
2090
2091fn read_encoded_lexical_index(
2092 reader: &mut impl Read,
2093 decoded: &DecodedHeader,
2094 consumed: &mut u64,
2095 deadline: Option<&OperationDeadline>,
2096) -> Result<Vec<u8>, SnapshotError> {
2097 let mut name_length = [0_u8; 1];
2098 read_payload_exact_with_deadline(
2099 reader,
2100 &mut name_length,
2101 consumed,
2102 decoded.payload_length,
2103 deadline,
2104 )?;
2105 let name_length_usize = usize::from(name_length[0]);
2106 if name_length_usize == 0 {
2107 return Err(SnapshotError::Invalid {
2108 reason: "invalid lexical-index name length",
2109 });
2110 }
2111 let mut name_and_counts = vec![0_u8; name_length_usize + 2];
2112 read_payload_exact_with_deadline(
2113 reader,
2114 &mut name_and_counts,
2115 consumed,
2116 decoded.payload_length,
2117 deadline,
2118 )?;
2119 if name_and_counts[name_length_usize] != 1 {
2120 return Err(SnapshotError::Invalid {
2121 reason: "unsupported lexical-index record version",
2122 });
2123 }
2124 let field_count = usize::from(name_and_counts[name_length_usize + 1]);
2125 if field_count == 0 || field_count > MAX_LEXICAL_FIELDS {
2126 return Err(SnapshotError::Invalid {
2127 reason: "invalid lexical-index field count",
2128 });
2129 }
2130 let mut encoded = Vec::new();
2131 encoded.push(name_length[0]);
2132 encoded.extend_from_slice(&name_and_counts);
2133 for _ in 0..field_count {
2134 check_snapshot_deadline(deadline)?;
2135 let mut segment_count = [0_u8; 1];
2136 read_payload_exact_with_deadline(
2137 reader,
2138 &mut segment_count,
2139 consumed,
2140 decoded.payload_length,
2141 deadline,
2142 )?;
2143 let segment_count_usize = usize::from(segment_count[0]);
2144 if segment_count_usize == 0 || segment_count_usize > MAX_LEXICAL_PATH_SEGMENTS {
2145 return Err(SnapshotError::Invalid {
2146 reason: "invalid lexical-index path",
2147 });
2148 }
2149 encoded.push(segment_count[0]);
2150 for _ in 0..segment_count_usize {
2151 check_snapshot_deadline(deadline)?;
2152 let mut length = [0_u8; 2];
2153 read_payload_exact_with_deadline(
2154 reader,
2155 &mut length,
2156 consumed,
2157 decoded.payload_length,
2158 deadline,
2159 )?;
2160 let length_usize = usize::from(u16::from_le_bytes(length));
2161 if length_usize == 0 || length_usize > MAX_LEXICAL_PATH_SEGMENT_BYTES {
2162 return Err(SnapshotError::Invalid {
2163 reason: "invalid lexical-index segment length",
2164 });
2165 }
2166 let mut segment = vec![0_u8; length_usize];
2167 read_payload_exact_with_deadline(
2168 reader,
2169 &mut segment,
2170 consumed,
2171 decoded.payload_length,
2172 deadline,
2173 )?;
2174 encoded.extend_from_slice(&length);
2175 encoded.extend_from_slice(&segment);
2176 }
2177 let mut weight = [0_u8; 4];
2178 read_payload_exact_with_deadline(
2179 reader,
2180 &mut weight,
2181 consumed,
2182 decoded.payload_length,
2183 deadline,
2184 )?;
2185 encoded.extend_from_slice(&weight);
2186 }
2187 Ok(encoded)
2188}
2189
2190fn encode_receipt(receipt: &CommitReceipt) -> [u8; RECEIPT_LENGTH] {
2191 let mut encoded = [0_u8; RECEIPT_LENGTH];
2192 encoded[..16].copy_from_slice(receipt.transaction_id.as_bytes());
2193 encoded[16..24].copy_from_slice(&receipt.commit_sequence.to_le_bytes());
2194 encoded[24..56].copy_from_slice(&receipt.commit_digest);
2195 encoded[56..88].copy_from_slice(&receipt.transaction_digest);
2196 encoded
2197}
2198
2199fn decode_snapshot_receipt(encoded: &[u8; RECEIPT_LENGTH]) -> CommitReceipt {
2200 CommitReceipt {
2201 transaction_id: uuid::Uuid::from_bytes(copy_array(&encoded[..16])),
2202 commit_sequence: u64::from_le_bytes(copy_array(&encoded[16..24])),
2203 commit_digest: copy_array(&encoded[24..56]),
2204 transaction_digest: copy_array(&encoded[56..88]),
2205 }
2206}
2207
2208fn write_receipt(
2209 writer: &mut impl Write,
2210 hasher: &mut blake3::Hasher,
2211 receipt: &CommitReceipt,
2212) -> Result<(), SnapshotError> {
2213 let encoded = encode_receipt(receipt);
2214 writer.write_all(&encoded)?;
2215 hasher.update(&encoded);
2216 Ok(())
2217}
2218
2219fn read_payload_exact(
2220 reader: &mut impl Read,
2221 buffer: &mut [u8],
2222 consumed: &mut u64,
2223 payload_length: u64,
2224) -> Result<(), SnapshotError> {
2225 let length = u64::try_from(buffer.len()).map_err(|_| SnapshotError::Invalid {
2226 reason: "payload length overflow",
2227 })?;
2228 let next = consumed.checked_add(length).ok_or(SnapshotError::Invalid {
2229 reason: "payload length overflow",
2230 })?;
2231 if next > payload_length {
2232 return Err(SnapshotError::Invalid {
2233 reason: "entry exceeds payload",
2234 });
2235 }
2236 read_exact_or_invalid(reader, buffer, "truncated payload")?;
2237 *consumed = next;
2238 Ok(())
2239}
2240
2241fn read_payload_exact_with_deadline(
2242 reader: &mut impl Read,
2243 buffer: &mut [u8],
2244 consumed: &mut u64,
2245 payload_length: u64,
2246 deadline: Option<&OperationDeadline>,
2247) -> Result<(), SnapshotError> {
2248 for chunk in buffer.chunks_mut(COPY_BUFFER_LENGTH) {
2249 check_snapshot_deadline(deadline)?;
2250 read_payload_exact(reader, chunk, consumed, payload_length)?;
2251 }
2252 check_snapshot_deadline(deadline)?;
2253 Ok(())
2254}
2255
2256fn read_exact_or_invalid(
2257 reader: &mut impl Read,
2258 buffer: &mut [u8],
2259 reason: &'static str,
2260) -> Result<(), SnapshotError> {
2261 reader.read_exact(buffer).map_err(|source| {
2262 if source.kind() == io::ErrorKind::UnexpectedEof {
2263 SnapshotError::Invalid { reason }
2264 } else {
2265 SnapshotError::Io(source)
2266 }
2267 })
2268}
2269
2270#[cfg(unix)]
2271fn sync_directory(path: &Path) -> Result<(), SnapshotError> {
2272 File::open(path)?.sync_all()?;
2273 Ok(())
2274}
2275
2276fn copy_array<const N: usize>(source: &[u8]) -> [u8; N] {
2277 let mut output = [0_u8; N];
2278 output.copy_from_slice(source);
2279 output
2280}
2281
2282struct TemporaryFileGuard {
2283 path: PathBuf,
2284 armed: bool,
2285}
2286
2287impl TemporaryFileGuard {
2288 fn new(path: PathBuf) -> Self {
2289 Self { path, armed: true }
2290 }
2291
2292 fn disarm(&mut self) {
2293 self.armed = false;
2294 }
2295}
2296
2297impl Drop for TemporaryFileGuard {
2298 fn drop(&mut self) {
2299 if self.armed {
2300 let _ignored = std::fs::remove_file(&self.path);
2301 }
2302 }
2303}
2304
2305#[cfg(test)]
2306mod tests {
2307 use std::{
2308 error::Error,
2309 fs::{self, OpenOptions},
2310 io::{Cursor, Seek, SeekFrom, Write},
2311 path::{Path, PathBuf},
2312 time::Duration,
2313 };
2314
2315 use super::{
2316 CHECKSUM_PREFIX_LENGTH, DIGEST_PREFIX_LENGTH, DISK_FORMAT_VERSION, DecodedHeader,
2317 ENTRY_HEADER_LENGTH, HEADER_LENGTH, MAGIC, OperationDeadline, SnapshotError,
2318 SnapshotReadLimits, SnapshotRecordVisitor, V2_COUNTS_LENGTH, encode_entry_header,
2319 read_encoded_lexical_index, read_encoded_vector, read_snapshot_records_with_policy,
2320 };
2321 use crate::{CommitReceipt, test_support::TestDirectory};
2322
2323 struct TransientMutationVisitor {
2324 path: PathBuf,
2325 second_value_offset: u64,
2326 seen: Vec<(Vec<u8>, Vec<u8>)>,
2327 }
2328
2329 struct GrowingSnapshotVisitor {
2330 path: PathBuf,
2331 grew: bool,
2332 }
2333
2334 impl TransientMutationVisitor {
2335 fn overwrite_second_value_prefix(&self, byte: u8) -> Result<(), SnapshotError> {
2336 let mut file = OpenOptions::new().read(true).write(true).open(&self.path)?;
2337 file.seek(SeekFrom::Start(self.second_value_offset))?;
2338 file.write_all(&[byte])?;
2339 file.sync_all()?;
2340 Ok(())
2341 }
2342 }
2343
2344 impl SnapshotRecordVisitor for TransientMutationVisitor {
2345 fn put(&mut self, key: &[u8], value: &[u8]) -> Result<(), SnapshotError> {
2346 self.seen.push((key.to_vec(), value.to_vec()));
2347 match self.seen.len() {
2348 1 => self.overwrite_second_value_prefix(b'X')?,
2349 2 => self.overwrite_second_value_prefix(b't')?,
2350 _ => {}
2351 }
2352 Ok(())
2353 }
2354
2355 fn receipt(&mut self, _receipt: &CommitReceipt) -> Result<(), SnapshotError> {
2356 Ok(())
2357 }
2358 }
2359
2360 impl SnapshotRecordVisitor for GrowingSnapshotVisitor {
2361 fn put(&mut self, _key: &[u8], _value: &[u8]) -> Result<(), SnapshotError> {
2362 if !self.grew {
2363 let mut file = OpenOptions::new().append(true).open(&self.path)?;
2364 file.write_all(b"x")?;
2365 file.sync_all()?;
2366 self.grew = true;
2367 }
2368 Ok(())
2369 }
2370
2371 fn receipt(&mut self, _receipt: &CommitReceipt) -> Result<(), SnapshotError> {
2372 Ok(())
2373 }
2374 }
2375
2376 fn write_two_entry_snapshot(path: &Path) -> Result<(Vec<u8>, u64), SnapshotError> {
2377 let entries: [(&[u8], &[u8]); 2] = [(b"a", b"one"), (b"b", b"two")];
2378 let mut payload = vec![0_u8; V2_COUNTS_LENGTH];
2379 let mut second_value_offset = 0_u64;
2380 for (index, (key, value)) in entries.into_iter().enumerate() {
2381 let entry_header = encode_entry_header(key, value)?;
2382 payload.extend_from_slice(&entry_header);
2383 payload.extend_from_slice(key);
2384 if index == 1 {
2385 second_value_offset =
2386 u64::try_from(HEADER_LENGTH + payload.len()).map_err(|_| {
2387 SnapshotError::Invalid {
2388 reason: "test snapshot offset overflow",
2389 }
2390 })?;
2391 }
2392 payload.extend_from_slice(value);
2393 }
2394
2395 let mut header = [0_u8; HEADER_LENGTH];
2396 header[..8].copy_from_slice(&MAGIC);
2397 header[8..10].copy_from_slice(&DISK_FORMAT_VERSION.to_le_bytes());
2398 header[12..20].copy_from_slice(&1_u64.to_le_bytes());
2399 header[20..52].copy_from_slice(&[7_u8; 32]);
2400 header[52..60].copy_from_slice(&2_u64.to_le_bytes());
2401 header[68..76].copy_from_slice(
2402 &u64::try_from(payload.len())
2403 .map_err(|_| SnapshotError::Invalid {
2404 reason: "test snapshot length overflow",
2405 })?
2406 .to_le_bytes(),
2407 );
2408 let checksum =
2409 crc32c::crc32c_append(crc32c::crc32c(&header[..CHECKSUM_PREFIX_LENGTH]), &payload);
2410 header[76..80].copy_from_slice(&checksum.to_le_bytes());
2411 let mut hasher = blake3::Hasher::new();
2412 hasher.update(&header[..DIGEST_PREFIX_LENGTH]);
2413 hasher.update(&payload);
2414 header[80..112].copy_from_slice(hasher.finalize().as_bytes());
2415
2416 let mut bytes = Vec::with_capacity(HEADER_LENGTH + payload.len());
2417 bytes.extend_from_slice(&header);
2418 bytes.extend_from_slice(&payload);
2419 fs::write(path, &bytes)?;
2420 let expected_second_value_offset = u64::try_from(
2421 HEADER_LENGTH
2422 + V2_COUNTS_LENGTH
2423 + ENTRY_HEADER_LENGTH
2424 + entries[0].0.len()
2425 + entries[0].1.len()
2426 + ENTRY_HEADER_LENGTH
2427 + entries[1].0.len(),
2428 )
2429 .map_err(|_| SnapshotError::Invalid {
2430 reason: "test snapshot offset overflow",
2431 })?;
2432 debug_assert_eq!(second_value_offset, expected_second_value_offset);
2433 Ok((bytes, second_value_offset))
2434 }
2435
2436 #[test]
2437 fn visitor_pass_rejects_transient_in_place_payload_mutation() -> Result<(), Box<dyn Error>> {
2438 let temporary = TestDirectory::new("snapshot-visitor-authentication")?;
2439 let snapshot_path = temporary.path().join("transient-mutation.hysnap");
2440 let (original_bytes, second_value_offset) = write_two_entry_snapshot(&snapshot_path)?;
2441 let mut visitor = TransientMutationVisitor {
2442 path: snapshot_path.clone(),
2443 second_value_offset,
2444 seen: Vec::new(),
2445 };
2446 let deadline = OperationDeadline::new(Duration::from_secs(5));
2447
2448 let result = read_snapshot_records_with_policy(
2449 &snapshot_path,
2450 &mut visitor,
2451 &SnapshotReadLimits::default(),
2452 &deadline,
2453 );
2454
2455 assert!(matches!(
2456 result,
2457 Err(SnapshotError::Invalid {
2458 reason: "CRC32C mismatch"
2459 })
2460 ));
2461 assert_eq!(
2462 visitor.seen,
2463 vec![
2464 (b"a".to_vec(), b"one".to_vec()),
2465 (b"b".to_vec(), b"Xwo".to_vec())
2466 ]
2467 );
2468 assert_eq!(fs::read(snapshot_path)?, original_bytes);
2469 Ok(())
2470 }
2471
2472 #[test]
2473 fn visitor_pass_rejects_same_handle_file_growth() -> Result<(), Box<dyn Error>> {
2474 let temporary = TestDirectory::new("snapshot-visitor-growth")?;
2475 let snapshot_path = temporary.path().join("growing.hysnap");
2476 let (original_bytes, _) = write_two_entry_snapshot(&snapshot_path)?;
2477 let mut visitor = GrowingSnapshotVisitor {
2478 path: snapshot_path.clone(),
2479 grew: false,
2480 };
2481 let deadline = OperationDeadline::new(Duration::from_secs(5));
2482
2483 let result = read_snapshot_records_with_policy(
2484 &snapshot_path,
2485 &mut visitor,
2486 &SnapshotReadLimits::default(),
2487 &deadline,
2488 );
2489
2490 assert!(matches!(
2491 result,
2492 Err(SnapshotError::Invalid {
2493 reason: "snapshot changed while being read"
2494 })
2495 ));
2496 assert!(visitor.grew);
2497 assert_eq!(
2498 fs::metadata(snapshot_path)?.len(),
2499 u64::try_from(original_bytes.len())? + 1
2500 );
2501 Ok(())
2502 }
2503
2504 #[test]
2505 fn snapshot_reader_rejects_non_regular_paths() -> Result<(), Box<dyn Error>> {
2506 let temporary = TestDirectory::new("snapshot-regular-file")?;
2507 assert!(matches!(
2508 super::open_snapshot_file(temporary.path()),
2509 Err(SnapshotError::Invalid {
2510 reason: "snapshot is not a regular file"
2511 })
2512 ));
2513 Ok(())
2514 }
2515
2516 #[test]
2517 fn vector_and_lexical_payload_helpers_observe_the_shared_deadline() {
2518 let decoded = DecodedHeader {
2519 disk_format_version: DISK_FORMAT_VERSION,
2520 checkpoint_sequence: 1,
2521 checkpoint_digest: Some([1; 32]),
2522 entry_count: 0,
2523 receipt_count: 0,
2524 payload_length: 1,
2525 expected_checksum: 0,
2526 expected_digest: [0; 32],
2527 };
2528 let deadline = OperationDeadline::new(Duration::ZERO);
2529
2530 let mut vector_payload = Cursor::new(vec![0_u8]);
2531 let mut consumed = 0;
2532 let vector = read_encoded_vector(
2533 &mut vector_payload,
2534 &decoded,
2535 &mut consumed,
2536 Some(&deadline),
2537 );
2538 assert!(matches!(vector, Err(source) if source.is_timeout()));
2539
2540 let mut lexical_payload = Cursor::new(vec![0_u8]);
2541 let mut consumed = 0;
2542 let lexical = read_encoded_lexical_index(
2543 &mut lexical_payload,
2544 &decoded,
2545 &mut consumed,
2546 Some(&deadline),
2547 );
2548 assert!(matches!(lexical, Err(source) if source.is_timeout()));
2549 }
2550}