From d9dbd5fe5d4956fb78825d5c1e4b4bb18d1b81a9 Mon Sep 17 00:00:00 2001
From: Wind-Explorer <windexplorer@outlook.com>
Date: Fri, 5 Dec 2025 00:38:35 +0800
Subject: [PATCH] ws auth

---
 src-tauri/src/core/services/auth.rs |  1 +
 src-tauri/src/core/services/ws.rs   | 67 ++++++++++++++++++++++-------
 src-tauri/src/core/state.rs         | 26 +++--------
 3 files changed, 59 insertions(+), 35 deletions(-)

diff --git a/src-tauri/src/core/services/auth.rs b/src-tauri/src/core/services/auth.rs
index 7677abe..690d555 100644
--- a/src-tauri/src/core/services/auth.rs
+++ b/src-tauri/src/core/services/auth.rs
@@ -586,6 +586,7 @@ where
                             error!("Failed to save auth pass: {}", e);
                         } else {
                             info!("Authentication successful!");
+                            crate::core::services::ws::init_ws_client().await;
                             on_success();
                         }
                     }
diff --git a/src-tauri/src/core/services/ws.rs b/src-tauri/src/core/services/ws.rs
index d81ba85..955e4db 100644
--- a/src-tauri/src/core/services/ws.rs
+++ b/src-tauri/src/core/services/ws.rs
@@ -1,11 +1,11 @@
 use rust_socketio::{ClientBuilder, Payload, RawClient};
 use serde_json::json;
 use tauri::async_runtime;
-use tracing::error;
+use tracing::{error, info};
 
 use crate::{
     core::{models::app_config::AppConfig, state::FDOLL},
-    lock_r,
+    lock_r, lock_w,
     services::cursor::CursorPosition,
 };
 
@@ -45,19 +45,54 @@ pub async fn report_cursor_data(cursor_position: CursorPosition) {
     }
 }
 
-pub fn build_ws_client(app_config: &AppConfig) -> rust_socketio::client::Client {
-    let client = match ClientBuilder::new(
-        app_config
-            .api_base_url
-            .as_ref()
-            .expect("Missing API base URL"),
-    )
-    .namespace("/")
-    .on("pong", on_pong)
-    .connect()
-    {
-        Ok(c) => c,
-        Err(_) => todo!("TODO error handling"),
+pub async fn init_ws_client() {
+    let app_config = {
+        let guard = lock_r!(FDOLL);
+        guard.app_config.clone()
     };
-    client
+
+    let ws_client = build_ws_client(&app_config).await;
+
+    let mut guard = lock_w!(FDOLL);
+    if let Some(clients) = guard.clients.as_mut() {
+        clients.ws_client = Some(ws_client);
+    }
+    info!("WebSocket client initialized after authentication");
+}
+
+pub async fn build_ws_client(app_config: &AppConfig) -> rust_socketio::client::Client {
+    let token = crate::core::services::auth::get_access_token()
+        .await
+        .expect("No access token available for WebSocket connection");
+
+    info!("Building WebSocket client with authentication");
+
+    let api_base_url = app_config
+        .api_base_url
+        .clone()
+        .expect("Missing API base URL");
+
+    let client = async_runtime::spawn_blocking(move || {
+        ClientBuilder::new(api_base_url)
+            .namespace("/")
+            .on("pong", on_pong)
+            .auth(json!({ "token": token }))
+            .connect()
+    })
+    .await
+    .expect("Failed to spawn blocking task");
+
+    match client {
+        Ok(c) => {
+            info!("WebSocket client connected successfully");
+            c
+        }
+        Err(e) => {
+            error!("Failed to connect WebSocket: {:?}", e);
+            panic!(
+                "TODO: better error handling for WebSocket connection - {}",
+                e
+            );
+        }
+    }
 }
diff --git a/src-tauri/src/core/state.rs b/src-tauri/src/core/state.rs
index 8d1c23a..833a61b 100644
--- a/src-tauri/src/core/state.rs
+++ b/src-tauri/src/core/state.rs
@@ -2,10 +2,7 @@
 use crate::{
     core::{
         models::app_config::{AppConfig, AuthConfig},
-        services::{
-            auth::{load_auth_pass, AuthPass},
-            ws::build_ws_client,
-        },
+        services::auth::{load_auth_pass, AuthPass},
     },
     lock_w,
 };
@@ -78,24 +75,15 @@ pub fn init_fdoll_state() {
         });
         info!("Initialized HTTP client");
 
-        // Clone app_config for async task
-        let app_config = guard.app_config.clone();
+        let has_auth = guard.auth_pass.is_some();
 
-        // Drop the write lock before spawning async task
         drop(guard);
 
-        // Initialize WebSocket client in a blocking task to avoid runtime conflicts
-        async_runtime::spawn(async move {
-            let ws_client = async_runtime::spawn_blocking(move || build_ws_client(&app_config))
-                .await
-                .expect("Failed to initialize WebSocket client");
-
-            let mut guard = lock_w!(FDOLL);
-            if let Some(clients) = guard.clients.as_mut() {
-                clients.ws_client = Some(ws_client);
-            }
-            info!("Initialized FDOLL state with WebSocket client");
-        });
+        if has_auth {
+            async_runtime::spawn(async move {
+                crate::core::services::ws::init_ws_client().await;
+            });
+        }
 
         info!("Initialized FDOLL state (WebSocket client initializing asynchronously)");
     }