Strip 'Cookie' from Vary header if the Vary header contains other entries (#48)

The feature to strip cookies will not be triggered if the Vary header contains something other than "Cookie", e.g. Vary: Language, Cookie. This change properly inspects and rebuilds the Vary header in such cases.

Author: vsalvino

testproject/home/tests.py

@@ -342,6 +342,14 @@ def test_client_tracking_cookies_ignore(self):
         self.client.cookies["_dataminer"] = "precious data"
         self.get_hit(self.page_cachedpage.get_url())
 
+    @override_settings(WAGTAIL_CACHE_IGNORE_COOKIES=True)
+    def test_vary_header_parse(self):
+        self.get_miss(reverse("vary_view"))
+        r = self.get_hit(reverse("vary_view"))
+        # Cookie should have been stripped from the Vary header while preserving
+        # case and order of the other items.
+        self.assertEqual(r["Vary"], "A, B, C")
+
     def test_page_restricted(self):
         auth_url = "/_util/authenticate_with_password/%d/%d/" % (
             self.view_restriction.id,

testproject/home/views.py

@@ -10,3 +10,9 @@ def cached_view(request):
 @nocache_page
 def nocached_view(request):
     return HttpResponse("Hello, World!")
+
+
+def vary_view(request):
+    r = HttpResponse("Variety is the spice of life.")
+    r.headers["Vary"] = "A, B, Cookie, C"
+    return r

testproject/testproject/urls.py

@@ -12,8 +12,9 @@
     path("django-admin/", admin.site.urls),
     path("admin/", include(wagtailadmin_urls)),
     path("documents/", include(wagtaildocs_urls)),
-    path("views/cached-view/", views.cached_view, name="cached_view"),
-    path("views/nocache-view/", views.nocached_view, name="nocached_view"),
+    path("views/cached/", views.cached_view, name="cached_view"),
+    path("views/nocache/", views.nocached_view, name="nocached_view"),
+    path("views/vary/", views.vary_view, name="vary_view"),
     # For anything not caught by a more specific rule above, hand over to
     # Wagtail's page serving mechanism. This should be the last pattern in
     # the list:

wagtailcache/cache.py

@@ -13,6 +13,7 @@
 from django.http.response import HttpResponse
 from django.template.response import SimpleTemplateResponse
 from django.utils.cache import (
+    cc_delim_re,
     get_cache_key,
     get_max_age,
     has_vary_header,
@@ -45,6 +46,9 @@ class Status(Enum):
 
 
 def _patch_header(response: HttpResponse, status: Status) -> None:
+    """
+    Adds our Cache Control status to the response headers.
+    """
     # Patch cache-control with no-cache if it is not already set.
     if status == Status.SKIP and not response.get("Cache-Control", None):
         response["Cache-Control"] = CacheControl.NOCACHE.value
@@ -53,6 +57,31 @@ def _patch_header(response: HttpResponse, status: Status) -> None:
         response[wagtailcache_settings.WAGTAIL_CACHE_HEADER] = status.value
 
 
+def _delete_vary_cookie(response: HttpResponse) -> None:
+    """
+    Deletes the ``Vary: Cookie`` header while keeping other items of the
+    Vary header in tact. Inspired by ``django.utils.cache.patch_vary_headers``.
+    """
+    if not response.has_header("Vary"):
+        return
+    # Parse the value of Vary header.
+    vary_headers = cc_delim_re.split(response["Vary"])
+    # Build a lowercase-keyed dict to preserve the original case.
+    vhdict = {}
+    for item in vary_headers:
+        vhdict.update({item.lower(): item})
+    # Delete "Cookie".
+    if "cookie" in vhdict:
+        del vhdict["cookie"]
+        # Delete the header if it's now empty.
+        if not vhdict:
+            del response["Vary"]
+            return
+        # Else patch the header.
+        vary_headers = [vhdict[k] for k in vhdict]
+        response["Vary"] = ", ".join(vary_headers)
+
+
 def _chop_querystring(r: WSGIRequest) -> WSGIRequest:
     """
     Given a request object, remove any of our ignored querystrings from it.
@@ -103,15 +132,15 @@ def _chop_response_vary(r: WSGIRequest, s: HttpResponse) -> HttpResponse:
     if (
         not s.has_header("Set-Cookie")
         and s.has_header("Vary")
-        and s["Vary"].lower() == "cookie"
+        and has_vary_header(s, "Cookie")
         and not (
             settings.CSRF_COOKIE_NAME in s.cookies
             or settings.CSRF_COOKIE_NAME in r.COOKIES
             or settings.SESSION_COOKIE_NAME in s.cookies
             or settings.SESSION_COOKIE_NAME in r.COOKIES
         )
     ):
-        del s["Vary"]
+        _delete_vary_cookie(s)
     return s