From 10263130903a176e2bf1042448106a76aaac674f Mon Sep 17 00:00:00 2001
From: Viranch Mehta <email@viranch.me>
Date: Wed, 29 Nov 2017 15:18:13 +0530
Subject: [PATCH] Add support for brotli compression

---
 hyper/http11/response.py |  3 +++
 hyper/http20/response.py |  2 ++
 setup.py                 |  2 +-
 test/test_http11.py      | 11 +++++++++++
 test/test_hyper.py       | 10 ++++++++++
 5 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/hyper/http11/response.py b/hyper/http11/response.py
index 31d2266d..7ff7a523 100644
--- a/hyper/http11/response.py
+++ b/hyper/http11/response.py
@@ -9,6 +9,7 @@
 import logging
 import weakref
 import zlib
+import brotli
 
 from ..common.decoder import DeflateDecoder
 from ..common.exceptions import ChunkedDecodeError, InvalidResponseError
@@ -88,6 +89,8 @@ def __init__(self, code, reason, headers, sock, connection=None,
         # http://stackoverflow.com/a/2695466/1401686
         if b'gzip' in self.headers.get(b'content-encoding', []):
             self._decompressobj = zlib.decompressobj(16 + zlib.MAX_WBITS)
+        elif b'br' in self.headers.get(b'content-encoding', []):
+            self._decompressobj = brotli.Decompressor()
         elif b'deflate' in self.headers.get(b'content-encoding', []):
             self._decompressobj = DeflateDecoder()
         else:
diff --git a/hyper/http20/response.py b/hyper/http20/response.py
index 280ffbb2..29437cb9 100644
--- a/hyper/http20/response.py
+++ b/hyper/http20/response.py
@@ -8,6 +8,7 @@
 """
 import logging
 import zlib
+import brotli
 
 from ..common.decoder import DeflateDecoder
 from ..common.headers import HTTPHeaderMap
@@ -31,6 +32,7 @@ def strip_headers(headers):
 
 decompressors = {
     b'gzip': lambda: zlib.decompressobj(16 + zlib.MAX_WBITS),
+    b'br': brotli.Decompressor,
     b'deflate': DeflateDecoder
 }
 
diff --git a/setup.py b/setup.py
index a2578a6b..94cd8d21 100644
--- a/setup.py
+++ b/setup.py
@@ -77,7 +77,7 @@ def run_tests(self):
         'Programming Language :: Python :: Implementation :: CPython',
     ],
     install_requires=[
-        'h2>=2.4,<3.0,!=2.5.0', 'hyperframe>=3.2,<4.0', 'rfc3986>=1.1.0,<2.0'
+        'h2>=2.4,<3.0,!=2.5.0', 'hyperframe>=3.2,<4.0', 'rfc3986>=1.1.0,<2.0', 'brotlipy>=0.7.0,<1.0'
     ],
     tests_require=['pytest', 'requests', 'mock'],
     cmdclass={'test': PyTest},
diff --git a/test/test_http11.py b/test/test_http11.py
index 4a9280c9..9f3fd3d0 100644
--- a/test/test_http11.py
+++ b/test/test_http11.py
@@ -7,6 +7,7 @@
 """
 import os
 import zlib
+import brotli
 
 from collections import namedtuple
 from io import BytesIO, StringIO
@@ -644,6 +645,16 @@ def test_response_transparently_decrypts_gzip(self):
 
         assert r.read() == b'this is test data'
 
+    def test_response_transparently_decrypts_brotli(self):
+        d = DummySocket()
+        headers = {b'content-encoding': [b'br'], b'connection': [b'close']}
+        r = HTTP11Response(200, 'OK', headers, d, None)
+
+        body = brotli.compress(b'this is test data')
+        d._buffer = BytesIO(body)
+
+        assert r.read() == b'this is test data'
+
     def test_response_transparently_decrypts_real_deflate(self):
         d = DummySocket()
         headers = {
diff --git a/test/test_hyper.py b/test/test_hyper.py
index ae4584a8..ca4696cd 100644
--- a/test/test_hyper.py
+++ b/test/test_hyper.py
@@ -26,6 +26,7 @@
 import pytest
 import socket
 import zlib
+import brotli
 from io import BytesIO
 
 TEST_DIR = os.path.abspath(os.path.dirname(__file__))
@@ -1088,6 +1089,15 @@ def test_response_transparently_decrypts_gzip(self):
 
         assert resp.read() == b'this is test data'
 
+    def test_response_transparently_decrypts_brotli(self):
+        headers = HTTPHeaderMap(
+            [(':status', '200'), ('content-encoding', 'br')]
+        )
+        body = brotli.compress(b'this is test data')
+        resp = HTTP20Response(headers, DummyStream(body))
+
+        assert resp.read() == b'this is test data'
+
     def test_response_transparently_decrypts_real_deflate(self):
         headers = HTTPHeaderMap(
             [(':status', '200'), ('content-encoding', 'deflate')]