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

#include <ctype.h>
#include <gtest/gtest.h>
#include <openssl/aes.h>
#include <openssl/err.h>
#include <openssl/hmac.h>
#include <openssl/rand.h>
#include <openssl/rsa.h>
#include <openssl/sha.h>
#include <openssl/x509.h>
#include <stdint.h>

#include <algorithm>
#include <chrono>
#include <functional>
#include <iostream>
#include <map>
#include <string>
#include <utility>
#include <vector>

// TODO(b/253779846) Change it to include a header instead
#include "oemcrypto_test.cpp"

using ::testing::Bool;
using ::testing::Combine;
using ::testing::Range;
using ::testing::tuple;
using ::testing::Values;
using ::testing::WithParamInterface;
using namespace std;

namespace wvoec {

/// @addtogroup security
/// @{

/** Test that OEMCrypto_FreeSecureBuffer fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoClientTest,
       OEMCryptoMemoryAllocateSecureBufferForHugeBufferSize) {
  Session s;
  s.open();
  auto oemcrypto_function = [&s](size_t buffer_size) {
    OEMCrypto_DestBufferDesc output_descriptor;
    int secure_fd;
    OEMCryptoResult sts = OEMCrypto_AllocateSecureBuffer(
        s.session_id(), buffer_size, &output_descriptor, &secure_fd);
    if (sts == OEMCrypto_SUCCESS) {
      OEMCrypto_FreeSecureBuffer(s.session_id(), &output_descriptor, secure_fd);
    }
    return sts;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
  s.close();
}

/** Test that OEMCrypto_WrapKeyboxOrOEMCert fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoClientTest,
       OEMCryptoMemoryWrapKeyboxOrOEMCertForHugeKeyboxLength) {
  auto oemcrypto_function = [](size_t keybox_length) {
    vector<uint8_t> keybox_buffer(keybox_length);
    size_t wrapped_keybox_length = keybox_length + 50;
    vector<uint8_t> wrapped_keybox_buffer(wrapped_keybox_length);
    vector<uint8_t> transport_key_buffer(20);
    memcpy(keybox_buffer.data(), &kTestKeybox, sizeof(kTestKeybox));
    return OEMCrypto_WrapKeyboxOrOEMCert(
        keybox_buffer.data(), keybox_length, wrapped_keybox_buffer.data(),
        &wrapped_keybox_length, transport_key_buffer.data(),
        transport_key_buffer.size());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, sizeof(kTestKeybox),
                                kHugeInputBufferLength, kCheckStatus);
}

/** Test that OEMCrypto_WrapKeyboxOrOEMCert fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoClientTest,
       OEMCryptoMemoryWrapKeyboxOrOEMCertForHugeWrappedKeyboxLength) {
  auto oemcrypto_function = [](size_t buffer_length) {
    size_t wrapped_keybox_length = buffer_length;
    vector<uint8_t> wrapped_keybox_buffer(wrapped_keybox_length);
    vector<uint8_t> transport_key_buffer(20);
    return OEMCrypto_WrapKeyboxOrOEMCert(
        reinterpret_cast<const uint8_t*>(&kTestKeybox), sizeof(kTestKeybox),
        wrapped_keybox_buffer.data(), &wrapped_keybox_length,
        transport_key_buffer.data(), transport_key_buffer.size());
  };
  // API expects keybox length and wrapped keybox length to be equal. We would
  // like to test for various values of wrapped keybox lengths, hence skipping
  // status check.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_WrapKeyboxOrOEMCert fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoClientTest,
       OEMCryptoMemoryWrapKeyboxOrOEMCertForHugeTransportKey) {
  auto oemcrypto_function = [](size_t transport_key_length) {
    size_t wrapped_keybox_length = sizeof(&kTestKeybox) + 50;
    vector<uint8_t> wrapped_keybox_buffer(wrapped_keybox_length);
    vector<uint8_t> transport_key_buffer(transport_key_length);
    return OEMCrypto_WrapKeyboxOrOEMCert(
        reinterpret_cast<const uint8_t*>(&kTestKeybox), sizeof(kTestKeybox),
        wrapped_keybox_buffer.data(), &wrapped_keybox_length,
        transport_key_buffer.data(), transport_key_buffer.size());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_WrapKeyboxOrOEMCert fails gracefully on a huge buffer.
 */
TEST_F(
    OEMCryptoClientTest,
    OEMCryptoMemoryWrapKeyboxOrOEMCertForHugeKeyboxLengthStartingFromLength1) {
  auto oemcrypto_function = [](size_t keybox_length) {
    vector<uint8_t> keybox_buffer(keybox_length);
    size_t wrapped_keybox_length = keybox_length + 50;
    vector<uint8_t> wrapped_keybox_buffer(wrapped_keybox_length);
    vector<uint8_t> transport_key_buffer(20);
    return OEMCrypto_WrapKeyboxOrOEMCert(
        keybox_buffer.data(), keybox_length, wrapped_keybox_buffer.data(),
        &wrapped_keybox_length, transport_key_buffer.data(),
        transport_key_buffer.size());
  };
  // Cannot check status as keybox will not be valid for this test.
  // We still want to test what happens if buffer lengths is less that keybox
  // length.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_SetSandbox fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoClientTest, OEMCryptoMemorySetSandboxForHugeSandboxIdLength) {
  auto oemcrypto_function = [](size_t buffer_length) {
    vector<uint8_t> buffer(buffer_length);
    return OEMCrypto_SetSandbox(buffer.data(), buffer.size());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_CopyBuffer fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoClientTest, OEMCryptoMemoryCopyBufferForHugeBufferLengths) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  vector<uint8_t> input_buffer;
  OEMCrypto_DestBufferDesc dest_buffer_descriptor;
  dest_buffer_descriptor.type = OEMCrypto_BufferType_Secure;

  auto oemcrypto_function = [&s, &dest_buffer_descriptor,
                             &input_buffer](size_t buffer_length) {
    input_buffer.resize(buffer_length);
    int secure_fd;
    OEMCryptoResult sts = OEMCrypto_AllocateSecureBuffer(
        s.session_id(), buffer_length, &dest_buffer_descriptor, &secure_fd);
    if (sts != OEMCrypto_SUCCESS) {
      LOGI("Secure buffers are not supported.");
      return sts;
    }

    dest_buffer_descriptor.buffer.secure.secure_buffer_length = buffer_length;
    OEMCryptoResult status = OEMCrypto_CopyBuffer(
        s.session_id(), input_buffer.data(), buffer_length,
        &dest_buffer_descriptor,
        OEMCrypto_FirstSubsample | OEMCrypto_LastSubsample);
    OEMCrypto_FreeSecureBuffer(s.session_id(), &dest_buffer_descriptor,
                               secure_fd);
    return status;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** @ingroup security
 * Test that OEMCrypto_CopyBuffer fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoClientTest,
       OEMCryptoMemoryCopyBufferDirectForHugeBufferLengths) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  vector<uint8_t> input_buffer;
  OEMCrypto_DestBufferDesc dest_buffer_descriptor;
  dest_buffer_descriptor.type = OEMCrypto_BufferType_Direct;
  dest_buffer_descriptor.buffer.direct.is_video = false;

  auto oemcrypto_function = [&s, &dest_buffer_descriptor,
                             &input_buffer](size_t buffer_length) {
    input_buffer.resize(buffer_length);
    OEMCryptoResult status = OEMCrypto_CopyBuffer(
        s.session_id(), input_buffer.data(), buffer_length,
        &dest_buffer_descriptor,
        OEMCrypto_FirstSubsample | OEMCrypto_LastSubsample);
    return status;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** @ingroup security
 * Test that OEMCrypto_CopyBuffer fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoClientTest, OEMCryptoMemoryCopyBufferForOutOfRangeOffset) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  vector<uint8_t> input_buffer;
  OEMCrypto_DestBufferDesc dest_buffer_descriptor;
  dest_buffer_descriptor.type = OEMCrypto_BufferType_Secure;

  size_t buffer_length = KiB;
  input_buffer.resize(buffer_length);
  int secure_fd;
  if (OEMCrypto_AllocateSecureBuffer(s.session_id(), buffer_length,
                                     &dest_buffer_descriptor,
                                     &secure_fd) != OEMCrypto_SUCCESS) {
    LOGI("Secure buffers are not supported.");
    return;
  }

  dest_buffer_descriptor.buffer.secure.secure_buffer_length = buffer_length;
  auto oemcrypto_function = [&s, &dest_buffer_descriptor, &input_buffer,
                             &buffer_length](size_t offset) {
    dest_buffer_descriptor.buffer.secure.offset = offset;
    return OEMCrypto_CopyBuffer(
        s.session_id(), input_buffer.data(), buffer_length,
        &dest_buffer_descriptor,
        OEMCrypto_FirstSubsample | OEMCrypto_LastSubsample);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
  OEMCrypto_FreeSecureBuffer(s.session_id(), &dest_buffer_descriptor,
                             secure_fd);
}

/******** Dangerous Tests - DO NOT RUN ***********/
/*The following tests try to test InstallKeybox API with random buffers of
varying length in order to catch any overflow issues. These tests override the
actual keybox on the device. Remove the if and endif statement to run these
tests on a device ONLY IF YOU ARE ABLE TO RECOVER THE KEYBOX on the device.*/
#if 0
TEST_F(OEMCryptoKeyboxTest, OEMCryptoMemoryInstallKeyboxForHugeKeyboxBuffer) {
  auto f = [](size_t keybox_length) {
    vector<uint8_t> keybox(keybox_length);
    memcpy(keybox.data(), &kTestKeybox, sizeof(kTestKeybox));
    return OEMCrypto_InstallKeyboxOrOEMCert(keybox.data(), keybox.size());
  };
  // Starting at sizeof(kTestKeybox) as we are copying valid keybox
  // at beginning of generated buffers.
  TestHugeLengthDoesNotCrashAPI(f, sizeof(kTestKeybox), kHugeInputBufferLength,
                                kCheckStatus);
}

TEST_F(OEMCryptoKeyboxTest,
       OEMCryptoMemoryInstallKeyboxForHugeKeyboxBufferStartingFromLength1) {
  auto f = [](size_t keybox_length) {
    vector<uint8_t> keybox(keybox_length);
    return OEMCrypto_InstallKeyboxOrOEMCert(keybox.data(), keybox.size());
  };
  // We are testing for keybox lengths starting from 1 which would return error,
  // hence skipping status check.
  TestHugeLengthDoesNotCrashAPI(f, !kCheckStatus);
}
#endif

/** @ingroup security
 * Test that OEMCrypto_LoadTestKeybox fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoKeyboxTest, OEMCryptoMemoryLoadTestKeyBoxForHugeKeyboxBuffer) {
  auto f = [](size_t keybox_length) {
    vector<uint8_t> keybox(keybox_length);
    memcpy(keybox.data(), &kTestKeybox, sizeof(kTestKeybox));
    return OEMCrypto_LoadTestKeybox(keybox.data(), keybox.size());
  };
  // Starting at sizeof(kTestKeybox) as we are copying valid keybox
  // at beginning of generated buffers.
  TestHugeLengthDoesNotCrashAPI(f, sizeof(kTestKeybox), kHugeInputBufferLength,
                                kCheckStatus);
}

/** @ingroup security
 * Test that OEMCrypto_LoadTestKeybox fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoKeyboxTest,
       OEMCryptoMemoryLoadTestKeyBoxForHugeKeyboxBufferStartingFromLength1) {
  auto f = [](size_t keybox_length) {
    vector<uint8_t> keybox(keybox_length);
    return OEMCrypto_LoadTestKeybox(keybox.data(), keybox.size());
  };
  // We are testing for keybox lengths starting from 1 which would return error,
  // hence skipping status check.
  TestHugeLengthDoesNotCrashAPI(f, !kCheckStatus);
}

/** Test that OEMCrypto_GetDeviceID fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoKeyboxTest, OEMCryptoMemoryGetDeviceIdForHugeIdLength) {
  auto oemcrypto_function = [](size_t input_length) {
    size_t device_id_length = input_length;
    vector<uint8_t> device_id(device_id_length);
    return OEMCrypto_GetDeviceID(device_id.data(), &device_id_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_GetKeyData fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoKeyboxTest, OEMCryptoMemoryGetKeyIdForHugeIdLength) {
  auto oemcrypto_function = [](size_t input_length) {
    size_t key_data_length = input_length;
    vector<uint8_t> key_data(key_data_length);
    return OEMCrypto_GetKeyData(key_data.data(), &key_data_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_GenerateDerivedKeys fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoKeyboxTest,
       OEMCryptoMemoryGenerateDerivedKeysForHugeMacContextLength) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  vector<uint8_t> mac_context;
  vector<uint8_t> enc_context;
  s.FillDefaultContext(&mac_context, &enc_context);

  auto oemcrypto_function = [&s, &mac_context,
                             &enc_context](size_t buffer_length) {
    mac_context.resize(buffer_length);
    return OEMCrypto_GenerateDerivedKeys(s.session_id(), mac_context.data(),
                                         mac_context.size(), enc_context.data(),
                                         enc_context.size());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_GenerateDerivedKeys fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoKeyboxTest,
       OEMCryptoMemoryGenerateDerivedKeysForHugeEncContextLength) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  vector<uint8_t> mac_context;
  vector<uint8_t> enc_context;
  s.FillDefaultContext(&mac_context, &enc_context);

  auto oemcrypto_function = [&s, &mac_context,
                             &enc_context](size_t buffer_length) {
    enc_context.resize(buffer_length);
    return OEMCrypto_GenerateDerivedKeys(s.session_id(), mac_context.data(),
                                         mac_context.size(), enc_context.data(),
                                         enc_context.size());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_GetOEMPublicCertificate fails gracefully on a huge
 * buffer.
 */
TEST_F(OEMCryptoProv30Test, OEMCryptoMemoryGetOEMPublicCertForHugeCertLength) {
  if (wrapped_rsa_key_.size() == 0) {
    // If we don't have a wrapped key yet, create one.
    // This wrapped key will be shared by all sessions in the test.
    ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
  }
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  // Install the DRM Cert's RSA key.
  ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_rsa_key_));
  ASSERT_NO_FATAL_FAILURE(s.SetTestRsaPublicKey());

  auto oemcrypto_function = [](size_t input_length) {
    size_t public_cert_length = input_length;
    vector<uint8_t> public_cert(public_cert_length);
    return OEMCrypto_GetOEMPublicCertificate(public_cert.data(),
                                             &public_cert_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_CreateUsageTableHeader fails gracefully on a huge
 * buffer.
 */
TEST_F(OEMCryptoSessionTests,
       OEMCryptoMemoryCreateUsageTableHeaderForHugeHeaderBufferLength) {
  auto oemcrypto_function = [](size_t buffer_length) {
    size_t header_buffer_length = buffer_length;
    vector<uint8_t> usage_table_header(header_buffer_length);
    return OEMCrypto_CreateUsageTableHeader(usage_table_header.data(),
                                            &header_buffer_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_PrepAndSignRenewalRequest fails gracefully on a huge
 * buffer.
 */
TEST_F(OEMCryptoMemoryLicenseTest,
       OEMCryptoMemoryPrepareRenewalRequestForHugeBufferLength) {
  RenewalRoundTrip renewal_messages(&license_messages_);
  auto oemcrypto_function = [&renewal_messages](size_t buffer_length) {
    renewal_messages.set_message_size(buffer_length);
    return renewal_messages.SignAndCreateRequestWithCustomBufferLengths();
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_PrepAndSignRenewalRequest fails gracefully on a huge
 * buffer.
 */
TEST_F(OEMCryptoMemoryLicenseTest,
       OEMCryptoMemoryPrepareRenewalRequestForHugeSignatureLength) {
  RenewalRoundTrip renewal_messages(&license_messages_);
  auto oemcrypto_function = [&renewal_messages](size_t buffer_length) {
    renewal_messages.set_request_signature_size(buffer_length);
    return renewal_messages.SignAndCreateRequestWithCustomBufferLengths();
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_PrepAndSignRenewalRequest fails gracefully on a huge
 * buffer.
 */
TEST_F(OEMCryptoMemoryLicenseTest,
       OEMCryptoMemoryPrepareRenewalRequestForHugeCoreMessageLength) {
  RenewalRoundTrip renewal_messages(&license_messages_);
  auto oemcrypto_function = [&renewal_messages](size_t buffer_length) {
    renewal_messages.set_core_message_size(buffer_length);
    return renewal_messages.SignAndCreateRequestWithCustomBufferLengths();
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that loading entitled content keys fails gracefully on a huge
 * buffer.
 */
TEST_F(OEMCryptoMemoryLicenseTest,
       OEMCryptoMemoryLoadEntitledKeysForHugeBufferLength) {
  auto oemcrypto_function = [&](size_t buffer_length) {
    size_t entitled_key_data_length =
        entitled_message_.entitled_key_data_size();
    vector<uint8_t> message(entitled_key_data_length);
    memcpy(message.data(), entitled_message_.entitled_key_data(),
           entitled_key_data_length);
    message.resize(buffer_length);
    return entitled_message_.LoadKeys(message);
  };
  // We are not constructing a valid message for load entitled content keys.
  // Hence skipping status check.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_LoadLicense fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoSessionTests,
       OEMCryptoMemoryLoadLicenseForHugeSignatureLength) {
  auto oemcrypto_function = [&](size_t signature_size) {
    Session s;
    LicenseRoundTrip license_messages(&s);
    s.open();
    InstallTestDrmKey(&s);
    license_messages.SignAndVerifyRequest();
    license_messages.CreateDefaultResponse();
    license_messages.EncryptAndSignResponse();
    vector<uint8_t> signature(signature_size);
    OEMCryptoResult result = OEMCrypto_LoadLicense(
        s.session_id(), license_messages.encrypted_response_buffer().data(),
        license_messages.encrypted_response_buffer().size(),
        license_messages.serialized_core_message().size(), signature.data(),
        signature_size);
    s.close();
    return result;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_LoadRenewal fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoSessionTests, OEMCryptoMemoryLoadRenewalForHugeResponseLength) {
  auto oemcrypto_function = [&](size_t message_size) {
    Session s;
    LicenseRoundTrip license_messages(&s);
    s.open();
    LoadLicense(s, license_messages);

    RenewalRoundTrip renewal_messages(&license_messages);
    renewal_messages.SignAndVerifyRequest();
    renewal_messages.CreateDefaultResponse();
    renewal_messages.set_message_size(message_size);
    renewal_messages.EncryptAndSignResponse();
    OEMCryptoResult result = renewal_messages.LoadResponse();
    s.close();
    return result;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_LoadRenewal fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoSessionTests,
       OEMCryptoMemoryLoadRenewalForHugeSignatureLength) {
  auto oemcrypto_function = [&](size_t signature_size) {
    Session s;
    LicenseRoundTrip license_messages(&s);
    s.open();
    LoadLicense(s, license_messages);

    RenewalRoundTrip renewal_messages(&license_messages);
    renewal_messages.SignAndVerifyRequest();
    renewal_messages.CreateDefaultResponse();
    renewal_messages.EncryptAndSignResponse();
    vector<uint8_t> signature(signature_size);
    OEMCryptoResult result = OEMCrypto_LoadRenewal(
        s.session_id(), renewal_messages.encrypted_response_buffer().data(),
        renewal_messages.encrypted_response_buffer().size(),
        renewal_messages.serialized_core_message().size(), signature.data(),
        signature_size);
    s.close();
    return result;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_QueryKeyControl fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoSessionTests,
       OEMCryptoMemoryLoadRenewalForHugeCoreMessageLength) {
  auto oemcrypto_function = [&](size_t core_message_size) {
    Session s;
    LicenseRoundTrip license_messages(&s);
    s.open();
    LoadLicense(s, license_messages);

    RenewalRoundTrip renewal_messages(&license_messages);
    renewal_messages.SignAndVerifyRequest();
    renewal_messages.CreateDefaultResponse();
    renewal_messages.set_core_message_size(core_message_size);
    renewal_messages.EncryptAndSignResponse();
    OEMCryptoResult result = renewal_messages.LoadResponse();
    s.close();
    return result;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_QueryKeyControl fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoSessionTests,
       OEMCryptoMemoryQueryKeyControlForHugeKeyIdLength) {
  Session session;
  LicenseRoundTrip license_messages(&session);
  session.open();
  LoadLicense(session, license_messages);
  OEMCrypto_SESSION session_id = session.session_id();
  vector<uint8_t> valid_key_id(
      license_messages.response_data().keys[0].key_id,
      license_messages.response_data().keys[0].key_id + kTestKeyIdMaxLength);
  auto oemcrypto_function = [&session_id,
                             &valid_key_id](size_t additional_key_id_length) {
    vector<uint8_t> key_id(valid_key_id);
    key_id.resize(valid_key_id.size() + additional_key_id_length);
    KeyControlBlock block;
    size_t size = sizeof(block);
    return OEMCrypto_QueryKeyControl(session_id, key_id.data(), key_id.size(),
                                     reinterpret_cast<uint8_t*>(&block), &size);
  };
  // We do not want to stop as soon as API results in an error as it would
  // return error on first iteration as key_id is invalid.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test OEMCrypto_QueryKeyControl doesn't crash for huge key_control_block
 * length.
 */
TEST_F(OEMCryptoSessionTests,
       OEMCryptoMemoryQueryKeyControlForHugeKeyControlBlockLength) {
  Session session;
  LicenseRoundTrip license_messages(&session);
  session.open();
  LoadLicense(session, license_messages);
  OEMCrypto_SESSION session_id = session.session_id();
  uint8_t* key_id = license_messages.response_data().keys[0].key_id;
  size_t key_id_length = license_messages.response_data().keys[0].key_id_length;
  auto oemcrypto_function = [&session_id, &key_id,
                             &key_id_length](size_t buffer_length) {
    size_t key_control_block_length = buffer_length;
    vector<uint8_t> key_control_block(key_control_block_length);
    return OEMCrypto_QueryKeyControl(session_id, key_id, key_id_length,
                                     key_control_block.data(),
                                     &key_control_block_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** This test verifies that OEMCrypto_SetDecryptHash doesn't crash for a very
    large hash buffer.
*/
TEST_F(OEMCryptoMemoryLicenseTest,
       OEMCryptoMemoryDecryptHashForHugeHashBuffer) {
  uint32_t session_id = session_.session_id();
  auto f = [session_id]() {
    const uint32_t frame_number = 1;
    const uint32_t crc32 = 0;
    return OEMCrypto_SetDecryptHash(session_id, frame_number, crc32);
  };
  TestHugeLengthDoesNotCrashAPI(f, kCheckStatus);
}

/** Test Decrypt fails gracefully for huge input. */
TEST_P(OEMCryptoSessionTestsDecryptTests,
       OEMCryptoMemoryDecryptCENCForHugeNumberOfSubSamples) {
  auto oemcrypto_function = [&](size_t number_of_subsamples) {
    std::vector<SubsampleSize> subsample_sizes;
    while (number_of_subsamples-- > 0) {
      subsample_sizes.push_back({1, 1});
    }
    SetSubsampleSizes(subsample_sizes);
    LoadLicense();
    MakeBuffers();
    EncryptData();
    OEMCryptoResult result = DecryptCENC();
    FreeSecureBuffers();
    // Closing the session and opening it for next iteration.
    // If it is last iteration, session will be closed in teardown method of
    // class.
    session_.close();
    session_.open();
    InstallTestDrmKey(&session_);
    return result;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, 1, 2 * MiB, kCheckStatus);

  // Avoid double free when test teardown calls FreeSecureBuffers()
  MakeBuffers();
}

/** Test Decrypt fails gracefully for huge input. */
TEST_P(OEMCryptoSessionTestsDecryptTests,
       OEMCryptoMemoryDecryptCENCForHugeNumberOfSamples) {
  auto oemcrypto_function = [&](size_t number_of_samples) {
    std::vector<std::vector<SubsampleSize>> samples;
    std::vector<SubsampleSize> subsample_sizes;
    subsample_sizes.push_back({1, 1});
    while (number_of_samples-- > 0) {
      samples.push_back(subsample_sizes);
    }
    SetSampleSizes(samples);
    LoadLicense();
    MakeBuffers();
    EncryptData();
    OEMCryptoResult result = DecryptCENC();
    FreeSecureBuffers();
    // Closing the session and opening it for next iteration.
    // If it is last iteration, session will be closed in teardown method of
    // class.
    session_.close();
    session_.open();
    InstallTestDrmKey(&session_);
    return result;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, 1, 2 * MiB, kCheckStatus);

  // Avoid double free when test teardown calls FreeSecureBuffers()
  MakeBuffers();
}

/** Test that OEMCrypto_LoadProvisioning fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoLoadsCertificate,
       OEMCryptoMemoryLoadProvisioningForHugeSignatureLength) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (!global_features.loads_certificate ||
      global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  auto oemcrypto_function = [&](size_t signature_size) {
    Session s;
    ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
    provisioning_messages.PrepareSession(keybox_);
    provisioning_messages.SignAndVerifyRequest();
    provisioning_messages.CreateDefaultResponse();
    provisioning_messages.EncryptAndSignResponse();
    vector<uint8_t> signature(signature_size);
    size_t wrapped_private_key_length = 0;
    // Find wrapped_private_key_length.
    OEMCrypto_LoadProvisioning(
        s.session_id(),
        provisioning_messages.encrypted_response_buffer().data(),
        provisioning_messages.encrypted_response_buffer().size(),
        provisioning_messages.serialized_core_message().size(),
        signature.data(), signature_size, nullptr, &wrapped_private_key_length);
    std::vector<uint8_t> wrapped_rsa_key(wrapped_private_key_length);
    OEMCryptoResult result = OEMCrypto_LoadProvisioning(
        s.session_id(),
        provisioning_messages.encrypted_response_buffer().data(),
        provisioning_messages.encrypted_response_buffer().size(),
        provisioning_messages.serialized_core_message().size(),
        signature.data(), signature_size, wrapped_rsa_key.data(),
        &wrapped_private_key_length);
    s.close();
    return result;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_LoadProvisioning fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoLoadsCertificate,
       OEMCryptoMemoryLoadProvisioningForHugeWrappedRsaKeyLength) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (!global_features.loads_certificate ||
      global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  auto oemcrypto_function = [&](size_t buffer_length) {
    Session s;
    ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
    provisioning_messages.PrepareSession(keybox_);
    provisioning_messages.SignAndVerifyRequest();
    provisioning_messages.CreateDefaultResponse();
    provisioning_messages.EncryptAndSignResponse();
    size_t wrapped_private_key_length = buffer_length;
    vector<uint8_t> wrapped_private_key(wrapped_private_key_length);
    OEMCryptoResult result = OEMCrypto_LoadProvisioning(
        s.session_id(),
        provisioning_messages.encrypted_response_buffer().data(),
        provisioning_messages.encrypted_response_buffer().size(),
        provisioning_messages.serialized_core_message().size(),
        provisioning_messages.response_signature().data(),
        provisioning_messages.response_signature().size(),
        wrapped_private_key.data(), &wrapped_private_key_length);
    s.close();
    return result;
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_LoadDRMPrivateKey fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoLoadsCertificate,
       OEMCryptoMemoryLoadDrmPrivateKeyForHugeWrappedRsaKeyLength) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (!global_features.loads_certificate ||
      global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
  auto oemcrypto_function = [&](size_t wrapped_rsa_key_length) {
    Session s;
    s.open();
    vector<uint8_t> wrapped_rsa_key_buffer = wrapped_rsa_key_;
    wrapped_rsa_key_buffer.resize(wrapped_rsa_key_length);
    OEMCryptoResult result = OEMCrypto_LoadDRMPrivateKey(
        s.session_id(), OEMCrypto_RSA_Private_Key,
        wrapped_rsa_key_buffer.data(), wrapped_rsa_key_buffer.size());
    s.close();
    return result;
  };
  // It is hard to generate varying length valid wrapped rsa key with valid
  // signature. Hence we just call function with random data and do not check
  // status to test API with varying length wrapped rsa key.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, wrapped_rsa_key_.size(),
                                kHugeInputBufferLength, !kCheckStatus);
}

/** Test that OEMCrypto_LoadDRMPrivateKey fails gracefully on a huge buffer.
 */
TEST_F(
    OEMCryptoLoadsCertificate,
    OEMCryptoMemoryLoadDrmPrivateKeyForHugeWrappedRsaKeyLengthStartingFromLength1) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (!global_features.loads_certificate ||
      global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
  auto oemcrypto_function = [&](size_t wrapped_rsa_key_length) {
    Session s;
    s.open();
    vector<uint8_t> wrapped_rsa_key_buffer(wrapped_rsa_key_length);
    OEMCryptoResult result = OEMCrypto_LoadDRMPrivateKey(
        s.session_id(), OEMCrypto_RSA_Private_Key,
        wrapped_rsa_key_buffer.data(), wrapped_rsa_key_buffer.size());
    s.close();
    return result;
  };
  // It is hard to generate varying length valid wrapped rsa key with valid
  // signature. Hence we just call function with random data and do not check
  // status to test API with varying length wrapped rsa key.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_LoadDRMPrivateKey fails gracefully on a huge buffer.
 */
TEST_F(OEMCryptoUsesCertificate,
       OEMCryptoMemoryDeriveKeysFromSessionKeyForHugeMacContext) {
  vector<uint8_t> session_key;
  vector<uint8_t> enc_session_key;
  ASSERT_TRUE(session_.GenerateSessionKey(&session_key, &enc_session_key));
  vector<uint8_t> mac_context;
  vector<uint8_t> enc_context;
  session_.FillDefaultContext(&mac_context, &enc_context);
  OEMCrypto_SESSION session_id = session_.session_id();
  auto oemcrypto_function = [&session_id, &enc_context, &mac_context,
                             &enc_session_key](size_t buffer_length) {
    mac_context.resize(buffer_length);
    return OEMCrypto_DeriveKeysFromSessionKey(
        session_id, enc_session_key.data(), enc_session_key.size(),
        mac_context.data(), mac_context.size(), enc_context.data(),
        enc_context.size());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_DeriveKeysFromSessionKey fails gracefully on a huge
 * buffer. */
TEST_F(OEMCryptoUsesCertificate,
       OEMCryptoMemoryDeriveKeysFromSessionKeyForHugeEncContext) {
  vector<uint8_t> session_key;
  vector<uint8_t> enc_session_key;
  ASSERT_TRUE(session_.GenerateSessionKey(&session_key, &enc_session_key));
  vector<uint8_t> mac_context;
  vector<uint8_t> enc_context;
  session_.FillDefaultContext(&mac_context, &enc_context);
  OEMCrypto_SESSION session_id = session_.session_id();
  auto oemcrypto_function = [&session_id, &enc_context, &mac_context,
                             &enc_session_key](size_t buffer_length) {
    enc_context.resize(buffer_length);
    return OEMCrypto_DeriveKeysFromSessionKey(
        session_id, enc_session_key.data(), enc_session_key.size(),
        mac_context.data(), mac_context.size(), enc_context.data(),
        enc_context.size());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that OEMCrypto_DeriveKeysFromSessionKey fails gracefully on a huge
 * buffer. */
TEST_F(OEMCryptoUsesCertificate,
       OEMCryptoMemoryDeriveKeysFromSessionKeyForHugeEncSessionKey) {
  vector<uint8_t> session_key;
  vector<uint8_t> enc_session_key;
  ASSERT_TRUE(session_.GenerateSessionKey(&session_key, &enc_session_key));
  vector<uint8_t> mac_context;
  vector<uint8_t> enc_context;
  session_.FillDefaultContext(&mac_context, &enc_context);
  OEMCrypto_SESSION session_id = session_.session_id();
  auto oemcrypto_function = [&session_id, &session_key, &enc_context,
                             &mac_context,
                             &enc_session_key](size_t buffer_length) {
    session_key.resize(buffer_length);
    return OEMCrypto_DeriveKeysFromSessionKey(
        session_id, enc_session_key.data(), enc_session_key.size(),
        mac_context.data(), mac_context.size(), enc_context.data(),
        enc_context.size());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that OEMCrypto_GenerateRSASignature fails gracefully on a huge
 * buffer. */
TEST_F(OEMCryptoLoadsCertificateAlternates,
       OEMCryptoMemoryGenerateRSASignatureForHugeBuffer) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (!global_features.loads_certificate ||
      global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  OEMCryptoResult sts;
  LoadCastCertificateKey(false);
  // If the device is a cast receiver, then this scheme is required.
  if (global_features.cast_receiver) {
    ASSERT_TRUE(key_loaded_);
  }
  if (key_loaded_) {
    // If the key did load, then it should be processed correctly.
    Session s;
    ASSERT_NO_FATAL_FAILURE(s.open());
    ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_rsa_key_));

    vector<uint8_t> message_buffer(10);
    size_t signature_length = 0;
    sts = OEMCrypto_GenerateRSASignature(s.session_id(), message_buffer.data(),
                                         message_buffer.size(), nullptr,
                                         &signature_length, kSign_PKCS1_Block1);
    if (sts == OEMCrypto_ERROR_NOT_IMPLEMENTED) return;
    ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, sts);
    ASSERT_NE(static_cast<size_t>(0), signature_length);
    vector<uint8_t> signature(signature_length);

    auto oemcrypto_function = [&](size_t buffer_length) {
      message_buffer.resize(buffer_length);
      return OEMCrypto_GenerateRSASignature(
          s.session_id(), message_buffer.data(), message_buffer.size(),
          signature.data(), &signature_length, kSign_PKCS1_Block1);
    };
    TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
    s.close();
  }
}

/** Test that OEMCrypto_GenerateRSASignature fails gracefully on a huge
 * buffer. */
TEST_F(OEMCryptoLoadsCertificateAlternates,
       OEMCryptoMemoryGenerateRSASignatureForHugeSignatureLength) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (!global_features.loads_certificate ||
      global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  LoadCastCertificateKey(false);
  // If the device is a cast receiver, then this scheme is required.
  if (global_features.cast_receiver) {
    ASSERT_TRUE(key_loaded_);
  }
  if (key_loaded_) {
    // If the key did load, then it should be processed correctly.
    Session s;
    ASSERT_NO_FATAL_FAILURE(s.open());
    ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_rsa_key_));

    vector<uint8_t> message_buffer(50);
    vector<uint8_t> signature;
    auto oemcrypto_function = [&](size_t signature_length) {
      signature.resize(signature_length);
      return OEMCrypto_GenerateRSASignature(
          s.session_id(), message_buffer.data(), message_buffer.size(),
          signature.data(), &signature_length, kSign_PKCS1_Block1);
    };
    TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
    s.close();
  }
}

/** Test that GetKeyHandle fails gracefully on a huge buffer. */
TEST_P(OEMCryptoGenericCryptoTest, OEMCryptoMemorySelectKeyForHugeKeyIdLength) {
  EncryptAndLoadKeys();
  OEMCrypto_SESSION session_id = session_.session_id();
  auto oemcrypto_function = [session_id](size_t key_id_length) {
    vector<uint8_t> key_id(key_id_length);
    vector<uint8_t> key_handle;
    return GetKeyHandleIntoVector(session_id, key_id.data(), key_id.size(),
                                  OEMCrypto_CipherMode_CENC, key_handle);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoGenericCryptoTest,
       OEMCryptoMemoryGenericKeyEncryptForHugeBuffer) {
  EncryptAndLoadKeys();
  unsigned int key_index = 0;
  vector<uint8_t> expected_encrypted;
  vector<uint8_t> key_handle;
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GetKeyHandleIntoVector(session_.session_id(),
                             session_.license().keys[key_index].key_id,
                             session_.license().keys[key_index].key_id_length,
                             OEMCrypto_CipherMode_CENC, key_handle));
  OEMCrypto_SESSION session_id = session_.session_id();
  auto& iv = iv_;
  auto oemcrypto_function = [&session_id, &iv](size_t buffer_length) mutable {
    vector<uint8_t> buffer(buffer_length);
    return OEMCrypto_Generic_Encrypt(
        key_handle.data(), key_handle.size(), buffer.data(), buffer.size(), iv,
        OEMCrypto_AES_CBC_128_NO_PADDING, buffer.data());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, 16, kHugeInputBufferLength,
                                kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoGenericCryptoTest,
       OEMCryptoMemoryGenericKeyDecryptForHugeBuffer) {
  EncryptAndLoadKeys();
  unsigned int key_index = 1;
  vector<uint8_t> key_handle;
  GetKeyHandleIntoVector(session_.session_id(),
                         session_.license().keys[key_index].key_id,
                         session_.license().keys[key_index].key_id_length,
                         OEMCrypto_CipherMode_CENC, key_handle);
  OEMCrypto_SESSION session_id = session_.session_id();
  auto iv = iv_;
  auto oemcrypto_function = [&session_id, &iv](size_t buffer_length) {
    vector<uint8_t> encrypted(buffer_length);
    vector<uint8_t> resultant(encrypted.size());

    return OEMCrypto_Generic_Decrypt(key_handle.data(), key_handle.size(),
                                     encrypted.data(), encrypted.size(), iv,
                                     OEMCrypto_AES_CBC_128_NO_PADDING,
                                     resultant.data());
  };
  // API expects length to be multiple of 16. Starting from 16.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, 16, kHugeInputBufferLength,
                                kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoGenericCryptoTest, OEMCryptoMemoryGenericKeySignForHugeBuffer) {
  EncryptAndLoadKeys();
  unsigned int key_index = 2;
  vector<uint8_t> key_handle;
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GetKeyHandleIntoVector(session_.session_id(),
                             session_.license().keys[key_index].key_id,
                             session_.license().keys[key_index].key_id_length,
                             OEMCrypto_CipherMode_CENC),
      key_handle);
  vector<uint8_t> signature(SHA256_DIGEST_LENGTH);
  size_t signature_length = signature.size();
  OEMCrypto_SESSION session_id = session_.session_id();
  auto oemcrypto_function = [&session_id, &signature,
                             &signature_length](size_t buffer_length) {
    vector<uint8_t> buffer(buffer_length);
    return OEMCrypto_Generic_Sign(
        key_handle.data(), key_handle.size(), buffer.data(), buffer.size(),
        OEMCrypto_HMAC_SHA256, signature.data(), &signature_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoGenericCryptoTest,
       OEMCryptoMemoryGenericKeySignForHugeSignatureLength) {
  EncryptAndLoadKeys();
  unsigned int key_index = 2;
  vector<uint8_t> key_handle;
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GetKeyHandleIntoVector(session_.session_id(),
                             session_.license().keys[key_index].key_id,
                             session_.license().keys[key_index].key_id_length,
                             OEMCrypto_CipherMode_CENC),
      key_handle);
  OEMCrypto_SESSION session_id = session_.session_id();
  auto clear_buffer = clear_buffer_;
  auto oemcrypto_function = [&session_id,
                             &clear_buffer](size_t signature_length) {
    vector<uint8_t> signature(signature_length);
    size_t gen_signature_length = signature_length;
    return OEMCrypto_Generic_Sign(key_handle.data(), key_handle.size(),
                                  clear_buffer.data(), clear_buffer.size(),
                                  OEMCrypto_HMAC_SHA256, signature.data(),
                                  &gen_signature_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoGenericCryptoTest,
       OEMCryptoMemoryGenericKeyVerifyForHugeBuffer) {
  EncryptAndLoadKeys();
  unsigned int key_index = 3;
  vector<uint8_t> key_handle;
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GetKeyHandleIntoVector(session_.session_id(),
                             session_.license().keys[key_index].key_id,
                             session_.license().keys[key_index].key_id_length,
                             OEMCrypto_CipherMode_CENC, key_handle));
  auto oemcrypto_function = [&](size_t buffer_length) {
    vector<uint8_t> buffer(buffer_length);
    vector<uint8_t> signature;
    SignBuffer(key_index, buffer, &signature);
    return GenericVerify(key_handle.data(), key_handle.size(), buffer.data(),
                         buffer.size(), OEMCrypto_HMAC_SHA256, signature.data(),
                         signature.size());
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoGenericCryptoTest,
       OEMCryptoMemoryGenericKeyVerifyForHugeSignatureLength) {
  EncryptAndLoadKeys();
  unsigned int key_index = 3;
  vector<uint8_t> signature;
  SignBuffer(key_index, clear_buffer_, &signature);

  vector<uint8_t> key_handle;
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GetKeyHandleIntoVector(session_.session_id(),
                             session_.license().keys[key_index].key_id,
                             session_.license().keys[key_index].key_id_length,
                             OEMCrypto_CipherMode_CENC, key_handle));
  OEMCrypto_SESSION session_id = session_.session_id();
  auto clear_buffer = clear_buffer_;
  auto oemcrypto_function = [&session_id, &clear_buffer,
                             &signature](size_t signature_length) {
    return OEMCrypto_Generic_Verify(key_handle.data(), key_handle.size(),
                                    clear_buffer.data(), clear_buffer.size(),
                                    OEMCrypto_HMAC_SHA256, signature.data(),
                                    signature_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoUsageTableTest,
       OEMCryptoMemoryUpdateUsageEntryForHugeHeaderBuffer) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Usage tables are not supported.";
  }
  auto oemcrypto_function = [&](size_t buffer_length) {
    if (buffer_length < encrypted_usage_header_.size()) {
      return OEMCrypto_ERROR_SHORT_BUFFER;
    }
    LicenseWithUsageEntry entry;
    entry.license_messages().set_api_version(license_api_version_);
    entry.MakeAndLoadOnline(this);
    Session& s = entry.session();
    size_t header_buffer_length = 0;
    size_t entry_buffer_length = 0;
    // Header buffer length varies as generation_numbers size changes on every
    // call. Hence, we need to call update usage entry in every loop to get
    // latest value of header_buffer_length.
    OEMCrypto_UpdateUsageEntry(s.session_id(), nullptr, &header_buffer_length,
                               nullptr, &entry_buffer_length);
    vector<uint8_t> encrypted_usage_entry(entry_buffer_length);
    vector<uint8_t> header_buffer(encrypted_usage_header_);
    header_buffer.resize(buffer_length);
    return OEMCrypto_UpdateUsageEntry(
        s.session_id(), header_buffer.data(), &buffer_length,
        encrypted_usage_entry.data(), &entry_buffer_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoUsageTableTest,
       OEMCryptoMemoryUpdateUsageEntryForHugeUsageEntryBuffer) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Usage tables are not supported.";
  }
  auto oemcrypto_function = [&](size_t buffer_length) {
    LicenseWithUsageEntry entry;
    entry.license_messages().set_api_version(license_api_version_);
    entry.MakeAndLoadOnline(this);
    Session& s = entry.session();
    size_t header_buffer_length = 0;
    size_t entry_buffer_length = 0;
    // Header buffer length varies as generation_numbers size changes on every
    // call. Hence, we need to call update usage entry in every loop to get
    // latest value of header_buffer_length.
    OEMCrypto_UpdateUsageEntry(s.session_id(), nullptr, &header_buffer_length,
                               nullptr, &entry_buffer_length);
    vector<uint8_t> header_buffer(encrypted_usage_header_);
    header_buffer.resize(header_buffer_length);
    vector<uint8_t> encrypted_usage_entry(buffer_length);
    return OEMCrypto_UpdateUsageEntry(
        s.session_id(), header_buffer.data(), &header_buffer_length,
        encrypted_usage_entry.data(), &buffer_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoUsageTableTest,
       OEMCryptoMemoryDeactivateUsageEntryForHugePstBuffer) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Usage tables are not supported.";
  }
  auto oemcrypto_function = [&](size_t buffer_length) {
    LicenseWithUsageEntry entry;
    std::string pst("pst");
    pst.resize(buffer_length);
    entry.license_messages().set_pst(pst);
    entry.license_messages().set_api_version(license_api_version_);
    entry.MakeAndLoadOnline(this);
    Session& s = entry.session();
    return OEMCrypto_DeactivateUsageEntry(
        s.session_id(), reinterpret_cast<const uint8_t*>(pst.c_str()),
        pst.length());
  };
  // The test setup assertions fails if pst length goes beyond kMaxPSTLength.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, 1, kMaxPSTLength,
                                kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoUsageTableTest,
       OEMCryptoMemoryLoadUsageTableHeaderForHugeHeader) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Usage tables are not supported.";
  }
  auto oemcrypto_function = [&](size_t buffer_length) {
    if (buffer_length < encrypted_usage_header_.size()) {
      return OEMCrypto_ERROR_SHORT_BUFFER;
    }

    LicenseWithUsageEntry entry;
    entry.license_messages().set_api_version(license_api_version_);
    entry.MakeOfflineAndClose(this);
    vector<uint8_t> header_buffer(encrypted_usage_header_);
    header_buffer.resize(buffer_length);
    return OEMCrypto_LoadUsageTableHeader(header_buffer.data(),
                                          header_buffer.size());
  };
  // We cannot generate an encrypted usage header of varying length with
  // valid signature. Hence irrespective of return status, we call API for
  // varying buffer lengths.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function,
                                encrypted_usage_header_.size(),
                                kHugeInputBufferLength, !kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(
    OEMCryptoUsageTableTest,
    OEMCryptoMemoryLoadUsageTableHeaderForHugeHeaderStartingHeaderLengthFrom1) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Usage tables are not supported.";
  }
  auto oemcrypto_function = [&](size_t buffer_length) {
    LicenseWithUsageEntry entry;
    entry.license_messages().set_api_version(license_api_version_);
    entry.MakeOfflineAndClose(this);
    vector<uint8_t> header_buffer(buffer_length);
    return OEMCrypto_LoadUsageTableHeader(header_buffer.data(),
                                          header_buffer.size());
  };
  // We cannot generate an encrypted usage header of varying length with
  // valid signature. Hence irrespective of return status, we call API for
  // varying buffer lengths.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoUsageTableTest,
       OEMCryptoMemoryLoadUsageEntryForHugeUsageEntryBuffer) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Usage tables are not supported.";
  }
  auto oemcrypto_function = [&](size_t buffer_length) {
    LicenseWithUsageEntry entry;
    entry.license_messages().set_api_version(license_api_version_);
    Session& s = entry.session();
    // Make first entry 0.
    entry.MakeOfflineAndClose(this);
    if (buffer_length < s.encrypted_usage_entry().size()) {
      return OEMCrypto_ERROR_SHORT_BUFFER;
    }
    Session s2;
    s2.open();
    InstallTestDrmKey(&s2);
    vector<uint8_t> encrypted_usage_entry(buffer_length);
    memcpy(encrypted_usage_entry.data(), s.encrypted_usage_entry().data(),
           s.encrypted_usage_entry().size());
    const uint32_t usage_entry_number = s.usage_entry_number();
    return OEMCrypto_LoadUsageEntry(s2.session_id(), usage_entry_number,
                                    encrypted_usage_entry.data(),
                                    encrypted_usage_entry.size());
  };
  // We cannot generate an encrypted usage enctry of varying length with
  // valid signature. Hence irrespective of return status, we call API for
  // varying buffer lengths.
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoUsageTableTest, OEMCryptoMemoryReportUsageForHugeReportBuffer) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Usage tables are not supported.";
  }
  auto oemcrypto_function = [&](size_t buffer_length) {
    LicenseWithUsageEntry entry;
    entry.license_messages().set_api_version(license_api_version_);
    entry.MakeAndLoadOnline(this);
    Session& s = entry.session();
    s.UpdateUsageEntry(&encrypted_usage_header_);
    size_t length = 0;
    OEMCrypto_ReportUsage(s.session_id(),
                          reinterpret_cast<const uint8_t*>(entry.pst().c_str()),
                          entry.pst().length(), nullptr, &length);
    if (ShouldGenerateCorpus()) {
      const std::string file_name =
          GetFileName("oemcrypto_report_usage_fuzz_seed_corpus");
      AppendToFile(file_name, reinterpret_cast<const char*>(&length),
                   sizeof(length));
      AppendToFile(file_name, entry.pst().c_str(), entry.pst().length());
    }
    vector<uint8_t> pst_report_buffer(buffer_length);
    return OEMCrypto_ReportUsage(
        s.session_id(), reinterpret_cast<const uint8_t*>(entry.pst().c_str()),
        entry.pst().length(), pst_report_buffer.data(), &buffer_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoUsageTableTest, OEMCryptoMemoryReportUsageForHugePstBuffer) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Usage tables are not supported.";
  }
  auto oemcrypto_function = [&](size_t buffer_length) {
    LicenseWithUsageEntry entry;
    entry.license_messages().set_api_version(license_api_version_);
    entry.MakeAndLoadOnline(this);
    Session& s = entry.session();
    s.UpdateUsageEntry(&encrypted_usage_header_);
    size_t length = 0;
    OEMCrypto_ReportUsage(s.session_id(),
                          reinterpret_cast<const uint8_t*>(entry.pst().c_str()),
                          entry.pst().length(), nullptr, &length);
    vector<uint8_t> pst_report_buffer(length);
    vector<uint8_t> pst(buffer_length);
    return OEMCrypto_ReportUsage(s.session_id(), pst.data(), pst.size(),
                                 pst_report_buffer.data(), &length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, !kCheckStatus);
}

/** Test that API fails gracefully on a huge buffer. */
TEST_P(OEMCryptoUsageTableTest,
       OEMCryptoMemoryShrinkUsageTableHeaderForHugeHeaderBufferLength) {
  if (!wvoec::global_features.usage_table) {
    GTEST_SKIP() << "Usage tables are not supported.";
  }
  LicenseWithUsageEntry entry0;
  entry0.set_pst("pst 0");
  entry0.MakeOfflineAndClose(this);
  auto oemcrypto_function = [&](size_t buffer_length) {
    LicenseWithUsageEntry entry1;
    entry1.set_pst("pst 1");
    entry1.MakeOfflineAndClose(this);
    size_t header_buffer_length = buffer_length;
    encrypted_usage_header_.resize(header_buffer_length);
    return OEMCrypto_ShrinkUsageTableHeader(1, encrypted_usage_header_.data(),
                                            &header_buffer_length);
  };
  TestHugeLengthDoesNotCrashAPI(oemcrypto_function, kCheckStatus);
}
/// @}

}  // namespace wvoec