diff --git a/rust/src/lzma.rs b/rust/src/lzma.rs index e347b83eb5..b2d62ad1ab 100644 --- a/rust/src/lzma.rs +++ b/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, + /// Records output-limit exhaustion before lzma-rs can wrap the writer error. + full: &'a Cell, + /// Records a failed C buffer expansion separately from other I/O errors. + alloc_failed: &'a Cell, +} + +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 { + 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 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); - let output = Cursor::new(output); + *output_len = 0; + 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() { - Ok(output) => { - *output_len = output.position() as usize; - LzmaStatus::LzmaOk - } + let status = match stream.finish() { + Ok(_) => LzmaStatus::LzmaOk, + Err(_) if output_alloc_failed.get() => LzmaStatus::LzmaOutputAllocError, + 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, + 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); + } +} diff --git a/src/util-file-swf-decompression.c b/src/util-file-swf-decompression.c index 2038d642fb..b9a4bb5a66 100644 --- a/src/util-file-swf-decompression.c +++ b/src/util-file-swf-decompression.c @@ -145,10 +145,8 @@ 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) - break; - - if (result == Z_OK && infstream.avail_out == 0) { + /* More output can be pending even after all input has been consumed. */ + if (result == Z_OK && infstream.avail_out == 0) { if (!FileSwfDecompressionBufferGrow(det_ctx, out_buffer, decompressed_data_limit)) { ret = 0; break; @@ -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,57 +180,46 @@ 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 = - FileSwfDecompressionOutputCapacity(out_buffer, decompressed_data_limit); - - ret = lzma_decompress(compressed_data, &inprocessed, out_buffer->buf + SWF_HEADER_LEN, - &outprocessed, MAX_SWF_DECOMPRESSED_LEN); - - if (ret == LzmaOk) { - *decompressed_data_produced = (uint32_t)outprocessed; - ret = 1; - break; - } - if (ret == LzmaOutputFull) { - *decompressed_data_produced = (uint32_t)outprocessed; - 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) { - DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_IO_ERROR); - } else if (ret == LzmaHeaderTooShortError) { - DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_HEADER_TOO_SHORT_ERROR); - } else if (ret == LzmaError) { - DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_DECODER_ERROR); - } else if (ret == LzmaMemoryError) { - DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_MEMLIMIT_ERROR); - } else if (ret == LzmaXzError) { - /* We should not see XZ compressed SWF files */ - DEBUG_VALIDATE_BUG_ON(ret == LzmaXzError); - DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_XZ_ERROR); - } else { - DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_UNKNOWN_ERROR); - } - ret = 0; - break; + size_t inprocessed = compressed_data_len; + size_t outprocessed = 0; + int ret = lzma_decompress(compressed_data, &inprocessed, out_buffer, SWF_HEADER_LEN, + decompressed_data_limit, &outprocessed, FileSwfLzmaOutputGrow, + MAX_SWF_DECOMPRESSED_LEN); + *decompressed_data_produced = (uint32_t)outprocessed; + + if (ret == LzmaOk || ret == LzmaOutputFull) { + DEBUG_VALIDATE_BUG_ON( + ret == LzmaOutputFull && *decompressed_data_produced != decompressed_data_limit); + return 1; } - return ret; + 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); + } else if (ret == LzmaError) { + DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_DECODER_ERROR); + } else if (ret == LzmaMemoryError) { + DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_MEMLIMIT_ERROR); + } else if (ret == LzmaXzError) { + /* We should not see XZ compressed SWF files */ + DEBUG_VALIDATE_BUG_ON(ret == LzmaXzError); + DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_XZ_ERROR); + } else { + DetectEngineSetEvent(det_ctx, FILE_DECODER_EVENT_LZMA_UNKNOWN_ERROR); + } + return 0; }