Add session cookie

src/controller/lobbyClient.py

@@ -12,6 +12,8 @@ from .auth import discord_process_code
 
 from controller import ClientController
 
+from .session_storage import generate_cookie, session_storage
+
 if TYPE_CHECKING:
     from controller import LobbyController
 
@@ -90,17 +92,15 @@ class LobbyClientController:
         """Send an error string to the connected client"""
         await self.send(Message(type=MessageType.ERROR, data_type="error", data=error))
 
-    @vibe_check(LoginMessage)
-    async def login_context(self, message: LoginMessage) -> None:
-        if message.type == MessageType.DISCORD_LOGIN:
-            user = discord_process_code(message.data)
-            if user is None:
-                return
-            if "id" in user:
+    async def login_user(self, user: dict):
         self.discord_uid = user["id"]
         self.username = user["username"]
         self.discord_avatar = user["avatar"]
         self.on_message = self.lobby_context
+        cookie = generate_cookie(user)
+        await self.send(
+            Message(type=MessageType.SET, data_type="session_cookie", data=cookie)
+        )
         await self.send(
             Message(
                 type=MessageType.CREATE,
@@ -115,6 +115,18 @@ class LobbyClientController:
                 data="lobby",
             )
         )
+
+    @vibe_check(LoginMessage)
+    async def login_context(self, message: LoginMessage) -> None:
+        if message.type == MessageType.DISCORD_LOGIN:
+            user = discord_process_code(message.data)
+            if user is None:
+                return
+            if "id" in user:
+                await self.login_user(user)
+        if message.type == MessageType.COOKIE_LOGIN:
+            if message.data in session_storage:
+                await self.login_user(session_storage[message.data])
         else:
             print("user not connected")
 

src/controller/session_storage.py

@@ -0,0 +1,12 @@
+import secrets
+from websockets.server import WebSocketServerProtocol
+
+from netcode.models import Message
+
+session_storage: dict[str, dict] = {}
+
+
+def generate_cookie(user: dict):
+    cookie = secrets.token_urlsafe(32)
+    session_storage[cookie] = user
+    return cookie

src/netcode/models.py

@@ -10,6 +10,7 @@ class MessageType(str, Enum):
     LOGIN = "login"
     READY = "ready"
     DISCORD_LOGIN = "discord_login"
+    COOKIE_LOGIN = "cookie_login"
     JOIN = "join"
 
     CONTEXT = "context"