// 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 WIDEVINE_ODK_TEST_FUZZING_ODK_FUZZ_HELPER_H_
#define WIDEVINE_ODK_TEST_FUZZING_ODK_FUZZ_HELPER_H_

#include <memory>
#include <string>

#include "core_message_features.h"
#include "core_message_serialize.h"
#include "fuzzing/odk_fuzz_structs.h"
#include "odk_attributes.h"
#include "odk_serialize.h"

namespace oemcrypto_core_message {
bool convert_byte_to_valid_boolean(const bool* in);

OEMCryptoResult odk_serialize_LicenseRequest(
    const void* in, uint8_t* out, size_t* size,
    const ODK_LicenseRequest& core_license_request,
    const ODK_NonceValues* nonce_values);

OEMCryptoResult odk_serialize_RenewalRequest(
    const void* in, uint8_t* out, size_t* size,
    const ODK_RenewalRequest& core_renewal, ODK_NonceValues* nonce_values);

OEMCryptoResult odk_serialize_ProvisioningRequest(
    const void* in, uint8_t* out, size_t* size,
    const ODK_ProvisioningRequest& core_provisioning,
    const ODK_NonceValues* nonce_values);

OEMCryptoResult odk_deserialize_LicenseResponse(const uint8_t* message,
                                                size_t core_message_length,
                                                ODK_ParseLicense_Args* a,
                                                ODK_NonceValues* nonce_values,
                                                ODK_ParsedLicense* parsed_lic);

OEMCryptoResult odk_deserialize_RenewalResponse(
    const uint8_t* buf, size_t len, ODK_ParseRenewal_Args* a,
    ODK_NonceValues* nonce_values, ODK_PreparedRenewalRequest* renewal_msg);

OEMCryptoResult odk_deserialize_ProvisioningResponse(
    const uint8_t* buf, size_t len, ODK_ParseProvisioning_Args* a,
    ODK_NonceValues* nonce_values, ODK_ParsedProvisioning* parsed_prov);

bool kdo_serialize_LicenseResponse(const ODK_ParseLicense_Args* args,
                                   const ODK_ParsedLicense& parsed_lic,
                                   std::string* oemcrypto_core_message);

bool kdo_serialize_RenewalResponse(
    const ODK_ParseRenewal_Args* args,
    const ODK_PreparedRenewalRequest& renewal_msg,
    std::string* oemcrypto_core_message);

bool kdo_serialize_ProvisioningResponse(
    const ODK_ParseProvisioning_Args* args,
    const ODK_ParsedProvisioning& parsed_prov,
    std::string* oemcrypto_core_message);

// Idea behind having three different functions is:
// Only ODK_ParseLicense structure had fields which needed additional
// procession. Having a single function with templated parameter T was
// failing during compile time because other two structures doesn't have
// fields that need additional processing. Hence to reduce code redundance and
// make us of common FuzzerMutateResponse across three response fuzzers,
// three independent functions were defined and renewal and provisioning
// functions would be empty as no additional processing is needed for them.
void ConvertDataToValidBools(ODK_ParsedLicense* t);

void ConvertDataToValidBools(ODK_PreparedRenewalRequest* t);

void ConvertDataToValidBools(ODK_ParsedProvisioning* t);

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

template <typename A, typename T, typename F, typename G>
size_t FuzzerMutateResponse(uint8_t* data, size_t size, size_t max_size,
                            const F& odk_deserialize_fun,
                            const G& kdo_serialize_fun) {
  const size_t kArgsSize = sizeof(A);
  const size_t kCoreResponseSize = sizeof(T);
  const size_t kTotalResponseSize = kArgsSize + kCoreResponseSize;

  // Deserializing data in order to make sure it deserializes properly.
  // Input byte array format: [function arguments][data to parse].
  std::shared_ptr<A> _args(new A());
  A* args = _args.get();
  memcpy(args, data, kArgsSize);
  ODK_NonceValues nonce_values = args->nonce_values;
  args->nonce_values.api_major_version = ODK_MAJOR_VERSION;
  const uint8_t* buf = data + kArgsSize;
  T t = {};
  OEMCryptoResult result =
      odk_deserialize_fun(buf, size - kArgsSize, args, &nonce_values, &t);

  // If data doesn't deserialize successfully, We copy random bytes into
  // T and serialize using kdo function
  // which will create a valid oemcrypto core message using
  // nonce and request hash from function args. OEMCrypto core message acts as
  // input to odk_kdo.
  if (result != OEMCrypto_SUCCESS) {
    if (max_size < kTotalResponseSize) {
      return 0;
    }
    // Initialize remaining bytes needed in data to zero.
    if (size < kTotalResponseSize) {
      memset(data + size, 0, kTotalResponseSize - size);
    }
    t = {};
    memcpy(&t, buf, kCoreResponseSize);
  }

  // Ask LLVM to run its usual mutations, hopefully giving us interesting
  // inputs. We copy deserialized data into pointer data, run mutations
  // and copy back the mutated data to args and t
  memcpy(data + kArgsSize, &t, kCoreResponseSize);
  LLVMFuzzerMutate(data, kTotalResponseSize, kTotalResponseSize);
  memcpy(args, data, kArgsSize);
  memcpy(&t, data + kArgsSize, kCoreResponseSize);
  // Convert boolean flags in parsed message to valid bytes to
  // avoid errors from msan. Only needed for parsed license.
  ConvertDataToValidBools(&t);
  // Serialize the data after mutation.
  std::string oemcrypto_core_message;
  if (!kdo_serialize_fun(args, t, &oemcrypto_core_message)) {
    return 0;
  }

  // Copy mutated and serialized oemcrypto_core_message to data
  // so that it acts as input to odk_kdo function.
  memcpy(data + kArgsSize, oemcrypto_core_message.data(),
         oemcrypto_core_message.size());
  return kArgsSize + oemcrypto_core_message.size();
}

/**
 * Template arguments:
 * A: struct holding function arguments
 * T: odk deserialize output/kdo serialize input structure
 * F: odk deserialize function
 * G: kdo serialize function
 *
 * raw bytes -> F deserialize -> struct T -> G serialize -> raw bytes
 */
template <typename A, typename T, typename F, typename G>
void odk_kdo(const F& odk_fun, const G& kdo_fun, const uint8_t* in,
             const size_t size, const size_t args_size, uint8_t* out UNUSED) {
  T t = {};
  // Input byte array format: [function arguments][data to parse]
  if (size < args_size) {
    return;
  }
  const uint8_t* buf = in + args_size;
  std::shared_ptr<A> _args(new A());
  A* args = _args.get();
  memcpy(args, in, args_size);
  args->nonce_values.api_major_version = ODK_MAJOR_VERSION;
  ODK_NonceValues nonce_values = args->nonce_values;

  OEMCryptoResult result =
      odk_fun(buf, size - args_size, args, &nonce_values, &t);
  if (result != OEMCrypto_SUCCESS) {
    return;
  }
  std::string oemcrypto_core_message;
  if (!kdo_fun(args, t, &oemcrypto_core_message)) {
    return;
  }
}

/**
 * Template arguments:
 * T: kdo deserialize output/odk serialize input structure
 * F: kdo deserialize function
 * G: odk serialize function
 *
 * raw bytes -> F deserialize -> struct T -> G serialize -> raw bytes
 */
template <typename T, typename F, typename G>
static void kdo_odk(const F& kdo_fun, const G& odk_fun, const uint8_t* in,
                    size_t size, const size_t clock_value_size, uint8_t* out) {
  if (size <= clock_value_size) {
    return;
  }
  // Input byte array format: [Clock Values][data to parse].
  // Only Renewal Request expects clock values to be present.
  std::string input(reinterpret_cast<const char*>(in) + clock_value_size,
                    size - clock_value_size);
  T t = {};
  if (!kdo_fun(input, &t)) {
    return;
  }
  ODK_NonceValues nonce_values = {t.api_minor_version, t.api_major_version,
                                  t.nonce, t.session_id};
  OEMCryptoResult err = odk_fun(in, out, &size, t, &nonce_values);
  if (OEMCrypto_SUCCESS != err) {
    return;
  }
}
}  // namespace oemcrypto_core_message
#endif  // WIDEVINE_ODK_TEST_FUZZING_ODK_FUZZ_HELPER_H_