Update Content-Length header in frontend monitoring middleware also a…

…dds waffle switch (#418)

* fix: update Content-Length header in fronted monitoring middleware

This adds the functionality to update Content-Length header if already set to avoid content trimming issue.

* feat: add waffle switch in frontend monitoring middleware

* chore: bump version to 15.14.2 for release

* fix: disable middleware if flag is not enabled

* fix: improve qulity and fix typo

CHANGELOG.rst

@@ -11,6 +11,13 @@ Change Log
 
 .. There should always be an "Unreleased" section for changes pending release.
 
+[5.14.2] - 2024-05-31
+---------------------
+Fixed
+~~~~~
+* FrontendMonitoringMiddleware now updates the Content-Length header if its already set.
+* FrontendMonitoringMiddleware will now be enabled once waffle switch named ``edx_django_utils.monitoring.enable_frontend_monitoring_middleware`` is enabled.
+
 [5.14.1] - 2024-05-22
 ---------------------
 Fixed

edx_django_utils/__init__.py

@@ -2,7 +2,7 @@
 EdX utilities for Django Application development..
 """
 
-__version__ = "5.14.1"
+__version__ = "5.14.2"
 
 default_app_config = (
     "edx_django_utils.apps.EdxDjangoUtilsConfig"

edx_django_utils/monitoring/README.rst

@@ -108,6 +108,7 @@ Frontend Monitoring Middleware
 --------------------------------
 
 This middleware ``FrontendMonitoringMiddleware`` inserts frontend monitoring related HTML script tags to the response, see docstring for details.
+In addition to adding the FrontendMonitoringMiddleware, you will need to enable a waffle switch ``edx_django_utils.monitoring.enable_frontend_monitoring_middleware`` to enable the frontend monitoring.
 
 Monitoring Memory Usage
 -----------------------

edx_django_utils/monitoring/internal/middleware.py

@@ -16,6 +16,7 @@
 import psutil
 import waffle  # pylint: disable=invalid-django-waffle-import
 from django.conf import settings
+from django.core.exceptions import MiddlewareNotUsed
 from django.utils.deprecation import MiddlewareMixin
 
 from edx_django_utils.cache import RequestCache
@@ -471,16 +472,24 @@ class FrontendMonitoringMiddleware:
     Middleware for adding the frontend monitoring scripts to the response.
     """
     def __init__(self, get_response):
+        # Disable the middleware if flag isn't enabled
+        if not self._is_enabled():
+            raise MiddlewareNotUsed
+
         self.get_response = get_response
 
     def __call__(self, request):
         response = self.get_response(request)
 
         content_type = response.headers.get('Content-Type', '')
+        content_disposition = response.headers.get('Content-Disposition')
 
         if response.status_code != 200 or not content_type.startswith('text/html'):
             return response
 
+        if content_disposition is not None and content_disposition.split(";")[0].strip().lower() == "attachment":
+            return response
+
         # .. setting_name: OPENEDX_TELEMETRY_FRONTEND_SCRIPTS
         # .. setting_default: None
         # .. setting_description: Scripts to inject to response for frontend monitoring, this can
@@ -497,7 +506,13 @@ def __call__(self, request):
             # Prevent a certain kind of easy mistake.
             raise Exception("OPENEDX_TELEMETRY_FRONTEND_SCRIPTS must be a string.")
 
+        original_content_len = len(response.content)
         response.content = self.inject_script(response.content, frontend_scripts)
+
+        # If HTML is added and Content-Length already set, make sure Content-Length header is updated.
+        # If not browsers can trim response, as we are adding HTML to the response.
+        if len(response.content) != original_content_len and response.headers.get("Content-Length"):
+            response.headers["Content-Length"] = str(len(response.content))
         return response
 
     def inject_script(self, content, script):
@@ -522,6 +537,12 @@ def insert_html_at_index(index):
         # Don't add the script if both head and body tag is missing.
         return content
 
+    def _is_enabled(self):
+        """
+        Returns whether this middleware is enabled.
+        """
+        return waffle.switch_is_active('edx_django_utils.monitoring.enable_frontend_monitoring_middleware')
+
 
 # This function should be cleaned up and made into a general logging utility, but it will first
 # need some work to make it able to handle multibyte characters.

edx_django_utils/monitoring/tests/test_middleware.py

@@ -7,6 +7,7 @@
 from unittest.mock import Mock, call, patch
 
 import ddt
+from django.core.exceptions import MiddlewareNotUsed
 from django.http import HttpRequest, HttpResponse, JsonResponse
 from django.test import TestCase
 from django.test.client import RequestFactory
@@ -335,6 +336,31 @@ def setUp(self):
         super().setUp()
         self.script = "<script>test script</script>"
 
+    @override_switch('edx_django_utils.monitoring.enable_frontend_monitoring_middleware', False)
+    def test_frontend_middleware_with_waffle_diasbled(self):
+        """
+        Test that middleware is disabled when waffle flag is not enabled.
+        """
+        original_html = '<head></head>'
+        with override_settings(OPENEDX_TELEMETRY_FRONTEND_SCRIPTS=self.script):
+            self.assertRaises(
+                MiddlewareNotUsed,
+                FrontendMonitoringMiddleware,
+                lambda r: HttpResponse(original_html, content_type='text/html')
+            )
+
+    @override_switch('edx_django_utils.monitoring.enable_frontend_monitoring_middleware', True)
+    def test_frontend_middleware_with_waffle_enabled(self):
+        """
+        Test that middleware works as expected when flag is enabled.
+        """
+        with override_settings(OPENEDX_TELEMETRY_FRONTEND_SCRIPTS=self.script):
+            middleware = FrontendMonitoringMiddleware(lambda r: HttpResponse('<head></head>', content_type='text/html'))
+            response = middleware(HttpRequest())
+        # Assert that the script is inserted into the response when flag is enabled
+        assert self.script.encode() in response.content
+
+    @override_switch('edx_django_utils.monitoring.enable_frontend_monitoring_middleware', True)
     @patch("edx_django_utils.monitoring.internal.middleware.FrontendMonitoringMiddleware.inject_script")
     def test_frontend_middleware_without_setting_variable(self, mock_inject_script):
         """
@@ -348,6 +374,7 @@ def test_frontend_middleware_without_setting_variable(self, mock_inject_script):
 
         mock_inject_script.assert_not_called()
 
+    @override_switch('edx_django_utils.monitoring.enable_frontend_monitoring_middleware', True)
     @patch("edx_django_utils.monitoring.internal.middleware.FrontendMonitoringMiddleware.inject_script")
     def test_frontend_middleware_for_json_requests(self, mock_inject_script):
         """
@@ -360,6 +387,7 @@ def test_frontend_middleware_for_json_requests(self, mock_inject_script):
 
         mock_inject_script.assert_not_called()
 
+    @override_switch('edx_django_utils.monitoring.enable_frontend_monitoring_middleware', True)
     @ddt.data(
         ('<html><body></body><html>', '<body>'),
         ('<html><head></head><body></body><html>', '</head>'),
@@ -379,6 +407,7 @@ def test_frontend_middleware_with_head_and_body_tag(self, original_html, expecte
         # Assert that the script is inserted at the right place
         assert f"{self.script}{expected_tag}".encode() in response.content
 
+    @override_switch('edx_django_utils.monitoring.enable_frontend_monitoring_middleware', True)
     @ddt.data(
         '<html></html>',
         '<center></center>',
@@ -392,3 +421,32 @@ def test_frontend_middleware_without_head_and_body_tag(self, original_html):
             response = middleware(HttpRequest())
         # Assert that the response content remains unchanged if no body tag is found
         assert response.content == original_html.encode()
+
+    @override_switch('edx_django_utils.monitoring.enable_frontend_monitoring_middleware', True)
+    def test_frontend_middleware_content_length_header_already_set(self):
+        """
+        Test that middleware updates the Content-Length header, when its already set.
+        """
+        original_html = '<head></head>'
+        with override_settings(OPENEDX_TELEMETRY_FRONTEND_SCRIPTS=self.script):
+            middleware = FrontendMonitoringMiddleware(lambda r: HttpResponse(
+                original_html, content_type='text/html', headers={'Content-Length': len(original_html)}))
+            response = middleware(HttpRequest())
+        # Assert that the response content contains script tag
+        assert self.script.encode() in response.content
+        # Assert that the Content-Length header is updated and script length is added.
+        assert response.headers.get('Content-Length') == str(len(original_html) + len(self.script))
+
+    @override_switch('edx_django_utils.monitoring.enable_frontend_monitoring_middleware', True)
+    def test_frontend_middleware_content_length_header_not_set(self):
+        """
+        Test that middleware doesn't set the Content-Length header when it's not already set.
+        """
+        original_html = '<head></head>'
+        with override_settings(OPENEDX_TELEMETRY_FRONTEND_SCRIPTS=self.script):
+            middleware = FrontendMonitoringMiddleware(lambda r: HttpResponse(original_html, content_type='text/html'))
+            response = middleware(HttpRequest())
+        # Assert that the response content contains script tag
+        assert self.script.encode() in response.content
+        # Assert that the Content-Length header isn't updated, when not set already
+        assert response.headers.get('Content-Length') is None