Check for ConnectionResetError when reading bytes from client

Author: Pithikos

websocket_server/websocket_server.py

@@ -190,6 +190,10 @@ def read_bytes(self, num):
     def read_next_message(self):
         try:
             b1, b2 = self.read_bytes(2)
+        except ConnectionResetError:
+            logger.info("Client closed connection.")
+            self.keep_alive = 0
+            return
         except ValueError as e:
             b1, b2 = 0, 0
 
@@ -198,10 +202,6 @@ def read_next_message(self):
         masked = b2 & MASKED
         payload_length = b2 & PAYLOAD_LEN
 
-        if not b1:
-            logger.info("Client closed connection.")
-            self.keep_alive = 0
-            return
         if opcode == OPCODE_CLOSE_CONN:
             logger.info("Client asked to close connection.")
             self.keep_alive = 0