nodedb_array/segment/
writer.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Streaming `SegmentWriter` — appends tile payloads, finalises with footer.
4//!
5//! Layout: `[ HEADER | framed_tile_0 | framed_tile_1 | ... | footer_block | trailer ]`.
6//! Each tile is encoded as zerompk + framed (length + CRC). Tile
7//! offsets and MBRs are recorded in-memory and flushed in the footer.
8
9use super::format::{
10    SegmentFooter, SegmentHeader, TileEntry, TileKind, framing::BlockFraming, header::HEADER_SIZE,
11};
12use crate::error::{ArrayError, ArrayResult};
13use crate::tile::dense_tile::DenseTile;
14use crate::tile::sparse_tile::SparseTile;
15use crate::types::TileId;
16
17pub struct SegmentWriter {
18    schema_hash: u64,
19    buf: Vec<u8>,
20    entries: Vec<TileEntry>,
21    finished: bool,
22}
23
24impl SegmentWriter {
25    pub fn new(schema_hash: u64) -> Self {
26        let mut buf = Vec::new();
27        SegmentHeader::new(schema_hash).encode_to(&mut buf);
28        Self {
29            schema_hash,
30            buf,
31            entries: Vec::new(),
32            finished: false,
33        }
34    }
35
36    pub fn append_sparse(&mut self, tile_id: TileId, tile: &SparseTile) -> ArrayResult<()> {
37        let mut payload = Vec::new();
38        crate::codec::tile_encode::encode_sparse_tile(tile, &mut payload)?;
39        self.append_framed(tile_id, TileKind::Sparse, &payload, tile.mbr.clone())
40    }
41
42    pub fn append_dense(&mut self, tile_id: TileId, tile: &DenseTile) -> ArrayResult<()> {
43        let payload = zerompk::to_msgpack_vec(tile).map_err(|e| ArrayError::SegmentCorruption {
44            detail: format!("dense tile encode failed: {e}"),
45        })?;
46        self.append_framed(tile_id, TileKind::Dense, &payload, tile.mbr.clone())
47    }
48
49    fn append_framed(
50        &mut self,
51        tile_id: TileId,
52        kind: TileKind,
53        payload: &[u8],
54        mbr: crate::tile::mbr::TileMBR,
55    ) -> ArrayResult<()> {
56        if let Some(last) = self.entries.last()
57            && tile_id <= last.tile_id
58        {
59            return Err(ArrayError::SegmentCorruption {
60                detail: format!(
61                    "tile IDs not strictly monotonic ascending: \
62                     prev=({},{}) new=({},{})",
63                    last.tile_id.hilbert_prefix,
64                    last.tile_id.system_from_ms,
65                    tile_id.hilbert_prefix,
66                    tile_id.system_from_ms,
67                ),
68            });
69        }
70        let offset = self.buf.len() as u64;
71        let framed_len = BlockFraming::encode(payload, &mut self.buf);
72        self.entries.push(TileEntry::new(
73            tile_id,
74            kind,
75            offset,
76            framed_len as u32,
77            mbr,
78        ));
79        Ok(())
80    }
81
82    pub fn tile_count(&self) -> usize {
83        self.entries.len()
84    }
85
86    /// Finalise the segment by appending the footer + trailer. Consumes
87    /// the writer and returns the complete byte vector.
88    ///
89    /// When `kek` is `Some`, the assembled plaintext buffer is wrapped in an
90    /// AES-256-GCM `SEGA` envelope before being returned. When `None`, the
91    /// raw `NDAS` segment bytes are returned.
92    pub fn finish(
93        mut self,
94        kek: Option<&nodedb_wal::crypto::WalEncryptionKey>,
95    ) -> ArrayResult<Vec<u8>> {
96        if self.finished {
97            return Err(ArrayError::SegmentCorruption {
98                detail: "SegmentWriter::finish called twice".into(),
99            });
100        }
101        let footer = SegmentFooter::new(self.schema_hash, std::mem::take(&mut self.entries));
102        footer.encode_to(&mut self.buf)?;
103        self.finished = true;
104
105        if let Some(key) = kek {
106            return super::encrypt::encrypt_segment(key, &self.buf);
107        }
108
109        Ok(self.buf)
110    }
111
112    pub fn header_size(&self) -> usize {
113        HEADER_SIZE
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use crate::schema::ArraySchemaBuilder;
121    use crate::schema::attr_spec::{AttrSpec, AttrType};
122    use crate::schema::dim_spec::{DimSpec, DimType};
123    use crate::tile::sparse_tile::{SparseRow, SparseTileBuilder};
124    use crate::types::cell_value::value::CellValue;
125    use crate::types::coord::value::CoordValue;
126    use crate::types::domain::{Domain, DomainBound};
127
128    /// Finish a writer without encryption — convenience for existing tests.
129    fn finish_plain(w: SegmentWriter) -> crate::error::ArrayResult<Vec<u8>> {
130        w.finish(None)
131    }
132
133    fn schema() -> crate::schema::ArraySchema {
134        ArraySchemaBuilder::new("g")
135            .dim(DimSpec::new(
136                "x",
137                DimType::Int64,
138                Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)),
139            ))
140            .dim(DimSpec::new(
141                "y",
142                DimType::Int64,
143                Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)),
144            ))
145            .attr(AttrSpec::new("v", AttrType::Int64, true))
146            .tile_extents(vec![4, 4])
147            .build()
148            .unwrap()
149    }
150
151    fn sparse_tile(s: &crate::schema::ArraySchema) -> SparseTile {
152        let mut b = SparseTileBuilder::new(s);
153        b.push_row(SparseRow::live(
154            &[CoordValue::Int64(1), CoordValue::Int64(2)],
155            &[CellValue::Int64(10)],
156            nodedb_types::Surrogate::new(1),
157            0,
158            nodedb_types::OPEN_UPPER,
159        ))
160        .unwrap();
161        b.push_row(SparseRow::live(
162            &[CoordValue::Int64(3), CoordValue::Int64(0)],
163            &[CellValue::Int64(20)],
164            nodedb_types::Surrogate::new(2),
165            0,
166            nodedb_types::OPEN_UPPER,
167        ))
168        .unwrap();
169        b.build()
170    }
171
172    #[test]
173    fn writer_round_trip_via_footer() {
174        let s = schema();
175        let mut w = SegmentWriter::new(0x1234);
176        w.append_sparse(TileId::snapshot(1), &sparse_tile(&s))
177            .unwrap();
178        w.append_sparse(TileId::snapshot(2), &sparse_tile(&s))
179            .unwrap();
180        let bytes = finish_plain(w).unwrap();
181        let footer = SegmentFooter::decode(&bytes).unwrap();
182        assert_eq!(footer.schema_hash, 0x1234);
183        assert_eq!(footer.tiles.len(), 2);
184        assert_eq!(footer.tiles[0].tile_id.hilbert_prefix, 1);
185        assert_eq!(footer.tiles[1].tile_id.hilbert_prefix, 2);
186    }
187
188    #[test]
189    fn writer_emits_header_first() {
190        let w = SegmentWriter::new(0xAA);
191        let bytes = finish_plain(w).unwrap();
192        let h = SegmentHeader::decode(&bytes).unwrap();
193        assert_eq!(h.schema_hash, 0xAA);
194    }
195
196    #[test]
197    fn writer_rejects_non_monotonic_tile_ids() {
198        let s = schema();
199        let mut w = SegmentWriter::new(0x1234);
200        w.append_sparse(TileId::new(2, 100), &sparse_tile(&s))
201            .unwrap();
202        let err = w
203            .append_sparse(TileId::new(1, 100), &sparse_tile(&s))
204            .unwrap_err();
205        assert!(
206            matches!(err, crate::error::ArrayError::SegmentCorruption { .. }),
207            "expected SegmentCorruption, got {err:?}"
208        );
209    }
210
211    #[test]
212    fn writer_allows_same_prefix_different_system_times() {
213        let s = schema();
214        let mut w = SegmentWriter::new(0x1234);
215        w.append_sparse(TileId::new(5, 100), &sparse_tile(&s))
216            .unwrap();
217        w.append_sparse(TileId::new(5, 200), &sparse_tile(&s))
218            .unwrap();
219        let bytes = finish_plain(w).unwrap();
220        let footer = SegmentFooter::decode(&bytes).unwrap();
221        assert_eq!(footer.tiles.len(), 2);
222        assert_eq!(footer.tiles[0].tile_id, TileId::new(5, 100));
223        assert_eq!(footer.tiles[1].tile_id, TileId::new(5, 200));
224    }
225
226    fn test_kek() -> nodedb_wal::crypto::WalEncryptionKey {
227        nodedb_wal::crypto::WalEncryptionKey::from_bytes(&[0x42u8; 32]).unwrap()
228    }
229
230    #[test]
231    fn array_segment_encrypted_at_rest() {
232        let s = schema();
233        let mut w = SegmentWriter::new(0xCAFE);
234        w.append_sparse(TileId::snapshot(1), &sparse_tile(&s))
235            .unwrap();
236        let kek = test_kek();
237        let bytes = w.finish(Some(&kek)).unwrap();
238        // Output must start with SEGA, not NDAS.
239        assert_eq!(&bytes[..4], b"SEGA");
240        // Must not start with plaintext NDAS magic.
241        assert_ne!(&bytes[..4], b"NDAS");
242    }
243
244    #[test]
245    fn array_segment_refuses_plaintext_with_kek() {
246        // Write an encrypted segment, then verify the reader rejects it
247        // when presented without a KEK.
248        let s = schema();
249        let mut w = SegmentWriter::new(0xBEEF);
250        w.append_sparse(TileId::snapshot(1), &sparse_tile(&s))
251            .unwrap();
252        let kek = test_kek();
253        let encrypted = w.finish(Some(&kek)).unwrap();
254        // Attempt to open without KEK — must return MissingKek.
255        let err = super::super::reader::SegmentReader::open_with_kek(&encrypted, None).unwrap_err();
256        assert!(
257            matches!(err, crate::error::ArrayError::MissingKek),
258            "expected MissingKek, got {err:?}"
259        );
260    }
261
262    #[test]
263    fn array_segment_refuses_encrypted_without_kek() {
264        // Write a plaintext segment, then verify the reader rejects it
265        // when presented with a KEK (encryption enforcement).
266        let s = schema();
267        let mut w = SegmentWriter::new(0xDEAD);
268        w.append_sparse(TileId::snapshot(1), &sparse_tile(&s))
269            .unwrap();
270        let plaintext = w.finish(None).unwrap();
271        let kek = test_kek();
272        // Attempt to open plaintext with KEK — must return KekRequired.
273        let err =
274            super::super::reader::SegmentReader::open_with_kek(&plaintext, Some(&kek)).unwrap_err();
275        assert!(
276            matches!(err, crate::error::ArrayError::KekRequired),
277            "expected KekRequired, got {err:?}"
278        );
279    }
280
281    #[test]
282    fn array_segment_handle_decrypts_into_owned_buffer() {
283        // Encrypt a multi-tile segment and round-trip it through the full
284        // SegmentReader::open_with_kek path.
285        let s = schema();
286        let mut w = SegmentWriter::new(0x1234);
287        w.append_sparse(TileId::new(1, 100), &sparse_tile(&s))
288            .unwrap();
289        w.append_sparse(TileId::new(2, 200), &sparse_tile(&s))
290            .unwrap();
291        let kek = test_kek();
292        let encrypted = w.finish(Some(&kek)).unwrap();
293        let owned = super::super::reader::OwnedSegmentReader::open_with_kek(&encrypted, Some(&kek))
294            .unwrap();
295        let reader = owned.reader();
296        assert_eq!(reader.tile_count(), 2);
297    }
298
299    #[test]
300    fn array_segment_encrypted_roundtrip_multiple_tiles() {
301        let s = schema();
302        let mut w = SegmentWriter::new(0xABCD);
303        for i in 1u64..=5 {
304            w.append_sparse(TileId::new(i, i as i64 * 100), &sparse_tile(&s))
305                .unwrap();
306        }
307        let kek = test_kek();
308        let encrypted = w.finish(Some(&kek)).unwrap();
309        let owned = super::super::reader::OwnedSegmentReader::open_with_kek(&encrypted, Some(&kek))
310            .unwrap();
311        let reader = owned.reader();
312        assert_eq!(reader.tile_count(), 5);
313        for idx in 0..5 {
314            let tile = reader.read_tile(idx).unwrap();
315            assert!(matches!(tile, super::super::reader::TilePayload::Sparse(_)));
316        }
317    }
318}
nodedb_array/segment/writer.rs

nodedb_array/segment/
writer.rs
