From b6baeb6cb065bc5ed071ca6caa4f9fec8a538d60 Mon Sep 17 00:00:00 2001
From: ibaker <ibaker@google.com>
Date: Wed, 30 Oct 2024 05:14:28 -0700
Subject: [PATCH] Support CEA-608 subtitles in Dolby Vision

Issue: androidx/media#1820

#cherrypick

PiperOrigin-RevId: 691378476
(cherry picked from commit 27371db225c80316cdb8e6abdb9220da737ced31)
---
 RELEASENOTES.md                               |  2 ++
 .../media3/container/NalUnitUtil.java         | 20 +++++++++++++++++++
 .../extractor/mp4/FragmentedMp4Extractor.java |  9 +++++++--
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/RELEASENOTES.md b/RELEASENOTES.md
index 59d4945021..b2a5db9b5a 100644
--- a/RELEASENOTES.md
+++ b/RELEASENOTES.md
@@ -21,6 +21,8 @@
     *   Fix CEA-608 subtitles in H.264 MPEG-TS streams not being output (this
         was broken in `1.5.0-alpha01` by
         https://github.com/androidx/media/commit/03a205f220ecf7681f85f8a752227e3986e257ff).
+    *   Support CEA-608 subtitles in Dolby Vision content
+        ([#1820](https://github.com/androidx/media/issues/1820)).
 *   Metadata:
 *   Image:
 *   DRM:
diff --git a/libraries/container/src/main/java/androidx/media3/container/NalUnitUtil.java b/libraries/container/src/main/java/androidx/media3/container/NalUnitUtil.java
index 080505e219..195566b2d9 100644
--- a/libraries/container/src/main/java/androidx/media3/container/NalUnitUtil.java
+++ b/libraries/container/src/main/java/androidx/media3/container/NalUnitUtil.java
@@ -15,6 +15,7 @@
  */
 package androidx.media3.container;
 
+import static androidx.media3.common.MimeTypes.containsCodecsCorrespondingToMimeType;
 import static com.google.common.math.DoubleMath.log2;
 import static java.lang.Math.max;
 import static java.lang.Math.min;
@@ -33,6 +34,7 @@ import java.math.RoundingMode;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Objects;
 
 /** Utility methods for handling H.264/AVC and H.265/HEVC NAL units. */
 @UnstableApi
@@ -614,6 +616,24 @@ public final class NalUnitUtil {
             && ((nalUnitHeaderFirstByte & 0x7E) >> 1) == H265_NAL_UNIT_TYPE_PREFIX_SEI);
   }
 
+  /**
+   * Returns whether the NAL unit with the specified header contains supplemental enhancement
+   * information.
+   *
+   * @param format The sample {@link Format}.
+   * @param nalUnitHeaderFirstByte The first byte of nal_unit().
+   * @return Whether the NAL unit with the specified header is an SEI NAL unit. False is returned if
+   *     the {@code MimeType} is {@code null}.
+   */
+  public static boolean isNalUnitSei(Format format, byte nalUnitHeaderFirstByte) {
+    return ((Objects.equals(format.sampleMimeType, MimeTypes.VIDEO_H264)
+                || containsCodecsCorrespondingToMimeType(format.codecs, MimeTypes.VIDEO_H264))
+            && (nalUnitHeaderFirstByte & 0x1F) == H264_NAL_UNIT_TYPE_SEI)
+        || ((Objects.equals(format.sampleMimeType, MimeTypes.VIDEO_H265)
+                || containsCodecsCorrespondingToMimeType(format.codecs, MimeTypes.VIDEO_H265))
+            && ((nalUnitHeaderFirstByte & 0x7E) >> 1) == H265_NAL_UNIT_TYPE_PREFIX_SEI);
+  }
+
   /**
    * Returns the type of the NAL unit in {@code data} that starts at {@code offset}.
    *
diff --git a/libraries/extractor/src/main/java/androidx/media3/extractor/mp4/FragmentedMp4Extractor.java b/libraries/extractor/src/main/java/androidx/media3/extractor/mp4/FragmentedMp4Extractor.java
index 18521ed7fd..db44bc77c6 100644
--- a/libraries/extractor/src/main/java/androidx/media3/extractor/mp4/FragmentedMp4Extractor.java
+++ b/libraries/extractor/src/main/java/androidx/media3/extractor/mp4/FragmentedMp4Extractor.java
@@ -1608,7 +1608,7 @@ public class FragmentedMp4Extractor implements Extractor {
           output.sampleData(nalPrefix, 1);
           processSeiNalUnitPayload =
               ceaTrackOutputs.length > 0
-                  && NalUnitUtil.isNalUnitSei(track.format.sampleMimeType, nalPrefixData[4]);
+                  && NalUnitUtil.isNalUnitSei(track.format, nalPrefixData[4]);
           sampleBytesWritten += 5;
           sampleSize += nalUnitLengthFieldLengthDiff;
           if (!isSampleDependedOn
@@ -1629,7 +1629,12 @@ public class FragmentedMp4Extractor implements Extractor {
             int unescapedLength =
                 NalUnitUtil.unescapeStream(nalBuffer.getData(), nalBuffer.limit());
             // 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.setPosition(
+                Objects.equals(track.format.sampleMimeType, MimeTypes.VIDEO_H265)
+                        || MimeTypes.containsCodecsCorrespondingToMimeType(
+                            track.format.codecs, MimeTypes.VIDEO_H265)
+                    ? 1
+                    : 0);
             nalBuffer.setLimit(unescapedLength);
 
             if (track.format.maxNumReorderSamples != Format.NO_VALUE