From fe7b3bb651a618b5a333ab4bf1001a7df794277e Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sun, 6 Sep 2026 20:04:25 -0600 Subject: [PATCH] test(downloader): cover session state carried into download children Assert that verify, params, cert and auth reach a spawned child, and that a mounted adapter the child cannot reproduce falls back to the single-process path. Documents --download-processes in the CLI reference. --- docs/guide/cli-reference.md | 2 +- tests/downloaders/test_multiprocess.py | 51 ++++++++- tests/downloaders/test_session_spec.py | 148 ++++++++++++++++++++++++- 3 files changed, 194 insertions(+), 7 deletions(-) diff --git a/docs/guide/cli-reference.md b/docs/guide/cli-reference.md index 5684c719..33866c55 100644 --- a/docs/guide/cli-reference.md +++ b/docs/guide/cli-reference.md @@ -206,7 +206,7 @@ order the service used, then renumbers the episodes into the order you asked for |---|---|---| | `--workers` | downloader default | Per-track download threads. | | `--adaptive-workers` | off | Scale per-track segment workers dynamically (up to the `--workers` cap) based on measured CDN throughput and errors. | -| `--download-processes` | `1` | Split large segment batches (24+) across this many download processes to exceed the single-process throughput cap on multi-gigabit connections. Ignored while a speed limit is set: the cap is one shared budget, which extra processes cannot share, so the download stays in a single process. | +| `--download-processes` | `1` | Split large segment batches (24+) across this many download processes to exceed the single-process throughput cap on multi-gigabit connections. Ignored while a speed limit is set: the cap is one shared budget, which extra processes cannot share, so the download stays in a single process. Also ignored when the service HTTP session carries state a child process cannot rebuild, such as a custom TLS adapter. | | `--continue-downloads` | off | Keep completed segment files across runs and resume a previously failed download. One-off enable of the [`continue_downloads`](../reference/configuration/download.md#continue_downloads) config option. | | `--downloads` | `1` | Number of tracks downloaded concurrently. | | `--speed-limit` | unlimited | Cap total download speed across all threads and tracks, e.g. `500k`, `5M`, `1.5G` or plain bytes/sec. Values are bytes, not bits (`5M` = 5.0 MB/s). `off` disables a configured limit. | diff --git a/tests/downloaders/test_multiprocess.py b/tests/downloaders/test_multiprocess.py index d5688063..3c56c00c 100644 --- a/tests/downloaders/test_multiprocess.py +++ b/tests/downloaders/test_multiprocess.py @@ -25,8 +25,10 @@ import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest +from requests import Session from unshackle.core.constants import DownloadCancelled +from unshackle.core.utils.sslciphers import SSLCiphers dl = importlib.import_module("unshackle.core.downloaders.requests") @@ -44,8 +46,15 @@ class _Handler(BaseHTTPRequestHandler): delay = getattr(server, "delay", 0.0) if delay: time.sleep(delay) + path, _, query = self.path.partition("?") + required = getattr(server, "require_query", "") + if required and query != required: + self.send_response(403) + self.send_header("Content-Length", "0") + self.end_headers() + return try: - index = int(self.path.rsplit("/", 1)[1].split(".")[0]) + index = int(path.rsplit("/", 1)[1].split(".")[0]) except ValueError: self.send_response(404) self.send_header("Content-Length", "0") @@ -72,6 +81,7 @@ class _Server(ThreadingHTTPServer): def server(): srv = _Server(("127.0.0.1", 0), _Handler) srv.delay = 0.0 + srv.require_query = "" thread = threading.Thread(target=srv.serve_forever, daemon=True) thread.start() yield srv @@ -180,3 +190,42 @@ def test_speed_limit_forces_single_process(server, tmp_path, monkeypatch): assert len(files) == n for f in files: assert f.read_bytes() == seg_body(int(f.stem.split("_")[1])) + + +def test_children_carry_session_params(server, tmp_path): + """Services keep their access token in session.params. A child that loses it gets a 403 per segment.""" + server.require_query = "at=auth-token" + session = Session() + session.params = {"at": "auth-token"} + n = dl.MP_MIN_SEGMENTS + urls = [{"url": url(server, f"/seg/{i}.bin")} for i in range(n)] + + for _ in dl.requests( + urls, output_dir=tmp_path, filename="seg_{i:04}.bin", session=session, max_workers=4, processes=2 + ): + pass + + files = sorted(tmp_path.glob("seg_*.bin")) + assert len(files) == n + for f in files: + assert f.read_bytes() == seg_body(int(f.stem.split("_")[1])) + + +def test_service_mounted_adapter_forces_single_process(server, tmp_path, monkeypatch): + """A child mounts a plain adapter, so a session carrying SSLCiphers has to stay in this process.""" + + def no_mp(**kwargs): + raise AssertionError("multiprocess fan-out engaged for a session with a service-mounted adapter") + + monkeypatch.setattr(dl, "download_multiprocess", no_mp) + session = Session() + session.mount("https://", SSLCiphers(security_level=2)) + n = dl.MP_MIN_SEGMENTS + urls = [{"url": url(server, f"/seg/{i}.bin")} for i in range(n)] + + for _ in dl.requests( + urls, output_dir=tmp_path, filename="seg_{i:04}.bin", session=session, max_workers=4, processes=2 + ): + pass + + assert len(list(tmp_path.glob("seg_*.bin"))) == n diff --git a/tests/downloaders/test_session_spec.py b/tests/downloaders/test_session_spec.py index cf41156c..ee42d08e 100644 --- a/tests/downloaders/test_session_spec.py +++ b/tests/downloaders/test_session_spec.py @@ -1,20 +1,29 @@ """Session spec fidelity across the multiprocess spawn boundary. ``build_session_spec``/``rebuild_session`` are how a spawned child gets the parent's session -state. The cross-process half needs a live CDN to observe, so these tests pin the two halves -that can be checked deterministically: the spec is picklable (spawn requires it), and a -rebuilt session carries the same cookie domains as the parent. A spec that flattens cookies -to name/value rebuilds every cookie scoped to localhost, so a cookie-authed CDN sees none of -them and 403s only when ``--download-processes`` is above 1. +state. The cross-process half needs a live CDN to observe, so these tests pin the halves that +can be checked deterministically: the spec is picklable (spawn requires it), and a rebuilt +session carries the parent's cookie domains, query params, and TLS settings. A spec that +flattens cookies to name/value rebuilds every cookie scoped to localhost, so a cookie-authed +CDN sees none of them and 403s only when ``--download-processes`` is above 1. State the child +cannot reproduce (a service-mounted adapter, an unpicklable auth callable) must give None, +which drops the caller back to a single process. """ import pickle +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import pytest from requests import Session +from requests.adapters import HTTPAdapter +from requests.auth import HTTPBasicAuth from unshackle.core.downloaders.requests import build_session_spec, rebuild_session +from unshackle.core.service import Service from unshackle.core.session import RnetSession from unshackle.core.session import session as make_session +from unshackle.core.utils.sslciphers import SSLCiphers def rnet_session_with_cookies() -> RnetSession: @@ -79,6 +88,50 @@ def test_rnet_spec_without_impersonate_is_not_rebuildable() -> None: assert build_session_spec(RnetSession()) is None +class _CookieHandler(BaseHTTPRequestHandler): + """Sets one cookie on /set and echoes the Cookie header it receives on /echo.""" + + def do_GET(self) -> None: + body = b"" + self.send_response(200) + if self.path == "/set": + self.send_header("Set-Cookie", "sid=server-set; Path=/") + else: + body = (self.headers.get("Cookie") or "").encode() + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): # keep the test output quiet + pass + + +@pytest.fixture +def cookie_server(): + srv = ThreadingHTTPServer(("127.0.0.1", 0), _CookieHandler) + srv.daemon_threads = True + thread = threading.Thread(target=srv.serve_forever, daemon=True) + thread.start() + yield f"http://127.0.0.1:{srv.server_address[1]}" + srv.shutdown() + thread.join(timeout=5) + + +def test_rnet_spec_carries_server_set_cookies(cookie_server) -> None: + """A server Set-Cookie lands only in the rnet client jar, and the spec has to carry it across.""" + ns = make_session("Chrome131") + ns.get(f"{cookie_server}/set") + + spec = build_session_spec(ns) + assert spec is not None + assert spec["origin_cookies"] == {cookie_server: "sid=server-set"} + + rebuilt = rebuild_session(pickle.loads(pickle.dumps(spec)), 1) + assert rebuilt is not None + + assert "sid=server-set" in rebuilt.get(f"{cookie_server}/echo").text + + def test_rebuilt_requests_session_pool_is_sized_to_worker_count() -> None: """A child rebuilding a requests session must size its pool, not inherit urllib3's default 10.""" spec = build_session_spec(Session()) @@ -123,3 +176,88 @@ def test_requests_spec_round_trips_cookie_domains() -> None: assert rebuilt.cookies.get_dict(domain="cdn.example.com") == {"cdn_token": "abc"} assert rebuilt.cookies.get_dict(domain="api.example.com") == {"api_token": "def"} + + +def test_requests_spec_carries_verify() -> None: + """A service that turns certificate verification off must not get it turned back on.""" + rs = Session() + rs.verify = False + + spec = build_session_spec(rs) + assert spec is not None + rebuilt = rebuild_session(pickle.loads(pickle.dumps(spec)), 1) + assert rebuilt is not None + + assert rebuilt.verify is False + + +def test_requests_spec_carries_params() -> None: + """Services hold an auth token in session.params; without it every segment request is unauthed.""" + rs = Session() + rs.params = {"at": "auth-token"} + + spec = build_session_spec(rs) + assert spec is not None + rebuilt = rebuild_session(pickle.loads(pickle.dumps(spec)), 1) + assert rebuilt is not None + + assert dict(rebuilt.params) == {"at": "auth-token"} + + +def test_requests_spec_carries_cert_and_auth() -> None: + rs = Session() + rs.cert = ("/tmp/client.pem", "/tmp/client.key") + rs.auth = HTTPBasicAuth("user", "pass") + + spec = build_session_spec(rs) + assert spec is not None + rebuilt = rebuild_session(pickle.loads(pickle.dumps(spec)), 1) + assert rebuilt is not None + + assert rebuilt.cert == ("/tmp/client.pem", "/tmp/client.key") + assert rebuilt.auth == HTTPBasicAuth("user", "pass") + + +def test_requests_spec_is_picklable() -> None: + rs = Session() + rs.verify = False + rs.params = {"at": "auth-token"} + rs.cookies.set("cdn_token", "abc", domain="cdn.example.com") + + spec = build_session_spec(rs) + assert spec is not None + + assert pickle.loads(pickle.dumps(spec))["params"] == {"at": "auth-token"} + + +def test_requests_spec_with_unpicklable_auth_is_not_rebuildable() -> None: + """auth takes any callable. One that spawn cannot pickle has to fall back, not kill the child.""" + rs = Session() + rs.auth = lambda request: request + + assert build_session_spec(rs) is None + + +def test_service_session_stays_rebuildable() -> None: + """The framework default adapter is reproducible. Failing it closed would disable the fan-out for every service.""" + spec = build_session_spec(Service.get_session()) + + assert spec is not None + + +def test_mounted_ssl_ciphers_is_not_rebuildable() -> None: + """SSLCiphers swaps the SSL context, so a child mounting a plain adapter negotiates a different cipher list.""" + rs = Session() + rs.mount("https://cdn.example.com", SSLCiphers(security_level=2)) + + assert build_session_spec(rs) is None + + +def test_rebuilt_requests_session_mounts_a_plain_adapter() -> None: + """The spec only ever describes a plain adapter, so the child must not appear to hold more than that.""" + spec = build_session_spec(Session()) + assert spec is not None + rebuilt = rebuild_session(spec, 4) + assert rebuilt is not None + + assert type(rebuilt.get_adapter("https://example.com")) is HTTPAdapter