From 12ed18c74cac660433a6ed796fd26e423adfcc19 Mon Sep 17 00:00:00 2001
From: aquilescanta <aquilescanta@google.com>
Date: Thu, 14 Feb 2019 15:42:16 +0000
Subject: [PATCH] Add encryption information to initialization segments

PiperOrigin-RevId: 233953493
---
 .../source/hls/playlist/HlsMediaPlaylist.java | 13 ++++-
 .../hls/playlist/HlsPlaylistParser.java       | 34 +++++++----
 .../playlist/HlsMediaPlaylistParserTest.java  | 57 +++++++++++++++++++
 3 files changed, 90 insertions(+), 14 deletions(-)

diff --git a/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/playlist/HlsMediaPlaylist.java b/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/playlist/HlsMediaPlaylist.java
index 81d4e7a818..5200de048b 100644
--- a/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/playlist/HlsMediaPlaylist.java
+++ b/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/playlist/HlsMediaPlaylist.java
@@ -88,8 +88,15 @@ public final class HlsMediaPlaylist extends HlsPlaylist {
      * @param uri See {@link #url}.
      * @param byterangeOffset See {@link #byterangeOffset}.
      * @param byterangeLength See {@link #byterangeLength}.
+     * @param fullSegmentEncryptionKeyUri See {@link #fullSegmentEncryptionKeyUri}.
+     * @param encryptionIV See {@link #encryptionIV}.
      */
-    public Segment(String uri, long byterangeOffset, long byterangeLength) {
+    public Segment(
+        String uri,
+        long byterangeOffset,
+        long byterangeLength,
+        String fullSegmentEncryptionKeyUri,
+        String encryptionIV) {
       this(
           uri,
           /* initializationSegment= */ null,
@@ -98,8 +105,8 @@ public final class HlsMediaPlaylist extends HlsPlaylist {
           /* relativeDiscontinuitySequence= */ -1,
           /* relativeStartTimeUs= */ C.TIME_UNSET,
           /* drmInitData= */ null,
-          /* fullSegmentEncryptionKeyUri= */ null,
-          /* encryptionIV= */ null,
+          fullSegmentEncryptionKeyUri,
+          encryptionIV,
           byterangeOffset,
           byterangeLength,
           /* hasGapTag= */ false);
diff --git a/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/playlist/HlsPlaylistParser.java b/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/playlist/HlsPlaylistParser.java
index 26e1e1a9a7..58d1597d95 100644
--- a/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/playlist/HlsPlaylistParser.java
+++ b/library/hls/src/main/java/com/google/android/exoplayer2/source/hls/playlist/HlsPlaylistParser.java
@@ -455,8 +455,8 @@ public final class HlsPlaylistParser implements ParsingLoadable.Parser<HlsPlayli
     boolean hasGapTag = false;
 
     DrmInitData playlistProtectionSchemes = null;
-    String encryptionKeyUri = null;
-    String encryptionIV = null;
+    String fullSegmentEncryptionKeyUri = null;
+    String fullSegmentEncryptionIV = null;
     TreeMap<String, SchemeData> currentSchemeDatas = new TreeMap<>();
     String encryptionScheme = null;
     DrmInitData cachedDrmInitData = null;
@@ -489,7 +489,19 @@ public final class HlsPlaylistParser implements ParsingLoadable.Parser<HlsPlayli
             segmentByteRangeOffset = Long.parseLong(splitByteRange[1]);
           }
         }
-        initializationSegment = new Segment(uri, segmentByteRangeOffset, segmentByteRangeLength);
+        if (fullSegmentEncryptionKeyUri != null && fullSegmentEncryptionIV == null) {
+          // See RFC 8216, Section 4.3.2.5.
+          throw new ParserException(
+              "The encryption IV attribute must be present when an initialization segment is "
+                  + "encrypted with METHOD=AES-128.");
+        }
+        initializationSegment =
+            new Segment(
+                uri,
+                segmentByteRangeOffset,
+                segmentByteRangeLength,
+                fullSegmentEncryptionKeyUri,
+                fullSegmentEncryptionIV);
         segmentByteRangeOffset = 0;
         segmentByteRangeLength = C.LENGTH_UNSET;
       } else if (line.startsWith(TAG_TARGET_DURATION)) {
@@ -521,17 +533,17 @@ public final class HlsPlaylistParser implements ParsingLoadable.Parser<HlsPlayli
         String method = parseStringAttr(line, REGEX_METHOD, variableDefinitions);
         String keyFormat =
             parseOptionalStringAttr(line, REGEX_KEYFORMAT, KEYFORMAT_IDENTITY, variableDefinitions);
-        encryptionKeyUri = null;
-        encryptionIV = null;
+        fullSegmentEncryptionKeyUri = null;
+        fullSegmentEncryptionIV = null;
         if (METHOD_NONE.equals(method)) {
           currentSchemeDatas.clear();
           cachedDrmInitData = null;
         } else /* !METHOD_NONE.equals(method) */ {
-          encryptionIV = parseOptionalStringAttr(line, REGEX_IV, variableDefinitions);
+          fullSegmentEncryptionIV = parseOptionalStringAttr(line, REGEX_IV, variableDefinitions);
           if (KEYFORMAT_IDENTITY.equals(keyFormat)) {
             if (METHOD_AES_128.equals(method)) {
               // The segment is fully encrypted using an identity key.
-              encryptionKeyUri = parseStringAttr(line, REGEX_URI, variableDefinitions);
+              fullSegmentEncryptionKeyUri = parseStringAttr(line, REGEX_URI, variableDefinitions);
             } else {
               // Do nothing. Samples are encrypted using an identity key, but this is not supported.
               // Hopefully, a traditional DRM alternative is also provided.
@@ -581,10 +593,10 @@ public final class HlsPlaylistParser implements ParsingLoadable.Parser<HlsPlayli
         hasEndTag = true;
       } else if (!line.startsWith("#")) {
         String segmentEncryptionIV;
-        if (encryptionKeyUri == null) {
+        if (fullSegmentEncryptionKeyUri == null) {
           segmentEncryptionIV = null;
-        } else if (encryptionIV != null) {
-          segmentEncryptionIV = encryptionIV;
+        } else if (fullSegmentEncryptionIV != null) {
+          segmentEncryptionIV = fullSegmentEncryptionIV;
         } else {
           segmentEncryptionIV = Long.toHexString(segmentMediaSequence);
         }
@@ -615,7 +627,7 @@ public final class HlsPlaylistParser implements ParsingLoadable.Parser<HlsPlayli
                 relativeDiscontinuitySequence,
                 segmentStartTimeUs,
                 cachedDrmInitData,
-                encryptionKeyUri,
+                fullSegmentEncryptionKeyUri,
                 segmentEncryptionIV,
                 segmentByteRangeOffset,
                 segmentByteRangeLength,
diff --git a/library/hls/src/test/java/com/google/android/exoplayer2/source/hls/playlist/HlsMediaPlaylistParserTest.java b/library/hls/src/test/java/com/google/android/exoplayer2/source/hls/playlist/HlsMediaPlaylistParserTest.java
index e7bf3c6324..9fc761badb 100644
--- a/library/hls/src/test/java/com/google/android/exoplayer2/source/hls/playlist/HlsMediaPlaylistParserTest.java
+++ b/library/hls/src/test/java/com/google/android/exoplayer2/source/hls/playlist/HlsMediaPlaylistParserTest.java
@@ -16,9 +16,11 @@
 package com.google.android.exoplayer2.source.hls.playlist;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.fail;
 
 import android.net.Uri;
 import com.google.android.exoplayer2.C;
+import com.google.android.exoplayer2.ParserException;
 import com.google.android.exoplayer2.source.hls.playlist.HlsMediaPlaylist.Segment;
 import com.google.android.exoplayer2.util.Util;
 import java.io.ByteArrayInputStream;
@@ -365,6 +367,61 @@ public class HlsMediaPlaylistParserTest {
     assertThat(segments.get(3).initializationSegment.url).isEqualTo("init2.ts");
   }
 
+  @Test
+  public void testEncryptedMapTag() throws IOException {
+    Uri playlistUri = Uri.parse("https://example.com/test3.m3u8");
+    String playlistString =
+        "#EXTM3U\n"
+            + "#EXT-X-VERSION:3\n"
+            + "#EXT-X-TARGETDURATION:5\n"
+            + "#EXT-X-MEDIA-SEQUENCE:10\n"
+            + "#EXT-X-KEY:METHOD=AES-128,"
+            + "URI=\"https://priv.example.com/key.php?r=2680\",IV=0x1566B\n"
+            + "#EXT-X-MAP:URI=\"init1.ts\""
+            + "#EXTINF:5.005,\n"
+            + "02/00/32.ts\n"
+            + "#EXT-X-KEY:METHOD=NONE\n"
+            + "#EXT-X-MAP:URI=\"init2.ts\""
+            + "#EXTINF:5.005,\n"
+            + "02/00/47.ts\n";
+    InputStream inputStream = new ByteArrayInputStream(Util.getUtf8Bytes(playlistString));
+    HlsMediaPlaylist playlist =
+        (HlsMediaPlaylist) new HlsPlaylistParser().parse(playlistUri, inputStream);
+
+    List<Segment> segments = playlist.segments;
+    Segment initSegment1 = segments.get(0).initializationSegment;
+    assertThat(initSegment1.fullSegmentEncryptionKeyUri)
+        .isEqualTo("https://priv.example.com/key.php?r=2680");
+    assertThat(initSegment1.encryptionIV).isEqualTo("0x1566B");
+    Segment initSegment2 = segments.get(1).initializationSegment;
+    assertThat(initSegment2.fullSegmentEncryptionKeyUri).isNull();
+    assertThat(initSegment2.encryptionIV).isNull();
+  }
+
+  @Test
+  public void testEncryptedMapTagWithNoIvFails() throws IOException {
+    Uri playlistUri = Uri.parse("https://example.com/test3.m3u8");
+    String playlistString =
+        "#EXTM3U\n"
+            + "#EXT-X-VERSION:3\n"
+            + "#EXT-X-TARGETDURATION:5\n"
+            + "#EXT-X-MEDIA-SEQUENCE:10\n"
+            + "#EXT-X-KEY:METHOD=AES-128,"
+            + "URI=\"https://priv.example.com/key.php?r=2680\"\n"
+            + "#EXT-X-MAP:URI=\"init1.ts\""
+            + "#EXTINF:5.005,\n"
+            + "02/00/32.ts\n";
+    InputStream inputStream = new ByteArrayInputStream(Util.getUtf8Bytes(playlistString));
+
+    try {
+      new HlsPlaylistParser().parse(playlistUri, inputStream);
+      fail();
+    } catch (ParserException e) {
+      // Expected because the initialization segment does not have a defined initialization vector,
+      // although it is affected by an EXT-X-KEY tag.
+    }
+  }
+
   @Test
   public void testMasterPlaylistAttributeInheritance() throws IOException {
     Uri playlistUri = Uri.parse("https://example.com/test3.m3u8");