Skip to main content

resopt/
optimizer.rs

1use anyhow::{Context, Result, bail, ensure};
2use serde::{Deserialize, Serialize};
3use std::io::Cursor;
4#[cfg(not(target_arch = "wasm32"))]
5use std::time::Duration;
6
7pub(crate) const MAX_INPUT: usize = 64 * 1024 * 1024;
8const MAX_DECODED: usize = 256 * 1024 * 1024;
9
10/// Lossless PNG policy. Unknown fields are rejected (including lossy settings).
11#[derive(Debug, Clone, Serialize, Deserialize)]
12#[serde(default, deny_unknown_fields)]
13pub struct Policy {
14    pub png_level: u8,
15    pub include_ignored: bool,
16    pub min_input_bytes: u64,
17    pub min_savings_bytes: u64,
18    pub min_savings_percent: f64,
19    /// Allow lossless bit-depth, color-type and palette reductions. These keep
20    /// every decoded RGBA sample but rewrite IHDR/PLTE/tRNS, so candidates are
21    /// verified on expanded pixels instead of raw buffers.
22    #[serde(skip_serializing_if = "std::ops::Not::not")]
23    pub reductions: bool,
24}
25
26impl Default for Policy {
27    fn default() -> Self {
28        Self {
29            png_level: 2,
30            include_ignored: false,
31            min_input_bytes: 50 * 1024,
32            min_savings_bytes: 1024,
33            min_savings_percent: 1.0,
34            reductions: false,
35        }
36    }
37}
38
39impl Policy {
40    pub fn validate(&self) -> Result<()> {
41        ensure!(
42            self.png_level <= 6,
43            "png_level must be 0..=6 (effort, not quality)"
44        );
45        ensure!(
46            self.min_savings_percent.is_finite()
47                && (0.0..=100.0).contains(&self.min_savings_percent),
48            "min_savings_percent must be 0..=100"
49        );
50        Ok(())
51    }
52}
53
54pub(crate) fn optimize(original: &[u8], policy: &Policy) -> Result<Vec<u8>> {
55    policy.validate()?;
56    ensure!(original.len() <= MAX_INPUT, "input exceeds 64 MiB limit");
57    let chunks = chunks(original)?;
58    ensure!(
59        !chunks.iter().any(|(kind, _)| kind == b"acTL"),
60        "animated PNG is not supported yet"
61    );
62    let strict = encode(original, policy.png_level, false)?;
63    verify(original, &strict, false)?;
64    if !policy.reductions {
65        return Ok(strict);
66    }
67    // A reduction that cannot be verified, or does not help, yields to strict.
68    let reduced = encode(original, policy.png_level, true)
69        .and_then(|candidate| verify(original, &candidate, true).map(|()| candidate));
70    Ok(match reduced {
71        Ok(reduced) if reduced.len() < strict.len() => reduced,
72        _ => strict,
73    })
74}
75
76fn encode(original: &[u8], level: u8, reductions: bool) -> Result<Vec<u8>> {
77    let mut options = oxipng::Options::from_preset(level);
78    options.optimize_alpha = false;
79    options.bit_depth_reduction = reductions;
80    options.color_type_reduction = reductions;
81    options.palette_reduction = reductions;
82    options.grayscale_reduction = reductions;
83    options.scale_16 = false;
84    options.interlace = None;
85    options.strip = oxipng::StripChunks::None;
86    options.max_decompressed_size = Some(MAX_DECODED);
87    #[cfg(not(target_arch = "wasm32"))]
88    {
89        options.timeout = Some(Duration::from_secs(30));
90    }
91    // Browser hosts cancel the Worker. std::time::Instant has no clock on
92    // wasm32-unknown-unknown; leave OxiPNG's native deadline disabled there.
93    #[cfg(target_arch = "wasm32")]
94    {
95        options.timeout = None;
96    }
97    Ok(oxipng::optimize_from_memory(original, &options)?)
98}
99
100/// Chunks a lossless reduction may rewrite; everything else must be identical.
101const REDUCIBLE_CHUNKS: [[u8; 4]; 3] = [*b"IHDR", *b"PLTE", *b"tRNS"];
102
103/// Independent decoding plus exact chunk comparison. Strict mode rejects any
104/// non-IDAT rewrite, even a harmless one. With `reductions`, IHDR/PLTE/tRNS may
105/// change as long as dimensions, interlacing and expanded RGBA samples do not.
106pub(crate) fn verify(original: &[u8], candidate: &[u8], reductions: bool) -> Result<()> {
107    ensure!(
108        original.len() <= MAX_INPUT && candidate.len() <= MAX_INPUT,
109        "input exceeds 64 MiB limit"
110    );
111    let original_chunks = chunks(original)?;
112    ensure!(
113        !original_chunks.iter().any(|(kind, _)| kind == b"acTL"),
114        "animated PNG is not supported yet"
115    );
116    let candidate_chunks = chunks(candidate)?;
117    if !reductions {
118        ensure!(
119            original_chunks == candidate_chunks,
120            "non-IDAT chunks changed; candidate rejected"
121        );
122        ensure!(
123            decode(original)? == decode(candidate)?,
124            "decoded pixels changed; candidate rejected"
125        );
126        return Ok(());
127    }
128    ensure!(
129        fixed_chunks(&original_chunks) == fixed_chunks(&candidate_chunks),
130        "chunks other than IHDR/PLTE/tRNS changed; candidate rejected"
131    );
132    ensure!(
133        layout(&original_chunks)? == layout(&candidate_chunks)?,
134        "dimensions or interlacing changed; candidate rejected"
135    );
136    crate::png_pixels::ensure_same_rgba(original, candidate, MAX_DECODED)
137}
138
139fn fixed_chunks<'a>(chunks: &[Chunk<'a>]) -> Vec<Chunk<'a>> {
140    chunks
141        .iter()
142        .filter(|(kind, _)| !REDUCIBLE_CHUNKS.contains(kind))
143        .copied()
144        .collect()
145}
146
147/// Width, height and the compression/filter/interlace bytes of IHDR.
148fn layout<'a>(chunks: &[Chunk<'a>]) -> Result<(&'a [u8], &'a [u8])> {
149    let (kind, header) = chunks.first().context("PNG has no chunks")?;
150    ensure!(kind == b"IHDR" && header.len() == 13, "invalid IHDR");
151    Ok((&header[..8], &header[10..]))
152}
153
154type Chunk<'a> = ([u8; 4], &'a [u8]);
155
156fn chunks(bytes: &[u8]) -> Result<Vec<Chunk<'_>>> {
157    ensure!(bytes.starts_with(b"\x89PNG\r\n\x1a\n"), "not a PNG");
158    let mut offset = 8;
159    let mut result = Vec::new();
160    let mut saw_idat = false;
161    while offset < bytes.len() {
162        ensure!(bytes.len() - offset >= 12, "truncated PNG chunk");
163        let length = u32::from_be_bytes(bytes[offset..offset + 4].try_into()?) as usize;
164        let end = offset
165            .checked_add(length)
166            .and_then(|end| end.checked_add(12))
167            .context("PNG chunk overflow")?;
168        ensure!(end <= bytes.len(), "truncated PNG chunk payload");
169        let kind: [u8; 4] = bytes[offset + 4..offset + 8].try_into()?;
170        if kind == *b"IDAT" && !saw_idat {
171            result.push((kind, &bytes[0..0]));
172            saw_idat = true;
173        } else if kind != *b"IDAT" {
174            result.push((kind, &bytes[offset + 8..end - 4]));
175        }
176        offset = end;
177        if kind == *b"IEND" {
178            ensure!(offset == bytes.len(), "data after IEND is not supported");
179            return Ok(result);
180        }
181    }
182    bail!("PNG has no IEND")
183}
184
185fn decode(bytes: &[u8]) -> Result<Vec<u8>> {
186    let mut decoder = png::Decoder::new(Cursor::new(bytes));
187    decoder.set_limits(png::Limits { bytes: MAX_DECODED });
188    let mut reader = decoder.read_info()?;
189    let size = reader
190        .output_buffer_size()
191        .context("PNG output buffer overflow")?;
192    ensure!(size <= MAX_DECODED, "decoded image exceeds 256 MiB limit");
193    let mut buffer = vec![0; size];
194    let frame = reader.next_frame(&mut buffer)?;
195    buffer.truncate(frame.buffer_size());
196    reader.finish()?;
197    Ok(buffer)
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    const SIDE: u32 = 64;
205
206    /// Opaque two-color RGBA: reducible to a 1-bit palette without pixel loss.
207    fn reducible_png() -> Vec<u8> {
208        let mut bytes = Vec::new();
209        {
210            let mut encoder = png::Encoder::new(&mut bytes, SIDE, SIDE);
211            encoder.set_color(png::ColorType::Rgba);
212            encoder.set_depth(png::BitDepth::Eight);
213            encoder
214                .add_text_chunk("Comment".into(), "keep me".into())
215                .unwrap();
216            let data: Vec<u8> = (0..SIDE * SIDE)
217                .flat_map(|i| {
218                    if (i / 8 + i / SIDE / 8).is_multiple_of(2) {
219                        [255, 0, 0, 255]
220                    } else {
221                        [0, 0, 255, 255]
222                    }
223                })
224                .collect();
225            encoder
226                .write_header()
227                .unwrap()
228                .write_image_data(&data)
229                .unwrap();
230        }
231        bytes
232    }
233
234    fn policy(reductions: bool) -> Policy {
235        Policy {
236            reductions,
237            ..Policy::default()
238        }
239    }
240
241    #[test]
242    fn reductions_shrink_further_and_need_reduced_verification() {
243        let original = reducible_png();
244        let strict = optimize(&original, &policy(false)).unwrap();
245        let reduced = optimize(&original, &policy(true)).unwrap();
246        assert!(reduced.len() < strict.len());
247        verify(&original, &reduced, true).unwrap();
248        let error = verify(&original, &reduced, false).unwrap_err();
249        assert!(error.to_string().contains("non-IDAT chunks changed"));
250    }
251
252    #[test]
253    fn strict_candidates_pass_both_verifications() {
254        let original = reducible_png();
255        let strict = optimize(&original, &policy(false)).unwrap();
256        verify(&original, &strict, false).unwrap();
257        verify(&original, &strict, true).unwrap();
258    }
259
260    #[test]
261    fn reductions_keep_ancillary_chunks() {
262        let original = reducible_png();
263        let reduced = optimize(&original, &policy(true)).unwrap();
264        let text = |bytes| {
265            chunks(bytes)
266                .unwrap()
267                .into_iter()
268                .find(|(kind, _)| kind == b"tEXt")
269                .map(|(_, data)| data.to_vec())
270        };
271        assert!(text(&original).is_some());
272        assert_eq!(text(&original), text(&reduced));
273    }
274
275    #[test]
276    fn reduced_verification_rejects_other_images() {
277        let original = reducible_png();
278        let mut other = Vec::new();
279        {
280            let mut encoder = png::Encoder::new(&mut other, SIDE, SIDE);
281            encoder.set_color(png::ColorType::Grayscale);
282            encoder.set_depth(png::BitDepth::Eight);
283            encoder
284                .add_text_chunk("Comment".into(), "keep me".into())
285                .unwrap();
286            encoder
287                .write_header()
288                .unwrap()
289                .write_image_data(&vec![0; (SIDE * SIDE) as usize])
290                .unwrap();
291        }
292        let error = verify(&original, &other, true).unwrap_err();
293        assert!(error.to_string().contains("expanded pixels changed"));
294    }
295
296    #[test]
297    fn reductions_never_produce_a_larger_file_than_strict_mode() {
298        // Tiny image: the palette costs more than it saves.
299        let mut original = Vec::new();
300        {
301            let mut encoder = png::Encoder::new(&mut original, 4, 4);
302            encoder.set_color(png::ColorType::Rgba);
303            encoder.set_depth(png::BitDepth::Eight);
304            let data: Vec<u8> = (0..16u8).flat_map(|i| [i, 0, 0, 255]).collect();
305            encoder
306                .write_header()
307                .unwrap()
308                .write_image_data(&data)
309                .unwrap();
310        }
311        let strict = optimize(&original, &policy(false)).unwrap();
312        let reduced = optimize(&original, &policy(true)).unwrap();
313        assert!(reduced.len() <= strict.len());
314    }
315
316    #[test]
317    fn default_policy_serializes_without_reductions() {
318        let json = serde_json::to_string(&Policy::default()).unwrap();
319        assert!(!json.contains("reductions"), "{json}");
320        let enabled = serde_json::to_string(&policy(true)).unwrap();
321        assert!(enabled.contains("\"reductions\":true"), "{enabled}");
322    }
323}