Support closing while sending a fragmented message.

On one hand, it will close the connection with an unfinished fragmented
message, which is less than ideal.

On the other hand, RFC 6455 implies that it should be legal and it's
probably best to let users close the connection if they want to close
the connection (rather than force them to call fail() instead).

src/websockets/protocol.py

@@ -354,12 +354,10 @@ def send_close(self, code: Optional[int] = None, reason: str = "") -> None:
             reason: close reason.
 
         Raises:
-            ProtocolError: If a fragmented message is being sent, if the code
-                isn't valid, or if a reason is provided without a code
+            ProtocolError: If the code isn't valid or if a reason is provided
+                without a code.
 
         """
-        if self.expect_continuation_frame:
-            raise ProtocolError("expected a continuation frame")
         if code is None:
             if reason != "":
                 raise ProtocolError("cannot send a reason without a code")

tests/test_protocol.py

@@ -1364,26 +1364,22 @@ def test_client_send_close_in_fragmented_message(self):
         client = Protocol(CLIENT)
         client.send_text(b"Spam", fin=False)
         self.assertFrameSent(client, Frame(OP_TEXT, b"Spam", fin=False))
-        # The spec says: "An endpoint MUST be capable of handling control
-        # frames in the middle of a fragmented message." However, since the
-        # endpoint must not send a data frame after a close frame, a close
-        # frame can't be "in the middle" of a fragmented message.
-        with self.assertRaises(ProtocolError) as raised:
-            client.send_close(CloseCode.GOING_AWAY)
-        self.assertEqual(str(raised.exception), "expected a continuation frame")
-        client.send_continuation(b"Eggs", fin=True)
+        with self.enforce_mask(b"\x3c\x3c\x3c\x3c"):
+            client.send_close()
+        self.assertEqual(client.data_to_send(), [b"\x88\x80\x3c\x3c\x3c\x3c"])
+        self.assertIs(client.state, CLOSING)
+        with self.assertRaises(InvalidState):
+            client.send_continuation(b"Eggs", fin=True)
 
     def test_server_send_close_in_fragmented_message(self):
-        server = Protocol(CLIENT)
+        server = Protocol(SERVER)
         server.send_text(b"Spam", fin=False)
         self.assertFrameSent(server, Frame(OP_TEXT, b"Spam", fin=False))
-        # The spec says: "An endpoint MUST be capable of handling control
-        # frames in the middle of a fragmented message." However, since the
-        # endpoint must not send a data frame after a close frame, a close
-        # frame can't be "in the middle" of a fragmented message.
-        with self.assertRaises(ProtocolError) as raised:
-            server.send_close(CloseCode.NORMAL_CLOSURE)
-        self.assertEqual(str(raised.exception), "expected a continuation frame")
+        server.send_close()
+        self.assertEqual(server.data_to_send(), [b"\x88\x00"])
+        self.assertIs(server.state, CLOSING)
+        with self.assertRaises(InvalidState):
+            server.send_continuation(b"Eggs", fin=True)
 
     def test_client_receive_close_in_fragmented_message(self):
         client = Protocol(CLIENT)
@@ -1392,10 +1388,6 @@ def test_client_receive_close_in_fragmented_message(self):
             client,
             Frame(OP_TEXT, b"Spam", fin=False),
         )
-        # The spec says: "An endpoint MUST be capable of handling control
-        # frames in the middle of a fragmented message." However, since the
-        # endpoint must not send a data frame after a close frame, a close
-        # frame can't be "in the middle" of a fragmented message.
         client.receive_data(b"\x88\x02\x03\xe8")
         self.assertIsInstance(client.parser_exc, ProtocolError)
         self.assertEqual(str(client.parser_exc), "incomplete fragmented message")
@@ -1410,10 +1402,6 @@ def test_server_receive_close_in_fragmented_message(self):
             server,
             Frame(OP_TEXT, b"Spam", fin=False),
         )
-        # The spec says: "An endpoint MUST be capable of handling control
-        # frames in the middle of a fragmented message." However, since the
-        # endpoint must not send a data frame after a close frame, a close
-        # frame can't be "in the middle" of a fragmented message.
         server.receive_data(b"\x88\x82\x00\x00\x00\x00\x03\xe9")
         self.assertIsInstance(server.parser_exc, ProtocolError)
         self.assertEqual(str(server.parser_exc), "incomplete fragmented message")