// 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 "log.h"
#include "oemcrypto_fuzz_helper.h"
#include "oemcrypto_fuzz_structs.h"
#include "oemcrypto_overflow.h"

namespace wvoec {
const size_t MAX_FUZZ_SAMPLE_SIZE = 5 * MB;
// Free dynamic memory allocated by fuzzer script.
void FreeOutputBuffers(OEMCrypto_SESSION session_id,
                       OEMCrypto_SampleDescription* sample_description,
                       size_t sample_index, int* secure_fd_array) {
  for (size_t i = 0; i < sample_index; i++) {
    OEMCrypto_DestBufferDesc fuzzed_output_descriptor =
        sample_description[i].buffers.output_descriptor;
    switch (fuzzed_output_descriptor.type) {
      case OEMCrypto_BufferType_Clear: {
        delete[] fuzzed_output_descriptor.buffer.clear.clear_buffer;
        break;
      }
      case OEMCrypto_BufferType_Secure: {
        OEMCrypto_FreeSecureBuffer(session_id, &fuzzed_output_descriptor,
                                   secure_fd_array[i]);
        break;
      }
      case OEMCrypto_BufferType_Direct: {
        break;
      }
    }
  }
}

// Function to initialize output buffer pointers by allocating memory.
// Limiting output buffer size to 5 MB as 4 MB is maximum size specified
// by resource rating tier documentation.
bool InitializeOutputBuffers(OEMCrypto_SESSION session_id,
                             OEMCrypto_DestBufferDesc& output_descriptor,
                             size_t sample_index,
                             vector<int>& secure_fd_array) {
  switch (output_descriptor.type) {
    case OEMCrypto_BufferType_Clear: {
      output_descriptor.buffer.clear.clear_buffer =
          new OEMCrypto_SharedMemory[std::min(
              MAX_FUZZ_SAMPLE_SIZE,
              output_descriptor.buffer.clear.clear_buffer_length)];
      return true;
    }
    case OEMCrypto_BufferType_Secure: {
      int* secure_fd;
      OEMCryptoResult sts = OEMCrypto_AllocateSecureBuffer(
          session_id,
          std::min(MAX_FUZZ_SAMPLE_SIZE,
                   output_descriptor.buffer.secure.secure_buffer_length),
          &output_descriptor, secure_fd);
      if (sts == OEMCrypto_SUCCESS) secure_fd_array[sample_index] = *secure_fd;
      return sts == OEMCrypto_SUCCESS;
    }
    case OEMCrypto_BufferType_Direct: {
      return true;
    }
  }
}

extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
  // Redirect printf and log statements from oemcrypto functions to a file to
  // reduce noise
  RedirectStdoutToFile();
  size_t samples_length;

  // Split data using separator.
  auto inputs = SplitInput(data, size);
  if (inputs.size() < 2) {
    return 0;
  }

  OEMCrypto_Decrypt_Cenc_Fuzz fuzzed_structure;
  if (inputs[0].size() < sizeof(fuzzed_structure)) {
    return 0;
  }
  // Copy OEMCrypto_Decrypt_Cenc_Fuzz from input data.
  memcpy(&fuzzed_structure, data, sizeof(fuzzed_structure));
  ConvertDataToValidEnum(OEMCrypto_CipherMode_MaxValue,
                         &fuzzed_structure.cipher_mode);

  size_t remaining_size_for_samples =
      inputs[0].size() - sizeof(fuzzed_structure);
  // Initialize FDP structures to read data using inbuilt functions.
  FuzzedDataProvider fuzzed_sample_data(data + sizeof(fuzzed_structure),
                                        remaining_size_for_samples);
  FuzzedDataProvider fuzzed_subsample_data(inputs[1].data(), inputs[1].size());

  // Read subsamples from fuzzed data.
  vector<OEMCrypto_SubSampleDescription> subsamples;
  while (fuzzed_subsample_data.remaining_bytes() >=
         sizeof(OEMCrypto_SubSampleDescription)) {
    OEMCrypto_SubSampleDescription subsample;
    fuzzed_subsample_data.ConsumeData(&subsample,
                                      sizeof(OEMCrypto_SubSampleDescription));
    subsamples.push_back(subsample);
  }
  if (subsamples.size() == 0) {
    return 0;
  }

  // Infer samples_length from fuzzed data.
  size_t sample_description_size = sizeof(OEMCrypto_SampleDescription);
  samples_length =
      fuzzed_sample_data.remaining_bytes() / sample_description_size;
  if (samples_length == 0) {
    return 0;
  }

  // Initialize sample_descriptions array.
  vector<OEMCrypto_SampleDescription> sample_descriptions(samples_length);
  // Create array to maintain secure_fd buffer values for secure buffers.
  vector<int> secure_fd_array(samples_length);

  OEMCryptoLicenseAPIFuzz license_api_fuzz;
  Session* session = license_api_fuzz.session();
  // Copy samples from fuzzed data.
  size_t input_subsample_index = 0;
  size_t total_input_data_length = 0;
  for (size_t i = 0; i < samples_length; i++) {
    fuzzed_sample_data.ConsumeData(&sample_descriptions[i],
                                   sample_description_size);
    ConvertDataToValidEnum(
        OEMCrypto_BufferType_MaxValue,
        &sample_descriptions[i].buffers.output_descriptor.type);

    // Copy random data into input sample data. Cap input data length at 5 MB,
    // 1 MB higher than that described by resource rating tier.
    total_input_data_length += std::min(
        MAX_FUZZ_SAMPLE_SIZE, sample_descriptions[i].buffers.input_data_length);

    // Copy sub sample data.
    sample_descriptions[i].subsamples = &subsamples[input_subsample_index];
    if (OPK_AddOverflowUX(input_subsample_index,
                          sample_descriptions[i].subsamples_length,
                          &input_subsample_index)) {
      return 0;
    }
    if (input_subsample_index > subsamples.size()) return 0;
  }  // Sample loop.

  // Allocate input/output buffers for each sample description.
  vector<OEMCrypto_SharedMemory> input_buffer(total_input_data_length);
  size_t input_buffer_index = 0;
  for (size_t i = 0; i < samples_length; i++) {
    sample_descriptions[i].buffers.input_data =
        &input_buffer[input_buffer_index];
    input_buffer_index += std::min(
        MAX_FUZZ_SAMPLE_SIZE, sample_descriptions[i].buffers.input_data_length);

    // Create output buffer pointers. If secure buffer is not supported, we
    // explicitly convert to clear buffer and fuzz.
    if (!InitializeOutputBuffers(
            session->session_id(),
            sample_descriptions[i].buffers.output_descriptor, i,
            secure_fd_array)) {
      LOGI(
          "[OEMCrypto decrypt CENC fuzz] Secure buffers are not supported. Use "
          "clear buffer instead.");
      sample_descriptions[i].buffers.output_descriptor.type =
          OEMCrypto_BufferType_Clear;
      InitializeOutputBuffers(session->session_id(),
                              sample_descriptions[i].buffers.output_descriptor,
                              i, secure_fd_array);
    }
  }

  // Load license and call decrypt_cenc API.
  license_api_fuzz.LoadLicense();
  OEMCrypto_SelectKey(session->session_id(), session->license().keys[0].key_id,
                      session->license().keys[0].key_id_length,
                      fuzzed_structure.cipher_mode);
  OEMCrypto_DecryptCENC(session->session_id(), sample_descriptions.data(),
                        samples_length, &fuzzed_structure.pattern);
  FreeOutputBuffers(session->session_id(), sample_descriptions.data(),
                    samples_length, secure_fd_array.data());
  return 0;
}
}  // namespace wvoec