Add getPsshInfo() API to MediaExtractorCompat

This method retrieves PSSH (Protection System Specific Header) data from the media as a map of UUID-to-bytes.

PiperOrigin-RevId: 697974134

libraries/exoplayer/src/androidTest/java/androidx/media3/exoplayer/MediaExtractorCompatTest.java

@@ -19,6 +19,7 @@ import static androidx.media3.common.C.PLAYREADY_UUID;
 import static androidx.media3.common.C.WIDEVINE_UUID;
 import static androidx.media3.common.MimeTypes.AUDIO_AAC;
 import static androidx.media3.common.MimeTypes.VIDEO_H264;
+import static androidx.media3.common.MimeTypes.VIDEO_MP4;
 import static androidx.media3.test.utils.TestUtil.buildTestData;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
@@ -32,6 +33,7 @@ import android.media.metrics.MediaMetricsManager;
 import android.media.metrics.PlaybackSession;
 import android.net.Uri;
 import android.os.PersistableBundle;
+import androidx.annotation.Nullable;
 import androidx.media3.common.C;
 import androidx.media3.common.DrmInitData;
 import androidx.media3.common.Format;
@@ -51,6 +53,7 @@ import androidx.media3.extractor.SeekMap.SeekPoints;
 import androidx.media3.extractor.SeekPoint;
 import androidx.media3.extractor.TrackOutput;
 import androidx.media3.extractor.mp4.Mp4Extractor;
+import androidx.media3.extractor.mp4.PsshAtomUtil;
 import androidx.media3.test.utils.TestUtil;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -63,6 +66,7 @@ import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.UUID;
 import java.util.concurrent.atomic.AtomicBoolean;
 import okhttp3.mockwebserver.MockResponse;
 import okhttp3.mockwebserver.MockWebServer;
@@ -631,16 +635,15 @@ public class MediaExtractorCompatTest {
   @Test
   public void getTrackFormat_withBothTrackAndSeekMapDurationsSet_prioritizesTrackDuration()
       throws IOException {
-    TrackOutput[] outputs = new TrackOutput[1];
     fakeExtractor.addReadAction(
         (input, seekPosition) -> {
-          outputs[0] = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
+          TrackOutput output = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
           extractorOutput.endTracks();
           extractorOutput.seekMap(
               new FakeSeekMap(
                   /* durationUs= */ 1_000_000L, (timeUs) -> new SeekPoints(SeekPoint.START)));
-          outputs[0].format(PLACEHOLDER_FORMAT_VIDEO);
+          output.format(PLACEHOLDER_FORMAT_VIDEO);
-          outputs[0].durationUs(2_000_000L);
+          output.durationUs(2_000_000L);
           return Extractor.RESULT_CONTINUE;
         });
     mediaExtractorCompat.setDataSource(PLACEHOLDER_URI, /* offset= */ 0);
@@ -655,15 +658,14 @@ public class MediaExtractorCompatTest {
   @Test
   public void getTrackFormat_withOnlySeekMapDurationSet_returnsSeekMapDuration()
       throws IOException {
-    TrackOutput[] outputs = new TrackOutput[1];
     fakeExtractor.addReadAction(
         (input, seekPosition) -> {
-          outputs[0] = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
+          TrackOutput output = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
           extractorOutput.endTracks();
           extractorOutput.seekMap(
               new FakeSeekMap(
                   /* durationUs= */ 1_000_000L, (timeUs) -> new SeekPoints(SeekPoint.START)));
-          outputs[0].format(PLACEHOLDER_FORMAT_VIDEO);
+          output.format(PLACEHOLDER_FORMAT_VIDEO);
           return Extractor.RESULT_CONTINUE;
         });
     mediaExtractorCompat.setDataSource(PLACEHOLDER_URI, /* offset= */ 0);
@@ -678,12 +680,11 @@ public class MediaExtractorCompatTest {
   @Test
   public void getTrackFormat_withNoTrackOrSeekMapDurationSet_returnsNoDuration()
       throws IOException {
-    TrackOutput[] outputs = new TrackOutput[1];
     fakeExtractor.addReadAction(
         (input, seekPosition) -> {
-          outputs[0] = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
+          TrackOutput output = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
           extractorOutput.endTracks();
-          outputs[0].format(
+          output.format(
               new Format.Builder()
                   .setSampleMimeType(MimeTypes.VIDEO_H264)
                   .setCodecs("avc.123")
@@ -729,11 +730,10 @@ public class MediaExtractorCompatTest {
 
   @Test
   public void getDrmInitData_withNoTracksHavingDrmInitData_returnsNull() throws IOException {
-    TrackOutput[] outputs = new TrackOutput[1];
     fakeExtractor.addReadAction(
         (input, seekPosition) -> {
-          outputs[0] = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
+          TrackOutput output = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
-          outputs[0].format(PLACEHOLDER_FORMAT_VIDEO);
+          output.format(PLACEHOLDER_FORMAT_VIDEO);
           extractorOutput.endTracks();
           return Extractor.RESULT_CONTINUE;
         });
@@ -746,16 +746,13 @@ public class MediaExtractorCompatTest {
   @Test
   public void getDrmInitData_withSingleTrackHavingDrmInitData_returnsDrmInitData()
       throws IOException {
-    TrackOutput[] outputs = new TrackOutput[1];
     DrmInitData.SchemeData schemeData =
-        new DrmInitData.SchemeData(
+        new DrmInitData.SchemeData(WIDEVINE_UUID, VIDEO_H264, buildTestData(128, /* seed= */ 1));
-            WIDEVINE_UUID, VIDEO_H264, buildTestData(128, 1 /* data seed */));
     DrmInitData drmInitData = new DrmInitData(schemeData);
     fakeExtractor.addReadAction(
         (input, seekPosition) -> {
-          outputs[0] = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
+          TrackOutput output = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
-          outputs[0].format(
+          output.format(PLACEHOLDER_FORMAT_VIDEO.buildUpon().setDrmInitData(drmInitData).build());
-              PLACEHOLDER_FORMAT_VIDEO.buildUpon().setDrmInitData(drmInitData).build());
           extractorOutput.endTracks();
           return Extractor.RESULT_CONTINUE;
         });
@@ -768,16 +765,15 @@ public class MediaExtractorCompatTest {
   @Test
   public void getDrmInitData_withMultipleTracksHavingDrmInitData_returnsFirstNonNullDrmInitData()
       throws IOException {
-    TrackOutput[] outputs = new TrackOutput[3];
     DrmInitData.SchemeData firstSchemeData =
-        new DrmInitData.SchemeData(WIDEVINE_UUID, AUDIO_AAC, buildTestData(128, 1 /* data seed */));
+        new DrmInitData.SchemeData(WIDEVINE_UUID, AUDIO_AAC, buildTestData(128, /* seed= */ 1));
     DrmInitData firstDrmInitData = new DrmInitData(firstSchemeData);
     DrmInitData.SchemeData secondSchemeData =
-        new DrmInitData.SchemeData(
+        new DrmInitData.SchemeData(PLAYREADY_UUID, AUDIO_AAC, buildTestData(128, /* seed= */ 2));
-            PLAYREADY_UUID, AUDIO_AAC, buildTestData(128, 2 /* data seed */));
     DrmInitData secondDrmInitData = new DrmInitData(secondSchemeData);
     fakeExtractor.addReadAction(
         (input, seekPosition) -> {
+          TrackOutput[] outputs = new TrackOutput[3];
           outputs[0] = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
           outputs[0].format(PLACEHOLDER_FORMAT_VIDEO);
           outputs[1] = extractorOutput.track(/* id= */ 1, C.TRACK_TYPE_AUDIO);
@@ -1025,6 +1021,48 @@ public class MediaExtractorCompatTest {
     assertThat(bundle.getInt(MediaExtractor.MetricsConstants.TRACKS)).isEqualTo(2);
   }
 
+  @Test
+  public void getPsshInfo_withMediaWithoutPsshData_returnsNull() throws IOException {
+    DrmInitData.SchemeData schemeData =
+        new DrmInitData.SchemeData(WIDEVINE_UUID, VIDEO_H264, buildTestData(128, /* seed= */ 1));
+    DrmInitData drmInitData = new DrmInitData(schemeData);
+    fakeExtractor.addReadAction(
+        (input, seekPosition) -> {
+          TrackOutput output = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
+          output.format(PLACEHOLDER_FORMAT_VIDEO.buildUpon().setDrmInitData(drmInitData).build());
+          extractorOutput.endTracks();
+          return Extractor.RESULT_CONTINUE;
+        });
+
+    mediaExtractorCompat.setDataSource(PLACEHOLDER_URI, /* offset= */ 0);
+
+    assertThat(mediaExtractorCompat.getPsshInfo()).isNull();
+  }
+
+  @Test
+  public void getPsshInfo_withMediaHavingPsshData_returnsCorrectPsshMap() throws IOException {
+    byte[] rawSchemeData = new byte[] {0, 1, 2, 3, 4, 5};
+    DrmInitData.SchemeData schemeData =
+        new DrmInitData.SchemeData(
+            WIDEVINE_UUID, VIDEO_MP4, PsshAtomUtil.buildPsshAtom(WIDEVINE_UUID, rawSchemeData));
+    DrmInitData drmInitData = new DrmInitData(schemeData);
+    fakeExtractor.addReadAction(
+        (input, seekPosition) -> {
+          TrackOutput output = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
+          output.format(PLACEHOLDER_FORMAT_VIDEO.buildUpon().setDrmInitData(drmInitData).build());
+          extractorOutput.endTracks();
+          return Extractor.RESULT_CONTINUE;
+        });
+    mediaExtractorCompat.setDataSource(PLACEHOLDER_URI, /* offset= */ 0);
+
+    @Nullable Map<UUID, byte[]> psshMap = mediaExtractorCompat.getPsshInfo();
+
+    assertThat(psshMap).isNotNull();
+    assertThat(psshMap).isNotEmpty();
+    assertThat(psshMap).hasSize(1);
+    assertThat(psshMap.get(WIDEVINE_UUID)).isEqualTo(rawSchemeData);
+  }
+
   // Internal methods.
 
   private void assertReadSample(int trackIndex, long timeUs, int size, byte... sampleData) {

libraries/exoplayer/src/main/java/androidx/media3/exoplayer/MediaExtractorCompat.java

@@ -72,6 +72,7 @@ import androidx.media3.extractor.SeekMap.SeekPoints;
 import androidx.media3.extractor.SeekPoint;
 import androidx.media3.extractor.TrackOutput;
 import androidx.media3.extractor.mp4.Mp4Extractor;
+import androidx.media3.extractor.mp4.PsshAtomUtil;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
@@ -84,9 +85,11 @@ import java.lang.annotation.RetentionPolicy;
 import java.nio.ByteBuffer;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 import org.checkerframework.checker.nullness.qual.EnsuresNonNullIf;
 
 /**
@@ -682,6 +685,36 @@ public final class MediaExtractorCompat {
     return bundle;
   }
 
+  /**
+   * Extracts PSSH data from the media, if present.
+   *
+   * @return A {@link Map} of UUID-to-byte[] pairs, where the {@link UUID} identifies the crypto
+   *     scheme, and the byte array contains the scheme-specific data. Returns {@code null} if no
+   *     PSSH data exists.
+   */
+  @Nullable
+  public Map<UUID, byte[]> getPsshInfo() {
+    @Nullable DrmInitData drmInitData = getDrmInitData();
+    if (drmInitData == null) {
+      return null;
+    }
+
+    Map<UUID, byte[]> psshDataMap = new HashMap<>();
+    for (int i = 0; i < drmInitData.schemeDataCount; i++) {
+      DrmInitData.SchemeData schemeData = drmInitData.get(i);
+      if (schemeData.data == null) {
+        continue;
+      }
+
+      @Nullable PsshAtomUtil.PsshAtom parsedPsshAtom = PsshAtomUtil.parsePsshAtom(schemeData.data);
+      if (parsedPsshAtom != null) {
+        psshDataMap.put(parsedPsshAtom.uuid, parsedPsshAtom.schemeData);
+      }
+    }
+
+    return psshDataMap.isEmpty() ? null : psshDataMap;
+  }
+
   @VisibleForTesting(otherwise = NONE)
   public Allocator getAllocator() {
     return allocator;