Fix websocket data parsing if content is larger than read BUFSIZE; Is… (

#184) * Fix websocket data parsing if content is larger than read BUFSIZE; Issue #181 * Add test for websocket response with request larger than read BUFSIZE
circuits · Nov 4, 2016 · 3ec4b32 · 3ec4b32 · ri0t · Dec 21, 2016
1 parent b0df3dd
commit 3ec4b32
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/circuits/protocols/websocket.py b/circuits/protocols/websocket.py
@@ -49,6 +49,7 @@ def __init__(self, sock=None, data=bytearray(), *args, **kwargs):
         self._pending_type = None
         self._close_received = False
         self._close_sent = False
+        self._buffer = bytearray()
 
         messages = self._parse_messages(bytearray(data))
         for message in messages:
@@ -90,6 +91,7 @@ def _parse_messages(self, data):
         msgs = []  # one chunk of bytes may result in several messages
         if self._close_received:
             return msgs
+        data = self._buffer + data
         while data:
             # extract final flag, opcode and masking
             final = bool(data[0] & 0x80 != 0)
@@ -111,7 +113,9 @@ def _parse_messages(self, data):
                 offset += 4
             # if not enough bytes available yet, retry after next read
             if len(data) - offset < payload_length:
+                self._buffer = data
                 break
+            self._buffer = bytearray()
             # rest of _buffer is payload
             msg = data[offset:offset + payload_length]
             if masking:  # unmask

diff --git a/tests/web/test_websockets.py b/tests/web/test_websockets.py
@@ -6,7 +6,7 @@
 
 from circuits import Component
 from circuits.web.controllers import Controller
-from circuits.net.sockets import close, write
+from circuits.net.sockets import close, write, BUFSIZE
 from circuits.web.websockets import WebSocketClient, WebSocketsDispatcher
 
 
@@ -81,6 +81,12 @@ def test(manager, watcher, webapp):
     assert watcher.wait("read", channel="ws")
     assert client.response == "Received: Hello!"
 
+    for size in (BUFSIZE, BUFSIZE + 1, BUFSZIE * 2):
+        data = "A" * (size + 1)
+        client.fire(write(data), "ws")
+        assert watcher.wait("read", channel="ws")
+        assert client.response == "Received: %s" % (data,)
+
     f = urlopen(webapp.server.http.base)
     s = f.read()
     assert s == b"Hello World!"