Merge pull request #8764 from Ryan-Everett-arm/threadsafe-key-wiping

Make key destruction thread safe
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index c67345b..10a23f6 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -198,6 +198,8 @@
  *
  * This function clears all data associated with the PSA layer,
  * including the whole key store.
+ * This function is not thread safe, it wipes every key slot regardless of
+ * state and reader count. It should only be called when no slot is in use.
  *
  * This is an Mbed TLS extension.
  */
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index beafc48..4a0666b 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1089,6 +1089,14 @@
         return status;
     }
 
+#if defined(MBEDTLS_THREADING_C)
+    /* We cannot unlock between setting the state to PENDING_DELETION
+     * and destroying the key in storage, as otherwise another thread
+     * could load the key into a new slot and the key will not be
+     * fully destroyed. */
+    PSA_THREADING_CHK_GOTO_EXIT(mbedtls_mutex_lock(
+                                    &mbedtls_threading_key_slot_mutex));
+#endif
     /* Set the key slot containing the key description's state to
      * PENDING_DELETION. This stops new operations from registering
      * to read the slot. Current readers can safely continue to access
@@ -1097,7 +1105,12 @@
      * If the key is persistent, we can now delete the copy of the key
      * from memory. If the key is opaque, we require the driver to
      * deal with the deletion. */
-    slot->state = PSA_SLOT_PENDING_DELETION;
+    status = psa_key_slot_state_transition(slot, PSA_SLOT_FULL,
+                                           PSA_SLOT_PENDING_DELETION);
+
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
 
     if (PSA_KEY_LIFETIME_IS_READ_ONLY(slot->attr.lifetime)) {
         /* Refuse the destruction of a read-only key (which may or may not work
@@ -1152,11 +1165,6 @@
         if (overall_status == PSA_SUCCESS) {
             overall_status = status;
         }
-
-        /* TODO: other slots may have a copy of the same key. We should
-         * invalidate them.
-         * https://github.com/ARMmbed/mbed-crypto/issues/214
-         */
     }
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 
@@ -1182,6 +1190,14 @@
     if (status != PSA_SUCCESS) {
         overall_status = status;
     }
+
+#if defined(MBEDTLS_THREADING_C)
+    /* Don't overwrite existing errors if the unlock fails. */
+    status = overall_status;
+    PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
+                              &mbedtls_threading_key_slot_mutex));
+#endif
+
     return overall_status;
 }
 
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 47ace35..dc38662 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -521,44 +521,78 @@
 
 psa_status_t psa_close_key(psa_key_handle_t handle)
 {
-    psa_status_t status;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
 
     if (psa_key_handle_is_null(handle)) {
         return PSA_SUCCESS;
     }
 
+#if defined(MBEDTLS_THREADING_C)
+    /* We need to set status as success, otherwise CORRUPTION_DETECTED
+     * would be returned if the lock fails. */
+    status = PSA_SUCCESS;
+    PSA_THREADING_CHK_RET(mbedtls_mutex_lock(
+                              &mbedtls_threading_key_slot_mutex));
+#endif
     status = psa_get_and_lock_key_slot_in_memory(handle, &slot);
     if (status != PSA_SUCCESS) {
         if (status == PSA_ERROR_DOES_NOT_EXIST) {
             status = PSA_ERROR_INVALID_HANDLE;
         }
-
+#if defined(MBEDTLS_THREADING_C)
+        PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
+                                  &mbedtls_threading_key_slot_mutex));
+#endif
         return status;
     }
+
     if (slot->registered_readers == 1) {
-        return psa_wipe_key_slot(slot);
+        status = psa_wipe_key_slot(slot);
     } else {
-        return psa_unregister_read(slot);
+        status = psa_unregister_read(slot);
     }
+#if defined(MBEDTLS_THREADING_C)
+    PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
+                              &mbedtls_threading_key_slot_mutex));
+#endif
+
+    return status;
 }
 
 psa_status_t psa_purge_key(mbedtls_svc_key_id_t key)
 {
-    psa_status_t status;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
 
+#if defined(MBEDTLS_THREADING_C)
+    /* We need to set status as success, otherwise CORRUPTION_DETECTED
+     * would be returned if the lock fails. */
+    status = PSA_SUCCESS;
+    PSA_THREADING_CHK_RET(mbedtls_mutex_lock(
+                              &mbedtls_threading_key_slot_mutex));
+#endif
     status = psa_get_and_lock_key_slot_in_memory(key, &slot);
     if (status != PSA_SUCCESS) {
+#if defined(MBEDTLS_THREADING_C)
+        PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
+                                  &mbedtls_threading_key_slot_mutex));
+#endif
         return status;
     }
 
     if ((!PSA_KEY_LIFETIME_IS_VOLATILE(slot->attr.lifetime)) &&
         (slot->registered_readers == 1)) {
-        return psa_wipe_key_slot(slot);
+        status = psa_wipe_key_slot(slot);
     } else {
-        return psa_unregister_read(slot);
+        status = psa_unregister_read(slot);
     }
+#if defined(MBEDTLS_THREADING_C)
+    PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
+                              &mbedtls_threading_key_slot_mutex));
+#endif
+
+    return status;
 }
 
 void mbedtls_psa_get_stats(mbedtls_psa_stats_t *stats)
diff --git a/library/psa_crypto_slot_management.h b/library/psa_crypto_slot_management.h
index 002429b..18a9144 100644
--- a/library/psa_crypto_slot_management.h
+++ b/library/psa_crypto_slot_management.h
@@ -92,6 +92,8 @@
 psa_status_t psa_initialize_key_slots(void);
 
 /** Delete all data from key slots in memory.
+ * This function is not thread safe, it wipes every key slot regardless of
+ * state and reader count. It should only be called when no slot is in use.
  *
  * This does not affect persistent storage. */
 void psa_wipe_all_key_slots(void);