Merge pull request #8937 from valeriosetti/issue8712

Clarify the documentation of mbedtls_pk_setup_opaque
diff --git a/docs/use-psa-crypto.md b/docs/use-psa-crypto.md
index 92d0985..f2983bd 100644
--- a/docs/use-psa-crypto.md
+++ b/docs/use-psa-crypto.md
@@ -75,13 +75,8 @@
 
 **Benefits:** isolation of long-term secrets, use of PSA Crypto drivers.
 
-**Limitations:** can only wrap a key pair, can only use it for private key
-operations. (That is, signature generation, and for RSA decryption too.)
-Note: for ECDSA, currently this uses randomized ECDSA while Mbed TLS uses
-deterministic ECDSA by default. The following operations are not supported
-with a context set this way, while they would be available with a normal
-context: `mbedtls_pk_check_pair()`, `mbedtls_pk_debug()`, all public key
-operations.
+**Limitations:** please refer to the documentation of `mbedtls_pk_setup_opaque()`
+for a full list of supported operations and limitations.
 
 **Use in X.509 and TLS:** opt-in. The application needs to construct the PK context
 using the new API in order to get the benefits; it can then pass the
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index fde302f..52f4cc6 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -359,32 +359,40 @@
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 /**
- * \brief           Initialize a PK context to wrap a PSA key.
+ * \brief Initialize a PK context to wrap a PSA key.
  *
- * \note            This function replaces mbedtls_pk_setup() for contexts
- *                  that wrap a (possibly opaque) PSA key instead of
- *                  storing and manipulating the key material directly.
+ * This function creates a PK context which wraps a PSA key. The PSA wrapped
+ * key must be an EC or RSA key pair (DH is not suported in the PK module).
  *
- * \param ctx       The context to initialize. It must be empty (type NONE).
- * \param key       The PSA key to wrap, which must hold an ECC or RSA key
- *                  pair (see notes below).
+ * Under the hood PSA functions will be used to perform the required
+ * operations and, based on the key type, used algorithms will be:
+ * * EC:
+ *     * verify, verify_ext, sign, sign_ext: ECDSA.
+ * * RSA:
+ *     * sign, decrypt: use the primary algorithm in the wrapped PSA key;
+ *     * sign_ext: RSA PSS if the pk_type is #MBEDTLS_PK_RSASSA_PSS, otherwise
+ *       it falls back to the sign() case;
+ *     * verify, verify_ext, encrypt: not supported.
  *
- * \note            The wrapped key must remain valid as long as the
- *                  wrapping PK context is in use, that is at least between
- *                  the point this function is called and the point
- *                  mbedtls_pk_free() is called on this context. The wrapped
- *                  key might then be independently used or destroyed.
+ * In order for the above operations to succeed, the policy of the wrapped PSA
+ * key must allow the specified algorithm.
  *
- * \note            This function is currently only available for ECC or RSA
- *                  key pairs (that is, keys containing private key material).
- *                  Support for other key types may be added later.
+ * Opaque PK contexts wrapping an EC keys also support \c mbedtls_pk_check_pair(),
+ * whereas RSA ones do not.
  *
- * \return          \c 0 on success.
- * \return          #MBEDTLS_ERR_PK_BAD_INPUT_DATA on invalid input
- *                  (context already used, invalid key identifier).
- * \return          #MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE if the key is not an
- *                  ECC key pair.
- * \return          #MBEDTLS_ERR_PK_ALLOC_FAILED on allocation failure.
+ * \warning The PSA wrapped key must remain valid as long as the wrapping PK
+ *          context is in use, that is at least between the point this function
+ *          is called and the point mbedtls_pk_free() is called on this context.
+ *
+ * \param ctx The context to initialize. It must be empty (type NONE).
+ * \param key The PSA key to wrap, which must hold an ECC or RSA key pair.
+ *
+ * \return    \c 0 on success.
+ * \return    #MBEDTLS_ERR_PK_BAD_INPUT_DATA on invalid input (context already
+ *            used, invalid key identifier).
+ * \return    #MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE if the key is not an ECC or
+ *            RSA key pair.
+ * \return    #MBEDTLS_ERR_PK_ALLOC_FAILED on allocation failure.
  */
 int mbedtls_pk_setup_opaque(mbedtls_pk_context *ctx,
                             const mbedtls_svc_key_id_t key);
diff --git a/library/pk.c b/library/pk.c
index 097777f..c29318d 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -1327,43 +1327,19 @@
     }
 
     if (mbedtls_pk_get_type(ctx) == MBEDTLS_PK_OPAQUE) {
-        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_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,
+        /* PSA_ALG_RSA_PSS() behaves the same as PSA_ALG_RSA_PSS_ANY_SALT() when
+         * performing a signature, but they are encoded differently. Instead of
+         * extracting the proper one from the wrapped key policy, just try both. */
+        status = psa_sign_hash(ctx->priv_id, PSA_ALG_RSA_PSS(psa_md_alg),
                                hash, hash_len,
                                sig, sig_size, sig_len);
+        if (status == PSA_ERROR_NOT_PERMITTED) {
+            status = psa_sign_hash(ctx->priv_id, PSA_ALG_RSA_PSS_ANY_SALT(psa_md_alg),
+                                   hash, hash_len,
+                                   sig, sig_size, sig_len);
+        }
         return PSA_PK_RSA_TO_MBEDTLS_ERR(status);
     }
 
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index ddcbd83..564bda8 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -2082,6 +2082,19 @@
     memset(hash, 0x2a, sizeof(hash));
     memset(sig, 0, sizeof(sig));
 
+#if defined(MBEDTLS_PKCS1_V21)
+    /* Check that trying to use the wrong pk_type in sign_ext() results in a failure.
+     * The PSA key was setup to use PKCS1 v1.5 signature algorithm, but here we try
+     * to use it for PSS (PKCS1 v2.1) and it should fail. */
+    if (key_pk_type == MBEDTLS_PK_RSA) {
+        TEST_EQUAL(mbedtls_pk_sign_ext(MBEDTLS_PK_RSASSA_PSS, &pk, md_alg, hash, hash_len,
+                                       sig, sizeof(sig), &sig_len,
+                                       mbedtls_test_rnd_std_rand, NULL),
+                   MBEDTLS_ERR_RSA_BAD_INPUT_DATA);
+    }
+#endif /* MBEDTLS_PKCS1_V21 */
+
+    /* Perform sign_ext() with the correct pk_type. */
     TEST_EQUAL(mbedtls_pk_sign_ext(key_pk_type, &pk, md_alg, hash, hash_len,
                                    sig, sizeof(sig), &sig_len,
                                    mbedtls_test_rnd_std_rand, NULL), 0);