/* Copyright 2020 Google LLC. All rights reserved. This file and proprietary */
/* source code may only be used and distributed under the Widevine Master */
/* 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_serialize.h"
#include "fuzzing/odk_fuzz_structs.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) {
  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_ */