1use percent_encoding::percent_decode_str;
7use sha2::{Digest, Sha256};
8use std::fs;
9use std::io::{self, Write};
10use std::path::{Path, PathBuf};
11
12#[derive(Debug, thiserror::Error)]
14pub enum FileUploadError {
15 #[error("File too large: {0} bytes (max: {1} bytes)")]
16 FileTooLarge(usize, usize),
17 #[error("Invalid file type: {0}")]
18 InvalidFileType(String),
19 #[error("IO error: {0}")]
20 Io(#[from] io::Error),
21 #[error("Upload error: {0}")]
22 Upload(String),
23 #[error("Checksum verification failed")]
24 ChecksumMismatch,
25 #[error("MIME type detection failed")]
26 MimeDetectionFailed,
27 #[error("Path traversal detected in filename")]
28 PathTraversal,
29}
30
31pub fn validate_safe_filename(filename: &str) -> Result<(), FileUploadError> {
37 if filename.is_empty() {
38 return Err(FileUploadError::Upload("Empty filename".to_string()));
39 }
40
41 let decoded = percent_decode_str(filename).decode_utf8_lossy();
44 for candidate in [filename, decoded.as_ref()] {
45 if candidate.contains('\0') {
46 return Err(FileUploadError::PathTraversal);
47 }
48 if candidate.contains("..") {
49 return Err(FileUploadError::PathTraversal);
50 }
51 if candidate.contains('/') || candidate.contains('\\') {
52 return Err(FileUploadError::PathTraversal);
53 }
54 if candidate.starts_with('/') || candidate.starts_with('\\') {
56 return Err(FileUploadError::PathTraversal);
57 }
58 if candidate.len() >= 2
59 && candidate.as_bytes()[0].is_ascii_alphabetic()
60 && candidate.as_bytes()[1] == b':'
61 {
62 return Err(FileUploadError::PathTraversal);
63 }
64 }
65 Ok(())
66}
67
68pub struct FileUploadHandler {
73 upload_dir: PathBuf,
74 max_size: usize,
75 allowed_extensions: Option<Vec<String>>,
76 verify_checksum: bool,
77 allowed_mime_types: Option<Vec<String>>,
78}
79
80impl FileUploadHandler {
81 pub fn new(upload_dir: PathBuf) -> Self {
97 Self {
98 upload_dir,
99 max_size: 10 * 1024 * 1024, allowed_extensions: None,
101 verify_checksum: false,
102 allowed_mime_types: None,
103 }
104 }
105
106 pub fn with_max_size(mut self, max_size: usize) -> Self {
119 self.max_size = max_size;
120 self
121 }
122
123 pub fn with_allowed_extensions(mut self, extensions: Vec<String>) -> Self {
135 self.allowed_extensions = Some(extensions);
136 self
137 }
138
139 pub fn with_checksum_verification(mut self, enabled: bool) -> Self {
151 self.verify_checksum = enabled;
152 self
153 }
154
155 pub fn with_allowed_mime_types(mut self, mime_types: Vec<String>) -> Self {
170 self.allowed_mime_types = Some(mime_types);
171 self
172 }
173
174 pub fn max_size(&self) -> usize {
176 self.max_size
177 }
178
179 pub fn upload_dir(&self) -> &Path {
181 &self.upload_dir
182 }
183
184 pub fn calculate_checksum(&self, content: &[u8]) -> String {
197 let mut hasher = Sha256::new();
198 hasher.update(content);
199 let result = hasher.finalize();
200 result.iter().map(|b| format!("{:02x}", b)).collect()
202 }
203
204 pub fn verify_file_checksum(
218 &self,
219 content: &[u8],
220 expected_checksum: &str,
221 ) -> Result<(), FileUploadError> {
222 let actual_checksum = self.calculate_checksum(content);
223 if actual_checksum == expected_checksum {
224 Ok(())
225 } else {
226 Err(FileUploadError::ChecksumMismatch)
227 }
228 }
229
230 pub fn detect_mime_type(&self, content: &[u8]) -> Option<String> {
251 if content.is_empty() {
252 return None;
253 }
254
255 if content.len() >= 8 && content[0..8] == [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A] {
257 return Some("image/png".to_string());
258 }
259
260 if content.len() >= 3 && content[0..3] == [0xFF, 0xD8, 0xFF] {
261 return Some("image/jpeg".to_string());
262 }
263
264 if content.len() >= 4 && content[0..4] == [0x47, 0x49, 0x46, 0x38] {
265 return Some("image/gif".to_string());
266 }
267
268 if content.len() >= 4 && content[0..4] == [0x25, 0x50, 0x44, 0x46] {
269 return Some("application/pdf".to_string());
270 }
271
272 if content.len() >= 4
273 && (content[0..4] == [0x50, 0x4B, 0x03, 0x04]
274 || content[0..4] == [0x50, 0x4B, 0x05, 0x06])
275 {
276 return Some("application/zip".to_string());
277 }
278
279 None
280 }
281
282 fn validate_mime_type(&self, content: &[u8]) -> Result<(), FileUploadError> {
284 if let Some(ref allowed) = self.allowed_mime_types {
285 let detected_mime = self
286 .detect_mime_type(content)
287 .ok_or(FileUploadError::MimeDetectionFailed)?;
288
289 if !allowed.contains(&detected_mime) {
290 return Err(FileUploadError::InvalidFileType(detected_mime));
291 }
292 }
293 Ok(())
294 }
295
296 pub fn handle_upload(
319 &self,
320 field_name: &str,
321 filename: &str,
322 content: &[u8],
323 ) -> Result<String, FileUploadError> {
324 validate_safe_filename(field_name)?;
326
327 if content.len() > self.max_size {
329 return Err(FileUploadError::FileTooLarge(content.len(), self.max_size));
330 }
331
332 if let Some(ref allowed) = self.allowed_extensions {
334 let extension = Path::new(filename)
335 .extension()
336 .and_then(|e| e.to_str())
337 .unwrap_or("");
338
339 if !allowed.iter().any(|ext| ext == extension) {
340 return Err(FileUploadError::InvalidFileType(extension.to_string()));
341 }
342 }
343
344 self.validate_mime_type(content)?;
346
347 fs::create_dir_all(&self.upload_dir)?;
349
350 let unique_filename = self.generate_unique_filename(field_name, filename);
352 let file_path = self.upload_dir.join(&unique_filename);
353
354 let mut file = fs::File::create(&file_path)?;
356 file.write_all(content)?;
357
358 Ok(unique_filename)
359 }
360
361 pub fn handle_upload_with_checksum(
383 &self,
384 field_name: &str,
385 filename: &str,
386 content: &[u8],
387 expected_checksum: &str,
388 ) -> Result<String, FileUploadError> {
389 if self.verify_checksum {
391 self.verify_file_checksum(content, expected_checksum)?;
392 }
393
394 self.handle_upload(field_name, filename, content)
396 }
397
398 fn generate_unique_filename(&self, field_name: &str, original_filename: &str) -> String {
405 let unique_id = uuid::Uuid::new_v4();
406
407 let basename = Path::new(original_filename)
409 .file_name()
410 .and_then(|n| n.to_str())
411 .unwrap_or(original_filename);
412
413 let extension = Path::new(basename)
414 .extension()
415 .and_then(|e| e.to_str())
416 .unwrap_or("");
417
418 if extension.is_empty() {
419 format!("{}_{}", field_name, unique_id)
420 } else {
421 format!("{}_{}.{}", field_name, unique_id, extension)
422 }
423 }
424
425 pub fn delete_upload(&self, filename: &str) -> Result<(), FileUploadError> {
438 validate_safe_filename(filename)?;
440 let file_path = self.upload_dir.join(filename);
441 fs::remove_file(file_path)?;
442 Ok(())
443 }
444}
445
446pub struct TemporaryFileUpload {
450 path: PathBuf,
451 auto_delete: bool,
452}
453
454impl TemporaryFileUpload {
455 pub fn new(path: PathBuf) -> Self {
467 Self {
468 path,
469 auto_delete: true,
470 }
471 }
472
473 pub fn with_content(path: PathBuf, content: &[u8]) -> Result<Self, FileUploadError> {
487 let mut file = fs::File::create(&path)?;
488 file.write_all(content)?;
489 Ok(Self {
490 path,
491 auto_delete: true,
492 })
493 }
494
495 pub fn keep(&mut self) {
508 self.auto_delete = false;
509 }
510
511 pub fn path(&self) -> &Path {
513 &self.path
514 }
515
516 pub fn auto_delete(&self) -> bool {
518 self.auto_delete
519 }
520
521 pub fn read_content(&self) -> Result<Vec<u8>, FileUploadError> {
537 Ok(fs::read(&self.path)?)
538 }
539}
540
541impl Drop for TemporaryFileUpload {
542 fn drop(&mut self) {
543 if self.auto_delete && self.path.exists() {
544 let _ = fs::remove_file(&self.path);
545 }
546 }
547}
548
549pub struct MemoryFileUpload {
554 filename: String,
555 content: Vec<u8>,
556 content_type: Option<String>,
557}
558
559impl MemoryFileUpload {
560 pub fn new(filename: String, content: Vec<u8>) -> Self {
575 Self {
576 filename,
577 content,
578 content_type: None,
579 }
580 }
581
582 pub fn with_content_type(filename: String, content: Vec<u8>, content_type: String) -> Self {
597 Self {
598 filename,
599 content,
600 content_type: Some(content_type),
601 }
602 }
603
604 pub fn filename(&self) -> &str {
606 &self.filename
607 }
608
609 pub fn content(&self) -> &[u8] {
611 &self.content
612 }
613
614 pub fn content_type(&self) -> Option<&str> {
616 self.content_type.as_deref()
617 }
618
619 pub fn size(&self) -> usize {
630 self.content.len()
631 }
632
633 pub fn is_empty(&self) -> bool {
647 self.content.is_empty()
648 }
649
650 pub fn save_to_disk(&self, path: PathBuf) -> Result<(), FileUploadError> {
663 let mut file = fs::File::create(path)?;
664 file.write_all(&self.content)?;
665 Ok(())
666 }
667}
668
669#[cfg(test)]
670mod tests {
671 use super::*;
672
673 #[test]
674 fn test_file_upload_handler_creation() {
675 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"));
676 assert_eq!(handler.max_size(), 10 * 1024 * 1024);
677 assert_eq!(handler.upload_dir(), Path::new("/tmp/uploads"));
678 }
679
680 #[test]
681 fn test_file_upload_handler_with_max_size() {
682 let handler =
683 FileUploadHandler::new(PathBuf::from("/tmp/uploads")).with_max_size(5 * 1024 * 1024);
684 assert_eq!(handler.max_size(), 5 * 1024 * 1024);
685 }
686
687 #[test]
688 fn test_file_upload_handler_size_validation() {
689 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads")).with_max_size(100);
690
691 let large_content = vec![0u8; 200];
692 let result = handler.handle_upload("test", "large.txt", &large_content);
693
694 assert!(result.is_err());
695 if let Err(FileUploadError::FileTooLarge(size, max)) = result {
696 assert_eq!(size, 200);
697 assert_eq!(max, 100);
698 } else {
699 panic!("Expected FileTooLarge error");
700 }
701 }
702
703 #[test]
704 fn test_file_upload_handler_extension_validation() {
705 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"))
706 .with_allowed_extensions(vec!["jpg".to_string(), "png".to_string()]);
707
708 let content = b"test content";
709 let result = handler.handle_upload("test", "document.pdf", content);
710
711 assert!(result.is_err());
712 if let Err(FileUploadError::InvalidFileType(ext)) = result {
713 assert_eq!(ext, "pdf");
714 } else {
715 panic!("Expected InvalidFileType error");
716 }
717 }
718
719 #[test]
720 fn test_temporary_file_upload_creation() {
721 let temp = TemporaryFileUpload::new(PathBuf::from("/tmp/test_temp.txt"));
722 assert_eq!(temp.path(), Path::new("/tmp/test_temp.txt"));
723 assert!(temp.auto_delete());
724 }
725
726 #[test]
727 fn test_temporary_file_upload_keep() {
728 let mut temp = TemporaryFileUpload::new(PathBuf::from("/tmp/test_keep.txt"));
729 temp.keep();
730 assert!(!temp.auto_delete());
731 }
732
733 #[test]
734 fn test_temporary_file_upload_with_content() {
735 let temp_path = PathBuf::from("/tmp/test_content_temp.txt");
736 let content = b"Test content";
737
738 let temp = TemporaryFileUpload::with_content(temp_path.clone(), content).unwrap();
739 assert!(temp_path.exists());
740
741 let read_content = temp.read_content().unwrap();
742 assert_eq!(read_content, content);
743
744 drop(temp);
745 assert!(!temp_path.exists());
746 }
747
748 #[test]
749 fn test_temporary_file_upload_auto_delete() {
750 let temp_path = PathBuf::from("/tmp/test_auto_delete.txt");
751 fs::write(&temp_path, b"test").unwrap();
752
753 {
754 let _temp = TemporaryFileUpload::new(temp_path.clone());
755 assert!(temp_path.exists());
756 }
757
758 assert!(!temp_path.exists());
759 }
760
761 #[test]
762 fn test_memory_file_upload_creation() {
763 let upload = MemoryFileUpload::new("test.txt".to_string(), vec![1, 2, 3, 4, 5]);
764
765 assert_eq!(upload.filename(), "test.txt");
766 assert_eq!(upload.content(), &[1, 2, 3, 4, 5]);
767 assert_eq!(upload.size(), 5);
768 assert!(!upload.is_empty());
769 }
770
771 #[test]
772 fn test_memory_file_upload_with_content_type() {
773 let upload = MemoryFileUpload::with_content_type(
774 "image.png".to_string(),
775 vec![0x89, 0x50, 0x4E, 0x47],
776 "image/png".to_string(),
777 );
778
779 assert_eq!(upload.filename(), "image.png");
780 assert_eq!(upload.content_type(), Some("image/png"));
781 }
782
783 #[test]
784 fn test_memory_file_upload_is_empty() {
785 let empty = MemoryFileUpload::new("empty.txt".to_string(), vec![]);
786 assert!(empty.is_empty());
787 assert_eq!(empty.size(), 0);
788
789 let non_empty = MemoryFileUpload::new("data.txt".to_string(), vec![1, 2, 3]);
790 assert!(!non_empty.is_empty());
791 assert_eq!(non_empty.size(), 3);
792 }
793
794 #[test]
795 fn test_memory_file_upload_save_to_disk() {
796 let temp_path = PathBuf::from("/tmp/test_memory_save.txt");
797 let upload = MemoryFileUpload::new("test.txt".to_string(), vec![1, 2, 3, 4, 5]);
798
799 let result = upload.save_to_disk(temp_path.clone());
800 assert!(result.is_ok());
801 assert!(temp_path.exists());
802
803 let content = fs::read(&temp_path).unwrap();
804 assert_eq!(content, vec![1, 2, 3, 4, 5]);
805
806 fs::remove_file(temp_path).unwrap();
807 }
808
809 #[rstest::rstest]
814 #[case("../../../etc/passwd")]
815 #[case("foo/../../bar")]
816 #[case("/etc/passwd")]
817 #[case("test\0file.txt")]
818 #[case("..%2f..%2fetc%2fpasswd")]
819 #[case("%2e%2e/%2e%2e/etc/passwd")]
820 fn test_delete_upload_rejects_path_traversal(#[case] filename: &str) {
821 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"));
823
824 let result = handler.delete_upload(filename);
826
827 assert!(
829 matches!(result, Err(FileUploadError::PathTraversal)),
830 "Expected PathTraversal error for filename: {}",
831 filename
832 );
833 }
834
835 #[rstest::rstest]
836 fn test_delete_upload_allows_safe_filenames() {
837 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"));
839
840 let result = handler.delete_upload("safe_file.txt");
843
844 assert!(
846 !matches!(result, Err(FileUploadError::PathTraversal)),
847 "Safe filename should not trigger path traversal error"
848 );
849 }
850
851 #[rstest::rstest]
852 #[case("normal.txt", true)]
853 #[case("my-file_123.jpg", true)]
854 #[case("report.pdf", true)]
855 #[case("image_2024.png", true)]
856 #[case("../../../etc/passwd", false)]
857 #[case("foo/../bar.txt", false)]
858 #[case("/absolute/path.txt", false)]
859 #[case("null\0byte.txt", false)]
860 #[case("", false)]
861 #[case("back\\slash.txt", false)]
862 #[case("C:\\Windows\\system32", false)]
863 #[case("..%2f..%2fetc%2fpasswd", false)]
864 #[case("%2e%2e%2f%2e%2e%2f", false)]
865 fn test_validate_safe_filename(#[case] filename: &str, #[case] should_pass: bool) {
866 let result = validate_safe_filename(filename);
868
869 assert_eq!(
871 result.is_ok(),
872 should_pass,
873 "validate_safe_filename({:?}) expected {} but got {}",
874 filename,
875 if should_pass { "Ok" } else { "Err" },
876 if result.is_ok() { "Ok" } else { "Err" },
877 );
878 }
879
880 #[rstest::rstest]
881 #[case("../malicious")]
882 #[case("foo/../../bar")]
883 #[case("..%2fmalicious")]
884 fn test_handle_upload_rejects_traversal_in_field_name(#[case] field_name: &str) {
885 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"));
887
888 let result = handler.handle_upload(field_name, "safe.txt", b"content");
890
891 assert!(
893 matches!(result, Err(FileUploadError::PathTraversal)),
894 "Expected PathTraversal error for field_name: {}",
895 field_name
896 );
897 }
898
899 #[rstest::rstest]
900 fn test_handle_upload_accepts_safe_field_name() {
901 let handler = FileUploadHandler::new(PathBuf::from("/tmp/reinhardt_upload_test"));
903
904 let result = handler.handle_upload("avatar", "photo.jpg", b"image data");
906
907 assert!(result.is_ok());
909
910 let _ = fs::remove_dir_all("/tmp/reinhardt_upload_test");
912 }
913}