grabcli_sqlite/
images.rs

1//! Content-addressed image store: `<cache dir>/images/ab/cd/<sha256>.<ext>`
2//! (files 0600, dirs 0700), indexed in the `images` table by URL. Several URLs
3//! with identical bytes share one file. When the total size of distinct files
4//! exceeds `[cache] image_max_mb` (default 500), the least recently used files
5//! are evicted.
6//!
7//! Images are fetched (by `grabcli_core::cache::CacheHandle::fetch_image`)
8//! without cookies or any session headers, and only from Grab's CDN hosts. Originals are stored as-is: Grab
9//! mostly serves AVIF/WebP, which the pure-Rust `image` crate can't decode
10//! (AVIF) or only encode losslessly (WebP), so there is no thumbnailing;
11//! width/height are recorded for PNG/JPEG only.
12
13use std::fs;
14use std::path::{Path, PathBuf};
15
16use grabcli_core::Error;
17use grabcli_core::cache::{CachedImage, hex_sha256};
18use rusqlite::{Connection, OptionalExtension, params};
19
20use crate::db::create_private_dirs;
21use crate::err::{DbError, Result};
22use crate::{ImageStats, SqliteCache};
23
24fn ext_for(content_type: &str) -> &'static str {
25    match content_type
26        .split(';')
27        .next()
28        .unwrap_or("")
29        .trim()
30        .to_ascii_lowercase()
31        .as_str()
32    {
33        "image/jpeg" | "image/jpg" => "jpg",
34        "image/png" => "png",
35        "image/webp" => "webp",
36        "image/avif" => "avif",
37        "image/gif" => "gif",
38        _ => "img",
39    }
40}
41
42fn file_path(root: &Path, hash: &str, content_type: &str) -> PathBuf {
43    let a = hash.get(0..2).unwrap_or("00");
44    let b = hash.get(2..4).unwrap_or("00");
45    root.join("images")
46        .join(a)
47        .join(b)
48        .join(format!("{hash}.{}", ext_for(content_type)))
49}
50
51fn dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
52    image::ImageReader::new(std::io::Cursor::new(bytes))
53        .with_guessed_format()
54        .ok()?
55        .into_dimensions()
56        .ok()
57}
58
59fn write_private_file(path: &Path, data: &[u8]) -> Result<()> {
60    use std::io::Write;
61    if let Some(dir) = path.parent() {
62        create_private_dirs(dir)?;
63    }
64    let tmp = path.with_extension(format!("tmp.{}", std::process::id()));
65    {
66        let mut opts = fs::OpenOptions::new();
67        opts.write(true).create(true).truncate(true);
68        #[cfg(unix)]
69        {
70            use std::os::unix::fs::OpenOptionsExt;
71            opts.mode(0o600);
72        }
73        let mut f = opts.open(&tmp)?;
74        f.write_all(data)?;
75    }
76    fs::rename(&tmp, path)?;
77    Ok(())
78}
79
80type ImageRow = (String, i64, Option<i64>, Option<i64>, String);
81
82fn row_for(conn: &Connection, url: &str) -> Result<Option<ImageRow>> {
83    Ok(conn
84        .query_row(
85            "SELECT hash, bytes, width, height, content_type FROM images WHERE url = ?1",
86            [url],
87            |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?, r.get(4)?)),
88        )
89        .optional()?)
90}
91
92fn to_cached(root: &Path, url: &str, row: ImageRow, cached: bool) -> CachedImage {
93    let (hash, bytes, w, h, ct) = row;
94    CachedImage {
95        url: url.to_string(),
96        path: Some(file_path(root, &hash, &ct)),
97        hash,
98        bytes: u64::try_from(bytes).unwrap_or(0),
99        width: w.and_then(|v| u32::try_from(v).ok()),
100        height: h.and_then(|v| u32::try_from(v).ok()),
101        content_type: ct,
102        cached,
103    }
104}
105
106pub(crate) fn stats(conn: &Connection, max_bytes: u64) -> Result<ImageStats> {
107    let (files, bytes): (i64, i64) = conn.query_row(
108        "SELECT COUNT(*), COALESCE(SUM(bytes), 0) FROM
109            (SELECT hash, MAX(bytes) AS bytes FROM images GROUP BY hash)",
110        [],
111        |r| Ok((r.get(0)?, r.get(1)?)),
112    )?;
113    let urls: i64 = conn.query_row("SELECT COUNT(*) FROM images", [], |r| r.get(0))?;
114    Ok(ImageStats {
115        files: u64::try_from(files).unwrap_or(0),
116        urls: u64::try_from(urls).unwrap_or(0),
117        bytes: u64::try_from(bytes).unwrap_or(0),
118        max_bytes,
119    })
120}
121
122/// Evict least-recently-used files until the store is within `max_bytes`.
123/// Returns `(files evicted, bytes freed)`.
124pub(crate) fn evict(conn: &Connection, root: &Path, max_bytes: u64) -> Result<(u64, u64)> {
125    let mut stmt = conn.prepare(
126        "SELECT hash, MAX(bytes), MAX(last_access), MIN(content_type)
127         FROM images GROUP BY hash ORDER BY MAX(last_access) ASC, hash ASC",
128    )?;
129    let files: Vec<(String, i64, i64, String)> = stmt
130        .query_map([], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)))?
131        .collect::<rusqlite::Result<_>>()?;
132    drop(stmt);
133    let mut total: u64 = files.iter().map(|f| u64::try_from(f.1).unwrap_or(0)).sum();
134    let (mut n, mut freed) = (0u64, 0u64);
135    for (hash, bytes, _, ct) in files {
136        if total <= max_bytes {
137            break;
138        }
139        let bytes = u64::try_from(bytes).unwrap_or(0);
140        let _ = fs::remove_file(file_path(root, &hash, &ct));
141        conn.execute("DELETE FROM images WHERE hash = ?1", [&hash])?;
142        conn.execute(
143            "UPDATE menu_items SET image_hash = NULL WHERE image_hash = ?1",
144            [&hash],
145        )?;
146        total = total.saturating_sub(bytes);
147        n += 1;
148        freed += bytes;
149    }
150    Ok((n, freed))
151}
152
153/// Remove every image row and file. Returns the number of URLs removed.
154pub(crate) fn clear(conn: &Connection, root: &Path) -> Result<u64> {
155    let n = conn.execute("DELETE FROM images", [])? as u64;
156    conn.execute("UPDATE menu_items SET image_hash = NULL", [])?;
157    let dir = root.join("images");
158    if dir.exists() {
159        fs::remove_dir_all(dir)?;
160    }
161    Ok(n)
162}
163
164impl SqliteCache {
165    /// The cached file for `url`, if present (and still on disk). Marks it
166    /// recently used.
167    pub fn cached_image_at(&self, url: &str, now: i64) -> Option<CachedImage> {
168        let conn = self.conn().ok()?;
169        let row = row_for(&conn, url).ok()??;
170        let img = to_cached(&self.root, url, row, true);
171        if !img.path.as_ref().is_some_and(|p| p.exists()) {
172            let _ = conn.execute("DELETE FROM images WHERE url = ?1", [url]);
173            return None;
174        }
175        let _ = conn.execute(
176            "UPDATE images SET last_access = ?1 WHERE hash = ?2",
177            params![now, img.hash],
178        );
179        Some(img)
180    }
181
182    /// Store downloaded image bytes for `url` (dedupes by content hash), then
183    /// evict over the size cap.
184    pub fn store_image_at(
185        &self,
186        url: &str,
187        bytes: &[u8],
188        content_type: &str,
189        now: i64,
190    ) -> Result<CachedImage> {
191        let hash = hex_sha256(bytes);
192        let path = file_path(&self.root, &hash, content_type);
193        if !path.exists() {
194            write_private_file(&path, bytes)?;
195        }
196        let dims = dimensions(bytes);
197        let conn = self.conn()?;
198        conn.execute(
199            "INSERT OR REPLACE INTO images
200                (url, hash, bytes, width, height, content_type, last_access, fetched_at)
201             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7)",
202            params![
203                url,
204                hash,
205                i64::try_from(bytes.len()).unwrap_or(i64::MAX),
206                dims.map(|d| d.0),
207                dims.map(|d| d.1),
208                content_type,
209                now
210            ],
211        )?;
212        // Touch every URL sharing this file so LRU sees the file as fresh.
213        conn.execute(
214            "UPDATE images SET last_access = ?1 WHERE hash = ?2",
215            params![now, hash],
216        )?;
217        conn.execute(
218            "UPDATE menu_items SET image_hash = ?1 WHERE image_url = ?2",
219            params![hash, url],
220        )?;
221        evict(&conn, &self.root, self.config.image_max_bytes())?;
222        let row = row_for(&conn, url)?
223            .ok_or_else(|| DbError(Error::invalid("image is larger than the image cache cap")))?;
224        Ok(to_cached(&self.root, url, row, false))
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use grabcli_core::cache::{CacheConfig, image_url_allowed};
232
233    fn exists(i: &CachedImage) -> bool {
234        i.path.as_ref().is_some_and(|p| p.exists())
235    }
236
237    #[test]
238    fn url_allowlist() {
239        assert!(image_url_allowed(
240            "https://food-cms.grab.com/compressed_webp/x.webp"
241        ));
242        assert!(image_url_allowed(
243            "https://huawei-food-cms.grab.com/a/b.avif"
244        ));
245        assert!(!image_url_allowed("http://food-cms.grab.com/x.png"));
246        assert!(!image_url_allowed("https://evil.com/x.png"));
247        assert!(!image_url_allowed("https://grab.com.evil.com/x.png"));
248        assert!(!image_url_allowed("https://user:pw@food-cms.grab.com/x"));
249        assert!(!image_url_allowed("https://food-cms.grab.com:8443/x"));
250        assert!(!image_url_allowed("not a url"));
251    }
252
253    fn png(w: u32, h: u32) -> Result<Vec<u8>> {
254        let img = image::RgbImage::new(w, h);
255        let mut out = std::io::Cursor::new(Vec::new());
256        img.write_to(&mut out, image::ImageFormat::Png)
257            .map_err(|e| DbError(Error::invalid(e.to_string())))?;
258        Ok(out.into_inner())
259    }
260
261    #[test]
262    fn store_dedupes_and_records_dimensions() -> Result<()> {
263        let dir = tempfile::tempdir()?;
264        let c = SqliteCache::new(dir.path(), CacheConfig::default());
265        let bytes = png(3, 2)?;
266        let a = c.store_image_at("https://food-cms.grab.com/a.png", &bytes, "image/png", 10)?;
267        let b = c.store_image_at("https://food-cms.grab.com/b.png", &bytes, "image/png", 11)?;
268        assert_eq!(a.path, b.path);
269        assert_eq!((a.width, a.height), (Some(3), Some(2)));
270        assert!(exists(&a));
271        assert!(
272            a.path
273                .clone()
274                .unwrap_or_default()
275                .to_string_lossy()
276                .contains(&format!(
277                    "images/{}/{}/{}.png",
278                    &a.hash[0..2],
279                    &a.hash[2..4],
280                    a.hash
281                ))
282        );
283        let s = c.stats_at(11)?;
284        assert_eq!((s.images.files, s.images.urls), (1, 2));
285        assert_eq!(s.images.bytes, bytes.len() as u64);
286        let hit = c
287            .cached_image_at("https://food-cms.grab.com/a.png", 12)
288            .ok_or_else(|| DbError(Error::not_found("a")))?;
289        assert!(hit.cached);
290        #[cfg(unix)]
291        {
292            use std::os::unix::fs::PermissionsExt;
293            assert_eq!(
294                fs::metadata(a.path.clone().unwrap_or_default())?
295                    .permissions()
296                    .mode()
297                    & 0o777,
298                0o600
299            );
300        }
301        Ok(())
302    }
303
304    #[test]
305    fn lru_eviction_by_total_size() -> Result<()> {
306        let dir = tempfile::tempdir()?;
307        // 1 MiB cap; each "image" is 400 KiB of distinct bytes.
308        let cfg = CacheConfig {
309            image_max_mb: Some(1),
310            ..CacheConfig::default()
311        };
312        let c = SqliteCache::new(dir.path(), cfg);
313        let blob = |n: u8| vec![n; 400 * 1024];
314        let url = |n: u8| format!("https://food-cms.grab.com/{n}.webp");
315        let first = c.store_image_at(&url(1), &blob(1), "image/webp", 100)?;
316        c.store_image_at(&url(2), &blob(2), "image/webp", 200)?;
317        // Touch #1 so #2 becomes the least recently used.
318        c.cached_image_at(&url(1), 300)
319            .ok_or_else(|| DbError(Error::not_found("1")))?;
320        c.store_image_at(&url(3), &blob(3), "image/webp", 400)?;
321        assert!(c.cached_image_at(&url(1), 401).is_some());
322        assert!(c.cached_image_at(&url(2), 401).is_none());
323        assert!(c.cached_image_at(&url(3), 401).is_some());
324        assert!(exists(&first));
325        let s = c.stats_at(400)?;
326        assert_eq!(s.images.files, 2);
327        assert!(s.images.bytes <= 1024 * 1024);
328
329        // A cap below one image: prune evicts down to it.
330        let tiny = SqliteCache::new(
331            dir.path(),
332            CacheConfig {
333                image_max_mb: Some(0),
334                ..CacheConfig::default()
335            },
336        );
337        let r = tiny.prune_at(500)?;
338        assert_eq!(r.images_evicted, 2);
339        assert!(!exists(&first));
340        Ok(())
341    }
342}
grabcli_sqlite/images.rs

grabcli_sqlite/
images.rs
