Dynamic key store: implementation

When MBEDTLS_PSA_KEY_STORE_DYNAMIC is enabled, key slots are now organized in
multiple slices. The slices are allocated on demand, which allows the key
store to grow. The size of slices grows exponentially, which allows reaching
a large number of slots with a small (static) number of slices without too
much overhead.

Maintain a linked list of free slots in each slice. This way, allocating a
slot takes O(1) time unless a slice needs to be allocated.

In this commit, slices are only ever freed when deinitializing the key
store. This should be improved in the future to free empty slices.

To avoid growing the persistent key cache without control, the persistent
key cache has a fixed size (reusing MBEDTLS_PSA_KEY_SLOT_COUNT to avoid
creating yet another option).

When MBEDTLS_PSA_KEY_STORE_DYNAMIC is disabled. no semantic change and
minimal changes to the code.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>

library/psa_crypto.c

@@ -1232,6 +1232,11 @@ psa_status_t psa_wipe_key_slot(psa_key_slot_t *slot)
             status = PSA_ERROR_CORRUPTION_DETECTED;
     }
 
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+    size_t slice_index = slot->slice_index;
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
+
     /* Multipart operations may still be using the key. This is safe
      * because all multipart operation objects are independent from
      * the key slot: if they need to access the key after the setup
@@ -1242,6 +1247,17 @@ psa_status_t psa_wipe_key_slot(psa_key_slot_t *slot)
      * zeroize because the metadata is not particularly sensitive.
      * This memset also sets the slot's state to PSA_SLOT_EMPTY. */
     memset(slot, 0, sizeof(*slot));
+
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+    /* If the slot is already corrupted, something went deeply wrong,
+     * like a thread still using the slot or a stray pointer leading
+     * to the slot's memory being used for another object. Let the slot
+     * leak rather than make the corruption worse. */
+    if (status == PSA_SUCCESS) {
+        status = psa_free_key_slot(slice_index, slot);
+    }
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
     return status;
 }
 
@@ -1753,8 +1769,6 @@ static psa_status_t psa_start_key_creation(
     psa_se_drv_table_entry_t **p_drv)
 {
     psa_status_t status;
-    psa_key_id_t volatile_key_id;
-    psa_key_slot_t *slot;
 
     (void) method;
     *p_drv = NULL;
@@ -1764,11 +1778,16 @@ static psa_status_t psa_start_key_creation(
         return status;
     }
 
+    int key_is_volatile = PSA_KEY_LIFETIME_IS_VOLATILE(attributes->lifetime);
+    psa_key_id_t volatile_key_id;
+
 #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);
+    status = psa_reserve_free_key_slot(
+        key_is_volatile ? &volatile_key_id : NULL,
+        p_slot);
 #if defined(MBEDTLS_THREADING_C)
     PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
                               &mbedtls_threading_key_slot_mutex));
@@ -1776,7 +1795,7 @@ static psa_status_t psa_start_key_creation(
     if (status != PSA_SUCCESS) {
         return status;
     }
-    slot = *p_slot;
+    psa_key_slot_t *slot = *p_slot;
 
     /* We're storing the declared bit-size of the key. It's up to each
      * creation mechanism to verify that this information is correct.
@@ -1787,7 +1806,7 @@ static psa_status_t psa_start_key_creation(
      * definition. */
 
     slot->attr = *attributes;
-    if (PSA_KEY_LIFETIME_IS_VOLATILE(slot->attr.lifetime)) {
+    if (key_is_volatile) {
 #if !defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
         slot->attr.id = volatile_key_id;
 #else

library/psa_crypto_core.h

@@ -80,10 +80,38 @@ typedef struct {
      * slots that are in a suitable state for the function.
      * For example, psa_get_and_lock_key_slot_in_memory, which finds a slot
      * containing a given key ID, will only check slots whose state variable is
-     * PSA_SLOT_FULL. */
+     * PSA_SLOT_FULL.
+     */
     psa_key_slot_state_t state;
 
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+    /* The index of the slice containing this slot.
+     * This field must be filled if the slot contains a key
+     * (including keys being created or destroyed), and can be either
+     * filled or 0 when the slot is free. */
+    uint8_t slice_index;
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
     union {
+        struct {
+            /* The index of the next slot in the free list for this
+             * slice, relative * to the next array element.
+             *
+             * That is, 0 means the next slot, 1 means the next slot
+             * but one, etc. -1 would mean the slot itself. -2 means
+             * the previous slot, etc.
+             *
+             * If this is beyond the array length, the free list ends with the
+             * current element.
+             *
+             * The reason for this strange encoding is that 0 means the next
+             * element. This way, when we allocate a slice and initialize it
+             * to all-zero, the slice is ready for use, with a free list that
+             * consists of all the slots in order.
+             */
+            int32_t next_free_relative_to_next;
+        } free;
+
         struct {
             /*
              * Number of functions registered as reading the material in the key slot.

library/psa_crypto_slot_management.c

@@ -58,13 +58,110 @@ MBEDTLS_STATIC_ASSERT(PSA_KEY_ID_VOLATILE_MAX < MBEDTLS_PSA_KEY_ID_BUILTIN_MIN |
 
 
 
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+
+/* Dynamic key store.
+ *
+ * The key store consists of multiple slices.
+ *
+ * The volatile keys are stored in variable-sized tables called slices.
+ * Slices are allocated on demand and deallocated when possible.
+ * The size of slices increases exponentially, so the average overhead
+ * (number of slots that are allocated but not used) is roughly
+ * proportional to the number of keys (with a factor that grows
+ * when the key store is fragmented).
+ *
+ * One slice is dedicated to the cache of persistent and built-in keys.
+ * For simplicity, they are separated from volatile keys. This cache
+ * slice has a fixed size and has the slice index KEY_SLOT_CACHE_SLICE_INDEX,
+ * located after the slices for volatile keys.
+ */
+
+/* Size of slice 0 containing the cache of persistent and built-in keys. */
+#define PERSISTENT_KEY_CACHE_COUNT MBEDTLS_PSA_KEY_SLOT_COUNT
+
+/* Volatile keys are stored in slices 1 through KEY_SLICE_COUNT inclusive.
+ * Each slice is twice the size of the previous slice.
+ * Volatile key identifiers encode the slice number as follows:
+ *     bits 30..31:  0b10 (mandated by the PSA Crypto specification).
+ *     bits 25..29:  slice index (0...KEY_SLOT_VOLATILE_SLICE_COUNT-1)
+ *     bits 0..24:   slot index in slice
+ */
+#define KEY_ID_SLOT_INDEX_WIDTH 25u
+#define KEY_ID_SLICE_INDEX_WIDTH 5u
+
+#define KEY_SLOT_VOLATILE_SLICE_BASE_LENGTH 16u
+#define KEY_SLOT_VOLATILE_SLICE_COUNT 22u
+#define KEY_SLICE_COUNT (KEY_SLOT_VOLATILE_SLICE_COUNT + 1u)
+#define KEY_SLOT_CACHE_SLICE_INDEX KEY_SLOT_VOLATILE_SLICE_COUNT
+
+#if KEY_ID_SLICE_INDEX_WIDTH + KEY_ID_SLOT_INDEX_WIDTH > 30
+#error "Not enough room in volatile key IDs for slice index and slot index"
+#endif
+#if KEY_SLICE_COUNT >= (1 << KEY_ID_SLICE_INDEX_WIDTH) - 1
+#error "Too many slices to fit the slice index in a volatile key ID"
+#endif
+#define KEY_SLICE_LENGTH_MAX                                            \
+    (KEY_SLOT_VOLATILE_SLICE_BASE_LENGTH << (KEY_SLOT_VOLATILE_SLICE_COUNT - 1))
+#if KEY_SLICE_LENGTH_MAX > 1 << KEY_ID_SLOT_INDEX_WIDTH
+#error "Not enough room in volatile key IDs for a slot index in the largest slice"
+#endif
+#if KEY_ID_SLICE_INDEX_WIDTH > 8
+#error "Slice index does not fit in uint8_t for psa_key_slot_t::slice_index"
+#endif
+
+
+/* Calculate the volatile key id to use for a given slot.
+ * This function assumes valid parameter values. */
+static psa_key_id_t volatile_key_id_of_index(size_t slice_idx,
+                                             size_t slot_idx)
+{
+    return 0x40000000u | (slice_idx << KEY_ID_SLOT_INDEX_WIDTH) | slot_idx;
+}
+
+/* Calculate the slice containing the given volatile key.
+ * This function assumes valid parameter values. */
+static size_t slice_index_of_volatile_key_id(psa_key_id_t key_id)
+{
+    size_t mask = (1LU << KEY_ID_SLICE_INDEX_WIDTH) - 1;
+    return (key_id >> KEY_ID_SLOT_INDEX_WIDTH) & mask;
+}
+
+/* Calculate the index of the slot containing the given volatile key.
+ * This function assumes valid parameter values. */
+static size_t slot_index_of_volatile_key_id(psa_key_id_t key_id)
+{
+    return key_id & ((1LU << KEY_ID_SLOT_INDEX_WIDTH) - 1);
+}
+
+/* In global_data.first_free_slot_index, use this special value to
+ * indicate that the slice is full. */
+#define FREE_SLOT_INDEX_NONE ((size_t) -1)
+
+#else /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
+/* Static key store.
+ *
+ * All the keys (volatile or persistent) are in a single slice.
+ * We only use slices as a concept to allow some differences between
+ * static and dynamic key store management to be buried in auxiliary
+ * functions.
+ */
+
 #define PERSISTENT_KEY_CACHE_COUNT MBEDTLS_PSA_KEY_SLOT_COUNT
 #define KEY_SLICE_COUNT 1u
 #define KEY_SLOT_CACHE_SLICE_INDEX 0
 
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
 
 typedef struct {
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+    psa_key_slot_t *key_slices[KEY_SLICE_COUNT];
+    size_t first_free_slot_index[KEY_SLOT_VOLATILE_SLICE_COUNT];
+#else /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
     psa_key_slot_t key_slots[MBEDTLS_PSA_KEY_SLOT_COUNT];
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
     uint8_t key_slots_initialized;
 } psa_global_data_t;
 
@@ -128,6 +225,46 @@ static inline psa_key_slot_t *get_persistent_key_slot(size_t slot_idx);
  */
 static inline psa_key_slot_t *get_key_slot(size_t slice_idx, size_t slot_idx);
 
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+
+static inline size_t key_slice_length(size_t slice_idx)
+{
+    if (slice_idx == KEY_SLOT_CACHE_SLICE_INDEX) {
+        return PERSISTENT_KEY_CACHE_COUNT;
+    } else {
+        return KEY_SLOT_VOLATILE_SLICE_BASE_LENGTH << slice_idx;
+    }
+}
+
+static inline psa_key_slot_t *get_volatile_key_slot(psa_key_id_t key_id)
+{
+    size_t slice_idx = slice_index_of_volatile_key_id(key_id);
+    if (slice_idx >= KEY_SLOT_VOLATILE_SLICE_COUNT) {
+        return NULL;
+    }
+    size_t slot_idx = slot_index_of_volatile_key_id(key_id);
+    if (slot_idx >= key_slice_length(slice_idx)) {
+        return NULL;
+    }
+    psa_key_slot_t *slice = global_data.key_slices[slice_idx];
+    if (slice == NULL) {
+        return NULL;
+    }
+    return &slice[slot_idx];
+}
+
+static inline psa_key_slot_t *get_persistent_key_slot(size_t slot_idx)
+{
+    return &global_data.key_slices[KEY_SLOT_CACHE_SLICE_INDEX][slot_idx];
+}
+
+static inline psa_key_slot_t *get_key_slot(size_t slice_idx, size_t slot_idx)
+{
+    return &global_data.key_slices[slice_idx][slot_idx];
+}
+
+#else /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
 static inline size_t key_slice_length(size_t slice_idx)
 {
     (void) slice_idx;
@@ -153,6 +290,9 @@ static inline psa_key_slot_t *get_key_slot(size_t slice_idx, size_t slot_idx)
     return &global_data.key_slots[slot_idx];
 }
 
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
+
 
 int psa_is_valid_key_id(mbedtls_svc_key_id_t key, int vendor_ok)
 {
@@ -219,8 +359,9 @@ static psa_status_t psa_get_and_lock_key_slot_in_memory(
 
         /* Check if both the PSA key identifier key_id and the owner
          * identifier of key match those of the key slot. */
-        if ((slot->state == PSA_SLOT_FULL) &&
-            (mbedtls_svc_key_id_equal(key, slot->attr.id))) {
+        if (slot != NULL &&
+            slot->state == PSA_SLOT_FULL &&
+            mbedtls_svc_key_id_equal(key, slot->attr.id)) {
             status = PSA_SUCCESS;
         } else {
             status = PSA_ERROR_DOES_NOT_EXIST;
@@ -254,11 +395,21 @@ static psa_status_t psa_get_and_lock_key_slot_in_memory(
 
 psa_status_t psa_initialize_key_slots(void)
 {
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+    global_data.key_slices[KEY_SLOT_CACHE_SLICE_INDEX] =
+        mbedtls_calloc(PERSISTENT_KEY_CACHE_COUNT,
+                       sizeof(*global_data.key_slices[KEY_SLOT_CACHE_SLICE_INDEX]));
+    if (global_data.key_slices[KEY_SLOT_CACHE_SLICE_INDEX] == NULL) {
+        return PSA_ERROR_INSUFFICIENT_MEMORY;
+    }
+#else /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
     /* Nothing to do: program startup and psa_wipe_all_key_slots() both
      * guarantee that the key slots are initialized to all-zero, which
      * means that all the key slots are in a valid, empty state. The global
      * data mutex is already held when calling this function, so no need to
      * lock it here, to set the flag. */
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
     global_data.key_slots_initialized = 1;
     return PSA_SUCCESS;
 }
@@ -266,17 +417,137 @@ psa_status_t psa_initialize_key_slots(void)
 void psa_wipe_all_key_slots(void)
 {
     for (size_t slice_idx = 0; slice_idx < KEY_SLICE_COUNT; slice_idx++) {
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+        if (global_data.key_slices[slice_idx] == NULL) {
+            continue;
+        }
+#endif  /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
         for (size_t slot_idx = 0; slot_idx < key_slice_length(slice_idx); slot_idx++) {
             psa_key_slot_t *slot = get_key_slot(slice_idx, slot_idx);
+            if (slot->state == PSA_SLOT_EMPTY) {
+                /* Don't call psa_wipe_key_slot() on an already-empty slot.
+                 * It rejects that case anyway, though we bypass it by setting
+                 * the slot state to PENDING_DELETION.
+                 * Also, when MBEDTLS_PSA_KEY_STORE_DYNAMIC is enabled,
+                 * psa_wipe_key_slot() needs to have a valid slice_index
+                 * field, but that value might not be correct in a
+                 * free slot. */
+                continue;
+            }
             slot->var.occupied.registered_readers = 1;
             slot->state = PSA_SLOT_PENDING_DELETION;
             (void) psa_wipe_key_slot(slot);
         }
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+        mbedtls_free(global_data.key_slices[slice_idx]);
+        global_data.key_slices[slice_idx] = NULL;
+#endif  /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
     }
+
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+    for (size_t slice_idx = 0; slice_idx < KEY_SLOT_VOLATILE_SLICE_COUNT; slice_idx++) {
+        global_data.first_free_slot_index[slice_idx] = 0;
+    }
+#endif  /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
     /* The global data mutex is already held when calling this function. */
     global_data.key_slots_initialized = 0;
 }
 
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+
+static psa_status_t psa_allocate_volatile_key_slot(psa_key_id_t *key_id,
+                                                   psa_key_slot_t **p_slot)
+{
+    size_t slice_idx;
+    for (slice_idx = 0; slice_idx < KEY_SLOT_VOLATILE_SLICE_COUNT; slice_idx++) {
+        if (global_data.first_free_slot_index[slice_idx] != FREE_SLOT_INDEX_NONE) {
+            break;
+        }
+    }
+    if (slice_idx == KEY_SLOT_VOLATILE_SLICE_COUNT) {
+        return PSA_ERROR_INSUFFICIENT_MEMORY;
+    }
+
+    if (global_data.key_slices[slice_idx] == NULL) {
+        global_data.key_slices[slice_idx] =
+            mbedtls_calloc(key_slice_length(slice_idx),
+                           sizeof(psa_key_slot_t));
+        if (global_data.key_slices[slice_idx] == NULL) {
+            return PSA_ERROR_INSUFFICIENT_MEMORY;
+        }
+    }
+    psa_key_slot_t *slice = global_data.key_slices[slice_idx];
+
+    size_t slot_idx = global_data.first_free_slot_index[slice_idx];
+    *key_id = volatile_key_id_of_index(slice_idx, slot_idx);
+
+    psa_key_slot_t *slot = &slice[slot_idx];
+    size_t next_free = slot_idx + 1 + slot->var.free.next_free_relative_to_next;
+    if (next_free >= key_slice_length(slice_idx)) {
+        next_free = FREE_SLOT_INDEX_NONE;
+    }
+    global_data.first_free_slot_index[slice_idx] = next_free;
+    /* The .next_free field is not meaningful when the slot is not free,
+     * so give it the same content as freshly initialized memory. */
+    slot->var.free.next_free_relative_to_next = 0;
+
+    psa_status_t status = psa_key_slot_state_transition(slot,
+                                                        PSA_SLOT_EMPTY,
+                                                        PSA_SLOT_FILLING);
+    if (status != PSA_SUCCESS) {
+        /* The only reason for failure is if the slot state was not empty.
+         * This indicates that something has gone horribly wrong.
+         * In this case, we leave the slot out of the free list, and stop
+         * modifying it. This minimizes any further corruption. The slot
+         * is a memory leak, but that's a lesser evil. */
+        return status;
+    }
+
+    *p_slot = slot;
+    slot->slice_index = slice_idx;
+    return PSA_SUCCESS;
+}
+
+psa_status_t psa_free_key_slot(size_t slice_idx,
+                               psa_key_slot_t *slot)
+{
+
+    if (slice_idx == KEY_SLOT_CACHE_SLICE_INDEX) {
+        /* This is a cache entry. We don't maintain a free list, so
+         * there's nothing to do. */
+        return PSA_SUCCESS;
+    }
+    if (slice_idx >= KEY_SLOT_VOLATILE_SLICE_COUNT) {
+        return PSA_ERROR_CORRUPTION_DETECTED;
+    }
+
+    psa_key_slot_t *slice = global_data.key_slices[slice_idx];
+    psa_key_slot_t *slice_end = slice + key_slice_length(slice_idx);
+    if (slot < slice || slot >= slice_end) {
+        /* The slot isn't actually in the slice! We can't detect that
+         * condition for sure, because the pointer comparison itself is
+         * undefined behavior in that case. That same condition makes the
+         * subtraction to calculate the slot index also UB.
+         * Give up now to avoid causing further corruption.
+         */
+        return PSA_ERROR_CORRUPTION_DETECTED;
+    }
+    size_t slot_idx = slot - slice;
+
+    size_t next_free = global_data.first_free_slot_index[slice_idx];
+    if (next_free >= key_slice_length(slice_idx)) {
+        /* The slot was full. The newly freed slot thus becomes the
+         * end of the free list. */
+        next_free = key_slice_length(slice_idx);
+    }
+    global_data.first_free_slot_index[slice_idx] = slot_idx;
+    slot->var.free.next_free_relative_to_next = next_free - slot_idx - 1;
+
+    return PSA_SUCCESS;
+}
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
 psa_status_t psa_reserve_free_key_slot(psa_key_id_t *volatile_key_id,
                                        psa_key_slot_t **p_slot)
 {
@@ -289,6 +560,17 @@ psa_status_t psa_reserve_free_key_slot(psa_key_id_t *volatile_key_id,
         goto error;
     }
 
+    if (volatile_key_id != NULL) {
+        *volatile_key_id = 0;
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+        return psa_allocate_volatile_key_slot(volatile_key_id, p_slot);
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+    }
+
+    /* With a dynamic key store, allocate an entry in the cache slice,
+     * applicable only to non-volatile keys that get cached in RAM.
+     * With a static key store, allocate an entry in the sole slice,
+     * applicable to all keys. */
     selected_slot = unused_persistent_key_slot = NULL;
     for (slot_idx = 0; slot_idx < PERSISTENT_KEY_CACHE_COUNT; slot_idx++) {
         psa_key_slot_t *slot = get_key_slot(KEY_SLOT_CACHE_SLICE_INDEX, slot_idx);
@@ -329,8 +611,18 @@ psa_status_t psa_reserve_free_key_slot(psa_key_id_t *volatile_key_id,
             goto error;
         }
 
-        *volatile_key_id = PSA_KEY_ID_VOLATILE_MIN +
-                           ((psa_key_id_t) (selected_slot - global_data.key_slots));
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+        selected_slot->slice_index = KEY_SLOT_CACHE_SLICE_INDEX;
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
+#if !defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+        if (volatile_key_id != NULL) {
+            /* Refresh slot_idx, for when the slot is not the original
+             * selected_slot but rather unused_persistent_key_slot.  */
+            slot_idx = selected_slot - global_data.key_slots;
+            *volatile_key_id = PSA_KEY_ID_VOLATILE_MIN + slot_idx;
+        }
+#endif
         *p_slot = selected_slot;
 
         return PSA_SUCCESS;
@@ -339,7 +631,6 @@ psa_status_t psa_reserve_free_key_slot(psa_key_id_t *volatile_key_id,
 
 error:
     *p_slot = NULL;
-    *volatile_key_id = 0;
 
     return status;
 }
@@ -498,9 +789,8 @@ psa_status_t psa_get_and_lock_key_slot(mbedtls_svc_key_id_t key,
     /* Loading keys from storage requires support for such a mechanism */
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) || \
     defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
-    psa_key_id_t volatile_key_id;
 
-    status = psa_reserve_free_key_slot(&volatile_key_id, p_slot);
+    status = psa_reserve_free_key_slot(NULL, p_slot);
     if (status != PSA_SUCCESS) {
 #if defined(MBEDTLS_THREADING_C)
         PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
@@ -760,15 +1050,20 @@ void mbedtls_psa_get_stats(mbedtls_psa_stats_t *stats)
     memset(stats, 0, sizeof(*stats));
 
     for (size_t slice_idx = 0; slice_idx < KEY_SLICE_COUNT; slice_idx++) {
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+        if (global_data.key_slices[slice_idx] == NULL) {
+            continue;
+        }
+#endif  /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
         for (size_t slot_idx = 0; slot_idx < key_slice_length(slice_idx); slot_idx++) {
             const psa_key_slot_t *slot = get_key_slot(slice_idx, slot_idx);
-            if (psa_key_slot_has_readers(slot)) {
-                ++stats->locked_slots;
-            }
             if (slot->state == PSA_SLOT_EMPTY) {
                 ++stats->empty_slots;
                 continue;
             }
+            if (psa_key_slot_has_readers(slot)) {
+                ++stats->locked_slots;
+            }
             if (PSA_KEY_LIFETIME_IS_VOLATILE(slot->attr.lifetime)) {
                 ++stats->volatile_slots;
             } else {

library/psa_crypto_slot_management.h

@@ -17,8 +17,10 @@
  *
  *  The first #MBEDTLS_PSA_KEY_SLOT_COUNT identifiers of the implementation
  *  range of key identifiers are reserved for volatile key identifiers.
- *  A volatile key identifier is equal to #PSA_KEY_ID_VOLATILE_MIN plus the
- *  index of the key slot containing the volatile key definition.
+ *
+ *  If \c id is a a volatile key identifier, #PSA_KEY_ID_VOLATILE_MIN - \c id
+ *  indicates the key slot containing the volatile key definition. See
+ *  psa_crypto_slot_management.c for details.
  */
 
 /** The minimum value for a volatile key identifier.
@@ -27,8 +29,12 @@
 
 /** The maximum value for a volatile key identifier.
  */
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+#define PSA_KEY_ID_VOLATILE_MAX (MBEDTLS_PSA_KEY_ID_BUILTIN_MIN - 1)
+#else /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
 #define PSA_KEY_ID_VOLATILE_MAX                                 \
     (PSA_KEY_ID_VOLATILE_MIN + MBEDTLS_PSA_KEY_SLOT_COUNT - 1)
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
 
 /** Test whether a key identifier is a volatile key identifier.
  *
@@ -113,13 +119,20 @@ void psa_wipe_all_key_slots(void);
  * 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] volatile_key_id   If null, reserve a cache slot for
+ *                               a persistent or built-in key.
+ *                               If non-null, allocate a slot for
+ *                               a volatile key.
+ *                               If non-null, on success, the volatile key
+ *                               identifier corresponding with the
+ *                               returned slot.
  * \param[out] p_slot            On success, a pointer to the slot.
  *
  * \retval #PSA_SUCCESS \emptydescription
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
  *         There were no free key slots.
+ *         When #MBEDTLS_PSA_KEY_STORE_DYNAMIC is enabled, there was not
+ *         enough memory to allocate more slots.
  * \retval #PSA_ERROR_BAD_STATE \emptydescription
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
  *         This function attempted to operate on a key slot which was in an
@@ -128,6 +141,29 @@ void psa_wipe_all_key_slots(void);
 psa_status_t psa_reserve_free_key_slot(psa_key_id_t *volatile_key_id,
                                        psa_key_slot_t **p_slot);
 
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+/** Return a key slot to the free list.
+ *
+ * Call this function when a slot obtained from psa_reserve_free_key_slot()
+ * is no longer in use.
+ *
+ * If multi-threading is enabled, the caller must hold the
+ * global key slot mutex.
+ *
+ * \param slice_idx             The slice containing the slot.
+ *                              This is `slot->slice_index` when the slot
+ *                              is obtained from psa_reserve_free_key_slot().
+ * \param slot                  The key slot.
+ *
+ * \retval #PSA_SUCCESS \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ *         This function attempted to operate on a key slot which was in an
+ *         unexpected state.
+ */
+psa_status_t psa_free_key_slot(size_t slice_idx,
+                               psa_key_slot_t *slot);
+#endif /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
+
 /** Change the state of a key slot.
  *
  * This function changes the state of the key slot from expected_state to