From 51c501ce3972f04033e4ab8fee134bd4d4c4254c Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Wed, 24 Jul 2024 14:11:41 +0300
Subject: [PATCH] Do not attempt to inject HTML into streaming responses

---
 cavalry/reporter.py                 | 16 +++++++++++-----
 cavalry_tests/tests/test_cavalry.py | 25 +++++++++++++++++--------
 cavalry_tests/urls.py               |  9 +++++++++
 3 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/cavalry/reporter.py b/cavalry/reporter.py
index 2dc542e..2d3c07b 100644
--- a/cavalry/reporter.py
+++ b/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:
diff --git a/cavalry_tests/tests/test_cavalry.py b/cavalry_tests/tests/test_cavalry.py
index cabbfb6..81be9f0 100644
--- a/cavalry_tests/tests/test_cavalry.py
+++ b/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)
diff --git a/cavalry_tests/urls.py b/cavalry_tests/urls.py
index 97ca30f..49da8b7 100644
--- a/cavalry_tests/urls.py
+++ b/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),
 ]