close websocket connections in an improved way

owrx/connection.py

@@ -20,13 +20,14 @@ class Client(object):
             self.conn.send(data)
         # these exception happen when the socket is closed
         except OSError:
+            logger.exception("OSError while sending data")
             self.close()
         except ValueError:
+            logger.exception("ValueError while sending data")
             self.close()
 
     def close(self):
         self.conn.close()
-        logger.debug("connection closed")
 
 
 class OpenWebRxReceiverClient(Client):
@@ -288,7 +289,3 @@ class WebSocketMessageHandler(object):
 
     def handleBinaryMessage(self, conn, data):
         logger.error("unsupported binary message, discarding")
-
-    def handleClose(self, conn):
-        if self.client:
-            self.client.close()

owrx/http.py

@@ -10,8 +10,8 @@ from owrx.controllers import (
     AprsSymbolsController,
 )
 from http.server import BaseHTTPRequestHandler
-import re
 from urllib.parse import urlparse, parse_qs
+import re
 
 import logging
 

owrx/websocket.py

@@ -1,16 +1,24 @@
 import base64
 import hashlib
 import json
+import os
+import select
 
 import logging
 
 logger = logging.getLogger(__name__)
 
 
+class IncompleteRead(Exception):
+    pass
+
+
 class WebSocketConnection(object):
     def __init__(self, handler, messageHandler):
         self.handler = handler
         self.messageHandler = messageHandler
+        (self.interruptPipeRecv, self.interruptPipeSend) = os.pipe()
+        self.open = True
         my_headers = self.handler.headers.items()
         my_header_keys = list(map(lambda x: x[0], my_headers))
         h_key_exists = lambda x: my_header_keys.count(x)
@@ -78,19 +86,31 @@ class WebSocketConnection(object):
         else:
             self.handler.wfile.flush()
 
+    def protected_read(self, num):
+        data = self.handler.rfile.read(num)
+        if len(data) != num:
+            raise IncompleteRead()
+        return data
+
+    def interrupt(self):
+        os.write(self.interruptPipeSend, bytes(0x00))
+
     def read_loop(self):
-        open = True
-        while open:
-            header = self.handler.rfile.read(2)
+        self.open = True
+        while self.open:
+            try:
+                (read, _, _) = select.select([self.interruptPipeRecv, self.handler.rfile], [], [])
+                if read[0] == self.handler.rfile:
+                    header = self.protected_read(2)
                     opcode = header[0] & 0x0F
                     length = header[1] & 0x7F
                     mask = (header[1] & 0x80) >> 7
                     if length == 126:
-                header = self.handler.rfile.read(2)
+                        header = self.protected_read(2)
                         length = (header[0] << 8) + header[1]
                     if mask:
-                masking_key = self.handler.rfile.read(4)
-            data = self.handler.rfile.read(length)
+                        masking_key = self.protected_read(4)
+                    data = self.protected_read(length)
                     if mask:
                         data = bytes([b ^ masking_key[index % 4] for (index, b) in enumerate(data)])
                     if opcode == 1:
@@ -99,12 +119,16 @@ class WebSocketConnection(object):
                     elif opcode == 2:
                         self.messageHandler.handleBinaryMessage(self, data)
                     elif opcode == 8:
-                open = False
-                self.messageHandler.handleClose(self)
+                        logger.debug("websocket close frame received; closing connection")
+                        self.open = False
                     else:
                         logger.warning("unsupported opcode: {0}".format(opcode))
+            except IncompleteRead:
+                logger.warning("incomplete websocket read; closing socket")
+                self.open = False
+
+        logger.debug("websocket loop ended; sending close frame")
 
-    def close(self):
         try:
             header = self.get_header(0, 8)
             self.handler.wfile.write(header)
@@ -114,11 +138,9 @@ class WebSocketConnection(object):
         except OSError:
             logger.exception("OSError while writing close frame:")
 
-        try:
-            self.handler.finish()
-            self.handler.connection.close()
-        except Exception:
-            logger.exception("while closing connection:")
+    def close(self):
+        self.open = False
+        self.interrupt()
 
 
 class WebSocketException(Exception):