From 6f831961dc4f04610160c4e03ed5bb67e1d620ff Mon Sep 17 00:00:00 2001 From: liris Date: Thu, 9 Oct 2014 07:40:44 +0900 Subject: [PATCH] - fixed #113 --- websocket/_core.py | 10 +++++++++- websocket/tests/test_websocket.py | 9 +++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/websocket/_core.py b/websocket/_core.py index 9b82938..6e7eaa9 100644 --- a/websocket/_core.py +++ b/websocket/_core.py @@ -352,6 +352,7 @@ class WebSocket(object): # These buffer over the build-up of a single frame. self._frame_buffer = _FrameBuffer() self._cont_data = None + self._recving_frames = None if enable_multithread: self.lock = threading.Lock() else: @@ -684,17 +685,24 @@ class WebSocket(object): # 'NoneType' object has no attribute 'opcode' raise WebSocketException("Not a valid frame %s" % frame) elif frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY, ABNF.OPCODE_CONT): - if not self.fire_cont_frame and frame.opcode == ABNF.OPCODE_CONT and not self._cont_data: + if frame.opcode == ABNF.OPCODE_CONT and not self._recving_frames: raise WebSocketException("Illegal frame") + if self._cont_data: self._cont_data[1] += frame.data else: + if frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY): + self._recving_frames = frame.opcode self._cont_data = [frame.opcode, frame.data] + if frame.fin: + self._recving_frames = None + if frame.fin or self.fire_cont_frame: data = self._cont_data self._cont_data = None return data + elif frame.opcode == ABNF.OPCODE_CLOSE: self.send_close() return (frame.opcode, frame.data) diff --git a/websocket/tests/test_websocket.py b/websocket/tests/test_websocket.py index 190d962..28e10f5 100644 --- a/websocket/tests/test_websocket.py +++ b/websocket/tests/test_websocket.py @@ -294,8 +294,8 @@ class WebSocketTest(unittest.TestCase): s = sock.sock = SockMock() # OPCODE=TEXT, FIN=0, MSG="Brevity is " s.add_packet(six.b("\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")) - # OPCODE=TEXT, FIN=0, MSG="Brevity is " - s.add_packet(six.b("\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")) + # OPCODE=CONT, FIN=0, MSG="Brevity is " + s.add_packet(six.b("\x00\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")) # OPCODE=CONT, FIN=1, MSG="the soul of wit" s.add_packet(six.b("\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")) @@ -306,6 +306,11 @@ class WebSocketTest(unittest.TestCase): _, data = sock.recv_data() self.assertEqual(data, six.b("the soul of wit")) + # OPCODE=CONT, FIN=0, MSG="Brevity is " + s.add_packet(six.b("\x80\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")) + + with self.assertRaises(ws.WebSocketException): + sock.recv_data() with self.assertRaises(ws.WebSocketConnectionClosedException): sock.recv()