1use std::collections::{BTreeMap, BTreeSet};
9
10use serde_json::Value;
11
12use super::{
13 AUDIT_ENTRY_LEN, BINARY_TRANSACTION_MAGIC, END_MARK, FRAME_HEADER_LEN, MAX_FRAME_PAYLOAD,
14 TEXT_SEGMENT_HEADER_LEN, TEXT_SEGMENT_MAGIC, TRANSACTION_MAGIC, TRANSACTION_MAGIC_PREFIX,
15};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub struct FrameRef {
20 pub offset: usize,
22 pub raw_len: u32,
24 pub comp_len: u32,
26 pub seq: u32,
28 pub crc: u32,
30}
31
32impl FrameRef {
33 pub fn end(&self) -> usize {
35 self.offset + FRAME_HEADER_LEN + self.comp_len as usize
36 }
37}
38
39#[derive(Debug, Clone, PartialEq, Eq)]
41pub enum SegmentEnd {
42 Marker,
44 ZeroHeader,
46 Exhausted,
48 TornTail { reason: &'static str },
51 Corrupt {
54 reason: &'static str,
55 committed_at: Option<usize>,
57 },
58}
59
60#[derive(Debug, Clone)]
62pub struct SegmentScan {
63 pub frames: Vec<FrameRef>,
64 pub end_offset: usize,
66 pub end: SegmentEnd,
67}
68
69fn word(data: &[u8], at: usize) -> u32 {
70 u32::from_le_bytes([data[at], data[at + 1], data[at + 2], data[at + 3]])
71}
72
73pub fn is_log_end(data: &[u8], at: usize) -> bool {
75 at.checked_add(FRAME_HEADER_LEN)
76 .is_some_and(|end| end <= data.len())
77 && word(data, at + 4) == 0
78 && word(data, at + 8) == 0
79 && ((word(data, at) == 0 && word(data, at + 12) == END_MARK)
80 || data[at..at + FRAME_HEADER_LEN].iter().all(|&b| b == 0))
81}
82
83fn is_end_marker(data: &[u8], at: usize) -> bool {
84 is_log_end(data, at) && word(data, at + 12) == END_MARK
85}
86
87pub fn decode_frame(data: &[u8], offset: usize) -> Result<(FrameRef, Vec<u8>), &'static str> {
92 let mut payload = Vec::new();
93 let frame = decode_frame_into(data, offset, &mut payload)?;
94 Ok((frame, payload))
95}
96
97fn decode_frame_into(
100 data: &[u8],
101 offset: usize,
102 out: &mut Vec<u8>,
103) -> Result<FrameRef, &'static str> {
104 let header_end = offset
105 .checked_add(FRAME_HEADER_LEN)
106 .filter(|&end| end <= data.len())
107 .ok_or("truncated-header")?;
108 let (crc, raw_len, comp_len, seq) = (
109 word(data, offset),
110 word(data, offset + 4),
111 word(data, offset + 8),
112 word(data, offset + 12),
113 );
114 if raw_len == 0 || comp_len == 0 || raw_len as usize > MAX_FRAME_PAYLOAD || comp_len > raw_len {
115 return Err("header-lengths");
116 }
117 let end = header_end
118 .checked_add(comp_len as usize)
119 .filter(|&end| end <= data.len())
120 .ok_or("truncated-payload")?;
121 let payload = &data[header_end..end];
122 if comp_len < raw_len {
123 if u64::from(raw_len) > u64::from(comp_len) * 255 {
124 return Err("implausible-ratio");
125 }
126 let raw_len = raw_len as usize;
129 if out.len() < raw_len {
130 out.resize(raw_len, 0);
131 }
132 let written =
133 lz4_flex::block::decompress_into(payload, &mut out[..raw_len]).map_err(|_| "lz4")?;
134 if written != raw_len {
136 return Err("raw-length");
137 }
138 out.truncate(written);
139 } else {
140 out.clear();
141 out.extend_from_slice(payload);
142 }
143 if crc32fast::hash(out) != crc {
144 return Err("crc32");
145 }
146 Ok(FrameRef {
147 offset,
148 raw_len,
149 comp_len,
150 seq,
151 crc,
152 })
153}
154
155pub(crate) fn last_nonzero(data: &[u8]) -> Option<usize> {
156 let (blocks, tail) = data.as_chunks::<4096>();
157 if let Some(at) = tail.iter().rposition(|&b| b != 0) {
158 return Some(blocks.len() * 4096 + at);
159 }
160 let block = blocks.iter().rposition(|bytes| bytes != &[0; 4096])?;
161 blocks[block]
162 .iter()
163 .rposition(|&b| b != 0)
164 .map(|at| block * 4096 + at)
165}
166
167fn committed_after(data: &[u8], from: usize, include_end_marker: bool) -> Option<usize> {
168 let last = last_nonzero(data)?;
169 let mut scratch = Vec::new();
171 (from..=last.saturating_sub(FRAME_HEADER_LEN - 1)).find(|&at| {
172 at + FRAME_HEADER_LEN <= data.len()
173 && ((include_end_marker && is_end_marker(data, at))
174 || decode_frame_into(data, at, &mut scratch).is_ok())
175 })
176}
177
178pub fn scan_segment(
182 data: &[u8],
183 is_final: bool,
184 mut visit: impl FnMut(&FrameRef, &[u8]),
185) -> SegmentScan {
186 let mut frames = Vec::new();
187 let mut payload = Vec::new();
188 let mut offset = 0usize;
189 let finish = |frames: Vec<FrameRef>, offset: usize, end: SegmentEnd| SegmentScan {
190 frames,
191 end_offset: offset,
192 end,
193 };
194 while offset + FRAME_HEADER_LEN <= data.len() {
195 if word(data, offset + 4) == 0 && word(data, offset + 8) == 0 {
196 if is_log_end(data, offset) {
197 if let Some(at) = committed_after(data, offset + FRAME_HEADER_LEN, false) {
200 return SegmentScan {
201 frames,
202 end_offset: offset,
203 end: SegmentEnd::Corrupt {
204 reason: "committed-after-end",
205 committed_at: Some(at),
206 },
207 };
208 }
209 let end = if is_end_marker(data, offset) {
210 SegmentEnd::Marker
211 } else {
212 SegmentEnd::ZeroHeader
213 };
214 return finish(frames, offset, end);
215 }
216 return fault(data, is_final, frames, offset, "invalid-end-header");
217 }
218 match decode_frame_into(data, offset, &mut payload) {
219 Ok(frame) => {
220 visit(&frame, &payload);
221 offset = frame.end();
222 frames.push(frame);
223 }
224 Err(reason) => return fault(data, is_final, frames, offset, reason),
225 }
226 }
227 finish(frames, offset, SegmentEnd::Exhausted)
228}
229
230fn fault(
231 data: &[u8],
232 is_final: bool,
233 frames: Vec<FrameRef>,
234 offset: usize,
235 reason: &'static str,
236) -> SegmentScan {
237 let committed_at = committed_after(data, offset + 1, true);
238 let end = if is_final && committed_at.is_none() {
239 SegmentEnd::TornTail { reason }
240 } else {
241 SegmentEnd::Corrupt {
242 reason,
243 committed_at,
244 }
245 };
246 SegmentScan {
247 frames,
248 end_offset: offset,
249 end,
250 }
251}
252
253#[derive(Debug, Clone, PartialEq, Eq)]
255pub enum PayloadKind<'a> {
256 Transaction(&'a [u8]),
258 BinaryTransaction(&'a [u8]),
261 OtherTransactionVersion(String),
263 Unrecognized,
265}
266
267pub fn classify_payload(payload: &[u8]) -> PayloadKind<'_> {
268 if let Some(json) = payload.strip_prefix(TRANSACTION_MAGIC) {
269 return PayloadKind::Transaction(json);
270 }
271 if let Some(body) = payload.strip_prefix(BINARY_TRANSACTION_MAGIC) {
272 return PayloadKind::BinaryTransaction(body);
273 }
274 if let Some(rest) = payload.strip_prefix(TRANSACTION_MAGIC_PREFIX) {
275 let digits = rest.iter().take_while(|b| b.is_ascii_digit()).count();
276 if digits > 0 && digits <= 9 && rest.get(digits) == Some(&0) {
277 return PayloadKind::OtherTransactionVersion(
278 String::from_utf8_lossy(&rest[..digits]).into_owned(),
279 );
280 }
281 }
282 PayloadKind::Unrecognized
283}
284
285pub const MUTATION_VARIANTS: [&str; 17] = [
287 "Vector",
288 "Atom",
289 "Composite",
290 "DeleteAtom",
291 "DeleteComposite",
292 "RefreshMeaningIndices",
293 "StoredTriple",
294 "Document",
295 "EncodeDocument",
296 "Derived",
297 "DeleteDocument",
298 "Delete",
299 "Triplet",
300 "Relation",
301 "RemoveRelation",
302 "RelationType",
303 "Rule",
304];
305
306const MAX_TRANSACTION_OPS: usize = 4096;
308const MAX_ID_BYTES: usize = u16::MAX as usize;
309const MAX_DOCUMENT_BYTES: u64 = 8 * 1024 * 1024;
310const MAX_CHUNK_BYTES: u64 = 64 * 1024;
311const MAX_CHUNKS: usize = 4096;
312const WHOLE_DOCUMENT_POSITIONS: u64 = 2;
314
315pub(crate) fn chunk_byte_limit(
320 whole_document: bool,
321 chunk_count: usize,
322 start_byte: u64,
323 lexical: bool,
324) -> Option<u64> {
325 if !whole_document {
326 return Some(MAX_CHUNK_BYTES);
327 }
328 (lexical && chunk_count == 1 && start_byte == 0).then_some(MAX_DOCUMENT_BYTES)
329}
330
331#[derive(Debug, Clone, Default)]
334pub struct Replay {
335 pub documents: BTreeMap<String, DocumentSummary>,
336 pub vectors: BTreeSet<String>,
337 pub derived: BTreeMap<String, Vec<String>>,
338 pub relations: BTreeSet<(String, String, String)>,
339 pub mutations: usize,
340}
341
342#[derive(Debug, Clone)]
343pub struct DocumentSummary {
344 pub chunks: usize,
345 pub indexed: usize,
348 pub version: String,
349 pub stamped: bool,
351}
352
353impl Replay {
354 pub fn vector_total(&self) -> usize {
356 self.vectors.len() + self.documents.values().map(|d| d.indexed).sum::<usize>()
357 }
358
359 fn drop_derived(&mut self, id: &str) {
360 if let Some(children) = self.derived.remove(id) {
361 for child in children {
362 self.vectors.remove(&child);
363 }
364 }
365 }
366
367 fn apply(&mut self, variant: &str, body: &Value) {
368 let id = body.get("id").and_then(Value::as_str).unwrap_or_default();
369 match variant {
370 "Vector" | "Composite" => {
371 if variant == "Vector" {
372 self.drop_derived(id);
373 }
374 self.vectors.insert(id.to_owned());
375 }
376 "Triplet" => {
377 self.drop_derived(id);
378 self.vectors.insert(id.to_owned());
379 }
380 "Document" => {
381 let version = body
382 .get("version")
383 .and_then(Value::as_str)
384 .unwrap_or_default();
385 let stamped = version.starts_with("mem1:")
386 || body
387 .get("metadata")
388 .and_then(Value::as_object)
389 .is_some_and(|m| m.contains_key("_mem"));
390 let chunks = body.get("chunks").and_then(Value::as_array);
391 self.documents.insert(
392 id.to_owned(),
393 DocumentSummary {
394 chunks: chunks.map_or(0, Vec::len),
395 indexed: chunks.map_or(0, |chunks| {
396 chunks
397 .iter()
398 .filter(|chunk| chunk.get("vector_state").is_none())
399 .count()
400 }),
401 version: version.to_owned(),
402 stamped,
403 },
404 );
405 }
406 "EncodeDocument" => {
407 let version = body
410 .get("version")
411 .and_then(Value::as_str)
412 .unwrap_or_default();
413 let vectors = body
414 .get("vectors")
415 .and_then(Value::as_array)
416 .map_or(0, Vec::len);
417 if let Some(document) = self.documents.get_mut(id) {
418 if document.version == version && document.chunks == vectors {
419 document.indexed = document.chunks;
420 }
421 }
422 }
423 "DeleteDocument" => {
424 self.documents.remove(id);
425 }
426 "Delete" => {
427 self.drop_derived(id);
428 self.vectors.remove(id);
429 }
430 "Derived" => {
431 let children = body
432 .get("children")
433 .and_then(Value::as_array)
434 .map(|c| {
435 c.iter()
436 .filter_map(Value::as_str)
437 .map(str::to_owned)
438 .collect()
439 })
440 .unwrap_or_default();
441 self.derived.insert(id.to_owned(), children);
442 }
443 "Relation" => {
444 self.relations
445 .insert(relation_key(body, "sourceId", "relationType", "targetId"));
446 }
447 "RemoveRelation" => {
448 self.relations
449 .remove(&relation_key(body, "source", "relation", "target"));
450 }
451 _ => {}
452 }
453 }
454}
455
456fn relation_key(body: &Value, a: &str, b: &str, c: &str) -> (String, String, String) {
457 let get = |k: &str| {
458 body.get(k)
459 .and_then(Value::as_str)
460 .unwrap_or_default()
461 .to_owned()
462 };
463 (get(a), get(b), get(c))
464}
465
466fn valid_id(value: Option<&Value>) -> bool {
467 value
468 .and_then(Value::as_str)
469 .is_some_and(|id| !id.is_empty() && id.len() <= MAX_ID_BYTES && !id.contains('\0'))
470}
471
472fn valid_vector(value: Option<&Value>, dimensions: Option<u64>) -> Result<(), String> {
473 let vector = value
474 .and_then(Value::as_object)
475 .ok_or("vector is not an object")?;
476 let dim = vector
477 .get("dim")
478 .and_then(Value::as_u64)
479 .ok_or("vector has no integer dim")?;
480 if dimensions.is_some_and(|d| d != dim) {
481 return Err(format!("vector dim {dim} differs from the store's"));
482 }
483 let indices = vector
484 .get("indices")
485 .and_then(Value::as_array)
486 .ok_or("vector has no indices array")?;
487 let mut previous: Option<u64> = None;
488 for index in indices {
489 let index = index.as_u64().ok_or("vector index is not an integer")?;
490 if index >= dim || previous.is_some_and(|p| p >= index) {
491 return Err("vector indices are not strictly ascending below dim".into());
492 }
493 previous = Some(index);
494 }
495 Ok(())
496}
497
498fn check_document(body: &Value, dimensions: Option<u64>) -> Result<(), String> {
499 if !valid_id(body.get("id")) {
500 return Err("document id is not 1..=65535 bytes without NUL".into());
501 }
502 let chunks = body
503 .get("chunks")
504 .and_then(Value::as_array)
505 .ok_or("document has no chunks array")?;
506 if chunks.is_empty() || chunks.len() > MAX_CHUNKS {
507 return Err(format!(
508 "document has {} chunks, expected 1..=4096",
509 chunks.len()
510 ));
511 }
512 for chunk in chunks {
513 if !valid_id(chunk.get("id")) {
514 return Err("chunk id is not 1..=65535 bytes without NUL".into());
515 }
516 valid_vector(chunk.get("vector"), dimensions)?;
517 let start = chunk.get("start_byte").and_then(Value::as_u64);
518 let end = chunk.get("end_byte").and_then(Value::as_u64);
519 let (Some(start), Some(end)) = (start, end) else {
520 return Err("chunk byte range is missing".into());
521 };
522 let whole_document = chunk
523 .get("terms")
524 .and_then(|terms| terms.get("$hms_positions"))
525 .and_then(|positions| positions.get("version"))
526 .and_then(Value::as_u64)
527 == Some(WHOLE_DOCUMENT_POSITIONS);
528 let lexical = chunk.get("vector_state").and_then(Value::as_str) == Some("lexical");
529 let Some(limit) = chunk_byte_limit(whole_document, chunks.len(), start, lexical) else {
530 return Err("whole-document postings require one lexical chunk from byte 0".into());
531 };
532 if start >= end || end > MAX_DOCUMENT_BYTES || end - start > limit {
533 return Err("chunk byte range is empty or over the limits".into());
534 }
535 if let Some(text) = chunk.get("text").and_then(Value::as_str) {
536 if text.len() as u64 != end - start {
537 return Err("chunk text length differs from its byte range".into());
538 }
539 }
540 if let Some(state) = chunk.get("vector_state") {
541 let known = matches!(state.as_str(), Some("pending" | "lexical"));
542 let empty = chunk
543 .get("vector")
544 .and_then(|v| v.get("indices"))
545 .and_then(Value::as_array)
546 .is_some_and(Vec::is_empty);
547 if !known || !empty {
548 return Err(
549 "chunk vector_state must be pending or lexical with an empty vector".into(),
550 );
551 }
552 }
553 }
554 Ok(())
555}
556
557fn check_encode_document(body: &Value, dimensions: Option<u64>) -> Result<(), String> {
560 if !valid_id(body.get("id")) {
561 return Err("document id is not 1..=65535 bytes without NUL".into());
562 }
563 let version_ok = body
564 .get("version")
565 .and_then(Value::as_str)
566 .is_some_and(|v| !v.is_empty() && v.len() <= 1024);
567 if !version_ok {
568 return Err("deferred encoding version is not 1..=1024 bytes".into());
569 }
570 let vectors = body
571 .get("vectors")
572 .and_then(Value::as_array)
573 .ok_or("deferred encoding has no vectors array")?;
574 if vectors.is_empty() || vectors.len() > MAX_CHUNKS {
575 return Err(format!(
576 "deferred encoding has {} vectors, expected 1..=4096",
577 vectors.len()
578 ));
579 }
580 for vector in vectors {
581 valid_vector(Some(vector), dimensions)?;
582 }
583 Ok(())
584}
585
586fn check_mutation(
588 value: &Value,
589 dimensions: Option<u64>,
590 replay: &mut Replay,
591) -> Result<(), String> {
592 let object = value.as_object().ok_or("mutation is not an object")?;
593 if object.len() != 1 {
594 return Err("mutation must have exactly one variant key".into());
595 }
596 let (variant, body) = object.iter().next().ok_or("empty mutation")?;
597 if !MUTATION_VARIANTS.contains(&variant.as_str()) {
598 return Err(format!("unknown mutation variant {variant:?}"));
599 }
600 match variant.as_str() {
601 "Vector" | "Atom" | "Composite" => {
602 if !valid_id(body.get("id")) {
603 return Err("vector id is not 1..=65535 bytes without NUL".into());
604 }
605 valid_vector(body.get("vector"), dimensions)?;
606 }
607 "Document" => check_document(body, dimensions)?,
608 "StoredTriple" => {
609 for key in ["subject_id", "relation_id", "object_id", "composite_id"] {
610 if !valid_id(body.get(key)) {
611 return Err(format!("stored triple {key} is not a valid id"));
612 }
613 }
614 if !body.get("deleted").is_some_and(Value::is_boolean) {
615 return Err("stored triple deleted is not a boolean".into());
616 }
617 }
618 "EncodeDocument" => check_encode_document(body, dimensions)?,
619 "Delete" | "DeleteDocument" | "DeleteAtom" | "DeleteComposite" => {
620 if !valid_id(body.get("id")) {
621 return Err("id is not 1..=65535 bytes without NUL".into());
622 }
623 }
624 "RefreshMeaningIndices" => {
625 if !body.as_object().is_some_and(|body| body.is_empty()) {
626 return Err("meaning index refresh body must be empty".into());
627 }
628 }
629 "Derived" => {
630 if !valid_id(body.get("id")) || !body.get("children").is_some_and(Value::is_array) {
631 return Err("derived record needs an id and children".into());
632 }
633 }
634 "Triplet" => {
635 for key in ["id", "subject", "relation", "object"] {
636 if !valid_id(body.get(key)) {
637 return Err(format!("triplet {key} is not a valid id"));
638 }
639 }
640 }
641 "Relation" => {
642 for key in ["sourceId", "relationType", "targetId"] {
643 if !valid_id(body.get(key)) {
644 return Err(format!("relation {key} is not a valid id"));
645 }
646 }
647 }
648 "RemoveRelation" => {
649 for key in ["source", "relation", "target"] {
650 if !valid_id(body.get(key)) {
651 return Err(format!("relation removal {key} is not a valid id"));
652 }
653 }
654 }
655 _ => {
656 if !body.is_object() {
657 return Err(format!("{variant} body is not an object"));
658 }
659 }
660 }
661 replay.apply(variant, body);
662 replay.mutations += 1;
663 Ok(())
664}
665
666pub fn transaction_json(payload: &[u8]) -> Option<Result<std::borrow::Cow<'_, [u8]>, String>> {
669 crate::core::engine::journal::payload_json(payload)
670 .map(|json| json.map_err(|error| format!("{error:#}")))
671}
672
673pub fn check_transaction(
676 json: &[u8],
677 dimensions: Option<u64>,
678 replay: &mut Replay,
679) -> Result<usize, String> {
680 check_transaction_with_schema(json, dimensions, super::STORE_SCHEMA_VERSION, replay)
681}
682
683pub fn check_transaction_with_schema(
685 json: &[u8],
686 dimensions: Option<u64>,
687 schema: u32,
688 replay: &mut Replay,
689) -> Result<usize, String> {
690 let value: Value = serde_json::from_slice(json).map_err(|e| format!("not JSON: {e}"))?;
691 let mutations = value.as_array().ok_or("transaction is not an array")?;
692 if schema < super::MEANING_MAINTENANCE_SCHEMA_VERSION
693 && mutations.iter().any(|mutation| {
694 ["DeleteAtom", "DeleteComposite", "RefreshMeaningIndices"]
695 .iter()
696 .any(|variant| mutation.get(variant).is_some())
697 })
698 {
699 return Err("scoped meaning maintenance requires store schema 5".into());
700 }
701 if mutations.is_empty() || mutations.len() > MAX_TRANSACTION_OPS {
702 return Err(format!(
703 "transaction has {} operations, expected 1..=4096",
704 mutations.len()
705 ));
706 }
707 for (n, mutation) in mutations.iter().enumerate() {
708 check_mutation(mutation, dimensions, replay).map_err(|e| format!("mutation {n}: {e}"))?;
709 }
710 Ok(mutations.len())
711}
712
713#[derive(Debug, Clone, Copy, PartialEq, Eq)]
715pub enum TextHeader {
716 Versioned { version: u8, flags: u8 },
718 Headerless,
720}
721
722pub fn parse_text_header(bytes: &[u8]) -> TextHeader {
723 if bytes.len() >= TEXT_SEGMENT_HEADER_LEN && &bytes[..4] == TEXT_SEGMENT_MAGIC {
724 TextHeader::Versioned {
725 version: bytes[4],
726 flags: bytes[5],
727 }
728 } else {
729 TextHeader::Headerless
730 }
731}
732
733#[derive(Debug, Clone, Copy, PartialEq, Eq)]
735pub struct AuditScan {
736 pub entries: usize,
737 pub trailing_bytes: usize,
739 pub unknown_ops: usize,
741}
742
743pub fn scan_audit(bytes: &[u8]) -> AuditScan {
744 let entries = bytes.len() / AUDIT_ENTRY_LEN;
745 let unknown_ops = (0..entries)
746 .filter(|i| !(1..=3).contains(&bytes[i * AUDIT_ENTRY_LEN + 8]))
747 .count();
748 AuditScan {
749 entries,
750 trailing_bytes: bytes.len() % AUDIT_ENTRY_LEN,
751 unknown_ops,
752 }
753}
754
755#[derive(Debug, Clone, PartialEq, Eq)]
757pub struct IndicesMeta {
758 pub checkpoint: (String, u64),
760 pub shard_count: u64,
761 pub shards: Vec<IndexShardMeta>,
762}
763
764#[derive(Debug, Clone, PartialEq, Eq)]
765pub struct IndexShardMeta {
766 pub nsg: bool,
767 pub ivf: bool,
768 pub nsg_crc: Option<u32>,
769 pub ivf_crc: Option<u32>,
770}
771
772pub fn parse_indices(bytes: &[u8]) -> Result<IndicesMeta, String> {
773 let value: Value = serde_json::from_slice(bytes).map_err(|e| e.to_string())?;
774 let checkpoint = value
775 .get("checkpoint")
776 .and_then(Value::as_array)
777 .filter(|c| c.len() == 2)
778 .and_then(|c| Some((c[0].as_str()?.to_owned(), c[1].as_u64()?)))
779 .ok_or("checkpoint is not [string, integer]")?;
780 let shard_count = value
781 .get("shardCount")
782 .and_then(Value::as_u64)
783 .ok_or("shardCount is not an integer")?;
784 let shards = value
785 .get("shards")
786 .and_then(Value::as_array)
787 .ok_or("shards is not an array")?
788 .iter()
789 .map(|s| IndexShardMeta {
790 nsg: s.get("nsg").and_then(Value::as_bool).unwrap_or(false),
791 ivf: s.get("ivf").and_then(Value::as_bool).unwrap_or(false),
792 nsg_crc: s
793 .get("nsgCrc")
794 .and_then(Value::as_u64)
795 .and_then(|c| u32::try_from(c).ok()),
796 ivf_crc: s
797 .get("ivfCrc")
798 .and_then(Value::as_u64)
799 .and_then(|c| u32::try_from(c).ok()),
800 })
801 .collect();
802 Ok(IndicesMeta {
803 checkpoint,
804 shard_count,
805 shards,
806 })
807}
808
809pub fn probe(bytes: &[u8]) -> usize {
812 let mut touched = 0usize;
813 let mut replay = Replay::default();
814 for is_final in [true, false] {
815 let scan = scan_segment(bytes, is_final, |_, payload| {
816 match classify_payload(payload) {
817 PayloadKind::Transaction(json) => {
818 touched += check_transaction(json, None, &mut replay).unwrap_or(0);
819 }
820 _ => touched += 1,
821 }
822 });
823 touched += scan.frames.len();
824 }
825 if let PayloadKind::Transaction(json) = classify_payload(bytes) {
826 touched += check_transaction(json, Some(1024), &mut replay).unwrap_or(0);
827 }
828 touched += usize::from(parse_text_header(bytes) != TextHeader::Headerless);
829 touched += scan_audit(bytes).entries;
830 touched += parse_indices(bytes).map_or(0, |m| m.shards.len());
831 touched += serde_json::from_slice::<super::super::storage::FormatManifest>(bytes)
832 .map_or(0, |m| m.version as usize);
833 #[cfg(feature = "experimental")]
834 {
835 touched += usize::from(crate::core::qgraph::ivf::QuantizedIvf::from_bytes(bytes).is_ok());
836 }
837 touched
838}
839
840#[cfg(test)]
841mod tests {
842 use super::*;
843
844 #[test]
845 fn last_nonzero_spans_padding_blocks() {
846 for len in [0, 1, 4095, 4096, 4097, 8191, 8192, 8193] {
847 let mut data = vec![0; len];
848 assert_eq!(last_nonzero(&data), None);
849 for at in [0, 4095, 4096, 8191, 8192]
850 .into_iter()
851 .filter(|&at| at < len)
852 {
853 data[at] = 1;
854 assert_eq!(last_nonzero(&data), Some(at));
855 data[at] = 0;
856 }
857 }
858 }
859
860 fn frame(payload: &[u8], seq: u32) -> Vec<u8> {
861 let compressed = lz4_flex::compress(payload);
862 let (body, comp_len) = if compressed.len() < payload.len() {
863 (compressed.clone(), compressed.len())
864 } else {
865 (payload.to_vec(), payload.len())
866 };
867 let mut out = Vec::new();
868 out.extend_from_slice(&crc32fast::hash(payload).to_le_bytes());
869 out.extend_from_slice(&(payload.len() as u32).to_le_bytes());
870 out.extend_from_slice(&(comp_len as u32).to_le_bytes());
871 out.extend_from_slice(&seq.to_le_bytes());
872 out.extend_from_slice(&body);
873 out
874 }
875
876 fn marker() -> [u8; 16] {
877 let mut m = [0u8; 16];
878 m[12..].copy_from_slice(&END_MARK.to_le_bytes());
879 m
880 }
881
882 #[test]
883 fn walks_frames_and_classifies_the_end() {
884 let mut segment = frame(&[b'a'; 300], 0);
885 segment.extend(frame(b"short", 1));
886 segment.extend(marker());
887 segment.resize(4096, 0);
888 let scan = scan_segment(&segment, true, |_, _| {});
889 assert_eq!(scan.frames.len(), 2);
890 assert_eq!(scan.end, SegmentEnd::Marker);
891 assert_eq!(scan.frames[1].end(), scan.end_offset);
892 assert!(scan.frames[0].comp_len < scan.frames[0].raw_len);
893
894 segment[scan.end_offset..scan.end_offset + 16].fill(0);
895 assert_eq!(
896 scan_segment(&segment, true, |_, _| {}).end,
897 SegmentEnd::ZeroHeader
898 );
899 }
900
901 #[test]
904 fn a_log_end_before_committed_frames_is_corruption() {
905 for end in [[0u8; 16], marker()] {
906 let mut segment = frame(b"first", 0);
907 segment.extend(end);
908 let hidden = segment.len();
909 segment.extend(frame(b"second", 1));
910 segment.extend(marker());
911 segment.resize(4096, 0);
912 for is_final in [true, false] {
913 let scan = scan_segment(&segment, is_final, |_, _| {});
914 assert_eq!(scan.frames.len(), 1);
915 assert!(
916 matches!(scan.end, SegmentEnd::Corrupt { committed_at: Some(at), .. } if at == hidden),
917 "{:?}",
918 scan.end
919 );
920 }
921 }
922 }
923
924 #[test]
927 fn a_frame_must_decode_to_its_declared_length() {
928 let mut segment = frame(&[b'a'; 1024], 0);
929 assert!(word(&segment, 8) < word(&segment, 4));
930 segment[4..8].copy_from_slice(&1025u32.to_le_bytes());
931 assert!(decode_frame(&segment, 0).is_err());
932 }
933
934 #[test]
937 fn tells_a_torn_tail_from_corruption() {
938 let first = frame(b"first frame payload", 0);
939 let second = frame(b"second frame payload", 1);
940 let mut log = first.clone();
941 log.extend(&second);
942 log.extend(marker());
943 log.resize(2048, 0);
944
945 let mut torn = log.clone();
946 let cut = first.len() + 20;
947 torn[cut..].fill(0);
948 let scan = scan_segment(&torn, true, |_, _| {});
949 assert!(
950 matches!(scan.end, SegmentEnd::TornTail { .. }),
951 "{:?}",
952 scan.end
953 );
954 assert_eq!(scan.end_offset, first.len());
955 assert!(matches!(
956 scan_segment(&torn, false, |_, _| {}).end,
957 SegmentEnd::Corrupt { .. }
958 ));
959
960 let mut flipped = log.clone();
961 flipped[16 + 3] ^= 0xff;
962 let scan = scan_segment(&flipped, true, |_, _| {});
963 assert!(matches!(
964 scan.end,
965 SegmentEnd::Corrupt {
966 committed_at: Some(_),
967 ..
968 }
969 ));
970 assert!(scan.frames.is_empty());
971
972 let mut last = log;
973 last[first.len() + 16 + 2] ^= 0xff;
974 assert!(matches!(
975 scan_segment(&last, true, |_, _| {}).end,
976 SegmentEnd::Corrupt {
977 committed_at: Some(_),
978 ..
979 }
980 ));
981 }
982
983 #[test]
984 fn rejects_hostile_lengths_without_allocating_them() {
985 let mut data = vec![0u8; 64];
986 data[4..8].copy_from_slice(&(MAX_FRAME_PAYLOAD as u32).to_le_bytes());
987 data[8..12].copy_from_slice(&16u32.to_le_bytes());
988 assert_eq!(decode_frame(&data, 0).unwrap_err(), "implausible-ratio");
989 data[4..8].copy_from_slice(&u32::MAX.to_le_bytes());
990 assert_eq!(decode_frame(&data, 0).unwrap_err(), "header-lengths");
991 data[4..8].copy_from_slice(&8u32.to_le_bytes());
992 data[8..12].copy_from_slice(&u32::MAX.to_le_bytes());
993 assert_eq!(decode_frame(&data, 0).unwrap_err(), "header-lengths");
994 }
995
996 #[test]
997 fn classifies_journal_versions() {
998 assert!(matches!(
999 classify_payload(b"HMS-TXN-2\0[]"),
1000 PayloadKind::Transaction(b"[]")
1001 ));
1002 assert_eq!(
1003 classify_payload(b"HMS-TXN-3\0[]"),
1004 PayloadKind::BinaryTransaction(b"[]")
1005 );
1006 assert_eq!(
1007 classify_payload(b"HMS-TXN-4\0"),
1008 PayloadKind::OtherTransactionVersion("4".into())
1009 );
1010 assert_eq!(classify_payload(b"HMS-TXN-\0"), PayloadKind::Unrecognized);
1011 assert_eq!(classify_payload(b"HMS-TXN-2"), PayloadKind::Unrecognized);
1012 }
1013
1014 #[test]
1015 fn replays_the_live_state() {
1016 let json = br#"[
1017 {"Document":{"id":"a","version":"mem1:{}","metadata":{},"chunks":[
1018 {"id":"hms:chunk:a:0","start_byte":0,"end_byte":3,"text":"abc",
1019 "vector":{"dim":8,"indices":[1,4]}}]}},
1020 {"Vector":{"id":"v","vector":{"dim":8,"indices":[]}}},
1021 {"Document":{"id":"b","version":"1","metadata":{"_mem":{}},"chunks":[
1022 {"id":"hms:chunk:b:0","start_byte":0,"end_byte":1,"vector":{"dim":8,"indices":[2]}}]}},
1023 {"DeleteDocument":{"id":"b"}}
1024 ]"#;
1025 let mut replay = Replay::default();
1026 assert_eq!(check_transaction(json, Some(8), &mut replay), Ok(4));
1027 assert_eq!(replay.documents.len(), 1);
1028 assert!(replay.documents["a"].stamped);
1029 assert_eq!(replay.vector_total(), 2);
1030 }
1031
1032 #[test]
1033 fn scoped_meaning_mutations_preserve_raw_vectors_and_require_schema_five() {
1034 let json = br#"[
1035 {"Vector":{"id":"alias","vector":{"dim":8,"indices":[1]}}},
1036 {"DeleteAtom":{"id":"alias"}},
1037 {"DeleteComposite":{"id":"alias"}},
1038 {"RefreshMeaningIndices":{}}
1039 ]"#;
1040 let mut replay = Replay::default();
1041 assert_eq!(
1042 check_transaction_with_schema(json, Some(8), 5, &mut replay),
1043 Ok(4)
1044 );
1045 assert_eq!(replay.vector_total(), 1);
1046 for schema in 2..5 {
1047 let mut rejected = Replay::default();
1048 assert!(check_transaction_with_schema(json, Some(8), schema, &mut rejected).is_err());
1049 assert_eq!(rejected.mutations, 0);
1050 }
1051 for invalid in [
1052 br#"[{"DeleteAtom":{"id":""}}]"#.as_slice(),
1053 br#"[{"DeleteComposite":{"id":"a\u0000b"}}]"#,
1054 br#"[{"RefreshMeaningIndices":{"unexpected":true}}]"#,
1055 ] {
1056 assert!(check_transaction(invalid, Some(8), &mut Replay::default()).is_err());
1057 }
1058 }
1059
1060 #[test]
1061 fn rejects_what_replay_rejects() {
1062 let bad = |json: &str| check_transaction(json.as_bytes(), Some(8), &mut Replay::default());
1063 assert!(bad("[]").is_err());
1064 assert!(bad(r#"[{"Mystery":{}}]"#).is_err());
1065 assert!(bad(r#"[{"Delete":{"id":""}}]"#).is_err());
1066 assert!(bad(r#"[{"Delete":{"id":"a\u0000b"}}]"#).is_err());
1067 assert!(bad(r#"[{"Vector":{"id":"v","vector":{"dim":8,"indices":[3,3]}}}]"#).is_err());
1068 assert!(bad(r#"[{"Vector":{"id":"v","vector":{"dim":8,"indices":[8]}}}]"#).is_err());
1069 assert!(bad(r#"[{"Vector":{"id":"v","vector":{"dim":9,"indices":[1]}}}]"#).is_err());
1070 assert!(bad(r#"[{"Vector":{"id":"v","vector":{"dim":8,"indices":[7]}}}]"#).is_ok());
1071 }
1072
1073 #[test]
1074 fn whole_document_chunks_reach_the_document_limit() {
1075 let document = |chunks: &str| {
1076 format!(r#"[{{"Document":{{"id":"w","version":"1","chunks":[{chunks}]}}}}]"#)
1077 };
1078 let chunk = |start: u64, end: u64, state: &str, version: u8| {
1079 format!(
1080 r#"{{"id":"hms:chunk:w:{start}","start_byte":{start},"end_byte":{end},"vector":{{"dim":8,"indices":[]}},"vector_state":"{state}","terms":{{"$hms_positions":{{"version":{version},"offsets":[0,1],"data":[0]}},"red":1}}}}"#
1081 )
1082 };
1083 let check = |json: String| {
1084 check_transaction(json.as_bytes(), Some(8), &mut Replay::default()).is_ok()
1085 };
1086 let over = MAX_CHUNK_BYTES + 1;
1087 assert!(check(document(&chunk(0, over, "lexical", 2))));
1088 assert!(check(document(&chunk(0, MAX_DOCUMENT_BYTES, "lexical", 2))));
1089 assert!(!check(document(&chunk(
1090 0,
1091 MAX_DOCUMENT_BYTES + 1,
1092 "lexical",
1093 2
1094 ))));
1095 assert!(!check(document(&chunk(1, over, "lexical", 2))));
1096 assert!(!check(document(&chunk(0, over, "pending", 2))));
1097 assert!(!check(document(&chunk(0, over, "lexical", 1))));
1098 assert!(check(document(&chunk(0, MAX_CHUNK_BYTES, "lexical", 1))));
1099 let pair = format!(
1100 "{},{}",
1101 chunk(0, 2, "lexical", 2),
1102 chunk(2, 3, "lexical", 2)
1103 );
1104 assert!(!check(document(&pair)));
1105 }
1106
1107 #[test]
1108 fn probe_survives_garbage() {
1109 let mut state = 0x9e37_79b9_7f4a_7c15u64;
1110 for len in [0usize, 1, 15, 16, 17, 100, 5000] {
1111 let bytes: Vec<u8> = (0..len)
1112 .map(|_| {
1113 state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
1114 (state >> 33) as u8
1115 })
1116 .collect();
1117 probe(&bytes);
1118 }
1119 }
1120}