From d6c223b0f6bad9333254ba0ed70b6b1f74a78a07 Mon Sep 17 00:00:00 2001
From: Vanilagy <1696106+Vanilagy@users.noreply.github.com>
Date: Fri, 5 Dec 2025 22:14:23 +0100
Subject: [PATCH] Fix keyframe sample index lookup logic for ISOBMFF in rare
 B-frame scenarios

---
 dev/demux.html                 | 10 ++++--
 src/isobmff/isobmff-demuxer.ts | 56 +++++++++++++++++++++++++---------
 src/media-sink.ts              |  5 ++-
 3 files changed, 52 insertions(+), 19 deletions(-)

diff --git a/dev/demux.html b/dev/demux.html
index 8b06ec0..cab4bd7 100644
--- a/dev/demux.html
+++ b/dev/demux.html
@@ -14,9 +14,13 @@
 			source: new Mediabunny.BlobSource(file),
 		});
 
-		const audioTrack = await input.getPrimaryAudioTrack();
-		const decoderConfig = await audioTrack.getDecoderConfig();
-		console.log(decoderConfig);
+		const videoTrack = await input.getPrimaryVideoTrack();
+		const sink = new Mediabunny.EncodedPacketSink(videoTrack);
+
+		for await (const packet of sink.packets()) {
+			console.log(packet);
+		}
+
 
 		/*
 		const videoTrack = await input.getPrimaryVideoTrack();
diff --git a/src/isobmff/isobmff-demuxer.ts b/src/isobmff/isobmff-demuxer.ts
index b76fd85..b524bb2 100644
--- a/src/isobmff/isobmff-demuxer.ts
+++ b/src/isobmff/isobmff-demuxer.ts
@@ -2489,11 +2489,8 @@ abstract class IsobmffTrackBacking implements InputTrackBacking {
 		const timestampInTimescale = this.mapTimestampIntoTimescale(timestamp);
 
 		const sampleTable = this.internalTrack.demuxer.getSampleTableForTrack(this.internalTrack);
-		const sampleIndex = getSampleIndexForTimestamp(sampleTable, timestampInTimescale);
-		const keyFrameSampleIndex = sampleIndex === -1
-			? -1
-			: getRelevantKeyframeIndexForSample(sampleTable, sampleIndex);
-		const regularPacket = await this.fetchPacketForSampleIndex(keyFrameSampleIndex, options);
+		const sampleIndex = getKeyframeSampleIndexForTimestamp(sampleTable, timestampInTimescale);
+		const regularPacket = await this.fetchPacketForSampleIndex(sampleIndex, options);
 
 		if (!sampleTableIsEmpty(sampleTable) || !this.internalTrack.demuxer.isFragmented) {
 			// Prefer the non-fragmented packet
@@ -2913,7 +2910,45 @@ const getSampleIndexForTimestamp = (sampleTable: SampleTable, timescaleUnits: nu
 
 		const entry = sampleTable.sampleTimingEntries[index]!;
 		return entry.startIndex
-			+ Math.min(Math.floor((timescaleUnits - entry.startDecodeTimestamp) / entry.delta), entry.count - 1);
+			+ Math.min(
+				Math.floor((timescaleUnits - entry.startDecodeTimestamp) / entry.delta),
+				entry.count - 1,
+			);
+	}
+};
+
+const getKeyframeSampleIndexForTimestamp = (sampleTable: SampleTable, timescaleUnits: number) => {
+	if (!sampleTable.keySampleIndices) {
+		// Every sample is a keyframe
+		return getSampleIndexForTimestamp(sampleTable, timescaleUnits);
+	}
+
+	if (sampleTable.presentationTimestamps) {
+		const index = binarySearchLessOrEqual(
+			sampleTable.presentationTimestamps,
+			timescaleUnits,
+			x => x.presentationTimestamp,
+		);
+		if (index === -1) {
+			return -1;
+		}
+
+		// Walk the samples in presentation order until we find one that's a keyframe
+		for (let i = index; i >= 0; i--) {
+			const sampleIndex = sampleTable.presentationTimestamps[i]!.sampleIndex;
+			const isKeyFrame = binarySearchExact(sampleTable.keySampleIndices, sampleIndex, x => x) !== -1;
+
+			if (isKeyFrame) {
+				return sampleIndex;
+			}
+		}
+
+		return -1;
+	} else {
+		const sampleIndex = getSampleIndexForTimestamp(sampleTable, timescaleUnits);
+
+		const index = binarySearchLessOrEqual(sampleTable.keySampleIndices, sampleIndex, x => x);
+		return sampleTable.keySampleIndices[index] ?? -1;
 	}
 };
 
@@ -3002,15 +3037,6 @@ const getSampleInfo = (sampleTable: SampleTable, sampleIndex: number): SampleInf
 	};
 };
 
-const getRelevantKeyframeIndexForSample = (sampleTable: SampleTable, sampleIndex: number) => {
-	if (!sampleTable.keySampleIndices) {
-		return sampleIndex;
-	}
-
-	const index = binarySearchLessOrEqual(sampleTable.keySampleIndices, sampleIndex, x => x);
-	return sampleTable.keySampleIndices[index] ?? -1;
-};
-
 const getNextKeyframeIndexForSample = (sampleTable: SampleTable, sampleIndex: number) => {
 	if (!sampleTable.keySampleIndices) {
 		return sampleIndex + 1;
diff --git a/src/media-sink.ts b/src/media-sink.ts
index 841d8f5..179b024 100644
--- a/src/media-sink.ts
+++ b/src/media-sink.ts
@@ -2127,11 +2127,14 @@ export class AudioBufferSink {
 
 	/** @internal */
 	_audioSampleToWrappedArrayBuffer(sample: AudioSample): WrappedAudioBuffer {
-		return {
+		const result: WrappedAudioBuffer = {
 			buffer: sample.toAudioBuffer(),
 			timestamp: sample.timestamp,
 			duration: sample.duration,
 		};
+
+		sample.close();
+		return result;
 	}
 
 	/**