From 3d8118d9dcae661fe2cc7d958d1a6ec8ee444c5c Mon Sep 17 00:00:00 2001
From: Ryan Everett <ryan.everett@arm.com>
Date: Tue, 30 Jan 2024 16:58:47 +0000
Subject: [PATCH] Revert psa_reserve_free_key_slot changes, lock in
 start_key_creation instead

This means we can hold the mutex around the call to reserve_free_key_slot
in get_and_lock_key_slot, avoiding inefficient rework.
(Changes to get_and_lock_key_slot are not in scope in this PR)

Signed-off-by: Ryan Everett <ryan.everett@arm.com>
---
 library/psa_crypto.c                 |  8 ++++++++
 library/psa_crypto_slot_management.c | 24 +++++++-----------------
 library/psa_crypto_slot_management.h |  3 +++
 3 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index a0e58a2712..5300126c36 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1679,7 +1679,15 @@ static psa_status_t psa_start_key_creation(
         return status;
     }
 
+#if defined(MBEDTLS_THREADING_C)
+    PSA_THREADING_CHK_RET(mbedtls_mutex_lock(
+                              &mbedtls_threading_key_slot_mutex));
+#endif
     status = psa_reserve_free_key_slot(&volatile_key_id, p_slot);
+#if defined(MBEDTLS_THREADING_C)
+    PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
+                              &mbedtls_threading_key_slot_mutex));
+#endif
     if (status != PSA_SUCCESS) {
         return status;
     }
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 07d7f35fc8..dc38662e1b 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -160,13 +160,9 @@ psa_status_t psa_reserve_free_key_slot(psa_key_id_t *volatile_key_id,
     size_t slot_idx;
     psa_key_slot_t *selected_slot, *unused_persistent_key_slot;
 
-#if defined(MBEDTLS_THREADING_C)
-    PSA_THREADING_CHK_GOTO_EXIT(mbedtls_mutex_lock(
-                                    &mbedtls_threading_key_slot_mutex));
-#endif
     if (!global_data.key_slots_initialized) {
         status = PSA_ERROR_BAD_STATE;
-        goto exit;
+        goto error;
     }
 
     selected_slot = unused_persistent_key_slot = NULL;
@@ -198,7 +194,7 @@ psa_status_t psa_reserve_free_key_slot(psa_key_id_t *volatile_key_id,
         psa_register_read(selected_slot);
         status = psa_wipe_key_slot(selected_slot);
         if (status != PSA_SUCCESS) {
-            goto exit;
+            goto error;
         }
     }
 
@@ -206,27 +202,21 @@ psa_status_t psa_reserve_free_key_slot(psa_key_id_t *volatile_key_id,
         status = psa_key_slot_state_transition(selected_slot, PSA_SLOT_EMPTY,
                                                PSA_SLOT_FILLING);
         if (status != PSA_SUCCESS) {
-            goto exit;
+            goto error;
         }
 
         *volatile_key_id = PSA_KEY_ID_VOLATILE_MIN +
                            ((psa_key_id_t) (selected_slot - global_data.key_slots));
         *p_slot = selected_slot;
 
-        goto exit;
+        return PSA_SUCCESS;
     }
     status = PSA_ERROR_INSUFFICIENT_MEMORY;
 
-exit:
-    if (status != PSA_SUCCESS) {
-        *p_slot = NULL;
-        *volatile_key_id = 0;
-    }
+error:
+    *p_slot = NULL;
+    *volatile_key_id = 0;
 
-#if defined(MBEDTLS_THREADING_C)
-    PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
-                              &mbedtls_threading_key_slot_mutex));
-#endif
     return status;
 }
 
diff --git a/library/psa_crypto_slot_management.h b/library/psa_crypto_slot_management.h
index 18a914496d..585de13184 100644
--- a/library/psa_crypto_slot_management.h
+++ b/library/psa_crypto_slot_management.h
@@ -107,6 +107,9 @@ void psa_wipe_all_key_slots(void);
  * It is the responsibility of the caller to change the slot's state to
  * PSA_SLOT_EMPTY/FULL once key creation has finished.
  *
+ * If multi-threading is enabled, the caller must hold the
+ * global key slot mutex.
+ *
  * \param[out] volatile_key_id   On success, volatile key identifier
  *                               associated to the returned slot.
  * \param[out] p_slot            On success, a pointer to the slot.