Skip to main content

grabcli_core/cache/
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 without cookies or any session headers, and only from
8//! Grab's CDN hosts ([`image_url_allowed`]). 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};
15use std::time::Duration;
16
17use rusqlite::{Connection, OptionalExtension, params};
18use serde::Serialize;
19
20use super::db::create_private_dirs;
21use super::{Cache, ImageStats, hex_sha256};
22use crate::error::{Error, ErrorKind, Result};
23use crate::jwt;
24
25/// Largest image we will download.
26const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024;
27
28/// A cached image file.
29#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
30pub struct CachedImage {
31    pub url: String,
32    pub path: PathBuf,
33    pub hash: String,
34    pub bytes: u64,
35    pub width: Option<u32>,
36    pub height: Option<u32>,
37    pub content_type: String,
38    /// True if served from the store without downloading.
39    pub cached: bool,
40}
41
42/// Only `https://` URLs on `grab.com` or a subdomain (the CDNs:
43/// `food-cms.grab.com`, `huawei-food-cms.grab.com`, `assets-mobile.grab.com`,
44/// …), with no credentials or custom port.
45pub fn image_url_allowed(url: &str) -> bool {
46    let Ok(u) = reqwest::Url::parse(url) else {
47        return false;
48    };
49    let host_ok = u
50        .host_str()
51        .map(str::to_ascii_lowercase)
52        .is_some_and(|h| h == "grab.com" || h.ends_with(".grab.com"));
53    u.scheme() == "https"
54        && host_ok
55        && u.username().is_empty()
56        && u.password().is_none()
57        && u.port().is_none()
58}
59
60fn ext_for(content_type: &str) -> &'static str {
61    match content_type
62        .split(';')
63        .next()
64        .unwrap_or("")
65        .trim()
66        .to_ascii_lowercase()
67        .as_str()
68    {
69        "image/jpeg" | "image/jpg" => "jpg",
70        "image/png" => "png",
71        "image/webp" => "webp",
72        "image/avif" => "avif",
73        "image/gif" => "gif",
74        _ => "img",
75    }
76}
77
78fn file_path(root: &Path, hash: &str, content_type: &str) -> PathBuf {
79    let a = hash.get(0..2).unwrap_or("00");
80    let b = hash.get(2..4).unwrap_or("00");
81    root.join("images")
82        .join(a)
83        .join(b)
84        .join(format!("{hash}.{}", ext_for(content_type)))
85}
86
87fn dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
88    image::ImageReader::new(std::io::Cursor::new(bytes))
89        .with_guessed_format()
90        .ok()?
91        .into_dimensions()
92        .ok()
93}
94
95fn write_private_file(path: &Path, data: &[u8]) -> Result<()> {
96    use std::io::Write;
97    if let Some(dir) = path.parent() {
98        create_private_dirs(dir)?;
99    }
100    let tmp = path.with_extension(format!("tmp.{}", std::process::id()));
101    {
102        let mut opts = fs::OpenOptions::new();
103        opts.write(true).create(true).truncate(true);
104        #[cfg(unix)]
105        {
106            use std::os::unix::fs::OpenOptionsExt;
107            opts.mode(0o600);
108        }
109        let mut f = opts.open(&tmp)?;
110        f.write_all(data)?;
111    }
112    fs::rename(&tmp, path)?;
113    Ok(())
114}
115
116type ImageRow = (String, i64, Option<i64>, Option<i64>, String);
117
118fn row_for(conn: &Connection, url: &str) -> Result<Option<ImageRow>> {
119    Ok(conn
120        .query_row(
121            "SELECT hash, bytes, width, height, content_type FROM images WHERE url = ?1",
122            [url],
123            |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?, r.get(4)?)),
124        )
125        .optional()?)
126}
127
128fn to_cached(root: &Path, url: &str, row: ImageRow, cached: bool) -> CachedImage {
129    let (hash, bytes, w, h, ct) = row;
130    CachedImage {
131        url: url.to_string(),
132        path: file_path(root, &hash, &ct),
133        hash,
134        bytes: u64::try_from(bytes).unwrap_or(0),
135        width: w.and_then(|v| u32::try_from(v).ok()),
136        height: h.and_then(|v| u32::try_from(v).ok()),
137        content_type: ct,
138        cached,
139    }
140}
141
142pub(crate) fn stats(conn: &Connection, max_bytes: u64) -> Result<ImageStats> {
143    let (files, bytes): (i64, i64) = conn.query_row(
144        "SELECT COUNT(*), COALESCE(SUM(bytes), 0) FROM
145            (SELECT hash, MAX(bytes) AS bytes FROM images GROUP BY hash)",
146        [],
147        |r| Ok((r.get(0)?, r.get(1)?)),
148    )?;
149    let urls: i64 = conn.query_row("SELECT COUNT(*) FROM images", [], |r| r.get(0))?;
150    Ok(ImageStats {
151        files: u64::try_from(files).unwrap_or(0),
152        urls: u64::try_from(urls).unwrap_or(0),
153        bytes: u64::try_from(bytes).unwrap_or(0),
154        max_bytes,
155    })
156}
157
158/// Evict least-recently-used files until the store is within `max_bytes`.
159/// Returns `(files evicted, bytes freed)`.
160pub(crate) fn evict(conn: &Connection, root: &Path, max_bytes: u64) -> Result<(u64, u64)> {
161    let mut stmt = conn.prepare(
162        "SELECT hash, MAX(bytes), MAX(last_access), MIN(content_type)
163         FROM images GROUP BY hash ORDER BY MAX(last_access) ASC, hash ASC",
164    )?;
165    let files: Vec<(String, i64, i64, String)> = stmt
166        .query_map([], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)))?
167        .collect::<rusqlite::Result<_>>()?;
168    drop(stmt);
169    let mut total: u64 = files.iter().map(|f| u64::try_from(f.1).unwrap_or(0)).sum();
170    let (mut n, mut freed) = (0u64, 0u64);
171    for (hash, bytes, _, ct) in files {
172        if total <= max_bytes {
173            break;
174        }
175        let bytes = u64::try_from(bytes).unwrap_or(0);
176        let _ = fs::remove_file(file_path(root, &hash, &ct));
177        conn.execute("DELETE FROM images WHERE hash = ?1", [&hash])?;
178        conn.execute(
179            "UPDATE menu_items SET image_hash = NULL WHERE image_hash = ?1",
180            [&hash],
181        )?;
182        total = total.saturating_sub(bytes);
183        n += 1;
184        freed += bytes;
185    }
186    Ok((n, freed))
187}
188
189/// Remove every image row and file. Returns the number of URLs removed.
190pub(crate) fn clear(conn: &Connection, root: &Path) -> Result<u64> {
191    let n = conn.execute("DELETE FROM images", [])? as u64;
192    conn.execute("UPDATE menu_items SET image_hash = NULL", [])?;
193    let dir = root.join("images");
194    if dir.exists() {
195        fs::remove_dir_all(dir)?;
196    }
197    Ok(n)
198}
199
200impl Cache {
201    /// The cached file for `url`, if present (and still on disk). Marks it
202    /// recently used.
203    pub fn cached_image(&self, url: &str) -> Option<CachedImage> {
204        self.cached_image_at(url, jwt::now())
205    }
206
207    fn cached_image_at(&self, url: &str, now: i64) -> Option<CachedImage> {
208        if !self.mode.reads() && !self.mode.writes() {
209            return None;
210        }
211        let conn = self.conn().ok()?;
212        let row = row_for(&conn, url).ok()??;
213        let img = to_cached(&self.root, url, row, true);
214        if !img.path.exists() {
215            let _ = conn.execute("DELETE FROM images WHERE url = ?1", [url]);
216            return None;
217        }
218        let _ = conn.execute(
219            "UPDATE images SET last_access = ?1 WHERE hash = ?2",
220            params![now, img.hash],
221        );
222        Some(img)
223    }
224
225    /// Store downloaded image bytes for `url` (dedupes by content hash), then
226    /// evict over the size cap.
227    pub fn store_image(&self, url: &str, bytes: &[u8], content_type: &str) -> Result<CachedImage> {
228        self.store_image_at(url, bytes, content_type, jwt::now())
229    }
230
231    pub fn store_image_at(
232        &self,
233        url: &str,
234        bytes: &[u8],
235        content_type: &str,
236        now: i64,
237    ) -> Result<CachedImage> {
238        if !self.mode.writes() {
239            return Err(Error::invalid("the cache is off; images are not stored"));
240        }
241        let hash = hex_sha256(bytes);
242        let path = file_path(&self.root, &hash, content_type);
243        if !path.exists() {
244            write_private_file(&path, bytes)?;
245        }
246        let dims = dimensions(bytes);
247        let conn = self.conn()?;
248        conn.execute(
249            "INSERT OR REPLACE INTO images
250                (url, hash, bytes, width, height, content_type, last_access, fetched_at)
251             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?7)",
252            params![
253                url,
254                hash,
255                i64::try_from(bytes.len()).unwrap_or(i64::MAX),
256                dims.map(|d| d.0),
257                dims.map(|d| d.1),
258                content_type,
259                now
260            ],
261        )?;
262        // Touch every URL sharing this file so LRU sees the file as fresh.
263        conn.execute(
264            "UPDATE images SET last_access = ?1 WHERE hash = ?2",
265            params![now, hash],
266        )?;
267        conn.execute(
268            "UPDATE menu_items SET image_hash = ?1 WHERE image_url = ?2",
269            params![hash, url],
270        )?;
271        evict(&conn, &self.root, self.config.image_max_bytes())?;
272        let row = row_for(&conn, url)?
273            .ok_or_else(|| Error::invalid("image is larger than the image cache cap"))?;
274        Ok(to_cached(&self.root, url, row, false))
275    }
276
277    /// The cached file for `url`, downloading it first if needed. Only Grab
278    /// CDN URLs are fetched ([`image_url_allowed`]); no cookies are sent.
279    pub async fn fetch_image(&self, url: &str) -> Result<CachedImage> {
280        if !image_url_allowed(url) {
281            return Err(Error::new(
282                ErrorKind::UrlNotAllowed,
283                "only https image URLs on grab.com hosts can be cached",
284            ));
285        }
286        if let Some(hit) = self.cached_image(url) {
287            return Ok(hit);
288        }
289        let client = reqwest::Client::builder()
290            .timeout(Duration::from_secs(30))
291            .redirect(reqwest::redirect::Policy::none())
292            .build()?;
293        let resp = client.get(url).send().await?;
294        let status = resp.status();
295        if !status.is_success() {
296            return Err(Error::new(
297                if status.as_u16() == 404 {
298                    ErrorKind::NotFound
299                } else {
300                    ErrorKind::Api
301                },
302                format!("{} fetching image", status.as_u16()),
303            ));
304        }
305        let content_type = resp
306            .headers()
307            .get(reqwest::header::CONTENT_TYPE)
308            .and_then(|v| v.to_str().ok())
309            .unwrap_or("")
310            .to_string();
311        if !content_type.to_ascii_lowercase().starts_with("image/") {
312            return Err(Error::new(
313                ErrorKind::Api,
314                format!("not an image (content-type `{content_type}`)"),
315            ));
316        }
317        if resp
318            .content_length()
319            .is_some_and(|l| l > MAX_IMAGE_BYTES as u64)
320        {
321            return Err(Error::invalid("image is too large"));
322        }
323        let bytes = resp.bytes().await?;
324        if bytes.len() > MAX_IMAGE_BYTES {
325            return Err(Error::invalid("image is too large"));
326        }
327        self.store_image(url, &bytes, &content_type)
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334    use crate::cache::{CacheConfig, CacheMode};
335
336    #[test]
337    fn url_allowlist() {
338        assert!(image_url_allowed(
339            "https://food-cms.grab.com/compressed_webp/x.webp"
340        ));
341        assert!(image_url_allowed(
342            "https://huawei-food-cms.grab.com/a/b.avif"
343        ));
344        assert!(!image_url_allowed("http://food-cms.grab.com/x.png"));
345        assert!(!image_url_allowed("https://evil.com/x.png"));
346        assert!(!image_url_allowed("https://grab.com.evil.com/x.png"));
347        assert!(!image_url_allowed("https://user:pw@food-cms.grab.com/x"));
348        assert!(!image_url_allowed("https://food-cms.grab.com:8443/x"));
349        assert!(!image_url_allowed("not a url"));
350    }
351
352    fn png(w: u32, h: u32) -> Result<Vec<u8>> {
353        let img = image::RgbImage::new(w, h);
354        let mut out = std::io::Cursor::new(Vec::new());
355        img.write_to(&mut out, image::ImageFormat::Png)
356            .map_err(|e| Error::invalid(e.to_string()))?;
357        Ok(out.into_inner())
358    }
359
360    #[test]
361    fn store_dedupes_and_records_dimensions() -> Result<()> {
362        let dir = tempfile::tempdir()?;
363        let c = Cache::new(dir.path(), CacheConfig::default(), CacheMode::Normal);
364        let bytes = png(3, 2)?;
365        let a = c.store_image_at("https://food-cms.grab.com/a.png", &bytes, "image/png", 10)?;
366        let b = c.store_image_at("https://food-cms.grab.com/b.png", &bytes, "image/png", 11)?;
367        assert_eq!(a.path, b.path);
368        assert_eq!((a.width, a.height), (Some(3), Some(2)));
369        assert!(a.path.exists());
370        assert!(a.path.to_string_lossy().contains(&format!(
371            "images/{}/{}/{}.png",
372            &a.hash[0..2],
373            &a.hash[2..4],
374            a.hash
375        )));
376        let s = c.stats_at(11)?;
377        assert_eq!((s.images.files, s.images.urls), (1, 2));
378        assert_eq!(s.images.bytes, bytes.len() as u64);
379        let hit = c
380            .cached_image("https://food-cms.grab.com/a.png")
381            .ok_or_else(|| Error::not_found("a"))?;
382        assert!(hit.cached);
383        #[cfg(unix)]
384        {
385            use std::os::unix::fs::PermissionsExt;
386            assert_eq!(fs::metadata(&a.path)?.permissions().mode() & 0o777, 0o600);
387        }
388        Ok(())
389    }
390
391    #[test]
392    fn lru_eviction_by_total_size() -> Result<()> {
393        let dir = tempfile::tempdir()?;
394        // 1 MiB cap; each "image" is 400 KiB of distinct bytes.
395        let cfg = CacheConfig {
396            image_max_mb: Some(1),
397            ..CacheConfig::default()
398        };
399        let c = Cache::new(dir.path(), cfg, CacheMode::Normal);
400        let blob = |n: u8| vec![n; 400 * 1024];
401        let url = |n: u8| format!("https://food-cms.grab.com/{n}.webp");
402        let first = c.store_image_at(&url(1), &blob(1), "image/webp", 100)?;
403        c.store_image_at(&url(2), &blob(2), "image/webp", 200)?;
404        // Touch #1 so #2 becomes the least recently used.
405        c.cached_image_at(&url(1), 300)
406            .ok_or_else(|| Error::not_found("1"))?;
407        c.store_image_at(&url(3), &blob(3), "image/webp", 400)?;
408        assert!(c.cached_image(&url(1)).is_some());
409        assert!(c.cached_image(&url(2)).is_none());
410        assert!(c.cached_image(&url(3)).is_some());
411        assert!(first.path.exists());
412        let s = c.stats_at(400)?;
413        assert_eq!(s.images.files, 2);
414        assert!(s.images.bytes <= 1024 * 1024);
415
416        // A cap below one image: prune evicts down to it.
417        let tiny = Cache::new(
418            dir.path(),
419            CacheConfig {
420                image_max_mb: Some(0),
421                ..CacheConfig::default()
422            },
423            CacheMode::Normal,
424        );
425        let r = tiny.prune_at(500)?;
426        assert_eq!(r.images_evicted, 2);
427        assert!(!first.path.exists());
428        Ok(())
429    }
430
431    #[test]
432    fn off_mode_stores_nothing() -> Result<()> {
433        let dir = tempfile::tempdir()?;
434        let c = Cache::new(dir.path(), CacheConfig::default(), CacheMode::Off);
435        assert!(
436            c.store_image("https://food-cms.grab.com/a.png", b"x", "image/png")
437                .is_err()
438        );
439        assert!(c.cached_image("https://food-cms.grab.com/a.png").is_none());
440        Ok(())
441    }
442}