Skip to main content

paperforge_pdf/
stream.rs

1use crate::error::PdfResult;
2
3#[derive(Debug, Clone, Copy, PartialEq)]
4pub enum StreamFilter {
5    ASCIIHex,
6    ASCII85,
7    Flate,
8    LZW,
9    RunLength,
10    CCITTFax,
11    JBIG2,
12    DCT,
13    JPX,
14}
15
16pub struct StreamDecoder;
17
18impl StreamDecoder {
19    pub fn new() -> Self {
20        Self
21    }
22
23    pub fn decode(&self, data: &[u8], filter: StreamFilter) -> PdfResult<Vec<u8>> {
24        self.decode_with_limit(data, filter, usize::MAX)
25    }
26
27    /// Decode with a hard cap on the produced bytes (guards decompression bombs).
28    pub fn decode_with_limit(
29        &self,
30        data: &[u8],
31        filter: StreamFilter,
32        limit: usize,
33    ) -> PdfResult<Vec<u8>> {
34        match filter {
35            StreamFilter::Flate => self.decode_flate(data, limit),
36            _ => Err(crate::error::PdfError::NotImplemented(format!(
37                "filter {:?} not yet implemented",
38                filter
39            ))),
40        }
41    }
42
43    fn decode_flate(&self, data: &[u8], limit: usize) -> PdfResult<Vec<u8>> {
44        use std::io::Read;
45        let mut decoder = flate2::read::ZlibDecoder::new(data);
46        let mut result = Vec::new();
47        decoder
48            .by_ref()
49            .take(limit.saturating_add(1) as u64)
50            .read_to_end(&mut result)?;
51        Ok(result)
52    }
53}
54
55impl Default for StreamDecoder {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61/// Flate (zlib) stream encoder.
62///
63/// The encoder keeps the underlying [`flate2::Compress`] state across calls and
64/// resets it per stream, so a single instance can compress many small streams
65/// (one per content stream, for instance) without paying the zlib state
66/// initialisation cost every time. This is roughly 4-6x faster than creating a
67/// fresh encoder per stream on typical tiny PDF content streams.
68pub struct StreamEncoder {
69    /// `(zlib level, compressor)`; the compressor is recreated only when the
70    /// requested level changes between calls.
71    compressor: Option<(u32, flate2::Compress)>,
72}
73
74impl StreamEncoder {
75    pub fn new() -> Self {
76        Self { compressor: None }
77    }
78
79    pub fn encode(&mut self, data: &[u8], filter: StreamFilter) -> PdfResult<Vec<u8>> {
80        self.encode_with_compression(data, filter, 6)
81    }
82
83    /// Like [`StreamEncoder::encode`] but with an explicit zlib level (0-9)
84    /// for the flate filter; other filters ignore the level.
85    pub fn encode_with_compression(
86        &mut self,
87        data: &[u8],
88        filter: StreamFilter,
89        level: u32,
90    ) -> PdfResult<Vec<u8>> {
91        let level = level.clamp(0, 9);
92        match filter {
93            StreamFilter::Flate => self.encode_flate(data, level),
94            _ => Err(crate::error::PdfError::NotImplemented(format!(
95                "filter {:?} not yet implemented",
96                filter
97            ))),
98        }
99    }
100
101    fn encode_flate(&mut self, data: &[u8], level: u32) -> PdfResult<Vec<u8>> {
102        if self
103            .compressor
104            .as_ref()
105            .map(|(l, _)| *l != level)
106            .unwrap_or(true)
107        {
108            self.compressor = Some((
109                level,
110                flate2::Compress::new(flate2::Compression::new(level), true),
111            ));
112        }
113        let (_, compressor) = self.compressor.as_mut().expect("initialized above");
114        compressor.reset();
115
116        // Pre-size the output; PDF content streams are small, so a single
117        // allocation usually suffices. `compress_vec` writes into the Vec's
118        // spare capacity (appending, never overwriting), so on a full buffer we
119        // simply reserve more space and call again until the stream finishes.
120        let mut out = Vec::with_capacity(data.len().saturating_add(data.len() / 2).max(128));
121        loop {
122            let before = compressor.total_out();
123            match compressor.compress_vec(data, &mut out, flate2::FlushCompress::Finish) {
124                Ok(flate2::Status::StreamEnd) => break,
125                Ok(_) | Err(_) => {
126                    // compress_vec only uses spare capacity; grow it and
127                    // continue. If the stream makes no progress at all, bail
128                    // out instead of looping forever.
129                    if compressor.total_out() == before {
130                        return Err(crate::error::PdfError::NotImplemented(
131                            "flate compression made no progress".to_string(),
132                        ));
133                    }
134                    out.reserve(data.len().max(64));
135                }
136            }
137        }
138        Ok(out)
139    }
140}
141
142impl Default for StreamEncoder {
143    fn default() -> Self {
144        Self::new()
145    }
146}