Skip to main content

max / multithreaded

8.2 KB · 274 lines History Blame Raw
1 //! S3 storage client for image uploads.
2 //! Delegates S3 operations to the shared `s3_storage` crate.
3
4 use uuid::Uuid;
5
6 use crate::config::S3Config;
7
8 /// S3 client wrapper for image storage.
9 #[derive(Clone)]
10 pub struct S3Storage {
11 inner: s3_storage::S3Client,
12 }
13
14 /// Maximum image size: 5 MB.
15 pub const MAX_IMAGE_SIZE: usize = 5 * 1024 * 1024;
16
17 /// Allowed image content types.
18 const ALLOWED_CONTENT_TYPES: &[&str] = &[
19 "image/png",
20 "image/jpeg",
21 "image/gif",
22 "image/webp",
23 ];
24
25 /// Allowed file extensions (lowercase).
26 #[cfg(test)]
27 const ALLOWED_EXTENSIONS: &[&str] = &["png", "jpg", "jpeg", "gif", "webp"];
28
29 impl S3Storage {
30 /// Create a new S3 client from configuration.
31 pub async fn new(config: &S3Config) -> Result<Self, String> {
32 let s3_config = s3_storage::S3Config {
33 endpoint: config.endpoint.clone(),
34 bucket: config.bucket.clone(),
35 access_key: config.access_key.clone(),
36 secret_key: config.secret_key.clone(),
37 region: config.region.clone(),
38 };
39
40 let inner = s3_storage::S3Client::new(&s3_config).await?;
41 Ok(Self { inner })
42 }
43
44 /// Upload bytes to S3.
45 #[tracing::instrument(skip_all)]
46 pub async fn upload(&self, s3_key: &str, content_type: &str, data: Vec<u8>) -> Result<(), String> {
47 self.inner.upload(s3_key, content_type, data, None).await
48 }
49
50 /// Download bytes from S3.
51 #[tracing::instrument(skip_all)]
52 pub async fn download(&self, s3_key: &str) -> Result<(Vec<u8>, String), String> {
53 self.inner.download(s3_key).await
54 }
55
56 /// Delete an object from S3.
57 #[tracing::instrument(skip_all)]
58 pub async fn delete(&self, s3_key: &str) -> Result<(), String> {
59 self.inner.delete(s3_key).await
60 }
61 }
62
63 /// Generate an S3 key for a forum image.
64 /// Format: `mt/{community_slug}/{uuid}.{ext}`
65 pub fn generate_image_key(community_slug: &str, ext: &str) -> String {
66 let id = Uuid::new_v4();
67 format!("mt/{community_slug}/{id}.{ext}")
68 }
69
70 /// Validate an uploaded file. Returns the sanitized extension and content type.
71 pub fn validate_image(
72 filename: &str,
73 content_type: &str,
74 size: usize,
75 ) -> Result<(&'static str, &'static str), &'static str> {
76 if size > MAX_IMAGE_SIZE {
77 return Err("Image exceeds 5 MB limit.");
78 }
79 if size == 0 {
80 return Err("Empty file.");
81 }
82
83 // Check extension
84 let ext = filename
85 .rsplit('.')
86 .next()
87 .map(|e| e.to_lowercase())
88 .unwrap_or_default();
89
90 let ext_str: &'static str = match ext.as_str() {
91 "png" => "png",
92 "jpg" | "jpeg" => "jpg",
93 "gif" => "gif",
94 "webp" => "webp",
95 _ => return Err("Allowed types: png, jpg, gif, webp."),
96 };
97
98 // Check content type
99 let ct: &'static str = ALLOWED_CONTENT_TYPES
100 .iter()
101 .find(|&&ct| ct == content_type)
102 .copied()
103 .ok_or("Invalid image content type.")?;
104
105 // Cross-validate: extension should match content type
106 let ext_matches = match ext_str {
107 "png" => ct == "image/png",
108 "jpg" => ct == "image/jpeg",
109 "gif" => ct == "image/gif",
110 "webp" => ct == "image/webp",
111 _ => false,
112 };
113 if !ext_matches {
114 return Err("File extension does not match content type.");
115 }
116
117 Ok((ext_str, ct))
118 }
119
120 /// Strip EXIF metadata from JPEG data.
121 /// Works by copying all JPEG segments except APP1 (EXIF) and APP13 (IPTC).
122 /// Returns the cleaned data, or the original data if parsing fails.
123 pub fn strip_exif_jpeg(data: &[u8]) -> Vec<u8> {
124 if data.len() < 4 || data[0] != 0xFF || data[1] != 0xD8 {
125 return data.to_vec(); // Not a valid JPEG
126 }
127
128 let mut out = Vec::with_capacity(data.len());
129 out.extend_from_slice(&[0xFF, 0xD8]); // SOI marker
130
131 let mut i = 2;
132 while i + 1 < data.len() {
133 if data[i] != 0xFF {
134 // Not a marker — copy rest as-is (image data)
135 out.extend_from_slice(&data[i..]);
136 break;
137 }
138
139 let marker = data[i + 1];
140
141 // SOS (Start of Scan) — copy the rest verbatim (compressed data follows)
142 if marker == 0xDA {
143 out.extend_from_slice(&data[i..]);
144 break;
145 }
146
147 // Markers without length (RST0-RST7, SOI, EOI, TEM)
148 if marker == 0x00 || marker == 0x01 || (0xD0..=0xD9).contains(&marker) {
149 out.extend_from_slice(&data[i..i + 2]);
150 i += 2;
151 continue;
152 }
153
154 // Read segment length
155 if i + 3 >= data.len() {
156 out.extend_from_slice(&data[i..]);
157 break;
158 }
159 let seg_len = ((data[i + 2] as usize) << 8) | (data[i + 3] as usize);
160 let total = 2 + seg_len; // marker (2) + length includes itself
161
162 if i + total > data.len() {
163 out.extend_from_slice(&data[i..]);
164 break;
165 }
166
167 // Skip APP1 (0xE1 = EXIF) and APP13 (0xED = IPTC/Photoshop)
168 if marker == 0xE1 || marker == 0xED {
169 i += total;
170 continue;
171 }
172
173 out.extend_from_slice(&data[i..i + total]);
174 i += total;
175 }
176
177 out
178 }
179
180 #[cfg(test)]
181 mod tests {
182 use super::*;
183
184 #[test]
185 fn validate_valid_png() {
186 let (ext, ct) = validate_image("photo.png", "image/png", 1024).unwrap();
187 assert_eq!(ext, "png");
188 assert_eq!(ct, "image/png");
189 }
190
191 #[test]
192 fn validate_valid_jpeg() {
193 let (ext, _) = validate_image("photo.jpg", "image/jpeg", 1024).unwrap();
194 assert_eq!(ext, "jpg");
195 }
196
197 #[test]
198 fn validate_rejects_oversized() {
199 let err = validate_image("big.png", "image/png", 6 * 1024 * 1024).unwrap_err();
200 assert!(err.contains("5 MB"));
201 }
202
203 #[test]
204 fn validate_rejects_bad_extension() {
205 let err = validate_image("file.exe", "application/octet-stream", 1024).unwrap_err();
206 assert!(err.contains("Allowed types"));
207 }
208
209 #[test]
210 fn validate_rejects_mismatched_type() {
211 let err = validate_image("photo.png", "image/jpeg", 1024).unwrap_err();
212 assert!(err.contains("does not match"));
213 }
214
215 #[test]
216 fn validate_rejects_empty() {
217 let err = validate_image("photo.png", "image/png", 0).unwrap_err();
218 assert!(err.contains("Empty"));
219 }
220
221 #[test]
222 fn strip_exif_preserves_non_jpeg() {
223 let data = b"not a jpeg";
224 let result = strip_exif_jpeg(data);
225 assert_eq!(result, data);
226 }
227
228 #[test]
229 fn strip_exif_minimal_jpeg() {
230 // Minimal JPEG: SOI + APP0 (JFIF) + SOS + EOI
231 let mut jpeg = vec![0xFF, 0xD8]; // SOI
232 // APP0 segment (marker + length + data)
233 jpeg.extend_from_slice(&[0xFF, 0xE0, 0x00, 0x04, 0x00, 0x00]); // APP0, len=4
234 // APP1 segment (EXIF — should be stripped)
235 jpeg.extend_from_slice(&[0xFF, 0xE1, 0x00, 0x04, 0xAA, 0xBB]); // APP1, len=4
236 // SOS + fake data
237 jpeg.extend_from_slice(&[0xFF, 0xDA, 0x00, 0x02]);
238 jpeg.extend_from_slice(&[0xFF, 0xD9]); // EOI
239
240 let result = strip_exif_jpeg(&jpeg);
241 // Should contain SOI, APP0, SOS+data — but NOT APP1
242 assert!(result.len() < jpeg.len(), "EXIF should be stripped");
243 // Check APP1 marker is absent
244 let has_app1 = result.windows(2).any(|w| w == [0xFF, 0xE1]);
245 assert!(!has_app1, "APP1 should be removed");
246 // Check APP0 is preserved
247 let has_app0 = result.windows(2).any(|w| w == [0xFF, 0xE0]);
248 assert!(has_app0, "APP0 should be preserved");
249 }
250
251 #[test]
252 fn generate_key_format() {
253 let key = generate_image_key("test-comm", "png");
254 assert!(key.starts_with("mt/test-comm/"));
255 assert!(key.ends_with(".png"));
256 }
257
258 #[test]
259 fn allowed_extensions_match() {
260 // Verify our ALLOWED_EXTENSIONS list matches validate_image behavior
261 for ext in ALLOWED_EXTENSIONS {
262 let filename = format!("test.{ext}");
263 let ct = match *ext {
264 "png" => "image/png",
265 "jpg" | "jpeg" => "image/jpeg",
266 "gif" => "image/gif",
267 "webp" => "image/webp",
268 _ => continue,
269 };
270 assert!(validate_image(&filename, ct, 1024).is_ok(), "Extension {ext} should be valid");
271 }
272 }
273 }
274