diff --git a/mrok/frontend/app.py b/mrok/frontend/app.py index 33d5cdb..278f1c4 100644 --- a/mrok/frontend/app.py +++ b/mrok/frontend/app.py @@ -6,7 +6,7 @@ from jinja2 import Environment, FileSystemLoader, select_autoescape from mrok.conf import get_settings -from mrok.frontend.utils import get_target_name, parse_accept_header +from mrok.frontend.utils import get_headers, get_target_name, parse_accept_header from mrok.proxy.app import ProxyAppBase from mrok.proxy.backend import AIOZitiNetworkBackend from mrok.proxy.exceptions import InvalidTargetError @@ -54,13 +54,11 @@ def setup_connection_pool( ) def get_upstream_base_url(self, scope: Scope) -> str: - target = get_target_name( - {k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {})} - ) + target = get_target_name(get_headers(scope)) if not target: raise InvalidTargetError() - return f"http://{target.lower()}" + return f"http://{target}" async def send_error_response( self, @@ -70,9 +68,7 @@ async def send_error_response( body: str, headers: list[tuple[bytes, bytes]] | None = None, ): - request_headers = { - k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {}) - } + request_headers = get_headers(scope) accept_header = request_headers.get("accept") if not (accept_header and str(http_status) in self._templates_by_error): return await super().send_error_response(scope, send, http_status, body) @@ -113,7 +109,7 @@ async def _render_error_template( return await template.render_async(context) def _extract_request_context(self, scope: Scope) -> dict[str, Any]: - headers = {k.decode("latin-1"): v.decode("latin-1") for k, v in scope.get("headers", [])} + headers = get_headers(scope) return { "method": scope.get("method"), diff --git a/mrok/frontend/logging.py b/mrok/frontend/logging.py new file mode 100644 index 0000000..b18e963 --- /dev/null +++ b/mrok/frontend/logging.py @@ -0,0 +1,17 @@ +import contextvars +import logging + +request_target: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "request_target", default=None +) + + +def setup_log_record_factory() -> None: + default_factory = logging.getLogRecordFactory() + + def record_factory(*args, **kwargs): + record = default_factory(*args, **kwargs) + record.request_target = request_target.get() or "-" + return record + + logging.setLogRecordFactory(record_factory) diff --git a/mrok/frontend/main.py b/mrok/frontend/main.py index a0ac8fa..de8cce7 100644 --- a/mrok/frontend/main.py +++ b/mrok/frontend/main.py @@ -7,7 +7,12 @@ from mrok.authentication import HTTPAuthManager from mrok.conf import get_settings from mrok.frontend.app import FrontendProxyApp -from mrok.frontend.middleware import ASGIAuthenticationMiddleware, HealthCheckMiddleware +from mrok.frontend.logging import setup_log_record_factory +from mrok.frontend.middleware import ( + ASGIAuthenticationMiddleware, + HealthCheckMiddleware, + TargetContextMiddleware, +) from mrok.logging import get_logging_config from mrok.proxy.asgi import ASGIAppWrapper @@ -45,6 +50,7 @@ def load(self): ASGIAuthenticationMiddleware, auth_manager=HTTPAuthManager(settings.frontend.auth), ) + app.add_middleware(TargetContextMiddleware) return app @@ -58,6 +64,8 @@ def run( max_keepalive_connections: int | None, keepalive_expiry: float | None, ): + setup_log_record_factory() + options = { "bind": f"{host}:{port}", "workers": workers, diff --git a/mrok/frontend/middleware.py b/mrok/frontend/middleware.py index ddb86ca..f47efbb 100644 --- a/mrok/frontend/middleware.py +++ b/mrok/frontend/middleware.py @@ -1,19 +1,33 @@ import json from mrok.authentication import HTTPAuthManager -from mrok.frontend.utils import get_target_name +from mrok.frontend.logging import request_target +from mrok.frontend.utils import get_headers, get_target_name from mrok.types.proxy import ASGIApp, ASGIReceive, ASGISend, Scope +class TargetContextMiddleware: + """Expose the target (extension/instance) of the current request to the log record factory.""" + + def __init__(self, app: ASGIApp): + self.app = app + + async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend): + target = get_target_name(get_headers(scope)) + token = request_target.set(target) + try: + await self.app(scope, receive, send) + finally: + request_target.reset(token) + + class HealthCheckMiddleware: def __init__(self, app: ASGIApp): self.app = app async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend): if scope["type"] == "http" and scope["path"] == "/healthcheck": - target = get_target_name( - {k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {})} - ) + target = get_target_name(get_headers(scope)) if not target: await send( diff --git a/mrok/frontend/utils.py b/mrok/frontend/utils.py index c7ab49b..a8b0064 100644 --- a/mrok/frontend/utils.py +++ b/mrok/frontend/utils.py @@ -1,6 +1,7 @@ import re from mrok.conf import get_settings +from mrok.types.proxy import Scope def parse_accept_header(accept: str | None) -> list[str]: @@ -79,5 +80,9 @@ def get_target_name(headers: dict[str, str]) -> str | None: re.fullmatch(settings.identifiers.extension.regex, target) or re.fullmatch(settings.identifiers.instance.regex, target) ): - return target + return target.lower() return None + + +def get_headers(scope: Scope) -> dict[str, str]: + return {k.decode("latin-1"): v.decode("latin-1") for k, v in scope.get("headers", [])} diff --git a/mrok/logging.py b/mrok/logging.py index b318aa6..1edd41f 100644 --- a/mrok/logging.py +++ b/mrok/logging.py @@ -3,6 +3,8 @@ from mrok.conf import Settings +ACCESS_LOG_FORMAT = "{asctime} {name} {levelname} (pid: {process}) [{request_target}] {message}" + class HealthCheckFilter(logging.Filter): def filter(self, record): @@ -12,6 +14,7 @@ def filter(self, record): def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict: log_level = "DEBUG" if settings.logging.debug else "INFO" handler = "rich" if settings.logging.rich else "console" + access_handler = "rich" if settings.logging.rich else "access" if cli_mode: mrok_handler = "cli" @@ -31,11 +34,16 @@ def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict: "style": "{", }, "plain": {"format": "%(message)s"}, + "access": { + "format": ACCESS_LOG_FORMAT, + "style": "{", + "defaults": {"request_target": "-"}, + }, }, "filters": { "healthcheck_filter": { "()": HealthCheckFilter, - } + }, }, "handlers": { "console": { @@ -55,6 +63,11 @@ def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict: "log_time_format": lambda x: x.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3], "rich_tracebacks": True, }, + "access": { + "class": "logging.StreamHandler", + "formatter": "access", + "stream": "ext://sys.stderr", + }, }, "root": { "handlers": [handler], @@ -62,7 +75,7 @@ def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict: }, "loggers": { "gunicorn.access": { - "handlers": [handler], + "handlers": [access_handler], "level": log_level, "propagate": False, "filters": ["healthcheck_filter"], @@ -73,7 +86,7 @@ def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict: "propagate": False, }, "uvicorn.access": { - "handlers": [handler], + "handlers": [access_handler], "level": log_level, "propagate": False, "filters": ["healthcheck_filter"], diff --git a/tests/frontend/test_logging.py b/tests/frontend/test_logging.py new file mode 100644 index 0000000..f7faece --- /dev/null +++ b/tests/frontend/test_logging.py @@ -0,0 +1,55 @@ +import logging + +import pytest +from pytest_mock import MockerFixture + +from mrok.frontend.logging import request_target, setup_log_record_factory +from mrok.logging import ACCESS_LOG_FORMAT, get_logging_config +from tests.types import SettingsFactory + + +def _install_factory(mocker: MockerFixture): + m_set_factory = mocker.patch("logging.setLogRecordFactory") + setup_log_record_factory() + return m_set_factory.call_args.args[0] + + +def test_log_record_factory_sets_request_target(mocker: MockerFixture): + factory = _install_factory(mocker) + + token = request_target.set("ext-1234-5678") + try: + record = factory("uvicorn.access", logging.INFO, __file__, 1, "msg", (), None) + finally: + request_target.reset(token) + + assert isinstance(record, logging.LogRecord) + assert record.request_target == "ext-1234-5678" + + +def test_log_record_factory_without_request_target(mocker: MockerFixture): + factory = _install_factory(mocker) + + record = factory("uvicorn.access", logging.INFO, __file__, 1, "msg", (), None) + + assert record.request_target == "-" + + +@pytest.mark.parametrize( + ("rich", "expected_handler"), + [(False, "access"), (True, "rich")], +) +def test_get_logging_config_access_handler( + settings_factory: SettingsFactory, rich: bool, expected_handler: str +): + settings = settings_factory(logging={"debug": True, "rich": rich}) + + config = get_logging_config(settings) + + assert config["loggers"]["gunicorn.access"]["handlers"] == [expected_handler] + assert config["loggers"]["uvicorn.access"]["handlers"] == [expected_handler] + assert config["formatters"]["access"] == { + "format": ACCESS_LOG_FORMAT, + "style": "{", + "defaults": {"request_target": "-"}, + } diff --git a/tests/frontend/test_main.py b/tests/frontend/test_main.py index 914927a..8d5a238 100644 --- a/tests/frontend/test_main.py +++ b/tests/frontend/test_main.py @@ -8,6 +8,7 @@ def test_run(mocker: MockerFixture): "mrok.frontend.main.get_logging_config", return_value={"logging": "config"}, ) + m_setup_factory = mocker.patch("mrok.frontend.main.setup_log_record_factory") mock_settings = mocker.MagicMock() mock_settings.controller.auth = {"backends": []} @@ -32,3 +33,4 @@ def test_run(mocker: MockerFixture): assert m_standalone_app.mock_calls[0].args[0]["mrok"]["keepalive_expiry"] == 99.5 m_app.run.assert_called_once() + m_setup_factory.assert_called_once() diff --git a/tests/frontend/test_middleware.py b/tests/frontend/test_middleware.py index e908c79..62f0bea 100644 --- a/tests/frontend/test_middleware.py +++ b/tests/frontend/test_middleware.py @@ -3,7 +3,8 @@ import pytest from pytest_mock import MockerFixture -from mrok.frontend.middleware import HealthCheckMiddleware +from mrok.frontend.logging import request_target +from mrok.frontend.middleware import HealthCheckMiddleware, TargetContextMiddleware @pytest.mark.asyncio @@ -58,3 +59,40 @@ async def test_healthcheck_passtrhough( ) m_app.assert_awaited_once_with(scope, m_receive, m_send) + + +@pytest.mark.asyncio +async def test_target_middleware_sets_target(mocker: MockerFixture): + mocker.patch("mrok.frontend.utils.get_frontend_domain", return_value=".extdomain") + seen: list[str | None] = [] + m_app = mocker.AsyncMock(side_effect=lambda *args: seen.append(request_target.get())) + m_receive = mocker.AsyncMock() + m_send = mocker.AsyncMock() + middleware = TargetContextMiddleware(m_app) + scope = { + "type": "http", + "path": "/", + "headers": [(b"host", b"EXT-1234-5678.extdomain:8000")], + } + + await middleware(scope, m_receive, m_send) + + m_app.assert_awaited_once_with(scope, m_receive, m_send) + assert seen == ["ext-1234-5678"] + assert request_target.get() is None + + +@pytest.mark.asyncio +async def test_target_middleware_without_target(mocker: MockerFixture): + mocker.patch("mrok.frontend.utils.get_frontend_domain", return_value=".extdomain") + seen: list[str | None] = [] + m_app = mocker.AsyncMock(side_effect=lambda *args: seen.append(request_target.get())) + m_receive = mocker.AsyncMock() + m_send = mocker.AsyncMock() + middleware = TargetContextMiddleware(m_app) + scope = {"type": "http", "path": "/", "headers": [(b"host", b"127.0.0.1")]} + + await middleware(scope, m_receive, m_send) + + m_app.assert_awaited_once_with(scope, m_receive, m_send) + assert seen == [None]