diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index e962e29..d70dccb 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1891,6 +1891,7 @@
     psa_key_slot_t *slot = NULL;
     psa_se_drv_table_entry_t *driver = NULL;
     size_t bits;
+    size_t storage_size = data_length;
 
     *key = MBEDTLS_SVC_KEY_ID_INIT;
 
@@ -1906,12 +1907,18 @@
         goto exit;
 
     /* In the case of a transparent key or an opaque key stored in local
-     * storage (thus not in the case of generating a key in a secure element
-     * or cryptoprocessor with storage), we have to allocate a buffer to
-     * hold the generated key material. */
+     * storage, we have to allocate a buffer to hold the generated key
+     * material. */
     if( slot->key.data == NULL )
     {
-        status = psa_allocate_buffer_to_slot( slot, data_length );
+        if( psa_key_lifetime_is_external( attributes->core.lifetime ) )
+        {
+            status = psa_driver_wrapper_get_key_buffer_size_from_key_data( attributes, data,
+                                           data_length , &storage_size );
+            if( status != PSA_SUCCESS )
+                goto exit;
+        }
+        status = psa_allocate_buffer_to_slot( slot, storage_size );
         if( status != PSA_SUCCESS )
             goto exit;
     }
@@ -4142,6 +4149,7 @@
 {
     uint8_t *data = NULL;
     size_t bytes = PSA_BITS_TO_BYTES( bits );
+    size_t storage_size = bytes;
     psa_status_t status;
 
     if( ! key_type_is_raw_bytes( slot->attr.type ) )
@@ -4160,15 +4168,21 @@
         psa_des_set_key_parity( data, bytes );
 #endif /* MBEDTLS_PSA_BUILTIN_KEY_TYPE_DES */
 
-    status = psa_allocate_buffer_to_slot( slot, bytes );
-    if( status != PSA_SUCCESS )
-        goto exit;
-
     slot->attr.bits = (psa_key_bits_t) bits;
     psa_key_attributes_t attributes = {
       .core = slot->attr
     };
 
+    if( psa_key_lifetime_is_external( attributes.core.lifetime ) )
+    {
+        status = psa_driver_wrapper_get_key_buffer_size( &attributes, &storage_size );
+        if( status != PSA_SUCCESS )
+            goto exit;
+    }
+    status = psa_allocate_buffer_to_slot( slot, storage_size );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
     status = psa_driver_wrapper_import_key( &attributes,
                                             data, bytes,
                                             slot->key.data,
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 38d0e30..2974d6f 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -380,8 +380,49 @@
     }
 }
 
+/** calculate the key buffer size required to store the key material of a key
+ *  associated with an opaque driver from input key data.
+ *
+ *
+ * \param[in] attributes        The key attributes
+ * \param[in] data              The input key data.
+ * \param[in] data_length       The input data length.
+ * \param[out] key_buffer_size  Minimum buffer size to contain the key material.
+ *
+ * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ */
+psa_status_t psa_driver_wrapper_get_key_buffer_size_from_key_data(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data,
+    size_t data_length,
+    size_t *key_buffer_size )
+{
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( attributes->core.lifetime );
+    psa_key_type_t key_type = attributes->core.type;
+
+    *key_buffer_size = 0;
+    switch( location )
+    {
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case PSA_CRYPTO_TEST_DRIVER_LOCATION:
+            *key_buffer_size = mbedtls_test_opaque_size_function( key_type,
+                                     PSA_BYTES_TO_BITS( data_length ) );
+            return( ( *key_buffer_size != 0 ) ?
+                    PSA_SUCCESS : PSA_ERROR_NOT_SUPPORTED );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+
+        default:
+            (void)key_type;
+            (void)data;
+            (void)data_length;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+}
+
 /** Get the key buffer size required to store the key material of a key
- *  associated with an opaque driver without storage.
+ *  associated with an opaque driver.
  *
  * \param[in] attributes  The key attributes.
  * \param[out] key_buffer_size  Minimum buffer size to contain the key material
@@ -389,11 +430,11 @@
  * \retval #PSA_SUCCESS
  *         The minimum size for a buffer to contain the key material has been
  *         returned successfully.
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         The size in bits of the key is not valid.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         The type and/or the size in bits of the key or the combination of
  *         the two is not supported.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The key is declared with a lifetime not known to us.
  */
 psa_status_t psa_driver_wrapper_get_key_buffer_size(
     const psa_key_attributes_t *attributes,
@@ -426,7 +467,7 @@
         default:
             (void)key_type;
             (void)key_bits;
-            return( PSA_ERROR_NOT_SUPPORTED );
+            return( PSA_ERROR_INVALID_ARGUMENT );
     }
 }
 
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index 38a6ee8..99455a8 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -85,6 +85,12 @@
     const psa_key_attributes_t *attributes,
     size_t *key_buffer_size );
 
+psa_status_t psa_driver_wrapper_get_key_buffer_size_from_key_data(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data,
+    size_t data_length,
+    size_t *key_buffer_size );
+
 psa_status_t psa_driver_wrapper_generate_key(
     const psa_key_attributes_t *attributes,
     uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length );