/* 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_key_table.h"

#include "stdint.h"
#include "string.h"

#include "assert_interface.h"
#include "logging_interface.h"

static KeyTable key_table;
static bool key_table_initialized = false;

OEMCryptoResult InitializeKeyTable(void) {
  ASSERT(MAX_NUMBER_OF_KEYS > 0, "MAX_NUMBER_OF_KEYS must be > 0");
  ASSERT(MAX_NUMBER_OF_KEYS < UINT32_MAX - 1,
         "MAX_NUMBER_OF_KEYS is too large");
  if (key_table_initialized) {
    return OEMCrypto_ERROR_INIT_FAILED;
  }
  key_table.size = MAX_NUMBER_OF_KEYS;
  key_table.first_free_key = 0;
  for (uint32_t i = 0; i < key_table.size; i++) {
    key_table.next_free_key[i] = i + 1;
    key_table.is_free[i] = true;
    memset(&key_table.keys[i], 0, sizeof(CryptoKey));
  }
  key_table_initialized = true;
  return OEMCrypto_SUCCESS;
}

uint32_t MaxNumberOfKeys(void) { return MAX_NUMBER_OF_KEYS; }

OEMCryptoResult NumberOfUsedKeys(uint32_t* num_used_keys) {
  ASSERT(num_used_keys != NULL, "num_used_keys is NULL");
  if (!key_table_initialized) {
    return OEMCrypto_ERROR_SYSTEM_INVALIDATED;
  }
  for (uint32_t i = 0; i < key_table.size; i++) {
    if (!key_table.is_free[i]) {
      (*num_used_keys)++;
    }
  }
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult GrabKey(uint32_t* index) {
  ASSERT(index != NULL, "index is NULL");
  if (!key_table_initialized) {
    return OEMCrypto_ERROR_SYSTEM_INVALIDATED;
  }
  if (key_table.first_free_key == MAX_NUMBER_OF_KEYS) {
    return OEMCrypto_ERROR_TOO_MANY_KEYS;
  }
  *index = key_table.first_free_key;
  key_table.first_free_key = key_table.next_free_key[*index];
  key_table.is_free[*index] = false;
  key_table.keys[*index].key_table_index = *index;
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult GetKey(uint32_t index, CryptoKey** key) {
  ASSERT(key != NULL, "key is NULL");
  if (!key_table_initialized) {
    return OEMCrypto_ERROR_SYSTEM_INVALIDATED;
  }
  if (index >= key_table.size || key_table.is_free[index]) {
    return OEMCrypto_ERROR_INVALID_CONTEXT;
  }
  *key = &key_table.keys[index];
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult CreateKey(CryptoKey** key, const uint8_t* serialized_bytes,
                          uint32_t serialized_bytes_length,
                          CryptoKeyType key_type,
                          CryptoKeyOperation key_operation,
                          CryptoKeySize key_size) {
  ASSERT(key != NULL, "key is NULL");
  OEMCryptoResult result;
  if (*key != NULL) {
    result = FreeKey(key);
    if (result != OEMCrypto_SUCCESS) return result;
  }
  uint32_t key_table_index = 0;
  result = GrabKey(&key_table_index);
  if (result != OEMCrypto_SUCCESS) return result;
  result = GetKey(key_table_index, key);
  if (result != OEMCrypto_SUCCESS) {
    FreeKey(key);
    return result;
  }
  result = InitializeCryptoKey(*key, serialized_bytes, serialized_bytes_length,
                               key_type, key_operation, key_size);
  if (result != OEMCrypto_SUCCESS) {
    FreeKey(key);
    return result;
  }
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult FreeKey(CryptoKey** key) {
  ASSERT(key != NULL, "key is NULL");
  if (*key == NULL) return OEMCrypto_SUCCESS;
  uint32_t index = (*key)->key_table_index;
  if (!key_table_initialized) {
    return OEMCrypto_ERROR_SYSTEM_INVALIDATED;
  }
  if (index >= key_table.size || key_table.is_free[index]) {
    return OEMCrypto_ERROR_INVALID_CONTEXT;
  }
  OEMCryptoResult result = FreeCryptoKey(&key_table.keys[index]);
  if (result != OEMCrypto_SUCCESS) {
    return result;
  }
  key_table.next_free_key[index] = key_table.first_free_key;
  key_table.is_free[index] = true;
  key_table.first_free_key = index;
  *key = NULL;
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult TerminateKeyTable(void) {
  if (!key_table_initialized) {
    return OEMCrypto_ERROR_TERMINATE_FAILED;
  }
  OEMCryptoResult result = OEMCrypto_SUCCESS;
  for (uint32_t i = 0; i < key_table.size; i++) {
    if (!key_table.is_free[i]) {
      result = OEMCrypto_ERROR_TERMINATE_FAILED;
      /* Attempt to free the key. */
      CryptoKey* key = &key_table.keys[i];
      OEMCryptoResult free_result = FreeKey(&key);
      if (free_result != OEMCrypto_SUCCESS) {
        LOGE("Could not free key at index %d with error: %d", i, free_result);
      }
    }
  }
  key_table_initialized = false;
  return result;
}