68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
from fastapi import APIRouter, WebSocket, Query, WebSocketDisconnect
|
|
from jose import JWTError, jwt
|
|
from sqlalchemy.orm import Session
|
|
import logging
|
|
|
|
from app.infrastructure.realtime.connection_manager import manager
|
|
from app.core.settings import settings
|
|
from typing import Optional
|
|
|
|
from app.core.ws_auth import authenticate_websocket
|
|
from app.db.database import SessionLocal
|
|
from app.modules.auth.models.user_model import User
|
|
from app.modules.documents.models.document_model import Project
|
|
from app.core.token_blacklist import TokenBlacklist
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
def _authenticate_ws(token: str, db: Session) -> Optional[User]:
|
|
"""
|
|
Kept as a thin delegation, not deleted.
|
|
|
|
Both WebSocket modules carried their own copy of this and drifted apart;
|
|
`app/core/ws_auth.py` is now the single implementation. The name stays
|
|
because the handshake below calls it and the probes resolve it by name --
|
|
the point was to remove the second set of *rules*, not the local symbol.
|
|
"""
|
|
return authenticate_websocket(token, db)
|
|
|
|
@router.websocket("/ws/{task_id}")
|
|
async def websocket_progress(websocket: WebSocket, task_id: str, token: str = Query(None)):
|
|
auth_token = token or websocket.cookies.get("docqube_access_token")
|
|
|
|
if auth_token and TokenBlacklist.is_token_blacklisted(auth_token):
|
|
await websocket.accept()
|
|
await websocket.close(code=4001, reason="Token revoked")
|
|
return
|
|
|
|
db = SessionLocal()
|
|
try:
|
|
user = _authenticate_ws(auth_token, db) if auth_token else None
|
|
if not user:
|
|
logger.warning(f"WebSocket progress auth failed for task {task_id}")
|
|
await websocket.close(code=4001, reason="Authentication failed")
|
|
return
|
|
|
|
project = db.query(Project).filter(Project.session_id == task_id, Project.user_id == user.id).first()
|
|
if not project:
|
|
logger.warning(f"Unauthorized WebSocket progress access: User {user.id} task {task_id}")
|
|
await websocket.close(code=4003, reason="Access denied to this task")
|
|
return
|
|
|
|
user_id = user.id
|
|
finally:
|
|
db.close()
|
|
|
|
await manager.connect_task(task_id, websocket)
|
|
|
|
try:
|
|
while True:
|
|
await websocket.receive_text()
|
|
except WebSocketDisconnect:
|
|
logger.info(f"WebSocket progress disconnected: user={user_id} task={task_id}")
|
|
except Exception as e:
|
|
logger.error(f"WebSocket progress error: user={user_id} task={task_id}: {e}")
|
|
finally:
|
|
manager.disconnect(websocket) |