use std::collections::HashMap; use std::fmt; use std::io::Read; use crate::bus::message::ContentBlock; const MAX_NATIVE_IMAGE_BYTES: u64 = 10 * 1024 * 1024; pub trait MediaHandler: Send + Sync { fn media_type(&self) -> &str; fn handle(&self, path: &str) -> Result, MediaHandlerError>; } #[derive(Debug)] pub enum MediaHandlerError { Io(std::io::Error), UnsupportedFormat(String), } impl fmt::Display for MediaHandlerError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { MediaHandlerError::Io(e) => write!(f, "I/O error: {}", e), MediaHandlerError::UnsupportedFormat(msg) => write!(f, "Unsupported format: {}", msg), } } } impl std::error::Error for MediaHandlerError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { match self { MediaHandlerError::Io(e) => Some(e), MediaHandlerError::UnsupportedFormat(_) => None, } } } pub struct ImageHandler; impl MediaHandler for ImageHandler { fn media_type(&self) -> &str { "image" } fn handle(&self, path: &str) -> Result, MediaHandlerError> { let (mime_type, base64_data) = encode_image_to_base64(path)?; let url = format!("data:{};base64,{}", mime_type, base64_data); Ok(vec![ContentBlock::image_url(url)]) } } fn encode_image_to_base64(path: &str) -> Result<(String, String), MediaHandlerError> { use base64::{Engine as _, engine::general_purpose::STANDARD}; let metadata = std::fs::metadata(path).map_err(MediaHandlerError::Io)?; if metadata.len() > MAX_NATIVE_IMAGE_BYTES { return Err(MediaHandlerError::UnsupportedFormat(format!( "image is too large: {} bytes (max {} bytes)", metadata.len(), MAX_NATIVE_IMAGE_BYTES ))); } let mime = mime_guess::from_path(path) .first_or_octet_stream() .to_string(); if !matches!( mime.as_str(), "image/png" | "image/jpeg" | "image/gif" | "image/webp" ) { return Err(MediaHandlerError::UnsupportedFormat(format!( "unsupported image MIME type: {mime}" ))); } let mut file = std::fs::File::open(path).map_err(MediaHandlerError::Io)?; let mut buffer = Vec::new(); file.read_to_end(&mut buffer) .map_err(MediaHandlerError::Io)?; if !valid_image_signature(&buffer, &mime) { return Err(MediaHandlerError::UnsupportedFormat(format!( "file content does not match image MIME type: {mime}" ))); } let encoded = STANDARD.encode(&buffer); Ok((mime, encoded)) } fn valid_image_signature(bytes: &[u8], mime: &str) -> bool { match mime { "image/png" => bytes.starts_with(b"\x89PNG\r\n\x1a\n"), "image/jpeg" => bytes.starts_with(&[0xff, 0xd8, 0xff]), "image/gif" => bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a"), "image/webp" => bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP", _ => false, } } pub struct MediaHandlerRegistry { handlers: HashMap>, } impl Default for MediaHandlerRegistry { fn default() -> Self { Self::new() } } impl MediaHandlerRegistry { pub fn new() -> Self { Self { handlers: HashMap::new(), } } pub fn register(&mut self, handler: Box) { self.handlers .insert(handler.media_type().to_string(), handler); } pub fn handle( &self, media_type: &str, path: &str, ) -> Result, MediaHandlerError> { match self.handlers.get(media_type) { Some(handler) => handler.handle(path), None => Err(MediaHandlerError::UnsupportedFormat(format!( "no handler for type: {}", media_type ))), } } pub fn supports(&self, media_type: &str) -> bool { self.handlers.contains_key(media_type) } pub fn with_defaults() -> Self { let mut reg = Self::new(); reg.register(Box::new(ImageHandler)); reg } }