feat(vaults): report bad keys to the vault that served them

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": "<kid_hex>", "key": "<key_hex>" }

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
+      server flags the pair in every local vault it holds, reports it to the vault
-      vault that served it, and the next licence for that KID skips the pair and
+      that served it, and the next licence for that KID skips the pair and reaches
-      reaches the CDM. A server with no local SQLite vault cannot store the flag. The
+      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

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

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 == []

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
+        neither happens. On failure the pair is flagged in the local vaults and reported
-        its source, the ciphertext is restored, and the licence runs again. Every vault skips a
+        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")

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:
+    def flag() -> None:
-            vault.flag_bad_key(vaults.service, UUID(hex=kid), served_key, source)
+        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"

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
+      not decode. The server flags the pair in its local vaults and reports it to the vault that
-      reaches the CDM when a local vault stores the flag. The server accepts only a pair it served
+      served it, so the next licence for that KID reaches the CDM. The server accepts only a pair
-      to this remote session.
+      it served to this remote session.
     parameters:
       - name: session_id
         in: path

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.

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