psa: Add ECP/RSA transparent test driver export_public_key entry point

Add ECP/RSA transparent test driver export_public_key
entry point and use it in the transparent test driver
supporting both ECP and RSA.

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/library/psa_crypto_ecp.c b/library/psa_crypto_ecp.c
index 4f20621..265d1ac 100644
--- a/library/psa_crypto_ecp.c
+++ b/library/psa_crypto_ecp.c
@@ -38,8 +38,20 @@
 #include <mbedtls/ecp.h>
 #include <mbedtls/error.h>
 
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) || \
+#if ( defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||  \
+      ( defined(PSA_CRYPTO_DRIVER_TEST) &&                   \
+        defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR) ) )
+#define BUILTIN_KEY_TYPE_ECC_KEY_PAIR    1
+#endif
+
+#if ( defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) ||  \
+      ( defined(PSA_CRYPTO_DRIVER_TEST) &&                   \
+        defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY) ) )
+#define BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY  1
+#endif
+
+#if defined(BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_ECDSA) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_DETERMINISTIC_ECDSA)
@@ -131,14 +143,14 @@
 
     return( status );
 }
-#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) ||
+#endif /* defined(BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
+        * defined(BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) ||
         * defined(MBEDTLS_PSA_BUILTIN_ALG_ECDSA) ||
         * defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) ||
         * defined(MBEDTLS_PSA_BUILTIN_ALG_DETERMINISTIC_ECDSA) */
 
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
+#if defined(BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
 psa_status_t mbedtls_psa_ecp_export_key( psa_key_type_t type,
                                          mbedtls_ecp_keypair *ecp,
                                          uint8_t *data,
@@ -190,7 +202,7 @@
     }
 }
 
-psa_status_t mbedtls_psa_ecp_export_public_key(
+static psa_status_t ecp_export_public_key(
     const psa_key_attributes_t *attributes,
     const uint8_t *key_buffer, size_t key_buffer_size,
     uint8_t *data, size_t data_size, size_t *data_length )
@@ -213,7 +225,43 @@
 
     return( status );
 }
+#endif /* defined(BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
+        * defined(BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) */
+
+#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
+
+psa_status_t mbedtls_psa_ecp_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    return( ecp_export_public_key( attributes, key_buffer, key_buffer_size,
+                                   data, data_size, data_length ) );
+}
+
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) */
 
+/*
+ * BEYOND THIS POINT, TEST DRIVER ENTRY POINTS ONLY.
+ */
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY)
+psa_status_t mbedtls_transparent_test_driver_ecp_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    return( ecp_export_public_key( attributes, key_buffer, key_buffer_size,
+                                   data, data_size, data_length ) );
+}
+#endif /* defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR) ||
+          defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY) */
+
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+
 #endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/library/psa_crypto_ecp.h b/library/psa_crypto_ecp.h
index c6d8437..1083c9b 100644
--- a/library/psa_crypto_ecp.h
+++ b/library/psa_crypto_ecp.h
@@ -82,4 +82,15 @@
     const uint8_t *key_buffer, size_t key_buffer_size,
     uint8_t *data, size_t data_size, size_t *data_length );
 
+/*
+ * BEYOND THIS POINT, TEST DRIVER ENTRY POINTS ONLY.
+ */
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+psa_status_t mbedtls_transparent_test_driver_ecp_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+
 #endif /* PSA_CRYPTO_ECP_H */
diff --git a/library/psa_crypto_rsa.c b/library/psa_crypto_rsa.c
index 3e5d1b4..1814773 100644
--- a/library/psa_crypto_rsa.c
+++ b/library/psa_crypto_rsa.c
@@ -39,12 +39,24 @@
 #include <mbedtls/pk.h>
 #include <mbedtls/pk_internal.h>
 
+#if ( defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||  \
+      ( defined(PSA_CRYPTO_DRIVER_TEST) &&                   \
+        defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR) ) )
+#define BUILTIN_KEY_TYPE_RSA_KEY_PAIR    1
+#endif
+
+#if ( defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) ||  \
+      ( defined(PSA_CRYPTO_DRIVER_TEST) &&                   \
+        defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY) ) )
+#define BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY  1
+#endif
+
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_CRYPT) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_SIGN) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_OAEP) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PSS) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
+    defined(BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
 
 /* Mbed TLS doesn't support non-byte-aligned key sizes (i.e. key sizes
  * that are not a multiple of 8) well. For example, there is only
@@ -121,11 +133,11 @@
         * defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_SIGN) ||
         * defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_OAEP) ||
         * defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PSS) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
+        * defined(BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
+        * defined(BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
 
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
+#if defined(BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
 psa_status_t mbedtls_psa_rsa_export_key( psa_key_type_t type,
                                          mbedtls_rsa_context *rsa,
                                          uint8_t *data,
@@ -182,7 +194,7 @@
 #endif /* MBEDTLS_PK_WRITE_C */
 }
 
-psa_status_t mbedtls_psa_rsa_export_public_key(
+static psa_status_t rsa_export_public_key(
     const psa_key_attributes_t *attributes,
     const uint8_t *key_buffer, size_t key_buffer_size,
     uint8_t *data, size_t data_size, size_t *data_length )
@@ -206,7 +218,43 @@
 
     return( status );
 }
+#endif /* defined(BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
+        * defined(BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
+
+#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
+
+psa_status_t mbedtls_psa_rsa_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    return( rsa_export_public_key( attributes, key_buffer, key_buffer_size,
+                                   data, data_size, data_length ) );
+}
+
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
 
+/*
+ * BEYOND THIS POINT, TEST DRIVER ENTRY POINTS ONLY.
+ */
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY)
+psa_status_t mbedtls_transparent_test_driver_rsa_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    return( rsa_export_public_key( attributes, key_buffer, key_buffer_size,
+                                   data, data_size, data_length ) );
+}
+#endif /* defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR) ||
+          defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY) */
+
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+
 #endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/library/psa_crypto_rsa.h b/library/psa_crypto_rsa.h
index 2fb5442..36756ed 100644
--- a/library/psa_crypto_rsa.h
+++ b/library/psa_crypto_rsa.h
@@ -82,4 +82,17 @@
     const uint8_t *key_buffer, size_t key_buffer_size,
     uint8_t *data, size_t data_size, size_t *data_length );
 
+/*
+ * BEYOND THIS POINT, TEST DRIVER ENTRY POINTS ONLY.
+ */
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+
+psa_status_t mbedtls_transparent_test_driver_rsa_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+
 #endif /* PSA_CRYPTO_RSA_H */
diff --git a/tests/src/drivers/key_management.c b/tests/src/drivers/key_management.c
index 76d2b3f..aef3b89 100644
--- a/tests/src/drivers/key_management.c
+++ b/tests/src/drivers/key_management.c
@@ -27,6 +27,8 @@
 #if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(PSA_CRYPTO_DRIVER_TEST)
 #include "psa/crypto.h"
 #include "psa_crypto_core.h"
+#include "psa_crypto_ecp.h"
+#include "psa_crypto_rsa.h"
 #include "mbedtls/ecp.h"
 #include "mbedtls/error.h"
 
@@ -265,9 +267,10 @@
     (void) data_length;
     return( PSA_ERROR_NOT_SUPPORTED );
 }
+
 psa_status_t test_transparent_export_public_key(
     const psa_key_attributes_t *attributes,
-    const uint8_t *key, size_t key_length,
+    const uint8_t *key_buffer, size_t key_buffer_size,
     uint8_t *data, size_t data_size, size_t *data_length )
 {
     ++test_driver_key_management_hooks.hits;
@@ -285,73 +288,39 @@
         return( PSA_SUCCESS );
     }
 
-    if( key == NULL || key_length == 0 )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-
-    psa_key_type_t keytype = psa_get_key_type( attributes );
-    (void) keytype;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_type_t key_type = psa_get_key_type( attributes );
 
 #if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR) || \
     defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY)
-    if( PSA_KEY_TYPE_IS_ECC( keytype ) )
+    if( PSA_KEY_TYPE_IS_ECC( key_type ) )
     {
-        if( !PSA_KEY_TYPE_IS_KEY_PAIR( keytype ) )
-            return( PSA_ERROR_INVALID_ARGUMENT );
-
-        /* Mostly copied from psa_crypto.c */
-        mbedtls_ecp_group_id grp_id = MBEDTLS_ECP_DP_NONE;
-        psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-        mbedtls_ecp_keypair ecp;
-        mbedtls_test_rnd_pseudo_info rnd_info;
-        memset( &rnd_info, 0x5A, sizeof( mbedtls_test_rnd_pseudo_info ) );
-
-        if( attributes->domain_parameters_size != 0 )
-            return( PSA_ERROR_NOT_SUPPORTED );
-
-        grp_id = mbedtls_ecc_group_of_psa( PSA_KEY_TYPE_ECC_GET_FAMILY( keytype ),
-                                           PSA_BITS_TO_BYTES( psa_get_key_bits( attributes ) ) );
-        if( grp_id == MBEDTLS_ECP_DP_NONE )
-            return( PSA_ERROR_NOT_SUPPORTED );
-
-        mbedtls_ecp_keypair_init( &ecp );
-
-        status = mbedtls_to_psa_error(
-                    mbedtls_ecp_group_load( &ecp.grp, grp_id ) );
-        if( status != PSA_SUCCESS )
-            goto ecp_exit;
-
-        status = mbedtls_to_psa_error(
-            mbedtls_ecp_read_key( ecp.grp.id,
-                                  &ecp,
-                                  key,
-                                  key_length ) );
-        if( status != PSA_SUCCESS )
-            goto ecp_exit;
-
-        /* Calculate the public key */
-        status = mbedtls_to_psa_error(
-            mbedtls_ecp_mul( &ecp.grp, &ecp.Q, &ecp.d, &ecp.grp.G,
-                             &mbedtls_test_rnd_pseudo_rand,
-                             &rnd_info ) );
-        if( status != PSA_SUCCESS )
-            goto ecp_exit;
-
-        status = mbedtls_to_psa_error(
-                    mbedtls_ecp_point_write_binary( &ecp.grp, &ecp.Q,
-                                                    MBEDTLS_ECP_PF_UNCOMPRESSED,
-                                                    data_length,
-                                                    data,
-                                                    data_size ) );
-        if( status != PSA_SUCCESS )
-            memset( data, 0, data_size );
-ecp_exit:
-        mbedtls_ecp_keypair_free( &ecp );
-        return( status );
+        status = mbedtls_transparent_test_driver_ecp_export_public_key(
+                      attributes,
+                      key_buffer, key_buffer_size,
+                      data, data_size, data_length );
     }
-#endif /* MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR ||
-        * MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY */
+    else
+#endif
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY)
+    if( PSA_KEY_TYPE_IS_RSA( key_type ) )
+    {
+        status = mbedtls_transparent_test_driver_rsa_export_public_key(
+                      attributes,
+                      key_buffer, key_buffer_size,
+                      data, data_size, data_length );
+    }
+    else
+#endif
+    {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        (void)key_buffer;
+        (void)key_buffer_size;
+        (void)key_type;
+    }
 
-    return( PSA_ERROR_NOT_SUPPORTED );
+    return( status );
 }
 
 psa_status_t test_opaque_export_public_key(