orchestrata_ultimate.gateway.auth_handler

Reference › orchestrata_ultimate › gateway › auth_handler

This module is the connection authentication flow of the gateway (SPEC-014 US-006). server is its only importer: it calls these functions for every new WebSocket before the socket is registered with the session connection manager. The authenticator classes it selects among live in auth; the narrative is in gateway.

create_authenticator is the single seam that maps GatewayConfig.auth_mode to an object: jwt gives TokenAuthenticator, static gives StaticTokenAuthenticator, hmac gives HmacAuthenticator. validate_via_mode then validates an AUTH payload against the selected mode, reading token for jwt and static, or client_id, timestamp, nonce and signature for hmac, and returns a claims dict with at least sub.

authenticate waits for the first frame under AUTH_TIMEOUT_SECONDS. A JSON-RPC 2.0 frame is passed to handle_jsonrpc_auth, which serves the auth.authenticate call used by the VizInt Command Center. Any other frame must be a WSMessage of type AUTH; the module validates it with validate_via_mode, replies with an authenticated message carrying the user id, and returns that id. A timeout or AuthenticationError sends an error through send_error, records the failure with the auth metrics in metrics, and returns None. try_connect_redis and try_connect_graph are lazy-connect helpers for the ring buffer and the audit writer.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/gateway/auth_handler.py

"""Authentication handlers for WebSocket gateway connections.

Extracted from server.py to keep WebSocketGateway under 300 lines.
Contains module-level functions for auth frame handling and
JSON-RPC 2.0 authentication dispatch.
"""

import asyncio
import json
import logging
from typing import Any, cast

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.gateway.auth import (
    HmacAuthenticator,
    StaticTokenAuthenticator,
    TokenAuthenticator,
)
from orchestrata_ultimate.gateway.config import GatewayConfig
from orchestrata_ultimate.gateway.exceptions import AuthenticationError
from orchestrata_ultimate.gateway.metrics import auth_failure, auth_success
from orchestrata_ultimate.gateway.models import (
    MessageType,
    WSMessage,
)

logger = logging.getLogger(__name__)

AUTH_TIMEOUT_SECONDS = 10

#: Any of the three shipped authenticators (SPEC-014 FR-005/006/007) --
#: which one a given ``authenticate()``/``handle_jsonrpc_auth()`` call
#: receives is determined by ``config.auth_mode`` at construction time
#: (``create_authenticator`` below).
Authenticator = TokenAuthenticator | StaticTokenAuthenticator | HmacAuthenticator


def create_authenticator(config: GatewayConfig) -> Authenticator:
    """Construct the authenticator selected by ``config.auth_mode``.

    The one seam that turns SPEC-014's three shipped auth strategies into a
    single object callers can hold onto: "jwt" (default) ->
    :class:`TokenAuthenticator`, "static" (FR-006) ->
    :class:`StaticTokenAuthenticator`, "hmac" (FR-007) ->
    :class:`HmacAuthenticator`. ``GatewayConfig.__post_init__`` already
    validates ``auth_mode`` is one of the three, so no ``else`` branch is
    reachable in production.
    """
    if config.auth_mode == "static":
        return StaticTokenAuthenticator(config)
    if config.auth_mode == "hmac":
        return HmacAuthenticator(config)
    return TokenAuthenticator(config)


@traces(
    spec="014",
    story="US-006",
    reqs=["REQ-014-091", "REQ-014-096"],
)
async def validate_via_mode(
    config: GatewayConfig,
    authenticator: Authenticator,
    payload: dict[str, Any],
) -> dict[str, Any]:
    """Validate AUTH-frame *payload* against the mode-selected authenticator.

    Dispatches on ``config.auth_mode`` (SPEC-014 FR-005/006/007): "jwt"
    (default) validates a JWT via ``TokenAuthenticator.validate_token`` --
    BYTE-IDENTICAL to the call this replaced, so default-mode behaviour is
    unchanged; "static" reads ``payload["token"]`` and validates it via
    :meth:`StaticTokenAuthenticator.authenticate` (FR-006); "hmac" reads
    ``client_id``/``timestamp``/``nonce``/``signature`` from *payload* and
    validates via :meth:`HmacAuthenticator.authenticate` (FR-007).

    Returns:
        Claims dict with at least ``sub``.

    Raises:
        AuthenticationError: On any validation failure.
    """
    if config.auth_mode == "static":
        token = payload.get("token", "")
        return await cast(StaticTokenAuthenticator, authenticator).authenticate(token)
    if config.auth_mode == "hmac":
        return await cast(HmacAuthenticator, authenticator).authenticate(
            client_id=payload.get("client_id", ""),
            timestamp=payload.get("timestamp", ""),
            nonce=payload.get("nonce", ""),
            signature=payload.get("signature", ""),
        )
    token = payload.get("token", "")
    return cast(TokenAuthenticator, authenticator).validate_token(token)


@traces(spec="014", story="US-006", reqs=["REQ-014-089"])
async def send_error(
    websocket: Any,
    error_msg: str,
    correlation_id: str,
) -> None:
    """Send an error frame on *websocket*.

    Args:
        websocket: Target WebSocket connection.
        error_msg: Error description text.
        correlation_id: ID of the originating message.
    """
    resp = WSMessage(
        msg_type=MessageType.ERROR,
        payload={"error": error_msg},
        correlation_id=correlation_id,
    )
    await websocket.send(json.dumps(resp.to_dict()))


@traces(
    spec="014",
    story="US-006",
    reqs=["REQ-014-089"],
)
async def authenticate(
    websocket: Any,
    conn_id: str,
    config: GatewayConfig,
    auth: Authenticator,
    connections: Any,
    ws_map: dict[str, Any],
    handle_jsonrpc_fn: Any,
) -> str | None:
    """Wait for an AUTH frame and validate the token.

    Supports both internal WSMessage AUTH frames and JSON-RPC 2.0
    auth.authenticate calls from the VizInt Command Center. The WSMessage
    AUTH frame path validates via ``config.auth_mode`` (SPEC-014
    FR-005/006/007 -- see :func:`validate_via_mode`), so it supports all
    three shipped auth strategies; the payload shape carries whatever the
    selected mode needs (``token`` for jwt/static, or
    ``client_id``/``timestamp``/``nonce``/``signature`` for hmac).

    Args:
        websocket: The websockets connection object.
        conn_id: Connection identifier.
        config: Gateway configuration.
        auth: The authenticator :func:`create_authenticator` built for
            ``config.auth_mode`` -- must match that mode.
        connections: Connection manager for registration.
        ws_map: WebSocket connection map.
        handle_jsonrpc_fn: Callback for JSON-RPC message handling.

    Returns:
        Authenticated user_id, or None on failure.
    """
    try:
        raw = await asyncio.wait_for(
            websocket.recv(),
            timeout=AUTH_TIMEOUT_SECONDS,
        )
        data = json.loads(raw)

        # JSON-RPC 2.0 auth from VizInt Command Center
        if "jsonrpc" in data:
            return await handle_jsonrpc_auth(
                websocket,
                conn_id,
                data,
                config,
                auth,
                connections,
                ws_map,
                handle_jsonrpc_fn,
            )

        # Internal WSMessage AUTH frame
        msg = WSMessage.from_dict(data)
        if msg.msg_type != MessageType.AUTH:
            await send_error(
                websocket,
                "Expected auth",
                msg.msg_id,
            )
            auth_failure()
            return None
        claims = await validate_via_mode(config, auth, msg.payload)
        user_id: str = str(claims["sub"])
        resp = WSMessage(
            msg_type=MessageType.AUTH,
            payload={
                "status": "authenticated",
                "user_id": user_id,
            },
            correlation_id=msg.msg_id,
        )
        await websocket.send(json.dumps(resp.to_dict()))
        auth_success()
        return user_id
    except (TimeoutError, AuthenticationError) as exc:
        await send_error(websocket, str(exc), "")
        auth_failure()
        return None


@traces(spec="014", story="US-006", reqs=["REQ-014-086", "REQ-014-087"])
async def handle_jsonrpc_auth(
    websocket: Any,
    conn_id: str,
    data: dict[str, Any],
    config: GatewayConfig,
    auth: Authenticator,
    connections: Any,
    ws_map: dict[str, Any],
    handle_jsonrpc_fn: Any,
) -> str | None:
    """Handle JSON-RPC 2.0 authentication message.

    Args:
        websocket: The websockets connection object.
        conn_id: Connection identifier.
        data: Parsed JSON-RPC message dict.
        config: Gateway configuration.
        auth: The authenticator :func:`create_authenticator` built for
            ``config.auth_mode`` -- must match that mode. This path
            previously fabricated a user id and returned authenticated=True
            without ever reading a token (CRITICAL bypass, audit
            2026-07-06).
        connections: Connection manager for registration.
        ws_map: WebSocket connection map.
        handle_jsonrpc_fn: Callback for JSON-RPC dispatch.

    Returns:
        Authenticated user_id, or None on failure.
    """
    method = data.get("method", "")
    msg_id = data.get("id", "")
    if method == "auth.authenticate":
        if config.auth_mode != "jwt":
            # FR-006/FR-007: static/hmac auth is validated over the
            # WSMessage AUTH frame (validate_via_mode), which carries the
            # richer payload shape hmac needs (client_id/timestamp/nonce/
            # signature) -- this VizInt-specific JSON-RPC path only ever
            # carried a single "token" field, so reject cleanly here
            # instead of calling TokenAuthenticator.validate_token on a
            # StaticTokenAuthenticator/HmacAuthenticator (AttributeError).
            await websocket.send(
                json.dumps(
                    {
                        "jsonrpc": "2.0",
                        "id": msg_id,
                        "result": {
                            "authenticated": False,
                            "error": (
                                f"auth_mode={config.auth_mode!r} is not "
                                "supported over JSON-RPC auth.authenticate; "
                                "use the WSMessage AUTH frame"
                            ),
                        },
                    }
                )
            )
            return None
        token = (data.get("params") or {}).get("token", "")
        if token:
            try:
                claims = cast(TokenAuthenticator, auth).validate_token(token)
            except AuthenticationError as exc:
                await websocket.send(
                    json.dumps(
                        {
                            "jsonrpc": "2.0",
                            "id": msg_id,
                            "result": {
                                "authenticated": False,
                                "error": str(exc),
                            },
                        }
                    )
                )
                auth_failure()
                return None
            user_id = str(claims["sub"])
        elif config.allow_anonymous:
            # Anonymous access is an EXPLICIT deployment opt-in, never the
            # silent default.
            user_id = f"vizint-{conn_id[:8]}"
        else:
            await websocket.send(
                json.dumps(
                    {
                        "jsonrpc": "2.0",
                        "id": msg_id,
                        "result": {
                            "authenticated": False,
                            "error": "Missing token",
                        },
                    }
                )
            )
            auth_failure()
            return None
        await websocket.send(
            json.dumps(
                {
                    "jsonrpc": "2.0",
                    "id": msg_id,
                    "result": {
                        "authenticated": True,
                        "client_id": user_id,
                        "scope": "read_write",
                    },
                }
            )
        )
        auth_success()
        return user_id
    # Non-auth JSON-RPC -- only allow if config permits
    if not config.allow_anonymous:
        logger.warning(
            "Anonymous access denied for method %r on %s",
            method,
            conn_id,
        )
        await websocket.close(
            4001,
            "Authentication required",
        )
        return None
    logger.warning(
        "Anonymous access granted for non-auth JSON-RPC method %r on connection %s",
        method,
        conn_id,
    )
    user_id = f"anon-{conn_id[:8]}"
    await connections.register(conn_id, user_id)
    ws_map[conn_id] = websocket
    await handle_jsonrpc_fn(websocket, conn_id, data)
    return user_id


@traces(spec="014", story="US-008", reqs=["REQ-014-121"])
async def try_connect_redis(ring_buffer: Any) -> None:
    """Attempt Redis connection, log warning on failure.

    Args:
        ring_buffer: Redis ring buffer instance.
    """
    try:
        # Bounded: an unresolvable host must not stall gateway boot — DNS
        # resolution of a dead name can hang for minutes (2026-07-15: LLMNR
        # fallback on silently-dropped multicast turned a fast NXDOMAIN into
        # a 120s+ hang and timed out the suite).
        await asyncio.wait_for(ring_buffer.connect(), timeout=10.0)
    except Exception:
        logger.warning("Redis unavailable")


async def try_connect_graph(audit: Any) -> None:
    """Attempt graph connection, log warning on failure.

    Args:
        audit: Audit writer instance.
    """
    try:
        # Bounded for the same reason as try_connect_redis above.
        await asyncio.wait_for(audit.connect(), timeout=10.0)
    except Exception:
        logger.warning("Graph unavailable")