handling edge cases in consume

1. too much data
2. not enough data

For #1 we make sure we read no more than `to_read` in any given consume
operation.

For #2 we make sure we have enough bytes on hand to unpack into a Frame,
buffering if we don't.

Signed-off-by: Jesse Jaggars <jjaggars@redhat.com>

Author: jhjaggars

receptor/messages/directive.py

@@ -1,6 +1,6 @@
 import datetime
-import logging
 import json
+import logging
 
 from ..exceptions import UnknownDirective
 from . import envelope
@@ -26,7 +26,7 @@ async def __call__(self, router, inner_env):
         serial = 0
         async for response in responses:
             serial += 1
-            enveloped_response = envelope.InnerEnvelope.make_response(
+            enveloped_response = envelope.Inner.make_response(
                 receptor=router.receptor,
                 recipient=inner_env.sender,
                 payload=response,

receptor/messages/envelope.py

@@ -1,8 +1,6 @@
 import asyncio
 import base64
 import datetime
-import io
-import itertools
 import json
 import logging
 import struct
@@ -61,29 +59,42 @@ class FramedBuffer:
     def __init__(self, loop=None):
         self.q = asyncio.Queue(loop=loop)
         self.header = None
+        self.framebuffer = bytearray()
         self.bb = bytearray()
         self.current_frame = None
         self.to_read = 0
 
     async def put(self, data):
+        logger.debug("put: %s ... %s", data[:16], data[-16:])
         if not self.to_read:
             return await self.handle_frame(data)
         await self.consume(data)
 
     async def handle_frame(self, data):
-        self.current_frame, rest = Frame.from_data(data)
+        try:
+            self.framebuffer += data
+            frame, rest = Frame.from_data(self.framebuffer)
+        except struct.error:
+            return  # We don't have enough data yet
+        else:
+            self.framebuffer = bytearray()
 
-        if self.current_frame.type not in Frame.Types:
+        if frame.type not in Frame.Types:
             raise Exception("Unknown Frame Type")
 
+        self.current_frame = frame
         self.to_read = self.current_frame.length
         await self.consume(rest)
 
     async def consume(self, data):
+        logger.debug("consuming %d bytes; to_read = %d bytes", len(data), self.to_read)
+        data, rest = data[:self.to_read], data[self.to_read:]
         self.to_read -= len(data)
         self.bb += data
         if self.to_read == 0:
             await self.finish()
+        if rest:
+            await self.handle_frame(rest)
 
     async def finish(self):
         if self.current_frame.type == Frame.Types.HEADER:
@@ -124,6 +135,9 @@ def __init__(self, type_, version, length, msg_id, id_):
         self.msg_id = msg_id
         self.id = id_
 
+    def __repr__(self):
+        return f"Frame({self.type}, {self.version}, {self.length}, {self.msg_id}, {self.id})"
+
     def serialize(self):
         return self.fmt.pack(
             bytes([self.type]), bytes([self.version]),
@@ -158,58 +172,6 @@ def join_uuid(hi, lo):
     return (hi << 64) | lo
 
 
-def glue(header, payload):
-    hf = Frame.wrap(header, type_=Frame.Types.HEADER)
-    pf = Frame.wrap(payload, msg_id=hf.msg_id)
-    return hf.serialize() + header + pf.serialize() + payload
-
-
-def gen_chunks(data, header, msg_id=None, chunksize=2 ** 8):
-    if msg_id is None:
-        msg_id = uuid.uuid4().int
-    seq = itertools.count()
-    buf = bytearray(chunksize)
-    bv = memoryview(buf)
-    header = json.dumps(header).encode("utf-8")
-    yield Frame(Frame.Types.HEADER, 1, len(header), msg_id, next(seq)).serialize() + header
-    yield Frame(Frame.Types.PAYLOAD, 1, len(data), msg_id, next(seq)).serialize()
-    buffer = io.BytesIO(data)
-    bytes_read = buffer.readinto(buf)
-    while bytes_read:
-        if bytes_read == chunksize:
-            yield bv.tobytes()
-        else:
-            yield bv[:bytes_read].tobytes()
-        bytes_read = buffer.readinto(buf)
-
-
-class OuterEnvelope:
-    def __init__(self, frame_id, sender, recipient, route_list, inner):
-        self.frame_id = frame_id
-        self.sender = sender
-        self.recipient = recipient
-        self.route_list = route_list
-        self.inner = inner
-        self.inner_obj = None
-
-    async def deserialize_inner(self, receptor):
-        self.inner_obj = await Inner.deserialize(receptor, self.inner)
-
-    @classmethod
-    def from_raw(cls, raw):
-        doc = json.loads(raw)
-        return cls(**doc)
-
-    def serialize(self):
-        return json.dumps(dict(
-            frame_id=self.frame_id,
-            sender=self.sender,
-            recipient=self.recipient,
-            route_list=self.route_list,
-            inner=self.inner
-        ))
-
-
 class Inner:
     def __init__(self, receptor, message_id, sender, recipient, message_type, timestamp,
                  raw_payload, directive=None, in_response_to=None, ttl=None, serial=1,

receptor/protocol.py

@@ -57,7 +57,9 @@ async def watch_queue(self):
                 continue
 
             try:
-                self.transport.write(msg + DELIM)
+                logger.debug("about to write %s ... %s", msg[:8], msg[-8:])
+                self.transport.write(msg)
+                logger.debug("written successfully")
             except Exception:
                 logger.exception("Error received trying to write to %s", self.id)
                 await buffer_obj.put(msg)
@@ -76,7 +78,7 @@ def connection_lost(self, exc):
         self.receptor.remove_connection(self)
 
     def data_received(self, data):
-        logger.debug(data)
+        logger.debug("recv: %s ... %s", data[:16], data[-16:])
         self.loop.create_task(self.incoming_buffer.put(data))
 
     async def wait_greeting(self):
@@ -95,10 +97,12 @@ async def wait_greeting(self):
             self.transport.close()
 
     def handle_handshake(self, data):
+        logger.debug("handle_handshake: %s", data)
         self.id = data["id"]
         self.meta = data.get("meta", {})
         self.receptor.add_connection(self)
         self.loop.create_task(self.watch_queue())
+        logger.debug("starting message_handler: %s", self.incoming_buffer)
         self.loop.create_task(self.receptor.message_handler(self.incoming_buffer))
 
     def send_handshake(self):
@@ -180,14 +184,15 @@ def emit_response(self, response):
 
     def _do_emit_callback(self, fut):
         res = fut.result()
-        self.transport.write(res.encode() + DELIM)
+        logger.debug("_do_emit_callback: %s", res)
+        self.transport.write(res + DELIM)
 
     def data_received(self, data):
         recipient, directive, payload = data.rstrip(DELIM).decode('utf8').split('\n', 2)
         message_id = str(uuid.uuid4())
         logger.info(f'{message_id}: Sending {directive} to {recipient}')
         sent_timestamp = datetime.datetime.utcnow()
-        inner_env = envelope.InnerEnvelope(
+        inner_env = envelope.Inner(
             receptor=self.receptor,
             message_id=message_id,
             sender=self.receptor.node_id,
@@ -212,7 +217,7 @@ def _data_received_callback(self, inner_env, fut):
         try:
             fut.result()
         except Exception as e:
-            err_resp = envelope.InnerEnvelope.make_response(
+            err_resp = envelope.Inner.make_response(
                 receptor=self.receptor,
                 recipient=inner_env.sender,
                 payload=str(e),

receptor/receptor.py

@@ -93,12 +93,18 @@ def update_connections(self, protocol_obj):
         self.update_connection_manifest(protocol_obj.id)
 
     async def message_handler(self, buf):
+        logger.debug("spawning message_handler")
         while True:
-            header, data = await buf.get()
-            if "cmd" in data and data["cmd"] == "ROUTE":
-                await self.handle_route_advertisement(data)
+            try:
+                data = await buf.get()
+            except Exception:
+                logger.exception("message_handler")
             else:
-                await self.handle_message(header, data)
+                logger.debug("message_handler: %s", data)
+                if "cmd" in data.header and data.header["cmd"] == "ROUTE":
+                    await self.handle_route_advertisement(data.header)
+                else:
+                    await self.handle_message(data)
 
     def add_connection(self, protocol_obj):
         self.update_connections(protocol_obj)
@@ -191,17 +197,17 @@ async def handle_response(self, inner):
         else:
             logger.warning(f'Received response to {in_response_to} but no record of sent message.')
 
-    async def handle_message(self, header, msg):
+    async def handle_message(self, msg):
         handlers = dict(
             directive=self.handle_directive,
             response=self.handle_response,
         )
         messages_received_counter.inc()
-        next_hop = self.router.next_hop(header.recipient)
+        next_hop = self.router.next_hop(msg.header["recipient"])
         if next_hop:
-            return await self.router.forward(header, msg, next_hop)
+            return await self.router.forward(msg, next_hop)
 
-        inner = await envelope.Inner.deserialize(self, msg)
+        inner = await envelope.Inner.deserialize(self, msg.payload)
 
         if inner.message_type not in handlers:
             raise exceptions.UnknownMessageType(

receptor/router.py

@@ -97,7 +97,7 @@ def get_nodes(self):
     async def ping_node(self, node_id, callback=log_ping):
         logger.info(f'Sending ping to node {node_id}')
         now = datetime.datetime.utcnow().isoformat()
-        ping_envelope = envelope.InnerEnvelope(
+        ping_envelope = envelope.Inner(
             receptor=self.receptor,
             message_id=str(uuid.uuid4()),
             sender=self.node_id,
@@ -137,17 +137,17 @@ def find_shortest_path(self, to_node_id):
                         mins[next_vertex] = next_total_cost
                         heapq.heappush(heap, (next_total_cost, next_vertex, path))
 
-    async def forward(self, outer_envelope, next_hop):
+    async def forward(self, msg, next_hop):
         """
         Forward a message on to the next hop closer to its destination
         """
         buffer_mgr = self.receptor.config.components_buffer_manager
         buffer_obj = buffer_mgr.get_buffer_for_node(next_hop, self.receptor)
-        outer_envelope.route_list.append(self.node_id)
-        logger.debug(f'Forwarding frame {outer_envelope.frame_id} to {next_hop}')
+        msg.header["route_list"].append(self.node_id)
+        logger.debug(f'Forwarding frame {msg.msg_id} to {next_hop}')
         try:
             route_counter.inc()
-            await buffer_obj.put(outer_envelope.serialize().encode("utf-8"))
+            await buffer_obj.put(msg.serialize())
         except ReceptorBufferError as e:
             logger.exception("Receptor Buffer Write Error forwarding message to {}: {}".format(next_hop, e))
             # TODO: Possible to find another route? This might be a hard failure
@@ -183,7 +183,7 @@ async def send(self, inner_envelope, expected_response=False):
             "recipient": inner_envelope.recipient,
             "route_list": [self.node_id]
         }
-        msg = envelope.FramedMessage(uuid.uuid4().int, header, signed)
+        msg = envelope.FramedMessage(msg_id=uuid.uuid4().int, header=header, payload=signed)
         logger.debug(f'Sending {inner_envelope.message_id} to {inner_envelope.recipient} via {next_node_id}')
         if expected_response and inner_envelope.message_type == 'directive':
             self.response_registry[inner_envelope.message_id] = dict(message_sent_time=inner_envelope.timestamp)

receptor/security/__init__.py

@@ -1,5 +1,6 @@
 import json
 import logging
+
 logger = logging.getLogger(__name__)
 
 
@@ -26,5 +27,4 @@ async def sign_response(self, inner_envelope):
              for attr in ['message_id', 'sender', 'recipient', 'message_type',
                           'timestamp', 'raw_payload', 'directive',
                           'in_response_to', 'ttl', 'serial', 'code']}
-        )
-
+        ).encode("utf-8")

receptor/tests/test_framedbuffer.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from receptor.messages.envelope import Frame, FramedBuffer, gen_chunks
+from receptor.messages.envelope import Frame, FramedBuffer, FramedMessage
 
 
 @pytest.yield_fixture
@@ -12,8 +12,8 @@ def msg_id():
 
 
 @pytest.yield_fixture
-def framed_buffer():
-    return FramedBuffer()
+def framed_buffer(event_loop):
+    return FramedBuffer(loop=event_loop)
 
 
 @pytest.mark.asyncio
@@ -37,21 +37,6 @@ async def test_framedbuffer(framed_buffer, msg_id):
     assert m.payload == payload + payload2
 
 
-@pytest.mark.asyncio
-async def test_gen_chunks():
-
-    b = FramedBuffer()
-
-    header = {"sender": "node1", "recipient": "node2", "route_list": []}
-    payload = b"this is a test with a buffer"
-    for chunk in gen_chunks(payload, header):
-        await b.put(chunk)
-
-    m = await b.get()
-    assert m.header == header
-    assert m.payload == payload
-
-
 @pytest.mark.asyncio
 async def test_hi(msg_id, framed_buffer):
     hi = json.dumps({"cmd": "hi"}).encode("utf-8")
@@ -94,3 +79,33 @@ async def test_command(framed_buffer, msg_id):
     m = await framed_buffer.get()
     assert m.header == cmd
     assert m.payload is None
+
+
+@pytest.mark.asyncio
+async def test_overfull(framed_buffer, msg_id):
+    header = {"foo": "bar"}
+    payload = b'this is a test'
+    msg = FramedMessage(header=header, payload=payload)
+
+    await framed_buffer.put(msg.serialize())
+
+    m = await framed_buffer.get()
+
+    assert m.header == header
+    assert m.payload == payload
+
+
+@pytest.mark.asyncio
+async def test_underfull(framed_buffer, msg_id):
+    header = {"foo": "bar"}
+    payload = b'this is a test'
+    msg = FramedMessage(header=header, payload=payload)
+    b = msg.serialize()
+
+    await framed_buffer.put(b[:10])
+    await framed_buffer.put(b[10:])
+
+    m = await framed_buffer.get()
+
+    assert m.header == header
+    assert m.payload == payload