kafrust_protocol/
record_batch.rs

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

kafrust_protocol/
record_batch.rs
