// Copyright 2020 Google LLC. All Rights Reserved.

#include "api/license_whitebox.h"

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

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

namespace widevine {

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

    TestLicenseBuilder builder;
    builder.Build(*public_key_, &license_);

    // We must make the default size large to hold the signature returned by
    // WB_License_SignLicenseRequest().
    signature_size_ = 256;
    signature_.resize(signature_size_);
  }

  std::vector<uint8_t> invalid_license_request_ = {
      0x1e, 0x70, 0xbd, 0xeb, 0x24, 0xf2, 0x9d, 0x05, 0xc5, 0xb5,
      0xf4, 0xca, 0xe6, 0x1d, 0x01, 0x97, 0x29, 0xf4, 0xe0, 0x7c,
      0xfd, 0xcc, 0x97, 0x8d, 0xc2, 0xbb, 0x2d, 0x9b, 0x6b, 0x45,
      0x06, 0xbd, 0x2c, 0x66, 0x10, 0x42, 0x73, 0x8d, 0x88, 0x9b,
      0x18, 0xcc, 0xcb, 0x7e, 0x43, 0x23, 0x06, 0xe9, 0x8f, 0x8f,
  };

  License license_;

  // These will be the output from each test case.
  size_t signature_size_;
  std::vector<uint8_t> signature_;
};

TEST_F(LicenseWhiteboxSignLicenseRequestTest, SuccessForInvalidLicenseRequest) {
  ASSERT_EQ(
      WB_License_SignLicenseRequest(whitebox_, invalid_license_request_.data(),
                                    invalid_license_request_.size(),
                                    signature_.data(), &signature_size_),
      WB_RESULT_OK);

  signature_.resize(signature_size_);

  ASSERT_TRUE(public_key_->VerifySignature(
      std::string(invalid_license_request_.begin(),
                  invalid_license_request_.end()),
      std::string(signature_.begin(), signature_.end())));
}

TEST_F(LicenseWhiteboxSignLicenseRequestTest, SuccessForValidLicenseRequest) {
  ASSERT_EQ(WB_License_SignLicenseRequest(whitebox_, license_.request.data(),
                                          license_.request.size(),
                                          signature_.data(), &signature_size_),
            WB_RESULT_OK);

  signature_.resize(signature_size_);

  ASSERT_TRUE(public_key_->VerifySignature(
      std::string(license_.request.begin(), license_.request.end()),
      std::string(signature_.begin(), signature_.end())));
}

TEST_F(LicenseWhiteboxSignLicenseRequestTest, InvalidParameterForNullWhitebox) {
  ASSERT_EQ(WB_License_SignLicenseRequest(nullptr, license_.request.data(),
                                          license_.request.size(),
                                          signature_.data(), &signature_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignLicenseRequestTest,
       InvalidParameterForNullLicenseRequest) {
  ASSERT_EQ(
      WB_License_SignLicenseRequest(whitebox_, nullptr, license_.request.size(),
                                    signature_.data(), &signature_size_),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignLicenseRequestTest,
       InvalidParameterForZeroLicenseRequestSize) {
  ASSERT_EQ(WB_License_SignLicenseRequest(whitebox_, license_.request.data(), 0,
                                          signature_.data(), &signature_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignLicenseRequestTest,
       InvalidParameterForNullSignature) {
  ASSERT_EQ(WB_License_SignLicenseRequest(whitebox_, license_.request.data(),
                                          license_.request.size(), nullptr,
                                          &signature_size_),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignLicenseRequestTest,
       InvalidParameterForNullSignatureSize) {
  ASSERT_EQ(WB_License_SignLicenseRequest(whitebox_, license_.request.data(),
                                          license_.request.size(),
                                          signature_.data(), nullptr),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(LicenseWhiteboxSignLicenseRequestTest, BufferTooSmall) {
  // We could test zero, but zero is just a subset of "smaller" and we want to
  // make sure that they zero is not the only check. So use something larger
  // than zero.
  signature_size_ = 1;

  ASSERT_EQ(WB_License_SignLicenseRequest(whitebox_, license_.request.data(),
                                          license_.request.size(),
                                          signature_.data(), &signature_size_),
            WB_RESULT_BUFFER_TOO_SMALL);

  // When WB_RESULT_BUFFER_TOO_SMALL is returned, the required buffer size
  // should be returned via |signature_size|. Since we don't know what it is, we
  // must rely on it being larger than the original "too small" size.
  ASSERT_GT(signature_size_, 1u);
}
}  // namespace widevine