pk: implement non-PSA mbedtls_pk_sign_ext()

This makes the function always available with its
its implementation depending on MBEDTLS_USE_PSA_CRYPTO.

Related dependencies and tests are updated as well.

Fixes #7583.

Signed-off-by: Tomi Fontanilles <129057597+tomi-font@users.noreply.github.com>
diff --git a/ChangeLog.d/non-psa-pk-implementation.txt b/ChangeLog.d/non-psa-pk-implementation.txt
new file mode 100644
index 0000000..88ed006
--- /dev/null
+++ b/ChangeLog.d/non-psa-pk-implementation.txt
@@ -0,0 +1,4 @@
+Changes
+   * mbedtls_pk_sign_ext() is now available even without PSA
+     (MBEDTLS_PSA_CRYPTO_C). This improves uniformity in the PK module;
+     it now only depends on MBEDTLS_USE_PSA_CRYPTO for its use of PSA.
diff --git a/include/mbedtls/config_adjust_legacy_crypto.h b/include/mbedtls/config_adjust_legacy_crypto.h
index e66d67a..74d037d 100644
--- a/include/mbedtls/config_adjust_legacy_crypto.h
+++ b/include/mbedtls/config_adjust_legacy_crypto.h
@@ -235,9 +235,9 @@
 #define MBEDTLS_PSA_CRYPTO_CLIENT
 #endif /* MBEDTLS_PSA_CRYPTO_C */
 
-/* The PK wrappers need pk_write functions to format RSA key objects
- * when they are dispatching to the PSA API. This happens under USE_PSA_CRYPTO,
- * and also even without USE_PSA_CRYPTO for mbedtls_pk_sign_ext(). */
+/* The PK wrappers need pk_write/pk_parse functions to format RSA key objects
+ * when they are dispatching to the PSA API. This happens under MBEDTLS_USE_PSA_CRYPTO,
+ * and even under just MBEDTLS_PSA_CRYPTO_C in psa_crypto_rsa.c. */
 #if defined(MBEDTLS_PSA_CRYPTO_C) && defined(MBEDTLS_RSA_C)
 #define MBEDTLS_PK_C
 #define MBEDTLS_PK_WRITE_C
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index 24b1188..355bf10 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -28,7 +28,7 @@
 #include "mbedtls/ecdsa.h"
 #endif
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO) || defined(MBEDTLS_PSA_CRYPTO_C)
+#if defined(MBEDTLS_PSA_CRYPTO_C)
 #include "psa/crypto.h"
 #endif
 
@@ -615,7 +615,6 @@
                     unsigned char *sig, size_t sig_size, size_t *sig_len,
                     int (*f_rng)(void *, unsigned char *, size_t), void *p_rng);
 
-#if defined(MBEDTLS_PSA_CRYPTO_C)
 /**
  * \brief           Make signature given a signature type.
  *
@@ -652,7 +651,6 @@
                         unsigned char *sig, size_t sig_size, size_t *sig_len,
                         int (*f_rng)(void *, unsigned char *, size_t),
                         void *p_rng);
-#endif /* MBEDTLS_PSA_CRYPTO_C */
 
 /**
  * \brief           Restartable version of \c mbedtls_pk_sign()
diff --git a/library/pk.c b/library/pk.c
index 957f64e..344d29f 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -579,7 +579,7 @@
 
         return PSA_PK_RSA_TO_MBEDTLS_ERR(status);
     } else
-#endif
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     {
         if (sig_len < mbedtls_pk_get_len(ctx)) {
             return MBEDTLS_ERR_RSA_VERIFY_FAILED;
@@ -672,7 +672,6 @@
                                        f_rng, p_rng, NULL);
 }
 
-#if defined(MBEDTLS_PSA_CRYPTO_C)
 /*
  * Make a signature given a signature type.
  */
@@ -684,11 +683,6 @@
                         int (*f_rng)(void *, unsigned char *, size_t),
                         void *p_rng)
 {
-#if defined(MBEDTLS_RSA_C)
-    psa_algorithm_t psa_md_alg;
-#endif /* MBEDTLS_RSA_C */
-    *sig_len = 0;
-
     if (ctx->pk_info == NULL) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
     }
@@ -702,8 +696,10 @@
                                sig, sig_size, sig_len, f_rng, p_rng);
     }
 
-#if defined(MBEDTLS_RSA_C)
-    psa_md_alg = mbedtls_md_psa_alg_from_type(md_alg);
+#if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_PKCS1_V21)
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    const psa_algorithm_t psa_md_alg = mbedtls_md_psa_alg_from_type(md_alg);
     if (psa_md_alg == 0) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
     }
@@ -720,12 +716,31 @@
     return mbedtls_pk_psa_rsa_sign_ext(PSA_ALG_RSA_PSS(psa_md_alg),
                                        ctx->pk_ctx, hash, hash_len,
                                        sig, sig_size, sig_len);
-#else /* MBEDTLS_RSA_C */
-    return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
-#endif /* !MBEDTLS_RSA_C */
+#else /* MBEDTLS_USE_PSA_CRYPTO */
 
+    if (sig_size < mbedtls_pk_get_len(ctx)) {
+        return MBEDTLS_ERR_PK_BUFFER_TOO_SMALL;
+    }
+
+    if (pk_hashlen_helper(md_alg, &hash_len) != 0) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
+    mbedtls_rsa_context *const rsa_ctx = mbedtls_pk_rsa(*ctx);
+
+    const int ret = mbedtls_rsa_rsassa_pss_sign(rsa_ctx, f_rng, p_rng, md_alg,
+                                                (unsigned int) hash_len, hash, sig);
+    if (ret == 0) {
+        *sig_len = rsa_ctx->len;
+    }
+    return ret;
+
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+#else
+    return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
+#endif /* MBEDTLS_RSA_C && MBEDTLS_PKCS1_V21 */
 }
-#endif /* MBEDTLS_PSA_CRYPTO_C */
 
 /*
  * Decrypt message
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 0fb3c42..b1fbd86 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -281,7 +281,7 @@
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
-#if defined(MBEDTLS_PSA_CRYPTO_C)
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
 int  mbedtls_pk_psa_rsa_sign_ext(psa_algorithm_t alg,
                                  mbedtls_rsa_context *rsa_ctx,
                                  const unsigned char *hash, size_t hash_len,
@@ -344,7 +344,7 @@
     }
     return ret;
 }
-#endif /* MBEDTLS_PSA_CRYPTO_C */
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 static int rsa_sign_wrap(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg,
diff --git a/library/pk_wrap.h b/library/pk_wrap.h
index 28c815a..a6237bf 100644
--- a/library/pk_wrap.h
+++ b/library/pk_wrap.h
@@ -131,6 +131,14 @@
 #endif
 #endif
 
+#if defined(MBEDTLS_RSA_C)
+int mbedtls_pk_psa_rsa_sign_ext(psa_algorithm_t psa_alg_md,
+                                mbedtls_rsa_context *rsa_ctx,
+                                const unsigned char *hash, size_t hash_len,
+                                unsigned char *sig, size_t sig_size,
+                                size_t *sig_len);
+#endif /* MBEDTLS_RSA_C */
+
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_PSA_CRYPTO_C)
@@ -143,14 +151,6 @@
 #endif /* PSA_WANT_KEY_TYPE_RSA_PUBLIC_KEY || PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_BASIC */
 #endif /* !MBEDTLS_DEPRECATED_REMOVED */
 
-#if defined(MBEDTLS_RSA_C)
-int mbedtls_pk_psa_rsa_sign_ext(psa_algorithm_t psa_alg_md,
-                                mbedtls_rsa_context *rsa_ctx,
-                                const unsigned char *hash, size_t hash_len,
-                                unsigned char *sig, size_t sig_size,
-                                size_t *sig_len);
-#endif /* MBEDTLS_RSA_C */
-
 #endif /* MBEDTLS_PSA_CRYPTO_C */
 
 #endif /* MBEDTLS_PK_WRAP_H */
diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index e697491..af1e20c 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -621,62 +621,62 @@
 depends_on:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V15:MBEDTLS_GENPRIME:MBEDTLS_PK_WRITE_C
 pk_psa_sign:1024:PSA_KEY_TYPE_RSA_KEY_PAIR:1024
 
-PK Sign ext:RSA2048,PK_RSA,MD_SHA256
+PK sign ext: RSA2048, PK_RSA, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA256
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA256
 
-PK Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA256
+PK sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA256
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA256
 
-PK Sign ext:RSA2048,PK_RSA,MD_SHA384
+PK sign ext: RSA2048, PK_RSA, MD_SHA384
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA384:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA384
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA384
 
-PK Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA384
+PK sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA384
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA384:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA384
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA384
 
-PK Sign ext:RSA2048,PK_RSA,MD_SHA512
+PK sign ext: RSA2048, PK_RSA, MD_SHA512
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA512:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA512
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA512
 
-PK Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA512
+PK sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA512
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA512:MBEDTLS_RSA_C:MBEDTLS_RSA_GEN_KEY_MIN_BITS <= 2048
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA512
+pk_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA512
 
-PK Sign ext:SECP256R1,PK_ECDSA,MD_SHA256
+PK sign ext: SECP256R1, PK_ECDSA, MD_SHA256
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP256R1:MBEDTLS_MD_CAN_SHA256
-pk_psa_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP256R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA256
+pk_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP256R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA256
 
-PK Sign ext:SECP384R1,PK_ECDSA,MD_SHA384
+PK sign ext: SECP384R1, PK_ECDSA, MD_SHA384
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP384R1:MBEDTLS_MD_CAN_SHA384
-pk_psa_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP384R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA384
+pk_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP384R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA384
 
-PK Sign ext:SECP521R1,PK_ECDSA,MD_SHA512
+PK sign ext: SECP521R1, PK_ECDSA, MD_SHA512
 depends_on:MBEDTLS_PK_CAN_ECDSA_SIGN:MBEDTLS_ECP_HAVE_SECP521R1:MBEDTLS_MD_CAN_SHA512
-pk_psa_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP521R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA512
+pk_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP521R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA512
 
-PK wrapped Sign ext:RSA2048,PK_RSA,MD_SHA256
+PSA wrapped sign ext: RSA2048, PK_RSA, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA256
 
-PK wrapped Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA256
+PSA wrapped sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA256
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA256:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA256
 
-PK wrapped Sign ext:RSA2048,PK_RSA,MD_SHA384
+PSA wrapped sign ext: RSA2048, PK_RSA, MD_SHA384
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA384:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA384
 
-PK wrapped Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA384
+PSA wrapped sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA384
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA384:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA384
 
-PK wrapped Sign ext:RSA2048,PK_RSA,MD_SHA512
+PSA wrapped sign ext: RSA2048, PK_RSA, MD_SHA512
 depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_MD_CAN_SHA512:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA512
 
-PK wrapped Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA512
+PSA wrapped sign ext: RSA2048, PK_RSASSA_PSS, MD_SHA512
 depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_MD_CAN_SHA512:MBEDTLS_RSA_C
 pk_psa_wrap_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA512
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index fa0b03b..5255feb 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -407,6 +407,16 @@
                                            buf, buf_size, &buf_size,
                                            NULL, NULL,
                                            NULL));
+    TEST_EQUAL(MBEDTLS_ERR_PK_BAD_INPUT_DATA,
+               mbedtls_pk_sign_ext(pk_type, &ctx, MBEDTLS_MD_NONE,
+                                   NULL, buf_size,
+                                   buf, buf_size, &buf_size,
+                                   NULL, NULL));
+    TEST_EQUAL(MBEDTLS_ERR_PK_BAD_INPUT_DATA,
+               mbedtls_pk_sign_ext(pk_type, &ctx, MBEDTLS_MD_SHA256,
+                                   NULL, 0,
+                                   buf, buf_size, &buf_size,
+                                   NULL, NULL));
 exit:
     mbedtls_pk_free(&ctx);
     USE_PSA_DONE();
@@ -443,14 +453,6 @@
                                             NULL) ==
                 MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 
-    TEST_ASSERT(mbedtls_pk_sign_restartable(&pk,
-                                            MBEDTLS_MD_NONE,
-                                            NULL, 0,
-                                            buf, sizeof(buf), &len,
-                                            mbedtls_test_rnd_std_rand, NULL,
-                                            NULL) ==
-                MBEDTLS_ERR_PK_BAD_INPUT_DATA);
-
     TEST_ASSERT(mbedtls_pk_sign(&pk,
                                 MBEDTLS_MD_NONE,
                                 NULL, 0,
@@ -458,6 +460,13 @@
                                 mbedtls_test_rnd_std_rand, NULL) ==
                 MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 
+    TEST_ASSERT(mbedtls_pk_sign_ext(MBEDTLS_PK_NONE, &pk,
+                                    MBEDTLS_MD_NONE,
+                                    NULL, 0,
+                                    buf, sizeof(buf), &len,
+                                    mbedtls_test_rnd_std_rand, NULL) ==
+                MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+
     TEST_ASSERT(mbedtls_pk_verify_restartable(&pk,
                                               MBEDTLS_MD_NONE,
                                               NULL, 0,
@@ -1175,22 +1184,31 @@
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
 
-    TEST_ASSERT(mbedtls_pk_setup(&pk,
-                                 mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)) == 0);
+    TEST_EQUAL(mbedtls_pk_setup(&pk,
+                                mbedtls_pk_info_from_type(MBEDTLS_PK_RSA)), 0);
 
 #if defined(MBEDTLS_PKCS1_V21)
-    TEST_ASSERT(mbedtls_pk_verify_ext(MBEDTLS_PK_RSASSA_PSS, NULL, &pk,
-                                      MBEDTLS_MD_NONE, hash, hash_len, sig, sig_len) ==
-                MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+    TEST_EQUAL(mbedtls_pk_verify_ext(MBEDTLS_PK_RSASSA_PSS, NULL, &pk,
+                                     MBEDTLS_MD_NONE, hash, hash_len, sig, sig_len),
+               MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 #endif /* MBEDTLS_PKCS1_V21 */
 
-    TEST_ASSERT(mbedtls_pk_verify(&pk, MBEDTLS_MD_NONE, hash, hash_len,
-                                  sig, sig_len) == MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+    TEST_EQUAL(mbedtls_pk_verify(&pk, MBEDTLS_MD_NONE, hash, hash_len,
+                                 sig, sig_len),
+               MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 
-    TEST_ASSERT(mbedtls_pk_sign(&pk, MBEDTLS_MD_NONE, hash, hash_len,
-                                sig, sizeof(sig), &sig_len,
-                                mbedtls_test_rnd_std_rand, NULL)
-                == MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+#if defined(MBEDTLS_PKCS1_V21)
+    TEST_EQUAL(mbedtls_pk_sign_ext(MBEDTLS_PK_RSASSA_PSS, &pk,
+                                   MBEDTLS_MD_NONE, hash, hash_len,
+                                   sig, sizeof(sig), &sig_len,
+                                   mbedtls_test_rnd_std_rand, NULL),
+               MBEDTLS_ERR_PK_BAD_INPUT_DATA);
+#endif /* MBEDTLS_PKCS1_V21 */
+
+    TEST_EQUAL(mbedtls_pk_sign(&pk, MBEDTLS_MD_NONE, hash, hash_len,
+                               sig, sizeof(sig), &sig_len,
+                               mbedtls_test_rnd_std_rand, NULL),
+               MBEDTLS_ERR_PK_BAD_INPUT_DATA);
 
 exit:
     mbedtls_pk_free(&pk);
@@ -1440,14 +1458,14 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_C:MBEDTLS_GENPRIME */
-void pk_psa_sign_ext(int pk_type, int parameter, int key_pk_type, int md_alg)
+/* BEGIN_CASE depends_on:MBEDTLS_GENPRIME */
+void pk_sign_ext(int pk_type, int parameter, int key_pk_type, int md_alg)
 {
     /* See the description of pk_genkey() for the description of the `parameter` argument. */
     mbedtls_pk_context pk;
     size_t sig_len;
     unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
-    unsigned char hash[PSA_HASH_MAX_SIZE];
+    unsigned char hash[MBEDTLS_MD_MAX_SIZE];
     size_t hash_len = mbedtls_md_get_size_from_type(md_alg);
     void const *options = NULL;
     mbedtls_pk_rsassa_pss_options rsassa_pss_options;
@@ -1455,16 +1473,15 @@
     memset(sig, 0, sizeof(sig));
 
     mbedtls_pk_init(&pk);
-    PSA_INIT();
+    MD_OR_USE_PSA_INIT();
 
-    TEST_ASSERT(mbedtls_pk_setup(&pk,
-                                 mbedtls_pk_info_from_type(pk_type)) == 0);
+    TEST_EQUAL(mbedtls_pk_setup(&pk,
+                                mbedtls_pk_info_from_type(pk_type)), 0);
+    TEST_EQUAL(pk_genkey(&pk, parameter), 0);
 
-    TEST_ASSERT(pk_genkey(&pk, parameter) == 0);
-
-    TEST_ASSERT(mbedtls_pk_sign_ext(key_pk_type, &pk, md_alg, hash, hash_len,
-                                    sig, sizeof(sig), &sig_len,
-                                    mbedtls_test_rnd_std_rand, NULL) == 0);
+    TEST_EQUAL(mbedtls_pk_sign_ext(key_pk_type, &pk, md_alg, hash, hash_len,
+                                   sig, sizeof(sig), &sig_len,
+                                   mbedtls_test_rnd_std_rand, NULL), 0);
 
     if (key_pk_type == MBEDTLS_PK_RSASSA_PSS) {
         rsassa_pss_options.mgf1_hash_id = md_alg;
@@ -1472,11 +1489,11 @@
         rsassa_pss_options.expected_salt_len = hash_len;
         options = (const void *) &rsassa_pss_options;
     }
-    TEST_ASSERT(mbedtls_pk_verify_ext(key_pk_type, options, &pk, md_alg,
-                                      hash, hash_len, sig, sig_len) == 0);
+    TEST_EQUAL(mbedtls_pk_verify_ext(key_pk_type, options, &pk, md_alg,
+                                     hash, hash_len, sig, sig_len), 0);
 exit:
     mbedtls_pk_free(&pk);
-    PSA_DONE();
+    MD_OR_USE_PSA_DONE();
 }
 /* END_CASE */