From b36db97ffd13cb1f373c1498077a3191fedfc31f Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Tue, 22 Sep 2026 20:23:22 -0600 Subject: [PATCH] fix(serve): stop repeating server vault pushes and retry failed ones --- tests/core/test_vault_push_retry.py | 65 +++++++++++++++++++++++++++++ unshackle/core/api/handlers.py | 17 +++++++- unshackle/vaults/HTTP.py | 7 +++- 3 files changed, 85 insertions(+), 4 deletions(-) create mode 100644 tests/core/test_vault_push_retry.py diff --git a/tests/core/test_vault_push_retry.py b/tests/core/test_vault_push_retry.py new file mode 100644 index 00000000..0ed902b2 --- /dev/null +++ b/tests/core/test_vault_push_retry.py @@ -0,0 +1,65 @@ +"""A failed vault push is reported as a failure, and a server does not store the same pair twice.""" + +from typing import Any +from uuid import UUID + +import pytest + +from unshackle.core.api import handlers +from unshackle.core.vaults import Vaults +from unshackle.vaults.HTTP import HTTP + +KID = UUID("11111111-2222-3333-4444-555555555555") +KEY = "a" * 32 + + +def test_an_unreachable_json_vault_fails_the_push(monkeypatch: pytest.MonkeyPatch) -> None: + vault = HTTP("json", "https://vault.invalid/rpc", "token", api_mode="json") + + def down(method: str, params: dict) -> dict: + raise ConnectionError("vault is down") + + monkeypatch.setattr(vault, "request", down) + vaults = Vaults("SVC") + vaults.vaults = [vault] + assert vaults.add_keys({KID: KEY}) == 0 + + +class FakeVault: + no_push = False + + def __init__(self, ok: bool) -> None: + self.ok = ok + + +def fake_server_vaults(pushes: list, vault_list: list[FakeVault]) -> Any: + class FakeVaults: + vaults = vault_list + + def __len__(self) -> int: + return len(vault_list) + + def add_keys(self, key_map: dict) -> int: + pushes.append(key_map) + return sum(vault.ok for vault in vault_list) + + return FakeVaults() + + +def test_server_does_not_push_a_pair_it_already_stored(monkeypatch: pytest.MonkeyPatch) -> None: + pushes: list = [] + monkeypatch.setattr(handlers, "load_server_vaults", lambda name: fake_server_vaults(pushes, [FakeVault(True)])) + monkeypatch.setattr(handlers, "CACHED_PAIRS", set()) + handlers.cache_to_vaults({KID.hex: KEY}, "SVC") + handlers.cache_to_vaults({KID.hex: KEY}, "SVC") + assert pushes == [{KID: KEY}] + + +def test_server_pushes_again_after_one_vault_failed(monkeypatch: pytest.MonkeyPatch) -> None: + pushes: list = [] + vault_list = [FakeVault(True), FakeVault(False)] + monkeypatch.setattr(handlers, "load_server_vaults", lambda name: fake_server_vaults(pushes, vault_list)) + monkeypatch.setattr(handlers, "CACHED_PAIRS", set()) + handlers.cache_to_vaults({KID.hex: KEY}, "SVC") + handlers.cache_to_vaults({KID.hex: KEY}, "SVC") + assert len(pushes) == 2 diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 1a960536..92cadeee 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -4075,17 +4075,30 @@ def check_vaults(kids: list, service_name: str) -> Optional[tuple[Dict[str, str] return None +CACHED_PAIRS: set[tuple[str, str, str]] = set() + + def cache_to_vaults(keys: Dict[str, str], service_name: str) -> None: - """Cache newly obtained keys to server vaults.""" + """Cache newly obtained keys to server vaults. + + A PlayReady licence always runs, and a Widevine licence runs when one PSSH KID has no content + key in a vault, so the same content keys come back on every request for a title. This process + does not send a pair again after every vault accepted it. + """ from uuid import UUID try: + new = {kid: key for kid, key in keys.items() if (service_name, kid, key) not in CACHED_PAIRS} + if not new: + return vaults = load_server_vaults(service_name) if not vaults.vaults: return - key_map = {UUID(hex=kid): key for kid, key in keys.items()} + key_map = {UUID(hex=kid): key for kid, key in new.items()} cached = vaults.add_keys(key_map) + if cached == len([vault for vault in vaults.vaults if not vault.no_push]): + CACHED_PAIRS.update((service_name, kid, key) for kid, key in new.items()) if cached: log.info(f"Cached {len(key_map)} key(s) to {cached}/{len(vaults)} server vault(s)") except (Exception, SystemExit) as e: diff --git a/unshackle/vaults/HTTP.py b/unshackle/vaults/HTTP.py index 2a0561ef..1b882c1c 100644 --- a/unshackle/vaults/HTTP.py +++ b/unshackle/vaults/HTTP.py @@ -281,6 +281,7 @@ class HTTP(Vault): inserted_count = 0 title = getattr(self, "current_title", None) + failures: list[Exception] = [] if self.api_mode == "json": if self.batch_insert: @@ -311,8 +312,10 @@ class HTTP(Vault): continue if response.get("inserted", False): inserted_count += 1 - except Exception: - continue + except Exception as e: + failures.append(e) + if failures and len(failures) == len(processed_kid_keys): + raise failures[-1] else: for kid, key in processed_kid_keys.items(): response = self.session.get(