From 5c571e6e9d24086abede757f5be5d275cd67a515 Mon Sep 17 00:00:00 2001
From: andrewlewis <andrewlewis@google.com>
Date: Wed, 15 Feb 2017 09:58:44 -0800
Subject: [PATCH] Handle H.265/HEVC SEI NAL units in FragmentedMp4Extractor.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=147609330
---
 .../extractor/mp4/FragmentedMp4Extractor.java   |  6 ++++--
 .../android/exoplayer2/util/NalUnitUtil.java    | 17 +++++++++++------
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/library/src/main/java/com/google/android/exoplayer2/extractor/mp4/FragmentedMp4Extractor.java b/library/src/main/java/com/google/android/exoplayer2/extractor/mp4/FragmentedMp4Extractor.java
index bc9b0fcad6..8144880338 100644
--- a/library/src/main/java/com/google/android/exoplayer2/extractor/mp4/FragmentedMp4Extractor.java
+++ b/library/src/main/java/com/google/android/exoplayer2/extractor/mp4/FragmentedMp4Extractor.java
@@ -1087,7 +1087,7 @@ public final class FragmentedMp4Extractor implements Extractor {
           // Write the NAL unit type byte.
           output.sampleData(nalPrefix, 1);
           processSeiNalUnitPayload = cea608TrackOutput != null
-              && NalUnitUtil.isNalUnitSei(nalPrefixData[4]);
+              && NalUnitUtil.isNalUnitSei(track.format.sampleMimeType, nalPrefixData[4]);
           sampleBytesWritten += 5;
           sampleSize += nalUnitLengthFieldLengthDiff;
         } else {
@@ -1100,7 +1100,9 @@ public final class FragmentedMp4Extractor implements Extractor {
             writtenBytes = sampleCurrentNalBytesRemaining;
             // Unescape and process the SEI NAL unit.
             int unescapedLength = NalUnitUtil.unescapeStream(nalBuffer.data, nalBuffer.limit());
-            nalBuffer.reset(unescapedLength);
+            // If the format is H.265/HEVC the NAL unit header has two bytes so skip one more byte.
+            nalBuffer.setPosition(MimeTypes.VIDEO_H265.equals(track.format.sampleMimeType) ? 1 : 0);
+            nalBuffer.setLimit(unescapedLength);
             CeaUtil.consume(fragment.getSamplePresentationTime(sampleIndex) * 1000L, nalBuffer,
                 cea608TrackOutput);
           } else {
diff --git a/library/src/main/java/com/google/android/exoplayer2/util/NalUnitUtil.java b/library/src/main/java/com/google/android/exoplayer2/util/NalUnitUtil.java
index a2643d5177..ab2fec0db7 100644
--- a/library/src/main/java/com/google/android/exoplayer2/util/NalUnitUtil.java
+++ b/library/src/main/java/com/google/android/exoplayer2/util/NalUnitUtil.java
@@ -103,8 +103,9 @@ public final class NalUnitUtil {
     2f
   };
 
-  private static final int NAL_UNIT_TYPE_SEI = 6; // Supplemental enhancement information
-  private static final int NAL_UNIT_TYPE_SPS = 7; // Sequence parameter set
+  private static final int H264_NAL_UNIT_TYPE_SEI = 6; // Supplemental enhancement information
+  private static final int H264_NAL_UNIT_TYPE_SPS = 7; // Sequence parameter set
+  private static final int H265_NAL_UNIT_TYPE_PREFIX_SEI = 39;
 
   private static final Object scratchEscapePositionsLock = new Object();
 
@@ -177,7 +178,7 @@ public final class NalUnitUtil {
     while (offset + 1 < length) {
       int value = data.get(offset) & 0xFF;
       if (consecutiveZeros == 3) {
-        if (value == 1 && (data.get(offset + 1) & 0x1F) == NAL_UNIT_TYPE_SPS) {
+        if (value == 1 && (data.get(offset + 1) & 0x1F) == H264_NAL_UNIT_TYPE_SPS) {
           // Copy from this NAL unit onwards to the start of the buffer.
           ByteBuffer offsetData = data.duplicate();
           offsetData.position(offset - 3);
@@ -202,11 +203,15 @@ public final class NalUnitUtil {
    * Returns whether the NAL unit with the specified header contains supplemental enhancement
    * information.
    *
-   * @param nalUnitHeader The header of the NAL unit (first byte of nal_unit()).
+   * @param mimeType The sample MIME type.
+   * @param nalUnitHeaderFirstByte The first byte of nal_unit().
    * @return Whether the NAL unit with the specified header is an SEI NAL unit.
    */
-  public static boolean isNalUnitSei(byte nalUnitHeader) {
-    return (nalUnitHeader & 0x1F) == NAL_UNIT_TYPE_SEI;
+  public static boolean isNalUnitSei(String mimeType, byte nalUnitHeaderFirstByte) {
+    return (MimeTypes.VIDEO_H264.equals(mimeType)
+        && (nalUnitHeaderFirstByte & 0x1F) == H264_NAL_UNIT_TYPE_SEI)
+        || (MimeTypes.VIDEO_H265.equals(mimeType)
+        && ((nalUnitHeaderFirstByte & 0x7E) >> 1) == H265_NAL_UNIT_TYPE_PREFIX_SEI);
   }
 
   /**