// 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 <vector>

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

namespace {

// Limit output buffer size to 5 MB as 4 MB is maximum size specified by
// resource rating tier documentation.
constexpr size_t MAX_FUZZ_SAMPLE_SIZE = 5 * wvoec::MB;

// Avoid calling non-trivial destructor.
wvoec::OEMCryptoLicenseAPIFuzz& license_api_fuzz =
    *new wvoec::OEMCryptoLicenseAPIFuzz;

}  // namespace

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

extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
  // Split data using separator.
  const std::vector<wvoec::FuzzedData> inputs =
      wvoec::SplitFuzzedData(data, size);
  if (inputs.size() < 3) {
    return 0;
  }

  // Read cipher mode and pattern from fuzzed data.
  wvoec::OEMCrypto_Decrypt_Cenc_Fuzz fuzzed_structure;
  if (inputs[0].size < sizeof(fuzzed_structure)) {
    return 0;
  }
  FuzzedDataProvider fuzzed_data(inputs[0].data, inputs[0].size);
  fuzzed_data.ConsumeData(&fuzzed_structure, sizeof(fuzzed_structure));
  wvoec::ConvertDataToValidEnum(OEMCrypto_CipherMode_MaxValue,
                                fuzzed_structure.cipher_mode);

  // Allocate sample descriptions.
  std::vector<OEMCrypto_SampleDescription> sample_descriptions(
      fuzzed_data.remaining_bytes() /
      sizeof(wvoec::OEMCrypto_SampleDescription_Fuzz));

  // Allocate input buffers for each sample description.
  std::vector<std::vector<OEMCrypto_SharedMemory>> input_buffers(
      sample_descriptions.size());

  // Allocate secure_fd values for secure buffers.
  std::vector<int> secure_fd_array(sample_descriptions.size());

  // Allocate subsamples for each sample description.
  std::vector<std::vector<OEMCrypto_SubSampleDescription>> subsamples(
      sample_descriptions.size());

  const uint32_t session_id = license_api_fuzz.session().session_id();

  // Free first given number of output buffers.
  const auto FreeOutputBuffers = [&sample_descriptions, session_id,
                                  &secure_fd_array](size_t num_buffers) {
    for (size_t i = 0; i < num_buffers; i++) {
      OEMCrypto_DestBufferDesc& output_descriptor =
          sample_descriptions[i].buffers.output_descriptor;
      switch (output_descriptor.type) {
        case OEMCrypto_BufferType_Clear:
          delete[] output_descriptor.buffer.clear.clear_buffer;
          break;

        case OEMCrypto_BufferType_Secure:
          OEMCrypto_FreeSecureBuffer(session_id, &output_descriptor,
                                     secure_fd_array[i]);
          break;

        case OEMCrypto_BufferType_Direct:
          break;
      }
    }
  };

  // Prepare each sample description.
  FuzzedDataProvider& sample_description_data = fuzzed_data;
  FuzzedDataProvider input_buffer_data(inputs[1].data, inputs[1].size);
  FuzzedDataProvider subsample_data(inputs[2].data, inputs[2].size);
  for (size_t i = 0; i < sample_descriptions.size(); i++) {
    // Read and normalize sample description fuzzed properties.
    wvoec::OEMCrypto_SampleDescription_Fuzz fuzzed_sample_description;
    sample_description_data.ConsumeData(&fuzzed_sample_description,
                                        sizeof(fuzzed_sample_description));
    fuzzed_sample_description.buffers.input_data_length %=
        MAX_FUZZ_SAMPLE_SIZE + 1;
    wvoec::ConvertDataToValidEnum(
        OEMCrypto_BufferType_MaxValue,
        fuzzed_sample_description.buffers.output_descriptor.type);
    fuzzed_sample_description.buffers.output_descriptor.buffer_config %=
        MAX_FUZZ_SAMPLE_SIZE + 1;

    // Read input data.
    if (fuzzed_sample_description.buffers.input_data_length >
        input_buffer_data.remaining_bytes()) {
      FreeOutputBuffers(i);
      return 0;
    }
    input_buffers[i] = input_buffer_data.ConsumeBytes<uint8_t>(
        fuzzed_sample_description.buffers.input_data_length);
    sample_descriptions[i].buffers.input_data = input_buffers[i].data();
    sample_descriptions[i].buffers.input_data_length = input_buffers[i].size();

    // Set subsample data.
    if (fuzzed_sample_description.subsamples_length >
        subsample_data.remaining_bytes() /
            sizeof(OEMCrypto_SubSampleDescription)) {
      FreeOutputBuffers(i);
      return 0;
    }
    if (fuzzed_sample_description.subsamples_length > 0) {
      subsamples[i].resize(fuzzed_sample_description.subsamples_length);
      subsample_data.ConsumeData(
          subsamples[i].data(),
          subsamples[i].size() * sizeof(OEMCrypto_SubSampleDescription));
    }
    sample_descriptions[i].subsamples = subsamples[i].data();
    sample_descriptions[i].subsamples_length = subsamples[i].size();

    // Set IV data.
    memcpy(sample_descriptions[i].iv, fuzzed_sample_description.iv,
           sizeof(sample_descriptions[i].iv));

    // Initialize output buffer.
    OEMCrypto_DestBufferDesc& output_descriptor =
        sample_descriptions[i].buffers.output_descriptor;
    const wvoec::OEMCrypto_DestBufferDesc_Fuzz& fuzzed_output_descriptor =
        fuzzed_sample_description.buffers.output_descriptor;
    output_descriptor.type = fuzzed_output_descriptor.type;
    switch (output_descriptor.type) {
      case OEMCrypto_BufferType_Clear:
        output_descriptor.buffer.clear.clear_buffer =
            new OEMCrypto_SharedMemory[fuzzed_output_descriptor.buffer_config];
        output_descriptor.buffer.clear.clear_buffer_length =
            fuzzed_output_descriptor.buffer_config;
        break;

      case OEMCrypto_BufferType_Secure:
        if (OEMCrypto_AllocateSecureBuffer(
                session_id, fuzzed_output_descriptor.buffer_config,
                &output_descriptor, &secure_fd_array[i]) != OEMCrypto_SUCCESS) {
          FreeOutputBuffers(i);
          return 0;
        }
        break;

      case OEMCrypto_BufferType_Direct:
        output_descriptor.buffer.direct.is_video =
            fuzzed_output_descriptor.buffer_config & 1;
        break;
    }
  }

  // Load license and call decrypt_cenc API.
  const wvoec::MessageKeyData& key =
      license_api_fuzz.session().license().keys[0];
  std::vector<uint8_t> key_handle;
  wvoec::GetKeyHandleIntoVector(session_id, key.key_id, key.key_id_length,
                                fuzzed_structure.cipher_mode, key_handle);
  OEMCrypto_DecryptCENC(key_handle.data(), key_handle.size(),
                        sample_descriptions.data(), sample_descriptions.size(),
                        &fuzzed_structure.pattern);

  // Free all output buffers.
  FreeOutputBuffers(sample_descriptions.size());

  return 0;
}