Code restructurization

src/auth/jwt.rs

@@ -1,9 +1,51 @@
 use crate::auth::types::AuthValidator;
-use crate::utils::jwt::{check_jwt, JWT_TOKEN};
+use ahash::AHasher;
+use base64::engine::general_purpose::URL_SAFE_NO_PAD;
+use base64::Engine;
+use jsonwebtoken::{Algorithm, DecodingKey, Validation};
+use moka::sync::Cache;
+use moka::Expiry;
 use pingora_proxy::Session;
+use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
+use std::env;
+use std::hash::{Hash, Hasher};
+use std::sync::{Arc, LazyLock};
+use std::time::{Duration, Instant, SystemTime};
 use urlencoding::decode;
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Claims {
+    pub master_key: String,
+    pub owner: String,
+    pub exp: u64,
+    pub random: Option<String>,
+}
+
+#[derive(Debug, Deserialize)]
+struct Expired {
+    exp: Option<u64>,
+}
+
+static JWT_VALIDATION: LazyLock<Validation> = LazyLock::new(|| Validation::new(Algorithm::HS256));
+
+pub static JWT_TOKEN: LazyLock<Option<Arc<str>>> = LazyLock::new(|| match env::var("JWT_KEY") {
+    Ok(key) if !key.is_empty() => Some(Arc::from(key.as_str())),
+    _ => None,
+});
+
+static JWT_CACHE: LazyLock<Cache<u64, u64>> = LazyLock::new(|| Cache::builder().max_capacity(100_000).expire_after(JwtExpiry).build());
+struct JwtExpiry;
+impl Expiry<u64, u64> for JwtExpiry {
+    fn expire_after_create(&self, _key: &u64, value: &u64, _current_time: Instant) -> Option<Duration> {
+        let now = SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs();
+        if *value > now {
+            Some(Duration::from_secs(value - now))
+        } else {
+            Some(Duration::ZERO)
+        }
+    }
+}
 pub struct JwtAuth();
 #[async_trait::async_trait]
 impl AuthValidator for JwtAuth {
@@ -40,3 +82,53 @@ fn get_query_param(session: &mut Session, key: &str) -> Option<String> {
         .collect();
     params.get(key).and_then(|v| decode(v).ok()).map(|s| s.to_string())
 }
+
+pub fn check_jwt(token: &str, secret: &str) -> bool {
+    let key = hash_token(token, secret);
+    let now = SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs();
+    if let Some(exp) = JWT_CACHE.get(&key) {
+        if exp < now {
+            return false;
+        }
+        return true;
+    }
+    match is_expired(token, now) {
+        Ok(true) => return false,
+        Ok(false) => {}
+        Err(_) => return false,
+    }
+
+    match jsonwebtoken::decode::<Claims>(token, &DecodingKey::from_secret(secret.as_ref()), &JWT_VALIDATION) {
+        Ok(data) => {
+            let now = SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs();
+            if data.claims.exp > now {
+                JWT_CACHE.insert(key, data.claims.exp);
+                true
+            } else {
+                false
+            }
+        }
+        Err(_) => false,
+    }
+}
+
+fn is_expired(token: &str, now: u64) -> Result<bool, Box<dyn std::error::Error>> {
+    let parts: Vec<&str> = token.split('.').collect();
+    if parts.len() != 3 {
+        return Err("Invalid JWT format".into());
+    }
+    let decoded = URL_SAFE_NO_PAD.decode(parts[1])?;
+    let claims: Expired = serde_json::from_slice(&decoded)?;
+    if let Some(exp) = claims.exp {
+        Ok(exp < now)
+    } else {
+        Ok(true)
+    }
+}
+
+fn hash_token(token: &str, secret: &str) -> u64 {
+    let mut hasher = AHasher::default();
+    token.hash(&mut hasher);
+    secret.hash(&mut hasher);
+    hasher.finish()
+}