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

#include "oec_decrypt_fallback_chain.h"

#include <stdint.h>
#include <string.h>

#include "oemcrypto_corpus_generator_helper.h"
#include "oemcrypto_types.h"
#include "string_conversions.h"

namespace {

void advance_dest_buffer(OEMCrypto_DestBufferDesc* dest_buffer, size_t bytes) {
  switch (dest_buffer->type) {
    case OEMCrypto_BufferType_Clear:
      dest_buffer->buffer.clear.clear_buffer += bytes;
      dest_buffer->buffer.clear.clear_buffer_length -= bytes;
      break;

    case OEMCrypto_BufferType_Secure:
      dest_buffer->buffer.secure.offset += bytes;
      break;

    case OEMCrypto_BufferType_Direct:
      // Nothing to do for this buffer type.
      break;
  }
}

void advance_iv_ctr(uint8_t (*subsample_iv)[wvoec::KEY_IV_SIZE], size_t bytes) {
  uint64_t counter;
  constexpr size_t half_iv_size = wvoec::KEY_IV_SIZE / 2;
  memcpy(&counter, &(*subsample_iv)[half_iv_size], half_iv_size);

  const size_t increment =
      bytes / wvoec::AES_128_BLOCK_SIZE;  // The truncation here is intentional
  counter = wvutil::htonll64(wvutil::ntohll64(counter) + increment);

  memcpy(&(*subsample_iv)[half_iv_size], &counter, half_iv_size);
}

}  // namespace

namespace wvoec {

// Decrypts the given array of samples. Handles fallback behavior correctly if
// the OEMCrypto implementation does not accept multiple samples.
OEMCryptoResult DecryptFallbackChain::Decrypt(
    OEMCrypto_SESSION session_id, const OEMCrypto_SampleDescription* samples,
    size_t samples_length, OEMCryptoCipherMode cipher_mode,
    const OEMCrypto_CENCEncryptPatternDesc* pattern) {
  OEMCryptoResult sts =
      OEMCrypto_DecryptCENC(session_id, samples, samples_length, pattern);

  // No need for a fallback. Abort early.
  if (sts != OEMCrypto_ERROR_BUFFER_TOO_LARGE) {
    if (ShouldGenerateCorpus()) {
      WriteDecryptCencCorpus(cipher_mode, samples, pattern, samples_length);
    }
    return sts;
  }

  // Fall back to decrypting individual samples.
  for (size_t i = 0; i < samples_length; ++i) {
    sts = DecryptSample(session_id, samples[i], cipher_mode, pattern);
    if (sts != OEMCrypto_SUCCESS) return sts;
  }

  return sts;
}

// Decrypts the given sample. Handles fallback behavior correctly if the
// OEMCrypto implementation does not accept full samples.
OEMCryptoResult DecryptFallbackChain::DecryptSample(
    OEMCrypto_SESSION session_id, const OEMCrypto_SampleDescription& sample,
    OEMCryptoCipherMode cipher_mode,
    const OEMCrypto_CENCEncryptPatternDesc* pattern) {
  OEMCryptoResult sts = OEMCrypto_DecryptCENC(session_id, &sample, 1, pattern);

  // No need for a fallback. Abort early.
  if (sts != OEMCrypto_ERROR_BUFFER_TOO_LARGE) {
    if (ShouldGenerateCorpus()) {
      WriteDecryptCencCorpus(cipher_mode, &sample, pattern, 1);
    }
    return sts;
  }

  // Fall back to decrypting individual subsamples.
  OEMCrypto_SampleDescription fake_sample = sample;
  for (size_t i = 0; i < sample.subsamples_length; ++i) {
    const OEMCrypto_SubSampleDescription& subsample = sample.subsamples[i];

    const size_t length =
        subsample.num_bytes_clear + subsample.num_bytes_encrypted;
    fake_sample.buffers.input_data_length = length;
    fake_sample.subsamples = &subsample;
    fake_sample.subsamples_length = 1;

    sts = DecryptSubsample(session_id, fake_sample, pattern, cipher_mode);
    if (sts != OEMCrypto_SUCCESS) return sts;

    fake_sample.buffers.input_data += length;
    advance_dest_buffer(&fake_sample.buffers.output_descriptor, length);
    if (cipher_mode == OEMCrypto_CipherMode_CENC) {
      advance_iv_ctr(&fake_sample.iv,
                     subsample.block_offset + subsample.num_bytes_encrypted);
    }
  }

  return sts;
}

// Decrypts the given subsample. Handles fallback behavior correctly if the
// OEMCrypto implementation does not accept full subsamples.
OEMCryptoResult DecryptFallbackChain::DecryptSubsample(
    OEMCrypto_SESSION session_id, const OEMCrypto_SampleDescription& sample,
    const OEMCrypto_CENCEncryptPatternDesc* pattern,
    OEMCryptoCipherMode cipher_mode) {
  OEMCryptoResult sts = OEMCrypto_DecryptCENC(session_id, &sample, 1, pattern);

  // No need for a fallback. Abort early.
  if (sts != OEMCrypto_ERROR_BUFFER_TOO_LARGE) {
    if (ShouldGenerateCorpus()) {
      WriteDecryptCencCorpus(cipher_mode, &sample, pattern, 1);
    }
    return sts;
  }

  // Fall back to decrypting individual subsample halves.
  const OEMCrypto_SubSampleDescription& subsample = sample.subsamples[0];
  OEMCrypto_SampleDescription fake_sample = sample;
  OEMCrypto_SubSampleDescription fake_subsample;
  fake_sample.subsamples = &fake_subsample;
  fake_sample.subsamples_length = 1;

  if (subsample.num_bytes_clear > 0) {
    fake_sample.buffers.input_data_length = subsample.num_bytes_clear;
    fake_subsample.num_bytes_clear = subsample.num_bytes_clear;
    fake_subsample.num_bytes_encrypted = 0;
    fake_subsample.block_offset = 0;

    fake_subsample.subsample_flags = 0;
    if (subsample.subsample_flags & OEMCrypto_FirstSubsample)
      fake_subsample.subsample_flags |= OEMCrypto_FirstSubsample;
    if (subsample.subsample_flags & OEMCrypto_LastSubsample &&
        subsample.num_bytes_encrypted == 0)
      fake_subsample.subsample_flags |= OEMCrypto_LastSubsample;

    sts = DecryptSubsampleHalf(session_id, fake_sample, pattern, cipher_mode);
    if (sts != OEMCrypto_SUCCESS) return sts;

    // Advance the buffers for the other half, in case they're needed.
    fake_sample.buffers.input_data += subsample.num_bytes_clear;
    advance_dest_buffer(&fake_sample.buffers.output_descriptor,
                        subsample.num_bytes_clear);
  }

  if (subsample.num_bytes_encrypted > 0) {
    fake_sample.buffers.input_data_length = subsample.num_bytes_encrypted;
    fake_subsample.num_bytes_clear = 0;
    fake_subsample.num_bytes_encrypted = subsample.num_bytes_encrypted;
    fake_subsample.block_offset = subsample.block_offset;

    fake_subsample.subsample_flags = 0;
    if (subsample.subsample_flags & OEMCrypto_FirstSubsample &&
        subsample.num_bytes_clear == 0)
      fake_subsample.subsample_flags |= OEMCrypto_FirstSubsample;
    if (subsample.subsample_flags & OEMCrypto_LastSubsample)
      fake_subsample.subsample_flags |= OEMCrypto_LastSubsample;

    sts = DecryptSubsampleHalf(session_id, fake_sample, pattern, cipher_mode);
    if (sts != OEMCrypto_SUCCESS) return sts;
  }

  return sts;
}

// Decrypts the given subsample half. There is no fallback behavior after this;
// an OEMCrypto_ERROR_BUFFER_TOO_LARGE produced here will be returned to the
// caller.
OEMCryptoResult DecryptFallbackChain::DecryptSubsampleHalf(
    OEMCrypto_SESSION session_id, const OEMCrypto_SampleDescription& sample,
    const OEMCrypto_CENCEncryptPatternDesc* pattern,
    OEMCryptoCipherMode cipher_mode) {
  if (ShouldGenerateCorpus()) {
    WriteDecryptCencCorpus(cipher_mode, &sample, pattern, 1);
  }
  return OEMCrypto_DecryptCENC(session_id, &sample, 1, pattern);
  // In a real CDM, you would want some fallback here to handle the case where
  // the buffer is too big for the OEMCrypto implementation. But in the case of
  // the tests, we won't be passing a buffer that's too big unless we are trying
  // to test that failure condition, so there's no need to handle that case
  // here.
}

// Used for OEMCrypto Fuzzing: Corpus format is as below, let | be separator.
// cipher_mode + pattern + sample_data for all samples |
// subsample_data for all samples
void WriteDecryptCencCorpus(
    OEMCryptoCipherMode cipher_mode,
    const OEMCrypto_SampleDescription* samples_description,
    const OEMCrypto_CENCEncryptPatternDesc* pattern, size_t samples_length) {
  const std::string file_name =
      GetFileName("oemcrypto_decrypt_cenc_fuzz_seed_corpus");
  OEMCrypto_Decrypt_Cenc_Fuzz decrypt_cenc_fuzz_struct;
  decrypt_cenc_fuzz_struct.cipher_mode = cipher_mode;
  decrypt_cenc_fuzz_struct.pattern = *pattern;
  // Cipher mode and Pattern.
  AppendToFile(file_name,
               reinterpret_cast<const char*>(&decrypt_cenc_fuzz_struct),
               sizeof(OEMCrypto_Decrypt_Cenc_Fuzz));

  // Sample data for all samples.
  for (size_t i = 0; i < samples_length; i++) {
    AppendToFile(file_name,
                 reinterpret_cast<const char*>(&samples_description[i]),
                 sizeof(OEMCrypto_SampleDescription));
  }
  AppendSeparator(file_name);

  // Subsample data for all samples.
  for (size_t i = 0; i < samples_length; i++) {
    for (size_t j = 0; j < samples_description[i].subsamples_length; j++) {
      AppendToFile(
          file_name,
          reinterpret_cast<const char*>(&samples_description[i].subsamples[j]),
          sizeof(OEMCrypto_SubSampleDescription));
    }
  }
}

}  // namespace wvoec