102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
"""WebSocket routes for real-time updates."""
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
|
|
from uuid import UUID
|
|
from typing import Optional
|
|
import logging
|
|
|
|
from app.database import SessionLocal
|
|
from app.websocket.connection_manager import manager
|
|
from app.services.map_share_service import check_map_access
|
|
from app.dependencies import get_user_from_token
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(tags=["websocket"])
|
|
|
|
|
|
@router.websocket("/ws/maps/{map_id}")
|
|
async def websocket_endpoint(
|
|
websocket: WebSocket,
|
|
map_id: UUID,
|
|
token: Optional[str] = Query(None),
|
|
share_token: Optional[str] = Query(None)
|
|
):
|
|
"""
|
|
WebSocket endpoint for real-time updates.
|
|
|
|
Clients can connect using:
|
|
- JWT token (authenticated users)
|
|
- Share token (guest access)
|
|
|
|
Example: ws://localhost:8000/ws/maps/{map_id}?token={jwt_token}
|
|
Example: ws://localhost:8000/ws/maps/{map_id}?share_token={share_token}
|
|
"""
|
|
# Accept the connection first
|
|
await websocket.accept()
|
|
|
|
# Create a temporary DB session just for authentication
|
|
# This session will be closed immediately after checking access
|
|
db = SessionLocal()
|
|
try:
|
|
# Verify access to the map
|
|
user = None
|
|
if token:
|
|
try:
|
|
user = get_user_from_token(token, db)
|
|
except Exception as e:
|
|
logger.error(f"Invalid token: {e}")
|
|
await websocket.close(code=1008, reason="Invalid token")
|
|
return
|
|
|
|
# Check map access
|
|
has_access, permission = check_map_access(db, map_id, user, share_token)
|
|
if not has_access:
|
|
await websocket.close(code=1008, reason="Access denied")
|
|
return
|
|
|
|
# Store permission for the connection
|
|
permission_value = permission.value
|
|
finally:
|
|
# CRITICAL: Close the DB session immediately after authentication
|
|
db.close()
|
|
|
|
# Add to connection manager (don't call accept again, connect method will accept)
|
|
# Note: We need to call connect but it will try to accept again, so we skip it
|
|
# Instead, manually add the connection
|
|
user_id = user.id if user else None
|
|
map_key = str(map_id)
|
|
|
|
if map_key not in manager.active_connections:
|
|
manager.active_connections[map_key] = set()
|
|
|
|
manager.active_connections[map_key].add((websocket, user_id))
|
|
manager.websocket_to_map[websocket] = map_key
|
|
|
|
user_info = f"user {user_id}" if user_id else "guest"
|
|
logger.info(f"Client ({user_info}) connected to map {map_id}. Total connections: {len(manager.active_connections[map_key])}")
|
|
|
|
try:
|
|
# Send initial connection message
|
|
await websocket.send_json({
|
|
"type": "connected",
|
|
"data": {
|
|
"map_id": str(map_id),
|
|
"permission": permission_value
|
|
}
|
|
})
|
|
|
|
# Keep connection alive and listen for messages
|
|
while True:
|
|
# Receive messages (for potential future use like cursor position, etc.)
|
|
data = await websocket.receive_json()
|
|
|
|
# Echo back for now (can add more features later)
|
|
logger.info(f"Received message from client: {data}")
|
|
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(websocket, map_id)
|
|
logger.info(f"Client disconnected from map {map_id}")
|
|
except Exception as e:
|
|
logger.error(f"WebSocket error: {e}")
|
|
manager.disconnect(websocket, map_id)
|