fix(api): reach every PlayReady header on all server CDM routes

The rebuild raised KIDNotFound for a header whose KID lives only in the manifest, the single-track route stripped the track's other headers, and the batch cache keyed on the first PSSH alone. Licensing failures now reach the client instead of only the server log.

tests/remote/unit/test_server_cdm_absorb_routes.py

@@ -0,0 +1,161 @@
+"""The server CDM must reach every WRM header on the routes that bypass the batch loop.
+
+Three regressions found by review of the absorb fix: the single-track route stripped the
+track's other PlayReady objects before the absorb could see them, the batch cache keyed on
+the first PSSH alone and so handed one track another track's narrower key set, and a track
+with no readable PSSH told the client nothing.
+"""
+
+import base64
+import json
+from types import SimpleNamespace
+from uuid import UUID
+
+import pyplayready.system.pssh as pr_pssh_mod
+import pytest
+
+import unshackle.core.cdm as cdm_mod
+import unshackle.core.cdm.detect as detect_mod
+import unshackle.core.drm as drm_mod
+from unshackle.core.api import handlers
+
+pytestmark = pytest.mark.unit
+
+KID_A = UUID("161219ec3df64e0eadbb3827c0ccc46d")
+KID_B = UUID("43de497c6d1945f49acc2951696ec628")
+KID_C = UUID("75982fac2cd8412a990a59c0ef2d4bb1")
+
+
+class _FakePRPSSH:
+    def __init__(self, data):
+        self.data = data
+
+
+class _FakePlayReady:
+    def __init__(self, pssh=None, pssh_b64=None, kid=None, **kwargs):
+        self.kids = [kid] if kid else [KID_A]
+        self.content_keys: dict = {}
+        self.pssh_b64 = pssh_b64
+        self.data = {"pssh_b64": pssh_b64} if pssh_b64 else {}
+
+    def absorb(self, *others):
+        for other in others:
+            if other is self:
+                continue
+            for kid in other.kids:
+                if kid not in self.kids:
+                    self.kids.append(kid)
+
+    def get_content_keys(self, cdm, certificate, licence):
+        self.content_keys = {kid: f"key_{kid.hex[:4]}" for kid in self.kids}
+
+
+@pytest.fixture
+def env(monkeypatch):
+    monkeypatch.setattr(pr_pssh_mod, "PSSH", _FakePRPSSH)
+    monkeypatch.setattr(drm_mod, "PlayReady", _FakePlayReady)
+    monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", lambda *a, **k: object())
+    monkeypatch.setattr(detect_mod, "is_playready_cdm", lambda cdm: True)
+    monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None)
+    monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev")
+    monkeypatch.setattr(handlers, "check_vaults", lambda kids, name: None)
+    monkeypatch.setattr(handlers, "cache_to_vaults", lambda keys, name: None)
+    monkeypatch.setattr(handlers, "server_cdm_allowed", lambda request, tag: True)
+    monkeypatch.setattr(handlers, "find_title_for_track", lambda tid, session: SimpleNamespace())
+    monkeypatch.setattr(handlers.config, "serve", {"users": {}}, raising=False)
+
+
+_FakePlayReady.__name__ = "PlayReady"  # the single-track route matches on class name, not type
+
+PSSH = base64.b64encode(b"A").decode()
+
+
+async def test_single_track_path_keeps_siblings(env, monkeypatch):
+    """`proxy_license` always sends `pssh`, so this route must not discard the other headers."""
+    track = SimpleNamespace(
+        id="vid",
+        drm=[_FakePlayReady(pssh_b64=PSSH, kid=KID_A), _FakePlayReady(pssh_b64=PSSH, kid=KID_B)],
+    )
+    session = SimpleNamespace(
+        service_tag="EX",
+        service_instance=SimpleNamespace(),
+        tracks=SimpleNamespace(get={"vid": track}.get),
+        log_buffer=None,
+    )
+
+    async def fake_get_session(sid, req):
+        return session
+
+    monkeypatch.setattr(handlers, "get_validated_session", fake_get_session)
+
+    resp = await handlers.session_license_handler(
+        {"track_id": "vid", "mode": "server_cdm", "drm_type": "playready", "pssh": PSSH},
+        "sess",
+        None,
+    )
+    keys = json.loads(resp.body)["keys"]
+    assert set(keys) == {KID_A.hex, KID_B.hex}, f"lost sibling headers: {sorted(keys)}"
+
+
+async def test_batch_cache_does_not_starve_second_track(env, monkeypatch):
+    """Two tracks sharing a first PSSH but holding different headers need separate licences."""
+    vid = SimpleNamespace(id="vid", drm=[_FakePlayReady(pssh_b64=PSSH, kid=KID_A)])
+    aud = SimpleNamespace(
+        id="aud",
+        drm=[
+            _FakePlayReady(pssh_b64=PSSH, kid=KID_A),
+            _FakePlayReady(pssh_b64=PSSH, kid=KID_B),
+            _FakePlayReady(pssh_b64=PSSH, kid=KID_C),
+        ],
+    )
+    tracks = {"vid": vid, "aud": aud}
+    session = SimpleNamespace(
+        service_tag="EX",
+        service_instance=SimpleNamespace(),
+        tracks=SimpleNamespace(get=tracks.get),
+        log_buffer=None,
+    )
+
+    async def fake_get_session(sid, req):
+        return session
+
+    monkeypatch.setattr(handlers, "get_validated_session", fake_get_session)
+    monkeypatch.setattr(handlers, "detect_cdm_type_for_service", lambda tag, cfg: "playready")
+    monkeypatch.setattr(handlers, "drm_preference_name", lambda track: None)
+    monkeypatch.setattr(handlers, "fetch_init_segment", lambda track, sess: None)
+
+    resp = await handlers.session_license_handler(
+        {"mode": "server_cdm", "track_ids": ["vid", "aud"], "drm_type": "playready"},
+        "sess",
+        None,
+    )
+    keys = json.loads(resp.body)["keys"]
+    assert set(keys["aud"]) == {KID_A.hex, KID_B.hex, KID_C.hex}, f"audio starved: {sorted(keys['aud'])}"
+
+
+async def test_no_pssh_tells_the_client_why(env, monkeypatch):
+    """A track with DRM but no readable PSSH must give the client a reason, not silence."""
+    from unshackle.core.api.session_log import SessionLogBuffer
+
+    buf = SessionLogBuffer()
+    track = SimpleNamespace(id="vid", drm=[object()])
+    session = SimpleNamespace(
+        service_tag="EX",
+        service_instance=SimpleNamespace(),
+        tracks=SimpleNamespace(get={"vid": track}.get),
+        log_buffer=buf,
+    )
+
+    async def fake_get_session(sid, req):
+        return session
+
+    monkeypatch.setattr(handlers, "get_validated_session", fake_get_session)
+    monkeypatch.setattr(handlers, "detect_cdm_type_for_service", lambda tag, cfg: "playready")
+    monkeypatch.setattr(handlers, "drm_preference_name", lambda track: None)
+    monkeypatch.setattr(handlers, "fetch_init_segment", lambda track, sess: None)
+    monkeypatch.setattr(handlers, "extract_pssh_from_track", lambda track, drm: None)
+
+    await handlers.session_license_handler(
+        {"mode": "server_cdm", "track_ids": ["vid"], "drm_type": "playready"}, "sess", None
+    )
+    assert buf.since(0), "client gets no reason why the track produced no keys"

tests/remote/unit/test_server_cdm_manifest_kid.py

@@ -0,0 +1,160 @@
+"""The server CDM must keep a KID that only the manifest names, and tell the client when a key is missing.
+
+Three regressions found by a second review of the absorb fix: the server rebuilt a
+`PlayReady` from the PSSH alone, so a header whose KID lives only in the manifest (a DASH
+`<kid>` or `default_KID`, an init-segment `tenc` box) raised `KIDNotFound` before the absorb
+could add it; the batch cache keyed on headers only, so two tracks that shared a header but
+named different KIDs got one licence; and a KID the licence did not cover was logged on the
+server alone.
+"""
+
+import base64
+import json
+from types import SimpleNamespace
+from uuid import UUID
+
+import pyplayready.system.pssh as pr_pssh_mod
+import pytest
+from pyplayready.system.pssh import PSSH as RealPSSH
+
+import unshackle.core.cdm as cdm_mod
+import unshackle.core.cdm.detect as detect_mod
+import unshackle.core.drm as drm_mod
+from unshackle.core.api import handlers
+from unshackle.core.api.session_log import SessionLogBuffer
+from unshackle.core.drm import PlayReady
+
+pytestmark = pytest.mark.unit
+
+KID_A = UUID("161219ec3df64e0eadbb3827c0ccc46d")
+KID_B = UUID("43de497c6d1945f49acc2951696ec628")
+
+
+def _kidless_pro_b64() -> str:
+    """A one-record PlayReady Object whose v4.2 header names no KID."""
+    xml = (
+        '<WRMHEADER xmlns="http://schemas.microsoft.com/DRM/2007/03/PlayReadyHeader" version="4.2.0.0">'
+        "<DATA><PROTECTINFO></PROTECTINFO><LA_URL>https://example.invalid/</LA_URL></DATA></WRMHEADER>"
+    )
+    record = xml.encode("utf-16-le")
+    pro = RealPSSH.PlayreadyHeader.build(
+        {"length": 10 + len(record), "records": [{"type": 1, "length": len(record), "data": record}]}
+    )
+    return base64.b64encode(pro).decode()
+
+
+class _FakePRPSSH:
+    def __init__(self, data):
+        self.data = data
+
+
+class _FakePlayReady:
+    def __init__(self, pssh=None, pssh_b64=None, kid=None, **kwargs):
+        self.kids = [kid] if kid else [KID_A]
+        self.content_keys: dict = {}
+        self.pssh_b64 = pssh_b64
+        self.data = {"pssh_b64": pssh_b64} if pssh_b64 else {}
+
+    @property
+    def kid(self):
+        return self.kids[0]
+
+    def absorb(self, *others):
+        for other in others:
+            if other is self:
+                continue
+            for kid in other.kids:
+                if kid not in self.kids:
+                    self.kids.append(kid)
+
+    def get_content_keys(self, cdm, certificate, licence):
+        self.content_keys = {kid: f"key_{kid.hex[:4]}" for kid in self.kids}
+
+
+_FakePlayReady.__name__ = "PlayReady"
+
+PSSH = base64.b64encode(b"A").decode()
+
+
+@pytest.fixture
+def server_env(monkeypatch):
+    monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", lambda *a, **k: object())
+    monkeypatch.setattr(detect_mod, "is_playready_cdm", lambda cdm: True)
+    monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None)
+    monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev")
+    monkeypatch.setattr(handlers, "check_vaults", lambda kids, name: None)
+    monkeypatch.setattr(handlers, "cache_to_vaults", lambda keys, name: None)
+    monkeypatch.setattr(handlers, "server_cdm_allowed", lambda request, tag: True)
+    monkeypatch.setattr(handlers, "find_title_for_track", lambda tid, session: SimpleNamespace())
+    monkeypatch.setattr(handlers, "detect_cdm_type_for_service", lambda tag, cfg: "playready")
+    monkeypatch.setattr(handlers, "drm_preference_name", lambda track: None)
+    monkeypatch.setattr(handlers, "fetch_init_segment", lambda track, sess: None)
+    monkeypatch.setattr(handlers.config, "serve", {"users": {}}, raising=False)
+
+
+@pytest.fixture
+def fake_playready(monkeypatch):
+    monkeypatch.setattr(pr_pssh_mod, "PSSH", _FakePRPSSH)
+    monkeypatch.setattr(drm_mod, "PlayReady", _FakePlayReady)
+
+
+def _session(tracks: dict, buf=None):
+    session = SimpleNamespace(
+        service_tag="EX",
+        service_instance=SimpleNamespace(),
+        tracks=SimpleNamespace(get=tracks.get),
+        log_buffer=buf,
+    )
+
+    async def fake_get_session(sid, req):
+        return session
+
+    return fake_get_session
+
+
+def test_manifest_kid_survives_the_rebuild(server_env, monkeypatch):
+    """A DASH `<kid>` reaches the local object through `kid=`; the server rebuild must not lose it."""
+    monkeypatch.setattr(
+        PlayReady,
+        "get_content_keys",
+        lambda self, cdm, certificate, licence: self.content_keys.update({k: "k" for k in self.kids}),
+    )
+    pro_b64 = _kidless_pro_b64()
+    track = SimpleNamespace(
+        id="vid", drm=[PlayReady(pssh=RealPSSH(base64.b64decode(pro_b64)), kid=KID_A, pssh_b64=pro_b64)]
+    )
+
+    keys = handlers.handle_single_server_cdm(SimpleNamespace(), SimpleNamespace(), track, pro_b64, "playready", None)
+
+    assert set(keys) == {KID_A.hex}
+
+
+async def test_batch_cache_splits_on_manifest_kid(server_env, fake_playready, monkeypatch):
+    """Two tracks with one shared header but their own manifest KIDs need their own licences."""
+    tracks = {
+        "vid": SimpleNamespace(id="vid", drm=[_FakePlayReady(pssh_b64=PSSH, kid=KID_A)]),
+        "aud": SimpleNamespace(id="aud", drm=[_FakePlayReady(pssh_b64=PSSH, kid=KID_B)]),
+    }
+    monkeypatch.setattr(handlers, "get_validated_session", _session(tracks))
+
+    resp = await handlers.session_license_handler(
+        {"mode": "server_cdm", "track_ids": ["vid", "aud"], "drm_type": "playready"}, "sess", None
+    )
+
+    keys = json.loads(resp.body)["keys"]
+    assert KID_B.hex in keys["aud"], f"audio got the video licence: {sorted(keys['aud'])}"
+
+
+async def test_uncovered_kid_reaches_the_client(server_env, fake_playready, monkeypatch):
+    """The client only sees `session.log_buffer`, so an uncovered track KID must land there."""
+    buf = SessionLogBuffer()
+    track = SimpleNamespace(id="vid", drm=[_FakePlayReady(pssh_b64=PSSH, kid=KID_A)], get_key_id=lambda init: KID_B)
+    monkeypatch.setattr(handlers, "get_validated_session", _session({"vid": track}, buf))
+    monkeypatch.setattr(handlers, "fetch_init_segment", lambda track, sess: b"init")
+    monkeypatch.setattr(handlers, "drm_from_init_segment", lambda track, session=None, init_data=None: [])
+
+    await handlers.session_license_handler(
+        {"mode": "server_cdm", "track_ids": ["vid"], "drm_type": "playready"}, "sess", None
+    )
+
+    assert any(KID_B.hex in r["message"] for r in buf.since(0)), "client gets no reason for the missing key"

unshackle/core/api/handlers.py

@@ -3685,10 +3685,10 @@ def handle_single_server_cdm(
         from unshackle.core.drm import PlayReady
 
         pr_pssh = PlayReadyPSSH(base64.b64decode(pssh_b64))
-        pr_drm = PlayReady(pssh=pr_pssh, pssh_b64=pssh_b64)
-
         siblings = [d for d in (getattr(track, "drm", None) or []) if isinstance(d, PlayReady)]
-        if len(siblings) > 1:
+        manifest_kid = next((getattr(d, "kid", None) for d in siblings if getattr(d, "kid", None)), None)
+        pr_drm = PlayReady(pssh=pr_pssh, pssh_b64=pssh_b64, kid=manifest_kid)
+        if siblings:
             pr_drm.absorb(*siblings)
 
         # Gate on the caller's CDM device first: no device, no keys from the vault or CDM.
@@ -3843,7 +3843,6 @@ async def session_license_handler(
     the challenge, get the license, and extract `KID:KEY` pairs. Supports
     batch (track_ids list) and single-track requests.
     """
-    import base64
 
     session = await get_validated_session(session_id, request)
 
@@ -3873,10 +3872,23 @@ async def session_license_handler(
 
         all_keys: Dict[str, Dict[str, str]] = {}
         drm_types: Dict[str, str] = {}
-        keys_by_pssh: Dict[str, Dict[str, str]] = {}
-        drm_type_by_pssh: Dict[str, str] = {}
+        keys_by_pssh: Dict[tuple, Dict[str, str]] = {}
+        drm_type_by_pssh: Dict[tuple, str] = {}
         actual_drm_type: Optional[str] = None
 
+        def pssh_set(track: Any) -> tuple:
+            """Every PSSH and KID the track carries, as the rest of the cache key.
+
+            Two tracks can share their first PSSH and still need different licences,
+            because the server licenses every header the track holds and asks for every
+            KID the manifest named on it.
+            """
+            drms = getattr(track, "drm", None) or []
+            return (
+                tuple(sorted(str(getattr(d, "pssh_b64", "") or "") for d in drms)),
+                tuple(sorted(str(kid) for d in drms for kid in (getattr(d, "kids", None) or []))),
+            )
+
         def warn(message: str) -> None:
             """Log on the server and copy into the remote session buffer the client drains.
 
@@ -3902,20 +3914,22 @@ async def session_license_handler(
                 if pssh_str:
                     break
             else:
+                warn(f"No PSSH on track {sanitize_log(str(track.id)[:12])} for {', '.join(candidates) or 'any CDM'}")
                 return {}, None, None
 
-            if pssh_str not in keys_by_pssh:
-                keys_by_pssh[pssh_str] = {}
+            cache_key = (pssh_str, pssh_set(track))
+            if cache_key not in keys_by_pssh:
+                keys_by_pssh[cache_key] = {}
                 try:
                     keys = handle_single_server_cdm(service, title, track, pssh_str, candidate, request)
                     if keys:
-                        keys_by_pssh[pssh_str] = keys
-                        drm_type_by_pssh[pssh_str] = candidate
+                        keys_by_pssh[cache_key] = keys
+                        drm_type_by_pssh[cache_key] = candidate
                 except SystemExit:
                     warn(f"Service exited while resolving keys for track {sanitize_log(str(track.id)[:12])}, skipping")
                 except (Exception, SystemExit) as e:
                     warn(f"Failed to resolve keys for track {sanitize_log(str(track.id)[:12])}: {e}")
-            return keys_by_pssh[pssh_str], drm_type_by_pssh.get(pssh_str), pssh_str
+            return keys_by_pssh[cache_key], drm_type_by_pssh.get(cache_key), pssh_str
 
         for tid in track_ids:
             track = session.tracks.get(tid)
@@ -3926,6 +3940,7 @@ async def session_license_handler(
             init_data = fetch_init_segment(track, svc_session)
             ensure_track_drm(track, svc_session, init_data)
             if not track.drm:
+                warn(f"Track {sanitize_log(tid[:12])} carries no DRM, so it has no keys to resolve")
                 continue
 
             title = find_title_for_track(tid, session)
@@ -3947,7 +3962,7 @@ async def session_license_handler(
                     candidates.remove(preferred)
                     candidates.insert(0, preferred)
                 else:
-                    log.warning(
+                    warn(
                         f"Track {sanitize_log(tid[:12])} wants {preferred} DRM "
                         "but the server has no device for it, using the configured DRM instead"
                     )
@@ -3973,7 +3988,7 @@ async def session_license_handler(
                     keys, track_drm_type = init_keys, init_drm_type
                 else:
                     track.drm = manifest_drm
-                    log.warning(f"No content key for KID {track_kid.hex} of track {sanitize_log(tid[:12])}")
+                    warn(f"No content key for KID {track_kid.hex} of track {sanitize_log(tid[:12])}")
 
             if keys:
                 all_keys[tid] = keys
@@ -4005,25 +4020,8 @@ async def session_license_handler(
 
         pssh_b64 = data.get("pssh")
         if pssh_b64:
-            cls_name = {"playready": "PlayReady", "widevine": "Widevine"}.get(drm_type)
-            track.drm = [d for d in (track.drm or []) if d.__class__.__name__ != cls_name]
             if drm_type == "playready":
                 track.pr_pssh = pssh_b64
-                from pyplayready.system.pssh import PSSH as PlayReadyPSSH
-
-                from unshackle.core.drm import PlayReady
-
-                pr_pssh = PlayReadyPSSH(base64.b64decode(pssh_b64))
-                pr_drm = PlayReady(pssh=pr_pssh, pssh_b64=pssh_b64)
-                track.drm.append(pr_drm)
-            elif drm_type == "widevine":
-                from pywidevine.pssh import PSSH as WidevinePSSH
-
-                from unshackle.core.drm import Widevine
-
-                wv_pssh = WidevinePSSH(pssh_b64)
-                wv_drm = Widevine(pssh=wv_pssh)
-                track.drm.append(wv_drm)
 
         if mode == "server_cdm":
             keys = handle_single_server_cdm(service, title, track, pssh_b64, drm_type, request)