diff --git a/aikido_zen/context/__init__.py b/aikido_zen/context/__init__.py index 82de90b61..9f1fe83ad 100644 --- a/aikido_zen/context/__init__.py +++ b/aikido_zen/context/__init__.py @@ -39,6 +39,9 @@ class Context: """ def __init__(self, context_obj=None, body=None, req=None, source=None): + # Initialize default attributes that may not be in context_obj + self.firewall_enforced = False + if context_obj: logger.debug("Creating Context instance based on dict object.") self.__dict__.update(context_obj) @@ -92,6 +95,7 @@ def __reduce__(self): "xml": self.xml, "outgoing_req_redirects": self.outgoing_req_redirects, "executed_middleware": self.executed_middleware, + "firewall_enforced": self.firewall_enforced, "route_params": self.route_params, }, None, diff --git a/aikido_zen/context/init_test.py b/aikido_zen/context/init_test.py index a08e4d502..ee3613972 100644 --- a/aikido_zen/context/init_test.py +++ b/aikido_zen/context/init_test.py @@ -70,6 +70,7 @@ def test_wsgi_context_1(): "xml": {}, "outgoing_req_redirects": [], "executed_middleware": False, + "firewall_enforced": False, "tracked_events": 0, "tracked_events_limit_warning_logged": False, "route_params": [], @@ -104,6 +105,7 @@ def test_wsgi_context_2(): "xml": {}, "outgoing_req_redirects": [], "executed_middleware": False, + "firewall_enforced": False, "tracked_events": 0, "tracked_events_limit_warning_logged": False, "route_params": [], diff --git a/aikido_zen/sources/starlette/starlette_applications.py b/aikido_zen/sources/starlette/starlette_applications.py index 980d85f20..d8a9fb92f 100644 --- a/aikido_zen/sources/starlette/starlette_applications.py +++ b/aikido_zen/sources/starlette/starlette_applications.py @@ -4,6 +4,7 @@ from ..functions.request_handler import request_handler from ...helpers.get_argument import get_argument from ...sinks import on_import, patch_function, before +from ...helpers.logging import logger @before @@ -16,6 +17,34 @@ def _call(func, instance, args, kwargs): new_context.set_as_current_context() request_handler(stage="init") + # Perform firewall enforcement at the application entry point to ensure + # mounted ASGI applications and custom ASGI callables are also protected + try: + pre_response_results = request_handler(stage="pre_response") + if pre_response_results: + response = create_starlette_response(pre_response_results) + if response: + # Block the request by returning the firewall response + return response + # Mark firewall as enforced to avoid duplicate checks in routing wrapper + new_context.firewall_enforced = True + new_context.set_as_current_context() + except Exception as e: + logger.debug("Exception occurred in pre_response stage at __call__: %s", e) + + +def create_starlette_response(pre_response): + """Tries to import PlainTextResponse and generates starlette plain text response""" + text, status_code = pre_response + try: + from starlette.responses import PlainTextResponse + except ImportError: + logger.info( + "Ensure `starlette` install is valid, failed to import starlette.responses" + ) + return None + return PlainTextResponse(text, status_code) + @on_import("starlette.applications", "starlette", version_requirement="0.16.0") def patch(m): diff --git a/aikido_zen/sources/starlette/starlette_routing.py b/aikido_zen/sources/starlette/starlette_routing.py index be88c496d..c953f6f6d 100644 --- a/aikido_zen/sources/starlette/starlette_routing.py +++ b/aikido_zen/sources/starlette/starlette_routing.py @@ -32,12 +32,23 @@ async def aikido_route_func(*args, **kwargs): if not req: return await extract_data_from_request(req) - pre_response_results = request_handler(stage="pre_response") - if pre_response_results: - response = create_starlette_response(pre_response_results) - if response: - # Make sure to not return when an error occurred or there is an invalid response - return response + + # Check if firewall enforcement already happened at __call__ stage + from aikido_zen.context import get_current_context + + context = get_current_context() + if not context or not context.firewall_enforced: + # Firewall not yet enforced, do it now + pre_response_results = request_handler(stage="pre_response") + if pre_response_results: + response = create_starlette_response(pre_response_results) + if response: + # Make sure to not return when an error occurred or there is an invalid response + return response + # Mark as enforced for any subsequent checks + if context: + context.firewall_enforced = True + context.set_as_current_context() except Exception as e: logger.debug("Exception occurred in pre_response stage starlette : %s", e) diff --git a/aikido_zen/test_utils/context_utils.py b/aikido_zen/test_utils/context_utils.py index 5fd8ede43..1e82cdb4e 100644 --- a/aikido_zen/test_utils/context_utils.py +++ b/aikido_zen/test_utils/context_utils.py @@ -57,4 +57,5 @@ def __init__(self): self.user = None self.rate_limit_group = None self.executed_middleware = False + self.firewall_enforced = False self.protection_forced_off = False