mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-15 21:45:10 +00:00
fix(vaults): stop querying sibling KIDs and skip unavailable vaults
The server CDM stub names only the manifest's KIDs; sibling keys stay content keys. API and HTTP vaults retry a 429, and a vault that is down or still rate limited costs one timeout per run.
This commit is contained in:
@@ -10,11 +10,14 @@ import pytest
|
||||
|
||||
from unshackle.commands import dl as dl_module
|
||||
from unshackle.commands.dl import dl
|
||||
from unshackle.core.vault import RateLimited, vault_session
|
||||
from unshackle.core.vaults import Vaults
|
||||
from unshackle.vaults.API import API
|
||||
from unshackle.vaults.API import API, Exceptions
|
||||
from unshackle.vaults.HTTP import HTTP
|
||||
from unshackle.vaults.SQLite import SQLite
|
||||
|
||||
KID = UUID("11111111-2222-3333-4444-555555555555")
|
||||
SIBLING = UUID("99999999-2222-3333-4444-555555555555")
|
||||
BAD = "aa" * 16
|
||||
GOOD = "bb" * 16
|
||||
|
||||
@@ -630,3 +633,100 @@ def test_api_vault_reports_the_pair_upstream(monkeypatch: pytest.MonkeyPatch, ca
|
||||
vaults.flag_bad_key(KID, BAD)
|
||||
vaults.unflag_bad_key(KID, BAD)
|
||||
assert not caplog.records
|
||||
|
||||
|
||||
def test_prefer_vault_keys_queries_each_named_kid_once() -> None:
|
||||
asked: list[UUID] = []
|
||||
cmd = dl.__new__(dl)
|
||||
cmd.vaults = SimpleNamespace(get_key=lambda kid: asked.append(kid) or (None, None))
|
||||
cmd._remote_service = SimpleNamespace(server_vault_keys={})
|
||||
video, audio = (SimpleNamespace(kids=[KID], content_keys={KID: GOOD, SIBLING: BAD}) for _ in range(2))
|
||||
|
||||
cmd.prefer_vault_keys(SimpleNamespace(tracks=[SimpleNamespace(drm=[video]), SimpleNamespace(drm=[audio])]))
|
||||
|
||||
assert asked == [KID]
|
||||
|
||||
|
||||
def test_cache_resolved_keys_skips_a_key_a_local_vault_holds(tmp_path: Path) -> None:
|
||||
local = SQLite("local", tmp_path / "local.db")
|
||||
local.add_key("SVC", SIBLING, BAD)
|
||||
cmd = dl.__new__(dl)
|
||||
cmd.vaults = Vaults("SVC")
|
||||
cmd.vaults.vaults = [local]
|
||||
cmd._remote_service = SimpleNamespace(server_vault_keys={})
|
||||
pushed: list[dict] = []
|
||||
cmd.flush_vault_writes = lambda writes: pushed.extend(w.args[0] for w in writes)
|
||||
dl.LICENSE_KEY_CACHE.clear()
|
||||
drm = SimpleNamespace(kids=[KID], content_keys={KID: GOOD, SIBLING: BAD})
|
||||
|
||||
cmd.cache_resolved_keys(SimpleNamespace(tracks=[SimpleNamespace(drm=[drm])]))
|
||||
|
||||
assert pushed == [{KID: GOOD}]
|
||||
|
||||
|
||||
def test_api_vault_does_not_shrink_the_batch_on_429(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
api = API("vault", "https://vault.example/", "tok")
|
||||
posts: list[str] = []
|
||||
limited = SimpleNamespace(status_code=429)
|
||||
monkeypatch.setattr(api.session, "post", lambda url, **kw: posts.append(url) or limited)
|
||||
|
||||
with pytest.raises(Exceptions.TooManyRequests):
|
||||
api.add_keys("SVC", {KID: GOOD})
|
||||
|
||||
assert len(posts) == 1
|
||||
|
||||
|
||||
def test_http_vault_does_not_fall_back_to_single_inserts_on_429(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
vault = HTTP("vault", "https://vault.example/", password="tok", api_mode="json")
|
||||
posts: list[str] = []
|
||||
limited = SimpleNamespace(status_code=429, ok=False)
|
||||
monkeypatch.setattr(vault.session, "post", lambda url, **kw: posts.append(url) or limited)
|
||||
|
||||
with pytest.raises(RateLimited):
|
||||
vault.add_keys("SVC", {KID: GOOD, SIBLING: BAD})
|
||||
|
||||
assert len(posts) == 1
|
||||
|
||||
|
||||
def test_vault_session_retries_429_but_not_a_dead_vault() -> None:
|
||||
retry = vault_session(10).get_adapter("https://vault.example").max_retries
|
||||
assert retry.status_forcelist == [429]
|
||||
assert (retry.connect, retry.read, retry.other) == (0, 0, 0)
|
||||
|
||||
|
||||
def test_an_unreachable_vault_costs_one_attempt_per_run(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import requests
|
||||
|
||||
def down(url: str, **kw: object) -> None:
|
||||
calls.append(url)
|
||||
raise requests.ConnectTimeout("no answer")
|
||||
|
||||
calls: list[str] = []
|
||||
api = API("api", "https://api.example/", "tok")
|
||||
http = HTTP("http", "https://http.example/", password="tok", api_mode="json")
|
||||
for vault in (api, http):
|
||||
monkeypatch.setattr(vault.session, "get", down)
|
||||
monkeypatch.setattr(vault.session, "post", down)
|
||||
with pytest.raises(requests.ConnectTimeout):
|
||||
vault.add_keys("SVC", {KID: GOOD, SIBLING: BAD})
|
||||
assert len(calls) == 2 # no batch step-down, no single-key fallback
|
||||
|
||||
calls.clear()
|
||||
vaults = Vaults("SVC")
|
||||
vaults.vaults = [api]
|
||||
assert vaults.get_key(KID) == (None, None)
|
||||
assert vaults.get_key(SIBLING) == (None, None)
|
||||
assert vaults.add_keys({KID: GOOD}) == 0
|
||||
assert calls == ["https://api.example/svc/" + KID.hex]
|
||||
|
||||
|
||||
def test_a_vault_still_rate_limited_is_not_used_again(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
calls: list[str] = []
|
||||
api = API("api", "https://api.example/", "tok")
|
||||
monkeypatch.setattr(api.session, "get", lambda url, **kw: calls.append(url) or SimpleNamespace(status_code=429))
|
||||
vaults = Vaults("SVC")
|
||||
vaults.vaults = [api]
|
||||
|
||||
assert vaults.get_key(KID) == (None, None)
|
||||
assert vaults.get_key(SIBLING) == (None, None)
|
||||
assert calls == ["https://api.example/svc/" + KID.hex]
|
||||
|
||||
@@ -170,3 +170,51 @@ def test_resolve_server_keys_does_not_warn_for_clear_tracks(monkeypatch):
|
||||
svc.resolve_server_keys(SimpleNamespace(tracks=Tracks([video_track, audio_track])))
|
||||
|
||||
assert warnings == ["Server CDM returned no content keys for track v-1"]
|
||||
|
||||
|
||||
def resolve(manifest_drm, drm_type="widevine"):
|
||||
"""Run resolve_server_keys for one track the server answers with its own key and a sibling."""
|
||||
svc = remote_service()
|
||||
svc.log = SimpleNamespace(warning=lambda *a, **k: None, debug=lambda *a, **k: None)
|
||||
svc._server_cdm = True
|
||||
svc.drain_server_logs = lambda: None
|
||||
keys = {CDM_KID.hex: "bb" * 16, KID.hex: "aa" * 16}
|
||||
svc.client = SimpleNamespace(post=lambda ep, data: {"keys": {"v-1": keys}, "drm_type": drm_type})
|
||||
track = SimpleNamespace(id="v-1", drm=manifest_drm)
|
||||
|
||||
class Tracks(list):
|
||||
videos = [track]
|
||||
audio: list = []
|
||||
|
||||
svc.resolve_server_keys(SimpleNamespace(tracks=Tracks([track])))
|
||||
assert set(track.drm[0].content_keys) == {KID, CDM_KID}
|
||||
return track.drm[0]
|
||||
|
||||
|
||||
def test_resolve_server_keys_names_only_the_manifest_kids():
|
||||
"""Sibling keys in the licence stay content keys but are not KIDs the track names."""
|
||||
from pywidevine.pssh import PSSH
|
||||
|
||||
from unshackle.core.drm import Widevine
|
||||
|
||||
widevine = Widevine(pssh=PSSH.new(system_id=PSSH.SystemId.Widevine, key_ids=[KID]))
|
||||
assert resolve([widevine]).kids == [KID]
|
||||
|
||||
playready = RemoteService.create_drm_stub("playready", [KID.hex])
|
||||
assert resolve([playready], "playready").kids == [KID]
|
||||
|
||||
|
||||
def test_resolve_server_keys_names_every_kid_without_manifest_drm():
|
||||
assert set(resolve(None).kids) == {KID, CDM_KID}
|
||||
|
||||
|
||||
def test_reconstruct_drm_keeps_the_default_kid_of_a_pssh_without_kids():
|
||||
from pywidevine.pssh import PSSH
|
||||
|
||||
from unshackle.core.drm import real_kids
|
||||
from unshackle.core.remote_service import reconstruct_drm
|
||||
|
||||
content_id_only = PSSH.new(system_id=PSSH.SystemId.Widevine, init_data=b"\x22\x03abc")
|
||||
(drm,) = reconstruct_drm([{"type": "widevine", "pssh": content_id_only.dumps(), "kids": [str(KID)]}])
|
||||
|
||||
assert real_kids(drm) == [KID]
|
||||
|
||||
@@ -33,6 +33,9 @@ class FakeVaults:
|
||||
def __len__(self):
|
||||
return 1
|
||||
|
||||
def held_locally(self, kid, key):
|
||||
return False
|
||||
|
||||
def get_key(self, kid):
|
||||
return self.held.get(kid), self.vault_used
|
||||
|
||||
|
||||
Reference in New Issue
Block a user