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

#ifndef OEMCRYPTO_FUZZ_HELPER_H_
#define OEMCRYPTO_FUZZ_HELPER_H_

#include <cstring>
#include <type_traits>
#include <vector>

#include "FuzzedDataProvider.h"
#include "oec_device_features.h"
#include "oemcrypto_corpus_generator_helper.h"
#include "oemcrypto_session_tests_helper.h"

// Forward-declare the libFuzzer's mutator callback. Mark it weak so that
// the program links successfully even outside of --config=asan-fuzzer
// (apparently the only config in which LLVM uses our custom mutator).
extern "C" size_t LLVMFuzzerMutate(uint8_t* Data, size_t Size, size_t MaxSize)
    __attribute__((weak));

namespace wvoec {

constexpr size_t KB = 1024;

// Default maximum length of fuzzing output parameters.
constexpr size_t MAX_FUZZ_OUTPUT_LENGTH = 5 * KB;

// Fuzzed data region.
struct FuzzedData {
  const uint8_t* data;
  size_t size;
};

// Initial setup to create a valid OEMCrypto state such as initializing crypto
// firmware/hardware, installing golden key box etc. in order to fuzz
// OEMCrypto APIs.
void InitializeFuzz(SessionUtil& session_util);

class SessionFuzz {
 public:
  void Initialize();

  void Terminate();

  void InstallTestDrmKey() {
    session_util_.InstallTestDrmKey(&session_);
  }

  Session& session() { return session_; }

  const Session& session() const { return session_; }

 private:
  SessionUtil session_util_;
  Session session_;
};

class OEMCryptoLicenseAPIFuzz {
 public:
  OEMCryptoLicenseAPIFuzz() : license_messages_(&session_fuzz_.session()) {}

  void Initialize();

  void Terminate();

  void LoadLicense() { LoadLicense(false); }

  void LoadLicenseWithGenericCryptoKeys() { LoadLicense(true); }

  LicenseRoundTrip& license_messages() { return license_messages_; }

  const LicenseRoundTrip& license_messages() const { return license_messages_; }

  Session& session() { return session_fuzz_.session(); }

  const Session& session() const { return session_fuzz_.session(); }

 private:
  void LoadLicense(bool generic_crypto_keys);

  SessionFuzz session_fuzz_;
  LicenseRoundTrip license_messages_;
};

class OEMCryptoProvisioningAPIFuzz {
 public:
  OEMCryptoProvisioningAPIFuzz()
      : provisioning_messages_(&session_, session_util_.encoded_rsa_key_) {}

  void Initialize();

  void Terminate();

  void LoadProvisioning();

  ProvisioningRoundTrip& provisioning_messages() {
    return provisioning_messages_;
  }

  const ProvisioningRoundTrip& provisioning_messages() const {
    return provisioning_messages_;
  }

  Session& session() { return session_; }

  const Session& session() const { return session_; }

 private:
  SessionUtil session_util_;
  Session session_;
  ProvisioningRoundTrip provisioning_messages_;
};

// Initial setup to create a valid state such as creating session, installing
// golden key box etc. in order to fuzz Load Renewal API.
class OEMCryptoRenewalAPIFuzz {
 public:
  OEMCryptoRenewalAPIFuzz()
      : renewal_messages_(&license_api_fuzz_.license_messages()) {}

  void Initialize() { license_api_fuzz_.Initialize(); }

  void Terminate() { license_api_fuzz_.Terminate(); }

  LicenseRoundTrip& license_messages() {
    return license_api_fuzz_.license_messages();
  }

  const LicenseRoundTrip& license_messages() const {
    return license_api_fuzz_.license_messages();
  }

  RenewalRoundTrip& renewal_messages() { return renewal_messages_; }

  const RenewalRoundTrip& renewal_messages() const { return renewal_messages_; }

 private:
  OEMCryptoLicenseAPIFuzz license_api_fuzz_;
  RenewalRoundTrip renewal_messages_;
};

class LicenseWithUsageEntryFuzz {
 public:
  LicenseWithUsageEntryFuzz() : license_messages_(&session_fuzz_.session()) {
    license_messages_.set_pst("my_pst");
  }

  void Initialize() { session_fuzz_.Initialize(); }

  void Terminate() { session_fuzz_.Terminate(); }

  void CreateUsageTableHeader();

  void InstallTestDrmKey() { session_fuzz_.InstallTestDrmKey(); }

  void LoadLicense();

  LicenseRoundTrip& license_messages() { return license_messages_; }

  const LicenseRoundTrip& license_messages() const { return license_messages_; }

  const std::vector<uint8_t>& encrypted_usage_header() const {
    return encrypted_usage_header_;
  }

  Session& session() { return session_fuzz_.session(); }

  const Session& session() const { return session_fuzz_.session(); }

 private:
  SessionFuzz session_fuzz_;
  LicenseRoundTrip license_messages_;
  std::vector<uint8_t> encrypted_usage_header_;
};

// Convert data from FuzzedDataProvider to valid enum value.
template <typename T>
T ConvertDataToValidEnum(FuzzedDataProvider& fuzzed_data, T max_enum_value) {
  using UnsignedT =
      typename std::make_unsigned<typename std::underlying_type<T>::type>::type;
  return static_cast<T>(fuzzed_data.ConsumeIntegralInRange<UnsignedT>(
      0, static_cast<UnsignedT>(max_enum_value)));
}

// Convert data to valid enum value in place.
template <typename T>
void ConvertDataToValidEnum(T max_enum_value, T& enum_data) {
  using UnsignedT =
      typename std::make_unsigned<typename std::underlying_type<T>::type>::type;
  UnsignedT data;
  std::memcpy(&data, &enum_data, sizeof(T));
  const auto max_value = static_cast<UnsignedT>(max_enum_value);
  if (data > max_value) {
    enum_data = static_cast<T>(data % (max_value + 1));
  }
}

// Redirect printf and log statements from oemcrypto functions to a file to
// reduce noise
void RedirectStdoutToFile();

// Split fuzzed data using delimiter "-_^_".
std::vector<FuzzedData> SplitFuzzedData(const uint8_t* data, size_t size);

// Check the status and exit fuzzer if arguments do not match. This is usually
// called to check status of APIs which are called to setup state for fuzzers.
void CheckStatusAndExitFuzzerOnFailure(OEMCryptoResult result,
                                       OEMCryptoResult expected_status);

}  // namespace wvoec

#endif  // OEMCRYPTO_FUZZ_HELPER_H_