xmtp_content_types/
compression.rs

1//! Compression for the bytes inside an `EncodedContent` envelope.
2
3use std::io::{self, Read, Write};
4
5use flate2::{
6    Compression as FlateCompression,
7    read::{DeflateDecoder, MultiGzDecoder, ZlibDecoder},
8    write::{GzEncoder, ZlibEncoder},
9};
10use xmtp_proto::xmtp::mls::message_contents::{Compression, EncodedContent};
11
12use crate::CodecError;
13
14pub const MAX_DECOMPRESSED_BYTES: usize = 16_777_216;
15pub const COMPRESSION_CHUNK_BYTES: usize = 64 * 1024;
16
17#[derive(Debug)]
18enum DecompressFailure {
19    Invalid(io::Error),
20    Limit,
21}
22
23/// Total decompressed bytes allowed while decoding one message and its nested content.
24#[derive(Debug)]
25pub struct DecompressionBudget {
26    remaining: usize,
27    peak_capacity: usize,
28}
29
30impl Default for DecompressionBudget {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36impl DecompressionBudget {
37    pub const fn new() -> Self {
38        Self {
39            remaining: MAX_DECOMPRESSED_BYTES,
40            peak_capacity: 0,
41        }
42    }
43
44    /// Largest output buffer capacity used by one decompression.
45    pub const fn peak_capacity(&self) -> usize {
46        self.peak_capacity
47    }
48
49    /// Total bytes successfully decompressed so far.
50    pub const fn used(&self) -> usize {
51        MAX_DECOMPRESSED_BYTES - self.remaining
52    }
53}
54
55fn read_bounded(
56    reader: &mut impl Read,
57    budget: &mut DecompressionBudget,
58) -> Result<Vec<u8>, DecompressFailure> {
59    let mut output = Vec::new();
60    let mut chunk = [0u8; COMPRESSION_CHUNK_BYTES];
61    loop {
62        let remaining = budget.remaining - output.len();
63        let request = remaining.saturating_add(1).min(COMPRESSION_CHUNK_BYTES);
64        let read = reader
65            .read(&mut chunk[..request])
66            .map_err(DecompressFailure::Invalid)?;
67        if read == 0 {
68            budget.remaining -= output.len();
69            return Ok(output);
70        }
71        if read > remaining {
72            return Err(DecompressFailure::Limit);
73        }
74        output.reserve_exact(read);
75        output.extend_from_slice(&chunk[..read]);
76        budget.peak_capacity = budget.peak_capacity.max(output.capacity());
77    }
78}
79
80/// Compresses content with a wire-compatible zlib or gzip frame.
81// implements: CTYPE-023
82pub fn compress(
83    mut content: EncodedContent,
84    algorithm: Compression,
85) -> Result<EncodedContent, CodecError> {
86    if content.compression.is_some() {
87        return Err(CodecError::Encode("content is already compressed".into()));
88    }
89    let bytes = match algorithm {
90        Compression::Deflate => {
91            let mut encoder = ZlibEncoder::new(Vec::new(), FlateCompression::default());
92            encoder
93                .write_all(&content.content)
94                .map_err(|e| CodecError::Encode(e.to_string()))?;
95            encoder
96                .finish()
97                .map_err(|e| CodecError::Encode(e.to_string()))?
98        }
99        Compression::Gzip => {
100            let mut encoder = GzEncoder::new(Vec::new(), FlateCompression::default());
101            encoder
102                .write_all(&content.content)
103                .map_err(|e| CodecError::Encode(e.to_string()))?;
104            encoder
105                .finish()
106                .map_err(|e| CodecError::Encode(e.to_string()))?
107        }
108    };
109    content.content = bytes;
110    content.compression = Some(algorithm as i32);
111    Ok(content)
112}
113
114/// Leaves content unchanged when the caller does not request compression.
115pub fn compress_if_requested(
116    content: EncodedContent,
117    algorithm: Option<Compression>,
118) -> Result<EncodedContent, CodecError> {
119    match algorithm {
120        Some(algorithm) => compress(content, algorithm),
121        None => Ok(content),
122    }
123}
124
125/// Decompresses before a caller selects a content codec.
126// implements: CTYPE-024, CTYPE-025
127pub fn decompress(content: EncodedContent) -> Result<EncodedContent, CodecError> {
128    decompress_with_budget(content, &mut DecompressionBudget::new())
129}
130
131/// Decompresses one envelope and charges its output to a shared message budget.
132pub fn decompress_with_budget(
133    mut content: EncodedContent,
134    budget: &mut DecompressionBudget,
135) -> Result<EncodedContent, CodecError> {
136    let Some(raw_algorithm) = content.compression else {
137        return Ok(content);
138    };
139    let algorithm = Compression::try_from(raw_algorithm)
140        .map_err(|_| CodecError::Decode(format!("unknown compression value {raw_algorithm}")))?;
141    let result = match algorithm {
142        Compression::Gzip => {
143            read_bounded(&mut MultiGzDecoder::new(content.content.as_slice()), budget)
144        }
145        Compression::Deflate => {
146            match read_bounded(&mut ZlibDecoder::new(content.content.as_slice()), budget) {
147                Ok(bytes) => Ok(bytes),
148                Err(DecompressFailure::Invalid(_)) => {
149                    read_bounded(&mut DeflateDecoder::new(content.content.as_slice()), budget)
150                }
151                Err(DecompressFailure::Limit) => Err(DecompressFailure::Limit),
152            }
153        }
154    };
155    content.content = result.map_err(|error| match error {
156        DecompressFailure::Invalid(error) => {
157            CodecError::Decode(format!("decompression failed: {error}"))
158        }
159        DecompressFailure::Limit => CodecError::Decode(format!(
160            "decompressed content exceeds {MAX_DECOMPRESSED_BYTES} bytes"
161        )),
162    })?;
163    content.compression = None;
164    Ok(content)
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170    use crate::{ContentCodec, text::TextCodec};
171
172    // Generated with Python zlib.compressobj(wbits=-15) for "hello compressed XMTP".
173    // verifies: CTYPE-024
174    #[xmtp_common::test(unwrap_try = true)]
175    async fn decompress_ios_raw_deflate_vector() {
176        let mut content = TextCodec::encode(String::new())?;
177        content.content = hex::decode("cb48cdc9c95748cecf2d284a2d2e4e4d5188f00d090000")?;
178        content.compression = Some(Compression::Deflate as i32);
179        assert_eq!(
180            TextCodec::decode(decompress(content)?)?,
181            "hello compressed XMTP"
182        );
183    }
184
185    // Generated with Python zlib.compress for "hello compressed XMTP".
186    // verifies: CTYPE-024
187    #[xmtp_common::test(unwrap_try = true)]
188    async fn decompress_kotlin_zlib_vector() {
189        let mut content = TextCodec::encode(String::new())?;
190        content.content =
191            hex::decode("789ccb48cdc9c95748cecf2d284a2d2e4e4d5188f00d090000599907d3")?;
192        content.compression = Some(Compression::Deflate as i32);
193        assert_eq!(
194            TextCodec::decode(decompress(content)?)?,
195            "hello compressed XMTP"
196        );
197    }
198
199    // Generated with macOS compression_encode_buffer using COMPRESSION_LZFSE
200    // (0x801) on eight copies of "hello compressed XMTP".
201    // verifies: CTYPE-024
202    #[xmtp_common::test(unwrap_try = true)]
203    async fn lzfse_labelled_gzip_is_decode_failure() {
204        let mut content = TextCodec::encode(String::new())?;
205        content.content = hex::decode(
206            "6276786ea800000026000000e00568656c6c6f20636f6d7072657373656420584d54503815f077e25450060000000000000062767824",
207        )?;
208        content.compression = Some(Compression::Gzip as i32);
209        assert!(matches!(decompress(content), Err(CodecError::Decode(_))));
210    }
211
212    // verifies: CTYPE-025
213    #[xmtp_common::test(unwrap_try = true)]
214    async fn decompression_stops_at_limit() {
215        struct CountedRead<R> {
216            inner: R,
217            bytes_read: usize,
218        }
219
220        impl<R: Read> Read for CountedRead<R> {
221            fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
222                let read = self.inner.read(buffer)?;
223                self.bytes_read += read;
224                Ok(read)
225            }
226        }
227
228        let source = vec![b'A'; MAX_DECOMPRESSED_BYTES + 2 * COMPRESSION_CHUNK_BYTES];
229        let mut encoder = ZlibEncoder::new(Vec::new(), FlateCompression::default());
230        encoder.write_all(&source)?;
231        let compressed = encoder.finish()?;
232        assert!(compressed.len() < 20_000);
233        let mut budget = DecompressionBudget::new();
234        let mut reader = CountedRead {
235            inner: ZlibDecoder::new(compressed.as_slice()),
236            bytes_read: 0,
237        };
238        let result = read_bounded(&mut reader, &mut budget);
239        assert!(matches!(result, Err(DecompressFailure::Limit)));
240        // One extra byte detects the limit without reading the rest of the source.
241        assert_eq!(reader.bytes_read, MAX_DECOMPRESSED_BYTES + 1);
242        assert!(budget.peak_capacity() <= MAX_DECOMPRESSED_BYTES + COMPRESSION_CHUNK_BYTES);
243        let mut content = TextCodec::encode(String::new())?;
244        content.content = compressed;
245        content.compression = Some(Compression::Deflate as i32);
246        assert!(matches!(decompress(content), Err(CodecError::Decode(_))));
247    }
248
249    // verifies: CTYPE-025
250    #[xmtp_common::test(unwrap_try = true)]
251    async fn stored_block_then_compressed_tail_stays_within_capacity_limit() {
252        use flate2::write::DeflateEncoder;
253
254        let first = vec![b'B'; 49_052];
255        let tail = vec![b'A'; MAX_DECOMPRESSED_BYTES + 1 - first.len()];
256        let mut raw_tail = DeflateEncoder::new(Vec::new(), FlateCompression::default());
257        raw_tail.write_all(&tail)?;
258        let raw_tail = raw_tail.finish()?;
259
260        // One non-final stored DEFLATE block, then a final compressed block.
261        let length = u16::try_from(first.len())?;
262        let mut compressed = vec![0x78, 0x9c, 0x00];
263        compressed.extend_from_slice(&length.to_le_bytes());
264        compressed.extend_from_slice(&(!length).to_le_bytes());
265        compressed.extend_from_slice(&first);
266        compressed.extend_from_slice(&raw_tail);
267        let mut s1 = 1u32;
268        let mut s2 = 0u32;
269        for byte in first.iter().chain(&tail) {
270            s1 = (s1 + u32::from(*byte)) % 65_521;
271            s2 = (s2 + s1) % 65_521;
272        }
273        compressed.extend_from_slice(&((s2 << 16) | s1).to_be_bytes());
274
275        let mut budget = DecompressionBudget::new();
276        let result = read_bounded(&mut ZlibDecoder::new(compressed.as_slice()), &mut budget);
277        assert!(matches!(result, Err(DecompressFailure::Limit)));
278        assert!(budget.peak_capacity() <= MAX_DECOMPRESSED_BYTES + COMPRESSION_CHUNK_BYTES);
279        let mut content = TextCodec::encode(String::new())?;
280        content.content = compressed;
281        content.compression = Some(Compression::Deflate as i32);
282        assert!(matches!(decompress(content), Err(CodecError::Decode(_))));
283    }
284
285    // verifies: CTYPE-024
286    #[xmtp_common::test(unwrap_try = true)]
287    async fn gzip_decodes_all_members() {
288        let mut content = TextCodec::encode(String::new())?;
289        let mut first = GzEncoder::new(Vec::new(), FlateCompression::default());
290        first.write_all(b"hello ")?;
291        let mut second = GzEncoder::new(Vec::new(), FlateCompression::default());
292        second.write_all(b"world")?;
293        content.content = first.finish()?;
294        content.content.extend(second.finish()?);
295        content.compression = Some(Compression::Gzip as i32);
296        assert_eq!(decompress(content)?.content, b"hello world");
297    }
298
299    // verifies: CTYPE-023, CTYPE-024
300    #[xmtp_common::test(unwrap_try = true)]
301    async fn compress_round_trip() {
302        let original = TextCodec::encode("hello compressed XMTP".into())?;
303        assert_eq!(compress_if_requested(original.clone(), None)?, original);
304        for (algorithm, magic) in [
305            (Compression::Deflate, &[0x78][..]),
306            (Compression::Gzip, &[0x1f, 0x8b][..]),
307        ] {
308            let compressed = compress_if_requested(original.clone(), Some(algorithm))?;
309            assert_eq!(compressed.compression, Some(algorithm as i32));
310            assert!(compressed.content.starts_with(magic));
311            assert_eq!(decompress(compressed)?, original);
312        }
313    }
314}
xmtp_content_types/compression.rs

xmtp_content_types/
compression.rs
