Commit

jy-tan · jy-tan · commit a1d3061d3797 · 2026-02-06T15:53:26.000-08:00
diff --git a/drift/instrumentation/django/html_utils.py b/drift/instrumentation/django/html_utils.py
@@ -62,6 +62,17 @@ def normalize_csrf_in_body(body: bytes | None) -> bytes | None:
             flags=re.IGNORECASE,
         )
 
+        # Pattern 3: JavaScript CSRF token assignment (e.g., DRF Swagger UI)
+        # request.headers["X-CSRFTOKEN"] = "ABC123...";
+        # Also handles single quotes and X-CSRFToken variants
+        csrf_js_pattern = r'(headers\[["\']X-CSRFTOKEN["\']\]\s*=\s*["\'])[^"\']+(["\'])'
+        body_str = re.sub(
+            csrf_js_pattern,
+            rf"\g<1>{CSRF_PLACEHOLDER}\2",
+            body_str,
+            flags=re.IGNORECASE,
+        )
+
         return body_str.encode("utf-8")
 
     except Exception as e:
diff --git a/tests/unit/test_html_utils.py b/tests/unit/test_html_utils.py
@@ -172,6 +172,68 @@ def test_handles_non_utf8_gracefully(self):
         result = normalize_csrf_in_body(html)
         assert result == html
 
+    def test_normalizes_js_csrf_header_assignment_double_quotes(self):
+        """JavaScript CSRF header assignment with double quotes should be normalized."""
+        html = b"""<script>
+    const requestInterceptor = (request) => {
+      request.headers["X-CSRFTOKEN"] = "7180AeP5yScQ72FIE7z8tCtrrQc32Sgaqpqgk8ktVXeBXw7a5qDgXoGgLiWF22Ew";
+      return request;
+    };
+    </script>"""
+        result = normalize_csrf_in_body(html)
+        assert result is not None
+        assert CSRF_PLACEHOLDER.encode() in result
+        assert b"7180AeP5yScQ72FIE7z8tCtrrQc32Sgaqpqgk8ktVXeBXw7a5qDgXoGgLiWF22Ew" not in result
+
+    def test_normalizes_js_csrf_header_assignment_single_quotes(self):
+        """JavaScript CSRF header assignment with single quotes should be normalized."""
+        html = b"""<script>
+    request.headers['X-CSRFTOKEN'] = 'someToken123';
+    </script>"""
+        result = normalize_csrf_in_body(html)
+        assert result is not None
+        assert CSRF_PLACEHOLDER.encode() in result
+        assert b"someToken123" not in result
+
+    def test_normalizes_js_csrf_header_case_insensitive(self):
+        """JavaScript CSRF header matching should be case-insensitive (X-CSRFToken)."""
+        html = b"""<script>
+    request.headers["X-CSRFToken"] = "myToken456";
+    </script>"""
+        result = normalize_csrf_in_body(html)
+        assert result is not None
+        assert CSRF_PLACEHOLDER.encode() in result
+        assert b"myToken456" not in result
+
+    def test_normalizes_js_csrf_header_preserves_surrounding_js(self):
+        """Surrounding JavaScript should be preserved when normalizing CSRF header."""
+        html = b"""<script>
+    const requestInterceptor = (request) => {
+      if (!["GET", undefined].includes(request.method)) {
+        request.headers["X-CSRFTOKEN"] = "tokenABC";
+      }
+      return request;
+    };
+    </script>"""
+        result = normalize_csrf_in_body(html)
+        assert result is not None
+        assert b"requestInterceptor" in result
+        assert b"return request" in result
+        assert CSRF_PLACEHOLDER.encode() in result
+        assert b"tokenABC" not in result
+
+    def test_normalizes_both_form_and_js_csrf_tokens(self):
+        """Both form hidden inputs and JS CSRF header assignments should be normalized."""
+        html = b"""<html>
+    <form><input name="csrfmiddlewaretoken" value="formToken123"></form>
+    <script>request.headers["X-CSRFTOKEN"] = "jsToken456";</script>
+    </html>"""
+        result = normalize_csrf_in_body(html)
+        assert result is not None
+        assert result.count(CSRF_PLACEHOLDER.encode()) == 2
+        assert b"formToken123" not in result
+        assert b"jsToken456" not in result
+
 
 class TestNormalizeHtmlBody:
     """Tests for normalize_html_body function."""
