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

mrok/agent/devtools/inspector/app.py

@@ -27,7 +27,7 @@ from textual.widgets.data_table import ColumnKey
 from textual.worker import get_current_worker
 
 from mrok import __version__
-from mrok.datastructures import Event, HTTPHeaders, HTTPResponse, Meta, WorkerMetrics
+from mrok.proxy.models import Event, HTTPHeaders, HTTPResponse, ServiceMetadata, WorkerMetrics
 
 
 def build_tree(node, data):
@@ -185,7 +185,7 @@ class InfoPanel(Static):
         #         mem=int(mean([m.process.mem for m in self.workers_metrics.values()])),
         #     )
 
-    def update_meta(self, meta: Meta) -> None:
+    def update_meta(self, meta: ServiceMetadata) -> None:
         table = self.query_one(DataTable)
         if len(table.rows) == 0:
             table.add_row("URL", f"https://{meta.extension}.{meta.domain}")

mrok/agent/sidecar/app.py

@@ -1,51 +1,77 @@
-import asyncio
 import logging
-from collections.abc import AsyncGenerator
-from contextlib import asynccontextmanager
 from pathlib import Path
+from typing import Literal
 
-from mrok.http.forwarder import ForwardAppBase
-from mrok.http.pool import ConnectionPool
-from mrok.http.types import Scope, StreamPair
+from httpcore import AsyncConnectionPool
+
+from mrok.proxy.app import ProxyAppBase
+from mrok.types.proxy import Scope
 
 logger = logging.getLogger("mrok.agent")
 
 
-class ForwardApp(ForwardAppBase):
+TargetType = Literal["tcp", "unix"]
+
+
+class SidecarProxyApp(ProxyAppBase):
     def __init__(
         self,
-        target_address: str | Path | tuple[str, int],
-        read_chunk_size: int = 65536,
-    ) -> None:
+        target: str | Path | tuple[str, int],
+        *,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
+        retries: int = 0,
+    ):
+        self._target = target
+        self._target_type, self._target_address = self._parse_target()
         super().__init__(
-            read_chunk_size=read_chunk_size,
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+            retries=retries,
         )
-        self._target_address = target_address
-        self._pool = ConnectionPool(
-            pool_name=str(self._target_address),
-            factory=self.connect,
-            initial_connections=5,
-            max_size=100,
-            idle_timeout=20.0,
-            reaper_interval=5.0,
+
+    def setup_connection_pool(
+        self,
+        max_connections: int | None,
+        max_keepalive_connections: int | None,
+        keepalive_expiry: float | None,
+        retries: int,
+    ) -> AsyncConnectionPool:
+        if self._target_type == "unix":
+            return AsyncConnectionPool(
+                max_connections=max_connections,
+                max_keepalive_connections=max_keepalive_connections,
+                keepalive_expiry=keepalive_expiry,
+                retries=retries,
+                uds=self._target_address,
+            )
+        return AsyncConnectionPool(
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+            retries=retries,
         )
 
-    async def connect(self) -> StreamPair:
-        if isinstance(self._target_address, tuple):
-            return await asyncio.open_connection(*self._target_address)
-        return await asyncio.open_unix_connection(str(self._target_address))
+    def get_upstream_base_url(self, scope: Scope) -> str:
+        if self._target_type == "unix":
+            return "http://localhost"
+        return f"http://{self._target_address}"
 
-    async def startup(self):
-        await self._pool.start()
+    def _parse_target(self) -> tuple[TargetType, str]:
+        if isinstance(self._target, Path) or (
+            isinstance(self._target, str) and ":" not in self._target
+        ):
+            return "unix", str(self._target)
 
-    async def shutdown(self):
-        await self._pool.stop()
+        if isinstance(self._target, str) and ":" in self._target:
+            host, port = str(self._target).split(":", 1)
+            host = host or "127.0.0.1"
+        elif isinstance(self._target, tuple) and len(self._target) == 2:
+            host = self._target[0]
+            port = str(self._target[1])
+        else:
+            raise Exception(f"Invalid target address: {self._target}")
 
-    @asynccontextmanager
-    async def select_backend(
-        self,
-        scope: Scope,
-        headers: dict[str, str],
-    ) -> AsyncGenerator[StreamPair, None]:
-        async with self._pool.acquire() as (reader, writer):
-            yield reader, writer
+        return "tcp", f"{host}:{port}"

mrok/agent/sidecar/main.py

@@ -1,8 +1,8 @@
 import logging
 from pathlib import Path
 
-from mrok.agent.sidecar.app import ForwardApp
-from mrok.master import MasterBase
+from mrok.agent.sidecar.app import SidecarProxyApp
+from mrok.proxy.master import MasterBase
 
 logger = logging.getLogger("mrok.proxy")
 
@@ -11,28 +11,49 @@ class SidecarAgent(MasterBase):
     def __init__(
         self,
         identity_file: str,
-        target_addr: str | Path | tuple[str, int],
+        target: str | Path | tuple[str, int],
         workers: int = 4,
+        events_enabled: bool = True,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
+        retries: int = 0,
         publishers_port: int = 50000,
         subscribers_port: int = 50001,
     ):
         super().__init__(
             identity_file,
-            workers,
-            False,
-            publishers_port,
-            subscribers_port,
+            workers=workers,
+            reload=False,
+            events_enabled=events_enabled,
+            events_pub_port=publishers_port,
+            events_sub_port=subscribers_port,
         )
-        self.target_address = target_addr
+        self._target = target
+        self._max_connections = max_connections
+        self._max_keepalive_connections = max_keepalive_connections
+        self._keepalive_expiry = keepalive_expiry
+        self._retries = retries
 
     def get_asgi_app(self):
-        return ForwardApp(self.target_address)
+        return SidecarProxyApp(
+            self._target,
+            max_connections=self._max_connections,
+            max_keepalive_connections=self._max_keepalive_connections,
+            keepalive_expiry=self._keepalive_expiry,
+            retries=self._retries,
+        )
 
 
 def run(
     identity_file: str,
     target_addr: str | Path | tuple[str, int],
     workers: int = 4,
+    events_enabled: bool = True,
+    max_connections: int | None = 10,
+    max_keepalive_connections: int | None = None,
+    keepalive_expiry: float | None = None,
+    retries: int = 0,
     publishers_port: int = 50000,
     subscribers_port: int = 50001,
 ):
@@ -40,6 +61,11 @@ def run(
         identity_file,
         target_addr,
         workers=workers,
+        events_enabled=events_enabled,
+        max_connections=max_connections,
+        max_keepalive_connections=max_keepalive_connections,
+        keepalive_expiry=keepalive_expiry,
+        retries=retries,
         publishers_port=publishers_port,
         subscribers_port=subscribers_port,
     )

mrok/agent/ziticorn.py

@@ -1,5 +1,5 @@
-from mrok.http.types import ASGIApp
-from mrok.master import MasterBase
+from mrok.proxy.master import MasterBase
+from mrok.types.proxy import ASGIApp
 
 
 class ZiticornAgent(MasterBase):
@@ -12,7 +12,13 @@ class ZiticornAgent(MasterBase):
         publishers_port: int = 50000,
         subscribers_port: int = 50001,
     ):
-        super().__init__(identity_file, workers, reload, publishers_port, subscribers_port)
+        super().__init__(
+            identity_file,
+            workers=workers,
+            reload=reload,
+            events_pub_port=publishers_port,
+            events_sub_port=subscribers_port,
+        )
         self.app = app
 
     def get_asgi_app(self):

mrok/cli/commands/__init__.py

@@ -1,8 +1,8 @@
-from mrok.cli.commands import admin, agent, controller, proxy
+from mrok.cli.commands import admin, agent, controller, frontend
 
 __all__ = [
     "admin",
     "agent",
     "controller",
-    "proxy",
+    "frontend",
 ]

mrok/cli/commands/admin/bootstrap.py

@@ -8,14 +8,15 @@ import typer
 
 from mrok.cli.commands.admin.utils import parse_tags
 from mrok.conf import Settings
-from mrok.ziti.api import TagsType, ZitiClientAPI, ZitiManagementAPI
+from mrok.types.ziti import Tags
+from mrok.ziti.api import ZitiClientAPI, ZitiManagementAPI
 from mrok.ziti.bootstrap import bootstrap_identity
 
 logger = logging.getLogger(__name__)
 
 
 async def bootstrap(
-    settings: Settings, forced: bool, tags: TagsType | None
+    settings: Settings, forced: bool, tags: Tags | None
 ) -> tuple[str, dict[str, Any] | None]:
     async with ZitiManagementAPI(settings) as mgmt_api, ZitiClientAPI(settings) as client_api:
         return await bootstrap_identity(

mrok/cli/commands/admin/utils.py

@@ -2,10 +2,10 @@ from datetime import datetime
 
 import typer
 
-from mrok.ziti.api import TagsType
+from mrok.types.ziti import Tags
 
 
-def parse_tags(pairs: list[str] | None) -> TagsType | None:
+def parse_tags(pairs: list[str] | None) -> Tags | None:
     if not pairs:
         return None
 

mrok/cli/commands/agent/run/sidecar.py

@@ -25,10 +25,55 @@ def register(app: typer.Typer) -> None:
             typer.Option(
                 "--workers",
                 "-w",
-                help=f"Number of workers. Default: {default_workers}",
+                help="Number of workers.",
                 show_default=True,
             ),
         ] = default_workers,
+        max_connections: Annotated[
+            int,
+            typer.Option(
+                "--max-pool-connections",
+                help=(
+                    "The maximum number of concurrent HTTP connections that "
+                    "the pool should allow. Any attempt to send a request on a pool that "
+                    "would exceed this amount will block until a connection is available."
+                ),
+                show_default=True,
+            ),
+        ] = 10,
+        max_keepalive_connections: Annotated[
+            int | None,
+            typer.Option(
+                "--max-pool-keepalive-connections",
+                help=(
+                    "The maximum number of idle HTTP connections "
+                    "that will be maintained in the pool."
+                ),
+                show_default=True,
+            ),
+        ] = None,
+        keepalive_expiry: Annotated[
+            float | None,
+            typer.Option(
+                "--max-pool-keepalive-expiry",
+                help=(
+                    "The duration in seconds that an idle HTTP connection "
+                    "may be maintained for before being expired from the pool."
+                ),
+                show_default=True,
+            ),
+        ] = None,
+        retries: Annotated[
+            int,
+            typer.Option(
+                "--max-pool-connect-retries",
+                help=(
+                    "The duration in seconds that an idle HTTP connection "
+                    "may be maintained for before being expired from the pool."
+                ),
+                show_default=True,
+            ),
+        ] = 0,
         publishers_port: Annotated[
             int,
             typer.Option(
@@ -53,6 +98,14 @@ def register(app: typer.Typer) -> None:
                 show_default=True,
             ),
         ] = 50001,
+        no_events: Annotated[
+            bool,
+            typer.Option(
+                "--no-events",
+                help="Disable events. Default: False",
+                show_default=True,
+            ),
+        ] = False,
     ):
         """Run a Sidecar Proxy to expose a web application through OpenZiti."""
         if ":" in str(target):
@@ -65,6 +118,11 @@ def register(app: typer.Typer) -> None:
             str(identity_file),
             target_addr,
             workers=workers,
+            events_enabled=not no_events,
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+            retries=retries,
             publishers_port=publishers_port,
             subscribers_port=subscribers_port,
         )

mrok/cli/commands/{proxy → frontend}/__init__.py

@@ -1,6 +1,6 @@
 import typer
 
-from mrok.cli.commands.proxy import run
+from mrok.cli.commands.frontend import run
 
 app = typer.Typer(help="mrok proxy commands.")
 run.register(app)

mrok/cli/commands/frontend/run.py

@@ -0,0 +1,91 @@
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from mrok import frontend
+from mrok.cli.utils import number_of_workers
+
+default_workers = number_of_workers()
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("run")
+    def run_frontend(
+        ctx: typer.Context,
+        identity_file: Path = typer.Argument(
+            ...,
+            help="Identity json file",
+        ),
+        host: Annotated[
+            str,
+            typer.Option(
+                "--host",
+                "-h",
+                help="Host to bind to. Default: 127.0.0.1",
+                show_default=True,
+            ),
+        ] = "127.0.0.1",
+        port: Annotated[
+            int,
+            typer.Option(
+                "--port",
+                "-P",
+                help="Port to bind to. Default: 8000",
+                show_default=True,
+            ),
+        ] = 8000,
+        workers: Annotated[
+            int,
+            typer.Option(
+                "--workers",
+                "-w",
+                help=f"Number of workers. Default: {default_workers}",
+                show_default=True,
+            ),
+        ] = default_workers,
+        max_connections: Annotated[
+            int,
+            typer.Option(
+                "--max-pool-connections",
+                help=(
+                    "The maximum number of concurrent HTTP connections that "
+                    "the pool should allow. Any attempt to send a request on a pool that "
+                    "would exceed this amount will block until a connection is available."
+                ),
+                show_default=True,
+            ),
+        ] = 1000,
+        max_keepalive_connections: Annotated[
+            int | None,
+            typer.Option(
+                "--max-pool-keepalive-connections",
+                help=(
+                    "The maximum number of idle HTTP connections "
+                    "that will be maintained in the pool."
+                ),
+                show_default=True,
+            ),
+        ] = 100,
+        keepalive_expiry: Annotated[
+            float | None,
+            typer.Option(
+                "--max-pool-keepalive-expiry",
+                help=(
+                    "The duration in seconds that an idle HTTP connection "
+                    "may be maintained for before being expired from the pool."
+                ),
+                show_default=True,
+            ),
+        ] = 300,
+    ):
+        """Run the mrok frontend with Gunicorn and Uvicorn workers."""
+        frontend.run(
+            identity_file,
+            host,
+            port,
+            workers,
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+        )

mrok/constants.py

@@ -2,5 +2,3 @@ import re
 
 RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
 RE_INSTANCE_ID = re.compile(r"(?i)INS-\d{4}-\d{4}-\d{4}")
-
-RE_SUBDOMAIN = re.compile(r"(?i)^(?:EXT-\d{4}-\d{4}|INS-\d{4}-\d{4}-\d{4})$")

mrok/controller/openapi/examples.py

@@ -13,6 +13,7 @@ INSTANCE_RESPONSE = {
     "name": "ins-1234-5678-0001.ext-1234-5678",
     "extension": {"id": "EXT-1234-5678"},
     "instance": {"id": "INS-1234-5678-0001"},
+    "status": "offline",
     "tags": {
         "account": "ACC-5555-3333",
         MROK_VERSION_TAG_NAME: "1.0",
@@ -25,6 +26,7 @@ INSTANCE_CREATE_RESPONSE = {
     "name": "ins-1234-5678-0001.ext-1234-5678",
     "extension": {"id": "EXT-1234-5678"},
     "instance": {"id": "INS-1234-5678-0001"},
+    "status": "online",
     "identity": {
         "ztAPI": "https://ziti.exts.platform.softwareone.com/edge/client/v1",
         "ztAPIs": None,
@@ -35,6 +37,17 @@ INSTANCE_CREATE_RESPONSE = {
             "ca": "pem:-----BEGIN CERTIFICATE-----\n...\n-----END CERTIFICATE-----\n",
         },
         "enableHa": None,
+        "mrok": {
+            "identity": "ins-0000-0000-0000.ext-0000-0000",
+            "extension": "ext-0000-0000",
+            "instance": "ins-0000-0000-0000",
+            "domain": "ext.s1.today",
+            "tags": {
+                "mrok-service": "ext-0000-0000",
+                "mrok-identity-type": "instance",
+                "mrok": "0.4.0",
+            },
+        },
     },
     "tags": {
         "account": "ACC-5555-3333",

mrok/controller/schemas.py

@@ -9,12 +9,12 @@ from pydantic import (
     computed_field,
 )
 
-from mrok.ziti.api import TagsType
+from mrok.types.ziti import Tags
 
 
 class BaseSchema(BaseModel):
     model_config = ConfigDict(from_attributes=True, extra="ignore")
-    tags: TagsType | None = None
+    tags: Tags | None = None
 
 
 class IdSchema(BaseModel):

mrok/frontend/__init__.py

@@ -0,0 +1,3 @@
+from mrok.frontend.main import run
+
+__all__ = ["run"]

mrok/frontend/app.py

@@ -0,0 +1,75 @@
+import re
+
+from httpcore import AsyncConnectionPool
+
+from mrok.conf import get_settings
+from mrok.proxy.app import ProxyAppBase
+from mrok.proxy.backend import AIOZitiNetworkBackend
+from mrok.proxy.exceptions import InvalidTargetError
+from mrok.types.proxy import Scope
+
+RE_SUBDOMAIN = re.compile(r"(?i)^(?:EXT-\d{4}-\d{4}|INS-\d{4}-\d{4}-\d{4})$")
+
+
+class FrontendProxyApp(ProxyAppBase):
+    def __init__(
+        self,
+        identity_file: str,
+        *,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
+        retries=0,
+    ):
+        self._identity_file = identity_file
+        self._proxy_domain = self._get_proxy_domain()
+        super().__init__(
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+            retries=retries,
+        )
+
+    def setup_connection_pool(
+        self,
+        max_connections: int | None,
+        max_keepalive_connections: int | None,
+        keepalive_expiry: float | None,
+        retries: int,
+    ) -> AsyncConnectionPool:
+        return AsyncConnectionPool(
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+            retries=retries,
+            network_backend=AIOZitiNetworkBackend(self._identity_file),
+        )
+
+    def get_upstream_base_url(self, scope: Scope) -> str:
+        target = self._get_target_name(
+            {k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {})}
+        )
+        return f"http://{target.lower()}"
+
+    def _get_proxy_domain(self):
+        settings = get_settings()
+        return (
+            settings.proxy.domain
+            if settings.proxy.domain[0] == "."
+            else f".{settings.proxy.domain}"
+        )
+
+    def _get_target_from_header(self, headers: dict[str, str], name: str) -> str | None:
+        header_value = headers.get(name, "")
+        if self._proxy_domain in header_value:
+            if ":" in header_value:
+                header_value, _ = header_value.split(":", 1)
+            return header_value[: -len(self._proxy_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 or not RE_SUBDOMAIN.fullmatch(target):
+            raise InvalidTargetError()
+        return target

mrok/{proxy → frontend}/main.py

@@ -6,9 +6,8 @@ from gunicorn.app.base import BaseApplication
 from uvicorn_worker import UvicornWorker
 
 from mrok.conf import get_settings
-from mrok.http.lifespan import LifespanWrapper
+from mrok.frontend.app import FrontendProxyApp
 from mrok.logging import get_logging_config
-from mrok.proxy.app import ProxyApp
 
 
 class MrokUvicornWorker(UvicornWorker):
@@ -39,20 +38,23 @@ def run(
     host: str,
     port: int,
     workers: int,
+    max_connections: int | None,
+    max_keepalive_connections: int | None,
+    keepalive_expiry: float | None,
 ):
-    proxy_app = ProxyApp(identity_file)
-
-    asgi_app = LifespanWrapper(
-        proxy_app,
-        proxy_app.startup,
-        proxy_app.shutdown,
+    app = FrontendProxyApp(
+        str(identity_file),
+        max_connections=max_connections,
+        max_keepalive_connections=max_keepalive_connections,
+        keepalive_expiry=keepalive_expiry,
     )
+
     options = {
         "bind": f"{host}:{port}",
         "workers": workers,
-        "worker_class": "mrok.proxy.main.MrokUvicornWorker",
+        "worker_class": "mrok.frontend.main.MrokUvicornWorker",
         "logconfig_dict": get_logging_config(get_settings()),
         "reload": False,
     }
 
-    StandaloneApplication(asgi_app, options).run()
+    StandaloneApplication(app, options).run()

mrok/proxy/__init__.py

@@ -1,3 +0,0 @@
-from mrok.proxy.main import run
-
-__all__ = ["run"]