Merge pull request #8616 from lpy4105/issue/8553/test-driver-only-rsa

Add test for driver-only RSA (crypto only)
diff --git a/ChangeLog.d/8357.txt b/ChangeLog.d/8357.txt
new file mode 100644
index 0000000..9cae396
--- /dev/null
+++ b/ChangeLog.d/8357.txt
@@ -0,0 +1,8 @@
+Features
+   * It is now possible to have AEADs support (CCM, GCM and ChaChaPoly) without
+     MBEDTLS_CIPHER_C. This holds both for the builtin suport (MBEDTLS_CCM_C,
+     MBEDTLS_GCM_C and MBEDTLS_CHACHAPOLY_C) as well as the PSA one
+     (PSA_WANT_ALG_CCM, PSA_WANT_ALG_GCM, PSA_WANT_ALG_CHACHA20_POLY1305).
+     On the PSA side this means that it is possible to enable
+     MBEDTLS_PSA_CRYPTO_C without MBEDTLS_CIPHER_C if none of the
+     non-authenticated ciphers is enabled.
diff --git a/configs/config-tfm.h b/configs/config-tfm.h
index 85b677b..197b808 100644
--- a/configs/config-tfm.h
+++ b/configs/config-tfm.h
@@ -45,6 +45,11 @@
 #undef MBEDTLS_PLATFORM_STD_EXIT_SUCCESS
 #undef MBEDTLS_PLATFORM_STD_EXIT_FAILURE
 
+/* CCM is the only cipher/AEAD enabled in TF-M configuration files, but it
+ * does not need CIPHER_C to be enabled, so we can disable it in order
+ * to reduce code size further. */
+#undef MBEDTLS_CIPHER_C
+
 /*
  * In order to get an example config that works cleanly out-of-the-box
  * for both baremetal and non-baremetal builds, we detect baremetal builds
diff --git a/docs/driver-only-builds.md b/docs/driver-only-builds.md
index 4bad2e8..2dcfe67 100644
--- a/docs/driver-only-builds.md
+++ b/docs/driver-only-builds.md
@@ -55,6 +55,7 @@
 - hashes: SHA-3, SHA-2, SHA-1, MD5, etc.
 - elliptic-curve cryptography (ECC): ECDH, ECDSA, EC J-PAKE, ECC key types.
 - finite-field Diffie-Hellman: FFDH algorithm, DH key types.
+- AEADs: GCM, CCM and ChachaPoly
 
 Supported means that when those are provided only by drivers, everything
 (including PK, X.509 and TLS if `MBEDTLS_USE_PSA_CRYPTO` is enabled) should
@@ -63,7 +64,7 @@
 below.
 
 In the near future (end of 2023), we are planning to also add support for
-ciphers (AES) and AEADs (GCM, CCM, ChachaPoly).
+ciphers (AES, ARIA, Camellia).
 
 Currently (mid-2023) we don't have plans to extend this to RSA. If
 you're interested in driver-only support for RSA, please let us know.
@@ -240,3 +241,34 @@
 ### Limitations
 Support for deterministic derivation of a DH keypair
 (i.e. `PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_DERIVE`) is not supported.
+
+AEADs
+-----
+
+[This section might contain incomplete data and it is going to be updated in
+#8358, i.e. the wrap-up task for accelerated ciphers and AEADs.]
+
+It is possible to have all AEADs operations provided only by a driver.
+
+More precisely you can:
+- enable desired PSA algorithm(s) and key type(s):
+  - `PSA_WANT_ALG_[CCM|GCM]` with `PSA_WANT_KEY_TYPE_[AES|ARIA|CAMELLIA]`
+  - `PSA_WANT_ALG_CHACHA20_POLY1305` with `PSA_WANT_KEY_TYPE_CHACHA20`;
+- enable `MBEDTLS_PSA_ACCEL_xxx` symbol(s) which correspond to the
+  `PSA_WANT_xxx` of the previous step (both for algorithms and key types);
+- disable builtin support of `MBEDTLS_[CCM|GCM|CHACHAPOLY|POLY1305]_C`
+  algorithms and key types `MBEDTLS_[AES|ARIA|CAMELLIA|CHACHA20]_C` for AEADs
+  which are accelerated.
+
+In a build in which all AEADs algorithms and related key types are accelerated
+all AEADs operations requested through the PSA Crypto API (including those in
+TLS and X.509) will be performed by the driver.
+Moreover if no unauthenticated cipher is required, it is also possible to
+disable all built-in block cipher's key types
+(i.e. `MBEDTLS_[AES|ARIA|CAMELLIA|CHACHA20]_C`) and `MBEDTLS_CIPHER_C`. This
+helps in further reducing code's footprint, but unfortunately it makes the
+following modules unavailable:
+- `MBEDTLS_PKCS[5|12]_C`
+- `MBEDTLS_CTR_DRBG_C`
+- `MBEDTLS_NIST_KW_C`
+
diff --git a/include/mbedtls/build_info.h b/include/mbedtls/build_info.h
index 44ecacf..7a70e25 100644
--- a/include/mbedtls/build_info.h
+++ b/include/mbedtls/build_info.h
@@ -62,6 +62,11 @@
 #define MBEDTLS_ARCH_IS_X86
 #endif
 
+#if !defined(MBEDTLS_PLATFORM_IS_WINDOWS_ON_ARM64) && \
+    (defined(_M_ARM64) || defined(_M_ARM64EC))
+#define MBEDTLS_PLATFORM_IS_WINDOWS_ON_ARM64
+#endif
+
 /* This is defined if the architecture is Armv8-A, or higher */
 #if !defined(MBEDTLS_ARCH_IS_ARMV8_A)
 #if defined(__ARM_ARCH) && defined(__ARM_ARCH_PROFILE)
diff --git a/include/mbedtls/check_config.h b/include/mbedtls/check_config.h
index 9b5b646..34ddcb1 100644
--- a/include/mbedtls/check_config.h
+++ b/include/mbedtls/check_config.h
@@ -153,7 +153,9 @@
 #endif /* not all curves accelerated */
 #endif /* some curve accelerated */
 
-#if defined(MBEDTLS_CTR_DRBG_C) && !defined(MBEDTLS_AES_C)
+#if defined(MBEDTLS_CTR_DRBG_C) && !(defined(MBEDTLS_AES_C) || \
+    (defined(MBEDTLS_PSA_CRYPTO_C) && defined(PSA_WANT_KEY_TYPE_AES) && \
+    defined(PSA_WANT_ALG_ECB_NO_PADDING)))
 #error "MBEDTLS_CTR_DRBG_C defined, but not all prerequisites"
 #endif
 
diff --git a/include/mbedtls/ctr_drbg.h b/include/mbedtls/ctr_drbg.h
index d1f19e6..c00756d 100644
--- a/include/mbedtls/ctr_drbg.h
+++ b/include/mbedtls/ctr_drbg.h
@@ -32,7 +32,14 @@
 
 #include "mbedtls/build_info.h"
 
+/* In case AES_C is defined then it is the primary option for backward
+ * compatibility purposes. If that's not available, PSA is used instead */
+#if defined(MBEDTLS_AES_C)
 #include "mbedtls/aes.h"
+#else
+#include "psa/crypto.h"
+#endif
+
 #include "entropy.h"
 
 #if defined(MBEDTLS_THREADING_C)
@@ -150,6 +157,13 @@
 #define MBEDTLS_CTR_DRBG_ENTROPY_NONCE_LEN (MBEDTLS_CTR_DRBG_ENTROPY_LEN + 1) / 2
 #endif
 
+#if !defined(MBEDTLS_AES_C)
+typedef struct mbedtls_ctr_drbg_psa_context {
+    mbedtls_svc_key_id_t key_id;
+    psa_cipher_operation_t operation;
+} mbedtls_ctr_drbg_psa_context;
+#endif
+
 /**
  * \brief          The CTR_DRBG context structure.
  */
@@ -175,7 +189,11 @@
                                                   * This is the maximum number of requests
                                                   * that can be made between reseedings. */
 
+#if defined(MBEDTLS_AES_C)
     mbedtls_aes_context MBEDTLS_PRIVATE(aes_ctx);        /*!< The AES context. */
+#else
+    mbedtls_ctr_drbg_psa_context MBEDTLS_PRIVATE(psa_ctx); /*!< The PSA context. */
+#endif
 
     /*
      * Callbacks (Entropy)
diff --git a/include/mbedtls/mbedtls_config.h b/include/mbedtls/mbedtls_config.h
index 96a3e43..758a514 100644
--- a/include/mbedtls/mbedtls_config.h
+++ b/include/mbedtls/mbedtls_config.h
@@ -2607,6 +2607,13 @@
  * The CTR_DRBG generator uses AES-256 by default.
  * To use AES-128 instead, enable \c MBEDTLS_CTR_DRBG_USE_128_BIT_KEY above.
  *
+ * AES support can either be achived through builtin (MBEDTLS_AES_C) or PSA.
+ * Builtin is the default option when MBEDTLS_AES_C is defined otherwise PSA
+ * is used.
+ *
+ * \warning When using PSA, the user should call `psa_crypto_init()` before
+ *          using any CTR_DRBG operation (except `mbedtls_ctr_drbg_init()`).
+ *
  * \note AES-128 will be used if \c MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH is set.
  *
  * \note To achieve a 256-bit security strength with CTR_DRBG,
@@ -2616,7 +2623,9 @@
  * Module:  library/ctr_drbg.c
  * Caller:
  *
- * Requires: MBEDTLS_AES_C
+ * Requires: MBEDTLS_AES_C or
+ *           (PSA_WANT_KEY_TYPE_AES and PSA_WANT_ALG_ECB_NO_PADDING and
+ *            MBEDTLS_PSA_CRYPTO_C)
  *
  * This module provides the CTR_DRBG AES random number generator.
  */
@@ -3155,8 +3164,7 @@
  *
  * Module:  library/psa_crypto.c
  *
- * Requires: MBEDTLS_CIPHER_C,
- *           either MBEDTLS_CTR_DRBG_C and MBEDTLS_ENTROPY_C,
+ * Requires: either MBEDTLS_CTR_DRBG_C and MBEDTLS_ENTROPY_C,
  *           or MBEDTLS_HMAC_DRBG_C and MBEDTLS_ENTROPY_C,
  *           or MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG.
  *
diff --git a/include/mbedtls/pkcs12.h b/include/mbedtls/pkcs12.h
index 42e8453..09f89a2 100644
--- a/include/mbedtls/pkcs12.h
+++ b/include/mbedtls/pkcs12.h
@@ -31,8 +31,8 @@
 #define MBEDTLS_PKCS12_DERIVE_IV        2   /**< initialization vector     */
 #define MBEDTLS_PKCS12_DERIVE_MAC_KEY   3   /**< integrity / MAC key       */
 
-#define MBEDTLS_PKCS12_PBE_DECRYPT      0
-#define MBEDTLS_PKCS12_PBE_ENCRYPT      1
+#define MBEDTLS_PKCS12_PBE_DECRYPT      MBEDTLS_DECRYPT
+#define MBEDTLS_PKCS12_PBE_ENCRYPT      MBEDTLS_ENCRYPT
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/include/mbedtls/pkcs5.h b/include/mbedtls/pkcs5.h
index e004f45..6cfe967 100644
--- a/include/mbedtls/pkcs5.h
+++ b/include/mbedtls/pkcs5.h
@@ -17,6 +17,7 @@
 
 #include "mbedtls/asn1.h"
 #include "mbedtls/md.h"
+#include "mbedtls/cipher.h"
 
 #include <stddef.h>
 #include <stdint.h>
@@ -30,8 +31,8 @@
 /** Given private key password does not allow for correct decryption. */
 #define MBEDTLS_ERR_PKCS5_PASSWORD_MISMATCH               -0x2e00
 
-#define MBEDTLS_PKCS5_DECRYPT      0
-#define MBEDTLS_PKCS5_ENCRYPT      1
+#define MBEDTLS_PKCS5_DECRYPT      MBEDTLS_DECRYPT
+#define MBEDTLS_PKCS5_ENCRYPT      MBEDTLS_ENCRYPT
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/include/mbedtls/ssl.h b/include/mbedtls/ssl.h
index 3c2696f..043988f 100644
--- a/include/mbedtls/ssl.h
+++ b/include/mbedtls/ssl.h
@@ -687,7 +687,6 @@
     MBEDTLS_SSL_SERVER_FINISHED,
     MBEDTLS_SSL_FLUSH_BUFFERS,
     MBEDTLS_SSL_HANDSHAKE_WRAPUP,
-
     MBEDTLS_SSL_NEW_SESSION_TICKET,
     MBEDTLS_SSL_SERVER_HELLO_VERIFY_REQUEST_SENT,
     MBEDTLS_SSL_HELLO_RETRY_REQUEST,
diff --git a/library/aesce.c b/library/aesce.c
index 279c5a2..eaaa5b5 100644
--- a/library/aesce.c
+++ b/library/aesce.c
@@ -23,7 +23,7 @@
  * By defining the macros ourselves we gain access to those declarations without
  * requiring -march on the command line.
  *
- * `arm_neon.h` could be included by any header file, so we put these defines
+ * `arm_neon.h` is included by common.h, so we put these defines
  * at the top of this file, before any includes.
  */
 #define __ARM_FEATURE_CRYPTO 1
@@ -76,12 +76,6 @@
 #    endif
 #endif
 
-#ifdef __ARM_NEON
-#include <arm_neon.h>
-#else
-#error "Target does not support NEON instructions"
-#endif
-
 #if !(defined(__ARM_FEATURE_CRYPTO) || defined(__ARM_FEATURE_AES)) || \
     defined(MBEDTLS_ENABLE_ARM_CRYPTO_EXTENSIONS_COMPILER_FLAG)
 #   if defined(__ARMCOMPILER_VERSION)
@@ -619,6 +613,6 @@
 #undef MBEDTLS_POP_TARGET_PRAGMA
 #endif
 
-#endif /* MBEDTLS_ARCH_IS_ARMV8_A */
+#endif /* MBEDTLS_AESCE_HAVE_CODE */
 
 #endif /* MBEDTLS_AESCE_C */
diff --git a/library/aesce.h b/library/aesce.h
index e2bf58a..a14d085 100644
--- a/library/aesce.h
+++ b/library/aesce.h
@@ -20,7 +20,8 @@
 #include "mbedtls/aes.h"
 
 
-#if defined(MBEDTLS_AESCE_C) && defined(MBEDTLS_ARCH_IS_ARMV8_A) && defined(__ARM_NEON) \
+#if defined(MBEDTLS_AESCE_C) \
+    && defined(MBEDTLS_ARCH_IS_ARMV8_A) && defined(MBEDTLS_HAVE_NEON_INTRINSICS) \
     && (defined(MBEDTLS_COMPILER_IS_GCC) || defined(__clang__) || defined(MSC_VER))
 
 /* MBEDTLS_AESCE_HAVE_CODE is defined if we have a suitable target platform, and a
@@ -129,7 +130,7 @@
 #error "AES hardware acceleration not supported on this platform / compiler"
 #endif
 
-#endif /* MBEDTLS_AESCE_C && MBEDTLS_ARCH_IS_ARMV8_A && __ARM_NEON &&
+#endif /* MBEDTLS_AESCE_C && MBEDTLS_ARCH_IS_ARMV8_A && MBEDTLS_HAVE_NEON_INTRINSICS &&
           (MBEDTLS_COMPILER_IS_GCC || __clang__ || MSC_VER) */
 
 #endif /* MBEDTLS_AESCE_H */
diff --git a/library/alignment.h b/library/alignment.h
index 4bca10e..4aab8e0 100644
--- a/library/alignment.h
+++ b/library/alignment.h
@@ -23,11 +23,16 @@
  * efficient when this is not defined.
  */
 #if defined(__ARM_FEATURE_UNALIGNED) \
-    || defined(__i386__) || defined(__amd64__) || defined(__x86_64__)
+    || defined(MBEDTLS_ARCH_IS_X86) || defined(MBEDTLS_ARCH_IS_X64) \
+    || defined(MBEDTLS_PLATFORM_IS_WINDOWS_ON_ARM64)
 /*
  * __ARM_FEATURE_UNALIGNED is defined where appropriate by armcc, gcc 7, clang 9
  * (and later versions) for Arm v7 and later; all x86 platforms should have
  * efficient unaligned access.
+ *
+ * https://learn.microsoft.com/en-us/cpp/build/arm64-windows-abi-conventions?view=msvc-170#alignment
+ * specifies that on Windows-on-Arm64, unaligned access is safe (except for uncached
+ * device memory).
  */
 #define MBEDTLS_EFFICIENT_UNALIGNED_ACCESS
 #endif
diff --git a/library/bignum.c b/library/bignum.c
index 1b7ff58..1869137 100644
--- a/library/bignum.c
+++ b/library/bignum.c
@@ -102,7 +102,8 @@
  * about whether the assignment was made or not.
  * (Leaking information about the respective sizes of X and Y is ok however.)
  */
-#if defined(_MSC_VER) && defined(_M_ARM64) && (_MSC_FULL_VER < 193131103)
+#if defined(_MSC_VER) && defined(MBEDTLS_PLATFORM_IS_WINDOWS_ON_ARM64) && \
+    (_MSC_FULL_VER < 193131103)
 /*
  * MSVC miscompiles this function if it's inlined prior to Visual Studio 2022 version 17.1. See:
  * https://developercommunity.visualstudio.com/t/c-compiler-miscompiles-part-of-mbedtls-library-on/1646989
diff --git a/library/common.h b/library/common.h
index bd5a0c3..e532777 100644
--- a/library/common.h
+++ b/library/common.h
@@ -21,7 +21,11 @@
 
 #if defined(__ARM_NEON)
 #include <arm_neon.h>
-#endif /* __ARM_NEON */
+#define MBEDTLS_HAVE_NEON_INTRINSICS
+#elif defined(MBEDTLS_PLATFORM_IS_WINDOWS_ON_ARM64)
+#include <arm64_neon.h>
+#define MBEDTLS_HAVE_NEON_INTRINSICS
+#endif
 
 
 #if defined(__GNUC__) && !defined(__ARMCC_VERSION) && !defined(__clang__) \
@@ -178,16 +182,16 @@
 {
     size_t i = 0;
 #if defined(MBEDTLS_EFFICIENT_UNALIGNED_ACCESS)
-#if defined(__ARM_NEON) && \
-    (!defined(MBEDTLS_COMPILER_IS_GCC) || \
-    (defined(MBEDTLS_COMPILER_IS_GCC) && MBEDTLS_GCC_VERSION >= 70300))
+#if defined(MBEDTLS_HAVE_NEON_INTRINSICS) && \
+    (!(defined(MBEDTLS_COMPILER_IS_GCC) && MBEDTLS_GCC_VERSION < 70300))
+    /* Old GCC versions generate a warning here, so disable the NEON path for these compilers */
     for (; (i + 16) <= n; i += 16) {
         uint8x16_t v1 = vld1q_u8(a + i);
         uint8x16_t v2 = vld1q_u8(b + i);
         uint8x16_t x = veorq_u8(v1, v2);
         vst1q_u8(r + i, x);
     }
-#elif defined(__amd64__) || defined(__x86_64__) || defined(__aarch64__)
+#elif defined(MBEDTLS_ARCH_IS_X64) || defined(MBEDTLS_ARCH_IS_ARM64)
     /* This codepath probably only makes sense on architectures with 64-bit registers */
     for (; (i + 8) <= n; i += 8) {
         uint64_t x = mbedtls_get_unaligned_uint64(a + i) ^ mbedtls_get_unaligned_uint64(b + i);
@@ -226,7 +230,7 @@
 {
     size_t i = 0;
 #if defined(MBEDTLS_EFFICIENT_UNALIGNED_ACCESS)
-#if defined(__amd64__) || defined(__x86_64__) || defined(__aarch64__)
+#if defined(MBEDTLS_ARCH_IS_X64) || defined(MBEDTLS_ARCH_IS_ARM64)
     /* This codepath probably only makes sense on architectures with 64-bit registers */
     for (; (i + 8) <= n; i += 8) {
         uint64_t x = mbedtls_get_unaligned_uint64(a + i) ^ mbedtls_get_unaligned_uint64(b + i);
diff --git a/library/constant_time.c b/library/constant_time.c
index c7077c3..d212ddf 100644
--- a/library/constant_time.c
+++ b/library/constant_time.c
@@ -21,19 +21,6 @@
 
 #include <string.h>
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO) && defined(MBEDTLS_SSL_SOME_SUITES_USE_MAC)
-#include "psa/crypto.h"
-/* Define a local translating function to save code size by not using too many
- * arguments in each translating place. */
-static int local_err_translation(psa_status_t status)
-{
-    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
-                                 ARRAY_LENGTH(psa_to_ssl_errors),
-                                 psa_generic_status_to_mbedtls);
-}
-#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
-#endif
-
 #if !defined(MBEDTLS_CT_ASM)
 /*
  * Define an object with the value zero, such that the compiler cannot prove that it
diff --git a/library/ctr_drbg.c b/library/ctr_drbg.c
index cf3816e..da34f95 100644
--- a/library/ctr_drbg.c
+++ b/library/ctr_drbg.c
@@ -24,15 +24,60 @@
 #include <stdio.h>
 #endif
 
+/* Using error translation functions from PSA to MbedTLS */
+#if !defined(MBEDTLS_AES_C)
+#include "psa_util_internal.h"
+#endif
+
 #include "mbedtls/platform.h"
 
+#if !defined(MBEDTLS_AES_C)
+static psa_status_t ctr_drbg_setup_psa_context(mbedtls_ctr_drbg_psa_context *psa_ctx,
+                                               unsigned char *key, size_t key_len)
+{
+    psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
+    psa_status_t status;
+
+    psa_set_key_usage_flags(&key_attr, PSA_KEY_USAGE_ENCRYPT);
+    psa_set_key_algorithm(&key_attr, PSA_ALG_ECB_NO_PADDING);
+    psa_set_key_type(&key_attr, PSA_KEY_TYPE_AES);
+    status = psa_import_key(&key_attr, key, key_len, &psa_ctx->key_id);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    status = psa_cipher_encrypt_setup(&psa_ctx->operation, psa_ctx->key_id, PSA_ALG_ECB_NO_PADDING);
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+exit:
+    psa_reset_key_attributes(&key_attr);
+    return status;
+}
+
+static void ctr_drbg_destroy_psa_contex(mbedtls_ctr_drbg_psa_context *psa_ctx)
+{
+    psa_cipher_abort(&psa_ctx->operation);
+    psa_destroy_key(psa_ctx->key_id);
+
+    psa_ctx->operation = psa_cipher_operation_init();
+    psa_ctx->key_id = MBEDTLS_SVC_KEY_ID_INIT;
+}
+#endif
+
 /*
  * CTR_DRBG context initialization
  */
 void mbedtls_ctr_drbg_init(mbedtls_ctr_drbg_context *ctx)
 {
     memset(ctx, 0, sizeof(mbedtls_ctr_drbg_context));
+#if defined(MBEDTLS_AES_C)
     mbedtls_aes_init(&ctx->aes_ctx);
+#else
+    ctx->psa_ctx.key_id = MBEDTLS_SVC_KEY_ID_INIT;
+    ctx->psa_ctx.operation = psa_cipher_operation_init();
+#endif
     /* Indicate that the entropy nonce length is not set explicitly.
      * See mbedtls_ctr_drbg_set_nonce_len(). */
     ctx->reseed_counter = -1;
@@ -56,7 +101,11 @@
         mbedtls_mutex_free(&ctx->mutex);
     }
 #endif
+#if defined(MBEDTLS_AES_C)
     mbedtls_aes_free(&ctx->aes_ctx);
+#else
+    ctr_drbg_destroy_psa_contex(&ctx->psa_ctx);
+#endif
     mbedtls_platform_zeroize(ctx, sizeof(mbedtls_ctr_drbg_context));
     ctx->reseed_interval = MBEDTLS_CTR_DRBG_RESEED_INTERVAL;
     ctx->reseed_counter = -1;
@@ -117,8 +166,17 @@
     unsigned char key[MBEDTLS_CTR_DRBG_KEYSIZE];
     unsigned char chain[MBEDTLS_CTR_DRBG_BLOCKSIZE];
     unsigned char *p, *iv;
-    mbedtls_aes_context aes_ctx;
     int ret = 0;
+#if defined(MBEDTLS_AES_C)
+    mbedtls_aes_context aes_ctx;
+#else
+    psa_status_t status;
+    size_t tmp_len;
+    mbedtls_ctr_drbg_psa_context psa_ctx;
+
+    psa_ctx.key_id = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_ctx.operation = psa_cipher_operation_init();
+#endif
 
     int i, j;
     size_t buf_len, use_len;
@@ -129,7 +187,6 @@
 
     memset(buf, 0, MBEDTLS_CTR_DRBG_MAX_SEED_INPUT +
            MBEDTLS_CTR_DRBG_BLOCKSIZE + 16);
-    mbedtls_aes_init(&aes_ctx);
 
     /*
      * Construct IV (16 bytes) and S in buffer
@@ -151,10 +208,20 @@
         key[i] = i;
     }
 
+#if defined(MBEDTLS_AES_C)
+    mbedtls_aes_init(&aes_ctx);
+
     if ((ret = mbedtls_aes_setkey_enc(&aes_ctx, key,
                                       MBEDTLS_CTR_DRBG_KEYBITS)) != 0) {
         goto exit;
     }
+#else
+    status = ctr_drbg_setup_psa_context(&psa_ctx, key, sizeof(key));
+    if (status != PSA_SUCCESS) {
+        ret = psa_generic_status_to_mbedtls(status);
+        goto exit;
+    }
+#endif
 
     /*
      * Reduce data to MBEDTLS_CTR_DRBG_SEEDLEN bytes of data
@@ -170,10 +237,19 @@
             use_len -= (use_len >= MBEDTLS_CTR_DRBG_BLOCKSIZE) ?
                        MBEDTLS_CTR_DRBG_BLOCKSIZE : use_len;
 
+#if defined(MBEDTLS_AES_C)
             if ((ret = mbedtls_aes_crypt_ecb(&aes_ctx, MBEDTLS_AES_ENCRYPT,
                                              chain, chain)) != 0) {
                 goto exit;
             }
+#else
+            status = psa_cipher_update(&psa_ctx.operation, chain, MBEDTLS_CTR_DRBG_BLOCKSIZE,
+                                       chain, MBEDTLS_CTR_DRBG_BLOCKSIZE, &tmp_len);
+            if (status != PSA_SUCCESS) {
+                ret = psa_generic_status_to_mbedtls(status);
+                goto exit;
+            }
+#endif
         }
 
         memcpy(tmp + j, chain, MBEDTLS_CTR_DRBG_BLOCKSIZE);
@@ -187,23 +263,46 @@
     /*
      * Do final encryption with reduced data
      */
+#if defined(MBEDTLS_AES_C)
     if ((ret = mbedtls_aes_setkey_enc(&aes_ctx, tmp,
                                       MBEDTLS_CTR_DRBG_KEYBITS)) != 0) {
         goto exit;
     }
+#else
+    ctr_drbg_destroy_psa_contex(&psa_ctx);
+
+    status = ctr_drbg_setup_psa_context(&psa_ctx, tmp, MBEDTLS_CTR_DRBG_KEYSIZE);
+    if (status != PSA_SUCCESS) {
+        ret = psa_generic_status_to_mbedtls(status);
+        goto exit;
+    }
+#endif
     iv = tmp + MBEDTLS_CTR_DRBG_KEYSIZE;
     p = output;
 
     for (j = 0; j < MBEDTLS_CTR_DRBG_SEEDLEN; j += MBEDTLS_CTR_DRBG_BLOCKSIZE) {
+#if defined(MBEDTLS_AES_C)
         if ((ret = mbedtls_aes_crypt_ecb(&aes_ctx, MBEDTLS_AES_ENCRYPT,
                                          iv, iv)) != 0) {
             goto exit;
         }
+#else
+        status = psa_cipher_update(&psa_ctx.operation, iv, MBEDTLS_CTR_DRBG_BLOCKSIZE,
+                                   iv, MBEDTLS_CTR_DRBG_BLOCKSIZE, &tmp_len);
+        if (status != PSA_SUCCESS) {
+            ret = psa_generic_status_to_mbedtls(status);
+            goto exit;
+        }
+#endif
         memcpy(p, iv, MBEDTLS_CTR_DRBG_BLOCKSIZE);
         p += MBEDTLS_CTR_DRBG_BLOCKSIZE;
     }
 exit:
+#if defined(MBEDTLS_AES_C)
     mbedtls_aes_free(&aes_ctx);
+#else
+    ctr_drbg_destroy_psa_contex(&psa_ctx);
+#endif
     /*
      * tidy up the stack
      */
@@ -236,6 +335,10 @@
     unsigned char *p = tmp;
     int i, j;
     int ret = 0;
+#if !defined(MBEDTLS_AES_C)
+    psa_status_t status;
+    size_t tmp_len;
+#endif
 
     memset(tmp, 0, MBEDTLS_CTR_DRBG_SEEDLEN);
 
@@ -252,10 +355,19 @@
         /*
          * Crypt counter block
          */
+#if defined(MBEDTLS_AES_C)
         if ((ret = mbedtls_aes_crypt_ecb(&ctx->aes_ctx, MBEDTLS_AES_ENCRYPT,
                                          ctx->counter, p)) != 0) {
             goto exit;
         }
+#else
+        status = psa_cipher_update(&ctx->psa_ctx.operation, ctx->counter, sizeof(ctx->counter),
+                                   p, MBEDTLS_CTR_DRBG_BLOCKSIZE, &tmp_len);
+        if (status != PSA_SUCCESS) {
+            ret = psa_generic_status_to_mbedtls(status);
+            goto exit;
+        }
+#endif
 
         p += MBEDTLS_CTR_DRBG_BLOCKSIZE;
     }
@@ -267,10 +379,20 @@
     /*
      * Update key and counter
      */
+#if defined(MBEDTLS_AES_C)
     if ((ret = mbedtls_aes_setkey_enc(&ctx->aes_ctx, tmp,
                                       MBEDTLS_CTR_DRBG_KEYBITS)) != 0) {
         goto exit;
     }
+#else
+    ctr_drbg_destroy_psa_contex(&ctx->psa_ctx);
+
+    status = ctr_drbg_setup_psa_context(&ctx->psa_ctx, tmp, MBEDTLS_CTR_DRBG_KEYSIZE);
+    if (status != PSA_SUCCESS) {
+        ret = psa_generic_status_to_mbedtls(status);
+        goto exit;
+    }
+#endif
     memcpy(ctx->counter, tmp + MBEDTLS_CTR_DRBG_KEYSIZE,
            MBEDTLS_CTR_DRBG_BLOCKSIZE);
 
@@ -447,10 +569,20 @@
                  good_nonce_len(ctx->entropy_len));
 
     /* Initialize with an empty key. */
+#if defined(MBEDTLS_AES_C)
     if ((ret = mbedtls_aes_setkey_enc(&ctx->aes_ctx, key,
                                       MBEDTLS_CTR_DRBG_KEYBITS)) != 0) {
         return ret;
     }
+#else
+    psa_status_t status;
+
+    status = ctr_drbg_setup_psa_context(&ctx->psa_ctx, key, MBEDTLS_CTR_DRBG_KEYSIZE);
+    if (status != PSA_SUCCESS) {
+        ret = psa_generic_status_to_mbedtls(status);
+        return status;
+    }
+#endif
 
     /* Do the initial seeding. */
     if ((ret = mbedtls_ctr_drbg_reseed_internal(ctx, custom, len,
@@ -531,10 +663,22 @@
         /*
          * Crypt counter block
          */
+#if defined(MBEDTLS_AES_C)
         if ((ret = mbedtls_aes_crypt_ecb(&ctx->aes_ctx, MBEDTLS_AES_ENCRYPT,
                                          ctx->counter, tmp)) != 0) {
             goto exit;
         }
+#else
+        psa_status_t status;
+        size_t tmp_len;
+
+        status = psa_cipher_update(&ctx->psa_ctx.operation, ctx->counter, sizeof(ctx->counter),
+                                   tmp, MBEDTLS_CTR_DRBG_BLOCKSIZE, &tmp_len);
+        if (status != PSA_SUCCESS) {
+            ret = psa_generic_status_to_mbedtls(status);
+            goto exit;
+        }
+#endif
 
         use_len = (output_len > MBEDTLS_CTR_DRBG_BLOCKSIZE)
             ? MBEDTLS_CTR_DRBG_BLOCKSIZE : output_len;
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 182d07f..0fb3c42 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -955,37 +955,34 @@
     return 0;
 }
 
-/* Common helper for ECDSA sign using PSA functions. */
+/* Common helper for ECDSA sign using PSA functions.
+ * Instead of extracting key's properties in order to check which kind of ECDSA
+ * signature it supports, we try both deterministic and non-deterministic.
+ */
 static int ecdsa_sign_psa(mbedtls_svc_key_id_t key_id, mbedtls_md_type_t md_alg,
                           const unsigned char *hash, size_t hash_len,
                           unsigned char *sig, size_t sig_size, size_t *sig_len)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_status_t status;
-    psa_algorithm_t psa_sig_md;
-    psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
-    psa_algorithm_t alg, alg2;
 
-    status = psa_get_key_attributes(key_id, &key_attr);
-    if (status != PSA_SUCCESS) {
+    status = psa_sign_hash(key_id,
+                           PSA_ALG_DETERMINISTIC_ECDSA(mbedtls_md_psa_alg_from_type(md_alg)),
+                           hash, hash_len, sig, sig_size, sig_len);
+    if (status == PSA_SUCCESS) {
+        goto done;
+    } else if (status != PSA_ERROR_NOT_PERMITTED) {
         return PSA_PK_ECDSA_TO_MBEDTLS_ERR(status);
     }
-    alg = psa_get_key_algorithm(&key_attr);
-    alg2 = psa_get_key_enrollment_algorithm(&key_attr);
-    psa_reset_key_attributes(&key_attr);
 
-    if (PSA_ALG_IS_DETERMINISTIC_ECDSA(alg) || PSA_ALG_IS_DETERMINISTIC_ECDSA(alg2)) {
-        psa_sig_md = PSA_ALG_DETERMINISTIC_ECDSA(mbedtls_md_psa_alg_from_type(md_alg));
-    } else {
-        psa_sig_md = PSA_ALG_ECDSA(mbedtls_md_psa_alg_from_type(md_alg));
-    }
-
-    status = psa_sign_hash(key_id, psa_sig_md, hash, hash_len,
-                           sig, sig_size, sig_len);
+    status = psa_sign_hash(key_id,
+                           PSA_ALG_ECDSA(mbedtls_md_psa_alg_from_type(md_alg)),
+                           hash, hash_len, sig, sig_size, sig_len);
     if (status != PSA_SUCCESS) {
         return PSA_PK_ECDSA_TO_MBEDTLS_ERR(status);
     }
 
+done:
     ret = pk_ecdsa_sig_asn1_from_psa(sig, sig_len, sig_size);
 
     return ret;
diff --git a/library/pkparse.c b/library/pkparse.c
index edebf92..18498e5 100644
--- a/library/pkparse.c
+++ b/library/pkparse.c
@@ -105,16 +105,21 @@
 {
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_usage_t flags;
     psa_status_t status;
 
     psa_set_key_type(&attributes, PSA_KEY_TYPE_ECC_KEY_PAIR(pk->ec_family));
-    psa_set_key_algorithm(&attributes, PSA_ALG_ECDH);
-    psa_key_usage_t flags = PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_DERIVE;
-    /* Montgomery allows only ECDH, others ECDSA too */
-    if (pk->ec_family != PSA_ECC_FAMILY_MONTGOMERY) {
-        flags |= PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_SIGN_MESSAGE;
-        psa_set_key_enrollment_algorithm(&attributes,
-                                         MBEDTLS_PK_PSA_ALG_ECDSA_MAYBE_DET(PSA_ALG_ANY_HASH));
+    if (pk->ec_family == PSA_ECC_FAMILY_MONTGOMERY) {
+        /* Do not set algorithm here because Montgomery keys cannot do ECDSA and
+         * the PK module cannot do ECDH. When the key will be used in TLS for
+         * ECDH, it will be exported and then re-imported with proper flags
+         * and algorithm. */
+        flags = PSA_KEY_USAGE_EXPORT;
+    } else {
+        psa_set_key_algorithm(&attributes,
+                              MBEDTLS_PK_PSA_ALG_ECDSA_MAYBE_DET(PSA_ALG_ANY_HASH));
+        flags = PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_SIGN_MESSAGE |
+                PSA_KEY_USAGE_EXPORT;
     }
     psa_set_key_usage_flags(&attributes, flags);
 
diff --git a/library/platform_util.c b/library/platform_util.c
index 6d2dd14..63643d2 100644
--- a/library/platform_util.c
+++ b/library/platform_util.c
@@ -243,10 +243,10 @@
 #include <time.h>
 #if !defined(_WIN32) && \
     (defined(unix) || defined(__unix) || defined(__unix__) || \
-    (defined(__APPLE__) && defined(__MACH__)))
+    (defined(__APPLE__) && defined(__MACH__)) || defined(__HAIKU__))
 #include <unistd.h>
-#endif /* !_WIN32 && (unix || __unix || __unix__ || (__APPLE__ && __MACH__)) */
-#if (defined(_POSIX_VERSION) && _POSIX_VERSION >= 199309L)
+#endif /* !_WIN32 && (unix || __unix || __unix__ || (__APPLE__ && __MACH__) || __HAIKU__) */
+#if (defined(_POSIX_VERSION) && _POSIX_VERSION >= 199309L) || defined(__HAIKU__)
 mbedtls_ms_time_t mbedtls_ms_time(void)
 {
     int ret;
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5455fa2..d61e118 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7720,6 +7720,11 @@
     }
     global_data.drivers_initialized = 1;
 
+    status = psa_initialize_key_slots();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+
     /* Initialize and seed the random generator. */
     mbedtls_psa_random_init(&global_data.rng);
     global_data.rng_state = RNG_INITIALIZED;
@@ -7729,11 +7734,6 @@
     }
     global_data.rng_state = RNG_SEEDED;
 
-    status = psa_initialize_key_slots();
-    if (status != PSA_SUCCESS) {
-        goto exit;
-    }
-
 #if defined(PSA_CRYPTO_STORAGE_HAS_TRANSACTIONS)
     status = psa_crypto_load_transaction();
     if (status == PSA_SUCCESS) {
diff --git a/library/sha256.c b/library/sha256.c
index 45ad6d8..8788981 100644
--- a/library/sha256.c
+++ b/library/sha256.c
@@ -28,7 +28,7 @@
  * By defining the macros ourselves we gain access to those declarations without
  * requiring -march on the command line.
  *
- * `arm_neon.h` could be included by any header file, so we put these defines
+ * `arm_neon.h` is included by common.h, so we put these defines
  * at the top of this file, before any includes.
  */
 #define __ARM_FEATURE_CRYPTO 1
@@ -62,9 +62,7 @@
 
 #  if defined(MBEDTLS_SHA256_USE_ARMV8_A_CRYPTO_IF_PRESENT) || \
     defined(MBEDTLS_SHA256_USE_ARMV8_A_CRYPTO_ONLY)
-#       ifdef __ARM_NEON
-#           include <arm_neon.h>
-#       else
+#       if !defined(MBEDTLS_HAVE_NEON_INTRINSICS)
 #           if defined(MBEDTLS_SHA256_USE_ARMV8_A_CRYPTO_IF_PRESENT)
 #               warning "Target does not support NEON instructions"
 #               undef MBEDTLS_SHA256_USE_ARMV8_A_CRYPTO_IF_PRESENT
@@ -126,12 +124,7 @@
 #      include <signal.h>
 #    endif
 #  endif
-#elif defined(_M_ARM64)
-#  if defined(MBEDTLS_SHA256_USE_ARMV8_A_CRYPTO_IF_PRESENT) || \
-    defined(MBEDTLS_SHA256_USE_ARMV8_A_CRYPTO_ONLY)
-#    include <arm64_neon.h>
-#  endif
-#else
+#elif !defined(MBEDTLS_PLATFORM_IS_WINDOWS_ON_ARM64)
 #  undef MBEDTLS_SHA256_USE_ARMV8_A_CRYPTO_ONLY
 #  undef MBEDTLS_SHA256_USE_ARMV8_A_CRYPTO_IF_PRESENT
 #endif
@@ -156,7 +149,7 @@
 {
     return 1;
 }
-#elif defined(_M_ARM64)
+#elif defined(MBEDTLS_PLATFORM_IS_WINDOWS_ON_ARM64)
 #define WIN32_LEAN_AND_MEAN
 #include <Windows.h>
 #include <processthreadsapi.h>
diff --git a/library/sha512.c b/library/sha512.c
index e7af121..6011254 100644
--- a/library/sha512.c
+++ b/library/sha512.c
@@ -19,7 +19,7 @@
  * By defining the macros ourselves we gain access to those declarations without
  * requiring -march on the command line.
  *
- * `arm_neon.h` could be included by any header file, so we put these defines
+ * `arm_neon.h` is included by common.h, so we put these defines
  * at the top of this file, before any includes.
  */
 #define __ARM_FEATURE_SHA512 1
@@ -48,9 +48,7 @@
 #  if defined(MBEDTLS_SHA512_USE_A64_CRYPTO_IF_PRESENT) || \
     defined(MBEDTLS_SHA512_USE_A64_CRYPTO_ONLY)
 /* *INDENT-OFF* */
-#   ifdef __ARM_NEON
-#       include <arm_neon.h>
-#   else
+#   if !defined(MBEDTLS_HAVE_NEON_INTRINSICS)
 #       error "Target does not support NEON instructions"
 #   endif
 /*
@@ -109,12 +107,7 @@
 #      include <signal.h>
 #    endif
 #  endif
-#elif defined(_M_ARM64)
-#  if defined(MBEDTLS_SHA512_USE_A64_CRYPTO_IF_PRESENT) || \
-    defined(MBEDTLS_SHA512_USE_A64_CRYPTO_ONLY)
-#    include <arm64_neon.h>
-#  endif
-#else
+#elif !defined(MBEDTLS_PLATFORM_IS_WINDOWS_ON_ARM64)
 #  undef MBEDTLS_SHA512_USE_A64_CRYPTO_ONLY
 #  undef MBEDTLS_SHA512_USE_A64_CRYPTO_IF_PRESENT
 #endif
@@ -142,7 +135,7 @@
                            NULL, 0);
     return ret == 0 && value != 0;
 }
-#elif defined(_M_ARM64)
+#elif defined(MBEDTLS_PLATFORM_IS_WINDOWS_ON_ARM64)
 /*
  * As of March 2022, there don't appear to be any PF_ARM_V8_* flags
  * available to pass to IsProcessorFeaturePresent() to check for
diff --git a/library/ssl_client.c b/library/ssl_client.c
index 55fe352..270db41 100644
--- a/library/ssl_client.c
+++ b/library/ssl_client.c
@@ -753,7 +753,8 @@
         session_negotiate->ticket != NULL) {
         mbedtls_ms_time_t now = mbedtls_ms_time();
         mbedtls_ms_time_t age = now - session_negotiate->ticket_reception_time;
-        if (age < 0 || age > session_negotiate->ticket_lifetime * 1000) {
+        if (age < 0 ||
+            age > (mbedtls_ms_time_t) session_negotiate->ticket_lifetime * 1000) {
             /* Without valid ticket, disable session resumption.*/
             MBEDTLS_SSL_DEBUG_MSG(
                 3, ("Ticket expired, disable session resumption"));
diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index eae192b..b9801a0 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -2113,6 +2113,7 @@
 
 #if defined(MBEDTLS_SSL_EARLY_DATA)
 int mbedtls_ssl_tls13_write_early_data_ext(mbedtls_ssl_context *ssl,
+                                           int in_new_session_ticket,
                                            unsigned char *buf,
                                            const unsigned char *end,
                                            size_t *out_len);
@@ -2795,6 +2796,13 @@
                                                    MBEDTLS_SSL_TLS1_3_TICKET_ALLOW_PSK_EPHEMERAL_RESUMPTION);
 }
 
+static inline unsigned int mbedtls_ssl_session_ticket_allow_early_data(
+    mbedtls_ssl_session *session)
+{
+    return !mbedtls_ssl_session_check_ticket_flags(session,
+                                                   MBEDTLS_SSL_TLS1_3_TICKET_ALLOW_EARLY_DATA);
+}
+
 static inline void mbedtls_ssl_session_set_ticket_flags(
     mbedtls_ssl_session *session, unsigned int flags)
 {
diff --git a/library/ssl_msg.c b/library/ssl_msg.c
index 48d0a03..6579c96 100644
--- a/library/ssl_msg.c
+++ b/library/ssl_msg.c
@@ -3627,7 +3627,7 @@
 static int ssl_handle_possible_reconnect(mbedtls_ssl_context *ssl)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t len;
+    size_t len = 0;
 
     if (ssl->conf->f_cookie_write == NULL ||
         ssl->conf->f_cookie_check == NULL) {
diff --git a/library/ssl_ticket.c b/library/ssl_ticket.c
index 8e7c145..cd730fb 100644
--- a/library/ssl_ticket.c
+++ b/library/ssl_ticket.c
@@ -510,7 +510,8 @@
         }
 #endif
 
-        mbedtls_ms_time_t ticket_lifetime = ctx->ticket_lifetime * 1000;
+        mbedtls_ms_time_t ticket_lifetime =
+            (mbedtls_ms_time_t) ctx->ticket_lifetime * 1000;
 
         if (ticket_age < 0 || ticket_age > ticket_lifetime) {
             ret = MBEDTLS_ERR_SSL_SESSION_TICKET_EXPIRED;
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 6678b71..4daf2e7 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -306,7 +306,7 @@
 {
     unsigned char *resized_buffer = mbedtls_calloc(1, len_new);
     if (resized_buffer == NULL) {
-        return -1;
+        return MBEDTLS_ERR_SSL_ALLOC_FAILED;
     }
 
     /* We want to copy len_new bytes when downsizing the buffer, and
@@ -2623,8 +2623,7 @@
 
 #if defined(MBEDTLS_SSL_CLI_C)
     if (session->endpoint == MBEDTLS_SSL_IS_CLIENT) {
-#if defined(MBEDTLS_SSL_SERVER_NAME_INDICATION) && \
-        defined(MBEDTLS_SSL_SESSION_TICKETS)
+#if defined(MBEDTLS_SSL_SERVER_NAME_INDICATION)
         size_t hostname_len;
         /* load host name */
         if (end - p < 2) {
@@ -2644,8 +2643,7 @@
             memcpy(session->hostname, p, hostname_len);
             p += hostname_len;
         }
-#endif /* MBEDTLS_SSL_SERVER_NAME_INDICATION &&
-          MBEDTLS_SSL_SESSION_TICKETS */
+#endif /* MBEDTLS_SSL_SERVER_NAME_INDICATION */
 
 #if defined(MBEDTLS_HAVE_TIME)
         if (end - p < 8) {
diff --git a/library/ssl_tls12_server.c b/library/ssl_tls12_server.c
index a07d0fb..923b093 100644
--- a/library/ssl_tls12_server.c
+++ b/library/ssl_tls12_server.c
@@ -2597,12 +2597,12 @@
     mbedtls_pk_context *pk;
     mbedtls_pk_type_t pk_type;
     psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    unsigned char buf[PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)];
+    size_t key_len;
 #if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     uint16_t tls_id = 0;
     psa_key_type_t key_type = PSA_KEY_TYPE_NONE;
-    size_t key_len;
     mbedtls_ecp_group_id grp_id;
-    unsigned char buf[PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)];
     mbedtls_ecp_keypair *key;
 #endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
 
@@ -2625,22 +2625,51 @@
                 return MBEDTLS_ERR_SSL_PK_TYPE_MISMATCH;
             }
 
-            ssl->handshake->xxdh_psa_privkey = pk->priv_id;
-
-            /* Key should not be destroyed in the TLS library */
-            ssl->handshake->xxdh_psa_privkey_is_external = 1;
-
-            status = psa_get_key_attributes(ssl->handshake->xxdh_psa_privkey,
-                                            &key_attributes);
+            /* Get the attributes of the key previously parsed by PK module in
+             * order to extract its type and length (in bits). */
+            status = psa_get_key_attributes(pk->priv_id, &key_attributes);
             if (status != PSA_SUCCESS) {
-                ssl->handshake->xxdh_psa_privkey = MBEDTLS_SVC_KEY_ID_INIT;
-                return PSA_TO_MBEDTLS_ERR(status);
+                ret = PSA_TO_MBEDTLS_ERR(status);
+                goto exit;
             }
-
             ssl->handshake->xxdh_psa_type = psa_get_key_type(&key_attributes);
             ssl->handshake->xxdh_psa_bits = psa_get_key_bits(&key_attributes);
 
-            psa_reset_key_attributes(&key_attributes);
+            if (pk_type == MBEDTLS_PK_OPAQUE) {
+                /* Opaque key is created by the user (externally from Mbed TLS)
+                 * so we assume it already has the right algorithm and flags
+                 * set. Just copy its ID as reference. */
+                ssl->handshake->xxdh_psa_privkey = pk->priv_id;
+                ssl->handshake->xxdh_psa_privkey_is_external = 1;
+            } else {
+                /* PK_ECKEY[_DH] and PK_ECDSA instead as parsed from the PK
+                 * module and only have ECDSA capabilities. Since we need
+                 * them for ECDH later, we export and then re-import them with
+                 * proper flags and algorithm. Of course We also set key's type
+                 * and bits that we just got above. */
+                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->xxdh_psa_type));
+                psa_set_key_bits(&key_attributes, ssl->handshake->xxdh_psa_bits);
+
+                status = psa_export_key(pk->priv_id, buf, sizeof(buf), &key_len);
+                if (status != PSA_SUCCESS) {
+                    ret = PSA_TO_MBEDTLS_ERR(status);
+                    goto exit;
+                }
+                status = psa_import_key(&key_attributes, buf, key_len,
+                                        &ssl->handshake->xxdh_psa_privkey);
+                if (status != PSA_SUCCESS) {
+                    ret = PSA_TO_MBEDTLS_ERR(status);
+                    goto exit;
+                }
+
+                /* Set this key as owned by the TLS library: it will be its duty
+                 * to clear it exit. */
+                ssl->handshake->xxdh_psa_privkey_is_external = 0;
+            }
 
             ret = 0;
             break;
@@ -2696,6 +2725,10 @@
             ret = MBEDTLS_ERR_SSL_PK_TYPE_MISMATCH;
     }
 
+exit:
+    psa_reset_key_attributes(&key_attributes);
+    mbedtls_platform_zeroize(buf, sizeof(buf));
+
     return ret;
 }
 #else /* MBEDTLS_USE_PSA_CRYPTO */
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 44814b9..ae11364 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -695,10 +695,8 @@
     mbedtls_ssl_session *session = ssl->session_negotiate;
     return ssl->handshake->resume &&
            session->tls_version == MBEDTLS_SSL_VERSION_TLS1_3 &&
-           (session->ticket_flags &
-            MBEDTLS_SSL_TLS1_3_TICKET_ALLOW_EARLY_DATA) &&
-           mbedtls_ssl_tls13_cipher_suite_is_offered(
-        ssl, session->ciphersuite);
+           mbedtls_ssl_session_ticket_allow_early_data(session) &&
+           mbedtls_ssl_tls13_cipher_suite_is_offered(ssl, session->ciphersuite);
 }
 #endif
 
@@ -1176,7 +1174,9 @@
     if (mbedtls_ssl_conf_tls13_some_psk_enabled(ssl) &&
         ssl_tls13_early_data_has_valid_ticket(ssl) &&
         ssl->conf->early_data_enabled == MBEDTLS_SSL_EARLY_DATA_ENABLED) {
-        ret = mbedtls_ssl_tls13_write_early_data_ext(ssl, p, end, &ext_len);
+
+        ret = mbedtls_ssl_tls13_write_early_data_ext(
+            ssl, 0, p, end, &ext_len);
         if (ret != 0) {
             return ret;
         }
@@ -1925,7 +1925,6 @@
 
     mbedtls_ssl_set_inbound_transform(ssl, handshake->transform_handshake);
     MBEDTLS_SSL_DEBUG_MSG(1, ("Switch to handshake keys for inbound traffic"));
-    ssl->session_negotiate->ciphersuite = handshake->ciphersuite_info->id;
     ssl->session_in = ssl->session_negotiate;
 
 cleanup:
@@ -2153,9 +2152,7 @@
     int ret;
     unsigned char *buf;
     size_t buf_len;
-#if defined(MBEDTLS_SSL_EARLY_DATA)
     mbedtls_ssl_handshake_params *handshake = ssl->handshake;
-#endif
 
     MBEDTLS_SSL_DEBUG_MSG(2, ("=> parse encrypted extensions"));
 
@@ -2185,11 +2182,19 @@
          * - The selected cipher suite
          * - The selected ALPN [RFC7301] protocol, if any
          *
-         * We check here that when early data is involved the server
-         * selected the cipher suite associated to the pre-shared key
-         * as it must have.
+         * The server has sent an early data extension in its Encrypted
+         * Extension message thus accepted to receive early data. We
+         * check here that the additional constraints on the handshake
+         * parameters, when early data are exchanged, are met,
+         * namely:
+         * - a PSK has been selected for the handshake
+         * - the selected PSK for the handshake was the first one proposed
+         *   by the client.
+         * - the selected ciphersuite for the handshake is the ciphersuite
+         *   associated with the selected PSK.
          */
-        if (handshake->selected_identity != 0 ||
+        if ((!mbedtls_ssl_tls13_key_exchange_mode_with_psk(ssl)) ||
+            handshake->selected_identity != 0 ||
             handshake->ciphersuite_info->id !=
             ssl->session_negotiate->ciphersuite) {
 
@@ -2203,6 +2208,18 @@
     }
 #endif
 
+    /*
+     * In case the client has proposed a PSK associated with a ticket,
+     * `ssl->session_negotiate->ciphersuite` still contains at this point the
+     * identifier of the ciphersuite associated with the ticket. This is that
+     * way because, if an exchange of early data is agreed upon, we need
+     * it to check that the ciphersuite selected for the handshake is the
+     * ticket ciphersuite (see above). This information is not needed
+     * anymore thus we can now set it to the identifier of the ciphersuite
+     * used in this session under negotiation.
+     */
+    ssl->session_negotiate->ciphersuite = handshake->ciphersuite_info->id;
+
     MBEDTLS_SSL_PROC_CHK(mbedtls_ssl_add_hs_msg_to_checksum(
                              ssl, MBEDTLS_SSL_HS_ENCRYPTED_EXTENSIONS,
                              buf, buf_len));
@@ -2647,6 +2664,37 @@
 
 #if defined(MBEDTLS_SSL_SESSION_TICKETS)
 
+#if defined(MBEDTLS_SSL_EARLY_DATA)
+/* From RFC 8446 section 4.2.10
+ *
+ * struct {
+ *     select (Handshake.msg_type) {
+ *         case new_session_ticket:   uint32 max_early_data_size;
+ *         ...
+ *     };
+ * } EarlyDataIndication;
+ */
+MBEDTLS_CHECK_RETURN_CRITICAL
+static int ssl_tls13_parse_new_session_ticket_early_data_ext(
+    mbedtls_ssl_context *ssl,
+    const unsigned char *buf,
+    const unsigned char *end)
+{
+    mbedtls_ssl_session *session = ssl->session;
+
+    MBEDTLS_SSL_CHK_BUF_READ_PTR(buf, end, 4);
+
+    session->max_early_data_size = MBEDTLS_GET_UINT32_BE(buf, 0);
+    mbedtls_ssl_session_set_ticket_flags(
+        session, MBEDTLS_SSL_TLS1_3_TICKET_ALLOW_EARLY_DATA);
+    MBEDTLS_SSL_DEBUG_MSG(
+        3, ("received max_early_data_size: %u",
+            (unsigned int) session->max_early_data_size));
+
+    return 0;
+}
+#endif /* MBEDTLS_SSL_EARLY_DATA */
+
 MBEDTLS_CHECK_RETURN_CRITICAL
 static int ssl_tls13_parse_new_session_ticket_exts(mbedtls_ssl_context *ssl,
                                                    const unsigned char *buf,
@@ -2680,15 +2728,12 @@
         switch (extension_type) {
 #if defined(MBEDTLS_SSL_EARLY_DATA)
             case MBEDTLS_TLS_EXT_EARLY_DATA:
-                if (extension_data_len != 4) {
-                    MBEDTLS_SSL_PEND_FATAL_ALERT(
-                        MBEDTLS_SSL_ALERT_MSG_DECODE_ERROR,
-                        MBEDTLS_ERR_SSL_DECODE_ERROR);
-                    return MBEDTLS_ERR_SSL_DECODE_ERROR;
-                }
-                if (ssl->session != NULL) {
-                    ssl->session->ticket_flags |=
-                        MBEDTLS_SSL_TLS1_3_TICKET_ALLOW_EARLY_DATA;
+                ret = ssl_tls13_parse_new_session_ticket_early_data_ext(
+                    ssl, p, p + extension_data_len);
+                if (ret != 0) {
+                    MBEDTLS_SSL_DEBUG_RET(
+                        1, "ssl_tls13_parse_new_session_ticket_early_data_ext",
+                        ret);
                 }
                 break;
 #endif /* MBEDTLS_SSL_EARLY_DATA */
diff --git a/library/ssl_tls13_generic.c b/library/ssl_tls13_generic.c
index cc77a94..fe2a2eb 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -1402,7 +1402,7 @@
  *
  * struct {
  *   select ( Handshake.msg_type ) {
- *     ...
+ *     case new_session_ticket:   uint32 max_early_data_size;
  *     case client_hello:         Empty;
  *     case encrypted_extensions: Empty;
  *   };
@@ -1410,20 +1410,37 @@
  */
 #if defined(MBEDTLS_SSL_EARLY_DATA)
 int mbedtls_ssl_tls13_write_early_data_ext(mbedtls_ssl_context *ssl,
+                                           int in_new_session_ticket,
                                            unsigned char *buf,
                                            const unsigned char *end,
                                            size_t *out_len)
 {
     unsigned char *p = buf;
-    *out_len = 0;
-    ((void) ssl);
 
-    MBEDTLS_SSL_CHK_BUF_PTR(p, end, 4);
+#if defined(MBEDTLS_SSL_SRV_C)
+    const size_t needed = in_new_session_ticket ? 8 : 4;
+#else
+    const size_t needed = 4;
+    ((void) in_new_session_ticket);
+#endif
+
+    *out_len = 0;
+
+    MBEDTLS_SSL_CHK_BUF_PTR(p, end, needed);
 
     MBEDTLS_PUT_UINT16_BE(MBEDTLS_TLS_EXT_EARLY_DATA, p, 0);
-    MBEDTLS_PUT_UINT16_BE(0, p, 2);
+    MBEDTLS_PUT_UINT16_BE(needed - 4, p, 2);
 
-    *out_len = 4;
+#if defined(MBEDTLS_SSL_SRV_C)
+    if (in_new_session_ticket) {
+        MBEDTLS_PUT_UINT32_BE(ssl->conf->max_early_data_size, p, 4);
+        MBEDTLS_SSL_DEBUG_MSG(
+            4, ("Sent max_early_data_size=%u",
+                (unsigned int) ssl->conf->max_early_data_size));
+    }
+#endif
+
+    *out_len = needed;
 
     mbedtls_ssl_tls13_set_hs_sent_ext_mask(ssl, MBEDTLS_TLS_EXT_EARLY_DATA);
 
diff --git a/library/ssl_tls13_server.c b/library/ssl_tls13_server.c
index d983a00..fe7a674 100644
--- a/library/ssl_tls13_server.c
+++ b/library/ssl_tls13_server.c
@@ -1845,6 +1845,13 @@
 
     }
 
+    if (!mbedtls_ssl_session_ticket_allow_early_data(ssl->session_negotiate)) {
+        MBEDTLS_SSL_DEBUG_MSG(
+            1,
+            ("EarlyData: rejected, early_data not allowed in ticket "
+             "permission bits."));
+        return;
+    }
 
     ssl->early_data_status = MBEDTLS_SSL_EARLY_DATA_STATUS_ACCEPTED;
 
@@ -1880,10 +1887,18 @@
 #if defined(MBEDTLS_SSL_EARLY_DATA)
     /* There is enough information, update early data state. */
     ssl_tls13_update_early_data_status(ssl);
+
+    if (ssl->early_data_status == MBEDTLS_SSL_EARLY_DATA_STATUS_ACCEPTED) {
+        ret = mbedtls_ssl_tls13_compute_early_transform(ssl);
+        if (ret != 0) {
+            MBEDTLS_SSL_DEBUG_RET(
+                1, "mbedtls_ssl_tls13_compute_early_transform", ret);
+            return ret;
+        }
+    }
 #endif /* MBEDTLS_SSL_EARLY_DATA */
 
     return 0;
-
 }
 
 /*
@@ -2509,7 +2524,8 @@
 
 #if defined(MBEDTLS_SSL_EARLY_DATA)
     if (ssl->early_data_status == MBEDTLS_SSL_EARLY_DATA_STATUS_ACCEPTED) {
-        ret = mbedtls_ssl_tls13_write_early_data_ext(ssl, p, end, &output_len);
+        ret = mbedtls_ssl_tls13_write_early_data_ext(
+            ssl, 0, p, end, &output_len);
         if (ret != 0) {
             return ret;
         }
@@ -2749,6 +2765,59 @@
 #endif /* MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_EPHEMERAL_ENABLED */
 
 /*
+ * RFC 8446 section A.2
+ *
+ *                                | Send ServerHello
+ *                                | K_send = handshake
+ *                                | Send EncryptedExtensions
+ *                                | [Send CertificateRequest]
+ * Can send                       | [Send Certificate + CertificateVerify]
+ * app data                       | Send Finished
+ * after   -->                    | K_send = application
+ * here                  +--------+--------+
+ *              No 0-RTT |                 | 0-RTT
+ *                       |                 |
+ *   K_recv = handshake  |                 | K_recv = early data
+ * [Skip decrypt errors] |    +------> WAIT_EOED -+
+ *                       |    |       Recv |      | Recv EndOfEarlyData
+ *                       |    | early data |      | K_recv = handshake
+ *                       |    +------------+      |
+ *                       |                        |
+ *                       +> WAIT_FLIGHT2 <--------+
+ *                                |
+ *                       +--------+--------+
+ *               No auth |                 | Client auth
+ *                       |                 |
+ *                       |                 v
+ *                       |             WAIT_CERT
+ *                       |        Recv |       | Recv Certificate
+ *                       |       empty |       v
+ *                       | Certificate |    WAIT_CV
+ *                       |             |       | Recv
+ *                       |             v       | CertificateVerify
+ *                       +-> WAIT_FINISHED <---+
+ *                                | Recv Finished
+ *
+ *
+ * The following function handles the state changes after WAIT_FLIGHT2 in the
+ * above diagram. We are not going to receive early data related messages
+ * anymore, prepare to receive the first handshake message of the client
+ * second flight.
+ */
+static void ssl_tls13_prepare_for_handshake_second_flight(
+    mbedtls_ssl_context *ssl)
+{
+    if (ssl->handshake->certificate_request_sent) {
+        mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_CLIENT_CERTIFICATE);
+    } else {
+        MBEDTLS_SSL_DEBUG_MSG(2, ("skip parse certificate"));
+        MBEDTLS_SSL_DEBUG_MSG(2, ("skip parse certificate verify"));
+
+        mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_CLIENT_FINISHED);
+    }
+}
+
+/*
  * Handler for MBEDTLS_SSL_SERVER_FINISHED
  */
 MBEDTLS_CHECK_RETURN_CRITICAL
@@ -2769,20 +2838,200 @@
         return ret;
     }
 
-    MBEDTLS_SSL_DEBUG_MSG(1, ("Switch to handshake keys for inbound traffic"));
+#if defined(MBEDTLS_SSL_EARLY_DATA)
+    if (ssl->early_data_status == MBEDTLS_SSL_EARLY_DATA_STATUS_ACCEPTED) {
+        /* See RFC 8446 section A.2 for more information */
+        MBEDTLS_SSL_DEBUG_MSG(
+            1, ("Switch to early keys for inbound traffic. "
+                "( K_recv = early data )"));
+        mbedtls_ssl_set_inbound_transform(
+            ssl, ssl->handshake->transform_earlydata);
+        mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_END_OF_EARLY_DATA);
+        return 0;
+    }
+#endif /* MBEDTLS_SSL_EARLY_DATA */
+    MBEDTLS_SSL_DEBUG_MSG(
+        1, ("Switch to handshake keys for inbound traffic "
+            "( K_recv = handshake )"));
     mbedtls_ssl_set_inbound_transform(ssl, ssl->handshake->transform_handshake);
 
-    if (ssl->handshake->certificate_request_sent) {
-        mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_CLIENT_CERTIFICATE);
-    } else {
-        MBEDTLS_SSL_DEBUG_MSG(2, ("skip parse certificate"));
-        MBEDTLS_SSL_DEBUG_MSG(2, ("skip parse certificate verify"));
-        mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_CLIENT_FINISHED);
-    }
+    ssl_tls13_prepare_for_handshake_second_flight(ssl);
 
     return 0;
 }
 
+#if defined(MBEDTLS_SSL_EARLY_DATA)
+/*
+ * Handler for MBEDTLS_SSL_END_OF_EARLY_DATA
+ */
+#define SSL_GOT_END_OF_EARLY_DATA      0
+#define SSL_GOT_EARLY_DATA             1
+/* Coordination:
+ * Deals with the ambiguity of not knowing if the next message is an
+ * EndOfEarlyData message or an application message containing early data.
+ * Returns a negative code on failure, or
+ * - SSL_GOT_END_OF_EARLY_DATA
+ * - SSL_GOT_EARLY_DATA
+ * indicating which message is received.
+ */
+MBEDTLS_CHECK_RETURN_CRITICAL
+static int ssl_tls13_end_of_early_data_coordinate(mbedtls_ssl_context *ssl)
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    if ((ret = mbedtls_ssl_read_record(ssl, 0)) != 0) {
+        MBEDTLS_SSL_DEBUG_RET(1, "mbedtls_ssl_read_record", ret);
+        return ret;
+    }
+    ssl->keep_current_message = 1;
+
+    if (ssl->in_msgtype == MBEDTLS_SSL_MSG_HANDSHAKE        &&
+        ssl->in_msg[0]  == MBEDTLS_SSL_HS_END_OF_EARLY_DATA) {
+        MBEDTLS_SSL_DEBUG_MSG(3, ("Received an end_of_early_data message."));
+        return SSL_GOT_END_OF_EARLY_DATA;
+    }
+
+    if (ssl->in_msgtype == MBEDTLS_SSL_MSG_APPLICATION_DATA) {
+        MBEDTLS_SSL_DEBUG_MSG(3, ("Received early data"));
+        return SSL_GOT_EARLY_DATA;
+    }
+
+    MBEDTLS_SSL_PEND_FATAL_ALERT(MBEDTLS_SSL_ALERT_MSG_UNEXPECTED_MESSAGE,
+                                 MBEDTLS_ERR_SSL_UNEXPECTED_MESSAGE);
+    return MBEDTLS_ERR_SSL_UNEXPECTED_MESSAGE;
+}
+
+MBEDTLS_CHECK_RETURN_CRITICAL
+static int ssl_tls13_parse_end_of_early_data(mbedtls_ssl_context *ssl,
+                                             const unsigned char *buf,
+                                             const unsigned char *end)
+{
+    /* RFC 8446 section 4.5
+     *
+     * struct {} EndOfEarlyData;
+     */
+    if (buf != end) {
+        MBEDTLS_SSL_PEND_FATAL_ALERT(MBEDTLS_SSL_ALERT_MSG_DECODE_ERROR,
+                                     MBEDTLS_ERR_SSL_DECODE_ERROR);
+        return MBEDTLS_ERR_SSL_DECODE_ERROR;
+    }
+    return 0;
+}
+
+MBEDTLS_CHECK_RETURN_CRITICAL
+static int ssl_tls13_process_early_application_data(mbedtls_ssl_context *ssl)
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    if ((ret = mbedtls_ssl_read_record(ssl, 0)) != 0) {
+        MBEDTLS_SSL_DEBUG_RET(1, "mbedtls_ssl_read_record", ret);
+        return ret;
+    }
+
+    /*
+     * Output early data
+     *
+     * For the time being, we print received data via debug message.
+     *
+     * TODO: Remove it when `mbedtls_ssl_read_early_data` is ready.
+     */
+    ssl->in_msg[ssl->in_msglen] = 0;
+    MBEDTLS_SSL_DEBUG_MSG(3, ("\n%s", ssl->in_msg));
+
+    /* RFC 8446 section 4.6.1
+     *
+     * A server receiving more than max_early_data_size bytes of 0-RTT data
+     * SHOULD terminate the connection with an "unexpected_message" alert.
+     *
+     * TODO: Add received data size check here.
+     */
+
+    return 0;
+}
+
+/*
+ * RFC 8446 section A.2
+ *
+ *                                | Send ServerHello
+ *                                | K_send = handshake
+ *                                | Send EncryptedExtensions
+ *                                | [Send CertificateRequest]
+ * Can send                       | [Send Certificate + CertificateVerify]
+ * app data                       | Send Finished
+ * after   -->                    | K_send = application
+ * here                  +--------+--------+
+ *              No 0-RTT |                 | 0-RTT
+ *                       |                 |
+ *   K_recv = handshake  |                 | K_recv = early data
+ * [Skip decrypt errors] |    +------> WAIT_EOED -+
+ *                       |    |       Recv |      | Recv EndOfEarlyData
+ *                       |    | early data |      | K_recv = handshake
+ *                       |    +------------+      |
+ *                       |                        |
+ *                       +> WAIT_FLIGHT2 <--------+
+ *                                |
+ *                       +--------+--------+
+ *               No auth |                 | Client auth
+ *                       |                 |
+ *                       |                 v
+ *                       |             WAIT_CERT
+ *                       |        Recv |       | Recv Certificate
+ *                       |       empty |       v
+ *                       | Certificate |    WAIT_CV
+ *                       |             |       | Recv
+ *                       |             v       | CertificateVerify
+ *                       +-> WAIT_FINISHED <---+
+ *                                | Recv Finished
+ *
+ * The function handles actions and state changes from 0-RTT to WAIT_FLIGHT2 in
+ * the above diagram.
+ */
+MBEDTLS_CHECK_RETURN_CRITICAL
+static int ssl_tls13_process_end_of_early_data(mbedtls_ssl_context *ssl)
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    MBEDTLS_SSL_DEBUG_MSG(2, ("=> ssl_tls13_process_end_of_early_data"));
+
+    MBEDTLS_SSL_PROC_CHK_NEG(ssl_tls13_end_of_early_data_coordinate(ssl));
+
+    if (ret == SSL_GOT_END_OF_EARLY_DATA) {
+        unsigned char *buf;
+        size_t buf_len;
+
+        MBEDTLS_SSL_PROC_CHK(mbedtls_ssl_tls13_fetch_handshake_msg(
+                                 ssl, MBEDTLS_SSL_HS_END_OF_EARLY_DATA,
+                                 &buf, &buf_len));
+
+        MBEDTLS_SSL_PROC_CHK(ssl_tls13_parse_end_of_early_data(
+                                 ssl, buf, buf + buf_len));
+
+        MBEDTLS_SSL_DEBUG_MSG(
+            1, ("Switch to handshake keys for inbound traffic"
+                "( K_recv = handshake )"));
+        mbedtls_ssl_set_inbound_transform(
+            ssl, ssl->handshake->transform_handshake);
+
+        MBEDTLS_SSL_PROC_CHK(mbedtls_ssl_add_hs_msg_to_checksum(
+                                 ssl, MBEDTLS_SSL_HS_END_OF_EARLY_DATA,
+                                 buf, buf_len));
+
+        ssl_tls13_prepare_for_handshake_second_flight(ssl);
+
+    } else if (ret == SSL_GOT_EARLY_DATA) {
+        MBEDTLS_SSL_PROC_CHK(ssl_tls13_process_early_application_data(ssl));
+    } else {
+        MBEDTLS_SSL_DEBUG_MSG(1, ("should never happen"));
+        ret = MBEDTLS_ERR_SSL_INTERNAL_ERROR;
+        goto cleanup;
+    }
+
+cleanup:
+    MBEDTLS_SSL_DEBUG_MSG(2, ("<= ssl_tls13_process_end_of_early_data"));
+    return ret;
+}
+#endif /* MBEDTLS_SSL_EARLY_DATA */
+
 /*
  * Handler for MBEDTLS_SSL_CLIENT_FINISHED
  */
@@ -2888,6 +3137,15 @@
     mbedtls_ssl_session_set_ticket_flags(
         session, ssl->handshake->tls13_kex_modes);
 #endif
+
+#if defined(MBEDTLS_SSL_EARLY_DATA)
+    if (ssl->conf->early_data_enabled == MBEDTLS_SSL_EARLY_DATA_ENABLED &&
+        ssl->conf->max_early_data_size > 0) {
+        mbedtls_ssl_session_set_ticket_flags(
+            session, MBEDTLS_SSL_TLS1_3_TICKET_ALLOW_EARLY_DATA);
+    }
+#endif /* MBEDTLS_SSL_EARLY_DATA */
+
     MBEDTLS_SSL_PRINT_TICKET_FLAGS(4, session->ticket_flags);
 
     /* Generate ticket_age_add */
@@ -2971,12 +3229,13 @@
  * The following fields are placed inside the ticket by the
  * f_ticket_write() function:
  *
- *  - creation time (start)
- *  - flags (flags)
+ *  - creation time (ticket_creation_time)
+ *  - flags (ticket_flags)
  *  - age add (ticket_age_add)
- *  - key (key)
- *  - key length (key_len)
+ *  - key (resumption_key)
+ *  - key length (resumption_key_len)
  *  - ciphersuite (ciphersuite)
+ *  - max_early_data_size (max_early_data_size)
  */
 MBEDTLS_CHECK_RETURN_CRITICAL
 static int ssl_tls13_write_new_session_ticket_body(mbedtls_ssl_context *ssl,
@@ -2991,6 +3250,7 @@
     mbedtls_ssl_session *session = ssl->session;
     size_t ticket_len;
     uint32_t ticket_lifetime;
+    unsigned char *p_extensions_len;
 
     *out_len = 0;
     MBEDTLS_SSL_DEBUG_MSG(2, ("=> write NewSessionTicket msg"));
@@ -3052,15 +3312,35 @@
 
     /* Ticket Extensions
      *
-     * Note: We currently don't have any extensions.
-     * Set length to zero.
+     * Extension extensions<0..2^16-2>;
      */
     ssl->handshake->sent_extensions = MBEDTLS_SSL_EXT_MASK_NONE;
 
     MBEDTLS_SSL_CHK_BUF_PTR(p, end, 2);
-    MBEDTLS_PUT_UINT16_BE(0, p, 0);
+    p_extensions_len = p;
     p += 2;
 
+#if defined(MBEDTLS_SSL_EARLY_DATA)
+    if (mbedtls_ssl_session_ticket_allow_early_data(session)) {
+        size_t output_len;
+
+        if ((ret = mbedtls_ssl_tls13_write_early_data_ext(
+                 ssl, 1, p, end, &output_len)) != 0) {
+            MBEDTLS_SSL_DEBUG_RET(
+                1, "mbedtls_ssl_tls13_write_early_data_ext", ret);
+            return ret;
+        }
+        p += output_len;
+    } else {
+        MBEDTLS_SSL_DEBUG_MSG(
+            4, ("early_data not allowed, "
+                "skip early_data extension in NewSessionTicket"));
+    }
+
+#endif /* MBEDTLS_SSL_EARLY_DATA */
+
+    MBEDTLS_PUT_UINT16_BE(p - p_extensions_len - 2, p_extensions_len, 0);
+
     *out_len = p - buf;
     MBEDTLS_SSL_DEBUG_BUF(4, "ticket", buf, *out_len);
     MBEDTLS_SSL_DEBUG_MSG(2, ("<= write new session ticket"));
@@ -3207,6 +3487,12 @@
             ret = ssl_tls13_write_server_finished(ssl);
             break;
 
+#if defined(MBEDTLS_SSL_EARLY_DATA)
+        case MBEDTLS_SSL_END_OF_EARLY_DATA:
+            ret = ssl_tls13_process_end_of_early_data(ssl);
+            break;
+#endif /* MBEDTLS_SSL_EARLY_DATA */
+
         case MBEDTLS_SSL_CLIENT_FINISHED:
             ret = ssl_tls13_process_client_finished(ssl);
             break;
diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index f6a6bb6..1b3dedb 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -52,7 +52,7 @@
 #define DFL_KEY_OPAQUE          0
 #define DFL_KEY_PWD             ""
 #define DFL_PSK                 ""
-#define DFL_EARLY_DATA          MBEDTLS_SSL_EARLY_DATA_DISABLED
+#define DFL_EARLY_DATA          ""
 #define DFL_PSK_OPAQUE          0
 #define DFL_PSK_IDENTITY        "Client_identity"
 #define DFL_ECJPAKE_PW          NULL
@@ -347,8 +347,9 @@
 
 #if defined(MBEDTLS_SSL_EARLY_DATA)
 #define USAGE_EARLY_DATA \
-    "    early_data=%%d        default: 0 (disabled)\n"      \
-    "                        options: 0 (disabled), 1 (enabled)\n"
+    "    early_data=%%s      The file path to read early data from\n" \
+    "                        default: \"\" (do nothing)\n"            \
+    "                        option: a file path\n"
 #else
 #define USAGE_EARLY_DATA ""
 #endif /* MBEDTLS_SSL_EARLY_DATA && MBEDTLS_SSL_PROTO_TLS1_3 */
@@ -543,7 +544,7 @@
     int reproducible;           /* make communication reproducible          */
     int skip_close_notify;      /* skip sending the close_notify alert      */
 #if defined(MBEDTLS_SSL_EARLY_DATA)
-    int early_data;             /* support for early data                   */
+    const char *early_data;     /* the path of the file to read early data from */
 #endif
     int query_config_mode;      /* whether to read config                   */
     int use_srtp;               /* Support SRTP                             */
@@ -741,6 +742,10 @@
     size_t cid_renego_len = 0;
 #endif
 
+#if defined(MBEDTLS_SSL_EARLY_DATA)
+    FILE *early_data_fp = NULL;
+#endif /* MBEDTLS_SSL_EARLY_DATA */
+
 #if defined(MBEDTLS_SSL_ALPN)
     const char *alpn_list[ALPN_LIST_SIZE];
 #endif
@@ -1196,15 +1201,7 @@
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
 #if defined(MBEDTLS_SSL_EARLY_DATA)
         else if (strcmp(p, "early_data") == 0) {
-            switch (atoi(q)) {
-                case 0:
-                    opt.early_data = MBEDTLS_SSL_EARLY_DATA_DISABLED;
-                    break;
-                case 1:
-                    opt.early_data = MBEDTLS_SSL_EARLY_DATA_ENABLED;
-                    break;
-                default: goto usage;
-            }
+            opt.early_data = q;
         }
 #endif /* MBEDTLS_SSL_EARLY_DATA */
 
@@ -1971,7 +1968,16 @@
     }
 
 #if defined(MBEDTLS_SSL_EARLY_DATA)
-    mbedtls_ssl_conf_early_data(&conf, opt.early_data);
+    int early_data_enabled = MBEDTLS_SSL_EARLY_DATA_DISABLED;
+    if (strlen(opt.early_data) > 0) {
+        if ((early_data_fp = fopen(opt.early_data, "rb")) == NULL) {
+            mbedtls_printf("failed\n  ! Cannot open '%s' for reading.\n",
+                           opt.early_data);
+            goto exit;
+        }
+        early_data_enabled = MBEDTLS_SSL_EARLY_DATA_ENABLED;
+    }
+    mbedtls_ssl_conf_early_data(&conf, early_data_enabled);
 #endif /* MBEDTLS_SSL_EARLY_DATA */
 
     if ((ret = mbedtls_ssl_setup(&ssl, &conf)) != 0) {
@@ -3029,6 +3035,12 @@
     mbedtls_ssl_config_free(&conf);
     mbedtls_ssl_session_free(&saved_session);
 
+#if defined(MBEDTLS_SSL_EARLY_DATA)
+    if (early_data_fp != NULL) {
+        fclose(early_data_fp);
+    }
+#endif
+
     if (session_data != NULL) {
         mbedtls_platform_zeroize(session_data, session_data_len);
     }
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index c96128b..e6ebd8e 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -556,6 +556,7 @@
     USAGE_GROUPS                                            \
     USAGE_SIG_ALGS                                          \
     USAGE_KEY_OPAQUE_ALGS                                   \
+    USAGE_EARLY_DATA                                        \
     "\n"
 
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
diff --git a/programs/test/benchmark.c b/programs/test/benchmark.c
index 3d751d0..85b7fc6 100644
--- a/programs/test/benchmark.c
+++ b/programs/test/benchmark.c
@@ -1057,16 +1057,16 @@
         for (i = 0; (size_t) i < sizeof(dhm_sizes) / sizeof(dhm_sizes[0]); i++) {
             mbedtls_dhm_init(&dhm);
 
-            if (mbedtls_mpi_read_binary(&dhm.P, dhm_P[i],
+            if (mbedtls_mpi_read_binary(&dhm.MBEDTLS_PRIVATE(P), dhm_P[i],
                                         dhm_P_size[i]) != 0 ||
-                mbedtls_mpi_read_binary(&dhm.G, dhm_G[i],
+                mbedtls_mpi_read_binary(&dhm.MBEDTLS_PRIVATE(G), dhm_G[i],
                                         dhm_G_size[i]) != 0) {
                 mbedtls_exit(1);
             }
 
-            n = mbedtls_mpi_size(&dhm.P);
+            n = mbedtls_mpi_size(&dhm.MBEDTLS_PRIVATE(P));
             mbedtls_dhm_make_public(&dhm, (int) n, buf, n, myrand, NULL);
-            if (mbedtls_mpi_copy(&dhm.GY, &dhm.GX) != 0) {
+            if (mbedtls_mpi_copy(&dhm.MBEDTLS_PRIVATE(GY), &dhm.MBEDTLS_PRIVATE(GX)) != 0) {
                 mbedtls_exit(1);
             }
 
diff --git a/programs/test/selftest.c b/programs/test/selftest.c
index 61dde5e..e132e4c 100644
--- a/programs/test/selftest.c
+++ b/programs/test/selftest.c
@@ -79,6 +79,7 @@
         if (verbose) {
             mbedtls_printf("  CALLOC(0,1): passed (same non-null)\n");
         }
+        empty2 = NULL;
     } else {
         if (verbose) {
             mbedtls_printf("  CALLOC(0,1): passed (distinct non-null)\n");
@@ -103,6 +104,7 @@
         if (verbose) {
             mbedtls_printf("  CALLOC(1,0): passed (same non-null)\n");
         }
+        empty2 = NULL;
     } else {
         if (verbose) {
             mbedtls_printf("  CALLOC(1,0): passed (distinct non-null)\n");
@@ -119,6 +121,7 @@
             mbedtls_printf("  CALLOC(1): failed (same buffer twice)\n");
         }
         ++failures;
+        buffer2 = NULL;
     } else {
         if (verbose) {
             mbedtls_printf("  CALLOC(1): passed\n");
diff --git a/scripts/assemble_changelog.py b/scripts/assemble_changelog.py
index d5f705c..07e6fc5 100755
--- a/scripts/assemble_changelog.py
+++ b/scripts/assemble_changelog.py
@@ -115,7 +115,7 @@
 class TextChangelogFormat(ChangelogFormat):
     """The traditional Mbed TLS changelog format."""
 
-    _unreleased_version_text = '= Mbed TLS x.x.x branch released xxxx-xx-xx'
+    _unreleased_version_text = '= {} x.x.x branch released xxxx-xx-xx'
     @classmethod
     def is_released_version(cls, title):
         # Look for an incomplete release date
@@ -123,6 +123,7 @@
 
     _top_version_re = re.compile(r'(?:\A|\n)(=[^\n]*\n+)(.*?\n)(?:=|$)',
                                  re.DOTALL)
+    _name_re = re.compile(r'=\s(.*)\s[0-9x]+\.', re.DOTALL)
     @classmethod
     def extract_top_version(cls, changelog_file_content):
         """A version section starts with a line starting with '='."""
@@ -131,9 +132,10 @@
         top_version_end = m.end(2)
         top_version_title = m.group(1)
         top_version_body = m.group(2)
+        name = re.match(cls._name_re, top_version_title).group(1)
         if cls.is_released_version(top_version_title):
             top_version_end = top_version_start
-            top_version_title = cls._unreleased_version_text + '\n\n'
+            top_version_title = cls._unreleased_version_text.format(name) + '\n\n'
             top_version_body = ''
         return (changelog_file_content[:top_version_start],
                 top_version_title, top_version_body,
@@ -245,6 +247,7 @@
         for category in STANDARD_CATEGORIES:
             self.categories[category] = ''
         offset = (self.header + self.top_version_title).count('\n') + 1
+
         self.add_categories_from_text(input_stream.name, offset,
                                       top_version_body, True)
 
diff --git a/scripts/generate_driver_wrappers.py b/scripts/generate_driver_wrappers.py
index 2fdc4cd..624ab81 100755
--- a/scripts/generate_driver_wrappers.py
+++ b/scripts/generate_driver_wrappers.py
@@ -108,17 +108,17 @@
         return json_data
 
 
-def load_schemas(mbedtls_root: str) -> Dict[str, Any]:
+def load_schemas(project_root: str) -> Dict[str, Any]:
     """
     Load schemas map
     """
     schema_file_paths = {
-        'transparent': os.path.join(mbedtls_root,
+        'transparent': os.path.join(project_root,
                                     'scripts',
                                     'data_files',
                                     'driver_jsons',
                                     'driver_transparent_schema.json'),
-        'opaque': os.path.join(mbedtls_root,
+        'opaque': os.path.join(project_root,
                                'scripts',
                                'data_files',
                                'driver_jsons',
@@ -131,13 +131,13 @@
     return driver_schema
 
 
-def read_driver_descriptions(mbedtls_root: str,
+def read_driver_descriptions(project_root: str,
                              json_directory: str,
                              jsondriver_list: str) -> list:
     """
     Merge driver JSON files into a single ordered JSON after validation.
     """
-    driver_schema = load_schemas(mbedtls_root)
+    driver_schema = load_schemas(project_root)
 
     with open(file=os.path.join(json_directory, jsondriver_list),
               mode='r',
@@ -163,11 +163,11 @@
     """
     Main with command line arguments.
     """
-    def_arg_mbedtls_root = build_tree.guess_mbedtls_root()
+    def_arg_project_root = build_tree.guess_project_root()
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('--mbedtls-root', default=def_arg_mbedtls_root,
-                        help='root directory of mbedtls source code')
+    parser.add_argument('--project-root', default=def_arg_project_root,
+                        help='root directory of repo source code')
     parser.add_argument('--template-dir',
                         help='directory holding the driver templates')
     parser.add_argument('--json-dir',
@@ -176,24 +176,27 @@
                         help='output file\'s location')
     args = parser.parse_args()
 
-    mbedtls_root = os.path.abspath(args.mbedtls_root)
+    project_root = os.path.abspath(args.project_root)
+
+    crypto_core_directory = build_tree.crypto_core_directory(project_root)
 
     output_directory = args.output_directory if args.output_directory is not None else \
-        os.path.join(mbedtls_root, 'library')
+        crypto_core_directory
+
     template_directory = args.template_dir if args.template_dir is not None else \
-        os.path.join(mbedtls_root,
+        os.path.join(project_root,
                      'scripts',
                      'data_files',
                      'driver_templates')
     json_directory = args.json_dir if args.json_dir is not None else \
-        os.path.join(mbedtls_root,
+        os.path.join(project_root,
                      'scripts',
                      'data_files',
                      'driver_jsons')
 
     try:
         # Read and validate list of driver jsons from driverlist.json
-        merged_driver_json = read_driver_descriptions(mbedtls_root,
+        merged_driver_json = read_driver_descriptions(project_root,
                                                       json_directory,
                                                       'driverlist.json')
     except DriverReaderException as e:
diff --git a/scripts/lcov.sh b/scripts/lcov.sh
index 9258ba7..0584a0a 100755
--- a/scripts/lcov.sh
+++ b/scripts/lcov.sh
@@ -32,8 +32,8 @@
 
 # Repository detection
 in_mbedtls_build_dir () {
-     test -d library
- }
+    test -d library
+}
 
 # Collect stats and build a HTML report.
 lcov_library_report () {
diff --git a/scripts/mbedtls_dev/build_tree.py b/scripts/mbedtls_dev/build_tree.py
index a657a51..ec67e4c 100644
--- a/scripts/mbedtls_dev/build_tree.py
+++ b/scripts/mbedtls_dev/build_tree.py
@@ -7,6 +7,7 @@
 
 import os
 import inspect
+from typing import Optional
 
 def looks_like_tf_psa_crypto_root(path: str) -> bool:
     """Whether the given directory looks like the root of the PSA Crypto source tree."""
@@ -21,9 +22,40 @@
 def looks_like_root(path: str) -> bool:
     return looks_like_tf_psa_crypto_root(path) or looks_like_mbedtls_root(path)
 
-def check_repo_path():
+def crypto_core_directory(root: Optional[str] = None, relative: Optional[bool] = False) -> str:
     """
-    Check that the current working directory is the project root, and throw
+    Return the path of the directory containing the PSA crypto core
+    for either TF-PSA-Crypto or Mbed TLS.
+
+    Returns either the full path or relative path depending on the
+    "relative" boolean argument.
+    """
+    if root is None:
+        root = guess_project_root()
+    if looks_like_tf_psa_crypto_root(root):
+        if relative:
+            return "core"
+        return os.path.join(root, "core")
+    elif looks_like_mbedtls_root(root):
+        if relative:
+            return "library"
+        return os.path.join(root, "library")
+    else:
+        raise Exception('Neither Mbed TLS nor TF-PSA-Crypto source tree found')
+
+def crypto_library_filename(root: Optional[str] = None) -> str:
+    """Return the crypto library filename for either TF-PSA-Crypto or Mbed TLS."""
+    if root is None:
+        root = guess_project_root()
+    if looks_like_tf_psa_crypto_root(root):
+        return "tfpsacrypto"
+    elif looks_like_mbedtls_root(root):
+        return "mbedcrypto"
+    else:
+        raise Exception('Neither Mbed TLS nor TF-PSA-Crypto source tree found')
+
+def check_repo_path():
+    """Check that the current working directory is the project root, and throw
     an exception if not.
     """
     if not all(os.path.isdir(d) for d in ["include", "library", "tests"]):
@@ -43,11 +75,10 @@
             return
     raise Exception('Mbed TLS source tree not found')
 
+def guess_project_root():
+    """Guess project source code directory.
 
-def guess_mbedtls_root():
-    """Guess mbedTLS source code directory.
-
-    Return the first possible mbedTLS root directory
+    Return the first possible project root directory.
     """
     dirs = set({})
     for frame in inspect.stack():
@@ -60,4 +91,30 @@
             dirs.add(d)
             if looks_like_root(d):
                 return d
-    raise Exception('Mbed TLS source tree not found')
+    raise Exception('Neither Mbed TLS nor TF-PSA-Crypto source tree found')
+
+def guess_mbedtls_root(root: Optional[str] = None) -> str:
+    """Guess Mbed TLS source code directory.
+
+    Return the first possible Mbed TLS root directory.
+    Raise an exception if we are not in Mbed TLS.
+    """
+    if root is None:
+        root = guess_project_root()
+    if looks_like_mbedtls_root(root):
+        return root
+    else:
+        raise Exception('Mbed TLS source tree not found')
+
+def guess_tf_psa_crypto_root(root: Optional[str] = None) -> str:
+    """Guess TF-PSA-Crypto source code directory.
+
+    Return the first possible TF-PSA-Crypto root directory.
+    Raise an exception if we are not in TF-PSA-Crypto.
+    """
+    if root is None:
+        root = guess_project_root()
+    if looks_like_tf_psa_crypto_root(root):
+        return root
+    else:
+        raise Exception('TF-PSA-Crypto source tree not found')
diff --git a/tests/include/test/drivers/cipher.h b/tests/include/test/drivers/cipher.h
index 950a174..2fe47e4 100644
--- a/tests/include/test/drivers/cipher.h
+++ b/tests/include/test/drivers/cipher.h
@@ -23,11 +23,17 @@
     /* If not PSA_SUCCESS, return this error code instead of processing the
      * function call. */
     psa_status_t forced_status;
+    psa_status_t forced_status_encrypt;
+    psa_status_t forced_status_set_iv;
     /* Count the amount of times one of the cipher driver functions is called. */
     unsigned long hits;
+    unsigned long hits_encrypt;
+    unsigned long hits_set_iv;
 } mbedtls_test_driver_cipher_hooks_t;
 
-#define MBEDTLS_TEST_DRIVER_CIPHER_INIT { NULL, 0, PSA_SUCCESS, 0 }
+#define MBEDTLS_TEST_DRIVER_CIPHER_INIT { NULL, 0, \
+                                          PSA_SUCCESS, PSA_SUCCESS, PSA_SUCCESS, \
+                                          0, 0, 0 }
 static inline mbedtls_test_driver_cipher_hooks_t
 mbedtls_test_driver_cipher_hooks_init(void)
 {
diff --git a/tests/include/test/drivers/key_management.h b/tests/include/test/drivers/key_management.h
index 9e2c898..526adbb 100644
--- a/tests/include/test/drivers/key_management.h
+++ b/tests/include/test/drivers/key_management.h
@@ -26,6 +26,8 @@
     /* Count the amount of times one of the key management driver functions
      * is called. */
     unsigned long hits;
+    /* Subset of hits which only counts key operations with EC key */
+    unsigned long hits_export_public_key;
     /* Location of the last key management driver called to import a key. */
     psa_key_location_t location;
 } mbedtls_test_driver_key_management_hooks_t;
@@ -34,7 +36,7 @@
  * sense that no PSA specification will assign a meaning to this location
  * (stated first in version 1.0.1 of the specification) and that it is not
  * used as a location of an opaque test drivers. */
-#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0, 0x800000 }
+#define MBEDTLS_TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0, 0, 0x800000 }
 static inline mbedtls_test_driver_key_management_hooks_t
 mbedtls_test_driver_key_management_hooks_init(void)
 {
diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 04b90b9..0b8c221 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -392,4 +392,27 @@
 #define MD_OR_USE_PSA_DONE() ((void) 0)
 #endif
 
+/** \def AES_PSA_INIT
+ *
+ * Call this macro to initialize the PSA subsystem if AES_C is not defined,
+ * so that CTR_DRBG uses PSA implementation to get AES-ECB.
+ *
+ * If the initialization fails, mark the test case as failed and jump to the
+ * \p exit label.
+ */
+/** \def AES_PSA_DONE
+ *
+ * Call this macro at the end of a test case if you called #AES_PSA_INIT.
+ *
+ * This is like #PSA_DONE except it does nothing under the same conditions as
+ * #AES_PSA_INIT.
+ */
+#if defined(MBEDTLS_AES_C)
+#define AES_PSA_INIT() ((void) 0)
+#define AES_PSA_DONE() ((void) 0)
+#else /* MBEDTLS_AES_C */
+#define AES_PSA_INIT()   PSA_INIT()
+#define AES_PSA_DONE()   PSA_DONE()
+#endif /* MBEDTLS_AES_C */
+
 #endif /* PSA_CRYPTO_HELPERS_H */
diff --git a/tests/opt-testcases/tls13-misc.sh b/tests/opt-testcases/tls13-misc.sh
index 9208384..c1682e3 100755
--- a/tests/opt-testcases/tls13-misc.sh
+++ b/tests/opt-testcases/tls13-misc.sh
@@ -261,9 +261,11 @@
 requires_any_configs_enabled MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_EPHEMERAL_ENABLED \
                              MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_ENABLED
 run_test    "TLS 1.3 m->G: EarlyData: basic check, good" \
-            "$G_NEXT_SRV -d 10 --priority=NORMAL:-VERS-ALL:+VERS-TLS1.3:+CIPHER-ALL:+ECDHE-PSK:+PSK --earlydata --disable-client-cert" \
-            "$P_CLI debug_level=4 early_data=1 reco_mode=1 reconnect=1 reco_delay=900" \
+            "$G_NEXT_SRV -d 10 --priority=NORMAL:-VERS-ALL:+VERS-TLS1.3:+CIPHER-ALL:+ECDHE-PSK:+PSK \
+                         --earlydata --maxearlydata 16384 --disable-client-cert" \
+            "$P_CLI debug_level=4 early_data=$EARLY_DATA_INPUT reco_mode=1 reconnect=1 reco_delay=900" \
             0 \
+            -c "received max_early_data_size: 16384" \
             -c "Reconnecting with saved session" \
             -c "NewSessionTicket: early_data(42) extension received." \
             -c "ClientHello: early_data(42) extension exists." \
@@ -285,7 +287,7 @@
                              MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_ENABLED
 run_test    "TLS 1.3 m->G: EarlyData: no early_data in NewSessionTicket, good" \
             "$G_NEXT_SRV -d 10 --priority=NORMAL:-VERS-ALL:+VERS-TLS1.3:+CIPHER-ALL:+ECDHE-PSK:+PSK --disable-client-cert" \
-            "$P_CLI debug_level=4 early_data=1 reco_mode=1 reconnect=1" \
+            "$P_CLI debug_level=4 early_data=$EARLY_DATA_INPUT reco_mode=1 reconnect=1" \
             0 \
             -c "Reconnecting with saved session" \
             -C "NewSessionTicket: early_data(42) extension received." \
@@ -489,34 +491,19 @@
 EARLY_DATA_INPUT_LEN=$(( $EARLY_DATA_INPUT_LEN_BLOCKS * 32 ))
 
 requires_gnutls_next
-requires_all_configs_enabled MBEDTLS_SSL_EARLY_DATA MBEDTLS_SSL_SESSION_TICKETS     \
-                             MBEDTLS_SSL_SRV_C MBEDTLS_DEBUG_C MBEDTLS_HAVE_TIME    \
-                             MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_EPHEMERAL_ENABLED \
-                             MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE
-requires_any_configs_enabled MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_ENABLED \
-                             MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_EPHEMERAL_ENABLED
-run_test "TLS 1.3 G->m: EarlyData: feature is disabled, fail." \
-         "$P_SRV force_version=tls13 debug_level=4 max_early_data_size=-1" \
-         "$G_NEXT_CLI localhost --priority=NORMAL:-VERS-ALL:+VERS-TLS1.3:+GROUP-ALL -d 10 -r --earlydata $EARLY_DATA_INPUT" \
-         1 \
-         -s "ClientHello: early_data(42) extension exists."                     \
-         -s "EncryptedExtensions: early_data(42) extension does not exist."    \
-         -s "NewSessionTicket: early_data(42) extension does not exist."        \
-         -s "Last error was: -29056 - SSL - Verification of the message MAC failed"
-
-requires_gnutls_next
 requires_all_configs_enabled MBEDTLS_SSL_EARLY_DATA MBEDTLS_SSL_SESSION_TICKETS \
                              MBEDTLS_SSL_SRV_C MBEDTLS_DEBUG_C MBEDTLS_HAVE_TIME \
                              MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_EPHEMERAL_ENABLED \
                              MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE
 requires_any_configs_enabled MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_ENABLED \
                              MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_EPHEMERAL_ENABLED
-run_test "TLS 1.3 G->m: EarlyData: feature is enabled, fail." \
+run_test "TLS 1.3 G->m: EarlyData: feature is enabled, good." \
          "$P_SRV force_version=tls13 debug_level=4 max_early_data_size=$EARLY_DATA_INPUT_LEN" \
          "$G_NEXT_CLI localhost --priority=NORMAL:-VERS-ALL:+VERS-TLS1.3:+GROUP-ALL:+KX-ALL \
                       -d 10 -r --earlydata $EARLY_DATA_INPUT " \
-         1 \
+         0 \
+         -s "NewSessionTicket: early_data(42) extension exists."            \
+         -s "Sent max_early_data_size=$EARLY_DATA_INPUT_LEN"                \
          -s "ClientHello: early_data(42) extension exists."                 \
          -s "EncryptedExtensions: early_data(42) extension exists."         \
-         -s "NewSessionTicket: early_data(42) extension does not exist."    \
-         -s "Last error was: -29056 - SSL - Verification of the message MAC failed"
+         -s "$( tail -1 $EARLY_DATA_INPUT )"
diff --git a/tests/scripts/all.sh b/tests/scripts/all.sh
index 4281557..10a76ec 100755
--- a/tests/scripts/all.sh
+++ b/tests/scripts/all.sh
@@ -2100,7 +2100,7 @@
     make test
 }
 
-component_test_valgrind_constant_flow () {
+component_release_test_valgrind_constant_flow () {
     # This tests both (1) everything that valgrind's memcheck usually checks
     # (heap buffer overflows, use of uninitialized memory, use-after-free,
     # etc.) and (2) branches or memory access depending on secret values,
@@ -2134,7 +2134,7 @@
     make memcheck
 }
 
-component_test_valgrind_constant_flow_psa () {
+component_release_test_valgrind_constant_flow_psa () {
     # This tests both (1) everything that valgrind's memcheck usually checks
     # (heap buffer overflows, use of uninitialized memory, use-after-free,
     # etc.) and (2) branches or memory access depending on secret values,
@@ -3265,6 +3265,10 @@
     # Check that p256m was built
     grep -q p256_ecdsa_ library/libmbedcrypto.a
 
+    # In "config-tfm.h" we disabled CIPHER_C tweaking TF-M's configuration
+    # files, so we want to ensure that it has not be re-enabled accidentally.
+    not grep mbedtls_cipher library/cipher.o
+
     # Run the tests
     msg "test: TF-M config + p256m driver + accel ECDH(E)/ECDSA"
     make test
@@ -3286,6 +3290,10 @@
     # Check that p256m was not built
     not grep p256_ecdsa_ library/libmbedcrypto.a
 
+    # In "config-tfm.h" we disabled CIPHER_C tweaking TF-M's configuration
+    # files, so we want to ensure that it has not be re-enabled accidentally.
+    not grep mbedtls_cipher library/cipher.o
+
     msg "test: TF-M config"
     make test
 }
@@ -3727,7 +3735,6 @@
     scripts/config.py unset MBEDTLS_PKCS5_C
     scripts/config.py unset MBEDTLS_PKCS12_C
 
-    scripts/config.py unset MBEDTLS_CTR_DRBG_C
     scripts/config.py unset MBEDTLS_NIST_KW_C
 }
 
@@ -5765,7 +5772,7 @@
     fi
 }
 
-component_test_valgrind () {
+component_release_test_valgrind () {
     msg "build: Release (clang)"
     # default config, in particular without MBEDTLS_USE_PSA_CRYPTO
     CC=clang cmake -D CMAKE_BUILD_TYPE:String=Release .
@@ -5793,7 +5800,7 @@
     fi
 }
 
-component_test_valgrind_psa () {
+component_release_test_valgrind_psa () {
     msg "build: Release, full (clang)"
     # full config, in particular with MBEDTLS_USE_PSA_CRYPTO
     scripts/config.py full
diff --git a/tests/scripts/check-generated-files.sh b/tests/scripts/check-generated-files.sh
index 67dedeb..3fe4e8c 100755
--- a/tests/scripts/check-generated-files.sh
+++ b/tests/scripts/check-generated-files.sh
@@ -22,8 +22,20 @@
     exit
 fi
 
-if [ -d library -a -d include -a -d tests ]; then :; else
-    echo "Must be run from Mbed TLS root" >&2
+in_mbedtls_repo () {
+    test -d include -a -d library -a -d programs -a -d tests
+}
+
+in_tf_psa_crypto_repo () {
+    test -d include -a -d core -a -d drivers -a -d programs -a -d tests
+}
+
+if in_mbedtls_repo; then
+    library_dir='library'
+elif in_tf_psa_crypto_repo; then
+    library_dir='core'
+else
+    echo "Must be run from Mbed TLS root or TF-PSA-Crypto root" >&2
     exit 1
 fi
 
@@ -114,16 +126,21 @@
 #   - **/CMakeLists.txt (to (re)build them with cmake)
 #   - scripts/make_generated_files.bat (to generate them under Windows)
 
-check scripts/generate_errors.pl library/error.c
-check scripts/generate_query_config.pl programs/test/query_config.c
-check scripts/generate_driver_wrappers.py library/psa_crypto_driver_wrappers.h library/psa_crypto_driver_wrappers_no_static.c
-check scripts/generate_features.pl library/version_features.c
-check scripts/generate_ssl_debug_helpers.py library/ssl_debug_helpers_generated.c
-# generate_visualc_files enumerates source files (library/*.c). It doesn't
-# care about their content, but the files must exist. So it must run after
-# the step that creates or updates these files.
-check scripts/generate_visualc_files.pl visualc/VS2013
+# These checks are common to Mbed TLS and TF-PSA-Crypto
 check scripts/generate_psa_constants.py programs/psa/psa_constant_names_generated.c
 check tests/scripts/generate_bignum_tests.py $(tests/scripts/generate_bignum_tests.py --list)
 check tests/scripts/generate_ecp_tests.py $(tests/scripts/generate_ecp_tests.py --list)
 check tests/scripts/generate_psa_tests.py $(tests/scripts/generate_psa_tests.py --list)
+check scripts/generate_driver_wrappers.py $library_dir/psa_crypto_driver_wrappers.h $library_dir/psa_crypto_driver_wrappers_no_static.c
+
+# Additional checks for Mbed TLS only
+if in_mbedtls_repo; then
+    check scripts/generate_errors.pl library/error.c
+    check scripts/generate_query_config.pl programs/test/query_config.c
+    check scripts/generate_features.pl library/version_features.c
+    check scripts/generate_ssl_debug_helpers.py library/ssl_debug_helpers_generated.c
+    # generate_visualc_files enumerates source files (library/*.c). It doesn't
+    # care about their content, but the files must exist. So it must run after
+    # the step that creates or updates these files.
+    check scripts/generate_visualc_files.pl visualc/VS2013
+fi
diff --git a/tests/scripts/test_psa_compliance.py b/tests/scripts/test_psa_compliance.py
index bed6d84..0d56ddf 100755
--- a/tests/scripts/test_psa_compliance.py
+++ b/tests/scripts/test_psa_compliance.py
@@ -50,12 +50,8 @@
 
     in_tf_psa_crypto_repo = build_tree.looks_like_tf_psa_crypto_root(root_dir)
 
-    if in_tf_psa_crypto_repo:
-        crypto_name = 'tfpsacrypto'
-        library_subdir = 'core'
-    else:
-        crypto_name = 'mbedcrypto'
-        library_subdir = 'library'
+    crypto_name = build_tree.crypto_library_filename(root_dir)
+    library_subdir = build_tree.crypto_core_directory(root_dir, relative=True)
 
     crypto_lib_filename = (library_build_dir + '/' +
                            library_subdir + '/' +
diff --git a/tests/src/drivers/test_driver_cipher.c b/tests/src/drivers/test_driver_cipher.c
index 678d8d5..2bc751a 100644
--- a/tests/src/drivers/test_driver_cipher.c
+++ b/tests/src/drivers/test_driver_cipher.c
@@ -41,6 +41,7 @@
     size_t *output_length)
 {
     mbedtls_test_driver_cipher_hooks.hits++;
+    mbedtls_test_driver_cipher_hooks.hits_encrypt++;
 
     if (mbedtls_test_driver_cipher_hooks.forced_output != NULL) {
         if (output_size < mbedtls_test_driver_cipher_hooks.forced_output_length) {
@@ -58,6 +59,9 @@
     if (mbedtls_test_driver_cipher_hooks.forced_status != PSA_SUCCESS) {
         return mbedtls_test_driver_cipher_hooks.forced_status;
     }
+    if (mbedtls_test_driver_cipher_hooks.forced_status_encrypt != PSA_SUCCESS) {
+        return mbedtls_test_driver_cipher_hooks.forced_status_encrypt;
+    }
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
     defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_CIPHER)
@@ -208,10 +212,14 @@
     size_t iv_length)
 {
     mbedtls_test_driver_cipher_hooks.hits++;
+    mbedtls_test_driver_cipher_hooks.hits_set_iv++;
 
     if (mbedtls_test_driver_cipher_hooks.forced_status != PSA_SUCCESS) {
         return mbedtls_test_driver_cipher_hooks.forced_status;
     }
+    if (mbedtls_test_driver_cipher_hooks.forced_status_set_iv != PSA_SUCCESS) {
+        return mbedtls_test_driver_cipher_hooks.forced_status_set_iv;
+    }
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
     defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_CIPHER)
diff --git a/tests/src/drivers/test_driver_key_management.c b/tests/src/drivers/test_driver_key_management.c
index 6442f22..d522ebf 100644
--- a/tests/src/drivers/test_driver_key_management.c
+++ b/tests/src/drivers/test_driver_key_management.c
@@ -529,6 +529,7 @@
     uint8_t *data, size_t data_size, size_t *data_length)
 {
     ++mbedtls_test_driver_key_management_hooks.hits;
+    ++mbedtls_test_driver_key_management_hooks.hits_export_public_key;
 
     if (mbedtls_test_driver_key_management_hooks.forced_status != PSA_SUCCESS) {
         return mbedtls_test_driver_key_management_hooks.forced_status;
diff --git a/tests/src/psa_crypto_helpers.c b/tests/src/psa_crypto_helpers.c
index d59a8f8..e1ea2b5 100644
--- a/tests/src/psa_crypto_helpers.c
+++ b/tests/src/psa_crypto_helpers.c
@@ -70,9 +70,20 @@
 
     mbedtls_psa_get_stats(&stats);
 
+#if defined(MBEDTLS_CTR_DRBG_C) && !defined(MBEDTLS_AES_C) && \
+    !defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG)
+    /* When AES_C is not defined and PSA does not have an external RNG,
+     * then CTR_DRBG uses PSA to perform AES-ECB. In this scenario 1 key
+     * slot is used internally from PSA to hold the AES key and it should
+     * not be taken into account when evaluating remaining open slots. */
+    if (stats.volatile_slots > 1) {
+        return "A volatile slot has not been closed properly.";
+    }
+#else
     if (stats.volatile_slots != 0) {
         return "A volatile slot has not been closed properly.";
     }
+#endif
     if (stats.persistent_slots != 0) {
         return "A persistent slot has not been closed properly.";
     }
diff --git a/tests/suites/test_suite_block_cipher.data b/tests/suites/test_suite_block_cipher.data
index cf321ae..097b567 100644
--- a/tests/suites/test_suite_block_cipher.data
+++ b/tests/suites/test_suite_block_cipher.data
@@ -182,56 +182,74 @@
 test_vec:MBEDTLS_CIPHER_ID_ARIA:"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f":"00112233445566778899aabbccddeeff":"f92bd7c79fb72e2f2b8f80c1972d24fc"
 
 Camellia-128-ECB Encrypt RFC3713 #1
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"0123456789abcdeffedcba9876543210":"0123456789abcdeffedcba9876543210":"67673138549669730857065648eabe43"
 
 Camellia-192-ECB Encrypt RFC3713 #1
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"0123456789abcdeffedcba98765432100011223344556677":"0123456789abcdeffedcba9876543210":"b4993401b3e996f84ee5cee7d79b09b9"
 
 Camellia-256-ECB Encrypt RFC3713 #1
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"0123456789abcdeffedcba987654321000112233445566778899aabbccddeeff":"0123456789abcdeffedcba9876543210":"9acc237dff16d76c20ef7c919e3a7509"
 
 Camellia-128-ECB Encrypt Perl EVP #1 [#1]
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"000102030405060708090A0B0C0D0E0F":"00112233445566778899AABBCCDDEEFF":"77CF412067AF8270613529149919546F"
 
 Camellia-192-ECB Encrypt Perl EVP #1 [#1]
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"000102030405060708090A0B0C0D0E0F1011121314151617":"00112233445566778899AABBCCDDEEFF":"B22F3C36B72D31329EEE8ADDC2906C68"
 
 Camellia-256-ECB Encrypt Perl EVP #1 [#1]
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F":"00112233445566778899AABBCCDDEEFF":"2EDF1F3418D53B88841FC8985FB1ECF2"
 
 Camellia-128-ECB Encrypt Perl EVP #1 [#2]
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"2B7E151628AED2A6ABF7158809CF4F3C":"6BC1BEE22E409F96E93D7E117393172A":"432FC5DCD628115B7C388D770B270C96"
 
 Camellia-128-ECB Encrypt Perl EVP #2
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"2B7E151628AED2A6ABF7158809CF4F3C":"AE2D8A571E03AC9C9EB76FAC45AF8E51":"0BE1F14023782A22E8384C5ABB7FAB2B"
 
 Camellia-128-ECB Encrypt Perl EVP #3
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"2B7E151628AED2A6ABF7158809CF4F3C":"30C81C46A35CE411E5FBC1191A0A52EF":"A0A1ABCD1893AB6FE0FE5B65DF5F8636"
 
 Camellia-128-ECB Encrypt Perl EVP #4
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"2B7E151628AED2A6ABF7158809CF4F3C":"F69F2445DF4F9B17AD2B417BE66C3710":"E61925E0D5DFAA9BB29F815B3076E51A"
 
 Camellia-192-ECB Encrypt Perl EVP #1 [#2]
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"8E73B0F7DA0E6452C810F32B809079E562F8EAD2522C6B7B":"6BC1BEE22E409F96E93D7E117393172A":"CCCC6C4E138B45848514D48D0D3439D3"
 
 Camellia-192-ECB Encrypt Perl EVP #2
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"8E73B0F7DA0E6452C810F32B809079E562F8EAD2522C6B7B":"AE2D8A571E03AC9C9EB76FAC45AF8E51":"5713C62C14B2EC0F8393B6AFD6F5785A"
 
 Camellia-192-ECB Encrypt Perl EVP #3
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"8E73B0F7DA0E6452C810F32B809079E562F8EAD2522C6B7B":"30C81C46A35CE411E5FBC1191A0A52EF":"B40ED2B60EB54D09D030CF511FEEF366"
 
 Camellia-192-ECB Encrypt Perl EVP #4
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"8E73B0F7DA0E6452C810F32B809079E562F8EAD2522C6B7B":"F69F2445DF4F9B17AD2B417BE66C3710":"909DBD95799096748CB27357E73E1D26"
 
 Camellia-256-ECB Encrypt Perl EVP #1 [#2]
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4":"6BC1BEE22E409F96E93D7E117393172A":"BEFD219B112FA00098919CD101C9CCFA"
 
 Camellia-256-ECB Encrypt Perl EVP #2
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4":"AE2D8A571E03AC9C9EB76FAC45AF8E51":"C91D3A8F1AEA08A9386CF4B66C0169EA"
 
 Camellia-256-ECB Encrypt Perl EVP #3
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4":"30C81C46A35CE411E5FBC1191A0A52EF":"A623D711DC5F25A51BB8A80D56397D28"
 
 Camellia-256-ECB Encrypt Perl EVP #4
+depends_on:MBEDTLS_CAMELLIA_C
 test_vec:MBEDTLS_CIPHER_ID_CAMELLIA:"603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4":"F69F2445DF4F9B17AD2B417BE66C3710":"7960109FB6DC42947FCFE59EA3C5EB6B"
 
diff --git a/tests/suites/test_suite_ctr_drbg.function b/tests/suites/test_suite_ctr_drbg.function
index c689699..066e70b 100644
--- a/tests/suites/test_suite_ctr_drbg.function
+++ b/tests/suites/test_suite_ctr_drbg.function
@@ -137,10 +137,12 @@
                                  data_t *result_string)
 {
     data_t empty = { 0, 0 };
+    AES_PSA_INIT();
     ctr_drbg_validate_internal(RESEED_NEVER, add_init,
                                entropy->len, entropy,
                                &empty, add1, add2,
                                result_string);
+    AES_PSA_DONE();
     goto exit; // goto is needed to avoid warning ( no test assertions in func)
 }
 /* END_CASE */
@@ -151,10 +153,12 @@
                           data_t *result_string)
 {
     data_t empty = { 0, 0 };
+    AES_PSA_INIT();
     ctr_drbg_validate_internal(RESEED_ALWAYS, add_init,
                                entropy->len / 3, entropy,
                                &empty, add1, add2,
                                result_string);
+    AES_PSA_DONE();
     goto exit; // goto is needed to avoid warning ( no test assertions in func)
 }
 /* END_CASE */
@@ -164,10 +168,12 @@
                                       data_t *add1, data_t *add_reseed,
                                       data_t *add2, data_t *result_string)
 {
+    AES_PSA_INIT();
     ctr_drbg_validate_internal(RESEED_SECOND, add_init,
                                entropy->len / 2, entropy,
                                add_reseed, add1, add2,
                                result_string);
+    AES_PSA_DONE();
     goto exit; // goto is needed to avoid warning ( no test assertions in func)
 }
 /* END_CASE */
@@ -177,10 +183,12 @@
                                     data_t *add1, data_t *add_reseed,
                                     data_t *add2, data_t *result_string)
 {
+    AES_PSA_INIT();
     ctr_drbg_validate_internal(RESEED_FIRST, add_init,
                                entropy->len / 2, entropy,
                                add_reseed, add1, add2,
                                result_string);
+    AES_PSA_DONE();
     goto exit; // goto is needed to avoid warning ( no test assertions in func)
 }
 /* END_CASE */
@@ -196,6 +204,8 @@
     size_t byte_strength = expected_bit_strength / 8;
 
     mbedtls_ctr_drbg_init(&ctx);
+
+    AES_PSA_INIT();
     test_offset_idx = 0;
     test_max_idx = sizeof(entropy);
     memset(entropy, 0, sizeof(entropy));
@@ -214,6 +224,7 @@
 
 exit:
     mbedtls_ctr_drbg_free(&ctx);
+    AES_PSA_DONE();
 }
 /* END_CASE */
 
@@ -228,6 +239,9 @@
     size_t expected_idx = 0;
 
     mbedtls_ctr_drbg_init(&ctx);
+
+    AES_PSA_INIT();
+
     test_offset_idx = 0;
     test_max_idx = sizeof(entropy);
     memset(entropy, 0, sizeof(entropy));
@@ -307,6 +321,7 @@
 
 exit:
     mbedtls_ctr_drbg_free(&ctx);
+    AES_PSA_DONE();
 }
 /* END_CASE */
 
@@ -317,6 +332,8 @@
 
     mbedtls_ctr_drbg_init(&ctx);
 
+    AES_PSA_INIT();
+
     TEST_ASSERT(mbedtls_ctr_drbg_seed(&ctx, mbedtls_test_rnd_std_rand,
                                       NULL, NULL, 0) == 0);
     TEST_ASSERT(mbedtls_ctr_drbg_write_seed_file(&ctx, path) == ret);
@@ -324,12 +341,15 @@
 
 exit:
     mbedtls_ctr_drbg_free(&ctx);
+    AES_PSA_DONE();
 }
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_SELF_TEST */
 void ctr_drbg_selftest()
 {
+    AES_PSA_INIT();
     TEST_ASSERT(mbedtls_ctr_drbg_self_test(1) == 0);
+    AES_PSA_DONE();
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index 1d96f72..032fa47 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -845,10 +845,10 @@
     psa_set_key_bits(&attributes, 0);
     psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_EXPORT);
 
-    mbedtls_test_driver_key_management_hooks.forced_status = force_status;
-
     PSA_ASSERT(psa_crypto_init());
 
+    mbedtls_test_driver_key_management_hooks.hits = 0;
+    mbedtls_test_driver_key_management_hooks.forced_status = force_status;
     actual_status = psa_import_key(&attributes, key_input->x, key_input->len, &key);
     TEST_EQUAL(mbedtls_test_driver_key_management_hooks.hits, 1);
     TEST_EQUAL(actual_status, expected_status);
@@ -906,6 +906,7 @@
     }
 
     mbedtls_test_driver_key_management_hooks.hits = 0;
+    mbedtls_test_driver_key_management_hooks.hits_export_public_key = 0;
     mbedtls_test_driver_key_management_hooks.forced_status = force_status;
 
     if (PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY(output_key_type)) {
@@ -923,7 +924,7 @@
 
     if (PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY(output_key_type) &&
         !PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY(input_key_type)) {
-        TEST_EQUAL(mbedtls_test_driver_key_management_hooks.hits, 1);
+        TEST_EQUAL(mbedtls_test_driver_key_management_hooks.hits_export_public_key, 1);
     }
 
     if (actual_status == PSA_SUCCESS) {
@@ -1059,9 +1060,11 @@
     PSA_ASSERT(psa_import_key(&attributes, key_data->x, key_data->len,
                               &key));
 
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+    mbedtls_test_driver_cipher_hooks.hits_encrypt = 0;
     PSA_ASSERT(psa_cipher_encrypt(key, alg, input->x, input->len, output1,
                                   output1_buffer_size, &output1_length));
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits_encrypt, 1);
     mbedtls_test_driver_cipher_hooks.hits = 0;
 
     PSA_ASSERT(psa_cipher_encrypt_setup(&operation, key, alg));
@@ -1161,6 +1164,7 @@
     PSA_ASSERT(psa_import_key(&attributes, key_data->x, key_data->len,
                               &key));
 
+    mbedtls_test_driver_cipher_hooks.hits = 0;
     PSA_ASSERT(psa_cipher_encrypt_setup(&operation, key, alg));
     TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
     mbedtls_test_driver_cipher_hooks.hits = 0;
@@ -1289,6 +1293,7 @@
     PSA_ASSERT(psa_import_key(&attributes, key_data->x, key_data->len,
                               &key));
 
+    mbedtls_test_driver_cipher_hooks.hits = 0;
     PSA_ASSERT(psa_cipher_decrypt_setup(&operation, key, alg));
     TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
     mbedtls_test_driver_cipher_hooks.hits = 0;
@@ -1414,6 +1419,7 @@
         mbedtls_test_driver_cipher_hooks.forced_output_length = expected_output->len;
     }
 
+    mbedtls_test_driver_cipher_hooks.hits = 0;
     status = psa_cipher_decrypt(key, alg, input, input_buffer_size, output,
                                 output_buffer_size, &output_length);
     TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
@@ -1468,23 +1474,26 @@
      * First test that if we don't force a driver error, encryption is
      * successful, then force driver error.
      */
+    mbedtls_test_driver_cipher_hooks.hits = 0;
+    mbedtls_test_driver_cipher_hooks.hits_encrypt = 0;
     status = psa_cipher_encrypt(
         key, alg, input->x, input->len,
         output, output_buffer_size, &function_output_length);
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits_encrypt, 1);
     TEST_EQUAL(status, PSA_SUCCESS);
     mbedtls_test_driver_cipher_hooks.hits = 0;
 
-    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
+    mbedtls_test_driver_cipher_hooks.forced_status_encrypt = PSA_ERROR_GENERIC_ERROR;
     /* Set the output buffer in a given state. */
     for (size_t i = 0; i < output_buffer_size; i++) {
         output[i] = 0xa5;
     }
 
+    mbedtls_test_driver_cipher_hooks.hits_encrypt = 0;
     status = psa_cipher_encrypt(
         key, alg, input->x, input->len,
         output, output_buffer_size, &function_output_length);
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits_encrypt, 1);
     TEST_EQUAL(status, PSA_ERROR_GENERIC_ERROR);
     /*
      * Check that the output buffer is still in the same state.
@@ -1545,17 +1554,18 @@
     TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 1);
     TEST_EQUAL(status, mbedtls_test_driver_cipher_hooks.forced_status);
     mbedtls_test_driver_cipher_hooks.hits = 0;
+    mbedtls_test_driver_cipher_hooks.hits_set_iv = 0;
 
-    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
+    mbedtls_test_driver_cipher_hooks.forced_status_set_iv = PSA_ERROR_GENERIC_ERROR;
     /* Set the output buffer in a given state. */
     for (size_t i = 0; i < 16; i++) {
         output[i] = 0xa5;
     }
 
     status = psa_cipher_generate_iv(&operation, output, 16, &function_output_length);
-    /* When generating the IV fails, it should call abort too */
-    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits, 2);
-    TEST_EQUAL(status, mbedtls_test_driver_cipher_hooks.forced_status);
+    TEST_EQUAL(mbedtls_test_driver_cipher_hooks.hits_set_iv, 1);
+    TEST_EQUAL(status, mbedtls_test_driver_cipher_hooks.forced_status_set_iv);
+    mbedtls_test_driver_cipher_hooks.forced_status_set_iv = PSA_SUCCESS;
     /*
      * Check that the output buffer is still in the same state.
      * This will fail if the output buffer is used by the core to pass the IV
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index b4f2d23..8564d35 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -746,19 +746,12 @@
              * MBEDTLS_SVC_KEY_ID_GET_KEY_ID( valid_handle ) is a volatile
              * key identifier as the imported key is a volatile key. Volatile
              * key identifiers are in the range from PSA_KEY_ID_VOLATILE_MIN
-             * to PSA_KEY_ID_VOLATILE_MAX included. Thus pick a key identifier
-             * in the range from PSA_KEY_ID_VOLATILE_MIN to
-             * PSA_KEY_ID_VOLATILE_MAX different from
-             * MBEDTLS_SVC_KEY_ID_GET_KEY_ID( valid_handle ) to build an
-             * unopened and thus invalid identifier.
+             * to PSA_KEY_ID_VOLATILE_MAX included. It is very unlikely that
+             * all IDs are used up to the last one, so pick
+             * PSA_KEY_ID_VOLATILE_MAX to build an unopened and thus invalid
+             * identifier.
              */
-
-            if (MBEDTLS_SVC_KEY_ID_GET_KEY_ID(valid_handle) ==
-                PSA_KEY_ID_VOLATILE_MIN) {
-                key_id = PSA_KEY_ID_VOLATILE_MIN + 1;
-            } else {
-                key_id = MBEDTLS_SVC_KEY_ID_GET_KEY_ID(valid_handle) - 1;
-            }
+            key_id = PSA_KEY_ID_VOLATILE_MAX;
 
             invalid_handle =
                 mbedtls_svc_key_id_make(0, key_id);
@@ -938,11 +931,16 @@
     mbedtls_svc_key_id_t persistent_key2 = MBEDTLS_SVC_KEY_ID_INIT;
     mbedtls_svc_key_id_t returned_key_id = MBEDTLS_SVC_KEY_ID_INIT;
     mbedtls_svc_key_id_t *keys = NULL;
+    mbedtls_psa_stats_t psa_key_slots_stats;
+    size_t available_key_slots = 0;
 
     TEST_ASSERT(MBEDTLS_PSA_KEY_SLOT_COUNT >= 1);
 
-    TEST_CALLOC(keys, MBEDTLS_PSA_KEY_SLOT_COUNT);
     PSA_ASSERT(psa_crypto_init());
+    mbedtls_psa_get_stats(&psa_key_slots_stats);
+    available_key_slots = psa_key_slots_stats.empty_slots;
+
+    TEST_CALLOC(keys, available_key_slots);
 
     psa_set_key_usage_flags(&attributes,
                             PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_COPY);
@@ -961,10 +959,10 @@
     TEST_ASSERT(mbedtls_svc_key_id_equal(returned_key_id, persistent_key));
 
     /*
-     * Create MBEDTLS_PSA_KEY_SLOT_COUNT volatile keys
+     * Create the maximum available number of volatile keys
      */
     psa_set_key_lifetime(&attributes, PSA_KEY_LIFETIME_VOLATILE);
-    for (i = 0; i < MBEDTLS_PSA_KEY_SLOT_COUNT; i++) {
+    for (i = 0; i < available_key_slots; i++) {
         PSA_ASSERT(psa_import_key(&attributes,
                                   (uint8_t *) &i, sizeof(i),
                                   &keys[i]));
@@ -983,12 +981,12 @@
      * Check we can export the volatile key created last and that it has the
      * expected value. Then, destroy it.
      */
-    PSA_ASSERT(psa_export_key(keys[MBEDTLS_PSA_KEY_SLOT_COUNT - 1],
+    PSA_ASSERT(psa_export_key(keys[available_key_slots - 1],
                               exported, sizeof(exported),
                               &exported_length));
-    i = MBEDTLS_PSA_KEY_SLOT_COUNT - 1;
+    i = available_key_slots - 1;
     TEST_MEMORY_COMPARE(exported, exported_length, (uint8_t *) &i, sizeof(i));
-    PSA_ASSERT(psa_destroy_key(keys[MBEDTLS_PSA_KEY_SLOT_COUNT - 1]));
+    PSA_ASSERT(psa_destroy_key(keys[available_key_slots - 1]));
 
     /*
      * Check that we can now access the persistent key again.
@@ -1011,7 +1009,7 @@
      * Check we can export the remaining volatile keys and that they have the
      * expected values.
      */
-    for (i = 0; i < (MBEDTLS_PSA_KEY_SLOT_COUNT - 1); i++) {
+    for (i = 0; i < (available_key_slots - 1); i++) {
         PSA_ASSERT(psa_export_key(keys[i],
                                   exported, sizeof(exported),
                                   &exported_length));
diff --git a/tests/suites/test_suite_random.function b/tests/suites/test_suite_random.function
index 58cddb7..155b8e7 100644
--- a/tests/suites/test_suite_random.function
+++ b/tests/suites/test_suite_random.function
@@ -26,7 +26,12 @@
     unsigned char output1[OUTPUT_SIZE];
     unsigned char output2[OUTPUT_SIZE];
 
+#if defined(MBEDTLS_AES_C)
     MD_PSA_INIT();
+#else
+    USE_PSA_INIT();
+#endif
+
 
     /* First round */
     mbedtls_entropy_init(&entropy);
@@ -56,7 +61,11 @@
 exit:
     mbedtls_ctr_drbg_free(&drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_AES_C)
     MD_PSA_DONE();
+#else
+    USE_PSA_DONE();
+#endif
 }
 /* END_CASE */
 
diff --git a/tests/suites/test_suite_x509write.function b/tests/suites/test_suite_x509write.function
index 4de9add..b59fd48 100644
--- a/tests/suites/test_suite_x509write.function
+++ b/tests/suites/test_suite_x509write.function
@@ -665,13 +665,15 @@
     mbedtls_x509write_cert ctx;
     uint8_t invalid_serial[MBEDTLS_X509_RFC5280_MAX_SERIAL_LEN + 1];
 
+#if defined(MBEDTLS_TEST_DEPRECATED) && defined(MBEDTLS_BIGNUM_C)
+    mbedtls_mpi serial_mpi;
+    mbedtls_mpi_init(&serial_mpi);
+#endif
+
     USE_PSA_INIT();
     memset(invalid_serial, 0x01, sizeof(invalid_serial));
 
 #if defined(MBEDTLS_TEST_DEPRECATED) && defined(MBEDTLS_BIGNUM_C)
-    mbedtls_mpi serial_mpi;
-
-    mbedtls_mpi_init(&serial_mpi);
     TEST_EQUAL(mbedtls_mpi_read_binary(&serial_mpi, invalid_serial,
                                        sizeof(invalid_serial)), 0);
     TEST_EQUAL(mbedtls_x509write_crt_set_serial(&ctx, &serial_mpi),