From 919e73f758fe8dae86342401ff2f69859827fe08 Mon Sep 17 00:00:00 2001
From: Kyle Barron <kylebarron2@gmail.com>
Date: Wed, 30 Sep 2020 10:36:15 -0600
Subject: [PATCH] Add gzip fallback (#13)

* Add gzip fallback

* Add tests

* Update README
---
 README.md               | 13 ++++++++++---
 brotli_asgi/__init__.py |  8 ++++++++
 tests.py                | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index a88aea7..9a7c145 100644
--- a/README.md
+++ b/README.md
@@ -1,11 +1,11 @@
-# brotli-middleware
+# brotli-asgi
 
 `BrotliMiddleware` adds [Brotli](https://github.com/google/brotli) response compression to ASGI applications (Starlette, FastAPI, Quart, etc.). It provides faster and more dense compression than GZip, and can be used as a drop in replacement for the `GZipMiddleware` shipped with Starlette.
 
 **Installation**
 
 ```bash
-pip install brotli-middleware
+pip install brotli-asgi
 ```
 
 ## Examples
@@ -49,7 +49,13 @@ def home() -> dict:
 
 ```python
 app.add_middleware(
-  BrotliMiddleware, quality=4, mode="text", lgwin=22, lgblock=0, minimum_size=400,
+  BrotliMiddleware,
+  quality=4,
+  mode="text",
+  lgwin=22,
+  lgblock=0,
+  minimum_size=400,
+  gzip_fallback=True
 )
 ```
 
@@ -60,6 +66,7 @@ app.add_middleware(
 - _(Optional)_ `lgwin`: Base 2 logarithm of the sliding window size. Range is 10 to 24.
 - _(Optional)_ `lgblock`: Base 2 logarithm of the maximum input block size. Range is 16 to 24. If set to 0, the value will be set based on the quality.
 - _(Optional)_ `minimum_size`: Only compress responses that are bigger than this value in bytes.
+- _(Optional)_ `gzip_fallback`: If `True`, uses gzip encoding if `br` is not in the Accept-Encoding header.
 
 ## Performance
 
diff --git a/brotli_asgi/__init__.py b/brotli_asgi/__init__.py
index 5eb8184..f86cbe1 100644
--- a/brotli_asgi/__init__.py
+++ b/brotli_asgi/__init__.py
@@ -7,6 +7,7 @@
 
 from brotli import MODE_FONT, MODE_GENERIC, MODE_TEXT, Compressor  # type: ignore
 from starlette.datastructures import Headers, MutableHeaders
+from starlette.middleware.gzip import GZipResponder
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
 
@@ -29,6 +30,7 @@ def __init__(
         lgwin: int = 22,
         lgblock: int = 0,
         minimum_size: int = 400,
+        gzip_fallback: bool = True,
     ) -> None:
         """
         Arguments.
@@ -45,6 +47,7 @@ def __init__(
             Range is 16 to 24. If set to 0, the value will be set based on the
             quality.
         minimum_size: Only compress responses that are bigger than this value in bytes.
+        gzip_fallback: If True, uses gzip encoding if br is not in the Accept-Encoding header.
         """
         self.app = app
         self.quality = quality
@@ -52,6 +55,7 @@ def __init__(
         self.minimum_size = minimum_size
         self.lgwin = lgwin
         self.lgblock = lgblock
+        self.gzip_fallback = gzip_fallback
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] == "http":
@@ -67,6 +71,10 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
                 )
                 await responder(scope, receive, send)
                 return
+            if self.gzip_fallback and "gzip" in headers.get("Accept-Encoding", ""):
+                responder = GZipResponder(self.app, self.minimum_size)
+                await responder(scope, receive, send)
+                return
         await self.app(scope, receive, send)
 
 
diff --git a/tests.py b/tests.py
index c820bf6..7ff744d 100644
--- a/tests.py
+++ b/tests.py
@@ -100,3 +100,37 @@ def homepage(request):
     client = TestClient(app)
     response = client.get("/", headers={"accept-encoding": "br"})
     assert response.status_code == 200
+
+
+def test_gzip_fallback():
+    app = Starlette()
+
+    app.add_middleware(BrotliMiddleware, gzip_fallback=True)
+
+    @app.route("/")
+    def homepage(request):
+        return PlainTextResponse("x" * 4000, status_code=200)
+
+    client = TestClient(app)
+    response = client.get("/", headers={"accept-encoding": "gzip"})
+    assert response.status_code == 200
+    assert response.text == "x" * 4000
+    assert response.headers["Content-Encoding"] == "gzip"
+    assert int(response.headers["Content-Length"]) < 4000
+
+
+def test_gzip_fallback_false():
+    app = Starlette()
+
+    app.add_middleware(BrotliMiddleware, gzip_fallback=False)
+
+    @app.route("/")
+    def homepage(request):
+        return PlainTextResponse("x" * 4000, status_code=200)
+
+    client = TestClient(app)
+    response = client.get("/", headers={"accept-encoding": "gzip"})
+    assert response.status_code == 200
+    assert response.text == "x" * 4000
+    assert "Content-Encoding" not in response.headers
+    assert int(response.headers["Content-Length"]) == 4000