apache_avro/
codec.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Logic for all supported compression codecs in Avro.
19
20use crate::{AvroResult, Error, error::Details, types::Value};
21use strum::{EnumIter, EnumString, IntoStaticStr};
22
23/// Settings for the `Deflate` codec.
24#[derive(Clone, Copy, Eq, PartialEq, Debug)]
25pub struct DeflateSettings {
26    compression_level: miniz_oxide::deflate::CompressionLevel,
27}
28
29impl DeflateSettings {
30    pub fn new(compression_level: miniz_oxide::deflate::CompressionLevel) -> Self {
31        DeflateSettings { compression_level }
32    }
33
34    /// Get the compression level as a `u8`, note that this means the [`miniz_oxide::deflate::CompressionLevel::DefaultCompression`] variant
35    /// will appear as `255`, this is normalized by the [`miniz_oxide`] crate later.
36    pub fn compression_level(&self) -> u8 {
37        self.compression_level as u8
38    }
39}
40
41impl Default for DeflateSettings {
42    /// Default compression level is `miniz_oxide::deflate::CompressionLevel::DefaultCompression`.
43    fn default() -> Self {
44        Self::new(miniz_oxide::deflate::CompressionLevel::DefaultCompression)
45    }
46}
47
48/// The compression codec used to compress blocks.
49#[derive(Clone, Copy, Debug, Eq, PartialEq, EnumIter, EnumString, IntoStaticStr)]
50#[strum(serialize_all = "kebab_case")]
51pub enum Codec {
52    /// The `Null` codec simply passes through data uncompressed.
53    Null,
54    /// The `Deflate` codec writes the data block using the deflate algorithm
55    /// as specified in RFC 1951, and typically implemented using the zlib library.
56    /// Note that this format (unlike the "zlib format" in RFC 1950) does not have a checksum.
57    Deflate(DeflateSettings),
58    #[cfg(feature = "snappy")]
59    /// The `Snappy` codec uses Google's [Snappy](http://google.github.io/snappy/)
60    /// compression library. Each compressed block is followed by the 4-byte, big-endian
61    /// CRC32 checksum of the uncompressed data in the block.
62    Snappy,
63    #[cfg(feature = "zstandard")]
64    /// The `Zstandard` codec uses Facebook's [Zstandard](https://facebook.github.io/zstd/)
65    Zstandard(zstandard::ZstandardSettings),
66    #[cfg(feature = "bzip")]
67    /// The `BZip2` codec uses [BZip2](https://sourceware.org/bzip2/)
68    /// compression library.
69    Bzip2(bzip::Bzip2Settings),
70    #[cfg(feature = "xz")]
71    /// The `Xz` codec uses [Xz utils](https://tukaani.org/xz/)
72    /// compression library.
73    Xz(xz::XzSettings),
74}
75
76impl From<Codec> for Value {
77    fn from(value: Codec) -> Self {
78        Self::Bytes(<&str>::from(value).as_bytes().to_vec())
79    }
80}
81
82impl Codec {
83    /// Compress a stream of bytes in-place.
84    pub fn compress(self, stream: &mut Vec<u8>) -> AvroResult<()> {
85        match self {
86            Codec::Null => (),
87            Codec::Deflate(settings) => {
88                let compressed =
89                    miniz_oxide::deflate::compress_to_vec(stream, settings.compression_level());
90                *stream = compressed;
91            }
92            #[cfg(feature = "snappy")]
93            Codec::Snappy => {
94                let mut encoded: Vec<u8> = vec![0; snap::raw::max_compress_len(stream.len())];
95                let compressed_size = snap::raw::Encoder::new()
96                    .compress(&stream[..], &mut encoded[..])
97                    .map_err(Details::SnappyCompress)?;
98
99                let checksum = crc_fast::crc32_iso_hdlc(&stream[..]);
100                let checksum_as_bytes = checksum.to_be_bytes();
101                let checksum_len = checksum_as_bytes.len();
102                encoded.truncate(compressed_size + checksum_len);
103                encoded[compressed_size..].copy_from_slice(&checksum_as_bytes);
104
105                *stream = encoded;
106            }
107            #[cfg(feature = "zstandard")]
108            Codec::Zstandard(settings) => {
109                use std::io::Write;
110                let mut encoder = zstd::Encoder::new(Vec::new(), settings.compression_level as i32)
111                    .map_err(Details::ZstdCompress)?;
112                encoder.write_all(stream).map_err(Details::ZstdCompress)?;
113                *stream = encoder.finish().map_err(Details::ZstdCompress)?;
114            }
115            #[cfg(feature = "bzip")]
116            Codec::Bzip2(settings) => {
117                use bzip2::read::BzEncoder;
118                use std::io::Read;
119
120                let mut encoder = BzEncoder::new(&stream[..], settings.compression());
121                let mut buffer = Vec::new();
122                encoder
123                    .read_to_end(&mut buffer)
124                    .unwrap_or_else(|_| unreachable!("No I/O errors possible with Vec<u8>"));
125                *stream = buffer;
126            }
127            #[cfg(feature = "xz")]
128            Codec::Xz(settings) => {
129                use liblzma::read::XzEncoder;
130                use std::io::Read;
131
132                let mut encoder = XzEncoder::new(&stream[..], settings.compression_level as u32);
133                let mut buffer = Vec::new();
134                encoder
135                    .read_to_end(&mut buffer)
136                    .unwrap_or_else(|_| unreachable!("No I/O errors possible with Vec<u8>"));
137                *stream = buffer;
138            }
139        };
140
141        Ok(())
142    }
143
144    /// Decompress a stream of bytes in-place.
145    pub fn decompress(self, stream: &mut Vec<u8>) -> AvroResult<()> {
146        // Cap the decompressed output at the configured allocation budget so a
147        // small compressed block cannot inflate to an enormous buffer (a
148        // "decompression bomb") and exhaust memory.
149        let max_bytes =
150            crate::util::max_allocation_bytes(crate::util::DEFAULT_MAX_ALLOCATION_BYTES);
151        *stream = match self {
152            Codec::Null => return Ok(()),
153            Codec::Deflate(_settings) => miniz_oxide::inflate::decompress_to_vec_with_limit(stream, max_bytes).map_err(|e| {
154                use std::io::ErrorKind;
155                use miniz_oxide::inflate::TINFLStatus;
156
157                let details = match e.status {
158                    TINFLStatus::FailedCannotMakeProgress | TINFLStatus::NeedsMoreInput => Details::DeflateDecompress(ErrorKind::UnexpectedEof.into()),
159                    TINFLStatus::Adler32Mismatch | TINFLStatus::Failed | TINFLStatus::BadParam => Details::DeflateDecompress(ErrorKind::InvalidData.into()),
160                    TINFLStatus::Done => Details::DeflateDecompress(std::io::Error::other("Unexpected error: miniz_oxide reported an error with a success status. Please report this to avro-rs developers.")),
161                    // Output is larger than max allocation allowed
162                    TINFLStatus::HasMoreOutput => Details::MemoryAllocation {
163                        desired: None,
164                        maximum: max_bytes,
165                    },
166                    other => Details::DeflateDecompress(std::io::Error::other(format!("Unexpected error: {other:?}")))
167                };
168                Error::new(details)
169            })?,
170            #[cfg(feature = "snappy")]
171            Codec::Snappy => {
172                // The block ends with a 4-byte CRC32; a truncated/corrupt block
173                // shorter than that must error rather than underflow the slice.
174                let data_end = stream
175                    .len()
176                    .checked_sub(4)
177                    .ok_or(Details::BadSnappyLength(stream.len()))?;
178                let decompressed_size = snap::raw::decompress_len(&stream[..data_end])
179                    .map_err(Details::GetSnappyDecompressLen)?;
180                // The decompressed size is taken from the (untrusted) block
181                // header, so bound it before allocating for it.
182                let decompressed_size = crate::util::safe_len(decompressed_size)?;
183                let mut decoded = vec![0; decompressed_size];
184                snap::raw::Decoder::new()
185                    .decompress(&stream[..data_end], &mut decoded[..])
186                    .map_err(Details::SnappyDecompress)?;
187
188                let mut last_four: [u8; 4] = [0; 4];
189                last_four.copy_from_slice(&stream[data_end..]);
190                let expected: u32 = u32::from_be_bytes(last_four);
191
192                let actual = crc_fast::crc32_iso_hdlc(&decoded);
193
194                if expected != actual {
195                    return Err(Details::SnappyCrc32{expected, actual}.into());
196                }
197                decoded
198            }
199            #[cfg(feature = "zstandard")]
200            Codec::Zstandard(_settings) => {
201                use std::io::{BufReader, Read};
202                use zstd::zstd_safe;
203
204                let mut decoded = Vec::new();
205                let buffer_size = zstd_safe::DCtx::in_size();
206                let buffer = BufReader::with_capacity(buffer_size, &stream[..]);
207                let decoder = zstd::Decoder::new(buffer).map_err(Details::ZstdDecompress)?;
208                // Read one byte past the budget so an output that exactly fills
209                // it is allowed, while a larger (bomb) output is detected.
210                decoder
211                    .take((max_bytes as u64).saturating_add(1))
212                    .read_to_end(&mut decoded)
213                    .map_err(Details::ZstdDecompress)?;
214                if decoded.len() > max_bytes {
215                    return Err(Details::MemoryAllocation { desired: None, maximum: max_bytes }.into());
216                }
217                decoded
218            }
219            #[cfg(feature = "bzip")]
220            Codec::Bzip2(_) => {
221                use bzip2::read::BzDecoder;
222                use std::io::Read;
223
224                let mut decoded = Vec::new();
225                BzDecoder::new(&stream[..])
226                    .take((max_bytes as u64).saturating_add(1))
227                    .read_to_end(&mut decoded)
228                    .map_err(Details::Bzip2Decompress)?;
229                if decoded.len() > max_bytes {
230                    return Err(Details::MemoryAllocation { desired: None, maximum: max_bytes }.into());
231                }
232                decoded
233            }
234            #[cfg(feature = "xz")]
235            Codec::Xz(_) => {
236                use liblzma::read::XzDecoder;
237                use std::io::Read;
238
239                let mut decoded: Vec<u8> = Vec::new();
240                XzDecoder::new(&stream[..])
241                    .take((max_bytes as u64).saturating_add(1))
242                    .read_to_end(&mut decoded)
243                    .map_err(Details::XzDecompress)?;
244                if decoded.len() > max_bytes {
245                    return Err(Details::MemoryAllocation { desired: None, maximum: max_bytes }.into());
246                }
247                decoded
248            }
249        };
250        Ok(())
251    }
252}
253
254#[cfg(feature = "bzip")]
255pub mod bzip {
256    use bzip2::Compression;
257
258    #[derive(Clone, Copy, Eq, PartialEq, Debug)]
259    pub struct Bzip2Settings {
260        pub compression_level: u8,
261    }
262
263    impl Bzip2Settings {
264        pub fn new(compression_level: u8) -> Self {
265            Self { compression_level }
266        }
267
268        pub(crate) fn compression(&self) -> Compression {
269            Compression::new(self.compression_level as u32)
270        }
271    }
272
273    impl Default for Bzip2Settings {
274        fn default() -> Self {
275            Bzip2Settings::new(Compression::best().level() as u8)
276        }
277    }
278}
279
280#[cfg(feature = "zstandard")]
281pub mod zstandard {
282    #[derive(Clone, Copy, Eq, PartialEq, Debug)]
283    pub struct ZstandardSettings {
284        pub compression_level: u8,
285    }
286
287    impl ZstandardSettings {
288        pub fn new(compression_level: u8) -> Self {
289            Self { compression_level }
290        }
291    }
292
293    impl Default for ZstandardSettings {
294        fn default() -> Self {
295            Self::new(0)
296        }
297    }
298}
299
300#[cfg(feature = "xz")]
301pub mod xz {
302    #[derive(Clone, Copy, Eq, PartialEq, Debug)]
303    pub struct XzSettings {
304        pub compression_level: u8,
305    }
306
307    impl XzSettings {
308        pub fn new(compression_level: u8) -> Self {
309            Self { compression_level }
310        }
311    }
312
313    impl Default for XzSettings {
314        fn default() -> Self {
315            XzSettings::new(9)
316        }
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323    use apache_avro_test_helper::TestResult;
324    use miniz_oxide::deflate::CompressionLevel;
325    use pretty_assertions::{assert_eq, assert_ne};
326
327    const INPUT: &[u8] = b"theanswertolifetheuniverseandeverythingis42theanswertolifetheuniverseandeverythingis4theanswertolifetheuniverseandeverythingis2";
328
329    #[test]
330    fn null_compress_and_decompress() -> TestResult {
331        let codec = Codec::Null;
332        let mut stream = INPUT.to_vec();
333        codec.compress(&mut stream)?;
334        assert_eq!(INPUT, stream.as_slice());
335        codec.decompress(&mut stream)?;
336        assert_eq!(INPUT, stream.as_slice());
337        Ok(())
338    }
339
340    #[test]
341    fn deflate_compress_and_decompress() -> TestResult {
342        compress_and_decompress(Codec::Deflate(DeflateSettings::new(
343            CompressionLevel::BestCompression,
344        )))
345    }
346
347    #[cfg(feature = "snappy")]
348    #[test]
349    fn snappy_compress_and_decompress() -> TestResult {
350        compress_and_decompress(Codec::Snappy)
351    }
352
353    #[cfg(feature = "snappy")]
354    #[test]
355    fn avro_rs_660_snappy_compress_writes_crc32_iso_hdlc_trailer() -> TestResult {
356        let mut stream = b"123456789".to_vec();
357        Codec::Snappy.compress(&mut stream)?;
358
359        assert_eq!(stream[stream.len() - 4..], 0xcbf4_3926_u32.to_be_bytes());
360
361        Ok(())
362    }
363
364    #[cfg(feature = "snappy")]
365    #[test]
366    fn snappy_decompress_short_block_errors_without_panicking() {
367        // A block shorter than the trailing 4-byte CRC must return an error
368        // rather than underflowing `stream.len() - 4` and panicking.
369        for len in 0..4usize {
370            let mut stream = vec![0u8; len];
371            let result = Codec::Snappy.decompress(&mut stream);
372            assert!(result.is_err(), "len={len} should error, got {result:?}");
373        }
374    }
375
376    #[cfg(feature = "zstandard")]
377    #[test]
378    fn zstd_compress_and_decompress() -> TestResult {
379        compress_and_decompress(Codec::Zstandard(zstandard::ZstandardSettings::default()))
380    }
381
382    #[cfg(feature = "bzip")]
383    #[test]
384    fn bzip_compress_and_decompress() -> TestResult {
385        compress_and_decompress(Codec::Bzip2(bzip::Bzip2Settings::default()))
386    }
387
388    #[cfg(feature = "xz")]
389    #[test]
390    fn xz_compress_and_decompress() -> TestResult {
391        compress_and_decompress(Codec::Xz(xz::XzSettings::default()))
392    }
393
394    fn compress_and_decompress(codec: Codec) -> TestResult {
395        let mut stream = INPUT.to_vec();
396        codec.compress(&mut stream)?;
397        assert_ne!(INPUT, stream.as_slice());
398        assert!(INPUT.len() > stream.len());
399        codec.decompress(&mut stream)?;
400        assert_eq!(INPUT, stream.as_slice());
401        Ok(())
402    }
403
404    #[test]
405    fn codec_to_str() {
406        assert_eq!(<&str>::from(Codec::Null), "null");
407        assert_eq!(
408            <&str>::from(Codec::Deflate(DeflateSettings::default())),
409            "deflate"
410        );
411
412        #[cfg(feature = "snappy")]
413        assert_eq!(<&str>::from(Codec::Snappy), "snappy");
414
415        #[cfg(feature = "zstandard")]
416        assert_eq!(
417            <&str>::from(Codec::Zstandard(zstandard::ZstandardSettings::default())),
418            "zstandard"
419        );
420
421        #[cfg(feature = "bzip")]
422        assert_eq!(
423            <&str>::from(Codec::Bzip2(bzip::Bzip2Settings::default())),
424            "bzip2"
425        );
426
427        #[cfg(feature = "xz")]
428        assert_eq!(<&str>::from(Codec::Xz(xz::XzSettings::default())), "xz");
429    }
430
431    #[test]
432    fn codec_from_str() {
433        use std::str::FromStr;
434
435        assert_eq!(Codec::from_str("null").unwrap(), Codec::Null);
436        assert_eq!(
437            Codec::from_str("deflate").unwrap(),
438            Codec::Deflate(DeflateSettings::default())
439        );
440
441        #[cfg(feature = "snappy")]
442        assert_eq!(Codec::from_str("snappy").unwrap(), Codec::Snappy);
443
444        #[cfg(feature = "zstandard")]
445        assert_eq!(
446            Codec::from_str("zstandard").unwrap(),
447            Codec::Zstandard(zstandard::ZstandardSettings::default())
448        );
449
450        #[cfg(feature = "bzip")]
451        assert_eq!(
452            Codec::from_str("bzip2").unwrap(),
453            Codec::Bzip2(bzip::Bzip2Settings::default())
454        );
455
456        #[cfg(feature = "xz")]
457        assert_eq!(
458            Codec::from_str("xz").unwrap(),
459            Codec::Xz(xz::XzSettings::default())
460        );
461
462        assert!(Codec::from_str("not a codec").is_err());
463    }
464}
apache_avro/codec.rs

apache_avro/
codec.rs
