diff options
Diffstat (limited to 'mitmproxy/net/http/encoding.py')
| -rw-r--r-- | mitmproxy/net/http/encoding.py | 24 |
1 files changed, 22 insertions, 2 deletions
diff --git a/mitmproxy/net/http/encoding.py b/mitmproxy/net/http/encoding.py index 8cb96e5c..16d399ca 100644 --- a/mitmproxy/net/http/encoding.py +++ b/mitmproxy/net/http/encoding.py @@ -9,6 +9,7 @@ from io import BytesIO import gzip import zlib import brotli +import zstandard as zstd from typing import Union, Optional, AnyStr # noqa @@ -52,7 +53,7 @@ def decode( decoded = custom_decode[encoding](encoded) except KeyError: decoded = codecs.decode(encoded, encoding, errors) - if encoding in ("gzip", "deflate", "br"): + if encoding in ("gzip", "deflate", "br", "zstd"): _cache = CachedDecode(encoded, encoding, errors, decoded) return decoded except TypeError: @@ -93,7 +94,7 @@ def encode(decoded: Optional[str], encoding: str, errors: str='strict') -> Optio encoded = custom_encode[encoding](decoded) except KeyError: encoded = codecs.encode(decoded, encoding, errors) - if encoding in ("gzip", "deflate", "br"): + if encoding in ("gzip", "deflate", "br", "zstd"): _cache = CachedDecode(encoded, encoding, errors, decoded) return encoded except TypeError: @@ -140,6 +141,23 @@ def encode_brotli(content: bytes) -> bytes: return brotli.compress(content) +def decode_zstd(content: bytes) -> bytes: + if not content: + return b"" + zstd_ctx = zstd.ZstdDecompressor() + try: + return zstd_ctx.decompress(content) + except zstd.ZstdError: + # If the zstd stream is streamed without a size header, + # try decoding with a 10MiB output buffer + return zstd_ctx.decompress(content, max_output_size=10 * 2**20) + + +def encode_zstd(content: bytes) -> bytes: + zstd_ctx = zstd.ZstdCompressor() + return zstd_ctx.compress(content) + + def decode_deflate(content: bytes) -> bytes: """ Returns decompressed data for DEFLATE. Some servers may respond with @@ -170,6 +188,7 @@ custom_decode = { "gzip": decode_gzip, "deflate": decode_deflate, "br": decode_brotli, + "zstd": decode_zstd, } custom_encode = { "none": identity, @@ -177,6 +196,7 @@ custom_encode = { "gzip": encode_gzip, "deflate": encode_deflate, "br": encode_brotli, + "zstd": encode_zstd, } __all__ = ["encode", "decode"] |