fix(dl): cache server-CDM keys to the local vaults

Batch and per-track remote licences read the vaults but never wrote back, so a --remote or import run kept nothing. Vault hits now replicate to the other vaults, and the server branch resolves keys under drm_lock like the local path.

docs/guide/vaults.md

@@ -30,6 +30,8 @@ When unshackle recovers new keys from a license, it pushes them to **all** confi
 - unshackle skips keys that already exist in a vault (it never overwrites or deletes existing data).
 - A vault marked `no_push: true` receives lookups, but unshackle never writes to it. This is useful for read-only or shared upstream vaults you do not own.
 
+unshackle pushes keys from every source, not only a local CDM. A run against a remote server that licenses with its own CDM (`--remote`) writes the keys the server returns, and `unshackle import` writes the keys it reads from an export file.
+
 !!! tip "A good two-vault setup"
     A common arrangement is a fast local **SQLite** vault plus a shared remote vault (**MySQL**, **HTTP**, or **API**). The local vault answers instantly and works offline; the remote vault lets you share keys across machines or with a group.
 

tests/remote/unit/test_server_cdm_vault_write.py

@@ -0,0 +1,173 @@
+"""Keys the server CDM returns must land in the local vaults.
+
+The server-CDM branch of prepare_drm reads the local vaults but used to return before
+the vault writes, so a remote run never kept the keys it paid for. Keys that came from a
+vault in the first place must not be pushed back.
+"""
+
+import threading
+import time
+from types import SimpleNamespace
+from uuid import UUID
+
+import pytest
+from rich.table import Table
+
+from unshackle.commands.dl import dl
+
+pytestmark = pytest.mark.unit
+
+NEW = UUID("161219ec3df64e0eadbb3827c0ccc46d")
+CACHED = UUID("75982fac2cd8412a990a59c0ef2d4bb1")
+
+
+class FakeVaults:
+    def __init__(self, held=None, vault_used=None):
+        self.held = held or {}
+        self.vault_used = vault_used
+        self.added: dict = {}
+        self.pushes: list = []
+        self.replicated: list = []
+
+    def __len__(self):
+        return 1
+
+    def get_key(self, kid):
+        return self.held.get(kid), self.vault_used
+
+    def add_keys(self, kid_keys):
+        self.added.update(kid_keys)
+        self.pushes.append(dict(kid_keys))
+        return 1
+
+    def add_key(self, kid, key, excluding=None):
+        self.replicated.append((kid, key, excluding))
+        return 1
+
+
+class FakeWidevine:
+    def __init__(self, kids):
+        self.kids = list(kids)
+        self.content_keys: dict = {}
+
+
+def make_cmd(vaults):
+    cmd = dl.__new__(dl)
+    cmd.log = SimpleNamespace(warning=print, info=print, debug=print)
+    cmd.service = "TEST"
+    cmd.vaults = vaults
+    cmd.vault_cache_tally = None
+    cmd._remote_service = SimpleNamespace(_server_cdm=True, _server_cdm_type="widevine")
+    dl.LICENSE_KEY_CACHE.clear()
+    dl.DRM_LOCKS.clear()
+    return cmd
+
+
+def prepare(cmd, drm, licence, track_kid, **kwargs):
+    table = Table.grid()
+    table.add_row("")
+    cmd.prepare_drm(
+        track=SimpleNamespace(drm=[drm]),
+        title=None,
+        drm=drm,
+        licence=licence,
+        certificate=None,
+        track_kid=track_kid,
+        table=table,
+        **kwargs,
+    )
+
+
+def run_prepare(vaults, kids, licence_keys, track_kid):
+    cmd = make_cmd(vaults)
+    drm = FakeWidevine(kids)
+
+    def licence(**kwargs):
+        drm.content_keys.update(licence_keys)
+
+    prepare(cmd, drm, licence, track_kid)
+    cmd.wait_vault_writes()
+    return drm
+
+
+def test_server_cdm_keys_are_cached_to_local_vaults():
+    vaults = FakeVaults()
+    run_prepare(vaults, [NEW], {NEW: "k_new"}, track_kid=NEW)
+    assert vaults.added == {NEW: "k_new"}
+
+
+def test_keys_read_from_a_vault_are_not_pushed_back():
+    vaults = FakeVaults({CACHED: "k_cached"})
+    run_prepare(vaults, [CACHED, NEW], {NEW: "k_new"}, track_kid=NEW)
+    assert vaults.added == {NEW: "k_new"}
+
+
+def test_batch_resolved_keys_are_cached_before_download():
+    """The batch licence fills track.drm before prepare_drm, so dl caches it at the call site."""
+    vaults = FakeVaults()
+    cmd = dl.__new__(dl)
+    cmd.log = SimpleNamespace(warning=print, info=print, debug=print)
+    cmd.service = "TEST"
+    cmd.vaults = vaults
+    cmd.vault_cache_tally = None
+    dl.LICENSE_KEY_CACHE.clear()
+
+    drm = FakeWidevine([NEW])
+    drm.content_keys = {NEW: "k_new"}
+    title = SimpleNamespace(tracks=[SimpleNamespace(drm=[drm]), SimpleNamespace(drm=None)])
+
+    cmd.cache_resolved_keys(title)
+    cmd.wait_vault_writes()
+
+    assert vaults.added == {NEW: "k_new"}
+
+
+def test_vault_hits_replicate_to_the_other_vaults():
+    """The local path pushes a vault hit to every other vault; the server path must too."""
+    vaults = FakeVaults({CACHED: "k_cached"}, vault_used="vault-a")
+    run_prepare(vaults, [CACHED], {}, track_kid=CACHED)
+    assert vaults.replicated == [(CACHED, "k_cached", "vault-a")]
+    assert vaults.added == {}
+
+
+def test_tracks_sharing_one_drm_licence_and_push_once():
+    """Two tracks on the same DRM run at the same time; the second must find the first one's keys."""
+    vaults = FakeVaults()
+    cmd = make_cmd(vaults)
+    drm = FakeWidevine([NEW])
+    calls = []
+
+    def licence(**kwargs):
+        calls.append(1)
+        time.sleep(0.2)
+        drm.content_keys.update({NEW: "k_new"})
+
+    threads = [threading.Thread(target=prepare, args=(cmd, drm, licence, NEW)) for _ in range(2)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+    cmd.wait_vault_writes()
+
+    assert len(calls) == 1
+    assert vaults.pushes == [{NEW: "k_new"}]
+
+
+def test_cdm_only_still_reuses_keys_from_the_run_cache():
+    """--cdm-only skips the vaults, not LICENSE_KEY_CACHE, so a second track with the same KID does not licence twice."""
+    vaults = FakeVaults()
+    cmd = make_cmd(vaults)
+    calls = []
+
+    for _ in range(2):
+        drm = FakeWidevine([NEW])
+
+        def licence(**kwargs):
+            calls.append(1)
+            drm.content_keys.update({NEW: "k_new"})
+
+        prepare(cmd, drm, licence, NEW, cdm_only=True)
+    cmd.wait_vault_writes()
+
+    assert len(calls) == 1
+    assert vaults.pushes == [{NEW: "k_new"}]

unshackle/commands/dl.py

@@ -2904,6 +2904,7 @@ class dl:
 
             if hasattr(service, "resolve_server_keys"):
                 service.resolve_server_keys(title)
+                self.cache_resolved_keys(title)
 
             dl_start_time = time.time()
 
@@ -3996,6 +3997,23 @@ class dl:
         self.vault_cache_tally = None
         self.log.info(f"Cached {keys} Key{'' if keys == 1 else 's'} to {successful_caches}/{len(self.vaults)} Vaults")
 
+    def cache_resolved_keys(self, title: Title_T) -> None:
+        """Cache the keys a batch licence or an import put on the tracks.
+
+        Both fill track.drm before any prepare_drm call, so this is the only point where those
+        keys can still reach the local vaults.
+        """
+        keys = {
+            kid: key
+            for track in title.tracks
+            for drm in (getattr(track, "drm", None) or [])
+            for kid, key in getattr(drm, "content_keys", {}).items()
+            if kid not in self.LICENSE_KEY_CACHE
+        }
+        if keys:
+            self.LICENSE_KEY_CACHE.update(keys)
+            self.flush_vault_writes([partial(self.cache_keys_to_vaults, keys)])
+
     def cache_keys_to_vaults(self, content_keys: dict[UUID, str]) -> None:
         successful_caches = self.vaults.add_keys(content_keys)
         keys, caches = self.vault_cache_tally or (0, successful_caches)
@@ -4054,29 +4072,44 @@ class dl:
         server_cdm = getattr(svc_for_cdm, "_server_cdm", getattr(self, "server_cdm", False))
 
         if server_cdm:
-            if not cdm_only:
+            with self.drm_lock(drm):
+                pending_vault_writes: list[Callable[[], Any]] = []
                 vault_kids = list(getattr(drm, "kids", None) or [])
                 if track_kid and track_kid not in vault_kids:
                     vault_kids.append(track_kid)
                 for kid in vault_kids:
                     if kid in drm.content_keys:
                         continue
-                    content_key = self.LICENSE_KEY_CACHE.get(kid) or self.vaults.get_key(kid)[0]
+                    content_key = self.LICENSE_KEY_CACHE.get(kid)
+                    if not content_key and not cdm_only:
+                        content_key, vault_used = self.vaults.get_key(kid)
+                        if content_key:
+                            pending_vault_writes.append(
+                                partial(self.vaults.add_key, kid, content_key, excluding=vault_used)
+                            )
                     if content_key:
                         drm.content_keys[kid] = content_key
                         self.LICENSE_KEY_CACHE[kid] = content_key
 
-            def missing_track_key() -> bool:
-                return not drm.content_keys or bool(track_kid and track_kid not in drm.content_keys)
+                known_keys = set(drm.content_keys)
 
-            if missing_track_key():
-                try:
-                    licence(
-                        drm_system="playready" if drm.__class__.__name__ == "PlayReady" else "widevine",
-                        challenge=b"",
-                    )
-                except Exception as e:
-                    self.log.debug(f"Server CDM licence with an empty challenge failed: {e!r}")
+                def missing_track_key() -> bool:
+                    return not drm.content_keys or bool(track_kid and track_kid not in drm.content_keys)
+
+                if missing_track_key():
+                    try:
+                        licence(
+                            drm_system="playready" if drm.__class__.__name__ == "PlayReady" else "widevine",
+                            challenge=b"",
+                        )
+                    except Exception as e:
+                        self.log.debug(f"Server CDM licence with an empty challenge failed: {e!r}")
+
+                new_keys = {kid: key for kid, key in drm.content_keys.items() if kid not in known_keys}
+                if new_keys:
+                    self.LICENSE_KEY_CACHE.update(new_keys)
+                    pending_vault_writes.append(partial(self.cache_keys_to_vaults, new_keys))
+                self.flush_vault_writes(pending_vault_writes)
 
             if not drm.content_keys:
                 self.log.warning("Server CDM did not resolve any keys for this track")