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#[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 #[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 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 #[cfg(target_arch = "wasm32")]
94 {
95 options.timeout = None;
96 }
97 Ok(oxipng::optimize_from_memory(original, &options)?)
98}
99
100const REDUCIBLE_CHUNKS: [[u8; 4]; 3] = [*b"IHDR", *b"PLTE", *b"tRNS"];
102
103pub(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
147fn 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 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 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}