cadmpeg_core/decode/
deflate.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Scoped DEFLATE output for detection.
3
4use super::{u64_from_index, DecodeContext, ScopedReservation, View};
5use crate::CodecError;
6use flate2::{Decompress, FlushDecompress, Status};
7
8impl DecodeContext<'_> {
9    /// Inflates one raw or zlib member up to `cap`. An oversized or malformed
10    /// member supplies no evidence. Resource refusals remain typed.
11    pub fn inflate_probe(
12        &self,
13        source: View<'_>,
14        cap: usize,
15        zlib: bool,
16    ) -> Result<Option<(Vec<u8>, ScopedReservation<'_>)>, CodecError> {
17        self.charge_work(
18            u64_from_index(source.window().len()),
19            "probe compressed input",
20        )?;
21        // miniz_oxide uses a 32 KiB dictionary and fixed Huffman tables. This
22        // 256 KiB bound includes the state and table storage for one decoder.
23        let _workspace = self.reserve_scoped(256 * 1024, "DEFLATE probe workspace")?;
24        let mut decoder = Decompress::new(zlib);
25        let mut storage = self.reserve_scoped(0, "DEFLATE probe output")?;
26        let mut output = Vec::new();
27        let mut chunk = [0_u8; 1024];
28        let mut offset = 0;
29        loop {
30            self.charge_work(u64_from_index(chunk.len()), "DEFLATE probe step")?;
31            let before_in = decoder.total_in();
32            let before_out = decoder.total_out();
33            let Ok(status) = decoder.decompress(
34                &source.window()[offset..],
35                &mut chunk,
36                FlushDecompress::None,
37            ) else {
38                return Ok(None);
39            };
40            let consumed = usize::try_from(decoder.total_in() - before_in)
41                .map_err(|_| CodecError::Malformed("DEFLATE probe input overflow".into()))?;
42            let produced = usize::try_from(decoder.total_out() - before_out)
43                .map_err(|_| CodecError::Malformed("DEFLATE probe output overflow".into()))?;
44            offset += consumed;
45            if cap
46                .checked_sub(output.len())
47                .is_none_or(|remaining| produced > remaining)
48            {
49                return Ok(None);
50            }
51            self.charge_work(u64_from_index(produced), "DEFLATE probe copy")?;
52            self.reserve_scoped_vec(&mut storage, &mut output, produced, "DEFLATE probe output")?;
53            output.extend_from_slice(&chunk[..produced]);
54            if status == Status::StreamEnd {
55                return Ok(Some((output, storage)));
56            }
57            if consumed == 0 && produced == 0 {
58                return Ok(Some((output, storage)));
59            }
60        }
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use crate::decode::{DecodeArena, DecodeContext, DecodePolicy, ResourceDimension};
67    use crate::CodecError;
68    use flate2::{write::DeflateEncoder, Compression};
69    use std::io::Write;
70
71    #[test]
72    fn deflate_probe_output_reservation_lives_with_returned_bytes() {
73        let mut encoder = DeflateEncoder::new(Vec::new(), Compression::default());
74        encoder.write_all(b"Document.xml").expect("encode");
75        let bytes = encoder.finish().expect("finish");
76        for release in [false, true] {
77            let arena = DecodeArena::new();
78            let mut policy = DecodePolicy::service();
79            policy.limits.max_materialized_bytes = 256 * 1024 + 12;
80            let (ctx, root) =
81                DecodeContext::from_root_bytes(&bytes, &arena, &policy).expect("root");
82            let probe = ctx
83                .inflate_probe(root, 12, false)
84                .expect("probe")
85                .expect("evidence");
86            if release {
87                drop(probe);
88            }
89            let result = ctx.reserve_scoped(
90                policy.limits.max_materialized_bytes,
91                "probe output lifetime",
92            );
93            if release {
94                assert!(result.is_ok());
95            } else {
96                assert!(
97                    matches!(result, Err(CodecError::ResourceLimit(limit)) if limit.dimension == ResourceDimension::MaterializedBytes && limit.used == 12)
98                );
99            }
100        }
101    }
102
103    #[test]
104    fn deflate_probe_propagates_work_and_scoped_refusals() {
105        let mut encoder = DeflateEncoder::new(Vec::new(), Compression::default());
106        encoder.write_all(b"Document.xml").expect("encode");
107        let bytes = encoder.finish().expect("finish");
108        for dimension in [
109            ResourceDimension::WorkUnits,
110            ResourceDimension::MaterializedBytes,
111            ResourceDimension::CollectionItems,
112        ] {
113            let arena = DecodeArena::new();
114            let mut policy = DecodePolicy::service();
115            match dimension {
116                ResourceDimension::WorkUnits => policy.limits.max_work_units = 0,
117                ResourceDimension::MaterializedBytes => policy.limits.max_materialized_bytes = 0,
118                ResourceDimension::CollectionItems => policy.limits.max_collection_items = 0,
119                _ => panic!("test dimension"),
120            }
121            let (ctx, root) =
122                DecodeContext::from_root_bytes(&bytes, &arena, &policy).expect("root");
123            assert!(
124                matches!(ctx.inflate_probe(root, 12, false), Err(CodecError::ResourceLimit(limit)) if limit.dimension == dimension)
125            );
126        }
127        let arena = DecodeArena::new();
128        let (ctx, root) =
129            DecodeContext::from_root_bytes(&bytes, &arena, &DecodePolicy::service()).expect("root");
130        let (output, storage) = ctx
131            .inflate_probe(root, 12, false)
132            .expect("probe")
133            .expect("evidence");
134        assert_eq!(output, b"Document.xml");
135        drop((output, storage));
136        assert!(ctx.inflate_probe(root, 11, false).expect("probe").is_none());
137    }
138}
cadmpeg_core/decode/deflate.rs

cadmpeg_core/decode/
deflate.rs
