feat: allow auth middleware to accept access_token query string

Author: deepworklab

chat-core/src/middlewares/auth.rs

@@ -1,5 +1,5 @@
 use axum::{
-    extract::{FromRequestParts, Request, State},
+    extract::{FromRequestParts, Query, Request, State},
     http::StatusCode,
     middleware::Next,
     response::{IntoResponse, Response},
@@ -8,39 +8,51 @@ use axum_extra::{
     TypedHeader,
     headers::{Authorization, authorization::Bearer},
 };
+use serde::{Deserialize, Serialize};
 use tracing::warn;
 
 use crate::middlewares::TokenVerify;
-
+#[derive(Debug, Deserialize, Serialize)]
+struct Params {
+    access_token: String,
+}
 pub async fn verify_token<T>(State(state): State<T>, request: Request, next: Next) -> Response
 where
     T: TokenVerify + Clone + Send + Sync + 'static,
 {
     let (mut parts, body) = request.into_parts();
-    let req =
+    let token =
         match TypedHeader::<Authorization<Bearer>>::from_request_parts(&mut parts, &state).await {
-            Ok(TypedHeader(Authorization(bearer))) => {
-                let token = bearer.token();
-                match state.verify(token) {
-                    Ok(claims) => {
-                        let mut req = Request::from_parts(parts, body);
-                        req.extensions_mut().insert(claims);
-                        req
-                    }
-                    Err(err) => {
-                        let msg = format!("verify token failed:{:?}", err);
-                        warn!(msg);
-                        return (StatusCode::FORBIDDEN, msg).into_response();
+            Ok(TypedHeader(Authorization(bearer))) => bearer.token().to_string(),
+            Err(err) => {
+                if err.is_missing() {
+                    match Query::<Params>::from_request_parts(&mut parts, &state).await {
+                        Ok(params) => params.access_token.to_string(),
+                        Err(err) => {
+                            let msg = format!("parse query failed:{}", err);
+                            warn!(msg);
+                            return (StatusCode::UNAUTHORIZED, msg).into_response();
+                        }
                     }
+                } else {
+                    let msg = format!("parse Authorization header failed: {}", err);
+                    warn!(msg);
+                    return (StatusCode::UNAUTHORIZED, msg).into_response();
                 }
             }
-            Err(err) => {
-                let msg = format!("parse Authorization header failed:{}", err);
-                warn!(msg);
-                return (StatusCode::UNAUTHORIZED, msg).into_response();
-            }
         };
-
+    let req = match state.verify(&token) {
+        Ok(claims) => {
+            let mut req = Request::from_parts(parts, body);
+            req.extensions_mut().insert(claims);
+            req
+        }
+        Err(err) => {
+            let msg = format!("verify token failed:{:?}", err);
+            warn!(msg);
+            return (StatusCode::FORBIDDEN, msg).into_response();
+        }
+    };
     next.run(req).await
 }
 
@@ -100,6 +112,11 @@ mod tests {
             .layer(from_fn_with_state(state.clone(), verify_token::<AppState>))
             .with_state(state);
 
+        // no token
+        let req = Request::builder().uri("/").body(Body::empty())?;
+        let res = app.clone().oneshot(req).await?;
+        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
+
         // good token
         let req = Request::builder()
             .uri("/")
@@ -108,17 +125,26 @@ mod tests {
         let res = app.clone().oneshot(req).await?;
         assert_eq!(res.status(), StatusCode::OK);
 
-        // no token
-        let req = Request::builder().uri("/").body(Body::empty())?;
-        let res = app.clone().oneshot(req).await?;
-        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
-
         // bad token
         let req = Request::builder()
             .uri("/")
             .header("Authorization", "Bearer bad-token")
             .body(Body::empty())?;
-        let res = app.oneshot(req).await?;
+        let res = app.clone().oneshot(req).await?;
+        assert_eq!(res.status(), StatusCode::FORBIDDEN);
+
+        // good token in query
+        let req = Request::builder()
+            .uri(format!("/?access_token={token}"))
+            .body(Body::empty())?;
+        let res = app.clone().oneshot(req).await?;
+        assert_eq!(res.status(), StatusCode::OK);
+
+        // bad token
+        let req = Request::builder()
+            .uri("/?access_token=bad-token".to_string())
+            .body(Body::empty())?;
+        let res = app.clone().oneshot(req).await?;
         assert_eq!(res.status(), StatusCode::FORBIDDEN);
 
         Ok(())

notify-server/index.html

@@ -23,7 +23,9 @@
 		<h1>Server Sent Events</h1>
 
 		<script lang="javascript">
-			var source = new EventSource("/events");
+			let token =
+				"eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJpYXQiOjE3Nzc3MTUwNTksImV4cCI6MTc3ODMxOTg1OSwibmJmIjoxNzc3NzE1MDU5LCJpc3MiOiJjaGF0X3NlcnZlciIsImF1ZCI6ImNoYXRfd2ViIiwiaWQiOjEsIndzX2lkIjoxLCJmdWxsbmFtZSI6IlR5ciBDaGVuIiwiZW1haWwiOiJ0Y2hlbkBhY21lLm9yZyIsImNyZWF0ZWRfYXQiOiIyMDI2LTA1LTAyVDA3OjI0OjQ5Ljc1Mjc3MVoifQ.TEMT5zhy-O0nx8fpll3FRCJXvkJi6QUXWU-N0Gh73El-Ka7kkf7oao-Zwd8URPB4Piv8g03QIRjAfGMH6tdYDA";
+			var source = new EventSource(`/events?access_token=${token}`);
 			source.onmessage = function (event) {
 				console.log("Got:", event.data);
 			};

notify-server/src/lib.rs

@@ -5,8 +5,15 @@ mod sse;
 use std::{ops::Deref, sync::Arc};
 
 use anyhow::Result;
-use axum::response::{Html, IntoResponse};
-use chat_core::{User, middlewares::TokenVerify, utils::DecodingKey};
+use axum::{
+    middleware::from_fn_with_state,
+    response::{Html, IntoResponse},
+};
+use chat_core::{
+    User,
+    middlewares::{TokenVerify, verify_token},
+    utils::DecodingKey,
+};
 use dashmap::DashMap;
 pub use notify::setup_pg_listener;
 use tokio::sync::broadcast;
@@ -33,7 +40,7 @@ pub fn get_router() -> (axum::Router, AppState) {
     let state = AppState::new(config);
     let app = axum::Router::new()
         .route("/events", axum::routing::get(sse::sse_handler))
-        //.layer(from_fn_with_state(state.clone(), verify_token::<AppState>))
+        .layer(from_fn_with_state(state.clone(), verify_token::<AppState>))
         .route("/index", axum::routing::get(index_handler))
         .with_state(state.clone());
     (app, state)

notify-server/src/notify.rs

@@ -34,7 +34,7 @@ struct ChatUpdated {
 // pg_notify('chat_message_created', row_to_json(NEW)::text);
 #[derive(Debug, Serialize, Deserialize)]
 struct ChatMessageCreated {
-    messages: Message,
+    message: Message,
     members: Vec<i64>,
 }
 pub async fn setup_pg_listener(state: AppState) -> anyhow::Result<()> {
@@ -94,7 +94,7 @@ impl Notification {
                 let user_ids = payload.members.iter().copied().collect();
                 Ok(Self {
                     user_ids,
-                    event: Arc::new(AppEvent::NewMessage(payload.messages)),
+                    event: Arc::new(AppEvent::NewMessage(payload.message)),
                 })
             }
             _ => Err(anyhow::anyhow!("Invalid notification type")),