From e2c882518cd6ba7f8b33174aa6101accb32e8cb5 Mon Sep 17 00:00:00 2001
From: Jerry Yu <jerry.h.yu@arm.com>
Date: Tue, 22 Mar 2022 21:18:52 +0800
Subject: [PATCH] Add pk_sign_ext unit tests

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
---
 tests/suites/test_suite_pk.data     | 40 ++++++++++++++++++++++++-----
 tests/suites/test_suite_pk.function |  5 ++--
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/tests/suites/test_suite_pk.data b/tests/suites/test_suite_pk.data
index 8a3a176cee..a83c66f436 100644
--- a/tests/suites/test_suite_pk.data
+++ b/tests/suites/test_suite_pk.data
@@ -341,11 +341,39 @@ PSA wrapped sign: BP512R1
 depends_on:MBEDTLS_ECP_DP_BP512R1_ENABLED
 pk_psa_sign:MBEDTLS_ECP_DP_BP512R1:PSA_ECC_FAMILY_BRAINPOOL_P_R1:512
 
-PK Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA512
-depends_on:MBEDTLS_X509_RSASSA_PSS_SUPPORT:MBEDTLS_SHA512_C
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA512
+PK Sign ext:RSA2048,PK_RSA,MD_SHA256
+depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_SHA256_C
+pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA256
 
-PK Sign ext:RSA2048,PK_RSA,MD_SHA512
-depends_on:MBEDTLS_X509_RSASSA_PSS_SUPPORT:MBEDTLS_SHA512_C:MBEDTLS_PKCS1_V15
-pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA512
+PK Sign ext:RSA2048,PK_RSASSA_PSS,MD_SHA256
+depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_SHA256_C
+pk_psa_sign_ext:MBEDTLS_PK_RSA:2048:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA256
+
+PK Sign ext:RSA3072,PK_RSA,MD_SHA384
+depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_SHA384_C
+pk_psa_sign_ext:MBEDTLS_PK_RSA:3072:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA384
+
+PK Sign ext:RSA3072,PK_RSASSA_PSS,MD_SHA384
+depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_SHA384_C
+pk_psa_sign_ext:MBEDTLS_PK_RSA:3072:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA384
+
+PK Sign ext:RSA4096,PK_RSA,MD_SHA512
+depends_on:MBEDTLS_PKCS1_V15:MBEDTLS_SHA512_C
+pk_psa_sign_ext:MBEDTLS_PK_RSA:4096:MBEDTLS_PK_RSA:MBEDTLS_MD_SHA512
+
+PK Sign ext:RSA4096,PK_RSASSA_PSS,MD_SHA512
+depends_on:MBEDTLS_PKCS1_V21:MBEDTLS_SHA512_C
+pk_psa_sign_ext:MBEDTLS_PK_RSA:4096:MBEDTLS_PK_RSASSA_PSS:MBEDTLS_MD_SHA512
+
+PK Sign ext:SECP256R1,PK_ECDSA,MD_SHA256
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_SHA256_C
+pk_psa_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP256R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA256
+
+PK Sign ext:SECP384R1,PK_ECDSA,MD_SHA384
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_SHA384_C
+pk_psa_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP384R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA384
+
+PK Sign ext:SECP521R1,PK_ECDSA,MD_SHA512
+depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_DP_SECP521R1_ENABLED:MBEDTLS_SHA512_C
+pk_psa_sign_ext:MBEDTLS_PK_ECDSA:MBEDTLS_ECP_DP_SECP521R1:MBEDTLS_PK_ECDSA:MBEDTLS_MD_SHA512
 
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index dceb477436..046517fc58 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -1095,7 +1095,8 @@ void pk_psa_sign_ext( int pk_type, int parameter, int key_pk_type, int md_alg )
     size_t sig_len;
     unsigned char sig[MBEDTLS_PK_SIGNATURE_MAX_SIZE];
     unsigned char hash[MBEDTLS_MD_MAX_SIZE];
-    size_t hash_len = sizeof( hash );
+    const mbedtls_md_info_t *md_info = mbedtls_md_info_from_type( md_alg );
+    size_t hash_len = mbedtls_md_get_size( md_info );
     void const *options = NULL;
 #if defined(MBEDTLS_X509_RSASSA_PSS_SUPPORT)
     mbedtls_pk_rsassa_pss_options rsassa_pss_options;
@@ -1125,7 +1126,7 @@ void pk_psa_sign_ext( int pk_type, int parameter, int key_pk_type, int md_alg )
     }
 #endif
     TEST_ASSERT( mbedtls_pk_verify_ext( key_pk_type, options, &pk, md_alg,
-                                    hash, sizeof( hash ), sig, sig_len ) == 0 );
+                                        hash, hash_len, sig, sig_len ) == 0 );
 exit:
     PSA_DONE( );
     mbedtls_pk_free( &pk );