1use std::borrow::Cow;
13use std::io::BufWriter;
14use std::path::Path;
15
16use image::codecs::jpeg::JpegEncoder;
17use image::codecs::png::PngEncoder;
18use image::imageops::FilterType;
19use image::{DynamicImage, ImageDecoder, ImageEncoder, ImageFormat, ImageReader, Limits};
20
21pub const MAX_SIDE: u32 = 16_384;
23pub const MAX_PIXELS: u64 = 200_000_000;
25pub const THUMB_EDGE: u32 = 256;
27pub const SEND_MAX_BYTES: u64 = 5 * 1024 * 1024;
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum Sniffed {
35 Image(ImageFormat),
36 Heic,
37 Avif,
38 Svg,
39 NotImage,
40}
41
42pub fn sniff(head: &[u8]) -> Sniffed {
43 if head.starts_with(&[0x89, b'P', b'N', b'G']) {
44 return Sniffed::Image(ImageFormat::Png);
45 }
46 if head.starts_with(&[0xFF, 0xD8, 0xFF]) {
47 return Sniffed::Image(ImageFormat::Jpeg);
48 }
49 if head.starts_with(b"GIF8") {
50 return Sniffed::Image(ImageFormat::Gif);
51 }
52 if head.len() >= 12 && &head[0..4] == b"RIFF" && &head[8..12] == b"WEBP" {
53 return Sniffed::Image(ImageFormat::WebP);
54 }
55 if head.starts_with(b"BM") && head.len() >= 14 {
56 return Sniffed::Image(ImageFormat::Bmp);
57 }
58 if head.starts_with(b"II*\0") || head.starts_with(b"MM\0*") {
59 return Sniffed::Image(ImageFormat::Tiff);
60 }
61 if head.len() >= 12 && &head[4..8] == b"ftyp" {
62 let brand = &head[8..12];
63 if matches!(
64 brand,
65 b"heic" | b"heix" | b"hevc" | b"heim" | b"heis" | b"hevm" | b"mif1" | b"msf1"
66 ) {
67 return Sniffed::Heic;
68 }
69 if matches!(brand, b"avif" | b"avis") {
70 return Sniffed::Avif;
71 }
72 }
73 let text = String::from_utf8_lossy(&head[..head.len().min(1024)]);
74 let trimmed = text.trim_start_matches('\u{feff}').trim_start();
75 if trimmed.starts_with("<svg") || (trimmed.starts_with("<?xml") && trimmed.contains("<svg")) {
76 return Sniffed::Svg;
77 }
78 Sniffed::NotImage
79}
80
81pub fn mime_of(format: ImageFormat) -> &'static str {
82 match format {
83 ImageFormat::Png => "image/png",
84 ImageFormat::Jpeg => "image/jpeg",
85 ImageFormat::Gif => "image/gif",
86 ImageFormat::WebP => "image/webp",
87 ImageFormat::Bmp => "image/bmp",
88 ImageFormat::Tiff => "image/tiff",
89 _ => "application/octet-stream",
90 }
91}
92
93pub fn ext_of(format: ImageFormat) -> &'static str {
94 match format {
95 ImageFormat::Png => "png",
96 ImageFormat::Jpeg => "jpg",
97 ImageFormat::Gif => "gif",
98 ImageFormat::WebP => "webp",
99 ImageFormat::Bmp => "bmp",
100 ImageFormat::Tiff => "tif",
101 _ => "bin",
102 }
103}
104
105#[derive(Debug, Clone, PartialEq, Eq)]
108pub struct ProcessError {
109 pub code: &'static str,
110 pub message: String,
111}
112
113impl ProcessError {
114 fn new(code: &'static str, message: impl Into<String>) -> Self {
115 Self {
116 code,
117 message: message.into(),
118 }
119 }
120}
121
122impl std::fmt::Display for ProcessError {
123 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124 f.write_str(&self.message)
125 }
126}
127
128#[derive(Debug, Clone)]
130pub struct Encoded {
131 pub bytes: Vec<u8>,
132 pub format: ImageFormat,
133 pub width: u32,
134 pub height: u32,
135}
136
137#[derive(Debug, Clone)]
139pub struct Derived {
140 pub width: u32,
141 pub height: u32,
142 pub first_frame_only: bool,
143 pub thumb: Encoded,
144 pub send: Encoded,
145}
146
147pub fn header_dimensions(path: &Path) -> Result<(u32, u32), ProcessError> {
150 let (w, h) = ImageReader::open(path)
151 .map_err(|e| ProcessError::new("io", e.to_string()))?
152 .with_guessed_format()
153 .map_err(|e| ProcessError::new("io", e.to_string()))?
154 .into_dimensions()
155 .map_err(|e| ProcessError::new("decode", format!("Couldn't read the image: {e}")))?;
156 check_dimensions(w, h)?;
157 Ok((w, h))
158}
159
160pub fn decode_cost_bytes(w: u32, h: u32) -> u64 {
163 (w as u64) * (h as u64) * 8
164}
165
166pub fn derive(
169 path: &Path,
170 format: ImageFormat,
171 send_max_edge: u32,
172) -> Result<Derived, ProcessError> {
173 header_dimensions(path)?;
175
176 let mut reader = ImageReader::open(path).map_err(|e| ProcessError::new("io", e.to_string()))?;
177 reader.set_format(format);
178 let mut limits = Limits::default();
179 limits.max_image_width = Some(MAX_SIDE);
180 limits.max_image_height = Some(MAX_SIDE);
181 limits.max_alloc = Some(MAX_PIXELS * 4 + 64 * 1024 * 1024);
183 reader.limits(limits);
184 let mut decoder = reader
185 .into_decoder()
186 .map_err(|e| ProcessError::new("decode", format!("Couldn't read the image: {e}")))?;
187 let orientation = decoder.orientation().ok();
188 let mut img = DynamicImage::from_decoder(decoder)
189 .map_err(|e| ProcessError::new("decode", format!("Couldn't decode the image: {e}")))?;
190 if let Some(o) = orientation {
191 img.apply_orientation(o);
192 }
193
194 let first_frame_only = format == ImageFormat::Gif && gif_is_animated(path);
195 let alpha = has_real_alpha(&img);
196 let lossless_source = matches!(
197 format,
198 ImageFormat::Png | ImageFormat::Bmp | ImageFormat::Gif | ImageFormat::Tiff
199 );
200
201 let thumb = encode_thumbnail(&img, alpha)?;
202 let send = encode_send_copy(&img, alpha, lossless_source, send_max_edge.max(64))?;
203
204 Ok(Derived {
205 width: img.width(),
206 height: img.height(),
207 first_frame_only,
208 thumb,
209 send,
210 })
211}
212
213pub fn check_dimensions(w: u32, h: u32) -> Result<(), ProcessError> {
214 if w == 0 || h == 0 {
215 return Err(ProcessError::new("decode", "The image has no pixels."));
216 }
217 if w > MAX_SIDE || h > MAX_SIDE || (w as u64) * (h as u64) > MAX_PIXELS {
218 return Err(ProcessError::new(
219 "too_large_dimensions",
220 format!(
221 "The image is {w}×{h}. The largest supported is {MAX_SIDE} px on a side and {} megapixels.",
222 MAX_PIXELS / 1_000_000
223 ),
224 ));
225 }
226 Ok(())
227}
228
229fn has_real_alpha(img: &DynamicImage) -> bool {
232 if !img.color().has_alpha() {
233 return false;
234 }
235 let probe = if img.width().max(img.height()) > 512 {
237 img.thumbnail(512, 512)
238 } else {
239 img.clone()
240 };
241 probe.to_rgba8().pixels().any(|p| p.0[3] < 255)
242}
243
244fn gif_is_animated(path: &Path) -> bool {
245 use image::AnimationDecoder;
246 let Ok(f) = std::fs::File::open(path) else {
247 return false;
248 };
249 let Ok(dec) = image::codecs::gif::GifDecoder::new(std::io::BufReader::new(f)) else {
250 return false;
251 };
252 dec.into_frames().take(2).count() > 1
253}
254
255fn fit(w: u32, h: u32, edge: u32) -> (u32, u32) {
256 let long = w.max(h);
257 if long <= edge {
258 return (w, h);
259 }
260 let scale = edge as f64 / long as f64;
261 (
262 ((w as f64 * scale).round() as u32).max(1),
263 ((h as f64 * scale).round() as u32).max(1),
264 )
265}
266
267fn encode_thumbnail(img: &DynamicImage, alpha: bool) -> Result<Encoded, ProcessError> {
268 let (tw, th) = fit(img.width(), img.height(), THUMB_EDGE);
269 let small: Cow<'_, DynamicImage> = if (tw, th) == (img.width(), img.height()) {
270 Cow::Borrowed(img)
271 } else {
272 Cow::Owned(img.thumbnail(tw, th))
273 };
274 if alpha {
275 encode_png(&small)
276 } else {
277 encode_jpeg(&small, 80)
278 }
279}
280
281fn encode_send_copy(
282 img: &DynamicImage,
283 alpha: bool,
284 lossless_source: bool,
285 max_edge: u32,
286) -> Result<Encoded, ProcessError> {
287 let mut edge = max_edge;
288 for _ in 0..5 {
291 let (sw, sh) = fit(img.width(), img.height(), edge);
292 let scaled: Cow<'_, DynamicImage> = if (sw, sh) == (img.width(), img.height()) {
295 Cow::Borrowed(img)
296 } else {
297 Cow::Owned(img.resize_exact(sw, sh, FilterType::CatmullRom))
298 };
299 let out = if alpha {
300 encode_png(&scaled)?
301 } else if lossless_source {
302 let png = encode_png(&scaled)?;
305 if png.bytes.len() as u64 <= SEND_MAX_BYTES {
306 png
307 } else {
308 encode_jpeg(&scaled, 85)?
309 }
310 } else {
311 encode_jpeg(&scaled, 85)?
312 };
313 if out.bytes.len() as u64 <= SEND_MAX_BYTES {
314 return Ok(out);
315 }
316 edge = ((sw.max(sh) as f64) * 0.8) as u32;
317 }
318 Err(ProcessError::new(
319 "decode",
320 "Couldn't make a small enough copy of this image to send.",
321 ))
322}
323
324fn encode_png(img: &DynamicImage) -> Result<Encoded, ProcessError> {
325 let rgba;
327 let rgb;
328 let (buf, color): (&[u8], image::ExtendedColorType) = if img.color().has_alpha() {
329 rgba = img.to_rgba8();
330 (rgba.as_raw(), image::ExtendedColorType::Rgba8)
331 } else {
332 rgb = img.to_rgb8();
333 (rgb.as_raw(), image::ExtendedColorType::Rgb8)
334 };
335 let mut out = Vec::new();
336 PngEncoder::new(BufWriter::new(&mut out))
337 .write_image(buf, img.width(), img.height(), color)
338 .map_err(|e| ProcessError::new("decode", format!("Couldn't encode PNG: {e}")))?;
339 Ok(Encoded {
340 bytes: out,
341 format: ImageFormat::Png,
342 width: img.width(),
343 height: img.height(),
344 })
345}
346
347fn encode_jpeg(img: &DynamicImage, quality: u8) -> Result<Encoded, ProcessError> {
348 let rgb = img.to_rgb8();
349 let mut out = Vec::new();
350 JpegEncoder::new_with_quality(&mut out, quality)
351 .write_image(
352 rgb.as_raw(),
353 img.width(),
354 img.height(),
355 image::ExtendedColorType::Rgb8,
356 )
357 .map_err(|e| ProcessError::new("decode", format!("Couldn't encode JPEG: {e}")))?;
358 Ok(Encoded {
359 bytes: out,
360 format: ImageFormat::Jpeg,
361 width: img.width(),
362 height: img.height(),
363 })
364}
365
366#[cfg(test)]
367mod tests {
368 use super::*;
369 use image::{Rgb, RgbImage, Rgba, RgbaImage};
370
371 fn write(dir: &Path, name: &str, img: &DynamicImage, fmt: ImageFormat) -> std::path::PathBuf {
372 let p = dir.join(name);
373 img.save_with_format(&p, fmt).unwrap();
374 p
375 }
376
377 #[test]
378 fn sniff_recognises_formats_by_content_not_name() {
379 assert_eq!(
380 sniff(&[0x89, b'P', b'N', b'G', 0, 0]),
381 Sniffed::Image(ImageFormat::Png)
382 );
383 assert_eq!(
384 sniff(&[0xFF, 0xD8, 0xFF, 0xE0]),
385 Sniffed::Image(ImageFormat::Jpeg)
386 );
387 assert_eq!(sniff(b"GIF89a......"), Sniffed::Image(ImageFormat::Gif));
388 assert_eq!(
389 sniff(b"RIFF\0\0\0\0WEBPVP8 "),
390 Sniffed::Image(ImageFormat::WebP)
391 );
392 assert_eq!(sniff(b"\0\0\0\x18ftypheic\0\0\0\0"), Sniffed::Heic);
393 assert_eq!(sniff(b"\0\0\0\x18ftypmif1\0\0\0\0"), Sniffed::Heic);
394 assert_eq!(sniff(b"\0\0\0\x1cftypavif\0\0\0\0"), Sniffed::Avif);
395 assert_eq!(sniff(b"\xef\xbb\xbf <svg xmlns=..."), Sniffed::Svg);
396 assert_eq!(sniff(b"<?xml version=\"1.0\"?><svg>"), Sniffed::Svg);
397 assert_eq!(sniff(b"fn main() {}"), Sniffed::NotImage);
398 assert_eq!(sniff(b""), Sniffed::NotImage);
399 }
400
401 #[test]
402 fn dimension_guard_rejects_bombs_before_decode() {
403 assert!(check_dimensions(4000, 3000).is_ok());
404 assert_eq!(
405 check_dimensions(20_000, 10).unwrap_err().code,
406 "too_large_dimensions"
407 );
408 assert_eq!(
409 check_dimensions(16_000, 16_000).unwrap_err().code,
410 "too_large_dimensions"
411 );
412 assert_eq!(check_dimensions(0, 10).unwrap_err().code, "decode");
413 }
414
415 #[test]
416 fn large_photo_is_downscaled_to_the_send_edge_as_jpeg() {
417 let dir = tempfile::tempdir().unwrap();
418 let img = DynamicImage::ImageRgb8(RgbImage::from_fn(3000, 1500, |x, y| {
419 Rgb([(x % 256) as u8, (y % 256) as u8, ((x + y) % 256) as u8])
420 }));
421 let p = write(dir.path(), "photo.jpg", &img, ImageFormat::Jpeg);
422 let d = derive(&p, ImageFormat::Jpeg, 2000).unwrap();
423 assert_eq!((d.width, d.height), (3000, 1500));
424 assert_eq!(d.send.format, ImageFormat::Jpeg);
425 assert_eq!((d.send.width, d.send.height), (2000, 1000));
426 assert!(d.send.bytes.len() as u64 <= SEND_MAX_BYTES);
427 assert_eq!(d.thumb.width.max(d.thumb.height), THUMB_EDGE);
428 }
429
430 #[test]
431 fn small_screenshot_stays_png_at_full_size() {
432 let dir = tempfile::tempdir().unwrap();
433 let img =
434 DynamicImage::ImageRgba8(RgbaImage::from_pixel(800, 600, Rgba([30, 30, 30, 255])));
435 let p = write(dir.path(), "shot.png", &img, ImageFormat::Png);
436 let d = derive(&p, ImageFormat::Png, 2000).unwrap();
437 assert_eq!(d.send.format, ImageFormat::Png);
438 assert_eq!((d.send.width, d.send.height), (800, 600));
439 assert_eq!(d.thumb.format, ImageFormat::Jpeg);
441 }
442
443 #[test]
444 fn transparency_is_kept_as_png() {
445 let dir = tempfile::tempdir().unwrap();
446 let img = DynamicImage::ImageRgba8(RgbaImage::from_fn(300, 300, |x, _| {
447 Rgba([255, 0, 0, if x < 150 { 0 } else { 255 }])
448 }));
449 let p = write(dir.path(), "logo.png", &img, ImageFormat::Png);
450 let d = derive(&p, ImageFormat::Png, 2000).unwrap();
451 assert_eq!(d.send.format, ImageFormat::Png);
452 assert_eq!(d.thumb.format, ImageFormat::Png);
453 }
454
455 #[test]
456 fn send_copy_carries_no_exif() {
457 let dir = tempfile::tempdir().unwrap();
458 let img = DynamicImage::ImageRgb8(RgbImage::from_pixel(64, 64, Rgb([1, 2, 3])));
459 let p = write(dir.path(), "a.jpg", &img, ImageFormat::Jpeg);
460 let d = derive(&p, ImageFormat::Jpeg, 2000).unwrap();
461 assert!(!d.send.bytes.windows(4).any(|w| w == b"Exif"));
463 }
464
465 #[test]
466 fn truncated_file_fails_cleanly() {
467 let dir = tempfile::tempdir().unwrap();
468 let img = DynamicImage::ImageRgb8(RgbImage::from_pixel(200, 200, Rgb([9, 9, 9])));
469 let p = write(dir.path(), "t.png", &img, ImageFormat::Png);
470 let bytes = std::fs::read(&p).unwrap();
471 std::fs::write(&p, &bytes[..bytes.len() / 3]).unwrap();
472 let e = derive(&p, ImageFormat::Png, 2000).unwrap_err();
473 assert_eq!(e.code, "decode");
474 }
475
476 #[test]
477 fn fit_keeps_aspect_and_never_upscales() {
478 assert_eq!(fit(4000, 2000, 2000), (2000, 1000));
479 assert_eq!(fit(100, 50, 2000), (100, 50));
480 assert_eq!(fit(1, 5000, 256), (1, 256));
481 }
482}