forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
6.4 KiB
1//! Avatar upload processing: validate, decode with limits, crop, resize and re-encode as PNG.
2//!
3//! Re-encoding to PNG drops all metadata (EXIF, GPS, comments) from the input.
4
5use std::io::Cursor;
6
7use image::{ImageFormat, ImageReader, Limits, imageops::FilterType};
8use sha2::{Digest, Sha256};
9
10use crate::error::{Error, Result};
11
12/// Largest accepted upload, in bytes.
13pub const MAX_UPLOAD_BYTES: usize = 5 * 1024 * 1024;
14
15/// Largest accepted width or height of the decoded source image, in pixels.
16const MAX_SOURCE_DIMENSION: u32 = 4096;
17
18/// Upper bound on decoder allocations (128 MiB).
19const MAX_ALLOC_BYTES: u64 = 128 * 1024 * 1024;
20
21/// Edge length of the stored avatar, in pixels.
22const AVATAR_SIZE: u32 = 256;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct ProcessedAvatar {
26 pub png: Vec<u8>,
27 pub hash: String,
28}
29
30/// Decode with limits (max 4096x4096, image::Limits), center-crop to a square, resize to 256x256
31/// (Lanczos3), re-encode as PNG. Metadata (EXIF/GPS) is dropped by re-encoding. SVG, unknown
32/// formats, too-large input or a decode failure give Error::Invalid.
33pub fn process_avatar(bytes: &[u8]) -> Result<ProcessedAvatar> {
34 if bytes.len() > MAX_UPLOAD_BYTES {
35 return Err(Error::Invalid(format!(
36 "avatar is larger than {MAX_UPLOAD_BYTES} bytes"
37 )));
38 }
39
40 // Format is decided from magic bytes only; the client-supplied type is never consulted.
41 let mut reader = ImageReader::new(Cursor::new(bytes))
42 .with_guessed_format()
43 .map_err(|_| Error::Invalid("unreadable image header".into()))?;
44 match reader.format() {
45 Some(ImageFormat::Png | ImageFormat::Jpeg | ImageFormat::WebP | ImageFormat::Gif) => {}
46 _ => {
47 return Err(Error::Invalid(
48 "avatar must be a PNG, JPEG, WebP or GIF image".into(),
49 ));
50 }
51 }
52
53 let mut limits = Limits::default();
54 limits.max_image_width = Some(MAX_SOURCE_DIMENSION);
55 limits.max_image_height = Some(MAX_SOURCE_DIMENSION);
56 limits.max_alloc = Some(MAX_ALLOC_BYTES);
57 reader.limits(limits);
58
59 let decoded = reader
60 .decode()
61 .map_err(|e| Error::Invalid(format!("could not decode avatar: {e}")))?;
62
63 // Center-crop to the largest square, then scale down to the avatar size.
64 let (width, height) = (decoded.width(), decoded.height());
65 let side = width.min(height);
66 let x = (width - side) / 2;
67 let y = (height - side) / 2;
68 let square = decoded.crop_imm(x, y, side, side);
69 let resized = square.resize_exact(AVATAR_SIZE, AVATAR_SIZE, FilterType::Lanczos3);
70
71 let mut png = Cursor::new(Vec::new());
72 resized
73 .write_to(&mut png, ImageFormat::Png)
74 .map_err(|e| Error::Invalid(format!("could not encode avatar: {e}")))?;
75 let png = png.into_inner();
76
77 let digest = Sha256::digest(&png);
78 let hash = hex::encode(&digest[..8]);
79
80 Ok(ProcessedAvatar { png, hash })
81}
82
83#[cfg(test)]
84mod tests {
85 use super::*;
86 use image::{DynamicImage, RgbImage};
87
88 const PNG_MAGIC: [u8; 8] = [0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A];
89
90 fn encode(img: &DynamicImage, format: ImageFormat) -> Vec<u8> {
91 let mut out = Cursor::new(Vec::new());
92 img.write_to(&mut out, format).expect("encode test image");
93 out.into_inner()
94 }
95
96 fn sample(width: u32, height: u32) -> DynamicImage {
97 DynamicImage::ImageRgb8(RgbImage::from_fn(width, height, |x, y| {
98 image::Rgb([(x % 256) as u8, (y % 256) as u8, 128])
99 }))
100 }
101
102 fn crc32(data: &[u8]) -> u32 {
103 let mut crc = 0xFFFF_FFFFu32;
104 for &byte in data {
105 crc ^= u32::from(byte);
106 for _ in 0..8 {
107 let mask = (crc & 1).wrapping_neg();
108 crc = (crc >> 1) ^ (0xEDB8_8320 & mask);
109 }
110 }
111 !crc
112 }
113
114 #[test]
115 fn png_is_cropped_and_resized_to_256() {
116 let input = encode(&sample(300, 200), ImageFormat::Png);
117 let out = process_avatar(&input).expect("png accepted");
118 assert_eq!(&out.png[..8], &PNG_MAGIC);
119 let decoded = image::load_from_memory_with_format(&out.png, ImageFormat::Png).unwrap();
120 assert_eq!((decoded.width(), decoded.height()), (256, 256));
121 }
122
123 #[test]
124 fn jpeg_is_accepted() {
125 let input = encode(&sample(300, 200), ImageFormat::Jpeg);
126 let out = process_avatar(&input).expect("jpeg accepted");
127 assert_eq!(&out.png[..8], &PNG_MAGIC);
128 }
129
130 #[test]
131 fn svg_is_rejected() {
132 let svg = br#"<svg xmlns="http://www.w3.org/2000/svg" width="10" height="10"></svg>"#;
133 assert!(matches!(process_avatar(svg), Err(Error::Invalid(_))));
134 }
135
136 #[test]
137 fn random_bytes_are_rejected() {
138 let junk: Vec<u8> = (0..4096u32)
139 .map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
140 .collect();
141 assert!(matches!(process_avatar(&junk), Err(Error::Invalid(_))));
142 assert!(matches!(process_avatar(&[]), Err(Error::Invalid(_))));
143 }
144
145 #[test]
146 fn oversized_input_is_rejected() {
147 let input = vec![0u8; MAX_UPLOAD_BYTES + 1];
148 assert!(matches!(process_avatar(&input), Err(Error::Invalid(_))));
149 }
150
151 #[test]
152 fn huge_declared_dimensions_are_rejected() {
153 // Take a valid tiny PNG and rewrite its IHDR to claim 50000x50000, fixing the CRC so
154 // only the dimension limit can reject it.
155 let mut png = encode(&sample(2, 2), ImageFormat::Png);
156 // Layout: 8-byte signature, 4-byte length, 4-byte "IHDR", then width and height.
157 assert_eq!(&png[12..16], b"IHDR");
158 png[16..20].copy_from_slice(&50_000u32.to_be_bytes());
159 png[20..24].copy_from_slice(&50_000u32.to_be_bytes());
160 let crc = crc32(&png[12..29]);
161 png[29..33].copy_from_slice(&crc.to_be_bytes());
162 assert!(matches!(process_avatar(&png), Err(Error::Invalid(_))));
163 }
164
165 #[test]
166 fn output_starts_with_png_magic() {
167 let input = encode(&sample(64, 64), ImageFormat::Png);
168 let out = process_avatar(&input).unwrap();
169 assert_eq!(&out.png[..8], &PNG_MAGIC);
170 }
171
172 #[test]
173 fn hash_is_stable_and_well_formed() {
174 let input = encode(&sample(120, 90), ImageFormat::Png);
175 let a = process_avatar(&input).unwrap();
176 let b = process_avatar(&input).unwrap();
177 assert_eq!(a.hash, b.hash);
178 assert_eq!(a.hash.len(), 16);
179 assert!(a.hash.bytes().all(|c| c.is_ascii_hexdigit()));
180 assert_eq!(a.png, b.png);
181 }
182}