1use 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
25const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024;
27
28#[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 pub cached: bool,
40}
41
42pub 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
158pub(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
189pub(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 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 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 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 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 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 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 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}