1use anyhow::{Context, Result, bail, ensure};
2use serde::{Deserialize, Serialize};
3use std::{io::Cursor, time::Duration};
4
5pub(crate) const MAX_INPUT: usize = 64 * 1024 * 1024;
6const MAX_DECODED: usize = 256 * 1024 * 1024;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
10#[serde(default, deny_unknown_fields)]
11pub struct Policy {
12 pub png_level: u8,
13 pub include_ignored: bool,
14 pub min_input_bytes: u64,
15 pub min_savings_bytes: u64,
16 pub min_savings_percent: f64,
17}
18
19impl Default for Policy {
20 fn default() -> Self {
21 Self {
22 png_level: 2,
23 include_ignored: false,
24 min_input_bytes: 50 * 1024,
25 min_savings_bytes: 1024,
26 min_savings_percent: 1.0,
27 }
28 }
29}
30
31impl Policy {
32 pub fn validate(&self) -> Result<()> {
33 ensure!(
34 self.png_level <= 6,
35 "png_level must be 0..=6 (effort, not quality)"
36 );
37 ensure!(
38 self.min_savings_percent.is_finite()
39 && (0.0..=100.0).contains(&self.min_savings_percent),
40 "min_savings_percent must be 0..=100"
41 );
42 Ok(())
43 }
44}
45
46pub(crate) fn optimize(original: &[u8], policy: &Policy) -> Result<Vec<u8>> {
47 policy.validate()?;
48 ensure!(original.len() <= MAX_INPUT, "input exceeds 64 MiB limit");
49 let chunks = chunks(original)?;
50 ensure!(
51 !chunks.iter().any(|(kind, _)| kind == b"acTL"),
52 "animated PNG is not supported yet"
53 );
54 let mut options = oxipng::Options::from_preset(policy.png_level);
55 options.optimize_alpha = false;
56 options.bit_depth_reduction = false;
57 options.color_type_reduction = false;
58 options.palette_reduction = false;
59 options.grayscale_reduction = false;
60 options.scale_16 = false;
61 options.interlace = None;
62 options.strip = oxipng::StripChunks::None;
63 options.max_decompressed_size = Some(MAX_DECODED);
64 options.timeout = Some(Duration::from_secs(30));
65 let candidate = oxipng::optimize_from_memory(original, &options)?;
66 verify(original, &candidate)?;
67 Ok(candidate)
68}
69
70pub(crate) fn verify(original: &[u8], candidate: &[u8]) -> Result<()> {
73 ensure!(
74 original.len() <= MAX_INPUT && candidate.len() <= MAX_INPUT,
75 "input exceeds 64 MiB limit"
76 );
77 let original_chunks = chunks(original)?;
78 ensure!(
79 !original_chunks.iter().any(|(kind, _)| kind == b"acTL"),
80 "animated PNG is not supported yet"
81 );
82 ensure!(
83 original_chunks == chunks(candidate)?,
84 "non-IDAT chunks changed; candidate rejected"
85 );
86 ensure!(
87 decode(original)? == decode(candidate)?,
88 "decoded pixels changed; candidate rejected"
89 );
90 Ok(())
91}
92
93type Chunk<'a> = ([u8; 4], &'a [u8]);
94
95fn chunks(bytes: &[u8]) -> Result<Vec<Chunk<'_>>> {
96 ensure!(bytes.starts_with(b"\x89PNG\r\n\x1a\n"), "not a PNG");
97 let mut offset = 8;
98 let mut result = Vec::new();
99 let mut saw_idat = false;
100 while offset < bytes.len() {
101 ensure!(bytes.len() - offset >= 12, "truncated PNG chunk");
102 let length = u32::from_be_bytes(bytes[offset..offset + 4].try_into()?) as usize;
103 let end = offset
104 .checked_add(length)
105 .and_then(|end| end.checked_add(12))
106 .context("PNG chunk overflow")?;
107 ensure!(end <= bytes.len(), "truncated PNG chunk payload");
108 let kind: [u8; 4] = bytes[offset + 4..offset + 8].try_into()?;
109 if kind == *b"IDAT" && !saw_idat {
110 result.push((kind, &bytes[0..0]));
111 saw_idat = true;
112 } else if kind != *b"IDAT" {
113 result.push((kind, &bytes[offset + 8..end - 4]));
114 }
115 offset = end;
116 if kind == *b"IEND" {
117 ensure!(offset == bytes.len(), "data after IEND is not supported");
118 return Ok(result);
119 }
120 }
121 bail!("PNG has no IEND")
122}
123
124fn decode(bytes: &[u8]) -> Result<Vec<u8>> {
125 let mut decoder = png::Decoder::new(Cursor::new(bytes));
126 decoder.set_limits(png::Limits { bytes: MAX_DECODED });
127 let mut reader = decoder.read_info()?;
128 let size = reader
129 .output_buffer_size()
130 .context("PNG output buffer overflow")?;
131 ensure!(size <= MAX_DECODED, "decoded image exceeds 256 MiB limit");
132 let mut buffer = vec![0; size];
133 let frame = reader.next_frame(&mut buffer)?;
134 buffer.truncate(frame.buffer_size());
135 reader.finish()?;
136 Ok(buffer)
137}