Skip to main content

hyphae_storage/
snapshot.rs

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