feat: add auth middleware

src/auth/jwt.rs

@@ -1,9 +1,10 @@
 use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
 use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::env;
+use once_cell::sync::Lazy;
 
-static ISSUER: &str = "https://auth.iceberg.black/realms/iceberg";
-
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct Claims {
     pub sub: String,
     pub preferred_username: Option<String>,
@@ -13,34 +14,52 @@ pub struct Claims {
     pub realm_access: Option<RealmAccess>,
 }
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct RealmAccess {
     pub roles: Vec<String>,
 }
 
-pub fn validate_token(token: &str, jwks: &serde_json::Value) -> Result<Claims, ()> {
-    let header = decode_header(token).map_err(|_| ())?;
+static ISSUER: Lazy<String> = Lazy::new(|| {
+    env::var("ISSUER").expect("ISSUER not set")
+});
 
-    let kid = header.kid.ok_or(())?;
+pub fn validate_token(token: &str, jwks: &Value) -> Result<Claims, String> {
+    // 1. Decode header
+    let header = decode_header(token).map_err(|_| "Invalid header")?;
 
-    let keys = jwks["keys"].as_array().ok_or(())?;
+    let kid = header.kid.ok_or("Missing kid")?;
+
+    // 2. Find matching key
+    let keys = jwks["keys"]
+        .as_array()
+        .ok_or("Invalid JWKS")?;
 
     let key = keys
         .iter()
         .find(|k| k["kid"] == kid)
-        .ok_or(())?;
+        .ok_or("Matching key not found")?;
 
-    let n = key["n"].as_str().ok_or(())?;
-    let e = key["e"].as_str().ok_or(())?;
+    // 3. Extract RSA components
+    let n = key["n"].as_str().ok_or("Missing n")?;
+    let e = key["e"].as_str().ok_or("Missing e")?;
 
     let decoding_key =
-        DecodingKey::from_rsa_components(n, e).map_err(|_| ())?;
+        DecodingKey::from_rsa_components(n, e)
+            .map_err(|_| "Invalid decoding key")?;
 
+    // 4. Setup validation rules
     let mut validation = Validation::new(Algorithm::RS256);
-    validation.set_issuer(&[ISSUER]);
 
+    validation.set_issuer(&[ISSUER.as_str()]);
+
+    // Optional but recommended:
+    validation.validate_exp = true;
+    validation.validate_aud = false; // depends on your Keycloak config
+
+    // 5. Decode & verify
     let token_data =
-        decode::<Claims>(token, &decoding_key, &validation).map_err(|_| ())?;
+        decode::<Claims>(token, &decoding_key, &validation)
+            .map_err(|_| "Token validation failed")?;
 
     Ok(token_data.claims)
 }

src/auth/keycloak.rs

@@ -14,10 +14,12 @@ struct JwksCache {
 static JWK_CACHE: Lazy<Arc<RwLock<Option<JwksCache>>>> =
     Lazy::new(|| Arc::new(RwLock::new(None)));
 
-async fn fetch_jwks() -> Result<Value, reqwest::Error> {
-    let url = env::var("JWKS_URL").expect("JWKS_URL not set");
+static JWKS_URL: Lazy<String> = Lazy::new(|| {
+    env::var("JWKS_URL").expect("JWKS_URL not set")
+});
 
-    let jwks = reqwest::get(url)
+async fn fetch_jwks() -> Result<Value, reqwest::Error> {
+    let jwks = reqwest::get(JWKS_URL.as_str())
         .await?
         .json::<Value>()
         .await?;
@@ -25,6 +27,19 @@ async fn fetch_jwks() -> Result<Value, reqwest::Error> {
     Ok(jwks)
 }
 
+pub async fn refresh_jwks() -> Result<Value, reqwest::Error> {
+    let jwks = fetch_jwks().await?;
+
+    let mut write = JWK_CACHE.write().await;
+
+    *write = Some(JwksCache {
+        jwks: jwks.clone(),
+        last_fetched: Instant::now(),
+    });
+
+    Ok(jwks)
+}
+
 pub async fn get_jwks() -> Result<Value, reqwest::Error> {
     let ttl = Duration::from_secs(3600); // 1 hour
 
@@ -40,14 +55,5 @@ pub async fn get_jwks() -> Result<Value, reqwest::Error> {
     }
 
     // Expired or empty → refresh
-    let jwks = fetch_jwks().await?;
-
-    let mut write = JWK_CACHE.write().await;
-
-    *write = Some(JwksCache {
-        jwks: jwks.clone(),
-        last_fetched: Instant::now(),
-    });
-
-    Ok(jwks)
+    refresh_jwks().await
 }

src/auth/middleware.rs

@@ -0,0 +1,54 @@
+use axum::{
+    extract::Request,
+    http::{StatusCode},
+    middleware::Next,
+    response::Response,
+};
+
+use crate::auth::{jwt::validate_token, keycloak::{get_jwks, refresh_jwks}};
+
+pub async fn auth_middleware(
+    mut request: Request,
+    next: Next,
+) -> Result<Response, StatusCode> {
+    dbg!("Middleware hit");
+
+    let headers = request.headers();
+
+    dbg!("Headers extracted");
+
+    let auth_header = headers
+        .get("authorization")
+        .and_then(|v| v.to_str().ok());
+
+    dbg!("Auth header: {:?}", auth_header);
+
+    let auth_header = auth_header.ok_or(StatusCode::UNAUTHORIZED)?;
+
+    let token = auth_header
+        .strip_prefix("Bearer ")
+        .ok_or(StatusCode::UNAUTHORIZED)?;
+
+    let jwks = get_jwks().await.map_err(|_| StatusCode::UNAUTHORIZED)?;
+
+    match validate_token(token, &jwks) {
+        Ok(claims) => {
+            dbg!("Token valid");
+
+            request.extensions_mut().insert(claims);
+
+            Ok(next.run(request).await)
+        }
+        Err(_) => {
+            let jwks = refresh_jwks().await.map_err(|_| StatusCode::UNAUTHORIZED)?;
+
+            match validate_token(token, &jwks) {
+                Ok(claims) => {
+                    request.extensions_mut().insert(claims);
+                    Ok(next.run(request).await)
+                }
+                Err(_) => Err(StatusCode::UNAUTHORIZED),
+            }
+        }
+    }
+}

src/auth/mod.rs

@@ -1,3 +1,3 @@
 pub mod jwt;
 pub mod keycloak;
-// pub mod middleware;
+pub mod middleware;