Add usage parameter to mbedtls_pk_can_do_ext()

Signed-off-by: Neil Armstrong <narmstrong@baylibre.com>
diff --git a/library/pk.c b/library/pk.c
index b95f957..ab443ea 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -236,8 +236,11 @@
 /*
  * Tell if a PK can do the operations of the given PSA algorithm
  */
-int mbedtls_pk_can_do_ext( const mbedtls_pk_context *ctx, psa_algorithm_t alg )
+int mbedtls_pk_can_do_ext( const mbedtls_pk_context *ctx, psa_algorithm_t alg,
+                           psa_key_usage_t usage )
 {
+    psa_key_usage_t key_usage;
+
     /* A context with null pk_info is not set up yet and can't do anything.
      * For backward compatibility, also accept NULL instead of a context
      * pointer. */
@@ -252,6 +255,12 @@
         PSA_ALG_IS_ECDH( alg ) == 0 )
         return( 0 );
 
+    /* Filter out non allowed usage flags */
+    if( ( usage & ~( PSA_KEY_USAGE_SIGN_HASH |
+                     PSA_KEY_USAGE_DECRYPT |
+                     PSA_KEY_USAGE_DERIVE ) ) != 0 )
+        return( 0 );
+
     /* Wildcard hash is not allowed */
     if( PSA_ALG_IS_SIGN_HASH( alg ) &&
         PSA_ALG_SIGN_GET_HASH( alg ) == PSA_ALG_ANY_HASH )
@@ -262,15 +271,32 @@
         mbedtls_pk_type_t type;
 
         if( PSA_ALG_IS_ECDSA( alg ) || PSA_ALG_IS_ECDH( alg ) )
+        {
             type = MBEDTLS_PK_ECKEY;
-        else if( PSA_ALG_IS_RSA_PKCS1V15_SIGN( alg ) ||
-                 alg == PSA_ALG_RSA_PKCS1V15_CRYPT )
+            key_usage = PSA_KEY_USAGE_SIGN_HASH |
+                        PSA_KEY_USAGE_DERIVE;
+        }
+        else if( PSA_ALG_IS_RSA_PKCS1V15_SIGN( alg ) )
+        {
             type = MBEDTLS_PK_RSA;
+            key_usage = PSA_KEY_USAGE_SIGN_HASH;
+        }
+        else if( alg == PSA_ALG_RSA_PKCS1V15_CRYPT )
+        {
+            type = MBEDTLS_PK_RSA;
+            key_usage = PSA_KEY_USAGE_DECRYPT;
+        }
         else if( PSA_ALG_IS_RSA_PSS( alg ) )
+        {
             type = MBEDTLS_PK_RSASSA_PSS;
+            key_usage = PSA_KEY_USAGE_SIGN_HASH;
+        }
         else
             return( 0 );
 
+        if( ( key_usage & usage ) != usage )
+            return( 0 );
+
         return( ctx->pk_info->can_do( type ) );
     }
 
@@ -285,8 +311,12 @@
 
     key_alg = psa_get_key_algorithm( &attributes );
     key_alg2 = psa_get_key_enrollment_algorithm( &attributes );
+    key_usage = psa_get_key_usage_flags( &attributes );
     psa_reset_key_attributes( &attributes );
 
+    if( ( key_usage & usage ) != usage )
+        return( 0 );
+
     /*
      * Common case: the key alg or alg2 only allows alg.
      * This will match PSA_ALG_RSA_PKCS1V15_CRYPT & PSA_ALG_IS_ECDH
@@ -299,7 +329,7 @@
 
     /*
      * If key alg or alg2 is a hash-and-sign with a wildcard for the hash,
-     * and alg is the same hash-and-sign family with any hash,
+     * and alg is the same hash-and-sign family with any hasor,
      * then alg is compliant with this key alg
      */
     if( PSA_ALG_IS_SIGN_HASH( alg ) )