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

#include "FuzzedDataProvider.h"
#include "OEMCryptoCENC.h"
#include "oec_session_util.h"
#include "oemcrypto_fuzz_helper.h"
#include "oemcrypto_fuzz_structs.h"
#include "oemcrypto_types.h"

namespace wvoec {

// Properties deserialized from fuzzed data.
struct FuzzedProperties {
  OEMCrypto_Generic_Api_Fuzz structure;
  std::vector<uint8_t> buffer;
  std::vector<uint8_t> signature;
};

// Contains value only if has_value is true.
struct OptionalFuzzedProperties {
  FuzzedProperties value;
  bool has_value;
};

OEMCryptoLicenseAPIFuzz license_api_fuzz;

OptionalFuzzedProperties DeserializeFuzzedData(const uint8_t* data,
                                               size_t size) {
  OptionalFuzzedProperties fuzzed_properties;
  const std::vector<FuzzedData> inputs = SplitFuzzedData(data, size);
  if (inputs.size() < 2 ||
      inputs[0].size < sizeof(fuzzed_properties.value.structure)) {
    fuzzed_properties.has_value = false;
    return fuzzed_properties;
  }
  FuzzedDataProvider fuzzed_data(inputs[0].data, inputs[0].size);
  fuzzed_data.ConsumeData(&fuzzed_properties.value.structure,
                          sizeof(fuzzed_properties.value.structure));
  ConvertDataToValidEnum(OEMCrypto_CipherMode_MaxValue,
                         &fuzzed_properties.value.structure.cipher_mode);
  ConvertDataToValidEnum(OEMCrypto_Algorithm_MaxValue,
                         &fuzzed_properties.value.structure.algorithm);
  fuzzed_properties.value.buffer = fuzzed_data.ConsumeRemainingBytes<uint8_t>();
  fuzzed_properties.value.signature.assign(inputs[1].data,
                                           inputs[1].data + inputs[1].size);
  fuzzed_properties.has_value = true;
  return fuzzed_properties;
}

extern "C" int LLVMFuzzerInitialize(int* argc, char*** argv) {
  RedirectStdoutToFile();
  license_api_fuzz.LoadLicense();
  return 0;
}

extern "C" size_t LLVMFuzzerCustomMutator(uint8_t* data, size_t size,
                                          size_t max_size, unsigned int seed) {
  // Deserialize fuzzed data.
  OptionalFuzzedProperties fuzzed_properties =
      DeserializeFuzzedData(data, size);
  if (!fuzzed_properties.has_value) {
    return 0;
  }

  // Get key handle for signing and verifying.
  Session* const session = license_api_fuzz.session();
  vector<uint8_t> key_handle;
  OEMCryptoResult result = GetKeyHandleIntoVector(
      session->session_id(), session->license().keys[0].key_id,
      session->license().keys[0].key_id_length,
      fuzzed_properties.value.structure.cipher_mode, key_handle);
  if (result == OEMCrypto_SUCCESS) {
    // Generate a new signature if verification fails.
    result =
        OEMCrypto_Generic_Verify(key_handle.data(), key_handle.size(),
                                 fuzzed_properties.value.buffer.data(),
                                 fuzzed_properties.value.buffer.size(),
                                 fuzzed_properties.value.structure.algorithm,
                                 fuzzed_properties.value.signature.data(),
                                 fuzzed_properties.value.signature.size());
    if (result != OEMCrypto_SUCCESS) {
      size_t signature_length = 0;
      OEMCrypto_Generic_Sign(key_handle.data(), key_handle.size(),
                             fuzzed_properties.value.buffer.data(),
                             fuzzed_properties.value.buffer.size(),
                             fuzzed_properties.value.structure.algorithm,
                             nullptr, &signature_length);
      fuzzed_properties.value.signature.resize(signature_length);
      OEMCrypto_Generic_Sign(key_handle.data(), key_handle.size(),
                             fuzzed_properties.value.buffer.data(),
                             fuzzed_properties.value.buffer.size(),
                             fuzzed_properties.value.structure.algorithm,
                             fuzzed_properties.value.signature.data(),
                             &signature_length);
      const size_t signature_offset =
          sizeof(fuzzed_properties.value.structure) +
          fuzzed_properties.value.buffer.size() + sizeof(kFuzzDataSeparator);
      size = signature_offset + signature_length;
      if (size > max_size) {
        return 0;
      }
      memcpy(data + signature_offset, fuzzed_properties.value.signature.data(),
             signature_length);
    }
  }

  return LLVMFuzzerMutate(data, size, max_size);
}

extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
  // Deserialize fuzzed data.
  const OptionalFuzzedProperties fuzzed_properties =
      DeserializeFuzzedData(data, size);
  if (!fuzzed_properties.has_value) {
    return 0;
  }

  // Select key and perform verification.
  Session* const session = license_api_fuzz.session();
  vector<uint8_t> key_handle;
  GetKeyHandleIntoVector(
      session->session_id(), session->license().keys[0].key_id,
      session->license().keys[0].key_id_length,
      fuzzed_properties.value.structure.cipher_mode, key_handle);
  OEMCrypto_Generic_Verify(key_handle.data(), key_handle.size(),
                           fuzzed_properties.value.buffer.data(),
                           fuzzed_properties.value.buffer.size(),
                           fuzzed_properties.value.structure.algorithm,
                           fuzzed_properties.value.signature.data(),
                           fuzzed_properties.value.signature.size());
  return 0;
}

}  // namespace wvoec