mrok 0.4.5__py3-none-any.whl → 0.5.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.datastructures import Event, HTTPHeaders, HTTPResponse, WorkerMetrics, ZitiMrokMeta
 
 
 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: ZitiMrokMeta) -> 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,29 +1,77 @@
-import asyncio
 import logging
 from pathlib import Path
+from typing import Literal
 
-from mrok.http.forwarder import ForwardAppBase
-from mrok.http.types import Scope, StreamReader, StreamWriter
+from httpcore import AsyncConnectionPool
+
+from mrok.proxy.app import ProxyAppBase
+from mrok.proxy.types 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=1000,
+        max_keepalive_connections=10,
+        keepalive_expiry=120,
+        retries=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
 
-    async def select_backend(
+    def setup_connection_pool(
         self,
-        scope: Scope,
-        headers: dict[str, str],
-    ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
-        if isinstance(self._target_address, tuple):
-            return await asyncio.open_connection(*self._target_address)
-        return await asyncio.open_unix_connection(str(self._target_address))
+        max_connections: int | None = 1000,
+        max_keepalive_connections: int | None = 10,
+        keepalive_expiry: float | None = 120.0,
+        retries: int = 0,
+    ) -> 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,
+        )
+
+    def get_upstream_base_url(self, scope: Scope) -> str:
+        if self._target_type == "unix":
+            return "http://localhost"
+        return f"http://{self._target_address}"
+
+    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)
+
+        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}")
+
+        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,7 +11,7 @@ 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,
         publishers_port: int = 50000,
         subscribers_port: int = 50001,
@@ -23,10 +23,10 @@ class SidecarAgent(MasterBase):
             publishers_port,
             subscribers_port,
         )
-        self.target_address = target_addr
+        self._target = target
 
     def get_asgi_app(self):
-        return ForwardApp(self.target_address)
+        return SidecarProxyApp(self._target)
 
 
 def run(

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.proxy.types import ASGIApp
 
 
 class ZiticornAgent(MasterBase):

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/register/extensions.py

@@ -1,5 +1,4 @@
 import asyncio
-import re
 from typing import Annotated
 
 import typer
@@ -7,11 +6,10 @@ from rich import print
 
 from mrok.cli.commands.admin.utils import parse_tags
 from mrok.conf import Settings
+from mrok.constants import RE_EXTENSION_ID
 from mrok.ziti.api import ZitiManagementAPI
 from mrok.ziti.services import register_service
 
-RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
-
 
 async def do_register(settings: Settings, extension_id: str, tags: list[str] | None):
     async with ZitiManagementAPI(settings) as api:
@@ -20,7 +18,7 @@ async def do_register(settings: Settings, extension_id: str, tags: list[str] | N
 
 def validate_extension_id(extension_id: str) -> str:
     if not RE_EXTENSION_ID.fullmatch(extension_id):
-        raise typer.BadParameter("ext_id must match EXT-xxxx-yyyy (case-insensitive)")
+        raise typer.BadParameter("it must match EXT-xxxx-yyyy (case-insensitive)")
     return extension_id
 
 

mrok/cli/commands/admin/register/instances.py

@@ -1,6 +1,5 @@
 import asyncio
 import json
-import re
 from pathlib import Path
 from typing import Annotated
 
@@ -8,11 +7,10 @@ import typer
 
 from mrok.cli.commands.admin.utils import parse_tags
 from mrok.conf import Settings
+from mrok.constants import RE_EXTENSION_ID, RE_INSTANCE_ID
 from mrok.ziti.api import ZitiClientAPI, ZitiManagementAPI
 from mrok.ziti.identities import register_identity
 
-RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
-
 
 async def do_register(
     settings: Settings, extension_id: str, instance_id: str, tags: list[str] | None
@@ -25,10 +23,16 @@ async def do_register(
 
 def validate_extension_id(extension_id: str):
     if not RE_EXTENSION_ID.fullmatch(extension_id):
-        raise typer.BadParameter("ext_id must match EXT-xxxx-yyyy (case-insensitive)")
+        raise typer.BadParameter("it must match EXT-xxxx-yyyy (case-insensitive)")
     return extension_id
 
 
+def validate_instance_id(instance_id: str):
+    if not RE_INSTANCE_ID.fullmatch(instance_id):
+        raise typer.BadParameter("it must match INS-xxxx-yyyy-zzzz (case-insensitive)")
+    return instance_id
+
+
 def register(app: typer.Typer) -> None:
     @app.command("instance")
     def register_instance(
@@ -36,7 +40,9 @@ def register(app: typer.Typer) -> None:
         extension_id: str = typer.Argument(
             ..., callback=validate_extension_id, help="Extension ID in format EXT-xxxx-yyyy"
         ),
-        instance_id: str = typer.Argument(..., help="Instance ID"),
+        instance_id: str = typer.Argument(
+            ..., callback=validate_instance_id, help="Instance ID in format INS-xxxx-yyyy-zzzz"
+        ),
         output: Path = typer.Argument(
             ...,
             file_okay=True,

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/{proxy → frontend}/run.py

@@ -3,7 +3,7 @@ from typing import Annotated
 
 import typer
 
-from mrok import proxy
+from mrok import frontend
 from mrok.cli.utils import number_of_workers
 
 default_workers = number_of_workers()
@@ -11,7 +11,7 @@ default_workers = number_of_workers()
 
 def register(app: typer.Typer) -> None:
     @app.command("run")
-    def run_proxy(
+    def run_frontend(
         ctx: typer.Context,
         identity_file: Path = typer.Argument(
             ...,
@@ -45,5 +45,5 @@ def register(app: typer.Typer) -> None:
             ),
         ] = default_workers,
     ):
-        """Run the mrok proxy with Gunicorn and Uvicorn workers."""
-        proxy.run(identity_file, host, port, workers)
+        """Run the mrok frontend with Gunicorn and Uvicorn workers."""
+        frontend.run(identity_file, host, port, workers)

mrok/constants.py

@@ -0,0 +1,4 @@
+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}")

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/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.proxy.types 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 = 1000,
+        max_keepalive_connections: int = 10,
+        keepalive_expiry: float = 120.0,
+        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 = 1000,
+        max_keepalive_connections: int | None = 100,
+        keepalive_expiry: float | None = 120.0,
+        retries: int = 0,
+    ) -> 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):
@@ -40,19 +39,14 @@ def run(
     port: int,
     workers: int,
 ):
-    proxy_app = ProxyApp(identity_file)
+    app = FrontendProxyApp(str(identity_file))
 
-    asgi_app = LifespanWrapper(
-        proxy_app,
-        proxy_app.startup,
-        proxy_app.shutdown,
-    )
     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"]

mrok/proxy/app.py

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