json web tokens (auth)

src/websockets.rs

@@ -1,6 +1,6 @@
 use axum::{
     extract::{
-        State, WebSocketUpgrade,
+        Query, State, WebSocketUpgrade,
         ws::{Message, WebSocket},
     },
     response::Response,
@@ -10,24 +10,31 @@ use futures_util::{
     SinkExt as _, StreamExt,
     stream::{SplitSink, SplitStream},
 };
+use serde::Deserialize;
 use tokio::sync::broadcast::Receiver;
 
-use crate::{message::ChatMessage, state::AppState};
+use crate::{auth::verify_token, message::ChatMessage, state::AppState};
+
+#[derive(Debug, Deserialize)]
+pub struct WebsocketHandlerParams {
+    token: String,
+}
 
 pub async fn websocket_handler(
     upgrade: WebSocketUpgrade,
+    Query(params): Query<WebsocketHandlerParams>,
     State(state): State<AppState>,
 ) -> Response {
-    upgrade.on_upgrade(move |websocket| handler(websocket, state))
+    upgrade.on_upgrade(move |websocket| handler(websocket, state, params))
 }
 
-async fn handler(mut socket: WebSocket, state: AppState) {
-    let mut next_client_id = state.next_client_id.lock().await;
+async fn handler(mut socket: WebSocket, state: AppState, params: WebsocketHandlerParams) {
+    let Some(claims) = verify_token(&params.token) else {
+        let _ = socket.close();
+        return;
+    };
 
-    *next_client_id += 1;
-    let _ = socket.send((*next_client_id).to_string().into()).await;
-
-    drop(next_client_id);
+    let _ = socket.send((claims.user_id).to_string().into()).await;
 
     let (sender, receiver) = socket.split();