Merge pull request #5350 from AndrzejKurek/psa-aead-invalid-tag-lengths-setup

Detect invalid tag lengths in psa_aead_setup
diff --git a/ChangeLog.d/chacha20_invalid_iv_len_fix.txt b/ChangeLog.d/chacha20_invalid_iv_len_fix.txt
new file mode 100644
index 0000000..af35e2a
--- /dev/null
+++ b/ChangeLog.d/chacha20_invalid_iv_len_fix.txt
@@ -0,0 +1,4 @@
+Default behavior changes
+   * mbedtls_cipher_set_iv will now fail with ChaCha20 and ChaCha20+Poly1305
+     for IV lengths other than 12. The library was silently overwriting this
+     length with 12, but did not inform the caller about it. Fixes #4301.
diff --git a/library/cipher.c b/library/cipher.c
index 03e84c6..4c7ca3f 100644
--- a/library/cipher.c
+++ b/library/cipher.c
@@ -386,6 +386,12 @@
 #if defined(MBEDTLS_CHACHA20_C)
     if ( ctx->cipher_info->type == MBEDTLS_CIPHER_CHACHA20 )
     {
+        /* Even though the actual_iv_size is overwritten with a correct value
+         * of 12 from the cipher info, return an error to indicate that
+         * the input iv_len is wrong. */
+        if( iv_len != 12 )
+            return( MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA );
+
         if ( 0 != mbedtls_chacha20_starts( (mbedtls_chacha20_context*)ctx->cipher_ctx,
                                            iv,
                                            0U ) ) /* Initial counter value */
@@ -393,6 +399,11 @@
             return( MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA );
         }
     }
+#if defined(MBEDTLS_CHACHAPOLY_C)
+    if ( ctx->cipher_info->type == MBEDTLS_CIPHER_CHACHA20_POLY1305 &&
+         iv_len != 12 )
+        return( MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA );
+#endif
 #endif
 
 #if defined(MBEDTLS_GCM_C)
diff --git a/tests/suites/test_suite_cipher.chacha20.data b/tests/suites/test_suite_cipher.chacha20.data
index 11de103..bcd0032 100644
--- a/tests/suites/test_suite_cipher.chacha20.data
+++ b/tests/suites/test_suite_cipher.chacha20.data
@@ -1,3 +1,7 @@
+Decrypt empty buffer
+depends_on:MBEDTLS_CHACHA20_C
+dec_empty_buf:MBEDTLS_CIPHER_CHACHA20:0:0
+
 Chacha20 RFC 7539 Test Vector #1
 depends_on:MBEDTLS_CHACHA20_C
 decrypt_test_vec:MBEDTLS_CIPHER_CHACHA20:-1:"0000000000000000000000000000000000000000000000000000000000000000":"000000000000000000000000":"76b8e0ada0f13d90405d6ae55386bd28bdd219b8a08ded1aa836efcc8b770dc7da41597c5157488d7724e03fb8d84a376a43b8f41518a11cc387b669b2ee6586":"00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000":"":"":0:0
@@ -109,3 +113,23 @@
 ChaCha20 Encrypt and decrypt 32 bytes in multiple parts
 depends_on:MBEDTLS_CHACHA20_C
 enc_dec_buf_multipart:MBEDTLS_CIPHER_CHACHA20:256:16:16:-1:16:16:16:16
+
+ChaCha20 IV Length 0
+depends_on:MBEDTLS_CHACHA20_C
+check_iv:MBEDTLS_CIPHER_CHACHA20:"CHACHA20":0:MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA
+
+ChaCha20 IV Length 11
+depends_on:MBEDTLS_CHACHA20_C
+check_iv:MBEDTLS_CIPHER_CHACHA20:"CHACHA20":11:MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA
+
+ChaCha20 IV Length 12
+depends_on:MBEDTLS_CHACHA20_C
+check_iv:MBEDTLS_CIPHER_CHACHA20:"CHACHA20":12:0
+
+ChaCha20 IV Length 13
+depends_on:MBEDTLS_CHACHA20_C
+check_iv:MBEDTLS_CIPHER_CHACHA20:"CHACHA20":13:MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA
+
+ChaCha20 IV Length 16
+depends_on:MBEDTLS_CHACHA20_C
+check_iv:MBEDTLS_CIPHER_CHACHA20:"CHACHA20":16:MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA
diff --git a/tests/suites/test_suite_cipher.chachapoly.data b/tests/suites/test_suite_cipher.chachapoly.data
index 8c246ad..908951a 100644
--- a/tests/suites/test_suite_cipher.chachapoly.data
+++ b/tests/suites/test_suite_cipher.chachapoly.data
@@ -121,3 +121,23 @@
 Chacha20+Poly1305 RFC 7539 Test Vector #1 (streaming)
 depends_on:MBEDTLS_CHACHAPOLY_C
 decrypt_test_vec:MBEDTLS_CIPHER_CHACHA20_POLY1305:-1:"1c9240a5eb55d38af333888604f6b5f0473917c1402b80099dca5cbc207075c0":"000000000102030405060708":"64a0861575861af460f062c79be643bd5e805cfd345cf389f108670ac76c8cb24c6cfc18755d43eea09ee94e382d26b0bdb7b73c321b0100d4f03b7f355894cf332f830e710b97ce98c8a84abd0b948114ad176e008d33bd60f982b1ff37c8559797a06ef4f0ef61c186324e2b3506383606907b6a7c02b0f9f6157b53c867e4b9166c767b804d46a59b5216cde7a4e99040c5a40433225ee282a1b0a06c523eaf4534d7f83fa1155b0047718cbc546a0d072b04b3564eea1b422273f548271a0bb2316053fa76991955ebd63159434ecebb4e466dae5a1073a6727627097a1049e617d91d361094fa68f0ff77987130305beaba2eda04df997b714d6c6f2c29a6ad5cb4022b02709b":"496e7465726e65742d4472616674732061726520647261667420646f63756d656e74732076616c696420666f722061206d6178696d756d206f6620736978206d6f6e74687320616e64206d617920626520757064617465642c207265706c616365642c206f72206f62736f6c65746564206279206f7468657220646f63756d656e747320617420616e792074696d652e20497420697320696e617070726f70726961746520746f2075736520496e7465726e65742d447261667473206173207265666572656e6365206d6174657269616c206f7220746f2063697465207468656d206f74686572207468616e206173202fe2809c776f726b20696e2070726f67726573732e2fe2809d":"f33388860000000000004e91":"eead9d67890cbb22392336fea1851f38":0:0
+
+ChaCha20+Poly1305 IV Length 0
+depends_on:MBEDTLS_CHACHAPOLY_C
+check_iv:MBEDTLS_CIPHER_CHACHA20_POLY1305:"CHACHA20-POLY1305":0:MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA
+
+ChaCha20+Poly1305 IV Length 11
+depends_on:MBEDTLS_CHACHAPOLY_C
+check_iv:MBEDTLS_CIPHER_CHACHA20_POLY1305:"CHACHA20-POLY1305":11:MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA
+
+ChaCha20+Poly1305 IV Length 12
+depends_on:MBEDTLS_CHACHAPOLY_C
+check_iv:MBEDTLS_CIPHER_CHACHA20_POLY1305:"CHACHA20-POLY1305":12:0
+
+ChaCha20+Poly1305 IV Length 13
+depends_on:MBEDTLS_CHACHAPOLY_C
+check_iv:MBEDTLS_CIPHER_CHACHA20_POLY1305:"CHACHA20-POLY1305":13:MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA
+
+ChaCha20+Poly1305 IV Length 16
+depends_on:MBEDTLS_CHACHAPOLY_C
+check_iv:MBEDTLS_CIPHER_CHACHA20_POLY1305:"CHACHA20-POLY1305":16:MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA
diff --git a/tests/suites/test_suite_cipher.function b/tests/suites/test_suite_cipher.function
index fd2985b..cd79ba4 100644
--- a/tests/suites/test_suite_cipher.function
+++ b/tests/suites/test_suite_cipher.function
@@ -442,6 +442,9 @@
     if( NULL != strstr( cipher_info->name, "CCM*-NO-TAG") )
         iv_len = 13; /* For CCM, IV length is expected to be between 7 and 13 bytes.
                       * For CCM*-NO-TAG, IV length must be exactly 13 bytes long. */
+    else if( cipher_info->type == MBEDTLS_CIPHER_CHACHA20 ||
+             cipher_info->type == MBEDTLS_CIPHER_CHACHA20_POLY1305 )
+        iv_len = 12;
     else
         iv_len = sizeof(iv);
 
@@ -568,7 +571,9 @@
                     int expected_finish_ret )
 {
     unsigned char key[32];
-    unsigned char iv[16];
+
+    unsigned char *iv = NULL;
+    size_t iv_len = 16;
 
     mbedtls_cipher_context_t ctx_dec;
     const mbedtls_cipher_info_t *cipher_info;
@@ -579,7 +584,6 @@
     size_t outlen = 0;
 
     memset( key, 0, 32 );
-    memset( iv , 0, 16 );
 
     mbedtls_cipher_init( &ctx_dec );
 
@@ -589,6 +593,14 @@
     /* Initialise context */
     cipher_info = mbedtls_cipher_info_from_type( cipher );
     TEST_ASSERT( NULL != cipher_info);
+
+    if( cipher_info->type == MBEDTLS_CIPHER_CHACHA20 ||
+        cipher_info->type == MBEDTLS_CIPHER_CHACHA20_POLY1305 )
+        iv_len = 12;
+
+    ASSERT_ALLOC( iv, iv_len );
+    memset( iv , 0, iv_len );
+
     TEST_ASSERT( sizeof(key) * 8 >= cipher_info->key_bitlen );
 
     TEST_ASSERT( 0 == mbedtls_cipher_setup( &ctx_dec, cipher_info ) );
@@ -597,7 +609,7 @@
                                              key, cipher_info->key_bitlen,
                                              MBEDTLS_DECRYPT ) );
 
-    TEST_ASSERT( 0 == mbedtls_cipher_set_iv( &ctx_dec, iv, 16 ) );
+    TEST_ASSERT( 0 == mbedtls_cipher_set_iv( &ctx_dec, iv, iv_len ) );
 
     TEST_ASSERT( 0 == mbedtls_cipher_reset( &ctx_dec ) );
 
@@ -627,6 +639,7 @@
     TEST_ASSERT( 0 == outlen );
 
 exit:
+    mbedtls_free( iv );
     mbedtls_cipher_free( &ctx_dec );
 }
 /* END_CASE */
@@ -689,6 +702,9 @@
     if( NULL != strstr( cipher_info->name, "CCM*-NO-TAG") )
         iv_len = 13; /* For CCM, IV length is expected to be between 7 and 13 bytes.
                       * For CCM*-NO-TAG, IV length must be exactly 13 bytes long. */
+    else if( cipher_info->type == MBEDTLS_CIPHER_CHACHA20 ||
+             cipher_info->type == MBEDTLS_CIPHER_CHACHA20_POLY1305 )
+        iv_len = 12;
     else
         iv_len = sizeof(iv);
 
@@ -1130,3 +1146,40 @@
         TEST_ASSERT( dlen == (size_t) dlen_check );
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void check_iv( int cipher_id, char * cipher_string,
+               int iv_len_val, int ret )
+{
+    size_t iv_len = iv_len_val;
+    unsigned char iv[16];
+
+    const mbedtls_cipher_info_t *cipher_info;
+    mbedtls_cipher_context_t ctx_dec;
+    mbedtls_cipher_context_t ctx_enc;
+
+    /*
+     * Prepare contexts
+     */
+    mbedtls_cipher_init( &ctx_dec );
+    mbedtls_cipher_init( &ctx_enc );
+
+    /* Check and get info structures */
+    cipher_info = mbedtls_cipher_info_from_type( cipher_id );
+    TEST_ASSERT( NULL != cipher_info );
+    TEST_ASSERT( mbedtls_cipher_info_from_string( cipher_string ) == cipher_info );
+    TEST_ASSERT( strcmp( mbedtls_cipher_info_get_name( cipher_info ),
+                         cipher_string ) == 0 );
+
+    /* Initialise enc and dec contexts */
+    TEST_ASSERT( 0 == mbedtls_cipher_setup( &ctx_dec, cipher_info ) );
+    TEST_ASSERT( 0 == mbedtls_cipher_setup( &ctx_enc, cipher_info ) );
+
+    TEST_ASSERT( ret == mbedtls_cipher_set_iv( &ctx_dec, iv, iv_len ) );
+    TEST_ASSERT( ret == mbedtls_cipher_set_iv( &ctx_enc, iv, iv_len ) );
+
+exit:
+    mbedtls_cipher_free( &ctx_dec );
+    mbedtls_cipher_free( &ctx_enc );
+}
+/* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index ea6c9b3..08a8b29 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -305,6 +305,22 @@
 depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
 aead_decrypt:PSA_KEY_TYPE_AES:"a0ec7b0052541d9e9c091fb7fc481409":PSA_ALG_GCM:"00e440846db73a490573deaf3728c94f":"a3cfcb832e935eb5bc3812583b3a1b2e82920c07fda3668a35d939d8f11379bb606d39e6416b2ef336fffb15aec3f47a71e191f4ff6c56ff15913562619765b26ae094713d60bab6ab82bfc36edaaf8c7ce2cf5906554dcc5933acdb9cb42c1d24718efdc4a09256020b024b224cfe602772bd688c6c8f1041a46f7ec7d51208":"3b6de52f6e582d317f904ee768895bd4d0790912efcf27b58651d0eb7eb0b2f07222c6ffe9f7e127d98ccb132025b098a67dc0ec0083235e9f83af1ae1297df4319547cbcb745cebed36abc1f32a059a05ede6c00e0da097521ead901ad6a73be20018bda4c323faa135169e21581e5106ac20853642e9d6b17f1dd925c872814365847fe0b7b7fbed325953df344a96":"5431d93278c35cfcd7ffa9ce2de5c6b922edffd5055a9eaa5b54cae088db007cf2d28efaf9edd1569341889073e87c0a88462d77016744be62132fd14a243ed6e30e12cd2f7d08a8daeec161691f3b27d4996df8745d74402ee208e4055615a8cb069d495cf5146226490ac615d7b17ab39fb4fdd098e4e7ee294d34c1312826":PSA_ERROR_INSUFFICIENT_MEMORY
 
+PSA MAC sign multipart, through driver: HMAC-SHA-224, parts: 0
+depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
+mac_sign_multipart:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"ba7d81028e07b30466b867d8fefaa52db111d56b45df5a0e1465bf39":0:PSA_SUCCESS
+
+PSA MAC sign multipart, through driver: HMAC-SHA-224, parts: 1
+depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
+mac_sign_multipart:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"896fb1128abbdf196832107cd49df33f47b4b1169912ba4f53684b22":1:PSA_SUCCESS
+
+PSA MAC sign multipart, through driver: HMAC-SHA-224, parts: 2
+depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
+mac_sign_multipart:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"896fb1128abbdf196832107cd49df33f47b4b1169912ba4f53684b22":2:PSA_SUCCESS
+
+PSA MAC sign multipart, through driver: HMAC-SHA-224, parts: 3
+depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
+mac_sign_multipart:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"896fb1128abbdf196832107cd49df33f47b4b1169912ba4f53684b22":3:PSA_SUCCESS
+
 PSA MAC sign, through driver: HMAC-SHA-224
 depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
 mac_sign:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"896fb1128abbdf196832107cd49df33f47b4b1169912ba4f53684b22":PSA_SUCCESS
@@ -329,6 +345,22 @@
 depends_on:PSA_WANT_ALG_CMAC:PSA_WANT_KEY_TYPE_AES
 mac_sign:PSA_KEY_TYPE_AES:"2b7e151628aed2a6abf7158809cf4f3c":PSA_ALG_CMAC:"6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411":"dfa66747de9ae63030ca32611497c827":PSA_ERROR_GENERIC_ERROR
 
+PSA MAC verify multipart, through driver: HMAC-SHA-224, parts: 0
+depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
+mac_verify_multipart:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"ba7d81028e07b30466b867d8fefaa52db111d56b45df5a0e1465bf39":0:PSA_SUCCESS
+
+PSA MAC verify multipart, through driver: HMAC-SHA-224, parts: 1
+depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
+mac_verify_multipart:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"896fb1128abbdf196832107cd49df33f47b4b1169912ba4f53684b22":1:PSA_SUCCESS
+
+PSA MAC verify multipart, through driver: HMAC-SHA-224, parts: 2
+depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
+mac_verify_multipart:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"896fb1128abbdf196832107cd49df33f47b4b1169912ba4f53684b22":2:PSA_SUCCESS
+
+PSA MAC verify multipart, through driver: HMAC-SHA-224, parts: 3
+depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
+mac_verify_multipart:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"896fb1128abbdf196832107cd49df33f47b4b1169912ba4f53684b22":3:PSA_SUCCESS
+
 PSA MAC verify, through driver: HMAC-SHA-224
 depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_224:PSA_WANT_KEY_TYPE_HMAC:MBEDTLS_PSA_ACCEL_ALG_HMAC
 mac_verify:PSA_KEY_TYPE_HMAC:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_HMAC(PSA_ALG_SHA_224):"4869205468657265":"896fb1128abbdf196832107cd49df33f47b4b1169912ba4f53684b22":PSA_SUCCESS
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 64adba9..aab81fb 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -1247,9 +1247,64 @@
     else
         TEST_EQUAL( forced_status, status );
 
-    if( mac_buffer_size > 0 )
-        memset( actual_mac, 0, mac_buffer_size );
+    PSA_ASSERT( psa_mac_abort( &operation ) );
+    TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
+
+    if( forced_status == PSA_SUCCESS )
+    {
+        ASSERT_COMPARE( expected_mac->x, expected_mac->len,
+                        actual_mac, mac_length );
+    }
+
+    mbedtls_free( actual_mac );
+    actual_mac = NULL;
+
+exit:
+    psa_mac_abort( &operation );
+    psa_destroy_key( key );
+    PSA_DONE( );
+    mbedtls_free( actual_mac );
     mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void mac_sign_multipart( int key_type_arg,
+                         data_t *key_data,
+                         int alg_arg,
+                         data_t *input,
+                         data_t *expected_mac,
+                         int fragments_count,
+                         int forced_status_arg )
+{
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_type_t key_type = key_type_arg;
+    psa_algorithm_t alg = alg_arg;
+    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    uint8_t *actual_mac = NULL;
+    size_t mac_buffer_size =
+        PSA_MAC_LENGTH( key_type, PSA_BYTES_TO_BITS( key_data->len ), alg );
+    size_t mac_length = 0;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t forced_status = forced_status_arg;
+    uint8_t *input_x = input->x;
+    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();
+
+    TEST_ASSERT( mac_buffer_size <= PSA_MAC_MAX_SIZE );
+    /* We expect PSA_MAC_LENGTH to be exact. */
+    TEST_ASSERT( expected_mac->len == mac_buffer_size );
+
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, key_type );
+
+    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
+                                &key ) );
+
+    ASSERT_ALLOC( actual_mac, mac_buffer_size );
     mbedtls_test_driver_mac_hooks.forced_status = forced_status;
 
     /*
@@ -1266,25 +1321,38 @@
     else
         TEST_EQUAL( forced_status, status );
 
-    status = psa_mac_update( &operation,
-                             input->x, input->len );
-    if( forced_status == PSA_SUCCESS )
-        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 2 );
-    else
-        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
-    if( forced_status == PSA_SUCCESS ||
-        forced_status == PSA_ERROR_NOT_SUPPORTED )
+    if ( fragments_count )
     {
-        PSA_ASSERT( status );
+        TEST_ASSERT( ( input->len / fragments_count ) > 0 );
     }
-    else
-        TEST_EQUAL( PSA_ERROR_BAD_STATE, status );
+
+    for ( int i = 0; i < fragments_count; i++)
+    {
+        int fragment_size = input->len / fragments_count;
+        if ( i == fragments_count - 1 )
+            fragment_size += ( input->len % fragments_count );
+
+        status = psa_mac_update( &operation,
+                                input_x, fragment_size );
+        if( forced_status == PSA_SUCCESS )
+            TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 2 + i );
+        else
+            TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
+        if( forced_status == PSA_SUCCESS ||
+            forced_status == PSA_ERROR_NOT_SUPPORTED )
+        {
+            PSA_ASSERT( status );
+        }
+        else
+            TEST_EQUAL( PSA_ERROR_BAD_STATE, status );
+        input_x += fragment_size;
+    }
 
     status = psa_mac_sign_finish( &operation,
                                   actual_mac, mac_buffer_size,
                                   &mac_length );
     if( forced_status == PSA_SUCCESS )
-        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 4 );
+        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 3 + fragments_count );
     else
         TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
 
@@ -1298,7 +1366,7 @@
 
     PSA_ASSERT( psa_mac_abort( &operation ) );
     if( forced_status == PSA_SUCCESS )
-        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 4 );
+        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 3 + fragments_count );
     else
         TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
 
@@ -1365,7 +1433,46 @@
     else
         TEST_EQUAL( forced_status, status );
 
+    PSA_ASSERT( psa_mac_abort( &operation ) );
+    TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
+exit:
+    psa_mac_abort( &operation );
+    psa_destroy_key( key );
+    PSA_DONE( );
     mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void mac_verify_multipart( int key_type_arg,
+                           data_t *key_data,
+                           int alg_arg,
+                           data_t *input,
+                           data_t *expected_mac,
+                           int fragments_count,
+                           int forced_status_arg )
+{
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_type_t key_type = key_type_arg;
+    psa_algorithm_t alg = alg_arg;
+    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
+    psa_status_t forced_status = forced_status_arg;
+    uint8_t *input_x = input->x;
+    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();
+
+    TEST_ASSERT( expected_mac->len <= PSA_MAC_MAX_SIZE );
+
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_VERIFY_HASH );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, key_type );
+
+    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
+                                &key ) );
+
     mbedtls_test_driver_mac_hooks.forced_status = forced_status;
 
     /*
@@ -1382,26 +1489,39 @@
     else
         TEST_EQUAL( forced_status, status );
 
-    status = psa_mac_update( &operation,
-                             input->x, input->len );
-    if( forced_status == PSA_SUCCESS )
-        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 2 );
-    else
-        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
-
-    if( forced_status == PSA_SUCCESS ||
-        forced_status == PSA_ERROR_NOT_SUPPORTED )
+    if ( fragments_count )
     {
-        PSA_ASSERT( status );
+        TEST_ASSERT( ( input->len / fragments_count ) > 0 );
     }
-    else
-        TEST_EQUAL( PSA_ERROR_BAD_STATE, status );
+
+    for ( int i = 0; i < fragments_count; i++)
+    {
+        int fragment_size = input->len / fragments_count;
+        if ( i == fragments_count - 1 )
+            fragment_size += ( input->len % fragments_count );
+
+        status = psa_mac_update( &operation,
+                                input_x, fragment_size );
+        if( forced_status == PSA_SUCCESS )
+            TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 2 + i );
+        else
+            TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
+
+        if( forced_status == PSA_SUCCESS ||
+            forced_status == PSA_ERROR_NOT_SUPPORTED )
+        {
+            PSA_ASSERT( status );
+        }
+        else
+            TEST_EQUAL( PSA_ERROR_BAD_STATE, status );
+        input_x += fragment_size;
+    }
 
     status = psa_mac_verify_finish( &operation,
                                     expected_mac->x,
                                     expected_mac->len );
     if( forced_status == PSA_SUCCESS )
-        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 4 );
+        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 3 + fragments_count );
     else
         TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
 
@@ -1416,7 +1536,7 @@
 
     PSA_ASSERT( psa_mac_abort( &operation ) );
     if( forced_status == PSA_SUCCESS )
-        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 4 );
+        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 3 + fragments_count );
     else
         TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );