fix(vaults): harden the KID window check and clear a stale flag on licence

docs/guide/vaults.md

@@ -34,11 +34,11 @@ unshackle pushes keys from every source, not only a local CDM. A run against a r
 
 ### Wrong keys from a poisoned vault
 
-A public vault can return the right KID with a wrong content key. The decrypters accept a wrong content key and write noise. unshackle therefore trusts a content key from a vault only after FFmpeg decodes the decrypted track with strict error detection. For a fragmented MP4, unshackle reads from the encrypted file which fragments each KID encrypts, and FFmpeg decodes up to three windows of each KID: from its first, middle and last fragment on through the next fragments of that KID, for up to 10 seconds. A window never runs into a fragment of another KID or a clear fragment. A clear lead or a lead under another KID therefore cannot hide a wrong content key, and only the KID that fails is flagged. A KID that is not in this map, and every KID of any other file, gets one check of the whole file. FFmpeg decodes the first seconds and the last seconds. A failure there flags each of those KIDs. Content keys from a licence skip this check. A file this FFmpeg build cannot judge counts as a pass, never as a wrong content key: a codec it has no decoder for, or a header it cannot read. Without FFmpeg there is no check at all.
+A public vault can return the right KID with a wrong content key. The decrypters accept a wrong content key and write noise. unshackle therefore trusts a content key from a vault only after FFmpeg decodes the decrypted track with strict error detection. For a fragmented MP4, unshackle reads from the encrypted file which fragments each KID encrypts, and FFmpeg decodes up to three windows of each KID: from its first, middle and last fragment on through the next fragments of that KID, for up to 4 seconds. A window never runs into a fragment of another KID or a clear fragment. A clear lead or a lead under another KID therefore cannot hide a wrong content key, and only the KID that fails is flagged. A KID that is not in this map, and every KID of any other file, gets one check of the whole file. FFmpeg decodes the first seconds and the last seconds. A failure there flags each of those KIDs. When a mapped KID fails, the KIDs that are not in the map get no verdict from that pass; the retry with the replaced key judges them. Content keys from a licence skip this check. A file this FFmpeg build cannot judge counts as a pass, never as a wrong content key: a codec it has no decoder for, or a header it cannot read. Without FFmpeg there is no check at all.
 
 When the decode fails, unshackle:
 
-1. Flags the KID:KEY in the `bad_keys` table of every local SQLite vault, with the service and the name of the vault it came from. The pair also goes out of the service table, except in a vault marked `no_push`, where only the flag is written. A flagged pair cannot go back into a local vault, from a later run or from `kv copy`. An `API` vault that served the pair also gets a report; `HTTP` and `MySQL` vaults have no report route. A pair from a local vault goes to no remote vault, even when a remote vault holds the same pair.
+1. Flags the KID:KEY in the `bad_keys` table of every local SQLite vault, with the service and the name of the vault it came from. The pair also goes out of the service table, except in a vault marked `no_push`, where only the flag is written. A flagged pair cannot go back into a local vault from another vault, a later run, or `kv copy`. A licence that returns the pair clears the flag and stores it. An `API` vault that served the pair also gets a report; `HTTP` and `MySQL` vaults have no report route. A pair from a local vault goes to no remote vault, even when a remote vault holds the same pair.
 2. Restores the ciphertext from a hard link it made before decryption, then asks for the content key again. Every vault lookup skips a flagged pair, so the request reaches the remaining vaults and then the CDM. A remote answer that arrived after the winning one is reused instead of queried again.
 3. Removes the flag when the CDM returns the same content key, because then the decode check was wrong, not the content key. A content key from another vault is checked in its turn, and flagged as well if it also fails.
 

tests/core/test_kid_windows.py

@@ -90,17 +90,37 @@ def test_each_kid_maps_to_its_own_fragments(tmp_path: Path) -> None:
     assert kid_windows(path) == KidMap({A: [(0.0, 2.0)], B: [(2.0, 6.0), (4.0, 6.0)], C: [(6.0, 8.0)]}, True)
 
 
-def test_windows_run_through_a_kid_but_stop_at_another_or_after_ten_seconds() -> None:
+def test_windows_run_through_a_kid_but_stop_at_another_or_after_the_longest() -> None:
     long_run = [(t * 2.0, t * 2.0 + 2) for t in range(20)]
-    assert pick_windows(long_run) == [(0.0, 10.0), (20.0, 30.0), (30.0, 40.0)]
+    assert pick_windows(long_run, longest=10.0) == [(0.0, 10.0), (20.0, 30.0), (30.0, 40.0)]
+    assert pick_windows(long_run) == [(0.0, 4.0), (20.0, 24.0), (36.0, 40.0)]
     # a gap is another KID or a clear stretch: no window crosses it, and the last one ends the last run
-    assert pick_windows([(0, 2), (2, 4), (10, 12), (12, 14), (14, 16)]) == [(0, 4), (10, 16)]
+    assert pick_windows([(0, 2), (2, 4), (10, 12), (12, 14), (14, 16)], longest=10.0) == [(0, 4), (10, 16)]
+    assert pick_windows([(0, 2), (2, 4), (10, 12), (12, 14), (14, 16)]) == [(0, 4), (10, 14.0), (12, 16)]
+
+
+def test_fragments_longer_than_the_longest_and_empty_spans() -> None:
+    assert pick_windows([(0.0, 60.0)]) == [(0.0, 4.0)]
+    assert pick_windows([(0, 15), (15, 30), (30, 45), (45, 60)]) == [(0, 4), (30, 34), (45, 49)]
+    assert pick_windows([(0.0, 5.0), (5.0, 5.0)]) == [(0.0, 4.0)]
+    assert pick_windows([(5.0, 5.0)]) == []
+
+
+def test_nested_box_bomb_and_oversized_sbgp(tmp_path: Path) -> None:
+    body = b""
+    for _ in range(5000):
+        body = box(b"trak", body)
+    (tmp_path / "bomb.mp4").write_bytes(box(b"moov", body))
+    assert kid_windows(tmp_path / "bomb.mp4") is None
+    huge = full(b"sbgp", b"seig" + struct.pack(">I", 0x00FFFFFF) + bytes(16))
+    (tmp_path / "sbgp.mp4").write_bytes(movie(encv(A)) + moof(1, 0, huge))
+    assert kid_windows(tmp_path / "sbgp.mp4") is None
 
 
 def test_first_middle_and_last_fragment(tmp_path: Path) -> None:
     path = tmp_path / "enc.mp4"
     path.write_bytes(movie(encv(A)) + b"".join(moof(1, t * 1000) for t in range(5)))
-    assert kid_windows(path) == KidMap({A: [(0.0, 7.0), (2.0, 7.0)]}, True)
+    assert kid_windows(path) == KidMap({A: [(0.0, 4.0), (2.0, 6.0), (3.0, 7.0)]}, True)
 
 
 def test_last_fragment_ends_after_its_sample_durations(tmp_path: Path) -> None:
@@ -145,7 +165,7 @@ def test_samples_past_the_sbgp_runs_take_the_sample_entry_kid(tmp_path: Path) ->
 def test_piff_tenc_and_tfxd(tmp_path: Path) -> None:
     path = tmp_path / "enc.ismv"
     path.write_bytes(movie(encv(A, piff=True)) + moof(1, 0, tfxd=True) + moof(1, 2000, tfxd=True))
-    assert kid_windows(path) == KidMap({A: [(0.0, 5.0), (2.0, 5.0)]}, True)
+    assert kid_windows(path) == KidMap({A: [(0.0, 4.0), (2.0, 5.0)]}, True)
 
 
 def test_oversized_box(tmp_path: Path) -> None:

tests/core/test_vault_bad_keys.py

@@ -316,6 +316,19 @@ def test_flag_moves_the_row_and_blocks_it_coming_back(tmp_path: Path) -> None:
     assert local.add_key("SVC", KID, GOOD)  # the KID itself is not blocked, only the pair
 
 
+def test_a_licence_key_clears_a_flag_from_an_earlier_run(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """A false flag must not outlive the run: the next licence returns the pair and it is stored again."""
+    local = SQLite("local", tmp_path / "local.db")
+    local.flag_bad_key("SVC", KID, GOOD, "poisoned")
+    vaults = Vaults("SVC")
+    vaults.vaults = [local]
+    cmd = make_cmd(monkeypatch, vaults)
+    cmd.vault_cache_tally = None
+    cmd.cache_keys_to_vaults({KID: GOOD})
+    assert not local.is_bad_key(KID, GOOD)
+    assert local.get_key(KID, "SVC") == GOOD
+
+
 def test_bulk_add_refuses_a_flagged_pair_in_any_case_or_kid_form(tmp_path: Path) -> None:
     local = SQLite("local", tmp_path / "local.db")
     local.flag_bad_key("SVC", KID, BAD, "poisoned")
@@ -404,6 +417,45 @@ def test_a_kid_absent_from_the_map_is_judged_by_the_fallback(tmp_path: Path, mon
     assert (None, 3, True) in calls  # the fallback knows the track is video
 
 
+def test_a_wrong_mapped_key_does_not_flag_a_kid_the_file_does_not_carry(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """The audio KID is in the licence and the vault but not in this video file: the whole-file
+    check fails on the wrong video key's noise, and that verdict must not reach the audio KID."""
+    audio = UUID(int=7)
+    local = SQLite("local", tmp_path / "local.db")
+    poisoned = Remote("poisoned", tmp_path / "poisoned.db")
+    poisoned.add_key("SVC", audio, GOOD)
+    poisoned.add_key("SVC", KID, BAD)
+    other = Remote("other", tmp_path / "other.db")
+    other.add_key("SVC", KID, GOOD)
+    vaults = Vaults("SVC")
+    vaults.vaults = [local, poisoned]
+    cmd = make_cmd(monkeypatch, vaults)
+    drm = FakeDRM({audio: vaults.get_key(audio)[0], KID: vaults.get_key(KID)[0]})
+    monkeypatch.setattr(
+        dl_module.verify, "kid_windows", lambda path: dl_module.verify.KidMap({KID: [(4.0, 6.0)]}, True)
+    )
+    monkeypatch.setattr(
+        dl_module,
+        "ffmpeg_decodes",
+        lambda path, start=None, seconds=3, video=None: drm.decrypted_with[-1][KID] == GOOD,
+    )
+    path = tmp_path / "video.mp4"
+    path.write_bytes(b"cipher")
+
+    def licence(drm: FakeDRM, track_kid: UUID) -> None:
+        vaults.vaults.append(other)
+        drm.content_keys[KID] = vaults.get_key(KID)[0]
+
+    cmd.decrypt_verified(drm, path, licence, KID)
+    cmd.wait_vault_writes()
+    assert drm.content_keys == {audio: GOOD, KID: GOOD}
+    assert local.is_bad_key(KID, BAD)
+    assert not local.is_bad_key(audio, GOOD)
+    assert local.get_key(audio, "SVC") == GOOD  # judged on the retry and copied
+
+
 def test_only_checked_kids_are_copied_to_other_vaults(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
     """Both the mapped KID and the absent one are proven before they reach another vault."""
     lead = UUID(int=7)

unshackle/commands/dl.py

@@ -4048,12 +4048,13 @@ class dl:
 
         backup = path.with_name(path.name + ".enc")
         backup.unlink(missing_ok=True)  # a killed run leaves one behind; its bytes may not match
-        if vault_kids() and decrypt and binaries.FFMPEG:
+        checking = bool(vault_kids() and decrypt and binaries.FFMPEG)
+        if checking:
             try:
                 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 vault_kids() and decrypt and binaries.FFMPEG else None
+        kid_map = 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
 
@@ -4064,11 +4065,13 @@ class dl:
             """The vault KIDs whose content key did not decode.
 
             A KID in the KID map fails when one of its windows fails. The KIDs that are not in the
-            map share the verdict of one start and end check of the whole file.
+            map share the verdict of one start and end check of the whole file, taken only when
+            every mapped KID passed: a wrong mapped key fails that check too and would flag them
+            for noise that is not theirs. The retry judges them again with the replaced key.
             """
             failed = {kid for kid in kids if kid in windows and not all(window_decodes(*w) for w in windows[kid])}
             absent = {kid for kid in kids if kid not in windows}
-            if absent and not ffmpeg_decodes(path, video=video):
+            if absent and not failed and not ffmpeg_decodes(path, video=video):
                 failed |= absent
             self.log.debug(
                 f"Key check on {path.name}: {len(kids) - len(absent)} KID(s) by fragment window, "
@@ -4198,6 +4201,11 @@ class dl:
             self.flush_vault_writes([partial(self.cache_keys_to_vaults, keys)])
 
     def cache_keys_to_vaults(self, content_keys: dict[UUID, str]) -> None:
+        """Store licence keys in every vault. A licence key clears a flag on its pair first: the
+        flag came from a decode check that was wrong, and a local vault refuses a flagged pair."""
+        for kid, key in content_keys.items():
+            if self.vaults.is_flagged(kid, key):
+                self.vaults.unflag_bad_key(kid, key)
         successful_caches = self.vaults.add_keys(content_keys)
         keys, caches = self.vault_cache_tally or (0, successful_caches)
         self.vault_cache_tally = (keys + len(content_keys), min(caches, successful_caches))

unshackle/core/drm/verify.py

@@ -103,6 +103,8 @@ def sbgp_runs(data: bytes, body: int) -> Optional[list[tuple[int, int]]]:
         return None
     pos = body + (12 if version == 1 else 8)
     count = struct.unpack_from(">I", data, pos)[0]
+    if count * 8 > len(data) - pos - 4:
+        raise ValueError("sbgp entry count is larger than the box")
     runs = struct.unpack_from(f">{count * 2}I", data, pos + 4)
     return list(zip(runs[::2], runs[1::2]))
 
@@ -202,7 +204,7 @@ def kid_windows(path: Path) -> Optional[KidMap]:
                                     moov_groups[track] = groups
                 elif kind == b"moof" and box:
                     fragments.extend(read_fragment(box, entries, defaults, moov_groups, clock))
-    except (OSError, struct.error, IndexError, ValueError, MemoryError, OverflowError):
+    except (OSError, struct.error, IndexError, ValueError, MemoryError, OverflowError, RecursionError):
         return None
 
     tracks: dict[int, list[Fragment]] = {}
@@ -232,15 +234,19 @@ def kid_windows(path: Path) -> Optional[KidMap]:
     return KidMap({kid: pick_windows(fragment_spans) for kid, fragment_spans in spans.items()}, video)
 
 
-def pick_windows(spans: list[tuple[float, float]], longest: float = 10.0) -> list[tuple[float, float]]:
+def pick_windows(spans: list[tuple[float, float]], longest: float = 4.0) -> list[tuple[float, float]]:
     """Pick up to three windows from one KID's fragment (start, end) spans: at its first, middle and last fragment.
 
     A window starts at a fragment and runs on through the next fragments of the same KID, up
     to ``longest`` seconds, but never into a fragment of another KID or a clear one. A short
     window can decode noise without an error: the decoder often needs several frames to find it.
+    A wrong key fails on the first inter frame in every measured case, so ``longest`` buys
+    decoder frames, not certainty; every second of it is decoded in full for a right key.
     The last window ends where the KID's last run of fragments ends.
     """
-    spans = sorted(spans)
+    spans = sorted(span for span in spans if span[1] > span[0])  # a duplicated tfdt gives an empty span
+    if not spans:
+        return []
     runs: list[list[float]] = []  # [start, end] of each run of back-to-back fragments
     for start, end in spans:
         if runs and abs(start - runs[-1][1]) <= 1e-3:
@@ -252,7 +258,7 @@ def pick_windows(spans: list[tuple[float, float]], longest: float = 10.0) -> lis
         return next(end for run_start, end in runs if run_start <= start < end)
 
     last_from = max(runs[-1][0], runs[-1][1] - longest)
-    last = next(i for i, (start, _) in enumerate(spans) if start >= last_from - 1e-3)
+    last = max(i for i, (start, _) in enumerate(spans) if start <= last_from + 1e-3)
     return [(spans[i][0], min(run_end(spans[i][0]), spans[i][0] + longest)) for i in sorted({0, len(spans) // 2, last})]
 
 

unshackle/core/utils/subprocess.py

@@ -66,7 +66,7 @@ def ffmpeg_decodes(path: Path, start: Optional[float] = None, seconds: float = 3
 
     With ``start``, the check decodes ``seconds`` from that time. An FFmpeg copy cuts the
     window first, so no sample after the window gets to the decoder. Without ``start``, the
-    check decodes the first seconds and the last 10 seconds, because a title can open with
+    check decodes the first seconds and the last 4 seconds, because a title can open with
     a clear or separately keyed lead.
 
     Every frame is decoded, because an HEVC keyframe decrypted with a wrong key can decode
@@ -97,8 +97,9 @@ def ffmpeg_decodes(path: Path, start: Optional[float] = None, seconds: float = 3
                     timeout=FFMPEG_CHECK_TIMEOUT,
                 )
             else:
+                # -ignore_editlist: the windows come from tfdt, and an elst would shift FFmpeg's clock off them
                 cut = subprocess.Popen(
-                    [ffmpeg, "-nostdin", "-v", "error", *window, "-i", str(path)]
+                    [ffmpeg, "-nostdin", "-v", "error", "-ignore_editlist", "1", *window, "-i", str(path)]
                     + ["-map", "0:v?", "-map", "0:a?", "-c", "copy", "-f", "nut", "pipe:1"],
                     stdout=subprocess.PIPE,
                     stderr=subprocess.DEVNULL,
@@ -131,7 +132,7 @@ def ffmpeg_decodes(path: Path, start: Optional[float] = None, seconds: float = 3
 
     if start is not None:
         return decode(["-ss", f"{start:.6f}", "-t", f"{max(seconds - 0.5, seconds / 2):.6f}"], keyframes=False)
-    return decode(["-t", str(seconds)], keyframes=False) and decode(["-sseof", "-10", "-t", "10"], keyframes=False)
+    return decode(["-t", str(seconds)], keyframes=False) and decode(["-sseof", "-4", "-t", "4"], keyframes=False)
 
 
 def ffprobe(uri: Union[bytes, Path]) -> dict:

unshackle/vaults/SQLite.py

@@ -113,9 +113,7 @@ class SQLite(Vault):
         self.create_bad_keys_table()
         bad = set(self.conn_factory.get().execute("SELECT lower(`kid`), lower(`key_`) FROM `bad_keys`").fetchall())
         kid_keys = {
-            kid: key_
-            for kid, key_ in kid_keys.items()
-            if ((kid.hex if isinstance(kid, UUID) else kid.replace("-", "").lower()), key_.lower()) not in bad
+            kid: key_ for kid, key_ in kid_keys.items() if (kid.replace("-", "").lower(), key_.lower()) not in bad
         }
 
         if not kid_keys: