//
// Copyright 2013 Google Inc. All Rights Reserved.
//

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

#include "gtest/gtest.h"
#include "media/stagefright/foundation/ABase.h"
#include "media/stagefright/foundation/AString.h"
#include "MockCDM.h"
#include "wv_cdm_constants.h"
#include "WVCryptoPlugin.h"

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

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

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

  static const uint32_t kDataSize = 64;
  uint8_t in[kDataSize];
  uint8_t out[kDataSize];

  static const uint32_t kSubSampleCount = 3;
  CryptoPlugin::SubSample subSamples[kSubSampleCount];

  virtual void SetUp() {
    FILE* fp = fopen("/dev/urandom", "r");
    fread(sessionId, sizeof(uint8_t), kSessionIdSize, fp);
    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);

    memset(out, 0, sizeof(out));

    memset(subSamples, 0, sizeof(subSamples));
    subSamples[0].mNumBytesOfEncryptedData = 16;
    subSamples[1].mNumBytesOfClearData = 16;
    subSamples[1].mNumBytesOfEncryptedData = 24;
    subSamples[2].mNumBytesOfEncryptedData = 8;

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

TEST_F(WVCryptoPluginTest, CorrectlyReportsSecureBuffers) {
  MockCDM cdm;
  WVCryptoPlugin plugin(sessionId, kSessionIdSize, &cdm);

  EXPECT_FALSE(plugin.requiresSecureDecoderComponent("video/mp4")) <<
      "WVCryptoPlugin incorrectly expects a secure video decoder";
  EXPECT_FALSE(plugin.requiresSecureDecoderComponent("audio/aac")) <<
      "WVCryptoPlugin incorrectly expects a secure audio decoder";
}

TEST_F(WVCryptoPluginTest, RejectsSecureDecode) {
  MockCDM cdm;
  WVCryptoPlugin plugin(sessionId, kSessionIdSize, &cdm);

  // Decrypt should not be called because we specified an unsupported
  // security level
  EXPECT_CALL(cdm, Decrypt(_, _, _, _, _, _, _, _))
      .Times(0);

  ssize_t res = plugin.decrypt(true, keyId, iv, CryptoPlugin::kMode_AES_CTR,
                               in, subSamples, kSubSampleCount, out, NULL);

  EXPECT_EQ(static_cast<ssize_t>(-EPERM), res) <<
      "WVCryptoPlugin allowed decryption to proceed despite being asked for an "
      "unsupported security level";
}

TEST_F(WVCryptoPluginTest, AttemptsToDecrypt) {
  MockCDM cdm;
  WVCryptoPlugin plugin(sessionId, kSessionIdSize, &cdm);

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

    EXPECT_CALL(cdm, Decrypt(ElementsAreArray(sessionId, kSessionIdSize), true,
                             ElementsAreArray(keyId, KEY_ID_SIZE), in, 16,
                             ElementsAreArray(iv, KEY_IV_SIZE), 0, out))
        .WillOnce(Return(wvcdm::NO_ERROR));

    EXPECT_CALL(cdm, Decrypt(ElementsAreArray(sessionId, kSessionIdSize), false,
                             ElementsAreArray(keyId, KEY_ID_SIZE), in + 16, 16,
                             ElementsAreArray(iv, KEY_IV_SIZE), 0, out + 16))
        .WillOnce(Return(wvcdm::NO_ERROR));

    EXPECT_CALL(cdm, Decrypt(ElementsAreArray(sessionId, kSessionIdSize), true,
                             ElementsAreArray(keyId, KEY_ID_SIZE), in + 32, 24,
                             ElementsAreArray(iv, KEY_IV_SIZE), 0, out + 32))
        .WillOnce(Return(wvcdm::NO_ERROR));

    EXPECT_CALL(cdm, Decrypt(ElementsAreArray(sessionId, kSessionIdSize), true,
                             ElementsAreArray(keyId, KEY_ID_SIZE), in + 56, 8,
                             ElementsAreArray(iv, KEY_IV_SIZE), 8, out + 56))
        .WillOnce(Return(wvcdm::NO_ERROR));
  }

  ssize_t res = plugin.decrypt(false, keyId, iv, CryptoPlugin::kMode_AES_CTR,
                               in, subSamples, kSubSampleCount, out, NULL);

  EXPECT_EQ(static_cast<ssize_t>(kDataSize), res) <<
      "WVCryptoPlugin decrypted the wrong number of bytes";
}