From e25af6fb500b1ada9082f7f7ffebb91f521fe3af Mon Sep 17 00:00:00 2001
From: lennoxlotl <admin@lennoxlotl.dev>
Date: Tue, 22 Oct 2024 01:47:20 +0200
Subject: [PATCH] feat(api): add ability to control img cache

---
 api/src/endpoint/{show.rs => index.rs} | 34 ++++++++++++++++++++++----
 api/src/endpoint/mod.rs                |  2 +-
 api/src/main.rs                        |  7 +++++-
 3 files changed, 36 insertions(+), 7 deletions(-)
 rename api/src/endpoint/{show.rs => index.rs} (64%)

diff --git a/api/src/endpoint/show.rs b/api/src/endpoint/index.rs
similarity index 64%
rename from api/src/endpoint/show.rs
rename to api/src/endpoint/index.rs
index 14f7507..bbdd8bd 100644
--- a/api/src/endpoint/show.rs
+++ b/api/src/endpoint/index.rs
@@ -2,23 +2,28 @@ use super::{
     fairing::{bucket::BucketGuard, database::PostgresDb},
     v1::{error::Error, UploaderResult},
 };
-use crate::{database::query::image::find_image_by_id, s3::bucket::BucketOperations};
+use crate::{database::query::image::find_image_by_id, s3::bucket::BucketOperations, GlobalConfig};
 use rocket::{
     get,
     http::ContentType,
     response::{self, Responder},
-    Request, Response,
+    Request, Response, State,
 };
 use std::{io::Cursor, str::FromStr};
 
 pub struct ImageShowResponse {
     data: Vec<u8>,
     content_type: String,
+    cache_time: usize,
 }
 
 impl ImageShowResponse {
-    pub fn new(data: Vec<u8>, content_type: String) -> Self {
-        Self { data, content_type }
+    pub fn new(data: Vec<u8>, content_type: String, cache_time: usize) -> Self {
+        Self {
+            data,
+            content_type,
+            cache_time,
+        }
     }
 }
 
@@ -27,16 +32,31 @@ impl<'r> Responder<'r, 'static> for ImageShowResponse {
     fn respond_to(self, _: &'r Request<'_>) -> response::Result<'static> {
         Response::build()
             .header(ContentType::from_str(&self.content_type).unwrap_or(ContentType::default()))
+            .raw_header(
+                "Cache-Control",
+                if self.cache_time > 0 {
+                    format!("max-age={}", self.cache_time)
+                } else {
+                    "no-cache".into()
+                },
+            )
             .streamed_body(Cursor::new(self.data))
             .ok()
     }
 }
 
+// TODO: come up with something better
+#[get("/")]
+pub async fn index() -> &'static str {
+    "hi :wave:"
+}
+
 #[get("/<id>")]
 pub async fn show_image(
     id: &str,
     database: PostgresDb,
     bucket: BucketGuard,
+    config: &State<GlobalConfig>,
 ) -> UploaderResult<ImageShowResponse> {
     let mut transaction = database.begin().await.map_err(|_| Error::DatabaseError)?;
     let image = find_image_by_id(&mut transaction, &id.to_string())
@@ -50,5 +70,9 @@ pub async fn show_image(
         .await
         .map_err(|_| Error::ImageConvertError)?
         .to_vec();
-    Ok(ImageShowResponse::new(image_bytes, image_type.to_string()))
+    Ok(ImageShowResponse::new(
+        image_bytes,
+        image_type.to_string(),
+        config.cache_length.unwrap_or(0),
+    ))
 }
diff --git a/api/src/endpoint/mod.rs b/api/src/endpoint/mod.rs
index 96d5739..9d7ae6b 100644
--- a/api/src/endpoint/mod.rs
+++ b/api/src/endpoint/mod.rs
@@ -1,3 +1,3 @@
 pub mod fairing;
-pub mod show;
+pub mod index;
 pub mod v1;
diff --git a/api/src/main.rs b/api/src/main.rs
index ac8a918..82c83ad 100644
--- a/api/src/main.rs
+++ b/api/src/main.rs
@@ -13,6 +13,8 @@ pub struct GlobalConfig {
     public_url: String,
     // Length of the image id used for "shwoing" the image
     image_id_length: usize,
+    // Defines a Cache-Control header, time is in seconds
+    cache_length: Option<usize>,
     // If not empty requires an authentication header containing this key for uploads
     auth_key: Option<String>,
 }
@@ -21,7 +23,10 @@ pub struct GlobalConfig {
 async fn main() -> Result<(), rocket::Error> {
     rocket::build()
         .mount("/api/v1/", create_v1_routes())
-        .mount("/", routes![endpoint::show::show_image])
+        .mount(
+            "/",
+            routes![endpoint::index::index, endpoint::index::show_image],
+        )
         .attach(AdHoc::config::<GlobalConfig>())
         .attach(BucketFairing::new())
         .attach(PostgresFairing::new())