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"
@@ -3060,6 +3097,7 @@ dependencies = [
  "base64",
  "bytes",
  "encoding_rs",
+ "futures-channel",
  "futures-core",
  "futures-util",
  "h2",
@@ -3672,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"
@@ -4082,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"

Cargo.toml

@@ -25,7 +25,7 @@ log = "0.4.29"
 futures = "0.3.32"
 notify = "9.0.0-rc.2"
 axum = { version = "0.8.8" }
-reqwest = { version = "0.13.2", features = ["json", "stream"] }
+reqwest = { version = "0.13.2", features = ["json", "stream", "blocking"] }
 serde_yml = "0.0.12"
 rand = "0.10.0"
 base64 = "0.22.1"
@@ -44,3 +44,5 @@ privdrop = "0.5.6"
 ctrlc = "3.5.2"
 serde_json = "1.0.149"
 subtle = "2.6.1"
+moka = { version = "0.12.1", features = ["sync"] }
+ahash = "0.8.12"

src/utils/auth.rs

@@ -1,21 +1,157 @@
 use crate::utils::jwt::check_jwt;
+// use reqwest::Client;
+use axum::http::StatusCode;
 use base64::engine::general_purpose::STANDARD;
 use base64::Engine;
 use pingora_proxy::Session;
 use std::collections::HashMap;
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 use subtle::ConstantTimeEq;
 use urlencoding::decode;
 
+// use pingora::http::{RequestHeader, ResponseHeader, StatusCode};
+use pingora::http::RequestHeader;
+// --------------------------------- //
+use pingora_core::connectors::http::Connector;
+use pingora_core::upstreams::peer::HttpPeer;
+use pingora_http::ResponseHeader;
+// --------------------------------- //
+
+#[async_trait::async_trait]
 trait AuthValidator {
-    fn validate(&self, session: &Session) -> bool;
+    async fn validate(&self, session: &mut Session) -> bool;
 }
 struct BasicAuth<'a>(&'a str);
 struct ApiKeyAuth<'a>(&'a str);
 struct JwtAuth<'a>(&'a str);
+struct ForwardAuth<'a>(&'a str);
 
+pub static AUTH_CONNECTOR: LazyLock<Connector> = LazyLock::new(|| Connector::new(None));
+
+#[async_trait::async_trait]
+impl AuthValidator for ForwardAuth<'_> {
+    async fn validate(&self, session: &mut Session) -> bool {
+        let method = match session.req_header().method.as_str() {
+            "HEAD" => "HEAD",
+            _ => "GET",
+        };
+
+        let auth_url = self.0;
+
+        let (plain, tls) = if let Some(p) = auth_url.strip_prefix("http://") {
+            (p, false)
+        } else if let Some(p) = auth_url.strip_prefix("https://") {
+            (p, true)
+        } else {
+            return false;
+        };
+
+        let (addr, uri) = if let Some(pos) = plain.find('/') {
+            (&plain[..pos], &plain[pos..])
+        } else {
+            (plain, "/")
+        };
+
+        let hp = match split_host_port(addr, tls) {
+            Some(hp) => hp,
+            None => return false,
+        };
+
+        let peer = HttpPeer::new((hp.0, hp.1), tls, hp.0.to_string());
+
+        let (mut http_session, _) = match AUTH_CONNECTOR.get_http_session(&peer).await {
+            Ok(s) => s,
+            Err(e) => {
+                log::warn!("ForwardAuth: connect failed: {}", e);
+                return false;
+            }
+        };
+
+        let mut auth_req = match RequestHeader::build(method, uri.as_bytes(), None) {
+            Ok(r) => r,
+            Err(e) => {
+                log::warn!("ForwardAuth: failed to build request: {}", e);
+                return false;
+            }
+        };
+
+        // auth_req.headers = session.req_header().headers.clone();
+        auth_req.insert_header("Host", addr).ok();
+        auth_req.insert_header("X-Forwarded-Uri", uri).ok();
+        auth_req.insert_header("X-Forwarded-Method", session.req_header().method.as_str()).ok();
+        if let Some(auth) = session.req_header().headers.get("authorization") {
+            auth_req.insert_header("Authorization", auth.clone()).ok();
+        }
+
+        if let Some(cookie) = session.req_header().headers.get("cookie") {
+            auth_req.insert_header("Cookie", cookie.clone()).ok();
+        }
+
+        if tls {
+            auth_req.insert_header("X-Forwarded-Proto", "https").ok();
+        } else {
+            auth_req.insert_header("X-Forwarded-Proto", "http").ok();
+        }
+
+        if let Err(e) = http_session.write_request_header(Box::new(auth_req)).await {
+            log::warn!("ForwardAuth: write failed: {}", e);
+            return false;
+        }
+
+        let status = match http_session.read_response_header().await {
+            Ok(_) => http_session.response_header().map(|r| r.status.as_u16()).unwrap_or(500),
+            Err(e) => {
+                log::warn!("ForwardAuth: read failed: {}", e);
+                return false;
+            }
+        };
+
+        let auth_headers_to_forward: Vec<(String, String)> = if let Some(resp_header) = http_session.response_header() {
+            resp_header
+                .headers
+                .iter()
+                .filter_map(|(name, value)| {
+                    let name_str = name.as_str();
+                    if name_str.starts_with("x-") || name_str.starts_with("remote-") || name_str.starts_with("locat") {
+                        value.to_str().ok().map(|v| (name_str.to_string(), v.to_string()))
+                    } else {
+                        None
+                    }
+                })
+                .collect()
+        } else {
+            Vec::new()
+        };
+
+        AUTH_CONNECTOR.release_http_session(http_session, &peer, None).await;
+
+        if (200..300).contains(&status) {
+            for (name, value) in auth_headers_to_forward {
+                session.req_header_mut().insert_header(name, value).ok();
+            }
+            true
+        } else if status == 302 || status == 301 {
+            let resp = ResponseHeader::build(StatusCode::MOVED_PERMANENTLY, None);
+            match resp {
+                Ok(mut r) => {
+                    for (name, value) in auth_headers_to_forward {
+                        r.insert_header(name, value).ok();
+                    }
+                    let _ = r.insert_header("Content-Length", "0");
+                    let _ = session.write_response_header(Box::new(r), true).await;
+                    true
+                }
+                Err(_) => return false,
+            }
+        } else {
+            false
+        }
+    }
+}
+
+#[async_trait::async_trait]
 impl AuthValidator for BasicAuth<'_> {
-    fn validate(&self, session: &Session) -> bool {
+    async fn validate(&self, session: &mut Session) -> bool {
         if let Some(header) = session.get_header("authorization") {
             if let Some(h) = header.to_str().ok() {
                 if let Some((_, val)) = h.split_once(' ') {
@@ -31,8 +167,9 @@ impl AuthValidator for BasicAuth<'_> {
     }
 }
 
+#[async_trait::async_trait]
 impl AuthValidator for ApiKeyAuth<'_> {
-    fn validate(&self, session: &Session) -> bool {
+    async fn validate(&self, session: &mut Session) -> bool {
         if let Some(header) = session.get_header("x-api-key") {
             if let Some(h) = header.to_str().ok() {
                 return h.as_bytes().ct_eq(self.0.as_bytes()).into();
@@ -42,8 +179,9 @@ impl AuthValidator for ApiKeyAuth<'_> {
     }
 }
 
+#[async_trait::async_trait]
 impl AuthValidator for JwtAuth<'_> {
-    fn validate(&self, session: &Session) -> bool {
+    async fn validate(&self, session: &mut Session) -> bool {
         let jwtsecret = self.0;
         if let Some(tok) = get_query_param(session, "araleztoken") {
             return check_jwt(tok.as_str(), jwtsecret);
@@ -61,11 +199,12 @@ impl AuthValidator for JwtAuth<'_> {
     }
 }
 
-pub fn authenticate(auth_type: &Arc<str>, credentials: &Arc<str>, session: &Session) -> bool {
+pub async fn authenticate(auth_type: &Arc<str>, credentials: &Arc<str>, session: &mut Session) -> bool {
     match &**auth_type {
-        "basic" => BasicAuth(credentials).validate(session),
+        "basic" => BasicAuth(credentials).validate(session).await,
-        "apikey" => ApiKeyAuth(credentials).validate(session),
+        "apikey" => ApiKeyAuth(credentials).validate(session).await,
-        "jwt" => JwtAuth(credentials).validate(session),
+        "jwt" => JwtAuth(credentials).validate(session).await,
+        "forward" => ForwardAuth(credentials).validate(session).await,
         _ => {
             log::warn!("Unsupported authentication mechanism : {}", auth_type);
             false
@@ -73,7 +212,7 @@ pub fn authenticate(auth_type: &Arc<str>, credentials: &Arc<str>, session: &Sess
     }
 }
 
-pub fn get_query_param(session: &Session, key: &str) -> Option<String> {
+pub fn get_query_param(session: &mut Session, key: &str) -> Option<String> {
     let query = session.req_header().uri.query()?;
 
     let params: HashMap<_, _> = query
@@ -87,3 +226,22 @@ pub fn get_query_param(session: &Session, key: &str) -> Option<String> {
         .collect();
     params.get(key).and_then(|v| decode(v).ok()).map(|s| s.to_string())
 }
+
+fn split_host_port(addr: &str, tls: bool) -> Option<(&str, u16, bool, &str)> {
+    match addr.split_once(':') {
+        Some((h, p)) => match p.parse::<u16>() {
+            Ok(port) => return Some((h, port, tls, h)),
+            Err(_) => {
+                log::warn!("ForwardAuth: invalid port in {}", addr);
+                return None;
+            }
+        },
+        None => {
+            if tls {
+                return Some((addr, 443u16, tls, addr));
+            } else {
+                return Some((addr, 80u16, tls, addr));
+            }
+        }
+    };
+}

src/utils/jwt.rs

@@ -1,16 +1,87 @@
+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>,
 }
-pub fn check_jwt(input: &str, secret: &str) -> bool {
+
-    let validation = Validation::new(Algorithm::HS256);
+#[derive(Debug, Deserialize)]
-    let token_data = decode::<Claims>(&input, &DecodingKey::from_secret(secret.as_ref()), &validation);
+struct Expired {
-    match token_data {
+    exp: Option<u64>,
-        Ok(_) => true,
+}
+
+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());
+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);
+    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 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()
+}

src/utils/structs.rs

@@ -76,7 +76,7 @@ pub struct HostConfig {
 pub struct Auth {
     #[serde(rename = "type")]
     pub auth_type: String,
-    #[serde(rename = "creds")]
+    #[serde(rename = "data")]
     pub auth_cred: String,
 }
 #[derive(Debug, Default, Serialize, Deserialize)]

src/web/proxyhttp.rs

@@ -89,7 +89,7 @@ impl ProxyHttp for LB {
                     None => return Ok(false),
                     Some(ref innermap) => {
                         if let Some(auth) = _ctx.extraparams.authentication.as_ref().or(innermap.authorization.as_ref()) {
-                            if !authenticate(&auth.auth_type, &auth.auth_cred, &session) {
+                            if !authenticate(&auth.auth_type, &auth.auth_cred, session).await {
                                 let _ = session.respond_error(401).await;
                                 warn!("Forbidden: {:?}, {}", session.client_addr(), session.req_header().uri.path());
                                 return Ok(true);

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) => {