Do not attempt to inject HTML into streaming responses

cavalry/reporter.py

@@ -4,7 +4,7 @@
 from typing import List, Optional
 
 from django.core.handlers.wsgi import WSGIRequest
-from django.http import HttpResponse
+from django.http import HttpResponse, StreamingHttpResponse
 
 from cavalry.policy import can_report_stacks
 from cavalry.stack import Stack
@@ -28,7 +28,7 @@
 def inject_html(request: WSGIRequest, response: HttpResponse, data: dict, summary_data: dict) -> None:
     try:
         body_closing_index = response.content.rindex(b"</body>")
-    except ValueError:
+    except (AttributeError, ValueError):
         return
     content = " | ".join(
         f"{key}={round(value, 3) if isinstance(value, float) else value}"
@@ -82,10 +82,16 @@ def build_log_command(query: dict) -> str:
 
 def inject_stats(request: WSGIRequest, response: HttpResponse, data: dict) -> None:
     summary_data = summarize_data(data)
-    content_type = response.get("content-type", "").lower()
-    if content_type.startswith("text/html") and ("charset" not in content_type or "utf-8" in content_type):
-        inject_html(request, response, data, summary_data)
     response["x-cavalry-data"] = json.dumps(summary_data)
+    if can_inject_html(response):
+        inject_html(request, response, data, summary_data)
+
+
+def can_inject_html(response: HttpResponse) -> bool:
+    if isinstance(response, StreamingHttpResponse):
+        return False
+    content_type = response.get("content-type", "").lower()
+    return content_type.startswith("text/html") and ("charset" not in content_type or "utf-8" in content_type)
 
 
 def summarize_data(data: dict) -> dict:

cavalry_tests/tests/test_cavalry.py

@@ -2,14 +2,16 @@
 
 import pytest
 import requests_mock
+from django.http import StreamingHttpResponse
 from django.test import Client
 
 
 @pytest.mark.django_db()
-@pytest.mark.parametrize("enable", [False, True])
-@pytest.mark.parametrize("as_admin", [False, True])
-@pytest.mark.parametrize("posting", [False, True])
-def test_cavalry(settings, as_admin, enable, posting, admin_user):
+@pytest.mark.parametrize("enable", [False, True], ids=("disabled", "enabled"))
+@pytest.mark.parametrize("as_admin", [False, True], ids=("as_user", "as_admin"))
+@pytest.mark.parametrize("posting", [False, True], ids=("nopost", "posting"))
+@pytest.mark.parametrize("streaming", [False, True], ids=("regular", "streaming"))
+def test_cavalry(settings, as_admin, enable, posting, admin_user, streaming):
     settings.CAVALRY_ENABLED = enable
     settings.CAVALRY_ELASTICSEARCH_URL_TEMPLATE = "http://localhost:59595/asdf/foo" if posting else None
     client = Client()
@@ -24,13 +26,20 @@ def test_cavalry(settings, as_admin, enable, posting, admin_user):
             status_code=201,
             json={"ok": True},
         )
-        content = client.get("/").content
+        resp = client.get("/streaming/" if streaming else "/")
+        if isinstance(resp, StreamingHttpResponse):
+            content = resp.getvalue()
+        else:
+            content = resp.content
 
     # Check precondition: the user seemed logged in
-    assert (admin_user.username.encode() in content) == as_admin
+    assert (f"Henlo {admin_user.username}".encode() in content) == as_admin
 
-    # Check that the injection div only appears for admins
-    assert (b"<div" in content) == (enable and as_admin)
+    should_expose_info = enable and as_admin
+    # Check that the injection div/header only appears for admins and when not streaming
+    assert (b"<div" in content) == (should_expose_info and not streaming)
+    # Check the header is present when it should be
+    assert ("x-cavalry-data" in resp.headers) == should_expose_info
 
     # Check that stats are posted only when posting is enabled
     assert bool(m.called) == (enable and posting)

cavalry_tests/urls.py

@@ -1,8 +1,17 @@
 from django.contrib import admin
+from django.http import StreamingHttpResponse
+from django.template.loader import render_to_string
 from django.urls import path
 from django.views.generic import TemplateView
 
+
+def streaming_view(request):
+    content = render_to_string("index.html", request=request)
+    return StreamingHttpResponse(iter(content), content_type="text/html")
+
+
 urlpatterns = [
     path("admin/", admin.site.urls),
     path("", TemplateView.as_view(template_name="index.html")),
+    path("streaming/", streaming_view),
 ]