From 0493ab56a424d886c2b63552dda5da21736ec939 Mon Sep 17 00:00:00 2001
From: Paul Elliott <paul.elliott@arm.com>
Date: Wed, 6 Mar 2024 12:26:58 +0000
Subject: [PATCH] Add PSA threaded init tests

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
---
 tests/suites/test_suite_psa_crypto_init.data  |   3 +
 .../test_suite_psa_crypto_init.function       | 115 ++++++++++++++++++
 2 files changed, 118 insertions(+)

diff --git a/tests/suites/test_suite_psa_crypto_init.data b/tests/suites/test_suite_psa_crypto_init.data
index 8c5b41d6cb..147d03fbed 100644
--- a/tests/suites/test_suite_psa_crypto_init.data
+++ b/tests/suites/test_suite_psa_crypto_init.data
@@ -10,6 +10,9 @@ deinit_without_init:0
 PSA deinit twice
 deinit_without_init:1
 
+PSA threaded init checks
+psa_threaded_init:100
+
 No random without init
 validate_module_init_generate_random:0
 
diff --git a/tests/suites/test_suite_psa_crypto_init.function b/tests/suites/test_suite_psa_crypto_init.function
index 7a434322ae..9ff33a6d84 100644
--- a/tests/suites/test_suite_psa_crypto_init.function
+++ b/tests/suites/test_suite_psa_crypto_init.function
@@ -1,6 +1,7 @@
 /* BEGIN_HEADER */
 #include <stdint.h>
 
+#include "psa_crypto_core.h"
 /* Some tests in this module configure entropy sources. */
 #include "psa_crypto_invasive.h"
 
@@ -112,6 +113,59 @@ static void custom_entropy_init(mbedtls_entropy_context *ctx)
 
 #endif /* !defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG) */
 
+#if defined MBEDTLS_THREADING_PTHREAD
+
+typedef struct {
+    int do_init;
+} thread_psa_init_ctx_t;
+
+static void *thread_psa_init_function(void *ctx)
+{
+    thread_psa_init_ctx_t *init_context = (thread_psa_init_ctx_t *) ctx;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    uint8_t random[10] = { 0 };
+
+    if (init_context->do_init) {
+        PSA_ASSERT(psa_crypto_init());
+    }
+
+    /* If this is a test only thread, then we can assume PSA is being started
+     * up on another thread and thus we cannot know whether the following tests
+     * will be successful or not. These checks are still useful, however even
+     * without checking the return codes as they may show up race conditions on
+     * the flags they check under TSAN.*/
+
+    /* Test getting if drivers are initialised. */
+    int can_do = psa_can_do_hash(PSA_ALG_NONE);
+
+    if (init_context->do_init) {
+        TEST_ASSERT(can_do == 1);
+    }
+
+#if !defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG)
+
+    /* Test getting global_data.rng_state. */
+    status = mbedtls_psa_crypto_configure_entropy_sources(NULL, NULL);
+
+    if (init_context->do_init) {
+        /* Bad state due to entropy sources already being setup in
+         * psa_crypto_init() */
+        TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
+    }
+#endif /* !defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG) */
+
+    /* Test using the PSA RNG ony if we know PSA is up and running. */
+    if (init_context->do_init) {
+        status = psa_generate_random(random, sizeof(random));
+
+        TEST_EQUAL(status, PSA_SUCCESS);
+    }
+
+exit:
+    return NULL;
+}
+#endif /* defined MBEDTLS_THREADING_PTHREAD */
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -154,6 +208,67 @@ void deinit_without_init(int count)
 }
 /* END_CASE */
 
+/* BEGIN_CASE depends_on:MBEDTLS_THREADING_PTHREAD */
+void psa_threaded_init(int arg_thread_count)
+{
+    thread_psa_init_ctx_t init_context;
+    thread_psa_init_ctx_t init_context_2;
+
+    size_t thread_count = (size_t) arg_thread_count;
+    mbedtls_test_thread_t *threads = NULL;
+
+    TEST_CALLOC(threads, sizeof(mbedtls_test_thread_t) * thread_count);
+
+    init_context.do_init = 1;
+
+    /* Test initialising PSA and testing certain protected globals on multiple
+     * threads. */
+    for (size_t i = 0; i < thread_count; i++) {
+        TEST_EQUAL(
+            mbedtls_test_thread_create(&threads[i],
+                                       thread_psa_init_function,
+                                       (void *) &init_context),
+            0);
+    }
+
+    for (size_t i = 0; i < thread_count; i++) {
+        TEST_EQUAL(mbedtls_test_thread_join(&threads[i]), 0);
+    }
+
+    PSA_DONE();
+
+    init_context_2.do_init = 0;
+
+    /* Test initialising PSA whilst also testing flags on other threads. */
+    for (size_t i = 0; i < thread_count; i++) {
+
+        if (i & 1) {
+
+            TEST_EQUAL(
+                mbedtls_test_thread_create(&threads[i],
+                                           thread_psa_init_function,
+                                           (void *) &init_context),
+                0);
+        } else {
+            TEST_EQUAL(
+                mbedtls_test_thread_create(&threads[i],
+                                           thread_psa_init_function,
+                                           (void *) &init_context_2),
+                0);
+        }
+    }
+
+    for (size_t i = 0; i < thread_count; i++) {
+        TEST_EQUAL(mbedtls_test_thread_join(&threads[i]), 0);
+    }
+exit:
+
+    PSA_DONE();
+
+    mbedtls_free(threads);
+}
+/* END_CASE */
+
 /* BEGIN_CASE */
 void validate_module_init_generate_random(int count)
 {