From d448ee8ff9b9d3204af5d714b06d568f6c0f2611 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Wed, 5 Aug 2026 17:56:33 -0600 Subject: [PATCH] fix(downloader): decode content-encoded bodies instead of writing raw Reading .raw skips urllib3's decoding, so brotli-served subtitles landed compressed and converted to empty files. Also drops a partial rather than Range-resuming an encoded body, since a Range indexes the encoded bytes. --- tests/downloaders/test_content_encoding.py | 110 +++++++++++++++++++++ unshackle/core/downloaders/requests.py | 33 ++++++- unshackle/core/utilities.py | 8 +- 3 files changed, 141 insertions(+), 10 deletions(-) create mode 100644 tests/downloaders/test_content_encoding.py diff --git a/tests/downloaders/test_content_encoding.py b/tests/downloaders/test_content_encoding.py new file mode 100644 index 00000000..fc1d97e9 --- /dev/null +++ b/tests/downloaders/test_content_encoding.py @@ -0,0 +1,110 @@ +import gzip +import socket +import threading + +import pytest +import requests as rq + +from unshackle.core.downloaders.requests import _is_content_encoded, download + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + (None, False), + ("", False), + (" ", False), + ("identity", False), + ("IDENTITY", False), + ("identity, identity", False), + ("gzip", True), + ("GZIP", True), + (" gzip ", True), + ("x-gzip", True), + ("br", True), + ("zstd", True), + ("deflate", True), + ("compress", True), + ("dcb", True), + ("gzip, br", True), + ("identity, gzip", True), + ("gzip, identity", True), + ("gzip,,", True), + ("some-future-coding", True), + ], +) +def test_is_content_encoded(value, expected): + assert _is_content_encoded(value) is expected + + +PLAIN = b"WEBVTT\n\n" + b"".join(f"{i:05d} cue text\n".encode() for i in range(2000)) + + +class _EncodingServer: + """Serves one gzip body, honouring Range against the encoded representation.""" + + def __init__(self): + self.encoded = gzip.compress(PLAIN) + self.sock = socket.socket() + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.sock.bind(("127.0.0.1", 0)) + self.sock.listen(8) + self.url = f"http://127.0.0.1:{self.sock.getsockname()[1]}/sub.vtt" + threading.Thread(target=self._serve, daemon=True).start() + + def _serve(self): + while True: + try: + conn, _ = self.sock.accept() + except OSError: + return + try: + request = conn.recv(65536).decode("latin1") + start = next( + ( + int(line.split("=", 1)[1].split("-")[0]) + for line in request.split("\r\n") + if line.lower().startswith("range:") + ), + None, + ) + if start is None: + body = self.encoded + head = f"HTTP/1.1 200 OK\r\nContent-Length: {len(body)}\r\n" + else: + body = self.encoded[start:] + head = ( + f"HTTP/1.1 206 Partial Content\r\nContent-Length: {len(body)}\r\n" + f"Content-Range: bytes {start}-{len(self.encoded) - 1}/{len(self.encoded)}\r\n" + ) + conn.sendall( + (head + "Content-Encoding: gzip\r\nAccept-Ranges: bytes\r\nConnection: close\r\n\r\n").encode() + + body + ) + except Exception: + pass + finally: + conn.close() + + def close(self): + self.sock.close() + + +@pytest.fixture +def encoding_server(): + server = _EncodingServer() + yield server + server.close() + + +def test_encoded_body_is_decoded_on_disk(tmp_path, encoding_server): + save_path = tmp_path / "sub.vtt" + list(download(url=encoding_server.url, save_path=save_path, session=rq.Session())) + assert save_path.read_bytes() == PLAIN + + +def test_encoded_body_discards_a_partial_instead_of_range_resuming(tmp_path, encoding_server): + save_path = tmp_path / "sub.vtt" + save_path.with_name("sub.vtt.!dev").write_bytes(PLAIN[:5000]) + list(download(url=encoding_server.url, save_path=save_path, session=rq.Session())) + assert save_path.read_bytes() == PLAIN diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 97aa3411..8f793b1c 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -161,6 +161,18 @@ def _is_rnet_session(session: Any) -> bool: return isinstance(session, RnetSession) +def _is_content_encoded(value: Optional[str]) -> bool: + """ + Check a Content-Encoding value for any coding other than absent or `identity`. + + Reading .raw skips urllib3's decoding, so an encoded body has to go through + iter_content or the compressed bytes land on disk as the payload. Codings + urllib3 cannot decode still count: an allowlist fails open on whatever token + a server sends next, and the cost of guessing wrong is a corrupt file. + """ + return any(coding.strip() not in ("", "identity") for coding in (value or "").lower().split(",")) + + def _probe_ranged(url: str, session: Any, **kwargs: Any) -> tuple[int, bool]: headers = {**(kwargs.get("headers") or {}), "Range": "bytes=0-0"} rest = {k: v for k, v in kwargs.items() if k != "headers"} @@ -175,8 +187,7 @@ def _probe_ranged(url: str, session: Any, **kwargs: Any) -> tuple[int, bool]: try: if resp.status_code != 206: return 0, False - ce = (resp.headers.get("Content-Encoding") or resp.headers.get("content-encoding") or "").lower() - if ce in ("gzip", "deflate", "br"): + if _is_content_encoded(resp.headers.get("Content-Encoding") or resp.headers.get("content-encoding")): return 0, False content_range = resp.headers.get("Content-Range") or resp.headers.get("content-range") or "" total = content_range.rsplit("/", 1)[-1].strip() @@ -306,7 +317,8 @@ def download( Download a file with optimized I/O. Supports both requests.Session and RnetSession for TLS fingerprinting. - Uses raw socket reads for requests.Session and native rnet streaming for RnetSession. + RnetSession streams natively; requests.Session reads the raw socket, except + on a content-encoded body, which has to go through iter_content to be decoded. Yields the following download status updates while chunks are downloading: @@ -393,16 +405,27 @@ def download( resume_offset = 0 if part_mode and stream.status_code != 206: raise IOError(f"expected 206 for ranged part, got {stream.status_code}") + content_encoded = False if use_rnet: content_length = stream.content_length or 0 else: + content_encoded = _is_content_encoded(stream.headers.get("Content-Encoding")) try: content_length = int(stream.headers.get("Content-Length", "0")) - if stream.headers.get("Content-Encoding", "").lower() in ["gzip", "deflate", "br"]: + if content_encoded: content_length = 0 except ValueError: content_length = 0 + if resumed and content_encoded: + try: + stream.close() + except Exception: + pass + tmp_file.unlink(missing_ok=True) + resume_offset = 0 + continue + limiter = _speed_limiter chunk_size = _adaptive_chunk_size(content_length) if limiter: @@ -434,7 +457,7 @@ def download( if use_rnet: chunks = stream.stream() - elif use_raw: + elif use_raw and not content_encoded: _read1 = getattr(stream.raw, "read1", None) if claimed is not None and RACER_READ1 else None if _read1 is not None: chunks = iter(lambda: _read1(chunk_size), b"") diff --git a/unshackle/core/utilities.py b/unshackle/core/utilities.py index 90c5e500..285a0520 100644 --- a/unshackle/core/utilities.py +++ b/unshackle/core/utilities.py @@ -454,9 +454,9 @@ def try_ensure_utf8(data: bytes) -> bytes: """ Try to ensure that the given data is encoded in UTF-8. - Automatically decompresses gzip/deflate/zlib data before encoding detection. - This handles cases where HTTP responses are saved with raw Content-Encoding - (e.g., when decode_content=False is used for performance). + Gzip or zlib compressed input is decompressed first, detected by magic bytes. + Callers hand in bytes read from files and from responses they fetched + themselves, so there is no Content-Encoding header left to consult. Parameters: data: Input data that may or may not yet be UTF-8 or another encoding. @@ -464,13 +464,11 @@ def try_ensure_utf8(data: bytes) -> bytes: Returns the input data encoded in UTF-8 if successful. If unable to detect the encoding of the input data, then the original data is returned as-received. """ - # Decompress gzip data (magic bytes: 1f 8b) if data[:2] == b"\x1f\x8b": try: data = gzip.decompress(data) except Exception: pass - # Decompress raw deflate/zlib data (common zlib headers: 78 01, 78 5e, 78 9c, 78 da) elif data[:1] == b"\x78" and len(data) > 1 and data[1:2] in (b"\x01", b"\x5e", b"\x9c", b"\xda"): try: data = zlib.decompress(data)