Add COMPRESS_EVALUATE_CONDITIONAL_REQUEST new behaviour

Author: jlopex

flask_compress/flask_compress.py

@@ -166,6 +166,7 @@ def init_app(self, app):
             ("COMPRESS_REGISTER", True),
             ("COMPRESS_STREAMS", True),
             ("COMPRESS_MUTATE_WEAK_ETAGS", True),
+            ("COMPRESS_EVALUATE_CONDITIONAL_REQUEST", False),
             ("COMPRESS_ALGORITHM", ["zstd", "br", "gzip", "deflate"]),
         ]
 
@@ -236,6 +237,13 @@ def after_request(self, response):
         if etag and (not is_weak or app.config["COMPRESS_MUTATE_WEAK_ETAGS"]):
             response.set_etag(f"{etag}:{chosen_algorithm}", weak=is_weak)
 
+        if (
+            app.config["COMPRESS_EVALUATE_CONDITIONAL_REQUEST"]
+            and request.method in ("GET", "HEAD")
+            and (not response.is_streamed or app.config["COMPRESS_STREAMS"])
+        ):
+            response.make_conditional(request)
+
         return response
 
     def compressed(self):

tests/test_flask_compress.py

@@ -95,6 +95,13 @@ 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)
 
+    def test_evaluate_conditional_request(self):
+        """Tests COMPRESS_EVALUATE_CONDITIONAL_REQUEST default value
+        is correctly set."""
+        self.assertEqual(
+            self.app.config["COMPRESS_EVALUATE_CONDITIONAL_REQUEST"], False
+        )
+
 
 class InitTests(unittest.TestCase):
     def setUp(self):
@@ -589,12 +596,24 @@ def strong():
             rv.set_etag("abc123", weak=False)
             return rv.make_conditional(request)
 
+        @self.app.route("/strong-compress-conditional/")
+        def strong_compress_conditional():
+            rv = make_response(render_template("large.html"))
+            rv.set_etag("abc123", weak=False)
+            return rv
+
         @self.app.route("/weak/")
         def weak():
             rv = make_response(render_template("large.html"))
             rv.set_etag("abc123", weak=True)
             return rv.make_conditional(request)
 
+        @self.app.route("/weak-compress-conditional/")
+        def weak_compress_conditional():
+            rv = make_response(render_template("large.html"))
+            rv.set_etag("abc123", weak=True)
+            return rv
+
     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")])
@@ -646,20 +665,89 @@ def test_conditional_get_uses_strong_compressed_representation(self):
 
         r2 = client.get(
             "/strong/",
-            headers=[("Accept-Encoding", "gzip"), ("If-None-Match", r1.headers["ETag"])],
+            headers=[
+                ("Accept-Encoding", "gzip"),
+                ("If-None-Match", r1.headers["ETag"]),
+            ],
         )
-        # This is the current behavior that breaks strong etags due rewrite
+        # This is the current behavior that breaks make_conditional
+        # strong etags due rewrite at after_request
+        # We would expect a 304 but it does not because of etag mismatch
         self.assertEqual(r2.status_code, 200)
 
     def test_conditional_get_uses_weak_compressed_representation(self):
+        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"]),
+            ],
+        )
+        # This is the current behavior that breaks make_conditional
+        # weak etags due rewrite at after_request
+        # We would expect a 304 but it does not because of etag mismatch
+        self.assertEqual(r2.status_code, 200)
+
+    def test_conditional_get_uses_strong_compressed_representation_evaluate_conditional(
+        self,
+    ):
+        self.app.config["COMPRESS_EVALUATE_CONDITIONAL_REQUEST"] = True
+        client = self.app.test_client()
+        r1 = client.get(
+            "/strong-compress-conditional/", headers=[("Accept-Encoding", "gzip")]
+        )
+        etag_header = r1.headers["ETag"]
+
+        r2 = client.get(
+            "/strong-compress-conditional/",
+            headers=[("Accept-Encoding", "gzip"), ("If-None-Match", etag_header)],
+        )
+        # This is the new behaviour we would expect after evaluating
+        # flask make_conditional at after_request
+        self.assertEqual(r2.status_code, 304)
+        self.assertEqual(r2.headers.get("ETag"), etag_header)
+        self.assertNotIn("Content-Encoding", r2.headers)
+        self.assertEqual(len(r2.get_data()), 0)
+
+    def test_conditional_get_uses_weak_compressed_representation_evaluate_conditional(
+        self,
+    ):
+        self.app.config["COMPRESS_EVALUATE_CONDITIONAL_REQUEST"] = True
+        client = self.app.test_client()
+        r1 = client.get(
+            "/weak-compress-conditional/", headers=[("Accept-Encoding", "gzip")]
+        )
+        etag_header = r1.headers["ETag"]
+
+        r2 = client.get(
+            "/weak-compress-conditional/",
+            headers=[("Accept-Encoding", "gzip"), ("If-None-Match", etag_header)],
+        )
+        # This is the new behaviour we would expect after evaluating
+        # flask make_conditional at after_request
+        self.assertEqual(r2.status_code, 304)
+        self.assertEqual(r2.headers.get("ETag"), etag_header)
+        self.assertNotIn("Content-Encoding", r2.headers)
+        self.assertEqual(len(r2.get_data()), 0)
+
+    def test_conditional_get_uses_weak_compressed_representation_no_mutate_etag(self):
         self.app.config["COMPRESS_MUTATE_WEAK_ETAGS"] = False
         client = self.app.test_client()
         r1 = client.get("/weak/", headers=[("Accept-Encoding", "gzip")])
+        etag_header = r1.headers["ETag"]
+
         r2 = client.get(
             "/weak/",
-            headers=[("Accept-Encoding", "gzip"), ("If-None-Match", r1.headers["ETag"])],
+            headers=[("Accept-Encoding", "gzip"), ("If-None-Match", etag_header)],
         )
+        # This is the new behaviour we would expect by not mutating
+        # the weak etags at after_request
         self.assertEqual(r2.status_code, 304)
+        self.assertEqual(r2.headers.get("ETag"), etag_header)
+        self.assertNotIn("Content-Encoding", r2.headers)
+        self.assertEqual(len(r2.get_data()), 0)
 
 
 if __name__ == "__main__":