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,6 +92,30 @@ class LobbyClientController:
         """Send an error string to the connected client"""
         await self.send(Message(type=MessageType.ERROR, data_type="error", data=error))
 
+    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,
+                data_type="object",
+                data=self.lobby_controller.lobby.serialize(),
+            )
+        )
+        await self.send(
+            Message(
+                type=MessageType.CONTEXT,
+                data_type="game_context_id",
+                data="lobby",
+            )
+        )
+
     @vibe_check(LoginMessage)
     async def login_context(self, message: LoginMessage) -> None:
         if message.type == MessageType.DISCORD_LOGIN:
@@ -97,24 +123,10 @@ class LobbyClientController:
             if user is None:
                 return
             if "id" in user:
-                self.discord_uid = user["id"]
+                await self.login_user(user)
-                self.username = user["username"]
+        if message.type == MessageType.COOKIE_LOGIN:
-                self.discord_avatar = user["avatar"]
+            if message.data in session_storage:
-                self.on_message = self.lobby_context
+                await self.login_user(session_storage[message.data])
-                await self.send(
-                    Message(
-                        type=MessageType.CREATE,
-                        data_type="object",
-                        data=self.lobby_controller.lobby.serialize(),
-                    )
-                )
-                await self.send(
-                    Message(
-                        type=MessageType.CONTEXT,
-                        data_type="game_context_id",
-                        data="lobby",
-                    )
-                )
         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"