From 1128348db6e5cfc9af0b3505a7d2e59dbd5732ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexis=20M=C3=A9taireau?= <alexis@notmyidea.org>
Date: Tue, 16 Apr 2024 18:11:18 +0200
Subject: [PATCH] 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.
---
 pyproject.toml |  1 +
 umap/ws.py     | 84 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+)
 create mode 100644 umap/ws.py

diff --git a/pyproject.toml b/pyproject.toml
index feb54f46..e8c59da7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,6 +35,7 @@ dependencies = [
   "django-sesame==3.2.2",
   "Pillow==10.3.0",
   "psycopg==3.1.19",
+  "pydantic==2.7.0",
   "requests==2.32.3",
   "rcssmin==1.1.2",
   "rjsmin==1.2.2",
diff --git a/umap/ws.py b/umap/ws.py
new file mode 100644
index 00000000..644032e6
--- /dev/null
+++ b/umap/ws.py
@@ -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())