completed implementation of #17

Cargo.lock

@@ -3060,6 +3060,7 @@ dependencies = [
  "base64",
  "bytes",
  "encoding_rs",
+ "futures-channel",
  "futures-core",
  "futures-util",
  "h2",

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),
-        "apikey" => ApiKeyAuth(credentials).validate(session),
-        "jwt" => JwtAuth(credentials).validate(session),
+        "basic" => BasicAuth(credentials).validate(session).await,
+        "apikey" => ApiKeyAuth(credentials).validate(session).await,
+        "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,42 @@
+use ahash::AHasher;
 use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
+use moka::sync::Cache;
 use serde::{Deserialize, Serialize};
+use std::hash::{Hash, Hasher};
+use std::sync::LazyLock;
 
 #[derive(Debug, Serialize, Deserialize)]
 pub(crate) struct Claims {
     pub(crate) user: String,
     pub(crate) exp: 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);
-    match token_data {
-        Ok(_) => true,
-        Err(_) => false,
-    }
+    token_data.is_ok()
+}
+*/
+
+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 result = decode::<Claims>(token, &DecodingKey::from_secret(secret.as_ref()), &JWT_VALIDATION).is_ok();
+    if result {
+        JWT_CACHE.insert(key, true);
+    }
+    result
+}
+
+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);