Merge branch 'feat/psa_its_custom_backend_v6.0' into 'release/v6.0'

Support custom storage backend for persistent PSA keys (v6.0)

See merge request espressif/esp-idf!48569

components/mbedtls/CMakeLists.txt

@@ -36,7 +36,8 @@ if(NOT ${IDF_TARGET} STREQUAL "linux")
     endif()
 endif()
 
-set(mbedtls_srcs "port/esp_mem.c")
+set(mbedtls_srcs "port/esp_mem.c"
+                 "port/psa_crypto_storage/esp_psa_key_file.c")
 set(mbedtls_include_dirs
     "port/include"
     "mbedtls/include"
@@ -57,6 +58,7 @@ endif()
 
 
 list(APPEND mbedtls_include_dirs "${COMPONENT_DIR}/port/psa_driver/include")
+list(APPEND mbedtls_include_dirs "${COMPONENT_DIR}/port/psa_crypto_storage/include")
 
 idf_component_register(SRCS "${mbedtls_srcs}"
     INCLUDE_DIRS "${mbedtls_include_dirs}"
@@ -240,6 +242,7 @@ if(NOT ${IDF_TARGET} STREQUAL "linux")
         target_link_libraries(tfpsacrypto PRIVATE "$<$<TARGET_EXISTS:idf::nvs_flash>:idf::nvs_flash>")
         # Define compile definition to indicate ESP-IDF PSA ITS implementation is available
         target_compile_definitions(tfpsacrypto PUBLIC "$<$<TARGET_EXISTS:idf::nvs_flash>:ESP_PSA_ITS_AVAILABLE>")
+        target_include_directories(tfpsacrypto PRIVATE "${COMPONENT_DIR}/port/psa_crypto_storage/include")
     else()
         # For v1: check if component is in build before adding source and linking
         idf_build_get_property(build_components BUILD_COMPONENTS)
@@ -248,6 +251,7 @@ if(NOT ${IDF_TARGET} STREQUAL "linux")
             idf_component_get_property(nvs_flash_lib nvs_flash COMPONENT_LIB)
             target_link_libraries(tfpsacrypto PRIVATE ${nvs_flash_lib})
             target_compile_definitions(tfpsacrypto PUBLIC ESP_PSA_ITS_AVAILABLE)
+            target_include_directories(tfpsacrypto PRIVATE "${COMPONENT_DIR}/port/psa_crypto_storage/include")
         endif()
     endif()
 endif()

components/mbedtls/Kconfig

@@ -30,6 +30,38 @@ menu "mbedTLS"
                 which is added through vfs component for ESP32 based targets or by
                 the host system when the target is Linux.
 
+        config MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+            bool "Enable custom storage backend for PSA ITS"
+            default n
+            help
+                Enable support for registering a custom storage backend that
+                handles PSA ITS operations for keys in a reserved UID range.
+                When enabled, users can call esp_psa_its_register_custom_backend()
+                to route storage operations for UIDs in the configured range to
+                their own implementation. UIDs outside the range continue using
+                the default NVS backend.
+
+        config MBEDTLS_PSA_ITS_CUSTOM_BACKEND_UID_MIN
+            hex "Minimum UID for custom backend range"
+            depends on MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+            default 0x30000000
+            range 0x00000001 0x3FFFFFFF
+            help
+                The minimum UID value (inclusive) that will be routed to the
+                custom storage backend. Must be within the PSA user key ID
+                range (0x00000001 - 0x3FFFFFFF).
+
+        config MBEDTLS_PSA_ITS_CUSTOM_BACKEND_UID_MAX
+            hex "Maximum UID for custom backend range"
+            depends on MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+            default 0x3FFFFFFF
+            range MBEDTLS_PSA_ITS_CUSTOM_BACKEND_UID_MIN 0x3FFFFFFF
+            help
+                The maximum UID value (inclusive) that will be routed to the
+                custom storage backend. Must be >= the minimum UID and within
+                the PSA user key ID range (0x00000001 - 0x3FFFFFFF); the lower
+                bound is enforced by Kconfig.
+
         config MBEDTLS_THREADING_C
             bool "Enable the threading abstraction layer"
             default y

components/mbedtls/port/psa_crypto_storage/esp_psa_its.c

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2025 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2025-2026 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Apache-2.0
  *
@@ -23,8 +23,28 @@
 #include "nvs_flash.h"
 #include "esp_log.h"
 
+#include "sdkconfig.h"
+
+#if CONFIG_MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+#include "esp_psa_its.h"
+#endif
+
 static const char *TAG = "esp_psa_its";
 
+#if CONFIG_MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+/* Single registered custom backend (NULL when none registered) */
+static const esp_psa_its_custom_ops_t *s_custom_ops = NULL;
+
+/**
+ * Check if a UID falls within the custom backend range.
+ */
+static inline bool uid_in_custom_range(psa_storage_uid_t uid)
+{
+    return (uid >= CONFIG_MBEDTLS_PSA_ITS_CUSTOM_BACKEND_UID_MIN &&
+            uid <= CONFIG_MBEDTLS_PSA_ITS_CUSTOM_BACKEND_UID_MAX);
+}
+#endif /* CONFIG_MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND */
+
 /* NVS namespace for PSA ITS */
 #define PSA_ITS_NVS_NAMESPACE "psa_its"
 
@@ -106,6 +126,15 @@ psa_status_t psa_its_get_info(psa_storage_uid_t uid,
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
+#if CONFIG_MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+    if (uid_in_custom_range(uid)) {
+        if (s_custom_ops == NULL || s_custom_ops->get_info == NULL) {
+            return PSA_ERROR_STORAGE_FAILURE;
+        }
+        return s_custom_ops->get_info(s_custom_ops->ctx, uid, p_info);
+    }
+#endif
+
     /* Convert UID to NVS key */
     uid_to_nvs_key(uid, nvs_key);
 
@@ -188,6 +217,16 @@ psa_status_t psa_its_get(psa_storage_uid_t uid,
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
+#if CONFIG_MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+    if (uid_in_custom_range(uid)) {
+        if (s_custom_ops == NULL || s_custom_ops->get == NULL) {
+            return PSA_ERROR_STORAGE_FAILURE;
+        }
+        return s_custom_ops->get(s_custom_ops->ctx, uid, data_offset,
+                                 data_length, p_data, p_data_length);
+    }
+#endif
+
     /* Convert UID to NVS key */
     uid_to_nvs_key(uid, nvs_key);
 
@@ -297,6 +336,16 @@ psa_status_t psa_its_set(psa_storage_uid_t uid,
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
+#if CONFIG_MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+    if (uid_in_custom_range(uid)) {
+        if (s_custom_ops == NULL || s_custom_ops->set == NULL) {
+            return PSA_ERROR_STORAGE_FAILURE;
+        }
+        return s_custom_ops->set(s_custom_ops->ctx, uid, data_length,
+                                 p_data, create_flags);
+    }
+#endif
+
     /* Convert UID to NVS key */
     uid_to_nvs_key(uid, nvs_key);
 
@@ -383,6 +432,15 @@ psa_status_t psa_its_remove(psa_storage_uid_t uid)
     psa_its_entry_t *existing_entry = NULL;
     psa_status_t status = PSA_ERROR_STORAGE_FAILURE;
 
+#if CONFIG_MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+    if (uid_in_custom_range(uid)) {
+        if (s_custom_ops == NULL || s_custom_ops->remove == NULL) {
+            return PSA_ERROR_STORAGE_FAILURE;
+        }
+        return s_custom_ops->remove(s_custom_ops->ctx, uid);
+    }
+#endif
+
     /* Convert UID to NVS key */
     uid_to_nvs_key(uid, nvs_key);
 
@@ -451,3 +509,30 @@ exit:
     nvs_close(handle);
     return status;
 }
+
+#if CONFIG_MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND
+psa_status_t esp_psa_its_register_custom_backend(const esp_psa_its_custom_ops_t *ops)
+{
+    if (ops == NULL || ops->set == NULL || ops->get == NULL ||
+        ops->get_info == NULL || ops->remove == NULL) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    if (s_custom_ops != NULL) {
+        return PSA_ERROR_NOT_PERMITTED;
+    }
+
+    s_custom_ops = ops;
+    return PSA_SUCCESS;
+}
+
+psa_status_t esp_psa_its_unregister_custom_backend(void)
+{
+    if (s_custom_ops == NULL) {
+        return PSA_ERROR_DOES_NOT_EXIST;
+    }
+
+    s_custom_ops = NULL;
+    return PSA_SUCCESS;
+}
+#endif

components/mbedtls/port/psa_crypto_storage/esp_psa_key_file.c

@@ -0,0 +1,81 @@
+/*
+ * SPDX-FileCopyrightText: 2026 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Thin wrapper around the internal tf-psa-crypto storage helpers
+ * (psa_format_key_data_for_storage / psa_parse_key_data_from_storage and
+ * the fork-side psa_persistent_key_storage_blob_size).
+ *
+ * The PSA persistent key file format is documented and stable per the
+ * Mbed TLS storage specification; the C functions that produce/consume it
+ * are not (they're in tf-psa-crypto/core/, not the public install root).
+ * This wrapper isolates user code from the internal-symbol coupling: if
+ * the upstream functions change shape or get removed, only this file
+ * needs updating — esp_psa_key_file_{size,pack,unpack} stay the same.
+ */
+
+#include "esp_psa_key_file.h"
+
+#include "psa/crypto.h"
+#include "psa_crypto_storage.h"
+
+size_t esp_psa_key_file_size(size_t material_len)
+{
+    return psa_persistent_key_storage_blob_size(material_len);
+}
+
+psa_status_t esp_psa_key_file_pack(const psa_key_attributes_t *attrs,
+                                   const uint8_t *material,
+                                   size_t material_len,
+                                   uint8_t *out_buf,
+                                   size_t out_buf_size,
+                                   size_t *out_len)
+{
+    if (attrs == NULL || out_buf == NULL || out_len == NULL ||
+        (material == NULL && material_len != 0)) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    const size_t total = psa_persistent_key_storage_blob_size(material_len);
+    if (out_buf_size < total) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+
+    psa_format_key_data_for_storage(material, material_len, attrs, out_buf);
+
+    *out_len = total;
+    return PSA_SUCCESS;
+}
+
+psa_status_t esp_psa_key_file_unpack(const uint8_t *blob,
+                                     size_t blob_len,
+                                     psa_key_attributes_t *attrs,
+                                     const uint8_t **material,
+                                     size_t *material_len)
+{
+    if (blob == NULL || attrs == NULL || material == NULL || material_len == NULL) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    uint8_t *upstream_material = NULL;
+    size_t upstream_material_len = 0;
+    psa_status_t status = psa_parse_key_data_from_storage(blob, blob_len,
+                                                          &upstream_material,
+                                                          &upstream_material_len,
+                                                          attrs);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    /* Upstream allocates and copies the material into a fresh buffer.
+     * Discard the copy and return a pointer into the caller's blob — the
+     * material section starts immediately after the fixed-size header, so
+     * the bytes are identical. Preserves the zero-copy public contract. */
+    *material = (upstream_material_len == 0) ? NULL
+                                              : blob + (blob_len - upstream_material_len);
+    *material_len = upstream_material_len;
+
+    psa_free_persistent_key_data(upstream_material, upstream_material_len);
+    return PSA_SUCCESS;
+}

components/mbedtls/port/psa_crypto_storage/include/esp_psa_its.h

@@ -0,0 +1,126 @@
+/*
+ * SPDX-FileCopyrightText: 2026 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * PSA ITS custom storage backend API.
+ *
+ * When CONFIG_MBEDTLS_PSA_ITS_CUSTOM_STORAGE_BACKEND is enabled, users can register
+ * a custom storage backend for a reserved range of PSA key IDs. UIDs within
+ * the configured range are routed to the registered backend; all other UIDs
+ * continue using the default NVS backend.
+ */
+
+#pragma once
+
+#include <stdint.h>
+#include <stddef.h>
+#include "psa/internal_trusted_storage.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @brief Custom storage backend operations for PSA ITS.
+ *
+ * Implement this structure to provide a custom storage backend for PSA ITS
+ * keys in the configured UID range. Callback signatures mirror the PSA ITS
+ * API with an additional user-provided context pointer.
+ *
+ * The user's implementation is opaque to the framework — it may internally
+ * route to any number of storage systems (FATFS, SPIFFS, secure elements,
+ * etc.) based on the UID or any other criteria.
+ */
+typedef struct {
+    /**
+     * @brief Store data for the given UID.
+     *
+     * @param ctx           User-provided context pointer
+     * @param uid           Storage UID (key ID)
+     * @param data_length   Length of the data to store
+     * @param p_data        Pointer to the data buffer
+     * @param create_flags  Storage flags (e.g., PSA_STORAGE_FLAG_WRITE_ONCE)
+     * @return PSA status code
+     */
+    psa_status_t (*set)(void *ctx,
+                        psa_storage_uid_t uid,
+                        uint32_t data_length,
+                        const void *p_data,
+                        psa_storage_create_flags_t create_flags);
+
+    /**
+     * @brief Retrieve data for the given UID.
+     *
+     * @param ctx           User-provided context pointer
+     * @param uid           Storage UID (key ID)
+     * @param data_offset   Byte offset within the stored data
+     * @param data_length   Number of bytes to retrieve
+     * @param p_data        Output buffer for the data
+     * @param p_data_length On success, set to the number of bytes written
+     * @return PSA status code
+     */
+    psa_status_t (*get)(void *ctx,
+                        psa_storage_uid_t uid,
+                        uint32_t data_offset,
+                        uint32_t data_length,
+                        void *p_data,
+                        size_t *p_data_length);
+
+    /**
+     * @brief Retrieve metadata for the given UID.
+     *
+     * @param ctx    User-provided context pointer
+     * @param uid    Storage UID (key ID)
+     * @param p_info Output structure for size and flags
+     * @return PSA status code
+     */
+    psa_status_t (*get_info)(void *ctx,
+                             psa_storage_uid_t uid,
+                             struct psa_storage_info_t *p_info);
+
+    /**
+     * @brief Remove data for the given UID.
+     *
+     * @param ctx  User-provided context pointer
+     * @param uid  Storage UID (key ID)
+     * @return PSA status code
+     */
+    psa_status_t (*remove)(void *ctx,
+                           psa_storage_uid_t uid);
+
+    /** User-provided context pointer, passed as the first argument to all callbacks. */
+    void *ctx;
+} esp_psa_its_custom_ops_t;
+
+/**
+ * @brief Register a custom storage backend for the configured UID range.
+ *
+ * Only one custom backend may be registered at a time. The UID range is
+ * determined by CONFIG_MBEDTLS_PSA_ITS_CUSTOM_BACKEND_UID_MIN (inclusive) and
+ * CONFIG_MBEDTLS_PSA_ITS_CUSTOM_BACKEND_UID_MAX (inclusive).
+ *
+ * @param ops Pointer to the operations structure. Must remain valid for
+ *            the lifetime of the registration. All four function pointers
+ *            (set, get, get_info, remove) must be non-NULL.
+ *
+ * @return PSA_SUCCESS on success
+ * @return PSA_ERROR_INVALID_ARGUMENT if ops or any callback is NULL
+ * @return PSA_ERROR_NOT_PERMITTED if a backend is already registered
+ */
+psa_status_t esp_psa_its_register_custom_backend(const esp_psa_its_custom_ops_t *ops);
+
+/**
+ * @brief Unregister the currently registered custom storage backend.
+ *
+ * After this call, UIDs in the custom range will return
+ * PSA_ERROR_STORAGE_FAILURE until a new backend is registered.
+ *
+ * @return PSA_SUCCESS on success
+ * @return PSA_ERROR_DOES_NOT_EXIST if no backend is registered
+ */
+psa_status_t esp_psa_its_unregister_custom_backend(void);
+
+#ifdef __cplusplus
+}
+#endif

components/mbedtls/port/psa_crypto_storage/include/esp_psa_key_file.h

@@ -0,0 +1,92 @@
+/*
+ * SPDX-FileCopyrightText: 2026 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Pack/unpack PSA persistent-key file blobs in the format documented at
+ * tf-psa-crypto/docs/architecture/mbed-crypto-storage-specification.md
+ * ("Key file format for Mbed TLS 2.25.0"), which is the stable storage
+ * contract for Mbed TLS 2.25 through < 4.
+ *
+ * Custom PSA ITS backends that synthesize blobs on read, or that strip the
+ * header on write to save space, use these helpers to convert between PSA
+ * attributes + raw key material and the documented blob format without
+ * depending on tf-psa-crypto internal functions.
+ */
+
+#pragma once
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include "psa/crypto.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @brief Total blob size produced by esp_psa_key_file_pack() for a given key
+ *        material length (i.e. fixed header size + @p material_len).
+ */
+size_t esp_psa_key_file_size(size_t material_len);
+
+/**
+ * @brief Pack key attributes and material into a key file blob.
+ *
+ * @param[in]  attrs         Source attributes (lifetime, type, bits, policy).
+ * @param[in]  material      Key material. For transparent keys, the
+ *                           psa_export_key() output. For opaque keys, the
+ *                           driver-specific opaque blob.
+ * @param[in]  material_len  Length of @p material in bytes.
+ * @param[out] out_buf       Destination buffer.
+ * @param[in]  out_buf_size  Size of @p out_buf. Must be at least
+ *                           esp_psa_key_file_size(material_len).
+ * @param[out] out_len       Set to the number of bytes written.
+ *
+ * @retval PSA_SUCCESS
+ * @retval PSA_ERROR_INVALID_ARGUMENT  Null pointer.
+ * @retval PSA_ERROR_BUFFER_TOO_SMALL  @p out_buf_size is too small.
+ */
+psa_status_t esp_psa_key_file_pack(const psa_key_attributes_t *attrs,
+                                   const uint8_t *material,
+                                   size_t material_len,
+                                   uint8_t *out_buf,
+                                   size_t out_buf_size,
+                                   size_t *out_len);
+
+/**
+ * @brief Parse a key file blob into key attributes and material.
+ *
+ * The returned @p material pointer points into @p blob and is valid only as
+ * long as @p blob itself is. No allocation occurs.
+ *
+ * Attribute fields described in the blob (lifetime, type, bits, usage,
+ * algorithm, enrollment algorithm) are set on @p attrs on success. Other
+ * fields (notably the key id) are left untouched, so the caller can set the
+ * id beforehand if needed.
+ *
+ * @param[in]  blob          Blob to parse.
+ * @param[in]  blob_len      Length of @p blob in bytes.
+ * @param[out] attrs         Filled with parsed attributes on success.
+ * @param[out] material      Set to point to the material section within @p blob,
+ *                           or NULL if the blob declares zero-length material.
+ * @param[out] material_len  Length of the material section in bytes.
+ *
+ * @retval PSA_SUCCESS
+ * @retval PSA_ERROR_INVALID_ARGUMENT  Null pointer, or @p blob_len smaller than
+ *                                     ESP_PSA_KEY_FILE_HEADER_SIZE.
+ * @retval PSA_ERROR_DATA_INVALID      Bad magic, unsupported version, or the
+ *                                     declared material length does not match
+ *                                     the blob size (the spec rejects trailing
+ *                                     data on load).
+ */
+psa_status_t esp_psa_key_file_unpack(const uint8_t *blob,
+                                     size_t blob_len,
+                                     psa_key_attributes_t *attrs,
+                                     const uint8_t **material,
+                                     size_t *material_len);
+
+#ifdef __cplusplus
+}
+#endif

components/mbedtls/port/psa_driver/esp_rsa_ds/psa_crypto_driver_esp_rsa_ds.c

@@ -349,7 +349,7 @@ void esp_rsa_ds_release_ds_lock(void)
     }
 }
 
-static int esp_rsa_ds_validate_opaque_key(const esp_rsa_ds_opaque_key_t *opaque_key)
+static psa_status_t esp_rsa_ds_validate_opaque_key(const esp_rsa_ds_opaque_key_t *opaque_key)
 {
     if (opaque_key == NULL) {
         return PSA_ERROR_INVALID_ARGUMENT;
@@ -394,6 +394,111 @@ static int esp_rsa_ds_validate_opaque_key(const esp_rsa_ds_opaque_key_t *opaque_
     return PSA_SUCCESS;
 }
 
+/**
+ * Serialize an already-validated opaque key into the persistent inline
+ * storage layout (eFuse or Key Manager, selected from key_recovery_info).
+ * Internal helper — does not validate inputs.
+ */
+static psa_status_t rsa_ds_format_persistent_key_buffer_internal(
+    const esp_rsa_ds_opaque_key_t *opaque_key,
+    uint8_t *buf, size_t buf_size, size_t *out_len)
+{
+#if SOC_KEY_MANAGER_SUPPORTED
+    if (opaque_key->key_recovery_info) {
+        if (buf_size < sizeof(esp_rsa_ds_km_key_storage_t)) {
+            return PSA_ERROR_BUFFER_TOO_SMALL;
+        }
+        esp_rsa_ds_km_key_storage_t *storage = (esp_rsa_ds_km_key_storage_t *)buf;
+        memset(storage, 0, sizeof(*storage));
+        storage->metadata.version = ESP_RSA_DS_KEY_STORAGE_VERSION_V1;
+        storage->metadata.key_source = ESP_RSA_DS_KEY_SOURCE_KEY_MGR;
+        storage->rsa_length_bits = opaque_key->ds_data_ctx->rsa_length_bits;
+        memcpy(&storage->ds_data, opaque_key->ds_data_ctx->esp_ds_data, sizeof(esp_ds_data_t));
+        memcpy(&storage->key_recovery_info, opaque_key->key_recovery_info,
+               sizeof(esp_key_mgr_key_recovery_info_t));
+        *out_len = sizeof(esp_rsa_ds_km_key_storage_t);
+        return PSA_SUCCESS;
+    }
+#endif /* SOC_KEY_MANAGER_SUPPORTED */
+
+    if (buf_size < sizeof(esp_rsa_ds_efuse_key_storage_t)) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+    esp_rsa_ds_efuse_key_storage_t *storage = (esp_rsa_ds_efuse_key_storage_t *)buf;
+    memset(storage, 0, sizeof(*storage));
+    storage->metadata.version = ESP_RSA_DS_KEY_STORAGE_VERSION_V1;
+    storage->metadata.key_source = ESP_RSA_DS_KEY_SOURCE_EFUSE;
+    storage->efuse_key_id = opaque_key->ds_data_ctx->efuse_key_id;
+    storage->rsa_length_bits = opaque_key->ds_data_ctx->rsa_length_bits;
+    memcpy(&storage->ds_data, opaque_key->ds_data_ctx->esp_ds_data, sizeof(esp_ds_data_t));
+    *out_len = sizeof(esp_rsa_ds_efuse_key_storage_t);
+    return PSA_SUCCESS;
+}
+
+size_t esp_rsa_ds_persistent_key_buffer_size(const esp_rsa_ds_opaque_key_t *opaque_key)
+{
+    if (opaque_key == NULL) {
+        return 0;
+    }
+    return esp_rsa_ds_get_storage_size(opaque_key, true);
+}
+
+psa_status_t esp_rsa_ds_format_persistent_key_buffer(const esp_rsa_ds_opaque_key_t *opaque_key,
+                                                 uint8_t *buf, size_t buf_size,
+                                                 size_t *out_len)
+{
+    if (opaque_key == NULL || buf == NULL || out_len == NULL) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    psa_status_t ret = esp_rsa_ds_validate_opaque_key(opaque_key);
+    if (ret != PSA_SUCCESS) {
+        return ret;
+    }
+
+    return rsa_ds_format_persistent_key_buffer_internal(opaque_key, buf, buf_size, out_len);
+}
+
+psa_status_t esp_rsa_ds_parse_persistent_key_buffer(const uint8_t *buf, size_t buf_len,
+                                                    esp_rsa_ds_opaque_key_t *out)
+{
+    if (buf == NULL || out == NULL || out->ds_data_ctx == NULL) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    esp_rsa_ds_key_source_t key_source = ESP_RSA_DS_KEY_SOURCE_EFUSE;
+    uint16_t rsa_length_bits = 0;
+    const esp_ds_data_t *ds_data = NULL;
+    hmac_key_id_t hmac_id = 0;
+#if SOC_KEY_MANAGER_SUPPORTED
+    esp_key_mgr_key_recovery_info_t *km_ri = NULL;
+#endif
+
+    psa_status_t status = esp_rsa_ds_extract_storage(
+        buf, buf_len, /* is_persistent = */ true,
+        &key_source, &rsa_length_bits, &ds_data, &hmac_id
+#if SOC_KEY_MANAGER_SUPPORTED
+        , &km_ri
+#endif
+    );
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    /* Aliases into buf — caller must not write through these pointers and
+     * must keep buf alive for as long as out is used. */
+    out->ds_data_ctx->esp_ds_data     = (esp_ds_data_t *)ds_data;
+    out->ds_data_ctx->efuse_key_id    = (uint8_t)hmac_id;
+    out->ds_data_ctx->rsa_length_bits = rsa_length_bits;
+#if SOC_KEY_MANAGER_SUPPORTED
+    out->key_recovery_info = km_ri;
+#else
+    (void)key_source;
+#endif
+
+    return PSA_SUCCESS;
+}
+
 psa_status_t esp_rsa_ds_opaque_sign_hash_start(
     esp_rsa_ds_opaque_sign_hash_operation_t *operation,
     const psa_key_attributes_t *attributes,
@@ -680,7 +785,7 @@ psa_status_t esp_rsa_ds_opaque_import_key(
     }
 
     const esp_rsa_ds_opaque_key_t *opaque_key = (const esp_rsa_ds_opaque_key_t *)data;
-    int ret = esp_rsa_ds_validate_opaque_key(opaque_key);
+    psa_status_t ret = esp_rsa_ds_validate_opaque_key(opaque_key);
     if (ret != PSA_SUCCESS) {
         return ret;
     }
@@ -711,36 +816,10 @@ psa_status_t esp_rsa_ds_opaque_import_key(
         *key_buffer_length = sizeof(esp_rsa_ds_volatile_key_storage_t);
     } else {
         /* Persistent: deep-copy all data into self-contained storage struct */
-#if SOC_KEY_MANAGER_SUPPORTED
-        if (key_source == ESP_RSA_DS_KEY_SOURCE_KEY_MGR) {
-            if (key_buffer_size < sizeof(esp_rsa_ds_km_key_storage_t)) {
-                return PSA_ERROR_BUFFER_TOO_SMALL;
-            }
-
-            esp_rsa_ds_km_key_storage_t *storage = (esp_rsa_ds_km_key_storage_t *)key_buffer;
-            storage->metadata.version = ESP_RSA_DS_KEY_STORAGE_VERSION_V1;
-            storage->metadata.key_source = ESP_RSA_DS_KEY_SOURCE_KEY_MGR;
-            storage->rsa_length_bits = opaque_key->ds_data_ctx->rsa_length_bits;
-            memcpy(&storage->ds_data, opaque_key->ds_data_ctx->esp_ds_data, sizeof(esp_ds_data_t));
-            memcpy(&storage->key_recovery_info, opaque_key->key_recovery_info,
-                   sizeof(esp_key_mgr_key_recovery_info_t));
-            *key_buffer_length = sizeof(esp_rsa_ds_km_key_storage_t);
-        } else
-#endif /* SOC_KEY_MANAGER_SUPPORTED */
-        {
-            if (key_buffer_size < sizeof(esp_rsa_ds_efuse_key_storage_t)) {
-                return PSA_ERROR_BUFFER_TOO_SMALL;
-            }
-
-            esp_rsa_ds_efuse_key_storage_t *storage = (esp_rsa_ds_efuse_key_storage_t *)key_buffer;
-            storage->metadata.version = ESP_RSA_DS_KEY_STORAGE_VERSION_V1;
-            storage->metadata.key_source = ESP_RSA_DS_KEY_SOURCE_EFUSE;
-            storage->efuse_key_id = opaque_key->ds_data_ctx->efuse_key_id;
-            storage->reserved = 0;
-            storage->rsa_length_bits = opaque_key->ds_data_ctx->rsa_length_bits;
-            memset(storage->reserved2, 0, sizeof(storage->reserved2));
-            memcpy(&storage->ds_data, opaque_key->ds_data_ctx->esp_ds_data, sizeof(esp_ds_data_t));
-            *key_buffer_length = sizeof(esp_rsa_ds_efuse_key_storage_t);
+        psa_status_t status = rsa_ds_format_persistent_key_buffer_internal(
+            opaque_key, key_buffer, key_buffer_size, key_buffer_length);
+        if (status != PSA_SUCCESS) {
+            return status;
         }
     }
 

components/mbedtls/port/psa_driver/include/psa_crypto_driver_esp_rsa_ds.h

@@ -38,6 +38,89 @@ extern "C" {
         PSA_KEY_PERSISTENCE_VOLATILE, \
         PSA_KEY_LOCATION_ESP_RSA_DS)
 
+/**
+ * @brief Buffer size needed to serialize an ESP-RSA DS key into the persistent
+ *        storage layout used by this driver.
+ *
+ * The size depends on the key source: eFuse-sourced keys and Key Manager-sourced
+ * keys (KM-capable SoCs only) use different inline storage structs. The source
+ * is inferred from @p opaque_key — the Key Manager layout is selected when
+ * @p opaque_key->key_recovery_info is non-NULL.
+ *
+ * @param opaque_key  User-facing opaque key (must be non-NULL).
+ * @return Storage buffer size in bytes, or 0 if @p opaque_key is NULL.
+ */
+size_t esp_rsa_ds_persistent_key_buffer_size(const esp_rsa_ds_opaque_key_t *opaque_key);
+
+/**
+ * @brief Serialize an ESP-RSA DS key into the persistent storage layout that
+ *        this driver expects when loading a persistent key from PSA storage.
+ *
+ * Custom PSA ITS backends that synthesize persistent RSA-DS key blobs at
+ * read time can use this helper to produce the @c key_data payload, then wrap
+ * it with esp_psa_its_pack_key_blob() to build the full PSA persistent key blob.
+ *
+ * The output layout (eFuse vs Key Manager) is selected automatically from
+ * @p opaque_key, mirroring the import path. Validation of @p opaque_key is
+ * performed before serialization.
+ *
+ * @param opaque_key  User-facing opaque key.
+ * @param buf         Output buffer (caller-allocated, sized via
+ *                    esp_rsa_ds_persistent_key_buffer_size()).
+ * @param buf_size    Size of @p buf in bytes.
+ * @param[out] out_len Bytes written to @p buf on success.
+ *
+ * @return PSA_SUCCESS on success
+ * @return PSA_ERROR_INVALID_ARGUMENT if any required input is NULL or the
+ *                                    opaque key fields are invalid
+ * @return PSA_ERROR_BUFFER_TOO_SMALL if @p buf_size is insufficient
+ */
+psa_status_t esp_rsa_ds_format_persistent_key_buffer(const esp_rsa_ds_opaque_key_t *opaque_key,
+                                                 uint8_t *buf, size_t buf_size,
+                                                 size_t *out_len);
+
+/**
+ * @brief Parse an ESP-RSA DS persistent key buffer.
+ *
+ * Inverse of esp_rsa_ds_format_persistent_key_buffer(): validates the
+ * key-storage metadata (version + source) in @p buf, then fills @p out
+ * with the same opaque-key shape the format path consumes.
+ *
+ * The caller owns the storage for @p out, including the @c esp_ds_data_ctx_t
+ * it points to (@p out->ds_data_ctx must be non-NULL before the call).
+ * On success, scalar fields (@c efuse_key_id, @c rsa_length_bits) are filled
+ * into the caller's @c esp_ds_data_ctx_t, and the @c esp_ds_data pointer
+ * — together with @c out->key_recovery_info on KM-capable SoCs — aliases
+ * into @p buf. Those pointers remain valid only for as long as @p buf is,
+ * and must not be written through.
+ *
+ * The key source (eFuse vs Key Manager) is conveyed implicitly: on success,
+ * @c out->key_recovery_info is non-NULL iff the buffer was produced from a
+ * Key-Manager-backed key. This mirrors how the format path discriminates
+ * via the same field.
+ *
+ * Custom PSA ITS backends that accept writes (translating PSA-formatted
+ * blobs handed to psa_its_set() into a native storage format) can use this
+ * helper after first stripping the PSA wrapper with esp_psa_its_unpack_key_blob().
+ *
+ * @param buf       Input buffer (as produced by esp_rsa_ds_format_persistent_key_buffer()
+ *                  or written by the driver's import path).
+ * @param buf_len   Length of @p buf in bytes.
+ * @param[in,out] out  Opaque key to fill. @p out->ds_data_ctx must point to a
+ *                  caller-allocated @c esp_ds_data_ctx_t.
+ *
+ * @return PSA_SUCCESS on success
+ * @return PSA_ERROR_INVALID_ARGUMENT if @p buf, @p out, or @p out->ds_data_ctx
+ *                                    is NULL, or @p buf_len is too small, or
+ *                                    the buffer is internally inconsistent
+ *                                    (e.g. invalid RSA length)
+ * @return PSA_ERROR_DATA_INVALID     if the metadata version or key source is
+ *                                    unrecognized, or the embedded esp_ds_data_t
+ *                                    does not match the declared key length
+ */
+psa_status_t esp_rsa_ds_parse_persistent_key_buffer(const uint8_t *buf, size_t buf_len,
+                                                    esp_rsa_ds_opaque_key_t *out);
+
 /**
  * @brief Start the RSA DS opaque sign hash operation
  *