From b3c86413d518ba59dd92df3a7ddf8afbd64c198e Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sat, 31 Oct 2026 12:42:43 -0600 Subject: [PATCH] fix(dl): take turns on licence requests so one licence serves every track Tracks with different KID sets still query the vaults at the same time, but wait for one licence lock and check the key cache again before they licence. --- tests/core/test_licence_turns.py | 140 +++++++++++++++++++++++++++++++ unshackle/commands/dl.py | 58 +++++++++++-- 2 files changed, 190 insertions(+), 8 deletions(-) create mode 100644 tests/core/test_licence_turns.py diff --git a/tests/core/test_licence_turns.py b/tests/core/test_licence_turns.py new file mode 100644 index 00000000..2913e8d1 --- /dev/null +++ b/tests/core/test_licence_turns.py @@ -0,0 +1,140 @@ +"""Licences take turns across DRMs with different KIDs, so one licence that carries every key is sent once.""" + +from __future__ import annotations + +import logging +import threading +import time +from types import SimpleNamespace +from typing import Any +from uuid import UUID + +import pytest +from pywidevine.pssh import PSSH +from rich.table import Table + +from unshackle.commands.dl import dl +from unshackle.core.cdm.detect import cdm_type_stub +from unshackle.core.constants import DOWNLOAD_CANCELLED, DownloadCancelled +from unshackle.core.drm import Widevine + +pytestmark = pytest.mark.unit + +KID_A, KID_B, KID_C = (UUID(c * 32) for c in "abc") +KEYS = {KID_A: "aa" * 16, KID_B: "bb" * 16, KID_C: "cc" * 16} + + +class NoVaults: + sources: dict = {} + + def __bool__(self) -> bool: + return False + + def get_key(self, kid: UUID) -> tuple[None, None]: + return None, None + + +@pytest.fixture +def licences(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Each licence returns every content key of the title, after a delay that lets the tracks overlap.""" + calls: list[str] = [] + + def get_content_keys(self: Widevine, **_: Any) -> None: + calls.append(threading.current_thread().name) + time.sleep(0.2) + self.content_keys = dict(KEYS) + + monkeypatch.setattr(Widevine, "get_content_keys", get_content_keys) + monkeypatch.setattr(dl, "LICENSE_KEY_CACHE", {}) + monkeypatch.setattr(dl, "DRM_LOCKS", {}) + DOWNLOAD_CANCELLED.clear() + yield calls + DOWNLOAD_CANCELLED.clear() + + +@pytest.fixture +def cmd() -> dl: + cmd = dl.__new__(dl) + cmd.log = logging.getLogger(__name__) + cmd.service = "TEST" + cmd.vaults = NoVaults() + cmd.vault_cache_tally = None + cmd._remote_service = None + cmd.debug_logger = None + cmd.cdm = cdm_type_stub("widevine") + return cmd + + +def widevine(*kids: UUID) -> Widevine: + return Widevine(pssh=PSSH.new(system_id=PSSH.SystemId.Widevine, key_ids=list(kids))) + + +def prepare(cmd: dl, drm: Widevine, track_kid: UUID, table: Table, **kwargs: Any) -> None: + cmd.prepare_drm( + drm, + track=SimpleNamespace(drm=[drm]), + title=None, + certificate=lambda **_: None, + licence=lambda **_: None, + table=table, + track_kid=track_kid, + **kwargs, + ) + + +def new_table() -> Table: + table = Table.grid() + table.add_row("") + return table + + +def run_together(cmd: dl, *jobs: tuple[Widevine, UUID]) -> tuple[Table, list[BaseException]]: + table = new_table() + errors: list[BaseException] = [] + + def run(drm: Widevine, kid: UUID) -> None: + try: + prepare(cmd, drm, kid, table) + except BaseException as e: + errors.append(e) + + threads = [threading.Thread(target=run, args=job, name=f"track-{i}") for i, job in enumerate(jobs)] + for t in threads: + t.start() + for t in threads: + t.join() + return table, errors + + +def tree_labels(table: Table, kid: UUID) -> list[str]: + trees = [cell for cell in table.columns[0].cells if hasattr(cell, "children")] + return [str(child.label) for tree in trees for child in tree.children if kid.hex in str(child.label)] + + +def test_one_licence_for_two_pssh_with_different_kids(cmd: dl, licences: list[str]) -> None: + first, second = widevine(KID_A, KID_B), widevine(KID_C) + + table, errors = run_together(cmd, (first, KID_A), (second, KID_C)) + + assert not errors + assert len(licences) == 1 + assert first.content_keys[KID_A] == KEYS[KID_A] and second.content_keys[KID_C] == KEYS[KID_C] + from_cache = [label for kid in (KID_A, KID_C) for label in tree_labels(table, kid) if label.endswith(" from cache")] + assert from_cache + + +def test_a_track_waiting_for_its_turn_sends_no_licence_after_a_failure(cmd: dl, licences: list[str]) -> None: + DOWNLOAD_CANCELLED.set() + + with pytest.raises(DownloadCancelled): + prepare(cmd, widevine(KID_C), KID_C, new_table()) + + assert licences == [] + + +def test_force_licenses_again_although_the_keys_are_cached(cmd: dl, licences: list[str]) -> None: + dl.LICENSE_KEY_CACHE.update(KEYS) + + prepare(cmd, widevine(KID_C), KID_C, new_table(), force=True) + + assert len(licences) == 1 diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 6578cc34..56222b37 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -14,6 +14,7 @@ from collections import Counter, defaultdict from collections.abc import Iterable from concurrent import futures from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack from copy import deepcopy from datetime import date, timedelta from functools import partial @@ -50,6 +51,7 @@ from unshackle.core.constants import ( DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, + DownloadCancelled, context_settings, ) from unshackle.core.credential import Credential @@ -1241,6 +1243,7 @@ class dl: DRM_LOCKS: dict[str, Lock] = {} DRM_LOCKS_GUARD = Lock() + LICENCE_LOCK = Lock() VAULT_WRITER = ThreadPoolExecutor(1, thread_name_prefix="vault-writer") EXPORT_LOCK = Lock() LICENSE_KEY_CACHE: dict[UUID, str] = {} @@ -4749,8 +4752,8 @@ class dl: """Return the lock for one DRM's set of content keys. Tracks that share KIDs still take turns, so the second one finds the content keys in - LICENSE_KEY_CACHE and skips the challenge. Tracks with different KIDs send a challenge - and query the key vaults at the same time. + LICENSE_KEY_CACHE and skips the challenge. Tracks with different KIDs query the key vaults + at the same time; their licences take turns under ``take_licence_turn``. Callers take the lock after get_drm_for_cdm(), which folds a track's sibling PlayReady objects into the returned one, so the KID set is complete for that track's DRM list @@ -4764,6 +4767,37 @@ class dl: with cls.DRM_LOCKS_GUARD: return cls.DRM_LOCKS.setdefault(key or type(drm).__name__, Lock()) + def take_licence_turn( + self, + turn: ExitStack, + drm: Union[ClearKeyCENC, Widevine, PlayReady, MonaLisa], + kids: list[UUID], + track: AnyTrack, + track_kid: Optional[UUID], + cek_tree: Optional[Tree], + ) -> None: + """Hold LICENCE_LOCK until ``turn`` closes, then take the content keys of ``kids`` an earlier licence cached. + + One licence can carry the keys of a DRM with other KIDs, so licences take turns. The caller + holds its ``drm_lock`` first, so the vault lookups before this still run at the same time. + """ + turn.enter_context(self.LICENCE_LOCK) + if DOWNLOAD_CANCELLED.is_set(): + raise DownloadCancelled() + for kid in kids: + key = self.LICENSE_KEY_CACHE.get(kid) + if not key or kid in drm.content_keys: + continue + drm.content_keys[kid] = key + if cek_tree is not None and not any(f"{kid.hex}:{key}" in str(x.label) for x in cek_tree.children): + cek_tree.add(f"[text2]{kid.hex}:{key}{['', '*'][kid == track_kid]} from cache") + log_event( + "license_cache_hit", + level="INFO", + service=self.service, + context={"kid": kid.hex, "content_key": key, "track": str(track), "drm_type": type(drm).__name__}, + ) + def service_licence(self, service: Service, drm_system: Optional[str] = None, **kwargs: Any) -> Any: """Call the service's licence function for the track's DRM system. @@ -4908,7 +4942,7 @@ class dl: licensed_locally = bool(client_licensed) and str(track.id) in client_licensed if server_cdm and not licensed_locally: - with self.drm_lock(drm): + with self.drm_lock(drm), ExitStack() as licence_turn: pending_vault_writes: list[Callable[[], Any]] = [] vault_kids = list(getattr(drm, "kids", None) or []) if track_kid and track_kid not in vault_kids: @@ -4923,11 +4957,13 @@ class dl: drm.content_keys[kid] = content_key self.LICENSE_KEY_CACHE[kid] = content_key - known_keys = set(drm.content_keys) - def missing_track_key() -> bool: return not drm.content_keys or bool(track_kid and track_kid not in drm.content_keys) + if missing_track_key(): + self.take_licence_turn(licence_turn, drm, vault_kids, track, track_kid, None) + known_keys = set(drm.content_keys) + if missing_track_key(): from unshackle.core.remote_service import RemoteError @@ -5079,7 +5115,7 @@ class dl: ) pending_vault_writes = [] - with self.drm_lock(drm): + with self.drm_lock(drm), ExitStack() as licence_turn: pssh_display = self.truncate_pssh_for_display(drm.pssh.dumps(), "Widevine") cek_tree = Tree(Text.assemble(("Widevine", "cyan"), (f"({pssh_display})", "text"), overflow="fold")) pre_existing_tree = next( @@ -5164,6 +5200,8 @@ class dl: if kid not in drm.content_keys and cdm_only: need_license = True + if need_license: + self.take_licence_turn(licence_turn, drm, [] if force else all_kids, track, track_kid, cek_tree) if need_license and all(kid in drm.content_keys for kid in all_kids): need_license = False @@ -5272,7 +5310,7 @@ class dl: ) pending_vault_writes = [] - with self.drm_lock(drm): + with self.drm_lock(drm), ExitStack() as licence_turn: pssh_display = self.truncate_pssh_for_display(drm.pssh_b64 or "", "PlayReady") cek_tree = Tree( Text.assemble( @@ -5364,6 +5402,8 @@ class dl: if kid not in drm.content_keys and cdm_only: need_license = True + if need_license: + self.take_licence_turn(licence_turn, drm, [] if force else all_kids, track, track_kid, cek_tree) if need_license and all(kid in drm.content_keys for kid in all_kids): need_license = False @@ -5432,7 +5472,7 @@ class dl: elif isinstance(drm, ClearKeyCENC): pending_vault_writes = [] - with self.drm_lock(drm): + with self.drm_lock(drm), ExitStack() as licence_turn: cek_tree = Tree(Text.assemble(("ClearKey", "cyan"), overflow="fold")) pre_existing_tree = next( (x for x in table.columns[0].cells if isinstance(x, Tree) and x.label == cek_tree.label), None @@ -5503,6 +5543,8 @@ class dl: if kid not in drm.content_keys and cdm_only: need_license = True + if need_license: + self.take_licence_turn(licence_turn, drm, [] if force else all_kids, track, track_kid, cek_tree) if need_license and all(kid in drm.content_keys for kid in all_kids): need_license = False