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

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}

@@ -3,18 +3,17 @@ 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.proxy.utils import must_capture_request, must_capture_response
+from mrok.types.proxy import (
     ASGIApp,
     ASGIReceive,
     ASGISend,
-    LifespanCallback,
     Message,
     ResponseCompleteCallback,
     Scope,
 )
-from mrok.proxy.utils import must_capture_request, must_capture_response
 
 logger = logging.getLogger("mrok.proxy")
 
@@ -99,7 +98,7 @@ class CaptureMiddleware:
 
 
 class MetricsMiddleware:
-    def __init__(self, app: ASGIApp, metrics: WorkerMetricsCollector):
+    def __init__(self, app: ASGIApp, metrics: MetricsCollector):
         self.app = app
         self.metrics = metrics
 
@@ -133,32 +132,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}

@@ -8,7 +8,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator
 from pydantic_core import core_schema
 
 
-class ZitiId(BaseModel):
+class X509Credentials(BaseModel):
     key: str
     cert: str
     ca: str
@@ -21,7 +21,7 @@ class ZitiId(BaseModel):
         return value
 
 
-class ZitiMrokMeta(BaseModel):
+class ServiceMetadata(BaseModel):
     model_config = ConfigDict(extra="ignore")
     identity: str
     extension: str
@@ -30,21 +30,21 @@ class ZitiMrokMeta(BaseModel):
     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
     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
+    mrok: ServiceMetadata | 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:
@@ -183,7 +183,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()

mrok/proxy/{config.py → 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.proxy.protocol import MrokHttpToolsProtocol
-from mrok.proxy.types import ASGIApp
+from mrok.types.proxy import ASGIApp
 
 logger = logging.getLogger("mrok.proxy")
 
-config.LIFESPAN["auto"] = "mrok.proxy.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/{proxy/types.py → types/proxy.py}

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable, Coroutine, MutableMapping
+from collections.abc import Awaitable, Callable, Coroutine, Mapping, MutableMapping
+from contextlib import AbstractAsyncContextManager
 from typing import Any, Never
 
-from mrok.proxy.datastructures import HTTPResponse
+from mrok.proxy.models import HTTPResponse
 
 Scope = MutableMapping[str, Any]
 Message = MutableMapping[str, Any]
@@ -11,5 +12,9 @@ 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...")
 

mrok/ziti/identities.py

@@ -5,8 +5,9 @@ from typing import Any
 import jwt
 
 from mrok.conf import Settings
+from mrok.types.ziti import Tags
 from mrok.ziti import pki
-from mrok.ziti.api import TagsType, ZitiClientAPI, ZitiManagementAPI
+from mrok.ziti.api import ZitiClientAPI, ZitiManagementAPI
 from mrok.ziti.constants import (
     MROK_IDENTITY_TYPE_TAG_NAME,
     MROK_IDENTITY_TYPE_TAG_VALUE_INSTANCE,
@@ -29,7 +30,7 @@ async def register_identity(
     client_api: ZitiClientAPI,
     service_external_id: str,
     identity_external_id: str,
-    tags: TagsType | None = None,
+    tags: Tags | None = None,
 ):
     service_name = service_external_id.lower()
     identity_tags = copy.copy(tags or {})
@@ -39,7 +40,7 @@ async def register_identity(
     if not service:
         raise ServiceNotFoundError(f"A service with name `{service_external_id}` does not exists.")
 
-    identity_name = f"{identity_external_id.lower()}.{service_name}"
+    identity_name = identity_external_id.lower()
     service_policy_name = f"{identity_name}:bind"
     self_service_policy_name = f"self.{service_policy_name}"
 
@@ -129,7 +130,7 @@ async def enroll_proxy_identity(
     mgmt_api: ZitiManagementAPI,
     client_api: ZitiClientAPI,
     identity_name: str,
-    tags: TagsType | None = None,
+    tags: Tags | None = None,
 ):
     identity = await mgmt_api.search_identity(identity_name)
     if identity:

mrok/ziti/services.py

@@ -2,7 +2,8 @@ import logging
 from typing import Any
 
 from mrok.conf import Settings
-from mrok.ziti.api import TagsType, ZitiManagementAPI
+from mrok.types.ziti import Tags
+from mrok.ziti.api import ZitiManagementAPI
 from mrok.ziti.errors import (
     ConfigTypeNotFoundError,
     ProxyIdentityNotFoundError,
@@ -14,7 +15,7 @@ logger = logging.getLogger(__name__)
 
 
 async def register_service(
-    settings: Settings, mgmt_api: ZitiManagementAPI, external_id: str, tags: TagsType | None
+    settings: Settings, mgmt_api: ZitiManagementAPI, external_id: str, tags: Tags | None
 ) -> dict[str, Any]:
     service_name = external_id.lower()
     registered = False

{mrok-0.5.0.dist-info → mrok-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mrok
-Version: 0.5.0
+Version: 0.6.0
 Summary: MPT Extensions OpenZiti Orchestrator
 Author: SoftwareOne AG
 License:                                  Apache License