/* 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. */

#ifndef OEMCRYPTO_TA_OEMCRYPTO_KEY_H_
#define OEMCRYPTO_TA_OEMCRYPTO_KEY_H_

#include "crypto_interface.h"
#include "oemcrypto_key_types.h"

typedef struct KeyControlBlock {
  bool valid;
  char verification[4];
  uint32_t duration;
  uint32_t nonce;
  uint32_t control_bits;
} KeyControlBlock;

typedef struct CryptoKey {
  /* Index into the global key table. */
  uint32_t key_table_index;
  uint64_t cookie;
  CryptoKeyType key_type;
  CryptoKeyOperation key_operation;
  /* key_handle is owned by the TEE. */
  TEE_Key_Handle key_handle;
  CryptoKeySize key_size;
  uint32_t allowed_schemes; /* For DRM RSA keys only */
  /* For entitlement or content keys only. */
  uint8_t key_id[KEY_ID_MAX_SIZE];
  uint8_t key_id_size;
  KeyControlBlock key_control_block;
  /* Index into either the content or entitlement key table in the session. */
  uint32_t session_key_index;
  /* For entitlement keys only. */
  uint32_t entitled_content_key_index; /* Index of entitled content key that
                                          this entitles. */
  bool has_entitled_content_key;
  /* For content keys only. */
  uint32_t entitlement_key_index; /* Index of entitlement key that entitles this
                                     key. */
  bool is_entitled_content_key;
  OEMCryptoCipherMode cipher_mode;
} CryptoKey;

/* Initializes the data fields in the |key| and sets |key|'s key_handle to
   the key handle formed from |serialized_bytes| and |size|.
   Returns the result of CreateKeyHandle if it fails and OEMCrypto_SUCCESS
   otherwise.
   |serialized_bytes_length| must be > 0 and |key_type| must be valid.
   Caller retains ownership of all pointers and they must not be NULL. */
OEMCryptoResult InitializeCryptoKey(CryptoKey* key,
                                    const uint8_t* serialized_bytes,
                                    uint32_t serialized_bytes_length,
                                    CryptoKeyType key_type,
                                    CryptoKeyOperation key_operation,
                                    CryptoKeySize key_size);

/* Initializes the data fields in the |key| and sets |key|'s key_handle to
   the key handle formed from |wrapped_key| and |size|.
   The data in |wrapped_key| was previously computed in WrapCryptoKey.
   Returns the result of CreateKeyHandle if it fails and OEMCrypto_SUCCESS
   otherwise.
   |serialized_bytes_length| must be > 0 and |key_type| must be valid.
   Caller retains ownership of all pointers and they must not be NULL. */
OEMCryptoResult InitializeCryptoKeyFromWrappedKey(
    CryptoKey* key, const uint8_t* wrapped_key, uint32_t wrapped_key_length,
    CryptoKeyType key_type, CryptoKeyOperation key_operation,
    CryptoKeySize key_size);

/* Wraps the key data into a buffer that can be saved to the file system. The
   wrapping must be device unique.  Caller retains ownership of |key| and
   |buffer| and they must not be NULL. Caller ensures that buffer_length is at
   least as big as the wrapped key size specified in
   oemcrypto_config_macros.h. */
OEMCryptoResult WrapCryptoKey(CryptoKey* key, uint8_t* buffer,
                              size_t buffer_length);

/* Frees the key handle associated with this key if it exists and then clears
   the key. Returns the result of freeing the key handle.
   Caller retains ownership of |key| and it must not be NULL. */
OEMCryptoResult FreeCryptoKey(CryptoKey* key);

/* Checks to make sure that |key| isn't NULL, its cookie is valid, its key_type
   matches |key_type|, and its key operation matches |key_operation|. */
bool CheckKey(CryptoKey* key, CryptoKeyType key_type,
              CryptoKeyOperation key_operation);

/* Parses the given |kcb| into a KeyControlBlock. If the verification fails, the
   key control block is marked invalid. Returns the parsed KeyControlBlock.
   Caller retains ownership of all pointers and they must not be NULL. */
KeyControlBlock ParseKeyControlBlock(const uint8_t* kcb);

#endif /* OEMCRYPTO_TA_OEMCRYPTO_KEY_H_ */