Adapt test_suite_ssl for psa crypto

Signed-off-by: Przemyslaw Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 7a0b1f7..acfc3a4 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -1196,6 +1196,14 @@
     mbedtls_cipher_info_t const *cipher_info;
     int ret = 0;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_key_type_t key_type;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_algorithm_t alg;
+    size_t key_bits;
+    psa_status_t status;
+#endif
+
     size_t keylen, maclen, ivlen;
     unsigned char *key0 = NULL, *key1 = NULL;
     unsigned char *md0 = NULL, *md1 = NULL;
@@ -1230,6 +1238,7 @@
     memset( key0, 0x1, keylen );
     memset( key1, 0x2, keylen );
 
+#if !defined(MBEDTLS_USE_PSA_CRYPTO)
     /* Setup cipher contexts */
     CHK( mbedtls_cipher_setup( &t_in->cipher_ctx_enc,  cipher_info ) == 0 );
     CHK( mbedtls_cipher_setup( &t_in->cipher_ctx_dec,  cipher_info ) == 0 );
@@ -1258,6 +1267,7 @@
                                 keylen << 3, MBEDTLS_ENCRYPT ) == 0 );
     CHK( mbedtls_cipher_setkey( &t_out->cipher_ctx_dec, key0,
                                 keylen << 3, MBEDTLS_DECRYPT ) == 0 );
+#endif
 
     /* Setup MAC contexts */
 #if defined(MBEDTLS_SSL_SOME_SUITES_USE_MAC)
@@ -1420,6 +1430,74 @@
     t_out->out_cid_len = cid0_len;
 #endif /* MBEDTLS_SSL_DTLS_CONNECTION_ID */
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    status = tls_mbedtls_cipher_to_psa( cipher_type,
+                            t_in->taglen,
+                            &alg,
+                            &key_type,
+                            &key_bits );
+
+    if ( status != PSA_SUCCESS)
+    {
+        ret = psa_status_to_mbedtls( status );
+        goto cleanup;
+    }
+
+    t_in->psa_alg = alg;
+    t_out->psa_alg = alg;
+
+    if ( alg != MBEDTLS_SSL_NULL_CIPHER )
+    {
+        psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
+        psa_set_key_algorithm( &attributes, alg );
+        psa_set_key_type( &attributes, key_type );
+
+        status = psa_import_key( &attributes,
+                        key0,
+                        PSA_BITS_TO_BYTES( key_bits ),
+                        &t_in->psa_key_enc );
+
+        if ( status != PSA_SUCCESS)
+        {
+            ret = psa_status_to_mbedtls( status );
+            goto cleanup;
+        }
+
+        status = psa_import_key( &attributes,
+                        key1,
+                        PSA_BITS_TO_BYTES( key_bits ),
+                        &t_in->psa_key_dec );
+
+        if ( status != PSA_SUCCESS)
+        {
+            ret = psa_status_to_mbedtls( status );
+            goto cleanup;
+        }
+
+        status = psa_import_key( &attributes,
+                        key1,
+                        PSA_BITS_TO_BYTES( key_bits ),
+                        &t_out->psa_key_enc );
+
+        if ( status != PSA_SUCCESS)
+        {
+            ret = psa_status_to_mbedtls( status );
+            goto cleanup;
+        }
+
+        status = psa_import_key( &attributes,
+                        key0,
+                        PSA_BITS_TO_BYTES( key_bits ),
+                        &t_out->psa_key_dec );
+
+        if ( status != PSA_SUCCESS)
+        {
+            ret = psa_status_to_mbedtls( status );
+            goto cleanup;
+        }
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
 cleanup:
 
     mbedtls_free( key0 );
@@ -3178,13 +3256,17 @@
     size_t const buflen = 512;
     mbedtls_record rec, rec_backup;
 
+    USE_PSA_INIT( );
+
     mbedtls_ssl_init( &ssl );
     mbedtls_ssl_transform_init( &t0 );
     mbedtls_ssl_transform_init( &t1 );
-    TEST_ASSERT( build_transforms( &t0, &t1, cipher_type, hash_id,
-                                   etm, tag_mode, ver,
-                                   (size_t) cid0_len,
-                                   (size_t) cid1_len ) == 0 );
+    ret = build_transforms( &t0, &t1, cipher_type, hash_id,
+                            etm, tag_mode, ver,
+                            (size_t) cid0_len,
+                            (size_t) cid1_len );
+
+    TEST_ASSERT(  ret == 0 );
 
     TEST_ASSERT( ( buf = mbedtls_calloc( 1, buflen ) ) != NULL );
 
@@ -3288,6 +3370,7 @@
     mbedtls_ssl_transform_free( &t1 );
 
     mbedtls_free( buf );
+    USE_PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3334,13 +3417,17 @@
     int seen_success; /* Indicates if in the current mode we've
                        * already seen a successful test. */
 
+    USE_PSA_INIT( );
+
     mbedtls_ssl_init( &ssl );
     mbedtls_ssl_transform_init( &t0 );
     mbedtls_ssl_transform_init( &t1 );
-    TEST_ASSERT( build_transforms( &t0, &t1, cipher_type, hash_id,
+    ret = build_transforms( &t0, &t1, cipher_type, hash_id,
                                    etm, tag_mode, ver,
                                    (size_t) cid0_len,
-                                   (size_t) cid1_len ) == 0 );
+                                   (size_t) cid1_len );
+
+    TEST_ASSERT(  ret == 0 );
 
     TEST_ASSERT( ( buf = mbedtls_calloc( 1, buflen ) ) != NULL );
 
@@ -3454,10 +3541,11 @@
     mbedtls_ssl_transform_free( &t1 );
 
     mbedtls_free( buf );
+    USE_PSA_DONE( );
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_AES_C:MBEDTLS_SSL_PROTO_TLS1_2 */
+/* BEGIN_CASE depends_on:MBEDTLS_CIPHER_MODE_CBC:MBEDTLS_AES_C:MBEDTLS_SSL_PROTO_TLS1_2:!MBEDTLS_USE_PSA_CRYPTO */
 void ssl_decrypt_non_etm_cbc( int cipher_type, int hash_id, int trunc_hmac,
                               int length_selector )
 {
@@ -3487,17 +3575,20 @@
     unsigned char add_data[13];
     unsigned char mac[MBEDTLS_MD_MAX_SIZE];
     int exp_ret;
+    int ret;
     const unsigned char pad_max_len = 255; /* Per the standard */
 
+    USE_PSA_INIT( );
+
     mbedtls_ssl_init( &ssl );
     mbedtls_ssl_transform_init( &t0 );
     mbedtls_ssl_transform_init( &t1 );
 
     /* Set up transforms with dummy keys */
-    TEST_ASSERT( build_transforms( &t0, &t1, cipher_type, hash_id,
+    ret = build_transforms( &t0, &t1, cipher_type, hash_id,
                                    0, trunc_hmac,
                                    MBEDTLS_SSL_MINOR_VERSION_3,
-                                   0 , 0 ) == 0 );
+                                   0 , 0 );
 
     /* Determine padding/plaintext length */
     TEST_ASSERT( length_selector >= -2 && length_selector <= 255 );
@@ -3666,6 +3757,7 @@
     mbedtls_ssl_transform_free( &t1 );
     mbedtls_free( buf );
     mbedtls_free( buf_save );
+    USE_PSA_DONE( );
 }
 /* END_CASE */
 
@@ -3964,6 +4056,8 @@
     size_t buf_len;
     int other_endpoint;
 
+    USE_PSA_INIT( );
+
     TEST_ASSERT( endpoint == MBEDTLS_SSL_IS_CLIENT ||
                  endpoint == MBEDTLS_SSL_IS_SERVER );
 
@@ -4039,6 +4133,7 @@
     mbedtls_free( buf );
     mbedtls_ssl_transform_free( &transform_send );
     mbedtls_ssl_transform_free( &transform_recv );
+    USE_PSA_DONE( );
 }
 /* END_CASE */