xlog_decoder/core/
compress.rs

1use std::io::{Read, Write};
2
3use flate2::{Compression, Decompress, FlushDecompress, Status};
4use thiserror::Error;
5
6#[derive(Debug, Error)]
7/// Errors raised by the compression and decompression helpers.
8pub enum CompressError {
9    /// Raw-deflate compression failed.
10    #[error("zlib compress failed: {0}")]
11    Zlib(String),
12    /// zstd compression failed.
13    #[error("zstd compress failed: {0}")]
14    Zstd(String),
15    /// Raw-deflate decompression failed.
16    #[error("zlib decompress failed: {0}")]
17    ZlibDecompress(String),
18    /// zstd frame decompression failed.
19    #[error("zstd decompress failed: {0}")]
20    ZstdDecompress(String),
21}
22
23/// Stateful compressor used by the async writer to emit one logical stream.
24pub trait StreamCompressor {
25    /// Compresses one input chunk and appends newly emitted bytes to `output`.
26    fn compress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<(), CompressError>;
27    /// Finalizes the stream and appends any remaining encoded bytes to `output`.
28    fn flush(&mut self, output: &mut Vec<u8>) -> Result<(), CompressError>;
29}
30
31/// Raw-deflate stream compressor compatible with Mars zlib settings.
32pub struct ZlibStreamCompressor {
33    inner: flate2::write::DeflateEncoder<Vec<u8>>,
34}
35
36impl ZlibStreamCompressor {
37    /// Creates a raw-deflate stream compressor with the requested compression level.
38    pub fn new(level: i32) -> Self {
39        let level = level.clamp(0, 9) as u32;
40        Self {
41            inner: flate2::write::DeflateEncoder::new(Vec::new(), Compression::new(level)),
42        }
43    }
44}
45
46impl Default for ZlibStreamCompressor {
47    fn default() -> Self {
48        Self::new(9)
49    }
50}
51
52impl StreamCompressor for ZlibStreamCompressor {
53    fn compress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<(), CompressError> {
54        self.inner
55            .write_all(input)
56            .map_err(|e| CompressError::Zlib(e.to_string()))?;
57        self.inner
58            .flush()
59            .map_err(|e| CompressError::Zlib(e.to_string()))?;
60        drain_encoded_bytes(self.inner.get_mut(), output);
61        Ok(())
62    }
63
64    fn flush(&mut self, output: &mut Vec<u8>) -> Result<(), CompressError> {
65        self.inner
66            .try_finish()
67            .map_err(|e| CompressError::Zlib(e.to_string()))?;
68        drain_encoded_bytes(self.inner.get_mut(), output);
69        Ok(())
70    }
71}
72
73/// zstd chunk compressor.
74///
75/// Each chunk is encoded as a zstd frame and can be decoded from concatenated frames.
76pub struct ZstdChunkCompressor {
77    level: i32,
78}
79
80impl ZstdChunkCompressor {
81    /// Creates a compressor that encodes each input chunk as a separate zstd frame.
82    pub fn new(level: i32) -> Self {
83        Self { level }
84    }
85}
86
87impl StreamCompressor for ZstdChunkCompressor {
88    fn compress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<(), CompressError> {
89        let compressed = zstd::bulk::compress(input, self.level)
90            .map_err(|e| CompressError::Zstd(e.to_string()))?;
91        output.extend_from_slice(&compressed);
92        Ok(())
93    }
94
95    fn flush(&mut self, _output: &mut Vec<u8>) -> Result<(), CompressError> {
96        Ok(())
97    }
98}
99
100/// Streaming zstd compressor compatible with Mars async path.
101///
102/// Uses one compression stream with `windowLog=16`, flushing each chunk and
103/// emitting the final frame epilogue in `flush()`.
104pub struct ZstdStreamCompressor {
105    inner: Option<zstd::stream::write::Encoder<'static, Vec<u8>>>,
106}
107
108impl ZstdStreamCompressor {
109    /// Creates a streaming zstd compressor configured to match Mars async settings.
110    pub fn new(level: i32) -> Result<Self, CompressError> {
111        let mut encoder = zstd::stream::write::Encoder::new(Vec::new(), level)
112            .map_err(|e| CompressError::Zstd(e.to_string()))?;
113        encoder
114            .window_log(16)
115            .map_err(|e| CompressError::Zstd(e.to_string()))?;
116        Ok(Self {
117            inner: Some(encoder),
118        })
119    }
120}
121
122impl StreamCompressor for ZstdStreamCompressor {
123    fn compress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<(), CompressError> {
124        let Some(encoder) = self.inner.as_mut() else {
125            return Err(CompressError::Zstd(
126                "zstd stream compressor already finished".to_string(),
127            ));
128        };
129        encoder
130            .write_all(input)
131            .map_err(|e| CompressError::Zstd(e.to_string()))?;
132        encoder
133            .flush()
134            .map_err(|e| CompressError::Zstd(e.to_string()))?;
135        drain_encoded_bytes(encoder.get_mut(), output);
136        Ok(())
137    }
138
139    fn flush(&mut self, output: &mut Vec<u8>) -> Result<(), CompressError> {
140        let Some(encoder) = self.inner.take() else {
141            return Ok(());
142        };
143        let mut encoded = encoder
144            .finish()
145            .map_err(|e| CompressError::Zstd(e.to_string()))?;
146        drain_encoded_bytes(&mut encoded, output);
147        Ok(())
148    }
149}
150
151fn drain_encoded_bytes(encoded: &mut Vec<u8>, output: &mut Vec<u8>) {
152    if encoded.is_empty() {
153        return;
154    }
155    output.extend_from_slice(encoded);
156    encoded.clear();
157}
158
159/// Decompresses raw-deflate data produced by the zlib stream compressor.
160///
161/// Mars xlog flushes async chunks using `Z_SYNC_FLUSH` without writing a final
162/// stream termination marker (`Z_FINISH`). Decompressing via streaming state machine
163/// ensures compatibility with flushed-but-not-ended deflate chunks across flate2/miniz versions.
164pub fn decompress_raw_zlib(input: &[u8]) -> Result<Vec<u8>, CompressError> {
165    if input.is_empty() {
166        return Ok(Vec::new());
167    }
168
169    let mut decompressor = Decompress::new(false);
170    let mut out = Vec::with_capacity(input.len().saturating_mul(2).max(1024));
171    let mut input_offset = 0;
172
173    while input_offset < input.len() {
174        if out.capacity() == out.len() {
175            out.reserve(out.len().max(4096));
176        }
177        let in_before = decompressor.total_in();
178        let out_before = decompressor.total_out();
179
180        let status = decompressor
181            .decompress_vec(&input[input_offset..], &mut out, FlushDecompress::Sync)
182            .map_err(|e| CompressError::ZlibDecompress(e.to_string()))?;
183
184        let consumed = (decompressor.total_in() - in_before) as usize;
185        let written = (decompressor.total_out() - out_before) as usize;
186        input_offset += consumed;
187
188        if status == Status::StreamEnd {
189            break;
190        }
191
192        if consumed == 0 && written == 0 {
193            break;
194        }
195    }
196
197    Ok(out)
198}
199
200/// Decompresses one or more concatenated zstd frames into a single byte buffer.
201pub fn decompress_zstd_frames(input: &[u8]) -> Result<Vec<u8>, CompressError> {
202    let mut reader = std::io::Cursor::new(input);
203    let mut decoder = zstd::stream::Decoder::new(&mut reader)
204        .map_err(|e| CompressError::ZstdDecompress(e.to_string()))?;
205    let mut out = Vec::new();
206    decoder
207        .read_to_end(&mut out)
208        .map_err(|e| CompressError::ZstdDecompress(e.to_string()))?;
209    Ok(out)
210}
211
212#[cfg(test)]
213mod tests {
214    use super::{
215        decompress_raw_zlib, decompress_zstd_frames, CompressError, StreamCompressor,
216        ZlibStreamCompressor, ZstdChunkCompressor, ZstdStreamCompressor,
217    };
218
219    #[test]
220    fn zlib_stream_roundtrips_incremental_chunks() {
221        let mut compressor = ZlibStreamCompressor::new(6);
222        let mut encoded = Vec::new();
223
224        compressor.compress_chunk(b"mars", &mut encoded).unwrap();
225        compressor.compress_chunk(b" xlog", &mut encoded).unwrap();
226        compressor.flush(&mut encoded).unwrap();
227
228        assert_eq!(decompress_raw_zlib(&encoded).unwrap(), b"mars xlog");
229    }
230
231    #[test]
232    fn stream_compressors_release_emitted_bytes_after_each_chunk() {
233        let mut zlib = ZlibStreamCompressor::new(6);
234        let mut zlib_out = Vec::new();
235        zlib.compress_chunk(b"mars", &mut zlib_out).unwrap();
236        assert!(zlib.inner.get_ref().is_empty());
237        zlib.compress_chunk(b" xlog", &mut zlib_out).unwrap();
238        assert!(zlib.inner.get_ref().is_empty());
239
240        let mut zstd = ZstdStreamCompressor::new(3).unwrap();
241        let mut zstd_out = Vec::new();
242        zstd.compress_chunk(b"mars", &mut zstd_out).unwrap();
243        assert!(zstd.inner.as_ref().unwrap().get_ref().is_empty());
244        zstd.compress_chunk(b" xlog", &mut zstd_out).unwrap();
245        assert!(zstd.inner.as_ref().unwrap().get_ref().is_empty());
246    }
247
248    #[test]
249    fn zstd_chunk_roundtrips_concatenated_frames() {
250        let mut compressor = ZstdChunkCompressor::new(3);
251        let mut encoded = Vec::new();
252
253        compressor.compress_chunk(b"mars", &mut encoded).unwrap();
254        compressor.compress_chunk(b" xlog", &mut encoded).unwrap();
255        compressor.flush(&mut encoded).unwrap();
256
257        assert_eq!(decompress_zstd_frames(&encoded).unwrap(), b"mars xlog");
258    }
259
260    #[test]
261    fn zstd_stream_rejects_compress_after_finish_and_double_flush_is_noop() {
262        let mut compressor = ZstdStreamCompressor::new(3).unwrap();
263        let mut encoded = Vec::new();
264
265        compressor.compress_chunk(b"hello", &mut encoded).unwrap();
266        compressor.flush(&mut encoded).unwrap();
267        let finished = encoded.clone();
268
269        let err = compressor
270            .compress_chunk(b"world", &mut encoded)
271            .unwrap_err();
272        assert!(
273            matches!(err, CompressError::Zstd(message) if message.contains("already finished"))
274        );
275
276        compressor.flush(&mut encoded).unwrap();
277        assert_eq!(encoded, finished);
278    }
279
280    #[test]
281    fn invalid_decompression_maps_to_specific_error_variants() {
282        assert!(matches!(
283            decompress_raw_zlib(b"not-a-deflate-stream"),
284            Err(CompressError::ZlibDecompress(_))
285        ));
286        assert!(matches!(
287            decompress_zstd_frames(b"not-a-zstd-frame"),
288            Err(CompressError::ZstdDecompress(_))
289        ));
290    }
291}
xlog_decoder/core/compress.rs

xlog_decoder/core/
compress.rs
