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
This commit is contained in:
rohks 2024-11-19 05:11:49 -08:00 committed by Copybara-Service
parent b9c9d95b90
commit d702e1d496
2 changed files with 94 additions and 23 deletions

View File

@ -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);
outputs[0].durationUs(2_000_000L);
output.format(PLACEHOLDER_FORMAT_VIDEO);
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);
outputs[0].format(PLACEHOLDER_FORMAT_VIDEO);
TrackOutput output = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_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(
WIDEVINE_UUID, VIDEO_H264, buildTestData(128, 1 /* data seed */));
new DrmInitData.SchemeData(WIDEVINE_UUID, VIDEO_H264, buildTestData(128, /* seed= */ 1));
DrmInitData drmInitData = new DrmInitData(schemeData);
fakeExtractor.addReadAction(
(input, seekPosition) -> {
outputs[0] = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO);
outputs[0].format(
PLACEHOLDER_FORMAT_VIDEO.buildUpon().setDrmInitData(drmInitData).build());
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;
});
@ -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(
PLAYREADY_UUID, AUDIO_AAC, buildTestData(128, 2 /* data seed */));
new DrmInitData.SchemeData(PLAYREADY_UUID, AUDIO_AAC, buildTestData(128, /* seed= */ 2));
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) {

View File

@ -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;