[Aikido] Enforce firewall decisions at Starlette app entry point to prevent bypass

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,

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": [],

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):

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)
 

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