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:
imSp4rky
2026-10-16 14:49:37 -06:00
parent ecb8296de8
commit dc169f81e1
10 changed files with 261 additions and 49 deletions
+6
View File
@@ -56,3 +56,9 @@ key_vaults:
Timeout in seconds for vault operations. Injected automatically into any backend whose
constructor accepts a `timeout` parameter (a per-vault `timeout` still wins). The `MySQL` backend
uses it as the default `connect_timeout`.
The `API` and `HTTP` backends send a request again up to 3 times when the vault answers
`429 Too Many Requests`. They wait the `Retry-After` time, but never longer than this timeout.
A connection failure or a timeout is not tried again. The run does not use a vault again after a
connection failure, a timeout, or a `429` that continues after the retries, so a vault that is
down costs one timeout per run.
+101 -1
View File
@@ -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
+15 -6
View File
@@ -55,7 +55,7 @@ from unshackle.core.constants import (
)
from unshackle.core.credential import Credential
from unshackle.core.downloaders import default_max_workers, format_speed, parse_speed_limit, set_speed_limit
from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, cbcs, own_kids, verify
from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine, cbcs, own_kids, real_kids, verify
from unshackle.core.events import events
from unshackle.core.export_name import ExportNamer, move_export, plural, season_episodes, title_label
from unshackle.core.providers.anilist import parse_anilist_ref
@@ -4555,13 +4555,21 @@ class dl:
poisoned row; the server key waits as the next candidate. A server CDM key is trusted
like a local CDM key: it goes to the run cache and the vaults, and a retry after a
flagged vault key takes it before any vault. A server vault key is unproven, so it
waits in ``Vaults.candidates`` and gets the same check when its turn comes.
waits in ``Vaults.candidates`` and gets the same check when its turn comes. Only the KIDs
the DRM names are looked up, each once, so sibling keys add no remote vault requests.
"""
server_vault_keys = getattr(self._remote_service, "server_vault_keys", {})
looked_up: dict[UUID, tuple[Optional[str], Any]] = {}
for track in title.tracks:
for drm in getattr(track, "drm", None) or []:
for kid, server_key in list(getattr(drm, "content_keys", {}).items()):
vault_key, vault = self.vaults.get_key(kid)
content_keys = getattr(drm, "content_keys", {})
named = set(real_kids(drm) or content_keys)
for kid, server_key in list(content_keys.items()):
if kid not in named:
continue
if kid not in looked_up:
looked_up[kid] = self.vaults.get_key(kid)
vault_key, vault = looked_up[kid]
if not vault_key or vault_key == server_key:
continue
drm.content_keys[kid] = vault_key
@@ -4579,10 +4587,10 @@ class dl:
"""Cache the keys a batch licence or an import put on the tracks.
Both fill track.drm before any prepare_drm call, so this is the only point where those
keys can still reach the local vaults.
keys can still reach the local vaults. A content key a local vault already holds is not pushed again.
"""
server_vault_keys = getattr(getattr(self, "_remote_service", None), "server_vault_keys", {})
keys = {
resolved = {
kid: key
for track in title.tracks
for drm in (getattr(track, "drm", None) or [])
@@ -4591,6 +4599,7 @@ class dl:
and server_vault_keys.get(kid) != key
and self.vaults.sources.get(kid, (None,))[0] != key
}
keys = {kid: key for kid, key in resolved.items() if not self.vaults.held_locally(kid, key)}
if keys:
self.LICENSE_KEY_CACHE.update(keys)
self.flush_vault_writes([partial(self.cache_keys_to_vaults, keys)])
+5 -4
View File
@@ -38,7 +38,7 @@ from unshackle.core.config import config
from unshackle.core.console import console, prompt_user
from unshackle.core.constants import AnyTrack
from unshackle.core.credential import Credential
from unshackle.core.drm import cbcs
from unshackle.core.drm import cbcs, real_kids
from unshackle.core.proxies.resolve import is_loopback
from unshackle.core.session import MAX_BACKOFF
from unshackle.core.titles import Title_T, Titles_T, remap_titles
@@ -303,7 +303,8 @@ def reconstruct_drm(drm_list: Optional[list]) -> list:
from unshackle.core.drm import Widevine
wv_pssh = WidevinePSSH(pssh_str)
result.append(Widevine(pssh=wv_pssh))
kids = drm_info.get("kids")
result.append(Widevine(pssh=wv_pssh, kid=kids[0] if kids and not wv_pssh.key_ids else None))
elif drm_type == "playready":
import base64 as b64
@@ -1542,9 +1543,9 @@ class RemoteService:
self.log.warning(f"Server CDM returned no content keys for track {track.id}")
continue
kid_list = list(track_keys.keys())
named = {kid.hex for drm in track.drm or [] for kid in real_kids(drm)} & track_keys.keys()
track_drm_type = drm_types_by_track.get(str(track.id), server_drm_type)
drm_obj = self.create_drm_stub(track_drm_type, kid_list)
drm_obj = self.create_drm_stub(track_drm_type, sorted(named) or list(track_keys))
for kid_hex, key_hex in track_keys.items():
drm_obj.content_keys[UUID(hex=kid_hex)] = key_hex
track.drm = [drm_obj]
+40
View File
@@ -2,6 +2,46 @@ from abc import ABCMeta, abstractmethod
from typing import Iterator, Optional, Union
from uuid import UUID
from requests import ConnectionError as RequestsConnectionError
from requests import Session, Timeout
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from unshackle.core import __version__
from unshackle.core.session import POOL_MAX_SIZE
class RateLimited(Exception):
"""The vault still answered 429 after the retries."""
UNREACHABLE = (RequestsConnectionError, Timeout, ConnectionError, TimeoutError)
def vault_session(timeout: float) -> Session:
"""An HTTP session for a remote vault that retries a 429 and nothing else.
A 429 means the request was not processed, so a POST is safe to send again. A connect or read
failure is not retried: a lookup waits on every remote vault, so a retry multiplies the timeout.
"""
retry = Retry(
total=3,
connect=0,
read=0,
other=0,
backoff_factor=0.5,
status_forcelist=[429],
allowed_methods=None,
retry_after_max=int(timeout),
raise_on_status=False,
)
session = Session()
adapter = HTTPAdapter(pool_connections=POOL_MAX_SIZE, pool_maxsize=POOL_MAX_SIZE, max_retries=retry)
session.mount("https://", adapter)
session.mount("http://", adapter)
session.headers.update({"User-Agent": f"unshackle v{__version__}"})
return session
class Vault(metaclass=ABCMeta):
# a vault on this machine; Vaults.get_key tries these first, before the network ones
+25 -6
View File
@@ -8,7 +8,7 @@ from uuid import UUID
from unshackle.core.config import config
from unshackle.core.utilities import get_debug_logger, import_module_by_path
from unshackle.core.vault import Vault
from unshackle.core.vault import UNREACHABLE, RateLimited, Vault
log = logging.getLogger(__name__)
@@ -29,6 +29,7 @@ class Vaults:
self.flagged: set[tuple[Union[UUID, str], str]] = set()
self.candidates: dict[Union[UUID, str], list[tuple[str, Vault]]] = {}
self.warned: set[tuple[str, Union[UUID, str], str]] = set()
self.unreachable: set[Vault] = set()
def __iter__(self) -> Iterator[Vault]:
return iter(self.vaults)
@@ -70,7 +71,7 @@ class Vaults:
and the first one with the content key wins.
"""
local = [v for v in self.vaults if v.local]
remote = [v for v in self.vaults if not v.local]
remote = [v for v in self.live() if not v.local]
for vault in local:
key = self.query_vault(vault, kid)
if key:
@@ -123,6 +124,7 @@ class Vaults:
return None
except Exception as e:
log.warning(f"Failed to get key from Vault '{vault.name}': {e}")
self.mark_unreachable(vault, e)
if dl:
dl.log_vault_query(
vault.name,
@@ -160,7 +162,7 @@ class Vaults:
"""
self.flagged.add((kid, key))
source = self.sources.pop(kid, None)
for vault in self.vaults:
for vault in self.live():
if not vault.local and (source is None or vault is not source[1]):
continue
try:
@@ -171,7 +173,7 @@ class Vaults:
def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None:
"""Remove a flag from every Vault."""
self.flagged.discard((kid, key))
for vault in self.vaults:
for vault in self.live():
try:
vault.unflag_bad_key(kid, key)
except Exception as e:
@@ -183,7 +185,7 @@ class Vaults:
This method pushes to all Vaults at the same time, so one unreachable Vault costs only its
own timeout and does not delay the others.
"""
vaults = [vault for vault in self.vaults if vault != excluding and not vault.no_push]
vaults = [v for v in self.live() if v != excluding and not v.no_push]
if not vaults:
return 0
with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool:
@@ -201,6 +203,7 @@ class Vaults:
return 0
except Exception as e:
log.warning(f"Failed to add key to Vault '{vault.name}': {e}")
self.mark_unreachable(vault, e)
if dl:
dl.log_vault_query(vault.name, "add_key", kid=str(kid), success=False, error=e)
return 0
@@ -214,7 +217,7 @@ class Vaults:
This method pushes to all Vaults at the same time, so one unreachable Vault costs only its
own timeout and does not delay the others.
"""
vaults = [vault for vault in self.vaults if not vault.no_push]
vaults = [vault for vault in self.live() if not vault.no_push]
if not vaults:
return 0
with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool:
@@ -229,7 +232,23 @@ class Vaults:
return 0
except Exception as e:
log.warning(f"Failed to add keys to Vault '{vault.name}': {e}")
self.mark_unreachable(vault, e)
return 0
def live(self) -> list[Vault]:
"""The vaults this run still uses."""
return [vault for vault in self.vaults if vault not in self.unreachable]
def held_locally(self, kid: Union[UUID, str], key: str) -> bool:
"""True when a local vault already holds this KID:KEY."""
return any(self.query_vault(vault, kid) == key for vault in self.vaults if vault.local)
def mark_unreachable(self, vault: Vault, error: Exception) -> None:
"""Skip a vault that is down or still rate limited for the rest of the run, so it costs one
timeout, not one per KID."""
if isinstance(error, (RateLimited, *UNREACHABLE)) and vault not in self.unreachable:
self.unreachable.add(vault)
log.warning(f"Vault '{vault.name}' is unavailable, so this run does not use it again")
__all__ = ("Vaults",)
+10 -13
View File
@@ -1,12 +1,9 @@
from typing import Iterator, Optional, Union
from uuid import UUID
from requests import Response, Session
from requests.adapters import HTTPAdapter
from requests import HTTPError, Response
from unshackle.core import __version__
from unshackle.core.session import POOL_MAX_SIZE
from unshackle.core.vault import Vault
from unshackle.core.vault import RateLimited, Vault, vault_session
class API(Vault):
@@ -24,11 +21,7 @@ class API(Vault):
super().__init__(name, no_push)
self.uri = uri.rstrip("/")
self.timeout = timeout
self.session = Session()
adapter = HTTPAdapter(pool_connections=POOL_MAX_SIZE, pool_maxsize=POOL_MAX_SIZE)
self.session.mount("https://", adapter)
self.session.mount("http://", adapter)
self.session.headers.update({"User-Agent": f"unshackle v{__version__}"})
self.session = vault_session(timeout)
self.session.headers.update({"Authorization": f"Bearer {token}"})
if headers:
@@ -44,6 +37,8 @@ class API(Vault):
timeout=self.timeout,
allow_redirects=False,
)
if response.status_code == 429:
raise Exceptions.TooManyRequests(f"Vault {self.name} rate limited the lookup")
response.raise_for_status()
data = response.json()
@@ -176,10 +171,12 @@ class API(Vault):
batch_size = 1
continue
if response.status_code == 429:
raise Exceptions.TooManyRequests(f"Vault {self.name} rate limited the push")
response.raise_for_status()
data = response.json()
except Exception:
# JSON decode error or connection issue - try smaller batch
except (ValueError, HTTPError):
# an HTTP error or a broken JSON reply - try smaller batch
if batch_size > 1:
if batch_size > 100:
batch_size -= 100
@@ -281,7 +278,7 @@ class Exceptions:
class AuthRejected(Exception):
"""Authentication Error Occurred, is your token valid? Do you have permission to make this call?"""
class TooManyRequests(Exception):
class TooManyRequests(RateLimited):
"""Rate Limited. Sent too many requests in a given amount of time."""
class PageInvalid(Exception):
+8 -19
View File
@@ -3,12 +3,7 @@ from enum import Enum
from typing import Iterator, Optional, Union
from uuid import UUID
from requests import Session
from requests.adapters import HTTPAdapter
from unshackle.core import __version__
from unshackle.core.session import POOL_MAX_SIZE
from unshackle.core.vault import Vault
from unshackle.core.vault import RateLimited, Vault, vault_session
class InsertResult(Enum):
@@ -61,11 +56,7 @@ class HTTP(Vault):
self.api_mode = api_mode.lower()
self.batch_insert = True
self.current_title = None
self.session = Session()
adapter = HTTPAdapter(pool_connections=POOL_MAX_SIZE, pool_maxsize=POOL_MAX_SIZE)
self.session.mount("https://", adapter)
self.session.mount("http://", adapter)
self.session.headers.update({"User-Agent": f"unshackle v{__version__}"})
self.session = vault_session(timeout)
self.api_session_id = None
if self.api_mode == "decrypt_labs":
@@ -97,6 +88,9 @@ class HTTP(Vault):
if r.status_code == 404:
return {"status": "not_found"}
if r.status_code == 429:
raise RateLimited(f"Vault {self.name} rate limited the request")
if not r.ok:
try:
detail = r.json().get("message") or r.reason.title()
@@ -124,11 +118,8 @@ class HTTP(Vault):
kid = kid.hex
if self.api_mode == "decrypt_labs":
response = self.session.post(self.url, json={"service": service.lower(), "kid": kid}, timeout=self.timeout)
try:
request_payload = {"service": service.lower(), "kid": kid}
response = self.session.post(self.url, json=request_payload, timeout=self.timeout)
if not response.ok:
return None
@@ -186,7 +177,7 @@ class HTTP(Vault):
elif isinstance(key_entry, dict):
if key_entry.get("kid") == kid:
return key_entry.get("key")
except Exception as e:
except ValueError as e:
print(f"Failed to get key ({e.__class__.__name__}: {e})")
return None
return None
@@ -300,8 +291,6 @@ class HTTP(Vault):
except ValueError as e:
if "unknown" in str(e).lower() and "method" in str(e).lower():
self.batch_insert = False
except Exception:
pass
for kid, key in processed_kid_keys.items():
try:
response = self.request(
@@ -317,7 +306,7 @@ class HTTP(Vault):
continue
if response.get("inserted", False):
inserted_count += 1
except Exception as e:
except ValueError as e:
failures.append(e)
if failures and len(failures) == len(processed_kid_keys):
raise failures[-1]