Adding additional logging, refactoring connection handler into separate class from server

Author: Ken

example/Containerfile

@@ -1,4 +1,4 @@
-FROM docker.io/library/haproxy:latest
+FROM docker.io/library/haproxy:2.3
 
 COPY haproxy/ /usr/local/etc/haproxy/
 

example/Pipfile

@@ -7,6 +7,7 @@ verify_ssl = true
 
 [packages]
 aiohttp = "*"
+haproxyspoa = {editable = true, path = "./.."}
 
 [requires]
 python_version = "3.8"

example/Pipfile.lock

@@ -3,6 +3,7 @@
 from haproxyspoa.payloads.ack import AckPayload
 from haproxyspoa.spoa_server import SpoaServer
 
+
 agent = SpoaServer()
 
 

example/haproxy_spoa.py renamed to example/agent.py

@@ -0,0 +1,15 @@
+import logging
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='[%(asctime)s] [%(levelname)s] %(message)s',
+    datefmt="%Y-%m-%dT%H:%M:%S%z"
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FlowIdLoggerAdapter(logging.LoggerAdapter):
+
+    def process(self, msg, kwargs):
+        return f"[{self.extra['flow_id']}] {msg}", kwargs

haproxyspoa/logging.py

@@ -19,6 +19,9 @@ def support_async(self):
         self.supported_list.add("async")
         return self
 
+    def __str__(self):
+        return ",".join(self.supported_list)
+
 
 class AgentHelloPayload:
     GENEROUS_MAX_FRAME_SIZE = 65536
@@ -32,7 +35,7 @@ def __init__(
     ):
         self.version = spop_version
         self.max_frame_size = max_frame_size
-        self.capabilities = ",".join(capabilities.supported_list) if capabilities is not None else ""
+        self.capabilities = str(capabilities) if capabilities is not None else ""
 
     def to_bytes(self) -> bytes:
         return write_kv_list({

haproxyspoa/payloads/agent_hello.py

@@ -1,9 +1,9 @@
 import asyncio
 import functools
-import logging
 from collections import defaultdict
 from typing import List
 
+from haproxyspoa.logging import logger, FlowIdLoggerAdapter
 from haproxyspoa.payloads.ack import AckPayload
 from haproxyspoa.payloads.agent_disconnect import DisconnectStatusCode, AgentDisconnectPayload
 from haproxyspoa.payloads.agent_hello import AgentHelloPayload, AgentCapabilities
@@ -12,91 +12,116 @@
 from haproxyspoa.payloads.notify import NotifyPayload
 from haproxyspoa.spoa_frame import Frame, AgentHelloFrame, FrameType
 
+import secrets
 
-class SpoaServer:
-
-    def __init__(self):
-        self.handlers = defaultdict(list)
-
-    def handler(self, message_key: str):
-        def _handler(fn):
-            @functools.wraps(fn)
-            def wrapper(*args, **kwargs):
-                return fn(*args, **kwargs)
-            self.handlers[message_key].append(wrapper)
-            return wrapper
-        return _handler
-
-    async def handle_connection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
-        haproxy_hello_frame = await Frame.read_frame(reader)
-        if not haproxy_hello_frame.headers.is_haproxy_hello():
-            await self.send_agent_disconnect(writer)
-            return
-        await self.handle_hello_handshake(haproxy_hello_frame, writer)
-
-        if HaproxyHelloPayload(haproxy_hello_frame.payload).healthcheck():
-            logging.info("It is just a health check, immediately disconnecting")
-            return
 
-        logging.info("Completed new handshake with Haproxy")
+class SpoaConnection:
+    
+    def __init__(self, writer: asyncio.StreamWriter, handlers):
+        self.logger = FlowIdLoggerAdapter(logger, {"flow_id": secrets.token_hex(4)})
+        self.handlers = handlers
+        self.writer = writer
 
-        while True:
-            frame = await Frame.read_frame(reader)
-
-            if frame.headers.is_haproxy_disconnect():
-                await self.handle_haproxy_disconnect(frame)
-                await self.send_agent_disconnect(writer)
-                return
-            elif frame.headers.is_haproxy_notify():
-                await self.handle_haproxy_notify(frame, writer)
-
-    async def handle_haproxy_notify(self, frame: Frame, writer: asyncio.StreamWriter):
+    async def handle_haproxy_notify(self, frame: Frame):
+        self.logger.debug("Incoming `notify` frame from HAProxy")
         notify_payload = NotifyPayload(frame.payload)
 
         response_futures = []
         for msg_key, msg_val in notify_payload.messages.items():
+            self.logger.info(f"Received request on key '{msg_key}'")
             for handler in self.handlers[msg_key]:
                 response_futures.append(handler(**notify_payload.messages[msg_key]))
 
+        self.logger.info(f"Found {len(response_futures)} matching handlers, awaiting response...")
         ack_payloads: List[AckPayload] = await asyncio.gather(*response_futures)
         ack = AckPayload.create_from_all(*ack_payloads)
+        payload = ack.to_bytes()
+
+        self.logger.info(f"Responding with combined payload of {len(payload.getbuffer())} bytes")
 
         ack_frame = Frame(
             frame_type=FrameType.ACK,
             stream_id=frame.headers.stream_id,
             frame_id=frame.headers.frame_id,
             flags=1,
-            payload=ack.to_bytes()
+            payload=payload
         )
-        await ack_frame.write_frame(writer)
+        await ack_frame.write_frame(self.writer)
 
-    async def send_agent_disconnect(self, writer: asyncio.StreamWriter):
+    async def send_agent_disconnect(self):
+        self.logger.info("Agent is now dropping connection")
         disconnect_frame = Frame(
             frame_type=FrameType.AGENT_DISCONNECT,
             flags=1,
             stream_id=0,
             frame_id=0,
             payload=AgentDisconnectPayload().to_buffer()
         )
-        await disconnect_frame.write_frame(writer)
+        await disconnect_frame.write_frame(self.writer)
 
     async def handle_haproxy_disconnect(self, frame: Frame):
         payload = HaproxyDisconnectPayload(frame.payload)
         if payload.status_code() != DisconnectStatusCode.NORMAL:
-            logging.info(f"Haproxy is disconnecting us with status code {payload.status_code()} - `{payload.message()}`")
+            self.logger.info(f"Haproxy is disconnecting us with status code {payload.status_code()} - `{payload.message()}`")
 
-    async def handle_hello_handshake(self, frame: Frame, writer: asyncio.StreamWriter):
+    async def handle_hello_handshake(self, frame: Frame):
+        capabilities = AgentCapabilities()
+        self.logger.info(f"Received `hello handshake`, responding with agent capabilities of: '{capabilities}'")
         agent_hello_frame = AgentHelloFrame(
             payload=AgentHelloPayload(
-                capabilities=AgentCapabilities()
+                capabilities=capabilities,
             ),
             stream_id=frame.headers.stream_id,
             frame_id=frame.headers.frame_id,
         )
-        await agent_hello_frame.write_frame(writer)
+        await agent_hello_frame.write_frame(self.writer)
+
+
+class SpoaServer:
+
+    def __init__(self):
+        self.handlers = defaultdict(list)
+
+    def handler(self, message_key: str):
+        def _handler(fn):
+            @functools.wraps(fn)
+            def wrapper(*args, **kwargs):
+                return fn(*args, **kwargs)
+            self.handlers[message_key].append(wrapper)
+            return wrapper
+        return _handler
+
+    async def handle_connection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
+        conn = SpoaConnection(writer, self.handlers)
+        
+        haproxy_hello_frame = await Frame.read_frame(reader)
+
+        if not haproxy_hello_frame.headers.is_haproxy_hello():
+            conn.logger.error(f"""
+                Expected a `hello` frame from HAProxy,
+                but received unexpected frame of type {haproxy_hello_frame.headers.frame_type}
+            """.strip())
+            await conn.send_agent_disconnect()
+            return
+        await conn.handle_hello_handshake(haproxy_hello_frame)
+
+        if HaproxyHelloPayload(haproxy_hello_frame.payload).healthcheck():
+            conn.logger.info("Health check, immediately disconnecting")
+            return
+
+        while True:
+            frame = await Frame.read_frame(reader)
+
+            if frame.headers.is_haproxy_disconnect():
+                await conn.handle_haproxy_disconnect(frame)
+                await conn.send_agent_disconnect()
+                return
+            elif frame.headers.is_haproxy_notify():
+                await conn.handle_haproxy_notify(frame)
 
     async def _run(self, host: str = "0.0.0.0", port: int = 9002):
         server = await asyncio.start_server(self.handle_connection, host=host, port=port, )
+        logger.info(f"HAProxy SPO Agent listening at {host}:{port}")
         await server.serve_forever()
 
     def run(self, *args, **kwargs):

haproxyspoa/spoa_server.py

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta:__legacy__"

pyproject.toml

@@ -16,11 +16,15 @@
     classifiers=[
         "License :: OSI Approved :: Apache Software License",
         "Programming Language :: Python :: 3",
-        # "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
+        "Development Status :: 2 - Pre-Alpha",
+        "Topic :: System :: Networking",
+        "Topic :: Internet :: WWW/HTTP",
+        "Framework :: AsyncIO",
     ],
-    packages=find_packages(exclude=("example")),
+    packages=find_packages(exclude=("example", "tests")),
     # include_package_data=True,
     install_requires=[],
 )