use super::types::{StorageError, StorageResult, UploadedFile};
#[derive(Debug, Clone, Default)]
pub struct MimeValidator {
strict: bool,
}
impl MimeValidator {
#[must_use]
pub const fn new() -> Self {
Self { strict: true }
}
#[must_use]
pub const fn permissive() -> Self {
Self { strict: false }
}
#[must_use]
pub fn detect_mime(&self, file: &UploadedFile) -> Option<&'static str> {
infer::get(&file.data).map(|kind| kind.mime_type())
}
pub fn validate_against_magic(
&self,
file: &UploadedFile,
allowed_types: &[&str],
) -> StorageResult<()> {
match self.detect_mime(file) {
Some(detected_type) => {
if !allowed_types.contains(&detected_type) {
return Err(StorageError::InvalidMimeType {
expected: allowed_types.iter().map(|s| (*s).to_string()).collect(),
actual: detected_type.to_string(),
});
}
Ok(())
}
None => {
if self.strict {
Err(StorageError::InvalidMimeType {
expected: allowed_types.iter().map(|s| (*s).to_string()).collect(),
actual: "unknown (could not detect from content)".to_string(),
})
} else {
file.validate_mime(allowed_types)
}
}
}
}
pub fn validate_header_matches_content(&self, file: &UploadedFile) -> StorageResult<()> {
match self.detect_mime(file) {
Some(detected_type) => {
if detected_type != file.content_type {
return Err(StorageError::InvalidMimeType {
expected: vec![file.content_type.clone()],
actual: detected_type.to_string(),
});
}
Ok(())
}
None => {
if self.strict {
Err(StorageError::InvalidMimeType {
expected: vec![file.content_type.clone()],
actual: "unknown (could not detect from content)".to_string(),
})
} else {
Ok(())
}
}
}
}
#[must_use]
pub fn is_image(&self, file: &UploadedFile) -> bool {
self.detect_mime(file)
.is_some_and(|mime| mime.starts_with("image/"))
}
#[must_use]
pub fn is_video(&self, file: &UploadedFile) -> bool {
self.detect_mime(file)
.is_some_and(|mime| mime.starts_with("video/"))
}
#[must_use]
pub fn is_document(&self, file: &UploadedFile) -> bool {
const DOCUMENT_TYPES: &[&str] = &[
"application/pdf",
"application/msword",
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
"application/vnd.ms-excel",
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
"application/vnd.ms-powerpoint",
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
];
self.detect_mime(file)
.is_some_and(|mime| DOCUMENT_TYPES.contains(&mime))
}
}
#[cfg(test)]
mod tests {
use super::*;
const JPEG_MAGIC: &[u8] = &[0xFF, 0xD8, 0xFF];
const PNG_MAGIC: &[u8] = &[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
const GIF_MAGIC: &[u8] = b"GIF89a";
const PDF_MAGIC: &[u8] = b"%PDF-1.4";
const ZIP_MAGIC: &[u8] = &[0x50, 0x4B, 0x03, 0x04];
#[test]
fn test_detect_jpeg() {
let file = UploadedFile::new("test.jpg", "image/jpeg", JPEG_MAGIC.to_vec());
let validator = MimeValidator::new();
assert_eq!(validator.detect_mime(&file), Some("image/jpeg"));
}
#[test]
fn test_detect_png() {
let file = UploadedFile::new("test.png", "image/png", PNG_MAGIC.to_vec());
let validator = MimeValidator::new();
assert_eq!(validator.detect_mime(&file), Some("image/png"));
}
#[test]
fn test_detect_gif() {
let file = UploadedFile::new("test.gif", "image/gif", GIF_MAGIC.to_vec());
let validator = MimeValidator::new();
assert_eq!(validator.detect_mime(&file), Some("image/gif"));
}
#[test]
fn test_detect_pdf() {
let file = UploadedFile::new("test.pdf", "application/pdf", PDF_MAGIC.to_vec());
let validator = MimeValidator::new();
assert_eq!(validator.detect_mime(&file), Some("application/pdf"));
}
#[test]
fn test_detect_unknown() {
let file = UploadedFile::new("test.txt", "text/plain", b"hello".to_vec());
let validator = MimeValidator::new();
assert_eq!(validator.detect_mime(&file), None);
}
#[test]
fn test_validate_against_magic_success() {
let file = UploadedFile::new("photo.jpg", "image/jpeg", JPEG_MAGIC.to_vec());
let validator = MimeValidator::new();
assert!(validator
.validate_against_magic(&file, &["image/jpeg", "image/png"])
.is_ok());
}
#[test]
fn test_validate_against_magic_failure() {
let file = UploadedFile::new("photo.jpg", "image/jpeg", JPEG_MAGIC.to_vec());
let validator = MimeValidator::new();
let result = validator.validate_against_magic(&file, &["image/png", "image/gif"]);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
StorageError::InvalidMimeType { .. }
));
}
#[test]
fn test_validate_against_magic_strict_unknown() {
let file = UploadedFile::new("test.txt", "text/plain", b"hello".to_vec());
let validator = MimeValidator::new(); let result = validator.validate_against_magic(&file, &["text/plain"]);
assert!(result.is_err()); }
#[test]
fn test_validate_against_magic_permissive_unknown() {
let file = UploadedFile::new("test.txt", "text/plain", b"hello".to_vec());
let validator = MimeValidator::permissive();
let result = validator.validate_against_magic(&file, &["text/plain"]);
assert!(result.is_ok()); }
#[test]
fn test_header_matches_content_honest() {
let file = UploadedFile::new("photo.png", "image/png", PNG_MAGIC.to_vec());
let validator = MimeValidator::new();
assert!(validator.validate_header_matches_content(&file).is_ok());
}
#[test]
fn test_header_matches_content_dishonest() {
let file = UploadedFile::new("fake.jpg", "image/jpeg", PNG_MAGIC.to_vec());
let validator = MimeValidator::new();
let result = validator.validate_header_matches_content(&file);
assert!(result.is_err());
}
#[test]
fn test_header_matches_content_malicious() {
let file = UploadedFile::new(
"malware.jpg",
"image/jpeg",
b"#!/bin/sh\nrm -rf /".to_vec(),
);
let validator = MimeValidator::new();
let result = validator.validate_header_matches_content(&file);
assert!(result.is_err());
}
#[test]
fn test_is_image() {
let validator = MimeValidator::new();
let jpeg = UploadedFile::new("photo.jpg", "image/jpeg", JPEG_MAGIC.to_vec());
assert!(validator.is_image(&jpeg));
let png = UploadedFile::new("photo.png", "image/png", PNG_MAGIC.to_vec());
assert!(validator.is_image(&png));
let pdf = UploadedFile::new("doc.pdf", "application/pdf", PDF_MAGIC.to_vec());
assert!(!validator.is_image(&pdf));
}
#[test]
fn test_is_document() {
let validator = MimeValidator::new();
let pdf = UploadedFile::new("doc.pdf", "application/pdf", PDF_MAGIC.to_vec());
assert!(validator.is_document(&pdf));
let jpeg = UploadedFile::new("photo.jpg", "image/jpeg", JPEG_MAGIC.to_vec());
assert!(!validator.is_document(&jpeg));
}
#[test]
fn test_forged_extension() {
let file = UploadedFile::new(
"malware.jpg",
"image/jpeg",
ZIP_MAGIC.to_vec(), );
let validator = MimeValidator::new();
assert!(validator
.validate_against_magic(&file, &["image/jpeg"])
.is_err());
}
}