97 lines
2.8 KiB
Python
97 lines
2.8 KiB
Python
from collections import defaultdict
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.database import SessionLocal
|
|
from app.mongo import get_mongo_db
|
|
from app.services.document_service import get_document
|
|
from app.services.user_service import get_user_by_email
|
|
from app.utils.security import verify_token
|
|
|
|
router = APIRouter()
|
|
|
|
rooms: dict[str, set[WebSocket]] = defaultdict(set)
|
|
presence: dict[str, dict[int, str]] = defaultdict(dict)
|
|
|
|
|
|
async def _user_from_token(token: str):
|
|
subject = verify_token(token, expected_type="access")
|
|
if not subject:
|
|
return None
|
|
db: Session = SessionLocal()
|
|
try:
|
|
return get_user_by_email(db, subject)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
@router.websocket("/ws/collab/{doc_id}")
|
|
async def collab_socket(websocket: WebSocket, doc_id: str) -> None:
|
|
token = websocket.query_params.get("token", "")
|
|
user = await _user_from_token(token)
|
|
if user is None:
|
|
await websocket.close(code=4401)
|
|
return
|
|
|
|
try:
|
|
mongo = get_mongo_db()
|
|
except RuntimeError:
|
|
await websocket.close(code=1013)
|
|
return
|
|
|
|
doc = await get_document(mongo, doc_id, user.id)
|
|
if doc is None:
|
|
await websocket.close(code=4403)
|
|
return
|
|
|
|
await websocket.accept()
|
|
rooms[doc_id].add(websocket)
|
|
presence[doc_id][user.id] = user.full_name or user.username
|
|
|
|
async def broadcast(payload: dict, skip: WebSocket | None = None) -> None:
|
|
dead: list[WebSocket] = []
|
|
for client in rooms[doc_id]:
|
|
if client is skip:
|
|
continue
|
|
try:
|
|
await client.send_json(payload)
|
|
except Exception:
|
|
dead.append(client)
|
|
for client in dead:
|
|
rooms[doc_id].discard(client)
|
|
|
|
await broadcast(
|
|
{
|
|
"type": "presence",
|
|
"peers": [{"id": uid, "name": name} for uid, name in presence[doc_id].items()],
|
|
}
|
|
)
|
|
|
|
try:
|
|
while True:
|
|
message = await websocket.receive_json()
|
|
if message.get("type") == "content":
|
|
await broadcast(
|
|
{
|
|
"type": "content",
|
|
"content": message.get("content"),
|
|
"from": user.id,
|
|
},
|
|
skip=websocket,
|
|
)
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
rooms[doc_id].discard(websocket)
|
|
presence[doc_id].pop(user.id, None)
|
|
if rooms[doc_id]:
|
|
await broadcast(
|
|
{
|
|
"type": "presence",
|
|
"peers": [{"id": uid, "name": name} for uid, name in presence[doc_id].items()],
|
|
}
|
|
)
|
|
else:
|
|
rooms.pop(doc_id, None)
|
|
presence.pop(doc_id, None)
|