add token

qwen_server/database_server.py

@@ -76,26 +76,13 @@ def change_checkbox_state(text, cache_file):
     return {'result': 'changed'}
 
 
-class BasicAuthMiddleware(BaseHTTPMiddleware):
-    def __init__(self, app, token: str):
-        super().__init__(app)
-        self.required_credentials = token
-
-    async def dispatch(self, request: Request, call_next):
-        authorization: str = request.headers.get("Authorization")
-        if authorization:
-            try:
-                schema, credentials = authorization.split()
-                if credentials == self.required_credentials:
-                    return await call_next(request)
-            except ValueError:
-                pass
-
+@app.middleware("http")
+async def validate_token(request: Request, call_next):
+    access_token: str = request.headers.get("Authorization") or request.query_params.get("access_token")
+    if not access_token or os.getenv("ACCESS_TOKEN") != access_token:
         headers = {'WWW-Authenticate': 'Basic'}
         return Response(status_code=401, headers=headers)
-
-
-app.add_middleware(BasicAuthMiddleware, token=os.getenv("ACCESS_TOKEN"))
+    return await call_next(request)
 
 
 @app.get('/healthz')
@@ -105,8 +92,6 @@ async def healthz(request: Request):
 
 @app.get('/cachedata/{file_name}')
 async def cache_data(file_name: str, access_token: str):
-    if not access_token or os.getenv("ACCESS_TOKEN") != access_token:
-        raise HTTPException(401, "the access token is not valid")
     cache_file = os.path.join(server_config.path.cache_root, file_name)
     lines = []
     for line in jsonlines.open(cache_file):
@@ -117,9 +102,6 @@ async def cache_data(file_name: str, access_token: str):
 @app.post('/endpoint')
 async def web_listening(request: Request):
     data = await request.json()
-    access_token: str = request.headers.get("Authorization")
-    if not access_token or os.getenv("ACCESS_TOKEN") != access_token:
-        raise HTTPException(401, "the access token is not valid")
     msg_type = data['task']
 
     cache_file_popup_url = os.path.join(server_config.path.cache_root, 'popup_url.jsonl')