add unit test for getKeyIdsLinkedToSession

canalplus · Jun 19, 2024 · fa443d4 · fa443d4
1 parent 7064774
commit fa443d4
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 2 deletions.
diff --git a/src/main_thread/decrypt/__tests__/__global__/content_decryptor.test.ts b/src/main_thread/decrypt/__tests__/__global__/content_decryptor.test.ts
@@ -1,5 +1,7 @@
 import { describe, beforeEach, it, expect, vi } from "vitest";
-import { getMissingKeyIds } from "../../content_decryptor";
+import { getKeyIdsLinkedToSession, getMissingKeyIds } from "../../content_decryptor";
+import InitDataValuesContainer from "../../utils/init_data_values_container";
+import KeySessionRecord from "../../utils/key_session_record";
 
 describe("content_decryptor - blacklist missing key Ids", () => {
   beforeEach(() => {
@@ -43,3 +45,62 @@ describe("content_decryptor - blacklist missing key Ids", () => {
     expect(result).toEqual([new Uint8Array([2]), new Uint8Array([4])]);
   });
 });
+
+describe("content_decryptor - getKeyIdsLinkedToSession", () => {
+  it("should return the whitelisted and blacklisted depending if keys are usable or not", () => {
+    const initDataValues = new InitDataValuesContainer([
+      {
+        systemId: "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed" /* Widevine */,
+        data: new Uint8Array([1, 2, 3]),
+      },
+    ]);
+
+    const keySessionRecord = new KeySessionRecord({
+      type: "cenc",
+      values: initDataValues,
+    });
+    const { whitelisted, blacklisted } = getKeyIdsLinkedToSession(
+      /* InitData */ {
+        type: "cenc",
+        values: initDataValues,
+      },
+      /* KeySessionRecord */ keySessionRecord,
+      /* SingleLicensePer */ "init-data",
+      /* isCurrentLicense */ true,
+      /* usableKeys */ [new Uint8Array([1]), new Uint8Array([2])],
+      /* unusableKeys */ [new Uint8Array([3])],
+    );
+
+    expect(whitelisted).toStrictEqual([new Uint8Array([1]), new Uint8Array([2])]);
+    expect(blacklisted).toStrictEqual([new Uint8Array([3])]);
+  });
+
+  it("should have 1 blacklisted keyIds because it's in the initData but it's not in the usableKeys list", () => {
+    const initDataValues = new InitDataValuesContainer([
+      {
+        systemId: "edef8ba9-79d6-4ace-a3c8-27dcd51d21ed" /* Widevine */,
+        data: new Uint8Array([1, 2, 3]),
+      },
+    ]);
+
+    const keySessionRecord = new KeySessionRecord({
+      type: "cenc",
+      values: initDataValues,
+    });
+    const { whitelisted, blacklisted } = getKeyIdsLinkedToSession(
+      /* InitData */ {
+        type: "cenc",
+        values: initDataValues,
+        keyIds: [new Uint8Array([1]), new Uint8Array([2]), new Uint8Array([3])],
+      },
+      /* KeySessionRecord */ keySessionRecord,
+      /* SingleLicensePer */ "init-data",
+      /* isCurrentLicense */ true,
+      /* usableKeys */ [new Uint8Array([1]), new Uint8Array([2])],
+      /* unusableKeys */ [],
+    );
+
+    expect(whitelisted).toStrictEqual([new Uint8Array([1]), new Uint8Array([2])]);
+    expect(blacklisted).toStrictEqual([new Uint8Array([3])]); // KeyId 3 is not in usableKeys list
+  });
+});
diff --git a/src/main_thread/decrypt/content_decryptor.ts b/src/main_thread/decrypt/content_decryptor.ts
@@ -998,7 +998,7 @@ export function getMissingInitDataKeyIds(
  *   - `blacklisted`: Array of key ids for keys that are considered unusable.
  *     The qualities linked to those keys should not be played.
  */
-function getKeyIdsLinkedToSession(
+export function getKeyIdsLinkedToSession(
   initializationData: IProcessedProtectionData,
   keySessionRecord: KeySessionRecord,
   singleLicensePer: undefined | "init-data" | "content" | "periods",