From 2aeb103cf0beec402fcfa786a727274b073450f7 Mon Sep 17 00:00:00 2001
From: Adam Johnson <me@adamj.eu>
Date: Mon, 7 Oct 2024 21:44:01 +0100
Subject: [PATCH] Send binary responses for any content-encoding (#526)

Fixes the issue reported in #496.
---
 CHANGELOG.rst             |  4 ++++
 README.rst                |  9 +++++----
 src/apig_wsgi/__init__.py | 14 ++++++-------
 tests/test_apig_wsgi.py   | 42 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 57 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index ad93412..fc35352 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -2,6 +2,10 @@
 Changelog
 =========
 
+* Send binary responses if the 'content-encoding' header is set to any value, rather than just 'gzip'.
+
+  Thanks to Zoe Guillen for the report in `PR #496 <https://github.com/adamchainz/apig-wsgi/pull/496>`__.
+
 * Support Python 3.13.
 
 * Drop Python 3.8 support.
diff --git a/README.rst b/README.rst
index d9aff60..4af7a3e 100644
--- a/README.rst
+++ b/README.rst
@@ -89,12 +89,13 @@ responses:
   header as well, which WSGI applications often ignore. You may need to delete
   and recreate your stages for this value to be copied over.
 
-Note that binary responses aren't sent if your response has a 'Content-Type'
-starting 'text/', 'application/json' or 'application/vnd.api+json' - this
-is to support sending larger text responses, since the base64 encoding would
+Note that binary responses aren't sent if your response has no
+'content-encoding' header and a 'content-type' header starting
+'text/', 'application/json' or 'application/vnd.api+json'. This behaviour is to
+support sending larger text responses, since the base64 encoding would
 otherwise inflate the content length. To avoid base64 encoding other content
 types, you can set ``non_binary_content_type_prefixes`` to a list or tuple of
-content type prefixes of your choice (which replaces the default list).
+content type prefixes of your choice, which replaces the default list.
 
 If the event from API Gateway contains the ``requestContext`` key, for example
 on format version 2 or from custom request authorizers, this will be available
diff --git a/src/apig_wsgi/__init__.py b/src/apig_wsgi/__init__.py
index b5d63ca..bb8fbdc 100644
--- a/src/apig_wsgi/__init__.py
+++ b/src/apig_wsgi/__init__.py
@@ -274,20 +274,18 @@ def _should_send_binary(self) -> bool:
         if not self.binary_support:
             return False
 
-        content_type = self._get_content_type()
-        if not content_type.startswith(self.non_binary_content_type_prefixes):
+        if self._get_content_encoding() > "":
             return True
 
-        content_encoding = self._get_content_encoding()
-        # Content type is non-binary but the content encoding might be.
-        return "gzip" in content_encoding.lower()
-
-    def _get_content_type(self) -> str:
-        return self._get_header("content-type") or ""
+        content_type = self._get_content_type()
+        return not content_type.startswith(self.non_binary_content_type_prefixes)
 
     def _get_content_encoding(self) -> str:
         return self._get_header("content-encoding") or ""
 
+    def _get_content_type(self) -> str:
+        return self._get_header("content-type") or ""
+
     def _get_header(self, header_name: str) -> str | None:
         header_name = header_name.lower()
         matching_headers = [v for k, v in self.headers if k.lower() == header_name]
diff --git a/tests/test_apig_wsgi.py b/tests/test_apig_wsgi.py
index fd2626d..c906007 100644
--- a/tests/test_apig_wsgi.py
+++ b/tests/test_apig_wsgi.py
@@ -197,6 +197,27 @@ def test_get_binary_support_default_text_content_types(
             "body": "Hello World\n",
         }
 
+    @parametrize_default_text_content_type
+    def test_get_binary_support_default_text_content_types_encoded(
+        self, simple_app: App, text_content_type: str
+    ) -> None:
+        simple_app.handler = make_lambda_handler(simple_app, binary_support=True)
+        simple_app.headers = [
+            ("Content-Type", text_content_type),
+            ("Content-Encoding", "brotli"),
+        ]
+
+        response = simple_app.handler(make_v1_event(), None)
+        assert response == {
+            "statusCode": 200,
+            "multiValueHeaders": {
+                "Content-Type": [text_content_type],
+                "Content-Encoding": ["brotli"],
+            },
+            "isBase64Encoded": True,
+            "body": b64encode(b"Hello World\n").decode("utf-8"),
+        }
+
     @parametrize_custom_text_content_type
     def test_get_binary_support_custom_text_content_types(
         self, simple_app: App, text_content_type: str
@@ -802,6 +823,27 @@ def test_get_binary_support_default_text_content_types(
             "body": "Hello World\n",
         }
 
+    @parametrize_default_text_content_type
+    def test_get_binary_support_default_text_content_types_encoded(
+        self, simple_app: App, text_content_type: str
+    ) -> None:
+        simple_app.handler = make_lambda_handler(simple_app, binary_support=True)
+        simple_app.headers = [
+            ("Content-Type", text_content_type),
+            ("Content-Encoding", "brotli"),
+        ]
+
+        response = simple_app.handler(make_v1_event(), None)
+        assert response == {
+            "statusCode": 200,
+            "multiValueHeaders": {
+                "Content-Type": [text_content_type],
+                "Content-Encoding": ["brotli"],
+            },
+            "isBase64Encoded": True,
+            "body": b64encode(b"Hello World\n").decode("utf-8"),
+        }
+
     @parametrize_custom_text_content_type
     def test_get_binary_support_custom_text_content_types(
         self, simple_app: App, text_content_type: str