Fix Basic.Return content frame leak

src/serena/channel.py

@@ -12,8 +12,9 @@
 )
 
 import anyio
+import attr
 import outcome
-from anyio import CancelScope, ClosedResourceError, EndOfStream, Event, Lock
+from anyio import CancelScope, ClosedResourceError, EndOfStream, Event, Lock, WouldBlock
 from anyio.abc import TaskStatus
 from anyio.lowlevel import checkpoint
 from outcome import Error, Value
@@ -29,7 +30,7 @@
 from serena.frame import BodyFrame, Frame
 from serena.message import AMQPEnvelope, AMQPMessage
 from serena.mixin import ChannelLike
-from serena.payloads.header import BasicHeader, ContentHeaderFrame, ContentHeaderPayload
+from serena.payloads.header import BasicHeader, ContentHeaderFrame
 from serena.payloads.method import (
     BasicAckPayload,
     BasicCancelPayload,
@@ -76,6 +77,47 @@
 PayloadType = TypeVar("PayloadType", bound=MethodPayload)
 
 
+@attr.s(slots=True, frozen=True)
+class _ContentMessage:
+    method: MethodFrame[MethodPayload] = attr.ib()
+    header: BasicHeader = attr.ib()
+    body: bytes = attr.ib()
+
+
+class _ContentAssembly:
+    def __init__(self, method: MethodFrame[MethodPayload]) -> None:
+        self.method = method
+        self._header: ContentHeaderFrame | None = None
+        self._body = bytearray()
+
+    def feed(self, frame: ContentHeaderFrame | BodyFrame) -> _ContentMessage | None:
+        if self._header is None:
+            if not isinstance(frame, ContentHeaderFrame):
+                raise AMQPStateError(f"Expected a header frame, got {frame} instead")
+
+            if frame.payload.class_id != self.method.payload.klass:
+                raise AMQPStateError(
+                    f"Class mismatch ({frame.payload.class_id} != {self.method.payload.klass})"
+                )
+
+            self._header = frame
+        else:
+            if not isinstance(frame, BodyFrame):
+                raise AMQPStateError(f"Expected a body frame, got {frame} instead")
+
+            self._body += frame.data
+
+        header = self._header
+        if len(self._body) >= header.payload.full_size:
+            return _ContentMessage(
+                method=self.method,
+                header=header.payload.payload,
+                body=bytes(self._body),
+            )
+
+        return None
+
+
 class Channel(ChannelLike):
     """
     A wrapper around an AMQP channel.
@@ -106,6 +148,7 @@ def __init__(self, channel_id: int, connection: AMQPConnection, stream_buffer_si
         self._delivery_send, self._delivery_receive = anyio.create_memory_object_stream[Frame](
             max_buffer_size=stream_buffer_size
         )
+        self._returned_content: dict[int, _ContentMessage] = {}
 
         self._close_info: ChannelClosePayload | None = None
         self._lock = Lock()
@@ -194,30 +237,30 @@ async def _enqueue_delivery(self, frame: Frame) -> None:
 
         await self._delivery_send.send(frame)
 
+    def _enqueue_return(self, message: _ContentMessage) -> None:
+        """
+        Enqueues a complete Basic.Return message.
+        """
+
+        self._returned_content[id(message.method)] = message
+        try:
+            self._enqueue_regular(message.method)
+        except WouldBlock:
+            self._returned_content.pop(id(message.method), None)
+            raise
+
+    def _pop_return_content(self, frame: MethodFrame[MethodPayload]) -> _ContentMessage | None:
+        return self._returned_content.pop(id(frame), None)
+
     async def _receive_delivery_message(self) -> AMQPMessage | None:
         """
         Receives a single delivery message. This will reassemble a full message into its constituent
         frames.
         """
 
-        method: MethodFrame[MethodPayload] | None = None
-        headers = None
-        body = b""
+        assembly: _ContentAssembly | None = None
 
         while True:
-            # check for successful reassembly first
-            if headers is not None and len(body) >= headers.payload.full_size:
-                assert method is not None, "reached reassembly with invalid method"
-
-                await checkpoint()
-                # hehe payload.payload
-                return AMQPMessage(
-                    channel=self,
-                    envelope=AMQPEnvelope.of(method.payload),
-                    header=headers.payload.payload,
-                    body=body,
-                )
-
             try:
                 next_frame = await self._delivery_receive.receive()
             except EndOfStream:
@@ -226,14 +269,14 @@ async def _receive_delivery_message(self) -> AMQPMessage | None:
 
                 raise UnexpectedCloseError.of(self._close_info) from None
 
-            if method is None:
+            if assembly is None:
                 if not isinstance(next_frame, MethodFrame):
                     raise AMQPStateError(f"Expected a method frame, got {next_frame} instead")
 
                 method_frame: MethodFrame[MethodPayload] = next_frame
 
                 if isinstance(method_frame.payload, BasicGetOkPayload | BasicDeliverPayload):
-                    method = method_frame
+                    assembly = _ContentAssembly(method_frame)
 
                 elif isinstance(method_frame.payload, BasicGetEmptyPayload):
                     return None
@@ -244,24 +287,19 @@ async def _receive_delivery_message(self) -> AMQPMessage | None:
                         f"got {method_frame.payload} instead"
                     )
 
-            elif headers is None:
-                if not isinstance(next_frame, ContentHeaderFrame):
-                    raise AMQPStateError(f"Expected a header frame, got {next_frame} instead")
-
-                # explicit type hint as pycharm incorrectly infers based on the previous if check
-                payload: ContentHeaderPayload = next_frame.payload
-                if payload.class_id != method.payload.klass:
-                    raise AMQPStateError(
-                        f"Class mismatch ({payload.class_id} != {method.payload.klass})"
-                    )
-
-                headers = next_frame
-
             else:
-                if not isinstance(next_frame, BodyFrame):
-                    raise AMQPStateError(f"Expected a body frame, got {next_frame} instead")
-
-                body += next_frame.data
+                if not isinstance(next_frame, ContentHeaderFrame | BodyFrame):
+                    raise AMQPStateError(f"Expected a content frame, got {next_frame} instead")
+
+                message = assembly.feed(next_frame)
+                if message is not None:
+                    await checkpoint()
+                    return AMQPMessage(
+                        channel=self,
+                        envelope=AMQPEnvelope.of(message.method.payload),
+                        header=message.header,
+                        body=message.body,
+                    )
 
     async def _receive_frame(self) -> MethodFrame[MethodPayload]:
         """
@@ -814,11 +852,14 @@ async def basic_publish(
                 logger.trace(f"C#{self.id}: Server ACKed published message")
 
             elif isinstance(payload, BasicReturnPayload):
+                returned_content = self._pop_return_content(response)
                 raise MessageReturnedError(
                     exchange=payload.exchange,
                     routing_key=payload.routing_key,
                     reply_code=payload.reply_code,
                     reply_text=payload.reply_text,
+                    header=returned_content.header if returned_content is not None else None,
+                    body=returned_content.body if returned_content is not None else b"",
                 )
 
             elif isinstance(payload, BasicNackPayload):

src/serena/connection.py

@@ -20,7 +20,7 @@
 from anyio.abc import ByteStream, TaskGroup
 from anyio.lowlevel import checkpoint
 
-from serena.channel import Channel
+from serena.channel import Channel, _ContentAssembly
 from serena.enums import ClassID, ReplyCode
 from serena.exc import (
     AMQPError,
@@ -29,15 +29,16 @@
     InvalidProtocolError,
     UnexpectedCloseError,
 )
-from serena.frame import Frame, FrameType
+from serena.frame import BodyFrame, Frame, FrameType
 from serena.frameparser import NEED_DATA, FrameParser
-from serena.payloads.header import BasicHeader
+from serena.payloads.header import BasicHeader, ContentHeaderFrame
 from serena.payloads.method import (
     BasicDeliverPayload,
     BasicGetEmptyPayload,
     BasicGetOkPayload,
     BasicQOSOkPayload,
     BasicQOSPayload,
+    BasicReturnPayload,
     ChannelCloseOkPayload,
     ChannelClosePayload,
     ChannelOpenPayload,
@@ -159,6 +160,7 @@ def __init__(
 
         # mapping of channel id -> Channel
         self._channel_channels: dict[int, Channel] = {}
+        self._return_assemblies: dict[int, _ContentAssembly] = {}
 
         self._write_lock = Lock()
 
@@ -492,6 +494,7 @@ def _remove_channel(self, channel_id: int, payload: ChannelClosePayload | None)
         """
 
         self._channels[channel_id] = False
+        self._return_assemblies.pop(channel_id, None)
         chan = self._channel_channels.pop(channel_id)
         chan._close(payload)
 
@@ -603,7 +606,10 @@ async def _listen_for_messages(self) -> None:
 
                     continue
 
-                if isinstance(
+                if isinstance(payload, BasicReturnPayload):
+                    self._return_assemblies[channel] = _ContentAssembly(method_frame)
+
+                elif isinstance(
                     payload, BasicDeliverPayload | BasicGetOkPayload | BasicGetEmptyPayload
                 ):
                     # requires special logic
@@ -623,6 +629,21 @@ async def _listen_for_messages(self) -> None:
                 channel = frame.channel_id
                 assert channel != 0, "header frame cannot happen on control channel"
                 channel_object = self._channel_channels[channel]
+
+                if assembly := self._return_assemblies.get(channel):
+                    assert isinstance(frame, ContentHeaderFrame | BodyFrame)
+                    content = assembly.feed(frame)
+                    if content is not None:
+                        self._return_assemblies.pop(channel, None)
+                        try:
+                            channel_object._enqueue_return(content)
+                        except WouldBlock:
+                            logger.warning(
+                                f"Channel #{channel} was not listening for returned message, "
+                                "dropping"
+                            )
+                    continue
+
                 await self._enqueue_frame(channel_object, frame)
 
     async def _listen_wrapper(self) -> None:

src/serena/exc.py

@@ -6,6 +6,7 @@
 from typing_extensions import override
 
 from serena.enums import ReplyCode
+from serena.payloads.header import BasicHeader
 from serena.payloads.method import (
     ChannelClosePayload,
     ConnectionClosePayload,
@@ -52,7 +53,16 @@ class MessageReturnedError(AMQPError):
     Thrown when a message is returned unexpectedly.
     """
 
-    def __init__(self, reply_code: ReplyCode, reply_text: str, exchange: str, routing_key: str):
+    def __init__(
+        self,
+        reply_code: ReplyCode,
+        reply_text: str,
+        exchange: str,
+        routing_key: str,
+        *,
+        header: BasicHeader | None = None,
+        body: bytes = b"",
+    ):
         #: The server-provided error code.
         self.reply_code = reply_code
         #: The server-provided error text.
@@ -64,6 +74,12 @@ def __init__(self, reply_code: ReplyCode, reply_text: str, exchange: str, routin
         #: The routing key the message was using.
         self.routing_key = routing_key
 
+        #: The returned message header.
+        self.header = header or BasicHeader()
+
+        #: The returned message body.
+        self.body = body
+
     @override
     def __str__(self) -> str:
         return (