kenzuya/unshackle
1
0
Fork 2
mirror of https://github.com/unshackle-dl/unshackle.git synced 2026-11-12 20:15:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(dl): take turns on licence requests so one licence serves every track

Browse Source 
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.
This commit is contained in:
imSp4rky
2026-10-31 12:42:43 -06:00
parent f2747076a6
commit b3c86413d5
2 changed files with 190 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/core/test_licence_turns.py
+140
View File
@@ -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
unshackle/commands/dl.py
+50 -8
View File
@@ -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