123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231 |
- """Convert between frames and higher-level AMQP methods"""
- # Copyright (C) 2007-2008 Barry Pederson <bp@barryp.org>
- #
- # This library is free software; you can redistribute it and/or
- # modify it under the terms of the GNU Lesser General Public
- # License as published by the Free Software Foundation; either
- # version 2.1 of the License, or (at your option) any later version.
- #
- # This library is distributed in the hope that it will be useful,
- # but WITHOUT ANY WARRANTY; without even the implied warranty of
- # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
- # Lesser General Public License for more details.
- #
- # You should have received a copy of the GNU Lesser General Public
- # License along with this library; if not, write to the Free Software
- # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301
- from __future__ import absolute_import
-
- from collections import defaultdict, deque
- from struct import pack, unpack
-
- from .basic_message import Message
- from .exceptions import AMQPError, UnexpectedFrame
- from .five import range, string
- from .serialization import AMQPReader
-
- __all__ = ['MethodReader']
-
- #
- # MethodReader needs to know which methods are supposed
- # to be followed by content headers and bodies.
- #
- _CONTENT_METHODS = [
- (60, 50), # Basic.return
- (60, 60), # Basic.deliver
- (60, 71), # Basic.get_ok
- ]
-
-
- class _PartialMessage(object):
- """Helper class to build up a multi-frame method."""
-
- def __init__(self, method_sig, args, channel):
- self.method_sig = method_sig
- self.args = args
- self.msg = Message()
- self.body_parts = []
- self.body_received = 0
- self.body_size = None
- self.complete = False
-
- def add_header(self, payload):
- class_id, weight, self.body_size = unpack('>HHQ', payload[:12])
- self.msg._load_properties(payload[12:])
- self.complete = (self.body_size == 0)
-
- def add_payload(self, payload):
- parts = self.body_parts
- self.body_received += len(payload)
- if self.body_received == self.body_size:
- if parts:
- parts.append(payload)
- self.msg.body = bytes().join(parts)
- else:
- self.msg.body = payload
- self.complete = True
- else:
- parts.append(payload)
-
-
- class MethodReader(object):
- """Helper class to receive frames from the broker, combine them if
- necessary with content-headers and content-bodies into complete methods.
-
- Normally a method is represented as a tuple containing
- (channel, method_sig, args, content).
-
- In the case of a framing error, an :exc:`ConnectionError` is placed
- in the queue.
-
- In the case of unexpected frames, a tuple made up of
- ``(channel, ChannelError)`` is placed in the queue.
-
- """
-
- def __init__(self, source):
- self.source = source
- self.queue = deque()
- self.running = False
- self.partial_messages = {}
- self.heartbeats = 0
- # For each channel, which type is expected next
- self.expected_types = defaultdict(lambda: 1)
- # not an actual byte count, just incremented whenever we receive
- self.bytes_recv = 0
- self._quick_put = self.queue.append
- self._quick_get = self.queue.popleft
-
- def _next_method(self):
- """Read the next method from the source, once one complete method has
- been assembled it is placed in the internal queue."""
- queue = self.queue
- put = self._quick_put
- read_frame = self.source.read_frame
- while not queue:
- try:
- frame_type, channel, payload = read_frame()
- except Exception as exc:
- #
- # Connection was closed? Framing Error?
- #
- put(exc)
- break
-
- self.bytes_recv += 1
-
- if frame_type not in (self.expected_types[channel], 8):
- put((
- channel,
- UnexpectedFrame(
- 'Received frame {0} while expecting type: {1}'.format(
- frame_type, self.expected_types[channel]))))
- elif frame_type == 1:
- self._process_method_frame(channel, payload)
- elif frame_type == 2:
- self._process_content_header(channel, payload)
- elif frame_type == 3:
- self._process_content_body(channel, payload)
- elif frame_type == 8:
- self._process_heartbeat(channel, payload)
-
- def _process_heartbeat(self, channel, payload):
- self.heartbeats += 1
-
- def _process_method_frame(self, channel, payload):
- """Process Method frames"""
- method_sig = unpack('>HH', payload[:4])
- args = AMQPReader(payload[4:])
-
- if method_sig in _CONTENT_METHODS:
- #
- # Save what we've got so far and wait for the content-header
- #
- self.partial_messages[channel] = _PartialMessage(
- method_sig, args, channel,
- )
- self.expected_types[channel] = 2
- else:
- self._quick_put((channel, method_sig, args, None))
-
- def _process_content_header(self, channel, payload):
- """Process Content Header frames"""
- partial = self.partial_messages[channel]
- partial.add_header(payload)
-
- if partial.complete:
- #
- # a bodyless message, we're done
- #
- self._quick_put((channel, partial.method_sig,
- partial.args, partial.msg))
- self.partial_messages.pop(channel, None)
- self.expected_types[channel] = 1
- else:
- #
- # wait for the content-body
- #
- self.expected_types[channel] = 3
-
- def _process_content_body(self, channel, payload):
- """Process Content Body frames"""
- partial = self.partial_messages[channel]
- partial.add_payload(payload)
- if partial.complete:
- #
- # Stick the message in the queue and go back to
- # waiting for method frames
- #
- self._quick_put((channel, partial.method_sig,
- partial.args, partial.msg))
- self.partial_messages.pop(channel, None)
- self.expected_types[channel] = 1
-
- def read_method(self):
- """Read a method from the peer."""
- self._next_method()
- m = self._quick_get()
- if isinstance(m, Exception):
- raise m
- if isinstance(m, tuple) and isinstance(m[1], AMQPError):
- raise m[1]
- return m
-
-
- class MethodWriter(object):
- """Convert AMQP methods into AMQP frames and send them out
- to the peer."""
-
- def __init__(self, dest, frame_max):
- self.dest = dest
- self.frame_max = frame_max
- self.bytes_sent = 0
-
- def write_method(self, channel, method_sig, args, content=None):
- write_frame = self.dest.write_frame
- payload = pack('>HH', method_sig[0], method_sig[1]) + args
-
- if content:
- # do this early, so we can raise an exception if there's a
- # problem with the content properties, before sending the
- # first frame
- body = content.body
- if isinstance(body, string):
- coding = content.properties.get('content_encoding', None)
- if coding is None:
- coding = content.properties['content_encoding'] = 'UTF-8'
-
- body = body.encode(coding)
- properties = content._serialize_properties()
-
- write_frame(1, channel, payload)
-
- if content:
- payload = pack('>HHQ', method_sig[0], 0, len(body)) + properties
-
- write_frame(2, channel, payload)
-
- chunk_size = self.frame_max - 8
- for i in range(0, len(body), chunk_size):
- write_frame(3, channel, body[i:i + chunk_size])
- self.bytes_sent += 1
|