From dbf864fd6911bccce99adfbc8f286b0b3d8cc93f Mon Sep 17 00:00:00 2001
From: Daniele Sparano <daniele.sparano@v-nova.com>
Date: Tue, 4 Feb 2025 10:50:16 +0000
Subject: [PATCH] Handle 4 byte start code prefix, ensure non empty samples,
 fix timestamp on very last sample on synthesized empty pusi

---
 .../androidx/media3/extractor/ts/H265Reader.java    | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/libraries/extractor/src/main/java/androidx/media3/extractor/ts/H265Reader.java b/libraries/extractor/src/main/java/androidx/media3/extractor/ts/H265Reader.java
index cddec8b0e4..cabdfc6be3 100644
--- a/libraries/extractor/src/main/java/androidx/media3/extractor/ts/H265Reader.java
+++ b/libraries/extractor/src/main/java/androidx/media3/extractor/ts/H265Reader.java
@@ -128,6 +128,7 @@ public final class H265Reader implements ElementaryStreamReader {
 
       // Scan the appended data, processing NAL units as they are encountered
       while (offset < limit) {
+        int prefixSize = 3;
         int nalUnitOffset = NalUnitUtil.findNalUnit(dataArray, offset, limit, prefixFlags);
 
         if (nalUnitOffset == limit) {
@@ -139,6 +140,13 @@ public final class H265Reader implements ElementaryStreamReader {
         // We've seen the start of a NAL unit of the following type.
         int nalUnitType = NalUnitUtil.getH265NalUnitType(dataArray, nalUnitOffset);
 
+        // Case of a 4 byte start code prefix 0x00000001, recoil NAL unit offset by one byte
+        // to avoid previous byte being assigned to the previous access unit.
+        if (nalUnitOffset > 0 && dataArray[nalUnitOffset - 1] == 0x00) {
+          nalUnitOffset--;
+          prefixSize = 4;
+        }
+
         // This is the number of bytes from the current offset to the start of the next NAL unit.
         // It may be negative if the NAL unit started in the previously consumed data.
         int lengthToNalUnit = nalUnitOffset - offset;
@@ -159,7 +167,7 @@ public final class H265Reader implements ElementaryStreamReader {
         // Indicate the start of the next NAL unit.
         startNalUnit(absolutePosition, bytesWrittenPastPosition, nalUnitType, pesTimeUs);
         // Continue scanning the data.
-        offset = nalUnitOffset + 3;
+        offset = nalUnitOffset + prefixSize;
       }
     }
   }
@@ -391,12 +399,13 @@ public final class H265Reader implements ElementaryStreamReader {
       // Output a final sample with the remaining NAL units up to the passed position
       samplePosition = nalUnitPosition;
       nalUnitPosition = position;
+      sampleTimeUs = nalUnitTimeUs;
       outputSample(/* offset= */ 0);
       readingSample = false;
     }
 
     private void outputSample(int offset) {
-      if (sampleTimeUs == C.TIME_UNSET) {
+      if (sampleTimeUs == C.TIME_UNSET || nalUnitPosition == samplePosition) {
         return;
       }
       @C.BufferFlags int flags = sampleIsKeyframe ? C.BUFFER_FLAG_KEY_FRAME : 0;