1use std::{
4 collections::BTreeMap,
5 fs::{File, OpenOptions},
6 io::{self, Read, Seek, SeekFrom, Write},
7 path::{Path, PathBuf},
8};
9
10use hyphae_core::{
11 DISK_FORMAT_VERSION, MIN_DISK_FORMAT_VERSION, Q15Vector, VectorMetric, VectorSpaceDefinition,
12 VectorSpaceName,
13};
14use hyphae_query::FieldPath;
15use hyphae_retrieval::{
16 LexicalField, LexicalIndexDefinition, MAX_LEXICAL_FIELDS, MAX_LEXICAL_PATH_SEGMENT_BYTES,
17 MAX_LEXICAL_PATH_SEGMENTS,
18};
19use thiserror::Error;
20
21use crate::{
22 CommitReceipt, MAX_KEY_BYTES, MaterializedIndexError, index::MaterializedIndex,
23 log::MAX_OPERATION_BYTES,
24};
25
26const MAGIC: [u8; 8] = *b"HYSNAP01";
27const HEADER_LENGTH: usize = 112;
28const HEADER_LENGTH_U64: u64 = 112;
29const CHECKSUM_PREFIX_LENGTH: usize = 76;
30const DIGEST_PREFIX_LENGTH: usize = 80;
31const ENTRY_HEADER_LENGTH: usize = 12;
32const ENTRY_HEADER_LENGTH_U64: u64 = 12;
33const RECEIPT_LENGTH: usize = 88;
34const RECEIPT_LENGTH_U64: u64 = 88;
35const V2_COUNTS_LENGTH: usize = 24;
36const V2_COUNTS_LENGTH_U64: u64 = 24;
37const VECTOR_SPACE_FIXED_LENGTH_U64: u64 = 5;
38const VECTOR_FIXED_LENGTH_U64: u64 = 7;
39const COPY_BUFFER_LENGTH: usize = 64 * 1024;
40const COPY_BUFFER_LENGTH_U64: u64 = 64 * 1024;
41
42#[derive(Clone, Debug, Eq, PartialEq)]
44pub struct SnapshotInfo {
45 pub path: PathBuf,
47 pub disk_format_version: u16,
49 pub checkpoint_sequence: u64,
51 pub checkpoint_digest: Option<[u8; 32]>,
53 pub entry_count: u64,
55 pub vector_space_count: u64,
57 pub vector_count: u64,
59 pub lexical_index_count: u64,
61 pub receipt_count: u64,
63 pub snapshot_digest: [u8; 32],
65 pub file_bytes: u64,
67}
68
69#[derive(Clone, Debug, Eq, PartialEq)]
72pub struct SnapshotReadLimits {
73 pub file_bytes: u64,
75 pub entries: u64,
77 pub decoded_bytes: u64,
79}
80
81impl Default for SnapshotReadLimits {
82 fn default() -> Self {
83 Self {
84 file_bytes: 512 * 1024 * 1024,
85 entries: 1_000_000,
86 decoded_bytes: 256 * 1024 * 1024,
87 }
88 }
89}
90
91#[derive(Clone, Debug, Eq, PartialEq)]
93pub struct SnapshotEntry {
94 pub key: Vec<u8>,
96 pub value: Vec<u8>,
98}
99
100#[derive(Clone, Debug, Eq, PartialEq)]
102pub struct SnapshotContents {
103 pub info: SnapshotInfo,
105 pub entries: Vec<SnapshotEntry>,
107 pub vector_spaces: Vec<VectorSpaceDefinition>,
109 pub vectors: Vec<SnapshotVectorEntry>,
111 pub lexical_indexes: Vec<LexicalIndexDefinition>,
113}
114
115#[derive(Clone, Debug, Eq, PartialEq)]
117pub struct SnapshotVectorEntry {
118 pub space: VectorSpaceName,
120 pub key: Vec<u8>,
122 pub vector: Q15Vector,
124}
125
126#[derive(Debug, Error)]
128pub enum SnapshotError {
129 #[error(transparent)]
131 Io(#[from] io::Error),
132
133 #[error("materialized index failure during snapshot: {source}")]
135 Index {
136 #[source]
138 source: Box<MaterializedIndexError>,
139 },
140
141 #[error("invalid snapshot: {reason}")]
143 Invalid {
144 reason: &'static str,
146 },
147
148 #[error("unsupported snapshot format {found}; supported format is {supported}")]
150 UnsupportedVersion {
151 found: u16,
153 supported: u16,
155 },
156
157 #[error("snapshot sequence {sequence} already exists for a different commit")]
159 CheckpointConflict {
160 sequence: u64,
162 },
163
164 #[error("snapshot file length {actual} exceeds verification limit {maximum}")]
166 FileLimitExceeded {
167 actual: u64,
169 maximum: u64,
171 },
172
173 #[error("snapshot entry count {actual} exceeds verification limit {maximum}")]
175 EntryLimitExceeded {
176 actual: u64,
178 maximum: u64,
180 },
181
182 #[error("snapshot decoded bytes exceed verification limit {maximum}")]
184 DecodedBytesLimitExceeded {
185 maximum: u64,
187 },
188}
189
190impl From<MaterializedIndexError> for SnapshotError {
191 fn from(source: MaterializedIndexError) -> Self {
192 Self::Index {
193 source: Box::new(source),
194 }
195 }
196}
197
198#[allow(clippy::too_many_lines)]
199pub(crate) fn create_snapshot(
200 index: &MaterializedIndex,
201 snapshots_directory: &Path,
202 temporary_directory: &Path,
203 disk_format_version: u16,
204) -> Result<SnapshotInfo, SnapshotError> {
205 let checkpoint = index.checkpoint()?;
206 if checkpoint.sequence == 0 && checkpoint.digest.is_some() {
207 return Err(SnapshotError::Invalid {
208 reason: "empty checkpoint has a digest",
209 });
210 }
211
212 let measurements = measure_payload(index, checkpoint.sequence, disk_format_version)?;
213 let final_path =
214 snapshots_directory.join(format!("snapshot-{:020}.hysnap", checkpoint.sequence));
215 if final_path.exists() {
216 let existing = verify_snapshot(&final_path)?;
217 if existing.checkpoint_digest != checkpoint.digest {
218 return Err(SnapshotError::CheckpointConflict {
219 sequence: checkpoint.sequence,
220 });
221 }
222 return Ok(existing);
223 }
224
225 let mut header = [0_u8; HEADER_LENGTH];
226 header[0..8].copy_from_slice(&MAGIC);
227 header[8..10].copy_from_slice(&disk_format_version.to_le_bytes());
228 header[10..12].copy_from_slice(&0_u16.to_le_bytes());
229 header[12..20].copy_from_slice(&checkpoint.sequence.to_le_bytes());
230 header[20..52].copy_from_slice(&checkpoint.digest.unwrap_or([0; 32]));
231 header[52..60].copy_from_slice(&measurements.entry_count.to_le_bytes());
232 header[60..68].copy_from_slice(&measurements.receipt_count.to_le_bytes());
233 header[68..76].copy_from_slice(&measurements.payload_length.to_le_bytes());
234
235 let mut checksum = crc32c::crc32c(&header[..CHECKSUM_PREFIX_LENGTH]);
236 if disk_format_version >= 2 {
237 checksum = crc32c::crc32c_append(checksum, &measurements.v2_counts());
238 }
239 let mut checksum_error = None;
240 index.for_each_entry(|key, value| {
241 if checksum_error.is_some() {
242 return;
243 }
244 match encode_entry_header(key, value) {
245 Ok(entry_header) => {
246 checksum = crc32c::crc32c_append(checksum, &entry_header);
247 checksum = crc32c::crc32c_append(checksum, key);
248 checksum = crc32c::crc32c_append(checksum, value);
249 }
250 Err(source) => checksum_error = Some(source),
251 }
252 })?;
253 if let Some(source) = checksum_error {
254 return Err(source);
255 }
256 let mut vector_checksum_error = None;
257 if disk_format_version >= 2 {
258 index.for_each_vector_space(|definition| {
259 if vector_checksum_error.is_none() {
260 match encode_vector_space(definition) {
261 Ok(encoded) => checksum = crc32c::crc32c_append(checksum, &encoded),
262 Err(source) => vector_checksum_error = Some(source),
263 }
264 }
265 })?;
266 index.for_each_vector(|space, key, vector| {
267 if vector_checksum_error.is_none() {
268 match encode_vector(space, key, vector) {
269 Ok(encoded) => checksum = crc32c::crc32c_append(checksum, &encoded),
270 Err(source) => vector_checksum_error = Some(source),
271 }
272 }
273 })?;
274 index.for_each_lexical_index(|definition| {
275 if vector_checksum_error.is_none() {
276 match encode_lexical_index(definition) {
277 Ok(encoded) => checksum = crc32c::crc32c_append(checksum, &encoded),
278 Err(source) => vector_checksum_error = Some(source),
279 }
280 }
281 })?;
282 }
283 if let Some(source) = vector_checksum_error {
284 return Err(source);
285 }
286 index.for_each_receipt(|receipt| {
287 checksum = crc32c::crc32c_append(checksum, &encode_receipt(receipt));
288 })?;
289 header[76..80].copy_from_slice(&checksum.to_le_bytes());
290
291 let temporary_path = temporary_directory.join(format!(
292 "snapshot-{:020}-{}.tmp",
293 checkpoint.sequence,
294 uuid::Uuid::now_v7()
295 ));
296 let mut file = OpenOptions::new()
297 .create_new(true)
298 .read(true)
299 .write(true)
300 .open(&temporary_path)?;
301 file.write_all(&header)?;
302 let mut hasher = blake3::Hasher::new();
303 hasher.update(&header[..DIGEST_PREFIX_LENGTH]);
304 if disk_format_version >= 2 {
305 let counts = measurements.v2_counts();
306 file.write_all(&counts)?;
307 hasher.update(&counts);
308 }
309 let mut write_error = None;
310 index.for_each_entry(|key, value| {
311 if write_error.is_none()
312 && let Err(source) = write_entry(&mut file, &mut hasher, key, value)
313 {
314 write_error = Some(source);
315 }
316 })?;
317 if let Some(source) = write_error {
318 return Err(source);
319 }
320 let mut vector_write_error = None;
321 if disk_format_version >= 2 {
322 index.for_each_vector_space(|definition| {
323 if vector_write_error.is_none()
324 && let Err(source) =
325 write_encoded(&mut file, &mut hasher, encode_vector_space(definition))
326 {
327 vector_write_error = Some(source);
328 }
329 })?;
330 index.for_each_vector(|space, key, vector| {
331 if vector_write_error.is_none()
332 && let Err(source) =
333 write_encoded(&mut file, &mut hasher, encode_vector(space, key, vector))
334 {
335 vector_write_error = Some(source);
336 }
337 })?;
338 index.for_each_lexical_index(|definition| {
339 if vector_write_error.is_none()
340 && let Err(source) =
341 write_encoded(&mut file, &mut hasher, encode_lexical_index(definition))
342 {
343 vector_write_error = Some(source);
344 }
345 })?;
346 }
347 if let Some(source) = vector_write_error {
348 return Err(source);
349 }
350 let mut receipt_write_error = None;
351 index.for_each_receipt(|receipt| {
352 if receipt_write_error.is_none()
353 && let Err(source) = write_receipt(&mut file, &mut hasher, receipt)
354 {
355 receipt_write_error = Some(source);
356 }
357 })?;
358 if let Some(source) = receipt_write_error {
359 return Err(source);
360 }
361 let snapshot_digest = *hasher.finalize().as_bytes();
362 file.seek(SeekFrom::Start(80))?;
363 file.write_all(&snapshot_digest)?;
364 file.sync_all()?;
365 drop(file);
366
367 let temporary_info = verify_snapshot(&temporary_path)?;
368 std::fs::rename(&temporary_path, &final_path)?;
369 #[cfg(unix)]
370 sync_directory(snapshots_directory)?;
371 Ok(SnapshotInfo {
372 path: final_path,
373 ..temporary_info
374 })
375}
376
377pub fn verify_snapshot(path: impl AsRef<Path>) -> Result<SnapshotInfo, SnapshotError> {
384 let path = path.as_ref();
385 let mut file = File::open(path)?;
386 let file_bytes = file.metadata()?.len();
387 let mut header = [0_u8; HEADER_LENGTH];
388 read_exact_or_invalid(&mut file, &mut header, "truncated header")?;
389 let decoded = decode_header(&header, file_bytes)?;
390 let (vector_space_count, vector_count, lexical_index_count) =
391 verify_payload(&mut file, &header, &decoded)?;
392
393 Ok(SnapshotInfo {
394 path: path.to_path_buf(),
395 disk_format_version: decoded.disk_format_version,
396 checkpoint_sequence: decoded.checkpoint_sequence,
397 checkpoint_digest: decoded.checkpoint_digest,
398 entry_count: decoded.entry_count,
399 vector_space_count,
400 vector_count,
401 lexical_index_count,
402 receipt_count: decoded.receipt_count,
403 snapshot_digest: decoded.expected_digest,
404 file_bytes,
405 })
406}
407
408pub fn load_snapshot(
420 path: impl AsRef<Path>,
421 limits: &SnapshotReadLimits,
422) -> Result<SnapshotContents, SnapshotError> {
423 let path = path.as_ref();
424 let mut collector = SnapshotCollector {
425 entries: Vec::new(),
426 vector_spaces: Vec::new(),
427 vectors: Vec::new(),
428 lexical_indexes: Vec::new(),
429 decoded_bytes: 0,
430 limits,
431 };
432 let info = read_snapshot_records_with_limits(path, &mut collector, Some(limits))?;
433 Ok(SnapshotContents {
434 info,
435 entries: collector.entries,
436 vector_spaces: collector.vector_spaces,
437 vectors: collector.vectors,
438 lexical_indexes: collector.lexical_indexes,
439 })
440}
441
442pub(crate) trait SnapshotRecordVisitor {
443 fn put(&mut self, key: &[u8], value: &[u8]) -> Result<(), SnapshotError>;
444 fn vector_space(&mut self, _definition: &VectorSpaceDefinition) -> Result<(), SnapshotError> {
445 Ok(())
446 }
447 fn vector(
448 &mut self,
449 _space: &VectorSpaceName,
450 _key: &[u8],
451 _vector: &Q15Vector,
452 ) -> Result<(), SnapshotError> {
453 Ok(())
454 }
455 fn lexical_index(&mut self, _definition: &LexicalIndexDefinition) -> Result<(), SnapshotError> {
456 Ok(())
457 }
458 fn receipt(&mut self, receipt: &CommitReceipt) -> Result<(), SnapshotError>;
459}
460
461pub(crate) fn read_snapshot_records(
462 path: &Path,
463 visitor: &mut impl SnapshotRecordVisitor,
464) -> Result<SnapshotInfo, SnapshotError> {
465 read_snapshot_records_with_limits(path, visitor, None)
466}
467
468fn read_snapshot_records_with_limits(
469 path: &Path,
470 visitor: &mut impl SnapshotRecordVisitor,
471 limits: Option<&SnapshotReadLimits>,
472) -> Result<SnapshotInfo, SnapshotError> {
473 let before = verify_snapshot(path)?;
474 if let Some(limits) = limits {
475 validate_read_limits(&before, limits)?;
476 }
477 let mut file = File::open(path)?;
478 let file_bytes = file.metadata()?.len();
479 let mut header = [0_u8; HEADER_LENGTH];
480 read_exact_or_invalid(&mut file, &mut header, "truncated header")?;
481 let decoded = decode_header(&header, file_bytes)?;
482 let mut consumed = 0_u64;
483 let (vector_space_count, vector_count, lexical_index_count) =
484 read_v2_counts(&mut file, &decoded, &mut consumed)?;
485
486 for _ in 0..decoded.entry_count {
487 let mut entry_header = [0_u8; ENTRY_HEADER_LENGTH];
488 read_payload_exact(
489 &mut file,
490 &mut entry_header,
491 &mut consumed,
492 decoded.payload_length,
493 )?;
494 let key_length = usize::try_from(u32::from_le_bytes(copy_array(&entry_header[..4])))
495 .map_err(|_| SnapshotError::Invalid {
496 reason: "key length overflow during restore",
497 })?;
498 let value_length = usize::try_from(u64::from_le_bytes(copy_array(&entry_header[4..12])))
499 .map_err(|_| SnapshotError::Invalid {
500 reason: "value length overflow during restore",
501 })?;
502 if key_length == 0 || key_length > MAX_KEY_BYTES || value_length > MAX_OPERATION_BYTES {
503 return Err(SnapshotError::Invalid {
504 reason: "record exceeds restore bounds",
505 });
506 }
507 let mut key = vec![0_u8; key_length];
508 let mut value = vec![0_u8; value_length];
509 read_payload_exact(&mut file, &mut key, &mut consumed, decoded.payload_length)?;
510 read_payload_exact(&mut file, &mut value, &mut consumed, decoded.payload_length)?;
511 visitor.put(&key, &value)?;
512 }
513 for _ in 0..vector_space_count {
514 let definition = read_vector_space(&mut file, &decoded, &mut consumed)?;
515 visitor.vector_space(&definition)?;
516 }
517 for _ in 0..vector_count {
518 let (space, key, vector) = read_vector(&mut file, &decoded, &mut consumed)?;
519 visitor.vector(&space, &key, &vector)?;
520 }
521 for _ in 0..lexical_index_count {
522 let definition = read_lexical_index(&mut file, &decoded, &mut consumed)?;
523 visitor.lexical_index(&definition)?;
524 }
525 for _ in 0..decoded.receipt_count {
526 let mut encoded = [0_u8; RECEIPT_LENGTH];
527 read_payload_exact(
528 &mut file,
529 &mut encoded,
530 &mut consumed,
531 decoded.payload_length,
532 )?;
533 visitor.receipt(&decode_snapshot_receipt(&encoded))?;
534 }
535 if consumed != decoded.payload_length {
536 return Err(SnapshotError::Invalid {
537 reason: "record counts do not consume payload during restore",
538 });
539 }
540
541 let after = verify_snapshot(path)?;
542 if before != after {
543 return Err(SnapshotError::Invalid {
544 reason: "snapshot changed during restore",
545 });
546 }
547 Ok(after)
548}
549
550fn validate_read_limits(
551 info: &SnapshotInfo,
552 limits: &SnapshotReadLimits,
553) -> Result<(), SnapshotError> {
554 if info.file_bytes > limits.file_bytes {
555 return Err(SnapshotError::FileLimitExceeded {
556 actual: info.file_bytes,
557 maximum: limits.file_bytes,
558 });
559 }
560 let logical_records = info
561 .entry_count
562 .checked_add(info.vector_space_count)
563 .and_then(|count| count.checked_add(info.vector_count))
564 .and_then(|count| count.checked_add(info.lexical_index_count))
565 .ok_or(SnapshotError::EntryLimitExceeded {
566 actual: u64::MAX,
567 maximum: limits.entries,
568 })?;
569 if logical_records > limits.entries {
570 return Err(SnapshotError::EntryLimitExceeded {
571 actual: logical_records,
572 maximum: limits.entries,
573 });
574 }
575 Ok(())
576}
577
578struct SnapshotCollector<'limits> {
579 entries: Vec<SnapshotEntry>,
580 vector_spaces: Vec<VectorSpaceDefinition>,
581 vectors: Vec<SnapshotVectorEntry>,
582 lexical_indexes: Vec<LexicalIndexDefinition>,
583 decoded_bytes: u64,
584 limits: &'limits SnapshotReadLimits,
585}
586
587impl SnapshotRecordVisitor for SnapshotCollector<'_> {
588 fn put(&mut self, key: &[u8], value: &[u8]) -> Result<(), SnapshotError> {
589 let next_entry_count = u64::try_from(self.entries.len())
590 .ok()
591 .and_then(|count| count.checked_add(1))
592 .ok_or(SnapshotError::EntryLimitExceeded {
593 actual: u64::MAX,
594 maximum: self.limits.entries,
595 })?;
596 if next_entry_count > self.limits.entries {
597 return Err(SnapshotError::EntryLimitExceeded {
598 actual: next_entry_count,
599 maximum: self.limits.entries,
600 });
601 }
602 let entry_bytes = u64::try_from(key.len())
603 .ok()
604 .and_then(|key_bytes| {
605 u64::try_from(value.len())
606 .ok()
607 .and_then(|value_bytes| key_bytes.checked_add(value_bytes))
608 })
609 .ok_or(SnapshotError::DecodedBytesLimitExceeded {
610 maximum: self.limits.decoded_bytes,
611 })?;
612 self.decoded_bytes = self.decoded_bytes.checked_add(entry_bytes).ok_or(
613 SnapshotError::DecodedBytesLimitExceeded {
614 maximum: self.limits.decoded_bytes,
615 },
616 )?;
617 if self.decoded_bytes > self.limits.decoded_bytes {
618 return Err(SnapshotError::DecodedBytesLimitExceeded {
619 maximum: self.limits.decoded_bytes,
620 });
621 }
622 self.entries.push(SnapshotEntry {
623 key: key.to_vec(),
624 value: value.to_vec(),
625 });
626 Ok(())
627 }
628
629 fn receipt(&mut self, _receipt: &CommitReceipt) -> Result<(), SnapshotError> {
630 Ok(())
631 }
632
633 fn vector_space(&mut self, definition: &VectorSpaceDefinition) -> Result<(), SnapshotError> {
634 self.add_decoded_bytes(definition.name.as_str().len())?;
635 self.vector_spaces.push(definition.clone());
636 Ok(())
637 }
638
639 fn vector(
640 &mut self,
641 space: &VectorSpaceName,
642 key: &[u8],
643 vector: &Q15Vector,
644 ) -> Result<(), SnapshotError> {
645 let vector_bytes = vector
646 .as_slice()
647 .len()
648 .checked_mul(2)
649 .and_then(|length| length.checked_add(space.as_str().len()))
650 .and_then(|length| length.checked_add(key.len()))
651 .ok_or(SnapshotError::DecodedBytesLimitExceeded {
652 maximum: self.limits.decoded_bytes,
653 })?;
654 self.add_decoded_bytes(vector_bytes)?;
655 self.vectors.push(SnapshotVectorEntry {
656 space: space.clone(),
657 key: key.to_vec(),
658 vector: vector.clone(),
659 });
660 Ok(())
661 }
662
663 fn lexical_index(&mut self, definition: &LexicalIndexDefinition) -> Result<(), SnapshotError> {
664 let encoded_length = encode_lexical_index(definition)?.len();
665 self.add_decoded_bytes(encoded_length)?;
666 self.lexical_indexes.push(definition.clone());
667 Ok(())
668 }
669}
670
671impl SnapshotCollector<'_> {
672 fn add_decoded_bytes(&mut self, bytes: usize) -> Result<(), SnapshotError> {
673 let bytes = u64::try_from(bytes).map_err(|_| SnapshotError::DecodedBytesLimitExceeded {
674 maximum: self.limits.decoded_bytes,
675 })?;
676 self.decoded_bytes = self.decoded_bytes.checked_add(bytes).ok_or(
677 SnapshotError::DecodedBytesLimitExceeded {
678 maximum: self.limits.decoded_bytes,
679 },
680 )?;
681 if self.decoded_bytes > self.limits.decoded_bytes {
682 return Err(SnapshotError::DecodedBytesLimitExceeded {
683 maximum: self.limits.decoded_bytes,
684 });
685 }
686 Ok(())
687 }
688}
689
690#[derive(Clone, Copy, Debug)]
691struct DecodedHeader {
692 disk_format_version: u16,
693 checkpoint_sequence: u64,
694 checkpoint_digest: Option<[u8; 32]>,
695 entry_count: u64,
696 receipt_count: u64,
697 payload_length: u64,
698 expected_checksum: u32,
699 expected_digest: [u8; 32],
700}
701
702fn decode_header(
703 header: &[u8; HEADER_LENGTH],
704 file_bytes: u64,
705) -> Result<DecodedHeader, SnapshotError> {
706 if header[0..8] != MAGIC {
707 return Err(SnapshotError::Invalid {
708 reason: "bad magic",
709 });
710 }
711 let version = u16::from_le_bytes(copy_array(&header[8..10]));
712 if !(MIN_DISK_FORMAT_VERSION..=DISK_FORMAT_VERSION).contains(&version) {
713 return Err(SnapshotError::UnsupportedVersion {
714 found: version,
715 supported: DISK_FORMAT_VERSION,
716 });
717 }
718 if u16::from_le_bytes(copy_array(&header[10..12])) != 0 {
719 return Err(SnapshotError::Invalid {
720 reason: "unsupported flags",
721 });
722 }
723
724 let checkpoint_sequence = u64::from_le_bytes(copy_array(&header[12..20]));
725 let raw_checkpoint_digest: [u8; 32] = copy_array(&header[20..52]);
726 let checkpoint_digest = if checkpoint_sequence == 0 {
727 if raw_checkpoint_digest != [0; 32] {
728 return Err(SnapshotError::Invalid {
729 reason: "empty checkpoint has a digest",
730 });
731 }
732 None
733 } else {
734 Some(raw_checkpoint_digest)
735 };
736 let entry_count = u64::from_le_bytes(copy_array(&header[52..60]));
737 let receipt_count = u64::from_le_bytes(copy_array(&header[60..68]));
738 if checkpoint_sequence == 0 && receipt_count != 0 {
739 return Err(SnapshotError::Invalid {
740 reason: "empty checkpoint has idempotency receipts",
741 });
742 }
743 let payload_length = u64::from_le_bytes(copy_array(&header[68..76]));
744 let expected_file_bytes =
745 HEADER_LENGTH_U64
746 .checked_add(payload_length)
747 .ok_or(SnapshotError::Invalid {
748 reason: "file length overflow",
749 })?;
750 if file_bytes != expected_file_bytes {
751 return Err(SnapshotError::Invalid {
752 reason: "file length mismatch",
753 });
754 }
755
756 Ok(DecodedHeader {
757 disk_format_version: version,
758 checkpoint_sequence,
759 checkpoint_digest,
760 entry_count,
761 receipt_count,
762 payload_length,
763 expected_checksum: u32::from_le_bytes(copy_array(&header[76..80])),
764 expected_digest: copy_array(&header[80..112]),
765 })
766}
767
768#[allow(clippy::too_many_lines)]
769fn verify_payload(
770 file: &mut File,
771 header: &[u8; HEADER_LENGTH],
772 decoded: &DecodedHeader,
773) -> Result<(u64, u64, u64), SnapshotError> {
774 let mut checksum = crc32c::crc32c(&header[..CHECKSUM_PREFIX_LENGTH]);
775 let mut hasher = blake3::Hasher::new();
776 hasher.update(&header[..DIGEST_PREFIX_LENGTH]);
777 let mut consumed = 0_u64;
778 let mut counts_bytes = [0_u8; V2_COUNTS_LENGTH];
779 let (vector_space_count, vector_count, lexical_index_count) =
780 if decoded.disk_format_version >= 2 {
781 read_payload_exact(
782 file,
783 &mut counts_bytes,
784 &mut consumed,
785 decoded.payload_length,
786 )?;
787 checksum = crc32c::crc32c_append(checksum, &counts_bytes);
788 hasher.update(&counts_bytes);
789 (
790 u64::from_le_bytes(copy_array(&counts_bytes[..8])),
791 u64::from_le_bytes(copy_array(&counts_bytes[8..16])),
792 u64::from_le_bytes(copy_array(&counts_bytes[16..24])),
793 )
794 } else {
795 (0, 0, 0)
796 };
797 let mut previous_key: Option<Vec<u8>> = None;
798 let mut buffer = vec![0_u8; COPY_BUFFER_LENGTH].into_boxed_slice();
799 for _ in 0..decoded.entry_count {
800 let mut entry_header = [0_u8; ENTRY_HEADER_LENGTH];
801 read_payload_exact(
802 file,
803 &mut entry_header,
804 &mut consumed,
805 decoded.payload_length,
806 )?;
807 checksum = crc32c::crc32c_append(checksum, &entry_header);
808 hasher.update(&entry_header);
809 let key_length = usize::try_from(u32::from_le_bytes(copy_array(&entry_header[..4])))
810 .map_err(|_| SnapshotError::Invalid {
811 reason: "key length overflow",
812 })?;
813 let value_length = u64::from_le_bytes(copy_array(&entry_header[4..12]));
814 if key_length == 0 || key_length > MAX_KEY_BYTES {
815 return Err(SnapshotError::Invalid {
816 reason: "invalid key length",
817 });
818 }
819
820 let mut key = vec![0_u8; key_length];
821 read_payload_exact(file, &mut key, &mut consumed, decoded.payload_length)?;
822 checksum = crc32c::crc32c_append(checksum, &key);
823 hasher.update(&key);
824 if previous_key
825 .as_ref()
826 .is_some_and(|previous| previous >= &key)
827 {
828 return Err(SnapshotError::Invalid {
829 reason: "keys are not strictly sorted",
830 });
831 }
832 previous_key = Some(key);
833
834 let mut remaining = value_length;
835 while remaining > 0 {
836 let chunk_length =
837 usize::try_from(remaining.min(COPY_BUFFER_LENGTH_U64)).map_err(|_| {
838 SnapshotError::Invalid {
839 reason: "value length overflow",
840 }
841 })?;
842 let chunk = &mut buffer[..chunk_length];
843 read_payload_exact(file, chunk, &mut consumed, decoded.payload_length)?;
844 checksum = crc32c::crc32c_append(checksum, chunk);
845 hasher.update(chunk);
846 remaining -= u64::try_from(chunk_length).map_err(|_| SnapshotError::Invalid {
847 reason: "value length overflow",
848 })?;
849 }
850 }
851 let mut definitions = BTreeMap::new();
852 let mut previous_space: Option<VectorSpaceName> = None;
853 for _ in 0..vector_space_count {
854 let encoded = read_encoded_vector_space(file, decoded, &mut consumed)?;
855 checksum = crc32c::crc32c_append(checksum, &encoded);
856 hasher.update(&encoded);
857 let definition = decode_vector_space(&encoded)?;
858 if previous_space
859 .as_ref()
860 .is_some_and(|previous| previous >= &definition.name)
861 {
862 return Err(SnapshotError::Invalid {
863 reason: "vector spaces are not strictly sorted",
864 });
865 }
866 previous_space = Some(definition.name.clone());
867 definitions.insert(definition.name.clone(), definition);
868 }
869 let mut previous_vector_identity: Option<(VectorSpaceName, Vec<u8>)> = None;
870 for _ in 0..vector_count {
871 let encoded = read_encoded_vector(file, decoded, &mut consumed)?;
872 checksum = crc32c::crc32c_append(checksum, &encoded);
873 hasher.update(&encoded);
874 let (space, key, vector) = decode_vector(&encoded)?;
875 if previous_vector_identity
876 .as_ref()
877 .is_some_and(|previous| previous >= &(space.clone(), key.clone()))
878 {
879 return Err(SnapshotError::Invalid {
880 reason: "vectors are not strictly sorted",
881 });
882 }
883 let definition = definitions.get(&space).ok_or(SnapshotError::Invalid {
884 reason: "vector references an undefined space",
885 })?;
886 definition
887 .validate_vector(&vector)
888 .map_err(|_| SnapshotError::Invalid {
889 reason: "vector dimension does not match its space",
890 })?;
891 previous_vector_identity = Some((space, key));
892 }
893 let mut previous_lexical_name: Option<VectorSpaceName> = None;
894 for _ in 0..lexical_index_count {
895 let encoded = read_encoded_lexical_index(file, decoded, &mut consumed)?;
896 checksum = crc32c::crc32c_append(checksum, &encoded);
897 hasher.update(&encoded);
898 let definition = decode_lexical_index(&encoded)?;
899 if previous_lexical_name
900 .as_ref()
901 .is_some_and(|previous| previous >= &definition.name)
902 {
903 return Err(SnapshotError::Invalid {
904 reason: "lexical indexes are not strictly sorted",
905 });
906 }
907 previous_lexical_name = Some(definition.name);
908 }
909 let mut previous_transaction_id = None;
910 for _ in 0..decoded.receipt_count {
911 let mut encoded = [0_u8; RECEIPT_LENGTH];
912 read_payload_exact(file, &mut encoded, &mut consumed, decoded.payload_length)?;
913 checksum = crc32c::crc32c_append(checksum, &encoded);
914 hasher.update(&encoded);
915
916 let transaction_id: [u8; 16] = copy_array(&encoded[..16]);
917 if previous_transaction_id
918 .as_ref()
919 .is_some_and(|previous| previous >= &transaction_id)
920 {
921 return Err(SnapshotError::Invalid {
922 reason: "transaction identifiers are not strictly sorted",
923 });
924 }
925 previous_transaction_id = Some(transaction_id);
926 let commit_sequence = u64::from_le_bytes(copy_array(&encoded[16..24]));
927 if commit_sequence == 0 || commit_sequence > decoded.checkpoint_sequence {
928 return Err(SnapshotError::Invalid {
929 reason: "idempotency receipt exceeds snapshot checkpoint",
930 });
931 }
932 }
933 if consumed != decoded.payload_length {
934 return Err(SnapshotError::Invalid {
935 reason: "record counts do not consume payload",
936 });
937 }
938 if checksum != decoded.expected_checksum {
939 return Err(SnapshotError::Invalid {
940 reason: "CRC32C mismatch",
941 });
942 }
943 let actual_digest = *hasher.finalize().as_bytes();
944 if actual_digest != decoded.expected_digest {
945 return Err(SnapshotError::Invalid {
946 reason: "BLAKE3 digest mismatch",
947 });
948 }
949 Ok((vector_space_count, vector_count, lexical_index_count))
950}
951
952#[derive(Clone, Copy, Debug)]
953struct Measurements {
954 entry_count: u64,
955 vector_space_count: u64,
956 vector_count: u64,
957 lexical_index_count: u64,
958 receipt_count: u64,
959 payload_length: u64,
960}
961
962impl Measurements {
963 fn v2_counts(self) -> [u8; V2_COUNTS_LENGTH] {
964 let mut encoded = [0_u8; V2_COUNTS_LENGTH];
965 encoded[..8].copy_from_slice(&self.vector_space_count.to_le_bytes());
966 encoded[8..16].copy_from_slice(&self.vector_count.to_le_bytes());
967 encoded[16..24].copy_from_slice(&self.lexical_index_count.to_le_bytes());
968 encoded
969 }
970}
971
972#[allow(clippy::too_many_lines)]
973fn measure_payload(
974 index: &MaterializedIndex,
975 checkpoint_sequence: u64,
976 disk_format_version: u16,
977) -> Result<Measurements, SnapshotError> {
978 if !(MIN_DISK_FORMAT_VERSION..=DISK_FORMAT_VERSION).contains(&disk_format_version) {
979 return Err(SnapshotError::UnsupportedVersion {
980 found: disk_format_version,
981 supported: DISK_FORMAT_VERSION,
982 });
983 }
984 let mut entry_count = Some(0_u64);
985 let mut payload_length = Some(0_u64);
986 let mut valid = true;
987 index.for_each_entry(|key, value| {
988 if key.is_empty() || key.len() > MAX_KEY_BYTES {
989 valid = false;
990 return;
991 }
992 let Ok(key_length) = u64::try_from(key.len()) else {
993 valid = false;
994 return;
995 };
996 let Ok(value_length) = u64::try_from(value.len()) else {
997 valid = false;
998 return;
999 };
1000 entry_count = entry_count.and_then(|count| count.checked_add(1));
1001 payload_length = payload_length.and_then(|length| {
1002 length
1003 .checked_add(ENTRY_HEADER_LENGTH_U64)
1004 .and_then(|length| length.checked_add(key_length))
1005 .and_then(|length| length.checked_add(value_length))
1006 });
1007 })?;
1008 let mut vector_space_count = Some(0_u64);
1009 let mut vector_count = Some(0_u64);
1010 let mut lexical_index_count = Some(0_u64);
1011 if disk_format_version >= 2 {
1012 payload_length = payload_length.and_then(|length| length.checked_add(V2_COUNTS_LENGTH_U64));
1013 index.for_each_vector_space(|definition| {
1014 let Ok(name_length) = u64::try_from(definition.name.as_str().len()) else {
1015 valid = false;
1016 return;
1017 };
1018 vector_space_count = vector_space_count.and_then(|count| count.checked_add(1));
1019 payload_length = payload_length.and_then(|length| {
1020 length
1021 .checked_add(VECTOR_SPACE_FIXED_LENGTH_U64)
1022 .and_then(|length| length.checked_add(name_length))
1023 });
1024 })?;
1025 index.for_each_vector(|space, key, vector| {
1026 let Ok(name_length) = u64::try_from(space.as_str().len()) else {
1027 valid = false;
1028 return;
1029 };
1030 let Ok(key_length) = u64::try_from(key.len()) else {
1031 valid = false;
1032 return;
1033 };
1034 let Ok(vector_bytes) = u64::try_from(vector.as_slice().len().saturating_mul(2)) else {
1035 valid = false;
1036 return;
1037 };
1038 vector_count = vector_count.and_then(|count| count.checked_add(1));
1039 payload_length = payload_length.and_then(|length| {
1040 length
1041 .checked_add(VECTOR_FIXED_LENGTH_U64)
1042 .and_then(|length| length.checked_add(name_length))
1043 .and_then(|length| length.checked_add(key_length))
1044 .and_then(|length| length.checked_add(vector_bytes))
1045 });
1046 })?;
1047 index.for_each_lexical_index(|definition| {
1048 let Ok(encoded) = encode_lexical_index(definition) else {
1049 valid = false;
1050 return;
1051 };
1052 let Ok(record_bytes) = u64::try_from(encoded.len()) else {
1053 valid = false;
1054 return;
1055 };
1056 lexical_index_count = lexical_index_count.and_then(|count| count.checked_add(1));
1057 payload_length = payload_length.and_then(|length| length.checked_add(record_bytes));
1058 })?;
1059 }
1060 let mut receipt_count = Some(0_u64);
1061 index.for_each_receipt(|receipt| {
1062 if receipt.commit_sequence == 0 || receipt.commit_sequence > checkpoint_sequence {
1063 valid = false;
1064 return;
1065 }
1066 receipt_count = receipt_count.and_then(|count| count.checked_add(1));
1067 payload_length = payload_length.and_then(|length| length.checked_add(RECEIPT_LENGTH_U64));
1068 })?;
1069 if !valid {
1070 return Err(SnapshotError::Invalid {
1071 reason: "index contains an invalid key or idempotency receipt",
1072 });
1073 }
1074 let Some(entry_count) = entry_count else {
1075 return Err(SnapshotError::Invalid {
1076 reason: "entry count overflow",
1077 });
1078 };
1079 let Some(payload_length) = payload_length else {
1080 return Err(SnapshotError::Invalid {
1081 reason: "payload length overflow",
1082 });
1083 };
1084 let Some(receipt_count) = receipt_count else {
1085 return Err(SnapshotError::Invalid {
1086 reason: "receipt count overflow",
1087 });
1088 };
1089 let Some(vector_space_count) = vector_space_count else {
1090 return Err(SnapshotError::Invalid {
1091 reason: "vector-space count overflow",
1092 });
1093 };
1094 let Some(vector_count) = vector_count else {
1095 return Err(SnapshotError::Invalid {
1096 reason: "vector count overflow",
1097 });
1098 };
1099 let Some(lexical_index_count) = lexical_index_count else {
1100 return Err(SnapshotError::Invalid {
1101 reason: "lexical-index count overflow",
1102 });
1103 };
1104 Ok(Measurements {
1105 entry_count,
1106 vector_space_count,
1107 vector_count,
1108 lexical_index_count,
1109 receipt_count,
1110 payload_length,
1111 })
1112}
1113
1114fn encode_entry_header(
1115 key: &[u8],
1116 value: &[u8],
1117) -> Result<[u8; ENTRY_HEADER_LENGTH], SnapshotError> {
1118 let key_length = u32::try_from(key.len()).map_err(|_| SnapshotError::Invalid {
1119 reason: "key length overflow",
1120 })?;
1121 let value_length = u64::try_from(value.len()).map_err(|_| SnapshotError::Invalid {
1122 reason: "value length overflow",
1123 })?;
1124 let mut entry_header = [0_u8; ENTRY_HEADER_LENGTH];
1125 entry_header[..4].copy_from_slice(&key_length.to_le_bytes());
1126 entry_header[4..].copy_from_slice(&value_length.to_le_bytes());
1127 Ok(entry_header)
1128}
1129
1130fn write_entry(
1131 writer: &mut impl Write,
1132 hasher: &mut blake3::Hasher,
1133 key: &[u8],
1134 value: &[u8],
1135) -> Result<(), SnapshotError> {
1136 let entry_header = encode_entry_header(key, value)?;
1137 for bytes in [&entry_header[..], key, value] {
1138 writer.write_all(bytes)?;
1139 hasher.update(bytes);
1140 }
1141 Ok(())
1142}
1143
1144fn encode_vector_space(definition: &VectorSpaceDefinition) -> Result<Vec<u8>, SnapshotError> {
1145 let name = definition.name.as_str().as_bytes();
1146 let name_length = u8::try_from(name.len()).map_err(|_| SnapshotError::Invalid {
1147 reason: "vector-space name length overflow",
1148 })?;
1149 let mut encoded = Vec::with_capacity(name.len() + 5);
1150 encoded.push(name_length);
1151 encoded.extend_from_slice(name);
1152 encoded.extend_from_slice(&definition.dimension.to_le_bytes());
1153 encoded.push(definition.metric as u8);
1154 encoded.push(1);
1155 Ok(encoded)
1156}
1157
1158fn decode_vector_space(encoded: &[u8]) -> Result<VectorSpaceDefinition, SnapshotError> {
1159 let name_length = usize::from(*encoded.first().ok_or(SnapshotError::Invalid {
1160 reason: "truncated vector-space record",
1161 })?);
1162 let expected_length = name_length.checked_add(5).ok_or(SnapshotError::Invalid {
1163 reason: "vector-space record length overflow",
1164 })?;
1165 if encoded.len() != expected_length || name_length == 0 {
1166 return Err(SnapshotError::Invalid {
1167 reason: "invalid vector-space record length",
1168 });
1169 }
1170 let name =
1171 std::str::from_utf8(&encoded[1..=name_length]).map_err(|_| SnapshotError::Invalid {
1172 reason: "invalid vector-space name",
1173 })?;
1174 let name = VectorSpaceName::new(name.to_owned()).map_err(|_| SnapshotError::Invalid {
1175 reason: "invalid vector-space name",
1176 })?;
1177 let dimension = u16::from_le_bytes(copy_array(&encoded[1 + name_length..3 + name_length]));
1178 if encoded[3 + name_length] != VectorMetric::Cosine as u8 || encoded[4 + name_length] != 1 {
1179 return Err(SnapshotError::Invalid {
1180 reason: "unsupported vector-space tags",
1181 });
1182 }
1183 VectorSpaceDefinition::cosine(name, dimension).map_err(|_| SnapshotError::Invalid {
1184 reason: "invalid vector-space dimension",
1185 })
1186}
1187
1188fn encode_vector(
1189 space: &VectorSpaceName,
1190 key: &[u8],
1191 vector: &Q15Vector,
1192) -> Result<Vec<u8>, SnapshotError> {
1193 if key.is_empty() || key.len() > MAX_KEY_BYTES {
1194 return Err(SnapshotError::Invalid {
1195 reason: "invalid vector object key",
1196 });
1197 }
1198 let space_name = space.as_str().as_bytes();
1199 let space_length = u8::try_from(space_name.len()).map_err(|_| SnapshotError::Invalid {
1200 reason: "vector-space name length overflow",
1201 })?;
1202 let key_length = u32::try_from(key.len()).map_err(|_| SnapshotError::Invalid {
1203 reason: "vector key length overflow",
1204 })?;
1205 let mut encoded =
1206 Vec::with_capacity(space_name.len() + key.len() + vector.as_slice().len() * 2 + 7);
1207 encoded.push(space_length);
1208 encoded.extend_from_slice(space_name);
1209 encoded.extend_from_slice(&key_length.to_le_bytes());
1210 encoded.extend_from_slice(key);
1211 encoded.extend_from_slice(&vector.dimension().to_le_bytes());
1212 for value in vector.as_slice() {
1213 encoded.extend_from_slice(&value.to_le_bytes());
1214 }
1215 Ok(encoded)
1216}
1217
1218fn decode_vector(encoded: &[u8]) -> Result<(VectorSpaceName, Vec<u8>, Q15Vector), SnapshotError> {
1219 let space_length = usize::from(*encoded.first().ok_or(SnapshotError::Invalid {
1220 reason: "truncated vector record",
1221 })?);
1222 let key_length_offset = 1_usize
1223 .checked_add(space_length)
1224 .ok_or(SnapshotError::Invalid {
1225 reason: "vector record length overflow",
1226 })?;
1227 let key_length_end = key_length_offset
1228 .checked_add(4)
1229 .ok_or(SnapshotError::Invalid {
1230 reason: "vector record length overflow",
1231 })?;
1232 let key_length = usize::try_from(u32::from_le_bytes(copy_array(
1233 encoded
1234 .get(key_length_offset..key_length_end)
1235 .ok_or(SnapshotError::Invalid {
1236 reason: "truncated vector record",
1237 })?,
1238 )))
1239 .map_err(|_| SnapshotError::Invalid {
1240 reason: "vector key length overflow",
1241 })?;
1242 if space_length == 0 || key_length == 0 || key_length > MAX_KEY_BYTES {
1243 return Err(SnapshotError::Invalid {
1244 reason: "invalid vector identity",
1245 });
1246 }
1247 let key_end = key_length_end
1248 .checked_add(key_length)
1249 .ok_or(SnapshotError::Invalid {
1250 reason: "vector record length overflow",
1251 })?;
1252 let dimension_end = key_end.checked_add(2).ok_or(SnapshotError::Invalid {
1253 reason: "vector record length overflow",
1254 })?;
1255 let dimension = usize::from(u16::from_le_bytes(copy_array(
1256 encoded
1257 .get(key_end..dimension_end)
1258 .ok_or(SnapshotError::Invalid {
1259 reason: "truncated vector record",
1260 })?,
1261 )));
1262 let expected_length = dimension
1263 .checked_mul(2)
1264 .and_then(|length| length.checked_add(dimension_end))
1265 .ok_or(SnapshotError::Invalid {
1266 reason: "vector record length overflow",
1267 })?;
1268 if encoded.len() != expected_length {
1269 return Err(SnapshotError::Invalid {
1270 reason: "invalid vector record length",
1271 });
1272 }
1273 let space = std::str::from_utf8(&encoded[1..key_length_offset]).map_err(|_| {
1274 SnapshotError::Invalid {
1275 reason: "invalid vector-space name",
1276 }
1277 })?;
1278 let space = VectorSpaceName::new(space.to_owned()).map_err(|_| SnapshotError::Invalid {
1279 reason: "invalid vector-space name",
1280 })?;
1281 let key = encoded[key_length_end..key_end].to_vec();
1282 let values = encoded[dimension_end..]
1283 .chunks_exact(2)
1284 .map(|chunk| i16::from_le_bytes(copy_array(chunk)))
1285 .collect::<Vec<_>>();
1286 let vector = Q15Vector::new(values).map_err(|_| SnapshotError::Invalid {
1287 reason: "invalid Q15 vector",
1288 })?;
1289 Ok((space, key, vector))
1290}
1291
1292fn encode_lexical_index(definition: &LexicalIndexDefinition) -> Result<Vec<u8>, SnapshotError> {
1293 let name = definition.name.as_str().as_bytes();
1294 let name_length = u8::try_from(name.len()).map_err(|_| SnapshotError::Invalid {
1295 reason: "lexical-index name length overflow",
1296 })?;
1297 let field_count =
1298 u8::try_from(definition.fields.len()).map_err(|_| SnapshotError::Invalid {
1299 reason: "lexical-index field count overflow",
1300 })?;
1301 let mut encoded = Vec::new();
1302 encoded.push(name_length);
1303 encoded.extend_from_slice(name);
1304 encoded.push(1);
1305 encoded.push(field_count);
1306 for field in &definition.fields {
1307 let segment_count =
1308 u8::try_from(field.path.segments().len()).map_err(|_| SnapshotError::Invalid {
1309 reason: "lexical-index segment count overflow",
1310 })?;
1311 encoded.push(segment_count);
1312 for segment in field.path.segments() {
1313 let segment_length =
1314 u16::try_from(segment.len()).map_err(|_| SnapshotError::Invalid {
1315 reason: "lexical-index segment length overflow",
1316 })?;
1317 encoded.extend_from_slice(&segment_length.to_le_bytes());
1318 encoded.extend_from_slice(segment.as_bytes());
1319 }
1320 encoded.extend_from_slice(&field.weight_micros.to_le_bytes());
1321 }
1322 Ok(encoded)
1323}
1324
1325#[allow(clippy::too_many_lines)]
1326fn decode_lexical_index(encoded: &[u8]) -> Result<LexicalIndexDefinition, SnapshotError> {
1327 let name_length = usize::from(*encoded.first().ok_or(SnapshotError::Invalid {
1328 reason: "truncated lexical-index record",
1329 })?);
1330 let name_end = 1_usize
1331 .checked_add(name_length)
1332 .ok_or(SnapshotError::Invalid {
1333 reason: "lexical-index record length overflow",
1334 })?;
1335 if name_length == 0 || encoded.get(name_end) != Some(&1) {
1336 return Err(SnapshotError::Invalid {
1337 reason: "invalid lexical-index record prefix",
1338 });
1339 }
1340 let name = std::str::from_utf8(encoded.get(1..name_end).ok_or(SnapshotError::Invalid {
1341 reason: "truncated lexical-index name",
1342 })?)
1343 .map_err(|_| SnapshotError::Invalid {
1344 reason: "invalid lexical-index name",
1345 })?;
1346 let name = VectorSpaceName::new(name.to_owned()).map_err(|_| SnapshotError::Invalid {
1347 reason: "invalid lexical-index name",
1348 })?;
1349 let mut cursor = name_end.checked_add(1).ok_or(SnapshotError::Invalid {
1350 reason: "lexical-index record length overflow",
1351 })?;
1352 let field_count = usize::from(*encoded.get(cursor).ok_or(SnapshotError::Invalid {
1353 reason: "truncated lexical-index field count",
1354 })?);
1355 cursor = cursor.checked_add(1).ok_or(SnapshotError::Invalid {
1356 reason: "lexical-index record length overflow",
1357 })?;
1358 if field_count == 0 || field_count > MAX_LEXICAL_FIELDS {
1359 return Err(SnapshotError::Invalid {
1360 reason: "invalid lexical-index field count",
1361 });
1362 }
1363 let mut fields = Vec::with_capacity(field_count);
1364 for _ in 0..field_count {
1365 let segment_count = usize::from(*encoded.get(cursor).ok_or(SnapshotError::Invalid {
1366 reason: "truncated lexical-index path",
1367 })?);
1368 cursor = cursor.checked_add(1).ok_or(SnapshotError::Invalid {
1369 reason: "lexical-index record length overflow",
1370 })?;
1371 if segment_count == 0 || segment_count > MAX_LEXICAL_PATH_SEGMENTS {
1372 return Err(SnapshotError::Invalid {
1373 reason: "invalid lexical-index path",
1374 });
1375 }
1376 let mut segments = Vec::with_capacity(segment_count);
1377 for _ in 0..segment_count {
1378 let length_end = cursor.checked_add(2).ok_or(SnapshotError::Invalid {
1379 reason: "lexical-index record length overflow",
1380 })?;
1381 let length = usize::from(u16::from_le_bytes(copy_array(
1382 encoded
1383 .get(cursor..length_end)
1384 .ok_or(SnapshotError::Invalid {
1385 reason: "truncated lexical-index segment length",
1386 })?,
1387 )));
1388 cursor = length_end;
1389 if length == 0 || length > MAX_LEXICAL_PATH_SEGMENT_BYTES {
1390 return Err(SnapshotError::Invalid {
1391 reason: "invalid lexical-index segment length",
1392 });
1393 }
1394 let segment_end = cursor.checked_add(length).ok_or(SnapshotError::Invalid {
1395 reason: "lexical-index record length overflow",
1396 })?;
1397 let segment = std::str::from_utf8(encoded.get(cursor..segment_end).ok_or(
1398 SnapshotError::Invalid {
1399 reason: "truncated lexical-index segment",
1400 },
1401 )?)
1402 .map_err(|_| SnapshotError::Invalid {
1403 reason: "invalid lexical-index segment",
1404 })?
1405 .to_owned();
1406 cursor = segment_end;
1407 segments.push(segment);
1408 }
1409 let weight_end = cursor.checked_add(4).ok_or(SnapshotError::Invalid {
1410 reason: "lexical-index record length overflow",
1411 })?;
1412 let weight_micros = u32::from_le_bytes(copy_array(encoded.get(cursor..weight_end).ok_or(
1413 SnapshotError::Invalid {
1414 reason: "truncated lexical-index field weight",
1415 },
1416 )?));
1417 cursor = weight_end;
1418 fields.push(LexicalField {
1419 path: FieldPath::new(segments),
1420 weight_micros,
1421 });
1422 }
1423 if cursor != encoded.len() {
1424 return Err(SnapshotError::Invalid {
1425 reason: "invalid lexical-index record length",
1426 });
1427 }
1428 LexicalIndexDefinition::new(name, fields).map_err(|_| SnapshotError::Invalid {
1429 reason: "invalid lexical-index definition",
1430 })
1431}
1432
1433fn write_encoded(
1434 writer: &mut impl Write,
1435 hasher: &mut blake3::Hasher,
1436 encoded: Result<Vec<u8>, SnapshotError>,
1437) -> Result<(), SnapshotError> {
1438 let encoded = encoded?;
1439 writer.write_all(&encoded)?;
1440 hasher.update(&encoded);
1441 Ok(())
1442}
1443
1444fn read_v2_counts(
1445 reader: &mut impl Read,
1446 decoded: &DecodedHeader,
1447 consumed: &mut u64,
1448) -> Result<(u64, u64, u64), SnapshotError> {
1449 if decoded.disk_format_version < 2 {
1450 return Ok((0, 0, 0));
1451 }
1452 let mut encoded = [0_u8; V2_COUNTS_LENGTH];
1453 read_payload_exact(reader, &mut encoded, consumed, decoded.payload_length)?;
1454 Ok((
1455 u64::from_le_bytes(copy_array(&encoded[..8])),
1456 u64::from_le_bytes(copy_array(&encoded[8..16])),
1457 u64::from_le_bytes(copy_array(&encoded[16..24])),
1458 ))
1459}
1460
1461fn read_encoded_vector_space(
1462 reader: &mut impl Read,
1463 decoded: &DecodedHeader,
1464 consumed: &mut u64,
1465) -> Result<Vec<u8>, SnapshotError> {
1466 let mut name_length = [0_u8; 1];
1467 read_payload_exact(reader, &mut name_length, consumed, decoded.payload_length)?;
1468 let remaining = usize::from(name_length[0])
1469 .checked_add(4)
1470 .ok_or(SnapshotError::Invalid {
1471 reason: "vector-space record length overflow",
1472 })?;
1473 let mut encoded = vec![name_length[0]];
1474 let mut tail = vec![0_u8; remaining];
1475 read_payload_exact(reader, &mut tail, consumed, decoded.payload_length)?;
1476 encoded.extend_from_slice(&tail);
1477 Ok(encoded)
1478}
1479
1480fn read_vector_space(
1481 reader: &mut impl Read,
1482 decoded: &DecodedHeader,
1483 consumed: &mut u64,
1484) -> Result<VectorSpaceDefinition, SnapshotError> {
1485 decode_vector_space(&read_encoded_vector_space(reader, decoded, consumed)?)
1486}
1487
1488fn read_encoded_vector(
1489 reader: &mut impl Read,
1490 decoded: &DecodedHeader,
1491 consumed: &mut u64,
1492) -> Result<Vec<u8>, SnapshotError> {
1493 let mut space_length = [0_u8; 1];
1494 read_payload_exact(reader, &mut space_length, consumed, decoded.payload_length)?;
1495 let space_length = usize::from(space_length[0]);
1496 let mut prefix_tail = vec![0_u8; space_length + 4];
1497 read_payload_exact(reader, &mut prefix_tail, consumed, decoded.payload_length)?;
1498 let key_length = usize::try_from(u32::from_le_bytes(copy_array(&prefix_tail[space_length..])))
1499 .map_err(|_| SnapshotError::Invalid {
1500 reason: "vector key length overflow",
1501 })?;
1502 if space_length == 0 || key_length == 0 || key_length > MAX_KEY_BYTES {
1503 return Err(SnapshotError::Invalid {
1504 reason: "invalid vector identity",
1505 });
1506 }
1507 let mut key_and_dimension = vec![0_u8; key_length + 2];
1508 read_payload_exact(
1509 reader,
1510 &mut key_and_dimension,
1511 consumed,
1512 decoded.payload_length,
1513 )?;
1514 let dimension = usize::from(u16::from_le_bytes(copy_array(
1515 &key_and_dimension[key_length..],
1516 )));
1517 let vector_bytes = dimension.checked_mul(2).ok_or(SnapshotError::Invalid {
1518 reason: "vector record length overflow",
1519 })?;
1520 let mut values = vec![0_u8; vector_bytes];
1521 read_payload_exact(reader, &mut values, consumed, decoded.payload_length)?;
1522 let mut encoded =
1523 Vec::with_capacity(1 + prefix_tail.len() + key_and_dimension.len() + values.len());
1524 encoded.push(
1525 u8::try_from(space_length).map_err(|_| SnapshotError::Invalid {
1526 reason: "vector-space name length overflow",
1527 })?,
1528 );
1529 encoded.extend_from_slice(&prefix_tail);
1530 encoded.extend_from_slice(&key_and_dimension);
1531 encoded.extend_from_slice(&values);
1532 Ok(encoded)
1533}
1534
1535fn read_vector(
1536 reader: &mut impl Read,
1537 decoded: &DecodedHeader,
1538 consumed: &mut u64,
1539) -> Result<(VectorSpaceName, Vec<u8>, Q15Vector), SnapshotError> {
1540 decode_vector(&read_encoded_vector(reader, decoded, consumed)?)
1541}
1542
1543fn read_encoded_lexical_index(
1544 reader: &mut impl Read,
1545 decoded: &DecodedHeader,
1546 consumed: &mut u64,
1547) -> Result<Vec<u8>, SnapshotError> {
1548 let mut name_length = [0_u8; 1];
1549 read_payload_exact(reader, &mut name_length, consumed, decoded.payload_length)?;
1550 let name_length_usize = usize::from(name_length[0]);
1551 if name_length_usize == 0 {
1552 return Err(SnapshotError::Invalid {
1553 reason: "invalid lexical-index name length",
1554 });
1555 }
1556 let mut name_and_counts = vec![0_u8; name_length_usize + 2];
1557 read_payload_exact(
1558 reader,
1559 &mut name_and_counts,
1560 consumed,
1561 decoded.payload_length,
1562 )?;
1563 if name_and_counts[name_length_usize] != 1 {
1564 return Err(SnapshotError::Invalid {
1565 reason: "unsupported lexical-index record version",
1566 });
1567 }
1568 let field_count = usize::from(name_and_counts[name_length_usize + 1]);
1569 if field_count == 0 || field_count > MAX_LEXICAL_FIELDS {
1570 return Err(SnapshotError::Invalid {
1571 reason: "invalid lexical-index field count",
1572 });
1573 }
1574 let mut encoded = Vec::new();
1575 encoded.push(name_length[0]);
1576 encoded.extend_from_slice(&name_and_counts);
1577 for _ in 0..field_count {
1578 let mut segment_count = [0_u8; 1];
1579 read_payload_exact(reader, &mut segment_count, consumed, decoded.payload_length)?;
1580 let segment_count_usize = usize::from(segment_count[0]);
1581 if segment_count_usize == 0 || segment_count_usize > MAX_LEXICAL_PATH_SEGMENTS {
1582 return Err(SnapshotError::Invalid {
1583 reason: "invalid lexical-index path",
1584 });
1585 }
1586 encoded.push(segment_count[0]);
1587 for _ in 0..segment_count_usize {
1588 let mut length = [0_u8; 2];
1589 read_payload_exact(reader, &mut length, consumed, decoded.payload_length)?;
1590 let length_usize = usize::from(u16::from_le_bytes(length));
1591 if length_usize == 0 || length_usize > MAX_LEXICAL_PATH_SEGMENT_BYTES {
1592 return Err(SnapshotError::Invalid {
1593 reason: "invalid lexical-index segment length",
1594 });
1595 }
1596 let mut segment = vec![0_u8; length_usize];
1597 read_payload_exact(reader, &mut segment, consumed, decoded.payload_length)?;
1598 encoded.extend_from_slice(&length);
1599 encoded.extend_from_slice(&segment);
1600 }
1601 let mut weight = [0_u8; 4];
1602 read_payload_exact(reader, &mut weight, consumed, decoded.payload_length)?;
1603 encoded.extend_from_slice(&weight);
1604 }
1605 Ok(encoded)
1606}
1607
1608fn read_lexical_index(
1609 reader: &mut impl Read,
1610 decoded: &DecodedHeader,
1611 consumed: &mut u64,
1612) -> Result<LexicalIndexDefinition, SnapshotError> {
1613 decode_lexical_index(&read_encoded_lexical_index(reader, decoded, consumed)?)
1614}
1615
1616fn encode_receipt(receipt: &CommitReceipt) -> [u8; RECEIPT_LENGTH] {
1617 let mut encoded = [0_u8; RECEIPT_LENGTH];
1618 encoded[..16].copy_from_slice(receipt.transaction_id.as_bytes());
1619 encoded[16..24].copy_from_slice(&receipt.commit_sequence.to_le_bytes());
1620 encoded[24..56].copy_from_slice(&receipt.commit_digest);
1621 encoded[56..88].copy_from_slice(&receipt.transaction_digest);
1622 encoded
1623}
1624
1625fn decode_snapshot_receipt(encoded: &[u8; RECEIPT_LENGTH]) -> CommitReceipt {
1626 CommitReceipt {
1627 transaction_id: uuid::Uuid::from_bytes(copy_array(&encoded[..16])),
1628 commit_sequence: u64::from_le_bytes(copy_array(&encoded[16..24])),
1629 commit_digest: copy_array(&encoded[24..56]),
1630 transaction_digest: copy_array(&encoded[56..88]),
1631 }
1632}
1633
1634fn write_receipt(
1635 writer: &mut impl Write,
1636 hasher: &mut blake3::Hasher,
1637 receipt: &CommitReceipt,
1638) -> Result<(), SnapshotError> {
1639 let encoded = encode_receipt(receipt);
1640 writer.write_all(&encoded)?;
1641 hasher.update(&encoded);
1642 Ok(())
1643}
1644
1645fn read_payload_exact(
1646 reader: &mut impl Read,
1647 buffer: &mut [u8],
1648 consumed: &mut u64,
1649 payload_length: u64,
1650) -> Result<(), SnapshotError> {
1651 let length = u64::try_from(buffer.len()).map_err(|_| SnapshotError::Invalid {
1652 reason: "payload length overflow",
1653 })?;
1654 let next = consumed.checked_add(length).ok_or(SnapshotError::Invalid {
1655 reason: "payload length overflow",
1656 })?;
1657 if next > payload_length {
1658 return Err(SnapshotError::Invalid {
1659 reason: "entry exceeds payload",
1660 });
1661 }
1662 read_exact_or_invalid(reader, buffer, "truncated payload")?;
1663 *consumed = next;
1664 Ok(())
1665}
1666
1667fn read_exact_or_invalid(
1668 reader: &mut impl Read,
1669 buffer: &mut [u8],
1670 reason: &'static str,
1671) -> Result<(), SnapshotError> {
1672 reader.read_exact(buffer).map_err(|source| {
1673 if source.kind() == io::ErrorKind::UnexpectedEof {
1674 SnapshotError::Invalid { reason }
1675 } else {
1676 SnapshotError::Io(source)
1677 }
1678 })
1679}
1680
1681#[cfg(unix)]
1682fn sync_directory(path: &Path) -> Result<(), SnapshotError> {
1683 File::open(path)?.sync_all()?;
1684 Ok(())
1685}
1686
1687fn copy_array<const N: usize>(source: &[u8]) -> [u8; N] {
1688 let mut output = [0_u8; N];
1689 output.copy_from_slice(source);
1690 output
1691}