mj_controller/
hel_image.rs

1//! Decode and constrain user supplied images for the controller's web APIs.
2
3use std::io::{self, Cursor, Write};
4
5use anyhow::{Context, Result, anyhow, bail};
6use image::codecs::jpeg::JpegEncoder;
7use image::codecs::png::{CompressionType, FilterType, PngEncoder};
8use image::metadata::Orientation;
9use image::{
10    ColorType, DynamicImage, GenericImageView, ImageDecoder, ImageFormat, ImageReader, Limits,
11    RgbaImage,
12};
13
14const MAX_INPUT_BYTES: usize = 64 * 1024 * 1024;
15const MAX_DECODED_BYTES: u64 = 256 * 1024 * 1024;
16const MAX_OUTPUT_BYTES: usize = hel::hel_attachment::MAX_IMAGE_BYTES;
17const MAX_JPEG_DIMENSION: u32 = 65_535;
18const JPEG_QUALITIES: [u8; 3] = [90, 85, 80];
19
20/// An image encoded in a bounded, browser-friendly representation.
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct OptimizedImage {
23    pub bytes: Vec<u8>,
24    pub mime_type: String,
25    pub width: u32,
26    pub height: u32,
27}
28
29/// Decode and optimize a JPEG, PNG, or WebP image.
30///
31/// Images larger than the output budget are recompressed and progressively resized. EXIF
32/// orientation is applied before deciding the output dimensions.
33pub fn optimize_image(bytes: &[u8]) -> Result<OptimizedImage> {
34    if bytes.len() > MAX_INPUT_BYTES {
35        bail!("image input exceeds the 64 MiB limit");
36    }
37
38    let mut reader = ImageReader::new(Cursor::new(bytes))
39        .with_guessed_format()
40        .context("guess image format")?;
41    let format = reader
42        .format()
43        .ok_or_else(|| anyhow!("unsupported image format"))?;
44    if !matches!(
45        format,
46        ImageFormat::Jpeg | ImageFormat::Png | ImageFormat::WebP
47    ) {
48        bail!("unsupported image format (expected JPEG, PNG, or WebP)");
49    }
50
51    let mut limits = Limits::default();
52    limits.max_alloc = Some(MAX_DECODED_BYTES);
53    reader.limits(limits);
54    let mut decoder = reader.into_decoder().context("create image decoder")?;
55    let orientation = decoder.orientation().context("read image orientation")?;
56    let (width, height) = decoder.dimensions();
57    let decoded_bytes = decoder.total_bytes();
58    if decoded_bytes > MAX_DECODED_BYTES {
59        bail!("decoded image exceeds the 256 MiB limit");
60    }
61
62    let mut image = DynamicImage::from_decoder(decoder).context("decode image")?;
63    image.apply_orientation(orientation);
64
65    // A normal JPEG or PNG can keep its original metadata and compression when it already fits.
66    if orientation == Orientation::NoTransforms
67        && matches!(format, ImageFormat::Jpeg | ImageFormat::Png)
68        && bytes.len() <= MAX_OUTPUT_BYTES
69    {
70        let (width, height) = image.dimensions();
71        return Ok(OptimizedImage {
72            bytes: clone_bytes(bytes)?,
73            mime_type: format.to_mime_type().to_owned(),
74            width,
75            height,
76        });
77    }
78
79    // Keep this check tied to the decoder's advertised representation. It catches dimensions
80    // whose multiplication would overflow before any image buffer is allocated.
81    let _ = checked_image_bytes(width, height, image.color())?;
82    optimize_dynamic_image(image)
83}
84
85/// Optimize an 8-bit RGBA image supplied as raw pixels.
86pub fn optimize_rgba(width: u32, height: u32, rgba: &[u8]) -> Result<OptimizedImage> {
87    if width == 0 || height == 0 {
88        bail!("image dimensions must be non-zero");
89    }
90    let expected_len = checked_rgba_len(width, height)?;
91    if rgba.len() != expected_len {
92        bail!("RGBA buffer length does not match image dimensions");
93    }
94
95    let mut pixels = Vec::new();
96    pixels
97        .try_reserve_exact(rgba.len())
98        .map_err(|error| anyhow!("allocate RGBA image: {error}"))?;
99    pixels.extend_from_slice(rgba);
100    let image = RgbaImage::from_raw(width, height, pixels)
101        .ok_or_else(|| anyhow!("invalid RGBA image dimensions"))?;
102    optimize_dynamic_image(DynamicImage::ImageRgba8(image))
103}
104
105fn optimize_dynamic_image(mut image: DynamicImage) -> Result<OptimizedImage> {
106    let transparent = contains_transparency(&image);
107
108    loop {
109        let (width, height) = image.dimensions();
110
111        // JPEG dimensions are limited to a 16-bit unsigned value. Resize before asking the
112        // encoder to handle an otherwise valid PNG/WebP with larger dimensions.
113        if !transparent && (width > MAX_JPEG_DIMENSION || height > MAX_JPEG_DIMENSION) {
114            let (next_width, next_height) =
115                fit_dimensions(width, height, MAX_JPEG_DIMENSION, MAX_JPEG_DIMENSION);
116            image = resize_checked(image, next_width, next_height)?;
117            continue;
118        }
119
120        if transparent {
121            if let Some(bytes) = encode_png(&image)? {
122                return Ok(OptimizedImage {
123                    bytes,
124                    mime_type: "image/png".to_owned(),
125                    width,
126                    height,
127                });
128            }
129        } else {
130            for quality in JPEG_QUALITIES {
131                if let Some(bytes) = encode_jpeg(&image, quality)? {
132                    return Ok(OptimizedImage {
133                        bytes,
134                        mime_type: "image/jpeg".to_owned(),
135                        width,
136                        height,
137                    });
138                }
139            }
140        }
141
142        let (next_width, next_height) = reduced_dimensions(width, height);
143        if (next_width, next_height) == (width, height) {
144            bail!("could not encode image within the 700 KiB output limit");
145        }
146        image = resize_checked(image, next_width, next_height)?;
147    }
148}
149
150fn encode_png(image: &DynamicImage) -> Result<Option<Vec<u8>>> {
151    let mut writer = LimitedWriter::new(MAX_OUTPUT_BYTES);
152    let result = image.write_with_encoder(PngEncoder::new_with_quality(
153        &mut writer,
154        CompressionType::Best,
155        FilterType::Adaptive,
156    ));
157    if writer.too_large {
158        return Ok(None);
159    }
160    result.context("encode PNG")?;
161    Ok(Some(writer.into_inner()))
162}
163
164fn encode_jpeg(image: &DynamicImage, quality: u8) -> Result<Option<Vec<u8>>> {
165    let mut writer = LimitedWriter::new(MAX_OUTPUT_BYTES);
166    let mut encoder = JpegEncoder::new_with_quality(&mut writer, quality);
167    let result = match image {
168        // Calling encode_image directly keeps RGBA inputs from making a second full-size RGB
169        // allocation. The encoder ignores alpha when it receives an opaque image.
170        DynamicImage::ImageLuma8(buffer) => encoder.encode_image(buffer),
171        DynamicImage::ImageLumaA8(buffer) => encoder.encode_image(buffer),
172        DynamicImage::ImageRgb8(buffer) => encoder.encode_image(buffer),
173        DynamicImage::ImageRgba8(buffer) => encoder.encode_image(buffer),
174        _ => image.write_with_encoder(encoder),
175    };
176    if writer.too_large {
177        return Ok(None);
178    }
179    result.context("encode JPEG")?;
180    Ok(Some(writer.into_inner()))
181}
182
183fn contains_transparency(image: &DynamicImage) -> bool {
184    match image {
185        DynamicImage::ImageLumaA8(buffer) => buffer.pixels().any(|pixel| pixel[1] != u8::MAX),
186        DynamicImage::ImageRgba8(buffer) => buffer.pixels().any(|pixel| pixel[3] != u8::MAX),
187        DynamicImage::ImageLumaA16(buffer) => buffer.pixels().any(|pixel| pixel[1] != u16::MAX),
188        DynamicImage::ImageRgba16(buffer) => buffer.pixels().any(|pixel| pixel[3] != u16::MAX),
189        DynamicImage::ImageRgba32F(buffer) => buffer.pixels().any(|pixel| pixel[3] != 1.0),
190        _ => false,
191    }
192}
193
194fn resize_checked(image: DynamicImage, width: u32, height: u32) -> Result<DynamicImage> {
195    if width == 0 || height == 0 {
196        bail!("image dimensions must be non-zero");
197    }
198    let current_bytes = checked_image_bytes(image.width(), image.height(), image.color())?;
199    let next_bytes = checked_image_bytes(width, height, image.color())?;
200    if current_bytes
201        .checked_add(next_bytes)
202        .is_none_or(|bytes| bytes > MAX_DECODED_BYTES.saturating_mul(2))
203    {
204        bail!("image resize would exceed the memory limit");
205    }
206    Ok(image.resize_exact(width, height, image::imageops::FilterType::Lanczos3))
207}
208
209fn reduced_dimensions(width: u32, height: u32) -> (u32, u32) {
210    let next_width = ((u64::from(width) * 4) / 5).max(1) as u32;
211    let next_height = ((u64::from(height) * 4) / 5).max(1) as u32;
212    if (next_width, next_height) == (width, height) {
213        if width > height {
214            (width.saturating_sub(1), height)
215        } else {
216            (width, height.saturating_sub(1))
217        }
218    } else {
219        (next_width, next_height)
220    }
221}
222
223fn fit_dimensions(width: u32, height: u32, max_width: u32, max_height: u32) -> (u32, u32) {
224    if width <= max_width && height <= max_height {
225        return (width, height);
226    }
227
228    let width_scale = u64::from(max_width) * u64::from(height);
229    let height_scale = u64::from(max_height) * u64::from(width);
230    let (scale_numerator, scale_denominator) = if width_scale <= height_scale {
231        (u64::from(max_width), u64::from(width))
232    } else {
233        (u64::from(max_height), u64::from(height))
234    };
235    let next_width = (u64::from(width) * scale_numerator / scale_denominator).max(1) as u32;
236    let next_height = (u64::from(height) * scale_numerator / scale_denominator).max(1) as u32;
237    (next_width, next_height)
238}
239
240fn checked_rgba_len(width: u32, height: u32) -> Result<usize> {
241    let bytes = u64::from(width)
242        .checked_mul(u64::from(height))
243        .and_then(|pixels| pixels.checked_mul(4))
244        .ok_or_else(|| anyhow!("RGBA image dimensions overflow"))?;
245    if bytes > MAX_DECODED_BYTES {
246        bail!("decoded image exceeds the 256 MiB limit");
247    }
248    usize::try_from(bytes).map_err(|_| anyhow!("RGBA image is too large for this platform"))
249}
250
251fn checked_image_bytes(width: u32, height: u32, color: ColorType) -> Result<u64> {
252    let bytes = u64::from(width)
253        .checked_mul(u64::from(height))
254        .and_then(|pixels| pixels.checked_mul(u64::from(color.bytes_per_pixel())))
255        .ok_or_else(|| anyhow!("image dimensions overflow"))?;
256    if bytes > MAX_DECODED_BYTES {
257        bail!("decoded image exceeds the 256 MiB limit");
258    }
259    Ok(bytes)
260}
261
262fn clone_bytes(bytes: &[u8]) -> Result<Vec<u8>> {
263    let mut cloned = Vec::new();
264    cloned
265        .try_reserve_exact(bytes.len())
266        .map_err(|error| anyhow!("allocate output image: {error}"))?;
267    cloned.extend_from_slice(bytes);
268    Ok(cloned)
269}
270
271struct LimitedWriter {
272    bytes: Vec<u8>,
273    limit: usize,
274    too_large: bool,
275}
276
277impl LimitedWriter {
278    fn new(limit: usize) -> Self {
279        Self {
280            bytes: Vec::new(),
281            limit,
282            too_large: false,
283        }
284    }
285
286    fn into_inner(self) -> Vec<u8> {
287        self.bytes
288    }
289}
290
291impl Write for LimitedWriter {
292    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
293        let Some(new_len) = self.bytes.len().checked_add(bytes.len()) else {
294            self.too_large = true;
295            return Err(io::Error::new(
296                io::ErrorKind::InvalidData,
297                "encoded image exceeds size limit",
298            ));
299        };
300        if new_len > self.limit {
301            self.too_large = true;
302            return Err(io::Error::new(
303                io::ErrorKind::InvalidData,
304                "encoded image exceeds size limit",
305            ));
306        }
307        self.bytes
308            .try_reserve(bytes.len())
309            .map_err(|error| io::Error::other(format!("allocate encoded image: {error}")))?;
310        self.bytes.extend_from_slice(bytes);
311        Ok(bytes.len())
312    }
313
314    fn flush(&mut self) -> io::Result<()> {
315        Ok(())
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    fn noise(width: u32, height: u32) -> Vec<u8> {
324        let mut state = 0x1234_5678_u32;
325        let mut bytes = vec![0; checked_rgba_len(width, height).unwrap()];
326        for byte in &mut bytes {
327            state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
328            *byte = (state >> 24) as u8;
329        }
330        for pixel in bytes.chunks_exact_mut(4) {
331            pixel[3] = u8::MAX;
332        }
333        bytes
334    }
335
336    #[test]
337    fn corrupt_input_is_rejected() {
338        assert!(optimize_image(b"not an image").is_err());
339    }
340
341    #[test]
342    fn oversized_input_is_rejected_before_decoding() {
343        let bytes = vec![0; MAX_INPUT_BYTES + 1];
344        let error = optimize_image(&bytes).unwrap_err().to_string();
345        assert!(error.contains("64 MiB"));
346    }
347
348    #[test]
349    fn transparent_rgba_stays_png() {
350        let mut rgba = vec![0; 2 * 2 * 4];
351        rgba.chunks_exact_mut(4).for_each(|pixel| pixel[3] = 128);
352        let optimized = optimize_rgba(2, 2, &rgba).unwrap();
353        assert_eq!(optimized.mime_type, "image/png");
354        assert!(optimized.bytes.len() <= MAX_OUTPUT_BYTES);
355        let decoder = ImageReader::new(Cursor::new(&optimized.bytes))
356            .with_guessed_format()
357            .unwrap()
358            .into_decoder()
359            .unwrap();
360        assert_eq!(decoder.dimensions(), (2, 2));
361        assert!(decoder.color_type().has_alpha());
362    }
363
364    #[test]
365    fn noisy_image_is_resized_to_fit_output_budget() {
366        let width = 1_600;
367        let height = 1_200;
368        let optimized = optimize_rgba(width, height, &noise(width, height)).unwrap();
369        assert_eq!(optimized.mime_type, "image/jpeg");
370        assert!(optimized.bytes.len() <= MAX_OUTPUT_BYTES);
371        assert!(optimized.width < width || optimized.height < height);
372    }
373
374    #[test]
375    fn rgba_length_must_match_dimensions() {
376        assert!(optimize_rgba(2, 2, &[0; 3]).is_err());
377    }
378}
mj_controller/hel_image.rs

mj_controller/
hel_image.rs
