feat(drm): license a track only when its own KID has no key

A vault key for the track's KID now stands in for the licence on the client and the server (PlayReady included). decrypt_verified licenses any other header KID the file uses.

docs/guide/vaults.md

@@ -25,6 +25,8 @@ A vault stores every content key as a **KID → KEY** pair, grouped by **service
 
 When unshackle needs a content key during a download it asks your vaults **in the order they appear in your config**. It returns the first non-null key it finds and stops searching. If a vault raises a permission error or is unreachable, unshackle skips it and continues the search with the next one.
 
+A track needs no licence when a vault holds the content key of the KID that encrypts it. The DRM header of the track can name more KIDs, for example those of other tracks. Before the decrypt, unshackle reads the downloaded file, and licenses each of those KIDs that the file uses and no vault holds. When unshackle cannot read which KIDs the file uses, it licenses all of them.
+
 When unshackle recovers new keys from a license, it pushes them to **all** configured vaults, so every vault gradually converges on the same set of keys. Two rules apply to pushing:
 
 - unshackle skips keys that already exist in a vault (it never overwrites or deletes existing data).

docs/reference/configuration/services.md

@@ -223,9 +223,10 @@ the remote client's `-q` when the remote session starts:
   error. The server checks the values the job runs with, after it applies the service defaults
   and the `serve:` download defaults.
 
-The limit protects the server's device, not the content keys. For a Widevine track, a server vault
-still supplies a content key at every height. A PlayReady track never reads the vault, so above
-the limit it always goes to the client's device. A licence for a lower track can carry the
+The limit protects the server's device, not the content keys. A server vault still supplies a
+content key at every height: for a Widevine track, the key of the track's own KID or of every KID in
+its header; for a PlayReady track, only the key of the track's own KID. A PlayReady track whose
+own KID the server does not know always goes to the client's device above the limit. A licence for a lower track can carry the
 content key of a higher track when the service shares one content key between them. With server
 accounts enabled, a client that reports a stronger device than it holds makes the server's
 account request those manifests; the licence server still refuses its challenge.
@@ -233,7 +234,8 @@ account request those manifests; the licence server still refuses its challenge.
 `server_vault` lets a remote session that the client's own device licenses take content keys
 from the server vault. Set it to `true` for every service, or to a list of service tags. It is
 `false` unless the entry sets it, and `server_cdm` for a service implies it. The server looks up
-every KID of the track in its vault and makes no live licence. A track that the vault does not
+the track's own KID in its vault, or every KID of the track when it does not know that KID, and
+makes no live licence. A track that the vault does not
 hold goes back to the client's device. When the service runs the remote session on the server's
 own device (a service option, such as a server-identity flag), the client's own vaults, and the
 server vault with this grant, are the only source, and a missing content key stops the download.

tests/core/test_licence_turns.py

@@ -72,7 +72,7 @@ def widevine(*kids: UUID) -> Widevine:
     return Widevine(pssh=PSSH.new(system_id=PSSH.SystemId.Widevine, key_ids=list(kids)))
 
 
-def prepare(cmd: dl, drm: Widevine, track_kid: UUID, table: Table, **kwargs: Any) -> None:
+def prepare(cmd: dl, drm: Widevine, track_kid: Optional[UUID], table: Table, **kwargs: Any) -> None:
     cmd.prepare_drm(
         drm,
         track=SimpleNamespace(drm=[drm]),
@@ -220,7 +220,25 @@ def test_a_refused_licence_without_the_track_key_reports_the_licence_error(
     assert "Licence request failed" in str(tree.children[-1].label)
 
 
-def test_a_refused_licence_keeps_the_known_track_key(
+def test_a_held_track_key_needs_no_licence(cmd: dl, licences: list[str], monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(dl, "LICENSE_KEY_CACHE", {KID_A: KEYS[KID_A]})
+    drm = widevine(KID_A, KID_B)
+
+    prepare(cmd, drm, KID_A, new_table())
+
+    assert licences == []
+    assert drm.content_keys == {KID_A: KEYS[KID_A]}
+
+
+def test_force_licenses_a_held_track_key_again(cmd: dl, licences: list[str], monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(dl, "LICENSE_KEY_CACHE", {KID_A: KEYS[KID_A]})
+
+    prepare(cmd, widevine(KID_A, KID_B), KID_A, new_table(), force=True)
+
+    assert len(licences) == 1
+
+
+def test_a_refused_licence_keeps_the_known_keys_when_the_track_kid_is_unknown(
     cmd: dl, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
 ) -> None:
     monkeypatch.setattr(Widevine, "get_content_keys", refuse_licence)
@@ -228,7 +246,7 @@ def test_a_refused_licence_keeps_the_known_track_key(
     monkeypatch.setattr(dl, "DRM_LOCKS", {})
     drm = widevine(KID_A, KID_B)
 
-    prepare(cmd, drm, KID_A, new_table())
+    prepare(cmd, drm, None, new_table())
 
     assert drm.content_keys == {KID_A: KEYS[KID_A]}
     assert [r.getMessage() for r in caplog.records if r.levelname == "WARNING"] == [

tests/core/test_mp4decrypt_track_key.py

@@ -131,7 +131,18 @@ def test_segment_decrypt_needs_a_key_for_a_kidless_track(
     drm.content_keys[key_kid] = KEY
 
     assert can_use(drm, "mp4decrypt", movie(two_sinf(A, piff_first=True))) is usable
-    assert can_use(drm, "mp4decrypt", movie(encv(A))) is True
+    assert can_use(drm, "mp4decrypt", movie(encv(A))) is usable
+
+
+def test_segment_decrypt_waits_for_a_key_to_every_kid_the_drm_names(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(config, "decrypt_segments", True, raising=False)
+    monkeypatch.setattr(binaries, "Mp4decrypt", tmp_path / "mp4decrypt")
+    drm = _widevine(A, B)
+    drm.content_keys[A] = KEY
+
+    assert can_use(drm, "mp4decrypt", movie(encv(A))) is False
 
 
 class StubDRM:

tests/core/test_vault_bad_keys.py

@@ -384,6 +384,32 @@ def test_stale_backup_from_a_killed_run_is_not_restored(tmp_path: Path, monkeypa
     assert path.read_bytes() == b"cipher"
 
 
+@pytest.mark.parametrize(
+    ("windows", "licensed"),
+    [({KID: [(0.0, 2.0)], SIBLING: [(2.0, 4.0)]}, [SIBLING]), ({KID: [(0.0, 2.0)]}, []), (None, [SIBLING])],
+    ids=["file-uses-it", "file-does-not", "unmapped"],
+)
+def test_an_unkeyed_kid_the_file_uses_is_licensed_before_the_decrypt(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, windows: Optional[dict], licensed: list[UUID]
+) -> None:
+    cmd = make_cmd(monkeypatch, Vaults("SVC"))
+    drm = FakeDRM({KID: GOOD})
+    drm.kids = [KID, SIBLING]
+    monkeypatch.setattr(
+        dl_module.verify,
+        "kid_windows",
+        lambda path: None if windows is None else dl_module.verify.KidMap(windows, True),
+    )
+    path = tmp_path / "video.mp4"
+    path.write_bytes(b"cipher")
+    licences: list[UUID] = []
+
+    cmd.decrypt_verified(drm, path, lambda drm, track_kid: licences.append(track_kid), KID)
+
+    assert licences == licensed
+    assert drm.decrypted_with == [{KID: GOOD}]
+
+
 def test_drm_without_content_keys_still_decrypts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
     """HLS AES-128 ClearKey has no KID:KEY map; the hook must pass it straight through."""
     cmd = make_cmd(monkeypatch, Vaults("SVC"))

tests/remote/unit/test_extract_pssh_strict.py

@@ -56,7 +56,15 @@ async def _run_batch(monkeypatch, tracks, cdm_type, calls):
         return session
 
     def fake_single(
-        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, vault_only=False, shared=None
+        service,
+        title,
+        track,
+        pssh_str,
+        drm_type,
+        request,
+        sources=None,
+        refusal=None,
+        **_,
     ):
         calls.append((pssh_str, drm_type))
         return {"00" * 16: "key"}

tests/remote/unit/test_remote_service_server_cdm_license.py

@@ -132,3 +132,12 @@ def test_a_refused_server_cdm_licence_logs_one_warning(caplog):
 
     assert svc.proxy_license(b"chal", track, "widevine") == b"chal"
     assert [r.levelname for r in caplog.records if r.levelname in ("WARNING", "ERROR")] == ["WARNING"]
+
+
+def test_the_kid_dl_needs_goes_to_the_server():
+    svc = _service("widevine", {"keys": {KID.hex: "k"}, "drm_type": "widevine"})
+    track = SimpleNamespace(id="vid", drm=[_Widevine()])
+
+    svc.get_widevine_license(challenge=b"", title=None, track=track, kid=KID)
+
+    assert svc.client.posts[0]["kid"] == KID.hex

tests/remote/unit/test_server_cdm_key_verification.py

@@ -115,12 +115,25 @@ async def test_batch_reports_vault_keys_from_the_init_segment_pssh(monkeypatch):
     session = SessionEntry(session_id="sess", service_tag="SVC", service_instance=SimpleNamespace())
     session.tracks = {"vid": track}
 
+    probed: list = []
+
     async def validated(sid, req):
         return session
 
     def fake_single(
-        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, vault_only=False, shared=None
+        service,
+        title,
+        track,
+        pssh_str,
+        drm_type,
+        request,
+        sources=None,
+        refusal=None,
+        vault_only=False,
+        shared=None,
+        track_kid=None,
     ):
+        probed.append(track_kid)
         if pssh_str == "init":
             sources[CDM_KID.hex] = "sqlite"
             return {CDM_KID.hex: "v"}
@@ -147,6 +160,66 @@ 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")}
+    assert probed == [CDM_KID, CDM_KID]
+
+
+async def test_batch_asks_again_when_a_cached_vault_answer_lacks_the_track_kid(monkeypatch):
+    """Renditions that share a header: a vault answer for one track's KID does not serve a track with another KID."""
+    kids = {"a": KID, "b": CDM_KID, "c": CDM_KID}
+    session = SessionEntry(session_id="sess", service_tag="SVC", service_instance=SimpleNamespace())
+    session.tracks = {
+        tid: SimpleNamespace(id=tid, drm=[_WV("shared", [KID, CDM_KID])], get_key_id=lambda init, k=kid: k)
+        for tid, kid in kids.items()
+    }
+    asked: list = []
+
+    async def validated(sid, req):
+        return session
+
+    def fake_single(
+        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, track_kid=None, **_
+    ):
+        asked.append(track_kid)
+        if track_kid == KID:
+            sources[KID.hex] = "sqlite"
+            return {KID.hex: "v"}
+        return {KID.hex: "v", CDM_KID.hex: "c"}
+
+    monkeypatch.setattr(handlers, "get_validated_session", validated)
+    monkeypatch.setattr(handlers, "server_cdm_allowed", lambda request, tag: True)
+    monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None)
+    monkeypatch.setattr(handlers, "find_title_for_track", lambda tid, session: SimpleNamespace())
+    monkeypatch.setattr(handlers, "detect_cdm_type_for_service", lambda tag, cfg: "widevine")
+    monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev")
+    monkeypatch.setattr(handlers, "fetch_init_segment", lambda track, sess: b"init")
+    monkeypatch.setattr(handlers, "handle_single_server_cdm", fake_single)
+    monkeypatch.setattr(handlers.config, "serve", {"users": {}}, raising=False)
+
+    resp = await handlers.session_license_handler(
+        {"mode": "server_cdm", "track_ids": list(kids), "drm_type": "widevine"}, "sess", None
+    )
+
+    payload = json.loads(resp.body)
+    assert asked == [KID, CDM_KID]
+    assert payload["keys"]["c"] == {KID.hex: "v", CDM_KID.hex: "c"}
+
+
+@pytest.mark.parametrize("drm", [["drm"], []], ids=["with-drm", "clear"])
+def test_the_server_reads_the_head_of_a_one_file_track_with_drm(drm):
+    from unshackle.core.tracks import Video
+
+    fetched: list = []
+    track = SimpleNamespace(
+        id="v",
+        data={},
+        descriptor=Video.Descriptor.URL,
+        drm=drm,
+        get_init_segment=lambda session: fetched.append(session) or b"head",
+    )
+
+    head = handlers.fetch_init_segment(track, "svc-session")
+
+    assert (head, fetched) == ((b"head", ["svc-session"]) if drm else (None, []))
 
 
 async def test_server_tells_a_remote_vault_only_when_it_served_the_pair(monkeypatch, session):

tests/remote/unit/test_server_cdm_manifest_kid.py

@@ -27,6 +27,8 @@ from unshackle.core.drm import PlayReady
 
 pytestmark = pytest.mark.unit
 
+REAL_CHECK_VAULTS = handlers.check_vaults
+
 KID_A = UUID("161219ec3df64e0eadbb3827c0ccc46d")
 KID_B = UUID("43de497c6d1945f49acc2951696ec628")
 
@@ -161,3 +163,59 @@ async def test_uncovered_kid_reaches_the_client(server_env, fake_playready, monk
     )
 
     assert any(KID_B.hex in r["message"] for r in buf.since(0)), "client gets no reason for the missing key"
+
+
+class _VaultsHolding:
+    """Server vaults that hold a content key for ``held`` only."""
+
+    def __init__(self, *held: UUID):
+        self.held = held
+        self.vaults = [SimpleNamespace(name="v1")]
+
+    def get_key(self, kid):
+        return ("k" + kid.hex[:4], self.vaults[0]) if kid in self.held else (None, None)
+
+
+@pytest.mark.parametrize(
+    ("track_kid", "found"), [(None, None), (KID_A, {KID_A.hex}), (KID_B, None)], ids=["every-kid", "held", "not-held"]
+)
+def test_check_vaults_answers_for_the_track_kid(monkeypatch, track_kid, found):
+    monkeypatch.setattr(handlers, "load_server_vaults", lambda name: _VaultsHolding(KID_A))
+
+    hit = handlers.check_vaults([KID_A, KID_B], "SVC", track_kid=track_kid)
+
+    assert (set(hit[0]) if hit else None) == found
+
+
+@pytest.mark.parametrize("track_kid", [KID_A, None], ids=["track-kid-held", "track-kid-unknown"])
+def test_a_playready_track_takes_the_vault_key_only_for_its_known_kid(server_env, monkeypatch, track_kid):
+    monkeypatch.setattr(handlers, "check_vaults", REAL_CHECK_VAULTS)
+    monkeypatch.setattr(handlers, "load_server_vaults", lambda name: _VaultsHolding(KID_A))
+    licensed: list = []
+    monkeypatch.setattr(
+        PlayReady,
+        "get_content_keys",
+        lambda self, cdm, certificate, licence: licensed.append(self) or self.content_keys.update({KID_A: "cdm"}),
+    )
+    pro_b64 = _kidless_pro_b64()
+    track = SimpleNamespace(
+        id="vid", drm=[PlayReady(pssh=RealPSSH(base64.b64decode(pro_b64)), kid=KID_A, pssh_b64=pro_b64)]
+    )
+    sources: dict = {}
+
+    keys = handlers.handle_single_server_cdm(
+        SimpleNamespace(), SimpleNamespace(), track, pro_b64, "playready", None, sources, track_kid=track_kid
+    )
+
+    assert bool(licensed) is (track_kid is None)
+    assert sources == ({KID_A.hex: "v1"} if track_kid else {})
+    assert set(keys) == {KID_A.hex}
+
+
+def test_the_client_kid_counts_only_when_the_track_drm_names_it():
+    track = SimpleNamespace(drm=[SimpleNamespace(kids=[KID_A])])
+
+    assert handlers.client_track_kid(KID_A.hex, track) == KID_A
+    assert handlers.client_track_kid(KID_B.hex, track) is None
+    assert handlers.client_track_kid("not-hex", track) is None
+    assert handlers.client_track_kid(None, track) is None

tests/remote/unit/test_server_cdm_max_height.py

@@ -397,7 +397,15 @@ def licence_env(monkeypatch):
     vault = {}
 
     def fake_single(
-        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, vault_only=False, shared=None
+        service,
+        title,
+        track,
+        pssh_str,
+        drm_type,
+        request,
+        sources=None,
+        refusal=None,
+        **_,
     ):
         calls.append((track.id, pssh_str))
         if pssh_str in vault:

tests/remote/unit/test_server_cdm_multikey.py

@@ -134,7 +134,15 @@ async def test_batch_shares_full_bundle_per_pssh(monkeypatch):
     bundle_b = {"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa": "k_b"}
 
     def fake_single(
-        service, title, track, pssh_str, drm_type, request, sources=None, refusal=None, vault_only=False, shared=None
+        service,
+        title,
+        track,
+        pssh_str,
+        drm_type,
+        request,
+        sources=None,
+        refusal=None,
+        **_,
     ):
         calls.append(pssh_str)
         return dict(bundle_a) if pssh_str == pssh_a else dict(bundle_b)

unshackle/commands/dl.py

@@ -4624,6 +4624,9 @@ class dl:
 
         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.
+
+        Before the decrypt, each KID the DRM names without a content key is licensed when the file
+        uses it, or when the file cannot be mapped, because ``prepare_drm`` licensed only for the track's KID.
         """
         keys: Any = getattr(drm, "content_keys", None)
         if keys is None:
@@ -4675,6 +4678,15 @@ class dl:
         if known_bad and licence:
             relicense({kid: drop(kid) for kid in known_bad}, {})
 
+        unkeyed = [kid for kid in getattr(drm, "kids", None) or [] if kid not in keys]
+        mapped = bool(unkeyed and decrypt and licence)
+        kid_map = verify.kid_windows(path) if mapped else None
+        if mapped and licence:
+            for kid in unkeyed:
+                if kid_map is None or kid in kid_map.windows:
+                    licence(drm, track_kid=kid)
+            keys = getattr(drm, "content_keys", keys)
+
         def vault_kids() -> dict[UUID, Vault]:
             """The KIDs whose key on the track came from a vault, with that vault."""
             found = {kid: self.vaults.sources.get(kid) for kid in keys}
@@ -4688,7 +4700,7 @@ class dl:
                 os.link(path, backup)
             except OSError as e:
                 self.log.debug(f"Cannot keep the ciphertext for a decrypt retry: {e!r}")
-        kid_map = verify.kid_windows(path) if checking else None
+        kid_map = (kid_map if mapped else verify.kid_windows(path)) if checking else None
         windows = kid_map.windows if kid_map else {}
         video = kid_map.video if kid_map else None
 
@@ -5005,6 +5017,8 @@ class dl:
     ) -> None:
         """Give a Widevine, PlayReady or ClearKey DRM its content keys: run cache, then key vaults, then ``get_keys``.
 
+        A held key for ``track_kid`` skips the licence, as the file is encrypted with that KID;
+        ``decrypt_verified`` licenses any other KID the file turns out to use.
         ``force`` skips the run cache, so the licence runs again for each KID no vault answers.
         A licence can answer a KID with a blank or null key; a vault key for that KID wins.
         """
@@ -5060,7 +5074,7 @@ class dl:
                 show_error(msg)
                 raise drm.Exceptions.CEKNotFound(msg)
 
-            need_license = bool(missing)
+            need_license = bool(missing) and (force or not track_kid or track_kid not in drm.content_keys)
             if need_license:
                 # force licenses again, so the turn takes no cached keys
                 self.take_licence_turn(licence_turn, drm, [] if force else all_kids, track, track_kid, cek_tree)
@@ -5315,6 +5329,7 @@ class dl:
                         licence(
                             drm_system="playready" if drm.__class__.__name__ == "PlayReady" else "widevine",
                             challenge=b"",
+                            kid=track_kid,
                         )
                     except click.ClickException:
                         raise

unshackle/core/api/handlers.py

@@ -4635,17 +4635,23 @@ def extract_pssh_from_track(track: Any, drm_type: str) -> Optional[str]:
 
 
 def fetch_init_segment(track: Any, session: Any = None) -> Optional[bytes]:
-    """The init segment of a DASH track; None for a track that is not DASH, or when the fetch fails."""
+    """The init segment of a DASH track, or the head of a one-file track with DRM; None otherwise, or when the fetch fails."""
     from unshackle.core.manifests import DASH as DASHManifest
 
-    dash = (getattr(track, "data", None) or {}).get("dash") or {}
-    if not all(dash.get(k) is not None for k in ("period", "adaptation_set", "representation", "manifest")):
-        return None
-
     session = getattr(track, "session", None) or session
     if session is None:
         return None
 
+    dash = (getattr(track, "data", None) or {}).get("dash") or {}
+    if not all(dash.get(k) is not None for k in ("period", "adaptation_set", "representation", "manifest")):
+        if getattr(track, "descriptor", None) is not Video.Descriptor.URL or not getattr(track, "drm", None):
+            return None
+        try:
+            return track.get_init_segment(session=session)
+        except Exception as e:
+            log.debug(f"Init segment fetch failed for {track.id}: {e!r}")
+            return None
+
     try:
         return DASHManifest.get_period_segments(
             period=dash["period"],
@@ -4856,28 +4862,29 @@ def once(shared: Dict[Any, Any], key: Any, make: Callable[[], Any]) -> Any:
 
 
 def check_vaults(
-    kids: list, service_name: str, shared: Optional[Dict[str, Any]] = None
+    kids: list, service_name: str, shared: Optional[Dict[str, Any]] = None, track_kid: Optional[uuid.UUID] = None
 ) -> Optional[tuple[Dict[str, str], Dict[str, str]]]:
-    """Examine the server vaults for existing content keys that match all KIDs.
+    """Examine the server vaults for existing content keys that match the KIDs.
 
-    Returns `(KID:KEY, KID:vault name)` if ALL KIDs are found, None otherwise. The vault
-    name travels to the client, which is the only side that can prove the content key decrypts.
+    Returns `(KID:KEY, KID:vault name)` when every KID is found, or with ``track_kid`` when that
+    KID is found, with the other KIDs the vaults hold; None otherwise. The vault name travels to
+    the client, which is the only side that can prove the content key decrypts.
     """
-    from uuid import UUID
-
     try:
         vaults = once({} if shared is None else shared, "vaults", lambda: load_server_vaults(service_name))
         if not vaults.vaults:
             return None
         keys: Dict[str, str] = {}
         sources: Dict[str, str] = {}
-        for kid in kids:
-            kid_uuid = kid if isinstance(kid, UUID) else UUID(hex=str(kid))
+        wanted = [kid if isinstance(kid, uuid.UUID) else uuid.UUID(hex=str(kid)) for kid in kids]
+        if track_kid:
+            wanted = [track_kid, *(kid for kid in wanted if kid != track_kid)]
+        for kid_uuid in wanted:
             content_key, vault_used = vaults.get_key(kid_uuid)
             if content_key:
                 keys[kid_uuid.hex] = content_key
                 sources[kid_uuid.hex] = vault_used.name if vault_used else "unknown"
-            else:
+            elif track_kid in (None, kid_uuid):
                 return None
         if keys:
             log.info(f"Vault hit: {len(keys)} key(s) from server vaults, skipping CDM")
@@ -4952,6 +4959,17 @@ def license_with_track_fallback(drm: Any, track: Any, **kwargs: Any) -> Dict[str
     return {kid: key for kid, key in keys.items() if not is_null_key(key)}
 
 
+def client_track_kid(kid_hex: Any, track: Any) -> Optional[uuid.UUID]:
+    """The KID a client asks the server CDM for, when the track's DRM names it; None for any other value."""
+    if not isinstance(kid_hex, str):
+        return None
+    try:
+        kid = uuid.UUID(hex=kid_hex)
+    except ValueError:
+        return None
+    return kid if any(kid in (getattr(d, "kids", None) or []) for d in getattr(track, "drm", None) or []) else None
+
+
 def pssh_kids(pssh_b64: str, drm_type: str) -> set:
     """The KIDs a base64 PSSH names, for either DRM system."""
     if drm_type == "playready":
@@ -5000,13 +5018,15 @@ def handle_single_server_cdm(
     refusal: Optional[APIError] = None,
     vault_only: bool = False,
     shared: Optional[Dict[str, Any]] = None,
+    track_kid: Optional[uuid.UUID] = None,
 ) -> Dict[str, str]:
     """Do the single-track server_cdm licensing with the DRM class get_content_keys() flow.
 
     ``sources`` is filled with the vault name for every returned content key a server vault
-    supplied. A content key the CDM licensed gets no entry. It raises ``refusal`` before a live licence:
-    after a Widevine vault miss, or at once for PlayReady, whose licence can carry keys its header does not name.
-    With ``vault_only`` no server device is involved: the header's KIDs come from the vault, else ``refusal``.
+    supplied. A content key the CDM licensed gets no entry. A vault answers in place of the licence
+    when it holds ``track_kid``, the KID the track is encrypted with, or else every KID of the header.
+    PlayReady without ``track_kid`` always licenses, as its licence can carry keys its header does not name.
+    It raises ``refusal`` before a live licence. With ``vault_only`` no server device is involved.
     ``shared`` keeps the vaults and CDMs one ``license_response`` call built, for its next track.
     """
     import base64
@@ -5023,12 +5043,16 @@ def handle_single_server_cdm(
     if not pssh_b64:
         raise APIError(APIErrorCode.INVALID_INPUT, "No PSSH available for server_cdm licensing")
 
+    def from_vaults(kids: Iterable[uuid.UUID]) -> Optional[Dict[str, str]]:
+        """The vault keys for ``kids`` (see ``check_vaults``), with their sources noted."""
+        hit = check_vaults(list(kids), service_name, shared=shared, track_kid=track_kid)
+        if hit and sources is not None:
+            sources.update(hit[1])
+        return hit[0] if hit else None
+
     if vault_only:
-        vault_hit = check_vaults(list(pssh_kids(pssh_b64, drm_type)), service_name, shared=shared)
-        if vault_hit:
-            if sources is not None:
-                sources.update(vault_hit[1])
-            return vault_hit[0]
+        if vault_keys := from_vaults(pssh_kids(pssh_b64, drm_type)):
+            return vault_keys
         raise refusal or APIError(APIErrorCode.NO_CONTENT, "The server vault holds no content key for this track")
 
     api_key = caller_key(request)
@@ -5049,6 +5073,8 @@ def handle_single_server_cdm(
         # Gate on the caller's CDM device first: no device, no keys from the vault or CDM.
         device_name = resolve_device_name(user_config, drm_type, service_name)
 
+        if track_kid and (vault_keys := from_vaults(pr_drm.kids)):
+            return vault_keys
         if refusal:
             raise refusal
 
@@ -5076,11 +5102,7 @@ def handle_single_server_cdm(
         # harvest server-side keys from the vault fallback below.
         device_name = resolve_device_name(user_config, drm_type, service_name)
 
-        vault_hit = check_vaults(wv_drm.kids, service_name, shared=shared)
-        if vault_hit:
-            vault_keys, vault_sources = vault_hit
-            if sources is not None:
-                sources.update(vault_sources)
+        if vault_keys := from_vaults(wv_drm.kids):
             return vault_keys
         if refusal:
             raise refusal
@@ -5353,12 +5375,13 @@ def license_response(
                 session.log_buffer.append(logging.WARNING, message)
 
         def license_track(
-            track: Any, title: Any, candidates: list
+            track: Any, title: Any, candidates: list, track_kid: Optional[uuid.UUID]
         ) -> tuple[Dict[str, str], Optional[str], Optional[str]]:
             """`(keys, drm_type, pssh)` for the track's current DRM, licensing once per unique PSSH.
 
-            A track the server must not license live goes into ``capped_tracks`` with no keys, and its PSSH
-            stays uncached, as a track under the cap can share it.
+            A cached vault answer that lacks ``track_kid`` is asked again for that KID. A track the server
+            must not license live goes into ``capped_tracks`` with no keys, and its PSSH stays uncached,
+            as a track under the cap can share it.
             """
             picked = pick_server_pssh(track, candidates)
             if not picked:
@@ -5367,9 +5390,11 @@ def license_response(
             candidate, pssh_str = picked
 
             cache_key = (pssh_str, pssh_set(track))
-            if cache_key not in keys_by_pssh:
-                keys_by_pssh[cache_key] = {}
-                sources_by_pssh[cache_key] = {}
+            cached = keys_by_pssh.get(cache_key)
+            from_vault = bool(cached) and set(cached or {}) <= set(sources_by_pssh.get(cache_key, {}))
+            if cached is None or (track_kid and track_kid.hex not in cached and from_vault):
+                keys_by_pssh.setdefault(cache_key, {})
+                sources_by_pssh.setdefault(cache_key, {})
                 try:
                     keys = handle_single_server_cdm(
                         service,
@@ -5382,9 +5407,10 @@ def license_response(
                         live_licence_refusal(session, track),
                         vault_only=vault_only,
                         shared=shared,
+                        track_kid=track_kid,
                     )
                     if keys:
-                        keys_by_pssh[cache_key] = keys
+                        keys_by_pssh[cache_key] = {**keys_by_pssh[cache_key], **keys}
                         drm_type_by_pssh[cache_key] = candidate
                 except SystemExit:
                     warn(f"Service exited while resolving keys for track {track_log_label(track)}, skipping")
@@ -5417,34 +5443,34 @@ def license_response(
                 else server_drm_candidates(service_tag, user_config, drm_type, track, warn)
             )
 
-            keys, track_drm_type, pssh_str = license_track(track, title, candidates)
+            probed: Optional[uuid.UUID] = None
+            if init_data:
+                try:
+                    probed = track.get_key_id(init_data)
+                except Exception as e:
+                    log.debug(f"KID probe failed for {sanitize_log(tid[:12])}: {e!r}")
+            keys, track_drm_type, pssh_str = license_track(track, title, candidates, probed)
             # a refused track still tries its init segment PSSH, where a server vault can hold the key
             refusal = capped_tracks.pop(tid, None)
 
-            track_kid = None
-            if init_data:
-                try:
-                    track_kid = track.get_key_id(init_data)
-                except Exception as e:
-                    log.debug(f"KID probe failed for {sanitize_log(tid[:12])}: {e!r}")
-            if track_kid and track_kid.hex not in keys:
+            if probed and probed.hex not in keys:
                 manifest_drm = track.drm
                 track.drm = drm_from_init_segment(track, init_data=init_data) or manifest_drm
-                init_keys, init_drm_type, init_pssh = license_track(track, title, candidates)
+                init_keys, init_drm_type, init_pssh = license_track(track, title, candidates, probed)
                 init_refusal = capped_tracks.pop(tid, None)
                 if init_pssh != pssh_str:
                     log.info(
-                        f"The manifest PSSH gave no content key for KID {track_kid.hex} of track "
+                        f"The manifest PSSH gave no content key for KID {probed.hex} of track "
                         f"{track_log_label(track)}, tried the init segment PSSH"
                     )
-                if track_kid.hex in init_keys:
+                if probed.hex in init_keys:
                     keys, track_drm_type, pssh_str = init_keys, init_drm_type, init_pssh
                     refusal = None
                 else:
                     track.drm = manifest_drm
                     refusal = refusal or init_refusal
                     if not refusal:
-                        warn(f"No content key for KID {track_kid.hex} of track {track_log_label(track)}")
+                        warn(f"No content key for KID {probed.hex} of track {track_log_label(track)}")
             if refusal:
                 capped_tracks[tid] = refusal
                 continue
@@ -5507,6 +5533,7 @@ def license_response(
                     f"No PSSH on track {track_log_label(track)} for {', '.join(candidates) or 'any CDM'}",
                 )
             server_drm_type, server_pssh = picked
+            asked_kid = client_track_kid(data.get("kid"), track)
             if pssh_b64 and server_drm_type != drm_type:
                 log.info(
                     f"Client asked for {sanitize_log(drm_type)} on track {track_log_label(track)}, "
@@ -5528,6 +5555,7 @@ def license_response(
                 key_sources,
                 live_licence_refusal(session, track),
                 vault_only=vault_only,
+                track_kid=asked_kid,
             )
             log.info(f"Server CDM resolved {len(keys)} key(s) for track {track_log_label(track)}")
             note_served_keys(session, keys, key_sources)

unshackle/core/drm/segment_decrypt.py

@@ -24,6 +24,7 @@ def can_use(drm: Any, decryption: str, init: Optional[bytes] = None) -> bool:
     """Whether unshackle can decrypt a track's segments one by one as they download.
 
     A track whose ``init`` mp4decrypt cannot key goes to the whole-file decrypt, which has the shaka fallback.
+    So does a track whose DRM names a KID with no content key, as only that decrypt can license a KID the file uses.
     """
     return (
         isinstance(drm, (Widevine, PlayReady, ClearKeyCENC))
@@ -31,6 +32,7 @@ def can_use(drm: Any, decryption: str, init: Optional[bytes] = None) -> bool:
         and config.decrypt_segments
         and bool(binaries.Mp4decrypt)
         and not (init and set(kidless_track_kids(init).values()) - set(drm.content_keys))
+        and set(getattr(drm, "kids", None) or []) <= set(drm.content_keys)
     )
 
 

unshackle/core/remote_service.py

@@ -2081,13 +2081,15 @@ class RemoteService:
         raw = self._chapters_by_title.get(title_id, [])
         return Chapters([Chapter(ch["timestamp"], ch.get("name")) for ch in raw])
 
-    def get_widevine_license(self, *, challenge: bytes, title: Title_T, track: AnyTrack) -> Optional[Union[bytes, str]]:
-        return self.proxy_license(challenge, track, "widevine")
+    def get_widevine_license(
+        self, *, challenge: bytes, title: Title_T, track: AnyTrack, kid: Optional[UUID] = None
+    ) -> Optional[Union[bytes, str]]:
+        return self.proxy_license(challenge, track, "widevine", kid)
 
     def get_playready_license(
-        self, *, challenge: bytes, title: Title_T, track: AnyTrack
+        self, *, challenge: bytes, title: Title_T, track: AnyTrack, kid: Optional[UUID] = None
     ) -> Optional[Union[bytes, str]]:
-        return self.proxy_license(challenge, track, "playready")
+        return self.proxy_license(challenge, track, "playready", kid)
 
     def get_clearkey_license(
         self, *, challenge: bytes, title: Title_T, track: AnyTrack
@@ -2115,8 +2117,13 @@ class RemoteService:
         )
         return base64.b64decode(resp["license"]) if resp.get("license") else None
 
-    def proxy_license(self, challenge: Union[bytes, str], track: AnyTrack, drm_type: str) -> bytes:
-        """Relay a licence challenge to the server. A failed relay raises ClickException, so dl's licence fallbacks see it."""
+    def proxy_license(
+        self, challenge: Union[bytes, str], track: AnyTrack, drm_type: str, kid: Optional[UUID] = None
+    ) -> bytes:
+        """Relay a licence challenge to the server. A failed relay raises ClickException, so dl's licence fallbacks see it.
+
+        ``kid`` is the KID the server CDM must answer; a server vault key for it stands in for the licence.
+        """
         if isinstance(challenge, str):
             challenge = challenge.encode("utf-8")
 
@@ -2136,6 +2143,8 @@ class RemoteService:
             if pssh_b64:
                 try:
                     request = {"drm_type": drm_type, "mode": "server_cdm", "pssh": pssh_b64}
+                    if kid:
+                        request["kid"] = kid.hex
                     with quiet():
                         resp = self.session_call(
                             lambda sid: self.client.post(