Cache for JWT tokens, to minimize crypto. BRAKING: Claims key "valid" renamed to "exp"

src/utils/jwt.rs

@@ -1,37 +1,82 @@
 use ahash::AHasher;
+use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
 use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
 use moka::sync::Cache;
+use moka::Expiry;
 use serde::{Deserialize, Serialize};
 use std::hash::{Hash, Hasher};
 use std::sync::LazyLock;
+use std::time::{Duration, Instant, SystemTime};
 
 #[derive(Debug, Serialize, Deserialize)]
-pub(crate) struct Claims {
-    pub(crate) user: String,
-    pub(crate) exp: u64,
+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_CACHE: LazyLock<Cache<u64, bool>> = LazyLock::new(|| Cache::builder().max_capacity(100_000).time_to_live(std::time::Duration::from_secs(60)).build());
 static JWT_VALIDATION: LazyLock<Validation> = LazyLock::new(|| Validation::new(Algorithm::HS256));
 
-/*
-pub fn check_jwt(input: &str, secret: &str) -> bool {
-    let validation = Validation::new(Algorithm::HS256);
-    let token_data = decode::<Claims>(&input, &DecodingKey::from_secret(secret.as_ref()), &validation);
-    token_data.is_ok()
+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 fn check_jwt(token: &str, secret: &str) -> bool {
     let key = hash_token(token, secret);
-    if let Some(v) = JWT_CACHE.get(&key) {
-        return v;
+    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;
     }
-    let result = decode::<Claims>(token, &DecodingKey::from_secret(secret.as_ref()), &JWT_VALIDATION).is_ok();
-    if result {
-        JWT_CACHE.insert(key, true);
+    match is_expired(token, now) {
+        Ok(true) => return false,
+        Ok(false) => {}
+        Err(_) => return false,
+    }
+
+    match 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)
     }
-    result
 }
 
 fn hash_token(token: &str, secret: &str) -> u64 {

src/web/webserver.rs

@@ -1,4 +1,6 @@
 use crate::utils::discovery::APIUpstreamProvider;
+// use std::net::SocketAddr;
+use crate::utils::jwt::Claims;
 use crate::utils::structs::{Config, Configuration, UpstreamsDashMap};
 use crate::utils::tools::{upstreams_liveness_json, upstreams_to_json};
 use axum::body::Body;
@@ -13,22 +15,14 @@ use futures::SinkExt;
 use jsonwebtoken::{encode, EncodingKey, Header};
 use log::{error, info, warn};
 use prometheus::{gather, Encoder, TextEncoder};
-use serde::{Deserialize, Serialize};
+use serde::Serialize;
 use std::collections::HashMap;
-// use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use subtle::ConstantTimeEq;
 use tokio::net::TcpListener;
 use tower_http::services::ServeDir;
 
-#[derive(Deserialize)]
-struct InputKey {
-    master_key: String,
-    owner: String,
-    valid: u64,
-}
-
 #[derive(Serialize, Debug)]
 struct OutToken {
     token: String,
@@ -119,15 +113,21 @@ async fn apply_config(content: &str, mut st: AppState) {
     }
 }
 
-async fn jwt_gen(State(state): State<AppState>, Json(payload): Json<InputKey>) -> (StatusCode, Json<OutToken>) {
+async fn jwt_gen(State(state): State<AppState>, Json(payload): Json<Claims>) -> (StatusCode, Json<OutToken>) {
     if payload.master_key == state.master_key {
-        let now = SystemTime::now() + Duration::from_secs(payload.valid * 60);
-        let a = now.duration_since(UNIX_EPOCH).unwrap().as_secs();
-        let claim = crate::utils::jwt::Claims { user: payload.owner, exp: a };
+        let now = SystemTime::now() + Duration::from_secs(payload.exp * 60);
+        let expire = now.duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
+
+        let claim = Claims {
+            master_key: String::new(),
+            owner: payload.owner,
+            exp: expire,
+            random: payload.random,
+        };
         match encode(&Header::default(), &claim, &EncodingKey::from_secret(payload.master_key.as_ref())) {
             Ok(t) => {
                 let tok = OutToken { token: t };
-                info!("Generating token: {:?}", tok);
+                info!("Generating token: {:?}", tok.token);
                 (StatusCode::CREATED, Json(tok))
             }
             Err(e) => {