From aa36d4c2250e62a139ca4951eb4e8db1508df6aa Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sun, 13 Sep 2026 20:46:38 -0600 Subject: [PATCH] feat(vaults): report bad keys to the vault that served them --- docs/dev/rest-api/endpoints.md | 2 +- docs/dev/rest-api/remote-sessions.md | 6 +-- tests/core/test_vault_bad_keys.py | 46 +++++++++++++++++++ .../unit/test_server_cdm_key_verification.py | 14 ++++++ unshackle/commands/dl.py | 7 +-- unshackle/core/api/handlers.py | 14 ++++-- unshackle/core/api/routes.py | 6 +-- unshackle/core/vaults.py | 20 ++++++-- unshackle/vaults/API.py | 42 ++++++++++++++++- 9 files changed, 139 insertions(+), 18 deletions(-) diff --git a/docs/dev/rest-api/endpoints.md b/docs/dev/rest-api/endpoints.md index 78b6ed08..35a043dd 100644 --- a/docs/dev/rest-api/endpoints.md +++ b/docs/dev/rest-api/endpoints.md @@ -1088,7 +1088,7 @@ Get the content keys for the DRM. The `mode` field selects one of two modes. ### `POST /api/session/{session_id}/keys/bad` -Report a content key from `vault_keys` that did not decrypt the track. The server flags the pair in its local vaults, under the name of the vault that served it, and the next licence for that KID reaches the CDM. The client never learns the server vault names. +Report a content key from `vault_keys` that did not decrypt the track. The server flags the pair in its local vaults and reports it to the vault that served it, and the next licence for that KID reaches the CDM. The client never learns the server vault names. ```json { "kid": "", "key": "" } diff --git a/docs/dev/rest-api/remote-sessions.md b/docs/dev/rest-api/remote-sessions.md index eae7017d..5ec73e61 100644 --- a/docs/dev/rest-api/remote-sessions.md +++ b/docs/dev/rest-api/remote-sessions.md @@ -239,9 +239,9 @@ There are two ways DRM keys get resolved, chosen by the client's `server_cdm` fl - A content key the server took from its own vault is unproven. The response lists its KID in `vault_keys`, the client decrypts and checks that the output decodes, and a wrong pair goes back to `POST /api/session/{id}/keys/bad`. The - server flags the pair in every local vault it holds, records the name of the - vault that served it, and the next licence for that KID skips the pair and - reaches the CDM. A server with no local SQLite vault cannot store the flag. The + server flags the pair in every local vault it holds, reports it to the vault + that served it, and the next licence for that KID skips the pair and reaches + the CDM. A server with no local SQLite vault cannot store the flag. The server accepts a report only for a pair it served to that remote session. The client never learns the server vault names. - A content key in the client's own vaults goes first, even when the server diff --git a/tests/core/test_vault_bad_keys.py b/tests/core/test_vault_bad_keys.py index f7091496..ccaf0734 100644 --- a/tests/core/test_vault_bad_keys.py +++ b/tests/core/test_vault_bad_keys.py @@ -1,3 +1,4 @@ +import logging import time from pathlib import Path from types import SimpleNamespace @@ -8,6 +9,7 @@ import pytest from unshackle.commands import dl as dl_module from unshackle.commands.dl import dl from unshackle.core.vaults import Vaults +from unshackle.vaults.API import API from unshackle.vaults.SQLite import SQLite KID = UUID("11111111-2222-3333-4444-555555555555") @@ -331,3 +333,47 @@ def test_a_sibling_verdict_during_the_decrypt_still_triggers_the_retry( assert path.read_bytes() == b"plain:" + GOOD.encode() assert drm.decrypted_with == [{KID: BAD}, {KID: GOOD}] assert licences == [GOOD] # the good key in the run cache survives the flag of the bad one + + +def test_api_vault_reports_the_pair_upstream(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None: + api = API("vault", "https://vault.example/", "tok") + calls: list[tuple[str, str, object]] = [] + ok = SimpleNamespace(status_code=200, raise_for_status=lambda: None, json=lambda: {"code": 0}) + monkeypatch.setattr(api.session, "post", lambda url, **kw: calls.append(("POST", url, kw.get("json"))) or ok) + monkeypatch.setattr(api.session, "delete", lambda url, **kw: calls.append(("DELETE", url, None)) or ok) + + vaults = Vaults("SVC") + vaults.vaults = [api] + vaults.sources[KID] = (BAD, api) + vaults.flag_bad_key(KID, BAD) + vaults.unflag_bad_key(KID, BAD) + + assert calls == [ + ("POST", f"https://vault.example/bad-keys/{KID.hex}", {"content_key": BAD, "service": "SVC"}), + ("DELETE", f"https://vault.example/bad-keys/{KID.hex}/{BAD}", None), + ] + + calls.clear() + vaults.sources[KID] = (BAD, SimpleNamespace(name="other")) + vaults.flag_bad_key(KID, BAD) + assert calls == [] + vaults.sources[KID] = (BAD, api) + + monkeypatch.setattr(api.session, "post", lambda url, **kw: (_ for _ in ()).throw(ConnectionError("down"))) + caplog.clear() + with caplog.at_level(logging.DEBUG, logger="unshackle.core.vaults"): + vaults.flag_bad_key(KID, BAD) + assert [r.levelname for r in caplog.records] == ["DEBUG"] + + gone = SimpleNamespace(status_code=404, raise_for_status=lambda: (_ for _ in ()).throw(RuntimeError("404"))) + monkeypatch.setattr(api.session, "post", lambda url, **kw: gone) + caplog.clear() + vaults.flag_bad_key(KID, BAD) + assert not caplog.records + + denied = SimpleNamespace(status_code=200, raise_for_status=lambda: None, json=lambda: {"code": 1, "message": "no"}) + monkeypatch.setattr(api.session, "post", lambda url, **kw: denied) + monkeypatch.setattr(api.session, "delete", lambda url, **kw: denied) + vaults.flag_bad_key(KID, BAD) + vaults.unflag_bad_key(KID, BAD) + assert not caplog.records diff --git a/tests/remote/unit/test_server_cdm_key_verification.py b/tests/remote/unit/test_server_cdm_key_verification.py index 1eaa817d..3104809b 100644 --- a/tests/remote/unit/test_server_cdm_key_verification.py +++ b/tests/remote/unit/test_server_cdm_key_verification.py @@ -132,3 +132,17 @@ async def test_batch_reports_vault_keys_from_the_init_segment_pssh(monkeypatch): assert payload["keys"] == {"vid": {CDM_KID.hex: "v"}} assert payload["vault_keys"] == [CDM_KID.hex] assert session.served_keys == {CDM_KID.hex: ("v", "sqlite")} + + +async def test_server_tells_a_remote_vault_only_when_it_served_the_pair(monkeypatch, session): + entry, local = session + bystander, source = _Vault(), _Vault() + bystander.local = source.local = False + bystander.name, source.name = "other", "sqlite" + monkeypatch.setattr( + handlers, "load_server_vaults", lambda name: SimpleNamespace(vaults=[local, bystander, source], service="SVC") + ) + await handlers.session_bad_key_handler({"kid": KID.hex, "key": "v"}, "sess") + assert local.flagged == [("SVC", KID, "v", "sqlite")] + assert source.flagged == [("SVC", KID, "v", "sqlite")] + assert bystander.flagged == [] diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 58f4866d..8aae1794 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -4006,11 +4006,11 @@ class dl: A poisoned vault returns the right KID with a wrong content key, and the decrypters accept it without complaint. A content key from a vault is therefore trusted, and only then copied to the other vaults, once FFmpeg can decode the result. Without FFmpeg - neither happens. On failure the pair is flagged in the local vaults with - its source, the ciphertext is restored, and the licence runs again. Every vault skips a + neither happens. On failure the pair is flagged in the local vaults and reported + to the vault that served it, the ciphertext is restored, and the licence runs again. Every vault skips a flagged pair, so each pass burns at most one vault; when the vaults run out the CDM answers, and a CDM key is not checked. A CDM that returns the flagged key clears the - flag, because then the decode check was wrong. + flag and stores the content key again, because then the decode check was wrong. With ``decrypt=False`` the downloader already decrypted the segments in place, so a failure can only flag and raise; the next run skips the flagged pair. @@ -4094,6 +4094,7 @@ class dl: source = self.vaults.sources.get(kid) if source is None: self.vaults.unflag_bad_key(kid, key) + self.flush_vault_writes([partial(self.vaults.add_key, kid, key)]) elif kid in kids and source[1] is kids[kid]: raise ValueError(f"{key} from {kids[kid].name} was bad and no other source has the key") raise ValueError("No vault or CDM produced a content key that decrypts the track") diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 3fd95ed9..802a6826 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -4369,9 +4369,17 @@ async def session_bad_key_handler( raise APIError(APIErrorCode.INVALID_INPUT, "This session was not served that KID:KEY pair") vaults = load_server_vaults(session.service_instance.__class__.__name__) - for vault in vaults.vaults: - if vault.local: - vault.flag_bad_key(vaults.service, UUID(hex=kid), served_key, source) + + def flag() -> None: + for vault in vaults.vaults: + if not vault.local and vault.name != source: + continue + try: + vault.flag_bad_key(vaults.service, UUID(hex=kid), served_key, source) + except Exception as e: + log.debug(f"Could not flag {kid} as bad in vault {sanitize_log(vault.name)}: {e!r}") + + await asyncio.to_thread(flag) session.served_keys.pop(kid, None) log.warning( f"Client proved {kid}:{served_key} from vault {sanitize_log(source)} wrong, flagged in the server vaults" diff --git a/unshackle/core/api/routes.py b/unshackle/core/api/routes.py index caf3d9c9..cf8e5342 100644 --- a/unshackle/core/api/routes.py +++ b/unshackle/core/api/routes.py @@ -1856,9 +1856,9 @@ async def session_bad_key(request: web.Request) -> web.Response: summary: Report a bad content key description: >- The client decrypted with a content key the server took from its vault and the output did - not decode. The server flags the pair in its local vaults, so the next licence for that KID - reaches the CDM when a local vault stores the flag. The server accepts only a pair it served - to this remote session. + not decode. The server flags the pair in its local vaults and reports it to the vault that + served it, so the next licence for that KID reaches the CDM. The server accepts only a pair + it served to this remote session. parameters: - name: session_id in: path diff --git a/unshackle/core/vaults.py b/unshackle/core/vaults.py index 4610b64d..d2fd2e06 100644 --- a/unshackle/core/vaults.py +++ b/unshackle/core/vaults.py @@ -148,19 +148,31 @@ class Vaults: return key if found else None def flag_bad_key(self, kid: Union[UUID, str], key: str) -> None: - """Flag a KID:KEY that failed decryption in every local Vault, naming the Vault it came from.""" + """Flag a KID:KEY that failed decryption, naming the Vault it came from. + + Every local vault records it. A remote vault is told only when it served the pair, + so it stops serving it too, and a vault that never held it is not bothered. A vault + that cannot take the flag is skipped with a debug line: the local flag already + protects this run, and a vault without the route is not a fault. + """ self.flagged.add((kid, key)) source = self.sources.pop(kid, None) for vault in self.vaults: - if vault.local: + if not vault.local and (source is None or vault is not source[1]): + continue + try: vault.flag_bad_key(self.service, kid, key, source[1].name if source else "unknown") + except Exception as e: + log.debug(f"Could not flag {key} as bad in Vault '{vault.name}': {e!r}") def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None: - """Remove a flag from every local Vault.""" + """Remove a flag from every Vault.""" self.flagged.discard((kid, key)) for vault in self.vaults: - if vault.local: + try: vault.unflag_bad_key(kid, key) + except Exception as e: + log.debug(f"Could not unflag {key} in Vault '{vault.name}': {e!r}") def add_key(self, kid: Union[UUID, str], key: str, excluding: Optional[Vault] = None) -> int: """Add a KID:KEY to all Vaults, optionally with an exclusion. diff --git a/unshackle/vaults/API.py b/unshackle/vaults/API.py index af74b62e..51c061c7 100644 --- a/unshackle/vaults/API.py +++ b/unshackle/vaults/API.py @@ -1,7 +1,7 @@ from typing import Iterator, Optional, Union from uuid import UUID -from requests import Session +from requests import Response, Session from unshackle.core import __version__ from unshackle.core.vault import Vault @@ -205,6 +205,46 @@ class API(Vault): return total_added + def flag_bad_key(self, service: str, kid: Union[UUID, str], key: str, source: str) -> None: + """Report a pair that failed to decrypt. The vault only accepts pairs it served.""" + if self.no_push: + return + if isinstance(kid, UUID): + kid = kid.hex + response = self.session.post( + url=f"{self.uri}/bad-keys/{kid}", + json={"content_key": key, "service": service}, + headers={"Accept": "application/json"}, + timeout=self.timeout, + allow_redirects=False, + ) + self._check(response, ignore=frozenset({1, 4})) + + def unflag_bad_key(self, kid: Union[UUID, str], key: str) -> None: + if self.no_push: + return + if isinstance(kid, UUID): + kid = kid.hex + response = self.session.delete( + url=f"{self.uri}/bad-keys/{kid}/{key}", + headers={"Accept": "application/json"}, + timeout=self.timeout, + allow_redirects=False, + ) + self._check(response, ignore=frozenset({1})) + + @staticmethod + def _check(response: Response, ignore: frozenset[int] = frozenset()) -> None: + if response.status_code in (404, 405, 501): + return + response.raise_for_status() + data = response.json() + code = int(data.get("code", 0)) + if not code or code in ignore: + return + error = {1: Exceptions.AuthRejected, 2: Exceptions.TooManyRequests}.get(code, ValueError) + raise error(f"{data.get('message')} ({code})") + def get_services(self) -> Iterator[str]: response = self.session.post( url=self.uri, headers={"Accept": "application/json"}, timeout=self.timeout, allow_redirects=False