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

#include "oemcrypto_session.h"

#include "string.h"

#include "assert_interface.h"
#include "clock_interface.h"
#include "oemcrypto_endianness.h"
#include "logging_interface.h"
#include "oemcrypto_key_table.h"
#include "oemcrypto_nonce_table.h"
#include "oemcrypto_overflow.h"
#include "oemcrypto_session_key_table.h"
#include "oemcrypto_usage_table.h"
#include "root_of_trust_interface.h"

static void clear_nonce_table(OEMCryptoSession* session) {
  for (int i = 0; i < NONCE_TABLE_SIZE; i++) {
    session->nonce_table.age[i] = 0;
    session->nonce_table.nonces[i] = 0;
    session->nonce_table.state[i] = NT_STATE_INVALID;
  }
}

OEMCryptoResult InitializeSession(OEMCryptoSession* session, uint32_t index) {
  ASSERT(session != NULL, "session is NULL");
  session->session_id = index;
  session->state = SESSION_INVALID;
  clear_nonce_table(session);
  session->drm_private_key = NULL;
  session->mac_key_client = NULL;
  session->mac_key_server = NULL;
  session->encryption_key = NULL;
  session->session_key = NULL;
  session->refresh_valid = false;
  session->license_type = OEMCrypto_ContentLicense;
  session->current_content_key_index = CONTENT_KEYS_PER_SESSION;
  for (int i = 0; i < CONTENT_KEYS_PER_SESSION; i++) {
    session->content_keys[i] = NULL;
  }
  session->num_content_keys = 0;
  for (int i = 0; i < ENTITLEMENT_KEYS_PER_SESSION; i++) {
    session->entitlement_keys[i] = NULL;
  }
  session->num_entitlement_keys = 0;
  session->valid_srm_version = false;
  session->timer_start = 0;
  session->compute_hash = false;
  session->current_hash = 0;
  session->given_hash = 0;
  session->current_frame_number = 0;
  session->bad_frame_number = 0;
  session->hash_error = OEMCrypto_SUCCESS;
  session->usage_entry_status = SESSION_HAS_NO_ENTRY;
  session->usage_entry_number = MAX_NUMBER_OF_USAGE_ENTRIES;
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult TerminateSession(OEMCryptoSession* session) {
  ASSERT(session != NULL, "session is NULL");
  OEMCryptoResult result = OEMCrypto_SUCCESS;

  result = FreeKey(&session->drm_private_key);

  OEMCryptoResult free_key_result = FreeKey(&session->mac_key_client);
  if (result == OEMCrypto_SUCCESS) result = free_key_result;

  free_key_result = FreeKey(&session->mac_key_server);
  if (result == OEMCrypto_SUCCESS) result = free_key_result;

  free_key_result = FreeKey(&session->encryption_key);
  if (result == OEMCrypto_SUCCESS) result = free_key_result;

  free_key_result = FreeKey(&session->session_key);
  if (result == OEMCrypto_SUCCESS) result = free_key_result;

  for (uint32_t i = 0; i < session->num_content_keys; i++) {
    free_key_result = FreeKey(&session->content_keys[i]);
    if (result == OEMCrypto_SUCCESS) result = free_key_result;
  }

  for (uint32_t i = 0; i < session->num_entitlement_keys; i++) {
    free_key_result = FreeKey(&session->entitlement_keys[i]);
    if (result == OEMCrypto_SUCCESS) result = free_key_result;
  }

  clear_nonce_table(session);

  if (session->usage_entry_status != SESSION_HAS_NO_ENTRY) {
    ReleaseEntry(session, session->usage_entry_number);
  }

  memset(session, 0, sizeof(OEMCryptoSession));
  return result;
}

OEMCryptoResult CheckStatePreCall(OEMCryptoSession* session,
                                  OEMCryptoSessionAPI api) {
  ASSERT(session != NULL, "session is NULL");
  switch (api) {
    case API_OPENSESSION:
      switch (session->state) {
        case SESSION_INVALID:
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_UNKNOWN_FAILURE;
      }
    case API_CLOSESESSION:
      return OEMCrypto_SUCCESS;
    case API_GETOEMPUBLICCERTIFICATE:
      switch (session->state) {
        case (SESSION_OPENED):
        case (SESSION_LOAD_OEM_RSA_KEY):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_GENERATEDERIVEDKEYS:
      switch (session->state) {
        case (SESSION_OPENED):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_GENERATENONCE:
      switch (session->state) {
        case (SESSION_OPENED):
        case (SESSION_DERIVED_KEYS):
        case (SESSION_DERIVED_KEYS_FROM_SESSION_KEY):
        case (SESSION_KEYS_LOADED):
        case (SESSION_DECRYPT_KEY_SELECTED):
        case (SESSION_LOAD_OEM_RSA_KEY):
        case (SESSION_LOAD_DRM_RSA_KEY):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_GENERATERSASIGNATURE:
      switch (session->state) {
        case (SESSION_LOAD_OEM_RSA_KEY):
        case (SESSION_LOAD_DRM_RSA_KEY):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_REWRAPDEVICERSAKEY:
      switch (session->state) {
        case (SESSION_DERIVED_KEYS):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_REWRAPDEVICERSAKEY30:
      switch (session->state) {
        case (SESSION_LOAD_OEM_RSA_KEY):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_LOADDEVICERSAKEY:
      switch (session->state) {
        case (SESSION_LOAD_OEM_RSA_KEY):
        case (SESSION_OPENED):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_DERIVEKEYSFROMSESSIONKEY:
      switch (session->state) {
        case (SESSION_LOAD_DRM_RSA_KEY):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_LOADKEYS:
      switch (session->state) {
        case (SESSION_DERIVED_KEYS_FROM_SESSION_KEY):
        case (SESSION_KEYS_LOADED):
        case (SESSION_DECRYPT_KEY_SELECTED):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_REFRESHKEYS:
      switch (session->state) {
        case (SESSION_KEYS_LOADED):
        case (SESSION_DECRYPT_KEY_SELECTED):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_GENERATESIGNATURE:
      switch (session->state) {
        case (SESSION_DERIVED_KEYS):
        case (SESSION_DERIVED_KEYS_FROM_SESSION_KEY):
        case (SESSION_KEYS_LOADED):
        case (SESSION_DECRYPT_KEY_SELECTED):
          /* The next two are needed for license release in v15. */
        case (SESSION_OPENED):
        case (SESSION_LOAD_DRM_RSA_KEY):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_SELECTKEY:
    case API_LOADENTITLEDCONTENTKEYS:
    case API_QUERYKEYCONTROL:
      switch (session->state) {
        case (SESSION_KEYS_LOADED):
        case (SESSION_DECRYPT_KEY_SELECTED):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_DECRYPTCENC:
    case API_GENERICENCRYPT:
    case API_GENERICDECRYPT:
    case API_GENERICSIGN:
    case API_GENERICVERIFY:
    case API_GETHASHERRORCODE:
      switch (session->state) {
        case (SESSION_DECRYPT_KEY_SELECTED):
          return OEMCrypto_SUCCESS;
        default:
          return OEMCrypto_ERROR_INVALID_CONTEXT;
      }
    case API_SETDECRYPTHASH:
      switch (session->state) {
        case (SESSION_INVALID):
          return OEMCrypto_ERROR_INVALID_CONTEXT;
        default:
          return OEMCrypto_SUCCESS;
      }
    default:
      return OEMCrypto_ERROR_NOT_IMPLEMENTED;
  }
}

OEMCryptoResult SetStatePostCall(OEMCryptoSession* session,
                                 OEMCryptoSessionAPI api) {
  ASSERT(session != NULL, "session is NULL");
  switch (api) {
    case API_OPENSESSION:
      session->state = SESSION_OPENED;
      break;
    case API_GETOEMPUBLICCERTIFICATE:
      session->state = SESSION_LOAD_OEM_RSA_KEY;
      break;
    case API_GENERATEDERIVEDKEYS:
      session->state = SESSION_DERIVED_KEYS;
      break;
    case API_GENERATENONCE:
    case API_GENERATERSASIGNATURE:
    case API_REWRAPDEVICERSAKEY:
    case API_REWRAPDEVICERSAKEY30:
    case API_GENERATESIGNATURE:
    case API_REFRESHKEYS:
    case API_LOADENTITLEDCONTENTKEYS:
    case API_GENERICENCRYPT:
    case API_GENERICDECRYPT:
    case API_GENERICSIGN:
    case API_GENERICVERIFY:
    case API_SETDECRYPTHASH:
    case API_GETHASHERRORCODE:
    case API_QUERYKEYCONTROL:
      /* State does not change. */
      break;
    case API_LOADDEVICERSAKEY:
      session->state = SESSION_LOAD_DRM_RSA_KEY;
      break;
    case API_DERIVEKEYSFROMSESSIONKEY:
      session->state = SESSION_DERIVED_KEYS_FROM_SESSION_KEY;
      break;
    case API_LOADKEYS:
      session->state = SESSION_KEYS_LOADED;
      break;
    case API_SELECTKEY:
    case API_DECRYPTCENC:
      session->state = SESSION_DECRYPT_KEY_SELECTED;
      break;
    default:
      return OEMCrypto_ERROR_NOT_IMPLEMENTED;
  }
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult StartTimer(OEMCryptoSession* session) {
  ASSERT(session != NULL, "session is NULL");
  OEMCrypto_Clock_Security_Level clock_type;
  return GetSystemTime(&session->timer_start, &clock_type);
}

OEMCryptoResult CurrentTimer(const OEMCryptoSession* session, uint64_t* diff) {
  ASSERT(session != NULL, "session is NULL");
  ASSERT(diff != NULL, "diff is NULL");
  uint64_t current_time;
  OEMCrypto_Clock_Security_Level clock_type;
  OEMCryptoResult result = GetSystemTime(&current_time, &clock_type);
  if (result != OEMCrypto_SUCCESS) return result;
  if (SubOverflowU64(current_time, session->timer_start, diff)) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult LoadDRMRSAKey(OEMCryptoSession* session,
                              const uint8_t* pkcs8_rsa_key,
                              uint32_t rsa_key_length) {
  ASSERT(session != NULL, "session is NULL");
  ASSERT(pkcs8_rsa_key != NULL, "pkcs8_rsa_key is NULL");
  ASSERT(rsa_key_length != 0, "rsa_key_length is 0");
  if (rsa_key_length < 8) return OEMCrypto_ERROR_INVALID_RSA_KEY;
  /* Determine the padding scheme allowed. */
  uint32_t allowed_schemes;
  if ((memcmp(pkcs8_rsa_key, "SIGN", 4) == 0)) {
    uint32_t schemes_n;
    memcpy((uint8_t*)&schemes_n, pkcs8_rsa_key + 4, sizeof(uint32_t));
    allowed_schemes = HostToNetworkU32(schemes_n);
    pkcs8_rsa_key += 8;
    rsa_key_length -= 8;
  } else {
    allowed_schemes = kSign_RSASSA_PSS;
  }

  if ((GetRSAPaddingSchemes() & allowed_schemes) != allowed_schemes) {
    return OEMCrypto_ERROR_NOT_IMPLEMENTED;
  }

  /* Verify that the key is valid. */
  OEMCryptoResult result =
      CreateKey(&session->drm_private_key, pkcs8_rsa_key, rsa_key_length,
                DRM_RSA_PRIVATE_KEY, UNKNOWN_KEY_OPERATION, UNKNOWN_KEY_SIZE);
  if (result != OEMCrypto_SUCCESS) return result;
  session->drm_private_key->allowed_schemes = allowed_schemes;
  return OEMCrypto_SUCCESS;
}

typedef enum DeriveKeyType {
  GENERIC_DERIVE_KEY_TYPE = 0x24514123,
  KEYBOX_DERIVE_KEY_TYPE = 0x0be5a960,
  DEVICE_KEY_DERIVE_KEY_TYPE = 0x4933dcdd,
} DeriveKeyType;

static OEMCryptoResult DeriveKey(const CryptoKey* master_key, uint8_t counter,
                                 const uint8_t* context,
                                 uint32_t context_length, uint8_t* out,
                                 DeriveKeyType derive_key_type) {
  switch (derive_key_type) {
    case GENERIC_DERIVE_KEY_TYPE:
      return DeriveKeyFromKeyHandle(master_key->key_handle, counter, context,
                                    context_length, out);
    case DEVICE_KEY_DERIVE_KEY_TYPE:
      return DeriveKeyFromDeviceKey(counter, context, context_length, out);
    case KEYBOX_DERIVE_KEY_TYPE:
      return DeriveKeyFromKeybox(counter, context, context_length, out);
    default:
      /* Unimplemented or invalid. */
      return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
}

static OEMCryptoResult DeriveMacAndEncryptionKeys(
    OEMCryptoSession* session, const CryptoKey* master_key,
    const uint8_t* mac_key_context, uint32_t mac_key_context_length,
    const uint8_t* enc_key_context, uint32_t enc_key_context_length,
    DeriveKeyType derive_key_type) {
  ASSERT(
      session != NULL && mac_key_context != NULL &&
          mac_key_context_length != 0 && enc_key_context != NULL &&
          enc_key_context_length != 0 &&
          !(master_key == NULL && derive_key_type == GENERIC_DERIVE_KEY_TYPE),
      "Parameters are NULL or 0");

  /* Generate derived keys for mac keys. */
  uint8_t mac_key_server[MAC_KEY_SIZE];
  uint8_t mac_key_client[MAC_KEY_SIZE];

  OEMCryptoResult result =
      DeriveKey(master_key, 1, mac_key_context, mac_key_context_length,
                mac_key_server, derive_key_type);
  if (result != OEMCrypto_SUCCESS) return result;
  result = DeriveKey(master_key, 2, mac_key_context, mac_key_context_length,
                     mac_key_server + KEY_SIZE_128, derive_key_type);
  if (result != OEMCrypto_SUCCESS) return result;

  result = DeriveKey(master_key, 3, mac_key_context, mac_key_context_length,
                     mac_key_client, derive_key_type);
  if (result != OEMCrypto_SUCCESS) return result;
  result = DeriveKey(master_key, 4, mac_key_context, mac_key_context_length,
                     mac_key_client + KEY_SIZE_128, derive_key_type);
  if (result != OEMCrypto_SUCCESS) return result;

  /* Generate derived key for encryption key. */
  uint8_t encryption_key[KEY_SIZE_128];
  result = DeriveKey(master_key, 1, enc_key_context, enc_key_context_length,
                     encryption_key, derive_key_type);
  if (result != OEMCrypto_SUCCESS) return result;

  result = CreateKey(&session->mac_key_server, mac_key_server, MAC_KEY_SIZE,
                     MAC_KEY_SERVER, MAC_KEY_SERVER_VERIFY, MAC_KEY_SIZE);
  if (result != OEMCrypto_SUCCESS) return result;
  result = CreateKey(&session->mac_key_client, mac_key_client, MAC_KEY_SIZE,
                     MAC_KEY_CLIENT, MAC_KEY_CLIENT_SIGN, MAC_KEY_SIZE);
  if (result != OEMCrypto_SUCCESS) {
    FreeKey(&session->mac_key_server);
    return result;
  }
  result = CreateKey(&session->encryption_key, encryption_key, KEY_SIZE_128,
                     ENCRYPTION_KEY, ENCRYPTION_KEY_ENCRYPT, KEY_SIZE_128);
  if (result != OEMCrypto_SUCCESS) {
    FreeKey(&session->mac_key_server);
    FreeKey(&session->mac_key_client);
  }
  return result;
}

OEMCryptoResult DeriveMacAndEncryptionKeysFromCryptoKey(
    OEMCryptoSession* session, const CryptoKey* master_key,
    const uint8_t* mac_key_context, uint32_t mac_key_context_length,
    const uint8_t* enc_key_context, uint32_t enc_key_context_length) {
  return DeriveMacAndEncryptionKeys(
      session, master_key, mac_key_context, mac_key_context_length,
      enc_key_context, enc_key_context_length, GENERIC_DERIVE_KEY_TYPE);
}

OEMCryptoResult DeriveMacAndEncryptionKeysFromDeviceKey(
    OEMCryptoSession* session, const uint8_t* mac_key_context,
    uint32_t mac_key_context_length, const uint8_t* enc_key_context,
    uint32_t enc_key_context_length) {
  return DeriveMacAndEncryptionKeys(
      session, NULL, mac_key_context, mac_key_context_length, enc_key_context,
      enc_key_context_length, DEVICE_KEY_DERIVE_KEY_TYPE);
}

OEMCryptoResult DeriveMacAndEncryptionKeysFromKeybox(
    OEMCryptoSession* session, const uint8_t* mac_key_context,
    uint32_t mac_key_context_length, const uint8_t* enc_key_context,
    uint32_t enc_key_context_length) {
  if (GetProvisioningMethod() != OEMCrypto_Keybox) {
    return OEMCrypto_ERROR_INVALID_CONTEXT;
  }
  OEMCryptoResult result = ValidateKeybox();
  if (result != OEMCrypto_SUCCESS) return result;
  return DeriveMacAndEncryptionKeys(
      session, NULL, mac_key_context, mac_key_context_length, enc_key_context,
      enc_key_context_length, KEYBOX_DERIVE_KEY_TYPE);
}

OEMCryptoResult VerifySignatureWithMacKeyServer(OEMCryptoSession* session,
                                                const uint8_t* message,
                                                uint32_t message_length,
                                                const uint8_t* signature) {
  ASSERT(session != NULL && message != NULL && message_length != 0 &&
             signature != NULL,
         "Parameters are NULL or 0");
  if (!CheckKey(session->mac_key_server, MAC_KEY_SERVER,
                MAC_KEY_SERVER_VERIFY)) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  uint8_t computed_signature[SHA256_DIGEST_LENGTH];
  OEMCryptoResult result =
      HMAC_SHA256(session->mac_key_server->key_handle, message, message_length,
                  computed_signature);
  if (result != OEMCrypto_SUCCESS) return result;
  if (memcmp(signature, computed_signature, SHA256_DIGEST_LENGTH) != 0) {
    return OEMCrypto_ERROR_SIGNATURE_FAILURE;
  }
  return OEMCrypto_SUCCESS;
}

static OEMCryptoResult CheckStatusOnline(OEMCryptoSession* session,
                                         uint32_t nonce, uint32_t control) {
  if (!session) return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  if (!(control & CONTROL_NONCE_ENABLED)) {
    /* TODO(b/154764983): fix this. */
    /* Server error. Continue, and assume nonce required. */
    LOGE("Server error: nonce not enabled.");
  }
  if (!CheckNonce(&session->nonce_table, nonce)) {
    return OEMCrypto_ERROR_INVALID_NONCE;
  }
  switch (session->usage_entry_status) {
    default: /* Invalid status. */
    case SESSION_HAS_NO_ENTRY:
      return OEMCrypto_ERROR_INVALID_CONTEXT;
    case SESSION_HAS_LOADED_ENTRY:
      /* Cannot reload usage entry for online license. */
      return OEMCrypto_ERROR_INVALID_CONTEXT;
    case SESSION_HAS_NEW_ENTRY:
      return OEMCrypto_SUCCESS;
  }
}

static OEMCryptoResult CheckStatusOffline(OEMCryptoSession* session,
                                          uint32_t nonce, uint32_t control) {
  if (!session) return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  if (!(control & CONTROL_NONCE_ENABLED)) {
    /* TODO(b/154764983): fix this. */
    /* Server error. Continue, and assume nonce required. */
    LOGE("Server error: nonce not enabled.");
  }
  switch (session->usage_entry_status) {
    default: /* Invalid status. */
    case SESSION_HAS_NO_ENTRY:
      return OEMCrypto_ERROR_INVALID_CONTEXT;
    case SESSION_HAS_LOADED_ENTRY:
      /* We do not check the nonce. Instead, LoadKeys will verify the pst and
       * mac keys match. */
      return OEMCrypto_SUCCESS;
    case SESSION_HAS_NEW_ENTRY:
      if (!CheckNonce(&session->nonce_table, nonce)) {
        return OEMCrypto_ERROR_INVALID_NONCE;
      }
      return OEMCrypto_SUCCESS;
  }
}

static OEMCryptoResult check_nonce_or_entry(OEMCryptoSession* session,
                                            KeyControlBlock key_control_block) {
  /* Note: we only check the nonce if the bit is enabled. */
  uint8_t replay_bit =
      (uint8_t)(key_control_block.control_bits & CONTROL_REPLAY_MASK);
  switch (replay_bits) {
    case CONTROL_NONCE_REQUIRED:
      /* Online license. Nonce always required. */
      return CheckStatusOnline(session, key_control_block.nonce,
                               key_control_block.control_bits);
      break;
    case CONTROL_NONCE_OR_ENTRY:
      /* Offline license. Nonce required on first load. */
      return CheckStatusOffline(session, key_control_block.nonce,
                                key_control_block.control_bits);
      break;
    default:
      if ((key_control_block.control_bits & CONTROL_NONCE_ENABLED) &&
          !CheckNonce(&session->nonce_table, key_control_block.nonce)) {
        return OEMCrypto_ERROR_INVALID_NONCE;
      }
      return OEMCrypto_SUCCESS;
  }
}

OEMCryptoResult InstallKey(OEMCryptoSession* session, const uint8_t* key_id,
                           uint32_t key_id_length, const uint8_t* key_data,
                           uint32_t key_data_length, const uint8_t* key_data_iv,
                           const uint8_t* key_control,
                           const uint8_t* key_control_iv) {
  ASSERT(session != NULL && key_id != NULL && key_id_length != 0 &&
             key_data != NULL && key_data_length != 0 && key_data_iv != NULL &&
             key_control != NULL && key_control_iv != NULL,
         "Parameters are NULL or 0");
  uint8_t raw_key[KEY_SIZE_256];
  uint8_t raw_key_control[KEY_CONTROL_SIZE];

  CryptoKey** current_key_ptr;
  uint32_t current_key_index = 0;
  CryptoKeyType key_type;
  CryptoKeyOperation key_operation;
  if (session->license_type == OEMCrypto_ContentLicense) {
    if (key_data_length != KEY_SIZE_128 && key_data_length != KEY_SIZE_256) {
      /* Generic crypto allows both key sizes. */
      return OEMCrypto_ERROR_INVALID_CONTEXT;
    }
    current_key_index = session->num_content_keys;
    if (current_key_index == CONTENT_KEYS_PER_SESSION) {
      return OEMCrypto_ERROR_INSUFFICIENT_RESOURCES;
    }
    current_key_ptr = &session->content_keys[current_key_index];
    key_type = CONTENT_KEY;
    key_operation = CONTENT_KEY_DECRYPT;
  } else {
    if (key_data_length != KEY_SIZE_256) return OEMCrypto_ERROR_INVALID_CONTEXT;
    current_key_index = session->num_entitlement_keys;
    if (current_key_index == ENTITLEMENT_KEYS_PER_SESSION) {
      return OEMCrypto_ERROR_INSUFFICIENT_RESOURCES;
    }
    current_key_ptr = &session->entitlement_keys[current_key_index];
    key_type = ENTITLEMENT_KEY;
    key_operation = ENTITLEMENT_KEY_DECRYPT;
  }

  if (session->encryption_key != NULL) {
    session->encryption_key->key_operation = ENCRYPTION_KEY_DECRYPT;
  }
  if (!CheckKey(session->encryption_key, ENCRYPTION_KEY,
                ENCRYPTION_KEY_DECRYPT)) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  OEMCryptoResult result = AESCBCDecrypt(
      session->encryption_key->key_handle, session->encryption_key->key_size,
      key_data, key_data_length, key_data_iv, raw_key);
  if (result != OEMCrypto_SUCCESS) return result;

  result = CreateKey(current_key_ptr, raw_key, key_data_length, key_type,
                     key_operation, key_data_length);
  if (result != OEMCrypto_SUCCESS) return result;

  CryptoKey* current_key = *current_key_ptr;
  memcpy(current_key->key_id, key_id, key_id_length);
  current_key->key_id_size = key_id_length;
  current_key->session_key_index = current_key_index;

  if (!CheckKey(current_key, key_type, key_operation)) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  /* We use the first 16 bytes regardless of the license type to decrypt the key
     control. */
  result = AESCBCDecrypt(current_key->key_handle, KEY_SIZE_128, key_control,
                         KEY_CONTROL_SIZE, key_control_iv, raw_key_control);
  if (result != OEMCrypto_SUCCESS) goto cleanup;

  current_key->key_control_block = ParseKeyControlBlock(raw_key_control);
  if (!(current_key->key_control_block.valid)) {
    result = OEMCrypto_ERROR_INVALID_CONTEXT;
    goto cleanup;
  }
  if ((current_key->key_control_block.control_bits &
       CONTROL_REQUIRE_ANTI_ROLLBACK_HARDWARE) &&
      !IsAntiRollbackHWPresent()) {
    result = OEMCrypto_ERROR_UNKNOWN_FAILURE;
    goto cleanup;
  }
  uint8_t minimum_patch_level = (current_key->key_control_block.control_bits &
                                 CONTROL_SECURITY_PATCH_LEVEL_MASK) >>
                                CONTROL_SECURITY_PATCH_LEVEL_SHIFT;
  if (minimum_patch_level > OEMCrypto_Security_Patch_Level()) {
    result = OEMCrypto_ERROR_UNKNOWN_FAILURE;
    goto cleanup;
  }
  result = check_nonce_or_entry(session, current_key->key_control_block);
  if (result != OEMCrypto_SUCCESS) goto cleanup;
  if (current_key->key_control_block.control_bits &
      CONTROL_SRM_VERSION_REQUIRED) {
    if (!session->valid_srm_version) {
      /* Require local display only. */
      current_key->key_control_block.control_bits |= CONTROL_HDCP_VERSION_MASK;
    }
  }
  if (session->license_type == OEMCrypto_ContentLicense) {
    current_key->is_entitled_content_key = false;
  } else {
    /* Entitlement keys will have entitled content keys loaded in later. */
    current_key->has_entitled_content_key = false;
  }

cleanup:
  if (result != OEMCrypto_SUCCESS) FreeKey(current_key_ptr);
  return result;
}

OEMCryptoResult UpdateMacKeys(OEMCryptoSession* session,
                              const uint8_t* enc_mac_keys,
                              const uint8_t* mac_keys_iv) {
  ASSERT(session != NULL && enc_mac_keys != NULL && mac_keys_iv != NULL,
         "Parameters are NULL");
  uint8_t raw_mac_keys[2 * MAC_KEY_SIZE];
  if (!CheckKey(session->encryption_key, ENCRYPTION_KEY,
                ENCRYPTION_KEY_DECRYPT)) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  OEMCryptoResult result = AESCBCDecrypt(
      session->encryption_key->key_handle, session->encryption_key->key_size,
      enc_mac_keys, 2 * MAC_KEY_SIZE, mac_keys_iv, raw_mac_keys);
  if (result != OEMCrypto_SUCCESS) return result;
  OEMCryptoResult free_mac_key_server_result =
      FreeKey(&session->mac_key_server);
  OEMCryptoResult free_mac_key_client_result =
      FreeKey(&session->mac_key_client);
  if (free_mac_key_server_result != OEMCrypto_SUCCESS) {
    return free_mac_key_server_result;
  }
  if (free_mac_key_client_result != OEMCrypto_SUCCESS) {
    return free_mac_key_client_result;
  }
  result = CreateKey(&session->mac_key_server, raw_mac_keys, MAC_KEY_SIZE,
                     MAC_KEY_SERVER, MAC_KEY_SERVER_VERIFY, MAC_KEY_SIZE);
  if (result != OEMCrypto_SUCCESS) return result;
  result = CreateKey(&session->mac_key_client, raw_mac_keys + MAC_KEY_SIZE,
                     MAC_KEY_SIZE, MAC_KEY_CLIENT, MAC_KEY_CLIENT_SIGN,
                     MAC_KEY_SIZE);
  if (result != OEMCrypto_SUCCESS) FreeKey(&session->mac_key_server);
  return result;
}

OEMCryptoResult RefreshKey(OEMCryptoSession* session, const uint8_t* key_id,
                           uint32_t key_id_length, const uint8_t* key_control,
                           const uint8_t* key_control_iv) {
  ASSERT(session != NULL && key_control != NULL,
         "session or key_control is NULL");

  KeyControlBlock key_control_block;

  if (key_id == NULL) {
    key_control_block = ParseKeyControlBlock(key_control);
    if (!key_control_block.valid) {
      return OEMCrypto_ERROR_INVALID_CONTEXT;
    }
    if ((key_control_block.control_bits & CONTROL_NONCE_ENABLED) &&
        !CheckNonce(&session->nonce_table, key_control_block.nonce)) {
      return OEMCrypto_ERROR_INVALID_NONCE;
    }
    UpdateDurationForAllKeys(session, key_control_block);
    return OEMCrypto_SUCCESS;
  }

  bool is_content_key = session->license_type == OEMCrypto_ContentLicense;
  CryptoKey* key =
      FindKeyFromTable(session, is_content_key, key_id, key_id_length);

  if (key == NULL) return OEMCrypto_ERROR_NO_CONTENT_KEY;

  uint8_t raw_key_control[KEY_CONTROL_SIZE];
  if (key_control_iv != NULL) {
    CryptoKeyType key_type = UNKNOWN_KEY_TYPE;
    CryptoKeyOperation key_operation = UNKNOWN_KEY_OPERATION;
    if (is_content_key) {
      key_type = CONTENT_KEY;
      key_operation = CONTENT_KEY_DECRYPT;
    } else {
      key_type = ENTITLEMENT_KEY;
      key_operation = ENTITLEMENT_KEY_DECRYPT;
    }
    if (!CheckKey(key, key_type, key_operation)) {
      return OEMCrypto_ERROR_UNKNOWN_FAILURE;
    }
    /* Decrypt using only the first 128 bits of the key. */
    OEMCryptoResult result =
        AESCBCDecrypt(key->key_handle, KEY_SIZE_128, key_control,
                      KEY_CONTROL_SIZE, key_control_iv, raw_key_control);
    if (result != OEMCrypto_SUCCESS) return result;
  } else {
    memcpy(raw_key_control, key_control, KEY_CONTROL_SIZE);
  }

  key_control_block = ParseKeyControlBlock(raw_key_control);

  if (!key_control_block.valid) return OEMCrypto_ERROR_INVALID_CONTEXT;
  if ((key_control_block.control_bits & CONTROL_NONCE_ENABLED) &&
      !CheckNonce(&session->nonce_table, key_control_block.nonce)) {
    return OEMCrypto_ERROR_INVALID_NONCE;
  }

  key->key_control_block.duration = key_control_block.duration;
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult CheckCurrentContentKeyUsage(const OEMCryptoSession* session,
                                            uint32_t use_type,
                                            OEMCryptoBufferType buffer_type) {
  ASSERT(session != NULL && (buffer_type == OEMCrypto_BufferType_Clear ||
                             buffer_type == OEMCrypto_BufferType_Secure),
         "session is NULL or invalid buffer_type");
  if (session->current_content_key_index >= session->num_content_keys) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  CryptoKey* current_content_key =
      session->content_keys[session->current_content_key_index];
  if (!CheckKey(current_content_key, CONTENT_KEY, CONTENT_KEY_DECRYPT)) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  KeyControlBlock control = current_content_key->key_control_block;
  if (current_content_key->is_entitled_content_key) {
    if (current_content_key->entitlement_key_index >=
        session->num_entitlement_keys) {
      return OEMCrypto_ERROR_UNKNOWN_FAILURE;
    }
    CryptoKey* entitlement_key =
        session->entitlement_keys[current_content_key->entitlement_key_index];
    if (!CheckKey(entitlement_key, ENTITLEMENT_KEY, ENTITLEMENT_KEY_DECRYPT)) {
      return OEMCrypto_ERROR_UNKNOWN_FAILURE;
    }
    control = entitlement_key->key_control_block;
  }
  if (use_type && (!(control.control_bits & use_type))) {
    /* Could not use this key for the given use type. */
    return OEMCrypto_ERROR_INVALID_CONTEXT;
  }
  OEMCryptoResult result;
  if (control.control_bits & CONTROL_DATA_PATH_SECURE) {
    if (!IsClosedPlatform() && buffer_type == OEMCrypto_BufferType_Clear) {
      return OEMCrypto_ERROR_DECRYPT_FAILED;
    }
  }
  switch (session->usage_entry_status) {
    case SESSION_HAS_NO_ENTRY:
      /* No entry, don't do anything. */
      break;
    case SESSION_HAS_DEACTIVATED_ENTRY:
      return OEMCrypto_ERROR_KEY_EXPIRED;
    case SESSION_HAS_NEW_ENTRY:
    case SESSION_HAS_LOADED_ENTRY:
      /* TODO: this is not very efficient to call on every decrypt. This will be
       * more optimized when we mix with the ODK library. */
      result = UpdateLastPlaybackTime(session);
      if (result != OEMCrypto_SUCCESS) return result;
      break;
    default:
      return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  if (control.duration > 0) {
    uint64_t time_elapsed = 0;
    result = CurrentTimer(session, &time_elapsed);
    if (result != OEMCrypto_SUCCESS) return result;
    if (control.duration < time_elapsed) {
      return OEMCrypto_ERROR_KEY_EXPIRED;
    }
  }
  OEMCrypto_HDCP_Capability capability = CurrentHDCPCapability();
  if (capability != HDCP_NO_DIGITAL_OUTPUT &&
      (control.control_bits & CONTROL_HDCP_REQUIRED)) {
    /* Check to see if HDCP requirements are satisfied. */
    uint8_t required_hdcp =
        (control.control_bits & CONTROL_HDCP_VERSION_MASK) >>
        CONTROL_HDCP_VERSION_SHIFT;
    if (required_hdcp > capability || capability == HDCP_NONE) {
      return OEMCrypto_ERROR_INSUFFICIENT_HDCP;
    }
  }
  /* We can't control analog output if the output is clear. Similarly, if
     analog is not disabled, we should fail. */
  if (control.control_bits & CONTROL_DISABLE_ANALOG_OUTPUT) {
    if (buffer_type == OEMCrypto_BufferType_Clear ||
        (IsAnalogDisplayActive() && !DisableAnalogDisplay())) {
      return OEMCrypto_ERROR_ANALOG_OUTPUT;
    }
  }
  /* If CGMS-A is active, we must turn off analog output. */
  if (control.control_bits & CONTROL_CGMS_MASK) {
    if (buffer_type == OEMCrypto_BufferType_Clear) {
      /* Similar to the above, we can't control CGMS if output is clear. */
      return OEMCrypto_ERROR_ANALOG_OUTPUT;
    }
    if (!IsCGMS_AActive() && IsAnalogDisplayActive() &&
        !DisableAnalogDisplay()) {
      /* CGMS must be active if analog output is active. */
      return OEMCrypto_ERROR_ANALOG_OUTPUT;
    }
  }
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult DecryptCBC(const OEMCryptoSession* session,
                           const uint8_t* initial_iv,
                           const OEMCrypto_CENCEncryptPatternDesc* pattern,
                           const uint8_t* cipher_data,
                           uint32_t cipher_data_length, uint8_t* clear_data) {
  ASSERT(session != NULL && initial_iv != NULL && pattern != NULL &&
             cipher_data != NULL && cipher_data_length != 0 &&
             clear_data != NULL,
         "Parameters are NULL or 0");
  if (session->current_content_key_index >= session->num_content_keys) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  CryptoKey* current_content_key =
      session->content_keys[session->current_content_key_index];
  if (!CheckKey(current_content_key, CONTENT_KEY, CONTENT_KEY_DECRYPT)) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  uint8_t iv[AES_BLOCK_SIZE];
  uint8_t next_iv[AES_BLOCK_SIZE];
  memcpy(iv, &initial_iv[0], AES_BLOCK_SIZE);

  uint32_t l = 0;
  uint32_t pattern_offset = pattern->offset;
  while (l < cipher_data_length) {
    uint32_t size = cipher_data_length - l;
    if (size > AES_BLOCK_SIZE) size = AES_BLOCK_SIZE;
    uint32_t pattern_length = pattern->encrypt + pattern->skip;
    bool skip_block =
        (pattern_offset >= pattern->encrypt) && (pattern_length > 0);
    if (pattern_length > 0) {
      pattern_offset = (pattern_offset + 1) % pattern_length;
    }
    if (skip_block || (size < AES_BLOCK_SIZE)) {
      memmove(&clear_data[l], &cipher_data[l], size);
    } else {
      uint8_t aes_output[AES_BLOCK_SIZE];
      /* Save the iv for the next block, in case cipher_data is in the same
         buffer as clear_data. */
      memcpy(next_iv, &cipher_data[l], AES_BLOCK_SIZE);
      OEMCryptoResult result = AESDecrypt(current_content_key->key_handle,
                                          &cipher_data[l], aes_output);
      if (result != OEMCrypto_SUCCESS) return result;
      for (uint32_t n = 0; n < AES_BLOCK_SIZE; n++) {
        clear_data[l + n] = aes_output[n] ^ iv[n];
      }
      memcpy(iv, next_iv, AES_BLOCK_SIZE);
    }
    l += size;
  }
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult DecryptCTR(const OEMCryptoSession* session,
                           const uint8_t* initial_iv, uint32_t block_offset,
                           const OEMCrypto_CENCEncryptPatternDesc* pattern,
                           const uint8_t* cipher_data,
                           uint32_t cipher_data_length, uint8_t* clear_data) {
  ASSERT(session != NULL && initial_iv != NULL && pattern != NULL &&
             cipher_data != NULL && cipher_data_length != 0 &&
             clear_data != NULL,
         "Parameters are NULL or 0");
  if (session->current_content_key_index >= session->num_content_keys) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  CryptoKey* current_content_key =
      session->content_keys[session->current_content_key_index];
  if (!CheckKey(current_content_key, CONTENT_KEY, CONTENT_KEY_DECRYPT)) {
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
  }
  uint8_t iv[AES_BLOCK_SIZE];
  memcpy(iv, &initial_iv[0], AES_BLOCK_SIZE);

  uint32_t l = 0;
  uint32_t pattern_offset = pattern->offset;
  while (l < cipher_data_length) {
    uint32_t size = cipher_data_length - l;
    if (size > AES_BLOCK_SIZE - block_offset)
      size = AES_BLOCK_SIZE - block_offset;
    uint32_t pattern_length = pattern->encrypt + pattern->skip;
    bool skip_block =
        (pattern_offset >= pattern->encrypt) && (pattern_length > 0);
    if (pattern_length > 0) {
      pattern_offset = (pattern_offset + 1) % pattern_length;
    }
    if (skip_block) {
      memmove(&clear_data[l], &cipher_data[l], size);
    } else {
      uint8_t aes_output[AES_BLOCK_SIZE];
      AESEncrypt(current_content_key->key_handle, iv, aes_output);
      for (uint32_t n = 0; n < size; n++) {
        clear_data[l + n] = aes_output[n + block_offset] ^ cipher_data[l + n];
      }
      /* Increment the lower 8 bytes of the iv only. */
      for (int n = 15; n > 7; n--) {
        if (++iv[n] != 0) break;
      };
    }
    l += size;
    block_offset = 0;
  }
  return OEMCrypto_SUCCESS;
}