diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index f9ea369..7809c0c 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -302,6 +302,19 @@
     /* Set mapped lifetime on the attributes */
     psa_set_key_lifetime( &attributes, lifetime );
 
+    /* Get the full key attributes from the driver in order to be able to
+     * calculate the required buffer size. */
+    status = psa_driver_wrapper_get_builtin_key(
+                slot_number, &attributes,
+                NULL, 0, NULL );
+    if( status != PSA_ERROR_BUFFER_TOO_SMALL )
+    {
+        /* Builtin keys cannot be defined by the attributes alone */
+        if( status == PSA_SUCCESS )
+            status = PSA_ERROR_CORRUPTION_DETECTED;
+        goto exit;
+    }
+
     /* If the key should exist according to the platform, then ask the driver
      * what its expected size is. */
     status = psa_driver_wrapper_get_key_buffer_size( &attributes,
@@ -310,7 +323,7 @@
         return( status );
 
     /* Allocate a buffer of the required size and load the builtin key directly
-     * into the slot buffer. */
+     * into the (now properly sized) slot buffer. */
     status = psa_allocate_buffer_to_slot( slot, key_buffer_size );
     if( status != PSA_SUCCESS )
         return( status );
diff --git a/tests/src/drivers/key_management.c b/tests/src/drivers/key_management.c
index 5daec6b..a0626fb 100644
--- a/tests/src/drivers/key_management.c
+++ b/tests/src/drivers/key_management.c
@@ -343,9 +343,6 @@
     psa_key_attributes_t *attributes,
     uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length )
 {
-    if( key_buffer_size < sizeof( psa_drv_slot_number_t ) )
-        return( PSA_ERROR_BUFFER_TOO_SMALL );
-
     switch( slot_number )
     {
         case PSA_CRYPTO_TEST_DRIVER_BUILTIN_AES_KEY_SLOT:
@@ -358,6 +355,9 @@
                 PSA_KEY_USAGE_EXPORT );
             psa_set_key_algorithm( attributes, PSA_ALG_CTR );
 
+            if( key_buffer_size < sizeof( psa_drv_slot_number_t ) )
+                return( PSA_ERROR_BUFFER_TOO_SMALL );
+
             *( (psa_drv_slot_number_t*) key_buffer ) =
                 PSA_CRYPTO_TEST_DRIVER_BUILTIN_AES_KEY_SLOT;
             *key_buffer_length = sizeof( psa_drv_slot_number_t );
@@ -375,6 +375,9 @@
             psa_set_key_algorithm(
                 attributes, PSA_ALG_ECDSA( PSA_ALG_ANY_HASH ) );
 
+            if( key_buffer_size < sizeof( psa_drv_slot_number_t ) )
+                return( PSA_ERROR_BUFFER_TOO_SMALL );
+
             *( (psa_drv_slot_number_t*) key_buffer ) =
                 PSA_CRYPTO_TEST_DRIVER_BUILTIN_ECDSA_KEY_SLOT;
             *key_buffer_length = sizeof( psa_drv_slot_number_t );