// 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

#include "packager/media/crypto/encryption_handler.h"

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include "packager/media/base/aes_decryptor.h"
#include "packager/media/base/aes_pattern_cryptor.h"
#include "packager/media/base/fixed_key_source.h"
#include "packager/media/base/media_handler_test_base.h"
#include "packager/media/base/test/status_test_util.h"
#include "packager/media/codecs/video_slice_header_parser.h"
#include "packager/media/codecs/vpx_parser.h"

namespace shaka {
namespace media {
namespace {

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));
  MOCK_METHOD3(GetCryptoPeriodKey,
               Status(uint32_t crypto_period_index,
                      TrackType track_type,
                      EncryptionKey* key));
};

class MockVpxParser : public VPxParser {
 public:
  MOCK_METHOD3(Parse,
               bool(const uint8_t* data,
                    size_t data_size,
                    std::vector<VPxFrameInfo>* vpx_frames));
};

class MockVideoSliceHeaderParser : public VideoSliceHeaderParser {
 public:
  MOCK_METHOD1(Initialize,
               bool(const std::vector<uint8_t>& decoder_configuration));
  MOCK_METHOD1(GetHeaderSize, int64_t(const Nalu& nalu));
};

}  // namespace

class EncryptionHandlerTest : public MediaHandlerTestBase {
 public:
  void SetUp() override { SetUpEncryptionHandler(EncryptionOptions()); }

  void SetUpEncryptionHandler(const EncryptionOptions& encryption_options) {
    encryption_handler_.reset(
        new EncryptionHandler(encryption_options, &mock_key_source_));
    SetUpGraph(1 /* one input */, 1 /* one output */, encryption_handler_);
  }

  Status Process(std::unique_ptr<StreamData> stream_data) {
    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));
  }

  void InjectVideoSliceHeaderParserForTesting(
      std::unique_ptr<VideoSliceHeaderParser> header_parser) {
    encryption_handler_->InjectVideoSliceHeaderParserForTesting(
        std::move(header_parser));
  }

 protected:
  std::shared_ptr<EncryptionHandler> encryption_handler_;
  StrictMock<MockKeySource> mock_key_source_;
};

TEST_F(EncryptionHandlerTest, Initialize) {
  ASSERT_OK(encryption_handler_->Initialize());
}

TEST_F(EncryptionHandlerTest, OnlyOneOutput) {
  // Connecting another handler will fail.
  ASSERT_OK(encryption_handler_->AddHandler(some_handler()));
  ASSERT_EQ(error::INVALID_ARGUMENT,
            encryption_handler_->Initialize().error_code());
}

TEST_F(EncryptionHandlerTest, OnlyOneInput) {
  ASSERT_OK(some_handler()->AddHandler(encryption_handler_));
  ASSERT_EQ(error::INVALID_ARGUMENT,
            encryption_handler_->Initialize().error_code());
}

namespace {

const bool kVp9SubsampleEncryption = true;
const bool kIsKeyFrame = true;
const bool kIsSubsegment = true;
const bool kEncrypted = true;
const size_t kStreamIndex = 0;
const uint32_t kTimeScale = 1000;
const int64_t kSampleDuration = 1000;
const int64_t kSegmentDuration = 1000;

// 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
    0x30, 0x01, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09,
    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
    0x31, 0x25, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09,
    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, 0x47,
    // Third non-video-slice NALU for H264 or superframe index for VP9.
    0x06, 0x67, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
};

// H264 subsample information for the the above data.
const size_t kNaluLengthSize = 1u;
const size_t kNaluHeaderSize = 1u;
const size_t kSubsampleSize1 = 49u;
const size_t kSliceHeaderSize1 = 1u;
const size_t kSubsampleSize2 = 50u;
const size_t kSliceHeaderSize2 = 16u;
const size_t kSubsampleSize3 = 7u;
// VP9 frame information for the above data. It should match H264 subsample
// 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

inline bool operator==(const SubsampleEntry& lhs, const SubsampleEntry& rhs) {
  return lhs.clear_bytes == rhs.clear_bytes &&
         lhs.cipher_bytes == rhs.cipher_bytes;
}

class EncryptionHandlerEncryptionTest
    : public EncryptionHandlerTest,
      public WithParamInterface<std::tr1::tuple<FourCC, Codec, bool>> {
 public:
  void SetUp() override {
    protection_scheme_ = std::tr1::get<0>(GetParam());
    codec_ = std::tr1::get<1>(GetParam());
    vp9_subsample_encryption_ = std::tr1::get<2>(GetParam());
  }

  std::vector<VPxFrameInfo> GetMockVpxFrameInfo() {
    std::vector<VPxFrameInfo> vpx_frames;
    vpx_frames.resize(2);
    vpx_frames[0].frame_size = kVpxFrameSize1;
    vpx_frames[0].uncompressed_header_size = kUncompressedHeaderSize1;
    vpx_frames[1].frame_size = kVpxFrameSize2;
    vpx_frames[1].uncompressed_header_size = kUncompressedHeaderSize2;
    return vpx_frames;
  }

  // The subsamples values should match |GetMockVpxFrameInfo| above.
  std::vector<SubsampleEntry> GetExpectedSubsamples() {
    std::vector<SubsampleEntry> subsamples;
    if (codec_ == kCodecAAC ||
        (codec_ == kCodecVP9 && !vp9_subsample_encryption_)) {
      return subsamples;
    }
    if (protection_scheme_ == kAppleSampleAesProtectionScheme) {
      subsamples.emplace_back(static_cast<uint16_t>(kSampleAesClearSize1),
                              static_cast<uint32_t>(kSampleAesCipherSize1));
      subsamples.emplace_back(static_cast<uint16_t>(kSubsampleSize3), 0u);
    } else {
      if (codec_ == kCodecVP9 || protection_scheme_ == FOURCC_cbc1 ||
          protection_scheme_ == FOURCC_cens ||
          protection_scheme_ == FOURCC_cenc) {
        // Align the encrypted bytes to multiple of 16 bytes.
        subsamples.emplace_back(static_cast<uint16_t>(kAlignedClearSize1),
                                static_cast<uint32_t>(kAlignedCipherSize1));
        // Subsample 2 is already aligned.
      } else {
        subsamples.emplace_back(static_cast<uint16_t>(kClearSize1),
                                static_cast<uint32_t>(kCipherSize1));
      }
      subsamples.emplace_back(static_cast<uint16_t>(kClearSize2),
                              static_cast<uint32_t>(kCipherSize2));
      subsamples.emplace_back(static_cast<uint16_t>(kSubsampleSize3), 0u);
    }
    return subsamples;
  }

  std::unique_ptr<StreamData> GetMediaSampleStreamData(int stream_index,
                                                       int64_t timestamp,
                                                       int64_t duration) {
    std::unique_ptr<StreamData> stream_data(new StreamData);
    stream_data->stream_index = stream_index;
    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:
        if (vp9_subsample_encryption_) {
          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:
        aes_decryptor.reset(new AesCtrDecryptor);
        break;
      case FOURCC_cbc1:
        aes_decryptor.reset(new AesCbcDecryptor(kNoPadding));
        break;
      case FOURCC_cens:
        aes_decryptor.reset(new AesPatternCryptor(
            decrypt_config.crypt_byte_block(), decrypt_config.skip_byte_block(),
            AesPatternCryptor::kEncryptIfCryptByteBlockRemaining,
            AesCryptor::kDontUseConstantIv,
            std::unique_ptr<AesCryptor>(new AesCtrDecryptor())));
        break;
      case FOURCC_cbcs:
        aes_decryptor.reset(new AesPatternCryptor(
            decrypt_config.crypt_byte_block(), decrypt_config.skip_byte_block(),
            AesPatternCryptor::kEncryptIfCryptByteBlockRemaining,
            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.";
    }

    if (!aes_decryptor->InitializeWithIv(
            std::vector<uint8_t>(kKey, kKey + sizeof(kKey)),
            decrypt_config.iv())) {
      return false;
    }

    if (decrypt_config.subsamples().empty()) {
      // Sample not encrypted using subsample encryption. Decrypt whole.
      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;
      }
      return true;
    }

    // Subsample decryption.
    const std::vector<SubsampleEntry>& subsamples = decrypt_config.subsamples();
    uint8_t* current_ptr = data;
    const uint8_t* const buffer_end = data + data_size;
    for (const auto& subsample : subsamples) {
      if (current_ptr + subsample.clear_bytes + subsample.cipher_bytes >
          buffer_end) {
        LOG(ERROR) << "Subsamples overflow sample buffer.";
        return false;
      }
      current_ptr += subsample.clear_bytes;
      if (!aes_decryptor->Crypt(current_ptr, subsample.cipher_bytes,
                                current_ptr)) {
        LOG(ERROR) << "Error decrypting subsample buffer.";
        return false;
      }
      current_ptr += subsample.cipher_bytes;
    }
    return true;
  }

  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:
        return 0;
      case FOURCC_cens:
      case FOURCC_cbcs:
        return 1;
      default:
        return 0;
    }
  }

  uint8_t GetExpectedSkipByteBlock() {
    // Always use full sample encryption for audio.
    if (codec_ == kCodecAAC)
      return 0;
    switch (protection_scheme_) {
      case FOURCC_cenc:
      case FOURCC_cbc1:
        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_;
  bool vp9_subsample_encryption_;
};

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;
  encryption_options.vp9_subsample_encryption = vp9_subsample_encryption_;
  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;
  encryption_options.vp9_subsample_encryption = vp9_subsample_encryption_;
  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_;
  encryption_options.vp9_subsample_encryption = vp9_subsample_encryption_;
  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)));

  InjectCodecParser();

  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, kIsKeyFrame));

  ASSERT_OK(
      Process(GetMediaSampleStreamData(kStreamIndex, 0, kSampleDuration)));
  ASSERT_EQ(2u, GetOutputStreamDataVector().size());
  ASSERT_EQ(kStreamIndex, GetOutputStreamDataVector().back()->stream_index);
  ASSERT_EQ(StreamDataType::kMediaSample,
            GetOutputStreamDataVector().back()->stream_data_type);

  auto* media_sample = GetOutputStreamDataVector().back()->media_sample.get();
  auto* decrypt_config = media_sample->decrypt_config();
  EXPECT_EQ(std::vector<uint8_t>(kKeyId, kKeyId + sizeof(kKeyId)),
            decrypt_config->key_id());
  EXPECT_EQ(std::vector<uint8_t>(kIv, kIv + sizeof(kIv)), decrypt_config->iv());
  EXPECT_EQ(GetExpectedSubsamples(), decrypt_config->subsamples());
  EXPECT_EQ(protection_scheme_, decrypt_config->protection_scheme());
  EXPECT_EQ(GetExpectedCryptByteBlock(), decrypt_config->crypt_byte_block());
  EXPECT_EQ(GetExpectedSkipByteBlock(), decrypt_config->skip_byte_block());

  ASSERT_TRUE(Decrypt(*decrypt_config, media_sample->writable_data(),
                      media_sample->data_size()));
  EXPECT_EQ(
      std::vector<uint8_t>(kData, kData + sizeof(kData)),
      std::vector<uint8_t>(media_sample->data(),
                           media_sample->data() + media_sample->data_size()));
}

INSTANTIATE_TEST_CASE_P(
    CencProtectionSchemes,
    EncryptionHandlerEncryptionTest,
    Combine(Values(FOURCC_cenc, FOURCC_cens, FOURCC_cbc1, FOURCC_cbcs),
            Values(kCodecAAC, kCodecH264, kCodecVP9),
            Values(kVp9SubsampleEncryption, !kVp9SubsampleEncryption)));
INSTANTIATE_TEST_CASE_P(AppleSampleAes,
                        EncryptionHandlerEncryptionTest,
                        Combine(Values(kAppleSampleAesProtectionScheme),
                                Values(kCodecAAC, kCodecH264),
                                Values(kVp9SubsampleEncryption)));

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