// Copyright 2019 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine Master
// 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 <algorithm>
#include <chrono>
#include <future>
#include <string>
#include <thread>
#include <vector>

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

#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 {

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

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

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

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

}  // namespace

class ParallelCdmTest : public WvCdmTestBase,
                        public ::testing::WithParamInterface<int> {
 public:
  ParallelCdmTest()
      : dummy_engine_metrics_(new metrics::EngineMetrics()),
        cdm_engine_(&file_system_, dummy_engine_metrics_) {}

  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 std::string& url, const CdmKeyRequest& key_request,
                      std::string* key_response) {
    UrlRequest url_request(url);
    ASSERT_TRUE(url_request.is_connected());

    std::string http_response;
    url_request.PostRequest(key_request.message);
    ASSERT_TRUE(url_request.GetResponse(&http_response));
    int status_code = url_request.GetStatusCode(http_response);
    ASSERT_EQ(kHttpOk, status_code);

    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) {
    CdmDecryptionParameters params(&key_id, input.data(), input.size(), &iv, 0,
                                   output->data());
    params.is_secure = false;
    CdmResponseType status = cdm_engine_.Decrypt(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);

    const int session_count = std::min(kThreadCount, GetMaxNumberOfSessions());
    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));
    }
  }

  FileSystem file_system_;
  shared_ptr<metrics::EngineMetrics> dummy_engine_metrics_;
  CdmEngine cdm_engine_;
};

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;
    ASSERT_NO_FATAL_FAILURE(GetKeyResponse(url, key_request, &key_response));

    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(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(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_CASE_P(Repeated, ParallelCdmTest,
                        ::testing::Range(0, kRepetitions));

}  // namespace wvcdm