Improve full-key-store tests

Split the "many transient keys" test function in two: one that expects to
successfully create many keys, and one that expects to fill the key store.
This will make things easier when we add a dynamic key store where filling
the key store is not practical unless artificially limited.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.data b/tests/suites/test_suite_psa_crypto_slot_management.data
index 7d364ac..560350c 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.data
+++ b/tests/suites/test_suite_psa_crypto_slot_management.data
@@ -214,8 +214,23 @@
 invalid handle: huge
 invalid_handle:INVALID_HANDLE_HUGE:PSA_ERROR_INVALID_HANDLE
 
-Open many transient keys
-many_transient_keys:42
+Key slot count: less than maximum
+many_transient_keys:MBEDTLS_PSA_KEY_SLOT_COUNT - 1
+
+Key slot count: maximum
+many_transient_keys:MBEDTLS_PSA_KEY_SLOT_COUNT
+
+Key slot count: try to overfill, destroy first
+fill_key_store:0
+
+Key slot count: try to overfill, destroy second
+fill_key_store:1
+
+Key slot count: try to overfill, destroy next-to-last
+fill_key_store:-2
+
+Key slot count: try to overfill, destroy last
+fill_key_store:-1
 
 # Eviction from a key slot to be able to import a new persistent key.
 Key slot eviction to import a new persistent key
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index 94f26f6..013945e 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -98,6 +98,11 @@
     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. */
+#define MAX_VOLATILE_KEYS MBEDTLS_PSA_KEY_SLOT_COUNT
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -813,21 +818,19 @@
     psa_set_key_type(&attributes, PSA_KEY_TYPE_RAW_DATA);
 
     for (i = 0; i < max_keys; i++) {
+        mbedtls_test_set_step(i);
         status = psa_import_key(&attributes,
                                 (uint8_t *) &i, sizeof(i),
                                 &keys[i]);
-        if (status == PSA_ERROR_INSUFFICIENT_MEMORY) {
-            break;
-        }
         PSA_ASSERT(status);
         TEST_ASSERT(!mbedtls_svc_key_id_is_null(keys[i]));
         for (j = 0; j < i; j++) {
             TEST_ASSERT(!mbedtls_svc_key_id_equal(keys[i], keys[j]));
         }
     }
-    max_keys = i;
 
     for (i = 1; i < max_keys; i++) {
+        mbedtls_test_set_step(i);
         PSA_ASSERT(psa_close_key(keys[i - 1]));
         PSA_ASSERT(psa_export_key(keys[i],
                                   exported, sizeof(exported),
@@ -843,6 +846,97 @@
 }
 /* END_CASE */
 
+/* BEGIN_CASE depends_on:MAX_VOLATILE_KEYS */
+void fill_key_store(int key_to_destroy_arg)
+{
+    mbedtls_svc_key_id_t *keys = NULL;
+    size_t max_keys = MAX_VOLATILE_KEYS;
+    size_t i, j;
+    psa_status_t status;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    uint8_t exported[sizeof(size_t)];
+    size_t exported_length;
+
+    PSA_ASSERT(psa_crypto_init());
+
+    mbedtls_psa_stats_t stats;
+    mbedtls_psa_get_stats(&stats);
+    /* Account for any system-created volatile key, e.g. for the RNG. */
+    max_keys -= stats.volatile_slots;
+    TEST_CALLOC(keys, max_keys + 1);
+
+    psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_EXPORT);
+    psa_set_key_algorithm(&attributes, 0);
+    psa_set_key_type(&attributes, PSA_KEY_TYPE_RAW_DATA);
+
+    /* Fill the key store. */
+    for (i = 0; i < max_keys; i++) {
+        mbedtls_test_set_step(i);
+        status = psa_import_key(&attributes,
+                                (uint8_t *) &i, sizeof(i),
+                                &keys[i]);
+        PSA_ASSERT(status);
+        TEST_ASSERT(!mbedtls_svc_key_id_is_null(keys[i]));
+        for (j = 0; j < i; j++) {
+            TEST_ASSERT(!mbedtls_svc_key_id_equal(keys[i], keys[j]));
+        }
+    }
+
+    /* Attempt to overfill. */
+    mbedtls_test_set_step(max_keys);
+    status = psa_import_key(&attributes,
+                            (uint8_t *) &max_keys, sizeof(max_keys),
+                            &keys[max_keys]);
+    TEST_EQUAL(status, PSA_ERROR_INSUFFICIENT_MEMORY);
+    TEST_ASSERT(mbedtls_svc_key_id_is_null(keys[max_keys]));
+
+    /* Check that the keys are not corrupted. */
+    for (i = 0; i < max_keys; i++) {
+        mbedtls_test_set_step(i);
+        PSA_ASSERT(psa_export_key(keys[i],
+                                  exported, sizeof(exported),
+                                  &exported_length));
+        TEST_MEMORY_COMPARE(exported, exported_length,
+                            (uint8_t *) &i, sizeof(i));
+    }
+
+    /* Destroy one key and try again. */
+    size_t key_to_destroy = (key_to_destroy_arg >= 0 ?
+                             (size_t) key_to_destroy_arg :
+                             max_keys + key_to_destroy_arg);
+    mbedtls_svc_key_id_t reused_id = keys[key_to_destroy];
+    const uint8_t replacement_value[1] = { 0x64 };
+    PSA_ASSERT(psa_destroy_key(keys[key_to_destroy]));
+    keys[key_to_destroy] = MBEDTLS_SVC_KEY_ID_INIT;
+    status = psa_import_key(&attributes,
+                            replacement_value, sizeof(replacement_value),
+                            &keys[key_to_destroy]);
+    PSA_ASSERT(status);
+    TEST_ASSERT(mbedtls_svc_key_id_equal(reused_id, keys[key_to_destroy]));
+
+    /* Check that the keys are not corrupted and destroy them. */
+    for (i = 0; i < max_keys; i++) {
+        mbedtls_test_set_step(i);
+        PSA_ASSERT(psa_export_key(keys[i],
+                                  exported, sizeof(exported),
+                                  &exported_length));
+        if (i == key_to_destroy) {
+            TEST_MEMORY_COMPARE(exported, exported_length,
+                                replacement_value, sizeof(replacement_value));
+        } else {
+            TEST_MEMORY_COMPARE(exported, exported_length,
+                                (uint8_t *) &i, sizeof(i));
+        }
+        PSA_ASSERT(psa_destroy_key(keys[i]));
+        keys[i] = MBEDTLS_SVC_KEY_ID_INIT;
+    }
+
+exit:
+    PSA_DONE();
+    mbedtls_free(keys);
+}
+/* END_CASE */
+
 /* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C */
 void key_slot_eviction_to_import_new_key(int lifetime_arg)
 {