Update EncryptionHandler to support key rotation and SampleAes

This CL also removes EncryptionConfig stream data type and merges it
into StreamInfo/SegmentInfo instead.

Change-Id: Idb70ce503e61d3c951225cc78b6b15c084e16dcd

packager/media/base/encryption_config.h

@@ -0,0 +1,29 @@
+// Copyright 2017 Google Inc. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+#ifndef PACKAGER_MEDIA_BASE_ENCRYPTION_CONFIG_H_
+#define PACKAGER_MEDIA_BASE_ENCRYPTION_CONFIG_H_
+
+#include "packager/media/base/fourccs.h"
+#include "packager/media/base/protection_system_specific_info.h"
+
+namespace shaka {
+namespace media {
+
+struct EncryptionConfig {
+  FourCC protection_scheme = FOURCC_cenc;
+  uint8_t crypt_byte_block = 0;
+  uint8_t skip_byte_block = 0;
+  uint8_t per_sample_iv_size = 0;
+  std::vector<uint8_t> constant_iv;
+  std::vector<uint8_t> key_id;
+  std::vector<ProtectionSystemSpecificInfo> key_system_info;
+};
+
+}  // namespace media
+}  // namespace shaka
+
+#endif  // PACKAGER_MEDIA_BASE_ENCRYPTION_CONFIG_H_

packager/media/base/fourccs.h

@@ -25,6 +25,8 @@ enum FourCC : uint32_t {
   FOURCC_avcC = 0x61766343,
   FOURCC_bloc = 0x626C6F63,
   FOURCC_cbc1 = 0x63626331,
+  // This is a fake protection scheme fourcc code to indicate Apple Sample AES.
+  FOURCC_cbca = 0x63626361,
   FOURCC_cbcs = 0x63626373,
   FOURCC_cenc = 0x63656e63,
   FOURCC_cens = 0x63656e73,
@@ -142,6 +144,8 @@ enum FourCC : uint32_t {
   FOURCC_zacp = 0x7A616370,
 };
 
+const FourCC kAppleSampleAesProtectionScheme = FOURCC_cbca;
+
 const inline std::string FourCCToString(FourCC fourcc) {
   char buf[5];
   buf[0] = (fourcc >> 24) & 0xff;

packager/media/base/media_base.gyp

@@ -43,6 +43,7 @@
         'decrypt_config.h',
         'decryptor_source.cc',
         'decryptor_source.h',
+        'encryption_config.h',
         'fixed_key_source.cc',
         'fixed_key_source.h',
         'fourccs.h',

packager/media/base/media_handler.h

@@ -22,7 +22,6 @@ enum class StreamDataType {
   kUnknown,
   kPeriodInfo,
   kStreamInfo,
-  kEncryptionConfig,
   kMediaSample,
   kMediaEvent,
   kSegmentInfo,
@@ -30,13 +29,16 @@ enum class StreamDataType {
 
 // TODO(kqyang): Define these structures.
 struct PeriodInfo {};
-struct EncryptionConfig {};
 struct MediaEvent {};
 struct SegmentInfo {
   bool is_subsegment = false;
   bool is_encrypted = false;
   int64_t start_timestamp = -1;
   int64_t duration = 0;
+  // This is only available if key rotation is enabled. Note that we may have
+  // a |key_rotation_encryption_config| even if the segment is not encrypted,
+  // which is the case for clear lead.
+  std::shared_ptr<EncryptionConfig> key_rotation_encryption_config;
 };
 
 // TODO(kqyang): Should we use protobuf?
@@ -46,7 +48,6 @@ struct StreamData {
 
   std::shared_ptr<PeriodInfo> period_info;
   std::shared_ptr<StreamInfo> stream_info;
-  std::shared_ptr<EncryptionConfig> encryption_config;
   std::shared_ptr<MediaSample> media_sample;
   std::shared_ptr<MediaEvent> media_event;
   std::shared_ptr<SegmentInfo> segment_info;
@@ -129,17 +130,6 @@ class MediaHandler {
     return Dispatch(std::move(stream_data));
   }
 
-  /// Dispatch the encryption config to downstream handlers.
-  Status DispatchEncryptionConfig(
-      size_t stream_index,
-      std::unique_ptr<EncryptionConfig> encryption_config) {
-    std::unique_ptr<StreamData> stream_data(new StreamData);
-    stream_data->stream_index = stream_index;
-    stream_data->stream_data_type = StreamDataType::kEncryptionConfig;
-    stream_data->encryption_config = std::move(encryption_config);
-    return Dispatch(std::move(stream_data));
-  }
-
   /// Dispatch the media sample to downstream handlers.
   Status DispatchMediaSample(size_t stream_index,
                              std::shared_ptr<MediaSample> media_sample) {

packager/media/base/media_handler_test_base.cc

@@ -111,6 +111,21 @@ std::unique_ptr<StreamData> MediaHandlerTestBase::GetMediaSampleStreamData(
   return stream_data;
 }
 
+std::unique_ptr<StreamData> MediaHandlerTestBase::GetSegmentInfoStreamData(
+    int stream_index,
+    int64_t start_timestamp,
+    int64_t duration,
+    bool is_subsegment) {
+  std::unique_ptr<StreamData> stream_data(new StreamData);
+  stream_data->stream_index = stream_index;
+  stream_data->stream_data_type = StreamDataType::kSegmentInfo;
+  stream_data->segment_info.reset(new SegmentInfo);
+  stream_data->segment_info->start_timestamp = start_timestamp;
+  stream_data->segment_info->duration = duration;
+  stream_data->segment_info->is_subsegment = is_subsegment;
+  return stream_data;
+}
+
 void MediaHandlerTestBase::SetUpGraph(int num_inputs,
                                       int num_outputs,
                                       std::shared_ptr<MediaHandler> handler) {

packager/media/base/media_handler_test_base.h

@@ -7,6 +7,7 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
+#include "packager/base/strings/string_number_conversions.h"
 #include "packager/media/base/media_handler.h"
 
 namespace shaka {
@@ -15,62 +16,107 @@ namespace media {
 class FakeMediaHandler;
 
 MATCHER_P3(IsStreamInfo, stream_index, time_scale, encrypted, "") {
+  *result_listener << "which is (" << stream_index << "," << time_scale << ","
+                   << (encrypted ? "encrypted" : "not encrypted") << ")";
   return arg->stream_index == stream_index &&
          arg->stream_data_type == StreamDataType::kStreamInfo &&
          arg->stream_info->time_scale() == time_scale &&
          arg->stream_info->is_encrypted() == encrypted;
 }
 
-MATCHER_P4(IsSegmentInfo, stream_index, timestamp, duration, subsegment, "") {
+MATCHER_P5(IsSegmentInfo,
+           stream_index,
+           start_timestamp,
+           duration,
+           subsegment,
+           encrypted,
+           "") {
+  *result_listener << "which is (" << stream_index << "," << start_timestamp
+                   << "," << duration << ","
+                   << (subsegment ? "subsegment" : "not subsegment") << ","
+                   << (encrypted ? "encrypted" : "not encrypted") << ")";
   return arg->stream_index == stream_index &&
          arg->stream_data_type == StreamDataType::kSegmentInfo &&
-         arg->segment_info->start_timestamp == timestamp &&
+         arg->segment_info->start_timestamp == start_timestamp &&
          arg->segment_info->duration == duration &&
-         arg->segment_info->is_subsegment == subsegment;
+         arg->segment_info->is_subsegment == subsegment &&
+         arg->segment_info->is_encrypted == encrypted;
 }
 
-MATCHER_P3(IsMediaSample, stream_index, timestamp, duration, "") {
+MATCHER_P6(MatchEncryptionConfig,
+           protection_scheme,
+           crypt_byte_block,
+           skip_byte_block,
+           per_sample_iv_size,
+           constant_iv,
+           key_id,
+           "") {
+  *result_listener << "which is (" << FourCCToString(protection_scheme) << ","
+                   << static_cast<int>(crypt_byte_block) << ","
+                   << static_cast<int>(skip_byte_block) << ","
+                   << static_cast<int>(per_sample_iv_size) << ","
+                   << base::HexEncode(constant_iv.data(), constant_iv.size())
+                   << "," << base::HexEncode(key_id.data(), key_id.size())
+                   << ")";
+  return arg.protection_scheme == protection_scheme &&
+         arg.crypt_byte_block == crypt_byte_block &&
+         arg.skip_byte_block == skip_byte_block &&
+         arg.per_sample_iv_size == per_sample_iv_size &&
+         arg.constant_iv == constant_iv && arg.key_id == key_id;
+}
+
+MATCHER_P4(IsMediaSample, stream_index, timestamp, duration, encrypted, "") {
+  *result_listener << "which is (" << stream_index << "," << timestamp << ","
+                   << duration << ","
+                   << (encrypted ? "encrypted" : "not encrypted") << ")";
   return arg->stream_index == stream_index &&
          arg->stream_data_type == StreamDataType::kMediaSample &&
          arg->media_sample->dts() == timestamp &&
-         arg->media_sample->duration() == duration;
+         arg->media_sample->duration() == duration &&
+         arg->media_sample->is_encrypted() == encrypted;
 }
 
 class MediaHandlerTestBase : public ::testing::Test {
  public:
   MediaHandlerTestBase();
 
-  /// Return a stream data with mock stream info.
+  /// @return a stream data with mock stream info.
   std::unique_ptr<StreamData> GetStreamInfoStreamData(int stream_index,
                                                       Codec codec,
                                                       uint32_t time_scale);
 
-  /// Return a stream data with mock video stream info.
+  /// @return a stream data with mock video stream info.
   std::unique_ptr<StreamData> GetVideoStreamInfoStreamData(
       int stream_index,
       uint32_t time_scale) {
     return GetStreamInfoStreamData(stream_index, kCodecVP9, time_scale);
   }
 
-  /// Return a stream data with mock audio stream info.
+  /// @return a stream data with mock audio stream info.
   std::unique_ptr<StreamData> GetAudioStreamInfoStreamData(
       int stream_index,
       uint32_t time_scale) {
     return GetStreamInfoStreamData(stream_index, kCodecAAC, time_scale);
   }
 
-  /// Return a stream data with mock media sample.
+  /// @return a stream data with mock media sample.
   std::unique_ptr<StreamData> GetMediaSampleStreamData(int stream_index,
                                                        int64_t timestamp,
                                                        int64_t duration,
                                                        bool is_keyframe);
 
+  /// @return a stream data with mock segment info.
+  std::unique_ptr<StreamData> GetSegmentInfoStreamData(int stream_index,
+                                                       int64_t start_timestamp,
+                                                       int64_t duration,
+                                                       bool is_subsegment);
+
   /// Setup a graph using |handler| with |num_inputs| and |num_outputs|.
   void SetUpGraph(int num_inputs,
                   int num_outputs,
                   std::shared_ptr<MediaHandler> handler);
 
-  /// Return the output stream data vector from handler.
+  /// @return the output stream data vector from handler.
   const std::vector<std::unique_ptr<StreamData>>& GetOutputStreamDataVector()
       const;
 

packager/media/base/stream_info.h

@@ -10,6 +10,8 @@
 #include <string>
 #include <vector>
 
+#include "packager/media/base/encryption_config.h"
+
 namespace shaka {
 namespace media {
 
@@ -80,6 +82,9 @@ class StreamInfo {
   const std::vector<uint8_t>& codec_config() const { return codec_config_; }
   const std::string& language() const { return language_; }
   bool is_encrypted() const { return is_encrypted_; }
+  const EncryptionConfig& encryption_config() const {
+    return encryption_config_;
+  }
 
   void set_duration(int duration) { duration_ = duration; }
   void set_codec(Codec codec) { codec_ = codec; }
@@ -89,6 +94,9 @@ class StreamInfo {
   }
   void set_language(const std::string& language) { language_ = language; }
   void set_is_encrypted(bool is_encrypted) { is_encrypted_ = is_encrypted; }
+  void set_encryption_config(const EncryptionConfig& encryption_config) {
+    encryption_config_ = encryption_config;
+  }
 
  private:
   // Whether the stream is Audio or Video.
@@ -105,6 +113,7 @@ class StreamInfo {
   // Note that in a potentially encrypted stream, individual buffers
   // can be encrypted or not encrypted.
   bool is_encrypted_;
+  EncryptionConfig encryption_config_;
   // Optional byte data required for some audio/video decoders such as Vorbis
   // codebooks.
   std::vector<uint8_t> codec_config_;

packager/media/base/video_stream_info.h

@@ -65,9 +65,9 @@ class VideoStreamInfo : public StreamInfo {
   uint32_t pixel_height_;
   int16_t trick_play_rate_;  // Non-zero for trick-play streams.
 
-  // Specifies the normalized size of the NAL unit length field. Can be 1, 2 or
+  // Specifies the size of the NAL unit length field. Can be 1, 2 or 4 bytes, or
-  // 4 bytes, or 0 if the size if unknown or the stream is not a AVC stream
+  // 0 if the stream is not a NAL structured video stream or if it is an AnnexB
-  // (H.264).
+  // byte stream.
   uint8_t nalu_length_size_;
 
   // Container-specific data used by CDM to generate a license request:

packager/media/chunking/chunking_handler_unittest.cc

@@ -67,23 +67,24 @@ TEST_F(ChunkingHandlerTest, AudioNoSubsegmentsThenFlush) {
                                                kDuration1, kKeyFrame)));
     // One output stream_data except when i == 3, which also has SegmentInfo.
     if (i == 3) {
-      EXPECT_THAT(
+      EXPECT_THAT(GetOutputStreamDataVector(),
-          GetOutputStreamDataVector(),
+                  ElementsAre(IsSegmentInfo(kStreamIndex0, 0, kDuration1 * 3,
-          ElementsAre(
+                                            !kIsSubsegment, !kEncrypted),
-              IsSegmentInfo(kStreamIndex0, 0, kDuration1 * 3, !kIsSubsegment),
+                              IsMediaSample(kStreamIndex0, i * kDuration1,
-              IsMediaSample(kStreamIndex0, i * kDuration1, kDuration1)));
+                                            kDuration1, !kEncrypted)));
     } else {
       EXPECT_THAT(GetOutputStreamDataVector(),
                   ElementsAre(IsMediaSample(kStreamIndex0, i * kDuration1,
-                                            kDuration1)));
+                                            kDuration1, !kEncrypted)));
     }
   }
 
   ClearOutputStreamDataVector();
   ASSERT_OK(OnFlushRequest(kStreamIndex0));
-  EXPECT_THAT(GetOutputStreamDataVector(),
+  EXPECT_THAT(
-              ElementsAre(IsSegmentInfo(kStreamIndex0, kDuration1 * 3,
+      GetOutputStreamDataVector(),
-                                        kDuration1 * 2, !kIsSubsegment)));
+      ElementsAre(IsSegmentInfo(kStreamIndex0, kDuration1 * 3, kDuration1 * 2,
+                                !kIsSubsegment, !kEncrypted)));
 }
 
 TEST_F(ChunkingHandlerTest, AudioWithSubsegments) {
@@ -101,13 +102,16 @@ TEST_F(ChunkingHandlerTest, AudioWithSubsegments) {
       GetOutputStreamDataVector(),
       ElementsAre(
           IsStreamInfo(kStreamIndex0, kTimeScale0, !kEncrypted),
-          IsMediaSample(kStreamIndex0, 0, kDuration1),
+          IsMediaSample(kStreamIndex0, 0, kDuration1, !kEncrypted),
-          IsMediaSample(kStreamIndex0, kDuration1, kDuration1),
+          IsMediaSample(kStreamIndex0, kDuration1, kDuration1, !kEncrypted),
-          IsSegmentInfo(kStreamIndex0, 0, kDuration1 * 2, kIsSubsegment),
+          IsSegmentInfo(kStreamIndex0, 0, kDuration1 * 2, kIsSubsegment,
-          IsMediaSample(kStreamIndex0, 2 * kDuration1, kDuration1),
+                        !kEncrypted),
-          IsSegmentInfo(kStreamIndex0, 0, kDuration1 * 3, !kIsSubsegment),
+          IsMediaSample(kStreamIndex0, 2 * kDuration1, kDuration1, !kEncrypted),
-          IsMediaSample(kStreamIndex0, 3 * kDuration1, kDuration1),
+          IsSegmentInfo(kStreamIndex0, 0, kDuration1 * 3, !kIsSubsegment,
-          IsMediaSample(kStreamIndex0, 4 * kDuration1, kDuration1)));
+                        !kEncrypted),
+          IsMediaSample(kStreamIndex0, 3 * kDuration1, kDuration1, !kEncrypted),
+          IsMediaSample(kStreamIndex0, 4 * kDuration1, kDuration1,
+                        !kEncrypted)));
 }
 
 TEST_F(ChunkingHandlerTest, VideoAndSubsegmentAndNonzeroStart) {
@@ -131,22 +135,22 @@ TEST_F(ChunkingHandlerTest, VideoAndSubsegmentAndNonzeroStart) {
           IsStreamInfo(kStreamIndex0, kTimeScale1, !kEncrypted),
           // The first samples @ kStartTimestamp is discarded - not key frame.
           IsMediaSample(kStreamIndex0, kVideoStartTimestamp + kDuration1,
-                        kDuration1),
+                        kDuration1, !kEncrypted),
           IsMediaSample(kStreamIndex0, kVideoStartTimestamp + kDuration1 * 2,
-                        kDuration1),
+                        kDuration1, !kEncrypted),
           // The next segment boundary 13245 / 1000 != 12645 / 1000.
           IsSegmentInfo(kStreamIndex0, kVideoStartTimestamp + kDuration1,
-                        kDuration1 * 2, !kIsSubsegment),
+                        kDuration1 * 2, !kIsSubsegment, !kEncrypted),
           IsMediaSample(kStreamIndex0, kVideoStartTimestamp + kDuration1 * 3,
-                        kDuration1),
+                        kDuration1, !kEncrypted),
           IsMediaSample(kStreamIndex0, kVideoStartTimestamp + kDuration1 * 4,
-                        kDuration1),
+                        kDuration1, !kEncrypted),
           // The subsegment has duration kDuration1 * 2 since it can only
           // terminate before key frame.
           IsSegmentInfo(kStreamIndex0, kVideoStartTimestamp + kDuration1 * 3,
-                        kDuration1 * 2, kIsSubsegment),
+                        kDuration1 * 2, kIsSubsegment, !kEncrypted),
           IsMediaSample(kStreamIndex0, kVideoStartTimestamp + kDuration1 * 5,
-                        kDuration1)));
+                        kDuration1, !kEncrypted)));
 }
 
 TEST_F(ChunkingHandlerTest, AudioAndVideo) {
@@ -182,56 +186,56 @@ TEST_F(ChunkingHandlerTest, AudioAndVideo) {
       ElementsAre(
           // The first samples @ kStartTimestamp is discarded - not key frame.
           IsMediaSample(kStreamIndex0, kAudioStartTimestamp + kDuration0,
-                        kDuration0),
+                        kDuration0, !kEncrypted),
           IsMediaSample(kStreamIndex1, kVideoStartTimestamp + kDuration1,
-                        kDuration1),
+                        kDuration1, !kEncrypted),
           IsMediaSample(kStreamIndex0, kAudioStartTimestamp + kDuration0 * 2,
-                        kDuration0),
+                        kDuration0, !kEncrypted),
           IsMediaSample(kStreamIndex1, kVideoStartTimestamp + kDuration1 * 2,
-                        kDuration1),
+                        kDuration1, !kEncrypted),
           IsMediaSample(kStreamIndex0, kAudioStartTimestamp + kDuration0 * 3,
-                        kDuration0),
+                        kDuration0, !kEncrypted),
           // The audio segment is terminated together with video stream.
           IsSegmentInfo(kStreamIndex0, kAudioStartTimestamp + kDuration0,
-                        kDuration0 * 3, !kIsSubsegment),
+                        kDuration0 * 3, !kIsSubsegment, !kEncrypted),
           // The next segment boundary 13245 / 1000 != 12645 / 1000.
           IsSegmentInfo(kStreamIndex1, kVideoStartTimestamp + kDuration1,
-                        kDuration1 * 2, !kIsSubsegment),
+                        kDuration1 * 2, !kIsSubsegment, !kEncrypted),
           IsMediaSample(kStreamIndex1, kVideoStartTimestamp + kDuration1 * 3,
-                        kDuration1),
+                        kDuration1, !kEncrypted),
           IsMediaSample(kStreamIndex0, kAudioStartTimestamp + kDuration0 * 4,
-                        kDuration0),
+                        kDuration0, !kEncrypted),
           IsMediaSample(kStreamIndex1, kVideoStartTimestamp + kDuration1 * 4,
-                        kDuration1)));
+                        kDuration1, !kEncrypted)));
   ClearOutputStreamDataVector();
 
   // The side comments below show the equivalent timestamp in video timescale.
   // The audio and video are made ~aligned.
-  ASSERT_OK(Process(GetMediaSampleStreamData(kStreamIndex0,
+  ASSERT_OK(Process(GetMediaSampleStreamData(
-                                             kAudioStartTimestamp + kDuration0 * 5,
+      kStreamIndex0, kAudioStartTimestamp + kDuration0 * 5, kDuration0,
-                                             kDuration0, true)));  // 13595
+      true)));  // 13595
-  ASSERT_OK(Process(GetMediaSampleStreamData(kStreamIndex1,
+  ASSERT_OK(Process(GetMediaSampleStreamData(
-                                             kVideoStartTimestamp + kDuration1 * 5,
+      kStreamIndex1, kVideoStartTimestamp + kDuration1 * 5, kDuration1,
-                                             kDuration1, true)));  // 13845
+      true)));  // 13845
-  ASSERT_OK(Process(GetMediaSampleStreamData(kStreamIndex0,
+  ASSERT_OK(Process(GetMediaSampleStreamData(
-                                             kAudioStartTimestamp + kDuration0 * 6,
+      kStreamIndex0, kAudioStartTimestamp + kDuration0 * 6, kDuration0,
-                                             kDuration0, true)));  // 13845
+      true)));  // 13845
   // This expectation are separated from the expectation above because
   // ElementsAre supports at most 10 elements.
   EXPECT_THAT(
       GetOutputStreamDataVector(),
       ElementsAre(
           IsMediaSample(kStreamIndex0, kAudioStartTimestamp + kDuration0 * 5,
-                        kDuration0),
+                        kDuration0, !kEncrypted),
           // Audio is terminated along with video below.
           IsSegmentInfo(kStreamIndex0, kAudioStartTimestamp + kDuration0 * 4,
-                        kDuration0 * 2, kIsSubsegment),
+                        kDuration0 * 2, kIsSubsegment, !kEncrypted),
           // The subsegment has duration kDuration1 * 2 since it can only
           // terminate before key frame.
           IsSegmentInfo(kStreamIndex1, kVideoStartTimestamp + kDuration1 * 3,
-                        kDuration1 * 2, kIsSubsegment),
+                        kDuration1 * 2, kIsSubsegment, !kEncrypted),
           IsMediaSample(kStreamIndex1, kVideoStartTimestamp + kDuration1 * 5,
-                        kDuration1)));
+                        kDuration1, !kEncrypted)));
 
   ClearOutputStreamDataVector();
   ASSERT_OK(OnFlushRequest(kStreamIndex0));
@@ -239,16 +243,16 @@ TEST_F(ChunkingHandlerTest, AudioAndVideo) {
       GetOutputStreamDataVector(),
       ElementsAre(
           IsMediaSample(kStreamIndex0, kAudioStartTimestamp + kDuration0 * 6,
-                        kDuration0),
+                        kDuration0, !kEncrypted),
           IsSegmentInfo(kStreamIndex0, kAudioStartTimestamp + kDuration0 * 4,
-                        kDuration0 * 3, !kIsSubsegment)));
+                        kDuration0 * 3, !kIsSubsegment, !kEncrypted)));
 
   ClearOutputStreamDataVector();
   ASSERT_OK(OnFlushRequest(kStreamIndex1));
   EXPECT_THAT(GetOutputStreamDataVector(),
-              ElementsAre(IsSegmentInfo(kStreamIndex1,
+              ElementsAre(IsSegmentInfo(
-                                        kVideoStartTimestamp + kDuration1 * 3,
+                  kStreamIndex1, kVideoStartTimestamp + kDuration1 * 3,
-                                        kDuration1 * 3, !kIsSubsegment)));
+                  kDuration1 * 3, !kIsSubsegment, !kEncrypted)));
 
   // Flush again will do nothing.
   ClearOutputStreamDataVector();

packager/media/crypto/encryption_handler.cc

@@ -26,6 +26,11 @@ namespace media {
 namespace {
 const size_t kCencBlockSize = 16u;
 
+// The default KID for key rotation is all 0s.
+const uint8_t kKeyRotationDefaultKeyId[] = {
+    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+};
+
 // Adds one or more subsamples to |*subsamples|.  This may add more than one
 // if one of the values overflows the integer in the subsample.
 void AddSubsample(uint64_t clear_bytes,
@@ -42,13 +47,6 @@ void AddSubsample(uint64_t clear_bytes,
     decrypt_config->AddSubsample(clear_bytes, cipher_bytes);
 }
 
-Codec GetVideoCodec(const StreamInfo& stream_info) {
-  if (stream_info.stream_type() != kStreamVideo) return kUnknownCodec;
-  const VideoStreamInfo& video_stream_info =
-      static_cast<const VideoStreamInfo&>(stream_info);
-  return video_stream_info.codec();
-}
-
 uint8_t GetNaluLengthSize(const StreamInfo& stream_info) {
   if (stream_info.stream_type() != kStreamVideo)
     return 0;
@@ -104,15 +102,21 @@ Status EncryptionHandler::Process(std::unique_ptr<StreamData> stream_data) {
     case StreamDataType::kStreamInfo:
       status = ProcessStreamInfo(stream_data->stream_info.get());
       break;
-    case StreamDataType::kSegmentInfo:
+    case StreamDataType::kSegmentInfo: {
-      if (!stream_data->segment_info->is_subsegment) {
+      SegmentInfo* segment_info = stream_data->segment_info.get();
-        new_segment_ = true;
+      segment_info->is_encrypted = remaining_clear_lead_ <= 0;
+
+      const bool key_rotation_enabled = crypto_period_duration_ != 0;
+      if (key_rotation_enabled)
+        segment_info->key_rotation_encryption_config = encryption_config_;
+      if (!segment_info->is_subsegment) {
+        if (key_rotation_enabled)
+          check_new_crypto_period_ = true;
         if (remaining_clear_lead_ > 0)
-          remaining_clear_lead_ -= stream_data->segment_info->duration;
+          remaining_clear_lead_ -= segment_info->duration;
-        else
-          stream_data->segment_info->is_encrypted = true;
       }
       break;
+    }
     case StreamDataType::kMediaSample:
       status = ProcessMediaSample(stream_data->media_sample.get());
       break;
@@ -135,12 +139,12 @@ Status EncryptionHandler::ProcessStreamInfo(StreamInfo* stream_info) {
   crypto_period_duration_ =
       encryption_options_.crypto_period_duration_in_seconds *
       stream_info->time_scale();
+  codec_ = stream_info->codec();
   nalu_length_size_ = GetNaluLengthSize(*stream_info);
-  video_codec_ = GetVideoCodec(*stream_info);
   track_type_ = GetTrackTypeForEncryption(
       *stream_info, encryption_options_.max_sd_pixels,
       encryption_options_.max_hd_pixels, encryption_options_.max_uhd1_pixels);
-  switch (video_codec_) {
+  switch (codec_) {
     case kCodecVP9:
       vpx_parser_.reset(new VP9Parser);
       break;
@@ -155,41 +159,43 @@ Status EncryptionHandler::ProcessStreamInfo(StreamInfo* stream_info) {
     default:
       // Other codecs should have nalu length size == 0.
       if (nalu_length_size_ > 0) {
-        LOG(WARNING) << "Unknown video codec '" << video_codec_ << "'";
+        LOG(WARNING) << "Unknown video codec '" << codec_ << "'";
         return Status(error::ENCRYPTION_FAILURE, "Unknown video codec.");
       }
   }
-  if (header_parser_ &&
+  if (header_parser_) {
-      !header_parser_->Initialize(stream_info->codec_config())) {
+    CHECK_NE(nalu_length_size_, 0u) << "AnnexB stream is not supported yet";
-    return Status(error::ENCRYPTION_FAILURE, "Fail to read SPS and PPS data.");
+    if (!header_parser_->Initialize(stream_info->codec_config())) {
+      return Status(error::ENCRYPTION_FAILURE,
+                    "Fail to read SPS and PPS data.");
+    }
   }
 
-  // Set up protection pattern.
+  Status status = SetupProtectionPattern(stream_info->stream_type());
-  if (encryption_options_.protection_scheme == FOURCC_cbcs ||
+  if (!status.ok())
-      encryption_options_.protection_scheme == FOURCC_cens) {
+    return status;
-    if (stream_info->stream_type() == kStreamVideo) {
+
-      // Use 1:9 pattern for video.
+  EncryptionKey encryption_key;
-      crypt_byte_block_ = 1u;
+  const bool key_rotation_enabled = crypto_period_duration_ != 0;
-      skip_byte_block_ = 9u;
+  if (key_rotation_enabled) {
+    check_new_crypto_period_ = true;
+    // Setup dummy key id and key to signal encryption for key rotation.
+    encryption_key.key_id.assign(
+        kKeyRotationDefaultKeyId,
+        kKeyRotationDefaultKeyId + sizeof(kKeyRotationDefaultKeyId));
+    // The key is not really used to encrypt any data. It is there just for
+    // convenience.
+    encryption_key.key = encryption_key.key_id;
   } else {
-      // Tracks other than video are protected using whole-block full-sample
+    status = key_source_->GetKey(track_type_, &encryption_key);
-      // encryption, which is essentially a pattern of 1:0. Note that this may
+    if (!status.ok())
-      // not be the same as the non-pattern based encryption counterparts, e.g.
+      return status;
-      // in 'cens' for full sample encryption, the whole sample is encrypted up
-      // to the last 16-byte boundary, see 23001-7:2016(E) 9.7; while in 'cenc'
-      // for full sample encryption, the last partial 16-byte block is also
-      // encrypted, see 23001-7:2016(E) 9.4.2. Another difference is the use of
-      // constant iv.
-      crypt_byte_block_ = 1u;
-      skip_byte_block_ = 0u;
-    }
-  } else {
-    // Not using pattern encryption.
-    crypt_byte_block_ = 0u;
-    skip_byte_block_ = 0u;
   }
+  if (!CreateEncryptor(encryption_key))
+    return Status(error::ENCRYPTION_FAILURE, "Failed to create encryptor");
 
   stream_info->set_is_encrypted(true);
+  stream_info->set_encryption_config(*encryption_config_);
   return Status::OK;
 }
 
@@ -202,59 +208,122 @@ Status EncryptionHandler::ProcessMediaSample(MediaSample* sample) {
       !vpx_parser_->Parse(sample->data(), sample->data_size(), &vpx_frames)) {
     return Status(error::ENCRYPTION_FAILURE, "Failed to parse vpx frame.");
   }
-  if (remaining_clear_lead_ > 0)
-    return Status::OK;
 
-  Status status;
+  // Need to setup the encryptor for new segments even if this segment does not
-  if (new_segment_) {
+  // need to be encrypted, so we can signal encryption metadata earlier to
-    EncryptionKey encryption_key;
+  // allows clients to prefetch the keys.
-    bool create_encryptor = false;
+  if (check_new_crypto_period_) {
-    if (crypto_period_duration_ != 0) {
     const int64_t current_crypto_period_index =
         sample->dts() / crypto_period_duration_;
     if (current_crypto_period_index != prev_crypto_period_index_) {
-        status = key_source_->GetCryptoPeriodKey(current_crypto_period_index,
+      EncryptionKey encryption_key;
-                                                 track_type_, &encryption_key);
+      Status status = key_source_->GetCryptoPeriodKey(
+          current_crypto_period_index, track_type_, &encryption_key);
       if (!status.ok())
         return status;
-        create_encryptor = true;
+      if (!CreateEncryptor(encryption_key))
-      }
-    } else if (!encryptor_) {
-      status = key_source_->GetKey(track_type_, &encryption_key);
-      if (!status.ok())
-        return status;
-      create_encryptor = true;
-    }
-    if (create_encryptor && !CreateEncryptor(&encryption_key))
         return Status(error::ENCRYPTION_FAILURE, "Failed to create encryptor");
-    new_segment_ = false;
+    }
+    check_new_crypto_period_ = false;
   }
 
+  if (remaining_clear_lead_ > 0)
+    return Status::OK;
+
   std::unique_ptr<DecryptConfig> decrypt_config(new DecryptConfig(
-      key_id_, encryptor_->iv(), std::vector<SubsampleEntry>(),
+      encryption_config_->key_id, encryptor_->iv(),
-      encryption_options_.protection_scheme, crypt_byte_block_,
+      std::vector<SubsampleEntry>(), encryption_options_.protection_scheme,
-      skip_byte_block_));
+      crypt_byte_block_, skip_byte_block_));
+  bool result = true;
   if (vpx_parser_) {
-    if (!EncryptVpxFrame(vpx_frames, sample, decrypt_config.get()))
+    result = EncryptVpxFrame(vpx_frames, sample, decrypt_config.get());
-      return Status(error::ENCRYPTION_FAILURE, "Failed to encrypt VPx frames.");
+    if (result) {
-    DCHECK_EQ(decrypt_config->GetTotalSizeOfSubsamples(), sample->data_size());
+      DCHECK_EQ(decrypt_config->GetTotalSizeOfSubsamples(),
-  } else if (nalu_length_size_ > 0) {
+                sample->data_size());
-    if (!EncryptNalFrame(sample, decrypt_config.get())) {
+    }
-      return Status(error::ENCRYPTION_FAILURE,
+  } else if (header_parser_) {
-                    "Failed to encrypt video frames.");
+    result = EncryptNalFrame(sample, decrypt_config.get());
+    if (result) {
+      DCHECK_EQ(decrypt_config->GetTotalSizeOfSubsamples(),
+                sample->data_size());
     }
-    DCHECK_EQ(decrypt_config->GetTotalSizeOfSubsamples(), sample->data_size());
   } else {
     DCHECK_LE(crypt_byte_block_, 1u);
     DCHECK_EQ(skip_byte_block_, 0u);
-    EncryptBytes(sample->writable_data(), sample->data_size());
+    if (sample->data_size() > leading_clear_bytes_size_) {
+      EncryptBytes(sample->writable_data() + leading_clear_bytes_size_,
+                   sample->data_size() - leading_clear_bytes_size_);
     }
+  }
+  if (!result)
+    return Status(error::ENCRYPTION_FAILURE, "Failed to encrypt samples.");
+  sample->set_is_encrypted(true);
   sample->set_decrypt_config(std::move(decrypt_config));
   encryptor_->UpdateIv();
   return Status::OK;
 }
 
-bool EncryptionHandler::CreateEncryptor(EncryptionKey* encryption_key) {
+Status EncryptionHandler::SetupProtectionPattern(StreamType stream_type) {
+  switch (encryption_options_.protection_scheme) {
+    case kAppleSampleAesProtectionScheme: {
+      const size_t kH264LeadingClearBytesSize = 32u;
+      const size_t kSmallNalUnitSize = 32u + 16u;
+      const size_t kAudioLeadingClearBytesSize = 16u;
+      switch (codec_) {
+        case kCodecH264:
+          // Apple Sample AES uses 1:9 pattern for video.
+          crypt_byte_block_ = 1u;
+          skip_byte_block_ = 9u;
+          leading_clear_bytes_size_ = kH264LeadingClearBytesSize;
+          min_protected_data_size_ = kSmallNalUnitSize + 1u;
+          break;
+        case kCodecAAC:
+          FALLTHROUGH_INTENDED;
+        case kCodecAC3:
+          // Audio is whole sample encrypted. We could not use a
+          // crypto_byte_block_ of 1 here as if there is one crypto block
+          // remaining, it need not be encrypted for video but it needs to be
+          // encrypted for audio.
+          crypt_byte_block_ = 0u;
+          skip_byte_block_ = 0u;
+          leading_clear_bytes_size_ = kAudioLeadingClearBytesSize;
+          min_protected_data_size_ = leading_clear_bytes_size_ + 1u;
+          break;
+        default:
+          return Status(error::ENCRYPTION_FAILURE,
+                        "Only AAC/AC3 and H264 are supported in Sample AES.");
+      }
+      break;
+    }
+    case FOURCC_cbcs:
+      FALLTHROUGH_INTENDED;
+    case FOURCC_cens:
+      if (stream_type == kStreamVideo) {
+        // Use 1:9 pattern for video.
+        crypt_byte_block_ = 1u;
+        skip_byte_block_ = 9u;
+      } else {
+        // Tracks other than video are protected using whole-block full-sample
+        // encryption, which is essentially a pattern of 1:0. Note that this may
+        // not be the same as the non-pattern based encryption counterparts,
+        // e.g. in 'cens' for full sample encryption, the whole sample is
+        // encrypted up to the last 16-byte boundary, see 23001-7:2016(E) 9.7;
+        // while in 'cenc' for full sample encryption, the last partial 16-byte
+        // block is also encrypted, see 23001-7:2016(E) 9.4.2. Another
+        // difference is the use of constant iv.
+        crypt_byte_block_ = 1u;
+        skip_byte_block_ = 0u;
+      }
+      break;
+    default:
+      // Not using pattern encryption.
+      crypt_byte_block_ = 0u;
+      skip_byte_block_ = 0u;
+  }
+  return Status::OK;
+}
+
+bool EncryptionHandler::CreateEncryptor(const EncryptionKey& encryption_key) {
   std::unique_ptr<AesCryptor> encryptor;
   switch (encryption_options_.protection_scheme) {
     case FOURCC_cenc:
@@ -277,22 +346,47 @@ bool EncryptionHandler::CreateEncryptor(EncryptionKey* encryption_key) {
           AesCryptor::kUseConstantIv,
           std::unique_ptr<AesCryptor>(new AesCbcEncryptor(kNoPadding))));
       break;
+    case kAppleSampleAesProtectionScheme:
+      if (crypt_byte_block_ == 0 && skip_byte_block_ == 0) {
+        encryptor.reset(
+            new AesCbcEncryptor(kNoPadding, AesCryptor::kUseConstantIv));
+      } else {
+        encryptor.reset(new AesPatternCryptor(
+            crypt_byte_block_, skip_byte_block_,
+            AesPatternCryptor::kSkipIfCryptByteBlockRemaining,
+            AesCryptor::kUseConstantIv,
+            std::unique_ptr<AesCryptor>(new AesCbcEncryptor(kNoPadding))));
+      }
+      break;
     default:
       LOG(ERROR) << "Unsupported protection scheme.";
       return false;
   }
 
-  if (encryption_key->iv.empty()) {
+  std::vector<uint8_t> iv = encryption_key.iv;
+  if (iv.empty()) {
     if (!AesCryptor::GenerateRandomIv(encryption_options_.protection_scheme,
-                                      &encryption_key->iv)) {
+                                      &iv)) {
       LOG(ERROR) << "Failed to generate random iv.";
       return false;
     }
   }
   const bool initialized =
-      encryptor->InitializeWithIv(encryption_key->key, encryption_key->iv);
+      encryptor->InitializeWithIv(encryption_key.key, iv);
   encryptor_ = std::move(encryptor);
-  key_id_ = encryption_key->key_id;
+
+  encryption_config_.reset(new EncryptionConfig);
+  encryption_config_->protection_scheme = encryption_options_.protection_scheme;
+  encryption_config_->crypt_byte_block = crypt_byte_block_;
+  encryption_config_->skip_byte_block = skip_byte_block_;
+  if (encryptor_->use_constant_iv()) {
+    encryption_config_->per_sample_iv_size = 0;
+    encryption_config_->constant_iv = iv;
+  } else {
+    encryption_config_->per_sample_iv_size = iv.size();
+  }
+  encryption_config_->key_id = encryption_key.key_id;
+  encryption_config_->key_system_info = encryption_key.key_system_info;
   return initialized;
 }
 
@@ -301,7 +395,6 @@ bool EncryptionHandler::EncryptVpxFrame(
     MediaSample* sample,
     DecryptConfig* decrypt_config) {
   uint8_t* data = sample->writable_data();
-  const bool is_superframe = vpx_frames.size() > 1;
   for (const VPxFrameInfo& frame : vpx_frames) {
     uint16_t clear_bytes =
         static_cast<uint16_t>(frame.uncompressed_header_size);
@@ -326,6 +419,7 @@ bool EncryptionHandler::EncryptVpxFrame(
     data += frame.frame_size;
   }
   // Add subsample for the superframe index if exists.
+  const bool is_superframe = vpx_frames.size() > 1;
   if (is_superframe) {
     size_t index_size = sample->data() + sample->data_size() - data;
     DCHECK_LE(index_size, 2 + vpx_frames.size() * 4);
@@ -339,8 +433,10 @@ bool EncryptionHandler::EncryptVpxFrame(
 
 bool EncryptionHandler::EncryptNalFrame(MediaSample* sample,
                                         DecryptConfig* decrypt_config) {
+  DCHECK_NE(nalu_length_size_, 0u);
+  DCHECK(header_parser_);
   const Nalu::CodecType nalu_type =
-      (video_codec_ == kCodecHVC1 || video_codec_ == kCodecHEV1) ? Nalu::kH265
+      (codec_ == kCodecHVC1 || codec_ == kCodecHEV1) ? Nalu::kH265
                                                      : Nalu::kH264;
   NaluReader reader(nalu_type, nalu_length_size_, sample->writable_data(),
                     sample->data_size());
@@ -352,20 +448,21 @@ bool EncryptionHandler::EncryptNalFrame(MediaSample* sample,
   Nalu nalu;
   NaluReader::Result result;
   while ((result = reader.Advance(&nalu)) == NaluReader::kOk) {
-    if (nalu.is_video_slice()) {
+    const uint64_t nalu_total_size = nalu.header_size() + nalu.payload_size();
+    if (nalu.is_video_slice() && nalu_total_size >= min_protected_data_size_) {
+      uint64_t current_clear_bytes = leading_clear_bytes_size_;
+      if (current_clear_bytes == 0) {
         // For video-slice NAL units, encrypt the video slice.  This skips
-      // the frame header.  If this is an unrecognized codec, the whole NAL unit
+        // the frame header.
-      // will be encrypted.
         const int64_t video_slice_header_size =
-          header_parser_ ? header_parser_->GetHeaderSize(nalu) : 0;
+            header_parser_->GetHeaderSize(nalu);
         if (video_slice_header_size < 0) {
           LOG(ERROR) << "Failed to read slice header.";
           return false;
         }
-
+        current_clear_bytes = nalu.header_size() + video_slice_header_size;
-      uint64_t current_clear_bytes =
+      }
-          nalu.header_size() + video_slice_header_size;
+      uint64_t cipher_bytes = nalu_total_size - current_clear_bytes;
-      uint64_t cipher_bytes = nalu.payload_size() - video_slice_header_size;
 
       // ISO/IEC 23001-7:2016 10.2 'cbc1' 10.3 'cens'
       // The BytesOfProtectedData size SHALL be a multiple of 16 bytes to
@@ -385,9 +482,8 @@ bool EncryptionHandler::EncryptNalFrame(MediaSample* sample,
           cipher_bytes, decrypt_config);
       accumulated_clear_bytes = 0;
     } else {
-      // For non-video-slice NAL units, don't encrypt.
+      // For non-video-slice or small NAL units, don't encrypt.
-      accumulated_clear_bytes +=
+      accumulated_clear_bytes += nalu_length_size_ + nalu_total_size;
-          nalu_length_size_ + nalu.header_size() + nalu.payload_size();
     }
   }
   if (result != NaluReader::kEOStream) {

packager/media/crypto/encryption_handler.h

@@ -68,7 +68,8 @@ class EncryptionHandler : public MediaHandler {
   // Processes media sample and encrypts it if needed.
   Status ProcessMediaSample(MediaSample* sample);
 
-  bool CreateEncryptor(EncryptionKey* encryption_key);
+  Status SetupProtectionPattern(StreamType stream_type);
+  bool CreateEncryptor(const EncryptionKey& encryption_key);
   bool EncryptVpxFrame(const std::vector<VPxFrameInfo>& vpx_frames,
                        MediaSample* sample,
                        DecryptConfig* decrypt_config);
@@ -83,26 +84,30 @@ class EncryptionHandler : public MediaHandler {
   const EncryptionOptions encryption_options_;
   KeySource* key_source_ = nullptr;
   KeySource::TrackType track_type_ = KeySource::TRACK_TYPE_UNKNOWN;
+  // Current encryption config and encryptor.
+  std::shared_ptr<EncryptionConfig> encryption_config_;
   std::unique_ptr<AesCryptor> encryptor_;
+  Codec codec_ = kUnknownCodec;
   // Specifies the size of NAL unit length in bytes. Can be 1, 2 or 4 bytes. 0
   // if it is not a NAL structured video.
   uint8_t nalu_length_size_ = 0;
-  Codec video_codec_ = kUnknownCodec;
+  // For Sample AES, 32 bytes for Video and 16 bytes for audio.
+  size_t leading_clear_bytes_size_ = 0;
+  // For Sample AES, 48+1 bytes for video NAL and 16+1 bytes for audio.
+  size_t min_protected_data_size_ = 0;
   // Remaining clear lead in the stream's time scale.
   int64_t remaining_clear_lead_ = 0;
   // Crypto period duration in the stream's time scale.
   uint64_t crypto_period_duration_ = 0;
   // Previous crypto period index if key rotation is enabled.
   int64_t prev_crypto_period_index_ = -1;
-  bool new_segment_ = true;
+  bool check_new_crypto_period_ = false;
 
   // Number of encrypted blocks (16-byte-block) in pattern based encryption.
   uint8_t crypt_byte_block_ = 0;
   /// Number of unencrypted blocks (16-byte-block) in pattern based encryption.
   uint8_t skip_byte_block_ = 0;
 
-  // Current key id.
-  std::vector<uint8_t> key_id_;
   // VPx parser for VPx streams.
   std::unique_ptr<VPxParser> vpx_parser_;
   // Video slice header parser for NAL strucutred streams.

packager/media/crypto/encryption_handler_unittest.cc

@@ -25,11 +25,32 @@ using ::testing::_;
 using ::testing::Combine;
 using ::testing::DoAll;
 using ::testing::ElementsAre;
+using ::testing::Mock;
 using ::testing::Return;
 using ::testing::SetArgPointee;
+using ::testing::StrictMock;
 using ::testing::Values;
+using ::testing::ValuesIn;
 using ::testing::WithParamInterface;
 
+const uint8_t kKeyId[]{
+    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+    0x08, 0x09, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
+};
+const uint8_t kKey[]{
+    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+    0x08, 0x09, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
+};
+const uint8_t kIv[]{
+    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+    0x08, 0x09, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
+};
+
+// The default KID for key rotation is all 0s.
+const uint8_t kKeyRotationDefaultKeyId[] = {
+    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+};
+
 class MockKeySource : public FixedKeySource {
  public:
   MOCK_METHOD2(GetKey, Status(TrackType track_type, EncryptionKey* key));
@@ -70,6 +91,14 @@ class EncryptionHandlerTest : public MediaHandlerTestBase {
     return encryption_handler_->Process(std::move(stream_data));
   }
 
+  EncryptionKey GetMockEncryptionKey() {
+    EncryptionKey encryption_key;
+    encryption_key.key_id.assign(kKeyId, kKeyId + sizeof(kKeyId));
+    encryption_key.key.assign(kKey, kKey + sizeof(kKey));
+    encryption_key.iv.assign(kIv, kIv + sizeof(kIv));
+    return encryption_key;
+  }
+
   void InjectVpxParserForTesting(std::unique_ptr<VPxParser> vpx_parser) {
     encryption_handler_->InjectVpxParserForTesting(std::move(vpx_parser));
   }
@@ -82,7 +111,7 @@ class EncryptionHandlerTest : public MediaHandlerTestBase {
 
  protected:
   std::shared_ptr<EncryptionHandler> encryption_handler_;
-  MockKeySource mock_key_source_;
+  StrictMock<MockKeySource> mock_key_source_;
 };
 
 TEST_F(EncryptionHandlerTest, Initialize) {
@@ -104,38 +133,68 @@ TEST_F(EncryptionHandlerTest, OnlyOneInput) {
 
 namespace {
 
-const size_t kStreamIndex = 0;
+const bool kIsKeyFrame = true;
+const bool kIsSubsegment = true;
 const bool kEncrypted = true;
+const size_t kStreamIndex = 0;
 const uint32_t kTimeScale = 1000;
-const uint32_t kMaxSdPixels = 100u;
+const int64_t kSampleDuration = 1000;
-const uint32_t kMaxHdPixels = 200u;
+const int64_t kSegmentDuration = 1000;
-const uint32_t kMaxUhd1Pixels = 300u;
 
 // The data is based on H264. The same data is also used to test audio, which
 // does not care the underlying data, and VP9, for which we will mock the
 // parser.
 const uint8_t kData[]{
     // First NALU
-    0x15, 0x01, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09,
+    0x30, 0x01, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09,
-    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19,
+    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x20, 0x21,
+    0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x30, 0x31, 0x32, 0x33,
+    0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45,
+    0x46,
     // Second NALU
-    0x13, 0x25, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09,
+    0x31, 0x25, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09,
-    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x20, 0x21,
-    // Third NALU
+    0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x30, 0x31, 0x32, 0x33,
+    0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x40, 0x41, 0x42, 0x43, 0x44, 0x45,
+    0x46, 0x47,
+    // Third non-video-slice NALU for H264 or superframe index for VP9.
     0x06, 0x67, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
 };
-const uint8_t kKeyId[]{
+
-    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+// H264 subsample information for the the above data.
-    0x08, 0x09, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
+const size_t kNaluLengthSize = 1u;
-};
+const size_t kNaluHeaderSize = 1u;
-const uint8_t kKey[]{
+const size_t kSubsampleSize1 = 49u;
-    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+const size_t kSliceHeaderSize1 = 1u;
-    0x08, 0x09, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
+const size_t kSubsampleSize2 = 50u;
-};
+const size_t kSliceHeaderSize2 = 16u;
-const uint8_t kIv[]{
+const size_t kSubsampleSize3 = 7u;
-    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+// VP9 frame information for the above data. It should match H264 subsample
-    0x08, 0x09, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15,
+// information.
-};
+const size_t kVpxFrameSize1 = kSubsampleSize1;
+const size_t kUncompressedHeaderSize1 =
+    kNaluLengthSize + kNaluHeaderSize + kSliceHeaderSize1;
+const size_t kVpxFrameSize2 = kSubsampleSize2;
+const size_t kUncompressedHeaderSize2 =
+    kNaluLengthSize + kNaluHeaderSize + kSliceHeaderSize2;
+// Subsample pairs for the above data.
+const size_t kClearSize1 = kUncompressedHeaderSize1;
+const size_t kCipherSize1 = kVpxFrameSize1 - kUncompressedHeaderSize1;
+const size_t kClearSize2 = kUncompressedHeaderSize2;
+const size_t kCipherSize2 = kVpxFrameSize2 - kUncompressedHeaderSize2;
+// Align cipher bytes for some protection schemes.
+const size_t kAesBlockSize = 16u;
+const size_t kAlignedClearSize1 = kClearSize1 + kCipherSize1 % kAesBlockSize;
+static_assert(kAlignedClearSize1 != kClearSize1,
+              "Clearsize 1 should not be aligned");
+const size_t kAlignedCipherSize1 = kCipherSize1 - kCipherSize1 % kAesBlockSize;
+// Apple Sample AES.
+const size_t kVideoLeadingClearBytesSize = 32u + kNaluLengthSize;
+// Subsample 1 is <= 48 bytes, so not encrypted and merged with subsample2.
+const size_t kSampleAesClearSize1 =
+    kSubsampleSize1 + kVideoLeadingClearBytesSize;
+const size_t kSampleAesCipherSize1 =
+    kSubsampleSize2 - kVideoLeadingClearBytesSize;
 
 }  // namespace
 
@@ -151,22 +210,15 @@ class EncryptionHandlerEncryptionTest
   void SetUp() override {
     protection_scheme_ = std::tr1::get<0>(GetParam());
     codec_ = std::tr1::get<1>(GetParam());
-
-    EncryptionOptions encryption_options;
-    encryption_options.protection_scheme = protection_scheme_;;
-    encryption_options.max_sd_pixels = kMaxSdPixels;
-    encryption_options.max_hd_pixels = kMaxHdPixels;
-    encryption_options.max_uhd1_pixels = kMaxUhd1Pixels;
-    SetUpEncryptionHandler(encryption_options);
   }
 
   std::vector<VPxFrameInfo> GetMockVpxFrameInfo() {
     std::vector<VPxFrameInfo> vpx_frames;
     vpx_frames.resize(2);
-    vpx_frames[0].frame_size = 22;
+    vpx_frames[0].frame_size = kVpxFrameSize1;
-    vpx_frames[0].uncompressed_header_size = 3;
+    vpx_frames[0].uncompressed_header_size = kUncompressedHeaderSize1;
-    vpx_frames[1].frame_size = 20;
+    vpx_frames[1].frame_size = kVpxFrameSize2;
-    vpx_frames[1].uncompressed_header_size = 4;
+    vpx_frames[1].uncompressed_header_size = kUncompressedHeaderSize2;
     return vpx_frames;
   }
 
@@ -175,29 +227,70 @@ class EncryptionHandlerEncryptionTest
     std::vector<SubsampleEntry> subsamples;
     if (codec_ == kCodecAAC)
       return subsamples;
+    if (protection_scheme_ == kAppleSampleAesProtectionScheme) {
+      subsamples.emplace_back(kSampleAesClearSize1, kSampleAesCipherSize1);
+      subsamples.emplace_back(kSubsampleSize3, 0u);
+    } else {
       if (codec_ == kCodecVP9 || protection_scheme_ == FOURCC_cbc1 ||
           protection_scheme_ == FOURCC_cens) {
         // Align the encrypted bytes to multiple of 16 bytes.
-      subsamples.emplace_back(6, 16);
+        subsamples.emplace_back(kAlignedClearSize1, kAlignedCipherSize1);
+        // Subsample 2 is already aligned.
       } else {
-      subsamples.emplace_back(3, 19);
+        subsamples.emplace_back(kClearSize1, kCipherSize1);
+      }
+      subsamples.emplace_back(kClearSize2, kCipherSize2);
+      subsamples.emplace_back(kSubsampleSize3, 0u);
     }
-    subsamples.emplace_back(4, 16);
-    subsamples.emplace_back(7, 0);
     return subsamples;
   }
 
-  EncryptionKey GetMockEncryptionKey() {
+  std::unique_ptr<StreamData> GetMediaSampleStreamData(int stream_index,
-    EncryptionKey encryption_key;
+                                                       int64_t timestamp,
-    encryption_key.key_id.assign(kKeyId, kKeyId + sizeof(kKeyId));
+                                                       int64_t duration) {
-    encryption_key.key.assign(kKey, kKey + sizeof(kKey));
+    std::unique_ptr<StreamData> stream_data(new StreamData);
-    encryption_key.iv.assign(kIv, kIv + sizeof(kIv));
+    stream_data->stream_index = stream_index;
-    return encryption_key;
+    stream_data->stream_data_type = StreamDataType::kMediaSample;
+    stream_data->media_sample.reset(
+        new MediaSample(kData, sizeof(kData), nullptr, 0, kIsKeyFrame));
+    stream_data->media_sample->set_dts(timestamp);
+    stream_data->media_sample->set_duration(duration);
+    return stream_data;
+  }
+
+  // Inject vpx parser / video slice header parser if needed.
+  void InjectCodecParser() {
+    switch (codec_) {
+      case kCodecVP9: {
+        std::unique_ptr<MockVpxParser> mock_vpx_parser(new MockVpxParser);
+        EXPECT_CALL(*mock_vpx_parser, Parse(_, sizeof(kData), _))
+            .WillRepeatedly(
+                DoAll(SetArgPointee<2>(GetMockVpxFrameInfo()), Return(true)));
+        InjectVpxParserForTesting(std::move(mock_vpx_parser));
+        break;
+      }
+      case kCodecH264: {
+        std::unique_ptr<MockVideoSliceHeaderParser> mock_header_parser(
+            new MockVideoSliceHeaderParser);
+        if (protection_scheme_ == kAppleSampleAesProtectionScheme) {
+          EXPECT_CALL(*mock_header_parser, GetHeaderSize(_)).Times(0);
+        } else {
+          EXPECT_CALL(*mock_header_parser, GetHeaderSize(_))
+              .WillOnce(Return(kSliceHeaderSize1))
+              .WillOnce(Return(kSliceHeaderSize2));
+        }
+        InjectVideoSliceHeaderParserForTesting(std::move(mock_header_parser));
+        break;
+      }
+      default:
+        break;
+    }
   }
 
   bool Decrypt(const DecryptConfig& decrypt_config,
                uint8_t* data,
                size_t data_size) {
+    size_t leading_clear_bytes_size = 0;
     std::unique_ptr<AesCryptor> aes_decryptor;
     switch (decrypt_config.protection_scheme()) {
       case FOURCC_cenc:
@@ -220,6 +313,24 @@ class EncryptionHandlerEncryptionTest
             AesCryptor::kUseConstantIv,
             std::unique_ptr<AesCryptor>(new AesCbcDecryptor(kNoPadding))));
         break;
+      case kAppleSampleAesProtectionScheme:
+        if (decrypt_config.crypt_byte_block() == 0 &&
+            decrypt_config.skip_byte_block() == 0) {
+          const size_t kAudioLeadingClearBytesSize = 16u;
+          // Only needed for audio; for video, it is already taken into
+          // consideration in subsamples.
+          leading_clear_bytes_size = kAudioLeadingClearBytesSize;
+          aes_decryptor.reset(
+              new AesCbcDecryptor(kNoPadding, AesCryptor::kUseConstantIv));
+        } else {
+          aes_decryptor.reset(new AesPatternCryptor(
+              decrypt_config.crypt_byte_block(),
+              decrypt_config.skip_byte_block(),
+              AesPatternCryptor::kSkipIfCryptByteBlockRemaining,
+              AesCryptor::kUseConstantIv,
+              std::unique_ptr<AesCryptor>(new AesCbcDecryptor(kNoPadding))));
+        }
+        break;
       default:
         LOG(FATAL) << "Not supposed to happen.";
     }
@@ -232,7 +343,9 @@ class EncryptionHandlerEncryptionTest
 
     if (decrypt_config.subsamples().empty()) {
       // Sample not encrypted using subsample encryption. Decrypt whole.
-      if (!aes_decryptor->Crypt(data, data_size, data)) {
+      if (!aes_decryptor->Crypt(data + leading_clear_bytes_size,
+                                data_size - leading_clear_bytes_size,
+                                data + leading_clear_bytes_size)) {
         LOG(ERROR) << "Error during bulk sample decryption.";
         return false;
       }
@@ -261,6 +374,13 @@ class EncryptionHandlerEncryptionTest
   }
 
   uint8_t GetExpectedCryptByteBlock() {
+    if (protection_scheme_ == kAppleSampleAesProtectionScheme) {
+      // Audio is whole sample encrypted. We could not use a
+      // crypto_byte_block_ of 1 for audio as if there is one crypto block
+      // remaining, it need not be encrypted for video but it needs to be
+      // encrypted for audio.
+      return codec_ == kCodecAAC ? 0 : 1;
+    }
     switch (protection_scheme_) {
       case FOURCC_cenc:
       case FOURCC_cbc1:
@@ -283,59 +403,171 @@ class EncryptionHandlerEncryptionTest
         return 0;
       case FOURCC_cens:
       case FOURCC_cbcs:
+      case kAppleSampleAesProtectionScheme:
         return 9;
       default:
         return 0;
     }
   }
 
+  uint8_t GetExpectedPerSampleIvSize() {
+    switch (protection_scheme_) {
+      case FOURCC_cenc:
+      case FOURCC_cens:
+      case FOURCC_cbc1:
+        return sizeof(kIv);
+      case FOURCC_cbcs:
+      case kAppleSampleAesProtectionScheme:
+        return 0;
+      default:
+        return 0;
+    }
+  }
+
+  std::vector<uint8_t> GetExpectedConstantIv() {
+    switch (protection_scheme_) {
+      case FOURCC_cbcs:
+      case kAppleSampleAesProtectionScheme:
+        return std::vector<uint8_t>(std::begin(kIv), std::end(kIv));
+      default:
+        return std::vector<uint8_t>();
+    }
+  }
+
  protected:
   FourCC protection_scheme_;
   Codec codec_;
 };
 
+TEST_P(EncryptionHandlerEncryptionTest, ClearLeadWithNoKeyRotation) {
+  const double kClearLeadInSeconds = 1.5 * kSegmentDuration / kTimeScale;
+  EncryptionOptions encryption_options;
+  encryption_options.protection_scheme = protection_scheme_;
+  encryption_options.clear_lead_in_seconds = kClearLeadInSeconds;
+  SetUpEncryptionHandler(encryption_options);
+
+  const EncryptionKey mock_encryption_key = GetMockEncryptionKey();
+  EXPECT_CALL(mock_key_source_, GetKey(_, _))
+      .WillOnce(
+          DoAll(SetArgPointee<1>(mock_encryption_key), Return(Status::OK)));
+  ASSERT_OK(Process(GetStreamInfoStreamData(kStreamIndex, codec_, kTimeScale)));
+  EXPECT_THAT(GetOutputStreamDataVector(),
+              ElementsAre(IsStreamInfo(kStreamIndex, kTimeScale, kEncrypted)));
+  const EncryptionConfig& encryption_config =
+      GetOutputStreamDataVector().back()->stream_info->encryption_config();
+  EXPECT_THAT(encryption_config,
+              MatchEncryptionConfig(
+                  protection_scheme_, GetExpectedCryptByteBlock(),
+                  GetExpectedSkipByteBlock(), GetExpectedPerSampleIvSize(),
+                  GetExpectedConstantIv(), mock_encryption_key.key_id));
+  ClearOutputStreamDataVector();
+  Mock::VerifyAndClearExpectations(&mock_key_source_);
+
+  InjectCodecParser();
+
+  // There are three segments. Only the third segment is encrypted.
+  for (int i = 0; i < 3; ++i) {
+    // Use single-frame segment for testing.
+    ASSERT_OK(Process(GetMediaSampleStreamData(
+        kStreamIndex, i * kSegmentDuration, kSegmentDuration)));
+    ASSERT_OK(Process(GetSegmentInfoStreamData(
+        kStreamIndex, i * kSegmentDuration, kSegmentDuration, !kIsSubsegment)));
+    const bool is_encrypted = i == 2;
+    const auto& output_stream_data = GetOutputStreamDataVector();
+    EXPECT_THAT(output_stream_data,
+                ElementsAre(IsMediaSample(kStreamIndex, i * kSegmentDuration,
+                                          kSegmentDuration, is_encrypted),
+                            IsSegmentInfo(kStreamIndex, i * kSegmentDuration,
+                                          kSegmentDuration, !kIsSubsegment,
+                                          is_encrypted)));
+    EXPECT_FALSE(output_stream_data.back()
+                     ->segment_info->key_rotation_encryption_config);
+    ClearOutputStreamDataVector();
+  }
+}
+
+TEST_P(EncryptionHandlerEncryptionTest, ClearLeadWithKeyRotation) {
+  const double kClearLeadInSeconds = 1.5 * kSegmentDuration / kTimeScale;
+  const double kCryptoPeriodDurationInSeconds = kSegmentDuration / kTimeScale;
+  EncryptionOptions encryption_options;
+  encryption_options.protection_scheme = protection_scheme_;
+  encryption_options.clear_lead_in_seconds = kClearLeadInSeconds;
+  encryption_options.crypto_period_duration_in_seconds =
+      kCryptoPeriodDurationInSeconds;
+  SetUpEncryptionHandler(encryption_options);
+
+  ASSERT_OK(Process(GetStreamInfoStreamData(kStreamIndex, codec_, kTimeScale)));
+  EXPECT_THAT(GetOutputStreamDataVector(),
+              ElementsAre(IsStreamInfo(kStreamIndex, kTimeScale, kEncrypted)));
+  const EncryptionConfig& encryption_config =
+      GetOutputStreamDataVector().back()->stream_info->encryption_config();
+  EXPECT_EQ(protection_scheme_, encryption_config.protection_scheme);
+  EXPECT_EQ(GetExpectedCryptByteBlock(), encryption_config.crypt_byte_block);
+  EXPECT_EQ(GetExpectedSkipByteBlock(), encryption_config.skip_byte_block);
+  EXPECT_EQ(std::vector<uint8_t>(std::begin(kKeyRotationDefaultKeyId),
+                                 std::end(kKeyRotationDefaultKeyId)),
+            encryption_config.key_id);
+  ClearOutputStreamDataVector();
+
+  InjectCodecParser();
+
+  // There are three segments. Only the third segment is encrypted.
+  // Crypto period duration is the same as segment duration, so there are three
+  // crypto periods, although only the last is encrypted.
+  for (int i = 0; i < 3; ++i) {
+    EXPECT_CALL(mock_key_source_, GetCryptoPeriodKey(i, _, _))
+        .WillOnce(DoAll(SetArgPointee<2>(GetMockEncryptionKey()),
+                        Return(Status::OK)));
+    // Use single-frame segment for testing.
+    ASSERT_OK(Process(GetMediaSampleStreamData(
+        kStreamIndex, i * kSegmentDuration, kSegmentDuration)));
+    ASSERT_OK(Process(GetSegmentInfoStreamData(
+        kStreamIndex, i * kSegmentDuration, kSegmentDuration, !kIsSubsegment)));
+    const bool is_encrypted = i == 2;
+    const auto& output_stream_data = GetOutputStreamDataVector();
+    EXPECT_THAT(output_stream_data,
+                ElementsAre(IsMediaSample(kStreamIndex, i * kSegmentDuration,
+                                          kSegmentDuration, is_encrypted),
+                            IsSegmentInfo(kStreamIndex, i * kSegmentDuration,
+                                          kSegmentDuration, !kIsSubsegment,
+                                          is_encrypted)));
+    EXPECT_THAT(*output_stream_data.back()
+                     ->segment_info->key_rotation_encryption_config,
+                MatchEncryptionConfig(
+                    protection_scheme_, GetExpectedCryptByteBlock(),
+                    GetExpectedSkipByteBlock(), GetExpectedPerSampleIvSize(),
+                    GetExpectedConstantIv(), GetMockEncryptionKey().key_id));
+    Mock::VerifyAndClearExpectations(&mock_key_source_);
+    ClearOutputStreamDataVector();
+  }
+}
+
 TEST_P(EncryptionHandlerEncryptionTest, Encrypt) {
+  EncryptionOptions encryption_options;
+  encryption_options.protection_scheme = protection_scheme_;
+  SetUpEncryptionHandler(encryption_options);
+
+  const EncryptionKey mock_encryption_key = GetMockEncryptionKey();
+  EXPECT_CALL(mock_key_source_, GetKey(_, _))
+      .WillOnce(
+          DoAll(SetArgPointee<1>(mock_encryption_key), Return(Status::OK)));
+
   ASSERT_OK(Process(GetStreamInfoStreamData(kStreamIndex, codec_, kTimeScale)));
   EXPECT_THAT(GetOutputStreamDataVector(),
               ElementsAre(IsStreamInfo(kStreamIndex, kTimeScale, kEncrypted)));
 
-  // Inject vpx parser / video slice header parser if needed.
+  InjectCodecParser();
-  switch (codec_) {
-    case kCodecVP9:{
-      std::unique_ptr<MockVpxParser> mock_vpx_parser(new MockVpxParser);
-      EXPECT_CALL(*mock_vpx_parser, Parse(_, sizeof(kData), _))
-          .WillOnce(
-              DoAll(SetArgPointee<2>(GetMockVpxFrameInfo()), Return(true)));
-      InjectVpxParserForTesting(std::move(mock_vpx_parser));
-      break;
-    }
-    case kCodecH264: {
-      std::unique_ptr<MockVideoSliceHeaderParser> mock_header_parser(
-          new MockVideoSliceHeaderParser);
-      // We want to return the same subsamples for VP9 and H264, so the return
-      // values here should match |GetMockVpxFrameInfo|.
-      EXPECT_CALL(*mock_header_parser, GetHeaderSize(_))
-          .WillOnce(Return(1))
-          .WillOnce(Return(2));
-      InjectVideoSliceHeaderParserForTesting(std::move(mock_header_parser));
-      break;
-    }
-    default:
-      break;
-  }
 
   std::unique_ptr<StreamData> stream_data(new StreamData);
   stream_data->stream_index = 0;
   stream_data->stream_data_type = StreamDataType::kMediaSample;
   stream_data->media_sample.reset(
-      new MediaSample(kData, sizeof(kData), nullptr, 0, true));
+      new MediaSample(kData, sizeof(kData), nullptr, 0, kIsKeyFrame));
 
-  EXPECT_CALL(mock_key_source_, GetKey(_, _))
+  ASSERT_OK(
-      .WillOnce(
+      Process(GetMediaSampleStreamData(kStreamIndex, 0, kSampleDuration)));
-          DoAll(SetArgPointee<1>(GetMockEncryptionKey()), Return(Status::OK)));
-  ASSERT_OK(Process(std::move(stream_data)));
   ASSERT_EQ(2u, GetOutputStreamDataVector().size());
-  ASSERT_EQ(0u, GetOutputStreamDataVector().back()->stream_index);
+  ASSERT_EQ(kStreamIndex, GetOutputStreamDataVector().back()->stream_index);
   ASSERT_EQ(StreamDataType::kMediaSample,
             GetOutputStreamDataVector().back()->stream_data_type);
 
@@ -358,12 +590,83 @@ TEST_P(EncryptionHandlerEncryptionTest, Encrypt) {
 }
 
 INSTANTIATE_TEST_CASE_P(
-    InstantiationName,
+    CencProtectionSchemes,
     EncryptionHandlerEncryptionTest,
     Combine(Values(FOURCC_cenc, FOURCC_cens, FOURCC_cbc1, FOURCC_cbcs),
             Values(kCodecAAC, kCodecH264, kCodecVP9)));
+INSTANTIATE_TEST_CASE_P(AppleSampleAes,
+                        EncryptionHandlerEncryptionTest,
+                        Combine(Values(kAppleSampleAesProtectionScheme),
+                                Values(kCodecAAC, kCodecH264)));
 
-// TODO(kqyang): Add more unit tests.
+namespace {
+
+const uint32_t kMaxSdPixels = 100u;
+const uint32_t kMaxHdPixels = 200u;
+const uint32_t kMaxUhd1Pixels = 300u;
+
+struct TrackTypeTestCase {
+  uint16_t width;
+  uint16_t height;
+  KeySource::TrackType track_type;
+};
+
+const TrackTypeTestCase kTrackTypeTestCases[] = {
+    TrackTypeTestCase{10, 10, KeySource::TRACK_TYPE_SD},
+    TrackTypeTestCase{11, 9, KeySource::TRACK_TYPE_SD},
+    TrackTypeTestCase{11, 10, KeySource::TRACK_TYPE_HD},
+    TrackTypeTestCase{20, 10, KeySource::TRACK_TYPE_HD},
+    TrackTypeTestCase{10, 20, KeySource::TRACK_TYPE_HD},
+    TrackTypeTestCase{19, 10, KeySource::TRACK_TYPE_HD},
+    TrackTypeTestCase{21, 10, KeySource::TRACK_TYPE_UHD1},
+    TrackTypeTestCase{29, 10, KeySource::TRACK_TYPE_UHD1},
+    TrackTypeTestCase{30, 10, KeySource::TRACK_TYPE_UHD1},
+    TrackTypeTestCase{20, 15, KeySource::TRACK_TYPE_UHD1},
+    TrackTypeTestCase{20, 16, KeySource::TRACK_TYPE_UHD2},
+    TrackTypeTestCase{1000, 1000, KeySource::TRACK_TYPE_UHD2},
+};
+
+}  // namespace
+
+class EncryptionHandlerTrackTypeTest
+    : public EncryptionHandlerTest,
+      public WithParamInterface<TrackTypeTestCase> {
+ public:
+  void SetUp() override {
+    EncryptionOptions encryption_options;
+    encryption_options.max_sd_pixels = kMaxSdPixels;
+    encryption_options.max_hd_pixels = kMaxHdPixels;
+    encryption_options.max_uhd1_pixels = kMaxUhd1Pixels;
+    SetUpEncryptionHandler(encryption_options);
+  }
+};
+
+TEST_F(EncryptionHandlerTrackTypeTest, AudioTrackType) {
+  EncryptionOptions encryption_options;
+  SetUpEncryptionHandler(encryption_options);
+  EXPECT_CALL(mock_key_source_, GetKey(KeySource::TRACK_TYPE_AUDIO, _))
+      .WillOnce(
+          DoAll(SetArgPointee<1>(GetMockEncryptionKey()), Return(Status::OK)));
+  ASSERT_OK(Process(GetAudioStreamInfoStreamData(kStreamIndex, kTimeScale)));
+}
+
+TEST_P(EncryptionHandlerTrackTypeTest, VideoTrackType) {
+  TrackTypeTestCase test_case = GetParam();
+  EXPECT_CALL(mock_key_source_, GetKey(test_case.track_type, _))
+      .WillOnce(
+          DoAll(SetArgPointee<1>(GetMockEncryptionKey()), Return(Status::OK)));
+  std::unique_ptr<StreamData> stream_data =
+      GetVideoStreamInfoStreamData(kStreamIndex, kTimeScale);
+  VideoStreamInfo* video_stream_info =
+      reinterpret_cast<VideoStreamInfo*>(stream_data->stream_info.get());
+  video_stream_info->set_width(test_case.width);
+  video_stream_info->set_height(test_case.height);
+  ASSERT_OK(Process(std::move(stream_data)));
+}
+
+INSTANTIATE_TEST_CASE_P(VideoResolutions,
+                        EncryptionHandlerTrackTypeTest,
+                        ValuesIn(kTrackTypeTestCases));
 
 }  // namespace media
 }  // namespace shaka