mirror of
https://github.com/umap-project/umap.git
synced 2025-04-28 19:42:36 +02:00
feat(WebSockets): Features a WebSocket server.
There is one "room" per map, and the server relays messages to all the other connected peers. Messages are checked for compliance with what's allowed as a security measure. They should also be checked in the clients to avoid potential attack vectors.
This commit is contained in:
parent
ee3dbb85ca
commit
1128348db6
2 changed files with 85 additions and 0 deletions
|
@ -35,6 +35,7 @@ dependencies = [
|
||||||
"django-sesame==3.2.2",
|
"django-sesame==3.2.2",
|
||||||
"Pillow==10.3.0",
|
"Pillow==10.3.0",
|
||||||
"psycopg==3.1.19",
|
"psycopg==3.1.19",
|
||||||
|
"pydantic==2.7.0",
|
||||||
"requests==2.32.3",
|
"requests==2.32.3",
|
||||||
"rcssmin==1.1.2",
|
"rcssmin==1.1.2",
|
||||||
"rjsmin==1.2.2",
|
"rjsmin==1.2.2",
|
||||||
|
|
84
umap/ws.py
Normal file
84
umap/ws.py
Normal file
|
@ -0,0 +1,84 @@
|
||||||
|
#!/usr/bin/env python
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from collections import defaultdict
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
|
import django
|
||||||
|
import websockets
|
||||||
|
from django.conf import settings
|
||||||
|
from pydantic import BaseModel
|
||||||
|
from pydantic.types import UUID4
|
||||||
|
from websockets.server import serve
|
||||||
|
|
||||||
|
# This needs to run before the django-specific imports
|
||||||
|
# See https://docs.djangoproject.com/en/5.0/topics/settings/#calling-django-setup-is-required-for-standalone-django-usage
|
||||||
|
from umap.settings import settings_as_dict
|
||||||
|
|
||||||
|
settings.configure(**settings_as_dict)
|
||||||
|
django.setup()
|
||||||
|
|
||||||
|
from sesame.utils import get_user # NOQA
|
||||||
|
from umap.models import Map, User # NOQA
|
||||||
|
|
||||||
|
# Contains the list of websocket connections handled by this process.
|
||||||
|
# It's a mapping of map_id to a set of the active websocket connections
|
||||||
|
CONNECTIONS = defaultdict(set)
|
||||||
|
|
||||||
|
|
||||||
|
class JoinMessage(BaseModel):
|
||||||
|
kind: str = "join"
|
||||||
|
token: str
|
||||||
|
map_id: UUID4
|
||||||
|
|
||||||
|
|
||||||
|
class OperationMessage(BaseModel):
|
||||||
|
kind: str = "operation"
|
||||||
|
subject: str = Literal["map", "layer", "feature"]
|
||||||
|
metadata: dict
|
||||||
|
data: dict
|
||||||
|
|
||||||
|
|
||||||
|
async def join_and_listen(map_id, websocket):
|
||||||
|
"""Join a "room" whith other connected peers.
|
||||||
|
|
||||||
|
New messages will be broadcasted to other connected peers.
|
||||||
|
"""
|
||||||
|
CONNECTIONS[map_id].add(websocket)
|
||||||
|
try:
|
||||||
|
async for raw_message in websocket:
|
||||||
|
# recompute the peers-list at the time of message-sending.
|
||||||
|
# as doing so beforehand would miss new connections
|
||||||
|
peers = CONNECTIONS[map_id] - {websocket}
|
||||||
|
|
||||||
|
# Only relay valid "operation" messages
|
||||||
|
OperationMessage.model_validate_json(raw_message)
|
||||||
|
websockets.broadcast(peers, raw_message)
|
||||||
|
finally:
|
||||||
|
CONNECTIONS[map_id].remove(websocket)
|
||||||
|
|
||||||
|
|
||||||
|
async def handler(websocket):
|
||||||
|
"""Main WebSocket handler.
|
||||||
|
|
||||||
|
If permissions are granted, let the peer enter a room.
|
||||||
|
"""
|
||||||
|
raw_message = await websocket.recv()
|
||||||
|
|
||||||
|
# The first event should always be 'join'
|
||||||
|
message: JoinMessage = JoinMessage.model_validate_json(raw_message)
|
||||||
|
|
||||||
|
user: User = await asyncio.to_thread(get_user, message.token)
|
||||||
|
map_obj: Map = await asyncio.to_thread(Map.objects.get, message.map_id)
|
||||||
|
|
||||||
|
if map_obj.can_edit(user):
|
||||||
|
await join_and_listen(message.map_id, websocket)
|
||||||
|
|
||||||
|
|
||||||
|
async def main():
|
||||||
|
print("WebSocket server waiting for connections")
|
||||||
|
async with serve(handler, "localhost", 8001):
|
||||||
|
await asyncio.Future() # run forever
|
||||||
|
|
||||||
|
|
||||||
|
asyncio.run(main())
|
Loading…
Reference in a new issue