nodedb_array/segment/
reader.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! `SegmentReader` — zero-copy view over segment bytes.
4//!
5//! Designed to wrap an `mmap`'d file in production. The reader only
6//! borrows the slice and lazily decodes tile payloads on demand; the
7//! footer is decoded eagerly at construction so subsequent tile lookups
8//! are bounded reads + zerompk decode (no rescan).
9
10use super::format::{
11    SegmentFooter, SegmentHeader, TileEntry, TileKind, framing::BlockFraming, header::HEADER_SIZE,
12};
13use crate::codec::tag::peek_tag;
14use crate::codec::tile_decode::decode_sparse_tile;
15use crate::error::{ArrayError, ArrayResult};
16use crate::tile::cell_payload::{CELL_GDPR_ERASURE_SENTINEL, CELL_TOMBSTONE_SENTINEL, CellPayload};
17use crate::tile::dense_tile::DenseTile;
18use crate::tile::sparse_tile::{RowKind, SparseTile};
19use crate::types::coord::value::CoordValue;
20
21/// Extract the encoded `CellPayload` bytes for a specific `coord` from a
22/// `SparseTile`.
23///
24/// Returns `None` when the coord is not present in the tile. Encodes the
25/// matched row's attrs + valid-time bounds back into the `CellPayload` wire
26/// format so the bytes can be fed directly to
27/// [`crate::query::ceiling::ceiling_resolve_cell`].
28///
29/// The returned `Vec<u8>` is always freshly allocated (no lifetime coupling to
30/// the tile).
31pub fn extract_cell_bytes(tile: &SparseTile, coord: &[CoordValue]) -> ArrayResult<Option<Vec<u8>>> {
32    let n = tile.row_count();
33    'rows: for row in 0..n {
34        // Reconstruct the coordinate for this row.
35        for (dim_idx, c) in coord.iter().enumerate() {
36            let dict =
37                tile.dim_dicts
38                    .get(dim_idx)
39                    .ok_or_else(|| ArrayError::SegmentCorruption {
40                        detail: format!("extract_cell_bytes: dim {dim_idx} out of range"),
41                    })?;
42            let entry_idx = *dict
43                .indices
44                .get(row)
45                .ok_or_else(|| ArrayError::SegmentCorruption {
46                    detail: format!("extract_cell_bytes: row {row} index out of range"),
47                })? as usize;
48            let stored =
49                dict.values
50                    .get(entry_idx)
51                    .ok_or_else(|| ArrayError::SegmentCorruption {
52                        detail: format!("extract_cell_bytes: dict entry {entry_idx} out of range"),
53                    })?;
54            if stored != c {
55                continue 'rows;
56            }
57        }
58        // Coord matched — check row kind before decoding payload.
59        let kind = tile.row_kind(row)?;
60        match kind {
61            RowKind::Tombstone => return Ok(Some(CELL_TOMBSTONE_SENTINEL.to_vec())),
62            RowKind::GdprErased => return Ok(Some(CELL_GDPR_ERASURE_SENTINEL.to_vec())),
63            RowKind::Live => {}
64        }
65        // Attribute columns hold live rows only: index by the live-row index.
66        let mut live_row = 0usize;
67        for earlier in 0..row {
68            if tile.row_kind(earlier)? == RowKind::Live {
69                live_row += 1;
70            }
71        }
72        let attrs: Vec<_> = tile
73            .attr_cols
74            .iter()
75            .map(|col| {
76                col.get(live_row)
77                    .cloned()
78                    .ok_or_else(|| ArrayError::SegmentCorruption {
79                        detail: format!(
80                            "extract_cell_bytes: attr col live row {live_row} out of range"
81                        ),
82                    })
83            })
84            .collect::<ArrayResult<Vec<_>>>()?;
85        let surrogate = tile.live_surrogate(row)?;
86        let valid_from_ms =
87            tile.valid_from_ms
88                .get(row)
89                .copied()
90                .ok_or_else(|| ArrayError::SegmentCorruption {
91                    detail: format!("extract_cell_bytes: valid_from_ms row {row} out of range"),
92                })?;
93        let valid_until_ms =
94            tile.valid_until_ms
95                .get(row)
96                .copied()
97                .ok_or_else(|| ArrayError::SegmentCorruption {
98                    detail: format!("extract_cell_bytes: valid_until_ms row {row} out of range"),
99                })?;
100        let payload = CellPayload {
101            valid_from_ms,
102            valid_until_ms,
103            attrs,
104            surrogate,
105        };
106        return payload.encode().map(Some);
107    }
108    Ok(None)
109}
110
111/// Dispatch sparse tile decoding by peeking the tag byte.
112///
113/// - Recognized tag (0 or 1): delegate to `decode_sparse_tile`.
114/// - Empty payload or unrecognized byte: return a corruption error.
115fn read_sparse_tile(payload: &[u8]) -> ArrayResult<SparseTile> {
116    match peek_tag(payload) {
117        Some(_) => decode_sparse_tile(payload),
118        None => Err(ArrayError::SegmentCorruption {
119            detail: format!(
120                "sparse tile has unrecognized tag byte {:#04x}",
121                payload.first().copied().unwrap_or(0)
122            ),
123        }),
124    }
125}
126
127/// Decoded tile payload.
128#[derive(Debug, Clone, PartialEq)]
129pub enum TilePayload {
130    Sparse(SparseTile),
131    Dense(DenseTile),
132}
133
134pub struct SegmentReader<'a> {
135    bytes: &'a [u8],
136    header: SegmentHeader,
137    footer: SegmentFooter,
138}
139
140impl<'a> SegmentReader<'a> {
141    pub fn open(bytes: &'a [u8]) -> ArrayResult<Self> {
142        if bytes.len() < HEADER_SIZE {
143            return Err(ArrayError::SegmentCorruption {
144                detail: format!("segment too small: {} bytes", bytes.len()),
145            });
146        }
147        let header = SegmentHeader::decode(&bytes[..HEADER_SIZE])?;
148        let footer = SegmentFooter::decode(bytes)?;
149        if header.schema_hash != footer.schema_hash {
150            return Err(ArrayError::SegmentCorruption {
151                detail: format!(
152                    "header/footer schema_hash mismatch: header={:x} footer={:x}",
153                    header.schema_hash, footer.schema_hash
154                ),
155            });
156        }
157        Ok(Self {
158            bytes,
159            header,
160            footer,
161        })
162    }
163
164    pub fn header(&self) -> &SegmentHeader {
165        &self.header
166    }
167
168    pub fn schema_hash(&self) -> u64 {
169        self.header.schema_hash
170    }
171
172    pub fn tiles(&self) -> &[TileEntry] {
173        &self.footer.tiles
174    }
175
176    pub fn tile_count(&self) -> usize {
177        self.footer.tiles.len()
178    }
179
180    /// Reverse-scan tile entries with the given `hilbert_prefix` and
181    /// `system_from_ms <= system_as_of`, returning the newest qualifying
182    /// tile version.
183    ///
184    /// Returns `Ok(None)` if no version exists at or before the cutoff.
185    ///
186    /// The `valid_at_ms` parameter is reserved for the query layer (Tier
187    /// 9.2). At the reader level, valid-time filtering is NOT applied —
188    /// the whole tile is returned. This keeps the reader cell-shape-agnostic.
189    pub fn read_tile_as_of(
190        &self,
191        hilbert_prefix: u64,
192        system_as_of: i64,
193        _valid_at_ms: Option<i64>,
194    ) -> ArrayResult<Option<TilePayload>> {
195        let tiles = &self.footer.tiles;
196
197        // Binary search for the first entry with hilbert_prefix.
198        let first = tiles.partition_point(|e| e.tile_id.hilbert_prefix < hilbert_prefix);
199        // Binary search for the first entry past the range:
200        // hilbert_prefix matches and system_from_ms <= system_as_of.
201        // Upper bound: first entry where prefix > hilbert_prefix.
202        let past_prefix = tiles.partition_point(|e| e.tile_id.hilbert_prefix <= hilbert_prefix);
203
204        // Slice of entries with matching prefix.
205        let candidates = &tiles[first..past_prefix];
206        if candidates.is_empty() {
207            return Ok(None);
208        }
209
210        // Within the prefix slice, entries are ordered by system_from_ms ascending.
211        // Find the rightmost entry with system_from_ms <= system_as_of.
212        let cutoff_pos = candidates.partition_point(|e| e.tile_id.system_from_ms <= system_as_of);
213        if cutoff_pos == 0 {
214            return Ok(None);
215        }
216
217        // The entry at cutoff_pos - 1 is the newest qualifying version.
218        let entry_idx = first + cutoff_pos - 1;
219        self.read_tile(entry_idx).map(Some)
220    }
221
222    /// Returns an iterator over all tile versions for `hilbert_prefix` whose
223    /// `system_from_ms <= system_as_of`, ordered **newest-first** by
224    /// `system_from_ms`.
225    ///
226    /// Callers supply this to [`nodedb_array::query::ceiling`] to resolve the
227    /// bitemporal ceiling for every coordinate in the prefix.
228    pub fn iter_tile_versions(
229        &self,
230        hilbert_prefix: u64,
231        system_as_of: i64,
232    ) -> ArrayResult<impl Iterator<Item = ArrayResult<(crate::types::TileId, TilePayload)>> + '_>
233    {
234        let tiles = &self.footer.tiles;
235
236        // Find the contiguous slice of entries whose hilbert_prefix matches.
237        let first = tiles.partition_point(|e| e.tile_id.hilbert_prefix < hilbert_prefix);
238        let past_prefix = tiles.partition_point(|e| e.tile_id.hilbert_prefix <= hilbert_prefix);
239
240        // Within [first..past_prefix], entries are ascending by system_from_ms.
241        // Restrict to those at or before the cutoff.
242        let cutoff_pos =
243            tiles[first..past_prefix].partition_point(|e| e.tile_id.system_from_ms <= system_as_of);
244        // Global index range: [first .. first+cutoff_pos).
245        let qualifying_start = first;
246        let qualifying_end = first + cutoff_pos;
247
248        // Iterate in reverse (newest-first).
249        let indices: Vec<usize> = (qualifying_start..qualifying_end).rev().collect();
250        Ok(indices.into_iter().map(move |idx| {
251            let tile_id = tiles[idx].tile_id;
252            self.read_tile(idx).map(|payload| (tile_id, payload))
253        }))
254    }
255
256    /// Decode tile #`idx`. CRC is checked by the framing layer.
257    pub fn read_tile(&self, idx: usize) -> ArrayResult<TilePayload> {
258        let entry = self
259            .footer
260            .tiles
261            .get(idx)
262            .ok_or_else(|| ArrayError::SegmentCorruption {
263                detail: format!(
264                    "tile index {idx} out of range (have {})",
265                    self.footer.tiles.len()
266                ),
267            })?;
268        let off = entry.offset as usize;
269        let len = entry.length as usize;
270        let end = off
271            .checked_add(len)
272            .ok_or_else(|| ArrayError::SegmentCorruption {
273                detail: "tile entry offset+length overflows".into(),
274            })?;
275        if end > self.bytes.len() {
276            return Err(ArrayError::SegmentCorruption {
277                detail: format!(
278                    "tile {idx} block out of bounds: off={off} len={len} \
279                     file_size={}",
280                    self.bytes.len()
281                ),
282            });
283        }
284        let (payload, _) = BlockFraming::decode(&self.bytes[off..end])?;
285        match entry.kind {
286            TileKind::Sparse => {
287                let t = read_sparse_tile(payload)?;
288                Ok(TilePayload::Sparse(t))
289            }
290            TileKind::Dense => {
291                let t: DenseTile =
292                    zerompk::from_msgpack(payload).map_err(|e| ArrayError::SegmentCorruption {
293                        detail: format!("dense tile decode failed: {e}"),
294                    })?;
295                Ok(TilePayload::Dense(t))
296            }
297        }
298    }
299}
300
301/// An owned segment reader that holds decrypted segment bytes.
302///
303/// Exists because `SegmentReader<'a>` borrows its byte slice; for
304/// encrypted segments the plaintext lives in a `Vec<u8>` that must
305/// outlive the reader. `OwnedSegmentReader` ties the two together.
306#[derive(Debug)]
307pub struct OwnedSegmentReader {
308    /// Decrypted plaintext segment bytes.
309    plaintext: Vec<u8>,
310    header: SegmentHeader,
311    footer: SegmentFooter,
312}
313
314impl OwnedSegmentReader {
315    /// Open a segment with optional at-rest decryption.
316    ///
317    /// - `kek = None` → requires a plaintext (`NDAS`) segment; returns
318    ///   `Err(MissingKek)` if the blob starts with `SEGA`.
319    /// - `kek = Some(key)` → requires an encrypted (`SEGA`) segment; decrypts
320    ///   the blob, then parses the inner plaintext. Returns `Err(KekRequired)`
321    ///   if the blob starts with `NDAS`.
322    pub fn open_with_kek(
323        blob: &[u8],
324        kek: Option<&nodedb_wal::crypto::WalEncryptionKey>,
325    ) -> ArrayResult<Self> {
326        use super::encrypt::{decrypt_segment, detect_encryption};
327        let is_encrypted = detect_encryption(blob)?;
328        let plaintext = match (is_encrypted, kek) {
329            (true, Some(key)) => decrypt_segment(key, blob)?,
330            (true, None) => return Err(ArrayError::MissingKek),
331            (false, Some(_)) => return Err(ArrayError::KekRequired),
332            (false, None) => blob.to_vec(),
333        };
334        let header = SegmentHeader::decode(&plaintext[..HEADER_SIZE.min(plaintext.len())])?;
335        let footer = SegmentFooter::decode(&plaintext)?;
336        if header.schema_hash != footer.schema_hash {
337            return Err(ArrayError::SegmentCorruption {
338                detail: format!(
339                    "header/footer schema_hash mismatch: header={:x} footer={:x}",
340                    header.schema_hash, footer.schema_hash
341                ),
342            });
343        }
344        Ok(Self {
345            plaintext,
346            header,
347            footer,
348        })
349    }
350
351    /// Borrow a `SegmentReader` over the owned plaintext bytes.
352    pub fn reader(&self) -> SegmentReader<'_> {
353        SegmentReader {
354            bytes: &self.plaintext,
355            header: self.header,
356            footer: self.footer.clone(),
357        }
358    }
359
360    /// Consume the owned reader and return the inner plaintext buffer.
361    pub fn into_plaintext(self) -> Vec<u8> {
362        self.plaintext
363    }
364}
365
366impl<'a> SegmentReader<'a> {
367    /// Validate a segment blob with optional KEK, returning `Err` on mismatch.
368    ///
369    /// - `kek = None` + encrypted blob → `Err(MissingKek)`
370    /// - `kek = Some` + plaintext blob → `Err(KekRequired)`
371    ///
372    /// On success the segment is fully parsed but the owned bytes are
373    /// discarded. Use [`OwnedSegmentReader`] when you need to keep the reader.
374    pub fn open_with_kek(
375        blob: &[u8],
376        kek: Option<&nodedb_wal::crypto::WalEncryptionKey>,
377    ) -> ArrayResult<OwnedSegmentReader> {
378        OwnedSegmentReader::open_with_kek(blob, kek)
379    }
380}
381
382#[cfg(test)]
383mod tests {
384    use super::*;
385    use crate::schema::ArraySchemaBuilder;
386    use crate::schema::attr_spec::{AttrSpec, AttrType};
387    use crate::schema::dim_spec::{DimSpec, DimType};
388    use crate::segment::writer::SegmentWriter;
389    use crate::tile::dense_tile::DenseTile;
390    use crate::tile::sparse_tile::SparseTileBuilder;
391    use crate::types::TileId;
392    use crate::types::cell_value::value::CellValue;
393    use crate::types::coord::value::CoordValue;
394    use crate::types::domain::{Domain, DomainBound};
395
396    fn schema() -> crate::schema::ArraySchema {
397        ArraySchemaBuilder::new("g")
398            .dim(DimSpec::new(
399                "x",
400                DimType::Int64,
401                Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)),
402            ))
403            .dim(DimSpec::new(
404                "y",
405                DimType::Int64,
406                Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)),
407            ))
408            .attr(AttrSpec::new("v", AttrType::Int64, true))
409            .tile_extents(vec![4, 4])
410            .build()
411            .unwrap()
412    }
413
414    fn make_sparse(s: &crate::schema::ArraySchema, base: i64) -> SparseTile {
415        let mut b = SparseTileBuilder::new(s);
416        b.push_row(crate::tile::sparse_tile::SparseRow::live(
417            &[CoordValue::Int64(base), CoordValue::Int64(base + 1)],
418            &[CellValue::Int64(base * 10)],
419            nodedb_types::Surrogate::new(base as u32 + 1),
420            0,
421            nodedb_types::OPEN_UPPER,
422        ))
423        .unwrap();
424        b.build()
425    }
426
427    #[test]
428    fn reader_round_trips_sparse_tiles() {
429        let s = schema();
430        let mut w = SegmentWriter::new(0xCAFE);
431        w.append_sparse(TileId::snapshot(1), &make_sparse(&s, 1))
432            .unwrap();
433        w.append_sparse(TileId::snapshot(2), &make_sparse(&s, 2))
434            .unwrap();
435        let bytes = w.finish(None).unwrap();
436        let r = SegmentReader::open(&bytes).unwrap();
437        assert_eq!(r.tile_count(), 2);
438        let t0 = r.read_tile(0).unwrap();
439        match t0 {
440            TilePayload::Sparse(t) => assert_eq!(t.nnz(), 1),
441            _ => panic!("expected sparse"),
442        }
443    }
444
445    #[test]
446    fn reader_round_trips_dense_tile() {
447        let s = schema();
448        let mut w = SegmentWriter::new(0xBEEF);
449        w.append_dense(TileId::snapshot(1), &DenseTile::empty(&s))
450            .unwrap();
451        let bytes = w.finish(None).unwrap();
452        let r = SegmentReader::open(&bytes).unwrap();
453        match r.read_tile(0).unwrap() {
454            TilePayload::Dense(t) => assert_eq!(t.cell_count(), 16),
455            _ => panic!("expected dense"),
456        }
457    }
458
459    #[test]
460    fn reader_rejects_mismatched_schema_hash() {
461        // Build a valid segment, then flip a byte in the header
462        // schema_hash and re-CRC manually-detected mismatch by way of
463        // header CRC failure. (We can't cheaply forge a valid header
464        // with a mismatched footer hash, so this exercises the header
465        // CRC path which guards the same invariant.)
466        let s = schema();
467        let mut w = SegmentWriter::new(0x1);
468        w.append_sparse(TileId::snapshot(1), &make_sparse(&s, 1))
469            .unwrap();
470        let mut bytes = w.finish(None).unwrap();
471        bytes[12] ^= 0xFF; // corrupt header schema_hash
472        assert!(SegmentReader::open(&bytes).is_err());
473    }
474
475    #[test]
476    fn reader_rejects_out_of_range_tile() {
477        let s = schema();
478        let mut w = SegmentWriter::new(0x1);
479        w.append_sparse(TileId::snapshot(1), &make_sparse(&s, 1))
480            .unwrap();
481        let bytes = w.finish(None).unwrap();
482        let r = SegmentReader::open(&bytes).unwrap();
483        assert!(r.read_tile(99).is_err());
484    }
485
486    #[test]
487    fn read_tile_as_of_returns_newest_at_cutoff() {
488        let s = schema();
489        let mut w = SegmentWriter::new(0xCAFE);
490        // Three versions of prefix=1 at system_from_ms 100, 200, 300.
491        w.append_sparse(TileId::new(1, 100), &make_sparse(&s, 1))
492            .unwrap();
493        w.append_sparse(TileId::new(1, 200), &make_sparse(&s, 2))
494            .unwrap();
495        w.append_sparse(TileId::new(1, 300), &make_sparse(&s, 3))
496            .unwrap();
497        let bytes = w.finish(None).unwrap();
498        let r = SegmentReader::open(&bytes).unwrap();
499        // Read at cutoff 250 — between v2 (200) and v3 (300). Should return v2.
500        let result = r.read_tile_as_of(1, 250, None).unwrap();
501        match result {
502            Some(TilePayload::Sparse(t)) => {
503                // v2 was built with base=2 — one non-zero entry.
504                assert_eq!(t.nnz(), 1);
505                // dim_dicts[0] holds the x-dim dictionary; values[0] = Int64(2).
506                assert_eq!(t.dim_dicts[0].values[0], CoordValue::Int64(2));
507            }
508            other => panic!("expected Some(Sparse), got {other:?}"),
509        }
510    }
511
512    #[test]
513    fn read_tile_as_of_returns_none_below_first_version() {
514        let s = schema();
515        let mut w = SegmentWriter::new(0xBEEF);
516        w.append_sparse(TileId::new(1, 100), &make_sparse(&s, 1))
517            .unwrap();
518        let bytes = w.finish(None).unwrap();
519        let r = SegmentReader::open(&bytes).unwrap();
520        // Read at cutoff before any version (50 < 100).
521        let result = r.read_tile_as_of(1, 50, None).unwrap();
522        assert!(result.is_none());
523    }
524
525    #[test]
526    fn extract_cell_bytes_finds_coord() {
527        let s = schema();
528        let sparse = make_sparse(&s, 3);
529        // make_sparse builds one entry at coord (base, base+1) = (3, 4).
530        let coord = vec![CoordValue::Int64(3), CoordValue::Int64(4)];
531        let bytes = extract_cell_bytes(&sparse, &coord).unwrap();
532        assert!(bytes.is_some(), "should find coord (3,4)");
533
534        let absent = vec![CoordValue::Int64(9), CoordValue::Int64(9)];
535        let none = extract_cell_bytes(&sparse, &absent).unwrap();
536        assert!(none.is_none(), "absent coord must return None");
537    }
538
539    #[test]
540    fn extract_cell_bytes_carries_valid_time_bounds() {
541        use crate::tile::cell_payload::CellPayload;
542        use crate::tile::sparse_tile::{SparseRow, SparseTileBuilder};
543        use nodedb_types::Surrogate;
544
545        let s = schema();
546        let mut b = SparseTileBuilder::new(&s);
547        b.push_row(SparseRow {
548            coord: &[CoordValue::Int64(1), CoordValue::Int64(2)],
549            attrs: &[CellValue::Int64(99)],
550            surrogate: Some(Surrogate::new(3)),
551            valid_from_ms: 100,
552            valid_until_ms: 200,
553            kind: crate::tile::sparse_tile::RowKind::Live,
554        })
555        .unwrap();
556        let tile = b.build();
557
558        let coord = vec![CoordValue::Int64(1), CoordValue::Int64(2)];
559        let bytes = extract_cell_bytes(&tile, &coord).unwrap().unwrap();
560        let payload = CellPayload::decode(&bytes).unwrap();
561        assert_eq!(payload.valid_from_ms, 100);
562        assert_eq!(payload.valid_until_ms, 200);
563    }
564
565    #[test]
566    fn iter_tile_versions_newest_first() {
567        let s = schema();
568        let mut w = SegmentWriter::new(0xCAFE);
569        w.append_sparse(TileId::new(1, 100), &make_sparse(&s, 1))
570            .unwrap();
571        w.append_sparse(TileId::new(1, 200), &make_sparse(&s, 2))
572            .unwrap();
573        w.append_sparse(TileId::new(1, 300), &make_sparse(&s, 3))
574            .unwrap();
575        let bytes = w.finish(None).unwrap();
576        let r = SegmentReader::open(&bytes).unwrap();
577        let versions: Vec<_> = r
578            .iter_tile_versions(1, i64::MAX)
579            .unwrap()
580            .map(|v| v.unwrap().0.system_from_ms)
581            .collect();
582        assert_eq!(versions, vec![300, 200, 100]);
583    }
584
585    #[test]
586    fn iter_tile_versions_respects_system_as_of() {
587        let s = schema();
588        let mut w = SegmentWriter::new(0xBEEF);
589        w.append_sparse(TileId::new(1, 100), &make_sparse(&s, 1))
590            .unwrap();
591        w.append_sparse(TileId::new(1, 200), &make_sparse(&s, 2))
592            .unwrap();
593        w.append_sparse(TileId::new(1, 300), &make_sparse(&s, 3))
594            .unwrap();
595        let bytes = w.finish(None).unwrap();
596        let r = SegmentReader::open(&bytes).unwrap();
597        let versions: Vec<_> = r
598            .iter_tile_versions(1, 250)
599            .unwrap()
600            .map(|v| v.unwrap().0.system_from_ms)
601            .collect();
602        // Cutoff 250 → v2 (200) and v1 (100); v3 (300) excluded.
603        assert_eq!(versions, vec![200, 100]);
604    }
605
606    #[test]
607    fn extract_cell_bytes_returns_tombstone_sentinel_for_tombstone_row() {
608        use crate::tile::cell_payload::{CELL_TOMBSTONE_SENTINEL, is_cell_tombstone};
609        use crate::tile::sparse_tile::{RowKind, SparseTileBuilder};
610
611        let s = schema();
612        let mut b = SparseTileBuilder::new(&s);
613        b.push_row(crate::tile::sparse_tile::SparseRow {
614            coord: &[CoordValue::Int64(7), CoordValue::Int64(8)],
615            attrs: &[],
616            surrogate: None,
617            valid_from_ms: 0,
618            valid_until_ms: nodedb_types::OPEN_UPPER,
619            kind: RowKind::Tombstone,
620        })
621        .unwrap();
622        let tile = b.build();
623        let coord = vec![CoordValue::Int64(7), CoordValue::Int64(8)];
624        let bytes = extract_cell_bytes(&tile, &coord).unwrap().unwrap();
625        assert!(
626            is_cell_tombstone(&bytes),
627            "expected CELL_TOMBSTONE_SENTINEL, got {bytes:?}"
628        );
629        assert_eq!(bytes, CELL_TOMBSTONE_SENTINEL);
630    }
631
632    /// A live row after a sentinel row reads its own attributes: the
633    /// attribute columns skip sentinel rows.
634    #[test]
635    fn extract_cell_bytes_reads_a_live_row_after_a_sentinel() {
636        use crate::tile::sparse_tile::{RowKind, SparseRow, SparseTileBuilder};
637        use nodedb_types::Surrogate;
638
639        let s = schema();
640        let mut b = SparseTileBuilder::new(&s);
641        b.push_row(SparseRow {
642            coord: &[CoordValue::Int64(1), CoordValue::Int64(1)],
643            attrs: &[],
644            surrogate: None,
645            valid_from_ms: 0,
646            valid_until_ms: nodedb_types::OPEN_UPPER,
647            kind: RowKind::Tombstone,
648        })
649        .unwrap();
650        b.push_row(SparseRow {
651            coord: &[CoordValue::Int64(2), CoordValue::Int64(2)],
652            attrs: &[CellValue::Int64(20)],
653            surrogate: Some(Surrogate::new(9)),
654            valid_from_ms: 5,
655            valid_until_ms: nodedb_types::OPEN_UPPER,
656            kind: RowKind::Live,
657        })
658        .unwrap();
659        let tile = b.build();
660        let coord = vec![CoordValue::Int64(2), CoordValue::Int64(2)];
661        let bytes = extract_cell_bytes(&tile, &coord).unwrap().unwrap();
662        let payload = CellPayload::decode(&bytes).unwrap();
663        assert_eq!(payload.attrs, vec![CellValue::Int64(20)]);
664        assert_eq!(payload.surrogate, Surrogate::new(9));
665    }
666
667    #[test]
668    fn extract_cell_bytes_returns_erasure_sentinel_for_erased_row() {
669        use crate::tile::cell_payload::{CELL_GDPR_ERASURE_SENTINEL, is_cell_gdpr_erasure};
670        use crate::tile::sparse_tile::{RowKind, SparseTileBuilder};
671
672        let s = schema();
673        let mut b = SparseTileBuilder::new(&s);
674        b.push_row(crate::tile::sparse_tile::SparseRow {
675            coord: &[CoordValue::Int64(4), CoordValue::Int64(5)],
676            attrs: &[],
677            surrogate: None,
678            valid_from_ms: 0,
679            valid_until_ms: nodedb_types::OPEN_UPPER,
680            kind: RowKind::GdprErased,
681        })
682        .unwrap();
683        let tile = b.build();
684        let coord = vec![CoordValue::Int64(4), CoordValue::Int64(5)];
685        let bytes = extract_cell_bytes(&tile, &coord).unwrap().unwrap();
686        assert!(
687            is_cell_gdpr_erasure(&bytes),
688            "expected CELL_GDPR_ERASURE_SENTINEL, got {bytes:?}"
689        );
690        assert_eq!(bytes, CELL_GDPR_ERASURE_SENTINEL);
691    }
692
693    #[test]
694    fn read_tile_as_of_finds_exact_match() {
695        let s = schema();
696        let mut w = SegmentWriter::new(0xDEAD);
697        w.append_sparse(TileId::new(1, 100), &make_sparse(&s, 1))
698            .unwrap();
699        w.append_sparse(TileId::new(1, 200), &make_sparse(&s, 2))
700            .unwrap();
701        w.append_sparse(TileId::new(1, 300), &make_sparse(&s, 3))
702            .unwrap();
703        let bytes = w.finish(None).unwrap();
704        let r = SegmentReader::open(&bytes).unwrap();
705        // Read at cutoff exactly equal to v2's system_from_ms.
706        let result = r.read_tile_as_of(1, 200, None).unwrap();
707        match result {
708            Some(TilePayload::Sparse(t)) => {
709                assert_eq!(t.nnz(), 1);
710                assert_eq!(t.dim_dicts[0].values[0], CoordValue::Int64(2));
711            }
712            other => panic!("expected Some(Sparse), got {other:?}"),
713        }
714    }
715
716    fn test_kek() -> nodedb_wal::crypto::WalEncryptionKey {
717        nodedb_wal::crypto::WalEncryptionKey::from_bytes(&[0xA1u8; 32]).unwrap()
718    }
719
720    fn write_plain(s: &crate::schema::ArraySchema, id: TileId) -> Vec<u8> {
721        let mut w = SegmentWriter::new(0xCAFE);
722        w.append_sparse(id, &make_sparse(s, 1)).unwrap();
723        w.finish(None).unwrap()
724    }
725
726    fn write_encrypted(s: &crate::schema::ArraySchema, id: TileId) -> Vec<u8> {
727        let kek = test_kek();
728        let mut w = SegmentWriter::new(0xCAFE);
729        w.append_sparse(id, &make_sparse(s, 1)).unwrap();
730        w.finish(Some(&kek)).unwrap()
731    }
732
733    #[test]
734    fn array_segment_refuses_plaintext_with_kek() {
735        let s = schema();
736        let plain = write_plain(&s, TileId::snapshot(1));
737        let kek = test_kek();
738        let err = OwnedSegmentReader::open_with_kek(&plain, Some(&kek)).unwrap_err();
739        assert!(
740            matches!(err, crate::error::ArrayError::KekRequired),
741            "expected KekRequired, got {err:?}"
742        );
743    }
744
745    #[test]
746    fn array_segment_refuses_encrypted_without_kek() {
747        let s = schema();
748        let encrypted = write_encrypted(&s, TileId::snapshot(1));
749        let err = OwnedSegmentReader::open_with_kek(&encrypted, None).unwrap_err();
750        assert!(
751            matches!(err, crate::error::ArrayError::MissingKek),
752            "expected MissingKek, got {err:?}"
753        );
754    }
755
756    #[test]
757    fn array_segment_tampered_ciphertext_rejected() {
758        let s = schema();
759        let mut encrypted = write_encrypted(&s, TileId::snapshot(1));
760        // Flip a byte in ciphertext after the current envelope preamble.
761        encrypted[nodedb_wal::crypto::SEGMENT_ENVELOPE_PREAMBLE_SIZE + 2] ^= 0xFF;
762        let kek = test_kek();
763        assert!(OwnedSegmentReader::open_with_kek(&encrypted, Some(&kek)).is_err());
764    }
765
766    #[test]
767    fn array_segment_encrypted_at_rest() {
768        let s = schema();
769        let encrypted = write_encrypted(&s, TileId::snapshot(1));
770        // Encrypted blob must start with SEGA.
771        assert_eq!(&encrypted[..4], b"SEGA");
772        let kek = test_kek();
773        let owned = OwnedSegmentReader::open_with_kek(&encrypted, Some(&kek)).unwrap();
774        let reader = owned.reader();
775        assert_eq!(reader.tile_count(), 1);
776    }
777
778    #[test]
779    fn array_segment_handle_decrypts_into_owned_buffer() {
780        let s = schema();
781        let kek = test_kek();
782        let mut w = SegmentWriter::new(0x1234);
783        w.append_sparse(TileId::new(1, 100), &make_sparse(&s, 1))
784            .unwrap();
785        w.append_sparse(TileId::new(2, 200), &make_sparse(&s, 2))
786            .unwrap();
787        let encrypted = w.finish(Some(&kek)).unwrap();
788        let owned = OwnedSegmentReader::open_with_kek(&encrypted, Some(&kek)).unwrap();
789        let reader = owned.reader();
790        assert_eq!(reader.tile_count(), 2);
791        assert_eq!(reader.tiles()[0].tile_id, TileId::new(1, 100));
792    }
793}
nodedb_array/segment/reader.rs

nodedb_array/segment/
reader.rs
