Reject invalid control frames when building

test/test_frame.py

@@ -95,6 +95,43 @@ def test_frame_too_large(self):
         f.payload_length = 1 << 63
         self.assertRaises(FrameTooLargeException, f.build)
 
+    def test_fragmented_control_frames_with_payload_are_rejected(self):
+        for opcode in (OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG):
+            f = Frame(opcode=opcode, body=b'control', fin=0)
+            self.assertRaises(ValueError, f.build)
+
+    def test_control_frames_cannot_exceed_125_bytes(self):
+        for opcode in (OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG):
+            for mask in (None, b'1234'):
+                f = Frame(opcode=opcode, body=b'*' * 126,
+                          masking_key=mask, fin=1)
+                self.assertRaises(FrameTooLargeException, f.build)
+
+    def test_control_frames_accept_empty_and_125_byte_payloads(self):
+        for opcode in (OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG):
+            for mask in (None, b'1234'):
+                for body in (b'', b'*' * 125):
+                    data = Frame(opcode=opcode, body=body,
+                                 masking_key=mask, fin=1).build()
+                    parsed = Frame()
+                    parsed.parser.send(data)
+                    self.assertEqual(parsed.opcode, opcode)
+                    self.assertEqual(parsed.payload_length, len(body))
+                    decoded = parsed.unmask(parsed.body) if mask else parsed.body
+                    self.assertEqual(decoded, body)
+
+    def test_control_frame_payload_limit(self):
+        for opcode in (OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG):
+            self.assertEqual(len(Frame(opcode=opcode, body=b'*' * 125, fin=1).build()), 127)
+            self.assertRaises(
+                FrameTooLargeException,
+                Frame(opcode=opcode, body=b'*' * 126, fin=1).build,
+            )
+
+    def test_control_frames_cannot_be_fragmented(self):
+        for opcode in (OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG):
+            self.assertRaises(ValueError, Frame(opcode=opcode, body=b'', fin=0).build)
+
     def test_passing_encoded_string(self):
         # once encoded the u'\xe9' character will be of length 2
         f = Frame(opcode=OPCODE_TEXT, body=u'\xe9trange'.encode('utf-8'), fin=1)
@@ -187,7 +224,8 @@ def test_fragmented_control_frame_is_invalid(self):
         self.assertRaises(ProtocolException, f.parser.send, b'0x9')
 
     def test_fragmented_control_frame_is_too_large(self):
-        bytes = Frame(opcode=OPCODE_PING, body=b'*'*65536, fin=1).build()
+        # Construct malformed peer input without the validating sender.
+        bytes = b'\x89\x7f' + pack('!Q', 65536) + b'*' * 65536
         f = Frame()
         self.assertRaises(FrameTooLargeException, f.parser.send, bytes)
 

test/test_stream.py

@@ -128,8 +128,9 @@ def test_invalid_encoded_bytes_on_continuation(self):
 
     def test_too_large_close_message(self):
         payload = struct.pack("!H", 1000) + b'*' * 330
-        f = Frame(opcode=OPCODE_CLOSE, body=payload,
-                  fin=1, masking_key=os.urandom(4)).build()
+        # A zero masking key leaves the malformed peer payload unchanged.
+        # Do not ask the validating sender to build this oversized control frame.
+        f = b'\x88\xfe' + struct.pack('!H', len(payload)) + b'\x00' * 4 + payload
         s = Stream()
         self.assertEqual(len(s.errors), 0)
         self.assertEqual(s.closing, None)

ws4py/framing.py

@@ -75,6 +75,12 @@ def build(self):
         if 0x3 <= self.opcode <= 0x7 or 0xB <= self.opcode:
             raise ValueError('Opcode cannot be a reserved opcode')
 
+        if self.opcode in (OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG):
+            if self.fin != 1:
+                raise ValueError('Control frames cannot be fragmented')
+            if self.payload_length > 125:
+                raise FrameTooLargeException()
+
         ## +-+-+-+-+-------+
         ## |F|R|R|R| opcode|
         ## |I|S|S|S|  (4)  |