Properly handle disconnections from remote end

Signed-off-by: Olivier Gayot <olivier.gayot@sigexec.com>

1 files changed, 4 insertions, 2 deletions

diff --git a/swiftstory/SwiftStory.py b/swiftstory/SwiftStory.py
index de3eab8..256c2be 100644
--- a/swiftstory/SwiftStory.py
+++ b/swiftstory/SwiftStory.py
@@ -2,6 +2,7 @@
 
 import argparse
 import asyncio
+import contextlib
 import json
 import websockets
 
@@ -62,8 +63,9 @@ async def message_received_handler(client, websocket, message):
 async def swiftstory(websocket, path):
     client = Client(websocket, game_manager)
 
-    async for message in websocket:
-        await message_received_handler(client, websocket, message)
+    with contextlib.suppress(websockets.exceptions.ConnectionClosed):
+        async for message in websocket:
+            await message_received_handler(client, websocket, message)
 
     await client.disconnect()