Merge pull request #9752 from gilles-peskine-arm/psa_init-20241104-dev

Fix MD_PSA_INIT called before initializing some data structures
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 66f52fe..b3a84b3 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -338,7 +338,11 @@
 endif()
 
 if (NOT EXISTS "${MBEDTLS_FRAMEWORK_DIR}/CMakeLists.txt")
-    message(FATAL_ERROR "${MBEDTLS_FRAMEWORK_DIR}/CMakeLists.txt not found. Run `git submodule update --init` from the source tree to fetch the submodule contents.")
+    if (EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/.git/")
+        message(FATAL_ERROR "${MBEDTLS_FRAMEWORK_DIR}/CMakeLists.txt not found (and does appear to be a git checkout). Run `git submodule update --init` from the source tree to fetch the submodule contents.")
+    else ()
+        message(FATAL_ERROR "${MBEDTLS_FRAMEWORK_DIR}/CMakeLists.txt not found (and does not appear to be a git checkout). Please ensure you have downloaded the right archive from the release page on GitHub.")
+    endif()
 endif()
 add_subdirectory(framework)
 
diff --git a/ChangeLog.d/mbedtls_psa_ecp_generate_key-no_public_key.txt b/ChangeLog.d/mbedtls_psa_ecp_generate_key-no_public_key.txt
new file mode 100644
index 0000000..69c00e1
--- /dev/null
+++ b/ChangeLog.d/mbedtls_psa_ecp_generate_key-no_public_key.txt
@@ -0,0 +1,3 @@
+Changes
+   * Improve performance of PSA key generation with ECC keys: it no longer
+     computes the public key (which was immediately discarded). Fixes #9732.
diff --git a/Makefile b/Makefile
index 71ef1be..1b480f5 100644
--- a/Makefile
+++ b/Makefile
@@ -6,11 +6,16 @@
     ifeq (,$(wildcard framework/exported.make))
         # Use the define keyword to get a multi-line message.
         # GNU make appends ".  Stop.", so tweak the ending of our message accordingly.
-        define error_message
-$(MBEDTLS_PATH)/framework/exported.make not found.
-Run `git submodule update --init` to fetch the submodule contents.
+        ifeq (,$(wildcard .git))
+            define error_message
+${MBEDTLS_PATH}/framework/exported.make not found (and does appear to be a git checkout). Run `git submodule update --init` from the source tree to fetch the submodule contents.
 This is a fatal error
-        endef
+            endef
+        else
+            define error_message
+${MBEDTLS_PATH}/framework/exported.make not found (and does not appear to be a git checkout). Please ensure you have downloaded the right archive from the release page on GitHub.
+            endef
+        endif
         $(error $(error_message))
     endif
     include framework/exported.make
diff --git a/framework b/framework
index d68446c..5560984 160000
--- a/framework
+++ b/framework
@@ -1 +1 @@
-Subproject commit d68446c9da02e536279a7aaa5a3c9850742ba30c
+Subproject commit 55609845504ce77f3714795785282456444967c8
diff --git a/tests/scripts/analyze_outcomes.py b/tests/scripts/analyze_outcomes.py
index f19d4b6..f6f7d87 100755
--- a/tests/scripts/analyze_outcomes.py
+++ b/tests/scripts/analyze_outcomes.py
@@ -409,6 +409,8 @@
     IGNORED_SUITES = [
         # Modules replaced by drivers
         'ecdsa', 'ecdh', 'ecjpake',
+        # Unit tests for the built-in implementation
+        'psa_crypto_ecp',
     ]
     IGNORED_TESTS = {
         'test_suite_config': [
@@ -449,6 +451,8 @@
     IGNORED_SUITES = [
         # Modules replaced by drivers
         'ecp', 'ecdsa', 'ecdh', 'ecjpake',
+        # Unit tests for the built-in implementation
+        'psa_crypto_ecp',
     ]
     IGNORED_TESTS = {
         'test_suite_config': [
@@ -489,6 +493,8 @@
         'ecp', 'ecdsa', 'ecdh', 'ecjpake',
         'bignum_core', 'bignum_random', 'bignum_mod', 'bignum_mod_raw',
         'bignum.generated', 'bignum.misc',
+        # Unit tests for the built-in implementation
+        'psa_crypto_ecp',
     ]
     IGNORED_TESTS = {
         'test_suite_config': [
@@ -534,6 +540,8 @@
         'ecp', 'ecdsa', 'ecdh', 'ecjpake', 'dhm',
         'bignum_core', 'bignum_random', 'bignum_mod', 'bignum_mod_raw',
         'bignum.generated', 'bignum.misc',
+        # Unit tests for the built-in implementation
+        'psa_crypto_ecp',
     ]
     IGNORED_TESTS = {
         'ssl-opt': [
@@ -604,6 +612,8 @@
         'ecp', 'ecdsa', 'ecdh', 'ecjpake',
         'bignum_core', 'bignum_random', 'bignum_mod', 'bignum_mod_raw',
         'bignum.generated', 'bignum.misc',
+        # Unit tests for the built-in implementation
+        'psa_crypto_ecp',
     ]
     IGNORED_TESTS = {
         'test_suite_config': [
diff --git a/tests/suites/test_suite_psa_crypto_ecp.data b/tests/suites/test_suite_psa_crypto_ecp.data
new file mode 100644
index 0000000..ffb7a7b
--- /dev/null
+++ b/tests/suites/test_suite_psa_crypto_ecp.data
@@ -0,0 +1,82 @@
+ECC generate: unknown family (0)
+generate_key:0:256:64:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: unknown family (0xff)
+generate_key:0xff:256:64:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: SECP_R1 bad bit-size (0)
+generate_key:PSA_ECC_FAMILY_SECP_R1:0:64:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: SECP_R1 bad bit-size (512)
+generate_key:PSA_ECC_FAMILY_SECP_R1:512:64:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: SECP_R1 bad bit-size (528)
+generate_key:PSA_ECC_FAMILY_SECP_R1:528:64:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: SECP_R1 256-bit not supported
+depends_on:!MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_256
+generate_key:PSA_ECC_FAMILY_SECP_R1:256:32:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: SECP_R1 384-bit not supported
+depends_on:!MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_384
+generate_key:PSA_ECC_FAMILY_SECP_R1:384:48:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: SECP_R1 521-bit not supported
+depends_on:!MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_521
+generate_key:PSA_ECC_FAMILY_SECP_R1:521:66:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: SECP_K1 256-bit not supported
+depends_on:!MBEDTLS_PSA_BUILTIN_ECC_SECP_K1_256
+generate_key:PSA_ECC_FAMILY_SECP_K1:256:32:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: Curve25519 not supported
+depends_on:!MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_255
+generate_key:PSA_ECC_FAMILY_MONTGOMERY:255:32:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: Curve448 not supported
+depends_on:!MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_448
+generate_key:PSA_ECC_FAMILY_MONTGOMERY:448:56:PSA_ERROR_NOT_SUPPORTED
+
+ECC generate: SECP_R1 256-bit, size=31, too small
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_256
+generate_key:PSA_ECC_FAMILY_SECP_R1:256:31:PSA_ERROR_BUFFER_TOO_SMALL
+
+ECC generate: SECP_R1 256-bit, size=32, ok
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_256
+generate_key:PSA_ECC_FAMILY_SECP_R1:256:32:PSA_SUCCESS
+
+ECC generate: SECP_R1 256-bit, size=33, ok
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_256
+generate_key:PSA_ECC_FAMILY_SECP_R1:256:33:PSA_SUCCESS
+
+ECC generate: SECP_R1 521-bit, size=65, too small
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_521
+generate_key:PSA_ECC_FAMILY_SECP_R1:521:65:PSA_ERROR_BUFFER_TOO_SMALL
+
+ECC generate: SECP_R1 521-bit, size=66, ok
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_SECP_R1_521
+generate_key:PSA_ECC_FAMILY_SECP_R1:521:66:PSA_SUCCESS
+
+ECC generate: Curve25519, size=31, too small
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_255
+generate_key:PSA_ECC_FAMILY_MONTGOMERY:255:31:PSA_ERROR_BUFFER_TOO_SMALL
+
+ECC generate: Curve25519, size=32, ok
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_255
+generate_key:PSA_ECC_FAMILY_MONTGOMERY:255:32:PSA_SUCCESS
+
+ECC generate: Curve25519, size=33, ok
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_255
+generate_key:PSA_ECC_FAMILY_MONTGOMERY:255:33:PSA_SUCCESS
+
+ECC generate: Curve448, size=55, too small
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_448
+generate_key:PSA_ECC_FAMILY_MONTGOMERY:448:55:PSA_ERROR_BUFFER_TOO_SMALL
+
+ECC generate: Curve448, size=56, ok
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_448
+generate_key:PSA_ECC_FAMILY_MONTGOMERY:448:56:PSA_SUCCESS
+
+ECC generate: Curve448, size=57, ok
+depends_on:MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_448
+generate_key:PSA_ECC_FAMILY_MONTGOMERY:448:57:PSA_SUCCESS
diff --git a/tests/suites/test_suite_psa_crypto_ecp.function b/tests/suites/test_suite_psa_crypto_ecp.function
new file mode 100644
index 0000000..5be7a2c
--- /dev/null
+++ b/tests/suites/test_suite_psa_crypto_ecp.function
@@ -0,0 +1,165 @@
+/* BEGIN_HEADER */
+/* Unit tests for internal functions for built-in ECC mechanisms. */
+#include <psa/crypto.h>
+
+#include "psa_crypto_ecp.h"
+
+#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR_GENERATE)
+/*
+ * Check if a buffer is all-0 bytes:
+ * return   1 if it is,
+ *          0 if it isn't.
+ *
+ * TODO: we use this in multiple test suites. Move it to tests/src.
+ */
+static int buffer_is_all_zero(const uint8_t *buf, size_t size)
+{
+    for (size_t i = 0; i < size; i++) {
+        if (buf[i] != 0) {
+            return 0;
+        }
+    }
+    return 1;
+}
+
+typedef struct {
+    unsigned bit_bot;           /* lowest non-forced bit */
+    unsigned bit_top;           /* highest non-forced bit */
+} ecc_private_key_stats_t;
+
+/* Do some sanity checks on an ECC private key. This is not intended to be
+ * a full validity check, just to catch some potential mistakes. */
+static int check_ecc_private_key(psa_ecc_family_t family, size_t bits,
+                                 const uint8_t *key, size_t key_length,
+                                 ecc_private_key_stats_t *stats)
+{
+    int ok = 0;
+
+    /* Check the expected length (same calculation for all curves). */
+    TEST_EQUAL(PSA_BITS_TO_BYTES(bits), key_length);
+
+    /* All-bits zero is invalid and means no key material was copied to the
+     * output buffer, or a grave RNG pluming failure. */
+    TEST_ASSERT(!buffer_is_all_zero(key, key_length));
+
+    /* Check the top byte of the value for non-byte-aligned curve sizes.
+     * This is a partial endianness check. */
+    if (bits % 8 != 0) {
+        /* All supported non-byte-aligned curve sizes are for Weierstrass
+         * curves with a big-endian representation. */
+        uint8_t top_byte = key[0];
+        uint8_t mask = 0xff << (bits & 8);
+        TEST_EQUAL(top_byte & mask, 0);
+    }
+
+    /* Check masked bits on Curve25519 and Curve448 scalars.
+     * See RFC 7748 \S4.1 (we expect the "decoded" form here). */
+#if defined(MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_255)
+    if (family == PSA_ECC_FAMILY_MONTGOMERY && bits == 255) {
+        TEST_EQUAL(key[0] & 0xf8, key[0]);
+        TEST_EQUAL(key[31] & 0xc0, 0x40);
+    }
+#endif /* MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_255 */
+#if defined(MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_448)
+    if (family == PSA_ECC_FAMILY_MONTGOMERY && bits == 448) {
+        TEST_EQUAL(key[0] & 0xfc, key[0]);
+        TEST_EQUAL(key[55] & 0x80, 0x80);
+    }
+#endif /* MBEDTLS_PSA_BUILTIN_ECC_MONTGOMERY_448 */
+
+    /* Don't bother to check that the value is in the exact permitted range
+     * (1 to p-1 for Weierstrass curves, 2^{n-1} to p-1 for Montgomery curves).
+     * We would need to bring in bignum machinery, and on most curves
+     * the probability of a number being out of range is negligible.
+     */
+
+    /* Collect statistics on random-valued bits */
+    /* Defaults for big-endian numbers */
+    uint8_t bit_bot_mask = 0x01;
+    size_t bit_bot_index = key_length - 1;
+    uint8_t bit_top_mask = (bits % 8 == 0 ? 0x80 : 1 << (bits % 8 - 1));
+    size_t bit_top_index = 0;
+    if (family == PSA_ECC_FAMILY_MONTGOMERY) {
+        bit_bot_index = 0;
+        bit_top_index = key_length - 1;
+        if (bits == 255) {
+            bit_bot_mask = 0x08;
+            bit_top_mask = 0x20;
+        } else {
+            bit_bot_mask = 0x04;
+            bit_top_mask = 0x40;
+        }
+    }
+    if (key[bit_bot_index] & bit_bot_mask) {
+        ++stats->bit_bot;
+    }
+    if (key[bit_top_index] & bit_top_mask) {
+        ++stats->bit_top;
+    }
+
+    ok = 1;
+exit:
+    return ok;
+}
+#endif
+
+/* END_HEADER */
+
+/* BEGIN_DEPENDENCIES
+ * depends_on:MBEDTLS_PSA_CRYPTO_C:MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY
+ * END_DEPENDENCIES
+ */
+
+/* BEGIN_CASE depends_on:MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR_GENERATE */
+void generate_key(int family_arg, int bits_arg,
+                  int output_size_arg,
+                  psa_status_t expected_status)
+{
+    psa_ecc_family_t family = family_arg;
+    size_t bits = bits_arg;
+    size_t output_size = output_size_arg;
+
+    uint8_t *output = NULL;
+    size_t output_length = SIZE_MAX;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_set_key_type(&attributes, PSA_KEY_TYPE_ECC_KEY_PAIR(family));
+    psa_set_key_bits(&attributes, bits);
+    ecc_private_key_stats_t stats = { 0, 0 };
+
+    PSA_INIT();
+    TEST_CALLOC(output, output_size);
+
+    /* In success cases, run multiple iterations so that we can make
+     * statistical observations. */
+    unsigned iteration_count = expected_status == PSA_SUCCESS ? 256 : 1;
+    for (unsigned i = 0; i < iteration_count; i++) {
+        mbedtls_test_set_step(i);
+        TEST_EQUAL(mbedtls_psa_ecp_generate_key(&attributes,
+                                                output, output_size,
+                                                &output_length),
+                   expected_status);
+        if (expected_status == PSA_SUCCESS) {
+            TEST_LE_U(output_length, output_size);
+            TEST_ASSERT(check_ecc_private_key(family, bits,
+                                              output, output_length,
+                                              &stats));
+        }
+    }
+
+    if (expected_status == PSA_SUCCESS) {
+        /* For selected bits, check that we saw the values 0 and 1 each
+         * at least some minimum number of times. The iteration count and
+         * the minimum are chosen so that a random failure is unlikely
+         * to more than cryptographic levels. */
+        unsigned const min_times = 10;
+        TEST_LE_U(min_times, stats.bit_bot);
+        TEST_LE_U(stats.bit_bot, iteration_count - min_times);
+        TEST_LE_U(min_times, stats.bit_top);
+        TEST_LE_U(stats.bit_top, iteration_count - min_times);
+    }
+
+exit:
+    PSA_DONE();
+    mbedtls_free(output);
+}
+/* END_CASE */
diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index 32a52de..5f44cc3 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -8098,6 +8098,112 @@
                                    key);
 }
 
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+static psa_status_t psa_generate_key_iop_abort_internal(
+    psa_generate_key_iop_t *operation)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    if (operation->id == 0) {
+        return PSA_SUCCESS;
+    }
+
+    status = mbedtls_psa_generate_key_iop_abort(&operation->ctx);
+
+    operation->id = 0;
+
+    return status;
+}
+#endif
+
+uint32_t psa_generate_key_iop_get_num_ops(
+    psa_generate_key_iop_t *operation)
+{
+    (void) operation;
+    return 0;
+}
+
+psa_status_t psa_generate_key_iop_setup(
+    psa_generate_key_iop_t *operation,
+    const psa_key_attributes_t *attributes)
+{
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_type_t type;
+
+    if (operation->id != 0 || operation->error_occurred) {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
+    type = psa_get_key_type(attributes);
+
+    if (!PSA_KEY_TYPE_IS_ECC(type)) {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        goto exit;
+    }
+
+    if (psa_get_key_bits(attributes) == 0) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    if (PSA_KEY_TYPE_IS_PUBLIC_KEY(type)) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    operation->attributes = *attributes;
+
+    operation->num_ops = 0;
+
+    /* We only support the builtin/Mbed TLS driver for now. */
+    operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
+
+    status = mbedtls_psa_generate_key_iop_setup(&operation->ctx, attributes);
+
+exit:
+    if (status != PSA_SUCCESS) {
+        operation->error_occurred = 1;
+        psa_generate_key_iop_abort_internal(operation);
+    }
+
+    return status;
+#else
+    (void) operation;
+    (void) attributes;
+    return PSA_ERROR_NOT_SUPPORTED;
+#endif
+}
+
+psa_status_t psa_generate_key_iop_complete(
+    psa_generate_key_iop_t *operation,
+    psa_key_id_t *key)
+{
+    (void) operation;
+    (void) key;
+
+    return PSA_ERROR_NOT_SUPPORTED;
+}
+
+psa_status_t psa_generate_key_iop_abort(
+    psa_generate_key_iop_t *operation)
+{
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    psa_status_t status;
+
+    status = psa_generate_key_iop_abort_internal(operation);
+
+    operation->error_occurred = 0;
+    operation->num_ops = 0;
+    return status;
+#else
+    (void) operation;
+    return PSA_SUCCESS;
+#endif
+}
+
+
 /****************************************************************/
 /* Module setup */
 /****************************************************************/
diff --git a/tf-psa-crypto/core/psa_crypto_core.h b/tf-psa-crypto/core/psa_crypto_core.h
index df0ee50..1753554 100644
--- a/tf-psa-crypto/core/psa_crypto_core.h
+++ b/tf-psa-crypto/core/psa_crypto_core.h
@@ -435,6 +435,40 @@
                                        size_t key_buffer_size,
                                        size_t *key_buffer_length);
 
+/**
+ * \brief Setup a new interruptible key generation operation.
+ *
+ *  \param[in] operation                 The \c mbedtls_psa_generate_key_iop_t to use.
+ *                                       This must be initialized first.
+ *  \param[in] attributes                The desired attributes of the generated key.
+ *
+ *  \retval #PSA_SUCCESS
+ *         The operation started successfully - call \c mbedtls_psa_generate_key_complete()
+ *         with the same operation to complete the operation.
+ * * \retval #PSA_ERROR_NOT_SUPPORTED
+ *           Either no internal interruptible operations are
+ *           currently supported, or the key attributes are not unsupported.
+ *  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ *         There was insufficient memory to load the key representation.
+ *
+ */
+psa_status_t mbedtls_psa_generate_key_iop_setup(
+    mbedtls_psa_generate_key_iop_t *operation,
+    const psa_key_attributes_t *attributes);
+
+/**
+ * \brief Abort a key generation operation.
+ *
+ * \param[in] operation               The \c mbedtls_psa_generate_key_iop_t to abort.
+ *
+ * \retval #PSA_SUCCESS
+ *         The operation was aborted successfully.
+ *
+ */
+psa_status_t mbedtls_psa_generate_key_iop_abort(
+    mbedtls_psa_generate_key_iop_t *operation);
+
+
 /** Sign a message with a private key. For hash-and-sign algorithms,
  *  this includes the hashing step.
  *
diff --git a/tf-psa-crypto/drivers/builtin/include/mbedtls/bignum.h b/tf-psa-crypto/drivers/builtin/include/mbedtls/bignum.h
index 8367cd3..22d5d84 100644
--- a/tf-psa-crypto/drivers/builtin/include/mbedtls/bignum.h
+++ b/tf-psa-crypto/drivers/builtin/include/mbedtls/bignum.h
@@ -238,6 +238,8 @@
 }
 mbedtls_mpi;
 
+#define MBEDTLS_MPI_INIT { 0, 0, 0 }
+
 /**
  * \brief           Initialize an MPI context.
  *
diff --git a/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h b/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h
index 7b0a80d..b340614 100644
--- a/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h
+++ b/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h
@@ -162,6 +162,8 @@
 }
 mbedtls_ecp_point;
 
+#define MBEDTLS_ECP_POINT_INIT { MBEDTLS_MPI_INIT, MBEDTLS_MPI_INIT, MBEDTLS_MPI_INIT }
+
 /**
  * \brief           The ECP group structure.
  *
@@ -250,6 +252,9 @@
 }
 mbedtls_ecp_group;
 
+#define MBEDTLS_ECP_GROUP_INIT { MBEDTLS_ECP_DP_NONE, MBEDTLS_MPI_INIT, MBEDTLS_MPI_INIT, \
+                                 MBEDTLS_MPI_INIT, MBEDTLS_ECP_POINT_INIT, MBEDTLS_MPI_INIT, \
+                                 0, 0, 0, NULL, NULL, NULL, NULL, NULL, 0 }
 /**
  * \name SECTION: Module settings
  *
@@ -419,6 +424,9 @@
 }
 mbedtls_ecp_keypair;
 
+#define MBEDTLS_ECP_KEYPAIR_INIT { MBEDTLS_ECP_GROUP_INIT, MBEDTLS_MPI_INIT, \
+                                   MBEDTLS_ECP_POINT_INIT }
+
 /**
  * The uncompressed point format for Short Weierstrass curves
  * (MBEDTLS_ECP_DP_SECP_XXX and MBEDTLS_ECP_DP_BP_XXX).
diff --git a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
index cc5b483..acb2482 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -321,38 +321,36 @@
     const psa_key_attributes_t *attributes,
     uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length)
 {
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-
     psa_ecc_family_t curve = PSA_KEY_TYPE_ECC_GET_FAMILY(
         attributes->type);
     mbedtls_ecp_group_id grp_id =
         mbedtls_ecc_group_from_psa(curve, attributes->bits);
-
-    const mbedtls_ecp_curve_info *curve_info =
-        mbedtls_ecp_curve_info_from_grp_id(grp_id);
-    mbedtls_ecp_keypair ecp;
-
-    if (grp_id == MBEDTLS_ECP_DP_NONE || curve_info == NULL) {
+    if (grp_id == MBEDTLS_ECP_DP_NONE) {
         return PSA_ERROR_NOT_SUPPORTED;
     }
 
+    mbedtls_ecp_keypair ecp;
     mbedtls_ecp_keypair_init(&ecp);
-    ret = mbedtls_ecp_gen_key(grp_id, &ecp,
-                              mbedtls_psa_get_random,
-                              MBEDTLS_PSA_RANDOM_STATE);
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    ret = mbedtls_ecp_group_load(&ecp.grp, grp_id);
     if (ret != 0) {
-        mbedtls_ecp_keypair_free(&ecp);
-        return mbedtls_to_psa_error(ret);
+        goto exit;
     }
 
-    status = mbedtls_to_psa_error(
-        mbedtls_ecp_write_key_ext(&ecp, key_buffer_length,
-                                  key_buffer, key_buffer_size));
+    ret = mbedtls_ecp_gen_privkey(&ecp.grp, &ecp.d,
+                                  mbedtls_psa_get_random,
+                                  MBEDTLS_PSA_RANDOM_STATE);
+    if (ret != 0) {
+        goto exit;
+    }
 
+    ret = mbedtls_ecp_write_key_ext(&ecp, key_buffer_length,
+                                    key_buffer, key_buffer_size);
+
+exit:
     mbedtls_ecp_keypair_free(&ecp);
-
-    return status;
+    return mbedtls_to_psa_error(ret);
 }
 #endif /* MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR_GENERATE */
 
@@ -596,41 +594,38 @@
 /* Interruptible ECC Key Generation */
 /****************************************************************/
 
-uint32_t psa_generate_key_iop_get_num_ops(
-    psa_generate_key_iop_t *operation)
-{
-    (void) operation;
-    return 0;
-}
+#if defined(MBEDTLS_ECP_RESTARTABLE)
 
-psa_status_t psa_generate_key_iop_setup(
-    psa_generate_key_iop_t *operation,
+psa_status_t mbedtls_psa_generate_key_iop_setup(
+    mbedtls_psa_generate_key_iop_t *operation,
     const psa_key_attributes_t *attributes)
 {
-    (void) operation;
-    (void) attributes;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    return PSA_ERROR_NOT_SUPPORTED;
+    mbedtls_ecp_keypair_init(&operation->ecp);
+
+    psa_ecc_family_t curve = PSA_KEY_TYPE_ECC_GET_FAMILY(
+        psa_get_key_type(attributes));
+    mbedtls_ecp_group_id grp_id =
+        mbedtls_ecc_group_from_psa(curve, psa_get_key_bits(attributes));
+    if (grp_id == MBEDTLS_ECP_DP_NONE) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    status = mbedtls_ecp_group_load(&operation->ecp.grp, grp_id);
+
+    return mbedtls_to_psa_error(status);
 }
 
-psa_status_t psa_generate_key_iop_complete(
-    psa_generate_key_iop_t *operation,
-    psa_key_id_t *key)
+psa_status_t mbedtls_psa_generate_key_iop_abort(
+    mbedtls_psa_generate_key_iop_t *operation)
 {
-    (void) operation;
-    (void) key;
-
-    return PSA_ERROR_NOT_SUPPORTED;
+    mbedtls_ecp_keypair_free(&operation->ecp);
+    operation->num_ops = 0;
+    return PSA_SUCCESS;
 }
 
-psa_status_t psa_generate_key_iop_abort(
-    psa_generate_key_iop_t *operation)
-{
-    (void) operation;
-
-    return PSA_ERROR_NOT_SUPPORTED;
-}
-
+#endif
 /****************************************************************/
 /* Interruptible ECC Key Agreement */
 /****************************************************************/
diff --git a/tf-psa-crypto/include/psa/crypto.h b/tf-psa-crypto/include/psa/crypto.h
index aa58033..58b6887 100644
--- a/tf-psa-crypto/include/psa/crypto.h
+++ b/tf-psa-crypto/include/psa/crypto.h
@@ -4360,8 +4360,9 @@
  *                              time. The only guarantee is that lower values
  *                              for \p max_ops means functions will block for a
  *                              lesser maximum amount of time. The functions
- *                              \c psa_sign_interruptible_get_num_ops() and
- *                              \c psa_verify_interruptible_get_num_ops() are
+ *                              \c psa_sign_interruptible_get_num_ops(),
+ *                              \c psa_verify_interruptible_get_num_ops() and
+ *                              \c psa_generate_key_iop_get_num_ops() are
  *                              provided to help with tuning this value.
  *
  * \note                        This value defaults to
diff --git a/tf-psa-crypto/include/psa/crypto_builtin_composites.h b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
index c14f5dd..c9c0c6b 100644
--- a/tf-psa-crypto/include/psa/crypto_builtin_composites.h
+++ b/tf-psa-crypto/include/psa/crypto_builtin_composites.h
@@ -211,4 +211,20 @@
 
 #define MBEDTLS_PSA_PAKE_OPERATION_INIT { { 0 } }
 
+typedef struct {
+#if defined(MBEDTLS_ECP_C)
+    mbedtls_ecp_keypair MBEDTLS_PRIVATE(ecp);
+    uint32_t num_ops;
+#else
+    /* Make the struct non-empty if algs not supported. */
+    unsigned MBEDTLS_PRIVATE(dummy);
+#endif
+} mbedtls_psa_generate_key_iop_t;
+
+#if defined(MBEDTLS_ECP_C)
+#define MBEDTLS_PSA_GENERATE_KEY_IOP_INIT { MBEDTLS_ECP_KEYPAIR_INIT, 0 }
+#else
+#define MBEDTLS_PSA_GENERATE_KEY_IOP_INIT { 0 }
+#endif
+
 #endif /* PSA_CRYPTO_BUILTIN_COMPOSITES_H */
diff --git a/tf-psa-crypto/include/psa/crypto_struct.h b/tf-psa-crypto/include/psa/crypto_struct.h
index 2eec948..76ef5c4 100644
--- a/tf-psa-crypto/include/psa/crypto_struct.h
+++ b/tf-psa-crypto/include/psa/crypto_struct.h
@@ -542,14 +542,18 @@
      * any driver (i.e. none of the driver contexts are active).
      */
     unsigned int MBEDTLS_PRIVATE(id);
-
+    mbedtls_psa_generate_key_iop_t MBEDTLS_PRIVATE(ctx);
+    psa_key_attributes_t MBEDTLS_PRIVATE(attributes);
+    unsigned int MBEDTLS_PRIVATE(error_occurred) : 1;
+    uint32_t MBEDTLS_PRIVATE(num_ops);
 #endif
 };
 
 #if defined(MBEDTLS_PSA_CRYPTO_CLIENT) && !defined(MBEDTLS_PSA_CRYPTO_C)
 #define PSA_GENERATE_KEY_IOP_INIT { 0 }
 #else
-#define PSA_GENERATE_KEY_IOP_INIT { 0 }
+#define PSA_GENERATE_KEY_IOP_INIT { 0, MBEDTLS_PSA_GENERATE_KEY_IOP_INIT, PSA_KEY_ATTRIBUTES_INIT, \
+                                    0, 0 }
 #endif
 
 static inline struct psa_generate_key_iop_s
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data
index 87fec19..dab2ee7 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.data
@@ -7496,12 +7496,44 @@
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_SECP_R1_256
 generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS:0
 
+PSA generate key: ECC, SECP384R1, good
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_SECP_R1_384
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):384:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS:0
+
+PSA generate key: ECC, SECP521R1, good
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_SECP_R1_521
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):521:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS:0
+
+PSA generate key: ECC, SECP192K1, good
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_SECP_K1_192
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):192:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS:0
+
+PSA generate key: ECC, SECP256K1, good
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_SECP_K1_256
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_K1):256:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS:0
+
+PSA generate key: ECC, brainpool256r1, good
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_BRAINPOOL_P_R1_256
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):256:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS:0
+
+PSA generate key: ECC, brainpool384r1, good
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_BRAINPOOL_P_R1_384
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_BRAINPOOL_P_R1):384:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS:0
+
 PSA generate key: ECC, SECP256R1, incorrect bit size
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_BASIC:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_IMPORT:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_EXPORT:PSA_WANT_ECC_SECP_R1_256
 # INVALID_ARGUMENT would make more sense, but our code as currently structured
 # doesn't fully relate the curve with its size.
 generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_ERROR_NOT_SUPPORTED:0
 
+PSA generate key: ECC, SECP256R1, zero bit size
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_BASIC:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_IMPORT:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_EXPORT:PSA_WANT_ECC_SECP_R1_256
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):0:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_ERROR_INVALID_ARGUMENT:0
+
+PSA generate key: ECC, SECP256R1, public key
+depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_BASIC:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_IMPORT:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_EXPORT:PSA_WANT_ECC_SECP_R1_256
+generate_key:PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_SECP_R1):0:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_ERROR_INVALID_ARGUMENT:0
+
 PSA generate key: ECC, Curve25519, good
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_MONTGOMERY_255
 generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_MONTGOMERY):255:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_SUCCESS:0
@@ -7534,6 +7566,9 @@
 depends_on:PSA_WANT_ALG_FFDH:PSA_WANT_KEY_TYPE_DH_KEY_PAIR_GENERATE
 generate_key:PSA_KEY_TYPE_DH_KEY_PAIR(PSA_DH_FAMILY_RFC7919):1024:PSA_KEY_USAGE_EXPORT:PSA_ALG_FFDH:PSA_ERROR_NOT_SUPPORTED:0
 
+PSA generate key interruptible object initializers zero properly
+generate_key_iop_init:
+
 PSA generate key custom: RSA, flags=1
 depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE
 generate_key_custom:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:0:1:"":PSA_ERROR_INVALID_ARGUMENT
diff --git a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
index b1c662f..eaafd90 100644
--- a/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
+++ b/tf-psa-crypto/tests/suites/test_suite_psa_crypto.function
@@ -10096,6 +10096,7 @@
     psa_status_t expected_status = expected_status_arg;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_attributes_t got_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_generate_key_iop_t operation = PSA_GENERATE_KEY_IOP_INIT;
 
     PSA_ASSERT(psa_crypto_init());
 
@@ -10111,21 +10112,46 @@
         TEST_ASSUME(status != PSA_ERROR_INSUFFICIENT_MEMORY);
     }
     TEST_EQUAL(status, expected_status);
-    if (expected_status != PSA_SUCCESS) {
-        goto exit;
+    if (expected_status == PSA_SUCCESS) {
+        /* Test the key information */
+        PSA_ASSERT(psa_get_key_attributes(key, &got_attributes));
+        TEST_EQUAL(psa_get_key_type(&got_attributes), type);
+        TEST_EQUAL(psa_get_key_bits(&got_attributes), bits);
+
+        /* Do something with the key according to its type and permitted usage. */
+        TEST_EQUAL(mbedtls_test_psa_exercise_key(key, usage, alg, 0), 1);
     }
 
-    /* Test the key information */
-    PSA_ASSERT(psa_get_key_attributes(key, &got_attributes));
-    TEST_EQUAL(psa_get_key_type(&got_attributes), type);
-    TEST_EQUAL(psa_get_key_bits(&got_attributes), bits);
+    /* Adjust expected_status for interruptible key generation.
+     * Interruptible key generation is only supported for ECC key pairs and even
+     * for those only when MBEDTLS_ECP_RESTARTABLE is on.
+     */
 
-    /* Do something with the key according to its type and permitted usage. */
-    if (!mbedtls_test_psa_exercise_key(key, usage, alg, 0)) {
-        goto exit;
+    if (!PSA_KEY_TYPE_IS_ECC(type)) {
+        expected_status = PSA_ERROR_NOT_SUPPORTED;
     }
 
+#if !defined(MBEDTLS_ECP_RESTARTABLE)
+    expected_status = PSA_ERROR_NOT_SUPPORTED;
+#endif
+
+    status = psa_generate_key_iop_setup(&operation, &attributes);
+    TEST_EQUAL(status, expected_status);
+
+    /* Test that calling setup 2 time consecutively fails */
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    status = psa_generate_key_iop_setup(&operation, &attributes);
+    TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
+#endif
+
+    PSA_ASSERT(psa_generate_key_iop_abort(&operation));
+
+    /* Test that after calling abort operation is reset to it's fresh state */
+    status = psa_generate_key_iop_setup(&operation, &attributes);
+    TEST_EQUAL(status, expected_status);
+
 exit:
+    psa_generate_key_iop_abort(&operation);
     /*
      * Key attributes may have been returned by psa_get_key_attributes()
      * thus reset them as required.
@@ -10138,6 +10164,21 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void generate_key_iop_init()
+{
+    psa_generate_key_iop_t init = PSA_GENERATE_KEY_IOP_INIT;
+    psa_generate_key_iop_t func = psa_generate_key_iop_init();
+    psa_generate_key_iop_t zero;
+
+    memset(&zero, 0, sizeof(zero));
+
+    PSA_ASSERT(psa_generate_key_iop_abort(&init));
+    PSA_ASSERT(psa_generate_key_iop_abort(&func));
+    PSA_ASSERT(psa_generate_key_iop_abort(&zero));
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void generate_key_custom(int type_arg,
                          int bits_arg,
                          int usage_arg,