From 2a18b04f7ffeaa8d151799352bae55cfb59c94ba Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Tue, 18 Mar 2025 01:46:12 +0800
Subject: [PATCH] Added authentication checks and token validation

- Added auth checks in health check logic
- Protected routes require authentication
- Validated token on app startup
- Added auth check in API interceptor
- Clear token on 401 unauthorized error
---
 lightrag_webui/src/App.tsx         |  5 +++--
 lightrag_webui/src/AppRouter.tsx   | 21 +++++++++++++++------
 lightrag_webui/src/api/lightrag.ts | 16 +++++++++++-----
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/lightrag_webui/src/App.tsx b/lightrag_webui/src/App.tsx
index 09ee6b31..102ff233 100644
--- a/lightrag_webui/src/App.tsx
+++ b/lightrag_webui/src/App.tsx
@@ -5,7 +5,7 @@ import MessageAlert from '@/components/MessageAlert'
 import ApiKeyAlert from '@/components/ApiKeyAlert'
 import StatusIndicator from '@/components/graph/StatusIndicator'
 import { healthCheckInterval } from '@/lib/constants'
-import { useBackendState } from '@/stores/state'
+import { useBackendState, useAuthStore } from '@/stores/state'
 import { useSettingsStore } from '@/stores/settings'
 import { useEffect } from 'react'
 import SiteHeader from '@/features/SiteHeader'
@@ -26,7 +26,8 @@ function App() {
 
   // Health check
   useEffect(() => {
-    if (!enableHealthCheck) return
+    const { isAuthenticated } = useAuthStore.getState();
+    if (!enableHealthCheck || !isAuthenticated) return
 
     // Check immediately
     useBackendState.getState().check()
diff --git a/lightrag_webui/src/AppRouter.tsx b/lightrag_webui/src/AppRouter.tsx
index d7a06fa8..a88dc5ea 100644
--- a/lightrag_webui/src/AppRouter.tsx
+++ b/lightrag_webui/src/AppRouter.tsx
@@ -1,5 +1,6 @@
-import { HashRouter as Router, Routes, Route } from 'react-router-dom'
-// import { useAuthStore } from '@/stores/state'
+import { HashRouter as Router, Routes, Route, Navigate } from 'react-router-dom'
+import { useEffect } from 'react'
+import { useAuthStore } from '@/stores/state'
 import { Toaster } from 'sonner'
 import App from './App'
 import LoginPage from '@/features/LoginPage'
@@ -10,16 +11,24 @@ interface ProtectedRouteProps {
 }
 
 const ProtectedRoute = ({ children }: ProtectedRouteProps) => {
-  // const { isAuthenticated } = useAuthStore()
+  const { isAuthenticated } = useAuthStore()
 
-  // if (!isAuthenticated) {
-  //   return <Navigate to="/login" replace />
-  // }
+  if (!isAuthenticated) {
+    return <Navigate to="/login" replace />
+  }
 
   return <>{children}</>
 }
 
 const AppRouter = () => {
+  // Check login at befor startup
+  useEffect(() => {
+    const token = localStorage.getItem('LIGHTRAG-API-TOKEN');
+    if (!token) {
+      useAuthStore.getState().logout();
+    }
+  }, []);
+
   return (
     <ThemeProvider>
       <Router>
diff --git a/lightrag_webui/src/api/lightrag.ts b/lightrag_webui/src/api/lightrag.ts
index f495cd28..8a71ab9a 100644
--- a/lightrag_webui/src/api/lightrag.ts
+++ b/lightrag_webui/src/api/lightrag.ts
@@ -142,10 +142,20 @@ const axiosInstance = axios.create({
   }
 })
 
-// Interceptor：add api key
+// Interceptor: add api key and check authentication
 axiosInstance.interceptors.request.use((config) => {
   const apiKey = useSettingsStore.getState().apiKey
   const token = localStorage.getItem('LIGHTRAG-API-TOKEN');
+
+  // Check authentication status for paths that require authentication
+  const authRequiredPaths = ['/documents', '/graphs', '/query', '/health']; // Add all paths that require authentication
+  const isAuthRequired = authRequiredPaths.some(path => config.url?.includes(path));
+
+  if (isAuthRequired && !token && config.url !== '/login') {
+    // Cancel the request and return a rejected Promise
+    return Promise.reject(new Error('Authentication required'));
+  }
+
   if (apiKey) {
     config.headers['X-API-Key'] = apiKey
   }
@@ -160,10 +170,6 @@ axiosInstance.interceptors.response.use(
   (response) => response,
   (error: AxiosError) => {
     if (error.response) {
-      interface ErrorResponse {
-        detail: string;
-      }
-
       if (error.response?.status === 401) {
         localStorage.removeItem('LIGHTRAG-API-TOKEN');
         sessionStorage.clear();