Apply review feedback

Signed-off-by: Steven Cooreman <steven.cooreman@silabs.com>
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index f30b75a..3797dea 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -247,43 +247,35 @@
     if( status != PSA_SUCCESS )
         goto exit;
 
-    if( psa_key_lifetime_is_external( slot->attr.lifetime ) )
-    {
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-        const psa_drv_se_t *drv;
-        psa_drv_se_context_t *drv_context;
-        if( psa_get_se_driver( slot->attr.lifetime, &drv, &drv_context ) )
-        {
-            psa_se_key_data_storage_t *data;
-            if( key_data_length != sizeof( *data ) )
-            {
-                status = PSA_ERROR_STORAGE_FAILURE;
-                goto exit;
-            }
-            data = (psa_se_key_data_storage_t *) key_data;
-            memcpy( &slot->data.se.slot_number, &data->slot_number,
-                    sizeof( slot->data.se.slot_number ) );
-        }
-        else
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-        {
-            /* A key that is successfully loaded from storage with an
-             * external lifetime, but doesn't belong to an SE driver,
-             * must be a PSA driver-associated key which we can just
-             * load like an internal key. */
-            if ( key_data == NULL )
-            {
-                status = PSA_ERROR_STORAGE_FAILURE;
-                goto exit;
-            }
-
-            status = psa_copy_key_material_into_slot( slot, key_data, key_data_length );
-        }
-    }
-    else
+    /* Special handling is required for loading keys associated with a
+     * dynamically registered SE interface. */
+    const psa_drv_se_t *drv;
+    psa_drv_se_context_t *drv_context;
+    if( psa_get_se_driver( slot->attr.lifetime, &drv, &drv_context ) )
     {
-        status = psa_copy_key_material_into_slot( slot, key_data, key_data_length );
+        psa_se_key_data_storage_t *data;
+        if( key_data_length != sizeof( *data ) )
+        {
+            status = PSA_ERROR_STORAGE_FAILURE;
+            goto exit;
+        }
+        data = (psa_se_key_data_storage_t *) key_data;
+        memcpy( &slot->data.se.slot_number, &data->slot_number,
+                sizeof( slot->data.se.slot_number ) );
+
+        status = PSA_SUCCESS;
+        goto exit;
     }
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
+    if ( key_data == NULL )
+    {
+        status = PSA_ERROR_STORAGE_FAILURE;
+        goto exit;
+    }
+
+    status = psa_copy_key_material_into_slot( slot, key_data, key_data_length );
 
 exit:
     psa_free_persistent_key_data( key_data, key_data_length );
@@ -360,31 +352,26 @@
     if ( psa_key_lifetime_is_external( lifetime ) )
     {
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+        /* Check whether a driver is registered against this lifetime */
         psa_se_drv_table_entry_t *driver = psa_get_se_driver_entry( lifetime );
-        if( driver == NULL )
-        {
-#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS)
-            /* Key location for external keys gets checked by the wrapper */
-            return( PSA_SUCCESS );
-#else
-            return( PSA_ERROR_INVALID_ARGUMENT );
-#endif
-        }
-        else
+        if( driver != NULL )
         {
             if (p_drv != NULL)
                 *p_drv = driver;
             return( PSA_SUCCESS );
         }
-#else
+#else /* MBEDTLS_PSA_CRYPTO_SE_C */
         (void) p_drv;
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
 #if defined(MBEDTLS_PSA_CRYPTO_DRIVERS)
         /* Key location for external keys gets checked by the wrapper */
         return( PSA_SUCCESS );
-#else
+#else /* MBEDTLS_PSA_CRYPTO_DRIVERS */
+        /* No support for external lifetimes at all, or dynamic interface
+         * did not find driver for requested lifetime. */
         return( PSA_ERROR_INVALID_ARGUMENT );
-#endif
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVERS */
     }
     else
         /* Local/internal keys are always valid */