mrok 0.2.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

mrok/controller/routes/extensions.py

@@ -14,8 +14,8 @@ from mrok.ziti.errors import (
     ServiceAlreadyRegisteredError,
     ServiceNotFoundError,
 )
-from mrok.ziti.identities import register_instance, unregister_instance
-from mrok.ziti.services import register_extension, unregister_extension
+from mrok.ziti.identities import register_identity, unregister_identity
+from mrok.ziti.services import register_service, unregister_service
 
 logger = logging.getLogger("mrok.controller")
 
@@ -83,7 +83,7 @@ async def create_extension(
     ],
 ):
     try:
-        service = await register_extension(settings, mgmt_api, data.extension.id, data.tags)
+        service = await register_service(settings, mgmt_api, data.extension.id, data.tags)
         return ExtensionRead(
             id=service["id"],
             name=service["name"],
@@ -149,7 +149,7 @@ async def delete_instance_by_id_or_extension_id(
     id_or_extension_id: str,
 ):
     try:
-        await unregister_extension(settings, mgmt_api, id_or_extension_id)
+        await unregister_service(settings, mgmt_api, id_or_extension_id)
     except ServiceNotFoundError:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
@@ -203,6 +203,7 @@ async def get_extensions(
     tags=["Instances"],
 )
 async def create_extension_instances(
+    settings: AppSettings,
     mgmt_api: ZitiManagementAPI,
     client_api: ZitiClientAPI,
     id_or_extension_id: str,
@@ -223,8 +224,8 @@ async def create_extension_instances(
     ],
 ):
     service = await fetch_extension_or_404(mgmt_api, id_or_extension_id)
-    identity, identity_file = await register_instance(
-        mgmt_api, client_api, service["name"], data.instance.id, data.tags
+    identity, identity_file = await register_identity(
+        settings, mgmt_api, client_api, service["name"], data.instance.id, data.tags
     )
     return InstanceRead(
         id=identity["id"],
@@ -299,10 +300,11 @@ async def get_instance_by_id_or_instance_id(
     tags=["Instances"],
 )
 async def delete_instance_by_id_or_instance_id(
+    settings: AppSettings,
     mgmt_api: ZitiManagementAPI,
     id_or_extension_id: str,
     id_or_instance_id: str,
 ):
     identity = await fetch_instance_or_404(mgmt_api, id_or_extension_id, id_or_instance_id)
     instance_id, extension_id = identity["name"].split(".")
-    await unregister_instance(mgmt_api, extension_id, instance_id)
+    await unregister_identity(settings, mgmt_api, extension_id, instance_id)

mrok/datastructures.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+from pydantic_core import core_schema
+
+
+class FixedSizeByteBuffer:
+    def __init__(self, max_size: int):
+        self._max_size = max_size
+        self._buf = bytearray()
+        self.overflow = False
+
+    def write(self, data: bytes) -> None:
+        if not data:
+            return
+
+        remaining = self._max_size - len(self._buf)
+        if remaining <= 0:
+            self.overflow = True
+            return
+
+        if len(data) > remaining:
+            self._buf.extend(data[:remaining])
+            self.overflow = True
+        else:
+            self._buf.extend(data)
+
+    def getvalue(self) -> bytes:
+        return bytes(self._buf)
+
+    def clear(self) -> None:
+        self._buf.clear()
+        self.overflow = False
+
+
+class HTTPHeaders(dict):
+    def __init__(self, initial=None):
+        super().__init__()
+        if initial:
+            for k, v in initial.items():
+                super().__setitem__(str(k).lower(), str(v))
+
+    def __getitem__(self, key: str) -> str:
+        return super().__getitem__(key.lower())
+
+    def __setitem__(self, key: str, value: str) -> None:
+        super().__setitem__(str(key).lower(), str(value))
+
+    def __delitem__(self, key: str) -> None:
+        super().__delitem__(key.lower())
+
+    def get(self, k, default=None):
+        return super().get(k.lower(), default)
+
+    @classmethod
+    def from_asgi(cls, items: list[tuple[bytes, bytes]]) -> HTTPHeaders:
+        d = {k.decode("latin-1"): v.decode("latin-1") for k, v in items}
+        return cls(d)
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source, handler):
+        """Provide a pydantic-core schema so Pydantic treats this as a mapping of str->str.
+
+        We generate the schema for `dict[str, str]` using the provided handler and wrap
+        it with a validator that converts the validated dict into `HTTPHeaders`.
+        """
+        # handler may be a callable or an object with `generate_schema`; handle both
+        try:
+            dict_schema = handler.generate_schema(dict[str, str])
+        except AttributeError:
+            dict_schema = handler(dict[str, str])
+
+        def _wrap(v, validator):
+            # `validator` will validate input according to `dict_schema` and return a dict
+            validated = validator(input_value=v)
+            if isinstance(validated, HTTPHeaders):
+                return validated
+            return cls(validated)
+
+        return core_schema.no_info_wrap_validator_function(
+            _wrap,
+            dict_schema,
+            serialization=core_schema.plain_serializer_function_ser_schema(lambda v: dict(v)),
+        )
+
+
+class HTTPRequest(BaseModel):
+    method: str
+    url: str
+    headers: HTTPHeaders
+    query_string: bytes
+    start_time: float
+    body: bytes | None = None
+    body_truncated: bool | None = None
+
+
+class HTTPResponse(BaseModel):
+    type: Literal["response"] = "response"
+    request: HTTPRequest
+    status: int
+    headers: HTTPHeaders
+    duration: float
+    body: bytes | None = None
+    body_truncated: bool | None = None
+
+
+class ProcessMetrics(BaseModel):
+    cpu: float
+    mem: float
+
+
+class DataTransferMetrics(BaseModel):
+    bytes_in: int
+    bytes_out: int
+
+
+class RequestsMetrics(BaseModel):
+    rps: int
+    total: int
+    successful: int
+    failed: int
+
+
+class ResponseTimeMetrics(BaseModel):
+    avg: float
+    min: int
+    max: int
+    p50: int
+    p90: int
+    p99: int
+
+
+class WorkerMetrics(BaseModel):
+    worker_id: str
+    data_transfer: DataTransferMetrics
+    requests: RequestsMetrics
+    response_time: ResponseTimeMetrics
+    process: ProcessMetrics
+
+
+class Meta(BaseModel):
+    identity: str
+    extension: str
+    instance: str
+    domain: str
+    tags: dict[str, str] | None = None
+
+
+class Status(BaseModel):
+    type: Literal["status"] = "status"
+    meta: Meta
+    metrics: WorkerMetrics
+
+
+class Event(BaseModel):
+    type: Literal["status", "response"]
+    data: Status | HTTPResponse = Field(discriminator="type")

mrok/http/config.py

@@ -8,21 +8,18 @@ from typing import Any
 import openziti
 from uvicorn import config
 
-from mrok.conf import get_settings
 from mrok.http.protocol import MrokHttpToolsProtocol
-from mrok.logging import setup_logging
+from mrok.http.types import ASGIApp
 
 logger = logging.getLogger("mrok.proxy")
 
 config.LIFESPAN["auto"] = "mrok.http.lifespan:MrokLifespan"
 
-ASGIApplication = config.ASGIApplication
-
 
 class MrokBackendConfig(config.Config):
     def __init__(
         self,
-        app: ASGIApplication | Callable[..., Any] | str,
+        app: ASGIApp | Callable[..., Any] | str,
         identity_file: str | Path,
         ziti_load_timeout_ms: int = 5000,
         backlog: int = 2048,
@@ -62,4 +59,4 @@ class MrokBackendConfig(config.Config):
         return sock
 
     def configure_logging(self) -> None:
-        setup_logging(get_settings())
+        return

mrok/http/constants.py

@@ -0,0 +1,22 @@
+MAX_REQUEST_BODY_BYTES = 2 * 1024 * 1024
+MAX_RESPONSE_BODY_BYTES = 5 * 1024 * 1024
+
+BINARY_CONTENT_TYPES = {
+    "application/octet-stream",
+    "application/pdf",
+}
+
+BINARY_PREFIXES = (
+    "image/",
+    "video/",
+    "audio/",
+)
+
+TEXTUAL_CONTENT_TYPES = {
+    "application/json",
+    "application/xml",
+    "application/javascript",
+    "application/x-www-form-urlencoded",
+}
+
+TEXTUAL_PREFIXES = ("text/",)

mrok/http/forwarder.py

@@ -1,14 +1,10 @@
 import abc
 import asyncio
 import logging
-from collections.abc import Awaitable, Callable
-from typing import Any
 
-logger = logging.getLogger("mrok.proxy")
+from mrok.http.types import ASGIReceive, ASGISend, Scope, StreamReader, StreamWriter
 
-Scope = dict[str, Any]
-ASGIReceive = Callable[[], Awaitable[dict[str, Any]]]
-ASGISend = Callable[[dict[str, Any]], Awaitable[None]]
+logger = logging.getLogger("mrok.proxy")
 
 
 class BackendNotFoundError(Exception):
@@ -24,24 +20,71 @@ class ForwardAppBase(abc.ABC):
     and streaming logic (requests and responses).
     """
 
-    def __init__(self, read_chunk_size: int = 65536) -> None:
-        # number of bytes to read per iteration when streaming bodies
+    def __init__(
+        self,
+        read_chunk_size: int = 65536,
+        lifespan_timeout: float = 10.0,
+    ) -> None:
         self._read_chunk_size: int = int(read_chunk_size)
+        self._lifespan_timeout = lifespan_timeout
+
+    async def handle_lifespan(self, receive: ASGIReceive, send: ASGISend) -> None:
+        while True:
+            event = await receive()
+            etype = event.get("type")
+
+            if etype == "lifespan.startup":
+                try:
+                    await asyncio.wait_for(self.startup(), self._lifespan_timeout)
+                except TimeoutError:
+                    logger.exception("Lifespan startup timed out")
+                    await send({"type": "lifespan.startup.failed", "message": "startup timeout"})
+                    continue
+                except Exception as e:
+                    logger.exception("Exception during lifespan startup")
+                    await send({"type": "lifespan.startup.failed", "message": str(e)})
+                    continue
+                await send({"type": "lifespan.startup.complete"})
+
+            elif etype == "lifespan.shutdown":
+                try:
+                    await asyncio.wait_for(self.shutdown(), self._lifespan_timeout)
+                except TimeoutError:
+                    logger.exception("Lifespan shutdown timed out")
+                    await send({"type": "lifespan.shutdown.failed", "message": "shutdown timeout"})
+                    return
+                except Exception as exc:
+                    logger.exception("Exception during lifespan shutdown")
+                    await send({"type": "lifespan.shutdown.failed", "message": str(exc)})
+                    return
+                await send({"type": "lifespan.shutdown.complete"})
+                return
 
     @abc.abstractmethod
     async def select_backend(
         self,
         scope: Scope,
         headers: dict[str, str],
-    ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter] | tuple[None, None]:
+    ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
         """Return (reader, writer) connected to the target backend."""
 
+    async def startup(self):
+        return
+
+    async def shutdown(self):
+        return
+
     async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend) -> None:
         """ASGI callable entry point.
 
         Delegates to smaller helper methods for readability. Subclasses only
         need to implement backend selection.
         """
+        scope_type = scope.get("type")
+        if scope_type == "lifespan":
+            await self.handle_lifespan(receive, send)
+            return
+
         if scope.get("type") != "http":
             await send({"type": "http.response.start", "status": 500, "headers": []})
             await send({"type": "http.response.body", "body": b"Unsupported"})
@@ -116,7 +159,7 @@ class ForwardAppBase(abc.ABC):
 
     async def write_request_line_and_headers(
         self,
-        writer: asyncio.StreamWriter,
+        writer: StreamWriter,
         method: str,
         path_qs: str,
         headers: list[tuple[bytes, bytes]],
@@ -130,7 +173,7 @@ class ForwardAppBase(abc.ABC):
         await writer.drain()
 
     async def stream_request_body(
-        self, receive: ASGIReceive, writer: asyncio.StreamWriter, use_chunked: bool
+        self, receive: ASGIReceive, writer: StreamWriter, use_chunked: bool
     ) -> None:
         if use_chunked:
             await self.stream_request_chunked(receive, writer)
@@ -138,9 +181,7 @@ class ForwardAppBase(abc.ABC):
 
         await self.stream_request_until_end(receive, writer)
 
-    async def stream_request_chunked(
-        self, receive: ASGIReceive, writer: asyncio.StreamWriter
-    ) -> None:
+    async def stream_request_chunked(self, receive: ASGIReceive, writer: StreamWriter) -> None:
         """Send request body to backend using HTTP/1.1 chunked encoding."""
         while True:
             event = await receive()
@@ -160,9 +201,7 @@ class ForwardAppBase(abc.ABC):
         writer.write(b"0\r\n\r\n")
         await writer.drain()
 
-    async def stream_request_until_end(
-        self, receive: ASGIReceive, writer: asyncio.StreamWriter
-    ) -> None:
+    async def stream_request_until_end(self, receive: ASGIReceive, writer: StreamWriter) -> None:
         """Send request body to backend when content length/transfer-encoding
         already provided (no chunking).
         """
@@ -180,7 +219,7 @@ class ForwardAppBase(abc.ABC):
                 return
 
     async def read_status_and_headers(
-        self, reader: asyncio.StreamReader, first_line: bytes
+        self, reader: StreamReader, first_line: bytes
     ) -> tuple[int, list[tuple[bytes, bytes]], dict[bytes, bytes]]:
         parts = first_line.decode(errors="ignore").split(" ", 2)
         status = int(parts[1]) if len(parts) >= 2 and parts[1].isdigit() else 502
@@ -217,14 +256,14 @@ class ForwardAppBase(abc.ABC):
         except Exception:
             return None
 
-    async def drain_trailers(self, reader: asyncio.StreamReader) -> None:
+    async def drain_trailers(self, reader: StreamReader) -> None:
         """Consume trailer header lines until an empty line is encountered."""
         while True:
             trailer = await reader.readline()
             if trailer in (b"\r\n", b"\n", b""):
                 break
 
-    async def stream_response_chunked(self, reader: asyncio.StreamReader, send: ASGISend) -> None:
+    async def stream_response_chunked(self, reader: StreamReader, send: ASGISend) -> None:
         """Read chunked-encoded response from reader, decode and forward to ASGI send."""
         while True:
             size_line = await reader.readline()
@@ -253,7 +292,7 @@ class ForwardAppBase(abc.ABC):
         await send({"type": "http.response.body", "body": b"", "more_body": False})
 
     async def stream_response_with_content_length(
-        self, reader: asyncio.StreamReader, send: ASGISend, content_length: int
+        self, reader: StreamReader, send: ASGISend, content_length: int
     ) -> None:
         """Read exactly content_length bytes and forward to ASGI send events."""
         remaining = content_length
@@ -272,7 +311,7 @@ class ForwardAppBase(abc.ABC):
         if not sent_final:
             await send({"type": "http.response.body", "body": b"", "more_body": False})
 
-    async def stream_response_until_eof(self, reader: asyncio.StreamReader, send: ASGISend) -> None:
+    async def stream_response_until_eof(self, reader: StreamReader, send: ASGISend) -> None:
         """Read from reader until EOF and forward chunks to ASGI send events."""
         while True:
             chunk = await reader.read(self._read_chunk_size)
@@ -282,7 +321,7 @@ class ForwardAppBase(abc.ABC):
         await send({"type": "http.response.body", "body": b"", "more_body": False})
 
     async def stream_response_body(
-        self, reader: asyncio.StreamReader, send: ASGISend, raw_headers: dict[bytes, bytes]
+        self, reader: StreamReader, send: ASGISend, raw_headers: dict[bytes, bytes]
     ) -> None:
         te = raw_headers.get(b"transfer-encoding", b"").lower()
         cl = raw_headers.get(b"content-length")

mrok/http/lifespan.py

@@ -1,10 +1,39 @@
 import logging
+from collections.abc import Awaitable, Callable
 
 from uvicorn.config import Config
 from uvicorn.lifespan.on import LifespanOn
 
+AsyncCallback = Callable[[], Awaitable[None]]
+
 
 class MrokLifespan(LifespanOn):
     def __init__(self, config: Config) -> None:
         super().__init__(config)
         self.logger = logging.getLogger("mrok.proxy")
+
+
+class LifespanWrapper:
+    def __init__(
+        self, app, on_startup: AsyncCallback | None = None, on_shutdown: AsyncCallback | 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:
+                        await self.on_startup()
+                    await send({"type": "lifespan.startup.complete"})
+
+                elif event["type"] == "lifespan.shutdown":
+                    if self.on_shutdown:
+                        await self.on_shutdown()
+                    await send({"type": "lifespan.shutdown.complete"})
+                    break
+        else:
+            await self.app(scope, receive, send)

mrok/http/middlewares.py

@@ -0,0 +1,143 @@
+import asyncio
+import inspect
+import logging
+import time
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+from mrok.datastructures import FixedSizeByteBuffer, HTTPHeaders, HTTPRequest, HTTPResponse
+from mrok.http.constants import MAX_REQUEST_BODY_BYTES, MAX_RESPONSE_BODY_BYTES
+from mrok.http.types import ASGIApp, ASGIReceive, ASGISend, Message, Scope
+from mrok.http.utils import must_capture_request, must_capture_response
+from mrok.metrics import WorkerMetricsCollector
+
+logger = logging.getLogger("mrok.proxy")
+
+ResponseCompleteCallback = Callable[[HTTPResponse], Coroutine[Any, Any, None] | None]
+
+
+class CaptureMiddleware:
+    def __init__(
+        self,
+        app: ASGIApp,
+        on_response_complete: ResponseCompleteCallback,
+    ):
+        self.app = app
+        self._on_response_complete = on_response_complete
+
+    async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend):
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        start_time = time.time()
+        method = scope["method"]
+        path = scope["path"]
+        query_string = scope.get("query_string", b"")
+        req_headers_raw = scope.get("headers", [])
+        req_headers = HTTPHeaders.from_asgi(req_headers_raw)
+
+        state = {}
+
+        req_buf = FixedSizeByteBuffer(MAX_REQUEST_BODY_BYTES)
+        capture_req_body = must_capture_request(method, req_headers)
+
+        request = HTTPRequest(
+            method=method,
+            url=path,
+            headers=req_headers,
+            query_string=query_string,
+            start_time=start_time,
+        )
+
+        # Response capture
+        resp_buf = FixedSizeByteBuffer(MAX_RESPONSE_BODY_BYTES)
+
+        async def receive_wrapper() -> Message:
+            msg = await receive()
+            if capture_req_body and msg["type"] == "http.request":
+                body = msg.get("body", b"")
+                req_buf.write(body)
+            return msg
+
+        async def send_wrapper(msg: Message):
+            if msg["type"] == "http.response.start":
+                state["status"] = msg["status"]
+                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":
+                body = msg.get("body", b"")
+                resp_buf.write(body)
+
+            await send(msg)
+
+        await self.app(scope, receive_wrapper, send_wrapper)
+
+        # Finalise request
+        request.body = req_buf.getvalue() if capture_req_body else None
+        request.body_truncated = req_buf.overflow if capture_req_body else None
+
+        # Finalise response
+        end_time = time.time()
+        duration = end_time - start_time
+
+        response = HTTPResponse(
+            request=request,
+            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,
+        )
+        await asyncio.create_task(self.handle_callback(response))
+
+    async def handle_callback(self, response: HTTPResponse):
+        try:
+            if inspect.iscoroutinefunction(self._on_response_complete):
+                await self._on_response_complete(response)
+            else:
+                await asyncio.get_running_loop().run_in_executor(
+                    None, self._on_response_complete, response
+                )
+        except Exception:
+            logger.exception("Error invoking callback")
+
+
+class MetricsMiddleware:
+    def __init__(self, app, metrics: WorkerMetricsCollector):
+        self.app = app
+        self.metrics = metrics
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        start_time = await self.metrics.on_request_start(scope)
+        status_code = 500  # default if app errors early
+
+        async def wrapped_receive():
+            msg = await receive()
+            if msg["type"] == "http.request" and msg.get("body"):
+                await self.metrics.on_request_body(len(msg["body"]))
+            return msg
+
+        async def wrapped_send(msg):
+            nonlocal status_code
+
+            if msg["type"] == "http.response.start":
+                status_code = msg["status"]
+                await self.metrics.on_response_start(status_code)
+
+            elif msg["type"] == "http.response.body":
+                body = msg.get("body", b"")
+                await self.metrics.on_response_chunk(len(body))
+
+            return await send(msg)
+
+        try:
+            await self.app(scope, wrapped_receive, wrapped_send)
+        finally:
+            await self.metrics.on_request_end(start_time, status_code)

mrok/http/types.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, MutableMapping
+from typing import Any, Protocol
+
+from mrok.datastructures import HTTPRequest, HTTPResponse
+
+Scope = MutableMapping[str, Any]
+Message = MutableMapping[str, Any]
+
+ASGIReceive = Callable[[], Awaitable[Message]]
+ASGISend = Callable[[Message], Awaitable[None]]
+ASGIApp = Callable[[Scope, ASGIReceive, ASGISend], Awaitable[None]]
+RequestCompleteCallback = Callable[[HTTPRequest], Awaitable | None]
+ResponseCompleteCallback = Callable[[HTTPResponse], Awaitable | None]
+
+
+class StreamReaderWrapper(Protocol):
+    async def read(self, n: int = -1) -> bytes: ...
+    async def readexactly(self, n: int) -> bytes: ...
+    async def readline(self) -> bytes: ...
+    def at_eof(self) -> bool: ...
+
+    @property
+    def underlying(self) -> asyncio.StreamReader: ...
+
+
+class StreamWriterWrapper(Protocol):
+    def write(self, data: bytes) -> None: ...
+    async def drain(self) -> None: ...
+    def close(self) -> None: ...
+    async def wait_closed(self) -> None: ...
+
+    @property
+    def transport(self): ...
+
+    @property
+    def underlying(self) -> asyncio.StreamWriter: ...
+
+
+StreamReader = StreamReaderWrapper | asyncio.StreamReader
+StreamWriter = StreamWriterWrapper | asyncio.StreamWriter