swf: grow lzma output in one decoder pass

The on-demand buffer growth introduced for SWF decompression called
lzma_decompress again from byte zero each time the output buffer filled.
Because lzma-rs can buffer up to a dictionary before writing, a single
SWF could be decoded repeatedly as the buffer grew, amplifying CPU use.

This adds a depth-bounded Rust writer that asks C to grow the inspection
buffer as decompressed bytes are emitted.
This keeps a single lzma-rs stream alive, so the input is decoded only once.

Ticket: #8829

rust/src/lzma.rs

@@ -19,13 +19,79 @@
 
 use lzma_rs::decompress::{Options, Stream};
 use lzma_rs::error::Error;
-use std::io::{Cursor, ErrorKind, Write};
+use std::cell::Cell;
+use std::io::{self, ErrorKind, Write};
+use std::os::raw::c_void;
+use std::ptr;
+
+type LzmaOutputGrowFn = unsafe extern "C" fn(output: *mut c_void, min_size: u32) -> *mut u8;
+
+struct LzmaOutput<'a> {
+    /// Opaque C InspectionBuffer passed back to the growth callback.
+    output: *mut c_void,
+    /// Bytes reserved before the decompressed data, for example the FWS header.
+    output_offset: usize,
+    /// Maximum decompressed bytes to retain, excluding output_offset.
+    output_limit: usize,
+    /// Ensures that the C inspection buffer has the requested total capacity.
+    output_grow: LzmaOutputGrowFn,
+    /// Number of decompressed bytes already written after output_offset.
+    position: &'a Cell<usize>,
+    /// Records output-limit exhaustion before lzma-rs can wrap the writer error.
+    full: &'a Cell<bool>,
+    /// Records a failed C buffer expansion separately from other I/O errors.
+    alloc_failed: &'a Cell<bool>,
+}
+
+impl Write for LzmaOutput<'_> {
+    /// Accept a decompressed chunk from lzma-rs, grow the C inspection buffer as needed, and
+    /// copy no more than output_limit. The external Cells preserve the result if lzma-rs wraps
+    /// a short write or consumes the writer after an error.
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        if buf.is_empty() {
+            return Ok(0);
+        }
+
+        let position = self.position.get();
+        let written = buf.len().min(self.output_limit - position);
+        if written == 0 {
+            self.full.set(true);
+            return Ok(0);
+        }
+
+        let required = self.output_offset + position + written;
+        debug_assert!(required <= u32::MAX as usize);
+        let output = unsafe { (self.output_grow)(self.output, required as u32) };
+        if output.is_null() {
+            self.alloc_failed.set(true);
+            return Err(io::Error::other("failed to grow LZMA output"));
+        }
+
+        unsafe {
+            ptr::copy_nonoverlapping(
+                buf.as_ptr(),
+                output.add(self.output_offset + position),
+                written,
+            );
+        }
+        self.position.set(position + written);
+        if written < buf.len() {
+            self.full.set(true);
+        }
+        Ok(written)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
 
 /// Propagate lzma crate errors
 #[repr(C)]
 pub enum LzmaStatus {
     LzmaOk,
     LzmaOutputFull,
+    LzmaOutputAllocError,
     LzmaIoError,
     LzmaHeaderTooShortError,
     LzmaError,
@@ -64,12 +130,25 @@ impl From<std::io::Error> for LzmaStatus {
 /// Use the lzma algorithm to decompress a chunk of data.
 #[no_mangle]
 pub unsafe extern "C" fn lzma_decompress(
-    input: *const u8, input_len: &mut usize, output: *mut u8, output_len: &mut usize,
+    input: *const u8, input_len: &mut usize, output: *mut c_void, output_offset: u32,
+    output_limit: u32, output_len: &mut usize,
+    output_grow: unsafe extern "C" fn(output: *mut c_void, min_size: u32) -> *mut u8,
     memlimit: usize,
 ) -> LzmaStatus {
     let input = std::slice::from_raw_parts(input, *input_len);
-    let output = std::slice::from_raw_parts_mut(output, *output_len);
+    *output_len = 0;
-    let output = Cursor::new(output);
+    let output_position = Cell::new(0);
+    let output_full = Cell::new(false);
+    let output_alloc_failed = Cell::new(false);
+    let output = LzmaOutput {
+        output,
+        output_offset: output_offset as usize,
+        output_limit: output_limit as usize,
+        output_grow,
+        position: &output_position,
+        full: &output_full,
+        alloc_failed: &output_alloc_failed,
+    };
 
     let options = Options {
         memlimit: Some(memlimit),
@@ -80,15 +159,80 @@ pub unsafe extern "C" fn lzma_decompress(
     let mut stream = Stream::new_with_options(&options, output);
 
     if let Err(e) = stream.write_all(input) {
-        return e.into();
+        *output_len = output_position.get();
+        return if output_alloc_failed.get() {
+            LzmaStatus::LzmaOutputAllocError
+        } else if output_full.get() {
+            LzmaStatus::LzmaOutputFull
+        } else {
+            e.into()
+        };
     }
 
-    match stream.finish() {
+    let status = match stream.finish() {
-        Ok(output) => {
+        Ok(_) => LzmaStatus::LzmaOk,
-            *output_len = output.position() as usize;
+        Err(_) if output_alloc_failed.get() => LzmaStatus::LzmaOutputAllocError,
-            LzmaStatus::LzmaOk
+        Err(_) if output_full.get() => LzmaStatus::LzmaOutputFull,
-        }
         Err(e) => e.into(),
-    }
+    };
+    *output_len = output_position.get();
+    status
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::Cursor;
+
+    struct TestOutput {
+        data: Vec<u8>,
+        grow_count: usize,
+    }
+
+    unsafe extern "C" fn test_output_grow(output: *mut c_void, min_size: u32) -> *mut u8 {
+        let output = &mut *(output as *mut TestOutput);
+        output.grow_count += 1;
+        output.data.resize(min_size as usize, 0);
+        output.data.as_mut_ptr()
+    }
+
+    fn assert_output_full(dict_size: u32) {
+        let input = vec![b'A'; 8192];
+        let mut compressed = Vec::new();
+        lzma_rs::lzma_compress(&mut Cursor::new(&input), &mut compressed).unwrap();
+        compressed[1..5].copy_from_slice(&dict_size.to_le_bytes());
+        let mut input_len = compressed.len();
+        let mut output = TestOutput {
+            data: Vec::new(),
+            grow_count: 0,
+        };
+        let mut output_len = 0;
+        let status = unsafe {
+            lzma_decompress(
+                compressed.as_ptr(),
+                &mut input_len,
+                &mut output as *mut TestOutput as *mut c_void,
+                8,
+                64,
+                &mut output_len,
+                test_output_grow,
+                50_000_000,
+            )
+        };
+
+        assert!(matches!(status, LzmaStatus::LzmaOutputFull));
+        assert_eq!(output_len, 64);
+        assert!(output.grow_count > 0);
+        assert_eq!(&output.data[8..8 + output_len], vec![b'A'; output_len]);
+    }
+
+    #[test]
+    fn output_full_during_finish() {
+        assert_output_full(8 * 1024 * 1024);
+    }
+
+    #[test]
+    fn output_full_during_write() {
+        assert_output_full(4096);
+    }
+}

src/util-file-swf-decompression.c

@@ -145,9 +145,7 @@ int FileSwfZlibDecompression(DetectEngineThreadCtx *det_ctx, const uint8_t *comp
         if (result == Z_STREAM_END || *decompressed_data_produced == decompressed_data_limit)
             break;
 
-        if (result == Z_OK && infstream.avail_in == 0)
+        /* More output can be pending even after all input has been consumed. */
-            break;
-
         if (result == Z_OK && infstream.avail_out == 0) {
             if (!FileSwfDecompressionBufferGrow(det_ctx, out_buffer, decompressed_data_limit)) {
                 ret = 0;
@@ -156,6 +154,10 @@ int FileSwfZlibDecompression(DetectEngineThreadCtx *det_ctx, const uint8_t *comp
             continue;
         }
 
+        /* A retry with no input returns Z_BUF_ERROR when no output was pending. */
+        if ((result == Z_OK || result == Z_BUF_ERROR) && infstream.avail_in == 0)
+            break;
+
         if (result == Z_DATA_ERROR) {
             DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_Z_DATA_ERROR);
         } else if (result == Z_STREAM_ERROR) {
@@ -178,40 +180,33 @@ int FileSwfZlibDecompression(DetectEngineThreadCtx *det_ctx, const uint8_t *comp
  * | 4 bytes         | 4 bytes    | 4 bytes        | 5 bytes    | n bytes   | 6 bytes         |
  * | 'ZWS' + version | script len | compressed len | LZMA props | LZMA data | LZMA end marker |
  */
+static uint8_t *FileSwfLzmaOutputGrow(void *output, uint32_t min_size)
+{
+    return SCInspectionBufferCheckAndExpand(output, min_size);
+}
+
 int FileSwfLzmaDecompression(DetectEngineThreadCtx *det_ctx, const uint8_t *compressed_data,
         uint32_t compressed_data_len, InspectionBuffer *out_buffer,
         uint32_t decompressed_data_limit, uint32_t *decompressed_data_produced)
 {
-    int ret = 0;
     *decompressed_data_produced = 0;
 
-    while (true) {
     size_t inprocessed = compressed_data_len;
-        size_t outprocessed =
+    size_t outprocessed = 0;
-                FileSwfDecompressionOutputCapacity(out_buffer, decompressed_data_limit);
+    int ret = lzma_decompress(compressed_data, &inprocessed, out_buffer, SWF_HEADER_LEN,
-
+            decompressed_data_limit, &outprocessed, FileSwfLzmaOutputGrow,
-        ret = lzma_decompress(compressed_data, &inprocessed, out_buffer->buf + SWF_HEADER_LEN,
+            MAX_SWF_DECOMPRESSED_LEN);
-                &outprocessed, MAX_SWF_DECOMPRESSED_LEN);
-
-        if (ret == LzmaOk) {
     *decompressed_data_produced = (uint32_t)outprocessed;
-            ret = 1;
+
-            break;
+    if (ret == LzmaOk || ret == LzmaOutputFull) {
-        }
+        DEBUG_VALIDATE_BUG_ON(
-        if (ret == LzmaOutputFull) {
+                ret == LzmaOutputFull && *decompressed_data_produced != decompressed_data_limit);
-            *decompressed_data_produced = (uint32_t)outprocessed;
+        return 1;
-            if (*decompressed_data_produced == decompressed_data_limit) {
-                ret = 1;
-                break;
-            }
-            if (!FileSwfDecompressionBufferGrow(det_ctx, out_buffer, decompressed_data_limit)) {
-                ret = 0;
-                break;
-            }
-            continue;
     }
 
-        if (ret == LzmaIoError) {
+    if (ret == LzmaOutputAllocError) {
+        DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_NO_MEM);
+    } else if (ret == LzmaIoError) {
         DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_IO_ERROR);
     } else if (ret == LzmaHeaderTooShortError) {
         DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_HEADER_TOO_SHORT_ERROR);
@@ -226,9 +221,5 @@ int FileSwfLzmaDecompression(DetectEngineThreadCtx *det_ctx, const uint8_t *comp
     } else {
         DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_UNKNOWN_ERROR);
     }
-        ret = 0;
+    return 0;
-        break;
-    }
-
-    return ret;
 }