Fix a race condition when receiving fragmented messages.

When get() was cancelled e.g. due to timing out in the same turn
of the event loop where put() received the next frame, it would
crash.

Also fix a similar issue in the trio implementation.

The sync implementation already handled the race condition.

Fix #1773.

Author: aaugustin

docs/project/changelog.rst

@@ -39,6 +39,12 @@ Improvements
 
 .. _uvloop: https://uvloop.readthedocs.io/
 
+Bug fixes
+.........
+
+* Fixed a race condition between receiving the next frame of a fragmented
+  message and canceling :meth:`~asyncio.connection.Connection.recv`.
+
 .. _17.2:
 
 17.2

src/websockets/asyncio/messages.py

@@ -3,7 +3,7 @@
 import asyncio
 import codecs
 import collections
-from collections.abc import AsyncIterator, Iterable
+from collections.abc import AsyncIterator, Sequence
 from typing import Any, Callable, Generic, Literal, TypeVar, overload
 
 from ..exceptions import ConcurrencyError
@@ -54,11 +54,10 @@ async def get(self, block: bool = True) -> T:
                 self.get_waiter = None
         return self.queue.popleft()
 
-    def reset(self, items: Iterable[T]) -> None:
-        """Put back items into an empty, idle queue."""
-        assert self.get_waiter is None, "cannot reset() while get() is running"
-        assert not self.queue, "cannot reset() while queue isn't empty"
-        self.queue.extend(items)
+    def put_back(self, items: Sequence[T]) -> None:
+        """Put back items at the front of the queue."""
+        assert self.get_waiter is None, "cannot put_back() while get() is running"
+        self.queue.extendleft(reversed(items))
 
     def abort(self) -> None:
         """Close the queue, raising EOFError in get() if necessary."""
@@ -167,9 +166,11 @@ async def get(self, decode: bool | None = None) -> Data:
                 try:
                     frame = await self.frames.get(not self.closed)
                 except asyncio.CancelledError:
-                    # Put frames already received back into the queue
-                    # so that future calls to get() can return them.
-                    self.frames.reset(frames)
+                    # Put frames already received back into the queue so that
+                    # future calls to get() can return them. The queue may be
+                    # non-empty if put() added frames in the same iteration of
+                    # the event loop where get() was canceled.
+                    self.frames.put_back(frames)
                     raise
                 self.maybe_resume()
                 assert frame.opcode is CONT

src/websockets/trio/messages.py

@@ -123,14 +123,18 @@ async def get(self, decode: bool | None = None) -> Data:
                 try:
                     frame = await self.recv_frames.receive()
                 except trio.Cancelled:
-                    # Put frames already received back into the queue
-                    # so that future calls to get() can return them.
-                    # Bypass the statistics() method for performance.
+                    # Put frames already received back into the queue so that
+                    # future calls to get() can return them. The queue may be
+                    # non-empty if put() added at least two frames before get()
+                    # was canceled: the first one goes to the waiting receive()
+                    # and the second one is buffered.
                     state = self.send_frames._state
                     assert not state.receive_tasks, "no task should receive"
-                    assert not state.data, "queue should be empty"
-                    for frame in frames:
-                        self.send_frames.send_nowait(frame)
+                    # Since the queue is unbounded, we don't need to assert that
+                    # state.send_tasks is empty. It's always empty.
+                    # Now, since no tasks are sending or receiving, it should be
+                    # safe to mutate the underlying deque directly.
+                    state.data.extendleft(reversed(frames))
                     raise
                 except trio.EndOfChannel:
                     raise EOFError("stream of frames ended")

tests/asyncio/test_messages.py

@@ -38,12 +38,19 @@ async def test_get_then_put(self):
         item = await getter_task
         self.assertEqual(item, 42)
 
-    async def test_reset(self):
-        """reset sets the content of the queue."""
-        self.queue.reset([42])
+    async def test_put_back(self):
+        """put_back sets the content of an empty queue."""
+        self.queue.put_back([42])
         item = await self.queue.get()
         self.assertEqual(item, 42)
 
+    async def test_put_back_non_empty_queue(self):
+        """put_back restores items at the start of a non-empty queue."""
+        self.queue.put(3)
+        self.queue.put_back([1, 2])
+        items = [await self.queue.get() for _ in range(3)]
+        self.assertEqual(items, [1, 2, 3])
+
     async def test_abort(self):
         """abort throws an exception in get."""
         getter_task = asyncio.create_task(self.queue.get())
@@ -212,6 +219,21 @@ async def test_cancel_get_after_first_frame(self):
         message = await self.assembler.get()
         self.assertEqual(message, "café")
 
+    async def test_cancel_get_after_first_frame_with_buffered_frames(self):
+        """get can be canceled safely when more frames were put meanwhile."""
+        self.assembler.put(Frame(TEXT, b"ca", fin=False))
+
+        getter_task = asyncio.create_task(self.assembler.get())
+        await asyncio.sleep(0)  # let the event loop start getter_task
+        self.assembler.put(Frame(CONT, b"f\xc3", fin=False))
+        self.assembler.put(Frame(CONT, b"\xa9"))
+        getter_task.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await getter_task
+
+        message = await self.assembler.get()
+        self.assertEqual(message, "café")
+
     # Test get_iter.
 
     async def test_get_iter_text_message_already_received(self):

tests/trio/test_messages.py

@@ -205,6 +205,20 @@ async def test_cancel_get_after_first_frame(self):
         message = await self.assembler.get()
         self.assertEqual(message, "café")
 
+    async def test_cancel_get_after_first_frame_with_buffered_frames(self):
+        """get can be canceled safely when more frames were put meanwhile."""
+        self.assembler.put(Frame(TEXT, b"ca", fin=False))
+
+        async with trio.open_nursery() as nursery:
+            nursery.start_soon(self.assembler.get)
+            await trio.testing.wait_all_tasks_blocked()
+            self.assembler.put(Frame(CONT, b"f\xc3", fin=False))
+            self.assembler.put(Frame(CONT, b"\xa9"))
+            nursery.cancel_scope.cancel()
+
+        message = await self.assembler.get()
+        self.assertEqual(message, "café")
+
     # Test get_iter.
 
     async def test_get_iter_text_message_already_received(self):