// Copyright 2019 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.

// These tests perform various end-to-end actions similar to what an application
// would, but in parallel, attempting to create as many collisions in the CDM
// code as possible.

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <chrono>
#include <future>
#include <map>
#include <string>
#include <thread>
#include <vector>

#include "cdm_engine.h"
#include "config_test_env.h"
#include "initialization_data.h"
#include "license_request.h"
#include "log.h"
#include "metrics_collections.h"
#include "test_base.h"
#include "test_printers.h"
#include "url_request.h"
#include "wv_cdm_constants.h"
#include "wv_cdm_types.h"

namespace wvcdm {

namespace {

using HttpHeaderFields = std::map<std::string, std::string>;

constexpr const int kHttpOk = 200;
const std::string kCencMimeType = "cenc";

constexpr const auto kMaxThreadExecutionTime = std::chrono::seconds(60);
constexpr const auto kMinimumWait = std::chrono::nanoseconds(1);

constexpr const int kRepetitions = 10;
constexpr const int kThreadCount = 24;

// Number of attempts to request a license key from the license server
// before failing.
constexpr size_t kMaxKeyRequestAttempts = 5;
// Wait time between failed key requests.
constexpr const auto kRequestRetryWait = std::chrono::milliseconds(10);

const std::vector<uint8_t> kKeyId =
    wvutil::a2b_hex("371ea35e1a985d75d198a7f41020dc23");
const std::vector<uint8_t> kIv =
    wvutil::a2b_hex("cedc47cccd6cb437af41325953c2e5e0");
const std::vector<uint8_t> kEncryptedData = wvutil::a2b_hex(
    "6274b3237fe5991ff570ca902e4b3306"
    "8be1e782cc97944686223fe6a2bc0936"
    "71644c1963ffc465b8b5731a9914ce26"
    "1fb8c2dc07a723755040011aafec883d"
    "8f77a6194674e61888803d0d0b5a6670"
    "9b92b79ee5a1ccaa5a6edd290b994657"
    "b201e2fc");
const std::vector<uint8_t> kClearData = wvutil::a2b_hex(
    "21fd60f6e690ff0b0cb5f89540380f92"
    "ca7a3c4326110261d1f88ab33af1e9a3"
    "dc12574b7f55bdac821ddbacfe2f2158"
    "62b8b1e9c3c7db4e49c6a8e9fa5a7780"
    "1bb11279670daf907e21432bc66fa21a"
    "5fe539268cc115829b69b613695c961a"
    "765e9820");

}  // namespace

class ParallelCdmTest : public WvCdmTestBaseWithEngine,
                        public ::testing::WithParamInterface<int> {
 public:
  void SetUp() override {
    WvCdmTestBase::SetUp();
    EnsureProvisioned();
  }

 protected:
  int GetMaxNumberOfSessions() {
    std::string max;
    CdmResponseType res = cdm_engine_.QueryStatus(
        kLevelDefault, QUERY_KEY_MAX_NUMBER_OF_SESSIONS, &max);
    EXPECT_EQ(NO_ERROR, res);
    if (res == NO_ERROR) {
      return std::stoi(max);
    } else {
      return 0;
    }
  }

  // Note that CdmEngine expects its caller to make sure |OpenSession| is not
  // called simultaneously with any other session functions. For this reason,
  // calling this function from a parallel thread is inadvisable.
  void OpenSession(CdmSessionId* session_id) {
    CdmResponseType status = cdm_engine_.OpenSession(
        config_.key_system(), nullptr, nullptr, session_id);
    ASSERT_EQ(NO_ERROR, status);
    ASSERT_TRUE(cdm_engine_.IsOpenSession(*session_id));
  }

  void GenerateKeyRequest(const CdmSessionId& session_id,
                          const InitializationData& init_data,
                          CdmKeyRequest* key_request) {
    CdmAppParameterMap empty_app_parameters;
    CdmKeySetId empty_key_set_id;

    CdmResponseType result = cdm_engine_.GenerateKeyRequest(
        session_id, empty_key_set_id, init_data, kLicenseTypeStreaming,
        empty_app_parameters, key_request);
    ASSERT_EQ(KEY_MESSAGE, result);
    ASSERT_EQ(kKeyRequestTypeInitial, key_request->type);
  }

  void GetKeyResponse(const CdmSessionId& session_id, const std::string& url,
                      const CdmKeyRequest& key_request,
                      std::string* key_response) {
    bool request_ok = false;
    std::string http_response;
    for (size_t attempt = 1; attempt <= kMaxKeyRequestAttempts; ++attempt) {
      UrlRequest url_request(url);
      ASSERT_TRUE(url_request.is_connected());
      url_request.PostRequest(key_request.message);
      if (url_request.GetResponse(&http_response)) {
        const int status_code = url_request.GetStatusCode(http_response);
        if (status_code == kHttpOk) {
          request_ok = true;
          break;
        }
        LOGW("License response error: code = %d, url = %s", status_code,
             url.c_str());
        HttpHeaderFields fields;
        if (url_request.GetDebugHeaderFields(http_response, &fields)) {
          for (auto field : fields) {
            LOGD("  %s: %s", field.first.c_str(), field.second.c_str());
          }
        }
      }
      LOGW("License request failed: sid = %s, attempt = %zu",
           session_id.c_str(), attempt);
      std::this_thread::sleep_for(kRequestRetryWait);
    }
    ASSERT_TRUE(request_ok);
    LicenseRequest license_request;
    license_request.GetDrmMessage(http_response, *key_response);
  }

  void AddKey(const CdmSessionId& session_id, const std::string& key_response) {
    CdmKeySetId ignored_key_set_id;
    CdmLicenseType license_type;
    CdmResponseType status = cdm_engine_.AddKey(
        session_id, key_response, &license_type, &ignored_key_set_id);
    ASSERT_EQ(KEY_ADDED, status);
    ASSERT_EQ(kLicenseTypeStreaming, license_type);
  }

  void Decrypt(const CdmSessionId& session_id, const KeyId& key_id,
               const std::vector<uint8_t>& input,
               const std::vector<uint8_t>& iv, std::vector<uint8_t>* output) {
    CdmDecryptionParametersV16 params(key_id);
    params.is_secure = false;
    CdmDecryptionSample sample(input.data(), output->data(), 0, input.size(),
                               iv);
    CdmDecryptionSubsample subsample(0, input.size());
    sample.subsamples.push_back(subsample);
    params.samples.push_back(sample);

    CdmResponseType status = cdm_engine_.DecryptV16(session_id, params);
    ASSERT_EQ(NO_ERROR, status);
  }

  // Note that CdmEngine expects its caller to make sure |CloseSession| is not
  // called simultaneously with any other session functions. For this reason,
  // calling this function from a parallel thread is inadvisable.
  void CloseSession(const CdmSessionId& session_id) {
    CdmResponseType status = cdm_engine_.CloseSession(session_id);
    ASSERT_EQ(NO_ERROR, status);
    ASSERT_FALSE(cdm_engine_.IsOpenSession(session_id));
  }

  template <class Function>
  void RunThreadsSimultaneously(Function do_work) {
    std::atomic<int> threads_waiting(0);
    std::vector<std::future<void>> threads;
    threads.reserve(kThreadCount);

    for (int thread_number = 0; thread_number < kThreadCount; ++thread_number) {
      threads.push_back(std::async(std::launch::async, [&] {
        ++threads_waiting;
        while (threads_waiting < kThreadCount) {
          std::this_thread::sleep_for(kMinimumWait);
        }

        do_work();
      }));
    }

    for (std::future<void>& future : threads) {
      EXPECT_EQ(std::future_status::ready,
                future.wait_for(kMaxThreadExecutionTime));
    }
  }

  template <class Function>
  void RunSessionThreadsSimultaneously(Function do_work) {
    std::atomic<int> threads_waiting(0);
    // The OEMCrypto V16 adapter makes use of one of the sessions,
    // reducing the maximum number of sessions available for the test.
    const int session_count =
        std::min(kThreadCount, GetMaxNumberOfSessions() - 1);
    LOGI("Running %d Threads", session_count);
    std::vector<CdmSessionId> sessions(session_count);
    std::vector<std::future<void>> threads;
    threads.reserve(sessions.size());

    for (int i = 0; i < session_count; ++i) {
      ASSERT_NO_FATAL_FAILURE(OpenSession(&sessions[i]));
    }

    for (const CdmSessionId& session_id : sessions) {
      threads.push_back(std::async(std::launch::async, [&, session_id] {
        ++threads_waiting;
        while (threads_waiting < session_count) {
          std::this_thread::sleep_for(kMinimumWait);
        }

        do_work(session_id);
      }));
    }

    for (std::future<void>& future : threads) {
      EXPECT_EQ(std::future_status::ready,
                future.wait_for(kMaxThreadExecutionTime));
    }

    for (const CdmSessionId& session_id : sessions) {
      ASSERT_NO_FATAL_FAILURE(CloseSession(session_id));
    }
  }
};

TEST_P(ParallelCdmTest, ParallelLicenseRequests) {
  const std::string url = config_.license_server() + config_.client_auth();
  const InitializationData init_data(kCencMimeType, binary_key_id());

  RunSessionThreadsSimultaneously([&](const CdmSessionId& session_id) {
    CdmKeyRequest key_request;
    ASSERT_NO_FATAL_FAILURE(
        GenerateKeyRequest(session_id, init_data, &key_request));

    std::string key_response;
    LOGD("Getting license request: sid = %s", session_id.c_str());
    ASSERT_NO_FATAL_FAILURE(
        GetKeyResponse(session_id, url, key_request, &key_response))
        << "SID: " << session_id;

    ASSERT_NO_FATAL_FAILURE(AddKey(session_id, key_response));
  });
}

TEST_P(ParallelCdmTest, ParallelDecryptSessions) {
  const std::string url = config_.license_server() + config_.client_auth();
  const InitializationData init_data(kCencMimeType, binary_key_id());
  const KeyId key_id(kKeyId.cbegin(), kKeyId.cend());

  RunSessionThreadsSimultaneously([&](const CdmSessionId& session_id) {
    CdmKeyRequest key_request;
    ASSERT_NO_FATAL_FAILURE(
        GenerateKeyRequest(session_id, init_data, &key_request));

    std::string key_response;
    ASSERT_NO_FATAL_FAILURE(
        GetKeyResponse(session_id, url, key_request, &key_response));

    ASSERT_NO_FATAL_FAILURE(AddKey(session_id, key_response));

    std::vector<uint8_t> output;
    output.resize(kEncryptedData.size());
    ASSERT_NO_FATAL_FAILURE(
        Decrypt(session_id, key_id, kEncryptedData, kIv, &output));
    EXPECT_EQ(kClearData, output);
  });
}

TEST_P(ParallelCdmTest, ParallelDecryptsInSameSession) {
  const std::string url = config_.license_server() + config_.client_auth();
  const InitializationData init_data(kCencMimeType, binary_key_id());
  const KeyId key_id(kKeyId.cbegin(), kKeyId.cend());

  CdmSessionId session_id;
  ASSERT_NO_FATAL_FAILURE(OpenSession(&session_id));

  CdmKeyRequest key_request;
  ASSERT_NO_FATAL_FAILURE(
      GenerateKeyRequest(session_id, init_data, &key_request));

  std::string key_response;
  ASSERT_NO_FATAL_FAILURE(
      GetKeyResponse(session_id, url, key_request, &key_response));

  ASSERT_NO_FATAL_FAILURE(AddKey(session_id, key_response));

  RunThreadsSimultaneously([&] {
    std::vector<uint8_t> output;
    output.resize(kEncryptedData.size());
    ASSERT_NO_FATAL_FAILURE(
        Decrypt(session_id, key_id, kEncryptedData, kIv, &output));
    EXPECT_EQ(kClearData, output);
  });

  ASSERT_NO_FATAL_FAILURE(CloseSession(session_id));
}

INSTANTIATE_TEST_SUITE_P(Repeated, ParallelCdmTest,
                         ::testing::Range(0, kRepetitions));

}  // namespace wvcdm