Add threading support to the mldsa allocator

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>

Author: bensze01

drivers/pqcp/src/mldsa_native_buffer_alloc.h

@@ -13,6 +13,10 @@
 #include <stdint.h>
 #include "src/sys.h"
 
+#if defined(MBEDTLS_THREADING_C)
+#include "threading_internal.h"
+#endif
+
 // Sufficient for signing with TF_PSA_CRYPTO_PQCP_MLDSA_87_ENABLED
 #define TF_PSA_CRYPTO_MLD_ALLOC_BUFFER_SIZE 123200
 
@@ -22,26 +26,47 @@ struct tf_psa_crypto_mld_context {
     int alloc_offset; // Negative values indicate allocator errors
 };
 
+#if defined(MBEDTLS_THREADING_C)
+#define TF_PSA_CRYPTO_MLD_ALLOC_LOCK() mbedtls_mutex_lock(&mbedtls_threading_pqcp_buffer_alloc_mutex)
+#define TF_PSA_CRYPTO_MLD_ALLOC_UNLOCK() mbedtls_mutex_unlock(&mbedtls_threading_pqcp_buffer_alloc_mutex)
+#else /* MBEDTLS_THREADING_C */
+#define TF_PSA_CRYPTO_MLD_ALLOC_LOCK() 0
+#define TF_PSA_CRYPTO_MLD_ALLOC_UNLOCK() 0
+#endif /* MBEDTLS_THREADING_C */
+
 #define TF_PSA_CRYPTO_MLD_CUSTOM_ALLOC(v, T, N, context) \
     T *(v) = NULL; \
     do { \
         /* Verify that the allocation would fit in the buffer by itself, avoiding overflows \
            This should be optimized away at compile-time */ \
         if ((N) > 0 && (N) <= TF_PSA_CRYPTO_MLD_ALLOC_BUFFER_SIZE / sizeof(T)) { \
+            if ((context).alloc_offset == 0) { \
+                (context).alloc_offset = TF_PSA_CRYPTO_MLD_ALLOC_LOCK(); \
+            } \
             if ((context).alloc_offset >= 0) { \
                 if ((size_t) (context).alloc_offset <= \
                     TF_PSA_CRYPTO_MLD_ALLOC_BUFFER_SIZE - MLD_ALIGN_UP(sizeof(T) * (N))) { \
                     (v) = (T *) (tf_psa_crypto_mld_alloc_buffer + (context).alloc_offset); \
                     (context).alloc_offset += MLD_ALIGN_UP(sizeof(T) * (N)); \
                 } else { \
-                    /* Fail all further allocations in this function -> goto cleanup */ \
+                    /* Fail all further allocations in this function -> goto cleanup
+                     * If we ended up here, that implies (alloc_offset != 0) \
+                     * thus we don't need to call UNLOCK */ \
                     (context).alloc_offset = PSA_ERROR_INSUFFICIENT_MEMORY; \
                 } \
             } \
         } \
     } while (0)
 
-#define TF_PSA_CRYPTO_MLD_CUSTOM_FREE(v, T, N, context)
+#define TF_PSA_CRYPTO_MLD_CUSTOM_FREE(v, T, N, context) \
+    do { \
+        if ((v) != NULL) { \
+            /* Only unlock after freeing the last allocation */ \
+            if ((uint8_t *) (v) == tf_psa_crypto_mld_alloc_buffer) { \
+                (void) TF_PSA_CRYPTO_MLD_ALLOC_UNLOCK(); \
+            } \
+        } \
+    } while (0)
 
 #endif /* TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC */
 

platform/threading.c

@@ -268,7 +268,10 @@ void mbedtls_threading_set_alt(
     mbedtls_mutex_init(&mbedtls_threading_key_slot_mutex);
     mbedtls_mutex_init(&mbedtls_threading_psa_globaldata_mutex);
     mbedtls_mutex_init(&mbedtls_threading_psa_rngdata_mutex);
+#if defined(TF_PSA_CRYPTO_PQCP_MLDSA_ENABLED) && defined(TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC)
+    mbedtls_mutext_init(&mbedtls_threading_pqcp_buffer_alloc_mutex)
 #endif
+#endif /* MBEDTLS_PSA_CRYPTO_C */
 }
 
 /*
@@ -286,7 +289,10 @@ void mbedtls_threading_free_alt(void)
     mbedtls_mutex_free(&mbedtls_threading_key_slot_mutex);
     mbedtls_mutex_free(&mbedtls_threading_psa_globaldata_mutex);
     mbedtls_mutex_free(&mbedtls_threading_psa_rngdata_mutex);
+#if defined(TF_PSA_CRYPTO_PQCP_MLDSA_ENABLED) && defined(TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC)
+    mbedtls_mutext_free(&mbedtls_threading_pqcp_buffer_alloc_mutex)
 #endif
+#endif /* MBEDTLS_PSA_CRYPTO_C */
 }
 #endif /* MBEDTLS_THREADING_ALT */
 
@@ -306,6 +312,9 @@ mbedtls_threading_mutex_t mbedtls_threading_gmtime_mutex MUTEX_INIT;
 mbedtls_threading_mutex_t mbedtls_threading_key_slot_mutex MUTEX_INIT;
 mbedtls_threading_mutex_t mbedtls_threading_psa_globaldata_mutex MUTEX_INIT;
 mbedtls_threading_mutex_t mbedtls_threading_psa_rngdata_mutex MUTEX_INIT;
+#if defined(TF_PSA_CRYPTO_PQCP_MLDSA_ENABLED) && defined(TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC)
+mbedtls_threading_mutex_t mbedtls_threading_pqcp_buffer_alloc_mutex MUTEX_INIT;
 #endif
+#endif /* MBEDTLS_PSA_CRYPTO_C */
 
 #endif /* MBEDTLS_THREADING_C */

platform/threading_internal.h

@@ -68,7 +68,11 @@ extern mbedtls_threading_mutex_t mbedtls_threading_psa_globaldata_mutex;
  * This mutex must be held when reading or writing to the PSA
  * global_data rng_state or rng struct members. */
 extern mbedtls_threading_mutex_t mbedtls_threading_psa_rngdata_mutex;
+
+#if defined(TF_PSA_CRYPTO_PQCP_MLDSA_ENABLED) && defined(TF_PSA_CRYPTO_PQCP_BUFFER_ALLOC)
+extern mbedtls_threading_mutex_t mbedtls_threading_pqcp_buffer_alloc_mutex;
 #endif
+#endif /* MBEDTLS_PSA_CRYPTO_C */
 
 #endif /* MBEDTLS_THREADING_C */