From 925341971dba653a811055af163cb9e75d5b0f1d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20P=C3=A9gouri=C3=A9-Gonnard?=
 <manuel.pegourie-gonnard@arm.com>
Date: Thu, 12 Feb 2026 11:52:19 +0100
Subject: [PATCH] PK: validate type upfront when copying from PSA
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The type was validated near the end of the function when importing, but
if makes more sense to validate upfront before we possibly allocate a
buffer, export the key to it etc.

This also guarantees a sensible error value without requiring a special
case when exporting on the stack.

Signed-off-by: Manuel Pégourié-Gonnard <manuel.pegourie-gonnard@arm.com>
---
 library/pk.c | 33 ++++++++++++++++++++++-----------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/library/pk.c b/library/pk.c
index 5047fa3b30..f9b05a3a51 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -608,7 +608,8 @@ int mbedtls_pk_get_psa_attributes(const mbedtls_pk_context *pk,
  * - using a macro in multiple places results in multiple copies of the code;
  * - this function only handles key types supported in PK.
  *
- * Return 0 on unexpected types. Callers need to check for that value.
+ * WARNING: callers need to ensure the type is supported before calling this
+ * function, possibly by calling is_valid_for_pk().
  */
 static size_t pk_export_max_size(psa_key_type_t key_type, size_t bits)
 {
@@ -937,6 +938,22 @@ int mbedtls_pk_import_into_psa(const mbedtls_pk_context *pk,
     }
 }
 
+static int is_valid_for_pk(psa_key_type_t key_type)
+{
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+    if (PSA_KEY_TYPE_IS_ECC(key_type)) {
+        return 1;
+    }
+#endif
+#if defined(MBEDTLS_RSA_C)
+    if (key_type == PSA_KEY_TYPE_RSA_PUBLIC_KEY ||
+        key_type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
+        return 1;
+    }
+#endif
+    return 0;
+}
+
 static int copy_from_psa(mbedtls_svc_key_id_t key_id,
                          mbedtls_pk_context *pk,
                          int public_only)
@@ -965,6 +982,10 @@ static int copy_from_psa(mbedtls_svc_key_id_t key_id,
     }
 
     key_type = psa_get_key_type(&key_attr);
+    if (!is_valid_for_pk(key_type)) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
     if (public_only) {
         key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(key_type);
     }
@@ -972,20 +993,10 @@ static int copy_from_psa(mbedtls_svc_key_id_t key_id,
 
 #if defined(PK_HAVE_KEYS_LARGER_THAN_ECC)
     exp_key_size = pk_export_max_size(key_type, key_bits);
-    if (exp_key_size == 0) {
-        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    }
     exp_key = mbedtls_calloc(1, exp_key_size);
     if (exp_key == NULL) {
         return MBEDTLS_ERR_PK_ALLOC_FAILED;
     }
-#else
-    /* In case we're passed non-ECC key (API misuse), return a sensible error
-     * now. Otherwise we might get BUFFER_TOO_SMALL when exporting below, which
-     * is unlikely to be helpful to the user as the buffer is internal. */
-    if (!PSA_KEY_TYPE_IS_ECC(key_type)) {
-        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    }
 #endif
 
     if (public_only) {