Merge 14f6c25eb6 into 83feb2c0e0

owrx/websocket.py

@@ -32,6 +32,8 @@ class Drained(WebSocketException):
 class WebSocketClosed(WebSocketException):
     pass
 
+class InterruptOrTimeout(WebSocketException):
+    pass
 
 class WebSocketConnection(object):
     connections = []
@@ -176,20 +178,25 @@ class WebSocketConnection(object):
         def protected_read(num):
             data = self.handler.rfile.read(num)
             if data is None:
-                raise Drained()
-            if len(data) != num:
+                data = bytes()
+            while self.open and len(data) < num:
+                    (read, _, _) = select.select([self.interruptPipeRecv, self.handler.rfile], [], [], 15)
+                    if self.handler.rfile in read:
+                        data += self.handler.rfile.read(num - len(data))
+                    else:
+                        if len(data) == 0:
+                            raise InterruptOrTimeout()
+                        else:
+                            raise IncompleteRead()
+            if len(data) < num:
                 raise IncompleteRead()
             return data
 
         self.open = True
         while self.open:
-            (read, _, _) = select.select([self.interruptPipeRecv, self.handler.rfile], [], [], 15)
-            if self.handler.rfile in read:
-                available = True
-                self.resetPing()
-                while self.open and available:
             try:
                 header = protected_read(2)
+                self.resetPing()
                 opcode = header[0] & 0x0F
                 length = header[1] & 0x7F
                 mask = (header[1] & 0x80) >> 7
@@ -198,6 +205,8 @@ class WebSocketConnection(object):
                     length = (header[0] << 8) + header[1]
                 if mask:
                     masking_key = protected_read(4)
+                data = None
+                if length > 0:
                     data = protected_read(length)
                     if mask:
                         data = bytes([b ^ masking_key[index % 4] for (index, b) in enumerate(data)])
@@ -222,8 +231,8 @@ class WebSocketConnection(object):
                     self.open = False
                 else:
                      logger.warning("unsupported opcode: {0}".format(opcode))
-                    except Drained:
-                        available = False
+            except InterruptOrTimeout:
+                pass
             except IncompleteRead:
                 logger.warning("incomplete read on websocket; closing connection")
                 self.open = False