1use std::io::Cursor;
2
3use anyhow::{Context, Result, bail};
4use image::codecs::gif::GifDecoder;
5use image::codecs::jpeg::JpegDecoder;
6use image::codecs::png::PngDecoder;
7use image::codecs::webp::WebPDecoder;
8use image::{AnimationDecoder, DynamicImage, ImageDecoder, ImageFormat, Limits, RgbaImage};
9
10use super::blocking;
11use super::validation::{validate_blob_size, validate_media_type};
12
13pub const MAX_IMAGE_EDGE: u32 = 16_384;
14pub const MAX_FRAME_PIXELS: u64 = 40_000_000;
15pub const MAX_ANIMATION_FRAMES: usize = 100;
16pub const MAX_ANIMATED_PIXELS: u64 = 100_000_000;
17pub const MAX_DECODER_ALLOCATION: u64 = 256 * 1024 * 1024;
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct ImageFacts {
21 pub media_type: String,
22 pub width: i64,
23 pub height: i64,
24 pub frame_count: usize,
25}
26
27#[derive(Debug)]
28pub struct ValidatedImage {
29 pub bytes: Vec<u8>,
30 pub facts: ImageFacts,
31}
32
33pub async fn validate_image(
34 bytes: Vec<u8>,
35 declared_media_type: Option<String>,
36) -> Result<ValidatedImage> {
37 blocking::run(move || validate_image_blocking(bytes, declared_media_type.as_deref())).await
38}
39
40pub fn validate_image_blocking(
41 bytes: Vec<u8>,
42 declared_media_type: Option<&str>,
43) -> Result<ValidatedImage> {
44 let (facts, _) = decode(&bytes, declared_media_type, false)?;
45 Ok(ValidatedImage { bytes, facts })
46}
47
48pub fn decode_first_frame(bytes: &[u8]) -> Result<RgbaImage> {
49 decode(bytes, None, true)?
50 .1
51 .ok_or_else(|| anyhow::anyhow!("error empty-attachment-image"))
52}
53
54fn decode(
55 bytes: &[u8],
56 declared_media_type: Option<&str>,
57 capture_first_frame: bool,
58) -> Result<(ImageFacts, Option<RgbaImage>)> {
59 validate_blob_size(bytes.len())?;
60 if let Some(media_type) = declared_media_type {
61 validate_media_type(media_type)?;
62 }
63 let format = image::guess_format(bytes).context("error invalid-attachment-image")?;
64 let media_type = media_type_for_format(format)
65 .ok_or_else(|| anyhow::anyhow!("error unsupported-attachment-image-format"))?;
66 if declared_media_type.is_some_and(|declared| declared != media_type) {
67 bail!(
68 "error attachment-media-type-mismatch declared={} detected={media_type}",
69 declared_media_type.unwrap_or_default()
70 );
71 }
72
73 let limits = image_limits();
74 let ((width, height), frame_count, first_frame) = match format {
75 ImageFormat::Png => {
76 let decoder = PngDecoder::with_limits(Cursor::new(bytes), limits)?;
77 let dimensions = decoder.dimensions();
78 validate_frame_dimensions(dimensions.0, dimensions.1)?;
79 if decoder.is_apng()? {
80 let (count, first) = validate_frames(
81 decoder.apng()?.into_frames(),
82 dimensions,
83 capture_first_frame,
84 )?;
85 (dimensions, count, first)
86 } else {
87 let image = DynamicImage::from_decoder(decoder)
88 .context("error invalid-attachment-image")?;
89 (
90 dimensions,
91 1,
92 capture_first_frame.then(|| image.into_rgba8()),
93 )
94 }
95 }
96 ImageFormat::Jpeg => {
97 let mut decoder = JpegDecoder::new(Cursor::new(bytes))?;
98 decoder.set_limits(limits)?;
99 let dimensions = decoder.dimensions();
100 validate_frame_dimensions(dimensions.0, dimensions.1)?;
101 let image =
102 DynamicImage::from_decoder(decoder).context("error invalid-attachment-image")?;
103 (
104 dimensions,
105 1,
106 capture_first_frame.then(|| image.into_rgba8()),
107 )
108 }
109 ImageFormat::Gif => {
110 let mut decoder = GifDecoder::new(Cursor::new(bytes))?;
111 decoder.set_limits(limits)?;
112 let dimensions = decoder.dimensions();
113 validate_frame_dimensions(dimensions.0, dimensions.1)?;
114 let (count, first) =
115 validate_frames(decoder.into_frames(), dimensions, capture_first_frame)?;
116 (dimensions, count, first)
117 }
118 ImageFormat::WebP => {
119 let mut decoder = WebPDecoder::new(Cursor::new(bytes))?;
120 decoder.set_limits(limits)?;
121 let dimensions = decoder.dimensions();
122 validate_frame_dimensions(dimensions.0, dimensions.1)?;
123 if decoder.has_animation() {
124 let (count, first) =
125 validate_frames(decoder.into_frames(), dimensions, capture_first_frame)?;
126 (dimensions, count, first)
127 } else {
128 let image = DynamicImage::from_decoder(decoder)
129 .context("error invalid-attachment-image")?;
130 (
131 dimensions,
132 1,
133 capture_first_frame.then(|| image.into_rgba8()),
134 )
135 }
136 }
137 _ => bail!("error unsupported-attachment-image-format"),
138 };
139
140 if frame_count == 0 {
141 bail!("error empty-attachment-image");
142 }
143 Ok((
144 ImageFacts {
145 media_type: media_type.to_string(),
146 width: i64::from(width),
147 height: i64::from(height),
148 frame_count,
149 },
150 first_frame,
151 ))
152}
153
154fn validate_frames(
155 frames: image::Frames<'_>,
156 dimensions: (u32, u32),
157 capture_first_frame: bool,
158) -> Result<(usize, Option<RgbaImage>)> {
159 let mut frame_count = 0_usize;
160 let mut cumulative_pixels = 0_u64;
161 let mut first_frame = None;
162 for frame in frames {
163 let frame = frame.context("error invalid-attachment-image")?;
164 frame_count += 1;
165 if frame_count > MAX_ANIMATION_FRAMES {
166 bail!("error attachment-animation-frame-limit");
167 }
168 let frame_dimensions = frame.buffer().dimensions();
169 if frame_dimensions != dimensions {
170 bail!("error invalid-attachment-frame-dimensions");
171 }
172 validate_frame_dimensions(frame_dimensions.0, frame_dimensions.1)?;
173 let pixels = u64::from(frame_dimensions.0) * u64::from(frame_dimensions.1);
174 cumulative_pixels = cumulative_pixels
175 .checked_add(pixels)
176 .ok_or_else(|| anyhow::anyhow!("error attachment-animation-pixel-limit"))?;
177 if cumulative_pixels > MAX_ANIMATED_PIXELS {
178 bail!("error attachment-animation-pixel-limit");
179 }
180 if capture_first_frame && first_frame.is_none() {
181 first_frame = Some(frame.into_buffer());
182 }
183 }
184 Ok((frame_count, first_frame))
185}
186
187fn image_limits() -> Limits {
188 let mut limits = Limits::default();
189 limits.max_image_width = Some(MAX_IMAGE_EDGE);
190 limits.max_image_height = Some(MAX_IMAGE_EDGE);
191 limits.max_alloc = Some(MAX_DECODER_ALLOCATION);
192 limits
193}
194
195fn validate_frame_dimensions(width: u32, height: u32) -> Result<()> {
196 if width == 0 || height == 0 || width > MAX_IMAGE_EDGE || height > MAX_IMAGE_EDGE {
197 bail!("error attachment-image-dimension-limit");
198 }
199 if u64::from(width) * u64::from(height) > MAX_FRAME_PIXELS {
200 bail!("error attachment-image-pixel-limit");
201 }
202 Ok(())
203}
204
205fn media_type_for_format(format: ImageFormat) -> Option<&'static str> {
206 match format {
207 ImageFormat::Png => Some("image/png"),
208 ImageFormat::Jpeg => Some("image/jpeg"),
209 ImageFormat::Gif => Some("image/gif"),
210 ImageFormat::WebP => Some("image/webp"),
211 _ => None,
212 }
213}
214
215#[cfg(test)]
216mod tests {
217 use std::io::Cursor;
218
219 use image::codecs::gif::GifEncoder;
220 use image::{DynamicImage, Frame, ImageFormat, Rgba, RgbaImage};
221
222 use super::*;
223
224 fn encoded(format: ImageFormat) -> Vec<u8> {
225 let image = DynamicImage::ImageRgba8(RgbaImage::from_pixel(3, 2, Rgba([1, 2, 3, 255])));
226 let mut bytes = Cursor::new(Vec::new());
227 image.write_to(&mut bytes, format).unwrap();
228 bytes.into_inner()
229 }
230
231 #[test]
232 fn detects_supported_formats_and_dimensions() {
233 for (format, media_type) in [
234 (ImageFormat::Png, "image/png"),
235 (ImageFormat::Jpeg, "image/jpeg"),
236 (ImageFormat::Gif, "image/gif"),
237 (ImageFormat::WebP, "image/webp"),
238 ] {
239 let image = validate_image_blocking(encoded(format), None).unwrap();
240 assert_eq!(image.facts.media_type, media_type);
241 assert_eq!((image.facts.width, image.facts.height), (3, 2));
242 }
243 }
244
245 #[test]
246 fn rejects_declared_media_type_mismatch() {
247 let error = validate_image_blocking(encoded(ImageFormat::Png), Some("image/jpeg"))
248 .unwrap_err()
249 .to_string();
250 assert!(error.contains("attachment-media-type-mismatch"));
251 }
252
253 #[test]
254 fn rejects_excessive_animation_frames() {
255 let frames = (0..=MAX_ANIMATION_FRAMES)
256 .map(|_| Frame::new(RgbaImage::from_pixel(1, 1, Rgba([1, 2, 3, 255]))));
257 let mut bytes = Vec::new();
258 GifEncoder::new(&mut bytes).encode_frames(frames).unwrap();
259 let error = validate_image_blocking(bytes, None)
260 .unwrap_err()
261 .to_string();
262 assert!(error.contains("attachment-animation-frame-limit"));
263 }
264
265 #[test]
266 fn rejects_malformed_and_unsupported_images() {
267 assert!(validate_image_blocking(b"not an image".to_vec(), None).is_err());
268 assert!(validate_image_blocking(b"BMunsupported".to_vec(), None).is_err());
269 }
270}