// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox.h"

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

#include "api/golden_data.h"
#include "api/license_whitebox_test_base.h"
#include "api/test_data.h"
#include "api/test_license_builder.h"
#include "base/logging.h"
#include "crypto_utils/crypto_util.h"
#include "crypto_utils/rsa_key.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace widevine {

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

    // Because we are going to use the same buffer for both tests, make sure it
    // will be large enough for either.
    masked_text_size_ = std::max(golden_data_.CBCContent().ciphertext.size(),
                                 golden_data_.CTRContent().ciphertext.size());
    masked_text_.resize(masked_text_size_);

    // We have no idea how big the secret string will be, but it should be safe
    // to assume it won't be larger than the plaintext.
    secret_string_size_ = masked_text_size_;
    secret_string_.resize(secret_string_size_);

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

  void LoadLicense(const std::vector<uint8_t>& padding) {
    TestLicenseBuilder 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_.CTRDecodeKey().level,
                          golden_data_.CTRDecodeKey().id,
                          golden_data_.CTRDecodeKey().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};

  size_t secret_string_size_;
  std::vector<uint8_t> secret_string_;

  size_t masked_text_size_;
  std::vector<uint8_t> masked_text_;

  std::vector<uint8_t> plaintext_;
};

TEST_F(LicenseWhiteboxMaskedDecryptTest, DecodeKeyWithCbcDataInCbcMode) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Returned data is masked, so it should be the correct size but not
  // match the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CBCDecodeKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CBCDecodeKey().content->plaintext);

  // Now unmask the data.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCDecodeKey().id.data(),
                                 golden_data_.CBCDecodeKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_EQ(plaintext_, golden_data_.CBCDecodeKey().content->plaintext);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, DecodeKeyWithCtrDataInCtrMode) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CTR, golden_data_.CTRDecodeKey().id.data(),
          golden_data_.CTRDecodeKey().id.size(),
          golden_data_.CTRDecodeKey().content->ciphertext.data(),
          golden_data_.CTRDecodeKey().content->ciphertext.size(),
          golden_data_.CTRDecodeKey().content->iv.data(),
          golden_data_.CTRDecodeKey().content->iv.size(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Returned data is masked, so it should be the correct size but not
  // match the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CTRDecodeKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CTRDecodeKey().content->plaintext);

  // Now unmask the data.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CTR,
                                 golden_data_.CTRDecodeKey().id.data(),
                                 golden_data_.CTRDecodeKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_EQ(plaintext_, golden_data_.CTRDecodeKey().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(LicenseWhiteboxMaskedDecryptTest, DecodeKeyWithCbcDataInCtrMode) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CTR, 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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Whatever is returned must not be the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CBCDecodeKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CBCDecodeKey().content->plaintext);

  // Now unmask the data. Still should not match.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CTR,
                                 golden_data_.CBCDecodeKey().id.data(),
                                 golden_data_.CBCDecodeKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_NE(masked_text_, golden_data_.CBCDecodeKey().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(LicenseWhiteboxMaskedDecryptTest, DecodeKeyWithCtrDataInCbcMode) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CBC, golden_data_.CTRDecodeKey().id.data(),
          golden_data_.CTRDecodeKey().id.size(),
          golden_data_.CTRDecodeKey().content->ciphertext.data(),
          golden_data_.CTRDecodeKey().content->ciphertext.size(),
          golden_data_.CTRDecodeKey().content->iv.data(),
          golden_data_.CTRDecodeKey().content->iv.size(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Whatever is returned must not be the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CTRDecodeKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CTRDecodeKey().content->plaintext);

  // Now unmask the data. Still should not match.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CTRDecodeKey().id.data(),
                                 golden_data_.CTRDecodeKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_NE(masked_text_, golden_data_.CTRDecodeKey().content->plaintext);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, CryptoKeyWithCbcDataInCbcMode) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CBC, golden_data_.CBCCryptoKey().id.data(),
          golden_data_.CBCCryptoKey().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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Returned data is masked, so it should be the correct size but not
  // match the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CBCDecodeKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CBCDecodeKey().content->plaintext);

  // Now unmask the data.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCCryptoKey().id.data(),
                                 golden_data_.CBCCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_EQ(plaintext_, golden_data_.CBCDecodeKey().content->plaintext);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, CryptoKeyWithCtrDataInCtrMode) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CTR, golden_data_.CTRCryptoKey().id.data(),
          golden_data_.CTRCryptoKey().id.size(),
          golden_data_.CTRDecodeKey().content->ciphertext.data(),
          golden_data_.CTRDecodeKey().content->ciphertext.size(),
          golden_data_.CTRDecodeKey().content->iv.data(),
          golden_data_.CTRDecodeKey().content->iv.size(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Returned data is masked, so it should be the correct size but not
  // match the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CTRDecodeKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CTRDecodeKey().content->plaintext);

  // Now unmask the data.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CTR,
                                 golden_data_.CTRCryptoKey().id.data(),
                                 golden_data_.CTRCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_EQ(plaintext_, golden_data_.CTRDecodeKey().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(LicenseWhiteboxMaskedDecryptTest, CryptoKeyWithCbcDataInCtrMode) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CTR, golden_data_.CBCCryptoKey().id.data(),
          golden_data_.CBCCryptoKey().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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Whatever is returned must not be the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CBCDecodeKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CBCDecodeKey().content->plaintext);

  // Now unmask the data. Still should not match.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CTR,
                                 golden_data_.CBCCryptoKey().id.data(),
                                 golden_data_.CBCCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_NE(masked_text_, golden_data_.CBCDecodeKey().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(LicenseWhiteboxMaskedDecryptTest, CryptoKeyWithCtrDataInCbcMode) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Whatever is returned must not be the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CTRCryptoKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CTRCryptoKey().content->plaintext);

  // Now unmask the data. Still should not match.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CTRCryptoKey().id.data(),
                                 golden_data_.CTRCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_NE(masked_text_, golden_data_.CTRCryptoKey().content->plaintext);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, CryptoKeyWithCbcDataAndPKCS8Padding) {
  LoadLicense(TestLicenseBuilder::PKSC8Padding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Returned data is masked, so it should be the correct size but not
  // match the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CBCCryptoKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CBCCryptoKey().content->plaintext);

  // Now unmask the data.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCCryptoKey().id.data(),
                                 golden_data_.CBCCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_EQ(plaintext_, golden_data_.CBCCryptoKey().content->plaintext);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, CryptoKeyWithCtrDataAndPKCS8Padding) {
  LoadLicense(TestLicenseBuilder::PKSC8Padding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Returned data is masked, so it should be the correct size but not
  // match the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CTRCryptoKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CTRCryptoKey().content->plaintext);

  // Now unmask the data.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CTR,
                                 golden_data_.CTRCryptoKey().id.data(),
                                 golden_data_.CTRCryptoKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  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.
//
// Since we have two CBC keys, try using the decode key and then the crypto
// key.
TEST_F(LicenseWhiteboxMaskedDecryptTest, SuccessWithMultipleKeys) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Returned data is masked, so it should be the correct size but not
  // match the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CBCDecodeKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CBCDecodeKey().content->plaintext);

  // Now unmask the data.
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCDecodeKey().id.data(),
                                 golden_data_.CBCDecodeKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_EQ(plaintext_, golden_data_.CBCDecodeKey().content->plaintext);

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

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CTR, golden_data_.CTRDecodeKey().id.data(),
          golden_data_.CTRDecodeKey().id.size(),
          golden_data_.CTRDecodeKey().content->ciphertext.data(),
          golden_data_.CTRDecodeKey().content->ciphertext.size(),
          golden_data_.CTRDecodeKey().content->iv.data(),
          golden_data_.CTRDecodeKey().content->iv.size(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);
  masked_text_.resize(masked_text_size_);

  // Returned data is masked, so it should be the correct size but not
  // match the original text.
  ASSERT_EQ(masked_text_.size(),
            golden_data_.CTRDecodeKey().content->plaintext.size());
  ASSERT_NE(masked_text_, golden_data_.CTRDecodeKey().content->plaintext);

  // Now unmask the data.
  secret_string_.clear();
  secret_string_size_ = masked_text_.size();
  secret_string_.resize(secret_string_size_);
  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CTR,
                                 golden_data_.CTRDecodeKey().id.data(),
                                 golden_data_.CTRDecodeKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  plaintext_.resize(masked_text_size_);
  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    plaintext_.data());

  ASSERT_EQ(plaintext_, golden_data_.CTRDecodeKey().content->plaintext);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, InvalidParameterForNullWhitebox) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          nullptr, 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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, InvalidParameterForInvalidCipherMode) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  // In order to trick the compiler into letting us pass an invalid enum value
  // to WB__License_MaskedDecrypt(), 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_MaskedDecrypt(
                whitebox_, invalid_mode, 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(),
                masked_text_.data(), &masked_text_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, InvalidParameterForNullKeyId) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_MaskedDecrypt(
                whitebox_, WB_CIPHER_MODE_CBC, nullptr,
                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(),
                masked_text_.data(), &masked_text_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, InvalidParameterForNullZeroKeyIdSize) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CBC, golden_data_.CBCDecodeKey().id.data(),
          0, golden_data_.CBCDecodeKey().content->ciphertext.data(),
          golden_data_.CBCDecodeKey().content->ciphertext.size(),
          golden_data_.CBCDecodeKey().content->iv.data(),
          golden_data_.CBCDecodeKey().content->iv.size(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, InvalidParameterForNullInputData) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CBC, golden_data_.CBCDecodeKey().id.data(),
          golden_data_.CBCDecodeKey().id.size(), nullptr,
          golden_data_.CBCDecodeKey().content->ciphertext.size(),
          golden_data_.CBCDecodeKey().content->iv.data(),
          golden_data_.CBCDecodeKey().content->iv.size(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_INVALID_PARAMETER);
}

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

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CBC, golden_data_.CBCDecodeKey().id.data(),
          golden_data_.CBCDecodeKey().id.size(),
          golden_data_.CBCDecodeKey().content->ciphertext.data(), 14,
          golden_data_.CBCDecodeKey().content->iv.data(),
          golden_data_.CBCDecodeKey().content->iv.size(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_INVALID_PARAMETER);
}

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

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          whitebox_, WB_CIPHER_MODE_CBC, golden_data_.CBCDecodeKey().id.data(),
          golden_data_.CBCDecodeKey().id.size(),
          golden_data_.CBCDecodeKey().content->ciphertext.data(), 0,
          golden_data_.CBCDecodeKey().content->iv.data(),
          golden_data_.CBCDecodeKey().content->iv.size(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, InvalidParameterForNullIV) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), nullptr,
          golden_data_.CBCDecodeKey().content->iv.size(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_INVALID_PARAMETER);
}

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

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), 9,
          masked_text_.data(), &masked_text_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, InvalidParameterForNullOutput) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), nullptr,
          &masked_text_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, InvalidParameterForNullOutputSize) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), masked_text_.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(LicenseWhiteboxMaskedDecryptTest, KeyUnavailableForMissingKeyId) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_MaskedDecrypt(
                whitebox_, WB_CIPHER_MODE_CBC, missing_key_id_.data(),
                missing_key_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(),
                masked_text_.data(), &masked_text_size_),
            WB_RESULT_KEY_UNAVAILABLE);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, KeyUnavailableForNonContentKey) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_MaskedDecrypt(
                whitebox_, WB_CIPHER_MODE_CBC, non_content_key_id_.data(),
                non_content_key_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(),
                masked_text_.data(), &masked_text_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(LicenseWhiteboxMaskedDecryptTest,
       InsufficientSecurityLevelForHardwareContentKey) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(WB_License_MaskedDecrypt(
                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(),
                masked_text_.data(), &masked_text_size_),
            WB_RESULT_INSUFFICIENT_SECURITY_LEVEL);
}

// 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.
TEST_F(LicenseWhiteboxMaskedDecryptTest, InvalidState) {
  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_INVALID_STATE);
}

TEST_F(LicenseWhiteboxMaskedDecryptTest, BufferTooSmall) {
  LoadLicense(TestLicenseBuilder::NoPadding());

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

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), masked_text_.data(),
          &masked_text_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(masked_text_size_,
            golden_data_.CBCDecodeKey().content->ciphertext.size());
}

// Check that the result of unmasking only a small portion of the data is the
// same as when we unmask the whole buffer.
TEST_F(LicenseWhiteboxMaskedDecryptTest, SuccessForSubRangeUnmask) {
  LoadLicense(TestLicenseBuilder::NoPadding());

  ASSERT_EQ(
      WB_License_MaskedDecrypt(
          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(), masked_text_.data(),
          &masked_text_size_),
      WB_RESULT_OK);

  ASSERT_EQ(
      WB_License_GetSecretString(whitebox_, WB_CIPHER_MODE_CBC,
                                 golden_data_.CBCDecodeKey().id.data(),
                                 golden_data_.CBCDecodeKey().id.size(),
                                 secret_string_.data(), &secret_string_size_),
      WB_RESULT_OK);
  secret_string_.resize(secret_string_size_);

  std::vector<uint8_t> full_unmask(masked_text_size_);
  std::vector<uint8_t> partial_unmask(3);

  WB_License_Unmask(masked_text_.data(), 0, masked_text_size_,
                    secret_string_.data(), secret_string_.size(),
                    full_unmask.data());
  WB_License_Unmask(masked_text_.data(), 4, partial_unmask.size(),
                    secret_string_.data(), secret_string_.size(),
                    partial_unmask.data());

  ASSERT_EQ(full_unmask[4], partial_unmask[0]);
  ASSERT_EQ(full_unmask[5], partial_unmask[1]);
  ASSERT_EQ(full_unmask[6], partial_unmask[2]);
}
}  // namespace widevine