refresh

src/main.rs

@@ -10,6 +10,7 @@ use axum::{
 };
 use dotenvy::dotenv;
 use futures::{sink::SinkExt, stream::StreamExt};
+use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
 use serde::{Deserialize, Serialize};
 use std::{collections::HashMap, sync::Arc};
 use tokio::sync::{RwLock, mpsc};
@@ -88,15 +89,22 @@ impl ChatState {
 // Application state
 struct AppState {
     chat: ChatState,
+    jwks: RwLock<HashMap<String, (String, String)>>,
 }
 
 #[tokio::main]
 async fn main() {
     dotenv().ok();
 
-    let app_state = Arc::new(AppState {
-        chat: ChatState::new(),
-    });
+    let chat = ChatState::new();
+    let jwks = RwLock::new(HashMap::new());
+
+    let app_state = Arc::new(AppState { chat, jwks });
+
+    // Initial JWKS fetch
+    if let Err(e) = fetch_jwks(&app_state).await {
+        eprintln!("Warning: Failed to fetch initial JWKS: {}", e);
+    }
 
     // Build application with routes
     let app = Router::new()
@@ -113,6 +121,34 @@ async fn main() {
     axum::serve(listener, app).await.unwrap();
 }
 
+async fn fetch_jwks(state: &AppState) -> Result<(), Box<dyn std::error::Error>> {
+    let resp: JwkSet = reqwest::get("https://idm.flegr.me/application/o/chat/jwks/")
+        .await?
+        .json()
+        .await?;
+
+    let mut keys = state.jwks.write().await;
+    for jwk in resp.keys {
+        if let (Some(kid), Some(n), Some(e)) = (jwk.kid, jwk.n, jwk.e) {
+            keys.insert(kid, (n, e));
+        }
+    }
+    println!("Fetched {} public keys from Authentik", keys.len());
+    Ok(())
+}
+
+#[derive(Deserialize)]
+struct JwkSet {
+    keys: Vec<Jwk>,
+}
+
+#[derive(Deserialize)]
+struct Jwk {
+    kid: Option<String>,
+    n: Option<String>,
+    e: Option<String>,
+}
+
 // Handlers
 
 async fn index() -> impl IntoResponse {
@@ -135,26 +171,59 @@ async fn websocket_handler(
     Query(params): Query<WsParams>,
     State(state): State<Arc<AppState>>,
 ) -> impl IntoResponse {
-    // For now, we'll verify the token by calling the userinfo endpoint.
-    // In a production app, you should verify the JWT signature locally using JWKS.
-    let client = reqwest::Client::new();
-    let user_info_resp = client
-        .get("https://idm.flegr.me/application/o/userinfo/")
-        .header("User-Agent", "axum-chat-app")
-        .bearer_auth(&params.token)
-        .send()
-        .await;
+    let header = match decode_header(&params.token) {
+        Ok(h) => h,
+        Err(_) => return StatusCode::BAD_REQUEST.into_response(),
+    };
 
-    match user_info_resp {
-        Ok(resp) if resp.status().is_success() => {
-            let user: User = match resp.json().await {
-                Ok(u) => u,
-                Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
-            };
-            ws.on_upgrade(move |socket| handle_websocket(socket, state, user))
+    let kid = match header.kid {
+        Some(k) => k,
+        None => return StatusCode::UNAUTHORIZED.into_response(),
+    };
+
+    let components = {
+        let jwks = state.jwks.read().await;
+        if let Some(c) = jwks.get(&kid).cloned() {
+            c
+        } else {
+            // Key not found, might need to refresh JWKS
+            drop(jwks);
+            if let Err(e) = fetch_jwks(&state).await {
+                eprintln!("Failed to refresh JWKS: {}", e);
+                return StatusCode::UNAUTHORIZED.into_response();
+            }
+            let jwks = state.jwks.read().await;
+            match jwks.get(&kid).cloned() {
+                Some(c) => c,
+                None => return StatusCode::UNAUTHORIZED.into_response(),
+            }
         }
-        _ => StatusCode::UNAUTHORIZED.into_response(),
-    }
+    };
+
+    let key = match DecodingKey::from_rsa_components(&components.0, &components.1) {
+        Ok(k) => k,
+        Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
+    };
+
+    let mut validation = Validation::new(Algorithm::RS256);
+    validation.set_audience(&["xLL8a23JBnBvEarjQUY6Jgd4zLnnqIi2y3oB6laM"]);
+    validation.set_issuer(&["https://idm.flegr.me/application/o/chat/"]);
+
+    let token_data = match decode::<Claims>(&params.token, &key, &validation) {
+        Ok(c) => c,
+        Err(e) => {
+            eprintln!("Token validation failed: {}", e);
+            return StatusCode::UNAUTHORIZED.into_response();
+        }
+    };
+
+    let user = User {
+        login: token_data.claims.preferred_username,
+        avatar_url: String::new(),
+    };
+
+    let state_clone = state.clone();
+    ws.on_upgrade(move |socket| handle_websocket(socket, state_clone, user))
 }
 
 // WebSocket connection handler