diff --git a/rust/compression/fsst/examples/benchmark.rs b/rust/compression/fsst/examples/benchmark.rs index c442243e112..f71abeefedd 100644 --- a/rust/compression/fsst/examples/benchmark.rs +++ b/rust/compression/fsst/examples/benchmark.rs @@ -56,7 +56,10 @@ fn benchmark(file_path: &str) { let mut decompression_out_bufs = vec![]; let mut decompression_out_offsets_bufs = vec![]; for _ in 0..TEST_NUM { - let this_decom_out_buf = vec![0u8; BUFFER_SIZE * 3]; + // `decompress` requires the output buffer to be at least 8x the compressed input (a 1-byte + // code can expand to an 8-byte symbol). The compressed buffer is at most `BUFFER_SIZE`, so + // `BUFFER_SIZE * 8` is a safe upper bound. + let this_decom_out_buf = vec![0u8; BUFFER_SIZE * 8]; let this_decom_out_offsets_buf = vec![0i32; BUFFER_SIZE * 3]; decompression_out_bufs.push(this_decom_out_buf); decompression_out_offsets_bufs.push(this_decom_out_offsets_buf); diff --git a/rust/compression/fsst/src/fsst.rs b/rust/compression/fsst/src/fsst.rs index 0a2bf1d03d7..d00a6ed806b 100644 --- a/rust/compression/fsst/src/fsst.rs +++ b/rust/compression/fsst/src/fsst.rs @@ -57,6 +57,12 @@ use std::ptr; #[inline] fn fsst_unaligned_load_unchecked(v: *const u8) -> u64 { + // SAFETY: the caller must guarantee that `v` points to at least 8 readable bytes. All callers + // uphold this: `compress_bulk` loads from a 520-byte stack buffer at an offset < 511 (leaving + // >= 8 bytes), `build_symbol_table` guards the load with `word.len() > 7 && curr < word.len() - 7`, + // `find_longest_symbol_from_char_slice` copies into a stack `[u8; 8]` before loading, and + // `FsstDecoder::init` reads symbols from a `symbol_table` buffer already validated to be + // exactly `FSST_SYMBOL_TABLE_SIZE` bytes. unsafe { ptr::read_unaligned(v as *const u64) } } @@ -812,12 +818,31 @@ fn decompress_bulk( ) -> io::Result<()> { let symbols = decoder.symbols; let lens = decoder.lens; + // SAFETY invariant shared by every `unsafe` block in this closure: + // - `out` is sized to at least 8x `compressed_strs` (checked in `FsstDecoder::init`, which the + // sole public entry point always runs before reaching this function). Each code advances + // `out_curr` by `lens[code]`, which is 1..=8 for a well-formed symbol table, and each + // consumed input byte yields at most 8 output bytes, so `out_curr + 8 <= out.len()` at every + // 8-byte write, including the final one. This is why we can `write_unaligned` a full 8-byte + // word per code and advance by only the length. + // NOTE: `lens` is loaded verbatim from the (untrusted) symbol table and is NOT re-validated + // to be <= 8 on decode, and offsets (below) are likewise trusted. A corrupted table or + // offset buffer can violate these bounds; callers must supply structures produced by + // `compress` (or otherwise trusted). Hardening the decoder against corrupt input is a + // separate concern, not addressed here. + // - The only unchecked read is `read_unaligned::`, gated by `in_curr + 4 <= in_end`; the + // scalar paths use bounds-checked indexing. `in_end` is a caller-provided offset into + // `compressed_strs`; the read is sound only if `in_end <= compressed_strs.len()`, which is a + // trusted precondition (holds for encoder-produced offsets; not validated here). let mut decompress = |mut in_curr: usize, in_end: usize, out_curr: &mut usize| { // Do SIMD operation here by 4 bytes while in_curr + 4 <= in_end { let next_block; let mut code; let mut len; + // SAFETY: the loop guard proves `in_curr + 4 <= in_end`. Per the closure-level + // invariant, `in_end <= compressed_strs.len()` is a trusted precondition (not checked + // here), so the 4-byte read is in bounds for well-formed input. unsafe { next_block = ptr::read_unaligned(compressed_strs.as_ptr().add(in_curr) as *const u32); @@ -983,6 +1008,10 @@ fn decompress_bulk( if in_curr < in_end { // last code cannot be an escape code let code = compressed_strs[in_curr] as usize; + // SAFETY: see the closure-level invariant. This is the final write and has no + // subsequent write to cover its slack, so it is the tightest case: `out_curr` is at + // most 8*(consumed_input_bytes - 1), and the 8-byte store lands within `out.len()` + // precisely because the caller sized `out` to 8x the input. unsafe { let src = symbols[code]; ptr::write_unaligned(out.as_mut_ptr().add(*out_curr) as *mut u64, src); @@ -1188,8 +1217,19 @@ impl FsstDecoder { } self.decoder_switch_on = (st_info & (1 << 24)) != 0; - // when decoder_switch_on is true, we make sure the out_buf is at least 3 times the size of the in_buf, - if self.decoder_switch_on && in_buf.len() * 3 > out_buf.len() { + // A single 1-byte code can decode to a symbol of up to MAX_SYMBOL_LENGTH (8) bytes, so the + // decoded output can be up to 8x the input. `decompress_bulk` also relies on this bound: it + // writes a full 8-byte word per code (advancing only by the symbol length), so the output + // buffer must be large enough that even the final write stays in bounds. Require out_buf to + // be at least 8x in_buf. `checked_mul` guards against `in_buf.len() * 8` wrapping on 32-bit + // targets (an input >= 512 MiB would otherwise bypass the check); treat overflow as too + // small. + if self.decoder_switch_on + && in_buf + .len() + .checked_mul(8) + .is_none_or(|needed| needed > out_buf.len()) + { return Err(io::Error::new( io::ErrorKind::InvalidInput, "output buffer too small for FSST decoder", @@ -1288,8 +1328,10 @@ pub fn compress( // the following 32 bits after FSST_MAGIC contains information about FSST encoding, such as decoder_switch_on, suffix_lim, terminator, n_symbols // when the decoder_switch_on is off in the in_buf header, `decompress` first make sure the out_buf is at least the same size as the in_buf, then simply copy the // input data to the output -// when the decoder_switch_on is on, `decompress` first make sure the out_buf is at least 3 times the size of the in_buf, then start decoding the -// data using the symbol table +// when the decoder_switch_on is on, `decompress` first make sure the out_buf is at least 8 times the size of the in_buf, then start decoding the +// data using the symbol table. The 8x bound is required for correctness: a 1-byte code can expand to +// an 8-byte symbol, and the decode loop writes a full 8-byte word per code, so a smaller buffer can +// be written out of bounds. // the out_offsets_buf should be at least the same size as the in_offsets_buf, otherwise an error is returned // the symbol_table is the same symbol table created by `compression` pub fn decompress( @@ -1644,4 +1686,59 @@ But exactly how the acquaintance and friendship came about, we cannot say."; ); } } + + // Build a genuinely FSST-compressed (decoder_switch_on) buffer to exercise the decode-side + // output-buffer size contract. Returns (symbol_table, compressed_bytes, compressed_offsets). + fn compress_paragraph() -> ([u8; FSST_SYMBOL_TABLE_SIZE], Vec, Vec) { + let test_input = TEST_PARAGRAPH.repeat((1024 * 1024) / TEST_PARAGRAPH.len()); + let lines_vec = test_input.lines().collect::>(); + let string_array = StringArray::from(lines_vec); + let mut compress_output_buf: Vec = vec![0; string_array.value_data().len()]; + let mut compress_offset_buf: Vec = vec![0; string_array.value_offsets().len()]; + let mut symbol_table = [0; FSST_SYMBOL_TABLE_SIZE]; + compress( + symbol_table.as_mut(), + string_array.value_data(), + string_array.value_offsets(), + &mut compress_output_buf, + &mut compress_offset_buf, + ) + .unwrap(); + (symbol_table, compress_output_buf, compress_offset_buf) + } + + // The decoder writes a full 8-byte word per code, so the output buffer must be at least 8x the + // compressed input. A buffer sized below 8x must be rejected rather than written out of bounds. + #[test_log::test(tokio::test)] + async fn test_decompress_rejects_undersized_output_buffer() { + let (symbol_table, compressed, compressed_offsets) = compress_paragraph(); + // Sanity check: this input actually engaged FSST compression (decoder_switch_on). + let st_info = u64::from_ne_bytes(symbol_table[..8].try_into().unwrap()); + assert!(st_info & (1 << 24) != 0, "expected decoder_switch_on input"); + + // One byte short of the 8x requirement must be rejected. + let mut too_small = vec![0u8; compressed.len() * 8 - 1]; + let mut out_offsets = vec![0i32; compressed_offsets.len()]; + let err = decompress( + &symbol_table, + &compressed, + &compressed_offsets, + &mut too_small, + &mut out_offsets, + ) + .unwrap_err(); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); + + // Exactly 8x is the tight bound and must succeed. + let mut exact = vec![0u8; compressed.len() * 8]; + let mut out_offsets = vec![0i32; compressed_offsets.len()]; + decompress( + &symbol_table, + &compressed, + &compressed_offsets, + &mut exact, + &mut out_offsets, + ) + .unwrap(); + } }