fix(decoding): box the FSE sequence tables so decode fits a 32 KiB stack

.github/workflows/ci.yml

@@ -396,6 +396,15 @@ jobs:
         # Corpus fixtures resolve from the manifest dir, so no
         # working-directory is needed.
         run: cargo nextest run --profile ci -p ffi-bench --features bench-internals,dict-builder
+      - name: Decode stack budget (32 KiB, panic = abort)
+        # `cargo test` always unwinds, and under unwind LLVM elides a copy of
+        # the frame state that an abort build (every `*-none` target) keeps in
+        # `decode_all`'s frame, so only a non-test binary built with
+        # panic = abort can see that frame grow past a kernel task's stack.
+        if: needs.changes.outputs.code == 'true' && runner.os == 'Linux'
+        env:
+          CARGO_PROFILE_RELEASE_PANIC: abort
+        run: cargo run --release -p ffi-bench --example stack_budget -- 32
 
   cross-i686:
     needs: lint

ffi-bench/Cargo.toml

@@ -266,6 +266,10 @@ required-features = ["dict-builder"]
 name = "ffi_loop_dict"
 path = "../zstd/examples/ffi_loop_dict.rs"
 
+[[example]]
+name = "stack_budget"
+path = "../zstd/examples/stack_budget.rs"
+
 [[example]]
 name = "profile_decode"
 path = "../zstd/examples/profile_decode.rs"

zstd/examples/stack_budget.rs

@@ -0,0 +1,92 @@
+//! Decodes on a thread whose stack is the budget a kernel task gets.
+//!
+//! `cargo test` always builds with `panic = "unwind"`, and under unwind LLVM
+//! elides a copy of the frame state that a `panic = "abort"` build (every
+//! `*-none` target) materialises inside `FrameDecoder::decode_all`. A stack
+//! test run by the harness therefore cannot see that regression. This example
+//! can: CI runs it in release with `CARGO_PROFILE_RELEASE_PANIC=abort`.
+//!
+//! A frame that does not fit hits the guard page and the process dies with
+//! "has overflowed its stack"; success prints one line per decode path.
+//!
+//! ```text
+//! CARGO_PROFILE_RELEASE_PANIC=abort \
+//!     cargo run --release -p ffi-bench --example stack_budget [KiB]
+//! ```
+
+use std::io::Read;
+use structured_zstd::decoding::{FrameDecoder, StreamingDecoder};
+use structured_zstd::encoding::{CompressionLevel, compress_to_vec};
+
+/// The task stack the decoder must fit in, including the caller's frames.
+const DEFAULT_BUDGET_KIB: usize = 32;
+
+/// Several blocks of each type: compressed (with sequences), RLE and Raw.
+fn mixed_input() -> Vec<u8> {
+    let mut data = Vec::new();
+    for i in 0..300_000u32 {
+        data.push((i * 7 % 251) as u8 ^ (i >> 9) as u8);
+    }
+    data.extend(core::iter::repeat_n(0xA5, 200_000));
+    let mut x = 0x9E37_79B9_7F4A_7C15u64;
+    for _ in 0..200_000 {
+        x ^= x << 13;
+        x ^= x >> 7;
+        x ^= x << 17;
+        data.push(x as u8);
+    }
+    data
+}
+
+fn on_budget(kib: usize, name: &str, f: impl FnOnce() + Send + 'static) {
+    std::thread::Builder::new()
+        .stack_size(kib * 1024)
+        .spawn(f)
+        .expect("spawn")
+        .join()
+        .unwrap_or_else(|_| panic!("{name} panicked"));
+    println!("ok: {name} fits {kib} KiB");
+}
+
+fn main() {
+    let kib = std::env::args()
+        .nth(1)
+        .map(|s| s.parse().expect("KiB"))
+        .unwrap_or(DEFAULT_BUDGET_KIB);
+
+    let data = mixed_input();
+    let frame = compress_to_vec(data.as_slice(), CompressionLevel::Default);
+
+    let (d, f) = (data.clone(), frame.clone());
+    on_budget(kib, "FrameDecoder::decode_all", move || {
+        let mut out = vec![0u8; d.len()];
+        let n = FrameDecoder::new()
+            .decode_all(&f, &mut out)
+            .expect("decode");
+        assert_eq!(&out[..n], d.as_slice());
+    });
+
+    on_budget(kib, "StreamingDecoder", move || {
+        let mut out = Vec::new();
+        StreamingDecoder::new(frame.as_slice())
+            .read_to_end(&mut out)
+            .expect("decode");
+        assert_eq!(out, data);
+    });
+
+    let dir = concat!(env!("CARGO_MANIFEST_DIR"), "/../zstd/dict_tests");
+    let read = |name: &str| {
+        let path = format!("{dir}/{name}");
+        std::fs::read(&path).unwrap_or_else(|e| panic!("read {path}: {e}"))
+    };
+    let dict = read("dictionary");
+    let frame = read("files/ModemManager.service.zst");
+    let expected = read("files/ModemManager.service");
+    on_budget(kib, "FrameDecoder::decode_all_with_dict_bytes", move || {
+        let mut out = vec![0u8; expected.len()];
+        let n = FrameDecoder::new()
+            .decode_all_with_dict_bytes(&frame, &mut out, &dict)
+            .expect("decode");
+        assert_eq!(&out[..n], expected.as_slice());
+    });
+}

zstd/src/decoding/dictionary.rs

@@ -1,3 +1,4 @@
+use alloc::boxed::Box;
 #[cfg(not(target_has_atomic = "ptr"))]
 use alloc::rc::Rc;
 #[cfg(target_has_atomic = "ptr")]
@@ -26,8 +27,10 @@ pub struct Dictionary {
     /// therefore still requires a non-zero one.
     pub id: u32,
     /// A dictionary can contain an entropy table, either FSE or
-    /// Huffman.
-    pub fse: FSEScratch,
+    /// Huffman. Boxed for the same reason as `DecoderScratch::fse`: the
+    /// ~12 KiB of inline tables must not ride the stack through
+    /// construction.
+    pub fse: Box<FSEScratch>,
     /// A dictionary can contain an entropy table, either FSE or
     /// Huffman.
     pub huf: HuffmanScratch,
@@ -70,10 +73,15 @@ pub const MAGIC_NUM: [u8; 4] = [0x37, 0xA4, 0x30, 0xEC];
 
 impl Dictionary {
     /// Heap bytes owned by this dictionary: the content plus the parsed
-    /// entropy tables' heap (the fixed-size FSE decode arrays are inline,
-    /// counted by `size_of::<Dictionary>()`).
+    /// entropy tables' heap. `fse` is a `Box`, so the fixed-size FSE decode
+    /// arrays are heap bytes rather than inline in `size_of::<Dictionary>()`;
+    /// charge that `size_of` here so a caller budgeting against the reported
+    /// figure sees them.
     pub fn heap_bytes(&self) -> usize {
-        self.dict_content.capacity() + self.fse.heap_bytes() + self.huf.heap_bytes()
+        self.dict_content.capacity()
+            + core::mem::size_of::<FSEScratch>()
+            + self.fse.heap_bytes()
+            + self.huf.heap_bytes()
     }
 
     /// Build a dictionary from raw content bytes (without entropy table sections).
@@ -98,7 +106,7 @@ impl Dictionary {
 
         Ok(Dictionary {
             id,
-            fse: FSEScratch::new(),
+            fse: Box::new(FSEScratch::new()),
             huf: HuffmanScratch::new(),
             dict_content,
             offset_hist: [1, 4, 8],
@@ -131,7 +139,7 @@ impl Dictionary {
             // caller asking for a dictionary that cannot exist.
             Ok(Dictionary {
                 id: 0,
-                fse: FSEScratch::new(),
+                fse: Box::new(FSEScratch::new()),
                 huf: HuffmanScratch::new(),
                 dict_content: Vec::new(),
                 offset_hist: [1, 4, 8],
@@ -186,7 +194,7 @@ impl Dictionary {
 
         let mut new_dict = Dictionary {
             id: 0,
-            fse: FSEScratch::new(),
+            fse: Box::new(FSEScratch::new()),
             huf: HuffmanScratch::new(),
             dict_content: Vec::new(),
             offset_hist: [1, 4, 8],

zstd/src/decoding/frame_decoder/tests.rs

@@ -3584,6 +3584,58 @@ fn a_zero_offset_sequence_is_refused_rather_than_executed() {
         .expect_err("decode_all_to_vec must refuse the frame");
 }
 
+/// A `FrameDecoder` and everything it owns inline must stay small enough to
+/// live in a kernel task's stack frame. With the FSE sequence tables held by
+/// value the decoder was ~14 KiB, and a `panic = "abort"` build (every
+/// `*-none` target) built a second copy of the state inside `decode_all`'s
+/// frame: ~42 KiB of stack between the caller and `decode_all`, which
+/// overflows a 32 KiB kernel stack. `cargo test` runs `panic = "unwind"`,
+/// where LLVM happens to elide the copy, so pin the root cause instead:
+/// the inline size of the types every decode entry point moves around.
+#[test]
+fn frame_decoder_stays_small_enough_for_a_kernel_stack() {
+    use crate::decoding::scratch::DecoderScratch;
+    const BUDGET: usize = 2 * 1024;
+    let decoder = core::mem::size_of::<FrameDecoder>();
+    let state = core::mem::size_of::<super::FrameDecoderState>();
+    let scratch = core::mem::size_of::<DecoderScratch>();
+    std::eprintln!("FrameDecoder={decoder} FrameDecoderState={state} DecoderScratch={scratch}");
+    assert!(decoder <= BUDGET, "FrameDecoder is {decoder} B inline");
+    assert!(state <= BUDGET, "FrameDecoderState is {state} B inline");
+    assert!(scratch <= BUDGET, "DecoderScratch is {scratch} B inline");
+}
+
+/// The boxed FSE tables must show up in the reported context size. They are
+/// ~12 KiB that moved from inline to the heap when `fse` became a `Box`, and
+/// a C caller budgets against this figure, so leaving them uncounted
+/// under-reports the decoder by more than the rest of a fresh scratch
+/// combined. Pins both owners: the decode scratch and a dictionary.
+#[test]
+fn boxed_fse_tables_are_reported_as_heap() {
+    use crate::decoding::Dictionary;
+    use crate::decoding::scratch::{DecoderScratch, FSEScratch};
+
+    let boxed = core::mem::size_of::<FSEScratch>();
+    assert!(
+        boxed > 4 * 1024,
+        "expected the FSE tables to be the bulk of the charge, got {boxed} B"
+    );
+
+    let scratch: DecoderScratch = DecoderScratch::new(1024);
+    assert!(
+        scratch.workspace_bytes() >= boxed,
+        "workspace_bytes() = {} omits the {boxed} B boxed FSE tables",
+        scratch.workspace_bytes()
+    );
+
+    let dict = Dictionary::from_raw_content(1, alloc::vec![0u8; 64]).expect("dictionary");
+    assert!(
+        dict.heap_bytes() >= boxed,
+        "Dictionary::heap_bytes() = {} omits the {boxed} B boxed FSE tables",
+        dict.heap_bytes()
+    );
+}
+
 /// RFC 8878 3.1.1.1.1: a frame whose descriptor has its Reserved_bit (bit 3)
 /// set is refused by every decode entry point, before any block is read, as
 /// upstream's `ZSTD_getFrameHeader_advanced` refuses it. The same frame

zstd/src/decoding/scratch.rs

@@ -6,6 +6,7 @@ use super::ringbuffer::RingBuffer;
 use crate::decoding::dictionary::{Dictionary, DictionaryHandle};
 use crate::fse::SeqFSETable;
 use crate::huff0::HuffmanTable;
+use alloc::boxed::Box;
 use alloc::vec::Vec;
 use core::ops::{Deref, DerefMut};
 
@@ -23,7 +24,14 @@ pub struct DecoderScratch<B: BufferBackend = RingBuffer> {
     /// The decoder used for Huffman blocks.
     pub huf: HuffmanScratch,
     /// The decoder used for FSE blocks.
-    pub fse: FSEScratch,
+    ///
+    /// Boxed so `DecoderScratch` -- and the `FrameDecoder` that holds one
+    /// inline -- stays small on the stack. The three sequence tables are
+    /// ~12 KiB of inline arrays; held by value, a `panic = "abort"` build
+    /// (every `*-none` kernel target) materialises one or two copies of the
+    /// whole state in `FrameDecoder::decode_all`'s frame, ~28 KiB, which
+    /// overflows a 32 KiB kernel task stack before a block is decoded.
+    pub fse: Box<FSEScratch>,
 
     pub buffer: DecodeBuffer<B>,
     pub offset_hist: [u32; 3],
@@ -140,7 +148,7 @@ impl<B: BufferBackend> DecoderScratch<B> {
                 table: HuffmanTable::new(),
                 table_source: TableSource::Local,
             },
-            fse: FSEScratch {
+            fse: Box::new(FSEScratch {
                 offsets: AlignedFSETable::new(MAX_OFFSET_CODE),
                 literal_lengths: AlignedFSETable::new(MAX_LITERAL_LENGTH_CODE),
                 match_lengths: AlignedFSETable::new(MAX_MATCH_LENGTH_CODE),
@@ -149,7 +157,7 @@ impl<B: BufferBackend> DecoderScratch<B> {
                 ll_source: SeqTableSource::Local,
                 of_source: SeqTableSource::Local,
                 ml_source: SeqTableSource::Local,
-            },
+            }),
             buffer: DecodeBuffer::new(window_size),
             offset_hist: [1, 4, 8],
 
@@ -162,11 +170,16 @@ impl<B: BufferBackend> DecoderScratch<B> {
     /// per-block literal and block-content buffers and the entropy tables. The
     /// window dominates and scales with the frame; the rest are bounded by the
     /// block maximum and the entropy alphabet.
+    ///
+    /// `fse` is a `Box`, so its fixed-size decode arrays are heap bytes rather
+    /// than inline in this struct: charge `size_of` for the boxed value on top
+    /// of the build-scratch vectors `FSEScratch::heap_bytes` sums.
     pub fn workspace_bytes(&self) -> usize {
         self.buffer.capacity()
             + self.literals_buffer.capacity()
             + self.block_content_buffer.capacity()
             + self.huf.heap_bytes()
+            + core::mem::size_of::<FSEScratch>()
             + self.fse.heap_bytes()
     }
 
@@ -463,9 +476,11 @@ pub struct FSEScratch {
 
 impl FSEScratch {
     /// Heap bytes owned by the three locally-built sequence FSE tables
-    /// (LL/ML/OF). The fixed-size decode arrays are inline (counted by
-    /// `size_of`); this sums their build-scratch vectors. `Dict`-sourced
-    /// tables read a shared handle and are not owned here.
+    /// (LL/ML/OF): this sums their build-scratch vectors. The fixed-size
+    /// decode arrays are part of `size_of::<FSEScratch>()` and so are NOT
+    /// counted here; every owner holds this struct in a `Box`, so each
+    /// charges that `size_of` itself. `Dict`-sourced tables read a shared
+    /// handle and are not owned here.
     pub fn heap_bytes(&self) -> usize {
         self.offsets.heap_bytes()
             + self.literal_lengths.heap_bytes()

zstd/src/encoding/frame_compressor/tests.rs

@@ -1391,7 +1391,7 @@ fn set_dictionary_from_bytes_matches_full_decode_byte_for_byte() {
 fn set_dictionary_accepts_a_dictionary_without_an_id() {
     let raw_content = crate::decoding::Dictionary {
         id: 0,
-        fse: crate::decoding::scratch::FSEScratch::new(),
+        fse: alloc::boxed::Box::new(crate::decoding::scratch::FSEScratch::new()),
         huf: crate::decoding::scratch::HuffmanScratch::new(),
         dict_content: vec![1, 2, 3],
         offset_hist: [1, 4, 8],
@@ -1409,7 +1409,7 @@ fn set_dictionary_accepts_a_dictionary_without_an_id() {
     let mut decoder = crate::decoding::FrameDecoder::new();
     let registered = decoder.add_dict(crate::decoding::Dictionary {
         id: 0,
-        fse: crate::decoding::scratch::FSEScratch::new(),
+        fse: alloc::boxed::Box::new(crate::decoding::scratch::FSEScratch::new()),
         huf: crate::decoding::scratch::HuffmanScratch::new(),
         dict_content: vec![1, 2, 3],
         offset_hist: [1, 4, 8],
@@ -1552,7 +1552,7 @@ fn set_dictionary_from_bytes_rejects_a_corrupt_serialized_dictionary() {
 fn set_dictionary_rejects_zero_repeat_offsets() {
     let invalid = crate::decoding::Dictionary {
         id: 1,
-        fse: crate::decoding::scratch::FSEScratch::new(),
+        fse: alloc::boxed::Box::new(crate::decoding::scratch::FSEScratch::new()),
         huf: crate::decoding::scratch::HuffmanScratch::new(),
         dict_content: vec![1, 2, 3],
         offset_hist: [0, 4, 8],