539 lines
21 KiB
Python
539 lines
21 KiB
Python
"""Real-time WebSocket gateway.
|
|
|
|
Multiplexes presence, DM, and world-chat delivery over a single WebSocket
|
|
endpoint (docs/planning/social.md US-2, US-4).
|
|
|
|
Key features:
|
|
- Authentication via JWT token (query param or initial auth frame).
|
|
- Scoped subscriptions to DM conversations and world channels.
|
|
- Server-derived presence heartbeat (Redis TTL refreshed periodically).
|
|
- Broadcast presence changes (online/offline) to mutual friends.
|
|
- Fan-out via Redis pub/sub.
|
|
- Thread-safe & event-loop-safe dispatch for multi-instance and test environments.
|
|
- Graceful reconnect without duplicate delivery or orphaned state.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from datetime import datetime, timezone
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status
|
|
from jose import JWTError, jwt
|
|
from sqlalchemy import or_, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
import redis.asyncio as aioredis
|
|
|
|
from app.config import settings
|
|
from app.database import async_session_factory, get_db
|
|
from app.models.chat import ConversationParticipant
|
|
from app.models.social import Friendship, FriendshipStatus
|
|
from app.models.user import User
|
|
from app.models.world import World, WorldVisibility
|
|
from app.services.presence import presence_store
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def get_db_session():
|
|
"""Helper to acquire an async DB session respecting test overrides if present."""
|
|
from app.main import app
|
|
|
|
override = app.dependency_overrides.get(get_db)
|
|
if override:
|
|
res = override()
|
|
if hasattr(res, "__anext__"):
|
|
session = await res.__anext__()
|
|
try:
|
|
yield session
|
|
finally:
|
|
try:
|
|
await res.__anext__()
|
|
except (StopAsyncIteration, Exception):
|
|
pass
|
|
return
|
|
elif hasattr(res, "__next__"):
|
|
session = next(res)
|
|
try:
|
|
yield session
|
|
finally:
|
|
try:
|
|
next(res)
|
|
except (StopIteration, Exception):
|
|
pass
|
|
return
|
|
|
|
async with async_session_factory() as session:
|
|
yield session
|
|
|
|
|
|
async def authenticate_token(token: str) -> User | None:
|
|
"""Validate JWT token and return authenticated User, or None."""
|
|
try:
|
|
payload = jwt.decode(token, settings.public_pem, algorithms=[settings.ALGORITHM])
|
|
user_id_raw = payload.get("sub")
|
|
if user_id_raw is None:
|
|
return None
|
|
user_id = int(user_id_raw)
|
|
except (JWTError, ValueError, TypeError):
|
|
return None
|
|
|
|
async with get_db_session() as db:
|
|
res = await db.execute(select(User).where(User.id == user_id))
|
|
user = res.scalar_one_or_none()
|
|
if user and user.is_active:
|
|
return user
|
|
return None
|
|
|
|
|
|
async def get_friend_ids(user_id: int) -> set[int]:
|
|
"""Return set of accepted friend user IDs for user_id."""
|
|
try:
|
|
async with get_db_session() as db:
|
|
res = await db.execute(
|
|
select(Friendship).where(
|
|
or_(Friendship.requester_id == user_id, Friendship.addressee_id == user_id),
|
|
Friendship.status == FriendshipStatus.ACCEPTED,
|
|
)
|
|
)
|
|
friendships = res.scalars().all()
|
|
return {
|
|
f.addressee_id if f.requester_id == user_id else f.requester_id
|
|
for f in friendships
|
|
}
|
|
except Exception as exc:
|
|
logger.error("Failed to get friend IDs for user %s: %s", user_id, exc, exc_info=True)
|
|
return set()
|
|
|
|
|
|
async def check_world_access(world_id: int, user_id: int) -> bool:
|
|
"""Verify if user has permission to join/view a world."""
|
|
async with get_db_session() as db:
|
|
res = await db.execute(select(World).where(World.id == world_id))
|
|
world = res.scalar_one_or_none()
|
|
if not world:
|
|
return False
|
|
if world.visibility == WorldVisibility.PUBLIC:
|
|
return True
|
|
if world.owner_id == user_id:
|
|
return True
|
|
if world.visibility == WorldVisibility.FRIENDS_ONLY:
|
|
res_fr = await db.execute(
|
|
select(Friendship.id).where(
|
|
or_(
|
|
and_(Friendship.requester_id == world.owner_id, Friendship.addressee_id == user_id),
|
|
and_(Friendship.requester_id == user_id, Friendship.addressee_id == world.owner_id),
|
|
),
|
|
Friendship.status == FriendshipStatus.ACCEPTED,
|
|
).limit(1)
|
|
)
|
|
return res_fr.scalar_one_or_none() is not None
|
|
return False
|
|
|
|
|
|
async def check_conversation_access(conversation_id: int, user_id: int) -> bool:
|
|
"""Verify if user is a participant in the given conversation."""
|
|
async with get_db_session() as db:
|
|
res = await db.execute(
|
|
select(ConversationParticipant.id).where(
|
|
ConversationParticipant.conversation_id == conversation_id,
|
|
ConversationParticipant.user_id == user_id,
|
|
).limit(1)
|
|
)
|
|
return res.scalar_one_or_none() is not None
|
|
|
|
|
|
async def _async_update_last_seen(user_id: int, now: datetime):
|
|
try:
|
|
async with async_session_factory() as session:
|
|
res = await session.execute(select(User).where(User.id == user_id))
|
|
db_user = res.scalar_one_or_none()
|
|
if db_user:
|
|
db_user.last_seen_at = now
|
|
await session.commit()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manages active WebSocket connections, subscriptions, presence, and fan-out."""
|
|
|
|
def __init__(self):
|
|
# ws -> User
|
|
self._connections: dict[WebSocket, User] = {}
|
|
# ws -> event loop id
|
|
self._ws_loops: dict[WebSocket, int] = {}
|
|
# user_id -> set of active WebSockets
|
|
self._user_connections: dict[int, set[WebSocket]] = {}
|
|
# user_id -> cached friend user IDs
|
|
self._user_friend_ids: dict[int, set[int]] = {}
|
|
# world_id -> set of WebSockets
|
|
self._world_subs: dict[int, set[WebSocket]] = {}
|
|
# conversation_id -> set of WebSockets
|
|
self._dm_subs: dict[int, set[WebSocket]] = {}
|
|
# ws -> asyncio heartbeat task
|
|
self._heartbeat_tasks: dict[WebSocket, asyncio.Task] = {}
|
|
# Per-event-loop Redis pub/sub listener tasks and events
|
|
self._redis_tasks: dict[int, asyncio.Task] = {}
|
|
self._subscribed_events: dict[int, asyncio.Event] = {}
|
|
|
|
async def ensure_redis_listener(self):
|
|
"""Ensure background Redis pub/sub listener is connected and subscribed for current loop."""
|
|
loop = asyncio.get_running_loop()
|
|
loop_key = id(loop)
|
|
task = self._redis_tasks.get(loop_key)
|
|
if task is None or task.done():
|
|
event = asyncio.Event()
|
|
self._subscribed_events[loop_key] = event
|
|
task = asyncio.create_task(self._redis_listener(loop_key, event))
|
|
self._redis_tasks[loop_key] = task
|
|
else:
|
|
event = self._subscribed_events.get(loop_key)
|
|
|
|
if event:
|
|
try:
|
|
await asyncio.wait_for(event.wait(), timeout=2.0)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
|
|
async def connect(self, websocket: WebSocket, user: User):
|
|
"""Register an authenticated connection."""
|
|
await self.ensure_redis_listener()
|
|
loop_key = id(asyncio.get_running_loop())
|
|
|
|
self._connections[websocket] = user
|
|
self._ws_loops[websocket] = loop_key
|
|
self._user_connections.setdefault(user.id, set()).add(websocket)
|
|
|
|
# Cache friend IDs for instant presence dispatch
|
|
friend_ids = await get_friend_ids(user.id)
|
|
self._user_friend_ids[user.id] = friend_ids
|
|
|
|
# Mark user online in Redis and DB
|
|
await presence_store.heartbeat(user.id)
|
|
now = datetime.now(timezone.utc)
|
|
asyncio.create_task(_async_update_last_seen(user.id, now))
|
|
|
|
# Start periodic heartbeat task (every 20 seconds)
|
|
self._heartbeat_tasks[websocket] = asyncio.create_task(
|
|
self._run_periodic_heartbeat(user.id)
|
|
)
|
|
|
|
# Broadcast online presence to friends via Redis
|
|
await self.broadcast_presence(user.id, user.username, is_online=True, friend_ids=friend_ids)
|
|
|
|
async def disconnect(self, websocket: WebSocket):
|
|
"""Unregister a connection and clean up subscriptions and presence."""
|
|
user = self._connections.pop(websocket, None)
|
|
self._ws_loops.pop(websocket, None)
|
|
task = self._heartbeat_tasks.pop(websocket, None)
|
|
if task:
|
|
task.cancel()
|
|
|
|
# Clean subscriptions
|
|
for world_id in list(self._world_subs.keys()):
|
|
self._world_subs[world_id].discard(websocket)
|
|
if not self._world_subs[world_id]:
|
|
self._world_subs.pop(world_id, None)
|
|
|
|
for conv_id in list(self._dm_subs.keys()):
|
|
self._dm_subs[conv_id].discard(websocket)
|
|
if not self._dm_subs[conv_id]:
|
|
self._dm_subs.pop(conv_id, None)
|
|
|
|
if user is None:
|
|
return
|
|
|
|
user_id = user.id
|
|
username = user.username
|
|
user_set = self._user_connections.get(user_id)
|
|
if user_set:
|
|
user_set.discard(websocket)
|
|
if not user_set:
|
|
self._user_connections.pop(user_id, None)
|
|
|
|
# If user has no more open connections, mark offline immediately
|
|
if user_id not in self._user_connections:
|
|
await presence_store.clear_heartbeat(user_id)
|
|
now = datetime.now(timezone.utc)
|
|
asyncio.create_task(_async_update_last_seen(user_id, now))
|
|
friend_ids = self._user_friend_ids.pop(user_id, set())
|
|
await self.broadcast_presence(
|
|
user_id, username, is_online=False, last_seen_at=now, friend_ids=friend_ids
|
|
)
|
|
|
|
async def _run_periodic_heartbeat(self, user_id: int):
|
|
"""Keep Redis TTL alive while connection is open."""
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(20)
|
|
await presence_store.heartbeat(user_id)
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
async def subscribe_world(self, websocket: WebSocket, world_id: int) -> bool:
|
|
user = self._connections.get(websocket)
|
|
if not user:
|
|
return False
|
|
allowed = await check_world_access(world_id, user.id)
|
|
if not allowed:
|
|
return False
|
|
self._world_subs.setdefault(world_id, set()).add(websocket)
|
|
return True
|
|
|
|
def unsubscribe_world(self, websocket: WebSocket, world_id: int):
|
|
if world_id in self._world_subs:
|
|
self._world_subs[world_id].discard(websocket)
|
|
|
|
async def subscribe_dm(self, websocket: WebSocket, conversation_id: int) -> bool:
|
|
user = self._connections.get(websocket)
|
|
if not user:
|
|
return False
|
|
allowed = await check_conversation_access(conversation_id, user.id)
|
|
if not allowed:
|
|
return False
|
|
self._dm_subs.setdefault(conversation_id, set()).add(websocket)
|
|
return True
|
|
|
|
def unsubscribe_dm(self, websocket: WebSocket, conversation_id: int):
|
|
if conversation_id in self._dm_subs:
|
|
self._dm_subs[conversation_id].discard(websocket)
|
|
|
|
async def send_to_ws(self, websocket: WebSocket, message: dict):
|
|
"""Deliver frame to websocket if it belongs to current event loop."""
|
|
cur_loop = id(asyncio.get_running_loop())
|
|
if self._ws_loops.get(websocket) != cur_loop:
|
|
return
|
|
try:
|
|
await websocket.send_text(json.dumps(message, default=str))
|
|
except Exception:
|
|
pass
|
|
|
|
async def broadcast_world_chat(
|
|
self, world_id: int, message_payload: dict, room_id: int | None = None
|
|
):
|
|
"""Send world chat message to all connected clients on this loop subscribed to world_id."""
|
|
targets = list(self._world_subs.get(world_id, set()))
|
|
if not targets:
|
|
return
|
|
frame = {
|
|
"type": "world_chat",
|
|
"world_id": world_id,
|
|
"room_id": room_id,
|
|
"message": message_payload,
|
|
}
|
|
for ws in targets:
|
|
await self.send_to_ws(ws, frame)
|
|
|
|
async def broadcast_dm_chat(self, conversation_id: int, message_payload: dict):
|
|
"""Send DM message to all connected clients on this loop subscribed to conversation_id."""
|
|
targets = list(self._dm_subs.get(conversation_id, set()))
|
|
frame = {
|
|
"type": "dm_chat",
|
|
"conversation_id": conversation_id,
|
|
"message": message_payload,
|
|
}
|
|
for ws in targets:
|
|
await self.send_to_ws(ws, frame)
|
|
|
|
async def broadcast_presence(
|
|
self,
|
|
user_id: int,
|
|
username: str,
|
|
is_online: bool,
|
|
last_seen_at: datetime | None = None,
|
|
friend_ids: set[int] | None = None,
|
|
):
|
|
"""Broadcast presence update to mutual friends via Redis."""
|
|
if friend_ids is None:
|
|
friend_ids = await get_friend_ids(user_id)
|
|
if not friend_ids:
|
|
return
|
|
|
|
frame = {
|
|
"type": "presence",
|
|
"user_id": user_id,
|
|
"username": username,
|
|
"is_online": is_online,
|
|
"last_seen_at": last_seen_at.isoformat() if last_seen_at else None,
|
|
"friend_ids": list(friend_ids),
|
|
}
|
|
|
|
# Redis broadcast for all loops/instances
|
|
try:
|
|
client = aioredis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True)
|
|
async with client:
|
|
await client.publish("presence:events", json.dumps(frame, default=str))
|
|
except Exception as exc:
|
|
logger.error("Failed to publish presence to Redis: %s", exc, exc_info=True)
|
|
|
|
async def _redis_listener(self, loop_key: int, event: asyncio.Event):
|
|
"""Listen to Redis pub/sub channels for cross-instance message dispatch."""
|
|
while True:
|
|
try:
|
|
client = aioredis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True)
|
|
pubsub = client.pubsub()
|
|
await pubsub.psubscribe("world:*:chat", "dm:*:chat", "presence:events")
|
|
event.set()
|
|
logger.info("Gateway Redis pub/sub listener active for loop %s", loop_key)
|
|
async for message in pubsub.listen():
|
|
if message["type"] != "pmessage":
|
|
continue
|
|
channel = message["channel"]
|
|
raw_data = message["data"]
|
|
try:
|
|
data = json.loads(raw_data)
|
|
except Exception:
|
|
continue
|
|
|
|
if channel.startswith("world:") and channel.endswith(":chat"):
|
|
parts = channel.split(":")
|
|
world_id = int(parts[1])
|
|
await self.broadcast_world_chat(world_id, data, room_id=data.get("room_id"))
|
|
elif channel.startswith("dm:") and channel.endswith(":chat"):
|
|
parts = channel.split(":")
|
|
conv_id = int(parts[1])
|
|
await self.broadcast_dm_chat(conv_id, data)
|
|
elif channel == "presence:events":
|
|
logger.info("Redis presence event received: %s. Active users: %s", data, list(self._user_connections.keys()))
|
|
for friend_id in data.get("friend_ids", []):
|
|
targets = list(self._user_connections.get(friend_id, set()))
|
|
logger.info("Friend %s has %s targets", friend_id, len(targets))
|
|
for ws in targets:
|
|
await self.send_to_ws(ws, data)
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as exc:
|
|
logger.debug("Redis listener encountered error, retrying in 2s: %s", exc)
|
|
event.clear()
|
|
await asyncio.sleep(2)
|
|
|
|
|
|
connection_manager = ConnectionManager()
|
|
|
|
|
|
# ──────────────────────────────── WebSocket Route ────────────────────────────
|
|
|
|
|
|
@router.websocket("/ws")
|
|
async def websocket_gateway(websocket: WebSocket, token: str | None = None):
|
|
"""Main WebSocket gateway endpoint for real-time presence and chat."""
|
|
user: User | None = None
|
|
# 1. Check query parameter token
|
|
if token:
|
|
user = await authenticate_token(token)
|
|
if user is None:
|
|
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
|
return
|
|
|
|
await websocket.accept()
|
|
|
|
# 2. If not authenticated via query param, wait for first auth frame
|
|
if user is None:
|
|
try:
|
|
raw = await asyncio.wait_for(websocket.receive_text(), timeout=3.0)
|
|
data = json.loads(raw)
|
|
if data.get("type") == "auth" and data.get("token"):
|
|
user = await authenticate_token(data["token"])
|
|
except (asyncio.TimeoutError, json.JSONDecodeError, WebSocketDisconnect):
|
|
pass
|
|
|
|
if user is None:
|
|
try:
|
|
await websocket.send_text(json.dumps({"type": "error", "message": "Authentication failed"}))
|
|
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
|
except Exception:
|
|
pass
|
|
return
|
|
|
|
# Mark user online in Redis immediately upon successful auth
|
|
await presence_store.heartbeat(user.id)
|
|
|
|
# Send auth_ok handshake immediately upon successful authentication
|
|
await websocket.send_text(
|
|
json.dumps({
|
|
"type": "auth_ok",
|
|
"user_id": user.id,
|
|
"username": user.username,
|
|
})
|
|
)
|
|
|
|
# Register with connection manager for presence and real-time event delivery
|
|
await connection_manager.connect(websocket, user)
|
|
|
|
try:
|
|
while True:
|
|
text = await websocket.receive_text()
|
|
try:
|
|
msg = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
await websocket.send_text(json.dumps({"type": "error", "message": "Invalid JSON"}))
|
|
continue
|
|
|
|
msg_type = msg.get("type")
|
|
if msg_type == "ping":
|
|
await websocket.send_text(json.dumps({"type": "pong"}))
|
|
elif msg_type == "heartbeat":
|
|
await presence_store.heartbeat(user.id)
|
|
await websocket.send_text(json.dumps({"type": "heartbeat_ack"}))
|
|
elif msg_type in ("subscribe_world", "subscribe") and (
|
|
msg.get("world_id") or msg.get("channel") == "world"
|
|
):
|
|
world_id = msg.get("world_id") or msg.get("id")
|
|
if isinstance(world_id, int):
|
|
ok = await connection_manager.subscribe_world(websocket, world_id)
|
|
if ok:
|
|
await websocket.send_text(
|
|
json.dumps({"type": "subscribed", "channel": "world", "id": world_id})
|
|
)
|
|
else:
|
|
await websocket.send_text(
|
|
json.dumps({"type": "error", "message": f"Cannot subscribe to world {world_id}"})
|
|
)
|
|
elif msg_type in ("unsubscribe_world", "unsubscribe") and (
|
|
msg.get("world_id") or msg.get("channel") == "world"
|
|
):
|
|
world_id = msg.get("world_id") or msg.get("id")
|
|
if isinstance(world_id, int):
|
|
connection_manager.unsubscribe_world(websocket, world_id)
|
|
await websocket.send_text(
|
|
json.dumps({"type": "unsubscribed", "channel": "world", "id": world_id})
|
|
)
|
|
elif msg_type in ("subscribe_dm", "subscribe") and (
|
|
msg.get("conversation_id") or msg.get("channel") == "dm"
|
|
):
|
|
conv_id = msg.get("conversation_id") or msg.get("id")
|
|
if isinstance(conv_id, int):
|
|
ok = await connection_manager.subscribe_dm(websocket, conv_id)
|
|
if ok:
|
|
await websocket.send_text(
|
|
json.dumps({"type": "subscribed", "channel": "dm", "id": conv_id})
|
|
)
|
|
else:
|
|
await websocket.send_text(
|
|
json.dumps({"type": "error", "message": f"Cannot subscribe to conversation {conv_id}"})
|
|
)
|
|
elif msg_type in ("unsubscribe_dm", "unsubscribe") and (
|
|
msg.get("conversation_id") or msg.get("channel") == "dm"
|
|
):
|
|
conv_id = msg.get("conversation_id") or msg.get("id")
|
|
if isinstance(conv_id, int):
|
|
connection_manager.unsubscribe_dm(websocket, conv_id)
|
|
await websocket.send_text(
|
|
json.dumps({"type": "unsubscribed", "channel": "dm", "id": conv_id})
|
|
)
|
|
else:
|
|
await websocket.send_text(json.dumps({"type": "error", "message": f"Unknown action: {msg_type}"}))
|
|
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
await connection_manager.disconnect(websocket)
|