Merge pull request #5708 from AndrzejKurek/timeless-struggles

Remove the dependency on MBEDTLS_TIME_H from the timing module
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index aac526e..cea12f1 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -52,10 +52,6 @@
 The list of maintained branches can be found in the [Current Branches section
 of BRANCHES.md](BRANCHES.md#current-branches).
 
-The only currently maintained LTS branch is:
-1. [mbedtls-2.16](https://github.com/Mbed-TLS/mbedtls/tree/mbedtls-2.16)
-
-
 Tests
 -----
 As mentioned, tests that show the correctness of the feature or bug fix should be added to the pull request, if no such tests exist.
diff --git a/ChangeLog.d/fix-parllel-cmake-build-fail.txt b/ChangeLog.d/fix-parllel-cmake-build-fail.txt
new file mode 100644
index 0000000..4746c7b
--- /dev/null
+++ b/ChangeLog.d/fix-parllel-cmake-build-fail.txt
@@ -0,0 +1,3 @@
+Bugfix
+   * Fix a race condition in out-of-source builds with CMake when generated data
+     files are already present. Fixes #5374
diff --git a/ChangeLog.d/mbedtls_ssl_hs_cb_t.txt b/ChangeLog.d/mbedtls_ssl_hs_cb_t.txt
new file mode 100644
index 0000000..28c3371
--- /dev/null
+++ b/ChangeLog.d/mbedtls_ssl_hs_cb_t.txt
@@ -0,0 +1,4 @@
+Features
+   * Introduce mbedtls_ssl_hs_cb_t typedef for use with
+     mbedtls_ssl_conf_cert_cb() and perhaps future callbacks
+     during TLS handshake.
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index 324612a..a0d4694 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -923,7 +923,8 @@
  *                  change or be removed at any time without notice.
  *
  * \note            Only ECDSA keys are supported so far. Signing with the
- *                  specified hash is the only allowed use of that key.
+ *                  specified hash & ECDH key agreement derivation operation
+ *                  are the only allowed use of that key.
  *
  * \param pk        Input: the EC key to import to a PSA key.
  *                  Output: a PK context wrapping that PSA key.
diff --git a/include/mbedtls/ssl.h b/include/mbedtls/ssl.h
index 9be083a..9566ff6 100644
--- a/include/mbedtls/ssl.h
+++ b/include/mbedtls/ssl.h
@@ -1218,6 +1218,25 @@
                                         const unsigned char server_random[32],
                                         mbedtls_tls_prf_types tls_prf_type );
 
+#if defined(MBEDTLS_SSL_SRV_C)
+/**
+ * \brief           Callback type: generic handshake callback
+ *
+ * \note            Callbacks may use user_data funcs to set/get app user data.
+ *                  See \c mbedtls_ssl_get_user_data_p()
+ *                      \c mbedtls_ssl_get_user_data_n()
+ *                      \c mbedtls_ssl_conf_get_user_data_p()
+ *                      \c mbedtls_ssl_conf_get_user_data_n()
+ *
+ * \param ssl       \c mbedtls_ssl_context on which the callback is run
+ *
+ * \return          The return value of the callback is 0 if successful,
+ *                  or a specific MBEDTLS_ERR_XXX code, which will cause
+ *                  the handshake to be aborted.
+ */
+typedef int (*mbedtls_ssl_hs_cb_t)( mbedtls_ssl_context *ssl );
+#endif
+
 /* A type for storing user data in a library structure.
  *
  * The representation of type may change in future versions of the library.
@@ -1477,7 +1496,7 @@
     mbedtls_ssl_user_data_t MBEDTLS_PRIVATE(user_data);
 
 #if defined(MBEDTLS_SSL_SRV_C)
-    int (*MBEDTLS_PRIVATE(f_cert_cb))(mbedtls_ssl_context *); /*!< certificate selection callback */
+    mbedtls_ssl_hs_cb_t MBEDTLS_PRIVATE(f_cert_cb);  /*!< certificate selection callback */
 #endif /* MBEDTLS_SSL_SRV_C */
 };
 
@@ -2278,19 +2297,15 @@
  *                  If set, the callback is always called for each handshake,
  *                  after `ClientHello` processing has finished.
  *
- *                  The callback has the following parameters:
- *                  - \c mbedtls_ssl_context*: The SSL context to which
- *                                             the operation applies.
- *                  The return value of the callback is 0 if successful,
- *                  or a specific MBEDTLS_ERR_XXX code, which will cause
- *                  the handshake to be aborted.
- *
  * \param conf      The SSL configuration to register the callback with.
  * \param f_cert_cb The callback for selecting server certificate after
  *                  `ClientHello` processing has finished.
  */
-void mbedtls_ssl_conf_cert_cb( mbedtls_ssl_config *conf,
-                               int (*f_cert_cb)(mbedtls_ssl_context *) );
+static inline void mbedtls_ssl_conf_cert_cb( mbedtls_ssl_config *conf,
+                                             mbedtls_ssl_hs_cb_t f_cert_cb )
+{
+    conf->MBEDTLS_PRIVATE(f_cert_cb) = f_cert_cb;
+}
 #endif /* MBEDTLS_SSL_SRV_C */
 
 /**
diff --git a/library/pk.c b/library/pk.c
index 7f4d5fe..3b42799 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -735,8 +735,10 @@
     /* prepare the key attributes */
     psa_set_key_type( &attributes, key_type );
     psa_set_key_bits( &attributes, bits );
-    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH |
+                                          PSA_KEY_USAGE_DERIVE);
     psa_set_key_algorithm( &attributes, PSA_ALG_ECDSA(hash_alg) );
+    psa_set_key_enrollment_algorithm( &attributes, PSA_ALG_ECDH );
 
     /* import private key into PSA */
     if( PSA_SUCCESS != psa_import_key( &attributes, d, d_len, key ) )
diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index 4964787..254627f 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -644,8 +644,9 @@
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO) || defined(MBEDTLS_SSL_PROTO_TLS1_3)
     psa_key_type_t ecdh_psa_type;
-    uint16_t ecdh_bits;
+    size_t ecdh_bits;
     mbedtls_svc_key_id_t ecdh_psa_privkey;
+    uint8_t ecdh_psa_privkey_is_external;
     unsigned char ecdh_psa_peerkey[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
     size_t ecdh_psa_peerkey_len;
 #endif /* MBEDTLS_USE_PSA_CRYPTO || MBEDTLS_SSL_PROTO_TLS1_3 */
@@ -2222,6 +2223,8 @@
             return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
         case PSA_ERROR_BAD_STATE:
             return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
+        case PSA_ERROR_BUFFER_TOO_SMALL:
+            return( MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL );
         default:
             return( MBEDTLS_ERR_PLATFORM_HW_ACCEL_FAILED );
     }
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 94cb776..3fc0701 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1313,14 +1313,6 @@
 }
 
 #if defined(MBEDTLS_SSL_SRV_C)
-void mbedtls_ssl_conf_cert_cb( mbedtls_ssl_config *conf,
-                               int (*f_cert_cb)(mbedtls_ssl_context *) )
-{
-    conf->f_cert_cb = f_cert_cb;
-}
-#endif /* MBEDTLS_SSL_SRV_C */
-
-#if defined(MBEDTLS_SSL_SRV_C)
 void mbedtls_ssl_conf_session_cache( mbedtls_ssl_config *conf,
                                      void *p_cache,
                                      mbedtls_ssl_cache_get_t *f_get_cache,
@@ -3204,7 +3196,8 @@
 
 #if defined(MBEDTLS_ECDH_C) && \
     ( defined(MBEDTLS_USE_PSA_CRYPTO) || defined(MBEDTLS_SSL_PROTO_TLS1_3) )
-    psa_destroy_key( handshake->ecdh_psa_privkey );
+    if( handshake->ecdh_psa_privkey_is_external == 0 )
+        psa_destroy_key( handshake->ecdh_psa_privkey );
 #endif /* MBEDTLS_ECDH_C && MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
diff --git a/library/ssl_tls12_client.c b/library/ssl_tls12_client.c
index 734d3a2..7771d38 100644
--- a/library/ssl_tls12_client.c
+++ b/library/ssl_tls12_client.c
@@ -1861,9 +1861,7 @@
     {
         return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
     }
-    if( ecdh_bits > 0xffff )
-        return( MBEDTLS_ERR_SSL_ILLEGAL_PARAMETER );
-    handshake->ecdh_bits = (uint16_t) ecdh_bits;
+    handshake->ecdh_bits = ecdh_bits;
 
     /* Keep a copy of the peer's public key */
     ecpoint_len = *(*p)++;
diff --git a/library/ssl_tls12_server.c b/library/ssl_tls12_server.c
index 486632e..9ecfdd2 100644
--- a/library/ssl_tls12_server.c
+++ b/library/ssl_tls12_server.c
@@ -2848,7 +2848,102 @@
 }
 #endif /* MBEDTLS_KEY_EXCHANGE_CERT_REQ_ALLOWED_ENABLED */
 
-#if defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) || \
+#if defined(MBEDTLS_USE_PSA_CRYPTO) &&                      \
+        ( defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) || \
+          defined(MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED) )
+static int ssl_get_ecdh_params_from_cert( mbedtls_ssl_context *ssl )
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    unsigned char buf[
+        PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)];
+    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    size_t ecdh_bits = 0;
+    size_t key_len;
+    mbedtls_pk_context *pk;
+    mbedtls_ecp_keypair *key;
+
+    pk = mbedtls_ssl_own_key( ssl );
+
+    if( pk == NULL )
+        return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
+
+    switch( mbedtls_pk_get_type( pk ) )
+    {
+    case MBEDTLS_PK_OPAQUE:
+        if( ! mbedtls_pk_can_do( pk, MBEDTLS_PK_ECKEY ) )
+            return( MBEDTLS_ERR_SSL_PK_TYPE_MISMATCH );
+
+        ssl->handshake->ecdh_psa_privkey =
+                *( (mbedtls_svc_key_id_t*) pk->pk_ctx );
+
+        /* Key should not be destroyed in the TLS library */
+        ssl->handshake->ecdh_psa_privkey_is_external = 1;
+
+        status = psa_get_key_attributes( ssl->handshake->ecdh_psa_privkey,
+                                         &key_attributes );
+        if( status != PSA_SUCCESS)
+        {
+            ssl->handshake->ecdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
+            return( psa_ssl_status_to_mbedtls( status ) );
+        }
+
+        ssl->handshake->ecdh_psa_type = psa_get_key_type( &key_attributes );
+        ssl->handshake->ecdh_bits = psa_get_key_bits( &key_attributes );
+
+        psa_reset_key_attributes( &key_attributes );
+
+        ret = 0;
+        break;
+    case MBEDTLS_PK_ECKEY:
+    case MBEDTLS_PK_ECKEY_DH:
+    case MBEDTLS_PK_ECDSA:
+        key = mbedtls_pk_ec( *pk );
+        if( key == NULL )
+            return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
+
+        /* Convert EC group to PSA key type. */
+        if( ( ssl->handshake->ecdh_psa_type =
+                    mbedtls_ecc_group_to_psa( key->grp.id,
+                                              &ecdh_bits ) ) == 0 )
+        {
+            return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
+        }
+
+        ssl->handshake->ecdh_bits = ecdh_bits;
+
+        key_attributes = psa_key_attributes_init();
+        psa_set_key_usage_flags( &key_attributes, PSA_KEY_USAGE_DERIVE );
+        psa_set_key_algorithm( &key_attributes, PSA_ALG_ECDH );
+        psa_set_key_type( &key_attributes,
+                PSA_KEY_TYPE_ECC_KEY_PAIR( ssl->handshake->ecdh_psa_type ) );
+        psa_set_key_bits( &key_attributes, ssl->handshake->ecdh_bits );
+
+        key_len = PSA_BITS_TO_BYTES( key->grp.pbits );
+        ret = mbedtls_ecp_write_key( key, buf, key_len );
+        if( ret != 0 )
+            goto cleanup;
+
+        status = psa_import_key( &key_attributes, buf, key_len,
+                                 &ssl->handshake->ecdh_psa_privkey );
+        if( status != PSA_SUCCESS )
+        {
+            ret = psa_ssl_status_to_mbedtls( status );
+            goto cleanup;
+        }
+
+        ret = 0;
+        break;
+    default:
+            ret = MBEDTLS_ERR_SSL_PK_TYPE_MISMATCH;
+    }
+
+cleanup:
+    mbedtls_platform_zeroize( buf, sizeof( buf ) );
+
+    return( ret );
+}
+#elif defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) || \
     defined(MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED)
 static int ssl_get_ecdh_params_from_cert( mbedtls_ssl_context *ssl )
 {
@@ -3085,12 +3180,12 @@
             handshake->ecdh_psa_type = mbedtls_psa_parse_tls_ecc_group(
                         (*curve)->tls_id, &ecdh_bits );
 
-            if( handshake->ecdh_psa_type == 0 || ecdh_bits > 0xffff )
+            if( handshake->ecdh_psa_type == 0 )
             {
                 MBEDTLS_SSL_DEBUG_MSG( 1, ( "Invalid ecc group parse." ) );
                 return( MBEDTLS_ERR_SSL_ILLEGAL_PARAMETER );
             }
-            handshake->ecdh_bits = (uint16_t) ecdh_bits;
+            handshake->ecdh_bits = ecdh_bits;
 
             key_attributes = psa_key_attributes_init();
             psa_set_key_usage_flags( &key_attributes, PSA_KEY_USAGE_DERIVE );
@@ -3832,9 +3927,13 @@
 #endif /* MBEDTLS_KEY_EXCHANGE_DHE_RSA_ENABLED */
 #if defined(MBEDTLS_USE_PSA_CRYPTO) &&                           \
         ( defined(MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED) ||     \
-          defined(MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED) )
+          defined(MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED) ||   \
+          defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) ||      \
+          defined(MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED) )
     if( ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA ||
-        ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA )
+        ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA ||
+        ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDH_RSA ||
+        ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA )
     {
         size_t data_len = (size_t)( *p++ );
         size_t buf_len = (size_t)( end - p );
@@ -3872,25 +3971,31 @@
         {
             ret = psa_ssl_status_to_mbedtls( status );
             MBEDTLS_SSL_DEBUG_RET( 1, "psa_raw_key_agreement", ret );
-            (void) psa_destroy_key( handshake->ecdh_psa_privkey );
+            if( handshake->ecdh_psa_privkey_is_external == 0 )
+                (void) psa_destroy_key( handshake->ecdh_psa_privkey );
             handshake->ecdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
             return( ret );
         }
 
-        status = psa_destroy_key( handshake->ecdh_psa_privkey );
-
-        if( status != PSA_SUCCESS )
+        if( handshake->ecdh_psa_privkey_is_external == 0 )
         {
-            ret = psa_ssl_status_to_mbedtls( status );
-            MBEDTLS_SSL_DEBUG_RET( 1, "psa_destroy_key", ret );
-            return( ret );
+            status = psa_destroy_key( handshake->ecdh_psa_privkey );
+
+            if( status != PSA_SUCCESS )
+            {
+                ret = psa_ssl_status_to_mbedtls( status );
+                MBEDTLS_SSL_DEBUG_RET( 1, "psa_destroy_key", ret );
+                return( ret );
+            }
         }
         handshake->ecdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
     }
     else
 #endif /* MBEDTLS_USE_PSA_CRYPTO &&
             ( MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED ||
-              MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED ) */
+              MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED ||
+              MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED ||
+              MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED ) */
 #if defined(MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED) ||                     \
     defined(MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED) ||                   \
     defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) ||                      \
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 3bb308c..8e1112d 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -229,9 +229,7 @@
         mbedtls_psa_parse_tls_ecc_group( named_group, &ecdh_bits ) ) == 0 )
             return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
 
-    if( ecdh_bits > 0xffff )
-        return( MBEDTLS_ERR_SSL_ILLEGAL_PARAMETER );
-    ssl->handshake->ecdh_bits = (uint16_t) ecdh_bits;
+    ssl->handshake->ecdh_bits = ecdh_bits;
 
     key_attributes = psa_key_attributes_init();
     psa_set_key_usage_flags( &key_attributes, PSA_KEY_USAGE_DERIVE );
@@ -262,12 +260,6 @@
 
     }
 
-    if( own_pubkey_len > (size_t)( end - buf ) )
-    {
-            MBEDTLS_SSL_DEBUG_MSG( 1, ( "No space in the buffer for ECDH public key." ) );
-        return( MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL );
-    }
-
     *out_len = own_pubkey_len;
 
     return( 0 );
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 2431e40..c1c9052 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -50,12 +50,18 @@
             ${CMAKE_CURRENT_SOURCE_DIR}/../include/psa/crypto_values.h
             ${CMAKE_CURRENT_SOURCE_DIR}/../include/psa/crypto_extra.h
     )
+
 else()
     foreach(file ${base_generated_data_files})
         link_to_source(${file})
     endforeach()
 endif()
-
+# CMake generates sub-makefiles for each target and calls them in subprocesses.
+# Without this command, cmake will generate rules in each sub-makefile. As a result,
+# they can cause race conditions in parallel builds.
+# With this line, only 4 sub-makefiles include the above command, that reduces
+# the risk of a race.
+add_custom_target(test_suite_generated_data DEPENDS ${generated_data_files})
 # Test suites caught by SKIP_TEST_SUITES are built but not executed.
 # "foo" as a skip pattern skips "test_suite_foo" and "test_suite_foo.bar"
 # but not "test_suite_foobar".
@@ -119,6 +125,7 @@
     )
 
     add_executable(test_suite_${data_name} test_suite_${data_name}.c $<TARGET_OBJECTS:mbedtls_test>)
+    add_dependencies(test_suite_${data_name} test_suite_generated_data)
     target_link_libraries(test_suite_${data_name} ${libs})
     # Include test-specific header files from ./include and private header
     # files (used by some invasive tests) from ../library. Public header
diff --git a/tests/ssl-opt.sh b/tests/ssl-opt.sh
index 71a3904..18fff9d 100755
--- a/tests/ssl-opt.sh
+++ b/tests/ssl-opt.sh
@@ -1164,7 +1164,8 @@
 
     analyze_test_commands "$@"
 
-    TIMES_LEFT=2
+    # One regular run and two retries
+    TIMES_LEFT=3
     while [ $TIMES_LEFT -gt 0 ]; do
         TIMES_LEFT=$(( $TIMES_LEFT - 1 ))
 
@@ -1583,6 +1584,24 @@
             -S "error" \
             -C "error"
 
+requires_config_enabled MBEDTLS_SSL_PROTO_TLS1_2
+requires_config_enabled MBEDTLS_USE_PSA_CRYPTO
+requires_config_enabled MBEDTLS_X509_CRT_PARSE_C
+requires_config_enabled MBEDTLS_ECDSA_C
+requires_config_enabled MBEDTLS_SHA256_C
+run_test    "Opaque key for server authentication (ECDH-)" \
+            "$P_SRV force_version=tls12 auth_mode=required key_opaque=1\
+             crt_file=data_files/server5.ku-ka.crt\
+             key_file=data_files/server5.key" \
+            "$P_CLI" \
+            0 \
+            -c "Verifying peer X.509 certificate... ok" \
+            -c "Ciphersuite is TLS-ECDH-" \
+            -s "key types: Opaque, none" \
+            -s "Ciphersuite is TLS-ECDH-" \
+            -S "error" \
+            -C "error"
+
 # Test using an opaque private key for client/server authentication
 requires_config_enabled MBEDTLS_SSL_PROTO_TLS1_2
 requires_config_enabled MBEDTLS_USE_PSA_CRYPTO
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index 210ab65..bf1e01b 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -780,3 +780,15 @@
 PSA encrypt transparent driver: fallback not available RSA OAEP-SHA-256
 depends_on:!MBEDTLS_PSA_BUILTIN_ALG_RSA_OAEP
 asymmetric_encrypt:PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256):"30818902818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc30203010001":"af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3":"874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1":"ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad":"":"":PSA_ERROR_NOT_SUPPORTED:PSA_ERROR_NOT_SUPPORTED
+
+PSA multi-part AEAD encrypt setup, AES-GCM, 128 bytes #1
+depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
+aead_encrypt_setup:PSA_KEY_TYPE_AES:"a0ec7b0052541d9e9c091fb7fc481409":PSA_ALG_GCM:"00e440846db73a490573deaf3728c94f":"a3cfcb832e935eb5bc3812583b3a1b2e82920c07fda3668a35d939d8f11379bb606d39e6416b2ef336fffb15aec3f47a71e191f4ff6c56ff15913562619765b26ae094713d60bab6ab82bfc36edaaf8c7ce2cf5906554dcc5933acdb9cb42c1d24718efdc4a09256020b024b224cfe602772bd688c6c8f1041a46f7ec7d51208":"5431d93278c35cfcd7ffa9ce2de5c6b922edffd5055a9eaa5b54cae088db007cf2d28efaf9edd1569341889073e87c0a88462d77016744be62132fd14a243ed6e30e12cd2f7d08a8daeec161691f3b27d4996df8745d74402ee208e4055615a8cb069d495cf5146226490ac615d7b17ab39fb4fdd098e4e7ee294d34c1312826":"3b6de52f6e582d317f904ee768895bd4d0790912efcf27b58651d0eb7eb0b2f07222c6ffe9f7e127d98ccb132025b098a67dc0ec0083235e9f83af1ae1297df4319547cbcb745cebed36abc1f32a059a05ede6c00e0da097521ead901ad6a73be20018bda4c323faa135169e21581e5106ac20853642e9d6b17f1dd925c87281":"4365847fe0b7b7fbed325953df344a96":PSA_SUCCESS:PSA_SUCCESS
+
+PSA multi-part AEAD encrypt setup, AES-GCM, 128 bytes #1, fallback
+depends_on:MBEDTLS_AES_C:MBEDTLS_GCM_C
+aead_encrypt_setup:PSA_KEY_TYPE_AES:"a0ec7b0052541d9e9c091fb7fc481409":PSA_ALG_GCM:"00e440846db73a490573deaf3728c94f":"a3cfcb832e935eb5bc3812583b3a1b2e82920c07fda3668a35d939d8f11379bb606d39e6416b2ef336fffb15aec3f47a71e191f4ff6c56ff15913562619765b26ae094713d60bab6ab82bfc36edaaf8c7ce2cf5906554dcc5933acdb9cb42c1d24718efdc4a09256020b024b224cfe602772bd688c6c8f1041a46f7ec7d51208":"5431d93278c35cfcd7ffa9ce2de5c6b922edffd5055a9eaa5b54cae088db007cf2d28efaf9edd1569341889073e87c0a88462d77016744be62132fd14a243ed6e30e12cd2f7d08a8daeec161691f3b27d4996df8745d74402ee208e4055615a8cb069d495cf5146226490ac615d7b17ab39fb4fdd098e4e7ee294d34c1312826":"3b6de52f6e582d317f904ee768895bd4d0790912efcf27b58651d0eb7eb0b2f07222c6ffe9f7e127d98ccb132025b098a67dc0ec0083235e9f83af1ae1297df4319547cbcb745cebed36abc1f32a059a05ede6c00e0da097521ead901ad6a73be20018bda4c323faa135169e21581e5106ac20853642e9d6b17f1dd925c87281":"4365847fe0b7b7fbed325953df344a96":PSA_ERROR_NOT_SUPPORTED:PSA_SUCCESS
+
+PSA multi-part AEAD encrypt setup, AES-GCM, 128 bytes #1, INSUFFICIENT_MEMORY
+depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
+aead_encrypt_setup:PSA_KEY_TYPE_AES:"a0ec7b0052541d9e9c091fb7fc481409":PSA_ALG_GCM:"00e440846db73a490573deaf3728c94f":"a3cfcb832e935eb5bc3812583b3a1b2e82920c07fda3668a35d939d8f11379bb606d39e6416b2ef336fffb15aec3f47a71e191f4ff6c56ff15913562619765b26ae094713d60bab6ab82bfc36edaaf8c7ce2cf5906554dcc5933acdb9cb42c1d24718efdc4a09256020b024b224cfe602772bd688c6c8f1041a46f7ec7d51208":"5431d93278c35cfcd7ffa9ce2de5c6b922edffd5055a9eaa5b54cae088db007cf2d28efaf9edd1569341889073e87c0a88462d77016744be62132fd14a243ed6e30e12cd2f7d08a8daeec161691f3b27d4996df8745d74402ee208e4055615a8cb069d495cf5146226490ac615d7b17ab39fb4fdd098e4e7ee294d34c1312826":"3b6de52f6e582d317f904ee768895bd4d0790912efcf27b58651d0eb7eb0b2f07222c6ffe9f7e127d98ccb132025b098a67dc0ec0083235e9f83af1ae1297df4319547cbcb745cebed36abc1f32a059a05ede6c00e0da097521ead901ad6a73be20018bda4c323faa135169e21581e5106ac20853642e9d6b17f1dd925c87281":"4365847fe0b7b7fbed325953df344a96":PSA_ERROR_INSUFFICIENT_MEMORY:PSA_ERROR_INSUFFICIENT_MEMORY
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 406432f..9e433bc 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -2328,3 +2328,123 @@
     PSA_DONE( );
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void aead_encrypt_setup( int key_type_arg, data_t *key_data,
+                         int alg_arg,
+                         data_t *nonce,
+                         data_t *additional_data,
+                         data_t *input_data,
+                         data_t *expected_ciphertext,
+                         data_t *expected_tag,
+                         int forced_status_arg,
+                         int expected_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;
+    size_t key_bits;
+    psa_status_t forced_status = forced_status_arg;
+    psa_status_t expected_status = expected_status_arg;
+    uint8_t *output_data = NULL;
+    size_t output_size = 0;
+    size_t output_length = 0;
+    size_t finish_output_length = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
+    size_t tag_length = 0;
+    uint8_t tag_buffer[PSA_AEAD_TAG_MAX_SIZE];
+
+    psa_aead_operation_t operation = psa_aead_operation_init();
+
+    mbedtls_test_driver_aead_hooks = mbedtls_test_driver_aead_hooks_init();
+
+    PSA_INIT( );
+
+    mbedtls_test_driver_aead_hooks.forced_status = forced_status;
+
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT );
+    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 ) );
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    key_bits = psa_get_key_bits( &attributes );
+
+    output_size = input_data->len + PSA_AEAD_TAG_LENGTH( key_type, key_bits,
+                                                         alg );
+
+    /* For all currently defined algorithms, PSA_AEAD_ENCRYPT_OUTPUT_SIZE
+     * should be exact. */
+    TEST_EQUAL( output_size,
+                PSA_AEAD_ENCRYPT_OUTPUT_SIZE( key_type, alg, input_data->len ) );
+    TEST_ASSERT( output_size <=
+                 PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE( input_data->len ) );
+    ASSERT_ALLOC( output_data, output_size );
+
+    status = psa_aead_encrypt_setup( &operation, key, alg );
+
+    TEST_EQUAL( status, expected_status );
+    TEST_EQUAL( mbedtls_test_driver_aead_hooks.hits_encrypt_setup, 1 );
+
+    if( status == PSA_SUCCESS )
+    {
+        /* Set the nonce. */
+        PSA_ASSERT( psa_aead_set_nonce( &operation, nonce->x, nonce->len ) );
+
+        TEST_EQUAL( mbedtls_test_driver_aead_hooks.hits_set_nonce,
+                    forced_status == PSA_SUCCESS ? 1 : 0 );
+
+        /* Check hooks hits and
+         * set length (additional data and data to encrypt) */
+        PSA_ASSERT( psa_aead_set_lengths( &operation, additional_data->len,
+                                          input_data->len ) );
+
+        TEST_EQUAL( mbedtls_test_driver_aead_hooks.hits_set_lengths,
+                    forced_status == PSA_SUCCESS ? 1 : 0 );
+
+        /* Pass the additional data */
+        PSA_ASSERT( psa_aead_update_ad( &operation, additional_data->x,
+                                        additional_data->len ) );
+
+        TEST_EQUAL( mbedtls_test_driver_aead_hooks.hits_update_ad,
+                    forced_status == PSA_SUCCESS ? 1 : 0 );
+
+        /* Pass the data to encrypt */
+        PSA_ASSERT( psa_aead_update( &operation, input_data->x, input_data->len,
+                                     output_data, output_size, &output_length ) );
+
+        TEST_EQUAL( mbedtls_test_driver_aead_hooks.hits_update,
+                    forced_status == PSA_SUCCESS ? 1 : 0 );
+
+        /* Finish the encryption operation */
+        PSA_ASSERT( psa_aead_finish( &operation, output_data + output_length,
+                                     output_size - output_length,
+                                     &finish_output_length, tag_buffer,
+                                     PSA_AEAD_TAG_MAX_SIZE, &tag_length ) );
+
+        TEST_EQUAL( mbedtls_test_driver_aead_hooks.hits_finish,
+                    forced_status == PSA_SUCCESS ? 1 : 0 );
+
+        TEST_EQUAL( mbedtls_test_driver_aead_hooks.hits_abort,
+                    forced_status == PSA_SUCCESS ? 1 : 0 );
+
+        /* Compare output_data and expected_ciphertext */
+        ASSERT_COMPARE( expected_ciphertext->x, expected_ciphertext->len,
+                        output_data, output_length );
+
+        TEST_EQUAL( output_length + finish_output_length, expected_ciphertext->len );
+
+        /* Compare tag and expected_tag */
+        ASSERT_COMPARE( expected_tag->x, expected_tag->len, tag_buffer, tag_length );
+    }
+
+exit:
+    /* Cleanup */
+    PSA_ASSERT( psa_destroy_key( key ) );
+    mbedtls_free( output_data );
+    PSA_DONE( );
+    mbedtls_test_driver_aead_hooks = mbedtls_test_driver_aead_hooks_init();
+}
+/* END_CASE */