From 91db1fa0b1fc37e69fa65e34a51bc76cc6259396 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Thu, 24 Sep 2026 10:18:03 -0600 Subject: [PATCH] feat(remote): keep slow remote session requests alive with a heartbeat Titles, tracks and licence stream a newline every 30s until the body is ready, so a slow service outlasts the client timeout and proxy idle limits; add a per-server timeout. --- docs/dev/rest-api/endpoints.md | 22 ++++- docs/dev/rest-api/errors.md | 15 ++- docs/dev/rest-api/remote-sessions.md | 22 ++++- docs/reference/configuration/services.md | 8 +- tests/remote/unit/test_extract_pssh_strict.py | 2 + tests/remote/unit/test_input_bridge.py | 9 ++ tests/remote/unit/test_remote_client.py | 40 ++++++++ tests/remote/unit/test_routes.py | 93 ++++++++++++++++++- tests/remote/unit/test_segment_filter.py | 2 + .../unit/test_server_cdm_absorb_routes.py | 4 + .../unit/test_server_cdm_manifest_kid.py | 2 + tests/remote/unit/test_server_cdm_multikey.py | 2 + .../unit/test_server_cdm_type_override.py | 2 + tests/remote/unit/test_session_store.py | 11 +++ unshackle/core/api/handlers.py | 36 +++++-- unshackle/core/api/input_bridge.py | 4 +- unshackle/core/api/routes.py | 63 ++++++++++++- unshackle/core/api/session_store.py | 16 +++- unshackle/core/api/stats.py | 12 ++- unshackle/core/remote_service.py | 50 +++++++--- unshackle/core/services.py | 4 +- 21 files changed, 377 insertions(+), 42 deletions(-) diff --git a/docs/dev/rest-api/endpoints.md b/docs/dev/rest-api/endpoints.md index 86fc3f8f..b40c9a6d 100644 --- a/docs/dev/rest-api/endpoints.md +++ b/docs/dev/rest-api/endpoints.md @@ -15,7 +15,7 @@ Start the server with [`unshackle serve`](../../guide/cli-reference.md). By defa - **Authentication.** Every request except `GET /api/health` requires the `X-Secret-Key` header when you configure an API key. `--no-key` disables the check entirely. See [Authentication](authentication.md). - **Content type.** Request bodies are JSON (`Content-Type: application/json`). All success responses are JSON unless this page documents a `204 No Content`. - **CORS.** Every response carries permissive CORS headers, and the server answers `OPTIONS` preflight requests automatically. -- **Compression.** JSON responses of 256 bytes or more are gzip-compressed when the client sends `Accept-Encoding: gzip`. +- **Compression.** JSON responses of 256 bytes or more are gzip-compressed when the client sends `Accept-Encoding: gzip`. A [slow response](#slow-responses) is not compressed. - **Service allowlist.** The effective allowlist for your API key filters the service-facing endpoints (the intersection of the global `serve.services` list and your per-key list). The server treats services you cannot use as unknown. !!! warning "Error responses have two shapes" @@ -883,6 +883,22 @@ POST /api/session/{id}/keys/bad → report a server-vault content key that did n DELETE /api/session/{id} → tear down, harvest updated cache ``` +### Slow responses + +The `titles`, `tracks` and `license` routes can take minutes on a large catalogue. When the +answer is not ready after 30 seconds, the route sends status `200` and the headers. It then +sends a newline every 30 seconds until the JSON body is ready. JSON parsers ignore the leading +newlines. This keeps the connection open through a reverse proxy that closes an idle request. + +!!! warning "A late failure has status `200`" + The server sent the status before the work ended. A failure after the first 30 seconds + arrives as the [standard error body](errors.md#the-standard-error-shape) with status `200`. + Check `status` in the body, not only the HTTP status. The body has no HTTP status field. + Use `error_code` and `retryable` to decide what to do. + +A slow response is not gzip-compressed. Gzip output cannot start before the body is ready, and +the newlines must go out at once. + ### `POST /api/session/create` Make a remote session for a service and title. Authentication runs asynchronously in the background. This call returns immediately with `status: "authenticating"`, and you then poll the prompt endpoint. The body accepts `service` and `title_id` (both required). It also accepts a broad set of optional keys, because the body allows `additionalProperties`. These are `credentials` (`{username, password, extra?}`), `cookies` (base64 of zlib-compressed Netscape cookie file), `proxy`, `no_proxy`, `profile`, `cache` (map of forwarded cache files, keyed by the file path relative to the service cache directory with `/` separators and no `.json` suffix, so a nested key such as `session_web/` lands in a subdirectory), `client_region`, `proxy_region`, `cdm_type`, `range_`, `vcodec`, `quality`, `best_available`, `dl_params`, `client`, and any service CLI options. `proxy_region` is the country code the client resolved `proxy` from; the server matches it against its own accounts. `client` is a freeform object the dashboard shows as sent (the CLI puts `version`, `code_hash`, `platform` and a redacted `argv` in it); the server ignores it above 4096 bytes of JSON. @@ -989,6 +1005,7 @@ The buffer keeps the newest 500 records. Only the service instance's own logger ### `GET /api/session/{session_id}/titles` Fetch the titles for an authenticated remote session. +A slow answer arrives as a [slow response](#slow-responses). === "Response `200`" @@ -1003,7 +1020,7 @@ Requires an authenticated remote session. Otherwise it returns `404 SESSION_NOT_ ### `POST /api/session/{session_id}/tracks` -Get tracks and chapters for a specific title within the remote session. Body: `{ "title_id": "..." }` (required). Unlike `list-tracks`, track objects here **include** download URLs, along with session headers/cookies and manifest data the client needs to download directly. +Get tracks and chapters for a specific title within the remote session. Body: `{ "title_id": "..." }` (required). Unlike `list-tracks`, track objects here **include** download URLs, along with session headers/cookies and manifest data the client needs to download directly. A slow answer arrives as a [slow response](#slow-responses). `manifests` holds each served DASH/ISM manifest once. A service can build an AdaptationSet of its own that the served manifest does not contain. `track_manifests` holds a one-AdaptationSet MPD for each such track, and the client re-parses that in place of the whole manifest. @@ -1071,6 +1088,7 @@ Get the HLS segments a service drops (ads, bumpers, dub cards). The server runs ### `POST /api/session/{session_id}/license` Get the content keys for the DRM. The `mode` field selects one of two modes. +A slow answer arrives as a [slow response](#slow-responses). **Proxy mode** (`mode: "proxy"`, the default) forwards the client's CDM `challenge` to the service's license endpoint and returns the raw license bytes: diff --git a/docs/dev/rest-api/errors.md b/docs/dev/rest-api/errors.md index 409906d5..3c8e9cd4 100644 --- a/docs/dev/rest-api/errors.md +++ b/docs/dev/rest-api/errors.md @@ -203,6 +203,18 @@ body that uses an integer `status` and no `error_code` or `timestamp`. (integer `status`, no `error_code`). The `GET /api/health` endpoint is exempt from authentication, so it never returns the key-error shape. +## Late failures with status `200` + +The remote session `titles`, `tracks` and `license` routes send status `200` and the headers +when the work takes longer than 30 seconds. They then send a newline every 30 seconds to keep +the connection open. See [Slow responses](endpoints.md#slow-responses). A failure after that +time arrives as the standard error body, but with status `200`. The body has no field for the +HTTP status that the error would otherwise have. `error_code` and `retryable` stay in the body. + +!!! warning "Check the body on these routes" + On these routes a `200` does not always mean success. Parse the body and treat + `"status": "error"` with an `error_code` as a failure. + ## Responses with no body Not every failure or completion has a JSON body. Some endpoints answer with @@ -257,7 +269,8 @@ For anyone writing a consumer of the API, a durable strategy is: !!! note "Reference consumer" unshackle's own remote-download client treats any status of `400` or above as - fatal: it logs `Server error []: ` and exits. It also wraps + fatal, and a `200` body with `"status": "error"` and an `error_code` too: it logs + `Server error []: ` and exits. It also wraps requests in a retry adapter (five attempts, backing off on `429`, `500`, `502`, `503`, and `504`). This is a reasonable template for your own client's error handling. diff --git a/docs/dev/rest-api/remote-sessions.md b/docs/dev/rest-api/remote-sessions.md index 3302791a..ebf5f7f3 100644 --- a/docs/dev/rest-api/remote-sessions.md +++ b/docs/dev/rest-api/remote-sessions.md @@ -277,8 +277,8 @@ generally never touch this directly, but it explains behaviour you might observe | Behavior | Value | Notes | |---|---|---| -| Idle session TTL | **300s** (5 min) default | Refreshed on every request to the session | -| Max concurrent sessions | **100** default | Oldest (least recently used) is evicted when full | +| Idle session TTL | **300s** (5 min) default | Refreshed on every request to the session and when a `titles` or `tracks` request ends. A remote session does not expire while a request works on it | +| Max concurrent sessions | **100** default | Oldest (least recently used) is evicted when full. A remote session that a request works on goes last | | Auth/input timeout | **600s** (10 min) | Sessions still authenticating or awaiting a prompt use this longer window instead of the TTL | | Cleanup sweep | every 60s | Expired sessions are removed and their input prompts cancelled | @@ -403,6 +403,9 @@ layer. - `InputBridge.request_input(prompt, timeout=600)` blocks the sync auth thread on a `threading.Event` until `submit_response()` or `cancel()` fires. - A timeout raises `TimeoutError` and marks the remote session `FAILED`. +- A prompt after authentication (in `get_titles()`, `get_tracks()` or licensing) raises + `RuntimeError` at once. The client polls for prompts only during authentication, so + no answer can come. - `AUTH_INPUT_TIMEOUT = 600.0` seconds. This is also the TTL granted to `AUTHENTICATING` / `PENDING_INPUT` sessions in the store. @@ -513,9 +516,18 @@ Source: `unshackle/core/remote_service.py`. This is the canonical consumer of th remote session API and a good template for any client. - **`RemoteClient.request`** sets `X-Secret-Key` and `User-Agent: unshackle/`, - uses a 120s timeout for `POST` and 30s for `GET`/`DELETE`, and treats any - `status_code >= 400` as fatal: it logs `Server error []: ` - and raises `SystemExit(1)`. + uses the server's `timeout` from `remote_services` (default 120s) as the read timeout, + and treats any `status_code >= 400` as fatal: it logs `Server error []: ` + and raises `SystemExit(1)`. A `200` body with `"status": "error"` is also fatal (see + the heartbeat below). +- **Heartbeat.** The `titles`, `tracks` and `license` routes send a newline every 30s while + the service works. The read timeout starts again at each byte, so a slow request stays + open for as long as the server works on it, and a reverse proxy does not close it as + idle. A route that answers within 30s sends the usual response. After 30s the status + is already `200`, so a failure after that time arrives as the usual error body with + status `200`. Your client must check `status` in the body. JSON parsers ignore the + leading newlines. The server does not gzip-compress a heartbeat response, because gzip + output cannot start before the body is ready. - **Retries.** The download-side HTTP session mounts an adapter with `Retry(total=5, backoff_factor=0.2, status_forcelist=[429, 500, 502, 503, 504])`. - **Flow.** `authenticate()` → `create` (+ poll `prompt` every 2s up to a 600s diff --git a/docs/reference/configuration/services.md b/docs/reference/configuration/services.md index 3778fa73..7be6f346 100644 --- a/docs/reference/configuration/services.md +++ b/docs/reference/configuration/services.md @@ -83,7 +83,7 @@ auth tokens in localStorage rather than in HTTP cookies need. Extraction is read Definitions of remote unshackle service servers, used by the `--remote` mode. Each entry is named by you (pick it with `--server`, or do not write that flag when you configure only one) and gives the server's `url` (required), an optional `api_key`, an optional `auth_headers` list, an -optional `server_cdm` boolean, and an optional `services` sub-dict of per-service local +optional `server_cdm` boolean, an optional `timeout` in seconds, and an optional `services` sub-dict of per-service local overrides such as `title_map`. Leave `server_cdm` unset and the client follows the server: it uses the server CDM for each @@ -97,6 +97,12 @@ keeps the one that works for the rest of the HTTP session. Names you give keep y are not repeated in the fallbacks, so unshackle tries `auth_headers: ["Authorization", "x-secret-key"]` as `Authorization`, `x-secret-key`, `X-Api-Key`. +`timeout` is the time in seconds that the client waits for data from the server. The +default is `120`. The server sends a heartbeat every 30 seconds while it gets titles, tracks +or licences, so a slow service does not cause a timeout. Increase `timeout` only when the +server runs an older version that sends no heartbeat. Give a number of seconds. With a +value below 30 the client can time out before the first heartbeat. `0` means the default. + In `--remote` mode unshackle turns the server's service list into synthetic CLI commands that operate against it, falling back to the tags in that `services` sub-dict when unshackle cannot fetch the list. Each synthetic command carries the server-side service's options and diff --git a/tests/remote/unit/test_extract_pssh_strict.py b/tests/remote/unit/test_extract_pssh_strict.py index 5b9a0f7c..2f371ef1 100644 --- a/tests/remote/unit/test_extract_pssh_strict.py +++ b/tests/remote/unit/test_extract_pssh_strict.py @@ -5,6 +5,7 @@ loop licensed a PlayReady-only track as Widevine by accident. A Widevine device license PlayReady-only content, but only as a deliberate step in the batch loop. """ +import asyncio import json from types import SimpleNamespace @@ -44,6 +45,7 @@ def test_extractor_is_type_strict(): async def _run_batch(monkeypatch, tracks, cdm_type, calls): session = SimpleNamespace( + lock=asyncio.Lock(), service_tag="EXAMPLE", service_instance=SimpleNamespace(), served_keys={}, diff --git a/tests/remote/unit/test_input_bridge.py b/tests/remote/unit/test_input_bridge.py index ef6b0cd4..bcbac4fa 100644 --- a/tests/remote/unit/test_input_bridge.py +++ b/tests/remote/unit/test_input_bridge.py @@ -107,3 +107,12 @@ def test_answered_only_after_an_accepted_response() -> None: assert bridge.submit_response("CODE") is True worker.join(timeout=2) assert bridge.answered is True + + +def test_request_input_after_authentication_fails_fast() -> None: + """The client stops polling for prompts once authenticated, so a later prompt must not wait for an answer.""" + bridge = InputBridge() + bridge.status = AuthStatus.AUTHENTICATED + with pytest.raises(RuntimeError, match="after authentication"): + bridge.request_input("Pick a profile", timeout=5) + assert bridge.status is AuthStatus.AUTHENTICATED diff --git a/tests/remote/unit/test_remote_client.py b/tests/remote/unit/test_remote_client.py index 62c2720d..2461d13a 100644 --- a/tests/remote/unit/test_remote_client.py +++ b/tests/remote/unit/test_remote_client.py @@ -168,3 +168,43 @@ def test_timeout_raises_systemexit(client: RemoteClient) -> None: ) with pytest.raises(SystemExit): client.get("/api/health") + + +@responses.activate +def test_late_error_after_heartbeat_is_fatal(client: RemoteClient) -> None: + """A heartbeat response carries a late failure as a 200 error body after leading newlines.""" + body = "\n\n" + json.dumps({"status": "error", "error_code": "SERVICE_ERROR", "message": "late"}) + responses.add(responses.GET, "http://srv:8786/api/session/s/titles", body=body, status=200) + with pytest.raises(SystemExit): + client.get("/api/session/s/titles") + assert client.request("get", "/api/session/s/titles", optional=True) == {} + + +@responses.activate +def test_heartbeat_body_parses_past_leading_newlines(client: RemoteClient) -> None: + responses.add(responses.GET, "http://srv:8786/api/session/s/titles", body='\n\n{"titles": []}', status=200) + assert client.get("/api/session/s/titles") == {"titles": []} + + +def test_read_timeout_in_body_reports_timeout(client: RemoteClient, caplog: pytest.LogCaptureFixture) -> None: + """requests re-raises a read timeout in the body as a ConnectionError; the client still names the timeout.""" + import requests + import urllib3 + + def stalled(*args: object, **kwargs: object) -> None: + raise requests.ConnectionError(urllib3.exceptions.ReadTimeoutError(None, "", "Read timed out.")) + + client.session.get = stalled # type: ignore[method-assign] + with pytest.raises(SystemExit): + client.get("/api/session/s/titles") + assert "timed out after 120" in caplog.text + + +@responses.activate +def test_proxy_error_page_logs_its_title(client: RemoteClient, caplog: pytest.LogCaptureFixture) -> None: + page = "example.com | 524: A timeout occurredx" + responses.add(responses.GET, "http://srv:8786/api/session/s/titles", body=page, status=524) + with pytest.raises(SystemExit): + client.get("/api/session/s/titles") + assert "HTTP 524: example.com | 524: A timeout occurred" in caplog.text + assert "DOCTYPE" not in caplog.text diff --git a/tests/remote/unit/test_routes.py b/tests/remote/unit/test_routes.py index 1f77291b..3fcffec3 100644 --- a/tests/remote/unit/test_routes.py +++ b/tests/remote/unit/test_routes.py @@ -7,12 +7,16 @@ stubbing the route table for selected paths. from __future__ import annotations +import asyncio +import json + import pytest from aiohttp import web +from unshackle.core.api import routes from unshackle.core.api.compression import compression_middleware from unshackle.core.api.errors import APIError, APIErrorCode -from unshackle.core.api.routes import api_handler, cors_middleware, setup_routes +from unshackle.core.api.routes import api_handler, cors_middleware, heartbeat, setup_routes pytestmark = pytest.mark.unit @@ -131,6 +135,93 @@ async def test_api_handler_maps_apierror(aiohttp_client) -> None: assert body["details"] == {"x": 1} +async def test_heartbeat_keeps_slow_request_alive(aiohttp_client, monkeypatch: pytest.MonkeyPatch) -> None: + """A fast route answers unchanged; a slow one streams newlines, then its body, and a late error arrives as 200.""" + monkeypatch.setattr(routes, "HEARTBEAT_INTERVAL", 0.05) + + @heartbeat + async def fast(request: web.Request) -> web.Response: + return web.json_response({"x": 1}, status=201) + + @heartbeat + async def slow(request: web.Request) -> web.Response: + await asyncio.sleep(0.2) + return web.json_response({"titles": [1]}) + + @heartbeat + async def slow_fail(request: web.Request) -> web.Response: + await asyncio.sleep(0.2) + raise APIError(APIErrorCode.SERVICE_ERROR, "late") + + app = web.Application() + app["debug_api"] = False + app.router.add_get("/fast", fast) + app.router.add_get("/slow", slow) + app.router.add_get("/slow_fail", slow_fail) + client = await aiohttp_client(app) + + resp = await client.get("/fast") + assert resp.status == 201 and await resp.json() == {"x": 1} + + resp = await client.get("/slow") + text = await resp.text() + assert resp.status == 200 and text.startswith("\n") and json.loads(text) == {"titles": [1]} + assert resp.headers["X-Accel-Buffering"] == "no" + + resp = await client.get("/slow_fail") + body = json.loads(await resp.text()) + assert resp.status == 200 and body["status"] == "error" and body["error_code"] == "SERVICE_ERROR" + + +async def test_heartbeat_late_error_counts_as_rejected(aiohttp_client, monkeypatch: pytest.MonkeyPatch) -> None: + """Stats count a late failure by its real status and the bytes the stream sent (with framing), not as an empty 200.""" + from unshackle.core.api.stats import stats, stats_middleware + + monkeypatch.setattr(routes, "HEARTBEAT_INTERVAL", 0.05) + monkeypatch.setattr(stats, "keys", {}) + + @heartbeat + async def slow_fail(request: web.Request) -> web.Response: + await asyncio.sleep(0.2) + raise APIError(APIErrorCode.SERVICE_ERROR, "late") + + app = web.Application(middlewares=[stats_middleware]) + app["debug_api"] = False + app.router.add_get("/slow_fail", slow_fail) + client = await aiohttp_client(app) + + resp = await client.get("/slow_fail") + sent = len(await resp.read()) + (entry,) = stats.keys.values() + assert resp.status == 200 and entry.rejected == 1 and entry.bytes_out >= sent > 0 + + +async def test_heartbeat_client_disconnect_is_quiet( + aiohttp_client, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """A client that leaves during the heartbeat logs no traceback, and the work runs to the end.""" + monkeypatch.setattr(routes, "HEARTBEAT_INTERVAL", 0.02) + finished = asyncio.Event() + + @heartbeat + async def slow(request: web.Request) -> web.Response: + await asyncio.sleep(0.3) + finished.set() + return web.json_response({}) + + app = web.Application() + app["debug_api"] = False + app.router.add_get("/slow", slow) + client = await aiohttp_client(app) + + resp = await client.get("/slow") + await resp.content.read(1) + resp.close() + await asyncio.wait_for(finished.wait(), 2) + await asyncio.sleep(0.1) + assert "Error handling request" not in caplog.text + + async def test_cors_preflight_returns_headers(make_app, aiohttp_client) -> None: app = make_app(remote_only=True) client = await aiohttp_client(app) diff --git a/tests/remote/unit/test_segment_filter.py b/tests/remote/unit/test_segment_filter.py index 67474b53..3275e059 100644 --- a/tests/remote/unit/test_segment_filter.py +++ b/tests/remote/unit/test_segment_filter.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json from types import SimpleNamespace @@ -29,6 +30,7 @@ BASE = "https://cdn.test/hls/media.m3u8" def make_session(monkeypatch, track): session = SimpleNamespace( + lock=asyncio.Lock(), tracks={"t1": track}, service_instance=SimpleNamespace(session=SimpleNamespace(get=lambda url: SimpleNamespace(text=PLAYLIST))), ) diff --git a/tests/remote/unit/test_server_cdm_absorb_routes.py b/tests/remote/unit/test_server_cdm_absorb_routes.py index d2abb979..ce2ff9ff 100644 --- a/tests/remote/unit/test_server_cdm_absorb_routes.py +++ b/tests/remote/unit/test_server_cdm_absorb_routes.py @@ -6,6 +6,7 @@ the first PSSH alone and so handed one track another track's narrower key set, a with no readable PSSH told the client nothing. """ +import asyncio import base64 import json from types import SimpleNamespace @@ -77,6 +78,7 @@ async def test_single_track_path_keeps_siblings(env, monkeypatch): drm=[_FakePlayReady(pssh_b64=PSSH, kid=KID_A), _FakePlayReady(pssh_b64=PSSH, kid=KID_B)], ) session = SimpleNamespace( + lock=asyncio.Lock(), service_tag="EX", service_instance=SimpleNamespace(), served_keys={}, @@ -111,6 +113,7 @@ async def test_batch_cache_does_not_starve_second_track(env, monkeypatch): ) tracks = {"vid": vid, "aud": aud} session = SimpleNamespace( + lock=asyncio.Lock(), service_tag="EX", service_instance=SimpleNamespace(), served_keys={}, @@ -142,6 +145,7 @@ async def test_no_pssh_tells_the_client_why(env, monkeypatch): buf = SessionLogBuffer() track = SimpleNamespace(id="vid", drm=[object()]) session = SimpleNamespace( + lock=asyncio.Lock(), service_tag="EX", service_instance=SimpleNamespace(), served_keys={}, diff --git a/tests/remote/unit/test_server_cdm_manifest_kid.py b/tests/remote/unit/test_server_cdm_manifest_kid.py index 04e91d8a..92c7a4d1 100644 --- a/tests/remote/unit/test_server_cdm_manifest_kid.py +++ b/tests/remote/unit/test_server_cdm_manifest_kid.py @@ -8,6 +8,7 @@ named different KIDs got one licence; and a KID the licence did not cover was lo server alone. """ +import asyncio import base64 import json from types import SimpleNamespace @@ -100,6 +101,7 @@ def fake_playready(monkeypatch): def _session(tracks: dict, buf=None): session = SimpleNamespace( + lock=asyncio.Lock(), service_tag="EX", service_instance=SimpleNamespace(), served_keys={}, diff --git a/tests/remote/unit/test_server_cdm_multikey.py b/tests/remote/unit/test_server_cdm_multikey.py index 309162ec..02f8b209 100644 --- a/tests/remote/unit/test_server_cdm_multikey.py +++ b/tests/remote/unit/test_server_cdm_multikey.py @@ -6,6 +6,7 @@ instead of short-circuiting on a vault hit for the PSSH KID alone, and the batch path must give each track the full bundle without borrowing another PSSH's keys. """ +import asyncio import base64 import json from types import SimpleNamespace @@ -118,6 +119,7 @@ async def test_batch_shares_full_bundle_per_pssh(monkeypatch): tracks = {tid: SimpleNamespace(id=tid, drm=[object()]) for tid in track_pssh} session = SimpleNamespace( + lock=asyncio.Lock(), service_tag="EXAMPLE", service_instance=SimpleNamespace(), served_keys={}, diff --git a/tests/remote/unit/test_server_cdm_type_override.py b/tests/remote/unit/test_server_cdm_type_override.py index 5b2ab564..7dd94498 100644 --- a/tests/remote/unit/test_server_cdm_type_override.py +++ b/tests/remote/unit/test_server_cdm_type_override.py @@ -8,6 +8,7 @@ client's choice decides, and the type the routes plan with must be the type of t they load. """ +import asyncio import json from types import SimpleNamespace from uuid import UUID @@ -74,6 +75,7 @@ def _session(track): get_playready_license=lambda **k: b"", ) return SimpleNamespace( + lock=asyncio.Lock(), service_tag="PCOK", service_instance=service, served_keys={}, diff --git a/tests/remote/unit/test_session_store.py b/tests/remote/unit/test_session_store.py index 4de3a695..8bdc0efe 100644 --- a/tests/remote/unit/test_session_store.py +++ b/tests/remote/unit/test_session_store.py @@ -76,6 +76,17 @@ async def test_cleanup_expired_drops_old_authenticated(store: SessionStore, monk assert store.session_count == 0 +async def test_session_with_work_in_progress_does_not_expire(store: SessionStore) -> None: + """A request that holds the lock for longer than the TTL keeps its session and cache directory.""" + from datetime import datetime, timedelta, timezone + + entry = await store.create("EXAMPLE", _FakeService(), session_id="busy") + entry.last_accessed = datetime.now(timezone.utc) - timedelta(seconds=store.ttl + 100) + async with entry.lock: + assert await store.cleanup_expired() == 0 + assert await store.get("busy") is entry + + async def test_cleanup_expired_keeps_pending_input_under_grace(store: SessionStore) -> None: """Sessions awaiting user input get a longer grace period (10 min) than authenticated TTL.""" entry = await store.create("EXAMPLE", _FakeService()) diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 346aba35..305475ed 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -3243,7 +3243,10 @@ async def session_titles_handler(session_id: str, request: Optional[web.Request] try: service_instance = session.service_instance async with session.lock: - titles = await asyncio.to_thread(service_instance.get_titles) + try: + titles = await asyncio.to_thread(service_instance.get_titles) + finally: + session.touch() session.titles = titles if hasattr(titles, "__iter__") and not isinstance(titles, str): @@ -3346,7 +3349,10 @@ async def session_tracks_handler( try: service_instance = session.service_instance async with session.lock: - tracks = await asyncio.to_thread(service_instance.get_tracks, title) + try: + tracks = await asyncio.to_thread(service_instance.get_tracks, title) + finally: + session.touch() title_tracks: Dict[str, Any] = {} for track in tracks.videos: @@ -4299,14 +4305,14 @@ def handle_proxy_license( # Surface it as a structured licence error, not an uncaught 500 the edge turns into a 502. try: if drm_type == "widevine": - license_response = service.get_widevine_license( + license_blob = service.get_widevine_license( **declared_kwargs( service.get_widevine_license, {"challenge": challenge_bytes, "title": title, "track": track} ) ) else: challenge_str = challenge_bytes.decode("utf-8", errors="replace") - license_response = service.get_playready_license( + license_blob = service.get_playready_license( **declared_kwargs( service.get_playready_license, {"challenge": challenge_str, "title": title, "track": track} ) @@ -4317,10 +4323,10 @@ def handle_proxy_license( log.exception(f"{sanitize_log(drm_type)} licence request failed for the proxied challenge") raise APIError(APIErrorCode.SERVICE_ERROR, f"Licence request failed: {exc}") from None - if isinstance(license_response, str): - license_response = license_response.encode("utf-8") + if isinstance(license_blob, str): + license_blob = license_blob.encode("utf-8") - return web.json_response({"license": base64.b64encode(license_response).decode("ascii")}) + return web.json_response({"license": base64.b64encode(license_blob).decode("ascii")}) async def session_segment_filter_handler( @@ -4350,7 +4356,7 @@ async def session_segment_filter_handler( if not callable(segment_filter) or not getattr(track, "url", None): return web.json_response({"unwanted": None}) - try: + def unwanted_segments() -> list[str]: import m3u8 import requests @@ -4360,7 +4366,11 @@ async def session_segment_filter_handler( if isinstance(response, requests.Response): response.encoding = response.encoding or "utf-8" playlist = m3u8.loads(response.text, uri=str(track.url)) - unwanted = [segment.absolute_uri for segment in playlist.segments if segment_filter(segment)] + return [segment.absolute_uri for segment in playlist.segments if segment_filter(segment)] + + try: + async with session.lock: + unwanted = await asyncio.to_thread(unwanted_segments) except (Exception, SystemExit): # `from None`: a debug-mode traceback would otherwise carry the service's exception chain log.exception(f"Error running the segment filter for track {sanitize_log(str(track_id))}") @@ -4385,6 +4395,14 @@ async def session_license_handler( session = await get_validated_session(session_id, request) require_authenticated(session) + async with session.lock: + return await asyncio.to_thread(license_response, data, session, session_id, request) + + +def license_response( + data: Dict[str, Any], session: Any, session_id: str, request: Optional[web.Request] +) -> web.Response: + """The licence response for :func:`session_license_handler`, made on a worker thread.""" track_id = data.get("track_id") track_ids = data.get("track_ids") challenge_b64 = data.get("challenge") diff --git a/unshackle/core/api/input_bridge.py b/unshackle/core/api/input_bridge.py index 928a169a..06359d85 100644 --- a/unshackle/core/api/input_bridge.py +++ b/unshackle/core/api/input_bridge.py @@ -57,11 +57,13 @@ class InputBridge: Raises: TimeoutError: If the bridge gets no response within *timeout*. - RuntimeError: If the server cancelled the bridge. + RuntimeError: If the server cancelled the bridge, or authentication is complete. """ with self._lock: if self._cancelled: raise RuntimeError("Session was cancelled") + if self._status == AuthStatus.AUTHENTICATED: + raise RuntimeError(f"Remote mode cannot relay a prompt after authentication: {prompt}") self._prompt = prompt self._response = None self._status = AuthStatus.PENDING_INPUT diff --git a/unshackle/core/api/routes.py b/unshackle/core/api/routes.py index 88326ca3..3cd5ff18 100644 --- a/unshackle/core/api/routes.py +++ b/unshackle/core/api/routes.py @@ -1,3 +1,4 @@ +import asyncio import functools import logging import re @@ -100,6 +101,53 @@ def api_handler(handler: Handler) -> Handler: return wrapper +HEARTBEAT_INTERVAL = 30.0 + + +def heartbeat(handler: Handler) -> Handler: + """Keep a slow request alive with a newline every ``HEARTBEAT_INTERVAL`` seconds. + + A response that is ready within one interval goes out unchanged. After that the route + sends status 200 and the headers, then a newline for each interval until the JSON body is + ready. JSON parsers ignore the leading whitespace. The status is already sent, so a late + failure arrives as the usual error body with ``"status": "error"`` and a 200 status. The + work continues if the client disconnects, so the remote session state stays complete. + """ + + @functools.wraps(handler) + async def wrapper(request: web.Request) -> web.StreamResponse: + task = asyncio.ensure_future(handler(request)) + done, _ = await asyncio.wait({task}, timeout=HEARTBEAT_INTERVAL) + if done: + return task.result() + + headers = {"Content-Type": "application/json", "X-Accel-Buffering": "no", **CORS_HEADERS} + response = web.StreamResponse(headers=headers) + await response.prepare(request) + try: + while not done: + await response.write(b"\n") + done, _ = await asyncio.wait({task}, timeout=HEARTBEAT_INTERVAL) + except ConnectionResetError: + log.info(f"Client left {request.path} before the response was ready; the work continues") + return response + debug_mode = request.app.get("debug_api", False) + try: + result = task.result() + except APIError as e: + result = build_error_response(e, debug_mode) + except Exception as e: + log.exception(f"Error in {request.path}") + result = handle_api_exception(e, debug_mode=debug_mode) + response["late_status"] = result.status + body = result.body if isinstance(result, web.Response) else None + await response.write(bytes(body) if isinstance(body, (bytes, bytearray)) else b"{}") + await response.write_eof() + return response + + return wrapper + + @api_handler async def health(request: web.Request) -> web.Response: """ @@ -1737,6 +1785,7 @@ async def session_create(request: web.Request) -> web.Response: @api_handler +@heartbeat async def session_titles(request: web.Request) -> web.Response: """ Get titles for an authenticated remote session. @@ -1751,7 +1800,10 @@ async def session_titles(request: web.Request) -> web.Response: type: string responses: '200': - description: List of titles + description: >- + List of titles. + After 30 s the route sends 200 and a newline every 30 s. A failure after that time + arrives as the error body with status 200. '404': description: Remote session not found """ @@ -1766,6 +1818,7 @@ async def session_titles(request: web.Request) -> web.Response: @api_handler +@heartbeat async def session_tracks(request: web.Request) -> web.Response: """ Get tracks and chapters for a specific title. @@ -1792,7 +1845,10 @@ async def session_tracks(request: web.Request) -> web.Response: description: ID of the title to get tracks for responses: '200': - description: Tracks and chapters for the title + description: >- + Tracks and chapters for the title. + After 30 s the route sends 200 and a newline every 30 s. A failure after that time + arrives as the error body with status 200. '404': description: Remote session or title not found """ @@ -1915,6 +1971,7 @@ async def session_segment_filter(request: web.Request) -> web.Response: @api_handler +@heartbeat async def session_license(request: web.Request) -> web.Response: """ Proxy DRM license through authenticated service. @@ -1955,6 +2012,8 @@ async def session_license(request: web.Request) -> web.Response: tracks, lists the content keys a server vault supplied, which the client has to prove before it trusts them. `clear_tracks`, absent when empty, lists the requested track ids that carry no DRM and so have no keys. + After 30 s the route sends 200 and a newline every 30 s. A failure after that time + arrives as the error body with status 200. '404': description: Remote session or track not found """ diff --git a/unshackle/core/api/session_store.py b/unshackle/core/api/session_store.py index df8fb575..5ea1ae14 100644 --- a/unshackle/core/api/session_store.py +++ b/unshackle/core/api/session_store.py @@ -62,6 +62,15 @@ class SessionEntry: """Update last_accessed timestamp.""" self.last_accessed = datetime.now(timezone.utc) + @property + def in_use(self) -> bool: + """True while a request holds the lock for service work; the store does not expire such a remote session. + + A request can run for longer than the TTL, and expiry would remove the cache directory + that the worker thread still writes to. + """ + return self.lock.locked() + def summary(self) -> Dict[str, Any]: """Dashboard view of the session: no service instance, no raw owner key.""" from unshackle.core.api.stats import mask_key @@ -145,7 +154,8 @@ class SessionStore: async with self._lock: max_sessions = self.max_sessions if max_sessions is not None and len(self._sessions) >= max_sessions: - oldest_id = min(self._sessions, key=lambda k: self._sessions[k].last_accessed) + sessions = self._sessions + oldest_id = min(sessions, key=lambda k: (sessions[k].in_use, sessions[k].last_accessed)) log.warning(f"Max sessions reached ({max_sessions}), evicting oldest: {oldest_id}") evicted = self._sessions.pop(oldest_id) if evicted.input_bridge: @@ -175,7 +185,7 @@ class SessionStore: if entry is None: return None - if entry.auth_status not in (AuthStatus.AUTHENTICATING, AuthStatus.PENDING_INPUT): + if not entry.in_use and entry.auth_status not in (AuthStatus.AUTHENTICATING, AuthStatus.PENDING_INPUT): elapsed = (datetime.now(timezone.utc) - entry.last_accessed).total_seconds() if elapsed > self.ttl: log.info(f"Session {sanitize_log(session_id)} expired (elapsed={elapsed:.0f}s, ttl={self.ttl}s)") @@ -215,6 +225,8 @@ class SessionStore: now = datetime.now(timezone.utc) expired = [] for sid, entry in self._sessions.items(): + if entry.in_use: + continue elapsed = (now - entry.last_accessed).total_seconds() if entry.auth_status in (AuthStatus.AUTHENTICATING, AuthStatus.PENDING_INPUT): if elapsed > AUTH_INPUT_TIMEOUT: diff --git a/unshackle/core/api/stats.py b/unshackle/core/api/stats.py index 865d60eb..e205a226 100644 --- a/unshackle/core/api/stats.py +++ b/unshackle/core/api/stats.py @@ -184,14 +184,16 @@ async def stats_middleware(request: web.Request, handler: Any) -> web.StreamResp stats.requests_total += 1 started = time.perf_counter() response = await handler(request) - if response.status == 401: + status = response.get("late_status", response.status) + bytes_out = getattr(response, "content_length", None) or response.body_length + if status == 401: stats.requests_rejected += 1 else: entry = stats.key_stats(configured_key(request_secret_key(request))) entry.requests += 1 entry.last_seen = time.time() - entry.bytes_out += getattr(response, "content_length", None) or 0 - if response.status >= 400: + entry.bytes_out += bytes_out + if status >= 400: entry.rejected += 1 session_id = request.match_info.get("session_id") if session_id and request.method != "OPTIONS" and not request.path.endswith("/logs"): @@ -204,10 +206,10 @@ async def stats_middleware(request: web.Request, handler: Any) -> web.StreamResp "method": request.method, "action": request.path.split(f"/{session_id}", 1)[-1].strip("/") or "info", "query": dict(request.query), - "status": response.status, + "status": status, "ms": round((time.perf_counter() - started) * 1000, 1), "bytes_in": request.content_length or 0, - "bytes_out": getattr(response, "content_length", None) or 0, + "bytes_out": bytes_out, }, ) return response diff --git a/unshackle/core/remote_service.py b/unshackle/core/remote_service.py index 1c0089da..b90c8781 100644 --- a/unshackle/core/remote_service.py +++ b/unshackle/core/remote_service.py @@ -25,6 +25,7 @@ from uuid import UUID import click import requests +import urllib3 from langcodes import Language from requests.adapters import HTTPAdapter, Retry from rich.padding import Padding @@ -61,8 +62,15 @@ def redact_secrets(text: str, data: Optional[Dict[str, Any]] = None) -> str: class RemoteClient: """HTTP client for the unshackle serve API.""" - def __init__(self, server_url: str, api_key: str, auth_headers: Optional[list[str]] = None) -> None: + def __init__( + self, + server_url: str, + api_key: str, + auth_headers: Optional[list[str]] = None, + timeout: Optional[float] = None, + ) -> None: self.server_url = server_url.rstrip("/") + self.timeout = timeout or 120 self.api_key = api_key self.auth_headers = list(auth_headers) if auth_headers else list(DEFAULT_AUTH_HEADERS) self._auth_header_index = 0 @@ -95,31 +103,45 @@ class RemoteClient: url = f"{self.server_url}{endpoint}" while True: try: - resp = getattr(self.session, method)(url, json=data, timeout=120) - except requests.ConnectionError: + resp = getattr(self.session, method)(url, json=data, timeout=self.timeout) + except requests.ConnectionError as e: + if isinstance(e.args[0] if e.args else None, urllib3.exceptions.ReadTimeoutError): + log.error( + f"Request to remote server timed out after {self.timeout}s: {endpoint}. " + "Raise 'timeout' for this server in remote_services." + ) + raise SystemExit(1) server_url = safe_display_url(self.server_url) log.error(f"Could not connect to remote server at {server_url}. Is it running? (unshackle serve)") raise SystemExit(1) except requests.Timeout: - log.error(f"Request to remote server timed out: {endpoint}") + log.error( + f"Request to remote server timed out after {self.timeout}s: {endpoint}. " + "Raise 'timeout' for this server in remote_services." + ) raise SystemExit(1) # servers differ on which header carries the key, so retry the rest before giving up if resp.status_code == 401 and self.next_auth_header(): continue break - if resp.status_code >= 400: - try: - detail = resp.json() - except ValueError: - detail = {} - error_msg = redact_secrets(str(detail.get("message", resp.text)), data) + try: + detail = resp.json() + except ValueError: + if resp.status_code < 400: + raise + detail = {} + late_error = isinstance(detail, dict) and detail.get("status") == "error" and "error_code" in detail + if resp.status_code >= 400 or late_error: + page_title = re.search(r"(.*?)", resp.text, re.S) if not detail else None + fallback = f"HTTP {resp.status_code}: {page_title.group(1).strip()}" if page_title else resp.text + error_msg = redact_secrets(str(detail.get("message", fallback)), data) error_code = detail.get("error_code", "UNKNOWN") if optional: log.debug(f"Optional endpoint {endpoint} unavailable [{error_code}]: {error_msg}") return {} log.error(f"Server error [{error_code}]: {error_msg}") raise SystemExit(1) - return resp.json() + return detail def post(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]: return self.request("post", endpoint, data) @@ -632,6 +654,7 @@ def resolve_server(server_name: Optional[str]) -> tuple[str, str, dict]: services = svc.get("services", {}) services["_server_cdm"] = svc.get("server_cdm") services["_auth_headers"] = resolve_auth_headers(svc, server_name) + services["_timeout"] = svc.get("timeout") return svc["url"], svc.get("api_key", ""), services if len(remote_services) == 1: @@ -642,6 +665,7 @@ def resolve_server(server_name: Optional[str]) -> tuple[str, str, dict]: services = svc.get("services", {}) services["_server_cdm"] = svc.get("server_cdm") services["_auth_headers"] = resolve_auth_headers(svc, name) + services["_timeout"] = svc.get("timeout") return svc["url"], svc.get("api_key", ""), services available = ", ".join(remote_services.keys()) @@ -842,7 +866,9 @@ class RemoteService: self.service_tag = service_tag self.title_id = title_id - self.client = RemoteClient(server_url, api_key, services_config.get("_auth_headers")) + self.client = RemoteClient( + server_url, api_key, services_config.get("_auth_headers"), services_config.get("_timeout") + ) self.ctx = ctx self._service_params = service_params or {} self.log = logging.getLogger(service_tag) diff --git a/unshackle/core/services.py b/unshackle/core/services.py index 2c935c3d..a1df35f4 100644 --- a/unshackle/core/services.py +++ b/unshackle/core/services.py @@ -452,7 +452,9 @@ class Services(click.Group): server_name = ctx.params.get("server") server_url, api_key, services_config = resolve_server(server_name) - client = RemoteClient(server_url, api_key, services_config.get("_auth_headers")) + client = RemoteClient( + server_url, api_key, services_config.get("_auth_headers"), services_config.get("_timeout") + ) result = client.get("/api/services") Services.remote_services_cache = result.get("services", []) return Services.remote_services_cache