smart-org/apps/api/app/interfaces/realtime/router.py

"""The WebSocket endpoint — one socket per browser, every live feed and every
RPC on it.

``GET /api/ws?token=<jwt>``

The token travels as a query parameter because the browser WebSocket API cannot
set an ``Authorization`` header. That is a deliberate, contained exception to the
header-only rule: the token is short-lived, the socket is TLS-terminated in any
real deployment, and the alternative (a cookie) would drag CSRF into a protocol
that has no CSRF story.

Protocol — small, symmetric, JSON in both directions. Server->client frames
share the ``type`` discriminator; client->server frames use ``action``::

    client -> {"action": "subscribe",   "channels": ["org:…"]}
    client -> {"action": "unsubscribe", "channels": ["task:…"]}
    client -> {"action": "ping"}
    client -> {"action": "rpc", "id": "…", "method": "POST",
               "path": "/api/organizations", "body": {"name": "…"}}
    server -> {"type": "hello",       "user_id": "…"|null}
    server -> {"type": "subscribed",  "accepted": [...], "rejected": [...]}
    server -> {"type": "event",       "event": {…envelope…}}
    server -> {"type": "rpc_result",  "id": "…", "status": 200, "body": …}
    server -> {"type": "pong"}
    server -> {"type": "error",       "message": "…"}

A connection may be **anonymous** (no token): that is the bootstrap state that
lets a brand-new visitor call ``/api/auth/login`` and redeem an invite over
RPC. Every other RPC and every ``subscribe`` requires an authenticated
connection; the token presented at connect time is re-injected as a bearer
header when the RPC is dispatched into the internal app, so the handlers' own
auth dependencies still run unchanged.

Authorization queries open their own short-lived session: a socket can stay open
for hours, and one SQLAlchemy session per open tab would quietly exhaust the
connection pool.
"""

from __future__ import annotations

import asyncio
import json
import logging

import httpx
from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect

from app.application.ports.realtime import dedupe_channels
from app.core.database import SessionFactory
from app.core.errors import AuthenticationError
from app.core.security import decode_access_token
from app.infrastructure.repositories.identity import UserRepository
from app.interfaces.realtime.authz import authorize_channel
from app.interfaces.realtime.manager import Connection, ConnectionManager

logger = logging.getLogger("agy.realtime")

router = APIRouter()

#: The manager instance this process fans out with. Created at import: the
#: endpoint and the lifecycle hooks can share it without a wiring ceremony.
_manager = ConnectionManager()


def realtime_manager() -> ConnectionManager:
    """The process-wide manager — attached to the bus by ``app.main``."""
    return _manager


#: What an anonymous (not-yet-logged-in) connection may call: login — the only
#: way in now that signup is closed — plus the invite redemption endpoints,
#: which authorize themselves with the invite token. Previewing an invite and
#: creating the account it names must both be possible before the visitor has
#: a session, so a fresh invitee is never stuck at the login screen.
_ANONYMOUS_RPC_PATHS = frozenset({"/api/auth/login"})
_ANONYMOUS_RPC_PREFIXES = ("/api/invites/",)


def _is_anonymous_rpc_allowed(path: str) -> bool:
    return path in _ANONYMOUS_RPC_PATHS or path.startswith(_ANONYMOUS_RPC_PREFIXES)

#: One in-process HTTP client that dispatches RPCs into the internal app. Built
#: lazily so importing this module never pulls the app into existence.
_rpc_client: httpx.AsyncClient | None = None


def _get_rpc_client() -> httpx.AsyncClient:
    global _rpc_client
    if _rpc_client is None:
        # The internal app holds every HTTP handler; it is never socket-bound, so
        # it is reached here via an in-process ASGI transport instead of the wire.
        from app.interfaces.realtime.rpc_app import get_rpc_app

        transport = httpx.ASGITransport(app=get_rpc_app())
        _rpc_client = httpx.AsyncClient(transport=transport, base_url="http://internal")
    return _rpc_client


@router.websocket("/ws")
async def realtime_socket(
    websocket: WebSocket,
    token: str | None = Query(default=None),
) -> None:
    connection = await _authenticate(websocket, token)
    if connection is None:
        return  # _authenticate already closed with a policy-violation code

    await _manager.register(connection)

    _send_frame(
        connection,
        {"type": "hello", "user_id": connection.user_id, "server_time": _now_iso()},
    )

    try:
        while True:
            raw = await websocket.receive_text()
            await _handle_message(connection, raw)
    except WebSocketDisconnect:
        pass
    except Exception:  # noqa: BLE001 - any read failure ends the socket cleanly
        logger.debug("realtime: socket loop ended", exc_info=True)
    finally:
        await _manager.unregister(connection)


async def _authenticate(websocket: WebSocket, token: str | None) -> Connection | None:
    """Accept an anonymous or an authenticated, active user.

    * No token -> an anonymous connection (bootstrap for login/signup).
    * A token that decodes to a real, active user -> an authenticated connection
      whose raw token is kept for RPC re-injection.
    * A presented-but-invalid token -> 1008, so the client can tell "bad token"
      from "network dropped" and stop retrying.
    """
    user = None
    raw_token: str | None = None
    if token:
        try:
            payload = decode_access_token(token)
            user_id = payload.get("sub")
            if user_id:
                # Own session: the socket must not depend on request scope.
                async with SessionFactory() as session:
                    user = await UserRepository(session).get(str(user_id))
                if user is not None and user.is_active:
                    raw_token = token
        except AuthenticationError:
            user = None

    if token and user is None:
        await websocket.close(code=1008, reason="unauthorized")
        return None

    await websocket.accept()
    return Connection(
        user_id=user.id if user else None,
        is_platform_owner=user.is_platform_owner if user else False,
        token=raw_token,
        websocket=websocket,
    )


async def _handle_message(connection: Connection, raw: str) -> None:
    try:
        message = json.loads(raw)
    except (TypeError, ValueError):
        _send_frame(connection, {"type": "error", "message": "malformed JSON frame"})
        return
    if not isinstance(message, dict):
        _send_frame(connection, {"type": "error", "message": "frame must be an object"})
        return

    action = message.get("action")

    if action == "ping":
        _send_frame(connection, {"type": "pong"})
        return

    if action == "rpc":
        # Run off the receive loop: a slow handler must not stall pings.
        asyncio.create_task(_handle_rpc(connection, message))
        return

    # subscribe / unsubscribe require a logged-in user.
    if connection.user_id is None:
        _send_frame(connection, {"type": "error", "message": "authentication required"})
        return

    channels = dedupe_channels(
        str(channel) for channel in (message.get("channels") or []) if isinstance(channel, str)
    )

    if action == "subscribe":
        accepted, rejected = await _subscribe(connection, channels)
        _send_frame(
            connection, {"type": "subscribed", "accepted": accepted, "rejected": rejected}
        )
        return

    if action == "unsubscribe":
        await _manager.unsubscribe(connection, list(channels))
        _send_frame(connection, {"type": "unsubscribed", "channels": channels})
        return

    _send_frame(connection, {"type": "error", "message": f"unknown action: {action!r}"})


async def _handle_rpc(connection: Connection, message: dict) -> None:
    """Dispatch one client RPC into the internal app and return the result."""
    req_id = message.get("id")
    if not req_id:
        _send_frame(connection, {"type": "error", "message": "rpc frame needs an id"})
        return

    method = str(message.get("method") or "GET").upper()
    path = message.get("path")
    if not isinstance(path, str) or not path:
        _send_rpc_result(
            connection,
            req_id,
            400,
            {"error": {"code": "bad_request", "message": "rpc frame needs a path"}},
        )
        return

    # Gate unauthenticated callers to the bootstrap RPCs only.
    if connection.user_id is None and not _is_anonymous_rpc_allowed(path):
        _send_rpc_result(
            connection,
            req_id,
            401,
            {"error": {"code": "unauthenticated", "message": "نیاز به ورود دارید."}},
        )
        return

    body = message.get("body")
    headers = {"X-Request-ID": str(req_id), "User-Agent": "realtime-rpc"}
    if connection.token:
        headers["Authorization"] = f"Bearer {connection.token}"

    try:
        client = _get_rpc_client()
        if isinstance(body, (dict, list)):
            response = await client.request(method, path, json=body, headers=headers)
        elif body is None:
            response = await client.request(method, path, headers=headers)
        else:
            response = await client.request(method, path, content=str(body), headers=headers)
    except Exception:  # noqa: BLE001 - reported to the caller as a 500 envelope
        logger.exception("realtime: rpc dispatch to internal app failed")
        _send_rpc_result(
            connection,
            req_id,
            500,
            {"error": {"code": "internal_error", "message": "خطای داخلی سرور."}},
        )
        return

    status = response.status_code
    try:
        payload = response.json()
    except Exception:
        payload = None
    _send_rpc_result(connection, req_id, status, payload)


def _send_rpc_result(
    connection: Connection, req_id: object, status: int, body: object
) -> None:
    _send_frame(
        connection,
        {"type": "rpc_result", "id": req_id, "status": status, "body": body},
    )


async def _subscribe(
    connection: Connection, channels: tuple[str, ...]
) -> tuple[list[str], list[str]]:
    if not channels:
        return [], []

    async def _authorize(conn: Connection, channel: str) -> bool:
        async with SessionFactory() as session:
            return await authorize_channel(
                session,
                user_id=conn.user_id or "",
                is_platform_owner=conn.is_platform_owner,
                channel=channel,
            )

    return await _manager.subscribe(connection, list(channels), _authorize)


def _now_iso() -> str:
    from app.domain.base import utcnow

    return utcnow().isoformat()


def _send_frame(connection: Connection, payload: dict) -> None:
    """Queue a frame on the connection's single-writer outbox.

    Every frame — control, event and rpc_result — goes through the same queue so
    a slow RPC task can never interleave bytes with the receive loop's own
    frames and corrupt the stream.
    """
    connection.offer(json.dumps(payload, ensure_ascii=False, default=str))