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