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

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

@@ -8,7 +8,7 @@ import time
 import psutil
 from hdrh.histogram import HdrHistogram
 
-from mrok.datastructures import (
+from mrok.proxy.models import (
     DataTransferMetrics,
     ProcessMetrics,
     RequestsMetrics,
@@ -20,11 +20,7 @@ logger = logging.getLogger("mrok.proxy")
 
 
 def _collect_process_usage(interval: float) -> ProcessMetrics:
-    try:
-        proc = psutil.Process(os.getpid())
-    except psutil.NoSuchProcess:
-        return ProcessMetrics(cpu=0.0, mem=0.0)
-
+    proc = psutil.Process(os.getpid())
     total_cpu = 0.0
     total_mem = 0.0
 
@@ -33,29 +29,28 @@ def _collect_process_usage(interval: float) -> ProcessMetrics:
     except Exception:
         total_cpu = 0.0
 
-    if interval and interval > 0:
+    if interval and interval > 0:  # pragma: no branch
         time.sleep(interval)
 
     try:
         total_cpu = proc.cpu_percent(None)
-    except Exception:
+    except Exception:  # pragma: no cover
         total_cpu = 0.0
 
     try:
         total_mem = proc.memory_percent()
-    except Exception:
+    except Exception:  # pragma: no cover
         total_mem = 0.0
 
     return ProcessMetrics(cpu=total_cpu, mem=total_mem)
 
 
-async def get_process_and_children_usage(interval: float = 0.1) -> ProcessMetrics:
+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):
-        # Request-level counters
         self.worker_id = worker_id
         self.total_requests = 0
         self.successful_requests = 0
@@ -63,14 +58,11 @@ class WorkerMetricsCollector:
         self.bytes_in = 0
         self.bytes_out = 0
 
-        # RPS
         self._tick_last = time.time()
         self._tick_requests = 0
 
-        # latency histogram
         self.hist = HdrHistogram(lowest, highest, sigfigs)
 
-        # async lock
         self._lock = asyncio.Lock()
 
     async def on_request_start(self, scope):
@@ -102,38 +94,34 @@ class WorkerMetricsCollector:
             self.hist.record_value(elapsed_ms)
 
     async def snapshot(self) -> WorkerMetrics:
-        try:
-            async with self._lock:
-                now = time.time()
-                delta = now - self._tick_last
-                rps = int(self._tick_requests / delta) if delta > 0 else 0
-                data = WorkerMetrics(
-                    worker_id=self.worker_id,
-                    process=await get_process_and_children_usage(),
-                    requests=RequestsMetrics(
-                        rps=rps,
-                        total=self.total_requests,
-                        successful=self.successful_requests,
-                        failed=self.failed_requests,
-                    ),
-                    data_transfer=DataTransferMetrics(
-                        bytes_in=self.bytes_in,
-                        bytes_out=self.bytes_out,
-                    ),
-                    response_time=ResponseTimeMetrics(
-                        avg=self.hist.get_mean_value(),
-                        min=self.hist.get_min_value(),
-                        max=self.hist.get_max_value(),
-                        p50=self.hist.get_value_at_percentile(50),
-                        p90=self.hist.get_value_at_percentile(90),
-                        p99=self.hist.get_value_at_percentile(99),
-                    ),
-                )
-
-                self._tick_last = now
-                self._tick_requests = 0
-
-                return data
-        except Exception:
-            logger.exception("Exception calculating snapshot")
-            raise
+        async with self._lock:
+            now = time.time()
+            delta = now - self._tick_last
+            rps = int(self._tick_requests / delta) if delta > 0 else 0
+            data = WorkerMetrics(
+                worker_id=self.worker_id,
+                process=await get_process_metrics(),
+                requests=RequestsMetrics(
+                    rps=rps,
+                    total=self.total_requests,
+                    successful=self.successful_requests,
+                    failed=self.failed_requests,
+                ),
+                data_transfer=DataTransferMetrics(
+                    bytes_in=self.bytes_in,
+                    bytes_out=self.bytes_out,
+                ),
+                response_time=ResponseTimeMetrics(
+                    avg=self.hist.get_mean_value(),
+                    min=self.hist.get_min_value(),
+                    max=self.hist.get_max_value(),
+                    p50=self.hist.get_value_at_percentile(50),
+                    p90=self.hist.get_value_at_percentile(90),
+                    p99=self.hist.get_value_at_percentile(99),
+                ),
+            )
+
+            self._tick_last = now
+            self._tick_requests = 0
+
+            return data

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

@@ -1,20 +1,22 @@
 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
+from mrok.proxy.constants import MAX_REQUEST_BODY_BYTES, MAX_RESPONSE_BODY_BYTES
+from mrok.proxy.metrics import MetricsCollector
+from mrok.proxy.models import FixedSizeByteBuffer, HTTPHeaders, HTTPRequest, HTTPResponse
+from mrok.proxy.utils import must_capture_request, must_capture_response
+from mrok.types.proxy import (
+    ASGIApp,
+    ASGIReceive,
+    ASGISend,
+    Message,
+    ResponseCompleteCallback,
+    Scope,
+)
 
 logger = logging.getLogger("mrok.proxy")
 
-ResponseCompleteCallback = Callable[[HTTPResponse], Coroutine[Any, Any, None] | None]
-
 
 class CaptureMiddleware:
     def __init__(
@@ -92,22 +94,11 @@ class CaptureMiddleware:
             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")
+        asyncio.create_task(self._on_response_complete(response))
 
 
 class MetricsMiddleware:
-    def __init__(self, app, metrics: WorkerMetricsCollector):
+    def __init__(self, app: ASGIApp, metrics: MetricsCollector):
         self.app = app
         self.metrics = metrics
 
@@ -116,11 +107,11 @@ class MetricsMiddleware:
             return await self.app(scope, receive, send)
 
         start_time = await self.metrics.on_request_start(scope)
-        status_code = 500  # default if app errors early
+        status_code = 500
 
         async def wrapped_receive():
             msg = await receive()
-            if msg["type"] == "http.request" and msg.get("body"):
+            if msg["type"] == "http.request" and msg.get("body"):  # pragma: no branch
                 await self.metrics.on_request_body(len(msg["body"]))
             return msg
 
@@ -131,7 +122,7 @@ class MetricsMiddleware:
                 status_code = msg["status"]
                 await self.metrics.on_response_start(status_code)
 
-            elif msg["type"] == "http.response.body":
+            elif msg["type"] == "http.response.body":  # pragma: no branch
                 body = msg.get("body", b"")
                 await self.metrics.on_response_chunk(len(body))
 

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

@@ -1,11 +1,52 @@
 from __future__ import annotations
 
+import json
+from pathlib import Path
 from typing import Literal
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 from pydantic_core import core_schema
 
 
+class X509Credentials(BaseModel):
+    key: str
+    cert: str
+    ca: str
+
+    @field_validator("key", "cert", "ca", mode="before")
+    @classmethod
+    def strip_pem_prefix(cls, value: str) -> str:
+        if isinstance(value, str) and value.startswith("pem:"):
+            return value[4:]
+        return value
+
+
+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 Identity(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    zt_api: str = Field(validation_alias="ztAPI")
+    id: X509Credentials
+    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: ServiceMetadata | None = None
+
+    @staticmethod
+    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 Identity.model_validate(data)
+
+
 class FixedSizeByteBuffer:
     def __init__(self, max_size: int):
         self._max_size = max_size
@@ -140,17 +181,9 @@ class WorkerMetrics(BaseModel):
     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
+    meta: ServiceMetadata
     metrics: WorkerMetrics
 
 

mrok/proxy/stream.py

@@ -0,0 +1,68 @@
+import asyncio
+import select
+import sys
+from typing import Any
+
+from httpcore import AsyncNetworkStream
+
+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):
+    def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        self._reader = reader
+        self._writer = writer
+
+    async def read(self, n: int, timeout: float | None = None) -> bytes:
+        return await asyncio.wait_for(self._reader.read(n), timeout)
+
+    async def write(self, data: bytes, timeout: float | None = None) -> None:
+        self._writer.write(data)
+        await asyncio.wait_for(self._writer.drain(), timeout)
+
+    async def aclose(self) -> None:
+        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):
+        self._receive = receive
+        self._more_body = True
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> bytes:
+        if not self._more_body:
+            raise StopAsyncIteration
+
+        msg = await self._receive()
+        if msg["type"] == "http.request":
+            chunk = msg.get("body", b"")
+            self._more_body = msg.get("more_body", False)
+            return chunk
+        elif msg["type"] == "http.disconnect":
+            raise Exception("Client disconnected.")
+
+        raise Exception("Unexpected asgi message.")

mrok/{http → proxy}/utils.py

@@ -1,6 +1,6 @@
 from collections.abc import Mapping
 
-from mrok.http.constants import (
+from mrok.proxy.constants import (
     BINARY_CONTENT_TYPES,
     BINARY_PREFIXES,
     MAX_REQUEST_BODY_BYTES,

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()

mrok/{http/config.py → proxy/ziticorn.py}

@@ -6,17 +6,40 @@ from pathlib import Path
 from typing import Any
 
 import openziti
-from uvicorn import config
+from uvicorn import config, server
+from uvicorn.lifespan.on import LifespanOn
+from uvicorn.protocols.http.httptools_impl import HttpToolsProtocol as UvHttpToolsProtocol
 
-from mrok.http.protocol import MrokHttpToolsProtocol
-from mrok.http.types import ASGIApp
+from mrok.types.proxy import ASGIApp
 
 logger = logging.getLogger("mrok.proxy")
 
-config.LIFESPAN["auto"] = "mrok.http.lifespan:MrokLifespan"
+config.LIFESPAN["auto"] = "mrok.proxy.ziticorn:Lifespan"
 
 
-class MrokBackendConfig(config.Config):
+class Lifespan(LifespanOn):
+    def __init__(self, lf_config: config.Config) -> None:
+        super().__init__(lf_config)
+        self.logger = logging.getLogger("mrok.proxy")
+
+
+class HttpToolsProtocol(UvHttpToolsProtocol):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.logger = logging.getLogger("mrok.proxy")
+        self.access_logger = logging.getLogger("mrok.access")
+        self.access_log = self.access_logger.hasHandlers()
+
+
+class Server(server.Server):
+    async def serve(self, sockets: list[socket.socket] | None = None) -> None:
+        if not sockets:
+            sockets = [self.config.bind_socket()]
+        with self.capture_signals():
+            await self._serve(sockets)
+
+
+class BackendConfig(config.Config):
     def __init__(
         self,
         app: ASGIApp | Callable[..., Any] | str,
@@ -32,7 +55,7 @@ class MrokBackendConfig(config.Config):
         super().__init__(
             app,
             loop="asyncio",
-            http=MrokHttpToolsProtocol,
+            http=HttpToolsProtocol,
             backlog=backlog,
         )
 

mrok/types/proxy.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Coroutine, Mapping, MutableMapping
+from contextlib import AbstractAsyncContextManager
+from typing import Any, Never
+
+from mrok.proxy.models import 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]]
+StatelessLifespan = Callable[[ASGIApp], AbstractAsyncContextManager[None]]
+StatefulLifespan = Callable[[ASGIApp], AbstractAsyncContextManager[Mapping[str, Any]]]
+Lifespan = StatelessLifespan | StatefulLifespan
+
+LifespanCallback = Callable[[], Awaitable[None]]
+ResponseCompleteCallback = Callable[[HTTPResponse], Coroutine[Any, Any, Never]]

mrok/types/ziti.py

@@ -0,0 +1 @@
+Tags = dict[str, str | bool | None]

mrok/ziti/api.py

@@ -11,12 +11,11 @@ from typing import Any, Literal
 import httpx
 
 from mrok.conf import Settings
+from mrok.types.ziti import Tags
 from mrok.ziti.constants import MROK_VERSION_TAG, MROK_VERSION_TAG_NAME
 
 logger = logging.getLogger(__name__)
 
-TagsType = dict[str, str | bool | None]
-
 
 class ZitiAPIError(Exception):
     pass
@@ -70,7 +69,7 @@ class BaseZitiAPI(ABC):
             ),
         )
 
-    async def create(self, endpoint: str, payload: dict[str, Any], tags: TagsType | None) -> str:
+    async def create(self, endpoint: str, payload: dict[str, Any], tags: Tags | None) -> str:
         payload["tags"] = self._merge_tags(tags)
         response: httpx.Response = await self.httpx_client.post(
             endpoint,
@@ -156,8 +155,8 @@ class BaseZitiAPI(ABC):
     ) -> None:
         return await self.httpx_client.__aexit__(exc_type, exc_val, exc_tb)
 
-    def _merge_tags(self, tags: TagsType | None) -> TagsType:
-        prepared_tags: TagsType = tags or {}
+    def _merge_tags(self, tags: Tags | None) -> Tags:
+        prepared_tags: Tags = tags or {}
         prepared_tags.update(MROK_VERSION_TAG)
         return prepared_tags
 
@@ -281,9 +280,7 @@ class ZitiManagementAPI(BaseZitiAPI):
     async def search_config(self, id_or_name) -> dict[str, Any] | None:
         return await self.search_by_id_or_name("/configs", id_or_name)
 
-    async def create_config(
-        self, name: str, config_type_id: str, tags: TagsType | None = None
-    ) -> str:
+    async def create_config(self, name: str, config_type_id: str, tags: Tags | None = None) -> str:
         return await self.create(
             "/configs",
             {
@@ -302,7 +299,7 @@ class ZitiManagementAPI(BaseZitiAPI):
     async def delete_config(self, config_id: str) -> None:
         return await self.delete("/configs", config_id)
 
-    async def create_config_type(self, name: str, tags: TagsType | None = None) -> str:
+    async def create_config_type(self, name: str, tags: Tags | None = None) -> str:
         return await self.create(
             "/config-types",
             {
@@ -316,7 +313,7 @@ class ZitiManagementAPI(BaseZitiAPI):
         self,
         name: str,
         config_id: str,
-        tags: TagsType | None = None,
+        tags: Tags | None = None,
     ) -> str:
         return await self.create(
             "/services",
@@ -332,7 +329,7 @@ class ZitiManagementAPI(BaseZitiAPI):
         self,
         name: str,
         service_id: str,
-        tags: TagsType | None = None,
+        tags: Tags | None = None,
     ) -> str:
         return await self.create(
             "/service-edge-router-policies",
@@ -351,7 +348,7 @@ class ZitiManagementAPI(BaseZitiAPI):
         self,
         name: str,
         identity_id: str,
-        tags: TagsType | None = None,
+        tags: Tags | None = None,
     ) -> str:
         return await self.create(
             "/edge-router-policies",
@@ -385,10 +382,10 @@ class ZitiManagementAPI(BaseZitiAPI):
     async def delete_service(self, service_id: str) -> None:
         return await self.delete("/services", service_id)
 
-    async def create_user_identity(self, name: str, tags: TagsType | None = None) -> str:
+    async def create_user_identity(self, name: str, tags: Tags | None = None) -> str:
         return await self._create_identity(name, "User", tags=tags)
 
-    async def create_device_identity(self, name: str, tags: TagsType | None = None) -> str:
+    async def create_device_identity(self, name: str, tags: Tags | None = None) -> str:
         return await self._create_identity(name, "Device", tags=tags)
 
     async def search_identity(self, id_or_name: str) -> dict[str, Any] | None:
@@ -412,12 +409,12 @@ class ZitiManagementAPI(BaseZitiAPI):
         return response.text
 
     async def create_dial_service_policy(
-        self, name: str, service_id: str, identity_id: str, tags: TagsType | None = None
+        self, name: str, service_id: str, identity_id: str, tags: Tags | None = None
     ) -> str:
         return await self._create_service_policy("Dial", name, service_id, identity_id, tags)
 
     async def create_bind_service_policy(
-        self, name: str, service_id: str, identity_id: str, tags: TagsType | None = None
+        self, name: str, service_id: str, identity_id: str, tags: Tags | None = None
     ) -> str:
         return await self._create_service_policy("Bind", name, service_id, identity_id, tags)
 
@@ -433,7 +430,7 @@ class ZitiManagementAPI(BaseZitiAPI):
         name: str,
         service_id: str,
         identity_id: str,
-        tags: TagsType | None = None,
+        tags: Tags | None = None,
     ) -> str:
         return await self.create(
             "/service-policies",
@@ -451,7 +448,7 @@ class ZitiManagementAPI(BaseZitiAPI):
         self,
         name: str,
         type: Literal["User", "Device", "Default"],
-        tags: TagsType | None = None,
+        tags: Tags | None = None,
     ) -> str:
         return await self.create(
             "/identities",

mrok/ziti/bootstrap.py

@@ -1,7 +1,8 @@
 import logging
 from typing import Any
 
-from mrok.ziti.api import TagsType, ZitiClientAPI, ZitiManagementAPI
+from mrok.types.ziti import Tags
+from mrok.ziti.api import ZitiClientAPI, ZitiManagementAPI
 from mrok.ziti.identities import enroll_proxy_identity
 
 logger = logging.getLogger(__name__)
@@ -13,7 +14,7 @@ async def bootstrap_identity(
     identity_name: str,
     mode: str,
     forced: bool,
-    tags: TagsType | None,
+    tags: Tags | None,
 ) -> tuple[str, dict[str, Any] | None]:
     logger.info(f"Bootstrapping '{identity_name}' identity...")