diff --git a/guard/__init__.py b/guard/__init__.py index 2511155f..05217854 100644 --- a/guard/__init__.py +++ b/guard/__init__.py @@ -31,6 +31,7 @@ sus_patterns_handler as sus_patterns_handler, ) +from guard.asgi import PureASGISecurityMiddleware as PureASGISecurityMiddleware from guard.middleware import SecurityMiddleware as SecurityMiddleware from guard.websocket import ( WS_CLOSE_CLIENT_ADDRESS_UNKNOWN as WS_CLOSE_CLIENT_ADDRESS_UNKNOWN, @@ -52,6 +53,7 @@ __all__ = [ "__version__", + "PureASGISecurityMiddleware", "SecurityMiddleware", "guard_websocket", "make_guard_websocket", diff --git a/guard/asgi.py b/guard/asgi.py new file mode 100644 index 00000000..ea4a77ce --- /dev/null +++ b/guard/asgi.py @@ -0,0 +1,289 @@ +"""Pure-ASGI variant of the security middleware. + +`PureASGISecurityMiddleware` runs the same request-screening pipeline as +`SecurityMiddleware` (initialization, CORS preflight, path exclusions, IP +extraction, ban lists, rate limiting, user-agent blocking, penetration +detection, behavioral usage rules) without inheriting from Starlette's +`BaseHTTPMiddleware`. + +Why it exists: `BaseHTTPMiddleware` wraps the downstream app in an anyio +task group whose cancel scope cancels in-flight downstream awaits (database +calls, long LLM streams) when a client disconnects mid-request. Framework +projects that ban `BaseHTTPMiddleware` for exactly this reason (open-webui +is one) can mount this variant instead and keep the same guard behavior. + +Response-side differences from `SecurityMiddleware`: + +- Security headers and CORS response headers are injected directly on the + `http.response.start` message; responses stream through unbuffered. +- Decorator-driven response rules (return patterns) and + `custom_response_modifier` need a materialized `Response` object and are + only applied by the `BaseHTTPMiddleware`-based middleware. + +The downstream app receives a replaying receive channel, so both guard and +the application can read the request body. +""" + +import time +from collections.abc import Awaitable, Callable +from typing import Any + +from fastapi import Response +from guard_core.exceptions import GuardRedisError +from guard_core.handlers.security_headers_handler import security_headers_manager +from guard_core.models import SecurityConfig +from guard_core.utils import extract_client_ip +from starlette.requests import Request +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from guard.adapters import StarletteGuardRequest, wrap_call_next +from guard.middleware import SecurityMiddleware + +_BOOL_TRUE = frozenset({"1", "true", "yes", "on"}) + + +class _BodyCachingReceive: + """Receive wrapper that caches request-body chunks for replay. + + Guard needs the body for penetration and custom-request checks; the + downstream application still expects to read it afterwards. Chunks are + buffered on the first pass and replayed to the app in order. + """ + + def __init__(self, receive: Receive) -> None: + self._receive = receive + self._cached: list[Message] = [] + self._body_done = False + + async def __call__(self) -> Message: + message = await self._receive() + if message["type"] == "http.request": + self._cached.append(message) + if not message.get("more_body", False): + self._body_done = True + return message + + def replay(self) -> Receive: + """Return a receive callable that yields the cached body, then + continues streaming from the original receive (disconnect waits).""" + messages = list(self._cached) + index = {"i": 0} + inner = self._receive + + async def replay_receive() -> Message: + i = index["i"] + if i < len(messages): + index["i"] = i + 1 + return messages[i] + return await inner() + + return replay_receive + + +class _HeaderInjectingSend: + """Send wrapper that merges security and CORS headers into the response + start message, leaving the body stream untouched.""" + + def __init__( + self, + send: Send, + *, + engine: SecurityMiddleware, + request: Request, + ) -> None: + self._send = send + self._engine = engine + self._request = request + self.status_code: int | None = None + + async def __call__(self, message: Message) -> None: + if message["type"] == "http.response.start": + self.status_code = message["status"] + headers: dict[bytes, bytes] = { + key: value for key, value in message["headers"] + } + extra = await security_headers_manager.get_headers( + request_path=self._request.url.path, config=self._engine.config + ) + for key, value in extra.items(): + headers.setdefault(key.encode("latin-1"), value.encode("latin-1")) + if self._engine._cors_handler is not None: + cors = self._engine._cors_handler.build_response_headers( + self._request.headers + ) + for key, value in cors.items(): + headers.setdefault(key.encode("latin-1"), value.encode("latin-1")) + message = {**message, "headers": list(headers.items())} + await self._send(message) + + +class PureASGISecurityMiddleware: + """Pure-ASGI security middleware sharing the pipeline of + `SecurityMiddleware`. + + Mount it like the classic middleware:: + + app.add_middleware(PureASGISecurityMiddleware, config=config) + + The class is disabled-by-default at the embedding layer the same way + the classic one is: construct it only when the operator enabled the + integration. WebSocket and lifespan scopes pass through untouched + (wrap websocket endpoints with `guard.websocket.make_guard_websocket` + for socket-level checks). + """ + + def __init__(self, app: ASGIApp, *, config: SecurityConfig) -> None: + self.app = app + self.config = config + # The classic middleware instance is used as the engine: its + # constructor builds the pipeline, handlers, rate limiter, event bus + # and CORS handler, and its screening helpers operate on plain + # Request objects, so none of that machinery is duplicated here. + self._engine = SecurityMiddleware(_NullApp(), config=config) + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + caching_receive = _BodyCachingReceive(receive) + request = Request(scope, receive=caching_receive) + wrapped_send = _HeaderInjectingSend(send, engine=self._engine, request=request) + start_time = time.time() + await self._screen(request, caching_receive, wrapped_send, scope, start_time) + + async def _screen( # noqa: PLR0915 + self, + request: Request, + caching_receive: _BodyCachingReceive, + wrapped_send: _HeaderInjectingSend, + scope: Scope, + start_time: float, + ) -> None: + engine = self._engine + try: + await engine._ensure_initialized(request) + except GuardRedisError as e: + engine.logger.error("Redis unavailable during initialization: %s", e) + response = await engine._redis_unavailable_response() + await self._send_response(response, request, wrapped_send._send) + return + + guard_request = StarletteGuardRequest(request) + engine._populate_guard_state(guard_request, request) + + preflight = await engine._handle_preflight(request, guard_request) + if preflight is not None: + await self._send_response(preflight, request, wrapped_send._send) + return + + call_next = self._build_call_next(scope, caching_receive, wrapped_send) + passthrough = await engine._handle_passthrough( + request, guard_request, wrap_call_next(call_next, request) + ) + if passthrough is not None: + await self._send_response(passthrough, request, wrapped_send._send) + return + + client_ip = await extract_client_ip( + guard_request, engine.config, engine.agent_handler + ) + route_config = engine.route_resolver.get_route_config(guard_request) + + bypass = await engine._handle_security_bypass( + request, + guard_request, + wrap_call_next(call_next, request), + route_config, + ) + if bypass is not None: + await self._send_response(bypass, request, wrapped_send._send) + return + + blocked = await engine._handle_pipeline_block(request, guard_request) + if blocked is not None: + await self._send_response(blocked, request, wrapped_send._send) + return + + if route_config and route_config.behavior_rules and client_ip: + await engine.behavioral_processor.process_usage_rules( + guard_request, client_ip, route_config + ) + + await self._call_downstream( + scope, caching_receive.replay(), wrapped_send, start_time + ) + + def _build_call_next( + self, + scope: Scope, + caching_receive: _BodyCachingReceive, + wrapped_send: _HeaderInjectingSend, + ) -> Callable[[Request], Awaitable[Response]]: + """Buffered downstream call for the bypass/passthrough helpers, which + expect a materialized `Response`. Only bypassed routes take this + path; the regular flow streams instead.""" + + async def call_next(request: Request) -> Response: + captured: dict[str, Any] = {"status": 500, "headers": [], "body": b""} + + async def capturing_send(message: Message) -> None: + if message["type"] == "http.response.start": + captured["status"] = message["status"] + captured["headers"] = list(message["headers"]) + elif message["type"] == "http.response.body": # pragma: no branch + captured["body"] += message.get("body", b"") + + await self._call_downstream( + scope, caching_receive.replay(), capturing_send, time.time() + ) + return Response( + status_code=captured["status"], + headers=dict( + (key.decode("latin-1"), value.decode("latin-1")) + for key, value in captured["headers"] + ), + content=captured["body"], + ) + + return call_next + + async def _call_downstream( + self, + scope: Scope, + receive: Receive, + send: Send | _HeaderInjectingSend, + start_time: float, + ) -> None: + await self.app(scope, receive, send) + self._engine.logger.debug( + "Request processed in %.4fs", time.time() - start_time + ) + + async def _send_response( + self, response: Response, request: Request, send: Send + ) -> None: + """Send a guard-generated blocking response, with CORS headers + applied the same way the classic dispatch applies them.""" + if self._engine._cors_handler is not None: + cors_headers = self._engine._cors_handler.build_response_headers( + request.headers + ) + for key, value in cors_headers.items(): + response.headers.setdefault(key, value) + body = getattr(response, "body", b"") or b"" + await send( + { + "type": "http.response.start", + "status": response.status_code, + "headers": list(response.raw_headers), + } + ) + await send({"type": "http.response.body", "body": body}) + + +class _NullApp: + """Placeholder downstream for the engine instance; never called because + only the screening helpers of the engine are used.""" + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + raise RuntimeError("PureASGISecurityMiddleware engine must not be called") diff --git a/tests/test_asgi_pure.py b/tests/test_asgi_pure.py new file mode 100644 index 00000000..7509c65f --- /dev/null +++ b/tests/test_asgi_pure.py @@ -0,0 +1,439 @@ +"""Tests for the pure-ASGI security middleware variant.""" + +import asyncio +import itertools +from typing import Any + +from fastapi import FastAPI, Request +from httpx import ASGITransport, AsyncClient + +from guard import PureASGISecurityMiddleware, SecurityConfig + +# Rate limiting and ban state in guard-core are process-wide, so every test +# drives its own TEST-NET client IP to stay hermetic. +_IP = itertools.count(1) + + +def _unique_ip() -> str: + n = next(_IP) + return f"198.51.{n // 250}.{(n % 250) + 1}" + + +def _build_app(config: SecurityConfig, *, with_echo: bool = False) -> FastAPI: + app = FastAPI() + app.add_middleware(PureASGISecurityMiddleware, config=config) + + @app.get("/ping") + async def ping(): + return {"ok": True} + + @app.get("/health") + async def health(): + return {"ok": "health"} + + if with_echo: + + @app.post("/echo") + async def echo(request: Request): + return {"received": await request.json()} + + return app + + +def _run(scenario, client_ip, config, with_echo=False): + async def runner(): + transport = ASGITransport( + app=_build_app(config, with_echo=with_echo), client=(client_ip, 50000) + ) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return await scenario(client) + + return asyncio.run(runner()) + + +def _base_config(**overrides): + kwargs = { + # In-memory state unless a Redis URL is configured: never implicitly + # depend on a Redis server being reachable at localhost. + "enable_redis": False, + "exclude_paths": ["/health"], + } + kwargs.update(overrides) + return SecurityConfig(**kwargs) + + +def test_blocked_ip_is_rejected(): + blocked = _unique_ip() + config = _base_config(blacklist=(blocked,)) + + async def scenario(client): + response = await client.get("/ping") + assert response.status_code == 403 + + _run(scenario, blocked, config) + + +def test_rate_limit_returns_429(): + client_ip = _unique_ip() + config = _base_config(rate_limit=2, rate_limit_window=60) + + async def scenario(client): + assert (await client.get("/ping")).status_code == 200 + assert (await client.get("/ping")).status_code == 200 + assert (await client.get("/ping")).status_code == 429 + + _run(scenario, client_ip, config) + + +def test_passive_mode_never_blocks(): + blocked = _unique_ip() + config = _base_config(passive_mode=True, blacklist=(blocked,)) + + async def scenario(client): + response = await client.get("/ping") + assert response.status_code == 200 + + _run(scenario, blocked, config) + + +def test_ip_lists_enforced_on_excluded_paths(): + blocked = _unique_ip() + config = _base_config(blacklist=(blocked,)) + + async def scenario(client): + response = await client.get("/health") + assert response.status_code == 403 + + _run(scenario, blocked, config) + + +def test_rate_limiting_applies_on_excluded_paths(): + client_ip = _unique_ip() + config = _base_config(rate_limit=1, rate_limit_window=60) + + async def scenario(client): + codes = [(await client.get("/health")).status_code for _ in range(3)] + assert codes == [200, 429, 429] + + _run(scenario, client_ip, config) + + +def test_security_headers_injected_on_response(): + client_ip = _unique_ip() + config = _base_config( + security_headers={ + "enabled": True, + "frame_options": "SAMEORIGIN", + "content_type_options": "nosniff", + } + ) + + async def scenario(client): + response = await client.get("/ping") + assert response.status_code == 200 + assert response.headers["x-frame-options"] == "SAMEORIGIN" + assert response.headers["x-content-type-options"] == "nosniff" + + _run(scenario, client_ip, config) + + +def test_body_reaches_downstream_after_guard_reads_it(): + client_ip = _unique_ip() + config = _base_config() + + async def scenario(client): + response = await client.post( + "/echo", + json={"message": "hello guard"}, + headers={"content-type": "application/json"}, + ) + assert response.status_code == 200 + assert response.json() == {"received": {"message": "hello guard"}} + + # inline Body import keeps the echo signature local to this test + _run(scenario, client_ip, config, with_echo=True) + + +def test_bypassed_route_skips_checks(): + from guard import SecurityDecorator + + blocked = "203.0.113.9" + config = _base_config(blacklist=(blocked,)) + app = FastAPI() + app.add_middleware(PureASGISecurityMiddleware, config=config) + decorator = SecurityDecorator(config) + app.state.guard_decorator = decorator + + @decorator.bypass(["all"]) + @app.get("/bypassed") + async def bypassed(): + return {"ok": True} + + async def runner(): + transport = ASGITransport(app=app, client=(blocked, 50000)) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return (await client.get("/bypassed")).status_code + + assert asyncio.run(runner()) == 200 + + +def test_unknown_client_fails_secure(): + config = _base_config() + app = FastAPI() + app.add_middleware(PureASGISecurityMiddleware, config=config) + + captured: dict[str, Any] = {} + + async def send(message): + captured[message["type"]] = message + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "GET", + "path": "/ping", + "raw_path": b"/ping", + "query_string": b"", + "root_path": "", + "headers": [(b"host", b"testserver")], + # no "client" key: guard cannot determine the source address + } + asyncio.run(app(scope, receive, send)) + assert captured["http.response.start"]["status"] == 403 + + +def test_null_app_raises_if_called(): + async def exercise(): + app = FastAPI() + middleware = PureASGISecurityMiddleware(app, config=_base_config()) + await middleware._engine.app({"type": "http"}, lambda: None, lambda _: None) + + try: + asyncio.run(exercise()) + raised = False + except RuntimeError: + raised = True + assert raised + + +def test_redis_unavailable_returns_503(): + from guard_core.exceptions import GuardRedisError + + client_ip = _unique_ip() + config = _base_config() + app = FastAPI() + middleware = PureASGISecurityMiddleware(app, config=config) + + async def raise_redis_error(request): + raise GuardRedisError(503, "connection refused") + + middleware._engine._ensure_initialized = raise_redis_error + + async def runner(): + transport = ASGITransport(app=middleware, client=(client_ip, 50000)) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return (await client.get("/ping")).status_code + + assert asyncio.run(runner()) == 503 + + +def test_cors_preflight_and_response_headers(): + client_ip = _unique_ip() + config = _base_config( + enable_cors=True, + cors_allow_origins=["http://good.origin"], + cors_allow_methods=["GET"], + cors_allow_headers=["*"], + ) + app = FastAPI() + app.add_middleware(PureASGISecurityMiddleware, config=config) + + @app.get("/ping") + async def ping(): + return {"ok": True} + + async def runner(): + transport = ASGITransport(app=app, client=(client_ip, 50000)) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + preflight = await client.options( + "/ping", + headers={ + "Origin": "http://good.origin", + "Access-Control-Request-Method": "GET", + }, + ) + response = await client.get( + "/ping", headers={"Origin": "http://good.origin"} + ) + return preflight, response + + preflight, response = asyncio.run(runner()) + assert preflight.status_code == 200 + assert response.status_code == 200 + assert response.headers.get("access-control-allow-origin") == ("http://good.origin") + + +def test_blocked_response_carries_cors_headers(): + blocked = _unique_ip() + config = _base_config( + enable_cors=True, + cors_allow_origins=["http://good.origin"], + blacklist=(blocked,), + ) + app = FastAPI() + app.add_middleware(PureASGISecurityMiddleware, config=config) + + @app.get("/ping") + async def ping(): + return {"ok": True} + + async def runner(): + transport = ASGITransport(app=app, client=(blocked, 50000)) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return await client.get("/ping", headers={"Origin": "http://good.origin"}) + + response = asyncio.run(runner()) + assert response.status_code == 403 + assert response.headers.get("access-control-allow-origin") == ("http://good.origin") + + +def test_behavioral_usage_rules_processed(): + from guard import SecurityDecorator + + client_ip = _unique_ip() + config = _base_config() + app = FastAPI() + app.add_middleware(PureASGISecurityMiddleware, config=config) + decorator = SecurityDecorator(config) + app.state.guard_decorator = decorator + + @decorator.usage_monitor(max_calls=100, window=3600, action="ban") + @app.get("/monitored") + async def monitored(): + return {"ok": True} + + async def runner(): + transport = ASGITransport(app=app, client=(client_ip, 50000)) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return (await client.get("/monitored")).status_code + + assert asyncio.run(runner()) == 200 + + +def test_body_caching_receive_replays_then_yields_inner(): + from guard.asgi import _BodyCachingReceive + + served = [ + {"type": "http.request", "body": b"chunk-one", "more_body": True}, + {"type": "http.request", "body": b"chunk-two", "more_body": False}, + ] + inner_calls: list[dict] = [] + + async def inner(): + message = served.pop(0) if served else {"type": "http.disconnect"} + inner_calls.append(message) + return message + + receive = _BodyCachingReceive(inner) + + async def drive(): + first = await receive() + second = await receive() + replayed = receive.replay() + replayed_messages = [await replayed() for _ in range(2)] + third = await replayed() + return first, second, replayed_messages, third + + first, second, replayed_messages, third = asyncio.run(drive()) + assert first["body"] == b"chunk-one" + assert second["body"] == b"chunk-two" + assert [m["body"] for m in replayed_messages] == [b"chunk-one", b"chunk-two"] + assert third["type"] == "http.disconnect" + # inner served the two body chunks, then the post-body disconnect wait + assert [m["body"] for m in inner_calls[:2]] == [b"chunk-one", b"chunk-two"] + assert inner_calls[-1]["type"] == "http.disconnect" + + +def test_chunked_body_is_cached_and_replayed(): + client_ip = _unique_ip() + config = _base_config(enable_penetration_detection=False) + app = FastAPI() + app.add_middleware(PureASGISecurityMiddleware, config=config) + bodies: list[bytes] = [] + + @app.post("/upload") + async def upload(request: Request): + async for chunk in request.stream(): + bodies.append(chunk) + return {"size": len(b"".join(bodies))} + + served = [ + {"type": "http.request", "body": b"part-one-", "more_body": True}, + {"type": "http.request", "body": b"part-two", "more_body": False}, + ] + + async def receive(): + return served.pop(0) + + captured: dict[str, Any] = {} + + async def send(message): + if message["type"] == "http.response.start": + captured["status"] = message["status"] + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "path": "/upload", + "raw_path": b"/upload", + "query_string": b"", + "root_path": "", + "headers": [(b"host", b"testserver"), (b"content-type", b"text/plain")], + "client": (client_ip, 50000), + } + asyncio.run(app(scope, receive, send)) + assert captured["status"] == 200 + assert b"".join(bodies) == b"part-one-part-two" + + +def test_receive_passes_disconnect_through_uncached(): + from guard.asgi import _BodyCachingReceive + + disconnect = {"type": "http.disconnect"} + + async def inner(): + return disconnect + + receive = _BodyCachingReceive(inner) + message = asyncio.run(receive()) + assert message is disconnect + + +def test_non_http_scopes_pass_through(): + calls: list[str] = [] + + class ProbeApp: + async def __call__(self, scope, receive, send): + calls.append(scope["type"]) + + middleware = PureASGISecurityMiddleware(ProbeApp(), config=_base_config()) + asyncio.run(middleware.__call__({"type": "lifespan"}, lambda: None, lambda _: None)) + assert calls == ["lifespan"] diff --git a/tests/test_reexports.py b/tests/test_reexports.py index 43eb0ab8..4f7a184e 100644 --- a/tests/test_reexports.py +++ b/tests/test_reexports.py @@ -15,6 +15,7 @@ def test_all_derives_from_guard_core_exports() -> None: assert set(guard.__all__) == { "__version__", + "PureASGISecurityMiddleware", "SecurityMiddleware", "guard_websocket", "make_guard_websocket",