Dynamic key store: make full-key-store tests work effectively

Add a practical way to fill the dynamic key store by artificially limiting
the slice length through a test hook.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/tf-psa-crypto/core/psa_crypto_slot_management.c b/tf-psa-crypto/core/psa_crypto_slot_management.c
index 5b4539c..d740960 100644
--- a/tf-psa-crypto/core/psa_crypto_slot_management.c
+++ b/tf-psa-crypto/core/psa_crypto_slot_management.c
@@ -138,6 +138,13 @@
  * indicate that the slice is full. */
 #define FREE_SLOT_INDEX_NONE ((size_t) -1)
 
+#if defined(MBEDTLS_TEST_HOOKS)
+size_t psa_key_slot_volatile_slice_count(void)
+{
+    return KEY_SLOT_VOLATILE_SLICE_COUNT;
+}
+#endif
+
 #else /* MBEDTLS_PSA_KEY_STORE_DYNAMIC */
 
 /* Static key store.
@@ -227,11 +234,20 @@
 
 #if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
 
+#if defined(MBEDTLS_TEST_HOOKS)
+size_t (*mbedtls_test_hook_psa_volatile_key_slice_length)(size_t slice_idx) = NULL;
+#endif
+
 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 {
+#if defined(MBEDTLS_TEST_HOOKS)
+        if (mbedtls_test_hook_psa_volatile_key_slice_length != NULL) {
+            return mbedtls_test_hook_psa_volatile_key_slice_length(slice_idx);
+        }
+#endif
         return KEY_SLOT_VOLATILE_SLICE_BASE_LENGTH << slice_idx;
     }
 }
diff --git a/tf-psa-crypto/core/psa_crypto_slot_management.h b/tf-psa-crypto/core/psa_crypto_slot_management.h
index 1e6e935..26b7396 100644
--- a/tf-psa-crypto/core/psa_crypto_slot_management.h
+++ b/tf-psa-crypto/core/psa_crypto_slot_management.h
@@ -100,6 +100,24 @@
  */
 psa_status_t psa_initialize_key_slots(void);
 
+#if defined(MBEDTLS_TEST_HOOKS) && defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+/* Allow test code to customize the key slice length. We use this in tests
+ * that exhaust the key store to reach a full key store in reasonable time
+ * and memory.
+ *
+ * The length of each slice must be between 1 and
+ * (1 << KEY_ID_SLOT_INDEX_WIDTH) inclusive.
+ *
+ * The length for a given slice index must not change while
+ * the key store is initialized.
+ */
+extern size_t (*mbedtls_test_hook_psa_volatile_key_slice_length)(
+    size_t slice_idx);
+
+/* The number of volatile key slices. */
+size_t psa_key_slot_volatile_slice_count(void);
+#endif
+
 /** 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.
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.data b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.data
index 1bf300a..f379dba 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.data
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.data
@@ -228,6 +228,11 @@
 Key slot count: maximum
 many_transient_keys:MBEDTLS_PSA_KEY_SLOT_COUNT - MBEDTLS_TEST_PSA_INTERNAL_KEYS
 
+Key slot count: dynamic: more than MBEDTLS_PSA_KEY_SLOT_COUNT
+depends_on:MBEDTLS_PSA_KEY_STORE_DYNAMIC
+# Check that MBEDTLS_PSA_KEY_SLOT_COUNT doesn't apply to volatile keys.
+many_transient_keys:MBEDTLS_PSA_KEY_SLOT_COUNT + 1
+
 Key slot count: try to overfill, destroy first
 fill_key_store:0
 
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.function b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.function
index b2d3f29..604c4bd 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -98,11 +98,27 @@
     return 0;
 }
 
-/* Currently, there is always a maximum number of volatile keys that can
- * realistically be reached in tests. When we add configurations where this
- * is not true, undefine the macro in such configurations. */
 #if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC)
+#if defined(MBEDTLS_TEST_HOOKS)
+/* Artificially restrictable dynamic key store */
+#define KEY_SLICE_1_LENGTH 4
+#define KEY_SLICE_2_LENGTH 10
+static size_t tiny_key_slice_length(size_t slice_idx)
+{
+    switch (slice_idx) {
+        case 1: return KEY_SLICE_1_LENGTH;
+        case 2: return KEY_SLICE_2_LENGTH;
+        default: return 1;
+    }
+}
+#define MAX_VOLATILE_KEYS                       \
+    (KEY_SLICE_1_LENGTH + KEY_SLICE_2_LENGTH +  \
+     psa_key_slot_volatile_slice_count() - 2)
+
+#else  /* Effectively unbounded dynamic key store */
 #undef MAX_VOLATILE_KEYS
+#endif
+
 #else  /* Static key store */
 #define MAX_VOLATILE_KEYS MBEDTLS_PSA_KEY_SLOT_COUNT
 #endif
@@ -871,6 +887,10 @@
     uint8_t exported[sizeof(size_t)];
     size_t exported_length;
 
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC) && defined(MBEDTLS_TEST_HOOKS)
+    mbedtls_test_hook_psa_volatile_key_slice_length = &tiny_key_slice_length;
+#endif
+
     PSA_ASSERT(psa_crypto_init());
 
     mbedtls_psa_stats_t stats;
@@ -953,6 +973,9 @@
 exit:
     PSA_DONE();
     mbedtls_free(keys);
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC) && defined(MBEDTLS_TEST_HOOKS)
+    mbedtls_test_hook_psa_volatile_key_slice_length = NULL;
+#endif
 }
 /* END_CASE */