From e112707fccc90dab3c7c17f67a3a729b9da18999 Mon Sep 17 00:00:00 2001
From: Phil Sar <lifippos13@icloud.com>
Date: Mon, 8 Aug 2022 17:13:07 +0300
Subject: [PATCH] Backend fixes for sockets staying open

---
 __init__.py             |  8 ++------
 src/game.py             | 13 ++++++++++++-
 src/server/backend.py   | 30 ++++++++++--------------------
 src/server/instances.py |  4 ----
 4 files changed, 24 insertions(+), 31 deletions(-)

diff --git a/__init__.py b/__init__.py
index 4ce4a4a..2600301 100644
--- a/__init__.py
+++ b/__init__.py
@@ -101,11 +101,10 @@
                     ]:
                         game.player.jump()
                     elif event.key == pygame.K_ESCAPE and not game.pause_menu:
-                        game.showing_gui = True
-                        game.showing_title = True
-                        game.pause_menu = False
                         if game.crashing:
                             src.game.game_crash.stop()
+                        game.reset_everything()
+                        game.client.stop()
                         game.gui.add(src.game.gui.Button((52, 90), "Online", game.start))
                         game.gui.add(
                             src.game.gui.Button((107, 90), "Private", game.start_private)
@@ -122,9 +121,6 @@
                         game.gui.add(
                             src.game.gui.EmojiButton((12, 10), "👥", game.join_with_code)
                         )
-                        game.pin_code = None
-                        game.private = False
-                        game.client.stop()
                         break
                     elif event.key == pygame.K_f and not game.crashing:
                         game.showing_gui = True
diff --git a/src/game.py b/src/game.py
index e4c637b..8edc34a 100644
--- a/src/game.py
+++ b/src/game.py
@@ -676,10 +676,21 @@ def __init__(self):
         self.ending_man = EndingIncrementManager(self)
         self.tile_timer = TimedTileToggler(self)
         self.sound = True
-        self.pin_code = None
+        self.pin_code = ""
         self.private = False
         mixer.play(-1)
 
+    def reset_everything(self):
+        """Reset everything!"""
+        self.showing_gui = True
+        self.showing_title = True
+        self.pause_menu = False
+        self.pin_code = ""
+        self.private = False
+        self.other_players.empty()
+        self.tiles = pygame.sprite.LayeredUpdates()
+        self.objects = pygame.sprite.LayeredUpdates(self.player)
+
     def quit(self):
         """Quit button event"""
         if self.client.running:
diff --git a/src/server/backend.py b/src/server/backend.py
index 57aa79f..06c932a 100644
--- a/src/server/backend.py
+++ b/src/server/backend.py
@@ -32,6 +32,7 @@ def _resource_path(file: str):
 manager = ConnectionManager()
 db = GameDatabase()
 anticheat = GameAntiCheat()
+players = set()
 
 
 async def new_game(player, private: bool = False):
@@ -83,16 +84,16 @@ async def join_with_code(player, code):
         await new_game(player)
 
 
-async def ping_pong(websocket):
+async def ping_pong(ws):
     """Trying to keep broadcast alive."""
-    while websocket in manager.active_broadcasts:
+    while ws is not None:
         t0 = time.perf_counter()
-        pong_waiter = await websocket.ping()
+        pong_waiter = await ws.ping()
         await pong_waiter
         t1 = time.perf_counter()
         latency = f"{t1-t0:.2f}"
         message = json.dumps({"type": "ping", "latency": latency})
-        await websocket.send(message)
+        await ws.send(message)
         await asyncio.sleep(0.5)
 
 
@@ -163,9 +164,6 @@ async def close_main(websocket, player):
     await manager.drop_main(websocket)
 
 
-players = set()
-
-
 async def close_broadcast(websocket, request_id):
     """Close broadcast websocket properly."""
     for play in players.copy():
@@ -224,16 +222,9 @@ async def handler(websocket):
             await manager.add_broadcast(websocket)
             for pl in players.copy():
                 if pl.unique_id == event["unique_id"]:
-                    while True:
-                        if isinstance(
-                            websocket, websockets.legacy.server.WebSocketServerProtocol
-                        ):
-                            logging.info("Attached websocket")
-                            pl.attach_broadcast(websocket)
-                            await ping_pong(websocket)
-                            break
-                        else:
-                            await asyncio.sleep(0.2)
+                    logging.info("Attached websocket")
+                    pl.broadcast = websocket
+                    await ping_pong(pl.broadcast)
 
     except websockets.exceptions.ConnectionClosedError:
         logging.info("Websocket closed with ConnectionClosedError")
@@ -243,10 +234,9 @@ async def handler(websocket):
 
     finally:
         # Drop websocket after figuring out its type
-        if websocket in manager.active_connections and player:
+        if websocket and player:
             await close_main(websocket, player)
-        elif websocket in manager.active_broadcasts:
-            await close_broadcast(websocket, event["unique_id"])
+            await close_broadcast(player.broadcast, player.unique_id)
 
 
 async def main(PORT):
diff --git a/src/server/instances.py b/src/server/instances.py
index 5e61381..53b2e9b 100644
--- a/src/server/instances.py
+++ b/src/server/instances.py
@@ -25,10 +25,6 @@ def data(self):
             "direction": self.direction,
         }
 
-    def attach_broadcast(self, websocket):
-        """Adds second websocket in Player"""
-        self.broadcast = websocket
-
 
 class GameInstance:
     """Logical Game Instance"""