//
// Copyright 2018 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine Master
// License Agreement.
//

#include <stdio.h>
#include <string>
#include <vector>

#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "media/stagefright/foundation/ABase.h"
#include "media/stagefright/foundation/AString.h"
#include "media/stagefright/MediaErrors.h"

#include "OEMCryptoCENC.h"
#include "wv_cdm_constants.h"
#include "wv_cdm_types.h"
#include "wv_content_decryption_module.h"
#include "WVCryptoPlugin.h"

using namespace android;
using namespace std;
using namespace testing;
using namespace wvcdm;
using namespace wvdrm;

namespace {
constexpr ssize_t kErrorUnsupportedCrypto = ERROR_DRM_VENDOR_MIN + 2;
}

class MockCDM : public WvContentDecryptionModule {
 public:
  MOCK_METHOD1(IsOpenSession, bool(const CdmSessionId&));

  MOCK_METHOD3(DecryptV16, CdmResponseType(const CdmSessionId&, bool,
                                           const CdmDecryptionParametersV16&));

  MOCK_METHOD2(QuerySessionStatus, CdmResponseType(const CdmSessionId&,
                                                   CdmQueryMap*));
};

class WVCryptoPluginTest : public Test {
 protected:
  static const uint32_t kSessionIdSize = 16;
  uint8_t sessionId[kSessionIdSize];

  virtual void SetUp() {
    FILE* fp = fopen("/dev/urandom", "r");
    fread(sessionId, sizeof(uint8_t), kSessionIdSize, fp);
    fclose(fp);

    // Set default CdmResponseType value for gMock
    DefaultValue<CdmResponseType>::Set(wvcdm::NO_ERROR);
  }
};

TEST_F(WVCryptoPluginTest, CorrectlyReportsSecureBuffers) {
  android::sp<StrictMock<MockCDM>> cdm = new StrictMock<MockCDM>();

  CdmQueryMap l1Map;
  l1Map[QUERY_KEY_SECURITY_LEVEL] = QUERY_VALUE_SECURITY_LEVEL_L1;
  CdmQueryMap l3Map;
  l3Map[QUERY_KEY_SECURITY_LEVEL] = QUERY_VALUE_SECURITY_LEVEL_L3;

  // Provide the expected behavior for IsOpenSession
  EXPECT_CALL(*cdm, IsOpenSession(_))
      .WillRepeatedly(Return(true));

  // Specify the expected calls to QuerySessionStatus
  EXPECT_CALL(*cdm, QuerySessionStatus(_, _))
      .WillOnce(DoAll(SetArgPointee<1>(l1Map),
                      Return(wvcdm::NO_ERROR)))
      .WillOnce(DoAll(SetArgPointee<1>(l3Map),
                      Return(wvcdm::NO_ERROR)));

  WVCryptoPlugin plugin(sessionId, kSessionIdSize, cdm.get());

  EXPECT_TRUE(plugin.requiresSecureDecoderComponent("video/mp4")) <<
      "WVCryptoPlugin incorrectly allows an insecure video decoder on L1";
  EXPECT_FALSE(plugin.requiresSecureDecoderComponent("video/mp4")) <<
      "WVCryptoPlugin incorrectly expects a secure video decoder on L3";
  EXPECT_FALSE(plugin.requiresSecureDecoderComponent("audio/aac")) <<
      "WVCryptoPlugin incorrectly expects a secure audio decoder";
}

// TODO(b/28295739):
// Add New MediaCrypto Unit Tests for CBC & Pattern Mode in cdmPatternDesc.

// Predicate that validates that the fields of a passed-in
// CdmDecryptionParametersV16 match the values it was given at construction
// time.
//
// This could be done with a huge pile of gMock matchers, but it is ugly and
// unmaintainable, particularly once you get into validating the subsamples. The
// logic here is complex enough to warrant a custom matcher for this one test.
class CDPMatcher {
 public:
  CDPMatcher(const uint8_t* keyId, bool isSecure, CryptoPlugin::Mode cipherMode,
             const CryptoPlugin::Pattern& pattern,
             const uint8_t* input, size_t inputLength,
             const uint8_t* output, size_t outputLength, const uint8_t* iv,
             const CryptoPlugin::SubSample* subsamples,
             size_t subsamplesLength)
    : mKeyId(keyId, keyId + KEY_ID_SIZE), mIsSecure(isSecure),
      mCipherMode(cipherMode), mPattern(pattern), mInput(input),
      mInputLength(inputLength), mOutput(output), mOutputLength(outputLength),
      mIv(iv, iv + KEY_IV_SIZE),
      mSubsamples(subsamples, subsamples + subsamplesLength) {}

  bool operator()(const CdmDecryptionParametersV16& params) const {
    if (mCipherMode == CryptoPlugin::kMode_AES_CTR &&
        params.cipher_mode != kCipherModeCtr) {
      return false;
    } else if (mCipherMode == CryptoPlugin::kMode_AES_CBC &&
        params.cipher_mode != kCipherModeCbc) {
      return false;
    }

    if (params.key_id != mKeyId ||
        params.is_secure != mIsSecure ||
        params.pattern.encrypt_blocks != mPattern.mEncryptBlocks ||
        params.pattern.skip_blocks != mPattern.mSkipBlocks ||
        params.samples.size() != 1) {
      return false;
    }

    const CdmDecryptionSample& sample = params.samples[0];
    if (sample.encrypt_buffer != mInput ||
        sample.encrypt_buffer_length != mInputLength ||
        sample.decrypt_buffer != mOutput ||
        sample.decrypt_buffer_size != mOutputLength ||
        sample.decrypt_buffer_offset != 0 ||
        sample.iv != mIv ||
        sample.subsamples.size() != mSubsamples.size()) {
      return false;
    }

    for (size_t i = 0; i < mSubsamples.size(); ++i) {
      const CryptoPlugin::SubSample& androidSubsample = mSubsamples[i];
      const CdmDecryptionSubsample& cdmSubsample = sample.subsamples[i];

      if (cdmSubsample.clear_bytes != androidSubsample.mNumBytesOfClearData||
          cdmSubsample.protected_bytes != androidSubsample.mNumBytesOfEncryptedData) {
        return false;
      }
    }

    return true;
  }

 private:
  const KeyId mKeyId;
  const bool mIsSecure;
  const CryptoPlugin::Mode mCipherMode;
  const CryptoPlugin::Pattern mPattern;
  const uint8_t* const mInput;
  const size_t mInputLength;
  const uint8_t* const mOutput;
  const size_t mOutputLength;
  const std::vector<uint8_t> mIv;
  const std::vector<CryptoPlugin::SubSample> mSubsamples;
};

TEST_F(WVCryptoPluginTest, AttemptsToDecrypt) {
  android::sp<StrictMock<MockCDM>> cdm = new StrictMock<MockCDM>();

  constexpr size_t kSubSampleCount = 6;
  CryptoPlugin::SubSample subSamples[kSubSampleCount];
  memset(subSamples, 0, sizeof(subSamples));
  subSamples[0].mNumBytesOfEncryptedData = 16;
  subSamples[1].mNumBytesOfClearData = 16;
  subSamples[1].mNumBytesOfEncryptedData = 16;
  subSamples[2].mNumBytesOfEncryptedData = 8;
  subSamples[3].mNumBytesOfClearData = 29;
  subSamples[3].mNumBytesOfEncryptedData = 24;
  subSamples[4].mNumBytesOfEncryptedData = 60;
  subSamples[5].mNumBytesOfEncryptedData = 16;

  uint8_t keyId[KEY_ID_SIZE];
  uint8_t iv[KEY_IV_SIZE];

  constexpr size_t kDataSize = 185;
  uint8_t inputData[kDataSize];
  uint8_t outputData[kDataSize];

  FILE* fp = fopen("/dev/urandom", "r");
  fread(keyId, sizeof(uint8_t), KEY_ID_SIZE, fp);
  fread(iv, sizeof(uint8_t), KEY_IV_SIZE, fp);
  fread(inputData, sizeof(uint8_t), kDataSize, fp);
  fclose(fp);

  android::CryptoPlugin::Pattern noPattern = { 0, 0 };

  // Provide the expected behavior for IsOpenSession
  EXPECT_CALL(*cdm, IsOpenSession(_))
      .WillRepeatedly(Return(true));

  // Specify the expected calls to Decrypt
  CDPMatcher paramsMatcher(keyId, false, CryptoPlugin::kMode_AES_CTR, noPattern,
                           inputData, kDataSize, outputData, kDataSize, iv,
                           subSamples, kSubSampleCount);

  EXPECT_CALL(*cdm, DecryptV16(ElementsAreArray(sessionId, kSessionIdSize),
                               true,
                               Truly(paramsMatcher)))
      .Times(1);

  WVCryptoPlugin plugin(sessionId, kSessionIdSize, cdm.get());
  AString errorDetailMessage;

  ssize_t res = plugin.decrypt(false, keyId, iv, CryptoPlugin::kMode_AES_CTR,
                               noPattern, inputData, subSamples,
                               kSubSampleCount, outputData,
                               &errorDetailMessage);

  EXPECT_EQ(static_cast<ssize_t>(kDataSize), res) <<
      "WVCryptoPlugin decrypted the wrong number of bytes";
  EXPECT_EQ(0u, errorDetailMessage.size()) <<
      "WVCryptoPlugin reported a detailed error message.";
}

TEST_F(WVCryptoPluginTest, RejectsCens) {
  android::sp<StrictMock<MockCDM>> cdm = new StrictMock<MockCDM>();

  constexpr size_t kSubSampleCount = 2;
  CryptoPlugin::SubSample subSamples[kSubSampleCount];
  memset(subSamples, 0, sizeof(subSamples));
  subSamples[0].mNumBytesOfEncryptedData = 16;
  subSamples[1].mNumBytesOfClearData = 16;
  subSamples[1].mNumBytesOfEncryptedData = 16;

  uint8_t keyId[KEY_ID_SIZE];
  uint8_t iv[KEY_IV_SIZE];

  constexpr size_t kDataSize = 48;
  uint8_t inputData[kDataSize];
  uint8_t outputData[kDataSize];

  FILE* fp = fopen("/dev/urandom", "r");
  fread(keyId, sizeof(uint8_t), KEY_ID_SIZE, fp);
  fread(iv, sizeof(uint8_t), KEY_IV_SIZE, fp);
  fread(inputData, sizeof(uint8_t), kDataSize, fp);
  fclose(fp);

  android::CryptoPlugin::Pattern recommendedPattern = { 1, 9 };

  // Provide the expected behavior for IsOpenSession
  EXPECT_CALL(*cdm, IsOpenSession(_))
      .WillRepeatedly(Return(true));

  // Refuse calls to Decrypt
  EXPECT_CALL(*cdm, DecryptV16(_, _, _))
      .Times(0);

  WVCryptoPlugin plugin(sessionId, kSessionIdSize, cdm.get());
  AString errorDetailMessage;

  ssize_t res = plugin.decrypt(false, keyId, iv, CryptoPlugin::kMode_AES_CTR,
                               recommendedPattern, inputData, subSamples,
                               kSubSampleCount, outputData,
                               &errorDetailMessage);

  EXPECT_EQ(res, kErrorUnsupportedCrypto) <<
      "WVCryptoPlugin did not return an error for 'cens'.";
  EXPECT_NE(errorDetailMessage.size(), 0u) <<
      "WVCryptoPlugin did not report a detailed error message for 'cens'.";
}

TEST_F(WVCryptoPluginTest, CommunicatesSecureBufferRequest) {
  android::sp<StrictMock<MockCDM>> cdm = new StrictMock<MockCDM>();

  uint8_t keyId[KEY_ID_SIZE];
  uint8_t iv[KEY_IV_SIZE];

  static const size_t kDataSize = 32;
  uint8_t in[kDataSize];
  uint8_t out[kDataSize];

  FILE* fp = fopen("/dev/urandom", "r");
  fread(keyId, sizeof(uint8_t), KEY_ID_SIZE, fp);
  fread(iv, sizeof(uint8_t), KEY_IV_SIZE, fp);
  fread(in, sizeof(uint8_t), kDataSize, fp);
  fclose(fp);

  static const uint32_t kSubSampleCount = 1;
  CryptoPlugin::SubSample subSamples[kSubSampleCount];
  memset(subSamples, 0, sizeof(subSamples));
  subSamples[0].mNumBytesOfClearData = 16;
  subSamples[0].mNumBytesOfEncryptedData = 16;

  // Provide the expected behavior for IsOpenSession
  EXPECT_CALL(*cdm, IsOpenSession(_))
      .WillRepeatedly(Return(true));

  // Specify the expected calls to Decrypt
  {
    InSequence calls;

    typedef CdmDecryptionParametersV16 CDP;

    EXPECT_CALL(*cdm, DecryptV16(_, _, Field(&CDP::is_secure, false)))
        .Times(1);

    EXPECT_CALL(*cdm, DecryptV16(_, _, Field(&CDP::is_secure, true)))
        .Times(1);
  }

  WVCryptoPlugin plugin(sessionId, kSessionIdSize, cdm.get());
  android::CryptoPlugin::Pattern noPattern = { 0, 0 };
  AString errorDetailMessage;

  ssize_t res = plugin.decrypt(false, keyId, iv, CryptoPlugin::kMode_AES_CTR,
                               noPattern, in, subSamples, kSubSampleCount,
                               out, &errorDetailMessage);
  ASSERT_GE(res, 0) <<
      "WVCryptoPlugin returned an error";
  EXPECT_EQ(0u, errorDetailMessage.size()) <<
      "WVCryptoPlugin reported a detailed error message.";

  res = plugin.decrypt(true, keyId, iv, CryptoPlugin::kMode_AES_CTR,
                       noPattern, in, subSamples, kSubSampleCount, out,
                       &errorDetailMessage);
  ASSERT_GE(res, 0) <<
      "WVCryptoPlugin returned an error";
  EXPECT_EQ(0u, errorDetailMessage.size()) <<
      "WVCryptoPlugin reported a detailed error message.";
}

TEST_F(WVCryptoPluginTest, AllowsSessionIdChanges) {
  android::sp<StrictMock<MockCDM>> cdm = new StrictMock<MockCDM>();

  uint8_t keyId[KEY_ID_SIZE];
  uint8_t iv[KEY_IV_SIZE];
  uint8_t sessionId2[kSessionIdSize];

  static const size_t kDataSize = 32;
  uint8_t in[kDataSize];
  uint8_t out[kDataSize];

  FILE* fp = fopen("/dev/urandom", "r");
  fread(keyId, sizeof(uint8_t), KEY_ID_SIZE, fp);
  fread(iv, sizeof(uint8_t), KEY_IV_SIZE, fp);
  fread(sessionId2, sizeof(uint8_t), kSessionIdSize, fp);
  fread(in, sizeof(uint8_t), kDataSize, fp);
  fclose(fp);

  static const uint32_t kSubSampleCount = 1;
  CryptoPlugin::SubSample subSamples[kSubSampleCount];
  memset(subSamples, 0, sizeof(subSamples));
  subSamples[0].mNumBytesOfClearData = 16;
  subSamples[0].mNumBytesOfEncryptedData = 16;

  Vector<uint8_t> sessionIdVector;
  sessionIdVector.appendArray(sessionId, kSessionIdSize);
  Vector<uint8_t> sessionId2Vector;
  sessionId2Vector.appendArray(sessionId2, kSessionIdSize);

  // Provide the expected behavior for IsOpenSession
  EXPECT_CALL(*cdm, IsOpenSession(_))
      .WillRepeatedly(Return(true));

  // Specify the expected calls to Decrypt
  {
    InSequence calls;

    EXPECT_CALL(*cdm,
                DecryptV16(ElementsAreArray(sessionId, kSessionIdSize), _, _))
        .Times(1);

    EXPECT_CALL(*cdm,
                DecryptV16(ElementsAreArray(sessionId2, kSessionIdSize), _, _))
        .Times(1);
  }

  uint8_t blank[1];  // Some compilers will not accept 0.
  WVCryptoPlugin plugin(blank, 0, cdm.get());
  android::CryptoPlugin::Pattern noPattern = { 0, 0 };
  AString errorDetailMessage;
  ssize_t res;

  res = plugin.setMediaDrmSession(sessionIdVector);
  EXPECT_EQ(android::NO_ERROR, res);
  res = plugin.decrypt(false, keyId, iv, CryptoPlugin::kMode_AES_CTR,
                       noPattern, in, subSamples, kSubSampleCount, out,
                       &errorDetailMessage);
  EXPECT_GE(res, 0) <<
      "WVCryptoPlugin returned an error";
  EXPECT_EQ(0u, errorDetailMessage.size()) <<
      "WVCryptoPlugin reported a detailed error message.";

  res = plugin.setMediaDrmSession(sessionId2Vector);
  EXPECT_EQ(android::NO_ERROR, res);
  res = plugin.decrypt(false, keyId, iv, CryptoPlugin::kMode_AES_CTR,
                       noPattern, in, subSamples, kSubSampleCount, out,
                       &errorDetailMessage);
  EXPECT_GE(res, 0) <<
      "WVCryptoPlugin returned an error";
  EXPECT_EQ(0u, errorDetailMessage.size()) <<
      "WVCryptoPlugin reported a detailed error message.";
}

TEST_F(WVCryptoPluginTest, DisallowsUnopenedSessionIdChanges) {
  android::sp<StrictMock<MockCDM>> cdm = new StrictMock<MockCDM>();

  uint8_t blank[1];  // Some compilers will not accept 0.
  Vector<uint8_t> sessionIdVector;
  sessionIdVector.appendArray(sessionId, kSessionIdSize);

  // Specify the expected calls to IsOpenSession
  {
    InSequence calls;

    EXPECT_CALL(*cdm, IsOpenSession(ElementsAreArray(blank, 0)))
        .WillOnce(Return(false));

    EXPECT_CALL(*cdm, IsOpenSession(ElementsAreArray(sessionId, kSessionIdSize)))
        .WillOnce(Return(false))
        .WillOnce(Return(true));
  }

  WVCryptoPlugin plugin(blank, 0, cdm.get());

  ssize_t res;
  res = plugin.setMediaDrmSession(sessionIdVector);
  EXPECT_EQ(android::ERROR_DRM_SESSION_NOT_OPENED, res);
  res = plugin.setMediaDrmSession(sessionIdVector);
  EXPECT_EQ(android::NO_ERROR, res);
}