SSL: fix test failures

1. Change USE_PSA_CRYPTO_INIT/DONE to MD_OR_USE.

2. Add missing occurrences - some of these were already necessary in
principle (in one form or another) but where missing and this was not
detected so far as `psa_hash` doesn't complain in case of a missing
init, but now MD makes it visible.

3. Add missing include in ssl_test_lib.h.

Signed-off-by: Manuel Pégourié-Gonnard <manuel.pegourie-gonnard@arm.com>
diff --git a/tests/src/test_helpers/ssl_helpers.c b/tests/src/test_helpers/ssl_helpers.c
index d248e29..b130edd 100644
--- a/tests/src/test_helpers/ssl_helpers.c
+++ b/tests/src/test_helpers/ssl_helpers.c
@@ -1757,7 +1757,7 @@
 #endif
     int expected_handshake_result = options->expected_handshake_result;
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
     mbedtls_platform_zeroize(&client, sizeof(client));
     mbedtls_platform_zeroize(&server, sizeof(server));
     mbedtls_test_ssl_message_queue server_queue, client_queue;
@@ -2119,7 +2119,7 @@
         mbedtls_free(context_buf);
     }
 #endif
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 #endif /* MBEDTLS_SSL_HANDSHAKE_WITH_CERT_ENABLED */
 
diff --git a/tests/suites/test_suite_debug.function b/tests/suites/test_suite_debug.function
index 8117c79..dad4a53 100644
--- a/tests/suites/test_suite_debug.function
+++ b/tests/suites/test_suite_debug.function
@@ -59,6 +59,8 @@
     mbedtls_ssl_config conf;
     struct buffer_data buffer;
 
+    MD_PSA_INIT();
+
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
     memset(buffer.buf, 0, 2000);
@@ -83,6 +85,7 @@
 exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    MD_PSA_DONE();
 }
 /* END_CASE */
 
@@ -94,6 +97,8 @@
     mbedtls_ssl_config conf;
     struct buffer_data buffer;
 
+    MD_PSA_INIT();
+
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
     memset(buffer.buf, 0, 2000);
@@ -115,6 +120,7 @@
 exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    MD_PSA_DONE();
 }
 /* END_CASE */
 
@@ -126,6 +132,8 @@
     mbedtls_ssl_config conf;
     struct buffer_data buffer;
 
+    MD_PSA_INIT();
+
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
     memset(buffer.buf, 0, 2000);
@@ -147,6 +155,7 @@
 exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    MD_PSA_DONE();
 }
 /* END_CASE */
 
@@ -159,6 +168,8 @@
     mbedtls_ssl_config conf;
     struct buffer_data buffer;
 
+    MD_PSA_INIT();
+
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
     mbedtls_x509_crt_init(&crt);
@@ -183,6 +194,7 @@
     mbedtls_x509_crt_free(&crt);
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    MD_PSA_DONE();
 }
 /* END_CASE */
 
@@ -195,6 +207,8 @@
     struct buffer_data buffer;
     mbedtls_mpi val;
 
+    MD_PSA_INIT();
+
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
     mbedtls_mpi_init(&val);
@@ -220,5 +234,6 @@
     mbedtls_mpi_free(&val);
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    MD_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index c8762d6..3ecd6c3 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -1081,6 +1081,8 @@
     mbedtls_ssl_context ssl;
     mbedtls_ssl_config conf;
 
+    MD_OR_USE_PSA_INIT();
+
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
 
@@ -1100,8 +1102,10 @@
     memcpy(ssl.in_ctr + 2, new->x, 6);
     TEST_ASSERT(mbedtls_ssl_dtls_replay_check(&ssl) == ret);
 
+exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1138,7 +1142,7 @@
     size_t const buflen = 512;
     mbedtls_record rec, rec_backup;
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
 
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
@@ -1245,7 +1249,7 @@
     mbedtls_ssl_transform_free(&t1);
 
     mbedtls_free(buf);
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1292,7 +1296,7 @@
     int seen_success; /* Indicates if in the current mode we've
                        * already seen a successful test. */
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
 
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
@@ -1410,7 +1414,7 @@
     mbedtls_ssl_transform_free(&t1);
 
     mbedtls_free(buf);
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1453,7 +1457,7 @@
     int ret;
     const unsigned char pad_max_len = 255; /* Per the standard */
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
 
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_transform_init(&t0);
@@ -1640,7 +1644,7 @@
     mbedtls_ssl_transform_free(&t1);
     mbedtls_free(buf);
     mbedtls_free(buf_save);
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -1968,7 +1972,7 @@
     size_t buf_len;
     int other_endpoint;
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(endpoint == MBEDTLS_SSL_IS_CLIENT ||
                 endpoint == MBEDTLS_SSL_IS_SERVER);
@@ -2045,7 +2049,7 @@
     mbedtls_free(buf);
     mbedtls_ssl_transform_free(&transform_send);
     mbedtls_ssl_transform_free(&transform_recv);
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2083,7 +2087,7 @@
         goto exit;
     }
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
 
     TEST_ASSERT(mbedtls_ssl_tls_prf(type, secret->x, secret->len,
                                     label, random->x, random->len,
@@ -2096,7 +2100,7 @@
 exit:
 
     mbedtls_free(output);
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2495,6 +2499,8 @@
     mbedtls_test_init_handshake_options(&options);
     options.pk_alg = MBEDTLS_PK_RSA;
 
+    MD_OR_USE_PSA_INIT();
+
     ret = mbedtls_test_ssl_endpoint_init(NULL, endpoint_type, &options,
                                          NULL, NULL, NULL, NULL);
     TEST_ASSERT(MBEDTLS_ERR_SSL_BAD_INPUT_DATA == ret);
@@ -2510,6 +2516,7 @@
 exit:
     mbedtls_test_ssl_endpoint_free(&ep, NULL);
     mbedtls_test_free_handshake_options(&options);
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2523,7 +2530,7 @@
     mbedtls_test_init_handshake_options(&options);
     options.pk_alg = MBEDTLS_PK_RSA;
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
     mbedtls_platform_zeroize(&base_ep, sizeof(base_ep));
     mbedtls_platform_zeroize(&second_ep, sizeof(second_ep));
 
@@ -2563,7 +2570,7 @@
     mbedtls_test_free_handshake_options(&options);
     mbedtls_test_ssl_endpoint_free(&base_ep, NULL);
     mbedtls_test_ssl_endpoint_free(&second_ep, NULL);
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2616,10 +2623,13 @@
 /* BEGIN_CASE depends_on:MBEDTLS_SSL_HANDSHAKE_WITH_CERT_ENABLED:MBEDTLS_PKCS1_V15:MBEDTLS_SSL_PROTO_TLS1_2:MBEDTLS_MD_CAN_SHA256 */
 void handshake_cipher(char *cipher, int pk_alg, int dtls)
 {
+    MD_OR_USE_PSA_INIT();
+
     test_handshake_psk_cipher(cipher, pk_alg, NULL, dtls);
 
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2683,10 +2693,15 @@
                   int expected_cli_fragments,
                   int expected_srv_fragments)
 {
+    MD_OR_USE_PSA_INIT();
+
     test_app_data(mfl, cli_msg_len, srv_msg_len, expected_cli_fragments,
                   expected_srv_fragments, 0);
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
+
+exit:
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2695,10 +2710,15 @@
                    int expected_cli_fragments,
                    int expected_srv_fragments)
 {
+    MD_OR_USE_PSA_INIT();
+
     test_app_data(mfl, cli_msg_len, srv_msg_len, expected_cli_fragments,
                   expected_srv_fragments, 1);
     /* The goto below is used to avoid an "unused label" warning.*/
     goto exit;
+
+exit:
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2833,7 +2853,7 @@
 
     mbedtls_ssl_config conf;
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
     mbedtls_ssl_config_init(&conf);
 
     TEST_ASSERT(mbedtls_ssl_conf_psk(&conf,
@@ -2848,7 +2868,7 @@
 
     mbedtls_ssl_config_free(&conf);
 
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2877,7 +2897,7 @@
 
     mbedtls_ssl_config conf;
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
     mbedtls_ssl_config_init(&conf);
 
     switch (mode) {
@@ -2931,7 +2951,7 @@
 exit:
 
     mbedtls_ssl_config_free(&conf);
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 
 }
 /* END_CASE */
@@ -2944,6 +2964,8 @@
     mbedtls_ssl_config conf;
     mbedtls_ssl_context ssl;
 
+    MD_OR_USE_PSA_INIT();
+
     mbedtls_ssl_config_init(&conf);
     mbedtls_ssl_init(&ssl);
 
@@ -2956,6 +2978,9 @@
 
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+
+exit:
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -2983,6 +3008,8 @@
 #endif
     mbedtls_ssl_conf_curves(&conf, curve_list);
 
+    MD_OR_USE_PSA_INIT();
+
     mbedtls_ssl_context ssl;
     mbedtls_ssl_init(&ssl);
     TEST_ASSERT(mbedtls_ssl_setup(&ssl, &conf) == 0);
@@ -2998,8 +3025,11 @@
         TEST_EQUAL(iana_tls_group_list[i], ssl.handshake->group_list[i]);
     }
 
+
+exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -3019,6 +3049,8 @@
 
     mbedtls_ssl_conf_groups(&conf, iana_tls_group_list);
 
+    MD_OR_USE_PSA_INIT();
+
     mbedtls_ssl_context ssl;
     mbedtls_ssl_init(&ssl);
     TEST_ASSERT(mbedtls_ssl_setup(&ssl, &conf) == 0);
@@ -3033,8 +3065,10 @@
         TEST_EQUAL(iana_tls_group_list[i], ssl.conf->group_list[i]);
     }
 
+exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -3054,7 +3088,7 @@
     options.srv_log_obj = &srv_pattern;
     options.srv_log_fun = mbedtls_test_ssl_log_analyzer;
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
     mbedtls_platform_zeroize(&client, sizeof(client));
     mbedtls_platform_zeroize(&server, sizeof(server));
 
@@ -3099,7 +3133,7 @@
     mbedtls_test_ssl_endpoint_free(&server, NULL);
     mbedtls_test_free_handshake_options(&options);
     mbedtls_debug_set_threshold(0);
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -3155,6 +3189,8 @@
 
     mbedtls_test_rnd_std_rand(NULL, own_cid, sizeof(own_cid));
 
+    MD_OR_USE_PSA_INIT();
+
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
 
@@ -3218,8 +3254,10 @@
 
     TEST_EQUAL(cid_enabled, MBEDTLS_SSL_CID_DISABLED);
 
+exit:
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -3234,7 +3272,7 @@
 
     uint16_t iana_tls_group_list[] = { MBEDTLS_SSL_IANA_TLS_GROUP_SECP256R1,
                                        MBEDTLS_SSL_IANA_TLS_GROUP_NONE };
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
     mbedtls_platform_zeroize(&client, sizeof(client));
     mbedtls_platform_zeroize(&server, sizeof(server));
 
@@ -3289,7 +3327,7 @@
     mbedtls_test_ssl_endpoint_free(&server, NULL);
     mbedtls_test_free_handshake_options(&options);
 
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 /* BEGIN_CASE depends_on:MBEDTLS_TEST_HOOKS:MBEDTLS_SSL_PROTO_TLS1_3:!MBEDTLS_SSL_PROTO_TLS1_2:MBEDTLS_SSL_CLI_C:MBEDTLS_SSL_SRV_C:MBEDTLS_SSL_HANDSHAKE_WITH_CERT_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED */
@@ -3308,7 +3346,7 @@
     /*
      * Test set-up
      */
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
     mbedtls_platform_zeroize(&client_ep, sizeof(client_ep));
     mbedtls_platform_zeroize(&server_ep, sizeof(server_ep));
 
@@ -3385,7 +3423,7 @@
     mbedtls_test_ssl_endpoint_free(&server_ep, NULL);
     mbedtls_test_free_handshake_options(&client_options);
     mbedtls_test_free_handshake_options(&server_options);
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -3403,7 +3441,7 @@
     size_t pwd_len = 0;
     int ret;
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
 
     mbedtls_ssl_init(&ssl);
 
@@ -3465,7 +3503,7 @@
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
 
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */
 
@@ -3475,7 +3513,7 @@
     psa_ecc_family_t psa_family;
     size_t psa_bits;
 
-    USE_PSA_INIT();
+    MD_OR_USE_PSA_INIT();
 
 #if defined(MBEDTLS_ECP_DP_SECP521R1_ENABLED) || defined(PSA_WANT_ECC_SECP_R1_521)
     TEST_AVAILABLE_ECC(25, MBEDTLS_ECP_DP_SECP521R1, PSA_ECC_FAMILY_SECP_R1, 521);
@@ -3543,6 +3581,6 @@
     TEST_UNAVAILABLE_ECC(30, MBEDTLS_ECP_DP_CURVE448, PSA_ECC_FAMILY_MONTGOMERY, 448);
 #endif
 
-    USE_PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */