Add pssh box parsing and expose them in resolveKey

docs/guide/reading-hls.md

@@ -352,7 +352,10 @@ using input = new Input({
 	formats: ALL_FORMATS,
 	formatOptions: {
 		isobmff: {
-			resolveKeyId: ({ keyId }) => {
+			resolveKeyId: ({ keyId, psshBoxes }) => {
+				// psshBoxes contains Protection System Specific Header boxes
+				// relevant to this key ID. They can be used to obtain a
+				// decryption key from a DRM license server.
 				const key = keyMap.get(keyId);
 				if (!key) {
 					throw new Error('Unknown key ID.');

src/hls/hls-segmented-input.ts

@@ -9,7 +9,16 @@
 import { AES_128_BLOCK_SIZE, createAes128CbcDecryptStream } from '../aes';
 import { ENCRYPTION_KEY_CACHE_GROUP, Input } from '../input';
 import { Segment, SegmentedInput, SegmentedInputTrackDeclaration, SegmentRetrievalOptions } from '../segmented-input';
-import { toDataView, joinPaths, last, assert, binarySearchLessOrEqual, arrayArgmin, wait } from '../misc';
+import {
+	toDataView,
+	joinPaths,
+	last,
+	assert,
+	binarySearchLessOrEqual,
+	arrayArgmin,
+	wait,
+	base64ToBytes,
+} from '../misc';
 import { readAllLines, readBytes, Reader } from '../reader';
 import { CustomPathedSource, ReadableStreamSource, SourceRef, SourceRequest } from '../source';
 import { HlsDemuxer } from './hls-demuxer';
@@ -27,9 +36,11 @@ import {
 	TAG_PROGRAM_DATE_TIME,
 	TAG_TARGETDURATION,
 } from './hls-misc';
-import { HlsInputFormat } from '../input-format';
+import { HlsInputFormat, type InputFormatOptions } from '../input-format';
+import { parsePsshBoxContents, psshBoxesAreEqual, type PsshBox } from '../isobmff/isobmff-misc';
 
 const IV_STRING_REGEX = /^0[xX][0-9a-fA-F]+$/;
+const BASE64_DATA_URI_REGEX = /^data:.*;base64,/i;
 
 export type HlsSegment = Segment & {
 	sequenceNumber: number | null;
@@ -47,6 +58,7 @@ export type HlsEncryptionInfo = {
 	keyFormat: string;
 } | {
 	method: 'SAMPLE-AES' | 'SAMPLE-AES-CTR';
+	psshBox: PsshBox | null;
 };
 
 export type HlsSegmentLocation = {
@@ -368,6 +380,11 @@ export class HlsSegmentedInput extends SegmentedInput {
 						keyFormat,
 					};
 				} else if (method === 'SAMPLE-AES' || method === 'SAMPLE-AES-CTR') {
+					const uri = attributes.get('uri');
+					if (!uri) {
+						throw new Error(`Invalid #EXT-X-KEY: ${method} requires a URI attribute.`);
+					}
+
 					const keyFormat = attributes.get('keyformat') ?? 'identity';
 					if (keyFormat === 'identity') {
 						throw new Error(
@@ -376,8 +393,26 @@ export class HlsSegmentedInput extends SegmentedInput {
 						);
 					}
 
+					let psshBox: PsshBox | null = null;
+					if (BASE64_DATA_URI_REGEX.test(uri)) {
+						const commaIndex = uri.indexOf(',');
+						const bytes = base64ToBytes(uri.slice(commaIndex + 1));
+
+						if (
+							bytes.length >= 8
+							&& bytes[4] === 0x70
+							&& bytes[5] === 0x73
+							&& bytes[6] === 0x73
+							&& bytes[7] === 0x68
+						) {
+							const size = toDataView(bytes).getUint32(0);
+							psshBox = parsePsshBoxContents(bytes.subarray(8, Math.min(size, bytes.length)));
+						}
+					}
+
 					currentKey = {
 						method,
+						psshBox,
 					};
 				} else {
 					throw new Error(
@@ -562,6 +597,38 @@ export class HlsSegmentedInput extends SegmentedInput {
 			initInput = this.getInputForSegment((hlsSegment.initSegment ?? hlsSegment.firstSegment)!);
 		}
 
+		const formatOptions: InputFormatOptions = {
+			...this.input._formatOptions,
+			isobmff: {
+				...this.input._formatOptions.isobmff,
+				// Intercept calls to resolveKeyId to inject our psshBox knowledge into it
+				resolveKeyId: this.input._formatOptions.isobmff?.resolveKeyId && ((options) => {
+					if (
+						!hlsSegment.encryption
+						|| !(
+							hlsSegment.encryption.method === 'SAMPLE-AES'
+							|| hlsSegment.encryption.method === 'SAMPLE-AES-CTR'
+						)
+						|| !hlsSegment.encryption.psshBox
+					) {
+						return this.input._formatOptions.isobmff!.resolveKeyId!(options);
+					}
+
+					let psshBoxes = options.psshBoxes;
+					const { psshBox } = hlsSegment.encryption;
+
+					if (
+						(psshBox.keyIds === null || psshBox.keyIds.includes(options.keyId))
+						&& !psshBoxes.some(x => psshBoxesAreEqual(x, psshBox))
+					) {
+						psshBoxes = [...psshBoxes, psshBox];
+					}
+
+					return this.input._formatOptions.isobmff!.resolveKeyId!({ ...options, psshBoxes });
+				}),
+			},
+		};
+
 		const input = new Input({
 			source: new CustomPathedSource(
 				hlsSegment.location.path,
@@ -638,7 +705,7 @@ export class HlsSegmentedInput extends SegmentedInput {
 			// Do not allow recursive HLS. Cool on paper, but allows for nasty infinite-depth request trees.
 			formats: this.input._formats.filter(x => !(x instanceof HlsInputFormat)),
 			initInput: initInput ?? undefined,
-			formatOptions: this.input._formatOptions,
+			formatOptions,
 		});
 
 		input._onFormatDetermined = (format) => {

src/index.ts

@@ -148,6 +148,7 @@ export {
 	EventListenerOptions,
 	FilePath,
 	MaybePromise,
+	PsshBox,
 	Rational,
 	Rectangle,
 	Rotation,

src/input-format.ts

@@ -9,6 +9,7 @@
 import { Demuxer } from './demuxer';
 import { Input } from './input';
 import { IsobmffDemuxer } from './isobmff/isobmff-demuxer';
+import type { PsshBox } from './isobmff/isobmff-misc';
 import {
 	EBMLId,
 	MAX_HEADER_SIZE,
@@ -738,7 +739,15 @@ export type IsobmffInputFormatOptions = {
 	resolveKeyId?: (options: {
 		/** The key ID that is to be resolved to a key. This is a 32-character lowercase hexadecimal string. */
 		keyId: string;
+		/**
+		 * Protection System Specific Header (pssh) boxes that apply to this key ID. Can be used to obtain a
+		 * description key from a DRM license server.
+		 */
+		psshBoxes: PsshBox[];
 	}) => MaybePromise<Uint8Array | string>;
+
+	/** @internal */
+	_suppressPsshParsing?: boolean;
 };
 
 export const validateInputFormatOptions = (options: InputFormatOptions, prefix: string) => {

src/isobmff/isobmff-demuxer.ts

@@ -64,7 +64,7 @@ import {
 	HEX_STRING_REGEX,
 } from '../misc';
 import { EncodedPacket, PLACEHOLDER_DATA } from '../packet';
-import { buildIsobmffMimeType } from './isobmff-misc';
+import { buildIsobmffMimeType, parsePsshBoxContents, psshBoxesAreEqual, PsshBox } from './isobmff-misc';
 import {
 	MAX_BOX_HEADER_SIZE,
 	MIN_BOX_HEADER_SIZE,
@@ -254,6 +254,7 @@ type Fragment = {
 	moofSize: number;
 	implicitBaseDataOffset: number;
 	trackData: Map<InternalTrack['id'], FragmentTrackData>;
+	psshBoxes: PsshBox[];
 };
 
 type TrackEncryptionInfo = {
@@ -301,6 +302,7 @@ export class IsobmffDemuxer extends Demuxer {
 
 	isFragmented = false;
 	fragmentTrackDefaults: FragmentTrackDefaults[] = [];
+	psshBoxes: PsshBox[] = [];
 	currentFragment: Fragment | null = null;
 	/**
 	 * Caches the last fragment that was read. Based on the assumption that there will be multiple reads to the
@@ -405,6 +407,7 @@ export class IsobmffDemuxer extends Demuxer {
 					this.metadataTags = initDemuxer.metadataTags;
 					this.isFragmented = true;
 					this.fragmentTrackDefaults = initDemuxer.fragmentTrackDefaults;
+					this.psshBoxes = initDemuxer.psshBoxes;
 
 					// Create tracks from the init input's tracks
 					for (const foreignTrack of initDemuxer.tracks) {
@@ -2041,6 +2044,7 @@ export class IsobmffDemuxer extends Demuxer {
 					moofSize: boxInfo.totalSize,
 					implicitBaseDataOffset: startPos,
 					trackData: new Map(),
+					psshBoxes: [],
 				};
 
 				this.readContiguousBoxes(slice.slice(contentStartPos, boxInfo.contentSize));
@@ -2108,6 +2112,20 @@ export class IsobmffDemuxer extends Demuxer {
 				}
 			}; break;
 
+			case 'pssh': {
+				if (this.input._formatOptions.isobmff?._suppressPsshParsing) {
+					break;
+				}
+
+				const psshBox = parsePsshBoxContents(readBytes(slice, boxInfo.contentSize));
+
+				if (this.currentFragment) {
+					this.currentFragment.psshBoxes.push(psshBox);
+				} else if (!this.currentTrack) {
+					this.psshBoxes.push(psshBox);
+				}
+			}; break;
+
 			case 'tfhd': {
 				assert(this.currentFragment);
 
@@ -3063,7 +3081,7 @@ abstract class IsobmffTrackBacking implements InputTrackBacking {
 				);
 
 				if (sampleIndex < entries.length) {
-					data = await decryptSample(this.internalTrack, entries[sampleIndex]!, data);
+					data = await decryptSample(this.internalTrack, entries[sampleIndex]!, data, null);
 				}
 			}
 		}
@@ -3110,7 +3128,7 @@ abstract class IsobmffTrackBacking implements InputTrackBacking {
 			data = readBytes(slice, fragmentSample.byteSize);
 
 			if (fragmentSample.encryption) {
-				data = await decryptSample(this.internalTrack, fragmentSample.encryption, data);
+				data = await decryptSample(this.internalTrack, fragmentSample.encryption, data, fragment);
 			}
 		}
 
@@ -3666,6 +3684,7 @@ const decryptSample = async (
 	track: InternalTrack,
 	sampleEncryption: SampleEncryptionInfo,
 	data: Uint8Array,
+	fragment: Fragment | null,
 ): Promise<Uint8Array> => {
 	assert(track.encryptionInfo);
 	const encryptionInfo = track.encryptionInfo;
@@ -3686,7 +3705,25 @@ const decryptSample = async (
 		}
 
 		const promise = (async () => {
-			const keyResult = await track.demuxer.input._formatOptions.isobmff!.resolveKeyId!({ keyId });
+			let psshBoxes = track.demuxer.psshBoxes;
+			if (fragment) {
+				psshBoxes = [
+					...psshBoxes,
+					...fragment.psshBoxes,
+				].filter(x => x.keyIds === null || x.keyIds.includes(keyId));
+
+				// Filter out duplicates
+				for (let i = 0; i < psshBoxes.length - 1; i++) {
+					for (let j = i + 1; j < psshBoxes.length; j++) {
+						if (psshBoxesAreEqual(psshBoxes[i]!, psshBoxes[j]!)) {
+							psshBoxes.splice(j, 1);
+							j--;
+						}
+					}
+				}
+			}
+
+			const keyResult = await track.demuxer.input._formatOptions.isobmff!.resolveKeyId!({ keyId, psshBoxes });
 
 			if (!(
 				(typeof keyResult === 'string' && keyResult.length === 32 && HEX_STRING_REGEX.test(keyResult))

src/isobmff/isobmff-misc.ts

@@ -6,6 +6,8 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/.
  */
 
+import { bytesToHexString, toDataView, uint8ArraysAreEqual } from '../misc';
+
 export const buildIsobmffMimeType = (info: {
 	isQuickTime: boolean;
 	hasVideo: boolean;
@@ -27,3 +29,62 @@ export const buildIsobmffMimeType = (info: {
 
 	return string;
 };
+
+/**
+ * Represents a Protection System Specific Header box as used by ISOBMFF Common Encryption. Contains
+ * DRM system-specific data that can be used to obtain a decryption key.
+ *
+ * @group Miscellaneous
+ * @public
+ */
+export type PsshBox = {
+	/** The system ID as a 32-bit lowercase hex string. */
+	systemId: string;
+	/**
+	 * The list of key IDs (32-bit lowercase hex strings) this box applies to, or `null` if it applies to all key IDs.
+	 */
+	keyIds: string[] | null;
+	/** The content protection system-specific data. */
+	data: Uint8Array;
+};
+
+export const parsePsshBoxContents = (contents: Uint8Array): PsshBox => {
+	const view = toDataView(contents);
+	let pos = 0;
+
+	const version = view.getUint8(pos);
+	pos += 1;
+
+	pos += 3; // Flags
+
+	const systemId = bytesToHexString(contents.subarray(pos, pos + 16));
+	pos += 16;
+
+	let keyIds: string[] | null = null;
+	if (version > 0) {
+		const kidCount = view.getUint32(pos);
+		pos += 4;
+
+		if (kidCount > 0) {
+			keyIds = [];
+			for (let i = 0; i < kidCount; i++) {
+				keyIds.push(bytesToHexString(contents.subarray(pos, pos + 16)));
+				pos += 16;
+			}
+		}
+	}
+
+	const dataSize = view.getUint32(pos);
+	pos += 4;
+
+	return {
+		systemId,
+		keyIds,
+		data: contents.slice(pos, pos + dataSize),
+	};
+};
+
+export const psshBoxesAreEqual = (a: PsshBox, b: PsshBox) => (
+	a.systemId === b.systemId
+	&& uint8ArraysAreEqual(a.data, b.data)
+);

src/misc.ts

@@ -1309,3 +1309,5 @@ export class ConcurrentRunner {
 		await Promise.all(this._queue);
 	}
 }
+
+export { PsshBox } from './isobmff/isobmff-misc';

test/node/hls-input.test.ts

@@ -897,7 +897,17 @@ test.concurrent('Widevine encryption (SAMPLE-AES-CTR) succeeds with string keys'
 		formats: ALL_FORMATS,
 		formatOptions: {
 			isobmff: {
-				resolveKeyId: ({ keyId }) => {
+				resolveKeyId: ({ keyId, psshBoxes }) => {
+					expect(psshBoxes).toHaveLength(1);
+					expect(psshBoxes[0]!.systemId).toBe('edef8ba979d64acea3c827dcd51d21ed');
+					expect(psshBoxes[0]!.keyIds).toBeNull();
+					expect(psshBoxes[0]!.data).toEqual(new Uint8Array([
+						34, 22, 115, 104, 97, 107, 97, 95,
+						99, 101, 99, 50, 102, 54, 52, 97,
+						97, 55, 56, 57, 48, 97, 49, 49,
+						72, 227, 220, 149, 155, 6,
+					]));
+
 					const key = keyMap.get(keyId);
 					assert(key);
 
@@ -949,6 +959,50 @@ test.concurrent('Widevine encryption (SAMPLE-AES-CTR) succeeds with buffer keys'
 	expect(lastPacket.timestamp + lastPacket.duration).toBe(60);
 });
 
+test.concurrent('Widevine HLS passes #EXT-X-KEY PSSH boxes to key resolver', async () => {
+	const keyMap = new Map([
+		['4d97930a3d7b55fa81d0028653f5e499', '429ec76475e7a952d224d8ef867f12b6'],
+		['d21373c0b8ab5ba9954742bcdfb5f48b', '150a6c7d7dee6a91b74dccfce5b31928'],
+		['6f1729072b4a5cd288c916e11846b89e', 'a84b4bd66901874556093454c075e2c6'],
+		['800aacaa522958ae888062b5695db6bf', '775dbf7289c4cc5847becd571f536ff2'],
+		['67b30c86756f57c5a0a38a23ac8c9178', 'efa2878c2ccf6dd47ab349fcf90e6259'],
+	]);
+
+	using input = new Input({
+		source: new UrlSource('https://storage.googleapis.com/shaka-demo-assets/angel-one-widevine-hls/hls.m3u8'),
+		formats: ALL_FORMATS,
+		formatOptions: {
+			isobmff: {
+				_suppressPsshParsing: true,
+				resolveKeyId: ({ keyId, psshBoxes }) => {
+					expect(psshBoxes).toHaveLength(1);
+					expect(psshBoxes[0]!.systemId).toBe('edef8ba979d64acea3c827dcd51d21ed');
+					expect(psshBoxes[0]!.keyIds).toBeNull();
+					expect(psshBoxes[0]!.data).toEqual(new Uint8Array([
+						34, 22, 115, 104, 97, 107, 97, 95,
+						99, 101, 99, 50, 102, 54, 52, 97,
+						97, 55, 56, 57, 48, 97, 49, 49,
+						72, 227, 220, 149, 155, 6,
+					]));
+
+					const key = keyMap.get(keyId);
+					assert(key);
+
+					return key;
+				},
+			},
+		},
+	});
+
+	const videoTrack = await input.getPrimaryVideoTrack();
+	assert(videoTrack);
+
+	const sink = new EncodedPacketSink(videoTrack);
+	const lastPacket = await sink.getPacket(Infinity);
+	assert(lastPacket);
+	expect(lastPacket.timestamp + lastPacket.duration).toBe(60);
+});
+
 test.concurrent('SourceRequest.isRoot', async () => {
 	using input = new Input({
 		source: new CustomPathedSource(