From bdc218f8e686a2936c4069b00acde66391eaaf4c Mon Sep 17 00:00:00 2001
From: Yeray Diaz Diaz <yeraydiazdiaz@gmail.com>
Date: Fri, 26 Apr 2019 20:09:57 +0100
Subject: [PATCH] Wrap decoding errors with custom exceptions

---
 httpcore/__init__.py   |  3 +++
 httpcore/decoders.py   | 37 +++++++++++++++++++++++++++++--------
 httpcore/exceptions.py | 24 ++++++++++++++++++++++++
 tests/test_decoding.py | 15 +++++++++++++++
 4 files changed, 71 insertions(+), 8 deletions(-)

diff --git a/httpcore/__init__.py b/httpcore/__init__.py
index 9824e854..7786c259 100644
--- a/httpcore/__init__.py
+++ b/httpcore/__init__.py
@@ -9,6 +9,9 @@ from .exceptions import (
     ResponseClosed,
     StreamConsumed,
     Timeout,
+    DeflateDecodingError,
+    GzipDecodingError,
+    BrotliDecodingError,
 )
 from .http11 import HTTP11Connection
 from .sync import SyncClient, SyncConnectionPool
diff --git a/httpcore/decoders.py b/httpcore/decoders.py
index b56745c4..dc56e131 100644
--- a/httpcore/decoders.py
+++ b/httpcore/decoders.py
@@ -12,6 +12,9 @@ except ImportError:  # pragma: nocover
     brotli = None
 
 
+import httpcore.exceptions
+
+
 class Decoder:
     def decode(self, data: bytes) -> bytes:
         raise NotImplementedError()  # pragma: nocover
@@ -39,10 +42,16 @@ class DeflateDecoder(Decoder):
         self.decompressor = zlib.decompressobj(-zlib.MAX_WBITS)
 
     def decode(self, data: bytes) -> bytes:
-        return self.decompressor.decompress(data)
+        try:
+            return self.decompressor.decompress(data)
+        except zlib.error as exc:
+            raise httpcore.exceptions.DeflateDecodingError from exc
 
     def flush(self) -> bytes:
-        return self.decompressor.flush()
+        try:
+            return self.decompressor.flush()
+        except zlib.error as exc:
+            raise httpcore.exceptions.DeflateDecodingError from exc
 
 
 class GZipDecoder(Decoder):
@@ -56,10 +65,16 @@ class GZipDecoder(Decoder):
         self.decompressor = zlib.decompressobj(zlib.MAX_WBITS | 16)
 
     def decode(self, data: bytes) -> bytes:
-        return self.decompressor.decompress(data)
+        try:
+            return self.decompressor.decompress(data)
+        except zlib.error as exc:
+            raise httpcore.exceptions.GzipDecodingError from exc
 
     def flush(self) -> bytes:
-        return self.decompressor.flush()
+        try:
+            return self.decompressor.flush()
+        except zlib.error as exc:
+            raise httpcore.exceptions.GzipDecodingError from exc
 
 
 class BrotliDecoder(Decoder):
@@ -77,16 +92,22 @@ class BrotliDecoder(Decoder):
         self.decompressor = brotli.Decompressor()
 
     def decode(self, data: bytes) -> bytes:
-        return self.decompressor.decompress(data)
+        try:
+            return self.decompressor.decompress(data)
+        except brotli.Error as exc:
+            raise httpcore.exceptions.BrotliDecodingError from exc
 
     def flush(self) -> bytes:
-        self.decompressor.finish()
-        return b""
+        try:
+            self.decompressor.finish()
+            return b""
+        except brotli.Error as exc:
+            raise httpcore.exceptions.BrotliDecodingError from exc
 
 
 class MultiDecoder(Decoder):
     """
-    Handle the case where mutliple encodings have been applied.
+    Handle the case where mutiple encodings have been applied.
     """
 
     def __init__(self, children: typing.Sequence[Decoder]) -> None:
diff --git a/httpcore/exceptions.py b/httpcore/exceptions.py
index 30814332..a85f128f 100644
--- a/httpcore/exceptions.py
+++ b/httpcore/exceptions.py
@@ -40,3 +40,27 @@ class ResponseClosed(Exception):
     Attempted to read or stream response content, but the request has been
     closed without loading the body.
     """
+
+
+class DecodingError(Exception):
+    """
+    Decoding of the response failed.
+    """
+
+
+class DeflateDecodingError(DecodingError):
+    """
+    Decoding of the response using deflate failed.
+    """
+
+
+class GzipDecodingError(DecodingError):
+    """
+    Decoding of the response using gzip failed.
+    """
+
+
+class BrotliDecodingError(DecodingError):
+    """
+    Decoding of the response using brotli failed.
+    """
diff --git a/tests/test_decoding.py b/tests/test_decoding.py
index f08c0abf..c4a17466 100644
--- a/tests/test_decoding.py
+++ b/tests/test_decoding.py
@@ -77,3 +77,18 @@ async def test_streaming():
     response = httpcore.Response(200, headers=headers, body=compress(body))
     assert not hasattr(response, "body")
     assert await response.read() == body
+
+
+@pytest.mark.parametrize(
+    'header_value, expected_exception',
+    [
+        (b"deflate", httpcore.DeflateDecodingError),
+        (b"gzip", httpcore.GzipDecodingError),
+        (b"br", httpcore.BrotliDecodingError),
+    ])
+def test_decoding_errors(header_value, expected_exception):
+    headers = [(b"Content-Encoding", header_value)]
+    body = b"test 123"
+    compressed_body = brotli.compress(body)[3:]
+    with pytest.raises(expected_exception):
+        response = httpcore.Response(200, headers=headers, body=compressed_body)