mrok 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

mrok/proxy/app.py

@@ -1,12 +1,12 @@
+import asyncio
 import logging
-import os
 from pathlib import Path
 
 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.ziti import ZitiConnectionManager
+from mrok.proxy.ziti import ZitiSocketCache
 
 logger = logging.getLogger("mrok.proxy")
 
@@ -21,8 +21,6 @@ class ProxyApp(ForwardAppBase):
         identity_file: str | Path,
         *,
         read_chunk_size: int = 65536,
-        ziti_connection_ttl_seconds: float = 60,
-        ziti_conn_cache_purge_interval_seconds: float = 10,
     ) -> None:
         super().__init__(read_chunk_size=read_chunk_size)
         self._identity_file = identity_file
@@ -32,40 +30,28 @@ class ProxyApp(ForwardAppBase):
             if settings.proxy.domain[0] == "."
             else f".{settings.proxy.domain}"
         )
-        self._conn_manager = ZitiConnectionManager(
-            identity_file,
-            ttl_seconds=ziti_connection_ttl_seconds,
-            purge_interval=ziti_conn_cache_purge_interval_seconds,
-        )
-
-    def get_target_from_header(self, name: str, headers: dict[str, str]) -> str:
-        header_value = headers.get(name)
-        if not header_value:
-            raise ProxyError(
-                f"Header {name} not found!",
-            )
-        if ":" in header_value:
-            header_value, _ = header_value.split(":", 1)
-        if not header_value.endswith(self._proxy_wildcard_domain):
-            raise ProxyError(f"Unexpected value for {name} header: `{header_value}`.")
+        self._ziti_socket_cache = ZitiSocketCache(self._identity_file)
 
-        return header_value[: -len(self._proxy_wildcard_domain)]
+    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:
-        try:
-            return self.get_target_from_header("x-forwared-for", headers)
-        except ProxyError as pe:
-            logger.warning(pe)
-            return self.get_target_from_header("host", headers)
+        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
 
     async def startup(self):
         setup_logging(get_settings())
-        await self._conn_manager.start()
-        logger.info(f"Proxy app startup completed: {os.getpid()}")
 
     async def shutdown(self):
-        await self._conn_manager.stop()
-        logger.info(f"Proxy app shutdown completed: {os.getpid()}")
+        await self._ziti_socket_cache.stop()
 
     async def select_backend(
         self,
@@ -73,5 +59,6 @@ class ProxyApp(ForwardAppBase):
         headers: dict[str, str],
     ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
         target_name = self.get_target_name(headers)
-
-        return await self._conn_manager.get(target_name)
+        sock = self._ziti_socket_cache.get_or_create(target_name)
+        reader, writer = await asyncio.open_connection(sock=sock)
+        return reader, writer

mrok/proxy/ziti.py

@@ -1,173 +1,103 @@
-"""Ziti-backed connection manager for the proxy.
-
-This manager owns creation of connections via an OpenZiti context, wraps
-streams to observe IO errors, evicts idle entries, and serializes creation
-per-key.
-"""
-
 import asyncio
-import logging
-import time
+import contextlib
+from asyncio import Task
 from pathlib import Path
 
-# typing imports intentionally minimized
 import openziti
+from aiocache import Cache
+from openziti.context import ZitiContext
+from openziti.zitisock import ZitiSocket
 
-from mrok.http.types import StreamReader, StreamWriter
-from mrok.proxy.dataclasses import CachedStreamEntry
-from mrok.proxy.streams import CachedStreamReader, CachedStreamWriter
-from mrok.proxy.types import CachedStream, ConnectionKey
-
-logger = logging.getLogger("mrok.proxy")
 
-
-class ZitiConnectionManager:
+class ZitiSocketCache:
     def __init__(
         self,
         identity_file: str | Path,
-        ziti_timeout_ms: int = 10000,
+        ziti_ctx_timeout_ms: int = 10_000,
         ttl_seconds: float = 60.0,
-        purge_interval: float = 10.0,
-    ):
+        cleanup_interval: float = 10.0,
+    ) -> None:
         self._identity_file = identity_file
-        self._ziti_ctx = None
-        self._ziti_timeout_ms = ziti_timeout_ms
-        self._ttl = float(ttl_seconds)
-        self._purge_interval = float(purge_interval)
-        self._cache: dict[ConnectionKey, CachedStreamEntry] = {}
-        self._lock = asyncio.Lock()
-        self._in_progress: dict[ConnectionKey, asyncio.Lock] = {}
-        self._purge_task: asyncio.Task | None = None
-
-    async def get(self, target: str) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
-        head, _, tail = target.partition(".")
-        terminator = target if head and tail else ""
-        service = tail if tail else head
-        r, w = await self._get_or_create_key((service, terminator))
-        return r, w
-
-    async def invalidate(self, key: ConnectionKey) -> None:
-        async with self._lock:
-            item = self._cache.pop(key, None)
-        if item is None:
-            return
-        await self._close_writer(item.writer)
+        self._ziti_ctx_timeout_ms = ziti_ctx_timeout_ms
+        self._ttl_seconds = ttl_seconds
+        self._cleanup_interval = cleanup_interval
 
-    async def start(self) -> None:
+        self._ziti_ctx: ZitiContext | None = None
+        self._cache = Cache(Cache.MEMORY)
+        self._active_sockets: dict[str, ZitiSocket] = {}
+        self._cleanup_task: Task | None = None
+
+    def _get_ziti_ctx(self) -> ZitiContext:
         if self._ziti_ctx is None:
-            ctx, err = openziti.load(str(self._identity_file), timeout=self._ziti_timeout_ms)
+            ctx, err = openziti.load(str(self._identity_file), timeout=self._ziti_ctx_timeout_ms)
             if err != 0:
                 raise Exception(f"Cannot create a Ziti context from the identity file: {err}")
             self._ziti_ctx = ctx
-        if self._purge_task is None:
-            self._purge_task = asyncio.create_task(self._purge_loop())
-            logger.info("Ziti connection manager started")
-
-    async def stop(self) -> None:
-        if self._purge_task is not None:
-            self._purge_task.cancel()
-            try:
-                await self._purge_task
-            except asyncio.CancelledError:
-                logger.debug("Purge task was cancelled")
-            except Exception as e:
-                logger.warning(f"An error occurred stopping the purge task: {e}")
-            self._purge_task = None
-            logger.info("Ziti connection manager stopped")
-
-        async with self._lock:
-            items = list(self._cache.items())
-            self._cache.clear()
-
-        for _, item in items:
-            await self._close_writer(item.writer)
-
-    async def _purge_loop(self) -> None:
+        return self._ziti_ctx
+
+    async def _create_socket(self, key: str):
+        return self._get_ziti_ctx().connect(key)
+
+    async def get_or_create(self, key: str):
+        sock = await self._cache.get(key)
+
+        if sock:
+            await self._cache.set(key, sock, ttl_seconds=self._ttl_seconds)
+            self._active_sockets[key] = sock
+            return sock
+
+        sock = await self._create_socket(key)
+        await self._cache.set(key, sock, ttl_seconds=self._ttl_seconds)
+        self._active_sockets[key] = sock
+        return sock
+
+    async def invalidate(self, key: str):
+        sock = await self._cache.get(key)
+        if sock:
+            await self._close_socket(sock)
+
+        await self._cache.delete(key)
+        self._active_sockets.pop(key, None)
+
+    async def start(self):
+        self._cleanup_task = asyncio.create_task(self._periodic_cleanup())
+        # Warmup ziti context
+        self._get_ziti_ctx()
+
+    async def stop(self):
+        """
+        Cleanup: stop background task + close all sockets.
+        """
+        self._cleanup_task.cancel()
+        with contextlib.suppress(Exception):
+            await self._cleanup_task
+
+        for sock in list(self._active_sockets.values()):
+            await self._close_socket(sock)
+
+        self._active_sockets.clear()
+        await self._cache.clear()
+
+    @staticmethod
+    async def _close_socket(sock: ZitiSocket):
+        with contextlib.suppress(Exception):
+            sock.close()
+
+    async def _periodic_cleanup(self):
         try:
             while True:
-                await asyncio.sleep(self._purge_interval)
-                await self._purge_once()
+                await asyncio.sleep(self._cleanup_interval)
+                await self._cleanup_once()
         except asyncio.CancelledError:
             return
 
-    async def _purge_once(self) -> None:
-        to_close: list[tuple[StreamReader, StreamWriter]] = []
-        async with self._lock:
-            now = time.time()
-            for key, item in list(self._cache.items()):
-                if now - item.last_access > self._ttl:
-                    to_close.append((item.reader, item.writer))
-                    del self._cache[key]
+    async def _cleanup_once(self):
+        keys_now = set(await self._cache.keys())
+        known_keys = set(self._active_sockets.keys())
 
-        for _, writer in to_close:
-            writer.close()
-            await self._close_writer(writer)
+        expired = known_keys - keys_now
 
-    def _is_writer_closed(self, writer: StreamWriter) -> bool:
-        return writer.transport.is_closing()
-
-    async def _close_writer(self, writer: StreamWriter) -> None:
-        writer.close()
-        try:
-            await writer.wait_closed()
-        except Exception as e:
-            logger.debug(f"Error closing writer: {e}")
-
-    async def _get_or_create_key(self, key: ConnectionKey) -> CachedStream:
-        """Internal: create or return a cached wrapped pair for the concrete key."""
-        await self._purge_once()
-        to_close = None
-        async with self._lock:
-            if key in self._cache:
-                now = time.time()
-                item = self._cache[key]
-                reader, writer = item.reader, item.writer
-                if not self._is_writer_closed(writer) and not reader.at_eof():
-                    self._cache[key] = CachedStreamEntry(reader, writer, now)
-                    return reader, writer
-                to_close = writer
-                del self._cache[key]
-
-            lock = self._in_progress.get(key)
-            if lock is None:
-                lock = asyncio.Lock()
-                self._in_progress[key] = lock
-
-        if to_close:
-            await self._close_writer(to_close)
-
-        async with lock:
-            try:
-                # # double-check cache after acquiring the per-key lock
-                # async with self._lock:
-                #     now = time.time()
-                #     if key in self._cache:
-                #         r, w, _ = self._cache[key]
-                #         if not self._is_writer_closed(w) and not r.at_eof():
-                #             self._cache[key] = (r, w, now)
-                #             return r, w
-
-                # perform creation via ziti context
-                extension, instance = key
-                logger.info(f"Create connection to {extension}: {instance}")
-                # loop = asyncio.get_running_loop()
-                # sock = await loop.run_in_executor(None, self._ziti_ctx.connect,
-                # extension, instance)
-                if instance:
-                    sock = self._ziti_ctx.connect(
-                        extension, terminator=instance
-                    )  # , terminator=instance)
-                else:
-                    sock = self._ziti_ctx.connect(extension)
-                orig_reader, orig_writer = await asyncio.open_connection(sock=sock)
-
-                reader = CachedStreamReader(orig_reader, key, self)
-                writer = CachedStreamWriter(orig_writer, key, self)
-
-                async with self._lock:
-                    self._cache[key] = CachedStreamEntry(reader, writer, time.time())
-
-                return reader, writer
-            finally:
-                async with self._lock:
-                    self._in_progress.pop(key, None)
+        for key in expired:
+            sock = self._active_sockets.pop(key, None)
+            if sock:
+                await self._close_socket(sock)

{mrok-0.4.1.dist-info → mrok-0.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mrok
-Version: 0.4.1
+Version: 0.4.3
 Summary: MPT Extensions OpenZiti Orchestrator
 Author: SoftwareOne AG
 License:                                  Apache License
@@ -206,6 +206,7 @@ License:                                  Apache License
            limitations under the License.
 License-File: LICENSE.txt
 Requires-Python: <4,>=3.12
+Requires-Dist: aiocache<0.13.0,>=0.12.3
 Requires-Dist: asn1crypto<2.0.0,>=1.5.1
 Requires-Dist: cryptography<46.0.0,>=45.0.7
 Requires-Dist: dynaconf<4.0.0,>=3.2.11

{mrok-0.4.1.dist-info → mrok-0.4.3.dist-info}/RECORD

@@ -71,12 +71,9 @@ mrok/http/server.py,sha256=Mj7C85fc-DXp-WTBWaOd7ag808oliLmFBH5bf-G2FHg,370
 mrok/http/types.py,sha256=XpNrvbfpANKvmjOBYtLF1FmDHoJF3z_MIMQHXoJlvmE,1302
 mrok/http/utils.py,sha256=sOixYu3R9-nNoMFYdifrreYvcFRIHYVtb6AAmtVzaLE,2125
 mrok/proxy/__init__.py,sha256=vWXyImroqM1Eq8e_oFPBup8VJ3reyp8SVjFTbLzRkI8,51
-mrok/proxy/app.py,sha256=YP4hM3BIRf_MFw9YNEnsFrrBz76zVegrNFcvcsKTFvI,2579
-mrok/proxy/dataclasses.py,sha256=DtX-Yuma-uOECOPefJnoQJhZMEtT6Za_27cd-lJE9Iw,237
+mrok/proxy/app.py,sha256=-xmHKSSqQhNdwISRXQCUVLI1n6HWs3d46cgKQ6_NC3A,2157
 mrok/proxy/main.py,sha256=ZXpticE6J4FABaslDB_8J5qklPsf3e7xIFSZmcPAAjQ,1588
-mrok/proxy/streams.py,sha256=6TMZwrQPbSyQqpqavsoTeyUmS2O026pJfiCnxLopPqg,3425
-mrok/proxy/types.py,sha256=dgWqAj6dFGVH_Q8-k8sU5h18yoUF_fTn-SRPIfEs_gA,308
-mrok/proxy/ziti.py,sha256=dKd6UzmEAFu9-gey871sPEDUZTkt4YVPyCYRzeA5mlA,6539
+mrok/proxy/ziti.py,sha256=Ce02EimVJNffyluFr1nhOMXh6kZ0kY2nt6Shg58wkBw,3206
 mrok/ziti/__init__.py,sha256=20OWMiexRhOovZOX19zlX87-V78QyWnEnSZfyAftUdE,263
 mrok/ziti/api.py,sha256=KvGiT9d4oSgC3JbFWLDQyuHcLX2HuZJoJ8nHmWtCDkY,16154
 mrok/ziti/bootstrap.py,sha256=QIDhlkIxPW2QRuumFq2D1WDbD003P5f3z24pAUsyeBI,2696
@@ -85,8 +82,8 @@ mrok/ziti/errors.py,sha256=yYCbVDwktnR0AYduqtynIjo73K3HOhIrwA_vQimvEd4,368
 mrok/ziti/identities.py,sha256=1BcwfqAJHMBhc3vRaf0aLaIkoHskj5Xe2Lsq2lO9Vs8,6735
 mrok/ziti/pki.py,sha256=o2tySqHC8-7bvFuI2Tqxg9vX6H6ZSxWxfP_9x29e19M,1954
 mrok/ziti/services.py,sha256=zR1PEBYwXVou20iJK4euh0ZZFAo9UB8PZk8f6SDmiUE,3194
-mrok-0.4.1.dist-info/METADATA,sha256=PuydR9Z4yuDSDMuJrn1FqWGb18fJTA6TFChAj-LwWZE,15796
-mrok-0.4.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mrok-0.4.1.dist-info/entry_points.txt,sha256=tloXwvU1uJicBJR2h-8HoVclPgwJWDwuREMHN8Zq-nU,38
-mrok-0.4.1.dist-info/licenses/LICENSE.txt,sha256=6PaICaoA3yNsZKLv5G6OKqSfLSoX7MakYqTDgJoTCBs,11346
-mrok-0.4.1.dist-info/RECORD,,
+mrok-0.4.3.dist-info/METADATA,sha256=rfG5MT0fAz27oU7Jzr8RFD3LLA9Syci51df1gHFSgLI,15836
+mrok-0.4.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mrok-0.4.3.dist-info/entry_points.txt,sha256=tloXwvU1uJicBJR2h-8HoVclPgwJWDwuREMHN8Zq-nU,38
+mrok-0.4.3.dist-info/licenses/LICENSE.txt,sha256=6PaICaoA3yNsZKLv5G6OKqSfLSoX7MakYqTDgJoTCBs,11346
+mrok-0.4.3.dist-info/RECORD,,

mrok/proxy/dataclasses.py

@@ -1,12 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-from mrok.http.types import StreamReader, StreamWriter
-
-
-@dataclass
-class CachedStreamEntry:
-    reader: StreamReader
-    writer: StreamWriter
-    last_access: float

mrok/proxy/streams.py

@@ -1,124 +0,0 @@
-import asyncio
-
-from mrok.proxy.types import ConnectionCache, ConnectionKey
-
-
-class CachedStreamReader:
-    def __init__(
-        self,
-        reader: asyncio.StreamReader,
-        key: ConnectionKey,
-        manager: ConnectionCache,
-    ):
-        self._reader = reader
-        self._key = key
-        self._manager = manager
-
-    async def read(self, n: int = -1) -> bytes:
-        try:
-            return await self._reader.read(n)
-        except (
-            asyncio.CancelledError,
-            asyncio.IncompleteReadError,
-            asyncio.LimitOverrunError,
-            BrokenPipeError,
-            ConnectionAbortedError,
-            ConnectionResetError,
-            RuntimeError,
-            TimeoutError,
-            UnicodeDecodeError,
-        ):
-            asyncio.create_task(self._manager.invalidate(self._key))
-            raise
-
-    async def readexactly(self, n: int) -> bytes:
-        try:
-            return await self._reader.readexactly(n)
-        except (
-            asyncio.CancelledError,
-            asyncio.IncompleteReadError,
-            asyncio.LimitOverrunError,
-            BrokenPipeError,
-            ConnectionAbortedError,
-            ConnectionResetError,
-            RuntimeError,
-            TimeoutError,
-            UnicodeDecodeError,
-        ):
-            asyncio.create_task(self._manager.invalidate(self._key))
-            raise
-
-    async def readline(self) -> bytes:
-        try:
-            return await self._reader.readline()
-        except (
-            asyncio.CancelledError,
-            asyncio.IncompleteReadError,
-            asyncio.LimitOverrunError,
-            BrokenPipeError,
-            ConnectionAbortedError,
-            ConnectionResetError,
-            RuntimeError,
-            TimeoutError,
-            UnicodeDecodeError,
-        ):
-            asyncio.create_task(self._manager.invalidate(self._key))
-            raise
-
-    def at_eof(self) -> bool:
-        return self._reader.at_eof()
-
-    @property
-    def underlying(self) -> asyncio.StreamReader:
-        return self._reader
-
-
-class CachedStreamWriter:
-    def __init__(
-        self,
-        writer: asyncio.StreamWriter,
-        key: ConnectionKey,
-        manager: ConnectionCache,
-    ):
-        self._writer = writer
-        self._key = key
-        self._manager = manager
-
-    def write(self, data: bytes) -> None:
-        try:
-            return self._writer.write(data)
-        except (RuntimeError, TypeError):
-            asyncio.create_task(self._manager.invalidate(self._key))
-            raise
-
-    async def drain(self) -> None:
-        try:
-            return await self._writer.drain()
-        except (
-            asyncio.CancelledError,
-            BrokenPipeError,
-            ConnectionAbortedError,
-            ConnectionResetError,
-            RuntimeError,
-            TimeoutError,
-        ):
-            asyncio.create_task(self._manager.invalidate(self._key))
-            raise
-
-    def close(self) -> None:
-        return self._writer.close()
-
-    async def wait_closed(self) -> None:
-        try:
-            return await self._writer.wait_closed()
-        except (ConnectionResetError, BrokenPipeError):
-            asyncio.create_task(self._manager.invalidate(self._key))
-            raise
-
-    @property
-    def transport(self):
-        return self._writer.transport
-
-    @property
-    def underlying(self) -> asyncio.StreamWriter:
-        return self._writer

mrok/proxy/types.py

@@ -1,12 +0,0 @@
-from __future__ import annotations
-
-from typing import Protocol
-
-from mrok.http.types import StreamReader, StreamWriter
-
-ConnectionKey = tuple[str, str | None]
-CachedStream = tuple[StreamReader, StreamWriter]
-
-
-class ConnectionCache(Protocol):
-    async def invalidate(self, key: ConnectionKey) -> None: ...