Add & splice in test driver for ECC keygen

Signed-off-by: Steven Cooreman <steven.cooreman@silabs.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 7da2ff6..03326f7 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5987,29 +5987,15 @@
     if( status != PSA_SUCCESS )
         goto exit;
 
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    if( driver != NULL )
-    {
-        const psa_drv_se_t *drv = psa_get_se_driver_methods( driver );
-        size_t pubkey_length = 0; /* We don't support this feature yet */
-        if( drv->key_management == NULL ||
-            drv->key_management->p_generate == NULL )
-        {
-            status = PSA_ERROR_NOT_SUPPORTED;
-            goto exit;
-        }
-        status = drv->key_management->p_generate(
-            psa_get_se_driver_context( driver ),
-            slot->data.se.slot_number, attributes,
-            NULL, 0, &pubkey_length );
-    }
-    else
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    {
-        status = psa_generate_key_internal(
-            slot, attributes->core.bits,
-            attributes->domain_parameters, attributes->domain_parameters_size );
-    }
+    status = psa_driver_wrapper_generate_key( attributes,
+                                              slot );
+    if( status != PSA_ERROR_NOT_SUPPORTED ||
+        psa_key_lifetime_is_external( attributes->core.lifetime ) )
+        goto exit;
+
+    status = psa_generate_key_internal(
+        slot, attributes->core.bits,
+        attributes->domain_parameters, attributes->domain_parameters_size );
 
 exit:
     if( status == PSA_SUCCESS )
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. */
diff --git a/tests/include/drivers/keygen.h b/tests/include/drivers/keygen.h
new file mode 100644
index 0000000..436df34
--- /dev/null
+++ b/tests/include/drivers/keygen.h
@@ -0,0 +1,49 @@
+/*
+ * Test driver for signature functions
+ */
+/*  Copyright (C) 2020, ARM Limited, All Rights Reserved
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ *  This file is part of mbed TLS (https://tls.mbed.org)
+ */
+
+#ifndef MBEDTLS_PSA_CRYPTO_TEST_DRIVERS_KEYGEN_H
+#define MBEDTLS_PSA_CRYPTO_TEST_DRIVERS_KEYGEN_H
+
+#if !defined(MBEDTLS_CONFIG_FILE)
+#include "mbedtls/config.h"
+#else
+#include MBEDTLS_CONFIG_FILE
+#endif
+
+#if defined(MBEDTLS_TEST_HOOKS)
+#include <psa/crypto_driver_common.h>
+
+extern void *test_driver_keygen_forced_output;
+extern size_t test_driver_keygen_forced_output_length;
+
+extern psa_status_t test_transparent_keygen_status;
+extern unsigned long test_transparent_keygen_hit;
+
+psa_status_t test_transparent_generate_key(
+    const psa_key_attributes_t *attributes,
+    uint8_t *key, size_t key_size, size_t *key_length );
+
+psa_status_t test_opaque_generate_key(
+    const psa_key_attributes_t *attributes,
+    uint8_t *key, size_t key_size, size_t *key_length );
+
+#endif /* MBEDTLS_TEST_HOOKS */
+#endif /* MBEDTLS_PSA_CRYPTO_TEST_DRIVERS_KEYGEN_H */
diff --git a/tests/include/drivers/test_driver.h b/tests/include/drivers/test_driver.h
index 5494674..fec305f 100644
--- a/tests/include/drivers/test_driver.h
+++ b/tests/include/drivers/test_driver.h
@@ -25,5 +25,6 @@
 #define MBEDTLS_PSA_CRYPTO_TEST_DRIVER_LIFETIME 0x7fffff
 
 #include "drivers/signature.h"
+#include "drivers/keygen.h"
 
 #endif /* MBEDTLS_PSA_CRYPTO_TEST_DRIVER_H */
diff --git a/tests/src/drivers/keygen.c b/tests/src/drivers/keygen.c
new file mode 100644
index 0000000..4c830c3
--- /dev/null
+++ b/tests/src/drivers/keygen.c
@@ -0,0 +1,129 @@
+/*
+ * Test driver for signature functions
+ */
+/*  Copyright (C) 2020, ARM Limited, All Rights Reserved
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ *  This file is part of mbed TLS (https://tls.mbed.org)
+ */
+
+#if !defined(MBEDTLS_CONFIG_FILE)
+#include "mbedtls/config.h"
+#else
+#include MBEDTLS_CONFIG_FILE
+#endif
+
+#if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(MBEDTLS_TEST_HOOKS)
+#include "psa/crypto.h"
+#include "mbedtls/ecp.h"
+#include "mbedtls/error.h"
+
+#include "drivers/keygen.h"
+
+#include "test/random.h"
+
+#include <string.h>
+
+/* If non-null, on success, copy this to the output. */
+void *test_driver_keygen_forced_output = NULL;
+size_t test_driver_keygen_forced_output_length = 0;
+
+psa_status_t test_transparent_keygen_status = PSA_ERROR_NOT_SUPPORTED;
+unsigned long test_transparent_keygen_hit = 0;
+
+psa_status_t test_transparent_generate_key(
+    const psa_key_attributes_t *attributes,
+    uint8_t *key, size_t key_size, size_t *key_length )
+{
+    ++test_transparent_keygen_hit;
+
+    if( test_transparent_keygen_status != PSA_SUCCESS )
+        return( test_transparent_keygen_status );
+
+    if( test_driver_keygen_forced_output != NULL )
+    {
+        if( test_driver_keygen_forced_output_length > key_size )
+            return( PSA_ERROR_BUFFER_TOO_SMALL );
+        memcpy( key, test_driver_keygen_forced_output,
+                test_driver_keygen_forced_output_length );
+        *key_length = test_driver_keygen_forced_output_length;
+        return( PSA_SUCCESS );
+    }
+
+    /* Copied from psa_crypto.c */
+#if defined(MBEDTLS_ECP_C)
+    if ( PSA_KEY_TYPE_IS_ECC( attributes->core.type ) && PSA_KEY_TYPE_IS_KEY_PAIR( attributes->core.type ) )
+    {
+        psa_ecc_family_t curve = PSA_KEY_TYPE_ECC_GET_FAMILY( attributes->core.type );
+        mbedtls_ecp_group_id grp_id =
+            mbedtls_ecc_group_of_psa( curve, PSA_BITS_TO_BYTES( attributes->core.bits ) );
+        const mbedtls_ecp_curve_info *curve_info =
+            mbedtls_ecp_curve_info_from_grp_id( grp_id );
+        mbedtls_ecp_keypair ecp;
+        mbedtls_test_rnd_pseudo_info rnd_info;
+        memset( &rnd_info, 0x5A, sizeof( mbedtls_test_rnd_pseudo_info ) );
+
+        int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+        if( attributes->domain_parameters_size != 0 )
+            return( PSA_ERROR_NOT_SUPPORTED );
+        if( grp_id == MBEDTLS_ECP_DP_NONE || curve_info == NULL )
+            return( PSA_ERROR_NOT_SUPPORTED );
+        if( curve_info->bit_size != attributes->core.bits )
+            return( PSA_ERROR_INVALID_ARGUMENT );
+        mbedtls_ecp_keypair_init( &ecp );
+        ret = mbedtls_ecp_gen_key( grp_id, &ecp,
+                                   &mbedtls_test_rnd_pseudo_rand,
+                                   &rnd_info );
+        if( ret != 0 )
+        {
+            mbedtls_ecp_keypair_free( &ecp );
+            return( mbedtls_to_psa_error( ret ) );
+        }
+
+        /* Make sure to use export representation */
+        size_t bytes = PSA_BITS_TO_BYTES( attributes->core.bits );
+        if( key_size < bytes )
+        {
+            mbedtls_ecp_keypair_free( &ecp );
+            return( PSA_ERROR_BUFFER_TOO_SMALL );
+        }
+        psa_status_t status = mbedtls_to_psa_error(
+            mbedtls_mpi_write_binary( &ecp.d, key, bytes ) );
+
+        if( status == PSA_SUCCESS )
+        {
+            *key_length = bytes;
+        }
+
+        mbedtls_ecp_keypair_free( &ecp );
+        return( status );
+    }
+    else
+#endif /* MBEDTLS_ECP_C */
+    return( PSA_ERROR_NOT_SUPPORTED );
+}
+
+psa_status_t test_opaque_generate_key(
+    const psa_key_attributes_t *attributes,
+    uint8_t *key, size_t key_size, size_t *key_length )
+{
+    (void) attributes;
+    (void) key;
+    (void) key_size;
+    (void) key_length;
+    return( PSA_ERROR_NOT_SUPPORTED );
+}
+
+#endif /* MBEDTLS_PSA_CRYPTO_DRIVERS && MBEDTLS_TEST_HOOKS */
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index 2d2c5c4..ddf283f 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -27,3 +27,15 @@
 
 verify_hash through transparent driver: error
 ecdsa_verify:PSA_ERROR_GENERIC_ERROR:1:PSA_ERROR_GENERIC_ERROR
+
+generate_key through transparent driver: fake
+generate_key:PSA_SUCCESS:1:PSA_SUCCESS
+
+generate_key through transparent driver: in-driver
+generate_key:PSA_SUCCESS:0:PSA_SUCCESS
+
+generate_key through transparent driver: fallback
+generate_key:PSA_ERROR_NOT_SUPPORTED:0:PSA_SUCCESS
+
+generate_key through transparent driver: error
+generate_key:PSA_ERROR_GENERIC_ERROR:0:PSA_ERROR_GENERIC_ERROR
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index a954446..9bb7943 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -171,3 +171,73 @@
     test_driver_forced_output_length = 0;
 }
 /* END_CASE */
+
+/* BEGIN_CASE depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED */
+void generate_key( int force_status_arg,
+                   int fake_output,
+                   int expected_status_arg )
+{
+    psa_status_t force_status = force_status_arg;
+    psa_status_t expected_status = expected_status_arg;
+    psa_key_handle_t handle = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_algorithm_t alg = PSA_ALG_ECDSA( PSA_ALG_SHA_256 );
+    const uint8_t *expected_output;
+    size_t expected_output_length;
+    psa_status_t actual_status;
+    uint8_t actual_output[sizeof(test_secp256r1_key_data)] = {0};
+    size_t actual_output_length;
+
+    psa_set_key_type( &attributes,
+                      PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_CURVE_SECP_R1 ) );
+    psa_set_key_bits( &attributes, 256 );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_EXPORT );
+    psa_set_key_algorithm( &attributes, alg );
+
+    if( fake_output )
+    {
+        expected_output = test_driver_keygen_forced_output = test_secp256r1_key_data;
+        expected_output_length = test_driver_keygen_forced_output_length =
+            sizeof( test_secp256r1_key_data );
+    }
+
+    test_transparent_keygen_hit = 0;
+    test_transparent_keygen_status = force_status;
+
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    actual_status = psa_generate_key( &attributes, &handle );
+    fprintf(stdout, "rteturn %d\n", actual_status);
+
+    TEST_EQUAL( test_transparent_keygen_hit, 1 );
+    TEST_EQUAL( actual_status, expected_status );
+
+    if( actual_status == PSA_SUCCESS )
+    {
+        psa_export_key( handle, actual_output, sizeof(actual_output), &actual_output_length );
+
+        if( fake_output )
+        {
+            ASSERT_COMPARE( actual_output, actual_output_length,
+                            expected_output, expected_output_length );
+        }
+        else
+        {
+            size_t zeroes = 0;
+            for( size_t i = 0; i < sizeof(actual_output); i++ )
+            {
+                if( actual_output[i] == 0)
+                    zeroes++;
+            }
+            TEST_ASSERT( zeroes != sizeof(actual_output) );
+        }
+    }
+exit:
+    psa_reset_key_attributes( &attributes );
+    psa_destroy_key( handle );
+    PSA_DONE( );
+    test_transparent_signature_sign_hash_status = PSA_ERROR_NOT_SUPPORTED;
+    test_driver_keygen_forced_output = NULL;
+    test_driver_keygen_forced_output_length = 0;
+}
+/* END_CASE */