From e2936f143e914c9b2dda1b5df95ed919b7fdf54f Mon Sep 17 00:00:00 2001
From: kane50613 <kane@yeecord.com>
Date: Fri, 6 Dec 2024 14:31:44 +0800
Subject: [PATCH] feat: Into Response for CacheResponse, create_cache_key
 accept u8 array

---
 src/db.rs                  | 14 ++++++++++++++
 src/handlers/api.rs        | 35 ++++++++++-------------------------
 src/handlers/invalidate.rs |  6 +++++-
 src/hash.rs                | 21 ++++++++-------------
 4 files changed, 37 insertions(+), 39 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index 62089a9..fa4df81 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -1,3 +1,5 @@
+use axum::response::Response;
+use reqwest::header::CONTENT_TYPE;
 use sqlx::Statement;
 use std::sync::LazyLock;
 
@@ -25,6 +27,18 @@ pub struct CacheResponse {
   pub content_type: Option<String>,
 }
 
+impl From<CacheResponse> for Response {
+  fn from(val: CacheResponse) -> Self {
+    let mut builder = Response::builder().status(val.status);
+
+    if let Some(content_type) = val.content_type {
+      builder = builder.header(CONTENT_TYPE, content_type);
+    }
+
+    builder.body(val.body.unwrap_or_default().into()).unwrap()
+  }
+}
+
 pub struct Database<'a> {
   connection: SqlitePool,
   queries: OnceCell<Queries<'a>>,
diff --git a/src/handlers/api.rs b/src/handlers/api.rs
index ecee7ad..a1519e7 100644
--- a/src/handlers/api.rs
+++ b/src/handlers/api.rs
@@ -12,7 +12,7 @@ use reqwest::{
 
 use crate::{
   db::{CacheResponse, DB},
-  hash::{create_cache_key, CacheKeyPayload},
+  hash::create_cache_key,
 };
 
 pub static HTTP: LazyLock<Client> = LazyLock::new(Client::default);
@@ -30,17 +30,14 @@ pub async fn api_handler(request: Request) -> Response {
       .unwrap();
   };
 
-  let cache_key = create_cache_key(CacheKeyPayload {
-    method: head.method.to_string(),
-    url: path.to_string(),
-    authorization: head
-      .headers
-      .get(AUTHORIZATION)
-      .map(|x| x.to_str().unwrap().to_string()),
-  });
+  let cache_key = create_cache_key(
+    head.method.as_str().as_bytes(),
+    path.as_str().as_bytes(),
+    head.headers.get(AUTHORIZATION).map(|x| x.as_bytes()),
+  );
 
   if let Some(cache_response) = DB.get(cache_key).await {
-    return response_from_cache(cache_response);
+    return cache_response.into();
   }
 
   let url = Uri::builder()
@@ -59,7 +56,7 @@ pub async fn api_handler(request: Request) -> Response {
     .await
     .unwrap();
 
-  let status = response.status();
+  let status = response.status().as_u16();
 
   let content_type = response.headers_mut().remove(CONTENT_TYPE);
   let body = response.bytes().await.unwrap();
@@ -68,22 +65,10 @@ pub async fn api_handler(request: Request) -> Response {
     key: cache_key,
     content_type: content_type.map(|x| x.to_str().unwrap().to_string()),
     body: Some(body.to_vec()),
-    status: status.as_u16(),
+    status,
   };
 
   DB.set(&cache_payload).await;
 
-  response_from_cache(cache_payload)
-}
-
-fn response_from_cache(cache_response: CacheResponse) -> Response {
-  let mut builder = Response::builder().status(cache_response.status);
-
-  if let Some(content_type) = &cache_response.content_type {
-    builder = builder.header(CONTENT_TYPE, content_type);
-  }
-
-  builder
-    .body(cache_response.body.unwrap_or_default().into())
-    .unwrap()
+  cache_payload.into()
 }
diff --git a/src/handlers/invalidate.rs b/src/handlers/invalidate.rs
index e958b9f..ac7feba 100644
--- a/src/handlers/invalidate.rs
+++ b/src/handlers/invalidate.rs
@@ -6,7 +6,11 @@ use crate::{
 };
 
 pub async fn invalidate_handler(body: Json<CacheKeyPayload>) {
-  let key = create_cache_key(body.0);
+  let key = create_cache_key(
+    body.method.as_bytes(),
+    body.url.as_bytes(),
+    body.authorization.as_ref().map(|header| header.as_bytes()),
+  );
 
   DB.delete(key).await;
 }
diff --git a/src/hash.rs b/src/hash.rs
index fc056a5..c52295d 100644
--- a/src/hash.rs
+++ b/src/hash.rs
@@ -8,21 +8,16 @@ pub struct CacheKeyPayload {
   pub authorization: Option<String>,
 }
 
-pub fn create_cache_key(payload: CacheKeyPayload) -> i64 {
-  let method_bytes = payload.method.as_bytes();
-  let url_bytes = payload.url.as_bytes();
-  let authorization_bytes = payload
-    .authorization
-    .as_deref()
-    .unwrap_or_default()
-    .as_bytes();
-
+pub fn create_cache_key(method: &[u8], url: &[u8], authorization: Option<&[u8]>) -> i64 {
   let mut buffer =
-    Vec::with_capacity(method_bytes.len() + url_bytes.len() + authorization_bytes.len());
+    Vec::with_capacity(method.len() + url.len() + authorization.map_or(0, |x| x.len()));
+
+  buffer.extend_from_slice(method);
+  buffer.extend_from_slice(url);
 
-  buffer.extend_from_slice(method_bytes);
-  buffer.extend_from_slice(url_bytes);
-  buffer.extend_from_slice(authorization_bytes);
+  if let Some(authorization) = authorization {
+    buffer.extend_from_slice(authorization);
+  }
 
   xxh3::xxh3_64(&buffer) as i64
 }