// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox.h"

#include <memory>
#include <string>
#include <vector>

#include "api/golden_data.h"
#include "api/license_builder.h"
#include "api/license_whitebox_test_base.h"
#include "api/test_data.h"
#include "crypto_utils/rsa_key.h"
#include "testing/include/gtest/gtest.h"

namespace widevine {

class LicenseWhiteboxDecryptTest : public LicenseWhiteboxTestBase {
 protected:
  void SetUp() override {
    LicenseWhiteboxTestBase::SetUp();

    // Because we are using the same buffer for both CTR and CBC, need to make
    // sure that it is large enough for either one.
    plaintext_size_ = std::max(golden_data_.CBCContent().ciphertext.size(),
                               golden_data_.CTRContent().ciphertext.size());
    plaintext_.resize(plaintext_size_);

    golden_data_.MakeKeyIdDifferent(&non_content_key_id_);
    golden_data_.MakeKeyIdDifferent(&missing_key_id_);
  }

  void LoadLicense(const std::vector<uint8_t>& padding) {
    LicenseBuilder builder;

    builder.AddContentKey(golden_data_.CBCCryptoKey().level,
                          golden_data_.CBCCryptoKey().id,
                          golden_data_.CBCCryptoKey().content->key);

    builder.AddContentKey(golden_data_.CTRCryptoKey().level,
                          golden_data_.CTRCryptoKey().id,
                          golden_data_.CTRCryptoKey().content->key);

    builder.AddContentKey(golden_data_.CBCDecodeKey().level,
                          golden_data_.CBCDecodeKey().id,
                          golden_data_.CBCDecodeKey().content->key);

    builder.AddContentKey(golden_data_.CBCHardwareKey().level,
                          golden_data_.CBCHardwareKey().id,
                          golden_data_.CBCHardwareKey().content->key);

    builder.AddOperatorSessionKey(non_content_key_id_);

    License license;
    builder.Build(*public_key_, &license);

    ASSERT_EQ(WB_License_ProcessLicenseResponse(
                  whitebox_, license.message.data(), license.message.size(),
                  license.signature.data(), license.signature.size(),
                  license.session_key.data(), license.session_key.size(),
                  license.request.data(), license.request.size()),
              WB_RESULT_OK);
  }

  // We need two special keys for this test, one that will be used for a
  // non-content key and one that will never be in the license.
  std::vector<uint8_t> non_content_key_id_ = {0, 0, 0};
  std::vector<uint8_t> missing_key_id_ = {1, 0, 0};

  // This is the buffer used to store the output of each decrypt call.
  size_t plaintext_size_;
  std::vector<uint8_t> plaintext_;
};

TEST_F(LicenseWhiteboxDecryptTest, CryptoKeyWithCbcDataInCbcMode) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_OK);
  plaintext_.resize(plaintext_size_);
  ASSERT_EQ(plaintext_, golden_data_.CBCCryptoKey().content->plaintext);
}

TEST_F(LicenseWhiteboxDecryptTest, CryptoKeyWithCtrDataInCtrMode) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CTR,
                         golden_data_.CTRCryptoKey().id.data(),
                         golden_data_.CTRCryptoKey().id.size(),
                         golden_data_.CTRCryptoKey().content->ciphertext.data(),
                         golden_data_.CTRCryptoKey().content->ciphertext.size(),
                         golden_data_.CTRCryptoKey().content->iv.data(),
                         golden_data_.CTRCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_OK);
  plaintext_.resize(plaintext_size_);
  ASSERT_EQ(plaintext_, golden_data_.CTRCryptoKey().content->plaintext);
}

// We try to decrypt CBC encrypted data in CTR mode. All operations should be
// successful, but the resulting plaintext should not match.
TEST_F(LicenseWhiteboxDecryptTest, CryptoKeyWithCbcDataInCtrMode) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CTR,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_OK);
  plaintext_.resize(plaintext_size_);
  ASSERT_NE(plaintext_, golden_data_.CBCCryptoKey().content->plaintext);
}

// We try to decrypt CTR encrypted data in CBC mode. All operations should be
// successful, but the resulting plaintext should not match.
TEST_F(LicenseWhiteboxDecryptTest, CryptoKeyWithCtrDataInCbcMode) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CTRCryptoKey().id.data(),
                         golden_data_.CTRCryptoKey().id.size(),
                         golden_data_.CTRCryptoKey().content->ciphertext.data(),
                         golden_data_.CTRCryptoKey().content->ciphertext.size(),
                         golden_data_.CTRCryptoKey().content->iv.data(),
                         golden_data_.CTRCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_OK);
  plaintext_.resize(plaintext_size_);
  ASSERT_NE(plaintext_, golden_data_.CTRCryptoKey().content->plaintext);
}

TEST_F(LicenseWhiteboxDecryptTest, CryptoKeyWithCbcDataAndPKCS8Padding) {
  LoadLicense(LicenseBuilder::PKSC8Padding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_OK);
  plaintext_.resize(plaintext_size_);
  ASSERT_EQ(plaintext_, golden_data_.CBCCryptoKey().content->plaintext);
}

TEST_F(LicenseWhiteboxDecryptTest, CryptoKeyWithCtrDataAndPKCS8Padding) {
  LoadLicense(LicenseBuilder::PKSC8Padding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CTR,
                         golden_data_.CTRCryptoKey().id.data(),
                         golden_data_.CTRCryptoKey().id.size(),
                         golden_data_.CTRCryptoKey().content->ciphertext.data(),
                         golden_data_.CTRCryptoKey().content->ciphertext.size(),
                         golden_data_.CTRCryptoKey().content->iv.data(),
                         golden_data_.CTRCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_OK);
  plaintext_.resize(plaintext_size_);
  ASSERT_EQ(plaintext_, golden_data_.CTRCryptoKey().content->plaintext);
}

// Try decrypting two different sets of content to make sure that two
// different keys can be used at the same time.
TEST_F(LicenseWhiteboxDecryptTest, SuccessWithMultipleKeys) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_OK);
  plaintext_.resize(plaintext_size_);
  ASSERT_EQ(plaintext_, golden_data_.CBCCryptoKey().content->plaintext);

  // Reset our output buffer.
  plaintext_.clear();
  plaintext_size_ = golden_data_.CTRDecodeKey().content->plaintext.size();
  plaintext_.resize(plaintext_size_);

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CTR,
                         golden_data_.CTRCryptoKey().id.data(),
                         golden_data_.CTRCryptoKey().id.size(),
                         golden_data_.CTRCryptoKey().content->ciphertext.data(),
                         golden_data_.CTRCryptoKey().content->ciphertext.size(),
                         golden_data_.CTRCryptoKey().content->iv.data(),
                         golden_data_.CTRCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_OK);

  plaintext_.resize(plaintext_size_);
  ASSERT_EQ(plaintext_, golden_data_.CTRCryptoKey().content->plaintext);
}

TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForNullWhitebox) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(nullptr, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForInvalidCipherMode) {
  LoadLicense(LicenseBuilder::NoPadding());

  // In order to trick the compiler into letting us pass an invalid enum value
  // to WB__License_Decrypt(), we need to cast it. If we don't do this, the
  // compiler tries to save us.
  const WB_CipherMode invalid_mode = static_cast<WB_CipherMode>(0xFF);

  ASSERT_EQ(WB_License_Decrypt(
                whitebox_, invalid_mode, golden_data_.CBCCryptoKey().id.data(),
                golden_data_.CBCCryptoKey().id.size(),
                golden_data_.CBCCryptoKey().content->ciphertext.data(),
                golden_data_.CBCCryptoKey().content->ciphertext.size(),
                golden_data_.CBCCryptoKey().content->iv.data(),
                golden_data_.CBCCryptoKey().content->iv.size(),
                plaintext_.data(), &plaintext_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForNullKeyId) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC, nullptr,
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForZeroKeyIdSize) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(), 0,
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForNullInputData) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(), nullptr,
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_INVALID_PARAMETER);
}

// AES CBC requires that the input be block aligned (multiple of 16). CTR does
// not care.
TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForInvalidCBCInputDataSize) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         14, golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_INVALID_PARAMETER);
}

// The white-box (using any cipher mode) should reject input with size zero.
TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForZeroInputDataSize) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         0, golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForNullIV) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(
          whitebox_, WB_CIPHER_MODE_CBC, golden_data_.CBCCryptoKey().id.data(),
          golden_data_.CBCCryptoKey().id.size(),
          golden_data_.CBCCryptoKey().content->ciphertext.data(),
          golden_data_.CBCCryptoKey().content->ciphertext.size(), nullptr,
          golden_data_.CBCCryptoKey().content->iv.size(), plaintext_.data(),
          &plaintext_size_),
      WB_RESULT_INVALID_PARAMETER);
}

// IV size should be 16. Any number other than 16 should fail.
TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForInvalidIVSize) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(), 9,
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForNullOutput) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         nullptr, &plaintext_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxDecryptTest, InvalidParameterForNullOutputSize) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), nullptr),
      WB_RESULT_INVALID_PARAMETER);
}

// For this test, "missing key id" specifically means a key id that was never
// in the license to start with. This is different than "non content key"
// and "dropped content key", as those keys were in the license but ignored.
TEST_F(LicenseWhiteboxDecryptTest, KeyUnavailableForMissingKeyId) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC, missing_key_id_.data(),
                         missing_key_id_.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_KEY_UNAVAILABLE);
}

TEST_F(LicenseWhiteboxDecryptTest, KeyUnavailableForNonContentKey) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         non_content_key_id_.data(), non_content_key_id_.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_KEY_UNAVAILABLE);
}

// Under normal circumstances, a hardware key should be dropped. The exception
// to this rule is on ChromeOS with a special license.
TEST_F(LicenseWhiteboxDecryptTest,
       InsufficientSecurityLevelForHardwareContentKey) {
  LoadLicense(LicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_Decrypt(
                whitebox_, WB_CIPHER_MODE_CBC,
                golden_data_.CBCHardwareKey().id.data(),
                golden_data_.CBCHardwareKey().id.size(),
                golden_data_.CBCHardwareKey().content->ciphertext.data(),
                golden_data_.CBCHardwareKey().content->ciphertext.size(),
                golden_data_.CBCHardwareKey().content->iv.data(),
                golden_data_.CBCHardwareKey().content->iv.size(),
                plaintext_.data(), &plaintext_size_),
            WB_RESULT_INSUFFICIENT_SECURITY_LEVEL);
}

TEST_F(LicenseWhiteboxDecryptTest, InsufficientSecurityLevelForDecodeKey) {
  LoadLicense(LicenseBuilder::NoPadding());

  // Use the software decode key as they are limited to
  // WB_License_Decrypt().
  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCDecodeKey().id.data(),
                         golden_data_.CBCDecodeKey().id.size(),
                         golden_data_.CBCDecodeKey().content->ciphertext.data(),
                         golden_data_.CBCDecodeKey().content->ciphertext.size(),
                         golden_data_.CBCDecodeKey().content->iv.data(),
                         golden_data_.CBCDecodeKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_INSUFFICIENT_SECURITY_LEVEL);
}

TEST_F(LicenseWhiteboxDecryptTest, BufferTooSmall) {
  LoadLicense(LicenseBuilder::NoPadding());

  // Our ciphertext will be large enough that we should not need to worry about
  // using a constant here.
  plaintext_size_ = 8;

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_BUFFER_TOO_SMALL);

  // We don't use padding so the reported plaintext size should be the same as
  // the cipher text size.
  ASSERT_EQ(plaintext_size_,
            golden_data_.CBCCryptoKey().content->ciphertext.size());
}

TEST_F(LicenseWhiteboxDecryptTest, InvalidState) {
  // Unlike the other tests, we do not call LoadLicense() as the criteria for
  // WB_RESULT_INVALID_STATE is that no key can be found and keys are provided
  // via a license.

  ASSERT_EQ(
      WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                         golden_data_.CBCCryptoKey().id.data(),
                         golden_data_.CBCCryptoKey().id.size(),
                         golden_data_.CBCCryptoKey().content->ciphertext.data(),
                         golden_data_.CBCCryptoKey().content->ciphertext.size(),
                         golden_data_.CBCCryptoKey().content->iv.data(),
                         golden_data_.CBCCryptoKey().content->iv.size(),
                         plaintext_.data(), &plaintext_size_),
      WB_RESULT_INVALID_STATE);
}

}  // namespace widevine