mrok 0.4.6__py3-none-any.whl → 0.6.0__py3-none-any.whl

mrok/proxy/app.py

@@ -1,101 +1,176 @@
-import asyncio
+import abc
 import logging
-from collections.abc import AsyncGenerator
-from contextlib import asynccontextmanager
-from pathlib import Path
 
-import openziti
-from openziti.context import ZitiContext
+from httpcore import AsyncConnectionPool, Request
 
-from mrok.conf import get_settings
-from mrok.constants import RE_SUBDOMAIN
-from mrok.http.forwarder import BackendUnavailableError, ForwardAppBase, InvalidBackendError
-from mrok.http.pool import ConnectionPool, PoolManager
-from mrok.http.types import Scope, StreamPair
-from mrok.logging import setup_logging
+from mrok.proxy.exceptions import ProxyError
+from mrok.proxy.stream import ASGIRequestBodyStream
+from mrok.types.proxy import ASGIReceive, ASGISend, Scope
 
 logger = logging.getLogger("mrok.proxy")
 
 
-class ProxyError(Exception):
-    pass
+HOP_BY_HOP_HEADERS = [
+    b"connection",
+    b"keep-alive",
+    b"proxy-authenticate",
+    b"proxy-authorization",
+    b"te",
+    b"trailers",
+    b"transfer-encoding",
+    b"upgrade",
+]
 
 
-class ProxyApp(ForwardAppBase):
+class ProxyAppBase(abc.ABC):
     def __init__(
         self,
-        identity_file: str | Path,
         *,
-        read_chunk_size: int = 65536,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
+        retries: int = 0,
     ) -> None:
-        super().__init__(read_chunk_size=read_chunk_size)
-        self._identity_file = identity_file
-        settings = get_settings()
-        self._proxy_wildcard_domain = (
-            settings.proxy.domain
-            if settings.proxy.domain[0] == "."
-            else f".{settings.proxy.domain}"
-        )
-        self._ziti_ctx: ZitiContext | None = None
-        self._pool_manager = PoolManager(self.build_connection_pool)
-
-    def get_target_from_header(self, headers: dict[str, str], name: str) -> str | None:
-        header_value = headers.get(name, "")
-        if self._proxy_wildcard_domain in header_value:
-            if ":" in header_value:
-                header_value, _ = header_value.split(":", 1)
-            return header_value[: -len(self._proxy_wildcard_domain)]
-
-    def get_target_name(self, headers: dict[str, str]) -> str:
-        target = self.get_target_from_header(headers, "x-forwarded-host")
-        if not target:
-            target = self.get_target_from_header(headers, "host")
-        if not target:
-            raise ProxyError("Neither Host nor X-Forwarded-Host contain a valid target name")
-        return target
-
-    def _get_ziti_ctx(self) -> ZitiContext:
-        if self._ziti_ctx is None:
-            ctx, err = openziti.load(str(self._identity_file), timeout=10_000)
-            if err != 0:
-                raise Exception(f"Cannot create a Ziti context from the identity file: {err}")
-            self._ziti_ctx = ctx
-        return self._ziti_ctx
-
-    async def startup(self):
-        setup_logging(get_settings())
-        self._get_ziti_ctx()
-
-    async def shutdown(self):
-        await self._pool_manager.shutdown()
-
-    async def build_connection_pool(self, key: str) -> ConnectionPool:
-        async def connect():
-            sock = self._get_ziti_ctx().connect(key)
-            reader, writer = await asyncio.open_connection(sock=sock)
-            return reader, writer
-
-        return ConnectionPool(
-            pool_name=key,
-            factory=connect,
-            initial_connections=5,
-            max_size=100,
-            idle_timeout=20.0,
-            reaper_interval=5.0,
+        self._pool = self.setup_connection_pool(
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+            retries=retries,
         )
 
-    @asynccontextmanager
-    async def select_backend(
+    @abc.abstractmethod
+    def setup_connection_pool(
         self,
-        scope: Scope,
-        headers: dict[str, str],
-    ) -> AsyncGenerator[StreamPair]:
-        target_name = self.get_target_name(headers)
-        if not target_name or not RE_SUBDOMAIN.fullmatch(target_name):
-            raise InvalidBackendError()
-        pool = await self._pool_manager.get_pool(target_name)
+        max_connections: int | None,
+        max_keepalive_connections: int | None,
+        keepalive_expiry: float | None,
+        retries: int,
+    ) -> AsyncConnectionPool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_upstream_base_url(self, scope: Scope) -> str:
+        raise NotImplementedError()
+
+    async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend) -> None:
+        if scope.get("type") == "lifespan":
+            return
+
+        if scope.get("type") != "http":
+            await self._send_error(send, 500, "Unsupported")
+            return
+
         try:
-            async with pool.acquire() as (reader, writer):
-                yield reader, writer
+            base_url = self.get_upstream_base_url(scope)
+            if base_url.endswith("/"):  # pragma: no cover
+                base_url = base_url[:-1]
+            full_path = self._format_path(scope)
+            url = f"{base_url}{full_path}"
+            method = scope.get("method", "GET").encode()
+            headers = self._prepare_headers(scope)
+
+            body_stream = ASGIRequestBodyStream(receive)
+
+            request = Request(
+                method=method,
+                url=url,
+                headers=headers,
+                content=body_stream,
+            )
+            response = await self._pool.handle_async_request(request)
+            logger.debug(f"connection pool status: {self._pool}")
+            response_headers = []
+            for k, v in response.headers:
+                if k.lower() not in HOP_BY_HOP_HEADERS:
+                    response_headers.append((k, v))
+
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": response.status,
+                    "headers": response_headers,
+                }
+            )
+
+            async for chunk in response.stream:  # type: ignore[union-attr]
+                await send(
+                    {
+                        "type": "http.response.body",
+                        "body": chunk,
+                        "more_body": True,
+                    }
+                )
+
+            await send({"type": "http.response.body", "body": b"", "more_body": False})
+            await response.aclose()
+
+        except ProxyError as pe:
+            await self._send_error(send, pe.http_status, pe.message)
+
         except Exception:
-            raise BackendUnavailableError()
+            logger.exception("Unexpected error in forwarder")
+            await self._send_error(send, 502, "Bad Gateway")
+
+    async def _send_error(self, send: ASGISend, http_status: int, body: str):
+        try:
+            await send({"type": "http.response.start", "status": http_status, "headers": []})
+            await send({"type": "http.response.body", "body": body.encode()})
+        except Exception as e:  # pragma: no cover
+            logger.error(f"Cannot send error response: {e}")
+
+    def _prepare_headers(self, scope: Scope) -> list[tuple[bytes, bytes]]:
+        headers: list[tuple[bytes, bytes]] = []
+        scope_headers = scope.get("headers", [])
+
+        for k, v in scope_headers:
+            if k.lower() not in HOP_BY_HOP_HEADERS:
+                headers.append((k, v))
+
+        self._merge_x_forwarded(headers, scope)
+
+        return headers
+
+    def _find_header(self, headers: list[tuple[bytes, bytes]], name: bytes) -> int | None:
+        """Return index of header `name` in `headers`, or None if missing."""
+        lname = name.lower()
+        for i, (k, _) in enumerate(headers):
+            if k.lower() == lname:
+                return i
+        return None
+
+    def _merge_x_forwarded(self, headers: list[tuple[bytes, bytes]], scope: Scope) -> None:
+        client = scope.get("client")
+        if client:
+            client_ip = client[0].encode()
+            idx = self._find_header(headers, b"x-forwarded-for")
+            if idx is None:
+                headers.append((b"x-forwarded-for", client_ip))
+            else:
+                k, v = headers[idx]
+                headers[idx] = (k, v + b", " + client_ip)
+
+        server = scope.get("server")
+        if server:
+            if self._find_header(headers, b"x-forwarded-host") is None:
+                headers.append((b"x-forwarded-host", server[0].encode()))
+            if server[1] and self._find_header(headers, b"x-forwarded-port") is None:
+                headers.append((b"x-forwarded-port", str(server[1]).encode()))
+
+        # Always set the protocol to https for upstream
+        idx_proto = self._find_header(headers, b"x-forwarded-proto")
+        if idx_proto is None:
+            headers.append((b"x-forwarded-proto", b"https"))
+        else:
+            k, _ = headers[idx_proto]
+            headers[idx_proto] = (k, b"https")
+
+    def _format_path(self, scope: Scope) -> str:
+        raw_path = scope.get("raw_path")
+        if raw_path:
+            return raw_path.decode()
+        q = scope.get("query_string", b"")
+        path = scope.get("path", "/")
+        path_qs = path
+        if q:
+            path_qs += "?" + q.decode()
+        return path_qs

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

@@ -0,0 +1,45 @@
+import asyncio
+from collections.abc import Iterable
+from pathlib import Path
+
+import openziti
+from httpcore import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
+from openziti.context import ZitiContext
+
+from mrok.proxy.exceptions import InvalidTargetError, TargetUnavailableError
+from mrok.proxy.stream import AIONetworkStream
+
+
+class AIOZitiNetworkBackend(AsyncNetworkBackend):
+    def __init__(self, identity_file: str | Path) -> None:
+        self._identity_file = identity_file
+        self._ziti_ctx: ZitiContext | None = None
+
+    def _get_ziti_ctx(self) -> ZitiContext:
+        if self._ziti_ctx is None:
+            ctx, err = openziti.load(str(self._identity_file), timeout=10_000)
+            if err != 0:
+                raise Exception(f"Cannot create a Ziti context from the identity file: {err}")
+            self._ziti_ctx = ctx
+        return self._ziti_ctx
+
+    async def connect_tcp(
+        self,
+        host: str,
+        port: int,
+        timeout: float | None = None,
+        local_address: str | None = None,
+        socket_options: Iterable[SOCKET_OPTION] | None = None,
+    ) -> AsyncNetworkStream:
+        ctx = self._get_ziti_ctx()
+        try:
+            sock = ctx.connect(host)
+            reader, writer = await asyncio.open_connection(sock=sock)
+            return AIONetworkStream(reader, writer)
+        except Exception as 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/exceptions.py

@@ -0,0 +1,22 @@
+from http import HTTPStatus
+
+from httpcore import ConnectError
+
+
+class ProxyError(Exception):
+    def __init__(self, http_status: HTTPStatus, message: str) -> None:
+        self.http_status: HTTPStatus = http_status
+        self.message: str = message
+
+
+class InvalidTargetError(ProxyError):
+    def __init__(self):
+        super().__init__(HTTPStatus.BAD_GATEWAY, "Bad Gateway: invalid target extension.")
+
+
+class TargetUnavailableError(ProxyError, ConnectError):
+    def __init__(self):
+        super().__init__(
+            HTTPStatus.SERVICE_UNAVAILABLE,
+            "Service Unavailable: the target extension is unavailable.",
+        )

mrok/{master.py → proxy/master.py}

@@ -1,6 +1,3 @@
-import asyncio
-import contextlib
-import json
 import logging
 import os
 import signal
@@ -11,21 +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.datastructures import Event, HTTPResponse, Meta, Status
-from mrok.http.config import MrokBackendConfig
-from mrok.http.lifespan import LifespanWrapper
-from mrok.http.middlewares import CaptureMiddleware, MetricsMiddleware
-from mrok.http.server import MrokServer
-from mrok.http.types import ASGIApp
 from mrok.logging import setup_logging
-from mrok.metrics import WorkerMetricsCollector
+from mrok.proxy.worker import Worker
+from mrok.types.proxy import ASGIApp
 
 logger = logging.getLogger("mrok.agent")
 
@@ -52,7 +42,7 @@ def start_events_router(events_pub_port: int, events_sub_port: int):
     try:
         logger.info(f"Events router process started: {os.getpid()}")
         zmq.proxy(frontend, backend)
-    except KeyboardInterrupt:
+    except KeyboardInterrupt:  # pragma: no cover
         pass
     finally:
         frontend.close()
@@ -62,78 +52,43 @@ def start_events_router(events_pub_port: int, events_sub_port: int):
 
 def start_uvicorn_worker(
     worker_id: str,
-    app: ASGIApp,
+    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())
-    if isinstance(app, str):
-        app = import_from_string(app)
 
-    setup_logging(get_settings())
-    identity = json.load(open(identity_file))
-    meta = Meta(**identity["mrok"])
-    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():
-        while True:
-            snap = await metrics.snapshot()
-            event = Event(type="status", data=Status(meta=meta, metrics=snap))
-            await pub.send_string(event.model_dump_json())
-            await asyncio.sleep(metrics_interval)
-
-    async def on_startup():  # noqa
-        nonlocal task
-        task = asyncio.create_task(status_sender())
-
-    async def on_shutdown():  # noqa
-        if task:
-            task.cancel()
-
-    async def on_response_complete(response: HTTPResponse):
-        event = Event(type="response", data=response)
-        await pub.send_string(event.model_dump_json())
-
-    config = MrokBackendConfig(
-        LifespanWrapper(
-            MetricsMiddleware(
-                CaptureMiddleware(
-                    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
@@ -154,7 +109,7 @@ class MasterBase(ABC):
 
     @abstractmethod
     def get_asgi_app(self):
-        pass
+        raise NotImplementedError()
 
     def setup_signals_handler(self):
         for sig in (signal.SIGINT, signal.SIGTERM):
@@ -173,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,
             ),
@@ -181,18 +137,6 @@ class MasterBase(ABC):
         logger.info(f"Worker {worker_id} [{p.pid}] started")
         return p
 
-    def start(self):
-        self.start_events_router()
-        self.start_workers()
-        self.monitor_thread.start()
-
-    def stop(self):
-        if self.monitor_thread.is_alive():
-            logger.debug("Wait for monitor worker to exit")
-            self.monitor_thread.join(timeout=MONITOR_THREAD_JOIN_TIMEOUT)
-        self.stop_workers()
-        self.stop_events_router()
-
     def start_events_router(self):
         self.zmq_pubsub_router_process = start_process(
             start_events_router,
@@ -204,21 +148,32 @@ class MasterBase(ABC):
             None,
         )
 
-    def stop_events_router(self):
-        self.zmq_pubsub_router_process.stop(sigint_timeout=5, sigkill_timeout=1)
-
     def start_workers(self):
         for i in range(self.workers):
             worker_id = self.worker_identifiers[i]
             p = self.start_worker(worker_id)
             self.worker_processes[worker_id] = p
 
+    def start(self):
+        self.start_events_router()
+        self.start_workers()
+        self.monitor_thread.start()
+
+    def stop_events_router(self):
+        self.zmq_pubsub_router_process.stop(sigint_timeout=5, sigkill_timeout=1)
+
     def stop_workers(self):
         for process in self.worker_processes.values():
-            if process.is_alive():
+            if process.is_alive():  # pragma: no branch
                 process.stop(sigint_timeout=5, sigkill_timeout=1)
         self.worker_processes.clear()
 
+    def stop(self):
+        if self.monitor_thread.is_alive():  # pragma: no branch
+            self.monitor_thread.join(timeout=MONITOR_THREAD_JOIN_TIMEOUT)
+        self.stop_workers()
+        self.stop_events_router()
+
     def restart(self):
         self.pause_event.set()
         self.stop_workers()