163 lines
6.2 KiB
Python
163 lines
6.2 KiB
Python
"""WebSocket API endpoints"""
|
|
|
|
import asyncio
|
|
from datetime import UTC, datetime
|
|
from typing import Optional
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
|
|
from jose import jwt, JWTError
|
|
|
|
from app.core.config import settings
|
|
from app.core.logging import get_logger
|
|
from app.core.time import to_iso8601_utc
|
|
from app.core.websocket.manager import manager
|
|
|
|
logger = get_logger(__name__, service="api")
|
|
router = APIRouter()
|
|
EARTH_UPDATES_CHANNEL = "earth_updates"
|
|
|
|
|
|
async def authenticate_token(token: str) -> Optional[dict]:
|
|
"""Authenticate WebSocket connection via token"""
|
|
try:
|
|
payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
|
|
if payload.get("type") != "access":
|
|
logger.warning_event(
|
|
"WebSocket auth failed: wrong token type",
|
|
event="auth.websocket.invalid_token_type",
|
|
)
|
|
return None
|
|
return payload
|
|
except JWTError as e:
|
|
logger.warning_event(
|
|
"WebSocket auth failed",
|
|
event="auth.websocket.decode_failed",
|
|
context={"error": str(e)},
|
|
)
|
|
return None
|
|
|
|
|
|
@router.websocket("/ws")
|
|
async def websocket_endpoint(
|
|
websocket: WebSocket,
|
|
token: str | None = Query(None),
|
|
):
|
|
"""WebSocket endpoint for real-time data"""
|
|
logger.info_event(
|
|
"WebSocket connection attempt",
|
|
event="auth.websocket.connection_attempt",
|
|
context={"token_preview": f"{token[:8]}..." if token else "anonymous"},
|
|
)
|
|
payload = await authenticate_token(token) if token else None
|
|
if token and payload is None:
|
|
logger.warning_event(
|
|
"WebSocket authentication failed, closing connection",
|
|
event="auth.websocket.connection_rejected",
|
|
)
|
|
await websocket.close(code=4001)
|
|
return
|
|
|
|
is_anonymous = payload is None
|
|
user_id = str(payload.get("sub")) if payload else f"anonymous:{id(websocket)}"
|
|
supported_channels = ["vessels", "earth_news", EARTH_UPDATES_CHANNEL] if is_anonymous else [
|
|
"gpu_clusters",
|
|
"submarine_cables",
|
|
"ixp_nodes",
|
|
"alerts",
|
|
"dashboard",
|
|
"datasource_tasks",
|
|
"vessels",
|
|
"earth_news",
|
|
EARTH_UPDATES_CHANNEL,
|
|
]
|
|
await manager.connect(websocket, user_id)
|
|
|
|
try:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "connection_established",
|
|
"data": {
|
|
"connection_id": f"conn_{user_id}",
|
|
"server_version": settings.VERSION,
|
|
"heartbeat_interval": 30,
|
|
"supported_channels": supported_channels,
|
|
},
|
|
}
|
|
)
|
|
|
|
while True:
|
|
try:
|
|
data = await asyncio.wait_for(websocket.receive_json(), timeout=30)
|
|
|
|
if data.get("type") == "heartbeat":
|
|
await websocket.send_json(
|
|
{
|
|
"type": "heartbeat",
|
|
"data": {"action": "pong", "timestamp": to_iso8601_utc(datetime.now(UTC))},
|
|
}
|
|
)
|
|
elif data.get("type") == "subscribe":
|
|
payload_data = data.get("data", {})
|
|
if not isinstance(payload_data, dict):
|
|
payload_data = {}
|
|
channels = payload_data.get("channels", [])
|
|
if isinstance(channels, str):
|
|
channels = [channels]
|
|
elif not isinstance(channels, list):
|
|
channels = []
|
|
channel = payload_data.get("channel")
|
|
if channel and channel not in channels:
|
|
channels = [*channels, channel]
|
|
if is_anonymous:
|
|
channels = [channel for channel in channels if channel in supported_channels]
|
|
vessel_subscription = None
|
|
if "vessels" in channels and "bbox" in payload_data:
|
|
try:
|
|
vessel_subscription = manager.subscribe_vessels(websocket, payload_data)
|
|
except ValueError as exc:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "subscription_error",
|
|
"data": {"channel": "vessels", "detail": str(exc)},
|
|
}
|
|
)
|
|
continue
|
|
channels = [channel for channel in channels if channel != "vessels"]
|
|
manager.subscribe(websocket, channels)
|
|
await websocket.send_json(
|
|
{
|
|
"type": "subscription_confirmed",
|
|
"data": {
|
|
"action": "subscribe",
|
|
"channels": [
|
|
*channels,
|
|
*(["vessels"] if vessel_subscription else []),
|
|
],
|
|
"vessels": vessel_subscription,
|
|
},
|
|
}
|
|
)
|
|
elif data.get("type") == "unsubscribe":
|
|
channels = data.get("data", {}).get("channels", [])
|
|
manager.unsubscribe(websocket, channels)
|
|
await websocket.send_json(
|
|
{
|
|
"type": "subscription_confirmed",
|
|
"data": {"action": "unsubscribe", "channels": channels},
|
|
}
|
|
)
|
|
elif data.get("type") == "control_frame":
|
|
await websocket.send_json(
|
|
{"type": "control_acknowledged", "data": {"received": True}}
|
|
)
|
|
else:
|
|
await websocket.send_json({"type": "ack", "data": {"received": True}})
|
|
|
|
except asyncio.TimeoutError:
|
|
await websocket.send_json({"type": "heartbeat", "data": {"action": "ping"}})
|
|
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
manager.disconnect(websocket, user_id)
|