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

Cargo.lock

@@ -120,6 +120,7 @@ checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c"
 name = "aralez"
 version = "0.9.2"
 dependencies = [
+ "ahash",
  "arc-swap",
  "async-trait",
  "axum",
@@ -132,6 +133,7 @@ dependencies = [
  "jsonwebtoken",
  "log",
  "mimalloc",
+ "moka",
  "notify",
  "pingora",
  "pingora-core",
@@ -645,6 +647,24 @@ dependencies = [
  "cfg-if",
 ]
 
+[[package]]
+name = "crossbeam-channel"
+version = "0.5.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
+dependencies = [
+ "crossbeam-utils",
+]
+
 [[package]]
 name = "crossbeam-queue"
 version = "0.3.12"
@@ -2016,6 +2036,23 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
+[[package]]
+name = "moka"
+version = "0.12.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "957228ad12042ee839f93c8f257b62b4c0ab5eaae1d4fa60de53b27c9d7c5046"
+dependencies = [
+ "crossbeam-channel",
+ "crossbeam-epoch",
+ "crossbeam-utils",
+ "equivalent",
+ "parking_lot",
+ "portable-atomic",
+ "smallvec",
+ "tagptr",
+ "uuid",
+]
+
 [[package]]
 name = "neli"
 version = "0.7.4"
@@ -3673,6 +3710,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "tagptr"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
+
 [[package]]
 name = "thiserror"
 version = "1.0.69"
@@ -4083,6 +4126,17 @@ version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
 
+[[package]]
+name = "uuid"
+version = "1.23.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5ac8b6f42ead25368cf5b098aeb3dc8a1a2c05a3eee8a9a1a68c640edbfc79d9"
+dependencies = [
+ "getrandom 0.4.2",
+ "js-sys",
+ "wasm-bindgen",
+]
+
 [[package]]
 name = "vcpkg"
 version = "0.2.15"

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 struct Claims {
-    pub(crate) user: String,
+    pub master_key: String,
-    pub(crate) exp: u64,
+    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));
 
-/*
+static JWT_CACHE: LazyLock<Cache<u64, u64>> = LazyLock::new(|| Cache::builder().max_capacity(100_000).expire_after(JwtExpiry).build());
-pub fn check_jwt(input: &str, secret: &str) -> bool {
+struct JwtExpiry;
-    let validation = Validation::new(Algorithm::HS256);
+impl Expiry<u64, u64> for JwtExpiry {
-    let token_data = decode::<Claims>(&input, &DecodingKey::from_secret(secret.as_ref()), &validation);
+    fn expire_after_create(&self, _key: &u64, value: &u64, _current_time: Instant) -> Option<Duration> {
-    token_data.is_ok()
+        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) {
+    let now = SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs();
-        return v;
+    if let Some(exp) = JWT_CACHE.get(&key) {
+        if exp < now {
+            return false;
         }
-    let result = decode::<Claims>(token, &DecodingKey::from_secret(secret.as_ref()), &JWT_VALIDATION).is_ok();
+        return true;
-    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 now = SystemTime::now() + Duration::from_secs(payload.exp * 60);
-        let a = now.duration_since(UNIX_EPOCH).unwrap().as_secs();
+        let expire = now.duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
-        let claim = crate::utils::jwt::Claims { user: payload.owner, exp: a };
+
+        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) => {