mrok 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

mrok/proxy/asgi.py

@@ -0,0 +1,96 @@
+from collections.abc import Iterator
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import Any, ParamSpec, Protocol
+
+from mrok.types.proxy import ASGIApp, ASGIReceive, ASGISend, Lifespan, Scope
+
+P = ParamSpec("P")
+
+
+class ASGIMiddleware(Protocol[P]):
+    def __call__(
+        self, app: ASGIApp, /, *args: P.args, **kwargs: P.kwargs
+    ) -> ASGIApp: ...  # pragma: no cover
+
+
+class Middleware:
+    def __init__(self, cls: ASGIMiddleware[P], *args: P.args, **kwargs: P.kwargs) -> None:
+        self.cls = cls
+        self.args = args
+        self.kwargs = kwargs
+
+    def __iter__(self) -> Iterator[Any]:
+        as_tuple = (self.cls, self.args, self.kwargs)
+        return iter(as_tuple)
+
+
+class ASGIAppWrapper:
+    def __init__(
+        self,
+        app: ASGIApp,
+        lifespan: Lifespan | None = None,
+    ) -> None:
+        self.app = app
+        self.lifespan = lifespan
+        self.middlware: list[Middleware] = []
+        self.middleare_stack: ASGIApp | None = None
+
+    def add_middleware(self, cls: ASGIMiddleware[P], *args: P.args, **kwargs: P.kwargs):
+        self.middlware.insert(0, Middleware(cls, *args, **kwargs))
+
+    def build_middleware_stack(self):
+        app = self.app
+        for cls, args, kwargs in reversed(self.middlware):
+            app = cls(app, *args, **kwargs)
+        return app
+
+    def get_starlette_lifespan(self):
+        router = getattr(self.app, "router", None)
+        if router is None:
+            return None
+        return getattr(router, "lifespan_context", None)
+
+    @asynccontextmanager
+    async def merge_lifespan(self, app: ASGIApp):
+        async with AsyncExitStack() as stack:
+            state: dict[Any, Any] = {}
+            if self.lifespan is not None:
+                outer_state = await stack.enter_async_context(self.lifespan(app))
+                state.update(outer_state or {})
+            starlette_lifespan = self.get_starlette_lifespan()
+            if starlette_lifespan is not None:
+                inner_state = await stack.enter_async_context(starlette_lifespan(app))
+                state.update(inner_state or {})
+            yield state
+
+    async def handle_lifespan(self, scope: Scope, receive: ASGIReceive, send: ASGISend) -> None:
+        started = False
+        app: Any = scope.get("app")
+        await receive()
+        try:
+            async with self.merge_lifespan(app) as state:
+                if state:
+                    if "state" not in scope:
+                        raise RuntimeError('"state" is unsupported by the current ASGI Server.')
+                    scope["state"].update(state)
+                await send({"type": "lifespan.startup.complete"})
+                started = True
+                await receive()
+        except Exception as e:  # pragma: no cover
+            if started:
+                await send({"type": "lifespan.shutdown.failed", "message": str(e)})
+            else:
+                await send({"type": "lifespan.startup.failed", "message": str(e)})
+            raise
+        else:
+            await send({"type": "lifespan.shutdown.complete"})
+
+    async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend) -> None:
+        if self.middleare_stack is None:  # pragma: no branch
+            self.middleware_stack = self.build_middleware_stack()
+        if scope["type"] == "lifespan":
+            scope["app"] = self
+            await self.handle_lifespan(scope, receive, send)
+            return
+
+        await self.middleware_stack(scope, receive, send)

mrok/proxy/backend.py

@@ -6,8 +6,8 @@ import openziti
 from httpcore import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
 from openziti.context import ZitiContext
 
-from mrok.proxy.exceptions import TargetUnavailableError
-from mrok.proxy.streams import AIONetworkStream
+from mrok.proxy.exceptions import InvalidTargetError, TargetUnavailableError
+from mrok.proxy.stream import AIONetworkStream
 
 
 class AIOZitiNetworkBackend(AsyncNetworkBackend):
@@ -37,7 +37,9 @@ class AIOZitiNetworkBackend(AsyncNetworkBackend):
             reader, writer = await asyncio.open_connection(sock=sock)
             return AIONetworkStream(reader, writer)
         except Exception as e:
-            raise TargetUnavailableError() from e
+            if e.args and e.args[0] == -24:  # the service exists but is not available
+                raise TargetUnavailableError() from e
+            raise InvalidTargetError() from e
 
     async def sleep(self, seconds: float) -> None:
         await asyncio.sleep(seconds)

mrok/proxy/event_publisher.py

@@ -0,0 +1,66 @@
+import asyncio
+import contextlib
+import logging
+
+import zmq
+import zmq.asyncio
+
+from mrok.proxy.asgi import ASGIAppWrapper
+from mrok.proxy.metrics import MetricsCollector
+from mrok.proxy.middleware import CaptureMiddleware, MetricsMiddleware
+from mrok.proxy.models import Event, HTTPResponse, ServiceMetadata, Status
+from mrok.types.proxy import ASGIApp
+
+logger = logging.getLogger("mrok.proxy")
+
+
+class EventPublisher:
+    def __init__(
+        self,
+        worker_id: str,
+        meta: ServiceMetadata | None = None,
+        event_publisher_port: int = 50000,
+        metrics_interval: float = 5.0,
+    ):
+        self._worker_id = worker_id
+        self._meta = meta
+        self._metrics_interval = metrics_interval
+        self.publisher_port = event_publisher_port
+        self._zmq_ctx = zmq.asyncio.Context()
+        self._publisher = self._zmq_ctx.socket(zmq.PUB)
+        self._metrics_collector = MetricsCollector(self._worker_id)
+        self._publish_task = None
+
+    async def on_startup(self):
+        self._publisher.connect(f"tcp://localhost:{self.publisher_port}")
+        self._publish_task = asyncio.create_task(self.publish_metrics_event())
+        logger.info(f"Events publishing for worker {self._worker_id} started")
+
+    async def on_shutdown(self):
+        self._publish_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self._publish_task
+        self._publisher.close()
+        self._zmq_ctx.term()
+        logger.info(f"Events publishing for worker {self._worker_id} stopped")
+
+    async def publish_metrics_event(self):
+        while True:
+            snap = await self._metrics_collector.snapshot()
+            event = Event(type="status", data=Status(meta=self._meta, metrics=snap))
+            await self._publisher.send_string(event.model_dump_json())
+            await asyncio.sleep(self._metrics_interval)
+
+    async def publish_response_event(self, response: HTTPResponse):
+        event = Event(type="response", data=response)
+        await self._publisher.send_string(event.model_dump_json())  # type: ignore[attr-defined]
+
+    def setup_middleware(self, app: ASGIAppWrapper):
+        app.add_middleware(CaptureMiddleware, self.publish_response_event)
+        app.add_middleware(MetricsMiddleware, self._metrics_collector)  # type: ignore
+
+    @contextlib.asynccontextmanager
+    async def lifespan(self, app: ASGIApp):
+        await self.on_startup()  # type: ignore
+        yield
+        await self.on_shutdown()  # type: ignore

mrok/proxy/master.py

@@ -1,5 +1,3 @@
-import asyncio
-import contextlib
 import logging
 import os
 import signal
@@ -10,20 +8,14 @@ from pathlib import Path
 from uuid import uuid4
 
 import zmq
-import zmq.asyncio
-from uvicorn.importer import import_from_string
 from watchfiles import watch
 from watchfiles.filters import PythonFilter
 from watchfiles.run import CombinedProcess, start_process
 
 from mrok.conf import get_settings
 from mrok.logging import setup_logging
-from mrok.proxy.config import MrokBackendConfig
-from mrok.proxy.datastructures import Event, HTTPResponse, Status, ZitiIdentity
-from mrok.proxy.metrics import WorkerMetricsCollector
-from mrok.proxy.middlewares import CaptureMiddleware, LifespanMiddleware, MetricsMiddleware
-from mrok.proxy.server import MrokServer
-from mrok.proxy.types import ASGIApp
+from mrok.proxy.worker import Worker
+from mrok.types.proxy import ASGIApp
 
 logger = logging.getLogger("mrok.agent")
 
@@ -62,76 +54,41 @@ def start_uvicorn_worker(
     worker_id: str,
     app: ASGIApp | str,
     identity_file: str,
+    events_enabled: bool,
     events_pub_port: int,
     metrics_interval: float = 5.0,
 ):
     import sys
 
     sys.path.insert(0, os.getcwd())
-    asgi_app = app if not isinstance(app, str) else import_from_string(app)
 
-    setup_logging(get_settings())
-    identity = ZitiIdentity.load_from_file(identity_file)
-    ctx = zmq.asyncio.Context()
-    pub = ctx.socket(zmq.PUB)
-    pub.connect(f"tcp://localhost:{events_pub_port}")
-    metrics = WorkerMetricsCollector(worker_id)
-
-    task = None
-
-    async def status_sender():  # pragma: no cover
-        while True:
-            snap = await metrics.snapshot()
-            event = Event(type="status", data=Status(meta=identity.mrok, metrics=snap))
-            await pub.send_string(event.model_dump_json())
-            await asyncio.sleep(metrics_interval)
-
-    async def on_startup():  # pragma: no cover
-        nonlocal task
-        await asyncio.sleep(0)
-        task = asyncio.create_task(status_sender())
-
-    async def on_shutdown():  # pragma: no cover
-        await asyncio.sleep(0)
-        if task:
-            task.cancel()
-
-    async def on_response_complete(response: HTTPResponse):  # pragma: no cover
-        event = Event(type="response", data=response)
-        await pub.send_string(event.model_dump_json())
-
-    config = MrokBackendConfig(
-        LifespanMiddleware(
-            MetricsMiddleware(
-                CaptureMiddleware(
-                    asgi_app,
-                    on_response_complete,
-                ),
-                metrics,
-            ),
-            on_startup=on_startup,
-            on_shutdown=on_shutdown,
-        ),
+    worker = Worker(
+        worker_id,
+        app,
         identity_file,
+        events_enabled=events_enabled,
+        event_publisher_port=events_pub_port,
+        metrics_interval=metrics_interval,
     )
-    server = MrokServer(config)
-    with contextlib.suppress(KeyboardInterrupt, asyncio.CancelledError):
-        server.run()
+    worker.run()
 
 
 class MasterBase(ABC):
     def __init__(
         self,
         identity_file: str,
-        workers: int,
-        reload: bool,
-        events_pub_port: int,
-        events_sub_port: int,
+        *,
+        workers: int = 4,
+        reload: bool = False,
+        events_enabled: bool = True,
+        events_pub_port: int = 50000,
+        events_sub_port: int = 50001,
         metrics_interval: float = 5.0,
     ):
         self.identity_file = identity_file
         self.workers = workers
         self.reload = reload
+        self.events_enabled = events_enabled
         self.events_pub_port = events_pub_port
         self.events_sub_port = events_sub_port
         self.metrics_interval = metrics_interval
@@ -171,6 +128,7 @@ class MasterBase(ABC):
                 worker_id,
                 self.get_asgi_app(),
                 self.identity_file,
+                self.events_enabled,
                 self.events_pub_port,
                 self.metrics_interval,
             ),

mrok/proxy/metrics.py

@@ -8,7 +8,7 @@ import time
 import psutil
 from hdrh.histogram import HdrHistogram
 
-from mrok.proxy.datastructures import (
+from mrok.proxy.models import (
     DataTransferMetrics,
     ProcessMetrics,
     RequestsMetrics,
@@ -49,7 +49,7 @@ async def get_process_metrics(interval: float = 0.1) -> ProcessMetrics:
     return await asyncio.to_thread(_collect_process_usage, interval)
 
 
-class WorkerMetricsCollector:
+class MetricsCollector:
     def __init__(self, worker_id: str, lowest=1, highest=60000, sigfigs=3):
         self.worker_id = worker_id
         self.total_requests = 0

mrok/proxy/{middlewares.py → middleware.py}

@@ -2,19 +2,19 @@ import asyncio
 import logging
 import time
 
-from mrok.proxy.constants import MAX_REQUEST_BODY_BYTES, MAX_RESPONSE_BODY_BYTES
-from mrok.proxy.datastructures import FixedSizeByteBuffer, HTTPHeaders, HTTPRequest, HTTPResponse
-from mrok.proxy.metrics import WorkerMetricsCollector
-from mrok.proxy.types import (
+from mrok.proxy.metrics import MetricsCollector
+from mrok.proxy.models import FixedSizeByteBuffer, HTTPHeaders, HTTPRequest, HTTPResponse
+from mrok.types.proxy import (
     ASGIApp,
     ASGIReceive,
     ASGISend,
-    LifespanCallback,
     Message,
     ResponseCompleteCallback,
     Scope,
 )
-from mrok.proxy.utils import must_capture_request, must_capture_response
+
+MAX_REQUEST_BODY_BYTES = 2 * 1024 * 1024
+MAX_RESPONSE_BODY_BYTES = 5 * 1024 * 1024
 
 logger = logging.getLogger("mrok.proxy")
 
@@ -43,7 +43,7 @@ class CaptureMiddleware:
         state = {}
 
         req_buf = FixedSizeByteBuffer(MAX_REQUEST_BODY_BYTES)
-        capture_req_body = must_capture_request(method, req_headers)
+        capture_req_body = method.upper() not in ("GET", "HEAD", "OPTIONS", "TRACE")
 
         request = HTTPRequest(
             method=method,
@@ -69,9 +69,7 @@ class CaptureMiddleware:
                 resp_headers = HTTPHeaders.from_asgi(msg.get("headers", []))
                 state["resp_headers_raw"] = resp_headers
 
-                state["capture_resp_body"] = must_capture_response(resp_headers)
-
-            if state["capture_resp_body"] and msg["type"] == "http.response.body":
+            if msg["type"] == "http.response.body":
                 body = msg.get("body", b"")
                 resp_buf.write(body)
 
@@ -92,14 +90,14 @@ class CaptureMiddleware:
             status=state["status"] or 0,
             headers=state["resp_headers_raw"],
             duration=duration,
-            body=resp_buf.getvalue() if state["capture_resp_body"] else None,
-            body_truncated=resp_buf.overflow if state["capture_resp_body"] else None,
+            body=resp_buf.getvalue(),
+            body_truncated=resp_buf.overflow,
         )
         asyncio.create_task(self._on_response_complete(response))
 
 
 class MetricsMiddleware:
-    def __init__(self, app: ASGIApp, metrics: WorkerMetricsCollector):
+    def __init__(self, app: ASGIApp, metrics: MetricsCollector):
         self.app = app
         self.metrics = metrics
 
@@ -133,32 +131,3 @@ class MetricsMiddleware:
             await self.app(scope, wrapped_receive, wrapped_send)
         finally:
             await self.metrics.on_request_end(start_time, status_code)
-
-
-class LifespanMiddleware:
-    def __init__(
-        self,
-        app,
-        on_startup: LifespanCallback | None = None,
-        on_shutdown: LifespanCallback | None = None,
-    ):
-        self.app = app
-        self.on_startup = on_startup
-        self.on_shutdown = on_shutdown
-
-    async def __call__(self, scope, receive, send):
-        if scope["type"] == "lifespan":
-            while True:
-                event = await receive()
-                if event["type"] == "lifespan.startup":
-                    if self.on_startup:  # pragma: no branch
-                        await self.on_startup()
-                    await send({"type": "lifespan.startup.complete"})
-
-                elif event["type"] == "lifespan.shutdown":
-                    if self.on_shutdown:  # pragma: no branch
-                        await self.on_shutdown()
-                    await send({"type": "lifespan.shutdown.complete"})
-                    break
-        else:
-            await self.app(scope, receive, send)

mrok/proxy/{datastructures.py → models.py}

@@ -1,14 +1,41 @@
 from __future__ import annotations
 
+import base64
+import binascii
 import json
+from collections.abc import Mapping
 from pathlib import Path
-from typing import Literal
+from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic.functional_serializers import PlainSerializer
+from pydantic.functional_validators import PlainValidator
 from pydantic_core import core_schema
 
 
-class ZitiId(BaseModel):
+def serialize_b64(v: bytes) -> str:
+    return base64.b64encode(v).decode("ascii")
+
+
+def deserialize_b64(v):
+    if isinstance(v, bytes):
+        return v
+    if isinstance(v, str):
+        try:
+            return base64.b64decode(v, validate=True)
+        except binascii.Error as e:  # pragma: no branch
+            raise ValueError("Invalid base64 data") from e
+    raise TypeError("Expected bytes or base64 string")  # pragma: no cover
+
+
+Base64Bytes = Annotated[
+    bytes,
+    PlainValidator(deserialize_b64),
+    PlainSerializer(serialize_b64, return_type=str),
+]
+
+
+class X509Credentials(BaseModel):
     key: str
     cert: str
     ca: str
@@ -16,35 +43,34 @@ class ZitiId(BaseModel):
     @field_validator("key", "cert", "ca", mode="before")
     @classmethod
     def strip_pem_prefix(cls, value: str) -> str:
-        if isinstance(value, str) and value.startswith("pem:"):
+        if isinstance(value, str) and value.startswith("pem:"):  # pragma: no branch
             return value[4:]
-        return value
+        return value  # pragma: no cover
 
 
-class ZitiMrokMeta(BaseModel):
+class ServiceMetadata(BaseModel):
     model_config = ConfigDict(extra="ignore")
-    identity: str
     extension: str
     instance: str
     domain: str | None = None
     tags: dict[str, str | bool | None] | None = None
 
 
-class ZitiIdentity(BaseModel):
+class Identity(BaseModel):
     model_config = ConfigDict(extra="ignore")
     zt_api: str = Field(validation_alias="ztAPI")
-    id: ZitiId
+    id: X509Credentials
+    mrok: ServiceMetadata
+    enable_ha: bool = Field(default=False, validation_alias="enableHa")
     zt_apis: str | None = Field(default=None, validation_alias="ztAPIs")
     config_types: str | None = Field(default=None, validation_alias="configTypes")
-    enable_ha: bool = Field(default=False, validation_alias="enableHa")
-    mrok: ZitiMrokMeta | None = None
 
     @staticmethod
-    def load_from_file(path: str | Path) -> ZitiIdentity:
+    def load_from_file(path: str | Path) -> Identity:
         path = Path(path)
         with path.open("r", encoding="utf-8") as f:
             data = json.load(f)
-        return ZitiIdentity.model_validate(data)
+        return Identity.model_validate(data)
 
 
 class FixedSizeByteBuffer:
@@ -77,7 +103,7 @@ class FixedSizeByteBuffer:
 
 
 class HTTPHeaders(dict):
-    def __init__(self, initial=None):
+    def __init__(self, initial: Mapping[Any, Any] | None = None):
         super().__init__()
         if initial:
             for k, v in initial.items():
@@ -131,9 +157,9 @@ class HTTPRequest(BaseModel):
     method: str
     url: str
     headers: HTTPHeaders
-    query_string: bytes
+    query_string: Base64Bytes | None = None
     start_time: float
-    body: bytes | None = None
+    body: Base64Bytes | None = None
     body_truncated: bool | None = None
 
 
@@ -143,7 +169,7 @@ class HTTPResponse(BaseModel):
     status: int
     headers: HTTPHeaders
     duration: float
-    body: bytes | None = None
+    body: Base64Bytes | None = None
     body_truncated: bool | None = None
 
 
@@ -183,7 +209,7 @@ class WorkerMetrics(BaseModel):
 
 class Status(BaseModel):
     type: Literal["status"] = "status"
-    meta: ZitiMrokMeta
+    meta: ServiceMetadata
     metrics: WorkerMetrics
 
 

mrok/proxy/{streams.py → stream.py}

@@ -1,8 +1,24 @@
 import asyncio
+import select
+import sys
+from typing import Any
 
 from httpcore import AsyncNetworkStream
 
-from mrok.proxy.types import ASGIReceive
+from mrok.types.proxy import ASGIReceive
+
+
+def is_readable(sock):  # pragma: no cover
+    # Stolen from
+    # https://github.com/python-trio/trio/blob/20ee2b1b7376db637435d80e266212a35837ddcc/trio/_socket.py#L471C1-L478C31
+
+    # use select.select on Windows, and select.poll everywhere else
+    if sys.platform == "win32":
+        rready, _, _ = select.select([sock], [], [], 0)
+        return bool(rready)
+    p = select.poll()
+    p.register(sock, select.POLLIN)
+    return bool(p.poll(0))
 
 
 class AIONetworkStream(AsyncNetworkStream):
@@ -21,6 +37,13 @@ class AIONetworkStream(AsyncNetworkStream):
         self._writer.close()
         await self._writer.wait_closed()
 
+    def get_extra_info(self, info: str) -> Any:
+        transport = self._writer.transport
+        if info == "is_readable":
+            sock = transport.get_extra_info("socket")
+            return is_readable(sock)
+        return transport.get_extra_info(info)
+
 
 class ASGIRequestBodyStream:
     def __init__(self, receive: ASGIReceive):

mrok/proxy/worker.py

@@ -0,0 +1,64 @@
+import asyncio
+import contextlib
+import logging
+from pathlib import Path
+
+from uvicorn.importer import import_from_string
+
+from mrok.conf import get_settings
+from mrok.logging import setup_logging
+from mrok.proxy.asgi import ASGIAppWrapper
+from mrok.proxy.event_publisher import EventPublisher
+from mrok.proxy.models import Identity
+from mrok.proxy.ziticorn import BackendConfig, Server
+from mrok.types.proxy import ASGIApp
+
+logger = logging.getLogger("mrok.proxy")
+
+
+class Worker:
+    def __init__(
+        self,
+        worker_id: str,
+        app: ASGIApp | str,
+        identity_file: str | Path,
+        *,
+        events_enabled: bool = True,
+        event_publisher_port: int = 50000,
+        metrics_interval: float = 5.0,
+    ):
+        self._worker_id = worker_id
+        self._identity_file = identity_file
+        self._identity = Identity.load_from_file(self._identity_file)
+        self._app = app
+
+        self._events_enabled = events_enabled
+        self._event_publisher = (
+            EventPublisher(
+                worker_id=worker_id,
+                meta=self._identity.mrok,
+                event_publisher_port=event_publisher_port,
+                metrics_interval=metrics_interval,
+            )
+            if events_enabled
+            else None
+        )
+
+    def setup_app(self):
+        app = ASGIAppWrapper(
+            self._app if not isinstance(self._app, str) else import_from_string(self._app),
+            lifespan=self._event_publisher.lifespan if self._events_enabled else None,
+        )
+
+        if self._events_enabled:
+            self._event_publisher.setup_middleware(app)
+        return app
+
+    def run(self):
+        setup_logging(get_settings())
+        app = self.setup_app()
+
+        config = BackendConfig(app, self._identity_file)
+        server = Server(config)
+        with contextlib.suppress(KeyboardInterrupt, asyncio.CancelledError):
+            server.run()