Squashed commit upgrading to mbedtls-3.6.0

Squash merging branch import/mbedtls-3.6.0

0fc9291f4 ("libmbedtls: bignum: restore mbedtls_mpi_exp_mod() from v3.5.2")
0ef87b1e6 ("libmbedtls: reset minimum rsa key size")
70b079496 ("libmbedtls: adjust use of rsa pk_wrap API")
6cf76464f ("libmbedtls: allow inclusion of arm_neon.h")
27df5c911 ("libmbedtls: fix cipher_wrap.c for NIST AES Key Wrap mode")
aa584f9ed ("libmbedtls: fix cipher_wrap.c for chacha20 and chachapoly")
523ae957e ("libmbedtls: add fault mitigation in mbedtls_rsa_rsassa_pkcs1_v15_verify()")
30bdb1bbf ("libmbedtls: add fault mitigation in mbedtls_rsa_rsassa_pss_verify_ext()")
e45cdab62 ("libmbedtls: add SM2 curve")
d2fda4fc2 ("libmbedtls: fix no CRT issue")
ab0eb5515 ("libmbedtls: add interfaces in mbedtls for context memory operation")
7925a6f26 ("libmedtls: mpi_miller_rabin: increase count limit")
8eaf69279 ("libmbedtls: add mbedtls_mpi_init_mempool()")
12e83fc8d ("libmbedtls: make mbedtls_mpi_mont*() available")
f9e261da5 ("mbedtls: configure mbedtls to reach for config")
7b6f378d7 ("mbedtls: remove default include/mbedtls/config.h")
c16331743 ("Import mbedtls-3.6.0")

Signed-off-by: Tom Van Eyck <tom.vaneyck@kuleuven.be>
Acked-by: Jerome Forissier <jerome.forissier@linaro.org>
diff --git a/lib/libmbedtls/mbedtls/library/pk.c b/lib/libmbedtls/mbedtls/library/pk.c
index 5e18ad2..097777f 100644
--- a/lib/libmbedtls/mbedtls/library/pk.c
+++ b/lib/libmbedtls/mbedtls/library/pk.c
@@ -2,19 +2,7 @@
  *  Public Key abstraction layer
  *
  *  Copyright The Mbed TLS Contributors
- *  SPDX-License-Identifier: Apache-2.0
- *
- *  Licensed under the Apache License, Version 2.0 (the "License"); you may
- *  not use this file except in compliance with the License.
- *  You may obtain a copy of the License at
- *
- *  http://www.apache.org/licenses/LICENSE-2.0
- *
- *  Unless required by applicable law or agreed to in writing, software
- *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
- *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *  See the License for the specific language governing permissions and
- *  limitations under the License.
+ *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
  */
 
 #include "common.h"
@@ -23,36 +11,34 @@
 #include "mbedtls/pk.h"
 #include "pk_wrap.h"
 #include "pkwrite.h"
-
-#include "hash_info.h"
+#include "pk_internal.h"
 
 #include "mbedtls/platform_util.h"
 #include "mbedtls/error.h"
 
 #if defined(MBEDTLS_RSA_C)
 #include "mbedtls/rsa.h"
+#include "rsa_internal.h"
 #endif
-#if defined(MBEDTLS_ECP_C)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 #include "mbedtls/ecp.h"
 #endif
 #if defined(MBEDTLS_ECDSA_C)
 #include "mbedtls/ecdsa.h"
 #endif
 
-#if defined(MBEDTLS_PSA_CRYPTO_C)
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
+#include "psa_util_internal.h"
 #include "mbedtls/psa_util.h"
-#define PSA_PK_TO_MBEDTLS_ERR(status) psa_pk_status_to_mbedtls(status)
-#define PSA_PK_RSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,     \
-                                                                  psa_to_pk_rsa_errors,            \
-                                                                  psa_pk_status_to_mbedtls)
-#define PSA_PK_ECDSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                                    psa_to_pk_ecdsa_errors,        \
-                                                                    psa_pk_status_to_mbedtls)
 #endif
 
 #include <limits.h>
 #include <stdint.h>
 
+#define PSA_EXPORT_KEY_PAIR_OR_PUBLIC_MAX_SIZE \
+    (PSA_EXPORT_KEY_PAIR_MAX_SIZE > PSA_EXPORT_PUBLIC_KEY_MAX_SIZE) ? \
+    PSA_EXPORT_KEY_PAIR_MAX_SIZE : PSA_EXPORT_PUBLIC_KEY_MAX_SIZE
+
 /*
  * Initialise a mbedtls_pk_context
  */
@@ -60,6 +46,15 @@
 {
     ctx->pk_info = NULL;
     ctx->pk_ctx = NULL;
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    ctx->priv_id = MBEDTLS_SVC_KEY_ID_INIT;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    memset(ctx->pub_raw, 0, sizeof(ctx->pub_raw));
+    ctx->pub_raw_len = 0;
+    ctx->ec_family = 0;
+    ctx->ec_bits = 0;
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 }
 
 /*
@@ -71,10 +66,18 @@
         return;
     }
 
-    if (ctx->pk_info != NULL) {
+    if ((ctx->pk_info != NULL) && (ctx->pk_info->ctx_free_func != NULL)) {
         ctx->pk_info->ctx_free_func(ctx->pk_ctx);
     }
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    /* The ownership of the priv_id key for opaque keys is external of the PK
+     * module. It's the user responsibility to clear it after use. */
+    if ((ctx->pk_info != NULL) && (ctx->pk_info->type != MBEDTLS_PK_OPAQUE)) {
+        psa_destroy_key(ctx->priv_id);
+    }
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+
     mbedtls_platform_zeroize(ctx, sizeof(mbedtls_pk_context));
 }
 
@@ -114,17 +117,17 @@
 #if defined(MBEDTLS_RSA_C)
         case MBEDTLS_PK_RSA:
             return &mbedtls_rsa_info;
-#endif
-#if defined(MBEDTLS_ECP_C)
+#endif /* MBEDTLS_RSA_C */
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
         case MBEDTLS_PK_ECKEY:
             return &mbedtls_eckey_info;
         case MBEDTLS_PK_ECKEY_DH:
             return &mbedtls_eckeydh_info;
-#endif
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
         case MBEDTLS_PK_ECDSA:
             return &mbedtls_ecdsa_info;
-#endif
+#endif /* MBEDTLS_PK_CAN_ECDSA_SOME */
         /* MBEDTLS_PK_RSA_ALT omitted on purpose */
         default:
             return NULL;
@@ -140,7 +143,8 @@
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
     }
 
-    if ((ctx->pk_ctx = info->ctx_alloc_func()) == NULL) {
+    if ((info->ctx_alloc_func != NULL) &&
+        ((ctx->pk_ctx = info->ctx_alloc_func()) == NULL)) {
         return MBEDTLS_ERR_PK_ALLOC_FAILED;
     }
 
@@ -158,7 +162,6 @@
 {
     const mbedtls_pk_info_t *info = NULL;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    mbedtls_svc_key_id_t *pk_ctx;
     psa_key_type_t type;
 
     if (ctx == NULL || ctx->pk_info != NULL) {
@@ -171,22 +174,19 @@
     type = psa_get_key_type(&attributes);
     psa_reset_key_attributes(&attributes);
 
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(type)) {
-        info = &mbedtls_pk_ecdsa_opaque_info;
-    } else if (type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
-        info = &mbedtls_pk_rsa_opaque_info;
+        info = &mbedtls_ecdsa_opaque_info;
+    } else
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+    if (type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
+        info = &mbedtls_rsa_opaque_info;
     } else {
         return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
     }
 
-    if ((ctx->pk_ctx = info->ctx_alloc_func()) == NULL) {
-        return MBEDTLS_ERR_PK_ALLOC_FAILED;
-    }
-
     ctx->pk_info = info;
-
-    pk_ctx = (mbedtls_svc_key_id_t *) ctx->pk_ctx;
-    *pk_ctx = key;
+    ctx->priv_id = key;
 
     return 0;
 }
@@ -315,18 +315,23 @@
         return (key_usage & usage) == usage;
     }
 
-    const mbedtls_svc_key_id_t *key = (const mbedtls_svc_key_id_t *) ctx->pk_ctx;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_algorithm_t key_alg, key_alg2;
     psa_status_t status;
 
-    status = psa_get_key_attributes(*key, &attributes);
+    status = psa_get_key_attributes(ctx->priv_id, &attributes);
     if (status != PSA_SUCCESS) {
         return 0;
     }
 
-    key_alg = psa_get_key_algorithm(&attributes);
-    key_alg2 = psa_get_key_enrollment_algorithm(&attributes);
+    psa_algorithm_t key_alg = psa_get_key_algorithm(&attributes);
+    /* Key's enrollment is available only when an Mbed TLS implementation of PSA
+     * Crypto is being used, i.e. when MBEDTLS_PSA_CRYPTO_C is defined.
+     * Even though we don't officially support using other implementations of PSA
+     * Crypto with TLS and X.509 (yet), we try to keep vendor's customizations
+     * separated. */
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+    psa_algorithm_t key_alg2 = psa_get_key_enrollment_algorithm(&attributes);
+#endif /* MBEDTLS_PSA_CRYPTO_C */
     key_usage = psa_get_key_usage_flags(&attributes);
     psa_reset_key_attributes(&attributes);
 
@@ -335,40 +340,657 @@
     }
 
     /*
-     * Common case: the key alg or alg2 only allows alg.
+     * Common case: the key alg [or alg2] only allows alg.
      * This will match PSA_ALG_RSA_PKCS1V15_CRYPT & PSA_ALG_IS_ECDH
      * directly.
      * This would also match ECDSA/RSA_PKCS1V15_SIGN/RSA_PSS with
-     * a fixed hash on key_alg/key_alg2.
+     * a fixed hash on key_alg [or key_alg2].
      */
-    if (alg == key_alg || alg == key_alg2) {
+    if (alg == key_alg) {
         return 1;
     }
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+    if (alg == key_alg2) {
+        return 1;
+    }
+#endif /* MBEDTLS_PSA_CRYPTO_C */
 
     /*
-     * If key_alg or key_alg2 is a hash-and-sign with a wildcard for the hash,
+     * If key_alg [or key_alg2] is a hash-and-sign with a wildcard for the hash,
      * and alg is the same hash-and-sign family with any hash,
      * then alg is compliant with this key alg
      */
     if (PSA_ALG_IS_SIGN_HASH(alg)) {
-
         if (PSA_ALG_IS_SIGN_HASH(key_alg) &&
             PSA_ALG_SIGN_GET_HASH(key_alg) == PSA_ALG_ANY_HASH &&
             (alg & ~PSA_ALG_HASH_MASK) == (key_alg & ~PSA_ALG_HASH_MASK)) {
             return 1;
         }
-
+#if defined(MBEDTLS_PSA_CRYPTO_C)
         if (PSA_ALG_IS_SIGN_HASH(key_alg2) &&
             PSA_ALG_SIGN_GET_HASH(key_alg2) == PSA_ALG_ANY_HASH &&
             (alg & ~PSA_ALG_HASH_MASK) == (key_alg2 & ~PSA_ALG_HASH_MASK)) {
             return 1;
         }
+#endif /* MBEDTLS_PSA_CRYPTO_C */
     }
 
     return 0;
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
+#if defined(MBEDTLS_RSA_C)
+static psa_algorithm_t psa_algorithm_for_rsa(const mbedtls_rsa_context *rsa,
+                                             int want_crypt)
+{
+    if (mbedtls_rsa_get_padding_mode(rsa) == MBEDTLS_RSA_PKCS_V21) {
+        if (want_crypt) {
+            mbedtls_md_type_t md_type = (mbedtls_md_type_t) mbedtls_rsa_get_md_alg(rsa);
+            return PSA_ALG_RSA_OAEP(mbedtls_md_psa_alg_from_type(md_type));
+        } else {
+            return PSA_ALG_RSA_PSS_ANY_SALT(PSA_ALG_ANY_HASH);
+        }
+    } else {
+        if (want_crypt) {
+            return PSA_ALG_RSA_PKCS1V15_CRYPT;
+        } else {
+            return PSA_ALG_RSA_PKCS1V15_SIGN(PSA_ALG_ANY_HASH);
+        }
+    }
+}
+#endif /* MBEDTLS_RSA_C */
+
+int mbedtls_pk_get_psa_attributes(const mbedtls_pk_context *pk,
+                                  psa_key_usage_t usage,
+                                  psa_key_attributes_t *attributes)
+{
+    mbedtls_pk_type_t pk_type = mbedtls_pk_get_type(pk);
+
+    psa_key_usage_t more_usage = usage;
+    if (usage == PSA_KEY_USAGE_SIGN_MESSAGE) {
+        more_usage |= PSA_KEY_USAGE_VERIFY_MESSAGE;
+    } else if (usage == PSA_KEY_USAGE_SIGN_HASH) {
+        more_usage |= PSA_KEY_USAGE_VERIFY_HASH;
+    } else if (usage == PSA_KEY_USAGE_DECRYPT) {
+        more_usage |= PSA_KEY_USAGE_ENCRYPT;
+    }
+    more_usage |= PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_COPY;
+
+    int want_private = !(usage == PSA_KEY_USAGE_VERIFY_MESSAGE ||
+                         usage == PSA_KEY_USAGE_VERIFY_HASH ||
+                         usage == PSA_KEY_USAGE_ENCRYPT);
+
+    switch (pk_type) {
+#if defined(MBEDTLS_RSA_C)
+        case MBEDTLS_PK_RSA:
+        {
+            int want_crypt = 0; /* 0: sign/verify; 1: encrypt/decrypt */
+            switch (usage) {
+                case PSA_KEY_USAGE_SIGN_MESSAGE:
+                case PSA_KEY_USAGE_SIGN_HASH:
+                case PSA_KEY_USAGE_VERIFY_MESSAGE:
+                case PSA_KEY_USAGE_VERIFY_HASH:
+                    /* Nothing to do. */
+                    break;
+                case PSA_KEY_USAGE_DECRYPT:
+                case PSA_KEY_USAGE_ENCRYPT:
+                    want_crypt = 1;
+                    break;
+                default:
+                    return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            /* Detect the presence of a private key in a way that works both
+             * in CRT and non-CRT configurations. */
+            mbedtls_rsa_context *rsa = mbedtls_pk_rsa(*pk);
+            int has_private = (mbedtls_rsa_check_privkey(rsa) == 0);
+            if (want_private && !has_private) {
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            psa_set_key_type(attributes, (want_private ?
+                                          PSA_KEY_TYPE_RSA_KEY_PAIR :
+                                          PSA_KEY_TYPE_RSA_PUBLIC_KEY));
+            psa_set_key_bits(attributes, mbedtls_pk_get_bitlen(pk));
+            psa_set_key_algorithm(attributes,
+                                  psa_algorithm_for_rsa(rsa, want_crypt));
+            break;
+        }
+#endif /* MBEDTLS_RSA_C */
+
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+        case MBEDTLS_PK_ECKEY:
+        case MBEDTLS_PK_ECKEY_DH:
+        case MBEDTLS_PK_ECDSA:
+        {
+            int sign_ok = (pk_type != MBEDTLS_PK_ECKEY_DH);
+            int derive_ok = (pk_type != MBEDTLS_PK_ECDSA);
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+            psa_ecc_family_t family = pk->ec_family;
+            size_t bits = pk->ec_bits;
+            int has_private = 0;
+            if (pk->priv_id != MBEDTLS_SVC_KEY_ID_INIT) {
+                has_private = 1;
+            }
+#else
+            const mbedtls_ecp_keypair *ec = mbedtls_pk_ec_ro(*pk);
+            int has_private = (ec->d.n != 0);
+            size_t bits = 0;
+            psa_ecc_family_t family =
+                mbedtls_ecc_group_to_psa(ec->grp.id, &bits);
+#endif
+            psa_algorithm_t alg = 0;
+            switch (usage) {
+                case PSA_KEY_USAGE_SIGN_MESSAGE:
+                case PSA_KEY_USAGE_SIGN_HASH:
+                case PSA_KEY_USAGE_VERIFY_MESSAGE:
+                case PSA_KEY_USAGE_VERIFY_HASH:
+                    if (!sign_ok) {
+                        return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+                    }
+#if defined(MBEDTLS_ECDSA_DETERMINISTIC)
+                    alg = PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH);
+#else
+                    alg = PSA_ALG_ECDSA(PSA_ALG_ANY_HASH);
+#endif
+                    break;
+                case PSA_KEY_USAGE_DERIVE:
+                    alg = PSA_ALG_ECDH;
+                    if (!derive_ok) {
+                        return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+                    }
+                    break;
+                default:
+                    return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            if (want_private && !has_private) {
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            psa_set_key_type(attributes, (want_private ?
+                                          PSA_KEY_TYPE_ECC_KEY_PAIR(family) :
+                                          PSA_KEY_TYPE_ECC_PUBLIC_KEY(family)));
+            psa_set_key_bits(attributes, bits);
+            psa_set_key_algorithm(attributes, alg);
+            break;
+        }
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+
+#if defined(MBEDTLS_PK_RSA_ALT_SUPPORT)
+        case MBEDTLS_PK_RSA_ALT:
+            return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
+#endif /* MBEDTLS_PK_RSA_ALT_SUPPORT */
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        case MBEDTLS_PK_OPAQUE:
+        {
+            psa_key_attributes_t old_attributes = PSA_KEY_ATTRIBUTES_INIT;
+            psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+            status = psa_get_key_attributes(pk->priv_id, &old_attributes);
+            if (status != PSA_SUCCESS) {
+                return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+            }
+            psa_key_type_t old_type = psa_get_key_type(&old_attributes);
+            switch (usage) {
+                case PSA_KEY_USAGE_SIGN_MESSAGE:
+                case PSA_KEY_USAGE_SIGN_HASH:
+                case PSA_KEY_USAGE_VERIFY_MESSAGE:
+                case PSA_KEY_USAGE_VERIFY_HASH:
+                    if (!(PSA_KEY_TYPE_IS_ECC_KEY_PAIR(old_type) ||
+                          old_type == PSA_KEY_TYPE_RSA_KEY_PAIR)) {
+                        return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+                    }
+                    break;
+                case PSA_KEY_USAGE_DECRYPT:
+                case PSA_KEY_USAGE_ENCRYPT:
+                    if (old_type != PSA_KEY_TYPE_RSA_KEY_PAIR) {
+                        return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+                    }
+                    break;
+                case PSA_KEY_USAGE_DERIVE:
+                    if (!(PSA_KEY_TYPE_IS_ECC_KEY_PAIR(old_type))) {
+                        return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+                    }
+                    break;
+                default:
+                    return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            psa_key_type_t new_type = old_type;
+            /* Opaque keys are always key pairs, so we don't need a check
+             * on the input if the required usage is private. We just need
+             * to adjust the type correctly if the required usage is public. */
+            if (!want_private) {
+                new_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(new_type);
+            }
+            more_usage = psa_get_key_usage_flags(&old_attributes);
+            if ((usage & more_usage) == 0) {
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            psa_set_key_type(attributes, new_type);
+            psa_set_key_bits(attributes, psa_get_key_bits(&old_attributes));
+            psa_set_key_algorithm(attributes, psa_get_key_algorithm(&old_attributes));
+            break;
+        }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+        default:
+            return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
+    psa_set_key_usage_flags(attributes, more_usage);
+    /* Key's enrollment is available only when an Mbed TLS implementation of PSA
+     * Crypto is being used, i.e. when MBEDTLS_PSA_CRYPTO_C is defined.
+     * Even though we don't officially support using other implementations of PSA
+     * Crypto with TLS and X.509 (yet), we try to keep vendor's customizations
+     * separated. */
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+    psa_set_key_enrollment_algorithm(attributes, PSA_ALG_NONE);
+#endif
+
+    return 0;
+}
+
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA) || defined(MBEDTLS_USE_PSA_CRYPTO)
+static psa_status_t export_import_into_psa(mbedtls_svc_key_id_t old_key_id,
+                                           const psa_key_attributes_t *attributes,
+                                           mbedtls_svc_key_id_t *new_key_id)
+{
+    unsigned char key_buffer[PSA_EXPORT_KEY_PAIR_MAX_SIZE];
+    size_t key_length = 0;
+    psa_status_t status = psa_export_key(old_key_id,
+                                         key_buffer, sizeof(key_buffer),
+                                         &key_length);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+    status = psa_import_key(attributes, key_buffer, key_length, new_key_id);
+    mbedtls_platform_zeroize(key_buffer, key_length);
+    return status;
+}
+
+static int copy_into_psa(mbedtls_svc_key_id_t old_key_id,
+                         const psa_key_attributes_t *attributes,
+                         mbedtls_svc_key_id_t *new_key_id)
+{
+    /* Normally, we prefer copying: it's more efficient and works even
+     * for non-exportable keys. */
+    psa_status_t status = psa_copy_key(old_key_id, attributes, new_key_id);
+    if (status == PSA_ERROR_NOT_PERMITTED /*missing COPY usage*/ ||
+        status == PSA_ERROR_INVALID_ARGUMENT /*incompatible policy*/) {
+        /* There are edge cases where copying won't work, but export+import
+         * might:
+         * - If the old key does not allow PSA_KEY_USAGE_COPY.
+         * - If the old key's usage does not allow what attributes wants.
+         *   Because the key was intended for use in the pk module, and may
+         *   have had a policy chosen solely for what pk needs rather than
+         *   based on a detailed understanding of PSA policies, we are a bit
+         *   more liberal than psa_copy_key() here.
+         */
+        /* Here we need to check that the types match, otherwise we risk
+         * importing nonsensical data. */
+        psa_key_attributes_t old_attributes = PSA_KEY_ATTRIBUTES_INIT;
+        status = psa_get_key_attributes(old_key_id, &old_attributes);
+        if (status != PSA_SUCCESS) {
+            return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+        }
+        psa_key_type_t old_type = psa_get_key_type(&old_attributes);
+        psa_reset_key_attributes(&old_attributes);
+        if (old_type != psa_get_key_type(attributes)) {
+            return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+        }
+        status = export_import_into_psa(old_key_id, attributes, new_key_id);
+    }
+    return PSA_PK_TO_MBEDTLS_ERR(status);
+}
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA || MBEDTLS_USE_PSA_CRYPTO */
+
+static int import_pair_into_psa(const mbedtls_pk_context *pk,
+                                const psa_key_attributes_t *attributes,
+                                mbedtls_svc_key_id_t *key_id)
+{
+    switch (mbedtls_pk_get_type(pk)) {
+#if defined(MBEDTLS_RSA_C)
+        case MBEDTLS_PK_RSA:
+        {
+            if (psa_get_key_type(attributes) != PSA_KEY_TYPE_RSA_KEY_PAIR) {
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            unsigned char key_buffer[
+                PSA_KEY_EXPORT_RSA_KEY_PAIR_MAX_SIZE(PSA_VENDOR_RSA_MAX_KEY_BITS)];
+            unsigned char *const key_end = key_buffer + sizeof(key_buffer);
+            unsigned char *key_data = key_end;
+            int ret = mbedtls_rsa_write_key(mbedtls_pk_rsa(*pk),
+                                            key_buffer, &key_data);
+            if (ret < 0) {
+                return ret;
+            }
+            size_t key_length = key_end - key_data;
+            ret = PSA_PK_TO_MBEDTLS_ERR(psa_import_key(attributes,
+                                                       key_data, key_length,
+                                                       key_id));
+            mbedtls_platform_zeroize(key_data, key_length);
+            return ret;
+        }
+#endif /* MBEDTLS_RSA_C */
+
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+        case MBEDTLS_PK_ECKEY:
+        case MBEDTLS_PK_ECKEY_DH:
+        case MBEDTLS_PK_ECDSA:
+        {
+            /* We need to check the curve family, otherwise the import could
+             * succeed with nonsensical data.
+             * We don't check the bit-size: it's optional in attributes,
+             * and if it's specified, psa_import_key() will know from the key
+             * data length and will check that the bit-size matches. */
+            psa_key_type_t to_type = psa_get_key_type(attributes);
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+            psa_ecc_family_t from_family = pk->ec_family;
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+            const mbedtls_ecp_keypair *ec = mbedtls_pk_ec_ro(*pk);
+            size_t from_bits = 0;
+            psa_ecc_family_t from_family = mbedtls_ecc_group_to_psa(ec->grp.id,
+                                                                    &from_bits);
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+            if (to_type != PSA_KEY_TYPE_ECC_KEY_PAIR(from_family)) {
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+            if (mbedtls_svc_key_id_is_null(pk->priv_id)) {
+                /* We have a public key and want a key pair. */
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            return copy_into_psa(pk->priv_id, attributes, key_id);
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+            if (ec->d.n == 0) {
+                /* Private key not set. Assume the input is a public key only.
+                 * (The other possibility is that it's an incomplete object
+                 * where the group is set but neither the public key nor
+                 * the private key. This is not possible through ecp.h
+                 * functions, so we don't bother reporting a more suitable
+                 * error in that case.) */
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            unsigned char key_buffer[PSA_BITS_TO_BYTES(PSA_VENDOR_ECC_MAX_CURVE_BITS)];
+            size_t key_length = 0;
+            int ret = mbedtls_ecp_write_key_ext(ec, &key_length,
+                                                key_buffer, sizeof(key_buffer));
+            if (ret < 0) {
+                return ret;
+            }
+            ret = PSA_PK_TO_MBEDTLS_ERR(psa_import_key(attributes,
+                                                       key_buffer, key_length,
+                                                       key_id));
+            mbedtls_platform_zeroize(key_buffer, key_length);
+            return ret;
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+        }
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        case MBEDTLS_PK_OPAQUE:
+            return copy_into_psa(pk->priv_id, attributes, key_id);
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+        default:
+            return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+}
+
+static int import_public_into_psa(const mbedtls_pk_context *pk,
+                                  const psa_key_attributes_t *attributes,
+                                  mbedtls_svc_key_id_t *key_id)
+{
+    psa_key_type_t psa_type = psa_get_key_type(attributes);
+
+#if defined(MBEDTLS_RSA_C) ||                                           \
+    (defined(MBEDTLS_PK_HAVE_ECC_KEYS) && !defined(MBEDTLS_PK_USE_PSA_EC_DATA)) || \
+    defined(MBEDTLS_USE_PSA_CRYPTO)
+    unsigned char key_buffer[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE];
+#endif
+    unsigned char *key_data = NULL;
+    size_t key_length = 0;
+
+    switch (mbedtls_pk_get_type(pk)) {
+#if defined(MBEDTLS_RSA_C)
+        case MBEDTLS_PK_RSA:
+        {
+            if (psa_type != PSA_KEY_TYPE_RSA_PUBLIC_KEY) {
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            unsigned char *const key_end = key_buffer + sizeof(key_buffer);
+            key_data = key_end;
+            int ret = mbedtls_rsa_write_pubkey(mbedtls_pk_rsa(*pk),
+                                               key_buffer, &key_data);
+            if (ret < 0) {
+                return ret;
+            }
+            key_length = (size_t) ret;
+            break;
+        }
+#endif /*MBEDTLS_RSA_C */
+
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+        case MBEDTLS_PK_ECKEY:
+        case MBEDTLS_PK_ECKEY_DH:
+        case MBEDTLS_PK_ECDSA:
+        {
+            /* We need to check the curve family, otherwise the import could
+             * succeed with nonsensical data.
+             * We don't check the bit-size: it's optional in attributes,
+             * and if it's specified, psa_import_key() will know from the key
+             * data length and will check that the bit-size matches. */
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+            if (psa_type != PSA_KEY_TYPE_ECC_PUBLIC_KEY(pk->ec_family)) {
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            key_data = (unsigned char *) pk->pub_raw;
+            key_length = pk->pub_raw_len;
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+            const mbedtls_ecp_keypair *ec = mbedtls_pk_ec_ro(*pk);
+            size_t from_bits = 0;
+            psa_ecc_family_t from_family = mbedtls_ecc_group_to_psa(ec->grp.id,
+                                                                    &from_bits);
+            if (psa_type != PSA_KEY_TYPE_ECC_PUBLIC_KEY(from_family)) {
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            int ret = mbedtls_ecp_write_public_key(
+                ec, MBEDTLS_ECP_PF_UNCOMPRESSED,
+                &key_length, key_buffer, sizeof(key_buffer));
+            if (ret < 0) {
+                return ret;
+            }
+            key_data = key_buffer;
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+            break;
+        }
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        case MBEDTLS_PK_OPAQUE:
+        {
+            psa_key_attributes_t old_attributes = PSA_KEY_ATTRIBUTES_INIT;
+            psa_status_t status =
+                psa_get_key_attributes(pk->priv_id, &old_attributes);
+            if (status != PSA_SUCCESS) {
+                return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+            }
+            psa_key_type_t old_type = psa_get_key_type(&old_attributes);
+            psa_reset_key_attributes(&old_attributes);
+            if (psa_type != PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(old_type)) {
+                return MBEDTLS_ERR_PK_TYPE_MISMATCH;
+            }
+            status = psa_export_public_key(pk->priv_id,
+                                           key_buffer, sizeof(key_buffer),
+                                           &key_length);
+            if (status != PSA_SUCCESS) {
+                return PSA_PK_TO_MBEDTLS_ERR(status);
+            }
+            key_data = key_buffer;
+            break;
+        }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+        default:
+            return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
+    return PSA_PK_TO_MBEDTLS_ERR(psa_import_key(attributes,
+                                                key_data, key_length,
+                                                key_id));
+}
+
+int mbedtls_pk_import_into_psa(const mbedtls_pk_context *pk,
+                               const psa_key_attributes_t *attributes,
+                               mbedtls_svc_key_id_t *key_id)
+{
+    /* Set the output immediately so that it won't contain garbage even
+     * if we error out before calling psa_import_key(). */
+    *key_id = MBEDTLS_SVC_KEY_ID_INIT;
+
+#if defined(MBEDTLS_PK_RSA_ALT_SUPPORT)
+    if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_RSA_ALT) {
+        return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
+    }
+#endif /* MBEDTLS_PK_RSA_ALT_SUPPORT */
+
+    int want_public = PSA_KEY_TYPE_IS_PUBLIC_KEY(psa_get_key_type(attributes));
+    if (want_public) {
+        return import_public_into_psa(pk, attributes, key_id);
+    } else {
+        return import_pair_into_psa(pk, attributes, key_id);
+    }
+}
+
+static int copy_from_psa(mbedtls_svc_key_id_t key_id,
+                         mbedtls_pk_context *pk,
+                         int public_only)
+{
+    psa_status_t status;
+    psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_type_t key_type;
+    psa_algorithm_t alg_type;
+    size_t key_bits;
+    /* Use a buffer size large enough to contain either a key pair or public key. */
+    unsigned char exp_key[PSA_EXPORT_KEY_PAIR_OR_PUBLIC_MAX_SIZE];
+    size_t exp_key_len;
+    int ret = MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+
+    if (pk == NULL) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
+    status = psa_get_key_attributes(key_id, &key_attr);
+    if (status != PSA_SUCCESS) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
+    if (public_only) {
+        status = psa_export_public_key(key_id, exp_key, sizeof(exp_key), &exp_key_len);
+    } else {
+        status = psa_export_key(key_id, exp_key, sizeof(exp_key), &exp_key_len);
+    }
+    if (status != PSA_SUCCESS) {
+        ret = PSA_PK_TO_MBEDTLS_ERR(status);
+        goto exit;
+    }
+
+    key_type = psa_get_key_type(&key_attr);
+    if (public_only) {
+        key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(key_type);
+    }
+    key_bits = psa_get_key_bits(&key_attr);
+    alg_type = psa_get_key_algorithm(&key_attr);
+
+#if defined(MBEDTLS_RSA_C)
+    if ((key_type == PSA_KEY_TYPE_RSA_KEY_PAIR) ||
+        (key_type == PSA_KEY_TYPE_RSA_PUBLIC_KEY)) {
+
+        ret = mbedtls_pk_setup(pk, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA));
+        if (ret != 0) {
+            goto exit;
+        }
+
+        if (key_type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
+            ret = mbedtls_rsa_parse_key(mbedtls_pk_rsa(*pk), exp_key, exp_key_len);
+        } else {
+            ret = mbedtls_rsa_parse_pubkey(mbedtls_pk_rsa(*pk), exp_key, exp_key_len);
+        }
+        if (ret != 0) {
+            goto exit;
+        }
+
+        mbedtls_md_type_t md_type = MBEDTLS_MD_NONE;
+        if (PSA_ALG_GET_HASH(alg_type) != PSA_ALG_ANY_HASH) {
+            md_type = mbedtls_md_type_from_psa_alg(alg_type);
+        }
+
+        if (PSA_ALG_IS_RSA_OAEP(alg_type) || PSA_ALG_IS_RSA_PSS(alg_type)) {
+            ret = mbedtls_rsa_set_padding(mbedtls_pk_rsa(*pk), MBEDTLS_RSA_PKCS_V21, md_type);
+        } else if (PSA_ALG_IS_RSA_PKCS1V15_SIGN(alg_type) ||
+                   alg_type == PSA_ALG_RSA_PKCS1V15_CRYPT) {
+            ret = mbedtls_rsa_set_padding(mbedtls_pk_rsa(*pk), MBEDTLS_RSA_PKCS_V15, md_type);
+        }
+        if (ret != 0) {
+            goto exit;
+        }
+    } else
+#endif /* MBEDTLS_RSA_C */
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+    if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(key_type) ||
+        PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY(key_type)) {
+        mbedtls_ecp_group_id grp_id;
+
+        ret = mbedtls_pk_setup(pk, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY));
+        if (ret != 0) {
+            goto exit;
+        }
+
+        grp_id = mbedtls_ecc_group_from_psa(PSA_KEY_TYPE_ECC_GET_FAMILY(key_type), key_bits);
+        ret = mbedtls_pk_ecc_set_group(pk, grp_id);
+        if (ret != 0) {
+            goto exit;
+        }
+
+        if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(key_type)) {
+            ret = mbedtls_pk_ecc_set_key(pk, exp_key, exp_key_len);
+            if (ret != 0) {
+                goto exit;
+            }
+            ret = mbedtls_pk_ecc_set_pubkey_from_prv(pk, exp_key, exp_key_len,
+                                                     mbedtls_psa_get_random,
+                                                     MBEDTLS_PSA_RANDOM_STATE);
+        } else {
+            ret = mbedtls_pk_ecc_set_pubkey(pk, exp_key, exp_key_len);
+        }
+        if (ret != 0) {
+            goto exit;
+        }
+    } else
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+    {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
+exit:
+    psa_reset_key_attributes(&key_attr);
+    mbedtls_platform_zeroize(exp_key, sizeof(exp_key));
+
+    return ret;
+}
+
+int mbedtls_pk_copy_from_psa(mbedtls_svc_key_id_t key_id,
+                             mbedtls_pk_context *pk)
+{
+    return copy_from_psa(key_id, pk, 0);
+}
+
+int mbedtls_pk_copy_public_from_psa(mbedtls_svc_key_id_t key_id,
+                                    mbedtls_pk_context *pk)
+{
+    return copy_from_psa(key_id, pk, 1);
+}
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
+
 /*
  * Helper for mbedtls_pk_sign and mbedtls_pk_verify
  */
@@ -378,7 +1000,7 @@
         return 0;
     }
 
-    *hash_len = mbedtls_hash_info_get_size(md_alg);
+    *hash_len = mbedtls_md_get_size_from_type(md_alg);
 
     if (*hash_len == 0) {
         return -1;
@@ -443,7 +1065,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 +1082,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);
 }
 
@@ -504,13 +1126,21 @@
         return mbedtls_pk_verify(ctx, md_alg, hash, hash_len, sig, sig_len);
     }
 
+    /* Ensure the PK context is of the right type otherwise mbedtls_pk_rsa()
+     * below would return a NULL pointer. */
+    if (mbedtls_pk_get_type(ctx) != MBEDTLS_PK_RSA) {
+        return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
+    }
+
 #if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_PKCS1_V21)
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     const mbedtls_pk_rsassa_pss_options *pss_opts;
 
+#if SIZE_MAX > UINT_MAX
     if (md_alg == MBEDTLS_MD_NONE && UINT_MAX < hash_len) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
     }
+#endif
 
     if (options == NULL) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
@@ -527,12 +1157,12 @@
         psa_status_t status = PSA_ERROR_DATA_CORRUPT;
         psa_status_t destruction_status = PSA_ERROR_DATA_CORRUPT;
 
-        psa_algorithm_t psa_md_alg = mbedtls_hash_info_psa_from_md(md_alg);
+        psa_algorithm_t psa_md_alg = mbedtls_md_psa_alg_from_type(md_alg);
         mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
         psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
         psa_algorithm_t psa_sig_alg = PSA_ALG_RSA_PSS_ANY_SALT(psa_md_alg);
         p = buf + sizeof(buf);
-        key_len = mbedtls_pk_write_pubkey(&p, buf, ctx);
+        key_len = mbedtls_rsa_write_pubkey(mbedtls_pk_rsa(*ctx), buf, &p);
 
         if (key_len < 0) {
             return key_len;
@@ -571,7 +1201,7 @@
 
         return PSA_PK_RSA_TO_MBEDTLS_ERR(status);
     } else
-#endif
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     {
         if (sig_len < mbedtls_pk_get_len(ctx)) {
             return MBEDTLS_ERR_RSA_VERIFY_FAILED;
@@ -626,7 +1256,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 +1275,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);
@@ -664,7 +1294,6 @@
                                        f_rng, p_rng, NULL);
 }
 
-#if defined(MBEDTLS_PSA_CRYPTO_C)
 /*
  * Make a signature given a signature type.
  */
@@ -676,11 +1305,6 @@
                         int (*f_rng)(void *, unsigned char *, size_t),
                         void *p_rng)
 {
-#if defined(MBEDTLS_RSA_C)
-    psa_algorithm_t psa_md_alg;
-#endif /* MBEDTLS_RSA_C */
-    *sig_len = 0;
-
     if (ctx->pk_info == NULL) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
     }
@@ -694,17 +1318,50 @@
                                sig, sig_size, sig_len, f_rng, p_rng);
     }
 
-#if defined(MBEDTLS_RSA_C)
-    psa_md_alg = mbedtls_hash_info_psa_from_md(md_alg);
+#if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_PKCS1_V21)
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    const psa_algorithm_t psa_md_alg = mbedtls_md_psa_alg_from_type(md_alg);
     if (psa_md_alg == 0) {
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
     }
 
     if (mbedtls_pk_get_type(ctx) == MBEDTLS_PK_OPAQUE) {
-        const mbedtls_svc_key_id_t *key = (const mbedtls_svc_key_id_t *) ctx->pk_ctx;
+        psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
+        psa_algorithm_t psa_alg, sign_alg;
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+        psa_algorithm_t psa_enrollment_alg;
+#endif /* MBEDTLS_PSA_CRYPTO_C */
         psa_status_t status;
 
-        status = psa_sign_hash(*key, PSA_ALG_RSA_PSS(psa_md_alg),
+        status = psa_get_key_attributes(ctx->priv_id, &key_attr);
+        if (status != PSA_SUCCESS) {
+            return PSA_PK_RSA_TO_MBEDTLS_ERR(status);
+        }
+        psa_alg = psa_get_key_algorithm(&key_attr);
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+        psa_enrollment_alg = psa_get_key_enrollment_algorithm(&key_attr);
+#endif /* MBEDTLS_PSA_CRYPTO_C */
+        psa_reset_key_attributes(&key_attr);
+
+        /* Since we're PK type is MBEDTLS_PK_RSASSA_PSS at least one between
+         * alg and enrollment alg should be of type RSA_PSS. */
+        if (PSA_ALG_IS_RSA_PSS(psa_alg)) {
+            sign_alg = psa_alg;
+        }
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+        else if (PSA_ALG_IS_RSA_PSS(psa_enrollment_alg)) {
+            sign_alg = psa_enrollment_alg;
+        }
+#endif /* MBEDTLS_PSA_CRYPTO_C */
+        else {
+            /* The opaque key has no RSA PSS algorithm associated. */
+            return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+        }
+        /* Adjust the hashing algorithm. */
+        sign_alg = (sign_alg & ~PSA_ALG_HASH_MASK) | PSA_ALG_GET_HASH(psa_md_alg);
+
+        status = psa_sign_hash(ctx->priv_id, sign_alg,
                                hash, hash_len,
                                sig, sig_size, sig_len);
         return PSA_PK_RSA_TO_MBEDTLS_ERR(status);
@@ -713,12 +1370,31 @@
     return mbedtls_pk_psa_rsa_sign_ext(PSA_ALG_RSA_PSS(psa_md_alg),
                                        ctx->pk_ctx, hash, hash_len,
                                        sig, sig_size, sig_len);
-#else /* MBEDTLS_RSA_C */
-    return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
-#endif /* !MBEDTLS_RSA_C */
+#else /* MBEDTLS_USE_PSA_CRYPTO */
 
+    if (sig_size < mbedtls_pk_get_len(ctx)) {
+        return MBEDTLS_ERR_PK_BUFFER_TOO_SMALL;
+    }
+
+    if (pk_hashlen_helper(md_alg, &hash_len) != 0) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+
+    mbedtls_rsa_context *const rsa_ctx = mbedtls_pk_rsa(*ctx);
+
+    const int ret = mbedtls_rsa_rsassa_pss_sign_no_mode_check(rsa_ctx, f_rng, p_rng, md_alg,
+                                                              (unsigned int) hash_len, hash, sig);
+    if (ret == 0) {
+        *sig_len = rsa_ctx->len;
+    }
+    return ret;
+
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+#else
+    return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
+#endif /* MBEDTLS_RSA_C && MBEDTLS_PKCS1_V21 */
 }
-#endif /* MBEDTLS_PSA_CRYPTO_C */
 
 /*
  * Decrypt message
@@ -736,7 +1412,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 +1432,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);
 }
 
@@ -786,12 +1462,15 @@
             return MBEDTLS_ERR_PK_TYPE_MISMATCH;
         }
     } else {
-        if (pub->pk_info != prv->pk_info) {
+        if ((prv->pk_info->type != MBEDTLS_PK_OPAQUE) &&
+            (pub->pk_info != prv->pk_info)) {
             return MBEDTLS_ERR_PK_TYPE_MISMATCH;
         }
     }
 
-    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 +1484,7 @@
         return 0;
     }
 
-    return ctx->pk_info->get_bitlen(ctx->pk_ctx);
+    return ctx->pk_info->get_bitlen((mbedtls_pk_context *) ctx);
 }
 
 /*
@@ -821,7 +1500,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;
 }
 
@@ -849,112 +1528,4 @@
     return ctx->pk_info->type;
 }
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
-/*
- * Load the key to a PSA key slot,
- * then turn the PK context into a wrapper for that key slot.
- *
- * Currently only works for EC & RSA private keys.
- */
-int mbedtls_pk_wrap_as_opaque(mbedtls_pk_context *pk,
-                              mbedtls_svc_key_id_t *key,
-                              psa_algorithm_t alg,
-                              psa_key_usage_t usage,
-                              psa_algorithm_t alg2)
-{
-#if !defined(MBEDTLS_ECP_C) && !defined(MBEDTLS_RSA_C)
-    ((void) pk);
-    ((void) key);
-    ((void) alg);
-    ((void) usage);
-    ((void) alg2);
-#else
-#if defined(MBEDTLS_ECP_C)
-    if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECKEY) {
-        const mbedtls_ecp_keypair *ec;
-        unsigned char d[MBEDTLS_ECP_MAX_BYTES];
-        size_t d_len;
-        psa_ecc_family_t curve_id;
-        psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-        psa_key_type_t key_type;
-        size_t bits;
-        int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-        psa_status_t status;
-
-        /* export the private key material in the format PSA wants */
-        ec = mbedtls_pk_ec(*pk);
-        d_len = PSA_BITS_TO_BYTES(ec->grp.nbits);
-        if ((ret = mbedtls_mpi_write_binary(&ec->d, d, d_len)) != 0) {
-            return ret;
-        }
-
-        curve_id = mbedtls_ecc_group_to_psa(ec->grp.id, &bits);
-        key_type = PSA_KEY_TYPE_ECC_KEY_PAIR(curve_id);
-
-        /* prepare the key attributes */
-        psa_set_key_type(&attributes, key_type);
-        psa_set_key_bits(&attributes, bits);
-        psa_set_key_usage_flags(&attributes, usage);
-        psa_set_key_algorithm(&attributes, alg);
-        if (alg2 != PSA_ALG_NONE) {
-            psa_set_key_enrollment_algorithm(&attributes, alg2);
-        }
-
-        /* import private key into PSA */
-        status = psa_import_key(&attributes, d, d_len, key);
-        if (status != PSA_SUCCESS) {
-            return PSA_PK_TO_MBEDTLS_ERR(status);
-        }
-
-        /* make PK context wrap the key slot */
-        mbedtls_pk_free(pk);
-        mbedtls_pk_init(pk);
-
-        return mbedtls_pk_setup_opaque(pk, *key);
-    } else
-#endif /* MBEDTLS_ECP_C */
-#if defined(MBEDTLS_RSA_C)
-    if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_RSA) {
-        unsigned char buf[MBEDTLS_PK_RSA_PRV_DER_MAX_BYTES];
-        psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-        int key_len;
-        psa_status_t status;
-
-        /* export the private key material in the format PSA wants */
-        key_len = mbedtls_pk_write_key_der(pk, buf, sizeof(buf));
-        if (key_len <= 0) {
-            return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-        }
-
-        /* prepare the key attributes */
-        psa_set_key_type(&attributes, PSA_KEY_TYPE_RSA_KEY_PAIR);
-        psa_set_key_bits(&attributes, mbedtls_pk_get_bitlen(pk));
-        psa_set_key_usage_flags(&attributes, usage);
-        psa_set_key_algorithm(&attributes, alg);
-        if (alg2 != PSA_ALG_NONE) {
-            psa_set_key_enrollment_algorithm(&attributes, alg2);
-        }
-
-        /* import private key into PSA */
-        status = psa_import_key(&attributes,
-                                buf + sizeof(buf) - key_len,
-                                key_len, key);
-
-        mbedtls_platform_zeroize(buf, sizeof(buf));
-
-        if (status != PSA_SUCCESS) {
-            return PSA_PK_TO_MBEDTLS_ERR(status);
-        }
-
-        /* make PK context wrap the key slot */
-        mbedtls_pk_free(pk);
-        mbedtls_pk_init(pk);
-
-        return mbedtls_pk_setup_opaque(pk, *key);
-    } else
-#endif /* MBEDTLS_RSA_C */
-#endif /* !MBEDTLS_ECP_C && !MBEDTLS_RSA_C */
-    return MBEDTLS_ERR_PK_TYPE_MISMATCH;
-}
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
 #endif /* MBEDTLS_PK_C */