1use crate::models::Media;
2use crate::services::image as img_service;
3use crate::Database;
4use anyhow::{bail, Result};
5use std::path::Path;
6use uuid::Uuid;
7
8pub const MAX_FILE_SIZE: usize = 50 * 1024 * 1024;
9
10pub const ALLOWED_MIME_TYPES: &[&str] = &[
11 "image/jpeg",
12 "image/png",
13 "image/gif",
14 "image/webp",
15 "application/pdf",
16 "video/mp4",
17 "video/webm",
18 "audio/mpeg",
19 "audio/ogg",
20];
21
22fn detect_mime_type(data: &[u8], claimed_mime: &str) -> Option<String> {
23 if let Some(kind) = infer::get(data) {
24 return Some(kind.mime_type().to_string());
25 }
26
27 if claimed_mime == "image/svg+xml" && data.len() > 5 {
28 let start = String::from_utf8_lossy(&data[..data.len().min(1000)]);
29 if start.contains("<svg") || start.contains("<?xml") {
30 return Some("image/svg+xml".to_string());
31 }
32 }
33
34 None
35}
36
37fn sanitize_svg(data: &[u8]) -> Result<Vec<u8>> {
38 let content = String::from_utf8_lossy(data);
39
40 let dangerous_patterns = [
41 "<script",
42 "javascript:",
43 "onload=",
44 "onerror=",
45 "onclick=",
46 "onmouseover=",
47 "onfocus=",
48 "onblur=",
49 "onchange=",
50 "onsubmit=",
51 "eval(",
52 "expression(",
53 "url(data:",
54 "xlink:href=\"javascript",
55 "xlink:href='javascript",
56 ];
57
58 let lower_content = content.to_lowercase();
59 for pattern in dangerous_patterns {
60 if lower_content.contains(pattern) {
61 bail!("SVG contains potentially dangerous content: {}", pattern);
62 }
63 }
64
65 Ok(data.to_vec())
66}
67
68fn get_safe_extension(detected_mime: &str) -> Option<&'static str> {
69 match detected_mime {
70 "image/jpeg" => Some("jpg"),
71 "image/png" => Some("png"),
72 "image/gif" => Some("gif"),
73 "image/webp" => Some("webp"),
74 "image/svg+xml" => Some("svg"),
75 "application/pdf" => Some("pdf"),
76 "video/mp4" => Some("mp4"),
77 "video/webm" => Some("webm"),
78 "audio/mpeg" => Some("mp3"),
79 "audio/ogg" => Some("ogg"),
80 _ => None,
81 }
82}
83
84pub fn upload_media(
85 db: &Database,
86 upload_dir: &Path,
87 original_name: &str,
88 mime_type: &str,
89 data: &[u8],
90 uploaded_by: Option<i64>,
91) -> Result<Media> {
92 if data.len() > MAX_FILE_SIZE {
93 bail!(
94 "File too large: {} bytes (max {} bytes)",
95 data.len(),
96 MAX_FILE_SIZE
97 );
98 }
99
100 let detected_mime = detect_mime_type(data, mime_type);
101 let actual_mime = detected_mime.as_deref().unwrap_or(mime_type);
102
103 let is_svg = actual_mime == "image/svg+xml";
104 if !ALLOWED_MIME_TYPES.contains(&actual_mime) && !is_svg {
105 bail!(
106 "File type not allowed: {}. Allowed types: {}",
107 actual_mime,
108 ALLOWED_MIME_TYPES.join(", ")
109 );
110 }
111
112 let final_data = if is_svg {
113 sanitize_svg(data)?
114 } else {
115 data.to_vec()
116 };
117
118 std::fs::create_dir_all(upload_dir)?;
119
120 let base_uuid = Uuid::new_v4();
121
122 let (filename, webp_filename, width, height, stored_data) =
123 if img_service::is_optimizable_image(actual_mime) {
124 match img_service::optimize_image(&final_data, actual_mime, None) {
125 Ok(optimized) => {
126 let ext = match optimized.original_format {
127 image::ImageFormat::Jpeg => "jpg",
128 image::ImageFormat::Png => "png",
129 image::ImageFormat::Gif => "gif",
130 image::ImageFormat::WebP => "webp",
131 _ => "bin",
132 };
133
134 let filename = format!("{}.{}", base_uuid, ext);
135 let webp_name = format!("{}.webp", base_uuid);
136
137 std::fs::write(upload_dir.join(&filename), &optimized.original)?;
138 std::fs::write(upload_dir.join(&webp_name), &optimized.webp)?;
139
140 if let Ok(thumb_data) =
141 img_service::generate_thumbnail(&optimized.original, None)
142 {
143 let thumb_name = format!("{}-thumb.webp", base_uuid);
144 std::fs::write(upload_dir.join(&thumb_name), thumb_data)?;
145 }
146
147 if let Ok(variants) =
149 img_service::generate_srcset_variants(&optimized.original)
150 {
151 for variant in variants {
152 let variant_name =
153 format!("{}{}.webp", base_uuid, variant.suffix);
154 let _ = std::fs::write(
155 upload_dir.join(&variant_name),
156 &variant.data,
157 );
158 }
159 }
160
161 (
162 filename,
163 Some(webp_name),
164 Some(optimized.width),
165 Some(optimized.height),
166 optimized.original,
167 )
168 }
169 Err(_) => {
170 let extension = get_safe_extension(actual_mime).unwrap_or("bin");
171 let filename = format!("{}.{}", base_uuid, extension);
172 std::fs::write(upload_dir.join(&filename), &final_data)?;
173 (filename, None, None, None, final_data.clone())
174 }
175 }
176 } else {
177 let extension = get_safe_extension(actual_mime).unwrap_or("bin");
178 let filename = format!("{}.{}", base_uuid, extension);
179 std::fs::write(upload_dir.join(&filename), &final_data)?;
180 (filename, None, None, None, final_data.clone())
181 };
182
183 let conn = db.get()?;
184 conn.execute(
185 "INSERT INTO media (filename, original_name, mime_type, size_bytes, uploaded_by, webp_filename, width, height) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
186 (&filename, original_name, actual_mime, stored_data.len() as i64, uploaded_by, &webp_filename, width, height),
187 )?;
188
189 let id = conn.last_insert_rowid();
190 let created_at: String =
191 conn.query_row("SELECT created_at FROM media WHERE id = ?", [id], |row| {
192 row.get(0)
193 })?;
194
195 Ok(Media {
196 id,
197 filename,
198 original_name: original_name.to_string(),
199 mime_type: actual_mime.to_string(),
200 size_bytes: stored_data.len() as i64,
201 alt_text: String::new(),
202 uploaded_by,
203 created_at,
204 })
205}
206
207pub fn count_media(db: &Database) -> Result<i64> {
208 let conn = db.get()?;
209 let count: i64 = conn.query_row("SELECT COUNT(*) FROM media", [], |row| row.get(0))?;
210 Ok(count)
211}
212
213pub fn list_media(db: &Database, limit: usize, offset: usize) -> Result<Vec<Media>> {
214 let conn = db.get()?;
215 let mut stmt = conn.prepare(
216 "SELECT id, filename, original_name, mime_type, size_bytes, alt_text, uploaded_by, created_at FROM media ORDER BY created_at DESC LIMIT ? OFFSET ?",
217 )?;
218 let media = stmt
219 .query_map((limit, offset), |row| {
220 Ok(Media {
221 id: row.get(0)?,
222 filename: row.get(1)?,
223 original_name: row.get(2)?,
224 mime_type: row.get(3)?,
225 size_bytes: row.get(4)?,
226 alt_text: row.get(5)?,
227 uploaded_by: row.get(6)?,
228 created_at: row.get(7)?,
229 })
230 })?
231 .collect::<Result<Vec<_>, _>>()?;
232 Ok(media)
233}
234
235pub fn get_media_by_filename(db: &Database, filename: &str) -> Result<Option<Media>> {
236 let conn = db.get()?;
237 let media = conn
238 .query_row(
239 "SELECT id, filename, original_name, mime_type, size_bytes, alt_text, uploaded_by, created_at FROM media WHERE filename = ?",
240 [filename],
241 |row| {
242 Ok(Media {
243 id: row.get(0)?,
244 filename: row.get(1)?,
245 original_name: row.get(2)?,
246 mime_type: row.get(3)?,
247 size_bytes: row.get(4)?,
248 alt_text: row.get(5)?,
249 uploaded_by: row.get(6)?,
250 created_at: row.get(7)?,
251 })
252 },
253 )
254 .ok();
255 Ok(media)
256}
257
258pub fn delete_media(db: &Database, upload_dir: &Path, id: i64) -> Result<()> {
259 let conn = db.get()?;
260
261 let (filename, webp_filename): (String, Option<String>) = conn.query_row(
262 "SELECT filename, webp_filename FROM media WHERE id = ?",
263 [id],
264 |row| Ok((row.get(0)?, row.get(1)?)),
265 )?;
266
267 let file_path = upload_dir.join(&filename);
268 if file_path.exists() {
269 std::fs::remove_file(file_path)?;
270 }
271
272 if let Some(webp) = webp_filename {
273 let webp_path = upload_dir.join(&webp);
274 if webp_path.exists() {
275 std::fs::remove_file(webp_path)?;
276 }
277 }
278
279 let base_name = filename
280 .rsplit_once('.')
281 .map(|(n, _)| n)
282 .unwrap_or(&filename);
283 let thumb_path = upload_dir.join(format!("{}-thumb.webp", base_name));
284 if thumb_path.exists() {
285 std::fs::remove_file(thumb_path)?;
286 }
287
288 conn.execute("DELETE FROM media WHERE id = ?", [id])?;
289 Ok(())
290}
291
292pub fn update_media_alt(db: &Database, id: i64, alt_text: &str) -> Result<()> {
293 let conn = db.get()?;
294 conn.execute("UPDATE media SET alt_text = ? WHERE id = ?", (alt_text, id))?;
295 Ok(())
296}