Crypto: Upgrade Mbed TLS to 2.25

Set the MBEDCRYPTO_VERSION to 2.25.0.

First three patches in existing v2.24 already applied in v2.25
and hence removed.

Replaced MBEDTLS_PSA_CRYPTO_KEY_FILE_ID_ENCODES_OWNER with
MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER in all configuration and
source as updated in v2.25 library.

Update all headers of psa/include as per mbedtls-v2.25 excluding
changes required to hide some implementation.

Update id field in the client_key_attributes structure to
psa_key_id_t.

Update Copyright year to 2021!

Removed patch 006 as not required in MbedTLS v2.25.0.

Update references of handle to key as per MbedTLS api changes.

Increase NUM_HANDLES to 32 to accommodate crypto api tests.

Added corresponding tfm implementation of psa_purge_key().

Signed-off-by: Maulik Patel <maulik.patel@arm.com>
Change-Id: I6a532da96735cf32996250c4a8733a8654c1f44e
diff --git a/secure_fw/partitions/crypto/crypto_aead.c b/secure_fw/partitions/crypto/crypto_aead.c
index 53b0848..d4aee11 100644
--- a/secure_fw/partitions/crypto/crypto_aead.c
+++ b/secure_fw/partitions/crypto/crypto_aead.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -38,7 +38,7 @@
 
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
     const struct tfm_crypto_aead_pack_input *aead_pack_input = &iov->aead_in;
-    psa_key_handle_t key_handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
     const uint8_t *nonce = aead_pack_input->nonce;
     size_t nonce_length = aead_pack_input->nonce_length;
@@ -48,20 +48,25 @@
     size_t ciphertext_size = out_vec[0].len;
     const uint8_t *additional_data = in_vec[2].base;
     size_t additional_data_length = in_vec[2].len;
+    mbedtls_svc_key_id_t encoded_key;
 
     /* Initialise ciphertext_length to zero. */
     out_vec[0].len = 0;
 
-    status = tfm_crypto_check_handle_owner(key_handle, NULL);
-    if (status == PSA_SUCCESS) {
-
-        status = psa_aead_encrypt(key_handle, alg, nonce, nonce_length,
-                                  additional_data, additional_data_length,
-                                  plaintext, plaintext_length,
-                                  ciphertext, ciphertext_size, &out_vec[0].len);
+    status = tfm_crypto_check_handle_owner(key_id, NULL);
+    if (status != PSA_SUCCESS) {
+        return status;
     }
 
-    return status;
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    return psa_aead_encrypt(encoded_key, alg, nonce, nonce_length,
+                            additional_data, additional_data_length,
+                            plaintext, plaintext_length,
+                            ciphertext, ciphertext_size, &out_vec[0].len);
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
 
@@ -83,7 +88,7 @@
 
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
     const struct tfm_crypto_aead_pack_input *aead_pack_input = &iov->aead_in;
-    psa_key_handle_t key_handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
     const uint8_t *nonce = aead_pack_input->nonce;
     size_t nonce_length = aead_pack_input->nonce_length;
@@ -93,20 +98,25 @@
     size_t plaintext_size = out_vec[0].len;
     const uint8_t *additional_data = in_vec[2].base;
     size_t additional_data_length = in_vec[2].len;
+    mbedtls_svc_key_id_t encoded_key;
 
     /* Initialise plaintext_length to zero. */
     out_vec[0].len = 0;
 
-    status = tfm_crypto_check_handle_owner(key_handle, NULL);
-    if (status == PSA_SUCCESS) {
-
-        status = psa_aead_decrypt(key_handle, alg, nonce, nonce_length,
-                                  additional_data, additional_data_length,
-                                  ciphertext, ciphertext_length,
-                                  plaintext, plaintext_size, &out_vec[0].len);
+    status = tfm_crypto_check_handle_owner(key_id, NULL);
+    if (status != PSA_SUCCESS) {
+        return status;
     }
 
-    return status;
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    return psa_aead_decrypt(encoded_key, alg, nonce, nonce_length,
+                            additional_data, additional_data_length,
+                            ciphertext, ciphertext_length,
+                            plaintext, plaintext_size, &out_vec[0].len);
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
 
diff --git a/secure_fw/partitions/crypto/crypto_asymmetric.c b/secure_fw/partitions/crypto/crypto_asymmetric.c
index 9f43eb4..c3a3909 100644
--- a/secure_fw/partitions/crypto/crypto_asymmetric.c
+++ b/secure_fw/partitions/crypto/crypto_asymmetric.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -35,19 +35,25 @@
     }
 
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_handle_t handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
     const uint8_t *hash = in_vec[1].base;
     size_t hash_length = in_vec[1].len;
     uint8_t *signature = out_vec[0].base;
     size_t signature_size = out_vec[0].len;
-    psa_status_t status = tfm_crypto_check_handle_owner(handle, NULL);
+    mbedtls_svc_key_id_t encoded_key;
 
+    psa_status_t status = tfm_crypto_check_handle_owner(key_id, NULL);
     if (status != PSA_SUCCESS) {
         return status;
     }
 
-    return psa_sign_hash(handle, alg, hash, hash_length,
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    return psa_sign_hash(encoded_key, alg, hash, hash_length,
                          signature, signature_size, &(out_vec[0].len));
 #endif /* TFM_CRYPTO_ASYMMETRIC_MODULE_DISABLED */
 }
@@ -68,19 +74,25 @@
 
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
 
-    psa_key_handle_t handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
     const uint8_t *hash = in_vec[1].base;
     size_t hash_length = in_vec[1].len;
     const uint8_t *signature = in_vec[2].base;
     size_t signature_length = in_vec[2].len;
-    psa_status_t status = tfm_crypto_check_handle_owner(handle, NULL);
+    mbedtls_svc_key_id_t encoded_key;
+    psa_status_t status = tfm_crypto_check_handle_owner(key_id, NULL);
 
     if (status != PSA_SUCCESS) {
         return status;
     }
 
-    return psa_verify_hash(handle, alg, hash, hash_length,
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    return psa_verify_hash(encoded_key, alg, hash, hash_length,
                            signature, signature_length);
 #endif /* TFM_CRYPTO_ASYMMETRIC_MODULE_DISABLED */
 }
@@ -102,7 +114,7 @@
     }
 
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_handle_t handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
     const uint8_t *input = in_vec[1].base;
     size_t input_length = in_vec[1].len;
@@ -113,13 +125,19 @@
     psa_key_type_t type;
     size_t key_bits;
     psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    mbedtls_svc_key_id_t encoded_key;
 
-    status = tfm_crypto_check_handle_owner(handle, NULL);
+    status = tfm_crypto_check_handle_owner(key_id, NULL);
     if (status != PSA_SUCCESS) {
         return status;
     }
 
-    status = psa_get_key_attributes(handle, &key_attributes);
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    status = psa_get_key_attributes(encoded_key, &key_attributes);
     if (status != PSA_SUCCESS) {
         return status;
     }
@@ -134,7 +152,7 @@
         return PSA_ERROR_BUFFER_TOO_SMALL;
     }
 
-    return psa_asymmetric_encrypt(handle, alg, input, input_length,
+    return psa_asymmetric_encrypt(encoded_key, alg, input, input_length,
                                   salt, salt_length,
                                   output, output_size, &(out_vec[0].len));
 #endif /* TFM_CRYPTO_ASYMMETRIC_MODULE_DISABLED */
@@ -156,7 +174,7 @@
     }
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
 
-    psa_key_handle_t handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
     const uint8_t *input = in_vec[1].base;
     size_t input_length = in_vec[1].len;
@@ -165,13 +183,19 @@
     uint8_t *output = out_vec[0].base;
     size_t output_size = out_vec[0].len;
     psa_status_t status;
+    mbedtls_svc_key_id_t encoded_key;
 
-    status = tfm_crypto_check_handle_owner(handle, NULL);
+    status = tfm_crypto_check_handle_owner(key_id, NULL);
     if (status != PSA_SUCCESS) {
         return status;
     }
 
-    return psa_asymmetric_decrypt(handle, alg, input, input_length,
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    return psa_asymmetric_decrypt(encoded_key, alg, input, input_length,
                                   salt, salt_length,
                                   output, output_size, &(out_vec[0].len));
 #endif /* TFM_CRYPTO_ASYMMETRIC_MODULE_DISABLED */
diff --git a/secure_fw/partitions/crypto/crypto_cipher.c b/secure_fw/partitions/crypto/crypto_cipher.c
index 6e47f61..03849df 100644
--- a/secure_fw/partitions/crypto/crypto_cipher.c
+++ b/secure_fw/partitions/crypto/crypto_cipher.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -133,10 +133,11 @@
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
     uint32_t handle = iov->op_handle;
     uint32_t *handle_out = out_vec[0].base;
-    psa_key_handle_t key_handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
+    mbedtls_svc_key_id_t encoded_key;
 
-    status = tfm_crypto_check_handle_owner(key_handle, NULL);
+    status = tfm_crypto_check_handle_owner(key_id, NULL);
     if (status != PSA_SUCCESS) {
         return status;
     }
@@ -148,10 +149,14 @@
     if (status != PSA_SUCCESS) {
         return status;
     }
-
     *handle_out = handle;
 
-    status = psa_cipher_encrypt_setup(operation, key_handle, alg);
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    status = psa_cipher_encrypt_setup(operation, encoded_key, alg);
     if (status != PSA_SUCCESS) {
         /* Release the operation context, ignore if the operation fails. */
         (void)tfm_crypto_operation_release(handle_out);
@@ -182,10 +187,11 @@
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
     uint32_t handle = iov->op_handle;
     uint32_t *handle_out = out_vec[0].base;
-    psa_key_handle_t key_handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
+    mbedtls_svc_key_id_t encoded_key;
 
-    status = tfm_crypto_check_handle_owner(key_handle, NULL);
+    status = tfm_crypto_check_handle_owner(key_id, NULL);
     if (status != PSA_SUCCESS) {
         return status;
     }
@@ -199,8 +205,12 @@
     }
 
     *handle_out = handle;
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
 
-    status = psa_cipher_decrypt_setup(operation, key_handle, alg);
+    status = psa_cipher_decrypt_setup(operation, encoded_key, alg);
     if (status != PSA_SUCCESS) {
         /* Release the operation context, ignore if the operation fails. */
         (void)tfm_crypto_operation_release(handle_out);
diff --git a/secure_fw/partitions/crypto/crypto_key.c b/secure_fw/partitions/crypto/crypto_key.c
index a2ca7d8..d07eef3 100644
--- a/secure_fw/partitions/crypto/crypto_key.c
+++ b/secure_fw/partitions/crypto/crypto_key.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -16,11 +16,11 @@
 #include <stdbool.h>
 
 #ifndef TFM_CRYPTO_MAX_KEY_HANDLES
-#define TFM_CRYPTO_MAX_KEY_HANDLES (16)
+#define TFM_CRYPTO_MAX_KEY_HANDLES (32)
 #endif
 struct tfm_crypto_handle_owner_s {
     int32_t owner;           /*!< Owner of the allocated handle */
-    psa_key_handle_t handle; /*!< Allocated handle */
+    psa_key_id_t key;        /*!< Allocated key */
     uint8_t in_use;          /*!< Flag to indicate if this in use */
 };
 
@@ -35,9 +35,9 @@
  */
 /*!@{*/
 psa_status_t tfm_crypto_key_attributes_from_client(
-                             const struct psa_client_key_attributes_s *client_key_attr,
-                             int32_t client_id,
-                             psa_key_attributes_t *key_attributes)
+                    const struct psa_client_key_attributes_s *client_key_attr,
+                    int32_t client_id,
+                    psa_key_attributes_t *key_attributes)
 {
     if (client_key_attr == NULL || key_attributes == NULL) {
         return PSA_ERROR_PROGRAMMER_ERROR;
@@ -60,8 +60,8 @@
 }
 
 psa_status_t tfm_crypto_key_attributes_to_client(
-                                   const psa_key_attributes_t *key_attributes,
-                                   struct psa_client_key_attributes_s *client_key_attr)
+                        const psa_key_attributes_t *key_attributes,
+                        struct psa_client_key_attributes_s *client_key_attr)
 {
     if (client_key_attr == NULL || key_attributes == NULL) {
         return PSA_ERROR_PROGRAMMER_ERROR;
@@ -83,7 +83,7 @@
     return PSA_SUCCESS;
 }
 
-psa_status_t tfm_crypto_check_handle_owner(psa_key_handle_t handle,
+psa_status_t tfm_crypto_check_handle_owner(psa_key_id_t key,
                                            uint32_t *index)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
@@ -99,7 +99,7 @@
     }
 
     for (i = 0; i < TFM_CRYPTO_MAX_KEY_HANDLES; i++) {
-        if (handle_owner[i].in_use && handle_owner[i].handle == handle) {
+        if (handle_owner[i].in_use && handle_owner[i].key == key) {
             if (handle_owner[i].owner == partition_id) {
                 if (index != NULL) {
                     *index = i;
@@ -115,6 +115,27 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
+psa_status_t tfm_crypto_encode_id_and_owner(psa_key_id_t key_id,
+                                            mbedtls_svc_key_id_t *enc_key_ptr)
+{
+    int32_t partition_id = 0;
+    psa_status_t status = tfm_crypto_get_caller_id(&partition_id);
+
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    /* If Null Pointer, return PSA_ERROR_PROGRAMMER_ERROR */
+    if (enc_key_ptr == NULL) {
+        return PSA_ERROR_PROGRAMMER_ERROR;
+    }
+
+    /* Use the client key id as the key_id and its partition id as the owner */
+    *enc_key_ptr = mbedtls_svc_key_id_make(partition_id, key_id);
+
+    return PSA_SUCCESS;
+}
+
 psa_status_t tfm_crypto_check_key_storage(uint32_t *index)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
@@ -134,7 +155,7 @@
 }
 
 psa_status_t tfm_crypto_set_key_storage(uint32_t index,
-                                        psa_key_handle_t key_handle)
+                                        psa_key_id_t key_handle)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
@@ -148,7 +169,7 @@
     }
 
     handle_owner[index].owner = partition_id;
-    handle_owner[index].handle = key_handle;
+    handle_owner[index].key = key_handle;
     handle_owner[index].in_use = TFM_CRYPTO_IN_USE;
 
     return PSA_SUCCESS;
@@ -194,16 +215,18 @@
 
     if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
         (in_vec[1].len != sizeof(struct psa_client_key_attributes_s)) ||
-        (out_vec[0].len != sizeof(psa_key_handle_t))) {
+        (out_vec[0].len != sizeof(psa_key_id_t))) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
     const struct psa_client_key_attributes_s *client_key_attr = in_vec[1].base;
     const uint8_t *data = in_vec[2].base;
     size_t data_length = in_vec[2].len;
-    psa_key_handle_t *key_handle = out_vec[0].base;
+    psa_key_id_t *psa_key = out_vec[0].base;
+
     psa_status_t status;
     psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
     uint32_t i = 0;
+    mbedtls_svc_key_id_t encoded_key;
     int32_t partition_id = 0;
     bool empty_found = false;
 
@@ -230,11 +253,13 @@
         return status;
     }
 
-    status = psa_import_key(&key_attributes, data, data_length, key_handle);
+    status = psa_import_key(&key_attributes, data, data_length, &encoded_key);
+    /* Update the imported key id */
+    *psa_key = encoded_key.key_id;
 
     if (status == PSA_SUCCESS) {
         handle_owner[i].owner = partition_id;
-        handle_owner[i].handle = *key_handle;
+        handle_owner[i].key = *psa_key;
         handle_owner[i].in_use = TFM_CRYPTO_IN_USE;
     }
 
@@ -254,15 +279,15 @@
     CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 2, out_len, 1, 1);
 
     if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (in_vec[1].len != sizeof(psa_app_key_id_t)) ||
-        (out_vec[0].len != sizeof(psa_key_handle_t))) {
+        (in_vec[1].len != sizeof(psa_key_id_t)) ||
+        (out_vec[0].len != sizeof(psa_key_id_t))) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    psa_app_key_id_t client_key_id = *((psa_app_key_id_t *)in_vec[1].base);
-    psa_key_handle_t *key_handle = out_vec[0].base;
+    psa_key_id_t client_key_id = *((psa_key_id_t *)in_vec[1].base);
+    psa_key_id_t *key = out_vec[0].base;
     psa_status_t status;
-    psa_key_id_t id;
+    mbedtls_svc_key_id_t encoded_key;
     int32_t partition_id;
     uint32_t i;
 
@@ -282,13 +307,14 @@
     }
 
     /* Use the client key id as the key_id and its partition id as the owner */
-    id = (psa_key_id_t){ .key_id = client_key_id, .owner = partition_id };
+    encoded_key = mbedtls_svc_key_id_make(partition_id, client_key_id);
 
-    status = psa_open_key(id, key_handle);
+    status = psa_open_key(encoded_key, &encoded_key);
+    *key = encoded_key.key_id;
 
     if (status == PSA_SUCCESS) {
         handle_owner[i].owner = partition_id;
-        handle_owner[i].handle = *key_handle;
+        handle_owner[i].key = *key;
         handle_owner[i].in_use = TFM_CRYPTO_IN_USE;
     }
 
@@ -313,19 +339,21 @@
     }
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
 
-    psa_key_handle_t key = iov->key_handle;
+    psa_key_id_t key = iov->key_id;
     uint32_t index;
-    psa_status_t status = tfm_crypto_check_handle_owner(key, &index);
+    mbedtls_svc_key_id_t encoded_key;
 
+    psa_status_t status = tfm_crypto_check_handle_owner(key, &index);
     if (status != PSA_SUCCESS) {
         return status;
     }
 
-    status = psa_close_key(key);
+    encoded_key = mbedtls_svc_key_id_make(handle_owner[index].owner, key);
+    status = psa_close_key(encoded_key);
 
     if (status == PSA_SUCCESS) {
         handle_owner[index].owner = 0;
-        handle_owner[index].handle = 0;
+        handle_owner[index].key = 0;
         handle_owner[index].in_use = TFM_CRYPTO_NOT_IN_USE;
     }
 
@@ -349,20 +377,21 @@
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    psa_key_handle_t key = iov->key_handle;
+    psa_key_id_t key = iov->key_id;
     uint32_t index;
-    psa_status_t status = tfm_crypto_check_handle_owner(key, &index);
+    mbedtls_svc_key_id_t encoded_key;
 
+    psa_status_t status = tfm_crypto_check_handle_owner(key, &index);
     if (status != PSA_SUCCESS) {
         return status;
     }
 
-    status = psa_destroy_key(key);
+    encoded_key = mbedtls_svc_key_id_make(handle_owner[index].owner, key);
 
+    status = psa_destroy_key(encoded_key);
     if (status == PSA_SUCCESS) {
         handle_owner[index].owner = 0;
-        handle_owner[index].handle = 0;
+        handle_owner[index].key = 0;
         handle_owner[index].in_use = TFM_CRYPTO_NOT_IN_USE;
     }
 
@@ -387,18 +416,23 @@
     }
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
 
-    psa_key_handle_t key = iov->key_handle;
+    psa_key_id_t key = iov->key_id;
     struct psa_client_key_attributes_s *client_key_attr = out_vec[0].base;
     psa_status_t status;
     psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    mbedtls_svc_key_id_t encoded_key;
 
     status = tfm_crypto_check_handle_owner(key, NULL);
     if (status != PSA_SUCCESS) {
         return status;
     }
 
-    status = psa_get_key_attributes(key, &key_attributes);
+    status = tfm_crypto_encode_id_and_owner(key, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
 
+    status = psa_get_key_attributes(encoded_key, &key_attributes);
     if (status == PSA_SUCCESS) {
         status = tfm_crypto_key_attributes_to_client(&key_attributes,
                                                      client_key_attr);
@@ -464,11 +498,21 @@
     }
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
 
-    psa_key_handle_t key = iov->key_handle;
+    psa_key_id_t key = iov->key_id;
     uint8_t *data = out_vec[0].base;
     size_t data_size = out_vec[0].len;
+    mbedtls_svc_key_id_t encoded_key;
+    uint32_t index;
 
-    return psa_export_key(key, data, data_size, &(out_vec[0].len));
+    psa_status_t status = tfm_crypto_check_handle_owner(key, &index);
+
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    encoded_key = mbedtls_svc_key_id_make(handle_owner[index].owner, key);
+    return psa_export_key(encoded_key, data, data_size,
+                          &(out_vec[0].len));
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
@@ -487,12 +531,61 @@
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    psa_key_handle_t key = iov->key_handle;
+    psa_key_id_t key = iov->key_id;
     uint8_t *data = out_vec[0].base;
     size_t data_size = out_vec[0].len;
+    mbedtls_svc_key_id_t encoded_key;
+    uint32_t index;
 
-    return psa_export_public_key(key, data, data_size, &(out_vec[0].len));
+    psa_status_t status = tfm_crypto_check_handle_owner(key, &index);
+
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    encoded_key = mbedtls_svc_key_id_make(handle_owner[index].owner, key);
+
+    return psa_export_public_key(encoded_key, data, data_size,
+                                 &(out_vec[0].len));
+#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
+}
+
+psa_status_t tfm_crypto_purge_key(psa_invec in_vec[],
+                                  size_t in_len,
+                                  psa_outvec out_vec[],
+                                  size_t out_len)
+{
+#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
+    return PSA_ERROR_NOT_SUPPORTED;
+#else
+    (void)out_vec;
+
+    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 0);
+
+    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
+        return PSA_ERROR_PROGRAMMER_ERROR;
+    }
+    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
+    psa_key_id_t key = iov->key_id;
+    uint32_t index;
+    mbedtls_svc_key_id_t encoded_key;
+
+    psa_status_t status = tfm_crypto_check_handle_owner(key, &index);
+
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    encoded_key = mbedtls_svc_key_id_make(handle_owner[index].owner, key);
+
+    status = psa_purge_key(encoded_key);
+    if (status == PSA_SUCCESS) {
+        handle_owner[index].owner = 0;
+        handle_owner[index].key = 0;
+        handle_owner[index].in_use = TFM_CRYPTO_NOT_IN_USE;
+    }
+
+    return status;
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
@@ -508,20 +601,22 @@
     CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 2, out_len, 1, 1);
 
     if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(psa_key_handle_t)) ||
+        (out_vec[0].len != sizeof(psa_key_id_t)) ||
         (in_vec[1].len != sizeof(struct psa_client_key_attributes_s))) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
 
-    psa_key_handle_t source_handle = iov->key_handle;
-    psa_key_handle_t *target_handle = out_vec[0].base;
+    psa_key_id_t source_key_id = iov->key_id;
+    psa_key_id_t *target_key_id = out_vec[0].base;
     const struct psa_client_key_attributes_s *client_key_attr = in_vec[1].base;
     psa_status_t status;
     psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
     uint32_t i = 0;
     int32_t partition_id = 0;
     bool empty_found = false;
+    mbedtls_svc_key_id_t target_key;
+    mbedtls_svc_key_id_t encoded_key;
 
     for (i = 0; i < TFM_CRYPTO_MAX_KEY_HANDLES; i++) {
         if (handle_owner[i].in_use == TFM_CRYPTO_NOT_IN_USE) {
@@ -546,11 +641,21 @@
         return status;
     }
 
-    status = psa_copy_key(source_handle, &key_attributes, target_handle);
+    status = tfm_crypto_check_handle_owner(source_key_id, NULL);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
 
+    status = tfm_crypto_encode_id_and_owner(source_key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    status = psa_copy_key(encoded_key, &key_attributes, &target_key);
+    *target_key_id = target_key.key_id;
     if (status == PSA_SUCCESS) {
         handle_owner[i].owner = partition_id;
-        handle_owner[i].handle = *target_handle;
+        handle_owner[i].key = *target_key_id;
         handle_owner[i].in_use = TFM_CRYPTO_IN_USE;
     }
 
@@ -571,16 +676,17 @@
 
     if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
         (in_vec[1].len != sizeof(struct psa_client_key_attributes_s)) ||
-        (out_vec[0].len != sizeof(psa_key_handle_t))) {
+        (out_vec[0].len != sizeof(psa_key_id_t))) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
-    psa_key_handle_t *key_handle = out_vec[0].base;
+    psa_key_id_t *key_handle = out_vec[0].base;
     const struct psa_client_key_attributes_s *client_key_attr = in_vec[1].base;
     psa_status_t status;
     psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
     uint32_t i = 0;
     int32_t partition_id = 0;
     bool empty_found = false;
+    mbedtls_svc_key_id_t encoded_key;
 
     for (i = 0; i < TFM_CRYPTO_MAX_KEY_HANDLES; i++) {
         if (handle_owner[i].in_use == TFM_CRYPTO_NOT_IN_USE) {
@@ -605,11 +711,12 @@
         return status;
     }
 
-    status = psa_generate_key(&key_attributes, key_handle);
+    status = psa_generate_key(&key_attributes, &encoded_key);
+    *key_handle = encoded_key.key_id;
 
     if (status == PSA_SUCCESS) {
         handle_owner[i].owner = partition_id;
-        handle_owner[i].handle = *key_handle;
+        handle_owner[i].key = *key_handle;
         handle_owner[i].in_use = TFM_CRYPTO_IN_USE;
     }
 
diff --git a/secure_fw/partitions/crypto/crypto_key_derivation.c b/secure_fw/partitions/crypto/crypto_key_derivation.c
index b2e09be..1d2cb7a 100644
--- a/secure_fw/partitions/crypto/crypto_key_derivation.c
+++ b/secure_fw/partitions/crypto/crypto_key_derivation.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -83,7 +83,7 @@
 static psa_status_t tfm_crypto_huk_derivation_output_key(
                                       const psa_key_attributes_t *attributes,
                                       psa_key_derivation_operation_t *operation,
-                                      psa_key_handle_t *handle)
+                                      mbedtls_svc_key_id_t *key_id)
 {
     enum tfm_plat_err_t err;
     size_t bytes = PSA_BITS_TO_BYTES(psa_get_key_bits(attributes));
@@ -103,7 +103,7 @@
     }
 
     return psa_import_key(attributes, operation->ctx.tls12_prf.output_block,
-                          bytes, handle);
+                          bytes, key_id);
 }
 
 static psa_status_t tfm_crypto_huk_derivation_abort(
@@ -333,13 +333,13 @@
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
     uint32_t handle = iov->op_handle;
-    psa_key_handle_t key_handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_key_derivation_step_t step = iov->step;
     psa_key_derivation_operation_t *operation = NULL;
+    mbedtls_svc_key_id_t encoded_key;
 
-    status = tfm_crypto_check_handle_owner(key_handle, NULL);
+    status = tfm_crypto_check_handle_owner(key_id, NULL);
     if (status != PSA_SUCCESS) {
         return status;
     }
@@ -352,7 +352,12 @@
         return status;
     }
 
-    return psa_key_derivation_input_key(operation, step, key_handle);
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    return psa_key_derivation_input_key(operation, step, encoded_key);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 
@@ -370,7 +375,7 @@
 
     if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
         (in_vec[1].len != sizeof(struct psa_client_key_attributes_s)) ||
-        (out_vec[0].len != sizeof(psa_key_handle_t))) {
+        (out_vec[0].len != sizeof(psa_key_id_t))) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
@@ -378,10 +383,11 @@
     uint32_t handle = iov->op_handle;
     const struct psa_client_key_attributes_s *client_key_attr = in_vec[1].base;
     psa_key_derivation_operation_t *operation = NULL;
-    psa_key_handle_t *key_handle = out_vec[0].base;
+    psa_key_id_t *key_handle = out_vec[0].base;
     psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
     int32_t partition_id;
     uint32_t index;
+    mbedtls_svc_key_id_t encoded_key;
 
     /* Look up the corresponding operation context */
     status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
@@ -410,11 +416,13 @@
 
     if (operation->alg == TFM_CRYPTO_ALG_HUK_DERIVATION) {
         status = tfm_crypto_huk_derivation_output_key(&key_attributes,
-                                                      operation, key_handle);
+                                                      operation, &encoded_key);
     } else {
         status = psa_key_derivation_output_key(&key_attributes, operation,
-                                               key_handle);
+                                               &encoded_key);
     }
+    *key_handle = encoded_key.key_id;
+
     if (status == PSA_SUCCESS) {
         status = tfm_crypto_set_key_storage(index, *key_handle);
     }
@@ -494,11 +502,12 @@
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
 
     uint32_t handle = iov->op_handle;
-    psa_key_handle_t private_key = iov->key_handle;
+    psa_key_id_t private_key = iov->key_id;
     const uint8_t *peer_key = in_vec[1].base;
     size_t peer_key_length = in_vec[1].len;
     psa_key_derivation_operation_t *operation = NULL;
     psa_key_derivation_step_t step = iov->step;
+    mbedtls_svc_key_id_t encoded_key;
 
     status = tfm_crypto_check_handle_owner(private_key, NULL);
     if (status != PSA_SUCCESS) {
@@ -513,8 +522,13 @@
         return status;
     }
 
+    status = tfm_crypto_encode_id_and_owner(private_key, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
     return psa_key_derivation_key_agreement(operation, step,
-                                            private_key,
+                                            encoded_key,
                                             peer_key,
                                             peer_key_length);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
@@ -559,11 +573,23 @@
     uint8_t *output = out_vec[0].base;
     size_t output_size = out_vec[0].len;
     psa_algorithm_t alg = iov->alg;
-    psa_key_handle_t private_key = iov->key_handle;
+    psa_key_id_t private_key = iov->key_id;
     const uint8_t *peer_key = in_vec[1].base;
     size_t peer_key_length = in_vec[1].len;
+    mbedtls_svc_key_id_t encoded_key;
 
-    return psa_raw_key_agreement(alg, private_key, peer_key, peer_key_length,
+    psa_status_t status = tfm_crypto_check_handle_owner(private_key, NULL);
+
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    status = tfm_crypto_encode_id_and_owner(private_key, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    return psa_raw_key_agreement(alg, encoded_key, peer_key, peer_key_length,
                                  output, output_size, &out_vec[0].len);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
diff --git a/secure_fw/partitions/crypto/crypto_mac.c b/secure_fw/partitions/crypto/crypto_mac.c
index c983b6a..e2f27c5 100644
--- a/secure_fw/partitions/crypto/crypto_mac.c
+++ b/secure_fw/partitions/crypto/crypto_mac.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -40,10 +40,11 @@
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
     uint32_t handle = iov->op_handle;
     uint32_t *handle_out = out_vec[0].base;
-    psa_key_handle_t key_handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
+    mbedtls_svc_key_id_t encoded_key;
 
-    status = tfm_crypto_check_handle_owner(key_handle, NULL);
+    status = tfm_crypto_check_handle_owner(key_id, NULL);
     if (status != PSA_SUCCESS) {
         return status;
     }
@@ -61,7 +62,12 @@
 
     *handle_out = handle;
 
-    status = psa_mac_sign_setup(operation, key_handle, alg);
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    status = psa_mac_sign_setup(operation, encoded_key, alg);
     if (status != PSA_SUCCESS) {
         /* Release the operation context, ignore if the operation fails. */
         (void)tfm_crypto_operation_release(handle_out);
@@ -92,10 +98,11 @@
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
     uint32_t handle = iov->op_handle;
     uint32_t *handle_out = out_vec[0].base;
-    psa_key_handle_t key_handle = iov->key_handle;
+    psa_key_id_t key_id = iov->key_id;
     psa_algorithm_t alg = iov->alg;
+    mbedtls_svc_key_id_t encoded_key;
 
-    status = tfm_crypto_check_handle_owner(key_handle, NULL);
+    status = tfm_crypto_check_handle_owner(key_id, NULL);
     if (status != PSA_SUCCESS) {
         return status;
     }
@@ -113,7 +120,12 @@
 
     *handle_out = handle;
 
-    status = psa_mac_verify_setup(operation, key_handle, alg);
+    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    status = psa_mac_verify_setup(operation, encoded_key, alg);
     if (status != PSA_SUCCESS) {
         /* Release the operation context, ignore if the operation fails. */
         (void)tfm_crypto_operation_release(handle_out);
diff --git a/secure_fw/partitions/crypto/crypto_spe.h b/secure_fw/partitions/crypto/crypto_spe.h
index 2ceaefd..507eda2 100644
--- a/secure_fw/partitions/crypto/crypto_spe.h
+++ b/secure_fw/partitions/crypto/crypto_spe.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -66,6 +66,8 @@
         PSA_FUNCTION_NAME(psa_export_key)
 #define psa_export_public_key \
         PSA_FUNCTION_NAME(psa_export_public_key)
+#define psa_purge_key \
+        PSA_FUNCTION_NAME(psa_purge_key)
 #define psa_copy_key \
         PSA_FUNCTION_NAME(psa_copy_key)
 #define psa_cipher_operation_init \
diff --git a/secure_fw/partitions/crypto/tfm_crypto.yaml b/secure_fw/partitions/crypto/tfm_crypto.yaml
index cc836c5..c97bf94 100644
--- a/secure_fw/partitions/crypto/tfm_crypto.yaml
+++ b/secure_fw/partitions/crypto/tfm_crypto.yaml
@@ -1,5 +1,5 @@
 #-------------------------------------------------------------------------------
-# Copyright (c) 2018-2020, Arm Limited. All rights reserved.
+# Copyright (c) 2018-2021, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -70,6 +70,13 @@
       "version_policy": "STRICT"
     },
     {
+      "name": "TFM_CRYPTO_PURGE_KEY",
+      "signal": "TFM_CRYPTO_PURGE_KEY",
+      "non_secure_clients": true,
+      "version": 1,
+      "version_policy": "STRICT"
+    },
+    {
       "name": "TFM_CRYPTO_COPY_KEY",
       "signal": "TFM_CRYPTO_COPY_KEY",
       "non_secure_clients": true,
diff --git a/secure_fw/partitions/crypto/tfm_crypto_api.h b/secure_fw/partitions/crypto/tfm_crypto_api.h
index fa6f136..3f3d9bf 100644
--- a/secure_fw/partitions/crypto/tfm_crypto_api.h
+++ b/secure_fw/partitions/crypto/tfm_crypto_api.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -79,9 +79,9 @@
  * \return Return values as described in \ref psa_status_t
  */
 psa_status_t tfm_crypto_key_attributes_from_client(
-                             const struct psa_client_key_attributes_s *client_key_attr,
-                             int32_t client_id,
-                             psa_key_attributes_t *key_attributes);
+                    const struct psa_client_key_attributes_s *client_key_attr,
+                    int32_t client_id,
+                    psa_key_attributes_t *key_attributes);
 
 /**
  * \brief Converts key attributes to client key attributes.
@@ -92,14 +92,14 @@
  * \return Return values as described in \ref psa_status_t
  */
 psa_status_t tfm_crypto_key_attributes_to_client(
-                                  const psa_key_attributes_t *key_attributes,
-                                  struct psa_client_key_attributes_s *client_key_attr);
+                        const psa_key_attributes_t *key_attributes,
+                        struct psa_client_key_attributes_s *client_key_attr);
 
 /**
  * \brief Checks that the requested handle belongs to the requesting
  *        partition
  *
- * \param[in]  handle Handle given as input
+ * \param[in]  key    key given as input
  * \param[out] index  Optionally, pointer to hold the internal index
  *                    corresponding to the input handle. Valid only
  *                    on PSA_SUCCESS, it's returned only if the input
@@ -107,7 +107,7 @@
  *
  * \return Return values as described in \ref psa_status_t
  */
-psa_status_t tfm_crypto_check_handle_owner(psa_key_handle_t handle,
+psa_status_t tfm_crypto_check_handle_owner(psa_key_id_t key,
                                            uint32_t *index);
 
 /**
@@ -130,7 +130,7 @@
  * \return Return values as described in \ref psa_status_t
  */
 psa_status_t tfm_crypto_set_key_storage(uint32_t index,
-                                        psa_key_handle_t key_handle);
+                                        psa_key_id_t key_handle);
 /**
  * \brief Allocate an operation context in the backend
  *
@@ -164,6 +164,16 @@
 psa_status_t tfm_crypto_operation_lookup(enum tfm_crypto_operation_type type,
                                          uint32_t handle,
                                          void **ctx);
+/**
+ * \brief Encodes the input key id and owner to output key
+ *
+ * \param[in]  key_id       Id of the key to encode
+ * \param[out] enc_key_ptr  Pointer to encoded key with id and owner
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_encode_id_and_owner(psa_key_id_t key_id,
+                                            mbedtls_svc_key_id_t *enc_key_ptr);
 
 #define LIST_TFM_CRYPTO_UNIFORM_SIGNATURE_API \
     X(tfm_crypto_get_key_attributes)          \
@@ -174,6 +184,7 @@
     X(tfm_crypto_destroy_key)                 \
     X(tfm_crypto_export_key)                  \
     X(tfm_crypto_export_public_key)           \
+    X(tfm_crypto_purge_key)                   \
     X(tfm_crypto_copy_key)                    \
     X(tfm_crypto_hash_compute)                \
     X(tfm_crypto_hash_compare)                \
diff --git a/secure_fw/partitions/crypto/tfm_crypto_secure_api.c b/secure_fw/partitions/crypto/tfm_crypto_secure_api.c
index 1db9be2..d7ae137 100644
--- a/secure_fw/partitions/crypto/tfm_crypto_secure_api.c
+++ b/secure_fw/partitions/crypto/tfm_crypto_secure_api.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -56,7 +56,7 @@
 }
 
 psa_status_t psa_open_key(psa_key_id_t id,
-                          psa_key_handle_t *handle)
+                          psa_key_id_t *key_id)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
@@ -70,7 +70,7 @@
         {.base = &id, .len = sizeof(psa_key_id_t)},
     };
     psa_outvec out_vec[] = {
-        {.base = handle, .len = sizeof(psa_key_handle_t)},
+        {.base = key_id, .len = sizeof(psa_key_id_t)},
     };
 
 #ifdef TFM_PSA_API
@@ -88,7 +88,7 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_close_key(psa_key_handle_t handle)
+psa_status_t psa_close_key(psa_key_id_t key_id)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
@@ -96,7 +96,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_CLOSE_KEY_SID,
-        .key_handle = handle,
+        .key_id = key_id,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -120,7 +120,7 @@
 psa_status_t psa_import_key(const psa_key_attributes_t *attributes,
                             const uint8_t *data,
                             size_t data_length,
-                            psa_key_handle_t *handle)
+                            psa_key_id_t *key_id)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
@@ -135,7 +135,7 @@
         {.base = data, .len = data_length}
     };
     psa_outvec out_vec[] = {
-        {.base = handle, .len = sizeof(psa_key_handle_t)}
+        {.base = key_id, .len = sizeof(psa_key_id_t)}
     };
 
 #ifdef TFM_PSA_API
@@ -152,7 +152,7 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_destroy_key(psa_key_handle_t handle)
+psa_status_t psa_destroy_key(psa_key_id_t key_id)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
@@ -160,7 +160,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_DESTROY_KEY_SID,
-        .key_handle = handle,
+        .key_id = key_id,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -180,7 +180,7 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_get_key_attributes(psa_key_handle_t handle,
+psa_status_t psa_get_key_attributes(psa_key_id_t key_id,
                                     psa_key_attributes_t *attributes)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
@@ -189,7 +189,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID,
-        .key_handle = handle,
+        .key_id = key_id,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -245,7 +245,7 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_export_key(psa_key_handle_t handle,
+psa_status_t psa_export_key(psa_key_id_t key_id,
                             uint8_t *data,
                             size_t data_size,
                             size_t *data_length)
@@ -256,7 +256,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_EXPORT_KEY_SID,
-        .key_handle = handle,
+        .key_id = key_id,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -282,7 +282,7 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_export_public_key(psa_key_handle_t handle,
+psa_status_t psa_export_public_key(psa_key_id_t key_id,
                                    uint8_t *data,
                                    size_t data_size,
                                    size_t *data_length)
@@ -293,7 +293,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID,
-        .key_handle = handle,
+        .key_id = key_id,
     };
 
     psa_invec in_vec[] = {
@@ -320,9 +320,37 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_copy_key(psa_key_handle_t source_handle,
+psa_status_t psa_purge_key(psa_key_id_t key_id)
+{
+#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
+    return PSA_ERROR_NOT_SUPPORTED;
+#else
+    psa_status_t status;
+    struct tfm_crypto_pack_iovec iov = {
+        .sfn_id = TFM_CRYPTO_PURGE_KEY_SID,
+        .key_id = key_id,
+    };
+    psa_invec in_vec[] = {
+        {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
+    };
+
+#ifdef TFM_PSA_API
+    PSA_CONNECT(TFM_CRYPTO);
+#endif
+
+    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_purge_key,
+                                    TFM_CRYPTO_PURGE_KEY);
+#ifdef TFM_PSA_API
+    PSA_CLOSE();
+#endif
+
+    return status;
+#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
+}
+
+psa_status_t psa_copy_key(psa_key_id_t source_key_id,
                           const psa_key_attributes_t *attributes,
-                          psa_key_handle_t *target_handle)
+                          psa_key_id_t *target_key_id)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
@@ -330,7 +358,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_COPY_KEY_SID,
-        .key_handle = source_handle,
+        .key_id = source_key_id,
     };
 
     psa_invec in_vec[] = {
@@ -339,7 +367,7 @@
     };
 
     psa_outvec out_vec[] = {
-        {.base = target_handle, .len = sizeof(psa_key_handle_t)},
+        {.base = target_key_id, .len = sizeof(psa_key_id_t)},
     };
 
 #ifdef TFM_PSA_API
@@ -431,7 +459,7 @@
 }
 
 psa_status_t psa_cipher_encrypt_setup(psa_cipher_operation_t *operation,
-                                      psa_key_handle_t handle,
+                                      psa_key_id_t key_id,
                                       psa_algorithm_t alg)
 {
 #ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
@@ -440,7 +468,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -467,7 +495,7 @@
 }
 
 psa_status_t psa_cipher_decrypt_setup(psa_cipher_operation_t *operation,
-                                      psa_key_handle_t handle,
+                                      psa_key_id_t key_id,
                                       psa_algorithm_t alg)
 {
 #ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
@@ -476,7 +504,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -908,7 +936,7 @@
 }
 
 psa_status_t psa_mac_sign_setup(psa_mac_operation_t *operation,
-                                psa_key_handle_t handle,
+                                psa_key_id_t key_id,
                                 psa_algorithm_t alg)
 {
 #ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
@@ -917,7 +945,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_MAC_SIGN_SETUP_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -944,7 +972,7 @@
 }
 
 psa_status_t psa_mac_verify_setup(psa_mac_operation_t *operation,
-                                  psa_key_handle_t handle,
+                                  psa_key_id_t key_id,
                                   psa_algorithm_t alg)
 {
 #ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
@@ -953,7 +981,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_MAC_VERIFY_SETUP_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -1121,7 +1149,7 @@
 #endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
 }
 
-psa_status_t psa_aead_encrypt(psa_key_handle_t handle,
+psa_status_t psa_aead_encrypt(psa_key_id_t key_id,
                               psa_algorithm_t alg,
                               const uint8_t *nonce,
                               size_t nonce_length,
@@ -1139,7 +1167,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_AEAD_ENCRYPT_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg,
         .aead_in = {.nonce = {0}, .nonce_length = nonce_length}
     };
@@ -1195,7 +1223,7 @@
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
 
-psa_status_t psa_aead_decrypt(psa_key_handle_t handle,
+psa_status_t psa_aead_decrypt(psa_key_id_t key_id,
                               psa_algorithm_t alg,
                               const uint8_t *nonce,
                               size_t nonce_length,
@@ -1213,7 +1241,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_AEAD_DECRYPT_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg,
         .aead_in = {.nonce = {0}, .nonce_length = nonce_length}
     };
@@ -1269,7 +1297,7 @@
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
 
-psa_status_t psa_asymmetric_sign(psa_key_handle_t handle,
+psa_status_t psa_asymmetric_sign(psa_key_id_t key_id,
                                  psa_algorithm_t alg,
                                  const uint8_t *hash,
                                  size_t hash_length,
@@ -1277,10 +1305,11 @@
                                  size_t signature_size,
                                  size_t *signature_length)
 {
-    return psa_sign_hash(handle, alg, hash, hash_length, signature, signature_size, signature_length);
+    return psa_sign_hash(key_id, alg, hash, hash_length, signature,
+                         signature_size, signature_length);
 }
 
-psa_status_t psa_sign_hash(psa_key_handle_t handle,
+psa_status_t psa_sign_hash(psa_key_id_t key_id,
                            psa_algorithm_t alg,
                            const uint8_t *hash,
                            size_t hash_length,
@@ -1294,7 +1323,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_SIGN_HASH_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg,
     };
 
@@ -1322,17 +1351,18 @@
 #endif /* TFM_CRYPTO_ASYMMETRIC_MODULE_DISABLED */
 }
 
-psa_status_t psa_asymmetric_verify(psa_key_handle_t handle,
+psa_status_t psa_asymmetric_verify(psa_key_id_t key_id,
                                    psa_algorithm_t alg,
                                    const uint8_t *hash,
                                    size_t hash_length,
                                    const uint8_t *signature,
                                    size_t signature_length)
 {
-    return psa_verify_hash(handle, alg, hash, hash_length, signature, signature_length);
+    return psa_verify_hash(key_id, alg, hash, hash_length,
+                           signature, signature_length);
 }
 
-psa_status_t psa_verify_hash(psa_key_handle_t handle,
+psa_status_t psa_verify_hash(psa_key_id_t key_id,
                              psa_algorithm_t alg,
                              const uint8_t *hash,
                              size_t hash_length,
@@ -1345,7 +1375,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_VERIFY_HASH_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg
     };
 
@@ -1369,7 +1399,7 @@
 #endif /* TFM_CRYPTO_ASYMMETRIC_MODULE_DISABLED */
 }
 
-psa_status_t psa_asymmetric_encrypt(psa_key_handle_t handle,
+psa_status_t psa_asymmetric_encrypt(psa_key_id_t key_id,
                                     psa_algorithm_t alg,
                                     const uint8_t *input,
                                     size_t input_length,
@@ -1385,7 +1415,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg
     };
 
@@ -1430,7 +1460,7 @@
 #endif /* TFM_CRYPTO_ASYMMETRIC_MODULE_DISABLED */
 }
 
-psa_status_t psa_asymmetric_decrypt(psa_key_handle_t handle,
+psa_status_t psa_asymmetric_decrypt(psa_key_id_t key_id,
                                     psa_algorithm_t alg,
                                     const uint8_t *input,
                                     size_t input_length,
@@ -1446,7 +1476,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .alg = alg
     };
 
@@ -1565,7 +1595,7 @@
 psa_status_t psa_key_derivation_input_key(
                                       psa_key_derivation_operation_t *operation,
                                       psa_key_derivation_step_t step,
-                                      psa_key_handle_t handle)
+                                      psa_key_id_t key_id)
 {
 #ifdef TFM_CRYPTO_GENERATOR_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
@@ -1573,7 +1603,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID,
-        .key_handle = handle,
+        .key_id = key_id,
         .step = step,
         .op_handle = operation->handle,
     };
@@ -1632,7 +1662,7 @@
 psa_status_t psa_key_derivation_key_agreement(
                                       psa_key_derivation_operation_t *operation,
                                       psa_key_derivation_step_t step,
-                                      psa_key_handle_t private_key,
+                                      psa_key_id_t private_key,
                                       const uint8_t *peer_key,
                                       size_t peer_key_length)
 {
@@ -1642,7 +1672,7 @@
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID,
-        .key_handle = private_key,
+        .key_id = private_key,
         .step = step,
         .op_handle = operation->handle,
     };
@@ -1710,7 +1740,7 @@
 }
 
 psa_status_t psa_generate_key(const psa_key_attributes_t *attributes,
-                              psa_key_handle_t *handle)
+                              psa_key_id_t *key_id)
 {
 #ifdef TFM_CRYPTO_GENERATOR_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
@@ -1726,7 +1756,7 @@
     };
 
     psa_outvec out_vec[] = {
-        {.base = handle, .len = sizeof(psa_key_handle_t)},
+        {.base = key_id, .len = sizeof(psa_key_id_t)},
     };
 
 #ifdef TFM_PSA_API
@@ -1817,7 +1847,7 @@
     return status;
 }
 
-psa_status_t psa_mac_compute(psa_key_handle_t handle,
+psa_status_t psa_mac_compute(psa_key_id_t key_id,
                              psa_algorithm_t alg,
                              const uint8_t *input,
                              size_t input_length,
@@ -1832,7 +1862,7 @@
     return status;
 }
 
-psa_status_t psa_mac_verify(psa_key_handle_t handle,
+psa_status_t psa_mac_verify(psa_key_id_t key_id,
                             psa_algorithm_t alg,
                             const uint8_t *input,
                             size_t input_length,
@@ -1846,7 +1876,7 @@
     return status;
 }
 
-psa_status_t psa_cipher_encrypt(psa_key_handle_t handle,
+psa_status_t psa_cipher_encrypt(psa_key_id_t key_id,
                                 psa_algorithm_t alg,
                                 const uint8_t *input,
                                 size_t input_length,
@@ -1861,7 +1891,7 @@
     return status;
 }
 
-psa_status_t psa_cipher_decrypt(psa_key_handle_t handle,
+psa_status_t psa_cipher_decrypt(psa_key_id_t key_id,
                                 psa_algorithm_t alg,
                                 const uint8_t *input,
                                 size_t input_length,
@@ -1877,7 +1907,7 @@
 }
 
 psa_status_t psa_raw_key_agreement(psa_algorithm_t alg,
-                                   psa_key_handle_t private_key,
+                                   psa_key_id_t private_key,
                                    const uint8_t *peer_key,
                                    size_t peer_key_length,
                                    uint8_t *output,
@@ -1891,7 +1921,7 @@
     struct tfm_crypto_pack_iovec iov = {
         .sfn_id = TFM_CRYPTO_RAW_KEY_AGREEMENT_SID,
         .alg = alg,
-        .key_handle = private_key
+        .key_id = private_key
     };
 
     psa_invec in_vec[] = {
@@ -2023,7 +2053,7 @@
 psa_status_t psa_key_derivation_output_key(
                                       const psa_key_attributes_t *attributes,
                                       psa_key_derivation_operation_t *operation,
-                                      psa_key_handle_t *handle)
+                                      psa_key_id_t *key_id)
 {
 #if (TFM_CRYPTO_GENERATOR_MODULE_DISABLED != 0)
     return PSA_ERROR_NOT_SUPPORTED;
@@ -2040,7 +2070,7 @@
     };
 
     psa_outvec out_vec[] = {
-        {.base = handle, .len = sizeof(psa_key_handle_t)}
+        {.base = key_id, .len = sizeof(psa_key_id_t)}
     };
 
 #ifdef TFM_PSA_API
@@ -2058,7 +2088,7 @@
 }
 
 psa_status_t psa_aead_encrypt_setup(psa_aead_operation_t *operation,
-                                    psa_key_handle_t handle,
+                                    psa_key_id_t key_id,
                                     psa_algorithm_t alg)
 {
     psa_status_t status;
@@ -2069,7 +2099,7 @@
 }
 
 psa_status_t psa_aead_decrypt_setup(psa_aead_operation_t *operation,
-                                    psa_key_handle_t handle,
+                                    psa_key_id_t key_id,
                                     psa_algorithm_t alg)
 {
     psa_status_t status;
diff --git a/secure_fw/partitions/crypto/tfm_mbedcrypto_include.h b/secure_fw/partitions/crypto/tfm_mbedcrypto_include.h
index 7bb7f46..e2b1423 100644
--- a/secure_fw/partitions/crypto/tfm_mbedcrypto_include.h
+++ b/secure_fw/partitions/crypto/tfm_mbedcrypto_include.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -10,6 +10,7 @@
 
 /* Define PSA_CRYPTO_SECURE to signal that we are compiling for the SPE */
 #define PSA_CRYPTO_SECURE 1
+
 /* Include the crypto_spe.h header before including the PSA Crypto header from
  * Mbed Crypto
  */
diff --git a/secure_fw/partitions/protected_storage/crypto/ps_crypto_interface.c b/secure_fw/partitions/protected_storage/crypto/ps_crypto_interface.c
index ecfa566..e4b66ff 100644
--- a/secure_fw/partitions/protected_storage/crypto/ps_crypto_interface.c
+++ b/secure_fw/partitions/protected_storage/crypto/ps_crypto_interface.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2017-2021, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -36,7 +36,7 @@
 typedef char PS_ERROR_NOT_AEAD_ALG[(PSA_ALG_IS_AEAD(PS_CRYPTO_ALG)) ? 1 : -1];
 
 static const uint8_t ps_key_label[] = "storage_key";
-static psa_key_handle_t ps_key_handle;
+static psa_key_id_t ps_key;
 static uint8_t ps_crypto_iv_buf[PS_IV_LEN_BYTES];
 
 psa_status_t ps_crypto_init(void)
@@ -74,7 +74,7 @@
     }
 
     /* Create the storage key from the key derivation operation */
-    status = psa_key_derivation_output_key(&attributes, &op, &ps_key_handle);
+    status = psa_key_derivation_output_key(&attributes, &op, &ps_key);
     if (status != PSA_SUCCESS) {
         goto err_release_op;
     }
@@ -88,7 +88,7 @@
     return PSA_SUCCESS;
 
 err_release_key:
-    (void)psa_destroy_key(ps_key_handle);
+    (void)psa_destroy_key(ps_key);
 
 err_release_op:
     (void)psa_key_derivation_abort(&op);
@@ -101,7 +101,7 @@
     psa_status_t status;
 
     /* Destroy the transient key */
-    status = psa_destroy_key(ps_key_handle);
+    status = psa_destroy_key(ps_key);
     if (status != PSA_SUCCESS) {
         return PSA_ERROR_GENERIC_ERROR;
     }
@@ -167,7 +167,7 @@
 {
     psa_status_t status;
 
-    status = psa_aead_encrypt(ps_key_handle, PS_CRYPTO_ALG,
+    status = psa_aead_encrypt(ps_key, PS_CRYPTO_ALG,
                               crypto->ref.iv, PS_IV_LEN_BYTES,
                               add, add_len,
                               in, in_len,
@@ -198,7 +198,7 @@
     (void)tfm_memcpy((in + in_len), crypto->ref.tag, PS_TAG_LEN_BYTES);
     in_len += PS_TAG_LEN_BYTES;
 
-    status = psa_aead_decrypt(ps_key_handle, PS_CRYPTO_ALG,
+    status = psa_aead_decrypt(ps_key, PS_CRYPTO_ALG,
                               crypto->ref.iv, PS_IV_LEN_BYTES,
                               add, add_len,
                               in, in_len,
@@ -217,7 +217,7 @@
     psa_status_t status;
     size_t out_len;
 
-    status = psa_aead_encrypt(ps_key_handle, PS_CRYPTO_ALG,
+    status = psa_aead_encrypt(ps_key, PS_CRYPTO_ALG,
                               crypto->ref.iv, PS_IV_LEN_BYTES,
                               add, add_len,
                               0, 0,
@@ -236,7 +236,7 @@
     psa_status_t status;
     size_t out_len;
 
-    status = psa_aead_decrypt(ps_key_handle, PS_CRYPTO_ALG,
+    status = psa_aead_decrypt(ps_key, PS_CRYPTO_ALG,
                               crypto->ref.iv, PS_IV_LEN_BYTES,
                               add, add_len,
                               crypto->ref.tag, PS_TAG_LEN_BYTES,