Update psa_get_and_lock_key_slot_X functions

Signed-off-by: Ryan Everett <ryan.everett@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 3c5bbbd..a27fd42 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -881,8 +881,9 @@
  * In case of a persistent key, the function loads the description of the key
  * into a key slot if not already done.
  *
- * On success, the returned key slot is locked. It is the responsibility of
- * the caller to unlock the key slot when it does not access it anymore.
+ * On success, the returned key slot has been registered for reading.
+ * It is the responsibility of the caller to call psa_unregister_read(slot)
+ * when they have finished reading the contents of the slot.
  */
 static psa_status_t psa_get_and_lock_key_slot_with_policy(
     mbedtls_svc_key_id_t key,
@@ -926,7 +927,7 @@
 
 error:
     *p_slot = NULL;
-    psa_unlock_key_slot(slot);
+    psa_unregister_read(slot);
 
     return status;
 }
@@ -941,8 +942,9 @@
  * psa_get_and_lock_key_slot_with_policy() when there is no opaque key support
  * for a cryptographic operation.
  *
- * On success, the returned key slot is locked. It is the responsibility of the
- * caller to unlock the key slot when it does not access it anymore.
+ * On success, the returned key slot has been registered for reading.
+ * It is the responsibility of the caller to call psa_unregister_read(slot)
+ * when they have finished reading the contents of the slot.
  */
 static psa_status_t psa_get_and_lock_transparent_key_slot_with_policy(
     mbedtls_svc_key_id_t key,
@@ -957,7 +959,7 @@
     }
 
     if (psa_key_lifetime_is_external((*p_slot)->attr.lifetime)) {
-        psa_unlock_key_slot(*p_slot);
+        psa_unregister_read(*p_slot);
         *p_slot = NULL;
         return PSA_ERROR_NOT_SUPPORTED;
     }
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 0f480fb..4846e33 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -108,7 +108,9 @@
 
         for (slot_idx = 0; slot_idx < MBEDTLS_PSA_KEY_SLOT_COUNT; slot_idx++) {
             slot = &global_data.key_slots[slot_idx];
-            if (mbedtls_svc_key_id_equal(key, slot->attr.id)) {
+            /* Only consider slots which are in a full state. */
+            if ((slot->state == PSA_SLOT_FULL) &&
+                (mbedtls_svc_key_id_equal(key, slot->attr.id))) {
                 break;
             }
         }
@@ -117,7 +119,7 @@
     }
 
     if (status == PSA_SUCCESS) {
-        status = psa_lock_key_slot(slot);
+        status = psa_register_read(slot);
         if (status == PSA_SUCCESS) {
             *p_slot = slot;
         }
@@ -367,7 +369,7 @@
     defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
     psa_key_id_t volatile_key_id;
 
-    status = psa_get_empty_key_slot(&volatile_key_id, p_slot);
+    status = psa_reserve_free_key_slot(&volatile_key_id, p_slot);
     if (status != PSA_SUCCESS) {
         return status;
     }
@@ -388,13 +390,24 @@
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 
     if (status != PSA_SUCCESS) {
+        /* Prepare the key slot to be wiped, and then wipe it.
+         * Don't overwrite status as a BAD_STATE error here
+         * can be reported in the psa_wipe_key_slot call. */
+        (*p_slot)->registered_readers = 1;
+        psa_key_slot_state_transition((*p_slot), PSA_SLOT_FILLING,
+                                      PSA_SLOT_PENDING_DELETION);
         psa_wipe_key_slot(*p_slot);
+
         if (status == PSA_ERROR_DOES_NOT_EXIST) {
             status = PSA_ERROR_INVALID_HANDLE;
         }
     } else {
         /* Add implicit usage flags. */
         psa_extend_key_usage_flags(&(*p_slot)->attr.policy.usage);
+
+        psa_key_slot_state_transition((*p_slot), PSA_SLOT_FILLING,
+                                      PSA_SLOT_FULL);
+        status = psa_register_read(*p_slot);
     }
 
     return status;
diff --git a/library/psa_crypto_slot_management.h b/library/psa_crypto_slot_management.h
index b2cf570..5858b18 100644
--- a/library/psa_crypto_slot_management.h
+++ b/library/psa_crypto_slot_management.h
@@ -54,8 +54,9 @@
  * In case of a persistent key, the function loads the description of the key
  * into a key slot if not already done.
  *
- * On success, the returned key slot is locked. It is the responsibility of
- * the caller to unlock the key slot when it does not access it anymore.
+ * On success, the returned key slot has been registered for reading.
+ * It is the responsibility of the caller to call psa_unregister_read(slot)
+ * when they have finished reading the contents of the slot.
  *
  * \param key           Key identifier to query.
  * \param[out] p_slot   On success, `*p_slot` contains a pointer to the
@@ -67,7 +68,9 @@
  *         description of the key identified by \p key.
  *         The key slot counter has been incremented.
  * \retval #PSA_ERROR_BAD_STATE
- *         The library has not been initialized.
+ *         The library has not been initialized. Or,
+ *         this call was operating on a key slot and found the slot in
+ *         an invalid state for the operation.
  * \retval #PSA_ERROR_INVALID_HANDLE
  *         \p key is not a valid key identifier.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY