diff --git a/tests/downloaders/test_content_encoding.py b/tests/downloaders/test_content_encoding.py index fc1d97e9..7aa58d7c 100644 --- a/tests/downloaders/test_content_encoding.py +++ b/tests/downloaders/test_content_encoding.py @@ -1,4 +1,5 @@ import gzip +import importlib import socket import threading @@ -108,3 +109,48 @@ def test_encoded_body_discards_a_partial_instead_of_range_resuming(tmp_path, enc 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 + + +class _FakeRnetResponse: + def __init__(self, status_code, headers, body): + self.status_code = status_code + self.headers = headers + self.content_length = len(body) + self._body = body + + def raise_for_status(self): + pass + + def stream(self): + yield self._body + + def close(self): + pass + + +class _FakeRnetSession: + """Duck-typed rnet stand-in: an encoded 206 for a Range request (the resume the + fail-closed guard must reject), an encoded 200 for the clean restart.""" + + def __init__(self, body): + self.body = body + self.range_requests = 0 + + def get(self, url, stream=True, **kwargs): + headers = kwargs.get("headers") or {} + if any(k.lower() == "range" for k in headers): + self.range_requests += 1 + return _FakeRnetResponse(206, {"Content-Encoding": "zstd"}, b"\x00mid-stream-slice\x00") + return _FakeRnetResponse(200, {"Content-Encoding": "zstd"}, self.body) + + +def test_rnet_encoded_body_discards_a_partial_instead_of_range_resuming(tmp_path, monkeypatch): + dl = importlib.import_module("unshackle.core.downloaders.requests") + monkeypatch.setattr(dl, "_is_rnet_session", lambda s: isinstance(s, _FakeRnetSession)) + session = _FakeRnetSession(PLAIN) + save_path = tmp_path / "sub.vtt" + save_path.with_name("sub.vtt.!dev").write_bytes(PLAIN[:5000]) + list(download(url="http://127.0.0.1:1/sub.vtt", save_path=save_path, session=session)) + # the encoded 206 must be dropped (never appended); the restart fetches the whole body + assert session.range_requests == 1 + assert save_path.read_bytes() == PLAIN diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 53bbc169..b7a6f40c 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -714,12 +714,11 @@ def download( # server ignored the slice's Range and sent the whole parent; writing that # as the segment would silently corrupt the merge, so fail the attempt raise IOError(f"expected 206 for byte-range segment, got {stream.status_code}") - content_encoded = False if use_rnet: - content_length = stream.content_length or 0 - ce = (stream.headers.get("Content-Encoding") or stream.headers.get("content-encoding") or "").lower() - if ce in ("gzip", "deflate", "br"): - content_length = 0 + content_encoded = _is_content_encoded( + stream.headers.get("Content-Encoding") or stream.headers.get("content-encoding") + ) + content_length = 0 if content_encoded else (stream.content_length or 0) else: content_encoded = _is_content_encoded(stream.headers.get("Content-Encoding")) try: