From c2f52d075919cf063986b2f63ced833ce3ee702d Mon Sep 17 00:00:00 2001
From: shaohuzhang1 <80892890+shaohuzhang1@users.noreply.github.com>
Date: Tue, 15 Jul 2025 13:49:33 +0800
Subject: [PATCH] fix: Access can only be granted after OpenAPI authentication
 (#3599)

---
 apps/common/init/init_doc.py                  |  6 +--
 .../middleware/doc_headers_middleware.py      | 42 +++++++++++++++----
 apps/maxkb/settings/base.py                   | 19 +++++++--
 apps/maxkb/urls.py                            |  7 +---
 apps/users/views/login.py                     |  2 +-
 5 files changed, 55 insertions(+), 21 deletions(-)
diff --git a/apps/common/init/init_doc.py b/apps/common/init/init_doc.py
index 559d4bbbc..e77a52cbf 100644
--- a/apps/common/init/init_doc.py
+++ b/apps/common/init/init_doc.py
@@ -18,15 +18,14 @@ chat_api_prefix = CONFIG.get_chat_path()[1:] + '/api/'
 
 def init_app_doc(system_urlpatterns):
     system_urlpatterns += [
-        path('schema/', SpectacularAPIView.as_view(), name='schema'),  # schema的配置文件的路由，下面两个ui也是根据这个配置文件来生成的
+        path('doc/schema/', SpectacularAPIView.as_view(), name='schema'),  # schema的配置文件的路由，下面两个ui也是根据这个配置文件来生成的
         path('doc/', SpectacularSwaggerView.as_view(url_name='schema'), name='swagger-ui'),  # swagger-ui的路由
-        path('redoc/', SpectacularRedocView.as_view(url_name='schema'), name='redoc'),  # redoc的路由
     ]
 
 
 def init_chat_doc(system_urlpatterns, chat_urlpatterns):
     system_urlpatterns += [
-        path('doc_chat_schema/',
+        path('doc_chat/schema/',
              SpectacularAPIView.as_view(patterns=[
                  URLPattern(pattern=f'{chat_api_prefix}{str(url.pattern)}', callback=url.callback,
                             default_args=url.default_args,
@@ -34,7 +33,6 @@ def init_chat_doc(system_urlpatterns, chat_urlpatterns):
                  ['chat', 'open', 'profile'].__contains__(url.name)]),
              name='chat_schema'),  # schema的配置文件的路由，下面两个ui也是根据这个配置文件来生成的
         path('doc_chat/', SpectacularSwaggerView.as_view(url_name='chat_schema'), name='swagger-ui'),  # swagger-ui的路由
-        path('redoc_chat/', SpectacularRedocView.as_view(url_name='chat_schema'), name='redoc'),  # redoc的路由
     ]
 
 
diff --git a/apps/common/middleware/doc_headers_middleware.py b/apps/common/middleware/doc_headers_middleware.py
index 4024e6ddb..67b32023e 100644
--- a/apps/common/middleware/doc_headers_middleware.py
+++ b/apps/common/middleware/doc_headers_middleware.py
@@ -6,9 +6,13 @@
     @date：2024/3/13 18:26
     @desc:
 """
+
 from django.http import HttpResponse
 from django.utils.deprecation import MiddlewareMixin
 
+from common.auth import TokenDetails, handles
+from maxkb.const import CONFIG
+
 content = """
 <!doctype html>
 <html lang="en">
@@ -18,18 +22,28 @@ content = """
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>Document</title>
     <script>
+    function setCookie(name, value, days) {
+    var expires = "";
+    if (days) {
+        var date = new Date();
+        date.setTime(date.getTime() + (days*2));
+        expires = "; expires=" + date.toUTCString();
+    }
+    document.cookie = name + "=" + (value || "")  + expires + "; path=/";
+}
       window.onload = () => {
         var xhr = new XMLHttpRequest()
-        xhr.open('GET', '/api/user', true)
+        xhr.open('GET', '/api/user/profile', true)
 
         xhr.setRequestHeader('Content-Type', 'application/json')
         const token = localStorage.getItem('token')
         const pathname = window.location.pathname
         if (token) {
-          xhr.setRequestHeader('Authorization', token)
+          xhr.setRequestHeader('Authorization', 'Bearer '+token)
           xhr.onreadystatechange = function () {
             if (xhr.readyState === 4) {
               if (xhr.status === 200) {
+                setCookie("Authorization",'Bearer '+token)
                 window.location.href = pathname
               }
               if (xhr.status === 401) {
@@ -48,15 +62,27 @@ content = """
   <body></body>
 </html>
 
-"""
+""".replace("/api/user/profile", CONFIG.get_admin_path() + '/api/user/profile').replace('/admin/login',
+                                                                                        CONFIG.get_admin_path() + '/login')
 
 
 class DocHeadersMiddleware(MiddlewareMixin):
     def process_response(self, request, response):
-        if request.path.startswith('/doc/') or request.path.startswith('/doc/chat/'):
-            HTTP_REFERER = request.META.get('HTTP_REFERER')
-            if HTTP_REFERER is None:
+        if request.path.startswith('/doc/') or request.path.startswith('/doc_chat/'):
+            auth = request.COOKIES.get('Authorization')
+            if auth is None:
                 return HttpResponse(content)
-            if HTTP_REFERER == request._current_scheme_host + request.path:
-                return response
+            else:
+                if not auth.startswith("Bearer "):
+                    return HttpResponse(content)
+                try:
+                    token = auth[7:]
+                    token_details = TokenDetails(token)
+                    for handle in handles:
+                        if handle.support(request, token, token_details.get_token_details):
+                            handle.handle(request, token, token_details.get_token_details)
+                            return response
+                    return HttpResponse(content)
+                except Exception as e:
+                    return HttpResponse(content)
         return response
diff --git a/apps/maxkb/settings/base.py b/apps/maxkb/settings/base.py
index 24a7a2bd3..0255827e3 100644
--- a/apps/maxkb/settings/base.py
+++ b/apps/maxkb/settings/base.py
@@ -100,6 +100,19 @@ TEMPLATES = [
          ],
      },
      },
+    {"NAME": "DOC",
+     'BACKEND': 'django.template.backends.django.DjangoTemplates',
+     'DIRS': ["apps/static/drf_spectacular_sidecar"],
+     'APP_DIRS': True,
+     'OPTIONS': {
+         'context_processors': [
+             'django.template.context_processors.debug',
+             'django.template.context_processors.request',
+             'django.contrib.auth.context_processors.auth',
+             'django.contrib.messages.context_processors.messages',
+         ],
+     },
+     },
 ]
 SPECTACULAR_SETTINGS = {
     'TITLE': 'MaxKB API',
@@ -107,9 +120,9 @@ SPECTACULAR_SETTINGS = {
     'VERSION': 'v2',
     'SERVE_INCLUDE_SCHEMA': False,
     # OTHER SETTINGS
-    'SWAGGER_UI_DIST': 'SIDECAR',  # shorthand to use the sidecar instead
-    'SWAGGER_UI_FAVICON_HREF': 'SIDECAR',
-    'REDOC_DIST': 'SIDECAR',
+    'SWAGGER_UI_DIST': '/doc/swagger-ui-dist',  # shorthand to use the sidecar instead
+    'SWAGGER_UI_FAVICON_HREF': '/doc/swagger-ui-dist/favicon-32x32.png',
+    'REDOC_DIST': '/doc/redoc',
     'SECURITY_DEFINITIONS': {
         'Bearer': {
             'type': 'apiKey',
diff --git a/apps/maxkb/urls.py b/apps/maxkb/urls.py
index b3fe01231..16e2bceed 100644
--- a/apps/maxkb/urls.py
+++ b/apps/maxkb/urls.py
@@ -49,15 +49,12 @@ urlpatterns = [
     path(f'{chat_ui_prefix[1:]}/', include('oss.retrieval_urls')),
 ]
 init_doc(urlpatterns, chat_urlpatterns)
-urlpatterns.append(
-    re_path(r'^static/(?P<path>.*)$', static.serve, {'document_root': settings.STATIC_ROOT}, name='static'),
-)
 
 
 def pro():
-    # 暴露静态主要是swagger资源
     urlpatterns.append(
-        re_path(r'^static/(?P<path>.*)$', static.serve, {'document_root': settings.STATIC_ROOT}, name='static'),
+        re_path(r'^doc/(?P<path>.*)$', static.serve,
+                {'document_root': os.path.join(settings.STATIC_ROOT, "drf_spectacular_sidecar")}, name='doc'),
     )
     # 暴露ui静态资源
     urlpatterns.append(
diff --git a/apps/users/views/login.py b/apps/users/views/login.py
index f8c7b1c51..4bdbba1f0 100644
--- a/apps/users/views/login.py
+++ b/apps/users/views/login.py
@@ -61,7 +61,7 @@ class Logout(APIView):
          get_operation_object=lambda r, k: {'name': r.user.username})
     def post(self, request: Request):
         version, get_key = Cache_Version.TOKEN.value
-        cache.delete(get_key(token=request.auth), version=version)
+        cache.delete(get_key(token=request.META.get('HTTP_AUTHORIZATION')[7:]), version=version)
         return result.success(True)
 
 
