Crypto: Align implementation to pass PSA API compliance

-- Enable the option to specify maximum supported key
   length and maximum number of key slots at build time
   for the key module
-- Enable the option to specify internal buffer size
   for scratch allocation at build time for the
   engine module
-- Make sure that MD-2 and MD-4 hashes are enabled and
   supported by the back end as they are tested by the
   PSA API compliance tests
-- Other alignment needed to pass PSA API compliance
   tests, as changes in return codes, more error
   checking, and documentation update when needed

Change-Id: I4bb78b06de2fa01580c4cbd361c946d32c614240
Signed-off-by: Jamie Fox <jamie.fox@arm.com>
Co-Authored-by: Antonio de Angelis <antonio.deangelis@arm.com>
diff --git a/secure_fw/services/crypto/CMakeLists.inc b/secure_fw/services/crypto/CMakeLists.inc
index 4b572eb..3cf8dc2 100644
--- a/secure_fw/services/crypto/CMakeLists.inc
+++ b/secure_fw/services/crypto/CMakeLists.inc
@@ -67,6 +67,21 @@
   #Inform the user about Crypto service features selected based on the Crypto service cmake flags
   message("The Crypto service compile configuration is as follows:")
   message("- CRYPTO_ENGINE_MBEDTLS: " ${CRYPTO_ENGINE_MBEDTLS})
+  if (NOT DEFINED CRYPTO_KEY_STORAGE_NUM)
+    message("- CRYPTO_KEY_STORAGE_NUM using default value")
+  else()
+    message("- CRYPTO_KEY_STORAGE_NUM: " ${CRYPTO_KEY_STORAGE_NUM})
+  endif()
+  if (NOT DEFINED CRYPTO_KEY_MAX_KEY_LENGTH)
+    message("- CRYPTO_KEY_MAX_KEY_LENGTH using default value")
+  else()
+    message("- CRYPTO_KEY_MAX_KEY_LENGTH: " ${CRYPTO_KEY_MAX_KEY_LENGTH})
+  endif()
+  if (NOT DEFINED CRYPTO_ENGINE_BUF_SIZE)
+    message("- CRYPTO_ENGINE_BUF_SIZE using default value")
+  else()
+    message("- CRYPTO_ENGINE_BUF_SIZE: " ${CRYPTO_ENGINE_BUF_SIZE})
+  endif()
 
 else()
   message(FATAL_ERROR "Build system currently doesn't support selectively disabling of a service.")
diff --git a/secure_fw/services/crypto/CMakeLists.txt b/secure_fw/services/crypto/CMakeLists.txt
index f7074c7..04bda24 100644
--- a/secure_fw/services/crypto/CMakeLists.txt
+++ b/secure_fw/services/crypto/CMakeLists.txt
@@ -18,11 +18,11 @@
 project(tfm_crypto LANGUAGES ASM C)
 embedded_project_fixup()
 
-###Some project global settings
+#Some project global settings
 set (CRYPTO_DIR "${CMAKE_CURRENT_LIST_DIR}")
 get_filename_component(TFM_ROOT_DIR "${CRYPTO_DIR}/../../.." ABSOLUTE)
 
-###Get the definition of what files we need to build
+#Get the definition of what files we need to build
 set (ENABLE_CRYPTO ON)
 if (NOT DEFINED CRYPTO_ENGINE_MBEDTLS)
     set (CRYPTO_ENGINE_MBEDTLS ON)
@@ -38,13 +38,31 @@
 
 include(CMakeLists.inc)
 
-###Configure how we build our target
+#Configure how we build our target
 if(DEFINED CORE_TEST)
   set (TFM_LVL 3)
 else()
   set (TFM_LVL 1)
 endif()
 
+#Create a list of the C defines
+list(APPEND TFM_CRYPTO_C_DEFINES_LIST __ARM_FEATURE_CMSE=3 __thumb2__ TFM_LVL=${TFM_LVL})
+
+if (CRYPTO_ENGINE_MBEDTLS)
+  list(APPEND TFM_CRYPTO_C_DEFINES_LIST TFM_CRYPTO_ENGINE_MBEDTLS MBEDTLS_CONFIG_FILE="platform/ext/common/tfm_mbedtls_config.h")
+endif()
+
+#Add module configuration parameters in case they are provided during cmake configuration step
+if (DEFINED CRYPTO_KEY_STORAGE_NUM)
+  list(APPEND TFM_CRYPTO_C_DEFINES_LIST TFM_CRYPTO_KEY_STORAGE_NUM=${CRYPTO_KEY_STORAGE_NUM})
+endif()
+if (DEFINED CRYPTO_KEY_MAX_KEY_LENGTH)
+  list(APPEND TFM_CRYPTO_C_DEFINES_LIST TFM_CRYPTO_KEY_MAX_KEY_LENGTH=${CRYPTO_KEY_MAX_KEY_LENGTH})
+endif()
+if (DEFINED CRYPTO_ENGINE_BUF_SIZE)
+  list(APPEND TFM_CRYPTO_C_DEFINES_LIST TFM_CRYPTO_ENGINE_BUF_SIZE=${CRYPTO_ENGINE_BUF_SIZE})
+endif()
+
 if (CRYPTO_ENGINE_MBEDTLS)
   #Set mbed TLS compiler flags
   set(MBEDTLS_C_FLAGS ${MBEDTLS_C_FLAGS_SERVICES})
@@ -61,15 +79,14 @@
 
 #Specify what we build (for the crypto service, build as a static library)
 add_library(tfm_crypto STATIC ${ALL_SRC_ASM} ${ALL_SRC_C})
+embedded_set_target_compile_defines(TARGET tfm_crypto LANGUAGE C DEFINES ${TFM_CRYPTO_C_DEFINES_LIST})
 if (CRYPTO_ENGINE_MBEDTLS)
-  embedded_set_target_compile_defines(TARGET tfm_crypto LANGUAGE C DEFINES __ARM_FEATURE_CMSE=3 __thumb2__ TFM_LVL=${TFM_LVL} MBEDTLS_CONFIG_FILE="platform/ext/common/tfm_mbedtls_config.h" TFM_CRYPTO_ENGINE_MBEDTLS)
   #Add a dependency on the mbed_tls_lib_install target.
   add_dependencies(tfm_crypto ${MBEDTLS_TARGET_NAME}_install)
   #Ask the compiler to merge the mbed TLS and the crypto libraries.
   compiler_merge_library(DEST tfm_crypto LIBS "${MBEDTLS_INSTALL_DIR}/lib/${CMAKE_STATIC_LIBRARY_PREFIX_C}mbedcrypto${CMAKE_STATIC_LIBRARY_SUFFIX_C}")
-else()
-  embedded_set_target_compile_defines(TARGET tfm_crypto LANGUAGE C DEFINES __ARM_FEATURE_CMSE=3 __thumb2__ TFM_LVL=${TFM_LVL})
 endif()
+
 #Set common compiler and linker flags
 config_setting_shared_compiler_flags(tfm_crypto)
 config_setting_shared_linker_flags(tfm_crypto)
diff --git a/secure_fw/services/crypto/crypto_aead.c b/secure_fw/services/crypto/crypto_aead.c
index ec8b3a1..1698fd0 100644
--- a/secure_fw/services/crypto/crypto_aead.c
+++ b/secure_fw/services/crypto/crypto_aead.c
@@ -16,6 +16,16 @@
 #include "tfm_crypto_api.h"
 #include "crypto_utils.h"
 
+/**
+ * \def CRYPTO_AEAD_MAX_KEY_LENGTH
+ *
+ * \brief Specifies the maximum key length supported by the
+ *        AEAD operations in this implementation
+ */
+#ifndef CRYPTO_AEAD_MAX_KEY_LENGTH
+#define CRYPTO_AEAD_MAX_KEY_LENGTH (32)
+#endif
+
 /*!
  * \defgroup public_psa Public functions, PSA
  *
@@ -36,12 +46,15 @@
 {
     psa_status_t status = PSA_SUCCESS;
     enum tfm_crypto_err_t err;
-    uint8_t key_data[TFM_CRYPTO_MAX_KEY_LENGTH];
+    uint8_t key_data[CRYPTO_AEAD_MAX_KEY_LENGTH];
     uint32_t key_size;
     psa_key_type_t key_type;
 
+    /* Initialise ciphertext_length to zero */
+    *ciphertext_length = 0;
+
     if (PSA_ALG_IS_AEAD(alg) == 0) {
-        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
+        return TFM_CRYPTO_ERR_PSA_ERROR_NOT_SUPPORTED;
     }
 
     if (PSA_AEAD_TAG_SIZE(alg) == 0) {
@@ -49,7 +62,7 @@
     }
 
     if (PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg, plaintext_length) > ciphertext_size) {
-        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
+        return TFM_CRYPTO_ERR_PSA_ERROR_BUFFER_TOO_SMALL;
     }
 
     if ((nonce_length == 0) ||
@@ -98,7 +111,12 @@
     /* Access the key data */
     err = tfm_crypto_get_key_information(key, &key_type, (size_t *)&key_size);
     if (err != TFM_CRYPTO_ERR_PSA_SUCCESS) {
-        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
+        return err;
+    }
+
+    /* Support only AES based AEAD */
+    if (key_type != PSA_KEY_TYPE_AES) {
+        return TFM_CRYPTO_ERR_PSA_ERROR_NOT_SUPPORTED;
     }
 
     /* Access the crypto service key module to retrieve key data */
@@ -106,7 +124,7 @@
                              PSA_KEY_USAGE_ENCRYPT,
                              alg,
                              key_data,
-                             TFM_CRYPTO_MAX_KEY_LENGTH,
+                             CRYPTO_AEAD_MAX_KEY_LENGTH,
                              (size_t *)&key_size);
     if (err != TFM_CRYPTO_ERR_PSA_SUCCESS) {
         return err;
@@ -126,6 +144,13 @@
                                             ciphertext,
                                             ciphertext_size,
                                             (uint32_t *)ciphertext_length);
+    if (status == PSA_SUCCESS) {
+        /* The ciphertext_length needs to take into account the tag length */
+        *ciphertext_length += PSA_AEAD_TAG_SIZE(alg);
+    } else {
+        /* In case of failure set the ciphertext_length to zero */
+        *ciphertext_length = 0;
+    }
 
     return PSA_STATUS_TO_TFM_CRYPTO_ERR(status);
 }
@@ -144,20 +169,24 @@
 {
     psa_status_t status = PSA_SUCCESS;
     enum tfm_crypto_err_t err;
-    uint8_t key_data[TFM_CRYPTO_MAX_KEY_LENGTH];
+    uint8_t key_data[CRYPTO_AEAD_MAX_KEY_LENGTH];
     uint32_t key_size;
     psa_key_type_t key_type;
 
+    /* Initialise plaintext_length to zero */
+    *plaintext_length = 0;
+
     if (PSA_ALG_IS_AEAD(alg) == 0) {
+        return TFM_CRYPTO_ERR_PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    if ((PSA_AEAD_TAG_SIZE(alg) == 0) ||
+        (ciphertext_length < PSA_AEAD_TAG_SIZE(alg))) {
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    if (PSA_AEAD_TAG_SIZE(alg) == 0) {
-        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
-    }
-
-    if (PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg, ciphertext_length) > plaintext_size) {
-        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
+    if (PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg,ciphertext_length) > plaintext_size) {
+        return TFM_CRYPTO_ERR_PSA_ERROR_BUFFER_TOO_SMALL;
     }
 
     if ((nonce_length == 0) ||
@@ -206,7 +235,12 @@
     /* Access the key data */
     err = tfm_crypto_get_key_information(key, &key_type, (size_t *)&key_size);
     if (err != TFM_CRYPTO_ERR_PSA_SUCCESS) {
-        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
+        return err;
+    }
+
+    /* Support only AES based AEAD */
+    if (key_type != PSA_KEY_TYPE_AES) {
+        return TFM_CRYPTO_ERR_PSA_ERROR_NOT_SUPPORTED;
     }
 
     /* Access the crypto service key module to retrieve key data */
@@ -214,7 +248,7 @@
                              PSA_KEY_USAGE_DECRYPT,
                              alg,
                              key_data,
-                             TFM_CRYPTO_MAX_KEY_LENGTH,
+                             CRYPTO_AEAD_MAX_KEY_LENGTH,
                              (size_t *)&key_size);
     if (err != TFM_CRYPTO_ERR_PSA_SUCCESS) {
         return err;
@@ -234,6 +268,9 @@
                                             plaintext,
                                             plaintext_size,
                                             (uint32_t *)plaintext_length);
+    if (status != PSA_SUCCESS) {
+        *plaintext_length = 0;
+    }
 
     return PSA_STATUS_TO_TFM_CRYPTO_ERR(status);
 }
diff --git a/secure_fw/services/crypto/crypto_alloc.c b/secure_fw/services/crypto/crypto_alloc.c
index 2a4d6bd..627dab7 100644
--- a/secure_fw/services/crypto/crypto_alloc.c
+++ b/secure_fw/services/crypto/crypto_alloc.c
@@ -120,6 +120,6 @@
         return TFM_CRYPTO_ERR_PSA_SUCCESS;
     }
 
-    return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
+    return TFM_CRYPTO_ERR_PSA_ERROR_BAD_STATE;
 }
 /*!@}*/
diff --git a/secure_fw/services/crypto/crypto_cipher.c b/secure_fw/services/crypto/crypto_cipher.c
index 84b5b78..a88f4a8 100644
--- a/secure_fw/services/crypto/crypto_cipher.c
+++ b/secure_fw/services/crypto/crypto_cipher.c
@@ -18,13 +18,23 @@
 #include "tfm_crypto_api.h"
 #include "crypto_utils.h"
 
+/**
+ * \def CRYPTO_CIPHER_MAX_KEY_LENGTH
+ *
+ * \brief Specifies the maximum key length supported by the
+ *        Cipher operations in this implementation
+ */
+#ifndef CRYPTO_CIPHER_MAX_KEY_LENGTH
+#define CRYPTO_CIPHER_MAX_KEY_LENGTH (32)
+#endif
+
 static enum tfm_crypto_err_t tfm_crypto_cipher_setup(
                                            psa_cipher_operation_t *operation,
                                            psa_key_slot_t key,
                                            psa_algorithm_t alg,
                                            enum engine_cipher_mode_t c_mode)
 {
-    uint8_t key_data[TFM_CRYPTO_MAX_KEY_LENGTH];
+    uint8_t key_data[CRYPTO_CIPHER_MAX_KEY_LENGTH];
     size_t key_size;
     psa_key_type_t key_type = PSA_KEY_TYPE_NONE;
     psa_status_t status = PSA_SUCCESS;
@@ -48,6 +58,16 @@
     /* Access the key module to retrieve key related information */
     err = tfm_crypto_get_key_information(key, &key_type, &key_size);
     if (err != TFM_CRYPTO_ERR_PSA_SUCCESS) {
+        return err;
+    }
+
+    /* Check if it's a raw data key type */
+    if (key_type == PSA_KEY_TYPE_RAW_DATA) {
+        return TFM_CRYPTO_ERR_PSA_ERROR_NOT_PERMITTED;
+    }
+
+    /* Check compatibility between key and algorithm */
+    if ((key_type == PSA_KEY_TYPE_ARC4) && (alg != PSA_ALG_ARC4)) {
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
@@ -101,7 +121,7 @@
                              usage,
                              alg,
                              key_data,
-                             TFM_CRYPTO_MAX_KEY_LENGTH,
+                             CRYPTO_CIPHER_MAX_KEY_LENGTH,
                              &key_size);
     if (err != TFM_CRYPTO_ERR_PSA_SUCCESS) {
         /* Release the operation context */
@@ -260,6 +280,9 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
+    /* Initialise the output length to zero */
+    *output_length = 0;
+
     /* Look up the corresponding operation context */
     err = tfm_crypto_operation_lookup(TFM_CRYPTO_CIPHER_OPERATION,
                                       operation->handle,
@@ -278,8 +301,6 @@
         /* This call is used to set the IV on the object */
         err = tfm_crypto_cipher_set_iv(operation, input, input_length);
 
-        *output_length = 0;
-
         return err;
     }
 
@@ -288,6 +309,13 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_BAD_STATE;
     }
 
+    /* FIXME: The implementation currently expects to work only on blocks
+     *        of input data whose length is equal to the block size
+     */
+    if (input_length > output_size) {
+        return TFM_CRYPTO_ERR_PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+
     /* Update the cipher output with the input chunk on the engine */
     status = tfm_crypto_engine_cipher_update(&(ctx->engine_ctx),
                                              input,
@@ -339,6 +367,13 @@
         return err;
     }
 
+    /* Check that the output buffer is large enough for up to one block size of
+     * output data.
+     */
+    if (output_size < ctx->block_size) {
+        return TFM_CRYPTO_ERR_PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+
     /* Finalise the operation on the crypto engine */
     status = tfm_crypto_engine_cipher_finish(&(ctx->engine_ctx),
                                              output,
diff --git a/secure_fw/services/crypto/crypto_engine.c b/secure_fw/services/crypto/crypto_engine.c
index f9fe896..dd09481 100644
--- a/secure_fw/services/crypto/crypto_engine.c
+++ b/secure_fw/services/crypto/crypto_engine.c
@@ -14,11 +14,19 @@
  */
 #include "crypto_engine.h"
 
+/**
+ * \brief Default value for the size of the static buffer used by the Engine
+ *        module as a scratch buffer for its own internal allocations
+ */
+#ifndef TFM_CRYPTO_ENGINE_BUF_SIZE
+#define TFM_CRYPTO_ENGINE_BUF_SIZE (1024)
+#endif
+
 #if defined(TFM_CRYPTO_ENGINE_MBEDTLS)
 /**
  * \brief Buffer size used by Mbed TLS for its allocations
  */
-#define TFM_CRYPTO_MBEDTLS_MEM_BUF_LEN (1024)
+#define TFM_CRYPTO_MBEDTLS_MEM_BUF_LEN (TFM_CRYPTO_ENGINE_BUF_SIZE)
 
 /**
  * \brief Static buffer to be used by Mbed TLS for memory allocations
@@ -221,8 +229,14 @@
         return PSA_SUCCESS;
     }
 
-    /* FIXME: For the time being map all errors to PSA_ERROR_UNKNOW_ERROR */
+    /* FIXME: Investigate all possible Mbed TLS errors and map them
+     *        to the the correct corresponding PSA status
+     */
     switch (ret) {
+    case MBEDTLS_ERR_CIPHER_FULL_BLOCK_EXPECTED:
+        return PSA_ERROR_INVALID_ARGUMENT;
+    case MBEDTLS_ERR_CIPHER_AUTH_FAILED:
+        return PSA_ERROR_INVALID_SIGNATURE;
     default:
         return PSA_ERROR_UNKNOWN_ERROR;
     }
diff --git a/secure_fw/services/crypto/crypto_key.c b/secure_fw/services/crypto/crypto_key.c
index 90638f2..e101420 100644
--- a/secure_fw/services/crypto/crypto_key.c
+++ b/secure_fw/services/crypto/crypto_key.c
@@ -5,6 +5,7 @@
  *
  */
 
+#include <stdbool.h>
 #include <stddef.h>
 
 #include "tfm_crypto_api.h"
@@ -14,10 +15,20 @@
 #include "secure_fw/core/secure_utilities.h"
 
 /**
- * \brief This value defines the maximum number of simultaneous key stores
- *        supported by this implementation.
+ * \brief This is the default value of maximum number of simultaneous
+ *        key stores supported.
  */
+#ifndef TFM_CRYPTO_KEY_STORAGE_NUM
 #define TFM_CRYPTO_KEY_STORAGE_NUM (4)
+#endif
+
+/**
+ * \brief This is the default value of the maximum supported key length
+ *        in bytes.
+ */
+#ifndef TFM_CRYPTO_MAX_KEY_LENGTH
+#define TFM_CRYPTO_MAX_KEY_LENGTH (32)
+#endif
 
 struct tfm_crypto_key_storage_s {
     uint8_t in_use;                 /*!< Indicates if the key store is in use */
@@ -30,6 +41,56 @@
 
 static struct tfm_crypto_key_storage_s
                                 key_storage[TFM_CRYPTO_KEY_STORAGE_NUM] = {{0}};
+
+/**
+ * \brief Get a pointer to the key store for the provided key slot.
+ *
+ * \param[in] key  Key slot
+ *
+ * \return Pointer to key store or NULL if key is not a valid key slot
+ */
+static struct tfm_crypto_key_storage_s *get_key_store(psa_key_slot_t key)
+{
+    if (key == 0 || key > TFM_CRYPTO_KEY_STORAGE_NUM) {
+        return NULL;
+    }
+
+    return &key_storage[key - 1];
+}
+
+/**
+ * \brief Check that the key type is supported and that key_length is a
+ *        supported key length for that key type.
+ *
+ * \param[in] type        Key type
+ * \param[in] key_length  Key data length in bytes
+ *
+ * \return True if the key type is supported and key_length is a supported
+ *         key length for that key type, false otherwise
+ */
+static bool key_type_is_supported(psa_key_type_t type, size_t key_length)
+{
+    if (key_length > TFM_CRYPTO_MAX_KEY_LENGTH) {
+        return false;
+    }
+
+    switch (type) {
+    case PSA_KEY_TYPE_RAW_DATA:
+    case PSA_KEY_TYPE_HMAC:
+    case PSA_KEY_TYPE_DERIVE:
+        return true; /* No further restictions on these key types */
+    case PSA_KEY_TYPE_AES:
+    case PSA_KEY_TYPE_CAMELLIA:
+        return (key_length == 16 || key_length == 24 || key_length == 32);
+    case PSA_KEY_TYPE_DES:
+        return (key_length == 8 || key_length == 16 || key_length == 24);
+    case PSA_KEY_TYPE_ARC4:
+        return key_length >= 1;
+    default:
+        return false; /* Other key types are not supported */
+    }
+}
+
 /*!
  * \defgroup public Public functions
  *
@@ -50,35 +111,37 @@
                                          size_t data_size,
                                          size_t *data_length)
 {
+    struct tfm_crypto_key_storage_s *key_store;
     size_t i;
 
-    if (key >= TFM_CRYPTO_KEY_STORAGE_NUM) {
+    key_store = get_key_store(key);
+    if (key_store == NULL) {
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    if (key_storage[key].in_use == TFM_CRYPTO_NOT_IN_USE) {
+    if (key_store->in_use == TFM_CRYPTO_NOT_IN_USE) {
         return TFM_CRYPTO_ERR_PSA_ERROR_EMPTY_SLOT;
     }
 
     /* Check that usage is permitted for this key */
-    if ((usage & key_storage[key].policy.usage) != usage) {
+    if ((usage & key_store->policy.usage) != usage) {
         return TFM_CRYPTO_ERR_PSA_ERROR_NOT_PERMITTED;
     }
 
     /* Check that alg is compatible with this key */
-    if (alg != 0 && alg != key_storage[key].policy.alg) {
+    if (alg != 0 && alg != key_store->policy.alg) {
         return TFM_CRYPTO_ERR_PSA_ERROR_NOT_PERMITTED;
     }
 
-    if (key_storage[key].data_length > data_size) {
+    if (key_store->data_length > data_size) {
         return TFM_CRYPTO_ERR_PSA_ERROR_BUFFER_TOO_SMALL;
     }
 
-    for (i = 0; i < key_storage[key].data_length; i++) {
-        data[i] = key_storage[key].data[i];
+    for (i = 0; i < key_store->data_length; i++) {
+        data[i] = key_store->data[i];
     }
 
-    *data_length = key_storage[key].data_length;
+    *data_length = key_store->data_length;
 
     return TFM_CRYPTO_ERR_PSA_SUCCESS;
 }
@@ -89,6 +152,7 @@
                                             size_t data_length)
 {
     enum tfm_crypto_err_t err;
+    struct tfm_crypto_key_storage_s *key_store;
     size_t i;
 
     err = tfm_crypto_memory_check((uint8_t *)data, data_length,
@@ -97,39 +161,42 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    if (key >= TFM_CRYPTO_KEY_STORAGE_NUM) {
+    key_store = get_key_store(key);
+    if (key_store == NULL) {
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    if (key_storage[key].in_use != TFM_CRYPTO_NOT_IN_USE) {
+    if (key_store->in_use != TFM_CRYPTO_NOT_IN_USE) {
         return TFM_CRYPTO_ERR_PSA_ERROR_OCCUPIED_SLOT;
     }
 
-    if (data_length > TFM_CRYPTO_MAX_KEY_LENGTH) {
+    if (!key_type_is_supported(type, data_length)) {
         return TFM_CRYPTO_ERR_PSA_ERROR_NOT_SUPPORTED;
     }
 
-    key_storage[key].in_use = TFM_CRYPTO_IN_USE;
-    key_storage[key].type = type;
+    key_store->in_use = TFM_CRYPTO_IN_USE;
+    key_store->type = type;
 
     for (i=0; i<data_length; i++) {
-        key_storage[key].data[i] = data[i];
+        key_store->data[i] = data[i];
     }
 
-    key_storage[key].data_length = data_length;
+    key_store->data_length = data_length;
 
     return TFM_CRYPTO_ERR_PSA_SUCCESS;
 }
 
 enum tfm_crypto_err_t tfm_crypto_destroy_key(psa_key_slot_t key)
 {
+    struct tfm_crypto_key_storage_s *key_store;
     uint32_t i;
 
-    if (key >= TFM_CRYPTO_KEY_STORAGE_NUM) {
+    key_store = get_key_store(key);
+    if (key_store == NULL) {
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    volatile uint8_t *p_mem = (uint8_t *) &(key_storage[key]);
+    volatile uint8_t *p_mem = (uint8_t *)key_store;
     uint32_t size_mem = sizeof(struct tfm_crypto_key_storage_s);
 
     /* memset the key_storage */
@@ -138,7 +205,7 @@
     }
 
     /* Set default values */
-    key_storage[key].in_use = TFM_CRYPTO_NOT_IN_USE;
+    key_store->in_use = TFM_CRYPTO_NOT_IN_USE;
 
     return TFM_CRYPTO_ERR_PSA_SUCCESS;
 }
@@ -148,6 +215,7 @@
                                                      size_t *bits)
 {
     enum tfm_crypto_err_t err;
+    struct tfm_crypto_key_storage_s *key_store;
 
     err = tfm_crypto_memory_check(type, sizeof(psa_key_type_t),
                                   TFM_MEMORY_ACCESS_RW);
@@ -160,17 +228,22 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    if (key >= TFM_CRYPTO_KEY_STORAGE_NUM) {
+    /* Initialise output parameters contents to zero */
+    *type = (psa_key_type_t) 0;
+    *bits = (size_t)0;
+
+    key_store = get_key_store(key);
+    if (key_store == NULL) {
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    if (key_storage[key].in_use == TFM_CRYPTO_NOT_IN_USE) {
+    if (key_store->in_use == TFM_CRYPTO_NOT_IN_USE) {
         return TFM_CRYPTO_ERR_PSA_ERROR_EMPTY_SLOT;
     }
 
     /* Get basic metadata */
-    *type = key_storage[key].type;
-    *bits = PSA_BYTES_TO_BITS(key_storage[key].data_length);
+    *type = key_store->type;
+    *bits = PSA_BYTES_TO_BITS(key_store->data_length);
 
     return TFM_CRYPTO_ERR_PSA_SUCCESS;
 }
@@ -294,6 +367,7 @@
                                                 const psa_key_policy_t *policy)
 {
     enum tfm_crypto_err_t err;
+    struct tfm_crypto_key_storage_s *key_store;
 
     err = tfm_crypto_memory_check((psa_key_policy_t *)policy,
                                   sizeof(psa_key_policy_t),
@@ -302,10 +376,6 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    if (key >= TFM_CRYPTO_KEY_STORAGE_NUM) {
-        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
-    }
-
     /* Check that the policy is valid */
     if (policy->usage & ~(PSA_KEY_USAGE_EXPORT
                           | PSA_KEY_USAGE_ENCRYPT
@@ -316,14 +386,19 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
+    key_store = get_key_store(key);
+    if (key_store == NULL) {
+        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     /* Changing the policy of an occupied slot is not permitted as
      * this is a requirement of the PSA Crypto API
      */
-    if (key_storage[key].in_use != TFM_CRYPTO_NOT_IN_USE) {
+    if (key_store->in_use != TFM_CRYPTO_NOT_IN_USE) {
         return TFM_CRYPTO_ERR_PSA_ERROR_OCCUPIED_SLOT;
     }
 
-    key_storage[key].policy = *policy;
+    key_store->policy = *policy;
 
     return TFM_CRYPTO_ERR_PSA_SUCCESS;
 }
@@ -332,6 +407,7 @@
                                                 psa_key_policy_t *policy)
 {
     enum tfm_crypto_err_t err;
+    struct tfm_crypto_key_storage_s *key_store;
 
     err = tfm_crypto_memory_check(policy, sizeof(psa_key_policy_t),
                                   TFM_MEMORY_ACCESS_RW);
@@ -339,11 +415,12 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    if (key >= TFM_CRYPTO_KEY_STORAGE_NUM) {
+    key_store = get_key_store(key);
+    if (key_store == NULL) {
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    *policy = key_storage[key].policy;
+    *policy = key_store->policy;
 
     return TFM_CRYPTO_ERR_PSA_SUCCESS;
 }
@@ -351,9 +428,7 @@
 enum tfm_crypto_err_t tfm_crypto_set_key_lifetime(psa_key_slot_t key,
                                                   psa_key_lifetime_t lifetime)
 {
-    if (key >= TFM_CRYPTO_KEY_STORAGE_NUM) {
-        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
-    }
+    struct tfm_crypto_key_storage_s *key_store;
 
     /* Check that the lifetime is valid */
     if (lifetime != PSA_KEY_LIFETIME_VOLATILE
@@ -362,10 +437,15 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
+    key_store = get_key_store(key);
+    if (key_store == NULL) {
+        return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     /* TF-M Crypto service does not support changing the lifetime of an occupied
      * slot.
      */
-    if (key_storage[key].in_use != TFM_CRYPTO_NOT_IN_USE) {
+    if (key_store->in_use != TFM_CRYPTO_NOT_IN_USE) {
         return TFM_CRYPTO_ERR_PSA_ERROR_OCCUPIED_SLOT;
     }
 
@@ -374,7 +454,7 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_NOT_SUPPORTED;
     }
 
-    key_storage[key].lifetime = lifetime;
+    key_store->lifetime = lifetime;
 
     return TFM_CRYPTO_ERR_PSA_SUCCESS;
 }
@@ -383,6 +463,7 @@
                                                   psa_key_lifetime_t *lifetime)
 {
     enum tfm_crypto_err_t err;
+    struct tfm_crypto_key_storage_s *key_store;
 
     err = tfm_crypto_memory_check(lifetime, sizeof(psa_key_lifetime_t),
                                   TFM_MEMORY_ACCESS_RW);
@@ -390,11 +471,12 @@
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    if (key >= TFM_CRYPTO_KEY_STORAGE_NUM) {
+    key_store = get_key_store(key);
+    if (key_store == NULL) {
         return TFM_CRYPTO_ERR_PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    *lifetime = key_storage[key].lifetime;
+    *lifetime = key_store->lifetime;
 
     return TFM_CRYPTO_ERR_PSA_SUCCESS;
 }
diff --git a/secure_fw/services/crypto/crypto_mac.c b/secure_fw/services/crypto/crypto_mac.c
index 25a98da..a13a00b 100644
--- a/secure_fw/services/crypto/crypto_mac.c
+++ b/secure_fw/services/crypto/crypto_mac.c
@@ -15,6 +15,16 @@
 #include "tfm_crypto_api.h"
 #include "crypto_utils.h"
 
+/**
+ * \def CRYPTO_HMAC_MAX_KEY_LENGTH
+ *
+ * \brief Specifies the maximum key length supported by the
+ *        HMAC operations in this implementation
+ */
+#ifndef CRYPTO_HMAC_MAX_KEY_LENGTH
+#define CRYPTO_HMAC_MAX_KEY_LENGTH (32)
+#endif
+
 static void mac_zeroize(void *data, size_t size)
 {
     tfm_memset(data, 0, size);
@@ -54,7 +64,7 @@
     enum tfm_crypto_err_t err;
     psa_key_type_t key_type;
     size_t key_size;
-    uint8_t key_data[TFM_CRYPTO_MAX_KEY_LENGTH];
+    uint8_t key_data[CRYPTO_HMAC_MAX_KEY_LENGTH];
     uint8_t hashed_key[PSA_HMAC_MAX_HASH_BLOCK_SIZE];
     size_t block_size;
     uint8_t ipad[PSA_HMAC_MAX_HASH_BLOCK_SIZE];
@@ -86,7 +96,7 @@
                              usage,
                              alg,
                              key_data,
-                             TFM_CRYPTO_MAX_KEY_LENGTH,
+                             CRYPTO_HMAC_MAX_KEY_LENGTH,
                              &key_size);
     if (err != TFM_CRYPTO_ERR_PSA_SUCCESS) {
         return err;