From 3af9bc18f3be7a3e3cc52aa5f08dd4c7b92ddcb6 Mon Sep 17 00:00:00 2001
From: Ryan Everett <ryan.everett@arm.com>
Date: Tue, 30 Jan 2024 17:21:57 +0000
Subject: [PATCH] Wrap get_and_lock_key_slot_in_memory calls in mutex

It is useful to do this for the call in get_and_lock_key_slot.
Documenting that get_and_lock_key_slot_in_memory requires the mutex
is not part of this PR

Signed-off-by: Ryan Everett <ryan.everett@arm.com>
---
 library/psa_crypto_slot_management.c | 31 +++++++++++++++++-----------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 599cc363b..f4c6ee005 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -527,26 +527,29 @@ psa_status_t psa_open_key(mbedtls_svc_key_id_t key, psa_key_handle_t *handle)
 
 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)
+    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 defined(MBEDTLS_THREADING_C)
-    PSA_THREADING_CHK_RET(mbedtls_mutex_lock(
-                              &mbedtls_threading_key_slot_mutex));
-#endif
     if (slot->registered_readers == 1) {
         status = psa_wipe_key_slot(slot);
     } else {
@@ -562,18 +565,22 @@ psa_status_t psa_close_key(psa_key_handle_t handle)
 
 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;
 
-    status = psa_get_and_lock_key_slot_in_memory(key, &slot);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
 #if defined(MBEDTLS_THREADING_C)
     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)) {
         status = psa_wipe_key_slot(slot);