Timothy J. Baek 10 meses atrás
pai
commit
ab270c1682
2 arquivos alterados com 9 adições e 6 exclusões
  1. 2 0
      backend/main.py
  2. 7 6
      backend/utils/utils.py

+ 2 - 0
backend/main.py

@@ -325,6 +325,7 @@ class ChatCompletionMiddleware(BaseHTTPMiddleware):
             data = json.loads(body_str) if body_str else {}
 
             user = get_current_user(
+                request,
                 get_http_authorization_cred(request.headers.get("Authorization"))
             )
 
@@ -558,6 +559,7 @@ class PipelineMiddleware(BaseHTTPMiddleware):
             data = json.loads(body_str) if body_str else {}
 
             user = get_current_user(
+                request,
                 get_http_authorization_cred(request.headers.get("Authorization"))
             )
 

+ 7 - 6
backend/utils/utils.py

@@ -25,7 +25,6 @@ ALGORITHM = "HS256"
 ##############
 
 bearer_security = HTTPBearer(auto_error=False)
-
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
 
@@ -79,15 +78,17 @@ def get_current_user(
     request: Request,
     auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),
 ):
-    # get token from cookie
-    token = request.cookies.get("token")
-
-    if auth_token is None and token is None:
-        raise HTTPException(status_code=403, detail="Not authenticated")
+    token = None
 
     if auth_token is not None:
         token = auth_token.credentials
 
+    if token is None and "token" in request.cookies:
+        token = request.cookies.get("token")
+
+    if token is None:
+        raise HTTPException(status_code=403, detail="Not authenticated")
+
     # auth by api key
     if token.startswith("sk-"):
         return get_current_user_by_api_key(token)