psa: export: Move registered SE support to the driver wrapper

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 766cf10..0cc89fe 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -454,6 +454,24 @@
     psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(
                                       psa_get_key_lifetime( attributes ) );
 
+    /* Try dynamically-registered SE interface first */
+#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( attributes->core.lifetime, &drv, &drv_context ) )
+    {
+        if( ( drv->key_management == NULL   ) ||
+            ( drv->key_management->p_export == NULL ) )
+            return( PSA_ERROR_NOT_SUPPORTED );
+
+        return( drv->key_management->p_export(
+                     drv_context,
+                     *( (psa_key_slot_number_t *)key_buffer ),
+                     data, data_size, data_length ) );
+    }
+#endif /* PSA_CRYPTO_SE_C */
+
     switch( location )
     {
         case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -492,6 +510,24 @@
     psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(
                                       psa_get_key_lifetime( attributes ) );
 
+    /* Try dynamically-registered SE interface first */
+#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( attributes->core.lifetime, &drv, &drv_context ) )
+    {
+        if( ( drv->key_management == NULL ) ||
+            ( drv->key_management->p_export_public == NULL ) )
+            return( PSA_ERROR_NOT_SUPPORTED );
+
+        return( drv->key_management->p_export_public(
+                    drv_context,
+                    *( (psa_key_slot_number_t *)key_buffer ),
+                    data, data_size, data_length ) );
+    }
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
     switch( location )
     {
         case PSA_KEY_LOCATION_LOCAL_STORAGE: