diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index acb39a1..dc383bc 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -640,14 +640,11 @@
             if (psa_is_dh_key_size_valid(PSA_BYTES_TO_BITS(data_length)) == 0) {
                 return PSA_ERROR_INVALID_ARGUMENT;
             }
-
-            /* Copy the key material. */
-            memcpy(key_buffer, data, data_length);
-            *key_buffer_length = data_length;
-            *bits = PSA_BYTES_TO_BITS(data_length);
-            (void) key_buffer_size;
-
-            return PSA_SUCCESS;
+            return mbedtls_psa_ffdh_import_key(attributes,
+                                               data, data_length,
+                                               key_buffer, key_buffer_size,
+                                               key_buffer_length,
+                                               bits);
         }
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_DH_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_DH_PUBLIC_KEY) */
diff --git a/library/psa_crypto_ffdh.c b/library/psa_crypto_ffdh.c
index db30a89..4550a72 100644
--- a/library/psa_crypto_ffdh.c
+++ b/library/psa_crypto_ffdh.c
@@ -134,7 +134,18 @@
     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 @@
     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 */
 
diff --git a/library/psa_crypto_ffdh.h b/library/psa_crypto_ffdh.h
index 62b05b2..5d7d951 100644
--- a/library/psa_crypto_ffdh.h
+++ b/library/psa_crypto_ffdh.h
@@ -112,4 +112,33 @@
     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 */