Merge pull request #1306 from mitmproxy/message-body-encoding

Improve Message Body Encoding

1 files changed, 63 insertions, 34 deletions

diff --git a/netlib/encoding.py b/netlib/encoding.py
index 98502451..8b67b543 100644
--- a/netlib/encoding.py
+++ b/netlib/encoding.py
@@ -1,39 +1,62 @@
 """
-    Utility functions for decoding response bodies.
+Utility functions for decoding response bodies.
 """
 from __future__ import absolute_import
+
+import codecs
 from io import BytesIO
 import gzip
 import zlib
 
+from typing import Union  # noqa
+
 
-ENCODINGS = {"identity", "gzip", "deflate"}
+def decode(obj, encoding, errors='strict'):
+    # type: (Union[str, bytes], str) -> Union[str, bytes]
+    """
+    Decode the given input object
 
+    Returns:
+        The decoded value
 
-def decode(e, content):
-    if not isinstance(content, bytes):
-        return None
-    encoding_map = {
-        "identity": identity,
-        "gzip": decode_gzip,
-        "deflate": decode_deflate,
-    }
-    if e not in encoding_map:
-        return None
-    return encoding_map[e](content)
+    Raises:
+        ValueError, if decoding fails.
+    """
+    try:
+        try:
+            return custom_decode[encoding](obj)
+        except KeyError:
+            return codecs.decode(obj, encoding, errors)
+    except Exception as e:
+        raise ValueError("{} when decoding {} with {}".format(
+            type(e).__name__,
+            repr(obj)[:10],
+            repr(encoding),
+        ))
+
+
+def encode(obj, encoding, errors='strict'):
+    # type: (Union[str, bytes], str) -> Union[str, bytes]
+    """
+    Encode the given input object
 
+    Returns:
+        The encoded value
 
-def encode(e, content):
-    if not isinstance(content, bytes):
-        return None
-    encoding_map = {
-        "identity": identity,
-        "gzip": encode_gzip,
-        "deflate": encode_deflate,
-    }
-    if e not in encoding_map:
-        return None
-    return encoding_map[e](content)
+    Raises:
+        ValueError, if encoding fails.
+    """
+    try:
+        try:
+            return custom_encode[encoding](obj)
+        except KeyError:
+            return codecs.encode(obj, encoding, errors)
+    except Exception as e:
+        raise ValueError("{} when encoding {} with {}".format(
+            type(e).__name__,
+            repr(obj)[:10],
+            repr(encoding),
+        ))
 
 
 def identity(content):
@@ -46,10 +69,7 @@ def identity(content):
 
 def decode_gzip(content):
     gfile = gzip.GzipFile(fileobj=BytesIO(content))
-    try:
-        return gfile.read()
-    except (IOError, EOFError):
-        return None
+    return gfile.read()
 
 
 def encode_gzip(content):
@@ -70,12 +90,9 @@ def decode_deflate(content):
         http://bugs.python.org/issue5784
     """
     try:
-        try:
-            return zlib.decompress(content)
-        except zlib.error:
-            return zlib.decompress(content, -15)
+        return zlib.decompress(content)
     except zlib.error:
-        return None
+        return zlib.decompress(content, -15)
 
 
 def encode_deflate(content):
@@ -84,4 +101,16 @@ def encode_deflate(content):
     """
     return zlib.compress(content)
 
-__all__ = ["ENCODINGS", "encode", "decode"]
+
+custom_decode = {
+    "identity": identity,
+    "gzip": decode_gzip,
+    "deflate": decode_deflate,
+}
+custom_encode = {
+    "identity": identity,
+    "gzip": encode_gzip,
+    "deflate": encode_deflate,
+}
+
+__all__ = ["encode", "decode"]