From 6b90afb2bba12487253e0fd683ac4b95aeeef08f 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 12:20:13 +0100
Subject: [PATCH] PK: ensure we test with keys on the stack or heap
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Manuel Pégourié-Gonnard <manuel.pegourie-gonnard@arm.com>
---
 library/pk.c                                  | 21 +------------------
 library/pk_internal.h                         | 21 +++++++++++++++++++
 ...test_suite_config.crypto_combinations.data |  8 +++++++
 tests/suites/test_suite_config.function       |  2 +-
 4 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/library/pk.c b/library/pk.c
index 55e14bf6fc..e34cba8937 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -35,26 +35,7 @@
 #include <limits.h>
 #include <stdint.h>
 
-/*
- * We're trying to statisfy two kinds of users:
- * - those who don't want to use the heap;
- * - those who can't afford large stack buffers.
- *
- * The current compromise is that if ECC is the only key type supported in PK,
- * then we export keys on the stack, and otherwise we use the heap.
- */
-#if !defined(MBEDTLS_RSA_C)
-#define PK_EXPORT_KEYS_ON_THE_STACK
-#endif
-
-#if defined(PK_EXPORT_KEYS_ON_THE_STACK)
-/* We know for ECC, pubkey are longer than privkeys, but double check */
-#define PK_EXPORT_KEY_STACK_BUFFER_SIZE  MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH
-#if MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH > PK_EXPORT_KEY_STACK_BUFFER_SIZE
-#undef PK_EXPORT_KEY_STACK_BUFFER_SIZE
-#define PK_EXPORT_KEY_STACK_BUFFER_SIZE  MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH
-#endif
-#else
+#if !defined(PK_EXPORT_KEYS_ON_THE_STACK)
 #include "mbedtls/platform.h" // for calloc/free
 #endif
 
diff --git a/library/pk_internal.h b/library/pk_internal.h
index e86a3a09d2..4d6081cbe2 100644
--- a/library/pk_internal.h
+++ b/library/pk_internal.h
@@ -44,6 +44,27 @@
 #define PEM_BEGIN_ENCRYPTED_PRIVATE_KEY_PKCS8 "-----BEGIN ENCRYPTED PRIVATE KEY-----"
 #define PEM_END_ENCRYPTED_PRIVATE_KEY_PKCS8   "-----END ENCRYPTED PRIVATE KEY-----"
 
+/*
+ * We're trying to statisfy two kinds of users:
+ * - those who don't want to use the heap;
+ * - those who can't afford large stack buffers.
+ *
+ * The current compromise is that if ECC is the only key type supported in PK,
+ * then we export keys on the stack, and otherwise we use the heap.
+ */
+#if !defined(MBEDTLS_RSA_C)
+#define PK_EXPORT_KEYS_ON_THE_STACK
+#endif
+
+#if defined(PK_EXPORT_KEYS_ON_THE_STACK)
+/* We know for ECC, pubkey are longer than privkeys, but double check */
+#define PK_EXPORT_KEY_STACK_BUFFER_SIZE  MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH
+#if MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH > PK_EXPORT_KEY_STACK_BUFFER_SIZE
+#undef PK_EXPORT_KEY_STACK_BUFFER_SIZE
+#define PK_EXPORT_KEY_STACK_BUFFER_SIZE  MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH
+#endif
+#endif
+
 #if defined(MBEDTLS_PK_HAVE_ECC_KEYS) && !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
 /**
  * Public function mbedtls_pk_ec() can be used to get direct access to the
diff --git a/tests/suites/test_suite_config.crypto_combinations.data b/tests/suites/test_suite_config.crypto_combinations.data
index 9cc09ec895..0f8585a519 100644
--- a/tests/suites/test_suite_config.crypto_combinations.data
+++ b/tests/suites/test_suite_config.crypto_combinations.data
@@ -12,3 +12,11 @@ pass:
 Config: ECC: Montgomery curves only
 depends_on:!MBEDTLS_ECP_SHORT_WEIERSTRASS_ENABLED:MBEDTLS_ECP_MONTGOMERY_ENABLED
 pass:
+
+Config: PK export keys on the stack (ECC is the only PK key type)
+depends_on:PK_EXPORT_KEYS_ON_THE_STACK
+pass:
+
+Config: PK export keys on the heap (ECC is not the only PK key type)
+depends_on:!PK_EXPORT_KEYS_ON_THE_STACK
+pass:
diff --git a/tests/suites/test_suite_config.function b/tests/suites/test_suite_config.function
index 9e9dd01990..8c68c0c5d3 100644
--- a/tests/suites/test_suite_config.function
+++ b/tests/suites/test_suite_config.function
@@ -1,5 +1,5 @@
 /* BEGIN_HEADER */
-
+#include "pk_internal.h"
 /* END_HEADER */
 
 /* BEGIN_CASE */