spoe-forge 0.0.1__py3-none-any.whl

spoe_forge/__init__.py

@@ -0,0 +1,15 @@
+"""SPOE Forge - A Python framework for building HAProxy SPOA agents."""
+
+from spoe_forge.spoe_forge import SpoeForge
+from spoe_forge.spop.spop_types import SetVarAction, UnsetVarAction, Action
+from spoe_forge.spop.constants import ActionScope
+from spoe_forge.agent.context import AgentContext
+
+__all__ = [
+    "SpoeForge",
+    "AgentContext",
+    "Action",
+    "ActionScope",
+    "SetVarAction",
+    "UnsetVarAction",
+]

spoe_forge/agent/context.py

@@ -0,0 +1,53 @@
+import logging
+
+from spoe_forge.spop.spop_types import SpoaDataType
+
+logger = logging.getLogger(__name__)
+
+
+class AgentContext:
+    """
+    Context object passed to message handlers containing message data.
+
+    Provides access to message name and arguments from HAProxy NOTIFY frames.
+    """
+
+    def __init__(self, message: str, args: dict[str, SpoaDataType]):
+        """
+        Create a new agent context.
+
+        :param str message: Message name from NOTIFY frame
+        :param dict[str, SpoaDataType] args: Message arguments
+        """
+        self._message = message
+        self._args = args
+
+    def get_arg(self, arg: str, default: SpoaDataType = None) -> SpoaDataType:
+        """
+        Retrieve a message argument by name.
+
+        :param str arg: Argument name to retrieve
+        :param default: Default value to return if argument is not present.
+        :return: Argument value, or None if not found
+        """
+        try:
+            return self._args[arg]
+        except KeyError:
+            return default
+
+    def get_args(self) -> dict[str, SpoaDataType]:
+        """
+        Retrieve all message arguments.
+
+        :return: Dictionary mapping argument names to values
+        """
+        return self._args
+
+    def has_arg(self, arg: str) -> bool:
+        """
+        Check if an argument exists in the message.
+
+        :param str arg: Argument name to check
+        :return: True if argument exists, False otherwise
+        """
+        return arg in self._args

spoe_forge/agent/exceptions.py

@@ -0,0 +1,7 @@
+from spoe_forge.exception import SpoeForgeError
+
+
+class SpoeAgentError(SpoeForgeError):
+    """Exception raised when agent message processing fails."""
+
+    pass

spoe_forge/agent/registry.py

@@ -0,0 +1,67 @@
+import asyncio
+import logging
+from typing import Callable
+
+from spoe_forge.agent.context import AgentContext
+from spoe_forge.agent.exceptions import SpoeAgentError
+from spoe_forge.spop.spop_types import SpoaDataType, Action
+
+logger = logging.getLogger(__name__)
+
+MessageHandlerFunc = Callable[[AgentContext], list[Action]]
+
+
+class AgentRegistry:
+    """
+    Internal registry mapping message names to handler functions.
+
+    Note: Users don't interact with this directly - it's used by the Agent class.
+    """
+
+    def __init__(self):
+        """Initialize a new agent registry."""
+        self._handlers: dict[str, MessageHandlerFunc] = {}
+        self._validation_cache: dict[str, bool] = {}
+
+    def register(self, message_name: str, handler: MessageHandlerFunc) -> None:
+        """
+        Register a handler for a message type.
+
+        :param str message_name: Message name to handle
+        :param MessageHandlerFunc handler: User-defined handler function
+        """
+        if message_name in self._handlers:
+            logger.warning(f"Overwriting existing handler for message '{message_name}'")
+
+        self._handlers[message_name] = handler
+        logger.debug(f"Registered handler for message '{message_name}'")
+
+    async def handle_message(
+        self, message: str, args: dict[str, SpoaDataType]
+    ) -> list[Action]:
+        """
+        Route a message to its registered handler.
+
+        Wraps synchronous handlers with asyncio.to_thread for async compatibility.
+
+        :param str message: Message name to handle
+        :param dict[str, SpoaDataType] args: Message arguments
+        :return: List of actions from handler (empty if no handler registered)
+        :raises SpoeAgentError: If handler raises an exception
+        """
+        handler = self._handlers.get(message)
+
+        if handler is None:
+            logger.warning(f"No handler registered for message '{message}'")
+            return []
+
+        ctx = AgentContext(message, args)
+        try:
+            return await asyncio.to_thread(handler, ctx)
+        except Exception as e:
+            logger.error(
+                f"Error in handler for message '{message}': {e}", exc_info=True
+            )
+            # Convert to our own error - will be caught in our own handlers to
+            # cause a graceful kill of the connection
+            raise SpoeAgentError(e)

spoe_forge/exception.py

@@ -0,0 +1,4 @@
+class SpoeForgeError(Exception):
+    """Base exception class for all SPOE Forge errors."""
+
+    pass

spoe_forge/log.py

@@ -0,0 +1,60 @@
+import logging
+import sys
+
+
+_default_handler = logging.StreamHandler(sys.stderr)
+_default_handler.setFormatter(
+    logging.Formatter("%(levelname)s | %(asctime)s | %(name)s | %(message)s")
+)
+
+
+def _has_level_handler(logger: logging.Logger) -> bool:
+    """
+    Check if there is a handler in the logging chain that will handle the
+    given logger's effective level.
+
+    This traverses the logger hierarchy to check if any parent logger
+    has a handler configured. Based on Flask's implementation.
+
+    :param logger: Logger to check
+    :return: True if a handler exists in the chain, False otherwise
+    """
+    level = logger.getEffectiveLevel()
+    current = logger
+
+    while current:
+        if any(handler.level <= level for handler in current.handlers):
+            return True
+
+        if not current.propagate:
+            break
+
+        current = current.parent
+
+    return False
+
+
+def create_logger(debug: bool = False) -> logging.Logger:
+    """
+    Get the SPOE Forge logger and configure it if needed.
+
+    This follows Flask's pattern:
+    - Uses 'spoe_forge' as the logger name
+    - Only adds default handler if no handler exists in the hierarchy
+    - Respects existing logging configuration
+
+    Internal function - called by spoe_forge.py at import time.
+
+    :return: Configured logger instance
+    """
+    logger = logging.getLogger("spoe_forge")
+
+    # Only add handler if none exists in the hierarchy
+    if not _has_level_handler(logger):
+        logger.addHandler(_default_handler)
+
+    # Set default level if not already set
+    if not logger.level:
+        logger.setLevel(logging.DEBUG if debug else logging.INFO)
+
+    return logger

spoe_forge/server/configuration.py

@@ -0,0 +1,135 @@
+import logging
+from math import floor
+
+from spoe_forge.server.constants import DEFAULT_MAX_FRAME_SIZE
+from spoe_forge.server.constants import SPOE_CAPABILITIES
+from spoe_forge.server.constants import SPOE_VERSION
+
+logger = logging.getLogger(__name__)
+
+
+class ServerConfiguration:
+    """
+    Manages SPOP protocol negotiation for a single connection.
+
+    Fresh instance created per connection to negotiate version, capabilities,
+    and max frame size with HAProxy.
+    """
+
+    _server_compatible: bool
+
+    _max_frame_size: int
+    _capabilities: list[str]
+    _version: str
+
+    def __init__(self, max_frame_size: int = DEFAULT_MAX_FRAME_SIZE) -> None:
+        """
+        Initialize server configuration with defaults.
+
+        :param int max_frame_size: Maximum frame size to negotiate with HAProxy
+        """
+        self._version = SPOE_VERSION
+        self._capabilities = []
+
+        self._max_frame_size = (
+            max_frame_size  # Set to default until negotiation complete
+        )
+        self._server_compatible = True
+
+    async def _check_version_compatibility(self, ha_versions: list[str]):
+        """
+        Check if server version is compatible with HAProxy versions.
+
+        :param list[str] ha_versions: Versions supported by HAProxy
+        """
+        float_ver = float(self._version)
+        for version in ha_versions:
+            float_ha_ver = float(version)
+            if floor(float_ha_ver) != floor(float_ha_ver):
+                continue
+
+            if float_ver <= float_ha_ver:
+                self._server_compatible = True
+                return
+
+        self._server_compatible = False
+        logger.error(
+            f"No compatible versions found for server compatibility. "
+            f"Agent supports {self._version}. "
+            f"HAProxy supports {', '.join(ha_versions)}."
+        )
+
+    async def _find_max_frame_size(self, ha_max_frame_size: int):
+        """
+        Negotiate maximum frame size with HAProxy.
+
+        Takes the minimum of server and HAProxy frame sizes.
+
+        :param int ha_max_frame_size: Maximum frame size supported by HAProxy
+        """
+        self._max_frame_size = min(ha_max_frame_size, self._max_frame_size)
+
+        if self.max_frame_size <= 0:
+            logger.error(
+                f"Frame size negotiation failed: negotiated size {self._max_frame_size} is not positive"
+            )
+            self._server_compatible = False
+            return
+
+    async def _find_common_capabilities(self, ha_capabilities: list[str]):
+        """
+        Find common capabilities between server and HAProxy.
+
+        :param list[str] ha_capabilities: Capabilities supported by HAProxy
+        """
+        self._capabilities = list(set(ha_capabilities) & set(SPOE_CAPABILITIES))
+
+        if "pipelining" not in self._capabilities:
+            logger.warning(
+                "Pipelining capability is not supported by HAProxy. "
+                "Req/Resp cycles may be slower than they could be."
+            )
+
+        # We don't set compatability here as pipelining is the only option available, and we support it w/o
+        # changes to the system - instead just log a warning that HAProxy won't take advantage of it and
+        # move on
+
+    async def negotiate_server_compatibility(
+        self,
+        supported_versions: list[str],
+        ha_max_frame_size: int,
+        ha_capabilities: list[str],
+    ) -> None:
+        """
+        Negotiate protocol compatibility with HAProxy.
+
+        Checks version, frame size, and capabilities to determine if connection
+        can proceed.
+
+        :param list[str] supported_versions: Versions supported by HAProxy
+        :param int ha_max_frame_size: Maximum frame size supported by HAProxy
+        :param list[str] ha_capabilities: Capabilities supported by HAProxy
+        """
+        await self._check_version_compatibility(supported_versions)
+        await self._find_max_frame_size(ha_max_frame_size)
+        await self._find_common_capabilities(ha_capabilities)
+
+    @property
+    def is_compatible(self) -> bool:
+        """Check if negotiation succeeded and connection is compatible."""
+        return self._server_compatible
+
+    @property
+    def version(self) -> str:
+        """Get negotiated SPOP protocol version."""
+        return self._version
+
+    @property
+    def max_frame_size(self) -> int:
+        """Get negotiated maximum frame size."""
+        return self._max_frame_size
+
+    @property
+    def capabilities(self) -> list[str]:
+        """Get negotiated capabilities shared by server and HAProxy."""
+        return self._capabilities

spoe_forge/server/constants.py

@@ -0,0 +1,33 @@
+from enum import IntEnum
+
+SPOE_VERSION = "2.0"
+"""Hardcoded to the version of SPOP we built around"""
+
+SPOE_CAPABILITIES = ["pipelining"]
+"""Pipelining is the only capability available in 2.0"""
+
+DEFAULT_MAX_FRAME_SIZE = 1024 * 4
+"""4kb max frame size - set to a comfortably low value"""
+
+
+class DisconnectCode(IntEnum):
+    # Predefined by HAProxy Protocol
+    NORMAL = 0
+    IO_ERROR = 1
+    TIMEOUT = 2
+    FRAME_TOO_BIG = 3
+    INVALID_FRAME_RECEIVED = 4
+    VERSION_NOT_FOUND = 5
+    MAX_FRAME_SIZE_NOT_FOUND = 6
+    CAPABILITY_NOT_FOUND = 7
+    UNSUPPORTED_VERSION = 8
+    MAX_FRAME_SIZE_OUT_OF_RANGE = 9
+    FRAGMENTATION_NOT_SUPPORTED = 10
+    INVALID_INTERLACED_FRAME = 11
+    FRAME_ID_NOT_FOUND = 12
+    RESOURCE_ALLOCATION_ERROR = 13
+    UNKNOWN_ERROR = 99
+
+    # Custom Error Codes
+    SERVER_INCOMPATIBLE = 101
+    PROTOCOL_ERROR = 102

spoe_forge/server/handler.py

@@ -0,0 +1,245 @@
+import logging
+from asyncio import StreamReader
+from asyncio import StreamWriter
+from typing import Callable
+from typing import Awaitable
+
+from spoe_forge.exception import SpoeForgeError
+from spoe_forge.server.configuration import ServerConfiguration
+from spoe_forge.server.constants import DisconnectCode
+from spoe_forge.spop.constants import FrameType
+from spoe_forge.spop.exception import SpopEOFError
+from spoe_forge.spop.frame import Disconnect
+from spoe_forge.spop.frame import Frame
+from spoe_forge.spop.frame import HaproxyHello
+from spoe_forge.spop.frame import Notify
+from spoe_forge.spop.spop_types import Action, Messages
+
+logger = logging.getLogger(__name__)
+
+
+class ForgeHandler:
+    """
+    Handles SPOP protocol lifecycle for a single connection.
+
+    Manages handshake, NOTIFY/ACK cycles, and disconnection for one HAProxy connection.
+    """
+
+    def __init__(
+        self,
+        notify_handler: Callable[[Messages], Awaitable[list[Action]]],
+        config: ServerConfiguration,
+        reader: StreamReader,
+        writer: StreamWriter,
+    ):
+        """
+        Initialize connection handler.
+
+        :param notify_handler: Callback from Forge to process Messages into Actions
+        :param ServerConfiguration config: Configuration for this connection
+        :param StreamReader reader: AsyncIO stream reader for connection
+        :param StreamWriter writer: AsyncIO stream writer for connection
+        """
+        self.notify_handler = notify_handler
+        self.config = config
+        self.reader = reader
+        self.writer = writer
+
+    async def close_connection(self):
+        """Close the connection stream and wait for it to close."""
+        if not self.writer.is_closing():
+            self.writer.close()
+
+        await self.writer.wait_closed()
+        logger.debug("Stream disconnected")
+
+    async def send_frame(self, frame: Frame) -> bool:
+        """
+        Encode and send a frame to HAProxy.
+
+        :param Frame frame: Frame to encode and send
+        :return: True if frame was sent successfully, False otherwise
+        """
+        if self.writer.is_closing():
+            logger.warning(
+                f"Could not send frame {frame.frame_type.name} - stream closed"
+            )
+            return False
+
+        try:
+            self.writer.write(await frame.encode(self.config.max_frame_size))
+        except SpoeForgeError as e:
+            logger.warning(
+                f"Failed to write frame to stream - {frame.frame_type.name} - {e}"
+            )
+            return False
+
+        await self.writer.drain()
+        return True
+
+    async def send_disconnect(self, status_code: DisconnectCode, message: str) -> bool:
+        """
+        Send AGENT_DISCONNECT frame to HAProxy.
+
+        :param DisconnectCode status_code: Disconnect reason code
+        :param str message: Human-readable disconnect message
+        :return: True if disconnect frame was sent successfully, False otherwise
+        """
+        err_frame = await Frame.construct(
+            FrameType.AGENT_DISCONNECT,
+            stream_id=0,
+            frame_id=0,
+            status_code=status_code,
+            message=message,
+        )
+
+        return await self.send_frame(err_frame)
+
+    async def send_disconnect_on_error(
+        self, status_code: DisconnectCode, message: str
+    ) -> bool:
+        """
+        Send AGENT_DISCONNECT frame and log as error.
+
+        :param DisconnectCode status_code: Disconnect reason code
+        :param str message: Human-readable disconnect message
+        :return: True if disconnect frame was sent successfully, False otherwise
+        """
+        logger.error(
+            f"SPOA server encountered an error, disconnecting: {status_code.name}: {message}"
+        )
+        return await self.send_disconnect(status_code, message)
+
+    async def handle_handshake(self) -> bool:
+        """
+        Handle SPOP handshake phase with HAProxy.
+
+        Receives HAPROXY_HELLO, negotiates compatibility, and sends AGENT_HELLO.
+        Closes connection immediately if healthcheck flag is set.
+
+        :return: True if handshake succeeded and processing should continue, False otherwise
+        """
+        frame = await Frame.decode(self.reader)
+        if not isinstance(frame, HaproxyHello):
+            await self.send_disconnect_on_error(
+                status_code=DisconnectCode.INVALID_FRAME_RECEIVED,
+                message=f"Expected HAPROXY-HELLO, received {frame.frame_type.name}",
+            )
+            return False
+
+        await self.config.negotiate_server_compatibility(
+            frame.supported_versions, frame.max_frame_size, frame.capabilities
+        )
+
+        if not self.config.is_compatible:
+            await self.send_disconnect_on_error(
+                status_code=DisconnectCode.SERVER_INCOMPATIBLE,
+                message="Handshake failed to find compatibility.",
+            )
+            return False
+
+        agent_hello = await Frame.construct(
+            FrameType.AGENT_HELLO,
+            stream_id=0,
+            frame_id=0,
+            version=self.config.version,
+            max_frame_size=self.config.max_frame_size,
+            capabilities=self.config.capabilities,
+        )
+
+        res = await self.send_frame(agent_hello)
+        if frame.healthcheck:
+            # Health check means we close the connection immediately after sending AGENT-HELLO
+            logger.debug("Healthcheck connection - closing after AGENT-HELLO")
+            return False
+
+        if res:
+            logger.debug(
+                f"Connection established - SPOP {self.config.version}, "
+                f"frame_size={self.config.max_frame_size}, "
+                f"capabilities={','.join(self.config.capabilities)}"
+            )
+
+        return res
+
+    async def handle_notify_cycle(self) -> bool:
+        """
+        Handle a single NOTIFY/ACK cycle or HAPROXY_DISCONNECT.
+
+        Receives NOTIFY frame, calls agent to process messages, and sends ACK with actions.
+        Also handles graceful HAPROXY_DISCONNECT frames and EOF conditions.
+
+        :return: True if cycle completed successfully, False if connection should close
+        """
+        try:
+            frame = await Frame.decode(self.reader)
+        except SpopEOFError:
+            logger.debug("Stream disconnected with EOF")
+            return False
+
+        # Handle graceful disconnect from HAProxy
+        if isinstance(frame, Disconnect):
+            if frame.status_code == DisconnectCode.NORMAL:
+                logger.debug("Connection closed gracefully by HAProxy")
+            else:
+                logger.warning(
+                    f"Received HAPROXY_DISCONNECT: {frame.message} (status: {frame.status_code})"
+                )
+
+            # Respond with AGENT_DISCONNECT
+            await self.send_disconnect(
+                status_code=DisconnectCode.NORMAL,
+                message="Disconnecting normally",
+            )
+            return False
+
+        if not isinstance(frame, Notify):
+            await self.send_disconnect_on_error(
+                status_code=DisconnectCode.INVALID_FRAME_RECEIVED,
+                message=f"Expected NOTIFY or HAPROXY_DISCONNECT, received {frame.frame_type.name}",
+            )
+            return False
+
+        actions = await self.notify_handler(frame.messages)
+
+        ack = await Frame.construct(
+            FrameType.ACK,
+            stream_id=frame.metadata.stream_id,
+            frame_id=frame.metadata.frame_id,
+            actions=actions,
+        )
+
+        return await self.send_frame(ack)
+
+    async def core_handler(self):
+        """
+        Main connection lifecycle handler.
+
+        Executes handshake followed by NOTIFY/ACK processing loop until connection closes.
+        Handles protocol errors and connection resets gracefully.
+        """
+        try:
+            if not await self.handle_handshake():
+                await self.close_connection()
+                return
+
+            while True:
+                if not await self.handle_notify_cycle():
+                    await self.close_connection()
+                    break
+
+        except SpoeForgeError as e:
+            if not await self.send_disconnect_on_error(
+                status_code=DisconnectCode.PROTOCOL_ERROR,
+                message=str(e),
+            ):
+                logger.error(
+                    f"Failed to send disconnect on error while handling: {str(e)}",
+                    exc_info=True,
+                )
+
+            await self.close_connection()
+
+        except ConnectionResetError:
+            # Expected case from HAProxy - we treat it as a graceful disconnect
+            logger.debug("Connection reset by HAProxy")