Add token verify in websocket request

Signed-off-by: YoungSoo Shin <shinys000114@gmail.com>

main/service/ws.c

@@ -6,6 +6,7 @@
 #include "esp_err.h"
 #include "esp_http_server.h"
 #include "esp_log.h"
+#include "freertos/FreeRTOS.h"
 #include "freertos/semphr.h"
 #include "nconfig.h"
 #include "pb.h"
@@ -203,13 +204,50 @@ static void uart_event_task(void* arg)
 
 static esp_err_t ws_handler(httpd_req_t* req)
 {
-    // esp_err_t err = api_auth_check(req);
-    // if (err != ESP_OK) {
-    //     return err;
-    // }
-
     if (req->method == HTTP_GET)
     {
+        ESP_LOGI(TAG, "WebSocket GET request received for URI: %s", req->uri);
+
+        char* query_str = NULL;
+        size_t query_len = httpd_req_get_url_query_len(req) + 1;
+        if (query_len > 1) {
+            query_str = malloc(query_len);
+            if (query_str == NULL) {
+                ESP_LOGE(TAG, "Failed to allocate memory for query string");
+                httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR, "Internal Server Error");
+                return ESP_FAIL;
+            }
+            if (httpd_req_get_url_query_str(req, query_str, query_len) != ESP_OK) {
+                ESP_LOGE(TAG, "Failed to get query string from URI: %s", req->uri);
+                free(query_str);
+                httpd_resp_send_err(req, HTTPD_500_INTERNAL_SERVER_ERROR, "Internal Server Error");
+                return ESP_FAIL;
+            }
+            ESP_LOGI(TAG, "Extracted query string: %s", query_str);
+        }
+
+        char token_str[TOKEN_LENGTH];
+        esp_err_t err = ESP_FAIL; // Default to fail
+
+        if (query_str) {
+            err = httpd_query_key_value(query_str, "token", token_str, sizeof(token_str));
+            free(query_str); // Free allocated query string
+        }
+
+        if (err == ESP_OK) {
+            ESP_LOGI(TAG, "Token extracted from query string, value: %s", token_str);
+            if (!auth_validate_token(token_str)) {
+                ESP_LOGW(TAG, "WebSocket connection attempt with invalid token for URI: %s", req->uri);
+                httpd_resp_send_err(req, HTTPD_401_UNAUTHORIZED, "Invalid or expired token");
+                return ESP_FAIL;
+            }
+            ESP_LOGD(TAG, "WebSocket token validated for URI: %s", req->uri);
+        } else {
+            ESP_LOGW(TAG, "Failed to extract token from query string or query string not found, error: %s", esp_err_to_name(err));
+            httpd_resp_send_err(req, HTTPD_401_UNAUTHORIZED, "Authorization token required");
+            return ESP_FAIL;
+        }
+
         ESP_LOGI(TAG, "Handshake done, the new connection was opened");
         return ESP_OK;
     }

page/src/main.js

@@ -238,14 +238,17 @@ function initialize() {
     // Always attach login form listener
     loginForm.addEventListener('submit', handleLogin);
 
-    if (checkAuth()) { // Check authentication status
-        // If authenticated, initialize main content
-        initializeMainAppContent();
-    } else {
-        // If not authenticated, show login form
+    if (!checkAuth()) { // If NOT authenticated
+        // Show login form
         loginContainer.style.setProperty('display', 'flex', 'important');
         mainContent.style.setProperty('display', 'none', 'important');
+        console.log('Not authenticated. Login form displayed. Main app content NOT initialized.');
+        return; // IMPORTANT: Stop execution here if not authenticated
     }
+
+    // If authenticated, initialize main content
+    console.log('Authenticated. Initializing main app content.');
+    initializeMainAppContent();
 }
 
 // --- Start Application ---

page/src/websocket.js

@@ -9,7 +9,7 @@
 export let websocket;
 
 // The WebSocket server address, derived from the current page's host (hostname + port).
-const gateway = `ws://${window.location.host}/ws`;
+const baseGateway = `ws://${window.location.host}/ws`;
 
 /**
  * Initializes the WebSocket connection and sets up event handlers.
@@ -19,6 +19,13 @@ const gateway = `ws://${window.location.host}/ws`;
  * @param {function} callbacks.onMessage - Called when a message is received from the server.
  */
 export function initWebSocket({onOpen, onClose, onMessage}) {
+    const token = localStorage.getItem('authToken');
+    let gateway = baseGateway;
+
+    if (token) {
+        gateway = `${baseGateway}?token=${token}`;
+    }
+
     console.log(`Trying to open a WebSocket connection to ${gateway}...`);
     websocket = new WebSocket(gateway);
     // Set binary type to arraybuffer to handle raw binary data from the UART.