1use std::cell::RefCell;
31use std::collections::HashMap;
32use std::fmt;
33
34use axum::body::Bytes;
35use serde::de::Error as _;
36use serde::{Deserialize, Deserializer, Serialize, Serializer};
37
38use crate::Result;
39use crate::crypto::random_token;
40use crate::storage::Storage;
41
42#[derive(Clone)]
44#[non_exhaustive]
45pub struct Upload {
46 file_name: String,
47 content_type: String,
48 bytes: Bytes,
49}
50
51impl fmt::Debug for Upload {
52 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53 f.debug_struct("Upload")
54 .field("file_name", &self.file_name)
55 .field("content_type", &self.content_type)
56 .field("size", &self.bytes.len())
57 .finish()
58 }
59}
60
61const ACTIVE_EXTENSIONS: &[&str] = &[
63 "html", "htm", "xhtml", "xht", "shtml", "mht", "mhtml", "xml", "xsl", "xslt", "js", "mjs",
64 "cjs", "php", "phtml", "asp", "aspx", "jsp", "cgi", "pl", "py", "sh", "swf", "hta", "htc",
65];
66
67impl Upload {
68 pub fn new(
71 file_name: impl Into<String>,
72 content_type: impl Into<String>,
73 bytes: impl Into<Bytes>,
74 ) -> Self {
75 Self {
76 file_name: file_name.into(),
77 content_type: content_type.into(),
78 bytes: bytes.into(),
79 }
80 }
81
82 pub fn file_name(&self) -> &str {
84 &self.file_name
85 }
86
87 pub fn content_type(&self) -> &str {
89 &self.content_type
90 }
91
92 pub fn bytes(&self) -> &Bytes {
94 &self.bytes
95 }
96
97 pub fn size(&self) -> usize {
99 self.bytes.len()
100 }
101
102 pub fn extension(&self) -> Option<String> {
104 let (_, ext) = self.file_name.rsplit_once('.')?;
105 let ext = ext.to_ascii_lowercase();
106 (!ext.is_empty() && ext.len() <= 10 && ext.chars().all(|c| c.is_ascii_alphanumeric()))
107 .then_some(ext)
108 }
109
110 pub fn sniffed_type(&self) -> Option<&'static str> {
112 let b = &self.bytes[..];
113 if b.starts_with(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) {
114 Some("image/png")
115 } else if b.starts_with(&[0xFF, 0xD8, 0xFF]) {
116 Some("image/jpeg")
117 } else if b.starts_with(b"GIF87a") || b.starts_with(b"GIF89a") {
118 Some("image/gif")
119 } else if b.len() >= 12 && &b[..4] == b"RIFF" && &b[8..12] == b"WEBP" {
120 Some("image/webp")
121 } else if b.starts_with(b"%PDF-") {
122 Some("application/pdf")
123 } else {
124 None
125 }
126 }
127
128 pub fn is_image(&self) -> bool {
130 self.sniffed_type().is_some_and(|t| t.starts_with("image/"))
131 }
132
133 pub fn dimensions(&self) -> Option<(u32, u32)> {
136 image_dimensions(&self.bytes)
137 }
138
139 fn safe_extension(&self) -> String {
145 match self.sniffed_type() {
146 Some("image/png") => "png".into(),
147 Some("image/jpeg") => "jpg".into(),
148 Some("image/gif") => "gif".into(),
149 Some("image/webp") => "webp".into(),
150 Some("application/pdf") => "pdf".into(),
151 _ => match self.extension() {
152 Some(ext) if ACTIVE_EXTENSIONS.contains(&ext.as_str()) => "txt".into(),
153 Some(ext) => ext,
154 None => "bin".into(),
155 },
156 }
157 }
158
159 fn key(&self, dir: &str) -> String {
160 let dir = dir.trim_matches('/');
161 let name = format!("{}.{}", &random_token()[..24], self.safe_extension());
162 if dir.is_empty() {
163 name
164 } else {
165 format!("{dir}/{name}")
166 }
167 }
168
169 pub async fn store(&self, storage: &Storage, dir: &str) -> Result<String> {
171 let key = self.key(dir);
172 storage.put(&key, self.bytes.clone()).await?;
173 Ok(key)
174 }
175
176 pub async fn store_public(&self, storage: &Storage, dir: &str) -> Result<String> {
178 let key = self.key(&format!("public/{}", dir.trim_matches('/')));
179 storage.put(&key, self.bytes.clone()).await?;
180 Ok(key)
181 }
182}
183
184impl Serialize for Upload {
186 fn serialize<S: Serializer>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> {
187 serializer.serialize_str(&self.file_name)
188 }
189}
190
191thread_local! {
192 static UPLOADS: RefCell<HashMap<String, Upload>> = RefCell::new(HashMap::new());
194}
195
196impl<'de> Deserialize<'de> for Upload {
197 fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
198 let token = String::deserialize(deserializer)?;
199 UPLOADS
200 .with(|uploads| uploads.borrow().get(&token).cloned())
201 .ok_or_else(|| D::Error::custom(NOT_A_FILE))
202 }
203}
204
205pub(crate) const NOT_A_FILE: &str = "expected an uploaded file";
206
207pub(crate) fn token(index: usize) -> String {
209 format!("\u{1}renox-upload:{index}")
210}
211
212pub(crate) fn with_uploads<R>(uploads: &HashMap<String, Upload>, f: impl FnOnce() -> R) -> R {
214 UPLOADS.with(|cell| *cell.borrow_mut() = uploads.clone());
215 let result = f();
216 UPLOADS.with(|cell| cell.borrow_mut().clear());
217 result
218}
219
220fn image_dimensions(b: &[u8]) -> Option<(u32, u32)> {
222 let be16 = |i: usize| Some(u16::from_be_bytes(b.get(i..i + 2)?.try_into().ok()?) as u32);
223 let le16 = |i: usize| Some(u16::from_le_bytes(b.get(i..i + 2)?.try_into().ok()?) as u32);
224 let le24 = |i: usize| {
225 let s = b.get(i..i + 3)?;
226 Some(s[0] as u32 | (s[1] as u32) << 8 | (s[2] as u32) << 16)
227 };
228 let size = if b.starts_with(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) {
229 if b.get(12..16)? != b"IHDR" {
230 return None;
231 }
232 let width = u32::from_be_bytes(b.get(16..20)?.try_into().ok()?);
233 let height = u32::from_be_bytes(b.get(20..24)?.try_into().ok()?);
234 (width, height)
235 } else if b.starts_with(b"GIF87a") || b.starts_with(b"GIF89a") {
236 (le16(6)?, le16(8)?)
237 } else if b.starts_with(&[0xFF, 0xD8]) {
238 let mut i = 2;
241 loop {
242 while *b.get(i)? != 0xFF {
243 i += 1;
244 }
245 while *b.get(i)? == 0xFF {
246 i += 1;
247 }
248 let marker = *b.get(i)?;
249 i += 1;
250 if matches!(marker, 0xD0..=0xD9 | 0x01) {
251 continue;
252 }
253 let length = be16(i)? as usize;
254 if (0xC0..=0xCF).contains(&marker) && !matches!(marker, 0xC4 | 0xC8 | 0xCC) {
255 break (be16(i + 5)?, be16(i + 3)?);
256 }
257 i += length;
258 }
259 } else if b.len() >= 30 && &b[..4] == b"RIFF" && &b[8..12] == b"WEBP" {
260 match &b[12..16] {
261 b"VP8 " if b.get(23..26)? == [0x9D, 0x01, 0x2A] => {
262 (le16(26)? & 0x3FFF, le16(28)? & 0x3FFF)
263 }
264 b"VP8L" if b[20] == 0x2F => {
265 let bits = u32::from_le_bytes(b.get(21..25)?.try_into().ok()?);
266 ((bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1)
267 }
268 b"VP8X" => (le24(24)? + 1, le24(27)? + 1),
269 _ => return None,
270 }
271 } else {
272 return None;
273 };
274 (size.0 > 0 && size.1 > 0).then_some(size)
275}
276
277#[cfg(test)]
278mod tests {
279 use super::*;
280
281 #[test]
282 fn reads_image_dimensions_from_headers() {
283 let mut png = b"\x89PNG\r\n\x1a\n\0\0\0\x0dIHDR".to_vec();
285 png.extend_from_slice(&3u32.to_be_bytes());
286 png.extend_from_slice(&2u32.to_be_bytes());
287 assert_eq!(image_dimensions(&png), Some((3, 2)));
288 let gif = b"GIF89a\x80\x02\xe0\x01";
290 assert_eq!(image_dimensions(gif), Some((640, 480)));
291 let jpeg = [
293 0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x04, 0x00, 0x00, 0xFF, 0xC0, 0x00, 0x11, 0x08, 0x00,
294 100, 0x00, 200,
295 ];
296 assert_eq!(image_dimensions(&jpeg), Some((200, 100)));
297 let mut webp = b"RIFF\0\0\0\0WEBPVP8X\x0a\0\0\0\0\0\0\0".to_vec();
299 webp.extend_from_slice(&[0x7F, 0x07, 0x00, 0x37, 0x04, 0x00]);
300 assert_eq!(image_dimensions(&webp), Some((1920, 1080)));
301 assert_eq!(image_dimensions(b"%PDF-1.7"), None);
302 assert_eq!(image_dimensions(b"\x89PNG\r\n\x1a\n"), None);
303 assert_eq!(image_dimensions(&[0xFF, 0xD8, 0xFF]), None);
304 }
305
306 fn upload(name: &str, bytes: &'static [u8]) -> Upload {
307 Upload::new(name, "application/octet-stream", Bytes::from_static(bytes))
308 }
309
310 #[test]
311 fn sniffs_content_instead_of_trusting_names() {
312 let png = upload("photo.txt", b"\x89PNG\r\n\x1a\nrest");
313 assert!(png.is_image());
314 assert_eq!(png.safe_extension(), "png");
315 let fake = upload("virus.jpg", b"MZ\x90\x00");
316 assert!(!fake.is_image());
317 assert_eq!(fake.safe_extension(), "jpg");
318 assert_eq!(upload("../../etc/passwd", b"x").extension(), None);
319 assert_eq!(upload("no-extension", b"x").safe_extension(), "bin");
320 }
321
322 #[test]
323 fn keys_are_random_and_contained() {
324 let a = upload("a.PNG", b"x").key("/products/");
325 assert!(a.starts_with("products/") && a.ends_with(".png"), "{a}");
326 assert_ne!(a, upload("a.png", b"x").key("products"));
327 }
328
329 fn webp(chunk: &[u8; 4], body: &[u8]) -> Vec<u8> {
333 let mut b = b"RIFF\0\0\0\0WEBP".to_vec();
334 b.extend_from_slice(chunk);
335 b.extend_from_slice(&[0, 0, 0, 0]);
336 b.extend_from_slice(body);
337 b.resize(40, 0);
338 b
339 }
340
341 #[test]
342 fn webp_lossy_and_lossless_headers_and_unknown_chunks() {
343 let mut lossy = vec![0u8; 3];
345 lossy.extend_from_slice(&[0x9D, 0x01, 0x2A]);
346 lossy.extend_from_slice(&640u16.to_le_bytes());
347 lossy.extend_from_slice(&480u16.to_le_bytes());
348 let file = Upload::new("a.webp", "image/webp", webp(b"VP8 ", &lossy));
349 assert_eq!(file.sniffed_type(), Some("image/webp"));
350 assert_eq!(file.dimensions(), Some((640, 480)));
351 let bits: u32 = (99) | (49 << 14);
353 let mut lossless = vec![0x2F];
354 lossless.extend_from_slice(&bits.to_le_bytes());
355 let file = Upload::new("b.webp", "image/webp", webp(b"VP8L", &lossless));
356 assert_eq!(file.dimensions(), Some((100, 50)));
357 let file = Upload::new("c.webp", "image/webp", webp(b"ALPH", &[]));
359 assert!(file.is_image());
360 assert_eq!(file.dimensions(), None);
361 }
362
363 #[test]
364 fn jpeg_headers_skip_fill_bytes_and_standalone_markers() {
365 let mut jpeg = vec![
368 0xFF, 0xD8, 0xFF, 0xFF, 0xD0, 0xFF, 0xE0, 0x00, 0x04, 0x00, 0x00,
369 ];
370 jpeg.extend_from_slice(&[0xFF, 0xC0, 0x00, 0x11, 0x08, 0x00, 0x14, 0x00, 0x1E]);
371 let file = Upload::new("photo.JPG", "image/jpeg", jpeg);
372 assert_eq!(file.sniffed_type(), Some("image/jpeg"));
373 assert_eq!(file.dimensions(), Some((30, 20)));
374 assert!(
375 file.key("photos").ends_with(".jpg"),
376 "{}",
377 file.key("photos")
378 );
379 let cut = Upload::new("cut.jpg", "image/jpeg", vec![0xFF, 0xD8, 0xFF, 0xE0, 0x00]);
381 assert_eq!(cut.dimensions(), None);
382 let odd = Upload::new(
384 "odd.png",
385 "image/png",
386 b"\x89PNG\r\n\x1a\n\0\0\0\x0dJUNK".to_vec(),
387 );
388 assert_eq!(odd.dimensions(), None);
389 }
390
391 #[test]
392 fn accessors_debug_and_old_input() {
393 let file = Upload::new("notes", "text/plain", b"hello".to_vec());
394 assert_eq!(file.content_type(), "text/plain");
395 assert_eq!(file.size(), 5);
396 assert_eq!(file.extension(), None);
397 assert!(file.key("").ends_with(".bin") && !file.key("").contains('/'));
399 let shown = format!("{file:?}");
401 assert!(
402 shown.contains("size: 5") && !shown.contains("hello"),
403 "{shown}"
404 );
405 assert_eq!(serde_json::to_value(&file).unwrap(), "notes");
407 }
408
409 #[tokio::test]
410 async fn store_keeps_the_file_private() {
411 let dir = tempfile::tempdir().unwrap();
412 let config = crate::Config {
413 storage_path: dir.path().to_path_buf(),
414 ..Default::default()
415 };
416 let storage = crate::storage::Storage::from_config(&config).unwrap();
417 let file = Upload::new("a.txt", "text/plain", b"hi".to_vec());
418 let key = file.store(&storage, "/notes/").await.unwrap();
419 assert!(key.starts_with("notes/") && key.ends_with(".txt"), "{key}");
420 assert_eq!(&storage.get(&key).await.unwrap().unwrap()[..], b"hi");
421 }
422
423 #[test]
426 fn sniffed_files_are_stored_under_their_real_extension() {
427 let mut webp = b"RIFF\0\0\0\0WEBPVP8 ".to_vec();
428 webp.extend_from_slice(&[0; 20]);
429 for (name, bytes, ext) in [
430 ("a.png", b"GIF89a\x01\x00\x01\x00".to_vec(), "gif"),
431 ("b.png", webp, "webp"),
432 ("c.txt", b"%PDF-1.7\n".to_vec(), "pdf"),
433 ] {
434 let file = Upload::new(name, "application/octet-stream", bytes);
435 let key = file.key("files");
436 assert!(key.ends_with(&format!(".{ext}")), "{name}: {key}");
437 }
438 let mut jpeg = vec![0xFF, 0xD8, 0x00, 0x13, 0xFF, 0xE0, 0x00, 0x02];
440 jpeg.extend_from_slice(&[0xFF, 0xC0, 0x00, 0x11, 0x08, 0x00, 0x0A, 0x00, 0x0B]);
441 let file = Upload::new("photo.jpg", "image/jpeg", jpeg);
442 assert_eq!(file.dimensions(), Some((11, 10)));
443 }
444}