Merge pull request #9651 from waleed-elmelegy-arm/add-iop-key-gen-com…

…plete

Add PSA interruptible key generation complete API

tf-psa-crypto/core/psa_crypto.c

@@ -8314,7 +8314,9 @@ static psa_status_t psa_generate_key_iop_abort_internal(
         return PSA_SUCCESS;
     }
 
-    status = mbedtls_psa_generate_key_iop_abort(&operation->ctx);
+    status = mbedtls_psa_ecp_generate_key_iop_abort(&operation->ctx);
+
+    psa_reset_key_attributes(&operation->attributes);
 
     operation->id = 0;
 
@@ -8366,7 +8368,7 @@ psa_status_t psa_generate_key_iop_setup(
     /* 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);
+    status = mbedtls_psa_ecp_generate_key_iop_setup(&operation->ctx, attributes);
 
 exit:
     if (status != PSA_SUCCESS) {
@@ -8386,10 +8388,42 @@ psa_status_t psa_generate_key_iop_complete(
     psa_generate_key_iop_t *operation,
     mbedtls_svc_key_id_t *key)
 {
+#if defined(MBEDTLS_ECP_RESTARTABLE)
+    psa_status_t status;
+    uint8_t key_data[PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)+1] = { 0 };
+    size_t key_len = 0;
+
+    if (operation->id == 0 || operation->error_occurred) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    status = mbedtls_psa_ecp_generate_key_iop_complete(&operation->ctx, key_data,
+                                                       sizeof(key_data), &key_len);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    status = psa_import_key(&operation->attributes,
+                            key_data + (sizeof(key_data) - key_len),
+                            key_len,
+                            key);
+
+exit:
+    if (status != PSA_OPERATION_INCOMPLETE) {
+        if (status != PSA_SUCCESS) {
+            operation->error_occurred = 1;
+        }
+        psa_generate_key_iop_abort_internal(operation);
+    }
+
+    mbedtls_platform_zeroize(key_data, sizeof(key_data));
+    return status;
+#else
     (void) operation;
     (void) key;
 
-    return PSA_ERROR_NOT_SUPPORTED;
+    return PSA_ERROR_BAD_STATE;
+#endif
 }
 
 psa_status_t psa_generate_key_iop_abort(

tf-psa-crypto/core/psa_crypto_core.h

@@ -435,40 +435,6 @@ psa_status_t psa_generate_key_internal(const psa_key_attributes_t *attributes,
                                        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.
  *

tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c

@@ -596,11 +596,11 @@ psa_status_t mbedtls_psa_key_agreement_ecdh(
 
 #if defined(MBEDTLS_ECP_RESTARTABLE)
 
-psa_status_t mbedtls_psa_generate_key_iop_setup(
+psa_status_t mbedtls_psa_ecp_generate_key_iop_setup(
     mbedtls_psa_generate_key_iop_t *operation,
     const psa_key_attributes_t *attributes)
 {
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    int status = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
     mbedtls_ecp_keypair_init(&operation->ecp);
 
@@ -617,7 +617,36 @@ psa_status_t mbedtls_psa_generate_key_iop_setup(
     return mbedtls_to_psa_error(status);
 }
 
-psa_status_t mbedtls_psa_generate_key_iop_abort(
+psa_status_t mbedtls_psa_ecp_generate_key_iop_complete(
+    mbedtls_psa_generate_key_iop_t *operation,
+    uint8_t *key_output,
+    size_t key_output_size,
+    size_t *key_len)
+{
+    *key_len = 0;
+    int status = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    *key_len = PSA_BITS_TO_BYTES(operation->ecp.grp.nbits);
+
+    if (*key_len > key_output_size) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+
+    status = mbedtls_ecp_gen_privkey(&operation->ecp.grp, &operation->ecp.d,
+                                     mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE);
+
+    if (status != 0) {
+        return mbedtls_to_psa_error(status);
+    }
+
+    operation->num_ops = 1;
+
+    status = mbedtls_mpi_write_binary(&operation->ecp.d, key_output, key_output_size);
+
+    return mbedtls_to_psa_error(status);
+}
+
+psa_status_t mbedtls_psa_ecp_generate_key_iop_abort(
     mbedtls_psa_generate_key_iop_t *operation)
 {
     mbedtls_ecp_keypair_free(&operation->ecp);

tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.h

@@ -143,6 +143,67 @@ psa_status_t mbedtls_psa_ecp_generate_key(
     const psa_key_attributes_t *attributes,
     uint8_t *key_buffer, 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_ecp_generate_key_iop_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_ecp_generate_key_iop_setup(
+    mbedtls_psa_generate_key_iop_t *operation,
+    const psa_key_attributes_t *attributes);
+
+/**
+ * \brief Continue and eventually complete a key generation operation.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       generate_key_complete entry point. This function behaves as a
+ *       generate_key_complete entry point as defined in the PSA driver
+ *       interface specification for transparent drivers.
+ *
+ * \param[in] operation                  The \c mbedtls_psa_generate_key_iop_t to use.
+ *                                       This must be initialized first and
+ *                                       had \c mbedtls_psa_ecp_generate_key_iop_setup()
+ *                                       called successfully.
+ * \param[out] key_output                The buffer to which the generated key
+ *                                       is to be written.
+ * \param[out] key_len                   On success, the number of bytes that make
+ *                                       up the returned key output.
+ * \retval #PSA_SUCCESS
+ *         The key was generated successfully.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT \emptydescription
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY \emptydescription
+ *
+ */
+psa_status_t mbedtls_psa_ecp_generate_key_iop_complete(
+    mbedtls_psa_generate_key_iop_t *operation,
+    uint8_t *key_output,
+    size_t key_output_size,
+    size_t *key_len);
+
+/**
+ * \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_ecp_generate_key_iop_abort(
+    mbedtls_psa_generate_key_iop_t *operation);
+
 /** Sign an already-calculated hash with ECDSA.
  *
  * \note The signature of this function is that of a PSA driver

tf-psa-crypto/tests/suites/test_suite_psa_crypto.function

@@ -10230,13 +10230,15 @@ void generate_key(int type_arg,
                   int is_large_key)
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    mbedtls_svc_key_id_t iop_key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_type_t type = type_arg;
     psa_key_usage_t usage = usage_arg;
     size_t bits = bits_arg;
     psa_algorithm_t alg = alg_arg;
     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_key_attributes_t iop_attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_generate_key_iop_t operation = PSA_GENERATE_KEY_IOP_INIT;
 
     PSA_ASSERT(psa_crypto_init());
@@ -10276,6 +10278,12 @@ void generate_key(int type_arg,
     expected_status = PSA_ERROR_NOT_SUPPORTED;
 #endif
 
+    /* Test calling complete() without calling setup() will fail. */
+    status = psa_generate_key_iop_complete(&operation, &iop_key);
+    TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
+
+    psa_generate_key_iop_abort(&operation);
+
     status = psa_generate_key_iop_setup(&operation, &attributes);
     TEST_EQUAL(status, expected_status);
 
@@ -10291,15 +10299,37 @@ void generate_key(int type_arg,
     status = psa_generate_key_iop_setup(&operation, &attributes);
     TEST_EQUAL(status, expected_status);
 
+    if (expected_status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    do {
+        status = psa_generate_key_iop_complete(&operation, &iop_key);
+    } while (status == PSA_OPERATION_INCOMPLETE);
+
+    TEST_EQUAL(status, PSA_SUCCESS);
+
+    PSA_ASSERT(psa_get_key_attributes(iop_key, &iop_attributes));
+    TEST_EQUAL(psa_get_key_type(&iop_attributes), type);
+    TEST_EQUAL(psa_get_key_bits(&iop_attributes), bits);
+
+    TEST_EQUAL(mbedtls_test_psa_exercise_key(iop_key, usage, alg, 0), 1);
+
+    /* Test calling complete() 2 times consecutively will fail. */
+    status = psa_generate_key_iop_complete(&operation, &iop_key);
+    TEST_EQUAL(status, PSA_ERROR_BAD_STATE);
+
 exit:
     psa_generate_key_iop_abort(&operation);
     /*
      * Key attributes may have been returned by psa_get_key_attributes()
      * thus reset them as required.
      */
     psa_reset_key_attributes(&got_attributes);
+    psa_reset_key_attributes(&iop_attributes);
 
     psa_destroy_key(key);
+    psa_destroy_key(iop_key);
     PSA_DONE();
 }
 /* END_CASE */