samkhya_core/
puffin.rs

1//! Iceberg Puffin sidecar format — reader and writer.
2//!
3//! Spec: <https://iceberg.apache.org/puffin-spec/>
4//!
5//! Layout:
6//! ```text
7//!  +-----------+
8//!  | Magic     |  4 bytes "PFA1"
9//!  +-----------+
10//!  | Blob 1    |  variable
11//!  +-----------+
12//!  | Blob 2    |  variable
13//!  +-----------+
14//!  | ...       |
15//!  +-----------+
16//!  | Footer    |  Magic + JSON payload + payload-len(LE u32)
17//!  |           |  + flags(LE u32) + Magic
18//!  +-----------+
19//! ```
20
21use std::collections::BTreeMap;
22use std::io::{Read, Seek, SeekFrom, Write};
23
24use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
25use serde::{Deserialize, Serialize};
26
27use crate::{Error, Result};
28
29const MAGIC: &[u8; 4] = b"PFA1";
30
31/// Puffin file property declaring the samkhya portability contract version.
32pub const SAMKHYA_SCHEMA_VERSION_PROPERTY: &str = "samkhya.schema-version";
33
34/// Current samkhya portability contract version.
35pub const SAMKHYA_SCHEMA_VERSION: &str = "1";
36
37/// Validate an optional samkhya file-schema marker.
38///
39/// Absence is accepted for v1 files written before the marker was introduced.
40/// An explicit future or malformed value fails closed so readers do not
41/// silently reinterpret a changed cross-engine contract.
42pub fn validate_samkhya_schema_version(version: Option<&str>) -> Result<()> {
43    match version {
44        None | Some(SAMKHYA_SCHEMA_VERSION) => Ok(()),
45        Some(other) => Err(Error::InvalidPuffin(format!(
46            "unsupported {} {other:?}; expected {SAMKHYA_SCHEMA_VERSION}",
47            SAMKHYA_SCHEMA_VERSION_PROPERTY
48        ))),
49    }
50}
51
52/// Maximum permitted size for a single blob payload, in bytes.
53///
54/// Caps untrusted `length` and decompressed-size values read from the
55/// JSON footer of a Puffin sidecar. A Bloom filter sized for 1B keys at
56/// 1% FPR fits comfortably under this limit (~1.2 GiB); legitimate
57/// HLL/CMS/histogram payloads are sub-MiB. Attacker-controlled sidecars
58/// declaring larger sizes are rejected before allocation, eliminating
59/// the trivial DoS surface where a small file claims a multi-EiB blob.
60const MAX_BLOB_LEN: u64 = 2 * 1024 * 1024 * 1024;
61
62/// Maximum permitted size for the JSON footer payload, in bytes.
63///
64/// The footer is a `serde_json` blob enumerating every per-blob metadata
65/// record. 16 MiB accommodates pathologically large numbers of blobs
66/// while still letting the reader reject attacker-controlled `payload_len`
67/// values that would force a multi-GiB allocation for the footer alone.
68const MAX_FOOTER_LEN: u64 = 16 * 1024 * 1024;
69
70/// Maximum number of blob entries enumerated in a single footer.
71///
72/// Caps the `blobs` array length after JSON parse so an attacker cannot
73/// stage a footer with millions of trivially-sized blob records that
74/// each individually pass the per-blob length cap. Legitimate samkhya
75/// sidecars carry one blob per stat per column — even a wide table with
76/// one HLL + one Bloom + one CMS + two histograms per column tops out
77/// well under this limit.
78const MAX_BLOB_COUNT: usize = 65_536;
79
80/// Compression codec applied to a blob payload before it is written to the file.
81///
82/// The codec name is recorded in the blob's `compression-codec` metadata so
83/// readers can decompress lazily without scanning unrelated blobs.
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum CompressionCodec {
86    /// Payload is stored verbatim.
87    None,
88    /// Payload is compressed with zstd (default level).
89    Zstd,
90}
91
92impl CompressionCodec {
93    /// Codec name as serialized in blob metadata (Puffin spec convention).
94    ///
95    /// # Examples
96    ///
97    /// ```
98    /// use samkhya_core::puffin::CompressionCodec;
99    ///
100    /// assert_eq!(CompressionCodec::None.as_str(), "none");
101    /// assert_eq!(CompressionCodec::Zstd.as_str(), "zstd");
102    /// ```
103    pub fn as_str(&self) -> &'static str {
104        match self {
105            CompressionCodec::None => "none",
106            CompressionCodec::Zstd => "zstd",
107        }
108    }
109
110    fn from_meta(meta: Option<&str>) -> Result<Self> {
111        match meta {
112            None | Some("none") => Ok(CompressionCodec::None),
113            Some("zstd") => Ok(CompressionCodec::Zstd),
114            Some(other) => Err(Error::InvalidPuffin(format!(
115                "unsupported blob compression codec {other:?}"
116            ))),
117        }
118    }
119}
120
121/// Footer payload (JSON-encoded inside the file).
122#[derive(Debug, Clone, Default, Serialize, Deserialize)]
123pub struct FooterPayload {
124    pub blobs: Vec<BlobMetadata>,
125    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
126    pub properties: BTreeMap<String, String>,
127}
128
129/// Per-blob metadata stored in the footer.
130#[derive(Debug, Clone, Serialize, Deserialize)]
131pub struct BlobMetadata {
132    #[serde(rename = "type")]
133    pub kind: String,
134    pub fields: Vec<i32>,
135    #[serde(
136        rename = "snapshot-id",
137        default,
138        skip_serializing_if = "Option::is_none"
139    )]
140    pub snapshot_id: Option<i64>,
141    #[serde(
142        rename = "sequence-number",
143        default,
144        skip_serializing_if = "Option::is_none"
145    )]
146    pub sequence_number: Option<i64>,
147    pub offset: u64,
148    pub length: u64,
149    #[serde(
150        rename = "compression-codec",
151        default,
152        skip_serializing_if = "Option::is_none"
153    )]
154    pub compression_codec: Option<String>,
155    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
156    pub properties: BTreeMap<String, String>,
157}
158
159/// A blob to be written to a Puffin file.
160///
161/// # Examples
162///
163/// ```
164/// use samkhya_core::puffin::Blob;
165///
166/// let payload = b"sketch bytes";
167/// let blob = Blob::new("samkhya.hll-v1", vec![7], payload);
168/// assert_eq!(blob.kind, "samkhya.hll-v1");
169/// assert_eq!(blob.fields, vec![7]);
170/// assert_eq!(blob.payload, payload);
171/// ```
172pub struct Blob<'a> {
173    pub kind: String,
174    pub fields: Vec<i32>,
175    pub payload: &'a [u8],
176    pub properties: BTreeMap<String, String>,
177}
178
179impl<'a> Blob<'a> {
180    /// Build a blob with empty properties. Add metadata afterwards by
181    /// inserting into `blob.properties`.
182    ///
183    /// # Examples
184    ///
185    /// ```
186    /// use samkhya_core::puffin::Blob;
187    ///
188    /// let blob = Blob::new("samkhya.bloom-v1", vec![1, 2], &[0u8, 1, 2, 3]);
189    /// assert!(blob.properties.is_empty());
190    /// ```
191    pub fn new(kind: impl Into<String>, fields: Vec<i32>, payload: &'a [u8]) -> Self {
192        Self {
193            kind: kind.into(),
194            fields,
195            payload,
196            properties: BTreeMap::new(),
197        }
198    }
199}
200
201/// Streaming writer for Puffin files.
202///
203/// Wraps any `Write + Seek` sink. The magic header is written lazily on the
204/// first blob (or finish); call [`PuffinWriter::finish`] to flush the footer
205/// and recover the inner writer.
206///
207/// # Examples
208///
209/// ```
210/// use std::io::Cursor;
211/// use samkhya_core::puffin::{Blob, PuffinReader, PuffinWriter};
212///
213/// let mut writer = PuffinWriter::new(Cursor::new(Vec::<u8>::new()));
214/// writer.add_blob(Blob::new("samkhya.test-v1", vec![0], b"hello")).unwrap();
215/// let cursor = writer.finish().unwrap();
216///
217/// // Round-trip: open the bytes back as a reader.
218/// let reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
219/// assert_eq!(reader.blobs().len(), 1);
220/// ```
221pub struct PuffinWriter<W: Write + Seek> {
222    inner: W,
223    blobs: Vec<BlobMetadata>,
224    properties: BTreeMap<String, String>,
225    pos: u64,
226    wrote_head: bool,
227}
228
229impl<W: Write + Seek> PuffinWriter<W> {
230    pub fn new(inner: W) -> Self {
231        let mut properties = BTreeMap::new();
232        properties.insert(
233            "created-by".to_owned(),
234            format!("samkhya-core version {}", env!("CARGO_PKG_VERSION")),
235        );
236        properties.insert(
237            SAMKHYA_SCHEMA_VERSION_PROPERTY.to_owned(),
238            SAMKHYA_SCHEMA_VERSION.to_owned(),
239        );
240        Self {
241            inner,
242            blobs: Vec::new(),
243            properties,
244            pos: 0,
245            wrote_head: false,
246        }
247    }
248
249    /// Set a file-level Puffin footer property.
250    pub fn with_file_property(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
251        self.properties.insert(key.into(), value.into());
252        self
253    }
254
255    fn ensure_head(&mut self) -> Result<()> {
256        if !self.wrote_head {
257            self.inner.write_all(MAGIC)?;
258            self.pos += MAGIC.len() as u64;
259            self.wrote_head = true;
260        }
261        Ok(())
262    }
263
264    /// Append a blob to the file.
265    ///
266    /// # Examples
267    ///
268    /// ```
269    /// use std::io::Cursor;
270    /// use samkhya_core::puffin::{Blob, PuffinWriter};
271    ///
272    /// let mut writer = PuffinWriter::new(Cursor::new(Vec::<u8>::new()));
273    /// writer.add_blob(Blob::new("samkhya.test-v1", vec![0], b"payload")).unwrap();
274    /// let _bytes = writer.finish().unwrap().into_inner();
275    /// ```
276    pub fn add_blob(&mut self, blob: Blob<'_>) -> Result<()> {
277        // Puffin v1 requires both fields even when the Iceberg snapshot is not
278        // known yet; -1 is the specification's unassigned sentinel.
279        self.add_blob_for_snapshot(blob, -1, -1)
280    }
281
282    /// Append an uncompressed blob with required Iceberg snapshot metadata.
283    pub fn add_blob_for_snapshot(
284        &mut self,
285        blob: Blob<'_>,
286        snapshot_id: i64,
287        sequence_number: i64,
288    ) -> Result<()> {
289        self.ensure_head()?;
290        let offset = self.pos;
291        self.inner.write_all(blob.payload)?;
292        let length = blob.payload.len() as u64;
293        self.pos += length;
294        self.blobs.push(BlobMetadata {
295            kind: blob.kind,
296            fields: blob.fields,
297            snapshot_id: Some(snapshot_id),
298            sequence_number: Some(sequence_number),
299            offset,
300            length,
301            compression_codec: None,
302            properties: blob.properties,
303        });
304        Ok(())
305    }
306
307    /// Append a blob with payload compressed under `codec`.
308    ///
309    /// `CompressionCodec::None` is equivalent to [`Self::add_blob`].
310    /// `CompressionCodec::Zstd` requires the `zstd` feature; otherwise this
311    /// returns an [`Error::InvalidPuffin`] explaining the missing feature.
312    pub fn add_blob_compressed(&mut self, blob: Blob<'_>, codec: CompressionCodec) -> Result<()> {
313        self.add_blob_compressed_for_snapshot(blob, codec, -1, -1)
314    }
315
316    /// Append a blob with compression and required snapshot metadata.
317    pub fn add_blob_compressed_for_snapshot(
318        &mut self,
319        blob: Blob<'_>,
320        codec: CompressionCodec,
321        snapshot_id: i64,
322        sequence_number: i64,
323    ) -> Result<()> {
324        match codec {
325            CompressionCodec::None => {
326                self.add_blob_for_snapshot(blob, snapshot_id, sequence_number)
327            }
328            CompressionCodec::Zstd => self.add_blob_zstd(blob, snapshot_id, sequence_number),
329        }
330    }
331
332    #[cfg(feature = "zstd")]
333    fn add_blob_zstd(
334        &mut self,
335        blob: Blob<'_>,
336        snapshot_id: i64,
337        sequence_number: i64,
338    ) -> Result<()> {
339        self.ensure_head()?;
340        let compressed = zstd::encode_all(blob.payload, 0)
341            .map_err(|e| Error::InvalidPuffin(format!("zstd encode: {e}")))?;
342        let offset = self.pos;
343        self.inner.write_all(&compressed)?;
344        let length = compressed.len() as u64;
345        self.pos += length;
346        self.blobs.push(BlobMetadata {
347            kind: blob.kind,
348            fields: blob.fields,
349            snapshot_id: Some(snapshot_id),
350            sequence_number: Some(sequence_number),
351            offset,
352            length,
353            compression_codec: Some(CompressionCodec::Zstd.as_str().to_string()),
354            properties: blob.properties,
355        });
356        Ok(())
357    }
358
359    #[cfg(not(feature = "zstd"))]
360    fn add_blob_zstd(
361        &mut self,
362        _blob: Blob<'_>,
363        _snapshot_id: i64,
364        _sequence_number: i64,
365    ) -> Result<()> {
366        Err(Error::InvalidPuffin(
367            "zstd codec requested but the `zstd` cargo feature is disabled".into(),
368        ))
369    }
370
371    /// Finalize the file: write the footer and return the inner writer.
372    pub fn finish(mut self) -> Result<W> {
373        self.ensure_head()?;
374        let footer = FooterPayload {
375            blobs: self.blobs,
376            properties: self.properties,
377        };
378        let payload = serde_json::to_vec(&footer)
379            .map_err(|e| Error::InvalidPuffin(format!("footer JSON encode: {e}")))?;
380        let payload_len = u32::try_from(payload.len())
381            .map_err(|_| Error::InvalidPuffin("footer payload exceeds u32".into()))?;
382
383        self.inner.write_all(MAGIC)?;
384        self.inner.write_all(&payload)?;
385        self.inner.write_u32::<LittleEndian>(payload_len)?;
386        self.inner.write_u32::<LittleEndian>(0)?; // flags
387        self.inner.write_all(MAGIC)?;
388        Ok(self.inner)
389    }
390}
391
392/// Reader for Puffin files — parses the footer once, lazily loads blob payloads.
393///
394/// # Examples
395///
396/// ```
397/// use std::io::Cursor;
398/// use samkhya_core::puffin::{Blob, PuffinReader, PuffinWriter};
399///
400/// // Write a file in memory, then read it back.
401/// let mut w = PuffinWriter::new(Cursor::new(Vec::<u8>::new()));
402/// w.add_blob(Blob::new("samkhya.hll-v1", vec![1], b"sketch bytes")).unwrap();
403/// let bytes = w.finish().unwrap().into_inner();
404///
405/// let mut reader = PuffinReader::open(Cursor::new(bytes)).unwrap();
406/// let (idx, meta) = reader.find_blob("samkhya.hll-v1").unwrap();
407/// assert_eq!(meta.kind, "samkhya.hll-v1");
408/// assert_eq!(reader.read_blob(idx).unwrap(), b"sketch bytes");
409/// ```
410pub struct PuffinReader<R: Read + Seek> {
411    inner: R,
412    footer: FooterPayload,
413    blob_data_end: u64,
414}
415
416impl<R: Read + Seek> PuffinReader<R> {
417    pub fn open(mut inner: R) -> Result<Self> {
418        let file_len = inner.seek(SeekFrom::End(0))?;
419        // 4 (head magic) + 4 (footer head magic) + 4 (payload len) + 4 (flags) + 4 (trailing magic)
420        if file_len < 20 {
421            return Err(Error::InvalidPuffin(format!(
422                "file too short: {file_len} bytes"
423            )));
424        }
425
426        // Trailing magic
427        inner.seek(SeekFrom::End(-4))?;
428        let mut trailing = [0u8; 4];
429        inner.read_exact(&mut trailing)?;
430        if &trailing != MAGIC {
431            return Err(Error::InvalidPuffin("trailing magic missing".into()));
432        }
433
434        // A compressed footer (bit 0) and all reserved flags are unsupported
435        // by this synchronous reader. Reject explicitly instead of trying to
436        // parse compressed bytes as JSON or silently accepting future flags.
437        inner.seek(SeekFrom::End(-8))?;
438        let flags = inner.read_u32::<LittleEndian>()?;
439        if flags != 0 {
440            let reason = if flags & 1 == 1 {
441                "compressed Puffin footers are not supported"
442            } else {
443                "Puffin footer contains unsupported reserved flags"
444            };
445            return Err(Error::InvalidPuffin(format!("{reason}: 0x{flags:08x}")));
446        }
447
448        // Payload length (4 bytes before flags)
449        inner.seek(SeekFrom::End(-12))?;
450        let payload_len = inner.read_u32::<LittleEndian>()? as u64;
451
452        if payload_len > MAX_FOOTER_LEN {
453            return Err(Error::InvalidPuffin(format!(
454                "footer payload length {payload_len} exceeds MAX_FOOTER_LEN {MAX_FOOTER_LEN}"
455            )));
456        }
457
458        let footer_total = 16u64 + payload_len; // head magic + payload + len + flags + trailing magic
459        if file_len < footer_total {
460            return Err(Error::InvalidPuffin(
461                "payload length exceeds file size".into(),
462            ));
463        }
464
465        // Payload bytes
466        let payload_start = file_len - 12 - payload_len;
467        inner.seek(SeekFrom::Start(payload_start))?;
468        let mut payload = vec![0u8; payload_len as usize];
469        inner.read_exact(&mut payload)?;
470
471        // Footer head magic (4 bytes before payload)
472        inner.seek(SeekFrom::Start(payload_start - 4))?;
473        let mut footer_head = [0u8; 4];
474        inner.read_exact(&mut footer_head)?;
475        if &footer_head != MAGIC {
476            return Err(Error::InvalidPuffin("footer head magic missing".into()));
477        }
478
479        // File head magic
480        inner.seek(SeekFrom::Start(0))?;
481        let mut head = [0u8; 4];
482        inner.read_exact(&mut head)?;
483        if &head != MAGIC {
484            return Err(Error::InvalidPuffin("file head magic missing".into()));
485        }
486
487        let footer: FooterPayload = serde_json::from_slice(&payload)
488            .map_err(|e| Error::InvalidPuffin(format!("footer JSON decode: {e}")))?;
489
490        if footer.blobs.len() > MAX_BLOB_COUNT {
491            return Err(Error::InvalidPuffin(format!(
492                "footer enumerates {} blobs; exceeds MAX_BLOB_COUNT {MAX_BLOB_COUNT}",
493                footer.blobs.len()
494            )));
495        }
496
497        Ok(Self {
498            inner,
499            footer,
500            blob_data_end: payload_start - 4,
501        })
502    }
503
504    pub fn footer(&self) -> &FooterPayload {
505        &self.footer
506    }
507
508    pub fn blobs(&self) -> &[BlobMetadata] {
509        &self.footer.blobs
510    }
511
512    /// Read a blob's payload by index.
513    pub fn read_blob(&mut self, idx: usize) -> Result<Vec<u8>> {
514        let meta = self
515            .footer
516            .blobs
517            .get(idx)
518            .ok_or_else(|| Error::InvalidPuffin(format!("blob index {idx} out of range")))?;
519        if meta.length > MAX_BLOB_LEN {
520            return Err(Error::InvalidPuffin(format!(
521                "blob {idx} length {} exceeds MAX_BLOB_LEN {MAX_BLOB_LEN}",
522                meta.length
523            )));
524        }
525        let blob_end = meta
526            .offset
527            .checked_add(meta.length)
528            .ok_or_else(|| Error::InvalidPuffin(format!("blob {idx} offset overflow")))?;
529        if meta.offset < MAGIC.len() as u64 || blob_end > self.blob_data_end {
530            return Err(Error::InvalidPuffin(format!(
531                "blob {idx} range {}..{blob_end} falls outside payload region {}..{}",
532                meta.offset,
533                MAGIC.len(),
534                self.blob_data_end
535            )));
536        }
537        self.inner.seek(SeekFrom::Start(meta.offset))?;
538        let mut buf = vec![0u8; meta.length as usize];
539        self.inner.read_exact(&mut buf)?;
540        Ok(buf)
541    }
542
543    /// Read a blob's payload by index and decompress it according to the
544    /// blob's `compression_codec` metadata.
545    ///
546    /// Falls back to the raw bytes when the codec is `None` / unset. Returns
547    /// an error if the recorded codec is unsupported in the current build
548    /// (e.g. `zstd` without the `zstd` feature).
549    pub fn read_blob_decompressed(&mut self, idx: usize) -> Result<Vec<u8>> {
550        let codec = {
551            let meta =
552                self.footer.blobs.get(idx).ok_or_else(|| {
553                    Error::InvalidPuffin(format!("blob index {idx} out of range"))
554                })?;
555            CompressionCodec::from_meta(meta.compression_codec.as_deref())?
556        };
557        let raw = self.read_blob(idx)?;
558        match codec {
559            CompressionCodec::None => Ok(raw),
560            CompressionCodec::Zstd => decode_zstd(&raw),
561        }
562    }
563
564    /// Find the first blob whose `type` (kind) matches `kind`.
565    ///
566    /// # Examples
567    ///
568    /// ```
569    /// use std::io::Cursor;
570    /// use samkhya_core::puffin::{Blob, PuffinReader, PuffinWriter};
571    ///
572    /// let mut w = PuffinWriter::new(Cursor::new(Vec::<u8>::new()));
573    /// w.add_blob(Blob::new("samkhya.hll-v1", vec![0], b"x")).unwrap();
574    /// let bytes = w.finish().unwrap().into_inner();
575    /// let reader = PuffinReader::open(Cursor::new(bytes)).unwrap();
576    /// assert!(reader.find_blob("samkhya.hll-v1").is_some());
577    /// assert!(reader.find_blob("absent.kind").is_none());
578    /// ```
579    pub fn find_blob(&self, kind: &str) -> Option<(usize, &BlobMetadata)> {
580        self.footer
581            .blobs
582            .iter()
583            .enumerate()
584            .find(|(_, b)| b.kind == kind)
585    }
586}
587
588/// Bounded streaming decompression helper shared across compression
589/// codecs.
590///
591/// Reads from `reader` into a fresh `Vec<u8>` whose final length is at
592/// most `max_bytes`. The `+1` sentinel byte distinguishes "decompressed
593/// exactly max_bytes" (legitimate boundary) from "would have produced
594/// more than max_bytes" (attacker-controlled). `codec_label` appears in
595/// the error message so the caller can identify which codec tripped.
596///
597/// Centralising this here means a future LZ4 / Brotli / Snappy codec
598/// only needs to construct a `Read` and call this function; the cap is
599/// not re-derived per codec. See `documents/SECURITY-REVIEW-2026-05-17.md`
600/// item M6. Currently used only by the zstd path; broaden the cfg gate
601/// when a new codec lands.
602#[cfg(feature = "zstd")]
603fn decompress_bounded<R: std::io::Read>(
604    mut reader: R,
605    max_bytes: usize,
606    codec_label: &'static str,
607) -> Result<Vec<u8>> {
608    use std::io::Read as _;
609
610    let mut out = Vec::with_capacity(max_bytes.min(256 * 1024));
611    let n = (&mut reader)
612        .take(max_bytes as u64 + 1)
613        .read_to_end(&mut out)
614        .map_err(|e| Error::InvalidPuffin(format!("{codec_label} decode: {e}")))?;
615    let _ = n; // silence unused-binding when read_to_end's return is not used
616    if out.len() > max_bytes {
617        return Err(Error::InvalidPuffin(format!(
618            "{codec_label} decompressed size exceeds {max_bytes}"
619        )));
620    }
621    Ok(out)
622}
623
624#[cfg(feature = "zstd")]
625fn decode_zstd(raw: &[u8]) -> Result<Vec<u8>> {
626    // Stream into a bounded buffer so a high-ratio compression bomb
627    // cannot force a multi-GiB allocation. The cap and the sentinel
628    // logic live in `decompress_bounded`.
629    let decoder = zstd::stream::Decoder::new(raw)
630        .map_err(|e| Error::InvalidPuffin(format!("zstd decoder init: {e}")))?;
631    decompress_bounded(decoder, MAX_BLOB_LEN as usize, "zstd")
632}
633
634#[cfg(not(feature = "zstd"))]
635fn decode_zstd(_raw: &[u8]) -> Result<Vec<u8>> {
636    Err(Error::InvalidPuffin(
637        "blob is zstd-compressed but the `zstd` cargo feature is disabled".into(),
638    ))
639}
640
641#[cfg(test)]
642mod tests {
643    use std::io::Cursor;
644
645    use super::*;
646    use crate::sketches::{HllSketch, Sketch};
647
648    #[test]
649    fn round_trip_empty() {
650        let writer = PuffinWriter::new(Cursor::new(Vec::new()));
651        let cursor = writer.finish().unwrap();
652        let reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
653        assert!(reader.blobs().is_empty());
654    }
655
656    #[test]
657    fn round_trip_single_blob() {
658        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
659        writer
660            .add_blob(Blob::new("samkhya.test-v1", vec![0], b"hello puffin"))
661            .unwrap();
662        let cursor = writer.finish().unwrap();
663
664        let mut reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
665        assert_eq!(reader.blobs().len(), 1);
666        assert_eq!(reader.blobs()[0].kind, "samkhya.test-v1");
667        assert_eq!(reader.blobs()[0].snapshot_id, Some(-1));
668        assert_eq!(reader.blobs()[0].sequence_number, Some(-1));
669        assert!(reader.footer().properties["created-by"].starts_with("samkhya-core version "));
670        assert_eq!(reader.read_blob(0).unwrap(), b"hello puffin");
671    }
672
673    #[test]
674    fn snapshot_aware_blob_metadata_round_trips() {
675        let mut writer =
676            PuffinWriter::new(Cursor::new(Vec::new())).with_file_property("deployment", "golden");
677        writer
678            .add_blob_for_snapshot(Blob::new("samkhya.test-v1", vec![17], b"payload"), 4_242, 7)
679            .unwrap();
680        let cursor = writer.finish().unwrap();
681
682        let reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
683        assert_eq!(reader.blobs()[0].snapshot_id, Some(4_242));
684        assert_eq!(reader.blobs()[0].sequence_number, Some(7));
685        assert_eq!(reader.footer().properties["deployment"], "golden");
686    }
687
688    #[test]
689    fn round_trip_multiple_blobs() {
690        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
691        writer
692            .add_blob(Blob::new("samkhya.hll-v1", vec![1], &[1, 2, 3, 4, 5]))
693            .unwrap();
694        writer
695            .add_blob(Blob::new("samkhya.bloom-v1", vec![2], &[10, 20, 30]))
696            .unwrap();
697        let cursor = writer.finish().unwrap();
698
699        let mut reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
700        assert_eq!(reader.blobs().len(), 2);
701        assert_eq!(reader.read_blob(0).unwrap(), vec![1, 2, 3, 4, 5]);
702        assert_eq!(reader.read_blob(1).unwrap(), vec![10, 20, 30]);
703        assert_eq!(
704            reader.find_blob("samkhya.bloom-v1").map(|(i, _)| i),
705            Some(1)
706        );
707        assert_eq!(reader.find_blob("absent.kind").map(|(i, _)| i), None);
708    }
709
710    #[test]
711    fn round_trip_hll_sketch_through_puffin() {
712        let mut hll = HllSketch::new(12).unwrap();
713        for i in 0..1000u32 {
714            hll.add(&i.to_le_bytes());
715        }
716        let payload = hll.to_bytes().unwrap();
717
718        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
719        writer
720            .add_blob(Blob::new(HllSketch::KIND, vec![7], &payload))
721            .unwrap();
722        let cursor = writer.finish().unwrap();
723
724        let mut reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
725        let (idx, meta) = reader.find_blob(HllSketch::KIND).unwrap();
726        assert_eq!(meta.fields, vec![7]);
727        let blob_bytes = reader.read_blob(idx).unwrap();
728        let hll2 = HllSketch::from_bytes(&blob_bytes).unwrap();
729        let err = (hll2.estimate() as f64 - 1000.0).abs() / 1000.0;
730        assert!(err < 0.1, "HLL estimate via Puffin off by {err}");
731    }
732
733    #[test]
734    fn rejects_too_short_file() {
735        let result = PuffinReader::open(Cursor::new(vec![0u8; 5]));
736        assert!(result.is_err());
737    }
738
739    #[test]
740    fn rejects_bad_trailing_magic() {
741        let mut buf = vec![0u8; 20];
742        buf[0..4].copy_from_slice(MAGIC);
743        // trailing 4 bytes are not magic
744        let result = PuffinReader::open(Cursor::new(buf));
745        assert!(result.is_err());
746    }
747
748    #[test]
749    fn read_blob_decompressed_no_codec_is_passthrough() {
750        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
751        writer
752            .add_blob(Blob::new("samkhya.test-v1", vec![0], b"plain payload"))
753            .unwrap();
754        let cursor = writer.finish().unwrap();
755
756        let mut reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
757        assert_eq!(reader.read_blob_decompressed(0).unwrap(), b"plain payload");
758    }
759
760    #[test]
761    fn read_blob_decompressed_rejects_unknown_codec() {
762        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
763        writer
764            .add_blob(Blob::new("samkhya.test-v1", vec![0], b"plain payload"))
765            .unwrap();
766        let cursor = writer.finish().unwrap();
767        let mut reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
768        reader.footer.blobs[0].compression_codec = Some("gzip".to_owned());
769
770        let error = reader.read_blob_decompressed(0).unwrap_err();
771        assert!(
772            error
773                .to_string()
774                .contains("unsupported blob compression codec")
775        );
776    }
777
778    #[test]
779    fn open_rejects_unsupported_footer_flags() {
780        let writer = PuffinWriter::new(Cursor::new(Vec::new()));
781        let mut bytes = writer.finish().unwrap().into_inner();
782        let flags_offset = bytes.len() - 8;
783        bytes[flags_offset..flags_offset + 4].copy_from_slice(&2_u32.to_le_bytes());
784
785        let error = PuffinReader::open(Cursor::new(bytes)).err().unwrap();
786        assert!(error.to_string().contains("unsupported reserved flags"));
787    }
788
789    #[test]
790    fn read_blob_rejects_oversized_length() {
791        // Build a minimal valid Puffin file then tamper with the JSON
792        // footer to claim a blob length of MAX_BLOB_LEN + 1. The reader
793        // must reject the read_blob call before allocating.
794        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
795        writer
796            .add_blob(Blob::new("samkhya.test-v1", vec![0], b"x"))
797            .unwrap();
798        let bytes = writer.finish().unwrap().into_inner();
799
800        // Locate the JSON footer payload (between footer-head magic and
801        // the payload-len trailer). The header is the first 4 bytes; the
802        // last 12 bytes are payload-len(4) + flags(4) + trailing-magic(4).
803        let file_len = bytes.len();
804        let payload_len =
805            u32::from_le_bytes(bytes[file_len - 12..file_len - 8].try_into().unwrap()) as usize;
806        let payload_start = file_len - 12 - payload_len;
807        let payload_end = file_len - 12;
808        let json = std::str::from_utf8(&bytes[payload_start..payload_end])
809            .unwrap()
810            .to_string();
811        // Replace the legitimate length:1 with an oversized value. The
812        // payload was `b"x"` so the field reads `"length":1`.
813        let bogus_len = MAX_BLOB_LEN + 1;
814        let tampered_json = json.replacen("\"length\":1", &format!("\"length\":{bogus_len}"), 1);
815        assert_ne!(tampered_json, json, "tamper failed — payload schema drift?");
816
817        // Rebuild the file with the tampered footer. Adjust payload_len
818        // to match the new JSON byte length.
819        let mut tampered = Vec::with_capacity(file_len + 32);
820        tampered.extend_from_slice(&bytes[..payload_start]);
821        tampered.extend_from_slice(tampered_json.as_bytes());
822        let new_payload_len = tampered_json.len() as u32;
823        tampered.write_u32::<LittleEndian>(new_payload_len).unwrap();
824        tampered.write_u32::<LittleEndian>(0).unwrap();
825        tampered.extend_from_slice(MAGIC);
826
827        let mut reader = PuffinReader::open(Cursor::new(tampered)).unwrap();
828        let err = reader.read_blob(0).unwrap_err();
829        match err {
830            Error::InvalidPuffin(msg) => assert!(
831                msg.contains("MAX_BLOB_LEN"),
832                "expected MAX_BLOB_LEN rejection, got: {msg}"
833            ),
834            other => panic!("expected InvalidPuffin, got {other:?}"),
835        }
836    }
837
838    #[test]
839    fn open_rejects_oversized_footer_payload_len() {
840        // Build a minimal file then rewrite its payload-len field to a
841        // value that exceeds MAX_FOOTER_LEN. Reader must reject before
842        // allocating a footer-sized buffer.
843        let writer = PuffinWriter::new(Cursor::new(Vec::new()));
844        let bytes = writer.finish().unwrap().into_inner();
845        let mut tampered = bytes.clone();
846        let oversized = (MAX_FOOTER_LEN as u32).saturating_add(1);
847        let len_offset = tampered.len() - 12;
848        tampered[len_offset..len_offset + 4].copy_from_slice(&oversized.to_le_bytes());
849        match PuffinReader::open(Cursor::new(tampered)) {
850            Ok(_) => panic!("expected MAX_FOOTER_LEN rejection, got Ok"),
851            Err(Error::InvalidPuffin(msg)) => assert!(
852                msg.contains("MAX_FOOTER_LEN"),
853                "expected MAX_FOOTER_LEN rejection, got: {msg}"
854            ),
855            Err(other) => panic!("expected InvalidPuffin, got {other:?}"),
856        }
857    }
858
859    #[cfg(not(feature = "zstd"))]
860    #[test]
861    fn requesting_zstd_without_feature_errors() {
862        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
863        let err = writer
864            .add_blob_compressed(
865                Blob::new("samkhya.test-v1", vec![0], b"x"),
866                CompressionCodec::Zstd,
867            )
868            .unwrap_err();
869        assert!(matches!(err, Error::InvalidPuffin(_)));
870    }
871}
872
873#[cfg(all(test, feature = "zstd"))]
874mod zstd_tests {
875    use std::io::Cursor;
876
877    use super::*;
878    use crate::sketches::{HllSketch, Sketch};
879
880    #[test]
881    fn round_trip_compressed_blob() {
882        // Payload with enough redundancy that zstd visibly shrinks it.
883        let payload = vec![0xABu8; 8192];
884
885        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
886        writer
887            .add_blob_compressed(
888                Blob::new("samkhya.test-v1", vec![0], &payload),
889                CompressionCodec::Zstd,
890            )
891            .unwrap();
892        let cursor = writer.finish().unwrap();
893
894        let mut reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
895        let meta = &reader.blobs()[0];
896        assert_eq!(meta.compression_codec.as_deref(), Some("zstd"));
897        // Compressed length should be strictly smaller than the original.
898        assert!((meta.length as usize) < payload.len());
899
900        let decoded = reader.read_blob_decompressed(0).unwrap();
901        assert_eq!(decoded, payload);
902    }
903
904    #[test]
905    fn round_trip_compressed_hll_sketch() {
906        let mut hll = HllSketch::new(14).unwrap();
907        for i in 0..5_000u32 {
908            hll.add(&i.to_le_bytes());
909        }
910        let bytes = hll.to_bytes().unwrap();
911
912        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
913        writer
914            .add_blob_compressed(
915                Blob::new(HllSketch::KIND, vec![1], &bytes),
916                CompressionCodec::Zstd,
917            )
918            .unwrap();
919        let cursor = writer.finish().unwrap();
920
921        let mut reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
922        let (idx, meta) = reader.find_blob(HllSketch::KIND).unwrap();
923        assert_eq!(meta.compression_codec.as_deref(), Some("zstd"));
924        let decoded = reader.read_blob_decompressed(idx).unwrap();
925        let hll2 = HllSketch::from_bytes(&decoded).unwrap();
926        let err = (hll2.estimate() as f64 - 5_000.0).abs() / 5_000.0;
927        assert!(err < 0.05, "HLL estimate off by {err}");
928    }
929
930    #[test]
931    fn none_codec_via_compressed_api_matches_plain() {
932        let mut writer = PuffinWriter::new(Cursor::new(Vec::new()));
933        writer
934            .add_blob_compressed(
935                Blob::new("samkhya.test-v1", vec![0], b"identity"),
936                CompressionCodec::None,
937            )
938            .unwrap();
939        let cursor = writer.finish().unwrap();
940        let mut reader = PuffinReader::open(Cursor::new(cursor.into_inner())).unwrap();
941        assert!(reader.blobs()[0].compression_codec.is_none());
942        assert_eq!(reader.read_blob_decompressed(0).unwrap(), b"identity");
943    }
944}
samkhya_core/puffin.rs

samkhya_core/
puffin.rs
