pk: add mbedtls_pk_copy_from_psa()

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index ff80290..6447a15 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -390,6 +390,47 @@
                             const mbedtls_svc_key_id_t key);
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
+/**
+ * \brief           Create a PK context starting from a key stored in PSA.
+ *                  This key:
+ *                  - must have PSA_KEY_USAGE_EXPORT attribute set and
+ *                  - must be a either a RSA or EC (DH is not managed in PK) and
+ *                  - must be either a key pair or a public key.
+ *
+ *                  The resulting PK object will be a transparent type:
+ *                  - MBEDTLS_PK_RSA for RSA keys or
+ *                  - MBEDTLS_PK_ECKEY for EC keys.
+ *                  Once this functions returns the PK object will be completely
+ *                  independent from the original PSA key that it was generated
+ *                  from.
+ *                  Calling `mbedtls_pk_sign`, `mbedtls_pk_verify`,
+ *                  `mbedtls_pk_encrypt`, `mbedtls_pk_decrypt` on the resulting
+ *                  PK context will perform an algorithm that is compatible with
+ *                  the PSA key's primary algorithm policy if that is a matching
+ *                  operation type (sign/verify, encrypt/decrypt), but with no
+ *                  restriction on the hash (as if the policy had
+ *                  `PSA_ALG_ANY_HASH` instead of a specific hash, and with
+ *                  `PSA_ALG_RSA_PKCS1V15_SIGN_RAW` merged with
+ *                  `PSA_ALG_RSA_PKCS1V15_SIGN(hash_alg)`).
+ *                  * For ECDSA, the choice of deterministic vs randomized will
+ *                    be based on the compile-time setting `MBEDTLS_ECDSA_DETERMINISTIC`.
+ *                  * For an RSA key, the output key will allow both encrypt/decrypt
+ *                    and sign/verify regardless of the original key's policy.
+ *                    The original key's policy determines the output key's padding
+ *                    mode.
+ *
+ * \param key_id    The ID of the key stored in PSA.
+ * \param pk        The PK context that will be filled. It must be initialized,
+ *                  but not setup.
+ *
+ * \return          0 on success.
+ * \return          MBEDTLS_ERR_PK_BAD_INPUT_DATA in case the provided input
+ *                  parameters are not correct.
+ */
+int mbedtls_pk_copy_from_psa(mbedtls_svc_key_id_t key_id, mbedtls_pk_context *pk);
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
+
 #if defined(MBEDTLS_PK_RSA_ALT_SUPPORT)
 /**
  * \brief           Initialize an RSA-alt context
diff --git a/library/pk.c b/library/pk.c
index 003ef4a..ffe350e 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -35,6 +35,10 @@
 #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
  */
@@ -1374,4 +1378,130 @@
     return ctx->pk_info->type;
 }
 
+#if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
+int mbedtls_pk_copy_from_psa(mbedtls_svc_key_id_t key_id, mbedtls_pk_context *pk)
+{
+    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 ((psa_get_key_usage_flags(&key_attr) & PSA_KEY_USAGE_EXPORT) != PSA_KEY_USAGE_EXPORT) {
+        ret = MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+        goto exit;
+    }
+
+    status = psa_export_key(key_id, exp_key, sizeof(exp_key), &exp_key_len);
+    if (status != PSA_SUCCESS) {
+        ret = psa_generic_status_to_mbedtls(status);
+        goto exit;
+    }
+
+    key_type = psa_get_key_type(&key_attr);
+    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)) {
+        if (!PSA_ALG_IS_RSA_PKCS1V15_SIGN(alg_type) &&
+            (alg_type != PSA_ALG_RSA_PKCS1V15_CRYPT) &&
+            !PSA_ALG_IS_RSA_OAEP(alg_type) &&
+            !PSA_ALG_IS_RSA_PSS(alg_type)) {
+            return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+        }
+
+        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 ((alg_type != PSA_ALG_RSA_PKCS1V15_CRYPT) &&
+            (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 {
+            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;
+
+        if (!PSA_ALG_IS_ECDSA(alg_type)) {
+            ret = MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+            goto exit;
+        }
+
+        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;
+}
+#endif /* MBEDTLS_PSA_CRYPTO_CLIENT */
+
 #endif /* MBEDTLS_PK_C */