tests: psa: Fix export related tests

Fix some export related tests that were
relying on the fact that the size of the
output buffer was checked after other
parameters.

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 3f2da76..307b2c1 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -244,7 +244,7 @@
 
 PSA import/export-public: cannot export-public a symmetric key
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C
-import_export_public_key:"2b7e151628aed2a6abf7158809cf4f3c":PSA_KEY_TYPE_AES:PSA_ALG_CBC_NO_PADDING:0:PSA_ERROR_INVALID_ARGUMENT:""
+import_export_public_key:"2b7e151628aed2a6abf7158809cf4f3c":PSA_KEY_TYPE_AES:PSA_ALG_CBC_NO_PADDING:0:PSA_ERROR_INVALID_ARGUMENT:"2b7e151628aed2a6abf7158809cf4f3c"
 
 PSA import/export EC secp256r1 public key: good
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 82514d6..a65c482 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1028,19 +1028,21 @@
 
     PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
 
+    exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(
+                        psa_get_key_type( &attributes ),
+                        psa_get_key_bits( &attributes ) );
+    ASSERT_ALLOC( exported, exported_size );
+
     if( ( usage & PSA_KEY_USAGE_EXPORT ) == 0 &&
         ! PSA_KEY_TYPE_IS_PUBLIC_KEY( psa_get_key_type( &attributes ) ) )
     {
-        TEST_EQUAL( psa_export_key( key, NULL, 0, &exported_length ),
+        TEST_EQUAL( psa_export_key( key, exported,
+                                    exported_size, &exported_length ),
                     PSA_ERROR_NOT_PERMITTED );
         ok = 1;
         goto exit;
     }
 
-    exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE( psa_get_key_type( &attributes ),
-                                                psa_get_key_bits( &attributes ) );
-    ASSERT_ALLOC( exported, exported_size );
-
     PSA_ASSERT( psa_export_key( key,
                                 exported, exported_size,
                                 &exported_length ) );
@@ -1071,9 +1073,16 @@
     PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
     if( ! PSA_KEY_TYPE_IS_ASYMMETRIC( psa_get_key_type( &attributes ) ) )
     {
-        TEST_EQUAL( psa_export_public_key( key, NULL, 0, &exported_length ),
+        exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(
+                            psa_get_key_type( &attributes ),
+                            psa_get_key_bits( &attributes ) );
+        ASSERT_ALLOC( exported, exported_size );
+
+        TEST_EQUAL( psa_export_public_key( key, exported,
+                                           exported_size, &exported_length ),
                     PSA_ERROR_INVALID_ARGUMENT );
-        return( 1 );
+        ok = 1;
+        goto exit;
     }
 
     public_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index b0c660b..d14dfbb 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -343,9 +343,9 @@
                         psa_get_key_type( &read_attributes ) );
             TEST_EQUAL( psa_get_key_bits( &attributes ),
                         psa_get_key_bits( &read_attributes ) );
+            ASSERT_ALLOC( reexported, key_data->len );
             if( usage_flags & PSA_KEY_USAGE_EXPORT )
             {
-                ASSERT_ALLOC( reexported, key_data->len );
                 PSA_ASSERT( psa_export_key( id, reexported, key_data->len,
                                             &reexported_length ) );
                 ASSERT_COMPARE( key_data->x, key_data->len,
@@ -353,7 +353,8 @@
             }
             else
             {
-                TEST_EQUAL( psa_export_key( id, NULL, 0, &reexported_length ),
+                TEST_EQUAL( psa_export_key( id, reexported,
+                                key_data->len, &reexported_length ),
                             PSA_ERROR_NOT_PERMITTED );
             }
             PSA_ASSERT( psa_close_key( handle ) );