From 791c0183c9f42ab38ad1ee71c666691fe56b5689 Mon Sep 17 00:00:00 2001
From: eternal-flame-AD <yume@yumechi.jp>
Date: Sat, 21 Dec 2024 20:00:28 -0600
Subject: [PATCH] set a more realistic cache-control for URL summary

Signed-off-by: eternal-flame-AD <yume@yumechi.jp>
---
 src/lib.rs | 29 ++++++++++++++++++-----------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 411e8ce..221e18f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -19,7 +19,7 @@ use axum::extract::ConnectInfo;
 use axum::{
     body::Body,
     extract::{Path, Query, State},
-    http::{self, HeaderMap, StatusCode},
+    http::{self, HeaderMap, HeaderValue, StatusCode},
     response::{IntoResponse, Redirect, Response},
     routing::get,
     Json, Router,
@@ -250,20 +250,27 @@ pub async fn set_cache_control<
     next: axum::middleware::Next,
 ) -> Response {
     let mut resp = next.run(request).await;
+    let ok = resp.status().is_success();
+    let headers = resp.headers_mut();
     if enabled {
-        if resp.status() == StatusCode::OK {
-            let headers = resp.headers_mut();
-            headers.insert(
-                "Cache-Control",
-                "public, max-age=31536000, immutable".parse().unwrap(),
-            );
+        if ok {
+            if headers.get("content-type").map_or(false, |x| {
+                [b"image/", b"video/", b"audio/"]
+                    .into_iter()
+                    .any(|y| x.as_bytes().starts_with(y))
+            }) {
+                headers.insert(
+                    "cache-control",
+                    HeaderValue::from_static("public, max-age=31536000, immutable"),
+                );
+            } else {
+                headers.insert("cache-control", HeaderValue::from_static("max-age=86400"));
+            }
         } else {
-            let headers = resp.headers_mut();
-            headers.insert("Cache-Control", "max-age=300".parse().unwrap());
+            headers.insert("cache-control", HeaderValue::from_static("max-age=300"));
         }
     } else {
-        let headers = resp.headers_mut();
-        headers.insert("Cache-Control", "no-store".parse().unwrap());
+        headers.insert("cache-control", HeaderValue::from_static("no-store"));
     }
     resp
 }