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

#include <gtest/gtest.h>

#include "oemcrypto_usage_table_test.h"

using ::testing::Range;

namespace wvoec {

/// @addtogroup generic
/// @{

TEST_P(OEMCryptoGenericCryptoTest, GenericKeyLoad) { EncryptAndLoadKeys(); }

/** Test that the Generic_Encrypt function works correctly. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyEncrypt) {
  EncryptAndLoadKeys();
  unsigned int key_index = 0;
  vector<uint8_t> expected_encrypted;
  EncryptBuffer(key_index, clear_buffer_, &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));
  vector<uint8_t> encrypted(clear_buffer_.size());
  ASSERT_EQ(OEMCrypto_SUCCESS,
            GenericEncrypt(key_handle.data(), key_handle.size(),
                           clear_buffer_.data(), clear_buffer_.size(), iv_,
                           OEMCrypto_AES_CBC_128_NO_PADDING, encrypted.data()));
  ASSERT_EQ(expected_encrypted, encrypted);
}

/** Test that the Generic_Encrypt function fails when not allowed. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyBadEncrypt) {
  EncryptAndLoadKeys();
  BadEncrypt(0, OEMCrypto_HMAC_SHA256, buffer_size_);
  // The buffer size must be a multiple of 16, so subtracting 10 is bad.
  BadEncrypt(0, OEMCrypto_AES_CBC_128_NO_PADDING, buffer_size_ - 10);
  BadEncrypt(1, OEMCrypto_AES_CBC_128_NO_PADDING, buffer_size_);
  BadEncrypt(2, OEMCrypto_AES_CBC_128_NO_PADDING, buffer_size_);
  BadEncrypt(3, OEMCrypto_AES_CBC_128_NO_PADDING, buffer_size_);
}

/** Test that the Generic_Encrypt works if the input and output buffers are the
 * same. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyEncryptSameBufferAPI12) {
  EncryptAndLoadKeys();
  unsigned int key_index = 0;
  vector<uint8_t> expected_encrypted;
  EncryptBuffer(key_index, clear_buffer_, &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));
  // Input and output are same buffer:
  vector<uint8_t> buffer = clear_buffer_;
  ASSERT_EQ(OEMCrypto_SUCCESS,
            GenericEncrypt(key_handle.data(), key_handle.size(), buffer.data(),
                           buffer.size(), iv_, OEMCrypto_AES_CBC_128_NO_PADDING,
                           buffer.data()));
  ASSERT_EQ(expected_encrypted, buffer);
}

TEST_P(
    OEMCryptoGenericCryptoTest,
    OEMCryptoMemoryGenericKeyEncryptForHugeBufferWithBufferLengthNotMultipleOf16) {
  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));
  vector<uint8_t> buffer(17);
  ASSERT_NO_FATAL_FAILURE(OEMCrypto_Generic_Encrypt(
      key_handle.data(), key_handle.size(), buffer.data(), buffer.size(), iv_,
      OEMCrypto_AES_CBC_128_NO_PADDING, buffer.data()));
}

/** Test Generic_Decrypt works correctly. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyDecrypt) {
  EncryptAndLoadKeys();
  unsigned int key_index = 1;
  vector<uint8_t> encrypted;
  EncryptBuffer(key_index, clear_buffer_, &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));
  vector<uint8_t> resultant(encrypted.size());
  ASSERT_EQ(OEMCrypto_SUCCESS,
            GenericDecrypt(key_handle.data(), key_handle.size(),
                           encrypted.data(), encrypted.size(), iv_,
                           OEMCrypto_AES_CBC_128_NO_PADDING, resultant.data()));
  ASSERT_EQ(clear_buffer_, resultant);
}

/** Test that Generic_Decrypt works correctly when the input and output buffers
 * are the same. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyDecryptSameBufferAPI12) {
  EncryptAndLoadKeys();
  unsigned int key_index = 1;
  vector<uint8_t> encrypted;
  EncryptBuffer(key_index, clear_buffer_, &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));
  vector<uint8_t> resultant(encrypted.size());
  ASSERT_EQ(OEMCrypto_SUCCESS,
            GenericDecrypt(key_handle.data(), key_handle.size(),
                           encrypted.data(), encrypted.size(), iv_,
                           OEMCrypto_AES_CBC_128_NO_PADDING, resultant.data()));
  ASSERT_EQ(clear_buffer_, resultant);
}

/** Test that Generic_Decrypt fails to decrypt to an insecure buffer if the key
 * requires a secure data path. */
TEST_P(OEMCryptoGenericCryptoTest, GenericSecureToClear) {
  license_messages_.set_control(wvoec::kControlObserveDataPath |
                                wvoec::kControlDataPathSecure);
  license_messages_.CreateResponseWithGenericCryptoKeys();
  EncryptAndLoadKeys();
  unsigned int key_index = 1;
  vector<uint8_t> encrypted;
  EncryptBuffer(key_index, clear_buffer_, &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));
  vector<uint8_t> resultant(encrypted.size());
  ASSERT_NE(OEMCrypto_SUCCESS,
            GenericDecrypt(key_handle.data(), key_handle.size(),
                           encrypted.data(), encrypted.size(), iv_,
                           OEMCrypto_AES_CBC_128_NO_PADDING, resultant.data()));
  ASSERT_NE(clear_buffer_, resultant);
}

/** Test that the Generic_Decrypt function fails when not allowed. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyBadDecrypt) {
  EncryptAndLoadKeys();
  BadDecrypt(1, OEMCrypto_HMAC_SHA256, buffer_size_);
  // The buffer size must be a multiple of 16, so subtracting 10 is bad.
  BadDecrypt(1, OEMCrypto_AES_CBC_128_NO_PADDING, buffer_size_ - 10);
  BadDecrypt(0, OEMCrypto_AES_CBC_128_NO_PADDING, buffer_size_);
  BadDecrypt(2, OEMCrypto_AES_CBC_128_NO_PADDING, buffer_size_);
  BadDecrypt(3, OEMCrypto_AES_CBC_128_NO_PADDING, buffer_size_);
}

TEST_P(OEMCryptoGenericCryptoTest, GenericKeySign) {
  EncryptAndLoadKeys();
  unsigned int key_index = 2;
  vector<uint8_t> expected_signature;
  SignBuffer(key_index, clear_buffer_, &expected_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));
  size_t gen_signature_length = 0;
  ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER,
            GenericSign(key_handle.data(), key_handle.size(),
                        clear_buffer_.data(), clear_buffer_.size(),
                        OEMCrypto_HMAC_SHA256, nullptr, &gen_signature_length));
  ASSERT_EQ(static_cast<size_t>(SHA256_DIGEST_LENGTH), gen_signature_length);
  vector<uint8_t> signature(SHA256_DIGEST_LENGTH);
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GenericSign(key_handle.data(), key_handle.size(), clear_buffer_.data(),
                  clear_buffer_.size(), OEMCrypto_HMAC_SHA256, signature.data(),
                  &gen_signature_length));
  ASSERT_EQ(expected_signature, signature);
}

/** Test that the Generic_Sign function fails when not allowed. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyBadSign) {
  EncryptAndLoadKeys();
  BadSign(0, OEMCrypto_HMAC_SHA256);             // Can't sign with encrypt key.
  BadSign(1, OEMCrypto_HMAC_SHA256);             // Can't sign with decrypt key.
  BadSign(3, OEMCrypto_HMAC_SHA256);             // Can't sign with verify key.
  BadSign(2, OEMCrypto_AES_CBC_128_NO_PADDING);  // Bad signing algorithm.
}

TEST_P(OEMCryptoGenericCryptoTest, GenericKeyVerify) {
  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));
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GenericVerify(key_handle.data(), key_handle.size(), clear_buffer_.data(),
                    clear_buffer_.size(), OEMCrypto_HMAC_SHA256,
                    signature.data(), signature.size()));
}

/** Test that the Generic_Verify function fails when not allowed. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyBadVerify) {
  EncryptAndLoadKeys();
  BadVerify(0, OEMCrypto_HMAC_SHA256, SHA256_DIGEST_LENGTH, false);
  BadVerify(1, OEMCrypto_HMAC_SHA256, SHA256_DIGEST_LENGTH, false);
  BadVerify(2, OEMCrypto_HMAC_SHA256, SHA256_DIGEST_LENGTH, false);
  BadVerify(3, OEMCrypto_HMAC_SHA256, SHA256_DIGEST_LENGTH, true);
  BadVerify(3, OEMCrypto_HMAC_SHA256, SHA256_DIGEST_LENGTH - 1, false);
  BadVerify(3, OEMCrypto_HMAC_SHA256, SHA256_DIGEST_LENGTH + 1, false);
  BadVerify(3, OEMCrypto_AES_CBC_128_NO_PADDING, SHA256_DIGEST_LENGTH, false);
}

/** Test Generic_Encrypt with the maximum buffer size. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyEncryptLargeBuffer) {
  ResizeBuffer(GetResourceValue(kMaxGenericBuffer));
  EncryptAndLoadKeys();
  unsigned int key_index = 0;
  vector<uint8_t> expected_encrypted;
  EncryptBuffer(key_index, clear_buffer_, &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));
  vector<uint8_t> encrypted(clear_buffer_.size());
  ASSERT_EQ(OEMCrypto_SUCCESS,
            GenericEncrypt(key_handle.data(), key_handle.size(),
                           clear_buffer_.data(), clear_buffer_.size(), iv_,
                           OEMCrypto_AES_CBC_128_NO_PADDING, encrypted.data()));
  ASSERT_EQ(expected_encrypted, encrypted);
}

/** Test Generic_Decrypt with the maximum buffer size. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyDecryptLargeBuffer) {
  // Some applications are known to pass in a block that is almost 400k.
  ResizeBuffer(GetResourceValue(kMaxGenericBuffer));
  EncryptAndLoadKeys();
  unsigned int key_index = 1;
  vector<uint8_t> encrypted;
  EncryptBuffer(key_index, clear_buffer_, &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));
  vector<uint8_t> resultant(encrypted.size());
  ASSERT_EQ(OEMCrypto_SUCCESS,
            GenericDecrypt(key_handle.data(), key_handle.size(),
                           encrypted.data(), encrypted.size(), iv_,
                           OEMCrypto_AES_CBC_128_NO_PADDING, resultant.data()));
  ASSERT_EQ(clear_buffer_, resultant);
}

/** Test Generic_Sign with the maximum buffer size. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeySignLargeBuffer) {
  ResizeBuffer(GetResourceValue(kMaxGenericBuffer));
  EncryptAndLoadKeys();
  unsigned int key_index = 2;
  vector<uint8_t> expected_signature;
  SignBuffer(key_index, clear_buffer_, &expected_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));
  size_t gen_signature_length = 0;
  ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER,
            GenericSign(key_handle.data(), key_handle.size(),
                        clear_buffer_.data(), clear_buffer_.size(),
                        OEMCrypto_HMAC_SHA256, nullptr, &gen_signature_length));
  ASSERT_EQ(static_cast<size_t>(SHA256_DIGEST_LENGTH), gen_signature_length);
  vector<uint8_t> signature(SHA256_DIGEST_LENGTH);
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GenericSign(key_handle.data(), key_handle.size(), clear_buffer_.data(),
                  clear_buffer_.size(), OEMCrypto_HMAC_SHA256, signature.data(),
                  &gen_signature_length));
  ASSERT_EQ(expected_signature, signature);
}

/** Test Generic_Verify with the maximum buffer size. */
TEST_P(OEMCryptoGenericCryptoTest, GenericKeyVerifyLargeBuffer) {
  ResizeBuffer(GetResourceValue(kMaxGenericBuffer));
  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));
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GenericVerify(key_handle.data(), key_handle.size(), clear_buffer_.data(),
                    clear_buffer_.size(), OEMCrypto_HMAC_SHA256,
                    signature.data(), signature.size()));
}

/** Test Generic_Encrypt when the key duration has expired. */
TEST_P(OEMCryptoGenericCryptoTest, KeyDurationEncrypt) {
  license_messages_.core_response()
      .timer_limits.total_playback_duration_seconds = kDuration;
  license_messages_.CreateResponseWithGenericCryptoKeys();
  EncryptAndLoadKeys();
  vector<uint8_t> expected_encrypted;
  EncryptBuffer(0, clear_buffer_, &expected_encrypted);
  unsigned int key_index = 0;
  vector<uint8_t> encrypted(clear_buffer_.size());

  vector<uint8_t> key_handle;
  // Should be valid key at the start.
  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));
  ASSERT_EQ(OEMCrypto_SUCCESS,
            GenericEncrypt(key_handle.data(), key_handle.size(),
                           clear_buffer_.data(), clear_buffer_.size(), iv_,
                           OEMCrypto_AES_CBC_128_NO_PADDING, encrypted.data()));
  ASSERT_EQ(expected_encrypted, encrypted);

  wvutil::TestSleep::Sleep(kLongSleep + kShortSleep);  // Should be expired key.
  encrypted.assign(clear_buffer_.size(), 0);
  OEMCryptoResult status = OEMCrypto_Generic_Encrypt(
      key_handle.data(), key_handle.size(), clear_buffer_.data(),
      clear_buffer_.size(), iv_, OEMCrypto_AES_CBC_128_NO_PADDING,
      encrypted.data());
  ASSERT_EQ(OEMCrypto_ERROR_KEY_EXPIRED, status);
  ASSERT_NE(encrypted, expected_encrypted);
  ASSERT_NO_FATAL_FAILURE(session_.TestGetKeyHandleExpired(key_index));
}

/** Test Generic_Decrypt when the key duration has expired. */
TEST_P(OEMCryptoGenericCryptoTest, KeyDurationDecrypt) {
  license_messages_.core_response()
      .timer_limits.total_playback_duration_seconds = kDuration;
  license_messages_.CreateResponseWithGenericCryptoKeys();
  EncryptAndLoadKeys();

  // Should be valid key at the start.
  unsigned int key_index = 1;
  vector<uint8_t> encrypted;
  EncryptBuffer(key_index, clear_buffer_, &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));
  vector<uint8_t> resultant(encrypted.size());
  ASSERT_EQ(OEMCrypto_SUCCESS,
            GenericDecrypt(key_handle.data(), key_handle.size(),
                           encrypted.data(), encrypted.size(), iv_,
                           OEMCrypto_AES_CBC_128_NO_PADDING, resultant.data()));
  ASSERT_EQ(clear_buffer_, resultant);

  wvutil::TestSleep::Sleep(kLongSleep + kShortSleep);  // Should be expired key.
  resultant.assign(encrypted.size(), 0);
  OEMCryptoResult status = GenericDecrypt(
      key_handle.data(), key_handle.size(), encrypted.data(), encrypted.size(),
      iv_, OEMCrypto_AES_CBC_128_NO_PADDING, resultant.data());
  ASSERT_EQ(OEMCrypto_ERROR_KEY_EXPIRED, status);
  ASSERT_NE(clear_buffer_, resultant);
  ASSERT_NO_FATAL_FAILURE(session_.TestGetKeyHandleExpired(key_index));
}

/** Test Generic_Sign when the key duration has expired. */
TEST_P(OEMCryptoGenericCryptoTest, KeyDurationSign) {
  license_messages_.core_response()
      .timer_limits.total_playback_duration_seconds = kDuration;
  license_messages_.CreateResponseWithGenericCryptoKeys();
  EncryptAndLoadKeys();

  unsigned int key_index = 2;
  vector<uint8_t> expected_signature;
  vector<uint8_t> signature(SHA256_DIGEST_LENGTH);
  size_t signature_length = signature.size();
  SignBuffer(key_index, clear_buffer_, &expected_signature);

  vector<uint8_t> key_handle;
  // Should be valid key at the start.
  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));
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GenericSign(key_handle.data(), key_handle.size(), clear_buffer_.data(),
                  clear_buffer_.size(), OEMCrypto_HMAC_SHA256, signature.data(),
                  &signature_length));
  ASSERT_EQ(expected_signature, signature);

  wvutil::TestSleep::Sleep(kLongSleep + kShortSleep);  // Should be expired key.
  signature.assign(SHA256_DIGEST_LENGTH, 0);
  OEMCryptoResult status =
      GenericSign(key_handle.data(), key_handle.size(), clear_buffer_.data(),
                  clear_buffer_.size(), OEMCrypto_HMAC_SHA256, signature.data(),
                  &signature_length);
  ASSERT_EQ(OEMCrypto_ERROR_KEY_EXPIRED, status);
  ASSERT_NE(expected_signature, signature);
  ASSERT_NO_FATAL_FAILURE(session_.TestGetKeyHandleExpired(key_index));
}

/** Test Generic_Verify when the key duration has expired. */
TEST_P(OEMCryptoGenericCryptoTest, KeyDurationVerify) {
  license_messages_.core_response()
      .timer_limits.total_playback_duration_seconds = kDuration;
  license_messages_.CreateResponseWithGenericCryptoKeys();
  EncryptAndLoadKeys();

  unsigned int key_index = 3;
  vector<uint8_t> signature;
  SignBuffer(key_index, clear_buffer_, &signature);

  vector<uint8_t> key_handle;
  // Should be valid key at the start.
  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));
  ASSERT_EQ(
      OEMCrypto_SUCCESS,
      GenericVerify(key_handle.data(), key_handle.size(), clear_buffer_.data(),
                    clear_buffer_.size(), OEMCrypto_HMAC_SHA256,
                    signature.data(), signature.size()));

  wvutil::TestSleep::Sleep(kLongSleep + kShortSleep);  // Should be expired key.
  OEMCryptoResult status = OEMCrypto_Generic_Verify(
      key_handle.data(), key_handle.size(), clear_buffer_.data(),
      clear_buffer_.size(), OEMCrypto_HMAC_SHA256, signature.data(),
      signature.size());
  ASSERT_EQ(OEMCrypto_ERROR_KEY_EXPIRED, status);
  ASSERT_NO_FATAL_FAILURE(session_.TestGetKeyHandleExpired(key_index));
}

const unsigned int kLongKeyId = 2;

// Test that short key ids are allowed.
class OEMCryptoGenericCryptoKeyIdLengthTest
    : public OEMCryptoGenericCryptoTest {
 protected:
  void SetUp() override {
    OEMCryptoGenericCryptoTest::SetUp();
    license_messages_.set_num_keys(5);
    license_messages_.set_control(wvoec::kControlAllowDecrypt);
    license_messages_.core_response()
        .timer_limits.total_playback_duration_seconds = kDuration;
    ASSERT_NO_FATAL_FAILURE(license_messages_.CreateDefaultResponse());
    SetUniformKeyIdLength(16);  // Start with all key ids being 16 bytes.
    // But, we are testing that the key ids do not have to have the same
    // length.
    //                                     12 bytes (common key id length).
    license_messages_.SetKeyId(0, "123456789012");
    license_messages_.SetKeyId(1, "12345");  // short key id.
    //                                      16 byte key id. (default)
    license_messages_.SetKeyId(2, "1234567890123456");
    license_messages_.SetKeyId(3, "12345678901234");  // 14 byte. (uncommon)
    license_messages_.SetKeyId(4, "1");               // very short key id.
    ASSERT_EQ(2u, kLongKeyId);
    ASSERT_NO_FATAL_FAILURE(license_messages_.FillCoreResponseSubstrings());
  }

  // Make all four keys have the same length.
  void SetUniformKeyIdLength(size_t key_id_length) {
    for (size_t i = 0; i < license_messages_.num_keys(); i++) {
      string key_id;
      key_id.resize(key_id_length, i + 'a');
      license_messages_.SetKeyId(i, key_id);
    }
    ASSERT_NO_FATAL_FAILURE(license_messages_.FillCoreResponseSubstrings());
  }

  void TestWithKey(unsigned int key_index) {
    ASSERT_LT(key_index, license_messages_.num_keys());
    EncryptAndLoadKeys();
    vector<uint8_t> encrypted;
    // To make sure OEMCrypto is not expecting the key_id to be zero padded,
    // we will create a buffer that is padded with 'Z'. Then, we use fill
    // the buffer with the longer of the three keys. If OEMCrypto is paying
    // attention to the key id length, it should pick out the correct key.
    vector<uint8_t> key_id_buffer(
        session_.license().keys[kLongKeyId].key_id_length + 5,
        'Z');  // Fill a bigger buffer with letter 'Z'.
    memcpy(key_id_buffer.data(), session_.license().keys[kLongKeyId].key_id,
           session_.license().keys[kLongKeyId].key_id_length);
    EncryptBuffer(key_index, clear_buffer_, &encrypted);
    vector<uint8_t> key_handle;
    ASSERT_EQ(
        OEMCrypto_SUCCESS,
        GetKeyHandleIntoVector(session_.session_id(), key_id_buffer.data(),
                               session_.license().keys[key_index].key_id_length,
                               OEMCrypto_CipherMode_CENC, key_handle));
    vector<uint8_t> resultant(encrypted.size());
    ASSERT_EQ(OEMCrypto_SUCCESS,
              OEMCrypto_Generic_Decrypt(key_handle.data(), key_handle.size(),
                                        encrypted.data(), encrypted.size(), iv_,
                                        OEMCrypto_AES_CBC_128_NO_PADDING,
                                        resultant.data()));
    ASSERT_EQ(clear_buffer_, resultant);
  }
};

TEST_P(OEMCryptoGenericCryptoKeyIdLengthTest, MediumKeyId) { TestWithKey(0); }

TEST_P(OEMCryptoGenericCryptoKeyIdLengthTest, ShortKeyId) { TestWithKey(1); }

TEST_P(OEMCryptoGenericCryptoKeyIdLengthTest, LongKeyId) { TestWithKey(2); }

TEST_P(OEMCryptoGenericCryptoKeyIdLengthTest, FourteenByteKeyId) {
  TestWithKey(3);
}

TEST_P(OEMCryptoGenericCryptoKeyIdLengthTest, VeryShortKeyId) {
  TestWithKey(4);
}

TEST_P(OEMCryptoGenericCryptoKeyIdLengthTest, UniformShortKeyId) {
  SetUniformKeyIdLength(5);
  TestWithKey(2);
}

TEST_P(OEMCryptoGenericCryptoKeyIdLengthTest, UniformLongKeyId) {
  SetUniformKeyIdLength(kTestKeyIdMaxLength);
  TestWithKey(2);
}

INSTANTIATE_TEST_SUITE_P(TestAll, OEMCryptoGenericCryptoTest,
                         Range<uint32_t>(kCoreMessagesAPI, kCurrentAPI + 1));

INSTANTIATE_TEST_SUITE_P(TestAll, OEMCryptoGenericCryptoKeyIdLengthTest,
                         Range<uint32_t>(kCoreMessagesAPI, kCurrentAPI + 1));

/// @}
}  // namespace wvoec