Add driver support for DH import key and export public key

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>

library/psa_crypto.c

@@ -640,14 +640,11 @@ psa_status_t psa_import_key_into_slot(
             if (psa_is_dh_key_size_valid(PSA_BYTES_TO_BITS(data_length)) == 0) {
                 return PSA_ERROR_INVALID_ARGUMENT;
             }
-
+            return mbedtls_psa_ffdh_import_key(attributes,
-            /* Copy the key material. */
+                                               data, data_length,
-            memcpy(key_buffer, data, data_length);
+                                               key_buffer, key_buffer_size,
-            *key_buffer_length = data_length;
+                                               key_buffer_length,
-            *bits = PSA_BYTES_TO_BITS(data_length);
+                                               bits);
-            (void) key_buffer_size;
-
-            return PSA_SUCCESS;
         }
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_DH_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_DH_PUBLIC_KEY) */

library/psa_crypto_ffdh.c

@@ -134,7 +134,18 @@ psa_status_t mbedtls_psa_export_ffdh_public_key(
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     mbedtls_mpi GX, G, X, P;
-    (void) attributes;
+    psa_key_type_t type = attributes->core.type;
+
+    if (PSA_KEY_TYPE_IS_PUBLIC_KEY(type)) {
+        if (key_buffer_size > data_size) {
+            return PSA_ERROR_BUFFER_TOO_SMALL;
+        }
+        memcpy(data, key_buffer, key_buffer_size);
+        memset(data + key_buffer_size, 0,
+               data_size - key_buffer_size);
+        *data_length = key_buffer_size;
+        return PSA_SUCCESS;
+    }
 
     mbedtls_mpi_init(&GX); mbedtls_mpi_init(&G);
     mbedtls_mpi_init(&X); mbedtls_mpi_init(&P);
@@ -199,6 +210,24 @@ cleanup:
     return status;
 }
 
+psa_status_t mbedtls_psa_ffdh_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits)
+{
+    (void) attributes;
+
+    if (key_buffer_size < data_length) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+    memcpy(key_buffer, data, data_length);
+    *key_buffer_length = data_length;
+    *bits = PSA_BYTES_TO_BITS(data_length);
+
+    return PSA_SUCCESS;
+}
+
 #endif /* MBEDTLS_PSA_BUILTIN_KEY_TYPE_DH_KEY_PAIR ||
           MBEDTLS_PSA_BUILTIN_KEY_TYPE_DH_PUBLIC_KEY */
 

library/psa_crypto_ffdh.h

@@ -112,4 +112,33 @@ psa_status_t mbedtls_psa_ffdh_generate_key(
     size_t key_buffer_size,
     size_t *key_buffer_length);
 
+/**
+ * \brief Import DH key.
+ *
+ * \note The signature of the function is that of a PSA driver import_key
+ *       entry point.
+ *
+ * \param[in]  attributes       The attributes for the key to import.
+ * \param[in]  data             The buffer containing the key data in import
+ *                              format.
+ * \param[in]  data_length      Size of the \p data buffer in bytes.
+ * \param[out] key_buffer       The buffer containing the key data in output
+ *                              format.
+ * \param[in]  key_buffer_size  Size of the \p key_buffer buffer in bytes. This
+ *                              size is greater or equal to \p data_length.
+ * \param[out] key_buffer_length  The length of the data written in \p
+ *                                key_buffer in bytes.
+ * \param[out] bits             The key size in number of bits.
+ *
+ * \retval #PSA_SUCCESS
+ *         The key was generated successfully.
+ * \retval #PSA_ERROR_BUFFER_TOO_SMALL
+ *         The size of \p key_buffer is too small.
+ */
+psa_status_t mbedtls_psa_ffdh_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits);
+
 #endif /* PSA_CRYPTO_FFDH_H */

tests/src/drivers/test_driver_key_management.c

@@ -321,9 +321,25 @@ psa_status_t mbedtls_test_transparent_import_key(
             data, data_length,
             key_buffer, key_buffer_size,
             key_buffer_length, bits);
+#endif
+    } else if (PSA_KEY_TYPE_IS_DH(type)) {
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+        (defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_KEY_TYPE_DH_KEY_PAIR) || \
+        defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_KEY_TYPE_DH_PUBLIC_KEY))
+        return libtestdriver1_mbedtls_psa_ffdh_import_key(
+            (const libtestdriver1_psa_key_attributes_t *) attributes,
+            data, data_length,
+            key_buffer, key_buffer_size,
+            key_buffer_length, bits);
+#elif defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
+        defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
+        return mbedtls_psa_ffdh_import_key(
+            attributes,
+            data, data_length,
+            key_buffer, key_buffer_size,
+            key_buffer_length, bits);
 #endif
     }
-
     (void) data;
     (void) data_length;
     (void) key_buffer;