fix route in sql comment and also guard against sql injection

Author: kdelee

ansible_base/lib/logging/context.py

@@ -5,7 +5,6 @@
 # Providing a default value is important so that they can be accessed
 # even when the context has not been explicitly set.
 trace_id_var = contextvars.ContextVar('trace_id', default=None)
-route_var = contextvars.ContextVar('route', default=None)
 origin_var = contextvars.ContextVar('origin', default=None)
 
 

ansible_base/lib/middleware/profiling/profile_request.py

@@ -6,12 +6,13 @@
 import time
 import uuid
 from typing import Optional, Union
+from urllib.parse import quote
 
 from django.conf import settings
 from django.db import connection
 from django.utils.translation import gettext_lazy as _
 
-from ansible_base.lib.logging.context import origin_var, route_var, trace_id_var
+from ansible_base.lib.logging.context import origin_var, trace_id_var
 from ansible_base.lib.utils.settings import get_function_from_setting, get_setting
 
 logger = logging.getLogger(__name__)
@@ -85,23 +86,50 @@ def __call__(self, request):
         return response
 
 
+# Define the maximum length for a value in a SQL comment
+SQL_COMMENT_MAX_LENGTH = 256
+
+
+def _sanitize_for_sql_comment(value: str) -> str:
+    """
+    Sanitizes a string for safe inclusion in a SQL comment.
+
+    - URL-encodes the value to handle special characters.
+    - Escapes the '%' character to prevent conflicts with database placeholders.
+    - Truncates the string to a maximum length.
+    """
+    # URL-encode the value
+    quoted_value = quote(str(value))
+    # Escape the '%' character for the database driver
+    sanitized_value = quoted_value.replace('%', '%%')
+    # Truncate to the maximum length
+    return sanitized_value[:SQL_COMMENT_MAX_LENGTH]
+
+
 class SQLQueryMetrics:
-    def __init__(self):
+    def __init__(self, request=None):
+        self.request = request
         self.query_count = 0
         self.query_time = 0.0
 
     def __call__(self, execute, sql, params, many, context):
         # Build the context comment
         context_items = []
+        # trace_id is already validated as a UUID, so it is safe
         if trace_id := trace_id_var.get():
-            context_items.append(f"trace_id={trace_id}")
-        if route := route_var.get():
-            context_items.append(f"route={route}")
+            context_items.append(f"trace_id='{trace_id}'")
+
+        # The route is only available after the URL resolver has run
+        if self.request and getattr(self.request, 'resolver_match', None):
+            if route := self.request.resolver_match.route:
+                context_items.append(f"route='{_sanitize_for_sql_comment(route)}'")
+
         if origin := origin_var.get():
-            context_items.append(f"origin={origin}")
+            context_items.append(f"origin='{_sanitize_for_sql_comment(origin)}'")
 
         if context_items:
-            sql = f"/* {', '.join(context_items)} */ {sql}"
+            comment = f"/* {', '.join(context_items)} */"
+            sql = f"{comment} {sql}"
 
         start_time = time.time()
         try:
@@ -126,7 +154,7 @@ def __call__(self, request):
                 "Please use the ObservabilityMiddleware instead of including profiling middleware individually."
             )
 
-        metrics = SQLQueryMetrics()
+        metrics = SQLQueryMetrics(request)
         with connection.execute_wrapper(metrics):
             response = self.get_response(request)
 

ansible_base/lib/middleware/request_context.py

@@ -1,6 +1,6 @@
 import uuid
 
-from ansible_base.lib.logging.context import origin_var, route_var, trace_id_var
+from ansible_base.lib.logging.context import origin_var, trace_id_var
 
 
 class _TraceContextMiddleware:
@@ -10,13 +10,25 @@ def __init__(self, get_response):
     def __call__(self, request):
         # Set the context for the request and store the tokens
         origin_token = origin_var.set('request')
-        # .get is case-insensitive, but we'll use lowercase for consistency
-        trace_id = request.headers.get('x-request-id', str(uuid.uuid4()))
-        trace_id_token = trace_id_var.set(trace_id)
 
-        route_token = None
-        if request.resolver_match:
-            route_token = route_var.set(request.resolver_match.route)
+        # Get the request ID from the header
+        header_trace_id = request.headers.get('x-request-id')
+        trace_id = None
+
+        if header_trace_id:
+            try:
+                # Validate that the provided header is a valid UUID
+                uuid.UUID(header_trace_id)
+                trace_id = header_trace_id
+            except ValueError:
+                # If it's not a valid UUID, discard it and we'll generate a new one
+                pass
+
+        # If no valid trace_id was found, generate a new one
+        if not trace_id:
+            trace_id = str(uuid.uuid4())
+
+        trace_id_token = trace_id_var.set(trace_id)
 
         try:
             response = self.get_response(request)
@@ -25,7 +37,5 @@ def __call__(self, request):
             # Reset the context variables to their previous state
             origin_var.reset(origin_token)
             trace_id_var.reset(trace_id_token)
-            if route_token:
-                route_var.reset(route_token)
 
         return response

test_app/tests/lib/middleware/test_profiling_middleware.py

@@ -137,14 +137,22 @@ def test_logs_warning_if_context_middleware_is_missing(self, mock_logger):
 
 class SQLQueryMetricsTest(TestCase):
     def test_sql_comment_injection(self):
-        from ansible_base.lib.logging.context import origin_var, route_var, trace_id_var
+        from django.test.client import RequestFactory
 
+        from ansible_base.lib.logging.context import origin_var, trace_id_var
+
+        # 1. Manually set the context, saving the tokens to reset it later.
         trace_id_token = trace_id_var.set("test-trace-id")
-        route_token = route_var.set("test/route")
         origin_token = origin_var.set("test-origin")
 
+        # 2. Create a mock request and manually set the resolver_match
+        factory = RequestFactory()
+        request = factory.get('/test-db/')
+        request.resolver_match = type('ResolverMatch', (), {'route': 'test/route'})
+
         try:
-            metrics = SQLQueryMetrics()
+            # 3. Instantiate our metrics class and call it directly.
+            metrics = SQLQueryMetrics(request)
             original_sql = "SELECT 1"
             modified_sql = ""
 
@@ -155,15 +163,16 @@ def mock_execute(sql, params, many, context):
 
             metrics(mock_execute, original_sql, [], False, {})
 
+            # 4. Assert that the SQL passed to our mock was correctly modified.
             self.assertIn("/*", modified_sql)
-            self.assertIn("trace_id=test-trace-id", modified_sql)
-            self.assertIn("route=test/route", modified_sql)
-            self.assertIn("origin=test-origin", modified_sql)
+            self.assertIn("trace_id='test-trace-id'", modified_sql)
+            self.assertIn("route='test/route'", modified_sql)
+            self.assertIn("origin='test-origin'", modified_sql)
             self.assertIn("*/", modified_sql)
             self.assertIn(original_sql, modified_sql)
         finally:
+            # 5. Reset the context variables to their previous state.
             trace_id_var.reset(trace_id_token)
-            route_var.reset(route_token)
             origin_var.reset(origin_token)
 
 
@@ -208,3 +217,26 @@ def test_observability_middleware_all_headers(self):
                 # 3. From _SQLProfilingMiddleware: Check SQL headers
                 self.assertIn('X-API-Query-Count', response)
                 self.assertIn('X-API-Query-Time', response)
+
+
+class SQLCommentSanitizationTest(TestCase):
+    def test_sanitization_escapes_disallowed_chars(self):
+        from ansible_base.lib.middleware.profiling.profile_request import _sanitize_for_sql_comment
+
+        malicious_string = "*/; DROP TABLE users; --"
+        sanitized = _sanitize_for_sql_comment(malicious_string)
+        self.assertEqual(sanitized, "%%2A/%%3B%%20DROP%%20TABLE%%20users%%3B%%20--")
+
+    def test_sanitization_allows_safe_chars(self):
+        from ansible_base.lib.middleware.profiling.profile_request import _sanitize_for_sql_comment
+
+        safe_string = "a-b_c.d/e123"
+        sanitized = _sanitize_for_sql_comment(safe_string)
+        self.assertEqual(sanitized, "a-b_c.d/e123")
+
+    def test_sanitization_truncates_long_strings(self):
+        from ansible_base.lib.middleware.profiling.profile_request import SQL_COMMENT_MAX_LENGTH, _sanitize_for_sql_comment
+
+        long_string = "a" * (SQL_COMMENT_MAX_LENGTH + 100)
+        sanitized = _sanitize_for_sql_comment(long_string)
+        self.assertEqual(len(sanitized), SQL_COMMENT_MAX_LENGTH)

test_app/tests/lib/middleware/test_request_context.py

@@ -86,3 +86,21 @@ def test_context_does_not_bleed_between_requests(self):
             UUID(trace_id_2, version=4)
         except ValueError:
             self.fail("trace_id_2 is not a valid UUID4")
+
+    def test_discards_invalid_uuid_in_header(self):
+        """
+        Test that the middleware discards an invalid UUID in the X-Request-ID
+        header and generates a new, valid one.
+        """
+        malicious_id = "not-a-uuid' --"
+        response = self.client.get('/context/', HTTP_X_REQUEST_ID=malicious_id)
+        self.assertEqual(response.status_code, 200)
+
+        # The trace_id in the response should be a new, valid UUID, not the malicious one.
+        new_trace_id = response.headers.get("X-Request-ID")
+        self.assertIsNotNone(new_trace_id)
+        self.assertNotEqual(new_trace_id, malicious_id)
+        try:
+            UUID(new_trace_id, version=4)
+        except ValueError:
+            self.fail("The new trace_id is not a valid UUID4")