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

mrok/http/forwarder.py

@@ -1,14 +1,10 @@
 import abc
 import asyncio
 import logging
-from collections.abc import Awaitable, Callable
-from typing import Any
 
-logger = logging.getLogger("mrok.proxy")
+from mrok.http.types import ASGIReceive, ASGISend, Scope, StreamReader, StreamWriter
 
-Scope = dict[str, Any]
-ASGIReceive = Callable[[], Awaitable[dict[str, Any]]]
-ASGISend = Callable[[dict[str, Any]], Awaitable[None]]
+logger = logging.getLogger("mrok.proxy")
 
 
 class BackendNotFoundError(Exception):
@@ -24,24 +20,71 @@ class ForwardAppBase(abc.ABC):
     and streaming logic (requests and responses).
     """
 
-    def __init__(self, read_chunk_size: int = 65536) -> None:
-        # number of bytes to read per iteration when streaming bodies
+    def __init__(
+        self,
+        read_chunk_size: int = 65536,
+        lifespan_timeout: float = 10.0,
+    ) -> None:
         self._read_chunk_size: int = int(read_chunk_size)
+        self._lifespan_timeout = lifespan_timeout
+
+    async def handle_lifespan(self, receive: ASGIReceive, send: ASGISend) -> None:
+        while True:
+            event = await receive()
+            etype = event.get("type")
+
+            if etype == "lifespan.startup":
+                try:
+                    await asyncio.wait_for(self.startup(), self._lifespan_timeout)
+                except TimeoutError:
+                    logger.exception("Lifespan startup timed out")
+                    await send({"type": "lifespan.startup.failed", "message": "startup timeout"})
+                    continue
+                except Exception as e:
+                    logger.exception("Exception during lifespan startup")
+                    await send({"type": "lifespan.startup.failed", "message": str(e)})
+                    continue
+                await send({"type": "lifespan.startup.complete"})
+
+            elif etype == "lifespan.shutdown":
+                try:
+                    await asyncio.wait_for(self.shutdown(), self._lifespan_timeout)
+                except TimeoutError:
+                    logger.exception("Lifespan shutdown timed out")
+                    await send({"type": "lifespan.shutdown.failed", "message": "shutdown timeout"})
+                    return
+                except Exception as exc:
+                    logger.exception("Exception during lifespan shutdown")
+                    await send({"type": "lifespan.shutdown.failed", "message": str(exc)})
+                    return
+                await send({"type": "lifespan.shutdown.complete"})
+                return
 
     @abc.abstractmethod
     async def select_backend(
         self,
         scope: Scope,
         headers: dict[str, str],
-    ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter] | tuple[None, None]:
+    ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
         """Return (reader, writer) connected to the target backend."""
 
+    async def startup(self):
+        return
+
+    async def shutdown(self):
+        return
+
     async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend) -> None:
         """ASGI callable entry point.
 
         Delegates to smaller helper methods for readability. Subclasses only
         need to implement backend selection.
         """
+        scope_type = scope.get("type")
+        if scope_type == "lifespan":
+            await self.handle_lifespan(receive, send)
+            return
+
         if scope.get("type") != "http":
             await send({"type": "http.response.start", "status": 500, "headers": []})
             await send({"type": "http.response.body", "body": b"Unsupported"})
@@ -116,7 +159,7 @@ class ForwardAppBase(abc.ABC):
 
     async def write_request_line_and_headers(
         self,
-        writer: asyncio.StreamWriter,
+        writer: StreamWriter,
         method: str,
         path_qs: str,
         headers: list[tuple[bytes, bytes]],
@@ -130,7 +173,7 @@ class ForwardAppBase(abc.ABC):
         await writer.drain()
 
     async def stream_request_body(
-        self, receive: ASGIReceive, writer: asyncio.StreamWriter, use_chunked: bool
+        self, receive: ASGIReceive, writer: StreamWriter, use_chunked: bool
     ) -> None:
         if use_chunked:
             await self.stream_request_chunked(receive, writer)
@@ -138,9 +181,7 @@ class ForwardAppBase(abc.ABC):
 
         await self.stream_request_until_end(receive, writer)
 
-    async def stream_request_chunked(
-        self, receive: ASGIReceive, writer: asyncio.StreamWriter
-    ) -> None:
+    async def stream_request_chunked(self, receive: ASGIReceive, writer: StreamWriter) -> None:
         """Send request body to backend using HTTP/1.1 chunked encoding."""
         while True:
             event = await receive()
@@ -160,9 +201,7 @@ class ForwardAppBase(abc.ABC):
         writer.write(b"0\r\n\r\n")
         await writer.drain()
 
-    async def stream_request_until_end(
-        self, receive: ASGIReceive, writer: asyncio.StreamWriter
-    ) -> None:
+    async def stream_request_until_end(self, receive: ASGIReceive, writer: StreamWriter) -> None:
         """Send request body to backend when content length/transfer-encoding
         already provided (no chunking).
         """
@@ -180,7 +219,7 @@ class ForwardAppBase(abc.ABC):
                 return
 
     async def read_status_and_headers(
-        self, reader: asyncio.StreamReader, first_line: bytes
+        self, reader: StreamReader, first_line: bytes
     ) -> tuple[int, list[tuple[bytes, bytes]], dict[bytes, bytes]]:
         parts = first_line.decode(errors="ignore").split(" ", 2)
         status = int(parts[1]) if len(parts) >= 2 and parts[1].isdigit() else 502
@@ -217,14 +256,14 @@ class ForwardAppBase(abc.ABC):
         except Exception:
             return None
 
-    async def drain_trailers(self, reader: asyncio.StreamReader) -> None:
+    async def drain_trailers(self, reader: StreamReader) -> None:
         """Consume trailer header lines until an empty line is encountered."""
         while True:
             trailer = await reader.readline()
             if trailer in (b"\r\n", b"\n", b""):
                 break
 
-    async def stream_response_chunked(self, reader: asyncio.StreamReader, send: ASGISend) -> None:
+    async def stream_response_chunked(self, reader: StreamReader, send: ASGISend) -> None:
         """Read chunked-encoded response from reader, decode and forward to ASGI send."""
         while True:
             size_line = await reader.readline()
@@ -253,7 +292,7 @@ class ForwardAppBase(abc.ABC):
         await send({"type": "http.response.body", "body": b"", "more_body": False})
 
     async def stream_response_with_content_length(
-        self, reader: asyncio.StreamReader, send: ASGISend, content_length: int
+        self, reader: StreamReader, send: ASGISend, content_length: int
     ) -> None:
         """Read exactly content_length bytes and forward to ASGI send events."""
         remaining = content_length
@@ -272,7 +311,7 @@ class ForwardAppBase(abc.ABC):
         if not sent_final:
             await send({"type": "http.response.body", "body": b"", "more_body": False})
 
-    async def stream_response_until_eof(self, reader: asyncio.StreamReader, send: ASGISend) -> None:
+    async def stream_response_until_eof(self, reader: StreamReader, send: ASGISend) -> None:
         """Read from reader until EOF and forward chunks to ASGI send events."""
         while True:
             chunk = await reader.read(self._read_chunk_size)
@@ -282,7 +321,7 @@ class ForwardAppBase(abc.ABC):
         await send({"type": "http.response.body", "body": b"", "more_body": False})
 
     async def stream_response_body(
-        self, reader: asyncio.StreamReader, send: ASGISend, raw_headers: dict[bytes, bytes]
+        self, reader: StreamReader, send: ASGISend, raw_headers: dict[bytes, bytes]
     ) -> None:
         te = raw_headers.get(b"transfer-encoding", b"").lower()
         cl = raw_headers.get(b"content-length")

mrok/http/lifespan.py

@@ -1,10 +1,39 @@
 import logging
+from collections.abc import Awaitable, Callable
 
 from uvicorn.config import Config
 from uvicorn.lifespan.on import LifespanOn
 
+AsyncCallback = Callable[[], Awaitable[None]]
+
 
 class MrokLifespan(LifespanOn):
     def __init__(self, config: Config) -> None:
         super().__init__(config)
         self.logger = logging.getLogger("mrok.proxy")
+
+
+class LifespanWrapper:
+    def __init__(
+        self, app, on_startup: AsyncCallback | None = None, on_shutdown: AsyncCallback | 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:
+                        await self.on_startup()
+                    await send({"type": "lifespan.startup.complete"})
+
+                elif event["type"] == "lifespan.shutdown":
+                    if self.on_shutdown:
+                        await self.on_shutdown()
+                    await send({"type": "lifespan.shutdown.complete"})
+                    break
+        else:
+            await self.app(scope, receive, send)

mrok/http/middlewares.py

@@ -0,0 +1,143 @@
+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
+
+logger = logging.getLogger("mrok.proxy")
+
+ResponseCompleteCallback = Callable[[HTTPResponse], Coroutine[Any, Any, None] | None]
+
+
+class CaptureMiddleware:
+    def __init__(
+        self,
+        app: ASGIApp,
+        on_response_complete: ResponseCompleteCallback,
+    ):
+        self.app = app
+        self._on_response_complete = on_response_complete
+
+    async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend):
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        start_time = time.time()
+        method = scope["method"]
+        path = scope["path"]
+        query_string = scope.get("query_string", b"")
+        req_headers_raw = scope.get("headers", [])
+        req_headers = HTTPHeaders.from_asgi(req_headers_raw)
+
+        state = {}
+
+        req_buf = FixedSizeByteBuffer(MAX_REQUEST_BODY_BYTES)
+        capture_req_body = must_capture_request(method, req_headers)
+
+        request = HTTPRequest(
+            method=method,
+            url=path,
+            headers=req_headers,
+            query_string=query_string,
+            start_time=start_time,
+        )
+
+        # Response capture
+        resp_buf = FixedSizeByteBuffer(MAX_RESPONSE_BODY_BYTES)
+
+        async def receive_wrapper() -> Message:
+            msg = await receive()
+            if capture_req_body and msg["type"] == "http.request":
+                body = msg.get("body", b"")
+                req_buf.write(body)
+            return msg
+
+        async def send_wrapper(msg: Message):
+            if msg["type"] == "http.response.start":
+                state["status"] = msg["status"]
+                resp_headers = HTTPHeaders.from_asgi(msg.get("headers", []))
+                state["resp_headers_raw"] = resp_headers
+
+                state["capture_resp_body"] = must_capture_response(resp_headers)
+
+            if state["capture_resp_body"] and msg["type"] == "http.response.body":
+                body = msg.get("body", b"")
+                resp_buf.write(body)
+
+            await send(msg)
+
+        await self.app(scope, receive_wrapper, send_wrapper)
+
+        # Finalise request
+        request.body = req_buf.getvalue() if capture_req_body else None
+        request.body_truncated = req_buf.overflow if capture_req_body else None
+
+        # Finalise response
+        end_time = time.time()
+        duration = end_time - start_time
+
+        response = HTTPResponse(
+            request=request,
+            status=state["status"] or 0,
+            headers=state["resp_headers_raw"],
+            duration=duration,
+            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")
+
+
+class MetricsMiddleware:
+    def __init__(self, app, metrics: WorkerMetricsCollector):
+        self.app = app
+        self.metrics = metrics
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        start_time = await self.metrics.on_request_start(scope)
+        status_code = 500  # default if app errors early
+
+        async def wrapped_receive():
+            msg = await receive()
+            if msg["type"] == "http.request" and msg.get("body"):
+                await self.metrics.on_request_body(len(msg["body"]))
+            return msg
+
+        async def wrapped_send(msg):
+            nonlocal status_code
+
+            if msg["type"] == "http.response.start":
+                status_code = msg["status"]
+                await self.metrics.on_response_start(status_code)
+
+            elif msg["type"] == "http.response.body":
+                body = msg.get("body", b"")
+                await self.metrics.on_response_chunk(len(body))
+
+            return await send(msg)
+
+        try:
+            await self.app(scope, wrapped_receive, wrapped_send)
+        finally:
+            await self.metrics.on_request_end(start_time, status_code)

mrok/http/types.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, MutableMapping
+from typing import Any, Protocol
+
+from mrok.datastructures import HTTPRequest, 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]]
+RequestCompleteCallback = Callable[[HTTPRequest], Awaitable | None]
+ResponseCompleteCallback = Callable[[HTTPResponse], Awaitable | None]
+
+
+class StreamReaderWrapper(Protocol):
+    async def read(self, n: int = -1) -> bytes: ...
+    async def readexactly(self, n: int) -> bytes: ...
+    async def readline(self) -> bytes: ...
+    def at_eof(self) -> bool: ...
+
+    @property
+    def underlying(self) -> asyncio.StreamReader: ...
+
+
+class StreamWriterWrapper(Protocol):
+    def write(self, data: bytes) -> None: ...
+    async def drain(self) -> None: ...
+    def close(self) -> None: ...
+    async def wait_closed(self) -> None: ...
+
+    @property
+    def transport(self): ...
+
+    @property
+    def underlying(self) -> asyncio.StreamWriter: ...
+
+
+StreamReader = StreamReaderWrapper | asyncio.StreamReader
+StreamWriter = StreamWriterWrapper | asyncio.StreamWriter

mrok/http/utils.py

@@ -0,0 +1,90 @@
+from collections.abc import Mapping
+
+from mrok.http.constants import (
+    BINARY_CONTENT_TYPES,
+    BINARY_PREFIXES,
+    MAX_REQUEST_BODY_BYTES,
+    MAX_RESPONSE_BODY_BYTES,
+    TEXTUAL_CONTENT_TYPES,
+    TEXTUAL_PREFIXES,
+)
+
+
+def is_binary(content_type: str) -> bool:
+    ct = content_type.lower()
+    if ct in BINARY_CONTENT_TYPES:
+        return True
+    if any(ct.startswith(p) for p in BINARY_PREFIXES):
+        return True
+    return False
+
+
+def is_textual(content_type: str) -> bool:
+    ct = content_type.lower()
+    if ct in TEXTUAL_CONTENT_TYPES:
+        return True
+    if any(ct.startswith(p) for p in TEXTUAL_PREFIXES):
+        return True
+    return False
+
+
+def must_capture_request(
+    method: str,
+    headers: Mapping,
+) -> bool:
+    method = method.upper()
+
+    # No body expected
+    if method in ("GET", "HEAD", "OPTIONS", "TRACE"):
+        return False
+
+    content_type = headers.get("content-type", "").lower()
+
+    content_length = None
+    if "content-length" in headers:
+        content_length = int(headers["content-length"])
+
+    if is_binary(content_type):
+        return False
+
+    if content_type.startswith("multipart/form-data"):
+        return False
+
+    if content_length is not None and content_length > MAX_REQUEST_BODY_BYTES:
+        return False
+
+    if is_textual(content_type):
+        return True
+
+    if content_length is None:
+        return True
+
+    return content_length <= MAX_REQUEST_BODY_BYTES
+
+
+def must_capture_response(
+    headers: Mapping,
+) -> bool:
+    content_type = headers.get("content-type", "").lower()
+    disposition = headers.get("content-disposition", "").lower()
+
+    content_length = None
+    if "content-length" in headers:
+        content_length = int(headers["content-length"])
+
+    if "attachment" in disposition:
+        return False
+
+    if is_binary(content_type):
+        return False
+
+    if content_length is not None and content_length > MAX_RESPONSE_BODY_BYTES:
+        return False
+
+    if is_textual(content_type):
+        return True
+
+    if content_length is None:
+        return True
+
+    return content_length <= MAX_RESPONSE_BODY_BYTES

mrok/logging.py

@@ -1,5 +1,9 @@
 import logging.config
 
+from rich.console import Console
+from rich.logging import RichHandler
+from textual_serve.server import LogHighlighter
+
 from mrok.conf import Settings
 
 
@@ -74,3 +78,21 @@ def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict:
 def setup_logging(settings: Settings, cli_mode: bool = False) -> None:
     logging_config = get_logging_config(settings, cli_mode)
     logging.config.dictConfig(logging_config)
+
+
+def setup_inspector_logging(console: Console) -> None:
+    logging.basicConfig(
+        level="WARNING",
+        format="%(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        handlers=[
+            RichHandler(
+                show_path=False,
+                show_time=True,
+                rich_tracebacks=True,
+                tracebacks_show_locals=True,
+                highlighter=LogHighlighter(),
+                console=console,
+            )
+        ],
+    )