fixed #43

2 files changed, 35 insertions, 6 deletions

diff --git a/test_websocket.py b/test_websocket.py
index 1e14cc6..0800a3c 100644
--- a/test_websocket.py
+++ b/test_websocket.py
@@ -23,15 +23,16 @@ import websocket as ws
 # [1]: https://tools.ietf.org/html/rfc6455#section-5.4
 #      "RFC6455: 5.4. Fragmentation"
 #
-TEST_FRAGMENTATION=False
+TEST_FRAGMENTATION = True
+
+TRACABLE = False
 
-TRACABLE=False
 
 def create_mask_key(n):
     return "abcd"
 
-class SockMock(object):
 
+class SockMock(object):
     def __init__(self):
         self.data = []
         self.sent = []
@@ -248,6 +249,14 @@ class WebSocketTest(unittest.TestCase):
             sock.recv()
 
     @unittest.skipUnless(TEST_FRAGMENTATION, "fragmentation not implemented")
+    def testRecvContFragmentation(self):
+        sock = ws.WebSocket()
+        s = sock.sock = SockMock()
+        # OPCODE=CONT, FIN=1, MSG="the soul of wit"
+        s.add_packet("\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
+        self.assertRaises(ws.WebSocketException, sock.recv)
+
+    @unittest.skipUnless(TEST_FRAGMENTATION, "fragmentation not implemented")
     def testRecvWithProlongedFragmentation(self):
         sock = ws.WebSocket()
         s = sock.sock = SockMock()
diff --git a/websocket.py b/websocket.py
index c2d5c9e..bdd6b45 100644
--- a/websocket.py
+++ b/websocket.py
@@ -228,6 +228,7 @@ class ABNF(object):
     """
 
     # operation code values.
+    OPCODE_CONT   = 0x0
     OPCODE_TEXT   = 0x1
     OPCODE_BINARY = 0x2
     OPCODE_CLOSE  = 0x8
@@ -235,11 +236,12 @@ class ABNF(object):
     OPCODE_PONG   = 0xa
 
     # available operation code value tuple
-    OPCODES = (OPCODE_TEXT, OPCODE_BINARY, OPCODE_CLOSE,
+    OPCODES = (OPCODE_CONT, OPCODE_TEXT, OPCODE_BINARY, OPCODE_CLOSE,
                 OPCODE_PING, OPCODE_PONG)
 
     # opcode human readable string
     OPCODE_MAP = {
+        OPCODE_CONT: "cont",
         OPCODE_TEXT: "text",
         OPCODE_BINARY: "binary",
         OPCODE_CLOSE: "close",
@@ -267,6 +269,11 @@ class ABNF(object):
         self.data = data
         self.get_mask_key = os.urandom
 
+    def __str__(self):
+        return "fin=" + str(self.fin) \
+                + " opcode=" + str(self.opcode) \
+                + " data=" + str(self.data)
+
     @staticmethod
     def create_frame(data, opcode):
         """
@@ -379,6 +386,7 @@ class WebSocket(object):
         self._frame_header = None
         self._frame_length = None
         self._frame_mask = None
+        self._cont_data = None
 
     def fileno(self):
         return self.sock.fileno()
@@ -552,11 +560,13 @@ class WebSocket(object):
         if self.get_mask_key:
             frame.get_mask_key = self.get_mask_key
         data = frame.format()
+        length = len(data)
         if traceEnabled:
             logger.debug("send: " + repr(data))
         while data:
             l = self._send(data)
             data = data[l:]
+        return length
 
     def send_binary(self, payload):
         return self.send(payload, ABNF.OPCODE_BINARY)
@@ -598,8 +608,18 @@ class WebSocket(object):
                 # handle error:
                 # 'NoneType' object has no attribute 'opcode'
                 raise WebSocketException("Not a valid frame %s" % frame)
-            elif frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY):
-                return (frame.opcode, frame.data)
+            elif frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY, ABNF.OPCODE_CONT):
+                if frame.opcode == ABNF.OPCODE_CONT and not self._cont_data:
+                    raise WebSocketException("Illegal frame")
+                if self._cont_data:
+                    self._cont_data[1] += frame.data
+                else:
+                    self._cont_data = [frame.opcode, frame.data]
+                
+                if frame.fin:
+                    data = self._cont_data
+                    self._cont_data = None
+                    return data
             elif frame.opcode == ABNF.OPCODE_CLOSE:
                 self.send_close()
                 return (frame.opcode, None)