Encapsulate the key cache in HlsChunkSource

PiperOrigin-RevId: 234773649

Encapsulate the key cache in HlsChunkSource
PiperOrigin-RevId: 234773649
e502672b · aquilescanta · Andrew Lewis · 0d24098c · e502672b · e502672b
Commit e502672b authored Feb 20, 2019 by aquilescanta Committed by Andrew Lewis Feb 20, 2019
Showing with 36 additions and 26 deletions
library/hls/src/main/java/com/google/android/exoplayer2/source/hls/HlsChunkSource.java
library/hls/src/main/java/com/google/android/exoplayer2/source/hls/HlsMediaChunk.java
--- a/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/HlsChunkSource.java
+++ b/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/HlsChunkSource.java
@@ -36,12 +36,12 @@ import com.google.android.exoplayer2.trackselection.TrackSelection;
 import com.google.android.exoplayer2.upstream.DataSource;
 import com.google.android.exoplayer2.upstream.DataSpec;
 import com.google.android.exoplayer2.upstream.TransferListener;
+import com.google.android.exoplayer2.util.Assertions;
 import com.google.android.exoplayer2.util.TimestampAdjuster;
 import com.google.android.exoplayer2.util.UriUtil;
 import com.google.android.exoplayer2.util.Util;
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -314,13 +314,13 @@ import java.util.Map;
    HlsMediaPlaylist.Segment segment = mediaPlaylist.segments.get(segmentIndexInPlaylist);
    // Check if the segment or its initialization segment are fully encrypted.
-    out.chunk =
+    Uri initSegmentKeyUri = getFullEncryptionKeyUri(mediaPlaylist, segment.initializationSegment);
-        maybeCreateEncryptionChunkFor(
+    out.chunk = maybeCreateEncryptionChunkFor(initSegmentKeyUri, selectedVariantIndex);
-            segment.initializationSegment, mediaPlaylist, selectedVariantIndex);
    if (out.chunk != null) {
      return;
    }
-    out.chunk = maybeCreateEncryptionChunkFor(segment, mediaPlaylist, selectedVariantIndex);
+    Uri mediaSegmentKeyUri = getFullEncryptionKeyUri(mediaPlaylist, segment);
+    out.chunk = maybeCreateEncryptionChunkFor(mediaSegmentKeyUri, selectedVariantIndex);
    if (out.chunk != null) {
      return;
    }
@@ -339,7 +339,8 @@ import java.util.Map;
            isTimestampMaster,
            timestampAdjusterProvider,
            previous,
-            keyCache.asUnmodifiable());
+            /* mediaSegmentKey= */ keyCache.get(mediaSegmentKeyUri),
+            /* initSegmentKey= */ keyCache.get(initSegmentKeyUri));
  }
  /**
@@ -485,12 +486,11 @@ import java.util.Map;
            : (mediaPlaylist.getEndTimeUs() - playlistTracker.getInitialStartTimeUs());
  }
-  private Chunk maybeCreateEncryptionChunkFor(
+  @Nullable
-      @Nullable Segment segment, HlsMediaPlaylist mediaPlaylist, int selectedVariantIndex) {
+  private Chunk maybeCreateEncryptionChunkFor(@Nullable Uri keyUri, int selectedVariantIndex) {
-    if (segment == null || segment.fullSegmentEncryptionKeyUri == null) {
+    if (keyUri == null) {
      return null;
    }
-    Uri keyUri = UriUtil.resolveToUri(mediaPlaylist.baseUri, segment.fullSegmentEncryptionKeyUri);
    if (keyCache.containsKey(keyUri)) {
      // The key is present in the key cache. We re-insert it to prevent it from being evicted by
      // the following key addition. Note that removal of the key is necessary to affect the
@@ -508,6 +508,14 @@ import java.util.Map;
        scratchSpace);
  }
+  @Nullable
+  private static Uri getFullEncryptionKeyUri(HlsMediaPlaylist playlist, @Nullable Segment segment) {
+    if (segment == null || segment.fullSegmentEncryptionKeyUri == null) {
+      return null;
+    }
+    return UriUtil.resolveToUri(playlist.baseUri, segment.fullSegmentEncryptionKeyUri);
+  }
  // Private classes.
  /**
@@ -640,21 +648,27 @@ import java.util.Map;
   */
  private static final class FullSegmentEncryptionKeyCache extends LinkedHashMap<Uri, byte[]> {
-    private final Map<Uri, byte[]> unmodifiableView;
    public FullSegmentEncryptionKeyCache() {
      super(
          /* initialCapacity= */ KEY_CACHE_SIZE * 2, /* loadFactor= */ 1, /* accessOrder= */ false);
-      unmodifiableView = Collections.unmodifiableMap(this);
    }
    @Override
-    protected boolean removeEldestEntry(Map.Entry<Uri, byte[]> entry) {
+    public byte[] get(Object keyUri) {
-      return size() > KEY_CACHE_SIZE;
+      if (keyUri == null) {
+        return null;
+      }
+      return super.get(keyUri);
+    }
+    @Override
+    public byte[] put(Uri keyUri, byte[] key) {
+      return super.put(keyUri, Assertions.checkNotNull(key));
    }
-    public Map<Uri, byte[]> asUnmodifiable() {
+    @Override
-      return unmodifiableView;
+    protected boolean removeEldestEntry(Map.Entry<Uri, byte[]> entry) {
+      return size() > KEY_CACHE_SIZE;
    }
  }
 }
--- a/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/HlsMediaChunk.java
+++ b/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/HlsMediaChunk.java
@@ -39,7 +39,6 @@ import java.io.EOFException;
 import java.io.IOException;
 import java.math.BigInteger;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 /**
@@ -64,7 +63,9 @@ import java.util.concurrent.atomic.AtomicInteger;
   * @param timestampAdjusterProvider The provider from which to obtain the {@link
   *     TimestampAdjuster}.
   * @param previousChunk The {@link HlsMediaChunk} that preceded this one. May be null.
-   * @param keyCache A map from encryption key URI to the corresponding encryption key.
+   * @param mediaSegmentKey The media segment decryption key, if fully encrypted. Null otherwise.
+   * @param initSegmentKey The initialization segment decryption key, if fully encrypted. Null
+   *     otherwise.
   */
  public static HlsMediaChunk createInstance(
      HlsExtractorFactory extractorFactory,
@@ -79,7 +80,8 @@ import java.util.concurrent.atomic.AtomicInteger;
      boolean isMasterTimestampSource,
      TimestampAdjusterProvider timestampAdjusterProvider,
      @Nullable HlsMediaChunk previousChunk,
-      Map<Uri, byte[]> keyCache) {
+      @Nullable byte[] mediaSegmentKey,
+      @Nullable byte[] initSegmentKey) {
    // Media segment.
    HlsMediaPlaylist.Segment mediaSegment = mediaPlaylist.segments.get(segmentIndexInPlaylist);
    DataSpec dataSpec =
@@ -88,9 +90,6 @@ import java.util.concurrent.atomic.AtomicInteger;
            mediaSegment.byterangeOffset,
            mediaSegment.byterangeLength,
            /* key= */ null);
-    byte[] mediaSegmentKey =
-        keyCache.get(
-            UriUtil.resolveToUri(mediaPlaylist.baseUri, mediaSegment.fullSegmentEncryptionKeyUri));
    boolean mediaSegmentEncrypted = mediaSegmentKey != null;
    byte[] mediaSegmentIv =
        mediaSegmentEncrypted ? getEncryptionIvArray(mediaSegment.encryptionIV) : null;
@@ -102,9 +101,6 @@ import java.util.concurrent.atomic.AtomicInteger;
    boolean initSegmentEncrypted = false;
    DataSource initDataSource = null;
    if (initSegment != null) {
-      byte[] initSegmentKey =
-          keyCache.get(
-              UriUtil.resolveToUri(mediaPlaylist.baseUri, initSegment.fullSegmentEncryptionKeyUri));
      initSegmentEncrypted = initSegmentKey != null;
      byte[] initSegmentIv =
          initSegmentEncrypted ? getEncryptionIvArray(initSegment.encryptionIV) : null;