Skip to main content

pebble_cms/services/
media.rs

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                    // Generate responsive srcset variants (400w, 800w, 1200w, 1600w)
148                    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}