android/libwvdrmengine/oemcrypto/test/oemcrypto_security_test.cpp

// 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_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_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_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_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_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 set sandbox doesn't crash for a large sandbox id leangth.
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_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);
}

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);
}

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

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);
}

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_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_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_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_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_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_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_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_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_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);
}

// This verifies that entitled content keys API does not crash for unreasonable
// input message buffer lengths.
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_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_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_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_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 OEMCrypto_QueryKeyControl doesn't crash for huge key_id_length.
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](size_t hash_length) {
    uint32_t frame_number = 1;
    vector<uint8_t> hash_buffer(hash_length);
    return OEMCrypto_SetDecryptHash(session_id, frame_number,
                                    hash_buffer.data(), hash_buffer.size());
  };
  TestHugeLengthDoesNotCrashAPI(f, kCheckStatus);
}

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_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_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_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_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_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_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_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_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_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);
    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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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