Add COMPRESS_MUTATE_WEAK_ETAGS new behaviour

Author: jlopex

flask_compress/flask_compress.py

@@ -165,6 +165,7 @@ def init_app(self, app):
             ("COMPRESS_CACHE_BACKEND", None),
             ("COMPRESS_REGISTER", True),
             ("COMPRESS_STREAMS", True),
+            ("COMPRESS_MUTATE_WEAK_ETAGS", True),
             ("COMPRESS_ALGORITHM", ["zstd", "br", "gzip", "deflate"]),
         ]
 
@@ -230,9 +231,10 @@ def after_request(self, response):
 
         # "123456789"   => "123456789:gzip"   - A strong ETag validator
         # W/"123456789" => W/"123456789:gzip" - A weak ETag validator
-        etag = response.headers.get("ETag")
-        if etag:
-            response.headers["ETag"] = f'{etag[:-1]}:{chosen_algorithm}"'
+        etag, is_weak = response.get_etag()
+
+        if etag and (not is_weak or app.config["COMPRESS_MUTATE_WEAK_ETAGS"]):
+            response.set_etag(f"{etag}:{chosen_algorithm}", weak=is_weak)
 
         return response
 

tests/test_flask_compress.py

@@ -3,7 +3,7 @@
 import tempfile
 import unittest
 
-from flask import Flask, render_template
+from flask import Flask, make_response, render_template, request
 from flask_caching import Cache
 
 from flask_compress import Compress, DictCache
@@ -91,6 +91,10 @@ def test_quality_level_default_zstd(self):
         """Tests COMPRESS_ZSTD_LEVEL default value is correctly set."""
         self.assertEqual(self.app.config["COMPRESS_ZSTD_LEVEL"], 3)
 
+    def test_mutate_weak_etags(self):
+        """Tests COMPRESS_MUTATE_WEAK_ETAGS default value is correctly set."""
+        self.assertEqual(self.app.config["COMPRESS_MUTATE_WEAK_ETAGS"], True)
+
 
 class InitTests(unittest.TestCase):
     def setUp(self):
@@ -570,5 +574,93 @@ def test_compression(self):
         self.assertEqual(self.cache_key_calls, 2)
 
 
+class ETagTests(unittest.TestCase):
+    def setUp(self):
+        self.app = Flask(__name__)
+        self.app.testing = True
+        self.app.config["COMPRESS_ALGORITHM"] = ["gzip"]
+        self.app.config["COMPRESS_MIN_SIZE"] = 1
+
+        Compress(self.app)
+
+        @self.app.route("/strong/")
+        def strong():
+            rv = make_response(render_template("large.html"))
+            rv.set_etag("abc123", weak=False)
+            return rv.make_conditional(request)
+
+        @self.app.route("/weak/")
+        def weak():
+            rv = make_response(render_template("large.html"))
+            rv.set_etag("abc123", weak=True)
+            return rv.make_conditional(request)
+
+    def test_strong_etag_is_mutated_with_suffix_and_remains_strong(self):
+        client = self.app.test_client()
+        r = client.get("/strong/", headers=[("Accept-Encoding", "gzip")])
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.headers.get("Content-Encoding"), "gzip")
+
+        tag, is_weak = r.get_etag()
+        self.assertFalse(is_weak)
+        self.assertEqual(tag, "abc123:gzip")
+        self.assertEqual(int(r.headers["Content-Length"]), len(r.data))
+
+    def test_weak_etag_is_mutated_when_flag_true(self):
+        client = self.app.test_client()
+        r = client.get("/weak/", headers=[("Accept-Encoding", "gzip")])
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.headers.get("Content-Encoding"), "gzip")
+
+        tag, is_weak = r.get_etag()
+        self.assertTrue(is_weak)
+        self.assertEqual(tag, "abc123:gzip")
+
+    def test_strong_etag_is_mutated_when_flag_false(self):
+        self.app.config["COMPRESS_MUTATE_WEAK_ETAGS"] = False
+        client = self.app.test_client()
+        r = client.get("/strong/", headers=[("Accept-Encoding", "gzip")])
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.headers.get("Content-Encoding"), "gzip")
+
+        tag, is_weak = r.get_etag()
+        self.assertFalse(is_weak)
+        # :gzip suffix when flag is False
+        self.assertEqual(tag, "abc123:gzip")
+
+    def test_weak_etag_is_preserved_when_flag_false(self):
+        self.app.config["COMPRESS_MUTATE_WEAK_ETAGS"] = False
+        client = self.app.test_client()
+        r = client.get("/weak/", headers=[("Accept-Encoding", "gzip")])
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.headers.get("Content-Encoding"), "gzip")
+
+        tag, is_weak = r.get_etag()
+        self.assertTrue(is_weak)
+        # No :gzip suffix when flag is False
+        self.assertEqual(tag, "abc123")
+
+    def test_conditional_get_uses_strong_compressed_representation(self):
+        client = self.app.test_client()
+        r1 = client.get("/strong/", headers=[("Accept-Encoding", "gzip")])
+
+        r2 = client.get(
+            "/strong/",
+            headers=[("Accept-Encoding", "gzip"), ("If-None-Match", r1.headers["ETag"])],
+        )
+        # This is the current behavior that breaks strong etags due rewrite
+        self.assertEqual(r2.status_code, 200)
+
+    def test_conditional_get_uses_weak_compressed_representation(self):
+        self.app.config["COMPRESS_MUTATE_WEAK_ETAGS"] = False
+        client = self.app.test_client()
+        r1 = client.get("/weak/", headers=[("Accept-Encoding", "gzip")])
+        r2 = client.get(
+            "/weak/",
+            headers=[("Accept-Encoding", "gzip"), ("If-None-Match", r1.headers["ETag"])],
+        )
+        self.assertEqual(r2.status_code, 304)
+
+
 if __name__ == "__main__":
     unittest.main()