1use std::io::{Read, Write};
2
3use crate::error::{Error, Result};
4
5const COMPRESSION_CODEC_MASK: i16 = 0x07;
6#[cfg(test)]
7const MAX_DECOMPRESSED_RECORD_BYTES: usize = 64 * 1024 * 1024;
8const XERIAL_SNAPPY_HEADER: [u8; 16] = [
9 0x82, b'S', b'N', b'A', b'P', b'P', b'Y', 0, 0, 0, 0, 1, 0, 0, 0, 1,
10];
11const XERIAL_SNAPPY_MAGIC: [u8; 8] = [0x82, b'S', b'N', b'A', b'P', b'P', b'Y', 0];
12const XERIAL_SNAPPY_BLOCK_BYTES: usize = 32 * 1024;
13const ZSTD_MAGIC: [u8; 4] = [0x28, 0xb5, 0x2f, 0xfd];
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum RecordBatchCompression {
17 None,
18 Gzip,
19 Snappy,
20 Lz4,
21 Zstd,
22}
23
24impl RecordBatchCompression {
25 pub fn from_attributes(attributes: i16) -> Result<Self> {
26 match attributes & COMPRESSION_CODEC_MASK {
27 0 => Ok(Self::None),
28 1 => Ok(Self::Gzip),
29 2 => Ok(Self::Snappy),
30 3 => Ok(Self::Lz4),
31 4 => Ok(Self::Zstd),
32 code => Err(Error::UnsupportedVersion {
33 kind: "record batch compression codec",
34 version: code,
35 }),
36 }
37 }
38
39 pub fn attributes(self) -> i16 {
40 match self {
41 Self::None => 0,
42 Self::Gzip => 1,
43 Self::Snappy => 2,
44 Self::Lz4 => 3,
45 Self::Zstd => 4,
46 }
47 }
48
49 pub fn name(self) -> &'static str {
50 match self {
51 Self::None => "none",
52 Self::Gzip => "gzip",
53 Self::Snappy => "snappy",
54 Self::Lz4 => "lz4",
55 Self::Zstd => "zstd",
56 }
57 }
58
59 pub fn is_compressed(self) -> bool {
60 self != Self::None
61 }
62}
63
64pub(crate) fn compress_record_batch_records(
65 compression: RecordBatchCompression,
66 records: &[u8],
67) -> Result<Vec<u8>> {
68 match compression {
69 RecordBatchCompression::None => Ok(records.to_vec()),
70 RecordBatchCompression::Gzip => gzip_compress(records),
71 RecordBatchCompression::Snappy => snappy_compress(records),
72 RecordBatchCompression::Lz4 => lz4_compress(records),
73 RecordBatchCompression::Zstd => zstd_compress(records),
74 }
75}
76
77#[cfg(test)]
78pub(crate) fn decompress_record_batch_records(
79 compression: RecordBatchCompression,
80 records: &[u8],
81) -> Result<Vec<u8>> {
82 decompress_record_batch_records_with_limit(compression, records, MAX_DECOMPRESSED_RECORD_BYTES)
83}
84
85pub(crate) fn decompress_record_batch_records_with_limit(
86 compression: RecordBatchCompression,
87 records: &[u8],
88 max_decompressed_bytes: usize,
89) -> Result<Vec<u8>> {
90 match compression {
91 RecordBatchCompression::None => {
92 ensure_decompressed_output_limit(records.len(), max_decompressed_bytes)?;
93 Ok(records.to_vec())
94 }
95 RecordBatchCompression::Gzip => gzip_decompress(records, max_decompressed_bytes),
96 RecordBatchCompression::Snappy => snappy_decompress(records, max_decompressed_bytes),
97 RecordBatchCompression::Lz4 => lz4_decompress_with_limit(records, max_decompressed_bytes),
98 RecordBatchCompression::Zstd => zstd_decompress_with_limit(records, max_decompressed_bytes),
99 }
100}
101
102fn gzip_compress(records: &[u8]) -> Result<Vec<u8>> {
103 let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
104 encoder
105 .write_all(records)
106 .map_err(|error| compression_error("gzip", error))?;
107 encoder
108 .finish()
109 .map_err(|error| compression_error("gzip", error))
110}
111
112fn gzip_decompress(records: &[u8], max_decompressed_bytes: usize) -> Result<Vec<u8>> {
113 let decoder = flate2::read::GzDecoder::new(records);
114 let read_limit = u64::try_from(max_decompressed_bytes)
115 .map_err(|_| Error::LengthOverflow("decompressed record batch"))?
116 .checked_add(1)
117 .ok_or(Error::LengthOverflow("decompressed record batch"))?;
118 let mut limited = decoder.take(read_limit);
119 let mut output = Vec::new();
120 limited
121 .read_to_end(&mut output)
122 .map_err(|error| compression_error("gzip", error))?;
123 ensure_decompressed_output_limit(output.len(), max_decompressed_bytes)?;
124 Ok(output)
125}
126
127fn snappy_compress(records: &[u8]) -> Result<Vec<u8>> {
128 let mut output = Vec::with_capacity(
129 XERIAL_SNAPPY_HEADER
130 .len()
131 .checked_add(records.len())
132 .ok_or(Error::LengthOverflow("snappy record batch"))?,
133 );
134 output.extend_from_slice(&XERIAL_SNAPPY_HEADER);
135
136 let mut encoder = snap::raw::Encoder::new();
137 for block in records.chunks(XERIAL_SNAPPY_BLOCK_BYTES) {
138 let compressed = encoder
139 .compress_vec(block)
140 .map_err(|error| snappy_error(error.to_string()))?;
141 let compressed_len = u32::try_from(compressed.len())
142 .map_err(|_| Error::LengthOverflow("snappy record batch block"))?;
143 output.extend_from_slice(&compressed_len.to_be_bytes());
144 output.extend_from_slice(&compressed);
145 }
146 Ok(output)
147}
148
149fn snappy_decompress(records: &[u8], max_decompressed_bytes: usize) -> Result<Vec<u8>> {
150 if records.starts_with(&XERIAL_SNAPPY_MAGIC) {
151 return snappy_xerial_decompress(records, max_decompressed_bytes);
152 }
153
154 ensure_snappy_output_limit(records, 0, max_decompressed_bytes)?;
155 snap::raw::Decoder::new()
156 .decompress_vec(records)
157 .map_err(|error| snappy_error(error.to_string()))
158}
159
160fn snappy_xerial_decompress(records: &[u8], max_decompressed_bytes: usize) -> Result<Vec<u8>> {
161 if !records.starts_with(&XERIAL_SNAPPY_HEADER) {
162 return Err(snappy_error(
163 "invalid or unsupported xerial framing header".to_owned(),
164 ));
165 }
166
167 let mut output = Vec::new();
168 let mut position = XERIAL_SNAPPY_HEADER.len();
169 let mut decoder = snap::raw::Decoder::new();
170 while position < records.len() {
171 let length_end = position
172 .checked_add(4)
173 .ok_or(Error::LengthOverflow("snappy record batch block"))?;
174 let length_bytes = records
175 .get(position..length_end)
176 .ok_or_else(|| snappy_error("truncated xerial block length".to_owned()))?;
177 let block_len = u32::from_be_bytes([
178 length_bytes[0],
179 length_bytes[1],
180 length_bytes[2],
181 length_bytes[3],
182 ]) as usize;
183 position = length_end;
184
185 let block_end = position
186 .checked_add(block_len)
187 .ok_or(Error::LengthOverflow("snappy record batch block"))?;
188 let block = records
189 .get(position..block_end)
190 .ok_or_else(|| snappy_error("xerial block extends past input".to_owned()))?;
191 ensure_snappy_output_limit(block, output.len(), max_decompressed_bytes)?;
192 let decompressed = decoder
193 .decompress_vec(block)
194 .map_err(|error| snappy_error(error.to_string()))?;
195 output.extend_from_slice(&decompressed);
196 position = block_end;
197 }
198 Ok(output)
199}
200
201fn ensure_snappy_output_limit(
202 block: &[u8],
203 already_decompressed: usize,
204 max_decompressed_bytes: usize,
205) -> Result<()> {
206 let block_len =
207 snap::raw::decompress_len(block).map_err(|error| snappy_error(error.to_string()))?;
208 let total_len = already_decompressed
209 .checked_add(block_len)
210 .ok_or(Error::LengthOverflow("decompressed record batch"))?;
211 ensure_decompressed_output_limit(total_len, max_decompressed_bytes)
212}
213
214fn lz4_compress(records: &[u8]) -> Result<Vec<u8>> {
215 let mut settings = lz_fear::CompressionSettings::default();
216 settings
217 .independent_blocks(true)
218 .block_checksums(false)
219 .content_checksum(false)
220 .block_size(64 * 1024);
221 let mut output = Vec::new();
222 settings
223 .compress(records, &mut output)
224 .map_err(|error| compression_reason("lz4", error.to_string()))?;
225 Ok(output)
226}
227
228fn lz4_decompress_with_limit(records: &[u8], max_decompressed_bytes: usize) -> Result<Vec<u8>> {
229 let decoder = lz_fear::LZ4FrameReader::new(records)
230 .map_err(|error| compression_reason("lz4", error.to_string()))?
231 .into_read();
232 let read_limit = u64::try_from(max_decompressed_bytes)
233 .map_err(|_| Error::LengthOverflow("decompressed record batch"))?
234 .checked_add(1)
235 .ok_or(Error::LengthOverflow("decompressed record batch"))?;
236 let mut limited = decoder.take(read_limit);
237 let mut output = Vec::new();
238 limited
239 .read_to_end(&mut output)
240 .map_err(|error| compression_error("lz4", error))?;
241 ensure_decompressed_output_limit(output.len(), max_decompressed_bytes)?;
242 Ok(output)
243}
244
245fn zstd_compress(records: &[u8]) -> Result<Vec<u8>> {
246 std::panic::catch_unwind(|| {
247 ruzstd::encoding::compress_to_vec(records, ruzstd::encoding::CompressionLevel::Fastest)
248 })
249 .map_err(|_| compression_reason("zstd", "encoder panicked".to_owned()))
250}
251
252fn zstd_decompress_with_limit(records: &[u8], max_decompressed_bytes: usize) -> Result<Vec<u8>> {
253 ensure_zstd_frame_limits(records, max_decompressed_bytes)?;
254 std::panic::catch_unwind(|| {
255 let decoder = ruzstd::decoding::StreamingDecoder::new(records)
256 .map_err(|error| compression_reason("zstd", error.to_string()))?;
257 let read_limit = u64::try_from(max_decompressed_bytes)
258 .map_err(|_| Error::LengthOverflow("decompressed record batch"))?
259 .checked_add(1)
260 .ok_or(Error::LengthOverflow("decompressed record batch"))?;
261 let mut limited = decoder.take(read_limit);
262 let mut output = Vec::new();
263 limited
264 .read_to_end(&mut output)
265 .map_err(|error| compression_error("zstd", error))?;
266 ensure_decompressed_output_limit(output.len(), max_decompressed_bytes)?;
267 Ok(output)
268 })
269 .map_err(|_| compression_reason("zstd", "decoder panicked".to_owned()))?
270}
271
272fn ensure_zstd_frame_limits(records: &[u8], max_decompressed_bytes: usize) -> Result<()> {
273 if !records.starts_with(&ZSTD_MAGIC) {
274 return Err(compression_reason("zstd", "invalid frame magic".to_owned()));
275 }
276
277 let descriptor = *records
278 .get(ZSTD_MAGIC.len())
279 .ok_or_else(|| compression_reason("zstd", "truncated frame header".to_owned()))?;
280 let single_segment = descriptor & 0x20 != 0;
281 let mut position = ZSTD_MAGIC.len() + 1;
282
283 let window_size = if single_segment {
284 None
285 } else {
286 let window_descriptor = *records
287 .get(position)
288 .ok_or_else(|| compression_reason("zstd", "truncated window descriptor".to_owned()))?;
289 position += 1;
290 let exponent = u64::from(window_descriptor >> 3);
291 let mantissa = u64::from(window_descriptor & 0x07);
292 let window_base = 1u64 << (10 + exponent);
293 Some(window_base + (window_base / 8) * mantissa)
294 };
295
296 let dictionary_id_bytes = match descriptor & 0x03 {
297 0 => 0,
298 1 => 1,
299 2 => 2,
300 3 => 4,
301 _ => unreachable!(),
302 };
303 position = position
304 .checked_add(dictionary_id_bytes)
305 .ok_or(Error::LengthOverflow("zstd frame header"))?;
306
307 let content_size_bytes = match descriptor >> 6 {
308 0 if single_segment => 1,
309 0 => 0,
310 1 => 2,
311 2 => 4,
312 3 => 8,
313 _ => unreachable!(),
314 };
315 let content_size = read_zstd_little_endian(records, position, content_size_bytes)?;
316 let content_size = if content_size_bytes == 2 {
317 content_size + 256
318 } else {
319 content_size
320 };
321 let required_window = window_size.unwrap_or(content_size);
322 let actual = required_window.max(content_size);
323 let max = u64::try_from(max_decompressed_bytes)
324 .map_err(|_| Error::LengthOverflow("decompressed record batch"))?;
325 if actual > max {
326 return Err(Error::LimitExceeded {
327 kind: "decompressed record batch bytes",
328 actual: usize::try_from(actual).unwrap_or(usize::MAX),
329 max: max_decompressed_bytes,
330 });
331 }
332 Ok(())
333}
334
335fn ensure_decompressed_output_limit(actual: usize, max: usize) -> Result<()> {
336 if actual > max {
337 return Err(Error::LimitExceeded {
338 kind: "decompressed record batch bytes",
339 actual,
340 max,
341 });
342 }
343 Ok(())
344}
345
346fn read_zstd_little_endian(records: &[u8], position: usize, length: usize) -> Result<u64> {
347 let end = position
348 .checked_add(length)
349 .ok_or(Error::LengthOverflow("zstd frame header"))?;
350 let bytes = records
351 .get(position..end)
352 .ok_or_else(|| compression_reason("zstd", "truncated frame header".to_owned()))?;
353 Ok(bytes.iter().enumerate().fold(0u64, |value, (index, byte)| {
354 value | (u64::from(*byte) << (index * 8))
355 }))
356}
357
358fn compression_error(codec: &'static str, error: std::io::Error) -> Error {
359 compression_reason(codec, error.to_string())
360}
361
362fn compression_reason(codec: &'static str, reason: String) -> Error {
363 Error::Compression { codec, reason }
364}
365
366fn snappy_error(reason: String) -> Error {
367 compression_reason("snappy", reason)
368}
369
370#[cfg(test)]
371#[allow(clippy::unwrap_used)]
372mod tests {
373 use super::{
374 compress_record_batch_records, decompress_record_batch_records,
375 decompress_record_batch_records_with_limit, lz4_decompress_with_limit,
376 zstd_decompress_with_limit, RecordBatchCompression, MAX_DECOMPRESSED_RECORD_BYTES,
377 XERIAL_SNAPPY_HEADER, ZSTD_MAGIC,
378 };
379 use crate::error::Error;
380
381 #[test]
382 fn snappy_xerial_roundtrip_spans_multiple_blocks() {
383 let records = vec![b'x'; 70 * 1024];
384
385 let compressed =
386 compress_record_batch_records(RecordBatchCompression::Snappy, &records).unwrap();
387 let decompressed =
388 decompress_record_batch_records(RecordBatchCompression::Snappy, &compressed).unwrap();
389
390 assert!(compressed.starts_with(&XERIAL_SNAPPY_HEADER));
391 assert_eq!(decompressed, records);
392 }
393
394 #[test]
395 fn snappy_decoder_accepts_raw_blocks() {
396 let records = b"kafka record batch";
397 let compressed = snap::raw::Encoder::new().compress_vec(records).unwrap();
398
399 assert_eq!(
400 decompress_record_batch_records(RecordBatchCompression::Snappy, &compressed).unwrap(),
401 records
402 );
403 }
404
405 #[test]
406 fn snappy_decoder_rejects_declared_output_over_limit() {
407 let mut declared_len = MAX_DECOMPRESSED_RECORD_BYTES + 1;
408 let mut hostile_block = Vec::new();
409 while declared_len >= 0x80 {
410 hostile_block.push((declared_len as u8) | 0x80);
411 declared_len >>= 7;
412 }
413 hostile_block.push(declared_len as u8);
414
415 assert_eq!(
416 decompress_record_batch_records(RecordBatchCompression::Snappy, &hostile_block)
417 .unwrap_err(),
418 Error::LimitExceeded {
419 kind: "decompressed record batch bytes",
420 actual: MAX_DECOMPRESSED_RECORD_BYTES + 1,
421 max: MAX_DECOMPRESSED_RECORD_BYTES,
422 }
423 );
424 }
425
426 #[test]
427 fn gzip_decoder_honors_custom_output_limit() {
428 let records = vec![b'x'; 1024];
429 let compressed =
430 compress_record_batch_records(RecordBatchCompression::Gzip, &records).unwrap();
431
432 assert_eq!(
433 decompress_record_batch_records_with_limit(
434 RecordBatchCompression::Gzip,
435 &compressed,
436 64,
437 )
438 .unwrap_err(),
439 Error::LimitExceeded {
440 kind: "decompressed record batch bytes",
441 actual: 65,
442 max: 64,
443 }
444 );
445 }
446
447 #[test]
448 fn snappy_decoder_honors_custom_output_limit() {
449 let records = vec![b'x'; 1024];
450 let compressed =
451 compress_record_batch_records(RecordBatchCompression::Snappy, &records).unwrap();
452
453 assert_eq!(
454 decompress_record_batch_records_with_limit(
455 RecordBatchCompression::Snappy,
456 &compressed,
457 64,
458 )
459 .unwrap_err(),
460 Error::LimitExceeded {
461 kind: "decompressed record batch bytes",
462 actual: 1024,
463 max: 64,
464 }
465 );
466 }
467
468 #[test]
469 fn snappy_decoder_rejects_truncated_xerial_block() {
470 let mut compressed = XERIAL_SNAPPY_HEADER.to_vec();
471 compressed.extend_from_slice(&10u32.to_be_bytes());
472 compressed.extend_from_slice(&[1, 2, 3]);
473
474 assert!(matches!(
475 decompress_record_batch_records(RecordBatchCompression::Snappy, &compressed),
476 Err(Error::Compression {
477 codec: "snappy",
478 ..
479 })
480 ));
481 }
482
483 #[test]
484 fn lz4_frame_roundtrips_with_kafka_magic() {
485 let records = vec![b'x'; 70 * 1024];
486
487 let compressed =
488 compress_record_batch_records(RecordBatchCompression::Lz4, &records).unwrap();
489 let decompressed =
490 decompress_record_batch_records(RecordBatchCompression::Lz4, &compressed).unwrap();
491
492 assert_eq!(&compressed[..4], &[0x04, 0x22, 0x4d, 0x18]);
493 assert_eq!(decompressed, records);
494 }
495
496 #[test]
497 fn lz4_decoder_rejects_output_over_limit() {
498 let records = vec![b'x'; 1024];
499 let compressed =
500 compress_record_batch_records(RecordBatchCompression::Lz4, &records).unwrap();
501
502 assert_eq!(
503 lz4_decompress_with_limit(&compressed, 64).unwrap_err(),
504 Error::LimitExceeded {
505 kind: "decompressed record batch bytes",
506 actual: 65,
507 max: 64,
508 }
509 );
510 }
511
512 #[test]
513 fn lz4_decoder_rejects_malformed_frame() {
514 assert!(matches!(
515 decompress_record_batch_records(RecordBatchCompression::Lz4, b"not an lz4 frame"),
516 Err(Error::Compression { codec: "lz4", .. })
517 ));
518 }
519
520 #[test]
521 fn zstd_frame_roundtrips_with_kafka_magic() {
522 let records = vec![b'x'; 140 * 1024];
523
524 let compressed =
525 compress_record_batch_records(RecordBatchCompression::Zstd, &records).unwrap();
526 let decompressed =
527 decompress_record_batch_records(RecordBatchCompression::Zstd, &compressed).unwrap();
528
529 assert_eq!(&compressed[..4], &ZSTD_MAGIC);
530 assert_eq!(decompressed, records);
531 }
532
533 #[test]
534 fn zstd_decoder_rejects_output_over_limit() {
535 let records = vec![b'x'; 1024];
536 let compressed =
537 compress_record_batch_records(RecordBatchCompression::Zstd, &records).unwrap();
538
539 assert!(matches!(
540 zstd_decompress_with_limit(&compressed, 64),
541 Err(Error::LimitExceeded {
542 kind: "decompressed record batch bytes",
543 max: 64,
544 ..
545 })
546 ));
547 }
548
549 #[test]
550 fn zstd_decoder_rejects_declared_window_over_limit() {
551 let frame = [ZSTD_MAGIC.as_slice(), &[0, 0]].concat();
552
553 assert!(matches!(
554 zstd_decompress_with_limit(&frame, 64),
555 Err(Error::LimitExceeded {
556 kind: "decompressed record batch bytes",
557 max: 64,
558 ..
559 })
560 ));
561 }
562
563 #[test]
564 fn zstd_decoder_rejects_malformed_frame() {
565 assert!(matches!(
566 decompress_record_batch_records(RecordBatchCompression::Zstd, b"not a zstd frame"),
567 Err(Error::Compression { codec: "zstd", .. })
568 ));
569 }
570}