Always access key slots through accessor functions

New functions psa_get_key_slot(), psa_get_empty_key_slot(),
psa_get_key_from_slot() to access a key slot object from a key slot
number. These functions perform all requisite validations:

* psa_get_key_slot() verifies that the key slot number is in range.
* psa_get_empty_key_slot() verifies that the slot is empty.
* psa_get_key_from_slot() verifies that the slot contains a key with
  a suitable policy.

Always use these functions so as to make sure that the requisite
validations are always performed.

library/psa_crypto.c

@@ -339,6 +339,72 @@ static psa_status_t mbedtls_to_psa_error( int ret )
     }
 }
 
+/* Retrieve a key slot, occupied or not. */
+static psa_status_t psa_get_key_slot( psa_key_slot_t key,
+                                      key_slot_t **p_slot )
+{
+    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    *p_slot = &global_data.key_slots[key];
+    return( PSA_SUCCESS );
+}
+
+/* Retrieve an empty key slot (slot with no key data, but possibly
+ * with some metadata such as a policy). */
+static psa_status_t psa_get_empty_key_slot( psa_key_slot_t key,
+                                            key_slot_t **p_slot )
+{
+    psa_status_t status;
+    key_slot_t *slot = NULL;
+
+    *p_slot = NULL;
+
+    status = psa_get_key_slot( key, &slot );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    if( slot->type != PSA_KEY_TYPE_NONE )
+        return( PSA_ERROR_OCCUPIED_SLOT );
+
+    *p_slot = slot;
+    return( status );
+}
+
+/* Retrieve a slot which must contain a key. The key must have allow all
+ * the usage flags set in \p usage. If \p alg is nonzero, the key must
+ * allow operations with this algorithm. */
+static psa_status_t psa_get_key_from_slot( psa_key_slot_t key,
+                                           key_slot_t **p_slot,
+                                           psa_key_usage_t usage,
+                                           psa_algorithm_t alg )
+{
+    psa_status_t status;
+    key_slot_t *slot = NULL;
+
+    *p_slot = NULL;
+
+    status = psa_get_key_slot( key, &slot );
+    if( status != PSA_SUCCESS )
+        return( status );
+    if( slot->type == PSA_KEY_TYPE_NONE )
+        return( PSA_ERROR_EMPTY_SLOT );
+
+    /* Enforce that usage policy for the key slot contains all the flags
+     * required by the usage parameter. There is one exception: public
+     * keys can always be exported, so we treat public key objects as
+     * if they had the export flag. */
+    if( PSA_KEY_TYPE_IS_PUBLIC_KEY( slot->type ) )
+        usage &= ~PSA_KEY_USAGE_EXPORT;
+    if( ( slot->policy.usage & usage ) != usage )
+        return( PSA_ERROR_NOT_PERMITTED );
+    if( alg != 0 && ( alg != slot->policy.alg ) )
+        return( PSA_ERROR_NOT_PERMITTED );
+
+    *p_slot = slot;
+    return( PSA_SUCCESS );
+}
+
 
 
 /****************************************************************/
@@ -481,16 +547,13 @@ psa_status_t psa_import_key( psa_key_slot_t key,
                              size_t data_length )
 {
     key_slot_t *slot;
-
+    psa_status_t status = PSA_SUCCESS;
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
+    status = psa_get_empty_key_slot( key, &slot );
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    if( status != PSA_SUCCESS )
-    slot = &global_data.key_slots[key];
+        return( status );
-    if( slot->type != PSA_KEY_TYPE_NONE )
-        return( PSA_ERROR_OCCUPIED_SLOT );
 
     if( key_type_is_raw_bytes( type ) )
     {
-        psa_status_t status;
         /* Ensure that a bytes-to-bit conversion won't overflow. */
         if( data_length > SIZE_MAX / 8 )
             return( PSA_ERROR_NOT_SUPPORTED );
@@ -510,7 +573,6 @@ psa_status_t psa_import_key( psa_key_slot_t key,
     {
         int ret;
         mbedtls_pk_context pk;
-        psa_status_t status = PSA_SUCCESS;
         mbedtls_pk_init( &pk );
         if( PSA_KEY_TYPE_IS_KEYPAIR( type ) )
             ret = mbedtls_pk_parse_key( &pk, data, data_length, NULL, 0 );
@@ -583,10 +645,12 @@ psa_status_t psa_import_key( psa_key_slot_t key,
 psa_status_t psa_destroy_key( psa_key_slot_t key )
 {
     key_slot_t *slot;
+    psa_status_t status;
+
+    status = psa_get_key_slot( key, &slot );
+    if( status != PSA_SUCCESS )
+        return( status );
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    slot = &global_data.key_slots[key];
     if( slot->type == PSA_KEY_TYPE_NONE )
     {
         /* No key material to clean, but do zeroize the slot below to wipe
@@ -629,16 +693,19 @@ psa_status_t psa_get_key_information( psa_key_slot_t key,
                                       size_t *bits )
 {
     key_slot_t *slot;
+    psa_status_t status;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
-        return( PSA_ERROR_EMPTY_SLOT );
-    slot = &global_data.key_slots[key];
     if( type != NULL )
-        *type = slot->type;
+        *type = 0;
     if( bits != NULL )
         *bits = 0;
+    status = psa_get_key_slot( key, &slot );
+    if( status != PSA_SUCCESS )
+        return( status );
     if( slot->type == PSA_KEY_TYPE_NONE )
         return( PSA_ERROR_EMPTY_SLOT );
+    if( type != NULL )
+        *type = slot->type;
 
     if( key_type_is_raw_bytes( slot->type ) )
     {
@@ -679,6 +746,13 @@ static  psa_status_t psa_internal_export_key( psa_key_slot_t key,
                                               int export_public_key )
 {
     key_slot_t *slot;
+    psa_status_t status;
+    /* Exporting a public key doesn't require a usage flag. If we're
+     * called by psa_export_public_key(), don't require the EXPORT flag.
+     * If we're called by psa_export_key(), do require the EXPORT flag;
+     * if the key turns out to be public key object, psa_get_key_from_slot()
+     * will ignore this flag. */
+    psa_key_usage_t usage = export_public_key ? 0 : PSA_KEY_USAGE_EXPORT;
 
     /* Set the key to empty now, so that even when there are errors, we always
      * set data_length to a value between 0 and data_size. On error, setting
@@ -686,20 +760,12 @@ static  psa_status_t psa_internal_export_key( psa_key_slot_t key,
      * unlikely to be accepted anywhere. */
     *data_length = 0;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
+    status = psa_get_key_from_slot( key, &slot, usage, 0 );
-        return( PSA_ERROR_EMPTY_SLOT );
+    if( status != PSA_SUCCESS )
-    slot = &global_data.key_slots[key];
+        return( status );
-    if( slot->type == PSA_KEY_TYPE_NONE )
-        return( PSA_ERROR_EMPTY_SLOT );
-
     if( export_public_key && ! PSA_KEY_TYPE_IS_ASYMMETRIC( slot->type ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    if( ! export_public_key &&
-        ! PSA_KEY_TYPE_IS_PUBLIC_KEY( slot->type ) &&
-        ( slot->policy.usage & PSA_KEY_USAGE_EXPORT ) == 0 )
-        return( PSA_ERROR_NOT_PERMITTED );
-
     if( key_type_is_raw_bytes( slot->type ) )
     {
         if( slot->data.raw.bytes > data_size )
@@ -1424,13 +1490,17 @@ psa_status_t psa_mac_start( psa_mac_operation_t *operation,
     if( status != PSA_SUCCESS )
         return( status );
 
-    slot = &global_data.key_slots[key];
+    status = psa_get_key_from_slot( key, &slot, 0, alg );
-    if( slot->type == PSA_KEY_TYPE_NONE )
+    if( status != PSA_SUCCESS )
-        return( PSA_ERROR_EMPTY_SLOT );
+        return( status );
 
+    /* Since this function is called identically for a sign or verify
+     * operation, we don't know yet whether the operation is permitted.
+     * Store the part of the key policy that we can't check in the
+     * operation structure. psa_mac_finish() or psa_mac_verify() will
+     * check that remaining part. */
     if( ( slot->policy.usage & PSA_KEY_USAGE_SIGN ) != 0 )
         operation->key_usage_sign = 1;
-
     if( ( slot->policy.usage & PSA_KEY_USAGE_VERIFY ) != 0 )
         operation->key_usage_verify = 1;
 
@@ -1919,27 +1989,14 @@ psa_status_t psa_asymmetric_sign( psa_key_slot_t key,
     (void) salt;
     (void) salt_length;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
+    status = psa_get_key_from_slot( key, &slot, PSA_KEY_USAGE_SIGN, alg );
-    {
+    if( status != PSA_SUCCESS )
-        status = PSA_ERROR_EMPTY_SLOT;
         goto exit;
-    }
-    slot = &global_data.key_slots[key];
-    if( slot->type == PSA_KEY_TYPE_NONE )
-    {
-        status = PSA_ERROR_EMPTY_SLOT;
-        goto exit;
-    }
     if( ! PSA_KEY_TYPE_IS_KEYPAIR( slot->type ) )
     {
         status = PSA_ERROR_INVALID_ARGUMENT;
         goto exit;
     }
-    if( ! ( slot->policy.usage & PSA_KEY_USAGE_SIGN ) )
-    {
-        status = PSA_ERROR_NOT_PERMITTED;
-        goto exit;
-    }
 
 #if defined(MBEDTLS_RSA_C)
     if( slot->type == PSA_KEY_TYPE_RSA_KEYPAIR )
@@ -1999,17 +2056,14 @@ psa_status_t psa_asymmetric_verify( psa_key_slot_t key,
                                     size_t signature_length )
 {
     key_slot_t *slot;
+    psa_status_t status;
 
     (void) salt;
     (void) salt_length;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
+    status = psa_get_key_from_slot( key, &slot, PSA_KEY_USAGE_VERIFY, alg );
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    if( status != PSA_SUCCESS )
-    slot = &global_data.key_slots[key];
+        return( status );
-    if( slot->type == PSA_KEY_TYPE_NONE )
-        return( PSA_ERROR_EMPTY_SLOT );
-    if( ! ( slot->policy.usage & PSA_KEY_USAGE_VERIFY ) )
-        return( PSA_ERROR_NOT_PERMITTED );
 
 #if defined(MBEDTLS_RSA_C)
     if( slot->type == PSA_KEY_TYPE_RSA_KEYPAIR ||
@@ -2054,19 +2108,17 @@ psa_status_t psa_asymmetric_encrypt( psa_key_slot_t key,
                                      size_t *output_length )
 {
     key_slot_t *slot;
+    psa_status_t status;
+
     (void) salt;
     (void) salt_length;
     *output_length = 0;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
+    status = psa_get_key_from_slot( key, &slot, PSA_KEY_USAGE_ENCRYPT, alg );
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    if( status != PSA_SUCCESS )
-    slot = &global_data.key_slots[key];
+        return( status );
-    if( slot->type == PSA_KEY_TYPE_NONE )
-        return( PSA_ERROR_EMPTY_SLOT );
     if( ! PSA_KEY_TYPE_IS_KEYPAIR( slot->type ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
-    if( ! ( slot->policy.usage & PSA_KEY_USAGE_ENCRYPT ) )
-        return( PSA_ERROR_NOT_PERMITTED );
 
 #if defined(MBEDTLS_RSA_C)
     if( slot->type == PSA_KEY_TYPE_RSA_KEYPAIR ||
@@ -2121,19 +2173,17 @@ psa_status_t psa_asymmetric_decrypt( psa_key_slot_t key,
                                      size_t *output_length )
 {
     key_slot_t *slot;
+    psa_status_t status;
+
     (void) salt;
     (void) salt_length;
     *output_length = 0;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
+    status = psa_get_key_from_slot( key, &slot, PSA_KEY_USAGE_DECRYPT, alg );
-        return( PSA_ERROR_EMPTY_SLOT );
+    if( status != PSA_SUCCESS )
-    slot = &global_data.key_slots[key];
+        return( status );
-    if( slot->type == PSA_KEY_TYPE_NONE )
-        return( PSA_ERROR_EMPTY_SLOT );
     if( ! PSA_KEY_TYPE_IS_KEYPAIR( slot->type ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
-    if( ! ( slot->policy.usage & PSA_KEY_USAGE_DECRYPT ) )
-        return( PSA_ERROR_NOT_PERMITTED );
 
 #if defined(MBEDTLS_RSA_C)
     if( slot->type == PSA_KEY_TYPE_RSA_KEYPAIR )
@@ -2216,6 +2266,9 @@ static psa_status_t psa_cipher_setup( psa_cipher_operation_t *operation,
     psa_key_type_t key_type;
     size_t key_bits;
     const mbedtls_cipher_info_t *cipher_info = NULL;
+    psa_key_usage_t usage = ( cipher_operation == MBEDTLS_ENCRYPT ?
+                              PSA_KEY_USAGE_ENCRYPT :
+                              PSA_KEY_USAGE_DECRYPT );
 
     status = psa_cipher_init( operation, alg );
     if( status != PSA_SUCCESS )
@@ -2224,7 +2277,9 @@ static psa_status_t psa_cipher_setup( psa_cipher_operation_t *operation,
     status = psa_get_key_information( key, &key_type, &key_bits );
     if( status != PSA_SUCCESS )
         return( status );
-    slot = &global_data.key_slots[key];
+    status = psa_get_key_from_slot( key, &slot, usage, alg);
+    if( status != PSA_SUCCESS )
+        return( status );
 
     cipher_info = mbedtls_cipher_info_from_psa( alg, key_type, key_bits, NULL );
     if( cipher_info == NULL )
@@ -2525,13 +2580,14 @@ psa_status_t psa_set_key_policy( psa_key_slot_t key,
                                  const psa_key_policy_t *policy )
 {
     key_slot_t *slot;
+    psa_status_t status;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT || policy == NULL )
+    if( policy == NULL )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    slot = &global_data.key_slots[key];
+    status = psa_get_empty_key_slot( key, &slot );
-    if( slot->type != PSA_KEY_TYPE_NONE )
+    if( status != PSA_SUCCESS )
-        return( PSA_ERROR_OCCUPIED_SLOT );
+        return( status );
 
     if( ( policy->usage & ~( PSA_KEY_USAGE_EXPORT |
                              PSA_KEY_USAGE_ENCRYPT |
@@ -2549,11 +2605,14 @@ psa_status_t psa_get_key_policy( psa_key_slot_t key,
                                  psa_key_policy_t *policy )
 {
     key_slot_t *slot;
+    psa_status_t status;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT || policy == NULL )
+    if( policy == NULL )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    slot = &global_data.key_slots[key];
+    status = psa_get_key_slot( key, &slot );
+    if( status != PSA_SUCCESS )
+        return( status );
 
     *policy = slot->policy;
 
@@ -2570,11 +2629,11 @@ psa_status_t psa_get_key_lifetime( psa_key_slot_t key,
                                    psa_key_lifetime_t *lifetime )
 {
     key_slot_t *slot;
+    psa_status_t status;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
+    status = psa_get_key_slot( key, &slot );
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    if( status != PSA_SUCCESS )
-
+        return( status );
-    slot = &global_data.key_slots[key];
 
     *lifetime = slot->lifetime;
 
@@ -2585,18 +2644,16 @@ psa_status_t psa_set_key_lifetime( psa_key_slot_t key,
                                    psa_key_lifetime_t lifetime )
 {
     key_slot_t *slot;
-
+    psa_status_t status;
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
-        return( PSA_ERROR_INVALID_ARGUMENT );
 
     if( lifetime != PSA_KEY_LIFETIME_VOLATILE &&
         lifetime != PSA_KEY_LIFETIME_PERSISTENT &&
         lifetime != PSA_KEY_LIFETIME_WRITE_ONCE)
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    slot = &global_data.key_slots[key];
+    status = psa_get_empty_key_slot( key, &slot );
-    if( slot->type != PSA_KEY_TYPE_NONE )
+    if( status != PSA_SUCCESS )
-        return( PSA_ERROR_OCCUPIED_SLOT );
+        return( status );
 
     if( lifetime != PSA_KEY_LIFETIME_VOLATILE )
         return( PSA_ERROR_NOT_SUPPORTED );
@@ -2639,18 +2696,15 @@ psa_status_t psa_aead_encrypt( psa_key_slot_t key,
     status = psa_get_key_information( key, &key_type, &key_bits );
     if( status != PSA_SUCCESS )
         return( status );
-    slot = &global_data.key_slots[key];
+    status = psa_get_key_from_slot( key, &slot, PSA_KEY_USAGE_ENCRYPT, alg );
-    if( slot->type == PSA_KEY_TYPE_NONE )
+    if( status != PSA_SUCCESS )
-        return( PSA_ERROR_EMPTY_SLOT );
+        return( status );
 
     cipher_info = mbedtls_cipher_info_from_psa( alg, key_type,
                                                 key_bits, &cipher_id );
     if( cipher_info == NULL )
         return( PSA_ERROR_NOT_SUPPORTED );
 
-    if( ( slot->policy.usage & PSA_KEY_USAGE_ENCRYPT ) == 0 )
-        return( PSA_ERROR_NOT_PERMITTED );
-
     if( ( key_type & PSA_KEY_TYPE_CATEGORY_MASK ) !=
         PSA_KEY_TYPE_CATEGORY_SYMMETRIC )
         return( PSA_ERROR_INVALID_ARGUMENT );
@@ -2787,18 +2841,15 @@ psa_status_t psa_aead_decrypt( psa_key_slot_t key,
     status = psa_get_key_information( key, &key_type, &key_bits );
     if( status != PSA_SUCCESS )
         return( status );
-    slot = &global_data.key_slots[key];
+    status = psa_get_key_from_slot( key, &slot, PSA_KEY_USAGE_DECRYPT, alg );
-    if( slot->type == PSA_KEY_TYPE_NONE )
+    if( status != PSA_SUCCESS )
-        return( PSA_ERROR_EMPTY_SLOT );
+        return( status );
 
     cipher_info = mbedtls_cipher_info_from_psa( alg, key_type,
                                                 key_bits, &cipher_id );
     if( cipher_info == NULL )
         return( PSA_ERROR_NOT_SUPPORTED );
 
-    if( !( slot->policy.usage & PSA_KEY_USAGE_DECRYPT ) )
-        return( PSA_ERROR_NOT_PERMITTED );
-
     if( ( key_type & PSA_KEY_TYPE_CATEGORY_MASK ) !=
         PSA_KEY_TYPE_CATEGORY_SYMMETRIC )
         return( PSA_ERROR_INVALID_ARGUMENT );
@@ -2903,19 +2954,18 @@ psa_status_t psa_generate_key( psa_key_slot_t key,
                                size_t parameters_size )
 {
     key_slot_t *slot;
+    psa_status_t status;
 
-    if( key == 0 || key > PSA_KEY_SLOT_COUNT )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    slot = &global_data.key_slots[key];
-    if( slot->type != PSA_KEY_TYPE_NONE )
-        return( PSA_ERROR_OCCUPIED_SLOT );
     if( parameters == NULL && parameters_size != 0 )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
+    status = psa_get_empty_key_slot( key, &slot );
+    if( status != PSA_SUCCESS )
+        return( status );
+
     if( key_type_is_raw_bytes( type ) )
     {
-        psa_status_t status = prepare_raw_data_slot( type, bits,
+        status = prepare_raw_data_slot( type, bits, &slot->data.raw );
-                                                     &slot->data.raw );
         if( status != PSA_SUCCESS )
             return( status );
         status = psa_generate_random( slot->data.raw.data,