systemprompt_files/services/upload/

validator.rs

use crate::config::FileUploadConfig;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum FileValidationError {
    #[error("File uploads are disabled")]
    UploadsDisabled,

    #[error("File size {size} bytes exceeds maximum allowed {max} bytes")]
    FileTooLarge { size: u64, max: u64 },

    #[error("File type '{mime_type}' is not allowed")]
    TypeNotAllowed { mime_type: String },

    #[error("File type '{mime_type}' is blocked for security reasons")]
    TypeBlocked { mime_type: String },

    #[error("File category '{category}' is disabled in configuration")]
    CategoryDisabled { category: String },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FileCategory {
    Image,
    Document,
    Audio,
    Video,
}

impl FileCategory {
    pub const fn storage_subdir(&self) -> &'static str {
        match self {
            Self::Image => "images",
            Self::Document => "documents",
            Self::Audio => "audio",
            Self::Video => "video",
        }
    }

    pub const fn display_name(&self) -> &'static str {
        match self {
            Self::Image => "image",
            Self::Document => "document",
            Self::Audio => "audio",
            Self::Video => "video",
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct FileValidator {
    config: FileUploadConfig,
}

impl FileValidator {
    const IMAGE_TYPES: &'static [&'static str] = &[
        "image/jpeg",
        "image/png",
        "image/gif",
        "image/webp",
        "image/svg+xml",
        "image/bmp",
        "image/tiff",
        "image/x-icon",
        "image/vnd.microsoft.icon",
    ];

    const DOCUMENT_TYPES: &'static [&'static 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",
        "text/plain",
        "text/csv",
        "text/markdown",
        "text/html",
        "application/json",
        "application/xml",
        "text/xml",
        "application/rtf",
    ];

    const AUDIO_TYPES: &'static [&'static str] = &[
        "audio/mpeg",
        "audio/mp3",
        "audio/wav",
        "audio/wave",
        "audio/x-wav",
        "audio/ogg",
        "audio/webm",
        "audio/aac",
        "audio/flac",
        "audio/mp4",
        "audio/x-m4a",
    ];

    const VIDEO_TYPES: &'static [&'static str] = &[
        "video/mp4",
        "video/webm",
        "video/ogg",
        "video/quicktime",
        "video/x-msvideo",
        "video/x-matroska",
    ];

    const BLOCKED_TYPES: &'static [&'static str] = &[
        "application/x-executable",
        "application/x-msdos-program",
        "application/x-msdownload",
        "application/x-sh",
        "application/x-shellscript",
        "application/x-csh",
        "application/x-bash",
        "application/bat",
        "application/x-bat",
        "application/x-msi",
        "application/vnd.microsoft.portable-executable",
        "application/x-dosexec",
        "application/x-python-code",
        "application/javascript",
        "text/javascript",
        "application/x-httpd-php",
        "application/x-php",
        "text/x-php",
        "application/x-perl",
        "text/x-perl",
        "application/x-ruby",
        "text/x-ruby",
        "application/java-archive",
        "application/x-java-class",
    ];

    pub const fn new(config: FileUploadConfig) -> Self {
        Self { config }
    }

    pub fn validate(
        &self,
        mime_type: &str,
        size_bytes: u64,
    ) -> Result<FileCategory, FileValidationError> {
        if !self.config.enabled {
            return Err(FileValidationError::UploadsDisabled);
        }

        if size_bytes > self.config.max_file_size_bytes {
            return Err(FileValidationError::FileTooLarge {
                size: size_bytes,
                max: self.config.max_file_size_bytes,
            });
        }

        let normalized_mime = mime_type.to_lowercase();

        if Self::BLOCKED_TYPES.contains(&normalized_mime.as_str()) {
            return Err(FileValidationError::TypeBlocked {
                mime_type: mime_type.to_string(),
            });
        }

        let category = Self::categorize_mime_type(&normalized_mime)?;

        if !self.is_category_allowed(&category) {
            return Err(FileValidationError::CategoryDisabled {
                category: category.display_name().to_string(),
            });
        }

        Ok(category)
    }

    fn categorize_mime_type(mime_type: &str) -> Result<FileCategory, FileValidationError> {
        if Self::IMAGE_TYPES.contains(&mime_type) {
            return Ok(FileCategory::Image);
        }

        if Self::DOCUMENT_TYPES.contains(&mime_type) {
            return Ok(FileCategory::Document);
        }

        if Self::AUDIO_TYPES.contains(&mime_type) {
            return Ok(FileCategory::Audio);
        }

        if Self::VIDEO_TYPES.contains(&mime_type) {
            return Ok(FileCategory::Video);
        }

        Err(FileValidationError::TypeNotAllowed {
            mime_type: mime_type.to_string(),
        })
    }

    const fn is_category_allowed(&self, category: &FileCategory) -> bool {
        match category {
            FileCategory::Image => self.config.allowed_types.images,
            FileCategory::Document => self.config.allowed_types.documents,
            FileCategory::Audio => self.config.allowed_types.audio,
            FileCategory::Video => self.config.allowed_types.video,
        }
    }

    pub fn get_extension(mime_type: &str, filename: Option<&str>) -> String {
        if let Some(name) = filename {
            if let Some(ext) = name.rsplit('.').next() {
                if !ext.is_empty()
                    && ext.len() <= 10
                    && ext != name
                    && ext.chars().all(|c| c.is_ascii_alphanumeric())
                {
                    return ext.to_lowercase();
                }
            }
        }

        let lower = mime_type.to_lowercase();
        MIME_EXTENSION_TABLE
            .iter()
            .find(|(mimes, _)| mimes.contains(&lower.as_str()))
            .map_or("bin", |(_, ext)| *ext)
            .to_string()
    }
}

const MIME_EXTENSION_TABLE: &[(&[&str], &str)] = &[
    (&["image/jpeg"], "jpg"),
    (&["image/png"], "png"),
    (&["image/gif"], "gif"),
    (&["image/webp"], "webp"),
    (&["image/svg+xml"], "svg"),
    (&["image/bmp"], "bmp"),
    (&["image/tiff"], "tiff"),
    (&["image/x-icon", "image/vnd.microsoft.icon"], "ico"),
    (&["application/pdf"], "pdf"),
    (&["text/plain"], "txt"),
    (&["text/csv"], "csv"),
    (&["text/markdown"], "md"),
    (&["text/html"], "html"),
    (&["application/json"], "json"),
    (&["application/xml", "text/xml"], "xml"),
    (&["application/rtf"], "rtf"),
    (&["application/msword"], "doc"),
    (
        &["application/vnd.openxmlformats-officedocument.wordprocessingml.document"],
        "docx",
    ),
    (&["application/vnd.ms-excel"], "xls"),
    (
        &["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"],
        "xlsx",
    ),
    (&["application/vnd.ms-powerpoint"], "ppt"),
    (
        &["application/vnd.openxmlformats-officedocument.presentationml.presentation"],
        "pptx",
    ),
    (&["audio/mpeg", "audio/mp3"], "mp3"),
    (&["audio/wav", "audio/wave", "audio/x-wav"], "wav"),
    (&["audio/ogg"], "ogg"),
    (&["audio/webm"], "weba"),
    (&["audio/aac"], "aac"),
    (&["audio/flac"], "flac"),
    (&["audio/mp4", "audio/x-m4a"], "m4a"),
    (&["video/mp4"], "mp4"),
    (&["video/webm"], "webm"),
    (&["video/ogg"], "ogv"),
    (&["video/quicktime"], "mov"),
    (&["video/x-msvideo"], "avi"),
    (&["video/x-matroska"], "mkv"),
];