Merge pull request #7461 from valeriosetti/issue7460-part1

Fixing USE_PSA_INIT/DONE in SSL/X509/PK test suites
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index a7bf198..3a8c5c6 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -24,7 +24,7 @@
 1. [Check for open issues](https://github.com/Mbed-TLS/mbedtls/issues) or [start a discussion](https://lists.trustedfirmware.org/mailman3/lists/mbed-tls.lists.trustedfirmware.org) around a feature idea or a bug.
 1. Fork the [Mbed TLS repository on GitHub](https://github.com/Mbed-TLS/mbedtls) to start making your changes. As a general rule, you should use the ["development" branch](https://github.com/Mbed-TLS/mbedtls/tree/development) as a basis.
 1. Write a test which shows that the bug was fixed or that the feature works as expected.
-1. Send a pull request (PR) and work with us until it gets merged and published. Contributions may need some modifications, so a few rounds of review and fixing may be necessary. We will include your name in the ChangeLog :)
+1. Send a pull request (PR) and work with us until it gets merged and published. Contributions may need some modifications, so a few rounds of review and fixing may be necessary. See our [review process guidelines](https://mbed-tls.readthedocs.io/en/latest/reviews/review-for-contributors/).
 1. For quick merging, the contribution should be short, and concentrated on a single feature or topic. The larger the contribution is, the longer it would take to review it and merge it.
 
 Backwards Compatibility
diff --git a/ChangeLog.d/mbedtls_ecdsa_can_do-unconditional-define.txt b/ChangeLog.d/mbedtls_ecdsa_can_do-unconditional-define.txt
new file mode 100644
index 0000000..22e8adb
--- /dev/null
+++ b/ChangeLog.d/mbedtls_ecdsa_can_do-unconditional-define.txt
@@ -0,0 +1,3 @@
+Bugfix
+   * Fix an error when MBEDTLS_ECDSA_SIGN_ALT is defined but not
+     MBEDTLS_ECDSA_VERIFY_ALT, causing ecdsa verify to fail. Fixes #7498.
diff --git a/ChangeLog.d/programs_psa_fix.txt b/ChangeLog.d/programs_psa_fix.txt
new file mode 100644
index 0000000..fe2099e
--- /dev/null
+++ b/ChangeLog.d/programs_psa_fix.txt
@@ -0,0 +1,3 @@
+Bugfix
+   * Fix missing PSA initialization in sample programs when
+     MBEDTLS_USE_PSA_CRYPTO is enabled.
diff --git a/include/mbedtls/check_config.h b/include/mbedtls/check_config.h
index 266eb9e..c81cd1c 100644
--- a/include/mbedtls/check_config.h
+++ b/include/mbedtls/check_config.h
@@ -66,13 +66,6 @@
 #error "MBEDTLS_HAVE_TIME_DATE without MBEDTLS_HAVE_TIME does not make sense"
 #endif
 
-#if defined(__aarch64__) && defined(__GNUC__)
-/* We don't do anything with MBEDTLS_AESCE_C on systems without ^ these two */
-#if defined(MBEDTLS_AESCE_C) && !defined(MBEDTLS_HAVE_ASM)
-#error "MBEDTLS_AESCE_C defined, but not all prerequisites"
-#endif
-#endif
-
 #if defined(MBEDTLS_CTR_DRBG_C) && !defined(MBEDTLS_AES_C)
 #error "MBEDTLS_CTR_DRBG_C defined, but not all prerequisites"
 #endif
diff --git a/include/mbedtls/mbedtls_config.h b/include/mbedtls/mbedtls_config.h
index ab4b83b..bc0f502 100644
--- a/include/mbedtls/mbedtls_config.h
+++ b/include/mbedtls/mbedtls_config.h
@@ -2077,12 +2077,15 @@
  * Module:  library/aesce.c
  * Caller:  library/aes.c
  *
- * Requires: MBEDTLS_HAVE_ASM, MBEDTLS_AES_C
+ * Requires: MBEDTLS_AES_C
  *
  * \warning Runtime detection only works on Linux. For non-Linux operating
  *          system, Armv8-A Cryptographic Extensions must be supported by
  *          the CPU when this option is enabled.
  *
+ * \note    Minimum compiler versions for this feature are Clang 4.0,
+ *          GCC 6.0 or MSVC 2019 version 16.11.2.
+ *
  * This module adds support for the AES Armv8-A Cryptographic Extensions on Aarch64 systems.
  */
 #define MBEDTLS_AESCE_C
diff --git a/library/aesce.c b/library/aesce.c
index fe056dc..ff8c2e0 100644
--- a/library/aesce.c
+++ b/library/aesce.c
@@ -48,22 +48,34 @@
 
 #if defined(MBEDTLS_HAVE_ARM64)
 
+/* Compiler version checks. */
+#if defined(__clang__)
+#   if __clang_major__ < 4
+#       error "Minimum version of Clang for MBEDTLS_AESCE_C is 4.0."
+#   endif
+#elif defined(__GNUC__)
+#   if __GNUC__ < 6
+#       error "Minimum version of GCC for MBEDTLS_AESCE_C is 6.0."
+#   endif
+#elif defined(_MSC_VER)
+/* TODO: We haven't verified MSVC from 1920 to 1928. If someone verified that,
+ *       please update this and document of `MBEDTLS_AESCE_C` in
+ *       `mbedtls_config.h`. */
+#   if _MSC_VER < 1929
+#       error "Minimum version of MSVC for MBEDTLS_AESCE_C is 2019 version 16.11.2."
+#   endif
+#endif
+
 #if !defined(__ARM_FEATURE_AES) || defined(MBEDTLS_ENABLE_ARM_CRYPTO_EXTENSIONS_COMPILER_FLAG)
 #   if defined(__clang__)
-#       if __clang_major__ < 4
-#           error "A more recent Clang is required for MBEDTLS_AESCE_C"
-#       endif
 #       pragma clang attribute push (__attribute__((target("crypto"))), apply_to=function)
 #       define MBEDTLS_POP_TARGET_PRAGMA
 #   elif defined(__GNUC__)
-#       if __GNUC__ < 6
-#           error "A more recent GCC is required for MBEDTLS_AESCE_C"
-#       endif
 #       pragma GCC push_options
 #       pragma GCC target ("arch=armv8-a+crypto")
 #       define MBEDTLS_POP_TARGET_PRAGMA
-#   else
-#       error "Only GCC and Clang supported for MBEDTLS_AESCE_C"
+#   elif defined(_MSC_VER)
+#       error "Required feature(__ARM_FEATURE_AES) is not enabled."
 #   endif
 #endif /* !__ARM_FEATURE_AES || MBEDTLS_ENABLE_ARM_CRYPTO_EXTENSIONS_COMPILER_FLAG */
 
@@ -295,12 +307,24 @@
  * Older compilers miss some intrinsic functions for `poly*_t`. We use
  * uint8x16_t and uint8x16x3_t as input/output parameters.
  */
+#if defined(__GNUC__) && !defined(__clang__)
+/* GCC reports incompatible type error without cast. GCC think poly64_t and
+ * poly64x1_t are different, that is different with MSVC and Clang. */
+#define MBEDTLS_VMULL_P64(a, b) vmull_p64((poly64_t) a, (poly64_t) b)
+#else
+/* MSVC reports `error C2440: 'type cast'` with cast. Clang does not report
+ * error with/without cast. And I think poly64_t and poly64x1_t are same, no
+ * cast for clang also. */
+#define MBEDTLS_VMULL_P64(a, b) vmull_p64(a, b)
+#endif
 static inline uint8x16_t pmull_low(uint8x16_t a, uint8x16_t b)
 {
+
     return vreinterpretq_u8_p128(
-        vmull_p64(
-            (poly64_t) vget_low_p64(vreinterpretq_p64_u8(a)),
-            (poly64_t) vget_low_p64(vreinterpretq_p64_u8(b))));
+        MBEDTLS_VMULL_P64(
+            vget_low_p64(vreinterpretq_p64_u8(a)),
+            vget_low_p64(vreinterpretq_p64_u8(b))
+            ));
 }
 
 static inline uint8x16_t pmull_high(uint8x16_t a, uint8x16_t b)
@@ -362,9 +386,14 @@
 static inline uint8x16_t poly_mult_reduce(uint8x16x3_t input)
 {
     uint8x16_t const ZERO = vdupq_n_u8(0);
-    /* use 'asm' as an optimisation barrier to prevent loading MODULO from memory */
+
     uint64x2_t r = vreinterpretq_u64_u8(vdupq_n_u8(0x87));
+#if defined(__GNUC__)
+    /* use 'asm' as an optimisation barrier to prevent loading MODULO from
+     * memory. It is for GNUC compatible compilers.
+     */
     asm ("" : "+w" (r));
+#endif
     uint8x16_t const MODULO = vreinterpretq_u8_u64(vshrq_n_u64(r, 64 - 8));
     uint8x16_t h, m, l; /* input high/middle/low 128b */
     uint8x16_t c, d, e, f, g, n, o;
diff --git a/library/aesce.h b/library/aesce.h
index 12ddc74..7048d77 100644
--- a/library/aesce.h
+++ b/library/aesce.h
@@ -30,11 +30,11 @@
 
 #include "mbedtls/aes.h"
 
-
-#if defined(MBEDTLS_HAVE_ASM) && defined(__GNUC__) && \
-    defined(__aarch64__) && !defined(MBEDTLS_HAVE_ARM64)
+#if !defined(MBEDTLS_HAVE_ARM64)
+#if defined(__aarch64__) || defined(_M_ARM64) || defined(_M_ARM64EC)
 #define MBEDTLS_HAVE_ARM64
 #endif
+#endif
 
 #if defined(MBEDTLS_HAVE_ARM64)
 
diff --git a/library/bignum_core.c b/library/bignum_core.c
index c6d92fb..b0ffa37 100644
--- a/library/bignum_core.c
+++ b/library/bignum_core.c
@@ -35,6 +35,23 @@
 
 size_t mbedtls_mpi_core_clz(mbedtls_mpi_uint a)
 {
+#if defined(__has_builtin)
+#if __has_builtin(__builtin_clz)
+    if (sizeof(mbedtls_mpi_uint) == sizeof(unsigned int)) {
+        return (size_t) __builtin_clz(a);
+    }
+#endif
+#if __has_builtin(__builtin_clzl)
+    if (sizeof(mbedtls_mpi_uint) == sizeof(unsigned long)) {
+        return (size_t) __builtin_clzl(a);
+    }
+#endif
+#if __has_builtin(__builtin_clzll)
+    if (sizeof(mbedtls_mpi_uint) == sizeof(unsigned long long)) {
+        return (size_t) __builtin_clzll(a);
+    }
+#endif
+#endif
     size_t j;
     mbedtls_mpi_uint mask = (mbedtls_mpi_uint) 1 << (biL - 1);
 
@@ -51,21 +68,17 @@
 
 size_t mbedtls_mpi_core_bitlen(const mbedtls_mpi_uint *A, size_t A_limbs)
 {
-    size_t i, j;
+    int i;
+    size_t j;
 
-    if (A_limbs == 0) {
-        return 0;
-    }
-
-    for (i = A_limbs - 1; i > 0; i--) {
+    for (i = ((int) A_limbs) - 1; i >= 0; i--) {
         if (A[i] != 0) {
-            break;
+            j = biL - mbedtls_mpi_core_clz(A[i]);
+            return (i * biL) + j;
         }
     }
 
-    j = biL - mbedtls_mpi_core_clz(A[i]);
-
-    return (i * biL) + j;
+    return 0;
 }
 
 /* Convert a big-endian byte array aligned to the size of mbedtls_mpi_uint
diff --git a/library/bignum_core.h b/library/bignum_core.h
index b3d05a3..158d2b3 100644
--- a/library/bignum_core.h
+++ b/library/bignum_core.h
@@ -102,9 +102,12 @@
 
 /** Count leading zero bits in a given integer.
  *
+ * \warning     The result is undefined if \p a == 0
+ *
  * \param a     Integer to count leading zero bits.
  *
- * \return      The number of leading zero bits in \p a.
+ * \return      The number of leading zero bits in \p a, if \p a != 0.
+ *              If \p a == 0, the result is undefined.
  */
 size_t mbedtls_mpi_core_clz(mbedtls_mpi_uint a);
 
diff --git a/library/ecdh.c b/library/ecdh.c
index b529af5..58ef881 100644
--- a/library/ecdh.c
+++ b/library/ecdh.c
@@ -20,7 +20,7 @@
 /*
  * References:
  *
- * SEC1 http://www.secg.org/index.php?action=secg,docs_secg
+ * SEC1 https://www.secg.org/sec1-v2.pdf
  * RFC 4492
  */
 
diff --git a/library/ecdsa.c b/library/ecdsa.c
index eb3c303..1faec16 100644
--- a/library/ecdsa.c
+++ b/library/ecdsa.c
@@ -20,7 +20,7 @@
 /*
  * References:
  *
- * SEC1 http://www.secg.org/index.php?action=secg,docs_secg
+ * SEC1 https://www.secg.org/sec1-v2.pdf
  */
 
 #include "common.h"
@@ -234,6 +234,19 @@
 }
 #endif /* ECDSA_DETERMINISTIC || !ECDSA_SIGN_ALT || !ECDSA_VERIFY_ALT */
 
+int mbedtls_ecdsa_can_do(mbedtls_ecp_group_id gid)
+{
+    switch (gid) {
+#ifdef MBEDTLS_ECP_DP_CURVE25519_ENABLED
+        case MBEDTLS_ECP_DP_CURVE25519: return 0;
+#endif
+#ifdef MBEDTLS_ECP_DP_CURVE448_ENABLED
+        case MBEDTLS_ECP_DP_CURVE448: return 0;
+#endif
+        default: return 1;
+    }
+}
+
 #if !defined(MBEDTLS_ECDSA_SIGN_ALT)
 /*
  * Compute ECDSA signature of a hashed message (SEC1 4.1.3)
@@ -373,19 +386,6 @@
     return ret;
 }
 
-int mbedtls_ecdsa_can_do(mbedtls_ecp_group_id gid)
-{
-    switch (gid) {
-#ifdef MBEDTLS_ECP_DP_CURVE25519_ENABLED
-        case MBEDTLS_ECP_DP_CURVE25519: return 0;
-#endif
-#ifdef MBEDTLS_ECP_DP_CURVE448_ENABLED
-        case MBEDTLS_ECP_DP_CURVE448: return 0;
-#endif
-        default: return 1;
-    }
-}
-
 /*
  * Compute ECDSA signature of a hashed message
  */
diff --git a/library/ecp.c b/library/ecp.c
index 5d13b8e..086acb3 100644
--- a/library/ecp.c
+++ b/library/ecp.c
@@ -24,9 +24,11 @@
  * GECC = Guide to Elliptic Curve Cryptography - Hankerson, Menezes, Vanstone
  * FIPS 186-3 http://csrc.nist.gov/publications/fips/fips186-3/fips_186-3.pdf
  * RFC 4492 for the related TLS structures and constants
+ * - https://www.rfc-editor.org/rfc/rfc4492
  * RFC 7748 for the Curve448 and Curve25519 curve definitions
+ * - https://www.rfc-editor.org/rfc/rfc7748
  *
- * [Curve25519] http://cr.yp.to/ecdh/curve25519-20060209.pdf
+ * [Curve25519] https://cr.yp.to/ecdh/curve25519-20060209.pdf
  *
  * [2] CORON, Jean-S'ebastien. Resistance against differential power analysis
  *     for elliptic curve cryptosystems. In : Cryptographic Hardware and
diff --git a/library/ecp_curves.c b/library/ecp_curves.c
index 1376f5d..1640107 100644
--- a/library/ecp_curves.c
+++ b/library/ecp_curves.c
@@ -4618,6 +4618,8 @@
 #endif
 #if defined(MBEDTLS_ECP_DP_SECP256K1_ENABLED)
 static int ecp_mod_p256k1(mbedtls_mpi *);
+MBEDTLS_STATIC_TESTABLE
+int mbedtls_ecp_mod_p256k1(mbedtls_mpi *);
 #endif
 
 #if defined(ECP_LOAD_GROUP)
@@ -5528,7 +5530,6 @@
                                   size_t adjust, size_t shift, mbedtls_mpi_uint mask)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t i;
     mbedtls_mpi M, R;
     mbedtls_mpi_uint Mp[P_KOBLITZ_MAX + P_KOBLITZ_R + 1];
 
@@ -5545,55 +5546,31 @@
     M.s = 1;
     M.p = Mp;
 
-    /* M = A1 */
-    M.n = N->n - (p_limbs - adjust);
-    if (M.n > p_limbs + adjust) {
-        M.n = p_limbs + adjust;
-    }
-    memset(Mp, 0, sizeof(Mp));
-    memcpy(Mp, N->p + p_limbs - adjust, M.n * sizeof(mbedtls_mpi_uint));
-    if (shift != 0) {
-        MBEDTLS_MPI_CHK(mbedtls_mpi_shift_r(&M, shift));
-    }
-    M.n += R.n; /* Make room for multiplication by R */
+    for (size_t pass = 0; pass < 2; pass++) {
+        /* M = A1 */
+        M.n = N->n - (p_limbs - adjust);
+        if (M.n > p_limbs + adjust) {
+            M.n = p_limbs + adjust;
+        }
+        memset(Mp, 0, sizeof(Mp));
+        memcpy(Mp, N->p + p_limbs - adjust, M.n * sizeof(mbedtls_mpi_uint));
+        if (shift != 0) {
+            MBEDTLS_MPI_CHK(mbedtls_mpi_shift_r(&M, shift));
+        }
+        M.n += R.n; /* Make room for multiplication by R */
 
-    /* N = A0 */
-    if (mask != 0) {
-        N->p[p_limbs - 1] &= mask;
-    }
-    for (i = p_limbs; i < N->n; i++) {
-        N->p[i] = 0;
-    }
+        /* N = A0 */
+        if (mask != 0) {
+            N->p[p_limbs - 1] &= mask;
+        }
+        for (size_t i = p_limbs; i < N->n; i++) {
+            N->p[i] = 0;
+        }
 
-    /* N = A0 + R * A1 */
-    MBEDTLS_MPI_CHK(mbedtls_mpi_mul_mpi(&M, &M, &R));
-    MBEDTLS_MPI_CHK(mbedtls_mpi_add_abs(N, N, &M));
-
-    /* Second pass */
-
-    /* M = A1 */
-    M.n = N->n - (p_limbs - adjust);
-    if (M.n > p_limbs + adjust) {
-        M.n = p_limbs + adjust;
+        /* N = A0 + R * A1 */
+        MBEDTLS_MPI_CHK(mbedtls_mpi_mul_mpi(&M, &M, &R));
+        MBEDTLS_MPI_CHK(mbedtls_mpi_add_abs(N, N, &M));
     }
-    memset(Mp, 0, sizeof(Mp));
-    memcpy(Mp, N->p + p_limbs - adjust, M.n * sizeof(mbedtls_mpi_uint));
-    if (shift != 0) {
-        MBEDTLS_MPI_CHK(mbedtls_mpi_shift_r(&M, shift));
-    }
-    M.n += R.n; /* Make room for multiplication by R */
-
-    /* N = A0 */
-    if (mask != 0) {
-        N->p[p_limbs - 1] &= mask;
-    }
-    for (i = p_limbs; i < N->n; i++) {
-        N->p[i] = 0;
-    }
-
-    /* N = A0 + R * A1 */
-    MBEDTLS_MPI_CHK(mbedtls_mpi_mul_mpi(&M, &M, &R));
-    MBEDTLS_MPI_CHK(mbedtls_mpi_add_abs(N, N, &M));
 
 cleanup:
     return ret;
@@ -5655,11 +5632,18 @@
 #endif /* MBEDTLS_ECP_DP_SECP224K1_ENABLED */
 
 #if defined(MBEDTLS_ECP_DP_SECP256K1_ENABLED)
+
+static int ecp_mod_p256k1(mbedtls_mpi *N)
+{
+    return mbedtls_ecp_mod_p256k1(N);
+}
+
 /*
  * Fast quasi-reduction modulo p256k1 = 2^256 - R,
  * with R = 2^32 + 2^9 + 2^8 + 2^7 + 2^6 + 2^4 + 1 = 0x01000003D1
  */
-static int ecp_mod_p256k1(mbedtls_mpi *N)
+MBEDTLS_STATIC_TESTABLE
+int mbedtls_ecp_mod_p256k1(mbedtls_mpi *N)
 {
     static mbedtls_mpi_uint Rp[] = {
         MBEDTLS_BYTES_TO_T_UINT_8(0xD1, 0x03, 0x00, 0x00, 0x01, 0x00, 0x00,
diff --git a/library/ecp_invasive.h b/library/ecp_invasive.h
index b7cd6e2..0d093f3 100644
--- a/library/ecp_invasive.h
+++ b/library/ecp_invasive.h
@@ -47,7 +47,7 @@
  *                  This is the bit-size of the key minus 1:
  *                  254 for Curve25519 or 447 for Curve448.
  * \param d         The randomly generated key. This is a number of size
- *                  exactly \p n_bits + 1 bits, with the least significant bits
+ *                  exactly \p high_bit + 1 bits, with the least significant bits
  *                  masked as specified in [Curve25519] and in [RFC7748] §5.
  * \param f_rng     The RNG function.
  * \param p_rng     The RNG context to be passed to \p f_rng.
@@ -55,7 +55,7 @@
  * \return          \c 0 on success.
  * \return          \c MBEDTLS_ERR_ECP_xxx or MBEDTLS_ERR_MPI_xxx on failure.
  */
-int mbedtls_ecp_gen_privkey_mx(size_t n_bits,
+int mbedtls_ecp_gen_privkey_mx(size_t high_bit,
                                mbedtls_mpi *d,
                                int (*f_rng)(void *, unsigned char *, size_t),
                                void *p_rng);
@@ -186,6 +186,13 @@
 
 #endif /* MBEDTLS_ECP_DP_SECP224K1_ENABLED */
 
+#if defined(MBEDTLS_ECP_DP_SECP256K1_ENABLED)
+
+MBEDTLS_STATIC_TESTABLE
+int mbedtls_ecp_mod_p256k1(mbedtls_mpi *N);
+
+#endif /* MBEDTLS_ECP_DP_SECP256K1_ENABLED */
+
 /** Initialise a modulus with hard-coded const curve data.
  *
  * \note            The caller is responsible for the \p N modulus' memory.
diff --git a/library/pk.c b/library/pk.c
index a93ed56..2516eed 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -443,7 +443,7 @@
             return ret;
         }
 
-        ret = ctx->pk_info->verify_rs_func(ctx->pk_ctx,
+        ret = ctx->pk_info->verify_rs_func(ctx,
                                            md_alg, hash, hash_len, sig, sig_len, rs_ctx->rs_ctx);
 
         if (ret != MBEDTLS_ERR_ECP_IN_PROGRESS) {
@@ -460,7 +460,7 @@
         return MBEDTLS_ERR_PK_TYPE_MISMATCH;
     }
 
-    return ctx->pk_info->verify_func(ctx->pk_ctx, md_alg, hash, hash_len,
+    return ctx->pk_info->verify_func(ctx, md_alg, hash, hash_len,
                                      sig, sig_len);
 }
 
@@ -626,7 +626,7 @@
             return ret;
         }
 
-        ret = ctx->pk_info->sign_rs_func(ctx->pk_ctx, md_alg,
+        ret = ctx->pk_info->sign_rs_func(ctx, md_alg,
                                          hash, hash_len,
                                          sig, sig_size, sig_len,
                                          f_rng, p_rng, rs_ctx->rs_ctx);
@@ -645,7 +645,7 @@
         return MBEDTLS_ERR_PK_TYPE_MISMATCH;
     }
 
-    return ctx->pk_info->sign_func(ctx->pk_ctx, md_alg,
+    return ctx->pk_info->sign_func(ctx, md_alg,
                                    hash, hash_len,
                                    sig, sig_size, sig_len,
                                    f_rng, p_rng);
@@ -736,7 +736,7 @@
         return MBEDTLS_ERR_PK_TYPE_MISMATCH;
     }
 
-    return ctx->pk_info->decrypt_func(ctx->pk_ctx, input, ilen,
+    return ctx->pk_info->decrypt_func(ctx, input, ilen,
                                       output, olen, osize, f_rng, p_rng);
 }
 
@@ -756,7 +756,7 @@
         return MBEDTLS_ERR_PK_TYPE_MISMATCH;
     }
 
-    return ctx->pk_info->encrypt_func(ctx->pk_ctx, input, ilen,
+    return ctx->pk_info->encrypt_func(ctx, input, ilen,
                                       output, olen, osize, f_rng, p_rng);
 }
 
@@ -791,7 +791,9 @@
         }
     }
 
-    return prv->pk_info->check_pair_func(pub->pk_ctx, prv->pk_ctx, f_rng, p_rng);
+    return prv->pk_info->check_pair_func((mbedtls_pk_context *) pub,
+                                         (mbedtls_pk_context *) prv,
+                                         f_rng, p_rng);
 }
 
 /*
@@ -805,7 +807,7 @@
         return 0;
     }
 
-    return ctx->pk_info->get_bitlen(ctx->pk_ctx);
+    return ctx->pk_info->get_bitlen((mbedtls_pk_context *) ctx);
 }
 
 /*
@@ -821,7 +823,7 @@
         return MBEDTLS_ERR_PK_TYPE_MISMATCH;
     }
 
-    ctx->pk_info->debug_func(ctx->pk_ctx, items);
+    ctx->pk_info->debug_func((mbedtls_pk_context *) ctx, items);
     return 0;
 }
 
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 57bfdca..6c9f97b 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -191,18 +191,18 @@
            type == MBEDTLS_PK_RSASSA_PSS;
 }
 
-static size_t rsa_get_bitlen(const void *ctx)
+static size_t rsa_get_bitlen(mbedtls_pk_context *pk)
 {
-    const mbedtls_rsa_context *rsa = (const mbedtls_rsa_context *) ctx;
+    const mbedtls_rsa_context *rsa = (const mbedtls_rsa_context *) pk->pk_ctx;
     return 8 * mbedtls_rsa_get_len(rsa);
 }
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-static int rsa_verify_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int rsa_verify_wrap(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg,
                            const unsigned char *hash, size_t hash_len,
                            const unsigned char *sig, size_t sig_len)
 {
-    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) ctx;
+    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) pk->pk_ctx;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
@@ -225,7 +225,7 @@
     /* mbedtls_pk_write_pubkey_der() expects a full PK context;
      * re-construct one to make it happy */
     key.pk_info = &mbedtls_rsa_info;
-    key.pk_ctx = ctx;
+    key.pk_ctx = rsa;
     key_len = mbedtls_pk_write_pubkey_der(&key, buf, sizeof(buf));
     if (key_len <= 0) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
@@ -260,12 +260,12 @@
     return ret;
 }
 #else
-static int rsa_verify_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int rsa_verify_wrap(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg,
                            const unsigned char *hash, size_t hash_len,
                            const unsigned char *sig, size_t sig_len)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) ctx;
+    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) pk->pk_ctx;
     size_t rsa_len = mbedtls_rsa_get_len(rsa);
 
     if (md_alg == MBEDTLS_MD_NONE && UINT_MAX < hash_len) {
@@ -354,7 +354,7 @@
 #endif /* MBEDTLS_PSA_CRYPTO_C */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-static int rsa_sign_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int rsa_sign_wrap(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
@@ -370,16 +370,16 @@
 
     return mbedtls_pk_psa_rsa_sign_ext(PSA_ALG_RSA_PKCS1V15_SIGN(
                                            psa_md_alg),
-                                       ctx, hash, hash_len,
+                                       pk->pk_ctx, hash, hash_len,
                                        sig, sig_size, sig_len);
 }
 #else
-static int rsa_sign_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int rsa_sign_wrap(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) ctx;
+    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) pk->pk_ctx;
 
     if (md_alg == MBEDTLS_MD_NONE && UINT_MAX < hash_len) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
@@ -397,12 +397,12 @@
 #endif
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-static int rsa_decrypt_wrap(void *ctx,
+static int rsa_decrypt_wrap(mbedtls_pk_context *pk,
                             const unsigned char *input, size_t ilen,
                             unsigned char *output, size_t *olen, size_t osize,
                             int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) ctx;
+    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) pk->pk_ctx;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
@@ -427,7 +427,7 @@
     /* mbedtls_pk_write_key_der() expects a full PK context;
      * re-construct one to make it happy */
     key.pk_info = &mbedtls_rsa_info;
-    key.pk_ctx = ctx;
+    key.pk_ctx = rsa;
     key_len = mbedtls_pk_write_key_der(&key, buf, sizeof(buf));
     if (key_len <= 0) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
@@ -466,12 +466,12 @@
     return ret;
 }
 #else
-static int rsa_decrypt_wrap(void *ctx,
+static int rsa_decrypt_wrap(mbedtls_pk_context *pk,
                             const unsigned char *input, size_t ilen,
                             unsigned char *output, size_t *olen, size_t osize,
                             int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) ctx;
+    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) pk->pk_ctx;
 
     if (ilen != mbedtls_rsa_get_len(rsa)) {
         return MBEDTLS_ERR_RSA_BAD_INPUT_DATA;
@@ -483,12 +483,12 @@
 #endif
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-static int rsa_encrypt_wrap(void *ctx,
+static int rsa_encrypt_wrap(mbedtls_pk_context *pk,
                             const unsigned char *input, size_t ilen,
                             unsigned char *output, size_t *olen, size_t osize,
                             int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) ctx;
+    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) pk->pk_ctx;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
@@ -513,7 +513,7 @@
     /* mbedtls_pk_write_pubkey_der() expects a full PK context;
      * re-construct one to make it happy */
     key.pk_info = &mbedtls_rsa_info;
-    key.pk_ctx = ctx;
+    key.pk_ctx = rsa;
     key_len = mbedtls_pk_write_pubkey_der(&key, buf, sizeof(buf));
     if (key_len <= 0) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
@@ -551,12 +551,12 @@
     return ret;
 }
 #else
-static int rsa_encrypt_wrap(void *ctx,
+static int rsa_encrypt_wrap(mbedtls_pk_context *pk,
                             const unsigned char *input, size_t ilen,
                             unsigned char *output, size_t *olen, size_t osize,
                             int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) ctx;
+    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) pk->pk_ctx;
     *olen = mbedtls_rsa_get_len(rsa);
 
     if (*olen > osize) {
@@ -568,14 +568,14 @@
 }
 #endif
 
-static int rsa_check_pair_wrap(const void *pub, const void *prv,
+static int rsa_check_pair_wrap(mbedtls_pk_context *pub, mbedtls_pk_context *prv,
                                int (*f_rng)(void *, unsigned char *, size_t),
                                void *p_rng)
 {
     (void) f_rng;
     (void) p_rng;
-    return mbedtls_rsa_check_pub_priv((const mbedtls_rsa_context *) pub,
-                                      (const mbedtls_rsa_context *) prv);
+    return mbedtls_rsa_check_pub_priv((const mbedtls_rsa_context *) pub->pk_ctx,
+                                      (const mbedtls_rsa_context *) prv->pk_ctx);
 }
 
 static void *rsa_alloc_wrap(void)
@@ -595,22 +595,24 @@
     mbedtls_free(ctx);
 }
 
-static void rsa_debug(const void *ctx, mbedtls_pk_debug_item *items)
+static void rsa_debug(mbedtls_pk_context *pk, mbedtls_pk_debug_item *items)
 {
 #if defined(MBEDTLS_RSA_ALT)
     /* Not supported */
-    (void) ctx;
+    (void) pk;
     (void) items;
 #else
+    mbedtls_rsa_context *rsa = (mbedtls_rsa_context *) pk->pk_ctx;
+
     items->type = MBEDTLS_PK_DEBUG_MPI;
     items->name = "rsa.N";
-    items->value = &(((mbedtls_rsa_context *) ctx)->N);
+    items->value = &(rsa->N);
 
     items++;
 
     items->type = MBEDTLS_PK_DEBUG_MPI;
     items->name = "rsa.E";
-    items->value = &(((mbedtls_rsa_context *) ctx)->E);
+    items->value = &(rsa->E);
 #endif
 }
 
@@ -649,9 +651,10 @@
            type == MBEDTLS_PK_ECDSA;
 }
 
-static size_t eckey_get_bitlen(const void *ctx)
+static size_t eckey_get_bitlen(mbedtls_pk_context *pk)
 {
-    return ((mbedtls_ecp_keypair *) ctx)->grp.pbits;
+    mbedtls_ecp_keypair *ecp = (mbedtls_ecp_keypair *) pk->pk_ctx;
+    return ecp->grp.pbits;
 }
 
 #if defined(MBEDTLS_PK_CAN_ECDSA_VERIFY)
@@ -716,11 +719,12 @@
     return 0;
 }
 
-static int ecdsa_verify_wrap(void *ctx_arg, mbedtls_md_type_t md_alg,
+static int ecdsa_verify_wrap(mbedtls_pk_context *pk,
+                             mbedtls_md_type_t md_alg,
                              const unsigned char *hash, size_t hash_len,
                              const unsigned char *sig, size_t sig_len)
 {
-    mbedtls_ecp_keypair *ctx = ctx_arg;
+    mbedtls_ecp_keypair *ctx = pk->pk_ctx;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
@@ -799,14 +803,14 @@
     return ret;
 }
 #else /* MBEDTLS_USE_PSA_CRYPTO */
-static int ecdsa_verify_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int ecdsa_verify_wrap(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg,
                              const unsigned char *hash, size_t hash_len,
                              const unsigned char *sig, size_t sig_len)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     ((void) md_alg);
 
-    ret = mbedtls_ecdsa_read_signature((mbedtls_ecdsa_context *) ctx,
+    ret = mbedtls_ecdsa_read_signature((mbedtls_ecdsa_context *) pk->pk_ctx,
                                        hash, hash_len, sig, sig_len);
 
     if (ret == MBEDTLS_ERR_ECP_SIG_LEN_MISMATCH) {
@@ -904,12 +908,12 @@
     return 0;
 }
 
-static int ecdsa_sign_wrap(void *ctx_arg, mbedtls_md_type_t md_alg,
+static int ecdsa_sign_wrap(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_ecp_keypair *ctx = ctx_arg;
+    mbedtls_ecp_keypair *ctx = pk->pk_ctx;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
@@ -974,12 +978,12 @@
     return ret;
 }
 #else /* MBEDTLS_USE_PSA_CRYPTO */
-static int ecdsa_sign_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int ecdsa_sign_wrap(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    return mbedtls_ecdsa_write_signature((mbedtls_ecdsa_context *) ctx,
+    return mbedtls_ecdsa_write_signature((mbedtls_ecdsa_context *) pk->pk_ctx,
                                          md_alg, hash, hash_len,
                                          sig, sig_size, sig_len,
                                          f_rng, p_rng);
@@ -989,12 +993,12 @@
 
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
 /* Forward declarations */
-static int ecdsa_verify_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int ecdsa_verify_rs_wrap(mbedtls_pk_context *ctx, mbedtls_md_type_t md_alg,
                                 const unsigned char *hash, size_t hash_len,
                                 const unsigned char *sig, size_t sig_len,
                                 void *rs_ctx);
 
-static int ecdsa_sign_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int ecdsa_sign_rs_wrap(mbedtls_pk_context *ctx, 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 (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
@@ -1041,7 +1045,7 @@
     mbedtls_free(ctx);
 }
 
-static int eckey_verify_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int eckey_verify_rs_wrap(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg,
                                 const unsigned char *hash, size_t hash_len,
                                 const unsigned char *sig, size_t sig_len,
                                 void *rs_ctx)
@@ -1056,10 +1060,10 @@
 
     /* set up our own sub-context if needed (that is, on first run) */
     if (rs->ecdsa_ctx.grp.pbits == 0) {
-        MBEDTLS_MPI_CHK(mbedtls_ecdsa_from_keypair(&rs->ecdsa_ctx, ctx));
+        MBEDTLS_MPI_CHK(mbedtls_ecdsa_from_keypair(&rs->ecdsa_ctx, pk->pk_ctx));
     }
 
-    MBEDTLS_MPI_CHK(ecdsa_verify_rs_wrap(&rs->ecdsa_ctx,
+    MBEDTLS_MPI_CHK(ecdsa_verify_rs_wrap(pk,
                                          md_alg, hash, hash_len,
                                          sig, sig_len, &rs->ecdsa_rs));
 
@@ -1067,7 +1071,7 @@
     return ret;
 }
 
-static int eckey_sign_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int eckey_sign_rs_wrap(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
@@ -1083,10 +1087,10 @@
 
     /* set up our own sub-context if needed (that is, on first run) */
     if (rs->ecdsa_ctx.grp.pbits == 0) {
-        MBEDTLS_MPI_CHK(mbedtls_ecdsa_from_keypair(&rs->ecdsa_ctx, ctx));
+        MBEDTLS_MPI_CHK(mbedtls_ecdsa_from_keypair(&rs->ecdsa_ctx, pk->pk_ctx));
     }
 
-    MBEDTLS_MPI_CHK(ecdsa_sign_rs_wrap(&rs->ecdsa_ctx, md_alg,
+    MBEDTLS_MPI_CHK(ecdsa_sign_rs_wrap(pk, md_alg,
                                        hash, hash_len, sig, sig_size, sig_len,
                                        f_rng, p_rng, &rs->ecdsa_rs));
 
@@ -1104,13 +1108,12 @@
  * - write the raw content of public key "pub" to a local buffer
  * - compare the two buffers
  */
-static int eckey_check_pair_psa(const mbedtls_ecp_keypair *pub,
-                                const mbedtls_ecp_keypair *prv)
+static int eckey_check_pair_psa(mbedtls_pk_context *pub, mbedtls_pk_context *prv)
 {
     psa_status_t status, destruction_status;
     psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
-    mbedtls_ecp_keypair *prv_ctx = (mbedtls_ecp_keypair *) prv;
-    mbedtls_ecp_keypair *pub_ctx = (mbedtls_ecp_keypair *) pub;
+    mbedtls_ecp_keypair *prv_ctx = prv->pk_ctx;
+    mbedtls_ecp_keypair *pub_ctx = pub->pk_ctx;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     /* We are using MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH for the size of this
      * buffer because it will be used to hold the private key at first and
@@ -1167,7 +1170,7 @@
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
-static int eckey_check_pair(const void *pub, const void *prv,
+static int eckey_check_pair(mbedtls_pk_context *pub, mbedtls_pk_context *prv,
                             int (*f_rng)(void *, unsigned char *, size_t),
                             void *p_rng)
 {
@@ -1176,8 +1179,8 @@
     (void) p_rng;
     return eckey_check_pair_psa(pub, prv);
 #elif defined(MBEDTLS_ECP_C)
-    return mbedtls_ecp_check_pub_priv((const mbedtls_ecp_keypair *) pub,
-                                      (const mbedtls_ecp_keypair *) prv,
+    return mbedtls_ecp_check_pub_priv((const mbedtls_ecp_keypair *) pub->pk_ctx,
+                                      (const mbedtls_ecp_keypair *) prv->pk_ctx,
                                       f_rng, p_rng);
 #else
     return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
@@ -1201,11 +1204,12 @@
     mbedtls_free(ctx);
 }
 
-static void eckey_debug(const void *ctx, mbedtls_pk_debug_item *items)
+static void eckey_debug(mbedtls_pk_context *pk, mbedtls_pk_debug_item *items)
 {
+    mbedtls_ecp_keypair *ecp = (mbedtls_ecp_keypair *) pk->pk_ctx;
     items->type = MBEDTLS_PK_DEBUG_ECP;
     items->name = "eckey.Q";
-    items->value = &(((mbedtls_ecp_keypair *) ctx)->Q);
+    items->value = &(ecp->Q);
 }
 
 const mbedtls_pk_info_t mbedtls_eckey_info = {
@@ -1279,7 +1283,7 @@
 }
 
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
-static int ecdsa_verify_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int ecdsa_verify_rs_wrap(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg,
                                 const unsigned char *hash, size_t hash_len,
                                 const unsigned char *sig, size_t sig_len,
                                 void *rs_ctx)
@@ -1288,7 +1292,7 @@
     ((void) md_alg);
 
     ret = mbedtls_ecdsa_read_signature_restartable(
-        (mbedtls_ecdsa_context *) ctx,
+        (mbedtls_ecdsa_context *) pk->pk_ctx,
         hash, hash_len, sig, sig_len,
         (mbedtls_ecdsa_restart_ctx *) rs_ctx);
 
@@ -1299,14 +1303,14 @@
     return ret;
 }
 
-static int ecdsa_sign_rs_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int ecdsa_sign_rs_wrap(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t), void *p_rng,
                               void *rs_ctx)
 {
     return mbedtls_ecdsa_write_signature_restartable(
-        (mbedtls_ecdsa_context *) ctx,
+        (mbedtls_ecdsa_context *) pk->pk_ctx,
         md_alg, hash, hash_len, sig, sig_size, sig_len, f_rng, p_rng,
         (mbedtls_ecdsa_restart_ctx *) rs_ctx);
 
@@ -1372,19 +1376,19 @@
     return type == MBEDTLS_PK_RSA;
 }
 
-static size_t rsa_alt_get_bitlen(const void *ctx)
+static size_t rsa_alt_get_bitlen(mbedtls_pk_context *pk)
 {
-    const mbedtls_rsa_alt_context *rsa_alt = (const mbedtls_rsa_alt_context *) ctx;
+    const mbedtls_rsa_alt_context *rsa_alt = pk->pk_ctx;
 
     return 8 * rsa_alt->key_len_func(rsa_alt->key);
 }
 
-static int rsa_alt_sign_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int rsa_alt_sign_wrap(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_rsa_alt_context *rsa_alt = (mbedtls_rsa_alt_context *) ctx;
+    mbedtls_rsa_alt_context *rsa_alt = pk->pk_ctx;
 
     if (UINT_MAX < hash_len) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
@@ -1402,12 +1406,12 @@
                               md_alg, (unsigned int) hash_len, hash, sig);
 }
 
-static int rsa_alt_decrypt_wrap(void *ctx,
+static int rsa_alt_decrypt_wrap(mbedtls_pk_context *pk,
                                 const unsigned char *input, size_t ilen,
                                 unsigned char *output, size_t *olen, size_t osize,
                                 int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_rsa_alt_context *rsa_alt = (mbedtls_rsa_alt_context *) ctx;
+    mbedtls_rsa_alt_context *rsa_alt = pk->pk_ctx;
 
     ((void) f_rng);
     ((void) p_rng);
@@ -1421,7 +1425,7 @@
 }
 
 #if defined(MBEDTLS_RSA_C)
-static int rsa_alt_check_pair(const void *pub, const void *prv,
+static int rsa_alt_check_pair(mbedtls_pk_context *pub, mbedtls_pk_context *prv,
                               int (*f_rng)(void *, unsigned char *, size_t),
                               void *p_rng)
 {
@@ -1436,14 +1440,14 @@
 
     memset(hash, 0x2a, sizeof(hash));
 
-    if ((ret = rsa_alt_sign_wrap((void *) prv, MBEDTLS_MD_NONE,
+    if ((ret = rsa_alt_sign_wrap(prv, MBEDTLS_MD_NONE,
                                  hash, sizeof(hash),
                                  sig, sizeof(sig), &sig_len,
                                  f_rng, p_rng)) != 0) {
         return ret;
     }
 
-    if (rsa_verify_wrap((void *) pub, MBEDTLS_MD_NONE,
+    if (rsa_verify_wrap(pub, MBEDTLS_MD_NONE,
                         hash, sizeof(hash), sig, sig_len) != 0) {
         return MBEDTLS_ERR_RSA_KEY_CHECK_FAILED;
     }
@@ -1515,9 +1519,9 @@
     mbedtls_free(ctx);
 }
 
-static size_t pk_opaque_get_bitlen(const void *ctx)
+static size_t pk_opaque_get_bitlen(mbedtls_pk_context *pk)
 {
-    const mbedtls_svc_key_id_t *key = (const mbedtls_svc_key_id_t *) ctx;
+    const mbedtls_svc_key_id_t *key = pk->pk_ctx;
     size_t bits;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
@@ -1542,13 +1546,13 @@
            type == MBEDTLS_PK_RSASSA_PSS;
 }
 
-static int pk_opaque_sign_wrap(void *ctx, mbedtls_md_type_t md_alg,
+static int pk_opaque_sign_wrap(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
 #if !defined(MBEDTLS_PK_CAN_ECDSA_SIGN) && !defined(MBEDTLS_RSA_C)
-    ((void) ctx);
+    ((void) pk);
     ((void) md_alg);
     ((void) hash);
     ((void) hash_len);
@@ -1559,7 +1563,7 @@
     ((void) p_rng);
     return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
 #else /* !MBEDTLS_PK_CAN_ECDSA_SIGN && !MBEDTLS_RSA_C */
-    const mbedtls_svc_key_id_t *key = (const mbedtls_svc_key_id_t *) ctx;
+    const mbedtls_svc_key_id_t *key = pk->pk_ctx;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_algorithm_t alg;
     psa_key_type_t type;
@@ -1641,12 +1645,12 @@
 };
 
 #if defined(PSA_WANT_KEY_TYPE_RSA_KEY_PAIR)
-static int pk_opaque_rsa_decrypt(void *ctx,
+static int pk_opaque_rsa_decrypt(mbedtls_pk_context *pk,
                                  const unsigned char *input, size_t ilen,
                                  unsigned char *output, size_t *olen, size_t osize,
                                  int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    const mbedtls_svc_key_id_t *key = (const mbedtls_svc_key_id_t *) ctx;
+    const mbedtls_svc_key_id_t *key = pk->pk_ctx;
     psa_status_t status;
 
     /* PSA has its own RNG */
diff --git a/library/pk_wrap.h b/library/pk_wrap.h
index 89bce40..91d240c 100644
--- a/library/pk_wrap.h
+++ b/library/pk_wrap.h
@@ -39,18 +39,18 @@
     const char *name;
 
     /** Get key size in bits */
-    size_t (*get_bitlen)(const void *);
+    size_t (*get_bitlen)(mbedtls_pk_context *pk);
 
     /** Tell if the context implements this type (e.g. ECKEY can do ECDSA) */
     int (*can_do)(mbedtls_pk_type_t type);
 
     /** Verify signature */
-    int (*verify_func)(void *ctx, mbedtls_md_type_t md_alg,
+    int (*verify_func)(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg,
                        const unsigned char *hash, size_t hash_len,
                        const unsigned char *sig, size_t sig_len);
 
     /** Make signature */
-    int (*sign_func)(void *ctx, mbedtls_md_type_t md_alg,
+    int (*sign_func)(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t),
@@ -58,13 +58,13 @@
 
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
     /** Verify signature (restartable) */
-    int (*verify_rs_func)(void *ctx, mbedtls_md_type_t md_alg,
+    int (*verify_rs_func)(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg,
                           const unsigned char *hash, size_t hash_len,
                           const unsigned char *sig, size_t sig_len,
                           void *rs_ctx);
 
     /** Make signature (restartable) */
-    int (*sign_rs_func)(void *ctx, mbedtls_md_type_t md_alg,
+    int (*sign_rs_func)(mbedtls_pk_context *pk, 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 (*f_rng)(void *, unsigned char *, size_t),
@@ -72,19 +72,19 @@
 #endif /* MBEDTLS_ECDSA_C && MBEDTLS_ECP_RESTARTABLE */
 
     /** Decrypt message */
-    int (*decrypt_func)(void *ctx, const unsigned char *input, size_t ilen,
+    int (*decrypt_func)(mbedtls_pk_context *pk, const unsigned char *input, size_t ilen,
                         unsigned char *output, size_t *olen, size_t osize,
                         int (*f_rng)(void *, unsigned char *, size_t),
                         void *p_rng);
 
     /** Encrypt message */
-    int (*encrypt_func)(void *ctx, const unsigned char *input, size_t ilen,
+    int (*encrypt_func)(mbedtls_pk_context *pk, const unsigned char *input, size_t ilen,
                         unsigned char *output, size_t *olen, size_t osize,
                         int (*f_rng)(void *, unsigned char *, size_t),
                         void *p_rng);
 
     /** Check public-private key pair */
-    int (*check_pair_func)(const void *pub, const void *prv,
+    int (*check_pair_func)(mbedtls_pk_context *pub, mbedtls_pk_context *prv,
                            int (*f_rng)(void *, unsigned char *, size_t),
                            void *p_rng);
 
@@ -103,7 +103,7 @@
 #endif /* MBEDTLS_ECDSA_C && MBEDTLS_ECP_RESTARTABLE */
 
     /** Interface with the debug module */
-    void (*debug_func)(const void *ctx, mbedtls_pk_debug_item *items);
+    void (*debug_func)(mbedtls_pk_context *pk, mbedtls_pk_debug_item *items);
 
 };
 #if defined(MBEDTLS_PK_RSA_ALT_SUPPORT)
diff --git a/programs/fuzz/fuzz_client.c b/programs/fuzz/fuzz_client.c
index 56a5efe..d4e1d74 100644
--- a/programs/fuzz/fuzz_client.c
+++ b/programs/fuzz/fuzz_client.c
@@ -78,6 +78,13 @@
     mbedtls_ctr_drbg_init(&ctr_drbg);
     mbedtls_entropy_init(&entropy);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, dummy_entropy, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
         goto exit;
@@ -175,6 +182,9 @@
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ssl_free(&ssl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_dtlsclient.c b/programs/fuzz/fuzz_dtlsclient.c
index a58f6f4..3659026 100644
--- a/programs/fuzz/fuzz_dtlsclient.c
+++ b/programs/fuzz/fuzz_dtlsclient.c
@@ -61,6 +61,13 @@
     mbedtls_ctr_drbg_init(&ctr_drbg);
     mbedtls_entropy_init(&entropy);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     srand(1);
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, dummy_entropy, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
@@ -119,6 +126,9 @@
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ssl_free(&ssl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_dtlsserver.c b/programs/fuzz/fuzz_dtlsserver.c
index cdd69c0..1632e9d 100644
--- a/programs/fuzz/fuzz_dtlsserver.c
+++ b/programs/fuzz/fuzz_dtlsserver.c
@@ -50,6 +50,20 @@
 
     mbedtls_ctr_drbg_init(&ctr_drbg);
     mbedtls_entropy_init(&entropy);
+#if defined(MBEDTLS_X509_CRT_PARSE_C) && defined(MBEDTLS_PEM_PARSE_C)
+    mbedtls_x509_crt_init(&srvcert);
+    mbedtls_pk_init(&pkey);
+#endif
+    mbedtls_ssl_init(&ssl);
+    mbedtls_ssl_config_init(&conf);
+    mbedtls_ssl_cookie_init(&cookie_ctx);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, dummy_entropy, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
@@ -58,8 +72,7 @@
 
     if (initialized == 0) {
 #if defined(MBEDTLS_X509_CRT_PARSE_C) && defined(MBEDTLS_PEM_PARSE_C)
-        mbedtls_x509_crt_init(&srvcert);
-        mbedtls_pk_init(&pkey);
+
         if (mbedtls_x509_crt_parse(&srvcert, (const unsigned char *) mbedtls_test_srv_crt,
                                    mbedtls_test_srv_crt_len) != 0) {
             return 1;
@@ -78,9 +91,6 @@
 
         initialized = 1;
     }
-    mbedtls_ssl_init(&ssl);
-    mbedtls_ssl_config_init(&conf);
-    mbedtls_ssl_cookie_init(&cookie_ctx);
 
     if (mbedtls_ssl_config_defaults(&conf,
                                     MBEDTLS_SSL_IS_SERVER,
@@ -154,9 +164,16 @@
 exit:
     mbedtls_ssl_cookie_free(&cookie_ctx);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_X509_CRT_PARSE_C) && defined(MBEDTLS_PEM_PARSE_C)
+    mbedtls_pk_free(&pkey);
+    mbedtls_x509_crt_free(&srvcert);
+#endif
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ssl_free(&ssl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_privkey.c b/programs/fuzz/fuzz_privkey.c
index 39c23e2..ce75624 100644
--- a/programs/fuzz/fuzz_privkey.c
+++ b/programs/fuzz/fuzz_privkey.c
@@ -30,13 +30,20 @@
 
     mbedtls_ctr_drbg_init(&ctr_drbg);
     mbedtls_entropy_init(&entropy);
+    mbedtls_pk_init(&pk);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, dummy_entropy, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
-        return 1;
+        goto exit;
     }
 
-    mbedtls_pk_init(&pk);
     ret = mbedtls_pk_parse_key(&pk, Data, Size, NULL, 0,
                                dummy_random, &ctr_drbg);
     if (ret == 0) {
@@ -83,7 +90,13 @@
             abort();
         }
     }
+exit:
+    mbedtls_entropy_free(&entropy);
+    mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_pk_free(&pk);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 #else
     (void) Data;
     (void) Size;
diff --git a/programs/fuzz/fuzz_pubkey.c b/programs/fuzz/fuzz_pubkey.c
index 7f5e4aa..9203b4e 100644
--- a/programs/fuzz/fuzz_pubkey.c
+++ b/programs/fuzz/fuzz_pubkey.c
@@ -11,6 +11,12 @@
     mbedtls_pk_context pk;
 
     mbedtls_pk_init(&pk);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     ret = mbedtls_pk_parse_public_key(&pk, Data, Size);
     if (ret == 0) {
 #if defined(MBEDTLS_RSA_C)
@@ -66,6 +72,10 @@
             abort();
         }
     }
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+exit:
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_pk_free(&pk);
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_server.c b/programs/fuzz/fuzz_server.c
index cd021e1..e767859 100644
--- a/programs/fuzz/fuzz_server.c
+++ b/programs/fuzz/fuzz_server.c
@@ -58,6 +58,21 @@
 
     mbedtls_ctr_drbg_init(&ctr_drbg);
     mbedtls_entropy_init(&entropy);
+#if defined(MBEDTLS_X509_CRT_PARSE_C) && defined(MBEDTLS_PEM_PARSE_C)
+    mbedtls_x509_crt_init(&srvcert);
+    mbedtls_pk_init(&pkey);
+#endif
+    mbedtls_ssl_init(&ssl);
+    mbedtls_ssl_config_init(&conf);
+#if defined(MBEDTLS_SSL_SESSION_TICKETS) && defined(MBEDTLS_SSL_TICKET_C)
+    mbedtls_ssl_ticket_init(&ticket_ctx);
+#endif
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, dummy_entropy, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
@@ -67,8 +82,6 @@
     if (initialized == 0) {
 
 #if defined(MBEDTLS_X509_CRT_PARSE_C) && defined(MBEDTLS_PEM_PARSE_C)
-        mbedtls_x509_crt_init(&srvcert);
-        mbedtls_pk_init(&pkey);
         if (mbedtls_x509_crt_parse(&srvcert, (const unsigned char *) mbedtls_test_srv_crt,
                                    mbedtls_test_srv_crt_len) != 0) {
             return 1;
@@ -92,11 +105,6 @@
 
         initialized = 1;
     }
-    mbedtls_ssl_init(&ssl);
-    mbedtls_ssl_config_init(&conf);
-#if defined(MBEDTLS_SSL_SESSION_TICKETS) && defined(MBEDTLS_SSL_TICKET_C)
-    mbedtls_ssl_ticket_init(&ticket_ctx);
-#endif
 
     if (mbedtls_ssl_config_defaults(&conf,
                                     MBEDTLS_SSL_IS_SERVER,
@@ -193,8 +201,14 @@
     mbedtls_entropy_free(&entropy);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_ssl_config_free(&conf);
+#if defined(MBEDTLS_X509_CRT_PARSE_C) && defined(MBEDTLS_PEM_PARSE_C)
+    mbedtls_x509_crt_free(&srvcert);
+    mbedtls_pk_free(&pkey);
+#endif
     mbedtls_ssl_free(&ssl);
-
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif
 #else
     (void) Data;
     (void) Size;
diff --git a/programs/fuzz/fuzz_x509crl.c b/programs/fuzz/fuzz_x509crl.c
index 6ff0c05..313540d 100644
--- a/programs/fuzz/fuzz_x509crl.c
+++ b/programs/fuzz/fuzz_x509crl.c
@@ -11,6 +11,12 @@
     unsigned char buf[4096];
 
     mbedtls_x509_crl_init(&crl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     ret = mbedtls_x509_crl_parse(&crl, Data, Size);
 #if !defined(MBEDTLS_X509_REMOVE_INFO)
     if (ret == 0) {
@@ -20,6 +26,11 @@
     ((void) ret);
     ((void) buf);
 #endif /* !MBEDTLS_X509_REMOVE_INFO */
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+exit:
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_x509_crl_free(&crl);
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_x509crt.c b/programs/fuzz/fuzz_x509crt.c
index 858c1ff..8442090 100644
--- a/programs/fuzz/fuzz_x509crt.c
+++ b/programs/fuzz/fuzz_x509crt.c
@@ -11,6 +11,12 @@
     unsigned char buf[4096];
 
     mbedtls_x509_crt_init(&crt);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     ret = mbedtls_x509_crt_parse(&crt, Data, Size);
 #if !defined(MBEDTLS_X509_REMOVE_INFO)
     if (ret == 0) {
@@ -20,6 +26,11 @@
     ((void) ret);
     ((void) buf);
 #endif /* !MBEDTLS_X509_REMOVE_INFO */
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+exit:
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_x509_crt_free(&crt);
 #else
     (void) Data;
diff --git a/programs/fuzz/fuzz_x509csr.c b/programs/fuzz/fuzz_x509csr.c
index 39fb4cb..395d3c2 100644
--- a/programs/fuzz/fuzz_x509csr.c
+++ b/programs/fuzz/fuzz_x509csr.c
@@ -11,6 +11,12 @@
     unsigned char buf[4096];
 
     mbedtls_x509_csr_init(&csr);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     ret = mbedtls_x509_csr_parse(&csr, Data, Size);
 #if !defined(MBEDTLS_X509_REMOVE_INFO)
     if (ret == 0) {
@@ -20,6 +26,11 @@
     ((void) ret);
     ((void) buf);
 #endif /* !MBEDTLS_X509_REMOVE_INFO */
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+exit:
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_x509_csr_free(&csr);
 #else
     (void) Data;
diff --git a/programs/pkey/gen_key.c b/programs/pkey/gen_key.c
index 029558d..9bee275 100644
--- a/programs/pkey/gen_key.c
+++ b/programs/pkey/gen_key.c
@@ -200,6 +200,15 @@
     mbedtls_ctr_drbg_init(&ctr_drbg);
     memset(buf, 0, sizeof(buf));
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -407,6 +416,9 @@
     mbedtls_pk_free(&key);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/pkey/key_app.c b/programs/pkey/key_app.c
index c80dcd0..cd16e33 100644
--- a/programs/pkey/key_app.c
+++ b/programs/pkey/key_app.c
@@ -99,6 +99,15 @@
     mbedtls_pk_init(&pk);
     memset(buf, 0, sizeof(buf));
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto cleanup;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P); mbedtls_mpi_init(&Q);
     mbedtls_mpi_init(&D); mbedtls_mpi_init(&E); mbedtls_mpi_init(&DP);
     mbedtls_mpi_init(&DQ); mbedtls_mpi_init(&QP);
@@ -305,8 +314,10 @@
 
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
-
     mbedtls_pk_free(&pk);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_mpi_free(&N); mbedtls_mpi_free(&P); mbedtls_mpi_free(&Q);
     mbedtls_mpi_free(&D); mbedtls_mpi_free(&E); mbedtls_mpi_free(&DP);
     mbedtls_mpi_free(&DQ); mbedtls_mpi_free(&QP);
diff --git a/programs/pkey/key_app_writer.c b/programs/pkey/key_app_writer.c
index 862c93f..e8f3e85 100644
--- a/programs/pkey/key_app_writer.c
+++ b/programs/pkey/key_app_writer.c
@@ -216,6 +216,15 @@
     memset(buf, 0, sizeof(buf));
 #endif
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     mbedtls_mpi_init(&N); mbedtls_mpi_init(&P); mbedtls_mpi_init(&Q);
     mbedtls_mpi_init(&D); mbedtls_mpi_init(&E); mbedtls_mpi_init(&DP);
     mbedtls_mpi_init(&DQ); mbedtls_mpi_init(&QP);
@@ -422,6 +431,9 @@
 
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/pkey/pk_decrypt.c b/programs/pkey/pk_decrypt.c
index 8862636..f60c946 100644
--- a/programs/pkey/pk_decrypt.c
+++ b/programs/pkey/pk_decrypt.c
@@ -67,6 +67,15 @@
 
     memset(result, 0, sizeof(result));
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 2) {
         mbedtls_printf("usage: mbedtls_pk_decrypt <key_file>\n");
 
@@ -139,6 +148,9 @@
     mbedtls_pk_free(&pk);
     mbedtls_entropy_free(&entropy);
     mbedtls_ctr_drbg_free(&ctr_drbg);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_ERROR_C)
     if (exit_code != MBEDTLS_EXIT_SUCCESS) {
diff --git a/programs/pkey/pk_encrypt.c b/programs/pkey/pk_encrypt.c
index eab3f08..04e5cc7 100644
--- a/programs/pkey/pk_encrypt.c
+++ b/programs/pkey/pk_encrypt.c
@@ -63,6 +63,15 @@
     mbedtls_entropy_init(&entropy);
     mbedtls_pk_init(&pk);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: mbedtls_pk_encrypt <key_file> <string of max 100 characters>\n");
 
@@ -140,6 +149,9 @@
     mbedtls_pk_free(&pk);
     mbedtls_entropy_free(&entropy);
     mbedtls_ctr_drbg_free(&ctr_drbg);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_ERROR_C)
     if (exit_code != MBEDTLS_EXIT_SUCCESS) {
diff --git a/programs/pkey/pk_sign.c b/programs/pkey/pk_sign.c
index 82cb6a1..57bd796 100644
--- a/programs/pkey/pk_sign.c
+++ b/programs/pkey/pk_sign.c
@@ -63,6 +63,15 @@
     mbedtls_ctr_drbg_init(&ctr_drbg);
     mbedtls_pk_init(&pk);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: mbedtls_pk_sign <key_file> <filename>\n");
 
@@ -140,6 +149,9 @@
     mbedtls_pk_free(&pk);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_ERROR_C)
     if (exit_code != MBEDTLS_EXIT_SUCCESS) {
diff --git a/programs/pkey/pk_verify.c b/programs/pkey/pk_verify.c
index 0c549e0..bca985b 100644
--- a/programs/pkey/pk_verify.c
+++ b/programs/pkey/pk_verify.c
@@ -55,6 +55,15 @@
 
     mbedtls_pk_init(&pk);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: mbedtls_pk_verify <key_file> <filename>\n");
 
@@ -114,6 +123,9 @@
 
 exit:
     mbedtls_pk_free(&pk);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #if defined(MBEDTLS_ERROR_C)
     if (exit_code != MBEDTLS_EXIT_SUCCESS) {
diff --git a/programs/pkey/rsa_sign_pss.c b/programs/pkey/rsa_sign_pss.c
index 03882cd..999669e 100644
--- a/programs/pkey/rsa_sign_pss.c
+++ b/programs/pkey/rsa_sign_pss.c
@@ -64,6 +64,15 @@
     mbedtls_pk_init(&pk);
     mbedtls_ctr_drbg_init(&ctr_drbg);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: rsa_sign_pss <key_file> <filename>\n");
 
@@ -153,6 +162,9 @@
     mbedtls_pk_free(&pk);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/pkey/rsa_verify_pss.c b/programs/pkey/rsa_verify_pss.c
index e21e927..8a1fb59 100644
--- a/programs/pkey/rsa_verify_pss.c
+++ b/programs/pkey/rsa_verify_pss.c
@@ -58,6 +58,15 @@
 
     mbedtls_pk_init(&pk);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc != 3) {
         mbedtls_printf("usage: rsa_verify_pss <key_file> <filename>\n");
 
@@ -129,6 +138,9 @@
 
 exit:
     mbedtls_pk_free(&pk);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/ssl/dtls_client.c b/programs/ssl/dtls_client.c
index 44a135f..e47715c 100644
--- a/programs/ssl/dtls_client.c
+++ b/programs/ssl/dtls_client.c
@@ -109,11 +109,21 @@
     mbedtls_ssl_config_init(&conf);
     mbedtls_x509_crt_init(&cacert);
     mbedtls_ctr_drbg_init(&ctr_drbg);
+    mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_printf("\n  . Seeding the random number generator...");
     fflush(stdout);
 
-    mbedtls_entropy_init(&entropy);
     if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                      (const unsigned char *) pers,
                                      strlen(pers))) != 0) {
@@ -324,12 +334,14 @@
 #endif
 
     mbedtls_net_free(&server_fd);
-
     mbedtls_x509_crt_free(&cacert);
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     /* Shell can not handle large exit numbers -> 1 for errors */
     if (ret < 0) {
diff --git a/programs/ssl/dtls_server.c b/programs/ssl/dtls_server.c
index 6f8c841..f218130 100644
--- a/programs/ssl/dtls_server.c
+++ b/programs/ssl/dtls_server.c
@@ -119,6 +119,16 @@
     mbedtls_entropy_init(&entropy);
     mbedtls_ctr_drbg_init(&ctr_drbg);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_debug_set_threshold(DEBUG_LEVEL);
 #endif
@@ -394,6 +404,9 @@
 #endif
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     /* Shell can not handle large exit numbers -> 1 for errors */
     if (ret < 0) {
diff --git a/programs/ssl/mini_client.c b/programs/ssl/mini_client.c
index 6dbbc6d..e8f4797 100644
--- a/programs/ssl/mini_client.c
+++ b/programs/ssl/mini_client.c
@@ -175,8 +175,16 @@
 #if defined(MBEDTLS_X509_CRT_PARSE_C)
     mbedtls_x509_crt_init(&ca);
 #endif
-
     mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                               (const unsigned char *) pers, strlen(pers)) != 0) {
         ret = ctr_drbg_seed_failed;
@@ -262,7 +270,6 @@
 
 exit:
     mbedtls_net_free(&server_fd);
-
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
@@ -270,6 +277,9 @@
 #if defined(MBEDTLS_X509_CRT_PARSE_C)
     mbedtls_x509_crt_free(&ca);
 #endif
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(ret);
 }
diff --git a/programs/ssl/ssl_client1.c b/programs/ssl/ssl_client1.c
index ea96a4d..259b8f9 100644
--- a/programs/ssl/ssl_client1.c
+++ b/programs/ssl/ssl_client1.c
@@ -91,11 +91,21 @@
     mbedtls_ssl_config_init(&conf);
     mbedtls_x509_crt_init(&cacert);
     mbedtls_ctr_drbg_init(&ctr_drbg);
+    mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_printf("\n  . Seeding the random number generator...");
     fflush(stdout);
 
-    mbedtls_entropy_init(&entropy);
+
     if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                      (const unsigned char *) pers,
                                      strlen(pers))) != 0) {
@@ -274,12 +284,14 @@
 #endif
 
     mbedtls_net_free(&server_fd);
-
     mbedtls_x509_crt_free(&cacert);
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/ssl/ssl_context_info.c b/programs/ssl/ssl_context_info.c
index 0ba0d2c..a5f0650 100644
--- a/programs/ssl/ssl_context_info.c
+++ b/programs/ssl/ssl_context_info.c
@@ -21,6 +21,7 @@
 
 #include "mbedtls/build_info.h"
 #include "mbedtls/debug.h"
+#include "mbedtls/platform.h"
 
 #include <stdio.h>
 #include <stdlib.h>
@@ -933,6 +934,15 @@
     size_t ssl_max_len = SSL_INIT_LEN;
     size_t ssl_len = 0;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        return MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     /* The 'b64_file' is opened when parsing arguments to check that the
      * file name is correct */
     parse_arguments(argc, argv);
@@ -1001,6 +1011,10 @@
         printf("Finished. No valid base64 code found\n");
     }
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     return 0;
 }
 
diff --git a/programs/ssl/ssl_fork_server.c b/programs/ssl/ssl_fork_server.c
index 7ee880d..4777ee0 100644
--- a/programs/ssl/ssl_fork_server.c
+++ b/programs/ssl/ssl_fork_server.c
@@ -105,6 +105,15 @@
     mbedtls_x509_crt_init(&srvcert);
     mbedtls_ctr_drbg_init(&ctr_drbg);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     signal(SIGCHLD, SIG_IGN);
 
     /*
@@ -366,13 +375,15 @@
 exit:
     mbedtls_net_free(&client_fd);
     mbedtls_net_free(&listen_fd);
-
     mbedtls_x509_crt_free(&srvcert);
     mbedtls_pk_free(&pkey);
     mbedtls_ssl_free(&ssl);
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/ssl/ssl_mail_client.c b/programs/ssl/ssl_mail_client.c
index 3b040aa..fb6f371 100644
--- a/programs/ssl/ssl_mail_client.c
+++ b/programs/ssl/ssl_mail_client.c
@@ -369,6 +369,16 @@
     mbedtls_x509_crt_init(&clicert);
     mbedtls_pk_init(&pkey);
     mbedtls_ctr_drbg_init(&ctr_drbg);
+    mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     if (argc < 2) {
 usage:
@@ -458,7 +468,6 @@
     mbedtls_printf("\n  . Seeding the random number generator...");
     fflush(stdout);
 
-    mbedtls_entropy_init(&entropy);
     if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                      (const unsigned char *) pers,
                                      strlen(pers))) != 0) {
@@ -796,6 +805,9 @@
     mbedtls_ssl_config_free(&conf);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/ssl/ssl_pthread_server.c b/programs/ssl/ssl_pthread_server.c
index 2b3baff..9416c3c 100644
--- a/programs/ssl/ssl_pthread_server.c
+++ b/programs/ssl/ssl_pthread_server.c
@@ -332,6 +332,16 @@
      */
     mbedtls_entropy_init(&entropy);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     /*
      * 1a. Seed the random number generator
      */
@@ -473,14 +483,14 @@
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
     mbedtls_ssl_config_free(&conf);
-
     mbedtls_net_free(&listen_fd);
-
     mbedtls_mutex_free(&debug_mutex);
-
 #if defined(MBEDTLS_MEMORY_BUFFER_ALLOC_C)
     mbedtls_memory_buffer_alloc_free();
 #endif
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(ret);
 }
diff --git a/programs/ssl/ssl_server.c b/programs/ssl/ssl_server.c
index 7dabda8..bb49155 100644
--- a/programs/ssl/ssl_server.c
+++ b/programs/ssl/ssl_server.c
@@ -104,6 +104,16 @@
     mbedtls_entropy_init(&entropy);
     mbedtls_ctr_drbg_init(&ctr_drbg);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        ret = MBEDTLS_ERR_SSL_HW_ACCEL_FAILED;
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_debug_set_threshold(DEBUG_LEVEL);
 #endif
@@ -343,7 +353,6 @@
 
     mbedtls_net_free(&client_fd);
     mbedtls_net_free(&listen_fd);
-
     mbedtls_x509_crt_free(&srvcert);
     mbedtls_pk_free(&pkey);
     mbedtls_ssl_free(&ssl);
@@ -353,6 +362,9 @@
 #endif
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(ret);
 }
diff --git a/programs/x509/cert_app.c b/programs/x509/cert_app.c
index a9656c6..51a79ec 100644
--- a/programs/x509/cert_app.c
+++ b/programs/x509/cert_app.c
@@ -153,6 +153,7 @@
     mbedtls_ssl_init(&ssl);
     mbedtls_ssl_config_init(&conf);
     mbedtls_x509_crt_init(&cacert);
+    mbedtls_entropy_init(&entropy);
 #if defined(MBEDTLS_X509_CRL_PARSE_C)
     mbedtls_x509_crl_init(&cacrl);
 #else
@@ -161,6 +162,15 @@
     memset(&cacrl, 0, sizeof(mbedtls_x509_crl));
 #endif
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -338,7 +348,6 @@
         mbedtls_printf("\n  . Seeding the random number generator...");
         fflush(stdout);
 
-        mbedtls_entropy_init(&entropy);
         if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                          (const unsigned char *) pers,
                                          strlen(pers))) != 0) {
@@ -448,6 +457,9 @@
 #endif
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/x509/cert_req.c b/programs/x509/cert_req.c
index 396aaf3..1772f87 100644
--- a/programs/x509/cert_req.c
+++ b/programs/x509/cert_req.c
@@ -179,6 +179,16 @@
     mbedtls_pk_init(&key);
     mbedtls_ctr_drbg_init(&ctr_drbg);
     memset(buf, 0, sizeof(buf));
+    mbedtls_entropy_init(&entropy);
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     if (argc < 2) {
 usage:
@@ -388,7 +398,6 @@
     mbedtls_printf("  . Seeding the random number generator...");
     fflush(stdout);
 
-    mbedtls_entropy_init(&entropy);
     if ((ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                      (const unsigned char *) pers,
                                      strlen(pers))) != 0) {
@@ -460,6 +469,9 @@
     mbedtls_pk_free(&key);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     cur = opt.san_list;
     while (cur != NULL) {
diff --git a/programs/x509/cert_write.c b/programs/x509/cert_write.c
index a822684..51b09f3 100644
--- a/programs/x509/cert_write.c
+++ b/programs/x509/cert_write.c
@@ -330,6 +330,15 @@
     memset(buf, 0, sizeof(buf));
     memset(serial, 0, sizeof(serial));
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -885,6 +894,9 @@
     mbedtls_pk_free(&loaded_issuer_key);
     mbedtls_ctr_drbg_free(&ctr_drbg);
     mbedtls_entropy_free(&entropy);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/x509/crl_app.c b/programs/x509/crl_app.c
index d74a488..6c671ff 100644
--- a/programs/x509/crl_app.c
+++ b/programs/x509/crl_app.c
@@ -70,6 +70,15 @@
      */
     mbedtls_x509_crl_init(&crl);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -125,6 +134,9 @@
 
 exit:
     mbedtls_x509_crl_free(&crl);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/programs/x509/load_roots.c b/programs/x509/load_roots.c
index 237bd7c..d024e98 100644
--- a/programs/x509/load_roots.c
+++ b/programs/x509/load_roots.c
@@ -123,6 +123,15 @@
     struct mbedtls_timing_hr_time timer;
     unsigned long ms;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc <= 1) {
         mbedtls_printf(USAGE);
         goto exit;
@@ -187,6 +196,9 @@
     exit_code = MBEDTLS_EXIT_SUCCESS;
 
 exit:
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     mbedtls_exit(exit_code);
 }
 #endif /* necessary configuration */
diff --git a/programs/x509/req_app.c b/programs/x509/req_app.c
index 83e2546..64b9f0b 100644
--- a/programs/x509/req_app.c
+++ b/programs/x509/req_app.c
@@ -70,6 +70,15 @@
      */
     mbedtls_x509_csr_init(&csr);
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_fprintf(stderr, "Failed to initialize PSA Crypto implementation: %d\n",
+                        (int) status);
+        goto exit;
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
     if (argc < 2) {
 usage:
         mbedtls_printf(USAGE);
@@ -125,6 +134,9 @@
 
 exit:
     mbedtls_x509_csr_free(&csr);
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    mbedtls_psa_crypto_free();
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     mbedtls_exit(exit_code);
 }
diff --git a/scripts/mbedtls_dev/ecp.py b/scripts/mbedtls_dev/ecp.py
index 94ecdfe..b7b66e4 100644
--- a/scripts/mbedtls_dev/ecp.py
+++ b/scripts/mbedtls_dev/ecp.py
@@ -617,3 +617,68 @@
     @property
     def is_valid(self) -> bool:
         return True
+
+
+class EcpP256K1Raw(bignum_common.ModOperationCommon,
+                   EcpTarget):
+    """Test cases for ECP P256 fast reduction."""
+    symbol = "-"
+    test_function = "ecp_mod_p256k1"
+    test_name = "ecp_mod_p256k1"
+    input_style = "fixed"
+    arity = 1
+    dependencies = ["MBEDTLS_ECP_DP_SECP256K1_ENABLED"]
+
+    moduli = ["fffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f"] # type: List[str]
+
+    input_values = [
+        "0", "1",
+
+        # Modulus - 1
+        "fffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2e",
+
+        # Modulus + 1
+        "fffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc30",
+
+        # 2^256 - 1
+        "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff",
+
+        # Maximum canonical P256 multiplication result
+        ("fffffffffffffffffffffffffffffffffffffffffffffffffffffffdfffff85c0"
+         "00000000000000000000000000000000000000000000001000007a4000e9844"),
+
+        # First 8 number generated by random.getrandbits(512) - seed(2,2)
+        ("4067c3584ee207f8da94e3e8ab73738fcf1822ffbc6887782b491044d5e34124"
+         "5c6e433715ba2bdd177219d30e7a269fd95bafc8f2a4d27bdcf4bb99f4bea973"),
+        ("82523e86feac7eb7dc38f519b91751dacdbd47d364be8049a372db8f6e405d93"
+         "ffed9235288bc781ae66267594c9c9500925e4749b575bd13653f8dd9b1f282e"),
+        ("e8624fab5186ee32ee8d7ee9770348a05d300cb90706a045defc044a09325626"
+         "e6b58de744ab6cce80877b6f71e1f6d2ef8acd128b4f2fc15f3f57ebf30b94fa"),
+        ("829a48d422fe99a22c70501e533c91352d3d854e061b90303b08c6e33c729578"
+         "2d6c797f8f7d9b782a1be9cd8697bbd0e2520e33e44c50556c71c4a66148a86f"),
+        ("e89204e2e8168561867e5e15bc01bfce6a27e0dfcbf8754472154e76e4c11ab2"
+         "fec3f6b32e8d4b8a8f54f8ceacaab39e83844b40ffa9b9f15c14bc4a829e07b0"),
+        ("bd143fa9b714210c665d7435c1066932f4767f26294365b2721dea3bf63f23d0"
+         "dbe53fcafb2147df5ca495fa5a91c89b97eeab64ca2ce6bc5d3fd983c34c769f"),
+        ("74667bffe202849da9643a295a9ac6decbd4d3e2d4dec9ef83f0be4e80371eb9"
+         "7f81375eecc1cb6347733e847d718d733ff98ff387c56473a7a83ee0761ebfd2"),
+        ("d08f1bb2531d6460f0caeef038c89b38a8acb5137c9260dc74e088a9b9492f25"
+         "8ebdbfe3eb9ac688b9d39cca91551e8259cc60b17604e4b4e73695c3e652c71a"),
+
+        # Next 2 number generated by random.getrandbits(256)
+        ("c5e2486c44a4a8f69dc8db48e86ec9c6e06f291b2a838af8d5c44a4eb3172062"),
+        ("d4c0dca8b4c9e755cc9c3adcf515a8234da4daeb4f3f87777ad1f45ae9500ec9"),
+    ]
+
+    @property
+    def arg_a(self) -> str:
+        hex_digits = bignum_common.hex_digits_for_limb(448 // self.bits_in_limb, self.bits_in_limb)
+        return super().format_arg('{:x}'.format(self.int_a)).zfill(hex_digits)
+
+    def result(self) -> List[str]:
+        result = self.int_a % self.int_n
+        return [self.format_result(result)]
+
+    @property
+    def is_valid(self) -> bool:
+        return True
diff --git a/tests/suites/test_suite_bignum_core.function b/tests/suites/test_suite_bignum_core.function
index e084b83..53aa002 100644
--- a/tests/suites/test_suite_bignum_core.function
+++ b/tests/suites/test_suite_bignum_core.function
@@ -309,6 +309,36 @@
 }
 /* END_CASE */
 
+
+/* BEGIN_CASE */
+void mpi_core_clz(int leading_zeros, int trailing_zeros)
+{
+    if ((size_t) (leading_zeros + trailing_zeros) >= (sizeof(mbedtls_mpi_uint) * 8)) {
+        // can't fit required number of leading and trailing zeros - skip test
+        goto exit;
+    }
+
+    // Construct a test input value where the count of leading zeros and
+    // trailing zeros is given in the test case, and we add ones to fill
+    // the gap.
+    mbedtls_mpi_uint x;
+    if ((leading_zeros + trailing_zeros) > 0) {
+        // some zero bits
+        uint32_t s = (sizeof(mbedtls_mpi_uint) * 8 - leading_zeros - trailing_zeros);
+        x = ((((mbedtls_mpi_uint) 1) << s) - 1) << trailing_zeros;
+    } else {
+        // all bits set
+        x = ~((mbedtls_mpi_uint) 0);
+    }
+
+    size_t n = mbedtls_mpi_core_clz(x);
+    TEST_EQUAL(n, leading_zeros);
+exit:
+    ;
+}
+/* END_CASE */
+
+
 /* BEGIN_CASE */
 void mpi_core_lt_ct(char *input_X, char *input_Y, int exp_ret)
 {
diff --git a/tests/suites/test_suite_bignum_core.misc.data b/tests/suites/test_suite_bignum_core.misc.data
index b61d708..ba86029 100644
--- a/tests/suites/test_suite_bignum_core.misc.data
+++ b/tests/suites/test_suite_bignum_core.misc.data
@@ -491,3 +491,35 @@
 Fill random core: 42 bytes, 5 missing limbs
 mpi_core_fill_random:42:0:-5:0:MBEDTLS_ERR_MPI_BAD_INPUT_DATA
 
+CLZ: 0 0: all ones
+mpi_core_clz:0:0
+
+CLZ: 1 0
+mpi_core_clz:1:0
+
+CLZ: 1 1
+mpi_core_clz:1:1
+
+CLZ: 4 5
+mpi_core_clz:4:5
+
+CLZ: 8 16
+mpi_core_clz:8:16
+
+CLZ: 31 0
+mpi_core_clz:31:0
+
+CLZ: 32 0
+mpi_core_clz:32:0
+
+CLZ: 33 0
+mpi_core_clz:33:0
+
+CLZ: 63 0
+mpi_core_clz:63:0
+
+CLZ: 64 0
+mpi_core_clz:64:0
+
+CLZ: 100000 0: skip overly long input
+mpi_core_clz:100000:0
diff --git a/tests/suites/test_suite_ecp.function b/tests/suites/test_suite_ecp.function
index 898240e..6d5ce9c 100644
--- a/tests/suites/test_suite_ecp.function
+++ b/tests/suites/test_suite_ecp.function
@@ -1433,6 +1433,45 @@
 }
 /* END_CASE */
 
+/* BEGIN_CASE depends_on:MBEDTLS_TEST_HOOKS:MBEDTLS_ECP_DP_SECP256K1_ENABLED */
+void ecp_mod_p256k1(char *input_N,
+                    char *input_X,
+                    char *result)
+{
+    mbedtls_mpi X;
+    mbedtls_mpi N;
+    mbedtls_mpi res;
+
+    mbedtls_mpi_init(&X);
+    mbedtls_mpi_init(&N);
+    mbedtls_mpi_init(&res);
+
+    TEST_EQUAL(mbedtls_test_read_mpi(&X,   input_X), 0);
+    TEST_EQUAL(mbedtls_test_read_mpi(&N,   input_N), 0);
+    TEST_EQUAL(mbedtls_test_read_mpi(&res, result),  0);
+
+    TEST_ASSERT(mbedtls_mpi_core_uint_le_mpi(0, X.p, X.n));
+    TEST_ASSERT(mbedtls_mpi_core_uint_le_mpi(0, N.p, N.n));
+    TEST_ASSERT(mbedtls_mpi_core_uint_le_mpi(0, res.p, res.n));
+
+    size_t limbs = N.n;
+    size_t bytes = limbs * sizeof(mbedtls_mpi_uint);
+
+    TEST_LE_U(X.n, 2 * limbs);
+    TEST_EQUAL(res.n, limbs);
+
+    TEST_EQUAL(mbedtls_ecp_mod_p256k1(&X), 0);
+    TEST_EQUAL(mbedtls_mpi_mod_mpi(&X, &X, &N), 0);
+    TEST_LE_U(mbedtls_mpi_core_bitlen(X.p, X.n), 256);
+    ASSERT_COMPARE(X.p, bytes, res.p, bytes);
+
+exit:
+    mbedtls_mpi_free(&X);
+    mbedtls_mpi_free(&N);
+    mbedtls_mpi_free(&res);
+}
+/* END_CASE */
+
 /* BEGIN_CASE depends_on:MBEDTLS_TEST_HOOKS */
 void ecp_mod_setup(char *input_A, int id, int ctype, int iret)
 {
diff --git a/tests/suites/test_suite_platform.data b/tests/suites/test_suite_platform.data
index 557b586..39b423f 100644
--- a/tests/suites/test_suite_platform.data
+++ b/tests/suites/test_suite_platform.data
@@ -7,6 +7,3 @@
 
 Time: delay milliseconds
 time_delay_milliseconds:1000
-
-Time: delay seconds
-time_delay_seconds:1
diff --git a/tests/suites/test_suite_platform.function b/tests/suites/test_suite_platform.function
index 54ddd42..7453c32 100644
--- a/tests/suites/test_suite_platform.function
+++ b/tests/suites/test_suite_platform.function
@@ -38,8 +38,6 @@
 
 /* END_DEPENDENCIES */
 
-
-
 /* BEGIN_CASE depends_on:MBEDTLS_HAVE_TIME */
 void time_get_milliseconds()
 {
@@ -76,6 +74,14 @@
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_HAVE_TIME */
+
+/*
+ * WARNING: DO NOT ENABLE THIS TEST. We keep the code here to document the
+ *          reason.
+ *
+ *          The test often failed on the CI. See #1517. CI failures cannot be
+ *          completely avoided due to out-of-sync clock sources.
+ */
 void time_delay_seconds(int delay_secs)
 {
     mbedtls_time_t current = mbedtls_time(NULL);
@@ -84,7 +90,23 @@
     sleep_ms(delay_secs * 1000);
 
     elapsed_secs = mbedtls_time(NULL) - current;
-    TEST_ASSERT(elapsed_secs >= delay_secs && elapsed_secs < 4 + delay_secs);
+
+    /*
+     * `mbedtls_time()` was defined as c99 function `time()`, returns the number
+     * of seconds since the Epoch. And it is affected by discontinuous changes
+     * from automatic drift adjustment or time setting system call. The POSIX.1
+     * specification for clock_settime says that discontinuous changes in
+     * CLOCK_REALTIME should not affect `nanosleep()`.
+     *
+     * If discontinuous changes occur during `nanosleep()`, we will get
+     * `elapsed_secs < delay_secs` for backward or `elapsed_secs > delay_secs`
+     * for forward.
+     *
+     * The following tolerance windows cannot be guaranteed.
+     * PLEASE DO NOT ENABLE IT IN CI TEST.
+     */
+    TEST_ASSERT(elapsed_secs - delay_secs >= -1 &&
+                elapsed_secs - delay_secs <   4);
     /* This goto is added to avoid warnings from the generated code. */
     goto exit;
 }
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 07aad73..f0b3574 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -3045,7 +3045,7 @@
 aead_encrypt_decrypt:PSA_KEY_TYPE_AES:"D7828D13B2B0BDC325A76236DF93CC6B":PSA_ALG_CTR:"000102030405060708090A0B0C0D0E0F":"":"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA AEAD encrypt/decrypt: invalid algorithm (ChaCha20)
-depends_on:MBEDTLS_CHACHA20_C
+depends_on:PSA_WANT_KEY_TYPE_CHACHA20
 aead_encrypt_decrypt:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_STREAM_CIPHER:"":"":"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA Multipart AEAD encrypt: AES - CCM, 23 bytes (lengths set)