fix(server): scope auth DB sessions to prevent connection-pool exhaustion (#6237)
This commit is contained in:
+14
-9
@@ -3,7 +3,7 @@ import secrets
|
||||
import uuid
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from db import get_db
|
||||
from db import SessionLocal
|
||||
from fastapi import Depends, HTTPException, Request
|
||||
from fastapi.security import APIKeyHeader, HTTPAuthorizationCredentials, HTTPBearer
|
||||
from jose import JWTError, jwt
|
||||
@@ -145,19 +145,24 @@ async def verify_auth(
|
||||
request: Request,
|
||||
credentials: HTTPAuthorizationCredentials | None = Depends(bearer_scheme),
|
||||
x_api_key: str | None = Depends(api_key_header),
|
||||
db: Session = Depends(get_db),
|
||||
) -> User | None:
|
||||
"""Authenticate via JWT, X-API-Key, or legacy ADMIN_API_KEY. Returns User or None."""
|
||||
"""Authenticate via JWT, X-API-Key, or legacy ADMIN_API_KEY. Returns User or None.
|
||||
|
||||
A short-lived session is opened only on the branches that query the DB, so no
|
||||
pooled connection is held for the lifetime of the (possibly long-running) request.
|
||||
"""
|
||||
if credentials is not None:
|
||||
_mark_auth_type(request, "bearer")
|
||||
return _resolve_user_from_jwt(credentials.credentials, db)
|
||||
with SessionLocal() as db:
|
||||
return _resolve_user_from_jwt(credentials.credentials, db)
|
||||
|
||||
if x_api_key is not None:
|
||||
if ADMIN_API_KEY and secrets.compare_digest(x_api_key, ADMIN_API_KEY):
|
||||
_mark_auth_type(request, "admin_api_key")
|
||||
return None
|
||||
_mark_auth_type(request, "api_key")
|
||||
return _resolve_user_from_api_key(x_api_key, db)
|
||||
with SessionLocal() as db:
|
||||
return _resolve_user_from_api_key(x_api_key, db)
|
||||
|
||||
if AUTH_DISABLED:
|
||||
_mark_auth_type(request, "disabled")
|
||||
@@ -173,12 +178,12 @@ async def verify_auth(
|
||||
async def require_auth(
|
||||
request: Request,
|
||||
user: User | None = Depends(verify_auth),
|
||||
db: Session = Depends(get_db),
|
||||
) -> User:
|
||||
"""Like verify_auth but guarantees a non-None User. Use for endpoints that require auth."""
|
||||
if user is None:
|
||||
if getattr(request.state, "auth_type", "none") in {"admin_api_key", "disabled"}:
|
||||
default_user = _get_default_user(db)
|
||||
with SessionLocal() as db:
|
||||
default_user = _get_default_user(db)
|
||||
if default_user is not None:
|
||||
return default_user
|
||||
raise HTTPException(status_code=401, detail="Authentication required.")
|
||||
@@ -193,7 +198,6 @@ _BOOTSTRAP_ADMIN = User(
|
||||
async def require_admin(
|
||||
request: Request,
|
||||
user: User | None = Depends(verify_auth),
|
||||
db: Session = Depends(get_db),
|
||||
) -> User:
|
||||
"""Like require_auth but also enforces admin role.
|
||||
|
||||
@@ -203,7 +207,8 @@ async def require_admin(
|
||||
auth_type = getattr(request.state, "auth_type", "none")
|
||||
if user is None:
|
||||
if auth_type in {"admin_api_key", "disabled"}:
|
||||
default_user = _get_default_user(db)
|
||||
with SessionLocal() as db:
|
||||
default_user = _get_default_user(db)
|
||||
if default_user is not None:
|
||||
if default_user.role != "admin":
|
||||
raise HTTPException(status_code=403, detail="Admin role required.")
|
||||
|
||||
Reference in New Issue
Block a user