haproxy-mcp/haproxy_mcp/haproxy_client.py

"""HAProxy Runtime API client functions."""

import socket
import subprocess
import select
import time

from .config import (
    HAPROXY_SOCKET,
    HAPROXY_CONTAINER,
    SOCKET_TIMEOUT,
    SOCKET_RECV_TIMEOUT,
    MAX_RESPONSE_SIZE,
    SUBPROCESS_TIMEOUT,
)
from .exceptions import HaproxyError


def haproxy_cmd(command: str) -> str:
    """Send command to HAProxy Runtime API.

    Args:
        command: The HAProxy runtime API command to execute

    Returns:
        The response from HAProxy

    Raises:
        HaproxyError: If connection fails, times out, or response exceeds size limit
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(SOCKET_TIMEOUT)
            s.connect(HAPROXY_SOCKET)
            s.sendall(f"{command}\n".encode())
            s.shutdown(socket.SHUT_WR)

            # Set socket to non-blocking for select-based recv loop
            s.setblocking(False)
            response = b""
            start_time = time.time()

            while True:
                # Check for overall timeout
                elapsed = time.time() - start_time
                if elapsed >= SOCKET_RECV_TIMEOUT:
                    raise HaproxyError(f"Response timeout after {SOCKET_RECV_TIMEOUT} seconds")

                # Wait for data with timeout (remaining time)
                remaining = SOCKET_RECV_TIMEOUT - elapsed
                ready, _, _ = select.select([s], [], [], min(remaining, 1.0))

                if ready:
                    data = s.recv(8192)
                    if not data:
                        break
                    response += data
                    if len(response) > MAX_RESPONSE_SIZE:
                        raise HaproxyError(f"Response exceeded {MAX_RESPONSE_SIZE} bytes limit")

            return response.decode().strip()
    except socket.timeout:
        raise HaproxyError("Connection timeout")
    except ConnectionRefusedError:
        raise HaproxyError("Connection refused - HAProxy not running?")
    except UnicodeDecodeError:
        raise HaproxyError("Invalid UTF-8 in response")
    except HaproxyError:
        raise
    except (OSError, BlockingIOError, BrokenPipeError) as e:
        raise HaproxyError(f"Socket error: {e}") from e


def haproxy_cmd_checked(command: str) -> str:
    """Send command to HAProxy and raise on error response.

    Args:
        command: HAProxy command to execute

    Returns:
        Command response

    Raises:
        HaproxyError: If HAProxy returns an error message
    """
    result = haproxy_cmd(command)
    _check_response_for_errors(result)
    return result


def _check_response_for_errors(response: str) -> None:
    """Check HAProxy response for error indicators and raise if found.

    HAProxy Runtime API returns plain text responses. Success responses are
    typically empty or contain requested data. Error responses contain
    specific keywords that indicate the command failed.

    Args:
        response: Response string from HAProxy Runtime API command.

    Raises:
        HaproxyError: If response contains any error indicator keyword.

    Error Indicators:
        - "No such": Resource doesn't exist (e.g., backend, server, map entry)
        - "not found": Similar to "No such", resource lookup failed
        - "error": General error in command execution
        - "failed": Operation could not be completed
        - "invalid": Malformed command or invalid parameter value
        - "unknown": Unrecognized command or parameter

    Examples:
        Successful responses (will NOT raise):
            - "" (empty string for successful set commands)
            - "1" (map entry ID after successful add)
            - Server state data (for show commands)

        Error responses (WILL raise HaproxyError):
            - "No such server." - Server doesn't exist in specified backend
            - "No such backend." - Backend name not found
            - "No such map." - Map file not loaded or doesn't exist
            - "Entry not found." - Map entry lookup failed
            - "Invalid server state." - Bad state value for set server state
            - "unknown keyword 'xyz'" - Unrecognized command parameter
            - "failed to allocate memory" - Resource allocation failure
            - "'set server' expects <addr>:<port>" - Invalid command syntax

    Note:
        The check is case-insensitive to catch variations like "Error:",
        "ERROR:", "error:" etc. that HAProxy may return.
    """
    error_indicators = ["No such", "not found", "error", "failed", "invalid", "unknown"]
    if response:
        response_lower = response.lower()
        for indicator in error_indicators:
            if indicator.lower() in response_lower:
                raise HaproxyError(f"HAProxy command failed: {response.strip()}")


def haproxy_cmd_batch(commands: list[str]) -> list[str]:
    """Send multiple commands to HAProxy.

    Note: HAProxy Runtime API only processes the first command when multiple
    commands are sent on a single connection. This function sends each command
    on a separate connection to ensure all commands are executed.

    Args:
        commands: List of HAProxy commands to execute

    Returns:
        List of responses for each command

    Raises:
        HaproxyError: If connection fails or any command returns an error
    """
    if not commands:
        return []

    if len(commands) == 1:
        return [haproxy_cmd_checked(commands[0])]

    # Send each command on separate connection (HAProxy limitation)
    responses = []
    for cmd in commands:
        try:
            resp = haproxy_cmd_checked(cmd)
            responses.append(resp)
        except HaproxyError:
            raise

    return responses


def reload_haproxy() -> tuple[bool, str]:
    """Validate and reload HAProxy configuration.

    Returns:
        Tuple of (success, message)
    """
    try:
        validate = subprocess.run(
            ["podman", "exec", HAPROXY_CONTAINER, "haproxy", "-c", "-f", "/usr/local/etc/haproxy/haproxy.cfg"],
            capture_output=True, text=True, timeout=SUBPROCESS_TIMEOUT
        )
        if validate.returncode != 0:
            return False, f"Config validation failed:\n{validate.stderr}"

        result = subprocess.run(
            ["podman", "kill", "--signal", "USR2", HAPROXY_CONTAINER],
            capture_output=True, text=True, timeout=SUBPROCESS_TIMEOUT
        )
        if result.returncode != 0:
            return False, f"Reload failed: {result.stderr}"
        return True, "OK"
    except subprocess.TimeoutExpired:
        return False, f"Command timed out after {SUBPROCESS_TIMEOUT} seconds"
    except FileNotFoundError:
        return False, "podman command not found"
    except OSError as e:
        return False, f"OS error: {e}"