// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox.h"

#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/crypto_util.h"
#include "crypto_utils/rsa_key.h"
#include "testing/include/gtest/gtest.h"

namespace widevine {

class LicenseWhiteboxProcessLicenseResponseTest
    : public LicenseWhiteboxTestBase {
 protected:
  void UseLicenseWithoutSigningKey() {
    LicenseBuilder builder;
    builder.AddStubbedContentKey();
    builder.Build(*public_key_, &license_);
  }

  void UseLicenseWithSigningKey(const std::vector<uint8_t>& padding) {
    LicenseBuilder builder;
    builder.AddSigningKey(LicenseBuilder::DefaultSigningKey(), padding);
    builder.AddStubbedContentKey();
    builder.Build(*public_key_, &license_);
  }

  License license_;
};

TEST_F(LicenseWhiteboxProcessLicenseResponseTest, SuccessWithoutSigningKey) {
  UseLicenseWithoutSigningKey();

  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);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseTest,
       SuccessWithSigningKeyNoPadding) {
  UseLicenseWithSigningKey(LicenseBuilder::NoPadding());

  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);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseTest,
       SuccessWithSigningKeyPKSC8Padding) {
  UseLicenseWithSigningKey(LicenseBuilder::PKSC8Padding());

  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);
}

class LicenseWhiteboxProcessLicenseResponseErrorTest
    : public LicenseWhiteboxProcessLicenseResponseTest {
 protected:
  void SetUp() override {
    LicenseWhiteboxProcessLicenseResponseTest::SetUp();

    // For these tests, we don't care what license we use, it just needs to be
    // a valid license.
    UseLicenseWithoutSigningKey();
  }
};

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidSignatureForModifedMessage) {
  Modify(&license_.message);

  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_INVALID_SIGNATURE);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidSignatureForModifedSignature) {
  Modify(&license_.signature);

  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_INVALID_SIGNATURE);
}

// The license request is used to derive the signing key. If the request was
// modified, then the wrong signing key should be generated.
TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidSignatureForModifedLicenseRequest) {
  Modify(&license_.request);

  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_INVALID_SIGNATURE);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForNullWhitebox) {
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                nullptr, 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_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForNullMessage) {
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, nullptr, 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_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForZeroMessageSize) {
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, license_.message.data(), 0,
                license_.signature.data(), license_.signature.size(),
                license_.session_key.data(), license_.session_key.size(),
                license_.request.data(), license_.request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForNullSignature) {
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, license_.message.data(), license_.message.size(),
                nullptr, license_.signature.size(), license_.session_key.data(),
                license_.session_key.size(), license_.request.data(),
                license_.request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForInvalidSignatureSize) {
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, license_.message.data(), license_.message.size(),
                license_.signature.data(), 5, license_.session_key.data(),
                license_.session_key.size(), license_.request.data(),
                license_.request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForNullSessionKey) {
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, license_.message.data(), license_.message.size(),
                license_.signature.data(), license_.signature.size(), nullptr,
                license_.session_key.size(), license_.request.data(),
                license_.request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForInvalidSessionKeySize) {
  ASSERT_EQ(WB_License_ProcessLicenseResponse(
                whitebox_, license_.message.data(), license_.message.size(),
                license_.signature.data(), license_.signature.size(),
                license_.session_key.data(), 5, license_.request.data(),
                license_.request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

// If the session key is modified, unwrapping it will fail. Therefore, we will
// know that the parameter is invalid (compared to a modified license request).
TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForModifedSessionKey) {
  Modify(&license_.session_key);

  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_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForNullLicenseRequest) {
  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(),
                nullptr, license_.request.size()),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxProcessLicenseResponseErrorTest,
       InvalidParameterForZeroLienseRequestSize) {
  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(), 0),
            WB_RESULT_INVALID_PARAMETER);
}

class LicenseWhiteboxMultiLicenseTest
    : public LicenseWhiteboxProcessLicenseResponseTest {
 protected:
  void SetUp() override {
    LicenseWhiteboxProcessLicenseResponseTest::SetUp();

    // For these tests, we don't care what license we use, it just needs to be
    // a valid license.
    UseLicenseWithoutSigningKey();
  }
};

// A whitebox can only process a license once. If it has loaded a license
// (successfully) it should reject later calls with WB_RESULT_INVALID_STATE.
TEST_F(LicenseWhiteboxMultiLicenseTest, InvalidState) {
  // Load the first license. This one is expected to succeed as the whitebox has
  // not loaded a license yet.
  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);

  // Attempt to load the same license again. This should fail as it already has
  // a license (even though it is the same 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_INVALID_STATE);
}

// Even though a whitebox can only load a license once, if it fails to load a
// license, it should still be able to try again.
TEST_F(LicenseWhiteboxMultiLicenseTest, SuccessAfterFailure) {
  // Force this one to fail my changing the request, this will cause an error
  // in key derivation which is a later step of license parsing.
  std::vector<uint8_t> bad_request = license_.request;
  Modify(&bad_request);
  ASSERT_NE(WB_License_ProcessLicenseResponse(
                whitebox_, license_.message.data(), license_.message.size(),
                license_.signature.data(), license_.signature.size(),
                license_.session_key.data(), license_.session_key.size(),
                bad_request.data(), bad_request.size()),
            WB_RESULT_OK);

  // Attempt to load the license again, but use the correct (unmodified)
  // request.
  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);
}

}  // namespace widevine