geoHeil commented on code in PR #12077:
URL: https://github.com/apache/arrow-rs/pull/12077#discussion_r4473035272
##########
arrow-ipc/src/compression.rs:
##########
@@ -274,6 +280,17 @@ impl CompressionCodec {
decompressed_size: usize,
context: &mut DecompressionContext,
) -> Result<Vec<u8>, ArrowError> {
+ let (max_expansion_ratio, codec_name) = match self {
+ CompressionCodec::Lz4Frame => (LZ4_MAX_EXPANSION_RATIO, "LZ4"),
+ CompressionCodec::Zstd(_) => (ZSTD_MAX_EXPANSION_RATIO, "ZSTD"),
+ };
+ // Refuse a size the input cannot expand to before allocating for it.
+ if decompressed_size > input.len().saturating_mul(max_expansion_ratio)
{
+ return Err(ArrowError::IpcError(format!(
+ "Advertised decompressed size of {decompressed_size} exceeds
what {} bytes of {codec_name} data can expand to",
+ input.len()
+ )));
+ }
Review Comment:
Thanks, that is simpler and it is the stronger bound. The ratio check only
tied the reservation to the input size: a 4 MiB ZSTD body could still ask for
128 GiB. Capping the reservation ties memory to the bytes actually
decompressed. Done in a4724f7, with the tests folded in 97f5965. The ratio
constants are gone:
- The reservation is `min(declared, MAX_PREALLOC_BYTES)`. That is the same
64 MiB cap `reader.rs` already uses for message bodies. Past the cap, the
buffer doubles, clamped at the declared length, only as the decoder fills it.
The existing `ret.len() != decompressed_size` check in `decompress` refuses a
length the data does not back.
- LZ4 already decodes through `lz4_flex::frame::FrameDecoder`, which is a
`Read`, so it grows freely.
- ZSTD needs a hybrid. `zstd::bulk::Decompressor::decompress` /
`decompress_to_buffer` is single-shot and writes into capacity reserved up
front, so it cannot grow. The bulk path stays for declared sizes up to the cap,
which makes it bounded and unchanged. Above the cap the buffer streams through
`zstd::stream::read::Decoder`, whose window allocation zstd caps at 2^27 bytes
by default.
Cost on the honest path:
- Buffers up to 64 MiB decode exactly as before.
- Larger buffers take about log2(n / 64 MiB) reallocations, and for ZSTD a
copy out of the stream window. That is the same trade-off already accepted for
message bodies.
- The final capacity equals the declared length.
Tests:
- `test_decompression_rejects_size_the_data_does_not_back` declares
`usize::MAX` over a few bytes of LZ4 and of ZSTD (a frame without a content
size). It expects the length-mismatch error. On main it panics with `capacity
overflow` for each codec.
- `test_decompression_beyond_prealloc` round-trips a 64 MiB + 1 buffer
through both codecs and checks the result has exactly its declared capacity.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
