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