146 lines
4.1 KiB
Rust
146 lines
4.1 KiB
Rust
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<Vec<ContentBlock>, 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<Vec<ContentBlock>, 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<String, Box<dyn MediaHandler>>,
|
|
}
|
|
|
|
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<dyn MediaHandler>) {
|
|
self.handlers
|
|
.insert(handler.media_type().to_string(), handler);
|
|
}
|
|
|
|
pub fn handle(
|
|
&self,
|
|
media_type: &str,
|
|
path: &str,
|
|
) -> Result<Vec<ContentBlock>, 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
|
|
}
|
|
}
|