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#[non_exhaustive]
14#[derive(Debug, thiserror::Error)]
15pub enum FileUploadError {
16 #[error("File too large: {0} bytes (max: {1} bytes)")]
18 FileTooLarge(usize, usize),
19 #[error("Invalid file type: {0}")]
21 InvalidFileType(String),
22 #[error("IO error: {0}")]
24 Io(#[from] io::Error),
25 #[error("Upload error: {0}")]
27 Upload(String),
28 #[error("Checksum verification failed")]
30 ChecksumMismatch,
31 #[error("MIME type detection failed")]
33 MimeDetectionFailed,
34 #[error("Path traversal detected in filename")]
36 PathTraversal,
37}
38
39pub fn validate_safe_filename(filename: &str) -> Result<(), FileUploadError> {
45 if filename.is_empty() {
46 return Err(FileUploadError::Upload("Empty filename".to_string()));
47 }
48
49 let decoded = percent_decode_str(filename).decode_utf8_lossy();
52 for candidate in [filename, decoded.as_ref()] {
53 if candidate.contains('\0') {
54 return Err(FileUploadError::PathTraversal);
55 }
56 if candidate.contains("..") {
57 return Err(FileUploadError::PathTraversal);
58 }
59 if candidate.contains('/') || candidate.contains('\\') {
60 return Err(FileUploadError::PathTraversal);
61 }
62 if candidate.starts_with('/') || candidate.starts_with('\\') {
64 return Err(FileUploadError::PathTraversal);
65 }
66 if candidate.len() >= 2
67 && candidate.as_bytes()[0].is_ascii_alphabetic()
68 && candidate.as_bytes()[1] == b':'
69 {
70 return Err(FileUploadError::PathTraversal);
71 }
72 }
73 Ok(())
74}
75
76pub struct FileUploadHandler {
81 upload_dir: PathBuf,
82 max_size: usize,
83 allowed_extensions: Option<Vec<String>>,
84 verify_checksum: bool,
85 allowed_mime_types: Option<Vec<String>>,
86}
87
88impl FileUploadHandler {
89 pub fn new(upload_dir: PathBuf) -> Self {
105 Self {
106 upload_dir,
107 max_size: 10 * 1024 * 1024, allowed_extensions: None,
109 verify_checksum: false,
110 allowed_mime_types: None,
111 }
112 }
113
114 pub fn with_max_size(mut self, max_size: usize) -> Self {
127 self.max_size = max_size;
128 self
129 }
130
131 pub fn with_allowed_extensions(mut self, extensions: Vec<String>) -> Self {
143 self.allowed_extensions = Some(extensions);
144 self
145 }
146
147 pub fn with_checksum_verification(mut self, enabled: bool) -> Self {
159 self.verify_checksum = enabled;
160 self
161 }
162
163 pub fn with_allowed_mime_types(mut self, mime_types: Vec<String>) -> Self {
178 self.allowed_mime_types = Some(mime_types);
179 self
180 }
181
182 pub fn max_size(&self) -> usize {
184 self.max_size
185 }
186
187 pub fn upload_dir(&self) -> &Path {
189 &self.upload_dir
190 }
191
192 pub fn calculate_checksum(&self, content: &[u8]) -> String {
205 let mut hasher = Sha256::new();
206 hasher.update(content);
207 let result = hasher.finalize();
208 result.iter().map(|b| format!("{:02x}", b)).collect()
210 }
211
212 pub fn verify_file_checksum(
226 &self,
227 content: &[u8],
228 expected_checksum: &str,
229 ) -> Result<(), FileUploadError> {
230 let actual_checksum = self.calculate_checksum(content);
231 if actual_checksum == expected_checksum {
232 Ok(())
233 } else {
234 Err(FileUploadError::ChecksumMismatch)
235 }
236 }
237
238 pub fn detect_mime_type(&self, content: &[u8]) -> Option<String> {
259 if content.is_empty() {
260 return None;
261 }
262
263 if content.len() >= 8 && content[0..8] == [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A] {
265 return Some("image/png".to_string());
266 }
267
268 if content.len() >= 3 && content[0..3] == [0xFF, 0xD8, 0xFF] {
269 return Some("image/jpeg".to_string());
270 }
271
272 if content.len() >= 4 && content[0..4] == [0x47, 0x49, 0x46, 0x38] {
273 return Some("image/gif".to_string());
274 }
275
276 if content.len() >= 4 && content[0..4] == [0x25, 0x50, 0x44, 0x46] {
277 return Some("application/pdf".to_string());
278 }
279
280 if content.len() >= 4
281 && (content[0..4] == [0x50, 0x4B, 0x03, 0x04]
282 || content[0..4] == [0x50, 0x4B, 0x05, 0x06])
283 {
284 return Some("application/zip".to_string());
285 }
286
287 None
288 }
289
290 fn validate_mime_type(&self, content: &[u8]) -> Result<(), FileUploadError> {
292 if let Some(ref allowed) = self.allowed_mime_types {
293 let detected_mime = self
294 .detect_mime_type(content)
295 .ok_or(FileUploadError::MimeDetectionFailed)?;
296
297 if !allowed.contains(&detected_mime) {
298 return Err(FileUploadError::InvalidFileType(detected_mime));
299 }
300 }
301 Ok(())
302 }
303
304 pub fn handle_upload(
327 &self,
328 field_name: &str,
329 filename: &str,
330 content: &[u8],
331 ) -> Result<String, FileUploadError> {
332 validate_safe_filename(field_name)?;
334
335 if content.len() > self.max_size {
337 return Err(FileUploadError::FileTooLarge(content.len(), self.max_size));
338 }
339
340 if let Some(ref allowed) = self.allowed_extensions {
342 let extension = Path::new(filename)
343 .extension()
344 .and_then(|e| e.to_str())
345 .unwrap_or("");
346
347 if !allowed.iter().any(|ext| ext == extension) {
348 return Err(FileUploadError::InvalidFileType(extension.to_string()));
349 }
350 }
351
352 self.validate_mime_type(content)?;
354
355 fs::create_dir_all(&self.upload_dir)?;
357
358 let unique_filename = self.generate_unique_filename(field_name, filename);
360 let file_path = self.upload_dir.join(&unique_filename);
361
362 let mut file = fs::File::create(&file_path)?;
364 file.write_all(content)?;
365
366 Ok(unique_filename)
367 }
368
369 pub fn handle_upload_with_checksum(
391 &self,
392 field_name: &str,
393 filename: &str,
394 content: &[u8],
395 expected_checksum: &str,
396 ) -> Result<String, FileUploadError> {
397 if self.verify_checksum {
399 self.verify_file_checksum(content, expected_checksum)?;
400 }
401
402 self.handle_upload(field_name, filename, content)
404 }
405
406 fn generate_unique_filename(&self, field_name: &str, original_filename: &str) -> String {
413 let unique_id = uuid::Uuid::new_v4();
414
415 let basename = Path::new(original_filename)
417 .file_name()
418 .and_then(|n| n.to_str())
419 .unwrap_or(original_filename);
420
421 let extension = Path::new(basename)
422 .extension()
423 .and_then(|e| e.to_str())
424 .unwrap_or("");
425
426 if extension.is_empty() {
427 format!("{}_{}", field_name, unique_id)
428 } else {
429 format!("{}_{}.{}", field_name, unique_id, extension)
430 }
431 }
432
433 pub fn delete_upload(&self, filename: &str) -> Result<(), FileUploadError> {
446 validate_safe_filename(filename)?;
448 let file_path = self.upload_dir.join(filename);
449 fs::remove_file(file_path)?;
450 Ok(())
451 }
452}
453
454pub struct TemporaryFileUpload {
458 path: PathBuf,
459 auto_delete: bool,
460}
461
462impl TemporaryFileUpload {
463 pub fn new(path: PathBuf) -> Self {
475 Self {
476 path,
477 auto_delete: true,
478 }
479 }
480
481 pub fn with_content(path: PathBuf, content: &[u8]) -> Result<Self, FileUploadError> {
495 let mut file = fs::File::create(&path)?;
496 file.write_all(content)?;
497 Ok(Self {
498 path,
499 auto_delete: true,
500 })
501 }
502
503 pub fn keep(&mut self) {
516 self.auto_delete = false;
517 }
518
519 pub fn path(&self) -> &Path {
521 &self.path
522 }
523
524 pub fn auto_delete(&self) -> bool {
526 self.auto_delete
527 }
528
529 pub fn read_content(&self) -> Result<Vec<u8>, FileUploadError> {
545 Ok(fs::read(&self.path)?)
546 }
547}
548
549impl Drop for TemporaryFileUpload {
550 fn drop(&mut self) {
551 if self.auto_delete && self.path.exists() {
552 let _ = fs::remove_file(&self.path);
553 }
554 }
555}
556
557pub struct MemoryFileUpload {
562 filename: String,
563 content: Vec<u8>,
564 content_type: Option<String>,
565}
566
567impl MemoryFileUpload {
568 pub fn new(filename: String, content: Vec<u8>) -> Self {
583 Self {
584 filename,
585 content,
586 content_type: None,
587 }
588 }
589
590 pub fn with_content_type(filename: String, content: Vec<u8>, content_type: String) -> Self {
605 Self {
606 filename,
607 content,
608 content_type: Some(content_type),
609 }
610 }
611
612 pub fn filename(&self) -> &str {
614 &self.filename
615 }
616
617 pub fn content(&self) -> &[u8] {
619 &self.content
620 }
621
622 pub fn content_type(&self) -> Option<&str> {
624 self.content_type.as_deref()
625 }
626
627 pub fn size(&self) -> usize {
638 self.content.len()
639 }
640
641 pub fn is_empty(&self) -> bool {
655 self.content.is_empty()
656 }
657
658 pub fn save_to_disk(&self, path: PathBuf) -> Result<(), FileUploadError> {
671 let mut file = fs::File::create(path)?;
672 file.write_all(&self.content)?;
673 Ok(())
674 }
675}
676
677#[cfg(test)]
678mod tests {
679 use super::*;
680
681 #[test]
682 fn test_file_upload_handler_creation() {
683 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"));
684 assert_eq!(handler.max_size(), 10 * 1024 * 1024);
685 assert_eq!(handler.upload_dir(), Path::new("/tmp/uploads"));
686 }
687
688 #[test]
689 fn test_file_upload_handler_with_max_size() {
690 let handler =
691 FileUploadHandler::new(PathBuf::from("/tmp/uploads")).with_max_size(5 * 1024 * 1024);
692 assert_eq!(handler.max_size(), 5 * 1024 * 1024);
693 }
694
695 #[test]
696 fn test_file_upload_handler_size_validation() {
697 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads")).with_max_size(100);
698
699 let large_content = vec![0u8; 200];
700 let result = handler.handle_upload("test", "large.txt", &large_content);
701
702 assert!(result.is_err());
703 if let Err(FileUploadError::FileTooLarge(size, max)) = result {
704 assert_eq!(size, 200);
705 assert_eq!(max, 100);
706 } else {
707 panic!("Expected FileTooLarge error");
708 }
709 }
710
711 #[test]
712 fn test_file_upload_handler_extension_validation() {
713 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"))
714 .with_allowed_extensions(vec!["jpg".to_string(), "png".to_string()]);
715
716 let content = b"test content";
717 let result = handler.handle_upload("test", "document.pdf", content);
718
719 assert!(result.is_err());
720 if let Err(FileUploadError::InvalidFileType(ext)) = result {
721 assert_eq!(ext, "pdf");
722 } else {
723 panic!("Expected InvalidFileType error");
724 }
725 }
726
727 #[test]
728 fn test_temporary_file_upload_creation() {
729 let temp = TemporaryFileUpload::new(PathBuf::from("/tmp/test_temp.txt"));
730 assert_eq!(temp.path(), Path::new("/tmp/test_temp.txt"));
731 assert!(temp.auto_delete());
732 }
733
734 #[test]
735 fn test_temporary_file_upload_keep() {
736 let mut temp = TemporaryFileUpload::new(PathBuf::from("/tmp/test_keep.txt"));
737 temp.keep();
738 assert!(!temp.auto_delete());
739 }
740
741 #[test]
742 fn test_temporary_file_upload_with_content() {
743 let temp_path = PathBuf::from("/tmp/test_content_temp.txt");
744 let content = b"Test content";
745
746 let temp = TemporaryFileUpload::with_content(temp_path.clone(), content).unwrap();
747 assert!(temp_path.exists());
748
749 let read_content = temp.read_content().unwrap();
750 assert_eq!(read_content, content);
751
752 drop(temp);
753 assert!(!temp_path.exists());
754 }
755
756 #[test]
757 fn test_temporary_file_upload_auto_delete() {
758 let temp_path = PathBuf::from("/tmp/test_auto_delete.txt");
759 fs::write(&temp_path, b"test").unwrap();
760
761 {
762 let _temp = TemporaryFileUpload::new(temp_path.clone());
763 assert!(temp_path.exists());
764 }
765
766 assert!(!temp_path.exists());
767 }
768
769 #[test]
770 fn test_memory_file_upload_creation() {
771 let upload = MemoryFileUpload::new("test.txt".to_string(), vec![1, 2, 3, 4, 5]);
772
773 assert_eq!(upload.filename(), "test.txt");
774 assert_eq!(upload.content(), &[1, 2, 3, 4, 5]);
775 assert_eq!(upload.size(), 5);
776 assert!(!upload.is_empty());
777 }
778
779 #[test]
780 fn test_memory_file_upload_with_content_type() {
781 let upload = MemoryFileUpload::with_content_type(
782 "image.png".to_string(),
783 vec![0x89, 0x50, 0x4E, 0x47],
784 "image/png".to_string(),
785 );
786
787 assert_eq!(upload.filename(), "image.png");
788 assert_eq!(upload.content_type(), Some("image/png"));
789 }
790
791 #[test]
792 fn test_memory_file_upload_is_empty() {
793 let empty = MemoryFileUpload::new("empty.txt".to_string(), vec![]);
794 assert!(empty.is_empty());
795 assert_eq!(empty.size(), 0);
796
797 let non_empty = MemoryFileUpload::new("data.txt".to_string(), vec![1, 2, 3]);
798 assert!(!non_empty.is_empty());
799 assert_eq!(non_empty.size(), 3);
800 }
801
802 #[test]
803 fn test_memory_file_upload_save_to_disk() {
804 let temp_path = PathBuf::from("/tmp/test_memory_save.txt");
805 let upload = MemoryFileUpload::new("test.txt".to_string(), vec![1, 2, 3, 4, 5]);
806
807 let result = upload.save_to_disk(temp_path.clone());
808 assert!(result.is_ok());
809 assert!(temp_path.exists());
810
811 let content = fs::read(&temp_path).unwrap();
812 assert_eq!(content, vec![1, 2, 3, 4, 5]);
813
814 fs::remove_file(temp_path).unwrap();
815 }
816
817 #[rstest::rstest]
822 #[case("../../../etc/passwd")]
823 #[case("foo/../../bar")]
824 #[case("/etc/passwd")]
825 #[case("test\0file.txt")]
826 #[case("..%2f..%2fetc%2fpasswd")]
827 #[case("%2e%2e/%2e%2e/etc/passwd")]
828 fn test_delete_upload_rejects_path_traversal(#[case] filename: &str) {
829 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"));
831
832 let result = handler.delete_upload(filename);
834
835 assert!(
837 matches!(result, Err(FileUploadError::PathTraversal)),
838 "Expected PathTraversal error for filename: {}",
839 filename
840 );
841 }
842
843 #[rstest::rstest]
844 fn test_delete_upload_allows_safe_filenames() {
845 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"));
847
848 let result = handler.delete_upload("safe_file.txt");
851
852 assert!(
854 !matches!(result, Err(FileUploadError::PathTraversal)),
855 "Safe filename should not trigger path traversal error"
856 );
857 }
858
859 #[rstest::rstest]
860 #[case("normal.txt", true)]
861 #[case("my-file_123.jpg", true)]
862 #[case("report.pdf", true)]
863 #[case("image_2024.png", true)]
864 #[case("../../../etc/passwd", false)]
865 #[case("foo/../bar.txt", false)]
866 #[case("/absolute/path.txt", false)]
867 #[case("null\0byte.txt", false)]
868 #[case("", false)]
869 #[case("back\\slash.txt", false)]
870 #[case("C:\\Windows\\system32", false)]
871 #[case("..%2f..%2fetc%2fpasswd", false)]
872 #[case("%2e%2e%2f%2e%2e%2f", false)]
873 fn test_validate_safe_filename(#[case] filename: &str, #[case] should_pass: bool) {
874 let result = validate_safe_filename(filename);
876
877 assert_eq!(
879 result.is_ok(),
880 should_pass,
881 "validate_safe_filename({:?}) expected {} but got {}",
882 filename,
883 if should_pass { "Ok" } else { "Err" },
884 if result.is_ok() { "Ok" } else { "Err" },
885 );
886 }
887
888 #[rstest::rstest]
889 #[case("../malicious")]
890 #[case("foo/../../bar")]
891 #[case("..%2fmalicious")]
892 fn test_handle_upload_rejects_traversal_in_field_name(#[case] field_name: &str) {
893 let handler = FileUploadHandler::new(PathBuf::from("/tmp/uploads"));
895
896 let result = handler.handle_upload(field_name, "safe.txt", b"content");
898
899 assert!(
901 matches!(result, Err(FileUploadError::PathTraversal)),
902 "Expected PathTraversal error for field_name: {}",
903 field_name
904 );
905 }
906
907 #[rstest::rstest]
908 fn test_handle_upload_accepts_safe_field_name() {
909 let handler = FileUploadHandler::new(PathBuf::from("/tmp/reinhardt_upload_test"));
911
912 let result = handler.handle_upload("avatar", "photo.jpg", b"image data");
914
915 assert!(result.is_ok());
917
918 let _ = fs::remove_dir_all("/tmp/reinhardt_upload_test");
920 }
921}