From a2dfc62e0825d9167e295476a1f595914dd4da81 Mon Sep 17 00:00:00 2001
From: vigneshv <vigneshv@google.com>
Date: Wed, 10 Aug 2016 14:06:47 -0700
Subject: [PATCH] MatroskaExtractor: Implement subsample encryption

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=129908648
---
 .../src/main/assets/sample_media.exolist.json |   4 +-
 .../mkv/subsample_encrypted_altref.webm       | Bin 0 -> 360 bytes
 .../subsample_encrypted_altref.webm.0.dump    |  34 ++++++
 .../mkv/subsample_encrypted_noaltref.webm     | Bin 0 -> 345 bytes
 .../subsample_encrypted_noaltref.webm.0.dump  |  34 ++++++
 .../extractor/mkv/MatroskaExtractorTest.java  |  18 +++
 .../extractor/mkv/MatroskaExtractor.java      | 110 ++++++++++--------
 7 files changed, 149 insertions(+), 51 deletions(-)
 create mode 100644 library/src/androidTest/assets/mkv/subsample_encrypted_altref.webm
 create mode 100644 library/src/androidTest/assets/mkv/subsample_encrypted_altref.webm.0.dump
 create mode 100644 library/src/androidTest/assets/mkv/subsample_encrypted_noaltref.webm
 create mode 100644 library/src/androidTest/assets/mkv/subsample_encrypted_noaltref.webm.0.dump

diff --git a/demo/src/main/assets/sample_media.exolist.json b/demo/src/main/assets/sample_media.exolist.json
index 18322ad135..cd442ed345 100644
--- a/demo/src/main/assets/sample_media.exolist.json
+++ b/demo/src/main/assets/sample_media.exolist.json
@@ -220,13 +220,13 @@
       },
       {
         "name": "WV: Secure Subsample (WebM, VP9 without altref)",
-        "uri": "https://storage.googleapis.com/widevine_test/vp9/sintel_1080p_vp9_altref_subsample/sintel_1080p_vp9_noaltref_subsample.mpd",
+        "uri": "https://storage.googleapis.com/widevine_test/vp9/sintel_1080p_vp9_noaltref_subsample/sintel_1080p_vp9_noaltref_subsample.mpd",
         "drm_scheme": "widevine",
         "drm_license_url": "https://widevine-proxy.appspot.com/proxy"
       },
       {
         "name": "WV: Secure Fullsample (WebM, VP9 without altref)",
-        "uri": "https://storage.googleapis.com/widevine_test/vp9/sintel_1080p_vp9_altref_subsample/sintel_1080p_vp9_noaltref_fullsample.mpd",
+        "uri": "https://storage.googleapis.com/widevine_test/vp9/sintel_1080p_vp9_noaltref_fullsample/sintel_1080p_vp9_noaltref_fullsample.mpd",
         "drm_scheme": "widevine",
         "drm_license_url": "https://widevine-proxy.appspot.com/proxy"
       }
diff --git a/library/src/androidTest/assets/mkv/subsample_encrypted_altref.webm b/library/src/androidTest/assets/mkv/subsample_encrypted_altref.webm
new file mode 100644
index 0000000000000000000000000000000000000000..a929e9545d1ed52bd2e0d1f231f0af28f1b436b5
GIT binary patch
literal 360
zcmb1gy}x+AQ(GgW({~{L)X3uWxsk)EsiizMDc7mJk;$pGkx3%BA)S!{1Q;y@eP@@f
z^4;AXyt+lyb7flan#Mc`Pb_3zHjuXz!jpI2Rtn_(f3p{?k6Z#|#U`!m8=Lu^99%kE
z>>C&ud>i_6GLt~g(>2gD(lgXE2zNvm0lDSsy35xa8H<m$Pmg<*aBIS)=0?W0*0A`n
z0L$El4M`4b0}LA(0*o3N1B`*#be{XeMppMn6BwD8Sy<WFIXDFbg@i>!#oV8_xW8y*
ue9*q3i7{hW;|GwxAVI;P`@E51@rp*E*O(chfCb250%8!&2E-gd3<UsMjAAta

literal 0
HcmV?d00001

diff --git a/library/src/androidTest/assets/mkv/subsample_encrypted_altref.webm.0.dump b/library/src/androidTest/assets/mkv/subsample_encrypted_altref.webm.0.dump
new file mode 100644
index 0000000000..1932ab78f7
--- /dev/null
+++ b/library/src/androidTest/assets/mkv/subsample_encrypted_altref.webm.0.dump
@@ -0,0 +1,34 @@
+seekMap:
+  isSeekable = false
+  duration = 1000
+  getPosition(0) = 0
+numberOfTracks = 1
+track 1:
+  format:
+    bitrate = -1
+    id = 1
+    containerMimeType = null
+    sampleMimeType = video/x-vnd.on2.vp9
+    maxInputSize = -1
+    width = 360
+    height = 240
+    frameRate = -1.0
+    rotationDegrees = -1
+    pixelWidthHeightRatio = 1.0
+    channelCount = -1
+    sampleRate = -1
+    pcmEncoding = -1
+    encoderDelay = -1
+    encoderPadding = -1
+    subsampleOffsetUs = 9223372036854775807
+    selectionFlags = 0
+    language = null
+    drmInitData = 1305012705
+    initializationData:
+  sample count = 1
+  sample 0:
+    time = 0
+    flags = 1073741824
+    data = length 39, hash B7FE77F4
+    encryption key = length 16, hash 4CE944CF
+tracksEnded = true
diff --git a/library/src/androidTest/assets/mkv/subsample_encrypted_noaltref.webm b/library/src/androidTest/assets/mkv/subsample_encrypted_noaltref.webm
new file mode 100644
index 0000000000000000000000000000000000000000..d4b357e5fbb196c3d7e6a7dced3d11217753382a
GIT binary patch
literal 345
zcmb1gy}x+AQ(GgW({~{L)X3uWxsk)EsiizMDc7mJk;$pGkx3%BA)S!{1Q<01eP@@f
z^4;AXyt+lyb7flan#Mc`Pb_3zHjuXz!jpI2Rtn_(f3p{?k6Z#|#U`!m8=Lu^99%kE
z>>C&ud>i_6GLt~g(>2gD(lgXE2zNvm0lDSsy35xa8H<m$-}7^ya>u8uxskE0H7q_X
zz%sXCLz2VV0K-Ox0Ha360AnCFo#+0rk=6at1V$!i7FITP4o(3<Az=|wG56;!?k^e{
mAGB|1V$9gp_yOcENKi1yJ#S=KJhKt#HD*Q_0CJe1egyz6zF{!{

literal 0
HcmV?d00001

diff --git a/library/src/androidTest/assets/mkv/subsample_encrypted_noaltref.webm.0.dump b/library/src/androidTest/assets/mkv/subsample_encrypted_noaltref.webm.0.dump
new file mode 100644
index 0000000000..8751c99b20
--- /dev/null
+++ b/library/src/androidTest/assets/mkv/subsample_encrypted_noaltref.webm.0.dump
@@ -0,0 +1,34 @@
+seekMap:
+  isSeekable = false
+  duration = 1000
+  getPosition(0) = 0
+numberOfTracks = 1
+track 1:
+  format:
+    bitrate = -1
+    id = 1
+    containerMimeType = null
+    sampleMimeType = video/x-vnd.on2.vp9
+    maxInputSize = -1
+    width = 360
+    height = 240
+    frameRate = -1.0
+    rotationDegrees = -1
+    pixelWidthHeightRatio = 1.0
+    channelCount = -1
+    sampleRate = -1
+    pcmEncoding = -1
+    encoderDelay = -1
+    encoderPadding = -1
+    subsampleOffsetUs = 9223372036854775807
+    selectionFlags = 0
+    language = null
+    drmInitData = 1305012705
+    initializationData:
+  sample count = 1
+  sample 0:
+    time = 0
+    flags = 1073741824
+    data = length 24, hash E58668B1
+    encryption key = length 16, hash 4CE944CF
+tracksEnded = true
diff --git a/library/src/androidTest/java/com/google/android/exoplayer2/extractor/mkv/MatroskaExtractorTest.java b/library/src/androidTest/java/com/google/android/exoplayer2/extractor/mkv/MatroskaExtractorTest.java
index a11ed20625..48eee69b50 100644
--- a/library/src/androidTest/java/com/google/android/exoplayer2/extractor/mkv/MatroskaExtractorTest.java
+++ b/library/src/androidTest/java/com/google/android/exoplayer2/extractor/mkv/MatroskaExtractorTest.java
@@ -33,4 +33,22 @@ public final class MatroskaExtractorTest extends InstrumentationTestCase {
     }, "mkv/sample.mkv", getInstrumentation());
   }
 
+  public void testWebmSubsampleEncryption() throws Exception {
+    TestUtil.assertOutput(new TestUtil.ExtractorFactory() {
+      @Override
+      public Extractor create() {
+        return new MatroskaExtractor();
+      }
+    }, "mkv/subsample_encrypted_noaltref.webm", getInstrumentation());
+  }
+
+  public void testWebmSubsampleEncryptionWithAltrefFrames() throws Exception {
+    TestUtil.assertOutput(new TestUtil.ExtractorFactory() {
+      @Override
+      public Extractor create() {
+        return new MatroskaExtractor();
+      }
+    }, "mkv/subsample_encrypted_altref.webm", getInstrumentation());
+  }
+
 }
diff --git a/library/src/main/java/com/google/android/exoplayer2/extractor/mkv/MatroskaExtractor.java b/library/src/main/java/com/google/android/exoplayer2/extractor/mkv/MatroskaExtractor.java
index efb76f2bb7..cafe8121a8 100644
--- a/library/src/main/java/com/google/android/exoplayer2/extractor/mkv/MatroskaExtractor.java
+++ b/library/src/main/java/com/google/android/exoplayer2/extractor/mkv/MatroskaExtractor.java
@@ -267,6 +267,11 @@ public final class MatroskaExtractor implements Extractor {
   // Sample reading state.
   private int sampleBytesRead;
   private boolean sampleEncodingHandled;
+  private boolean sampleSignalByteRead;
+  private boolean sampleInitializationVectorRead;
+  private boolean samplePartitionCountRead;
+  private byte sampleSignalByte;
+  private int samplePartitionCount;
   private int sampleCurrentNalBytesRemaining;
   private int sampleBytesWritten;
   private boolean sampleRead;
@@ -858,6 +863,11 @@ public final class MatroskaExtractor implements Extractor {
     sampleBytesWritten = 0;
     sampleCurrentNalBytesRemaining = 0;
     sampleEncodingHandled = false;
+    sampleSignalByteRead = false;
+    samplePartitionCountRead = false;
+    samplePartitionCount = 0;
+    sampleSignalByte = (byte) 0;
+    sampleInitializationVectorRead = false;
     sampleStrippedBytes.reset();
   }
 
@@ -901,40 +911,50 @@ public final class MatroskaExtractor implements Extractor {
         // If the sample is encrypted, read its encryption signal byte and set the IV size.
         // Clear the encrypted flag.
         blockFlags &= ~C.BUFFER_FLAG_ENCRYPTED;
-        input.readFully(scratch.data, 0, 1);
-        sampleBytesRead++;
-        if ((scratch.data[0] & 0x80) == 0x80) {
-          throw new ParserException("Extension bit is set in signal byte");
-        }
-        // Bits 1-5 are reserved and must be 0.
-        if ((scratch.data[0] & 0x7C) != 0) {
-          throw new ParserException("One of the reserved bits is set in signal byte");
-        }
-        if ((scratch.data[0] & 0x01) == 0x01) {
-          boolean hasSubsampleEncryption = false;
-          if ((scratch.data[0] & 0x02) == 0x02) {
-            hasSubsampleEncryption = true;
+        if (!sampleSignalByteRead) {
+          input.readFully(scratch.data, 0, 1);
+          sampleBytesRead++;
+          if ((scratch.data[0] & 0x80) == 0x80) {
+            throw new ParserException("Extension bit is set in signal byte");
           }
-          // Write the signal byte, containing the IV size and the subsample encryption flag.
-          scratch.data[0] = (byte) (ENCRYPTION_IV_SIZE | (hasSubsampleEncryption ? 0x80 : 0x00));
-          scratch.setPosition(0);
-          output.sampleData(scratch, 1);
-          sampleBytesWritten++;
+          sampleSignalByte = scratch.data[0];
+          sampleSignalByteRead = true;
+        }
+        boolean isEncrypted = (sampleSignalByte & 0x01) == 0x01;
+        if (isEncrypted) {
+          boolean hasSubsampleEncryption = (sampleSignalByte & 0x02) == 0x02;
           blockFlags |= C.BUFFER_FLAG_ENCRYPTED;
-          // Read and Write the IV.
-          input.readFully(encryptionInitializationVector.data, 0, ENCRYPTION_IV_SIZE);
-          sampleBytesRead += ENCRYPTION_IV_SIZE;
-          encryptionInitializationVector.setPosition(0);
-          output.sampleData(encryptionInitializationVector, ENCRYPTION_IV_SIZE);
-          sampleBytesWritten += ENCRYPTION_IV_SIZE;
-          if (hasSubsampleEncryption) {
-            input.readFully(scratch.data, 0, 1);
-            sampleBytesRead++;
+          if (!sampleInitializationVectorRead) {
+            input.readFully(encryptionInitializationVector.data, 0, ENCRYPTION_IV_SIZE);
+            sampleBytesRead += ENCRYPTION_IV_SIZE;
+            sampleInitializationVectorRead = true;
+            // Write the signal byte, containing the IV size and the subsample encryption flag.
+            scratch.data[0] = (byte) (ENCRYPTION_IV_SIZE | (hasSubsampleEncryption ? 0x80 : 0x00));
             scratch.setPosition(0);
-            int partitionCount = scratch.readUnsignedByte();
-            // First partition is at offset 0 and is implicit (not included in partitionCount).
-            int partitionOffset = 0;
-            short subsampleCount = (short) Math.ceil((partitionCount + 1) / 2.0);
+            output.sampleData(scratch, 1);
+            sampleBytesWritten++;
+            // Write the IV.
+            encryptionInitializationVector.setPosition(0);
+            output.sampleData(encryptionInitializationVector, ENCRYPTION_IV_SIZE);
+            sampleBytesWritten += ENCRYPTION_IV_SIZE;
+          }
+          if (hasSubsampleEncryption) {
+            if (!samplePartitionCountRead) {
+              input.readFully(scratch.data, 0, 1);
+              sampleBytesRead++;
+              scratch.setPosition(0);
+              samplePartitionCount = scratch.readUnsignedByte();
+              samplePartitionCountRead = true;
+            }
+            int samplePartitionDataSize = samplePartitionCount * 4;
+            if (scratch.limit() < samplePartitionDataSize) {
+              scratch.reset(new byte[samplePartitionDataSize], samplePartitionDataSize);
+            }
+            input.readFully(scratch.data, 0, samplePartitionDataSize);
+            sampleBytesRead += samplePartitionDataSize;
+            scratch.setPosition(0);
+            scratch.setLimit(samplePartitionDataSize);
+            short subsampleCount = (short) (1 + (samplePartitionCount / 2));
             int subsampleDataSize = 2 + 6 * subsampleCount;
             if (encryptionSubsampleDataBuffer == null
                 || encryptionSubsampleDataBuffer.capacity() < subsampleDataSize) {
@@ -942,26 +962,16 @@ public final class MatroskaExtractor implements Extractor {
             }
             encryptionSubsampleDataBuffer.position(0);
             encryptionSubsampleDataBuffer.putShort(subsampleCount);
-            // Loop through the partition offsets and write out the data in the way MediaCodec wants
-            // it (ISO 23001-7 Part 7):
+            // Loop through the partition offsets and write out the data in the way ExoPlayer
+            // wants it (ISO 23001-7 Part 7):
             //   2 bytes - sub sample count.
             //   for each sub sample:
             //     2 bytes - clear data size.
             //     4 bytes - encrypted data size.
-            for (int i = 0; i <= partitionCount; i++) {
+            int partitionOffset = 0;
+            for (int i = 0; i < samplePartitionCount; i++) {
               int previousPartitionOffset = partitionOffset;
-              if (i < partitionCount) {
-                input.readFully(scratch.data, 0, 4);
-                sampleBytesRead += 4;
-                scratch.setPosition(0);
-                partitionOffset = scratch.readUnsignedIntToInt();
-              } else {
-                // Offset for the last parition is the same as the size of the frame's data.
-                partitionOffset = size - sampleBytesRead;
-              }
-              if (partitionOffset < previousPartitionOffset) {
-                throw new ParserException("Parition offsets are not increasing.");
-              }
+              partitionOffset = scratch.readUnsignedIntToInt();
               if ((i % 2) == 0) {
                 encryptionSubsampleDataBuffer.putShort(
                     (short) (partitionOffset - previousPartitionOffset));
@@ -969,9 +979,11 @@ public final class MatroskaExtractor implements Extractor {
                 encryptionSubsampleDataBuffer.putInt(partitionOffset - previousPartitionOffset);
               }
             }
-            // If the loop above ran odd number of times (happens when an altref is present), insert
-            // a fake encrypted section of size 0.
-            if ((partitionCount + 1) % 2 == 1) {
+            int finalPartitionSize = size - sampleBytesRead - partitionOffset;
+            if ((samplePartitionCount % 2) == 1) {
+              encryptionSubsampleDataBuffer.putInt(finalPartitionSize);
+            } else {
+              encryptionSubsampleDataBuffer.putShort((short) finalPartitionSize);
               encryptionSubsampleDataBuffer.putInt(0);
             }
             encryptionSubsampleData.reset(encryptionSubsampleDataBuffer.array(), subsampleDataSize);