[nrf fromtree] psa: use static key slot buffers to store key material

This commit takes only relevant changes of PR
Mbed-TLS/mbedtls#9448 that was merged
in upstream Mbed TLS in the LTS branch "mbedtls-3.6".
Since the original PR was made of several commits, but most
of them were only affecting tests cases (not used in Zephyr),
only changes belonging to the "include" and "library" folders
were included here.

== IMPORTANT ==
Changes introduced in this commit will be automatically
part of Mbed TLS release 3.6.3, so by the time Zephyr's
Mbed TLS fork repo is bumbed to that official release,
this commit MUST be discarded.

This commit introduces the possibility to use static key slot
buffers in the PSA core instead of dynamically allocating them
when needed. This helps reducing heap memory usage as well as
potentially removing heap management ROM code if heap is not
used anywhere else in the Zephyr application.

Signed-off-by: Valerio Setti <vsetti@baylibre.com>
Signed-off-by: Frank Audun Kvamtrø <frank.kvamtro@nordicsemi.no>
(Cherry-picked from commit 4952e1328529ee549d412b498ea71c54f30aa3b1)

include/library/psa_crypto_core.h

@@ -155,7 +155,11 @@ typedef struct {
     /* Dynamically allocated key data buffer.
      * Format as specified in psa_export_key(). */
     struct key_data {
+#if defined(MBEDTLS_PSA_STATIC_KEY_SLOTS)
+        uint8_t data[MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE];
+#else
         uint8_t *data;
+#endif
         size_t bytes;
     } key;
 } psa_key_slot_t;

include/library/psa_crypto_storage.h

@@ -21,9 +21,16 @@ extern "C" {
 #include <stdint.h>
 #include <string.h>
 
-/* Limit the maximum key size in storage. This should have no effect
- * since the key size is limited in memory. */
+/* Limit the maximum key size in storage. */
+#if defined(MBEDTLS_PSA_STATIC_KEY_SLOTS)
+/* Reflect the maximum size for the key buffer. */
+#define PSA_CRYPTO_MAX_STORAGE_SIZE (MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE)
+#else
+/* Just set an upper boundary but it should have no effect since the key size
+ * is limited in memory. */
 #define PSA_CRYPTO_MAX_STORAGE_SIZE (PSA_BITS_TO_BYTES(PSA_MAX_KEY_BITS))
+#endif
+
 /* Sanity check: a file size must fit in 32 bits. Allow a generous
  * 64kB of metadata. */
 #if PSA_CRYPTO_MAX_STORAGE_SIZE > 0xffff0000

include/mbedtls/check_config.h

@@ -741,6 +741,11 @@
 #error "MBEDTLS_PSA_INJECT_ENTROPY is not compatible with MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG"
 #endif
 
+#if defined(MBEDTLS_PSA_KEY_STORE_DYNAMIC) &&           \
+    defined(MBEDTLS_PSA_STATIC_KEY_SLOTS)
+#error "MBEDTLS_PSA_KEY_STORE_DYNAMIC and MBEDTLS_PSA_STATIC_KEY_SLOTS cannot be defined simultaneously"
+#endif
+
 #if defined(MBEDTLS_PSA_ITS_FILE_C) && \
     !defined(MBEDTLS_FS_IO)
 #error "MBEDTLS_PSA_ITS_FILE_C defined, but not all prerequisites"

include/mbedtls/entropy.h

@@ -17,12 +17,13 @@
 
 #include "md.h"
 
-#if defined(MBEDTLS_MD_CAN_SHA512) && !defined(MBEDTLS_ENTROPY_FORCE_SHA256)
+#if (defined(MBEDTLS_MD_CAN_SHA512) || defined(PSA_WANT_ALG_SHA_512)) && \
+    !defined(MBEDTLS_ENTROPY_FORCE_SHA256)
 #define MBEDTLS_ENTROPY_SHA512_ACCUMULATOR
 #define MBEDTLS_ENTROPY_MD  MBEDTLS_MD_SHA512
 #define MBEDTLS_ENTROPY_BLOCK_SIZE      64      /**< Block size of entropy accumulator (SHA-512) */
 #else
-#if defined(MBEDTLS_MD_CAN_SHA256)
+#if (defined(MBEDTLS_MD_CAN_SHA256) || defined(PSA_WANT_ALG_SHA_256))
 #define MBEDTLS_ENTROPY_SHA256_ACCUMULATOR
 #define MBEDTLS_ENTROPY_MD  MBEDTLS_MD_SHA256
 #define MBEDTLS_ENTROPY_BLOCK_SIZE      32      /**< Block size of entropy accumulator (SHA-256) */

include/mbedtls/mbedtls_config.h

@@ -3268,6 +3268,26 @@
  */
 #define MBEDTLS_PSA_ITS_FILE_C
 
+/**
+ * \def MBEDTLS_PSA_STATIC_KEY_SLOTS
+ *
+ * Statically preallocate memory to store keys' material in PSA instead
+ * of allocating it dynamically when required. This allows builds without a
+ * heap, if none of the enabled cryptographic implementations or other features
+ * require it.
+ * This feature affects both volatile and persistent keys which means that
+ * it's not possible to persistently store a key which is larger than
+ * #MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE.
+ *
+ * \note This feature comes with a (potentially) higher RAM usage since:
+ *       - All the key slots are allocated no matter if they are used or not.
+ *       - Each key buffer's length is #MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE bytes.
+ *
+ * Requires: MBEDTLS_PSA_CRYPTO_C
+ *
+ */
+//#define MBEDTLS_PSA_STATIC_KEY_SLOTS
+
 /**
  * \def MBEDTLS_RIPEMD160_C
  *
@@ -4069,6 +4089,19 @@
  */
 //#define MBEDTLS_PSA_KEY_SLOT_COUNT 32
 
+/**
+ * \def MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE
+ *
+ * Define the size (in bytes) of each static key buffer when
+ * #MBEDTLS_PSA_STATIC_KEY_SLOTS is set. If not
+ * explicitly defined then it's automatically guessed from available PSA keys
+ * enabled in the build through PSA_WANT_xxx symbols.
+ * If required by the application this parameter can be set to higher values
+ * in order to store larger objects (ex: raw keys), but please note that this
+ * will increase RAM usage.
+ */
+//#define MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE       256
+
 /* RSA OPTIONS */
 //#define MBEDTLS_RSA_GEN_KEY_MIN_BITS            1024 /**<  Minimum RSA key size that can be generated in bits (Minimum possible value is 128 bits) */
 

include/psa/crypto_extra.h

@@ -32,6 +32,16 @@ extern "C" {
 #define MBEDTLS_PSA_KEY_SLOT_COUNT 32
 #endif
 
+/* If the size of static key slots is not explicitly defined by the user, then
+ * set it to the maximum between PSA_EXPORT_KEY_PAIR_OR_PUBLIC_MAX_SIZE and
+ * PSA_CIPHER_MAX_KEY_LENGTH.
+ * See mbedtls_config.h for the definition. */
+#if !defined(MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE)
+#define MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE  \
+    ((PSA_EXPORT_KEY_PAIR_OR_PUBLIC_MAX_SIZE > PSA_CIPHER_MAX_KEY_LENGTH) ? \
+     PSA_EXPORT_KEY_PAIR_OR_PUBLIC_MAX_SIZE : PSA_CIPHER_MAX_KEY_LENGTH)
+#endif /* !MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE*/
+
 /** \addtogroup attributes
  * @{
  */

include/psa/crypto_sizes.h

@@ -1038,6 +1038,10 @@
     PSA_KEY_EXPORT_FFDH_PUBLIC_KEY_MAX_SIZE(PSA_VENDOR_FFDH_MAX_KEY_BITS)
 #endif
 
+#define PSA_EXPORT_KEY_PAIR_OR_PUBLIC_MAX_SIZE \
+    ((PSA_EXPORT_KEY_PAIR_MAX_SIZE > PSA_EXPORT_PUBLIC_KEY_MAX_SIZE) ? \
+     PSA_EXPORT_KEY_PAIR_MAX_SIZE : PSA_EXPORT_PUBLIC_KEY_MAX_SIZE)
+
 /** Sufficient output buffer size for psa_raw_key_agreement().
  *
  * This macro returns a compile-time constant if its arguments are
@@ -1085,6 +1089,27 @@
 #define PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE    PSA_BITS_TO_BYTES(PSA_VENDOR_FFDH_MAX_KEY_BITS)
 #endif
 
+/** Maximum key length for ciphers.
+ *
+ * Since there is no additional PSA_WANT_xxx symbol to specifiy the size of
+ * the key once a cipher is enabled (as it happens for asymmetric keys for
+ * example), the maximum key length is taken into account for each cipher.
+ * The resulting value will be the maximum cipher's key length given depending
+ * on which ciphers are enabled.
+ *
+ * Note: max value for AES used below would be doubled if XTS were enabled, but
+ *       this mode is currently not supported in Mbed TLS implementation of PSA
+ *       APIs.
+ */
+#if (defined(PSA_WANT_KEY_TYPE_AES) || defined(PSA_WANT_KEY_TYPE_ARIA) || \
+    defined(PSA_WANT_KEY_TYPE_CAMELLIA) || defined(PSA_WANT_KEY_TYPE_CHACHA20))
+#define PSA_CIPHER_MAX_KEY_LENGTH       32u
+#elif defined(PSA_WANT_KEY_TYPE_DES)
+#define PSA_CIPHER_MAX_KEY_LENGTH       24u
+#else
+#define PSA_CIPHER_MAX_KEY_LENGTH       0u
+#endif
+
 /** The default IV size for a cipher algorithm, in bytes.
  *
  * The IV that is generated as part of a call to #psa_cipher_encrypt() is always

library/pk.c

@@ -35,10 +35,6 @@
 #include <limits.h>
 #include <stdint.h>
 
-#define PSA_EXPORT_KEY_PAIR_OR_PUBLIC_MAX_SIZE \
-    (PSA_EXPORT_KEY_PAIR_MAX_SIZE > PSA_EXPORT_PUBLIC_KEY_MAX_SIZE) ? \
-    PSA_EXPORT_KEY_PAIR_MAX_SIZE : PSA_EXPORT_PUBLIC_KEY_MAX_SIZE
-
 /*
  * Initialise a mbedtls_pk_context
  */

library/psa_crypto.c

@@ -710,6 +710,11 @@ MBEDTLS_STATIC_TESTABLE psa_status_t psa_mac_key_can_do(
 psa_status_t psa_allocate_buffer_to_slot(psa_key_slot_t *slot,
                                          size_t buffer_length)
 {
+#if defined(MBEDTLS_PSA_STATIC_KEY_SLOTS)
+    if (buffer_length > ((size_t) MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE)) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+#else
     if (slot->key.data != NULL) {
         return PSA_ERROR_ALREADY_EXISTS;
     }
@@ -718,6 +723,7 @@ psa_status_t psa_allocate_buffer_to_slot(psa_key_slot_t *slot,
     if (slot->key.data == NULL) {
         return PSA_ERROR_INSUFFICIENT_MEMORY;
     }
+#endif
 
     slot->key.bytes = buffer_length;
     return PSA_SUCCESS;
@@ -1186,11 +1192,18 @@ static psa_status_t psa_get_and_lock_transparent_key_slot_with_policy(
 
 psa_status_t psa_remove_key_data_from_memory(psa_key_slot_t *slot)
 {
+#if defined(MBEDTLS_PSA_STATIC_KEY_SLOTS)
+    if (slot->key.bytes > 0) {
+        mbedtls_platform_zeroize(slot->key.data, MBEDTLS_PSA_STATIC_KEY_SLOT_BUFFER_SIZE);
+    }
+#else
     if (slot->key.data != NULL) {
         mbedtls_zeroize_and_free(slot->key.data, slot->key.bytes);
     }
 
     slot->key.data = NULL;
+#endif /* MBEDTLS_PSA_STATIC_KEY_SLOTS */
+
     slot->key.bytes = 0;
 
     return PSA_SUCCESS;
@@ -2109,7 +2122,7 @@ psa_status_t psa_import_key(const psa_key_attributes_t *attributes,
      * storage ( thus not in the case of importing a key in a secure element
      * with storage ( MBEDTLS_PSA_CRYPTO_SE_C ) ),we have to allocate a
      * buffer to hold the imported key material. */
-    if (slot->key.data == NULL) {
+    if (slot->key.bytes == 0) {
         if (psa_key_lifetime_is_external(attributes->lifetime)) {
             status = psa_driver_wrapper_get_key_buffer_size_from_key_data(
                 attributes, data, data_length, &storage_size);
@@ -8026,7 +8039,7 @@ psa_status_t psa_generate_key_custom(const psa_key_attributes_t *attributes,
      * storage ( thus not in the case of generating a key in a secure element
      * with storage ( MBEDTLS_PSA_CRYPTO_SE_C ) ),we have to allocate a
      * buffer to hold the generated key material. */
-    if (slot->key.data == NULL) {
+    if (slot->key.bytes == 0) {
         if (PSA_KEY_LIFETIME_GET_LOCATION(attributes->lifetime) ==
             PSA_KEY_LOCATION_LOCAL_STORAGE) {
             status = psa_validate_key_type_and_size_for_key_generation(