use crate::error::{Result, ZeptoError};
use sha2::{Digest, Sha256};
use std::path::PathBuf;
use tokio::fs;
pub const MAX_IMAGE_SIZE: usize = 20 * 1024 * 1024;
pub const SUPPORTED_TYPES: &[&str] = &["image/jpeg", "image/png", "image/gif", "image/webp"];
pub struct MediaStore {
base_dir: PathBuf,
}
impl MediaStore {
pub fn new(base_dir: PathBuf) -> Self {
Self { base_dir }
}
pub async fn save(&self, data: &[u8], mime_type: &str) -> Result<String> {
let ext = mime_to_ext(mime_type);
let hash = sha256_prefix(data);
let filename = format!("{}.{}", hash, ext);
let rel_path = format!("media/{}", filename);
let abs_path = self.base_dir.join(&rel_path);
let media_dir = self.base_dir.join("media");
fs::create_dir_all(&media_dir).await?;
if abs_path.exists() {
return Ok(rel_path);
}
fs::write(&abs_path, data).await?;
Ok(rel_path)
}
pub async fn load(&self, rel_path: &str) -> Result<Vec<u8>> {
let abs_path = self.base_dir.join(rel_path);
let bytes = fs::read(&abs_path).await?;
Ok(bytes)
}
}
pub fn mime_to_ext(mime_type: &str) -> &'static str {
match mime_type {
"image/jpeg" => "jpg",
"image/png" => "png",
"image/gif" => "gif",
"image/webp" => "webp",
_ => "bin",
}
}
pub fn validate_image(data: &[u8], mime_type: &str, max_size: usize) -> Result<()> {
if data.len() > max_size {
return Err(ZeptoError::Tool(format!(
"Image size {} bytes exceeds the maximum allowed size of {} bytes",
data.len(),
max_size
)));
}
if !SUPPORTED_TYPES.contains(&mime_type) {
return Err(ZeptoError::Tool(format!(
"Unsupported image type '{}'. Supported types: {}",
mime_type,
SUPPORTED_TYPES.join(", ")
)));
}
Ok(())
}
fn sha256_prefix(data: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(data);
let digest = hasher.finalize();
hex::encode(digest)[..16].to_string()
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn test_save_media_creates_file_with_hash_name() {
let tmp = TempDir::new().unwrap();
let store = MediaStore::new(tmp.path().to_path_buf());
let data = b"fake jpeg data";
let path = store.save(data, "image/jpeg").await.unwrap();
assert!(path.starts_with("media/"));
assert!(path.ends_with(".jpg"));
assert!(tmp.path().join(&path).exists());
}
#[tokio::test]
async fn test_save_media_deduplicates() {
let tmp = TempDir::new().unwrap();
let store = MediaStore::new(tmp.path().to_path_buf());
let data = b"same data";
let path1 = store.save(data, "image/png").await.unwrap();
let path2 = store.save(data, "image/png").await.unwrap();
assert_eq!(path1, path2);
}
#[tokio::test]
async fn test_load_media_returns_bytes() {
let tmp = TempDir::new().unwrap();
let store = MediaStore::new(tmp.path().to_path_buf());
let data = b"image bytes here";
let path = store.save(data, "image/jpeg").await.unwrap();
let loaded = store.load(&path).await.unwrap();
assert_eq!(loaded, data);
}
#[test]
fn test_mime_to_extension() {
assert_eq!(mime_to_ext("image/jpeg"), "jpg");
assert_eq!(mime_to_ext("image/png"), "png");
assert_eq!(mime_to_ext("image/gif"), "gif");
assert_eq!(mime_to_ext("image/webp"), "webp");
assert_eq!(mime_to_ext("image/unknown"), "bin");
}
#[test]
fn test_validate_image_valid_types() {
assert!(validate_image(b"data", "image/jpeg", MAX_IMAGE_SIZE).is_ok());
assert!(validate_image(b"data", "image/png", MAX_IMAGE_SIZE).is_ok());
assert!(validate_image(b"data", "image/gif", MAX_IMAGE_SIZE).is_ok());
assert!(validate_image(b"data", "image/webp", MAX_IMAGE_SIZE).is_ok());
}
#[test]
fn test_validate_image_rejects_unsupported() {
assert!(validate_image(b"data", "image/tiff", MAX_IMAGE_SIZE).is_err());
assert!(validate_image(b"data", "application/pdf", MAX_IMAGE_SIZE).is_err());
}
#[test]
fn test_validate_image_rejects_oversized() {
let big = vec![0u8; 21 * 1024 * 1024];
assert!(validate_image(&big, "image/jpeg", MAX_IMAGE_SIZE).is_err());
}
#[tokio::test]
async fn test_save_then_load_round_trip_matches() {
let tmp = TempDir::new().unwrap();
let store = MediaStore::new(tmp.path().to_path_buf());
let mut data = vec![0xFF, 0xD8, 0xFF, 0xE0];
data.extend_from_slice(&[0u8; 1000]);
let path = store.save(&data, "image/jpeg").await.unwrap();
let loaded = store.load(&path).await.unwrap();
assert_eq!(data, loaded, "Loaded bytes must match saved bytes exactly");
assert!(path.starts_with("media/"));
assert!(path.ends_with(".jpg"));
}
}