Add & splice in test driver for ECC keygen

Signed-off-by: Steven Cooreman <steven.cooreman@silabs.com>
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 09a7ccd..0f6cfe9 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -23,6 +23,9 @@
 
 #include "psa_crypto_core.h"
 #include "psa_crypto_driver_wrappers.h"
+#include "mbedtls/platform.h"
+
+#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS)
 
 /* Include test driver definition when running tests */
 #if defined(MBEDTLS_TEST_HOOKS)
@@ -31,16 +34,18 @@
 #undef MBEDTLS_PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT
 #define MBEDTLS_PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT
 #include "drivers/test_driver.h"
-#endif
+#endif /* MBEDTLS_TEST_HOOKS */
 
+/* Include driver definition file for each registered driver here */
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVERS */
+
+/* Support the 'old' SE interface when asked to */
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
 #undef MBEDTLS_PSA_CRYPTO_DRIVER_PRESENT
 #define MBEDTLS_PSA_CRYPTO_DRIVER_PRESENT
 #include "psa_crypto_se.h"
 #endif
 
-/* Include driver definition file for each registered driver */
-
 /* Start delegation functions */
 psa_status_t psa_driver_wrapper_sign_hash( psa_key_slot_t *slot,
                                            psa_algorithm_t alg,
@@ -223,13 +228,153 @@
 #endif /* MBEDTLS_PSA_CRYPTO_DRIVER_PRESENT */
 }
 
+#if defined(MBEDTLS_PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+static psa_status_t get_expected_key_size( const psa_key_attributes_t *attributes,
+                                           size_t *expected_size )
+{
+    if( PSA_KEY_LIFETIME_GET_LOCATION( attributes->core.lifetime ) == PSA_KEY_LOCATION_LOCAL_STORAGE )
+    {
+        if( PSA_KEY_TYPE_IS_UNSTRUCTURED( attributes->core.type ) )
+        {
+            *expected_size = PSA_BITS_TO_BYTES( attributes->core.bits );
+            return PSA_SUCCESS;
+        }
+
+        if( PSA_KEY_TYPE_IS_ECC( attributes->core.type ) )
+        {
+            if( PSA_KEY_TYPE_IS_KEY_PAIR( attributes->core.type ) )
+            {
+                *expected_size = PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE( attributes->core.bits );
+                return PSA_SUCCESS;
+            }
+            else
+            {
+                *expected_size = PSA_KEY_EXPORT_ECC_PUBLIC_KEY_MAX_SIZE( attributes->core.bits );
+                return PSA_SUCCESS;
+            }
+        }
+
+        if( PSA_KEY_TYPE_IS_RSA( attributes->core.type ) )
+        {
+            if( PSA_KEY_TYPE_IS_KEY_PAIR( attributes->core.type ) )
+            {
+                *expected_size = PSA_KEY_EXPORT_RSA_KEY_PAIR_MAX_SIZE( attributes->core.bits );
+                return PSA_SUCCESS;
+            }
+            else
+            {
+                *expected_size = PSA_KEY_EXPORT_RSA_PUBLIC_KEY_MAX_SIZE( attributes->core.bits );
+                return PSA_SUCCESS;
+            }
+        }
+
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+    else
+    {
+        /* TBD: opaque driver support, need to calculate size through driver-defined size function */
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+}
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVER_PRESENT */
+
 psa_status_t psa_driver_wrapper_generate_key( const psa_key_attributes_t *attributes,
                                               psa_key_slot_t *slot )
 {
+#if defined(MBEDTLS_PSA_CRYPTO_DRIVER_PRESENT)
+    /* 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( slot->attr.lifetime, &drv, &drv_context ) )
+    {
+        size_t pubkey_length = 0; /* We don't support this feature yet */
+        if( drv->key_management == NULL ||
+            drv->key_management->p_generate == NULL )
+        {
+            /* Key is defined as being in SE, but we have no way to generate it */
+            return PSA_ERROR_NOT_SUPPORTED;
+        }
+        return( drv->key_management->p_generate(
+            drv_context,
+            slot->data.se.slot_number, attributes,
+            NULL, 0, &pubkey_length ) );
+    }
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
+    /* Then try accelerator API */
+#if defined(MBEDTLS_PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+    psa_status_t status = PSA_ERROR_INVALID_ARGUMENT;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(slot->attr.lifetime);
+    size_t export_size = 0;
+
+    status = get_expected_key_size( attributes, &export_size );
+    if( status != PSA_SUCCESS )
+        return status;
+
+    slot->data.key.data = mbedtls_calloc(1, export_size);
+    if( slot->data.key.data == NULL )
+        return PSA_ERROR_INSUFFICIENT_MEMORY;
+    slot->data.key.bytes = export_size;
+
+    switch( location )
+    {
+        case PSA_KEY_LOCATION_LOCAL_STORAGE:
+            /* Key is stored in the slot in export representation, so
+             * cycle through all known transparent accelerators */
+
+            /* Transparent drivers are limited to generating asymmetric keys */
+            if( ! PSA_KEY_TYPE_IS_ASYMMETRIC( slot->attr.type ) )
+            {
+                status = PSA_ERROR_NOT_SUPPORTED;
+                break;
+            }
+#if defined(MBEDTLS_TEST_HOOKS)
+            status = test_transparent_generate_key( attributes,
+                                                    slot->data.key.data,
+                                                    slot->data.key.bytes,
+                                                    &slot->data.key.bytes );
+            /* Declared with fallback == true */
+            if( status != PSA_ERROR_NOT_SUPPORTED )
+                break;
+#endif /* MBEDTLS_TEST_HOOKS */
+            /* Fell through, meaning no accelerator supports this operation */
+            status = PSA_ERROR_NOT_SUPPORTED;
+            break;
+        /* Add cases for opaque driver here */
+#if defined(MBEDTLS_TEST_HOOKS)
+        case MBEDTLS_PSA_CRYPTO_TEST_DRIVER_LIFETIME:
+            status = test_opaque_generate_key( attributes,
+                                               slot->data.key.data,
+                                               slot->data.key.bytes,
+                                               &slot->data.key.bytes );
+            break;
+#endif /* MBEDTLS_TEST_HOOKS */
+        default:
+            /* Key is declared with a lifetime not known to us */
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            break;
+    }
+
+    if( status != PSA_SUCCESS )
+    {
+        /* free allocated buffer */
+        mbedtls_free( slot->data.key.data );
+        slot->data.key.data = NULL;
+        slot->data.key.bytes = 0;
+    }
+
+    return( status );
+#else /* MBEDTLS_PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+    return PSA_ERROR_NOT_SUPPORTED;
+#endif /* MBEDTLS_PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+#else /* MBEDTLS_PSA_CRYPTO_DRIVER_PRESENT */
     (void) attributes;
     (void) slot;
 
     return PSA_ERROR_NOT_SUPPORTED;
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVER_PRESENT */
 }
 
 /* End of automatically generated file. */