python-platform 0.1.1__py3-none-any.whl

python_platform/fastapi/realtime/runtime.py

@@ -0,0 +1,423 @@
+"""Application-owned WebSocket registry、writer 与 shutdown runtime。"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import cast
+from uuid import UUID, uuid4
+
+from dishka import AsyncContainer
+from fastapi import WebSocket
+
+from ...application_services.invocation import _InvocationRuntime
+from ...authorization import AccessTokenAuthenticator, CurrentUser, UnauthenticatedError
+from ...options import Options
+from ...realtime import (
+    RealtimeDeliveryReport,
+    RealtimeEnvelope,
+    RealtimePublisher,
+    RealtimeSnapshot,
+)
+from ...realtime.messages import _encode_envelope
+from .connection import (
+    _PLATFORM_PROBLEM_MESSAGE_TYPE,
+    WebSocketConnection,
+    _WebSocketCloseCode,
+)
+from .options import FastApiRealtimeOptions
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class _TextFrame:
+    text: str
+
+
+@dataclass(frozen=True, slots=True)
+class _CloseFrame:
+    code: _WebSocketCloseCode
+
+
+@dataclass(slots=True)
+class _ConnectionState:
+    connection_id: UUID
+    websocket: WebSocket
+    authenticator: AccessTokenAuthenticator
+    access_token: str
+    current_user: CurrentUser
+    correlation_id: str
+    outbound: asyncio.Queue[_TextFrame | _CloseFrame]
+    endpoint_task: asyncio.Task[object]
+    closed: asyncio.Event
+    writer_task: asyncio.Task[None] | None = None
+    watchdog_task: asyncio.Task[None] | None = None
+    consecutive_invalid_messages: int = 0
+    closing: bool = False
+    peer_disconnected: bool = False
+    close_code: int = int(_WebSocketCloseCode.NORMAL)
+
+
+class _RealtimeAdmissionError(Exception):
+    def __init__(self, code: _WebSocketCloseCode) -> None:
+        self.code = code
+        super().__init__(code.name)
+
+
+class _FastApiRealtimeRuntime:
+    """每个 Application 独占的连接 owner；不存在进程级 registry。"""
+
+    __slots__ = (
+        "_accepting",
+        "_connections",
+        "_dropped",
+        "_enqueued",
+        "_invocations",
+        "_lock",
+        "_options",
+        "_user_connections",
+    )
+
+    def __init__(
+        self,
+        options: Options[FastApiRealtimeOptions],
+        invocations: _InvocationRuntime,
+    ) -> None:
+        self._options = options.value
+        self._invocations = invocations
+        self._accepting = False
+        self._connections: dict[UUID, _ConnectionState] = {}
+        self._user_connections: dict[UUID, set[UUID]] = {}
+        self._lock = asyncio.Lock()
+        self._enqueued = 0
+        self._dropped = 0
+
+    @property
+    def options(self) -> FastApiRealtimeOptions:
+        return self._options
+
+    async def start_accepting(self) -> None:
+        async with self._lock:
+            self._accepting = True
+
+    async def connect(
+        self,
+        websocket: WebSocket,
+        session: AsyncContainer,
+        authenticator: AccessTokenAuthenticator,
+        access_token: str,
+        current_user: CurrentUser,
+        correlation_id: str,
+    ) -> WebSocketConnection:
+        if not current_user.is_authenticated or current_user.user_id is None:
+            raise _RealtimeAdmissionError(_WebSocketCloseCode.POLICY_VIOLATION)
+        endpoint_task = asyncio.current_task()
+        if endpoint_task is None:
+            raise RuntimeError("managed WebSocket requires an asyncio Task")
+        state = _ConnectionState(
+            connection_id=uuid4(),
+            websocket=websocket,
+            authenticator=authenticator,
+            access_token=access_token,
+            current_user=current_user,
+            correlation_id=correlation_id,
+            # 多出的一个槽只保留给 close frame；业务容量仍严格等于 Options。
+            outbound=asyncio.Queue(maxsize=self._options.outbound_queue_capacity + 1),
+            endpoint_task=cast(asyncio.Task[object], endpoint_task),
+            closed=asyncio.Event(),
+        )
+        async with self._lock:
+            if not self._accepting:
+                raise _RealtimeAdmissionError(_WebSocketCloseCode.APPLICATION_SHUTDOWN)
+            user_connections = self._user_connections.setdefault(current_user.user_id, set())
+            if (
+                len(self._connections) >= self._options.max_connections
+                or len(user_connections) >= self._options.max_connections_per_user
+            ):
+                if not user_connections:
+                    self._user_connections.pop(current_user.user_id, None)
+                raise _RealtimeAdmissionError(_WebSocketCloseCode.TRY_AGAIN_LATER)
+            self._connections[state.connection_id] = state
+            user_connections.add(state.connection_id)
+        try:
+            await websocket.accept()
+        except BaseException:
+            await self.disconnect(state.connection_id)
+            raise
+        state.writer_task = asyncio.create_task(
+            self._writer(state),
+            name=f"python-platform-ws-writer-{state.connection_id}",
+        )
+        state.watchdog_task = asyncio.create_task(
+            self._watchdog(state),
+            name=f"python-platform-ws-watchdog-{state.connection_id}",
+        )
+        return WebSocketConnection(self, state, session)
+
+    async def disconnect(self, connection_id: UUID) -> None:
+        async with self._lock:
+            state = self._connections.pop(connection_id, None)
+            if state is None:
+                return
+            user_id = state.current_user.user_id
+            if user_id is not None:
+                user_connections = self._user_connections.get(user_id)
+                if user_connections is not None:
+                    user_connections.discard(connection_id)
+                    if not user_connections:
+                        self._user_connections.pop(user_id, None)
+        state.closing = True
+        current_task = asyncio.current_task()
+        tasks = tuple(
+            task
+            for task in (state.writer_task, state.watchdog_task)
+            if task is not None and task is not current_task and not task.done()
+        )
+        for task in tasks:
+            task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        state.access_token = ""
+        state.closed.set()
+
+    async def finish(
+        self,
+        connection_id: UUID,
+        *,
+        code: _WebSocketCloseCode = _WebSocketCloseCode.NORMAL,
+    ) -> None:
+        async with self._lock:
+            state = self._connections.get(connection_id)
+        if state is None:
+            return
+        if state.peer_disconnected:
+            await self.disconnect(connection_id)
+            return
+        await self._request_close(state, code)
+        writer = state.writer_task
+        if writer is not None and not writer.done():
+            try:
+                await asyncio.wait_for(
+                    asyncio.shield(writer),
+                    timeout=self._options.shutdown_timeout_seconds,
+                )
+            except TimeoutError:
+                writer.cancel()
+                await asyncio.gather(writer, return_exceptions=True)
+        await self.disconnect(connection_id)
+
+    async def send_to_connection(
+        self,
+        connection_id: UUID,
+        *,
+        message_type: str,
+        payload: Mapping[str, object],
+    ) -> RealtimeDeliveryReport:
+        async with self._lock:
+            state = self._connections.get(connection_id)
+        states = () if state is None else (state,)
+        return await self._send(states, message_type=message_type, payload=payload)
+
+    async def send_to_user(
+        self,
+        user_id: UUID,
+        *,
+        message_type: str,
+        payload: Mapping[str, object],
+    ) -> RealtimeDeliveryReport:
+        async with self._lock:
+            connection_ids = tuple(self._user_connections.get(user_id, ()))
+            states = tuple(
+                state
+                for connection_id in connection_ids
+                if (state := self._connections.get(connection_id)) is not None
+            )
+        return await self._send(states, message_type=message_type, payload=payload)
+
+    async def snapshot(self) -> RealtimeSnapshot:
+        async with self._lock:
+            return RealtimeSnapshot(
+                accepting=self._accepting,
+                active_connections=len(self._connections),
+                active_users=len(self._user_connections),
+                enqueued=self._enqueued,
+                dropped=self._dropped,
+            )
+
+    async def shutdown(self) -> None:
+        async with self._lock:
+            self._accepting = False
+            states = tuple(self._connections.values())
+        for state in states:
+            await self._request_close(state, _WebSocketCloseCode.APPLICATION_SHUTDOWN)
+        waiters = tuple(asyncio.create_task(state.closed.wait()) for state in states)
+        if waiters:
+            _, pending = await asyncio.wait(
+                waiters,
+                timeout=self._options.shutdown_timeout_seconds,
+            )
+            for waiter in pending:
+                waiter.cancel()
+            await asyncio.gather(*pending, return_exceptions=True)
+        remaining = tuple(state for state in states if not state.closed.is_set())
+        for state in remaining:
+            if not state.endpoint_task.done():
+                state.endpoint_task.cancel()
+        if remaining:
+            await asyncio.gather(
+                *(state.endpoint_task for state in remaining),
+                return_exceptions=True,
+            )
+            for state in remaining:
+                await self.disconnect(state.connection_id)
+
+    async def _send(
+        self,
+        states: tuple[_ConnectionState, ...],
+        *,
+        message_type: str,
+        payload: Mapping[str, object],
+    ) -> RealtimeDeliveryReport:
+        correlation_id = self._invocations.correlation_id() or str(uuid4())
+        text = _encode_envelope(
+            RealtimeEnvelope(
+                message_type=message_type,
+                correlation_id=correlation_id,
+                payload=payload,
+            )
+        )
+        if not states:
+            return RealtimeDeliveryReport(matched=0, enqueued=0, dropped=0)
+        enqueued = 0
+        dropped = 0
+        for state in states:
+            if state.closing:
+                dropped += 1
+                continue
+            try:
+                if state.outbound.qsize() >= self._options.outbound_queue_capacity:
+                    raise asyncio.QueueFull
+                state.outbound.put_nowait(_TextFrame(text))
+                enqueued += 1
+            except asyncio.QueueFull:
+                dropped += 1
+                await self._request_close(
+                    state,
+                    _WebSocketCloseCode.TRY_AGAIN_LATER,
+                    discard_pending=True,
+                )
+        async with self._lock:
+            self._enqueued += enqueued
+            self._dropped += dropped
+        return RealtimeDeliveryReport(
+            matched=len(states),
+            enqueued=enqueued,
+            dropped=dropped,
+        )
+
+    async def _enqueue_problem(
+        self,
+        state: _ConnectionState,
+        *,
+        code: str,
+        status: int,
+    ) -> RealtimeDeliveryReport:
+        return await self._send(
+            (state,),
+            message_type=_PLATFORM_PROBLEM_MESSAGE_TYPE,
+            payload={"code": code, "status": status},
+        )
+
+    async def _request_close(
+        self,
+        state: _ConnectionState,
+        code: _WebSocketCloseCode,
+        *,
+        discard_pending: bool = False,
+    ) -> None:
+        if state.closing:
+            return
+        state.closing = True
+        state.close_code = int(code)
+        evicted = 0
+        if discard_pending:
+            while True:
+                try:
+                    frame = state.outbound.get_nowait()
+                except asyncio.QueueEmpty:
+                    break
+                if isinstance(frame, _TextFrame):
+                    evicted += 1
+        state.outbound.put_nowait(_CloseFrame(code))
+        if evicted:
+            async with self._lock:
+                self._dropped += evicted
+
+    async def _refresh_identity(self, state: _ConnectionState) -> CurrentUser:
+        try:
+            refreshed = await state.authenticator.authenticate_access_token(state.access_token)
+        except Exception as error:
+            _log_connection_failure("authentication", error)
+            await self._request_close(state, _WebSocketCloseCode.POLICY_VIOLATION)
+            raise UnauthenticatedError from None
+        if (
+            not refreshed.is_authenticated
+            or refreshed.user_id != state.current_user.user_id
+            or refreshed.session_id != state.current_user.session_id
+        ):
+            await self._request_close(state, _WebSocketCloseCode.POLICY_VIOLATION)
+            raise UnauthenticatedError
+        state.current_user = refreshed
+        return refreshed
+
+    async def _writer(self, state: _ConnectionState) -> None:
+        try:
+            while True:
+                frame = await state.outbound.get()
+                if isinstance(frame, _TextFrame):
+                    await state.websocket.send_text(frame.text)
+                    continue
+                await state.websocket.close(code=int(frame.code))
+                return
+        except asyncio.CancelledError:
+            raise
+        except Exception as error:
+            _log_connection_failure("writer", error)
+            state.closing = True
+            if not state.peer_disconnected and not state.endpoint_task.done():
+                state.endpoint_task.cancel()
+
+    async def _watchdog(self, state: _ConnectionState) -> None:
+        try:
+            while not state.closing:
+                try:
+                    await asyncio.wait_for(
+                        state.closed.wait(),
+                        timeout=self._options.authentication_refresh_seconds,
+                    )
+                    return
+                except TimeoutError:
+                    pass
+                try:
+                    await self._refresh_identity(state)
+                except UnauthenticatedError:
+                    return
+        except asyncio.CancelledError:
+            raise
+
+
+def _realtime_publisher(runtime: _FastApiRealtimeRuntime) -> RealtimePublisher:
+    return runtime
+
+
+def _log_connection_failure(component: str, error: Exception) -> None:
+    _LOGGER.warning(
+        "managed WebSocket operation failed",
+        extra={
+            "realtime_component": component,
+            "realtime_exception_type": type(error).__name__,
+        },
+    )

python_platform/fastapi/request_context.py

@@ -0,0 +1,244 @@
+"""HTTP bearer authentication、task-local identity 与 request audit bridge。"""
+
+from __future__ import annotations
+
+import logging
+import time
+from datetime import UTC, datetime
+from typing import cast
+from uuid import uuid4
+
+from anyio import CancelScope
+from dishka import AsyncContainer
+from dishka import Scope as DIScope
+from dishka.exceptions import NoFactoryError
+from dishka.integrations.starlette import ContainerMiddleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from ..application import Application
+from ..application_services.invocation import (
+    _close_action,
+    _InvocationRuntime,
+    _raise_scope_failure,
+)
+from ..auditing import AuditSink, RequestAuditRecord
+from ..authorization import AccessTokenAuthenticator, CurrentUser, UnauthenticatedError
+from ..invocation.contracts import _is_safe_correlation_id
+from ..unit_of_work import UnitOfWorkManager
+from .errors import FastApiAdapterError
+
+_LOGGER = logging.getLogger(__name__)
+_CORRELATION_HEADER = b"x-correlation-id"
+_AUTHORIZATION_HEADER = b"authorization"
+
+
+class _RequestScopeMiddleware:
+    """唯一 lease 包住唯一 REQUEST；原生 ASGI 链在相同 task 解析与关闭资源。"""
+
+    def __init__(self, app: ASGIApp, *, application: Application) -> None:
+        self._app = app
+        self._application = application
+        self._other = ContainerMiddleware(app)
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self._other(scope, receive, send)
+            return
+        async with self._application._http_request_lease_for_integration():
+            request = Request(scope, receive, send)
+            container = self._application._container_for_integration()(
+                {Request: request}, scope=DIScope.REQUEST
+            )
+            request.state.dishka_container = container
+            primary: BaseException | None = None
+            try:
+                await self._app(scope, receive, send)
+            except BaseException as error:
+                primary = error
+            finally:
+                # AnyIO 的 level cancellation 不能反复打断资源 finalizer；asyncio 重复取消
+                # 仍由既有 scope close 监督处理。lease 要等所有清理完成后才释放。
+                with CancelScope(shield=True):
+                    primary, failures = await _close_action(container, primary)
+                request.state.dishka_container = None
+            _raise_scope_failure(primary, failures, False)
+
+
+class _RequestTaskMiddleware:
+    """原生 middleware 下游的可信 task 入口，不迁移任何活动数据库状态。"""
+
+    def __init__(self, app: ASGIApp) -> None:
+        self._app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self._app(scope, receive, send)
+            return
+        request = Request(scope)
+        container = cast(AsyncContainer, request.state.dishka_container)
+        runtime = await container.get(_InvocationRuntime)
+        manager = await container.get(UnitOfWorkManager)
+        if manager.current is not None:
+            raise FastApiAdapterError("middleware must finish its UnitOfWork before call_next")
+        user, correlation_id = request.state.python_platform_identity
+        with runtime.trusted_http_task(container, user, correlation_id, scope["method"]):
+            await self._app(scope, receive, send)
+
+
+class _RequestContextMiddleware:
+    __slots__ = ("_app", "_application")
+
+    def __init__(self, app: ASGIApp, *, application: Application) -> None:
+        self._app = app
+        self._application = application
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self._app(scope, receive, send)
+            return
+
+        correlation_id = _correlation_id(scope)
+        occurred_at = datetime.now(UTC)
+        started_ns = time.perf_counter_ns()
+        current_user = CurrentUser.anonymous()
+        status_code = 500
+        failure: BaseException | None = None
+
+        async def send_with_context(message: Message) -> None:
+            nonlocal status_code
+            if message["type"] == "http.response.start":
+                status_code = cast(int, message["status"])
+                headers = list(message.get("headers", ()))
+                if not any(name.lower() == _CORRELATION_HEADER for name, _ in headers):
+                    headers.append((_CORRELATION_HEADER, correlation_id.encode("ascii")))
+                message = {**message, "headers": headers}
+            await send(message)
+
+        # lease 必须包住下游 Dishka REQUEST scope；stop 会等到 response、两级 audit 与
+        # REQUEST finalizer 全部完成后才开始 Module shutdown/root close。
+        container = cast(AsyncContainer, Request(scope).state.dishka_container)
+        runtime = await container.get(_InvocationRuntime)
+        with runtime.call_scope(container), runtime.http_scope(scope["method"]):
+            try:
+                with runtime.authenticating():
+                    current_user = await self._authenticate(scope)
+                Request(scope).state.python_platform_identity = (current_user, correlation_id)
+                with runtime.trusted_http_task(
+                    container, current_user, correlation_id, scope["method"]
+                ):
+                    await self._app(scope, receive, send_with_context)
+            except UnauthenticatedError as original_error:
+                failure = original_error
+                status_code = 401
+                response = JSONResponse(
+                    status_code=401,
+                    content={"detail": "Not authenticated"},
+                    headers={"WWW-Authenticate": "Bearer"},
+                )
+                await response(scope, receive, send_with_context)
+            except BaseException as original_error:
+                failure = original_error
+                raise
+            finally:
+                await self._write_audit(
+                    scope,
+                    current_user=current_user,
+                    correlation_id=correlation_id,
+                    occurred_at=occurred_at,
+                    status_code=status_code,
+                    duration_ns=time.perf_counter_ns() - started_ns,
+                    failure=failure
+                    or getattr(Request(scope).state, "python_platform_transfer_failure", None),
+                )
+
+    async def _authenticate(self, scope: Scope) -> CurrentUser:
+        authorization = _header(scope, _AUTHORIZATION_HEADER)
+        if authorization is None:
+            return CurrentUser.anonymous()
+        scheme, separator, token = authorization.partition(" ")
+        if scheme.casefold() != "bearer" or not separator or not token or token != token.strip():
+            raise UnauthenticatedError
+        try:
+            authenticator = cast(
+                AccessTokenAuthenticator,
+                await cast(AsyncContainer, Request(scope).state.dishka_container).get(
+                    AccessTokenAuthenticator
+                ),
+            )
+        except NoFactoryError:
+            # Host 收到认证凭据却没有验证 authority 时必须 fail closed。
+            raise UnauthenticatedError from None
+        return await authenticator.authenticate_access_token(token)
+
+    async def _write_audit(
+        self,
+        scope: Scope,
+        *,
+        current_user: CurrentUser,
+        correlation_id: str,
+        occurred_at: datetime,
+        status_code: int,
+        duration_ns: int,
+        failure: BaseException | None,
+    ) -> None:
+        try:
+            container = self._application._container_for_integration()
+        except RuntimeError:
+            return
+        try:
+            sink = await container.get(AuditSink)
+        except NoFactoryError:
+            return
+        except Exception as audit_error:
+            _log_audit_failure(audit_error, route_template="<unresolved>")
+            return
+        route = scope.get("route")
+        route_path = getattr(route, "path", None)
+        route_template = route_path if isinstance(route_path, str) else "<unmatched>"
+        try:
+            await sink.write_request(
+                RequestAuditRecord(
+                    occurred_at=occurred_at,
+                    user_id=current_user.user_id,
+                    session_id=current_user.session_id,
+                    correlation_id=correlation_id,
+                    route_template=route_template,
+                    status_code=status_code,
+                    duration_ns=duration_ns,
+                    exception_type=None if failure is None else type(failure).__name__,
+                )
+            )
+        except Exception as audit_error:
+            _log_audit_failure(audit_error, route_template=route_template)
+
+
+def _correlation_id(scope: Scope) -> str:
+    supplied = _header(scope, _CORRELATION_HEADER)
+    if _is_safe_correlation_id(supplied):
+        return supplied
+    return str(uuid4())
+
+
+def _header(scope: Scope, name: bytes) -> str | None:
+    headers = cast(tuple[tuple[bytes, bytes], ...], scope.get("headers", ()))
+    for header_name, raw_value in headers:
+        if header_name.lower() == name:
+            try:
+                return raw_value.decode("latin-1")
+            except UnicodeDecodeError:
+                return None
+    return None
+
+
+def _log_audit_failure(error: Exception, *, route_template: str) -> None:
+    # 仅记录静态类型与 route template；禁止输出 provider exception 文本。
+    _LOGGER.error(
+        "request audit write failed",
+        exc_info=error,
+        extra={
+            "audit_exception_type": type(error).__name__,
+            "audit_route_template": route_template,
+        },
+    )