Generalize getPsshInfo to properly accomodate WebM.

- Rather than returning a map, return a DrmInitData object,
  with mapped and non-mapped implementations.
- Include a suitable mimeType to pass to the MediaDrm. Previously
  we were incorrectly passing the mimeType of the samples,
  where-as MediaDrm expects the container mimeType. Note that
  it doesn't matter whether the mimeType starts with "video" or
  "audio", hence using video mimeTypes everywhere.

library/src/main/java/com/google/android/exoplayer/MediaCodecTrackRenderer.java

@@ -16,6 +16,7 @@
 package com.google.android.exoplayer;
 
 import com.google.android.exoplayer.MediaCodecUtil.DecoderQueryException;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.drm.DrmSessionManager;
 import com.google.android.exoplayer.util.Assertions;
 import com.google.android.exoplayer.util.Util;
@@ -33,8 +34,6 @@ import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
-import java.util.UUID;
 
 /**
  * An abstract {@link TrackRenderer} that uses {@link MediaCodec} to decode samples for rendering.
@@ -164,7 +163,7 @@ public abstract class MediaCodecTrackRenderer extends TrackRenderer {
   protected final Handler eventHandler;
 
   private MediaFormat format;
-  private Map<UUID, byte[]> drmInitData;
+  private DrmInitData drmInitData;
   private MediaCodec codec;
   private boolean codecIsAdaptive;
   private ByteBuffer[] inputBuffers;
@@ -281,7 +280,7 @@ public abstract class MediaCodecTrackRenderer extends TrackRenderer {
         throw new ExoPlaybackException("Media requires a DrmSessionManager");
       }
       if (!openedDrmSession) {
-        drmSessionManager.open(drmInitData, mimeType);
+        drmSessionManager.open(drmInitData);
         openedDrmSession = true;
       }
       int drmSessionState = drmSessionManager.getState();

library/src/main/java/com/google/android/exoplayer/MediaFormatHolder.java

@@ -15,8 +15,7 @@
  */
 package com.google.android.exoplayer;
 
-import java.util.Map;
+import com.google.android.exoplayer.drm.DrmInitData;
-import java.util.UUID;
 
 /**
  * Holds a {@link MediaFormat} and corresponding drm scheme initialization data.
@@ -28,9 +27,8 @@ public final class MediaFormatHolder {
    */
   public MediaFormat format;
   /**
-   * Initialization data for each of the drm schemes supported by the media, keyed by scheme UUID.
+   * Initialization data for drm schemes supported by the media. Null if the media is not encrypted.
-   * Null if the media is not encrypted.
    */
-  public Map<UUID, byte[]> drmInitData;
+  public DrmInitData drmInitData;
 
 }

library/src/main/java/com/google/android/exoplayer/chunk/ChunkSampleSource.java

@@ -352,7 +352,7 @@ public class ChunkSampleSource implements SampleSource, Loader.Callback {
     if (mediaFormat != null && !mediaFormat.equals(downstreamMediaFormat, true)) {
       chunkSource.getMaxVideoDimensions(mediaFormat);
       formatHolder.format = mediaFormat;
-      formatHolder.drmInitData = mediaChunk.getPsshInfo();
+      formatHolder.drmInitData = mediaChunk.getDrmInitData();
       downstreamMediaFormat = mediaFormat;
       return FORMAT_READ;
     }

library/src/main/java/com/google/android/exoplayer/chunk/ContainerMediaChunk.java

@@ -19,14 +19,12 @@ import com.google.android.exoplayer.MediaFormat;
 import com.google.android.exoplayer.ParserException;
 import com.google.android.exoplayer.SampleHolder;
 import com.google.android.exoplayer.chunk.parser.Extractor;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.upstream.DataSource;
 import com.google.android.exoplayer.upstream.DataSpec;
 import com.google.android.exoplayer.upstream.NonBlockingInputStream;
 import com.google.android.exoplayer.util.Assertions;
 
-import java.util.Map;
-import java.util.UUID;
-
 /**
  * A {@link MediaChunk} extracted from a container.
  */
@@ -38,7 +36,7 @@ public final class ContainerMediaChunk extends MediaChunk {
 
   private boolean prepared;
   private MediaFormat mediaFormat;
-  private Map<UUID, byte[]> psshInfo;
+  private DrmInitData drmInitData;
 
   /**
    * @deprecated Use the other constructor, passing null as {@code psshInfo}.
@@ -60,8 +58,9 @@ public final class ContainerMediaChunk extends MediaChunk {
    * @param endTimeUs The end time of the media contained by the chunk, in microseconds.
    * @param nextChunkIndex The index of the next chunk, or -1 if this is the last chunk.
    * @param extractor The extractor that will be used to extract the samples.
-   * @param psshInfo Pssh data. May be null if pssh data is present within the stream, meaning it
+   * @param drmInitData DRM initialization data. May be null if DRM initialization data is present
-   *     can be obtained directly from {@code extractor}, or if no pssh data is required.
+   *     within the stream, meaning it can be obtained directly from {@code extractor}, or if no
+   *     DRM initialization data is required.
    * @param maybeSelfContained Set to true if this chunk might be self contained, meaning it might
    *     contain a moov atom defining the media format of the chunk. This parameter can always be
    *     safely set to true. Setting to false where the chunk is known to not be self contained may
@@ -70,12 +69,12 @@ public final class ContainerMediaChunk extends MediaChunk {
    */
   public ContainerMediaChunk(DataSource dataSource, DataSpec dataSpec, Format format,
       int trigger, long startTimeUs, long endTimeUs, int nextChunkIndex, Extractor extractor,
-      Map<UUID, byte[]> psshInfo, boolean maybeSelfContained, long sampleOffsetUs) {
+      DrmInitData drmInitData, boolean maybeSelfContained, long sampleOffsetUs) {
     super(dataSource, dataSpec, format, trigger, startTimeUs, endTimeUs, nextChunkIndex);
     this.extractor = extractor;
     this.maybeSelfContained = maybeSelfContained;
     this.sampleOffsetUs = sampleOffsetUs;
-    this.psshInfo = psshInfo;
+    this.drmInitData = drmInitData;
   }
 
   @Override
@@ -111,9 +110,9 @@ public final class ContainerMediaChunk extends MediaChunk {
       }
       if (prepared) {
         mediaFormat = extractor.getFormat();
-        Map<UUID, byte[]> extractorPsshInfo = extractor.getPsshInfo();
+        DrmInitData extractorDrmInitData = extractor.getDrmInitData();
-        if (extractorPsshInfo != null) {
+        if (extractorDrmInitData != null) {
-          psshInfo = extractorPsshInfo;
+          drmInitData = extractorDrmInitData;
         }
       }
     }
@@ -145,8 +144,8 @@ public final class ContainerMediaChunk extends MediaChunk {
   }
 
   @Override
-  public Map<UUID, byte[]> getPsshInfo() {
+  public DrmInitData getDrmInitData() {
-    return psshInfo;
+    return drmInitData;
   }
 
 }

library/src/main/java/com/google/android/exoplayer/chunk/MediaChunk.java

@@ -18,12 +18,10 @@ package com.google.android.exoplayer.chunk;
 import com.google.android.exoplayer.MediaFormat;
 import com.google.android.exoplayer.ParserException;
 import com.google.android.exoplayer.SampleHolder;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.upstream.DataSource;
 import com.google.android.exoplayer.upstream.DataSpec;
 
-import java.util.Map;
-import java.util.UUID;
-
 /**
  * An abstract base class for {@link Chunk}s that contain media samples.
  */
@@ -129,12 +127,12 @@ public abstract class MediaChunk extends Chunk {
   public abstract MediaFormat getMediaFormat();
 
   /**
-   * Returns the pssh information associated with the chunk.
+   * Returns the DRM initialization data associated with the chunk.
    * <p>
    * Should only be called after the chunk has been successfully prepared.
    *
-   * @return The pssh information.
+   * @return The DRM initialization data.
    */
-  public abstract Map<UUID, byte[]> getPsshInfo();
+  public abstract DrmInitData getDrmInitData();
 
 }

library/src/main/java/com/google/android/exoplayer/chunk/SingleSampleMediaChunk.java

@@ -17,14 +17,12 @@ package com.google.android.exoplayer.chunk;
 
 import com.google.android.exoplayer.MediaFormat;
 import com.google.android.exoplayer.SampleHolder;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.upstream.DataSource;
 import com.google.android.exoplayer.upstream.DataSpec;
 import com.google.android.exoplayer.upstream.NonBlockingInputStream;
 import com.google.android.exoplayer.util.Assertions;
 
-import java.util.Map;
-import java.util.UUID;
-
 /**
  * A {@link MediaChunk} containing a single sample.
  */
@@ -132,7 +130,7 @@ public class SingleSampleMediaChunk extends MediaChunk {
   }
 
   @Override
-  public Map<UUID, byte[]> getPsshInfo() {
+  public DrmInitData getDrmInitData() {
     return null;
   }
 

library/src/main/java/com/google/android/exoplayer/chunk/parser/Extractor.java

@@ -18,11 +18,9 @@ package com.google.android.exoplayer.chunk.parser;
 import com.google.android.exoplayer.MediaFormat;
 import com.google.android.exoplayer.ParserException;
 import com.google.android.exoplayer.SampleHolder;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.upstream.NonBlockingInputStream;
 
-import java.util.Map;
-import java.util.UUID;
-
 /**
  * Facilitates extraction of media samples from a container format.
  */
@@ -42,7 +40,7 @@ public interface Extractor {
   public static final int RESULT_READ_SAMPLE = 4;
   /**
    * Initialization data was read. The parsed data can be read using {@link #getFormat()} and
-   * {@link #getPsshInfo}.
+   * {@link #getDrmInitData()}.
    */
   public static final int RESULT_READ_INIT = 8;
   /**
@@ -79,12 +77,12 @@ public interface Extractor {
   public MediaFormat getFormat();
 
   /**
-   * Returns the pssh information parsed from the stream.
+   * Returns DRM initialization data parsed from the stream.
    *
-   * @return The pssh information. May be null if pssh data has yet to be parsed, or if the stream
+   * @return The DRM initialization data. May be null if the initialization data has yet to be
-   *     does not contain any pssh data.
+   *     parsed, or if the stream does not contain any DRM initialization data.
    */
-  public Map<UUID, byte[]> getPsshInfo();
+  public DrmInitData getDrmInitData();
 
   /**
    * Consumes data from a {@link NonBlockingInputStream}.

library/src/main/java/com/google/android/exoplayer/chunk/parser/mp4/FragmentedMp4Extractor.java

@@ -21,6 +21,7 @@ import com.google.android.exoplayer.ParserException;
 import com.google.android.exoplayer.SampleHolder;
 import com.google.android.exoplayer.chunk.parser.Extractor;
 import com.google.android.exoplayer.chunk.parser.SegmentIndex;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.mp4.Atom;
 import com.google.android.exoplayer.mp4.Atom.ContainerAtom;
 import com.google.android.exoplayer.mp4.Atom.LeafAtom;
@@ -28,6 +29,7 @@ import com.google.android.exoplayer.mp4.CommonMp4AtomParsers;
 import com.google.android.exoplayer.mp4.Mp4Util;
 import com.google.android.exoplayer.mp4.Track;
 import com.google.android.exoplayer.upstream.NonBlockingInputStream;
+import com.google.android.exoplayer.util.MimeTypes;
 import com.google.android.exoplayer.util.ParsableByteArray;
 import com.google.android.exoplayer.util.Util;
 
@@ -38,10 +40,8 @@ import android.media.MediaExtractor;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.Stack;
 import java.util.UUID;
@@ -145,7 +145,7 @@ public final class FragmentedMp4Extractor implements Extractor {
   private int lastSyncSampleIndex;
 
   // Data parsed from moov and sidx atoms
-  private final HashMap<UUID, byte[]> psshData;
+  private DrmInitData.Mapped drmInitData;
   private SegmentIndex segmentIndex;
   private Track track;
   private DefaultSampleValues extendsDefaults;
@@ -165,7 +165,6 @@ public final class FragmentedMp4Extractor implements Extractor {
     extendedTypeScratch = new byte[16];
     containerAtoms = new Stack<ContainerAtom>();
     fragmentRun = new TrackFragment();
-    psshData = new HashMap<UUID, byte[]>();
   }
 
   /**
@@ -179,8 +178,8 @@ public final class FragmentedMp4Extractor implements Extractor {
   }
 
   @Override
-  public Map<UUID, byte[]> getPsshInfo() {
+  public DrmInitData getDrmInitData() {
-    return psshData.isEmpty() ? null : psshData;
+    return drmInitData;
   }
 
   @Override
@@ -370,7 +369,10 @@ public final class FragmentedMp4Extractor implements Extractor {
         int dataSize = psshAtom.readInt();
         byte[] data = new byte[dataSize];
         psshAtom.readBytes(data, 0, dataSize);
-        psshData.put(uuid, data);
+        if (drmInitData == null) {
+          drmInitData = new DrmInitData.Mapped(MimeTypes.VIDEO_MP4);
+        }
+        drmInitData.put(uuid, data);
       }
     }
     ContainerAtom mvex = moov.getContainerAtomOfType(Atom.TYPE_mvex);

library/src/main/java/com/google/android/exoplayer/chunk/parser/webm/WebmExtractor.java

@@ -21,19 +21,18 @@ import com.google.android.exoplayer.ParserException;
 import com.google.android.exoplayer.SampleHolder;
 import com.google.android.exoplayer.chunk.parser.Extractor;
 import com.google.android.exoplayer.chunk.parser.SegmentIndex;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.upstream.NonBlockingInputStream;
 import com.google.android.exoplayer.util.LongArray;
 import com.google.android.exoplayer.util.MimeTypes;
 
+import android.annotation.SuppressLint;
 import android.media.MediaCodec;
 import android.media.MediaExtractor;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -111,7 +110,7 @@ public final class WebmExtractor implements Extractor {
 
   private final EbmlReader reader;
   private final byte[] simpleBlockTimecodeAndFlags = new byte[3];
-  private final HashMap<UUID, byte[]> psshInfo = new HashMap<UUID, byte[]>();
+  private DrmInitData.Universal drmInitData;
 
   private SampleHolder sampleHolder;
   private int readResults;
@@ -199,8 +198,8 @@ public final class WebmExtractor implements Extractor {
   }
 
   @Override
-  public Map<UUID, byte[]> getPsshInfo() {
+  public DrmInitData getDrmInitData() {
-    return psshInfo.isEmpty() ? null : psshInfo;
+    return drmInitData;
   }
 
   /* package */ int getElementType(int id) {
@@ -296,8 +295,7 @@ public final class WebmExtractor implements Extractor {
         if (encryptionKeyId == null) {
           throw new ParserException("Encrypted Track found but ContentEncKeyID was not found");
         }
-        // Widevine.
+        drmInitData = new DrmInitData.Universal(MimeTypes.VIDEO_WEBM, encryptionKeyId);
-        psshInfo.put(new UUID(0xEDEF8BA979D64ACEL, 0xA3C827DCD51D21EDL), encryptionKeyId);
         return true;
       case ID_AUDIO:
         isAudioTrack = true;
@@ -427,6 +425,7 @@ public final class WebmExtractor implements Extractor {
     return true;
   }
 
+  @SuppressLint("InlinedApi")
   /* package */ boolean onBinaryElement(
       int id, long elementOffsetBytes, int headerSizeBytes, int contentsSizeBytes,
       NonBlockingInputStream inputStream) throws ParserException {

library/src/main/java/com/google/android/exoplayer/dash/DashChunkSource.java

@@ -39,6 +39,7 @@ import com.google.android.exoplayer.dash.mpd.MediaPresentationDescription;
 import com.google.android.exoplayer.dash.mpd.Period;
 import com.google.android.exoplayer.dash.mpd.RangedUri;
 import com.google.android.exoplayer.dash.mpd.Representation;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.text.webvtt.WebvttParser;
 import com.google.android.exoplayer.upstream.DataSource;
 import com.google.android.exoplayer.upstream.DataSpec;
@@ -54,8 +55,6 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.UUID;
 
 /**
  * An {@link ChunkSource} for DASH streams.
@@ -96,7 +95,7 @@ public class DashChunkSource implements ChunkSource {
   private final ManifestFetcher<MediaPresentationDescription> manifestFetcher;
   private final int adaptationSetIndex;
   private final int[] representationIndices;
-  private final Map<UUID, byte[]> psshInfo;
+  private final DrmInitData drmInitData;
 
   private MediaPresentationDescription currentManifest;
   private boolean finishedCurrentManifest;
@@ -190,7 +189,7 @@ public class DashChunkSource implements ChunkSource {
     this.evaluation = new Evaluation();
     this.headerBuilder = new StringBuilder();
 
-    psshInfo = getPsshInfo(currentManifest, adaptationSetIndex);
+    drmInitData = getDrmInitData(currentManifest, adaptationSetIndex);
     Representation[] representations = getFilteredRepresentations(currentManifest,
         adaptationSetIndex, representationIndices);
     long periodDurationUs = (representations[0].periodDurationMs == TrackRenderer.UNKNOWN_TIME_US)
@@ -407,7 +406,7 @@ public class DashChunkSource implements ChunkSource {
     // Do nothing.
   }
 
-  private boolean mimeTypeIsWebm(String mimeType) {
+  private static boolean mimeTypeIsWebm(String mimeType) {
     return mimeType.startsWith(MimeTypes.VIDEO_WEBM) || mimeType.startsWith(MimeTypes.AUDIO_WEBM);
   }
 
@@ -475,8 +474,8 @@ public class DashChunkSource implements ChunkSource {
           startTimeUs, endTimeUs, nextAbsoluteSegmentNum, null, representationHolder.vttHeader);
     } else {
       return new ContainerMediaChunk(dataSource, dataSpec, representation.format, trigger,
-          startTimeUs, endTimeUs, nextAbsoluteSegmentNum, representationHolder.extractor, psshInfo,
+          startTimeUs, endTimeUs, nextAbsoluteSegmentNum, representationHolder.extractor,
-          false, presentationTimeOffsetUs);
+          drmInitData, false, presentationTimeOffsetUs);
     }
   }
 
@@ -529,19 +528,24 @@ public class DashChunkSource implements ChunkSource {
     }
   }
 
-  private static Map<UUID, byte[]> getPsshInfo(MediaPresentationDescription manifest,
+  private static DrmInitData getDrmInitData(MediaPresentationDescription manifest,
       int adaptationSetIndex) {
     AdaptationSet adaptationSet = manifest.periods.get(0).adaptationSets.get(adaptationSetIndex);
+    String drmInitMimeType = mimeTypeIsWebm(adaptationSet.representations.get(0).format.mimeType)
+        ? MimeTypes.VIDEO_WEBM : MimeTypes.VIDEO_MP4;
     if (adaptationSet.contentProtections.isEmpty()) {
       return null;
     } else {
-      Map<UUID, byte[]> psshInfo = new HashMap<UUID, byte[]>();
+      DrmInitData.Mapped drmInitData = null;
       for (ContentProtection contentProtection : adaptationSet.contentProtections) {
         if (contentProtection.uuid != null && contentProtection.data != null) {
-          psshInfo.put(contentProtection.uuid, contentProtection.data);
+          if (drmInitData == null) {
+            drmInitData = new DrmInitData.Mapped(drmInitMimeType);
+          }
+          drmInitData.put(contentProtection.uuid, contentProtection.data);
         }
       }
-      return psshInfo.isEmpty() ? null : psshInfo;
+      return drmInitData;
     }
   }
 

library/src/main/java/com/google/android/exoplayer/drm/DrmSessionManager.java

@@ -18,9 +18,6 @@ package com.google.android.exoplayer.drm;
 import android.annotation.TargetApi;
 import android.media.MediaCrypto;
 
-import java.util.Map;
-import java.util.UUID;
-
 /**
  * Manages a DRM session.
  */
@@ -36,7 +33,7 @@ public interface DrmSessionManager {
    */
   public static final int STATE_CLOSED = 1;
   /**
-   * The session is being opened (i.e. {@link #open(Map, String)} has been called, but the session
+   * The session is being opened (i.e. {@link #open(DrmInitData)} has been called, but the session
    * is not yet open).
    */
   public static final int STATE_OPENING = 2;
@@ -52,11 +49,9 @@ public interface DrmSessionManager {
   /**
    * Opens the session, possibly asynchronously.
    *
-   * @param drmInitData Initialization data for the drm schemes supported by the media, keyed by
+   * @param drmInitData DRM initialization data.
-   *     scheme UUID.
-   * @param mimeType The mimeType of the media.
    */
-  void open(Map<UUID, byte[]> drmInitData, String mimeType);
+  void open(DrmInitData drmInitData);
 
   /**
    * Closes the session.

library/src/main/java/com/google/android/exoplayer/drm/StreamingDrmSessionManager.java

@@ -31,7 +31,6 @@ import android.os.Looper;
 import android.os.Message;
 
 import java.util.HashMap;
-import java.util.Map;
 import java.util.UUID;
 
 /**
@@ -168,7 +167,7 @@ public class StreamingDrmSessionManager implements DrmSessionManager {
   }
 
   @Override
-  public void open(Map<UUID, byte[]> psshData, String mimeType) {
+  public void open(DrmInitData drmInitData) {
     if (++openCount != 1) {
       return;
     }
@@ -178,8 +177,8 @@ public class StreamingDrmSessionManager implements DrmSessionManager {
       postRequestHandler = new PostRequestHandler(requestHandlerThread.getLooper());
     }
     if (this.schemePsshData == null) {
-      this.mimeType = mimeType;
+      mimeType = drmInitData.mimeType;
-      schemePsshData = psshData.get(uuid);
+      schemePsshData = drmInitData.get(uuid);
       if (schemePsshData == null) {
         onError(new IllegalStateException("Media does not support uuid: " + uuid));
         return;

library/src/main/java/com/google/android/exoplayer/smoothstreaming/SmoothStreamingChunkSource.java

@@ -30,6 +30,7 @@ import com.google.android.exoplayer.chunk.MediaChunk;
 import com.google.android.exoplayer.chunk.parser.Extractor;
 import com.google.android.exoplayer.chunk.parser.mp4.FragmentedMp4Extractor;
 import com.google.android.exoplayer.chunk.parser.mp4.TrackEncryptionBox;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.mp4.Track;
 import com.google.android.exoplayer.smoothstreaming.SmoothStreamingManifest.ProtectionElement;
 import com.google.android.exoplayer.smoothstreaming.SmoothStreamingManifest.StreamElement;
@@ -38,6 +39,7 @@ import com.google.android.exoplayer.upstream.DataSource;
 import com.google.android.exoplayer.upstream.DataSpec;
 import com.google.android.exoplayer.util.CodecSpecificDataUtil;
 import com.google.android.exoplayer.util.ManifestFetcher;
+import com.google.android.exoplayer.util.MimeTypes;
 
 import android.net.Uri;
 import android.os.SystemClock;
@@ -48,8 +50,6 @@ import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
-import java.util.UUID;
 
 /**
  * An {@link ChunkSource} for SmoothStreaming.
@@ -71,7 +71,7 @@ public class SmoothStreamingChunkSource implements ChunkSource {
   private final int maxHeight;
 
   private final SparseArray<FragmentedMp4Extractor> extractors;
-  private final Map<UUID, byte[]> psshInfo;
+  private final DrmInitData drmInitData;
   private final SmoothStreamingFormat[] formats;
 
   private SmoothStreamingManifest currentManifest;
@@ -143,9 +143,11 @@ public class SmoothStreamingChunkSource implements ChunkSource {
       byte[] keyId = getKeyId(protectionElement.data);
       trackEncryptionBoxes = new TrackEncryptionBox[1];
       trackEncryptionBoxes[0] = new TrackEncryptionBox(true, INITIALIZATION_VECTOR_SIZE, keyId);
-      psshInfo = Collections.singletonMap(protectionElement.uuid, protectionElement.data);
+      DrmInitData.Mapped drmInitData = new DrmInitData.Mapped(MimeTypes.VIDEO_MP4);
+      drmInitData.put(protectionElement.uuid, protectionElement.data);
+      this.drmInitData = drmInitData;
     } else {
-      psshInfo = null;
+      drmInitData = null;
     }
 
     int trackCount = trackIndices != null ? trackIndices.length : streamElement.tracks.length;
@@ -299,7 +301,7 @@ public class SmoothStreamingChunkSource implements ChunkSource {
 
     Uri uri = streamElement.buildRequestUri(selectedFormat.trackIndex, chunkIndex);
     Chunk mediaChunk = newMediaChunk(selectedFormat, uri, null,
-        extractors.get(Integer.parseInt(selectedFormat.id)), psshInfo, dataSource,
+        extractors.get(Integer.parseInt(selectedFormat.id)), drmInitData, dataSource,
         currentAbsoluteChunkIndex, isLastChunk, chunkStartTimeUs, nextChunkStartTimeUs, 0);
     out.chunk = mediaChunk;
   }
@@ -365,7 +367,7 @@ public class SmoothStreamingChunkSource implements ChunkSource {
   }
 
   private static MediaChunk newMediaChunk(Format formatInfo, Uri uri, String cacheKey,
-      Extractor extractor, Map<UUID, byte[]> psshInfo, DataSource dataSource, int chunkIndex,
+      Extractor extractor, DrmInitData drmInitData, DataSource dataSource, int chunkIndex,
       boolean isLast, long chunkStartTimeUs, long nextChunkStartTimeUs, int trigger) {
     int nextChunkIndex = isLast ? -1 : chunkIndex + 1;
     long nextStartTimeUs = isLast ? -1 : nextChunkStartTimeUs;
@@ -374,7 +376,7 @@ public class SmoothStreamingChunkSource implements ChunkSource {
     // In SmoothStreaming each chunk contains sample timestamps relative to the start of the chunk.
     // To convert them the absolute timestamps, we need to set sampleOffsetUs to -chunkStartTimeUs.
     return new ContainerMediaChunk(dataSource, dataSpec, formatInfo, trigger, chunkStartTimeUs,
-        nextStartTimeUs, nextChunkIndex, extractor, psshInfo, false, -chunkStartTimeUs);
+        nextStartTimeUs, nextChunkIndex, extractor, drmInitData, false, -chunkStartTimeUs);
   }
 
   private static byte[] getKeyId(byte[] initData) {

library/src/main/java/com/google/android/exoplayer/source/FrameworkSampleExtractor.java

@@ -19,7 +19,9 @@ import com.google.android.exoplayer.MediaFormat;
 import com.google.android.exoplayer.SampleHolder;
 import com.google.android.exoplayer.SampleSource;
 import com.google.android.exoplayer.TrackRenderer;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.util.Assertions;
+import com.google.android.exoplayer.util.MimeTypes;
 import com.google.android.exoplayer.util.Util;
 
 import android.annotation.TargetApi;
@@ -141,8 +143,8 @@ public final class FrameworkSampleExtractor implements SampleExtractor {
   }
 
   @Override
-  public Map<UUID, byte[]> getDrmInitData(int track) {
+  public DrmInitData getDrmInitData(int track) {
-    return Util.SDK_INT >= 18 ? getPsshInfoV18() : null;
+    return Util.SDK_INT >= 18 ? getDrmInitDataV18() : null;
   }
 
   @Override
@@ -176,9 +178,15 @@ public final class FrameworkSampleExtractor implements SampleExtractor {
   }
 
   @TargetApi(18)
-  private Map<UUID, byte[]> getPsshInfoV18() {
+  private DrmInitData getDrmInitDataV18() {
+    // MediaExtractor only supports psshInfo for MP4, so it's ok to hard code the mimeType here.
     Map<UUID, byte[]> psshInfo = mediaExtractor.getPsshInfo();
-    return (psshInfo == null || psshInfo.isEmpty()) ? null : psshInfo;
+    if (psshInfo == null || psshInfo.isEmpty()) {
+      return null;
+    }
+    DrmInitData.Mapped drmInitData = new DrmInitData.Mapped(MimeTypes.VIDEO_MP4);
+    drmInitData.putAll(psshInfo);
+    return drmInitData;
   }
 
 }

library/src/main/java/com/google/android/exoplayer/source/SampleExtractor.java

@@ -19,10 +19,9 @@ import com.google.android.exoplayer.MediaFormat;
 import com.google.android.exoplayer.SampleHolder;
 import com.google.android.exoplayer.SampleSource;
 import com.google.android.exoplayer.TrackRenderer;
+import com.google.android.exoplayer.drm.DrmInitData;
 
 import java.io.IOException;
-import java.util.Map;
-import java.util.UUID;
 
 /**
  * Extractor for reading track metadata and samples stored in tracks.
@@ -79,7 +78,7 @@ public interface SampleExtractor {
   MediaFormat getMediaFormat(int track);
 
   /** Returns the DRM initialization data for {@code track}. */
-  Map<UUID, byte[]> getDrmInitData(int track);
+  DrmInitData getDrmInitData(int track);
 
   /**
    * Reads the next sample in the track at index {@code track} into {@code sampleHolder}, returning

library/src/test/java/com/google/android/exoplayer/chunk/parser/webm/WebmExtractorTest.java

@@ -20,10 +20,12 @@ import com.google.android.exoplayer.MediaFormat;
 import com.google.android.exoplayer.ParserException;
 import com.google.android.exoplayer.SampleHolder;
 import com.google.android.exoplayer.chunk.parser.SegmentIndex;
+import com.google.android.exoplayer.drm.DrmInitData;
 import com.google.android.exoplayer.upstream.ByteArrayNonBlockingInputStream;
 import com.google.android.exoplayer.upstream.NonBlockingInputStream;
 import com.google.android.exoplayer.util.MimeTypes;
 
+import android.annotation.SuppressLint;
 import android.media.MediaCodec;
 import android.media.MediaExtractor;
 import android.test.InstrumentationTestCase;
@@ -31,7 +33,6 @@ import android.test.InstrumentationTestCase;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
-import java.util.Map;
 import java.util.UUID;
 
 public class WebmExtractorTest extends InstrumentationTestCase {
@@ -56,6 +57,7 @@ public class WebmExtractorTest extends InstrumentationTestCase {
   private static final int TEST_VORBIS_BOOKS_SIZE = 4140;
   private static final byte[] TEST_ENCRYPTION_KEY_ID = { 0x00, 0x01, 0x02, 0x03 };
   private static final UUID WIDEVINE_UUID = new UUID(0xEDEF8BA979D64ACEL, 0xA3C827DCD51D21EDL);
+  private static final UUID ZERO_UUID = new UUID(0, 0);
   // First 8 bytes of IV come from the container, last 8 bytes are always initialized to 0.
   private static final byte[] TEST_INITIALIZATION_VECTOR = {
       0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
@@ -110,10 +112,10 @@ public class WebmExtractorTest extends InstrumentationTestCase {
     assertEquals(EXPECTED_INIT_RESULT, extractor.read(testInputStream, sampleHolder));
     assertFormat();
     assertIndex(new IndexPoint(0, 0, TEST_DURATION_US));
-    Map<UUID, byte[]> psshInfo = extractor.getPsshInfo();
+    DrmInitData drmInitData = extractor.getDrmInitData();
-    assertNotNull(psshInfo);
+    assertNotNull(drmInitData);
-    assertTrue(psshInfo.containsKey(WIDEVINE_UUID));
+    android.test.MoreAsserts.assertEquals(TEST_ENCRYPTION_KEY_ID, drmInitData.get(WIDEVINE_UUID));
-    android.test.MoreAsserts.assertEquals(TEST_ENCRYPTION_KEY_ID, psshInfo.get(WIDEVINE_UUID));
+    android.test.MoreAsserts.assertEquals(TEST_ENCRYPTION_KEY_ID, drmInitData.get(ZERO_UUID));
   }
 
   public void testPrepareThreeCuePoints() throws ParserException {
@@ -353,6 +355,7 @@ public class WebmExtractorTest extends InstrumentationTestCase {
     }
   }
 
+  @SuppressLint("InlinedApi")
   private void assertSample(
       MediaSegment mediaSegment, int timeUs, boolean keyframe, boolean invisible,
       boolean encrypted) {
@@ -695,7 +698,7 @@ public class WebmExtractorTest extends InstrumentationTestCase {
 
   }
 
-  /** Used by {@link createTracksElementWithVideo} to create a Track header with Encryption. */
+  /** Used by {@link #createTracksElementWithVideo} to create a Track header with Encryption. */
   private static final class ContentEncodingSettings {
 
     private final int order;