MatroskaExtractor: Implement subsample encryption

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=129908648

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"
       }

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

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

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());
+  }
+
 }

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);
+        if (!sampleSignalByteRead) {
-        sampleBytesRead++;
+          input.readFully(scratch.data, 0, 1);
-        if ((scratch.data[0] & 0x80) == 0x80) {
+          sampleBytesRead++;
-          throw new ParserException("Extension bit is set in signal byte");
+          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;
           }
-          // Write the signal byte, containing the IV size and the subsample encryption flag.
+          sampleSignalByte = scratch.data[0];
-          scratch.data[0] = (byte) (ENCRYPTION_IV_SIZE | (hasSubsampleEncryption ? 0x80 : 0x00));
+          sampleSignalByteRead = true;
-          scratch.setPosition(0);
+        }
-          output.sampleData(scratch, 1);
+        boolean isEncrypted = (sampleSignalByte & 0x01) == 0x01;
-          sampleBytesWritten++;
+        if (isEncrypted) {
+          boolean hasSubsampleEncryption = (sampleSignalByte & 0x02) == 0x02;
           blockFlags |= C.BUFFER_FLAG_ENCRYPTED;
-          // Read and Write the IV.
+          if (!sampleInitializationVectorRead) {
-          input.readFully(encryptionInitializationVector.data, 0, ENCRYPTION_IV_SIZE);
+            input.readFully(encryptionInitializationVector.data, 0, ENCRYPTION_IV_SIZE);
-          sampleBytesRead += ENCRYPTION_IV_SIZE;
+            sampleBytesRead += ENCRYPTION_IV_SIZE;
-          encryptionInitializationVector.setPosition(0);
+            sampleInitializationVectorRead = true;
-          output.sampleData(encryptionInitializationVector, ENCRYPTION_IV_SIZE);
+            // Write the signal byte, containing the IV size and the subsample encryption flag.
-          sampleBytesWritten += ENCRYPTION_IV_SIZE;
+            scratch.data[0] = (byte) (ENCRYPTION_IV_SIZE | (hasSubsampleEncryption ? 0x80 : 0x00));
-          if (hasSubsampleEncryption) {
-            input.readFully(scratch.data, 0, 1);
-            sampleBytesRead++;
             scratch.setPosition(0);
-            int partitionCount = scratch.readUnsignedByte();
+            output.sampleData(scratch, 1);
-            // First partition is at offset 0 and is implicit (not included in partitionCount).
+            sampleBytesWritten++;
-            int partitionOffset = 0;
+            // Write the IV.
-            short subsampleCount = (short) Math.ceil((partitionCount + 1) / 2.0);
+            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
+            // Loop through the partition offsets and write out the data in the way ExoPlayer
-            // it (ISO 23001-7 Part 7):
+            // 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) {
+              partitionOffset = scratch.readUnsignedIntToInt();
-                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.");
-              }
               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
+            int finalPartitionSize = size - sampleBytesRead - partitionOffset;
-            // a fake encrypted section of size 0.
+            if ((samplePartitionCount % 2) == 1) {
-            if ((partitionCount + 1) % 2 == 1) {
+              encryptionSubsampleDataBuffer.putInt(finalPartitionSize);
+            } else {
+              encryptionSubsampleDataBuffer.putShort((short) finalPartitionSize);
               encryptionSubsampleDataBuffer.putInt(0);
             }
             encryptionSubsampleData.reset(encryptionSubsampleDataBuffer.array(), subsampleDataSize);