psa: mgmt: Add key slot access counter

Add key slot access counter to be able to
state if a key slot containing the description
of a permanent key can be reset or reset
and re-used.

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 2c4878d..a510e3c 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1189,20 +1189,25 @@
 
 /** 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. */
+ * operations with this algorithm.
+ *
+ * On success, the access counter of the returned key slot is incremented by
+ * one. It is the responsibility of the caller to call
+ * psa_decrement_key_slot_access_count() when it does not access the key slot
+ * anymore.
+ */
 static psa_status_t psa_get_key_from_slot( mbedtls_svc_key_id_t key,
                                            psa_key_slot_t **p_slot,
                                            psa_key_usage_t usage,
                                            psa_algorithm_t alg )
 {
-    psa_status_t status;
-    psa_key_slot_t *slot = NULL;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_slot_t *slot;
 
-    *p_slot = NULL;
-
-    status = psa_get_key_slot( key, &slot );
+    status = psa_get_key_slot( key, p_slot );
     if( status != PSA_SUCCESS )
         return( status );
+    slot = *p_slot;
 
     /* Enforce that usage policy for the key slot contains all the flags
      * required by the usage parameter. There is one exception: public
@@ -1210,15 +1215,22 @@
      * if they had the export flag. */
     if( PSA_KEY_TYPE_IS_PUBLIC_KEY( slot->attr.type ) )
         usage &= ~PSA_KEY_USAGE_EXPORT;
+
+    status = PSA_ERROR_NOT_PERMITTED;
     if( ( slot->attr.policy.usage & usage ) != usage )
-        return( PSA_ERROR_NOT_PERMITTED );
+        goto error;
 
     /* Enforce that the usage policy permits the requested algortihm. */
     if( alg != 0 && ! psa_key_policy_permits( &slot->attr.policy, alg ) )
-        return( PSA_ERROR_NOT_PERMITTED );
+        goto error;
 
-    *p_slot = slot;
     return( PSA_SUCCESS );
+
+error:
+    *p_slot = NULL;
+    psa_decrement_key_slot_access_count( slot );
+
+    return( status );
 }
 
 /** Retrieve a slot which must contain a transparent key.
@@ -1228,6 +1240,11 @@
  *
  * This is a temporary function to use instead of psa_get_key_from_slot()
  * until secure element support is fully implemented.
+ *
+ * On success, the access counter of the returned key slot is incremented by
+ * one. It is the responsibility of the caller to call
+ * psa_decrement_key_slot_access_count() when it does not access the key slot
+ * anymore.
  */
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
 static psa_status_t psa_get_transparent_key( mbedtls_svc_key_id_t key,
@@ -1238,11 +1255,14 @@
     psa_status_t status = psa_get_key_from_slot( key, p_slot, usage, alg );
     if( status != PSA_SUCCESS )
         return( status );
+
     if( psa_key_slot_is_external( *p_slot ) )
     {
+        psa_decrement_key_slot_access_count( *p_slot );
         *p_slot = NULL;
         return( PSA_ERROR_NOT_SUPPORTED );
     }
+
     return( PSA_SUCCESS );
 }
 #else /* MBEDTLS_PSA_CRYPTO_SE_C */
@@ -1473,8 +1493,9 @@
 psa_status_t psa_get_key_attributes( mbedtls_svc_key_id_t key,
                                      psa_key_attributes_t *attributes )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
-    psa_status_t status;
 
     psa_reset_key_attributes( attributes );
 
@@ -1528,7 +1549,10 @@
 
     if( status != PSA_SUCCESS )
         psa_reset_key_attributes( attributes );
-    return( status );
+
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
@@ -1688,8 +1712,9 @@
                              size_t data_size,
                              size_t *data_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
-    psa_status_t status;
 
     /* 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
@@ -1703,8 +1728,11 @@
     status = psa_get_key_from_slot( key, &slot, PSA_KEY_USAGE_EXPORT, 0 );
     if( status != PSA_SUCCESS )
         return( status );
-    return( psa_internal_export_key( slot, data, data_size,
-                                     data_length, 0 ) );
+
+    status = psa_internal_export_key( slot, data, data_size, data_length, 0 );
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 psa_status_t psa_export_public_key( mbedtls_svc_key_id_t key,
@@ -1712,8 +1740,9 @@
                                     size_t data_size,
                                     size_t *data_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
-    psa_status_t status;
 
     /* 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
@@ -1725,8 +1754,11 @@
     status = psa_get_key_from_slot( key, &slot, 0, 0 );
     if( status != PSA_SUCCESS )
         return( status );
-    return( psa_internal_export_key( slot, data, data_size,
-                                     data_length, 1 ) );
+
+    status = psa_internal_export_key( slot, data, data_size, data_length, 1 );
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 #if defined(static_assert)
@@ -1833,6 +1865,11 @@
  * In case of failure at any step, stop the sequence and call
  * psa_fail_key_creation().
  *
+ * On success, the access counter of the returned key slot is incremented by
+ * one. It is the responsibility of the caller to call
+ * psa_decrement_key_slot_access_count() when it does not access the key slot
+ * anymore.
+ *
  * \param method            An identification of the calling function.
  * \param[in] attributes    Key attributes for the new key.
  * \param[out] key          On success, identifier of the key. Note that the
@@ -1943,7 +1980,6 @@
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
     *key = slot->attr.id;
-
     return( PSA_SUCCESS );
 }
 
@@ -2203,6 +2239,9 @@
         psa_fail_key_creation( slot, driver );
         *key = MBEDTLS_SVC_KEY_ID_INIT;
     }
+    else
+        status = psa_decrement_key_slot_access_count( slot );
+
     return( status );
 }
 
@@ -2233,9 +2272,10 @@
 
 exit:
     if( status != PSA_SUCCESS )
-    {
         psa_fail_key_creation( slot, driver );
-    }
+    else
+        status = psa_decrement_key_slot_access_count( slot );
+
     /* Registration doesn't keep the key in RAM. */
     psa_close_key( key );
     return( status );
@@ -2261,7 +2301,8 @@
                            const psa_key_attributes_t *specified_attributes,
                            mbedtls_svc_key_id_t *target_key )
 {
-    psa_status_t status;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *source_slot = NULL;
     psa_key_slot_t *target_slot = NULL;
     psa_key_attributes_t actual_attributes = *specified_attributes;
@@ -2308,7 +2349,12 @@
         psa_fail_key_creation( target_slot, driver );
         *target_key = MBEDTLS_SVC_KEY_ID_INIT;
     }
-    return( status );
+    else
+        status = psa_decrement_key_slot_access_count( target_slot );
+
+    decrement_status = psa_decrement_key_slot_access_count( source_slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 
@@ -3094,7 +3140,8 @@
                                    psa_algorithm_t alg,
                                    int is_sign )
 {
-    psa_status_t status;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
     size_t key_bits;
     psa_key_usage_t usage =
@@ -3203,7 +3250,10 @@
     {
         operation->key_set = 1;
     }
-    return( status );
+
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 psa_status_t psa_mac_sign_setup( psa_mac_operation_t *operation,
@@ -3700,8 +3750,9 @@
                             size_t signature_size,
                             size_t *signature_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
-    psa_status_t status;
 
     *signature_length = signature_size;
     /* Immediately reject a zero-length signature buffer. This guarantees
@@ -3807,7 +3858,10 @@
         memset( signature, '!', signature_size );
     /* If signature_size is 0 then we have nothing to do. We must not call
      * memset because signature may be NULL in this case. */
-    return( status );
+
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 psa_status_t psa_verify_hash( mbedtls_svc_key_id_t key,
@@ -3817,8 +3871,9 @@
                               const uint8_t *signature,
                               size_t signature_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
-    psa_status_t status;
 
     status = psa_get_key_from_slot( key, &slot,
                                     PSA_KEY_USAGE_VERIFY_HASH, alg );
@@ -3834,7 +3889,7 @@
                                              signature_length );
     if( status != PSA_ERROR_NOT_SUPPORTED ||
         psa_key_lifetime_is_external( slot->attr.lifetime ) )
-        return status;
+        goto exit;
 
 #if defined(MBEDTLS_RSA_C)
     if( PSA_KEY_TYPE_IS_RSA( slot->attr.type ) )
@@ -3846,7 +3901,7 @@
                                               slot->data.key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
-            return( status );
+            goto exit;
 
         status = psa_rsa_verify( rsa,
                                  alg,
@@ -3854,7 +3909,7 @@
                                  signature, signature_length );
         mbedtls_rsa_free( rsa );
         mbedtls_free( rsa );
-        return( status );
+        goto exit;
     }
     else
 #endif /* defined(MBEDTLS_RSA_C) */
@@ -3870,25 +3925,31 @@
                                                   slot->data.key.bytes,
                                                   &ecp );
             if( status != PSA_SUCCESS )
-                return( status );
+                goto exit;
             status = psa_ecdsa_verify( ecp,
                                        hash, hash_length,
                                        signature, signature_length );
             mbedtls_ecp_keypair_free( ecp );
             mbedtls_free( ecp );
-            return( status );
+            goto exit;
         }
         else
 #endif /* defined(MBEDTLS_PSA_BUILTIN_ALG_ECDSA) || defined(MBEDTLS_PSA_BUILTIN_ALG_DETERMINISTIC_ECDSA) */
         {
-            return( PSA_ERROR_INVALID_ARGUMENT );
+            status =  PSA_ERROR_INVALID_ARGUMENT;
+            goto exit;
         }
     }
     else
 #endif /* defined(MBEDTLS_ECP_C) */
     {
-        return( PSA_ERROR_NOT_SUPPORTED );
+        status = PSA_ERROR_NOT_SUPPORTED;
     }
+
+exit:
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 #if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_PKCS1_V21)
@@ -3912,8 +3973,9 @@
                                      size_t output_size,
                                      size_t *output_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
-    psa_status_t status;
 
     (void) input;
     (void) input_length;
@@ -3931,7 +3993,10 @@
         return( status );
     if( ! ( PSA_KEY_TYPE_IS_PUBLIC_KEY( slot->attr.type ) ||
             PSA_KEY_TYPE_IS_KEY_PAIR( slot->attr.type ) ) )
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
 
 #if defined(MBEDTLS_RSA_C)
     if( PSA_KEY_TYPE_IS_RSA( slot->attr.type ) )
@@ -3989,13 +4054,17 @@
 
         mbedtls_rsa_free( rsa );
         mbedtls_free( rsa );
-        return( status );
     }
     else
 #endif /* defined(MBEDTLS_RSA_C) */
     {
-        return( PSA_ERROR_NOT_SUPPORTED );
+        status = PSA_ERROR_NOT_SUPPORTED;
     }
+
+exit:
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 psa_status_t psa_asymmetric_decrypt( mbedtls_svc_key_id_t key,
@@ -4008,8 +4077,9 @@
                                      size_t output_size,
                                      size_t *output_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
-    psa_status_t status;
 
     (void) input;
     (void) input_length;
@@ -4026,7 +4096,10 @@
     if( status != PSA_SUCCESS )
         return( status );
     if( ! PSA_KEY_TYPE_IS_KEY_PAIR( slot->attr.type ) )
-        return( PSA_ERROR_INVALID_ARGUMENT );
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
 
 #if defined(MBEDTLS_RSA_C)
     if( slot->attr.type == PSA_KEY_TYPE_RSA_KEY_PAIR )
@@ -4037,7 +4110,7 @@
                                               slot->data.key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
-            return( status );
+            goto exit;
 
         if( input_length != mbedtls_rsa_get_len( rsa ) )
         {
@@ -4084,13 +4157,17 @@
 rsa_exit:
         mbedtls_rsa_free( rsa );
         mbedtls_free( rsa );
-        return( status );
     }
     else
 #endif /* defined(MBEDTLS_RSA_C) */
     {
-        return( PSA_ERROR_NOT_SUPPORTED );
+        status = PSA_ERROR_NOT_SUPPORTED;
     }
+
+exit:
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 
@@ -4104,8 +4181,9 @@
                                       psa_algorithm_t alg,
                                       mbedtls_operation_t cipher_operation )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     int ret = 0;
-    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
     psa_key_slot_t *slot;
     size_t key_bits;
     const mbedtls_cipher_info_t *cipher_info = NULL;
@@ -4249,7 +4327,10 @@
     }
     else
         psa_cipher_abort( operation );
-    return( status );
+
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 psa_status_t psa_cipher_encrypt_setup( psa_cipher_operation_t *operation,
@@ -4615,6 +4696,7 @@
     const mbedtls_cipher_info_t *cipher_info;
     union
     {
+        unsigned dummy; /* Make the union non-empty even with no supported algorithms. */
 #if defined(MBEDTLS_CCM_C)
         mbedtls_ccm_context ccm;
 #endif /* MBEDTLS_CCM_C */
@@ -4630,6 +4712,8 @@
     uint8_t tag_length;
 } aead_operation_t;
 
+#define AEAD_OPERATION_INIT {0, 0, {0}, 0, 0, 0}
+
 static void psa_aead_abort_internal( aead_operation_t *operation )
 {
     switch( operation->core_alg )
@@ -4645,6 +4729,8 @@
             break;
 #endif /* MBEDTLS_GCM_C */
     }
+
+    psa_decrement_key_slot_access_count( operation->slot );
 }
 
 static psa_status_t psa_aead_setup( aead_operation_t *operation,
@@ -4666,7 +4752,10 @@
         mbedtls_cipher_info_from_psa( alg, operation->slot->attr.type, key_bits,
                                       &cipher_id );
     if( operation->cipher_info == NULL )
-        return( PSA_ERROR_NOT_SUPPORTED );
+    {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        goto cleanup;
+    }
 
     switch( PSA_ALG_AEAD_WITH_TAG_LENGTH( alg, 0 ) )
     {
@@ -4678,7 +4767,10 @@
              * The call to mbedtls_ccm_encrypt_and_tag or
              * mbedtls_ccm_auth_decrypt will validate the tag length. */
             if( PSA_BLOCK_CIPHER_BLOCK_SIZE( operation->slot->attr.type ) != 16 )
-                return( PSA_ERROR_INVALID_ARGUMENT );
+            {
+                status = PSA_ERROR_INVALID_ARGUMENT;
+                goto cleanup;
+            }
             mbedtls_ccm_init( &operation->ctx.ccm );
             status = mbedtls_to_psa_error(
                 mbedtls_ccm_setkey( &operation->ctx.ccm, cipher_id,
@@ -4697,7 +4789,10 @@
              * The call to mbedtls_gcm_crypt_and_tag or
              * mbedtls_gcm_auth_decrypt will validate the tag length. */
             if( PSA_BLOCK_CIPHER_BLOCK_SIZE( operation->slot->attr.type ) != 16 )
-                return( PSA_ERROR_INVALID_ARGUMENT );
+            {
+                status = PSA_ERROR_INVALID_ARGUMENT;
+                goto cleanup;
+            }
             mbedtls_gcm_init( &operation->ctx.gcm );
             status = mbedtls_to_psa_error(
                 mbedtls_gcm_setkey( &operation->ctx.gcm, cipher_id,
@@ -4714,7 +4809,10 @@
             operation->full_tag_length = 16;
             /* We only support the default tag length. */
             if( alg != PSA_ALG_CHACHA20_POLY1305 )
-                return( PSA_ERROR_NOT_SUPPORTED );
+            {
+                status = PSA_ERROR_NOT_SUPPORTED;
+                goto cleanup;
+            }
             mbedtls_chachapoly_init( &operation->ctx.chachapoly );
             status = mbedtls_to_psa_error(
                 mbedtls_chachapoly_setkey( &operation->ctx.chachapoly,
@@ -4725,7 +4823,8 @@
 #endif /* MBEDTLS_CHACHAPOLY_C */
 
         default:
-            return( PSA_ERROR_NOT_SUPPORTED );
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto cleanup;
     }
 
     if( PSA_AEAD_TAG_LENGTH( alg ) > operation->full_tag_length )
@@ -4755,7 +4854,7 @@
                                size_t *ciphertext_length )
 {
     psa_status_t status;
-    aead_operation_t operation;
+    aead_operation_t operation = AEAD_OPERATION_INIT;
     uint8_t *tag;
 
     *ciphertext_length = 0;
@@ -4869,7 +4968,7 @@
                                size_t *plaintext_length )
 {
     psa_status_t status;
-    aead_operation_t operation;
+    aead_operation_t operation = AEAD_OPERATION_INIT;
     const uint8_t *tag = NULL;
 
     *plaintext_length = 0;
@@ -5409,6 +5508,9 @@
         psa_fail_key_creation( slot, driver );
         *key = MBEDTLS_SVC_KEY_ID_INIT;
     }
+    else
+        status = psa_decrement_key_slot_access_count( slot );
+
     return( status );
 }
 
@@ -5772,8 +5874,9 @@
     psa_key_derivation_step_t step,
     mbedtls_svc_key_id_t key )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
-    psa_status_t status;
 
     status = psa_get_transparent_key( key, &slot,
                                       PSA_KEY_USAGE_DERIVE, operation->alg );
@@ -5788,10 +5891,14 @@
     if( step == PSA_KEY_DERIVATION_INPUT_SECRET )
         operation->can_output_key = 1;
 
-    return( psa_key_derivation_input_internal( operation,
-                                               step, slot->attr.type,
-                                               slot->data.key.data,
-                                               slot->data.key.bytes ) );
+    status = psa_key_derivation_input_internal( operation,
+                                                step, slot->attr.type,
+                                                slot->data.key.data,
+                                                slot->data.key.bytes );
+
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 
@@ -5939,8 +6046,10 @@
                                                const uint8_t *peer_key,
                                                size_t peer_key_length )
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
-    psa_status_t status;
+
     if( ! PSA_ALG_IS_KEY_AGREEMENT( operation->alg ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
     status = psa_get_transparent_key( private_key, &slot,
@@ -5959,7 +6068,10 @@
         if( step == PSA_KEY_DERIVATION_INPUT_SECRET )
             operation->can_output_key = 1;
     }
-    return( status );
+
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 psa_status_t psa_raw_key_agreement( psa_algorithm_t alg,
@@ -5970,8 +6082,9 @@
                                     size_t output_size,
                                     size_t *output_length )
 {
-    psa_key_slot_t *slot;
-    psa_status_t status;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t decrement_status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_slot_t *slot = NULL;
 
     if( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) )
     {
@@ -6001,7 +6114,10 @@
         psa_generate_random( output, output_size );
         *output_length = output_size;
     }
-    return( status );
+
+    decrement_status = psa_decrement_key_slot_access_count( slot );
+
+    return( ( status == PSA_SUCCESS ) ? decrement_status : status );
 }
 
 
@@ -6250,6 +6366,9 @@
         psa_fail_key_creation( slot, driver );
         *key = MBEDTLS_SVC_KEY_ID_INIT;
     }
+    else
+        status = psa_decrement_key_slot_access_count( slot );
+
     return( status );
 }
 
diff --git a/library/psa_crypto_core.h b/library/psa_crypto_core.h
index 86d804b..32d1d60 100644
--- a/library/psa_crypto_core.h
+++ b/library/psa_crypto_core.h
@@ -36,6 +36,33 @@
 typedef struct
 {
     psa_core_key_attributes_t attr;
+
+    /*
+     * Number of on-going accesses, read and/or write, to the key slot by the
+     * library.
+     *
+     * This counter is incremented by one each time a library function
+     * retrieves through one of the dedicated internal API a pointer to the
+     * key slot.
+     *
+     * This counter is decremented by one each time a library function stops
+     * accessing to the key slot and states it by calling the
+     * psa_decrement_key_slot_access_count() API.
+     *
+     * This counter is used to prevent resetting the key slot while the library
+     * may access it. For example, such control is needed in the following
+     * scenarios:
+     * . In case of key slot starvation, all key slots contain the description
+     *   of a key, and the library asks for the description of a permanent
+     *   key not present in the key slots, the key slots currently accessed by
+     *   the library cannot be reclaimed to free a key slot to load the
+     *   permanent key.
+     * . In case of a multi-threaded application where one thread asks to close
+     *   or purge or destroy a key while it is in used by the library through
+     *   another thread.
+     */
+    size_t access_count;
+
     union
     {
         /* Dynamically allocated key data buffer.
@@ -74,6 +101,20 @@
     return( slot->attr.type != 0 );
 }
 
+/** Test whether a key slot is accessed.
+ *
+ * A key slot is accessed iff its access counter is strickly greater than
+ * 0.
+ *
+ * \param[in] slot  The key slot to test.
+ *
+ * \return 1 if the slot is accessed, 0 otherwise.
+ */
+static inline int psa_is_key_slot_accessed( const psa_key_slot_t *slot )
+{
+    return( slot->access_count > 0 );
+}
+
 /** Retrieve flags from psa_key_slot_t::attr::core::flags.
  *
  * \param[in] slot      The key slot to query.
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 7308f6f..e207477 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -88,6 +88,11 @@
  *  key with identifier key_id can only be stored in slot of index
  *  ( key_id - #PSA_KEY_ID_VOLATILE_MIN ).
  *
+ *  On success, the access counter of the returned key slot is incremented by
+ *  one. It is the responsibility of the caller to call
+ *  psa_decrement_key_slot_access_count() when it does not access the key slot
+ *  anymore.
+ *
  * \param key           Key identifier to query.
  * \param[out] p_slot   On success, `*p_slot` contains a pointer to the
  *                      key slot containing the description of the key
@@ -135,7 +140,10 @@
     }
 
     if( status == PSA_SUCCESS )
+    {
         *p_slot = slot;
+        psa_increment_key_slot_access_count( slot );
+    }
 
     return( status );
 }
@@ -177,9 +185,12 @@
             *volatile_key_id = PSA_KEY_ID_VOLATILE_MIN +
                                ( (psa_key_id_t)slot_idx ) - 1;
 
+            psa_increment_key_slot_access_count( *p_slot );
+
             return( PSA_SUCCESS );
         }
     }
+
     *p_slot = NULL;
     return( PSA_ERROR_INSUFFICIENT_MEMORY );
 }
@@ -232,6 +243,10 @@
     if( ! global_data.key_slots_initialized )
         return( PSA_ERROR_BAD_STATE );
 
+    /*
+     * On success, the pointer to the slot is passed directly to the caller
+     * thus no need to decrement the key slot access counter here.
+     */
     status = psa_search_key_in_slots( key, p_slot );
     if( status != PSA_ERROR_DOES_NOT_EXIST )
         return( status );
@@ -257,6 +272,36 @@
 
 }
 
+psa_status_t psa_decrement_key_slot_access_count( psa_key_slot_t *slot )
+{
+    if( slot == NULL )
+        return( PSA_SUCCESS );
+
+    if( slot->access_count > 0 )
+    {
+        slot->access_count--;
+        return( PSA_SUCCESS );
+    }
+
+    /*
+     * As the return error code may not be handled in case of multiple errors,
+     * do our best to report if the access counter is equal to zero: if
+     * available call MBEDTLS_PARAM_FAILED that may terminate execution (if
+     * called as part of the execution of a unit test suite this will stop the
+     * test suite execution) and if MBEDTLS_PARAM_FAILED does not terminate
+     * execution ouput an error message on standard error output.
+     */
+#ifdef MBEDTLS_CHECK_PARAMS
+    MBEDTLS_PARAM_FAILED( slot->access_count > 0 );
+#endif
+#ifdef MBEDTLS_PLATFORM_C
+    mbedtls_fprintf( stderr,
+        "\nFATAL psa_decrement_key_slot_access_count Decrementing a zero access counter.\n" );
+#endif
+
+    return( PSA_ERROR_CORRUPTION_DETECTED );
+}
+
 psa_status_t psa_validate_key_location( psa_key_lifetime_t lifetime,
                                         psa_se_drv_table_entry_t **p_drv )
 {
@@ -315,7 +360,7 @@
 
     *handle = key;
 
-    return( PSA_SUCCESS );
+    return( psa_decrement_key_slot_access_count( slot ) );
 
 #else /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
     (void) key;
@@ -349,7 +394,7 @@
         return( status );
 
     if( PSA_KEY_LIFETIME_IS_VOLATILE( slot->attr.lifetime ) )
-        return PSA_SUCCESS;
+        return( psa_decrement_key_slot_access_count( slot ) );
 
     return( psa_wipe_key_slot( slot ) );
 }
diff --git a/library/psa_crypto_slot_management.h b/library/psa_crypto_slot_management.h
index 2b90ce8..d22e343 100644
--- a/library/psa_crypto_slot_management.h
+++ b/library/psa_crypto_slot_management.h
@@ -70,6 +70,11 @@
  *  In case of a persistent key, the function loads the description of the key
  *  into a key slot if not already done.
  *
+ *  On success, the access counter of the returned key slot is incremented by
+ *  one. It is the responsibility of the caller to call
+ *  psa_decrement_key_slot_access_count() when it does not access the slot
+ *  anymore.
+ *
  * \param key           Key identifier to query.
  * \param[out] p_slot   On success, `*p_slot` contains a pointer to the
  *                      key slot containing the description of the key
@@ -110,7 +115,10 @@
 /** Find a free key slot.
  *
  * This function returns a key slot that is available for use and is in its
- * ground state (all-bits-zero).
+ * ground state (all-bits-zero). On success, the access counter of the
+ * returned key slot is incremented by one. It is the responsibility of the
+ * caller to call psa_decrement_key_slot_access_count() when it does not access
+ * the key slot anymore.
  *
  * \param[out] volatile_key_id   On success, volatile key identifier
  *                               associated to the returned slot.
@@ -123,6 +131,35 @@
 psa_status_t psa_get_empty_key_slot( psa_key_id_t *volatile_key_id,
                                      psa_key_slot_t **p_slot );
 
+/** Increment slot access counter.
+ *
+ * This function increments the slot access counter by one.
+ *
+ * \param[in] slot  The key slot.
+ */
+static inline void psa_increment_key_slot_access_count( psa_key_slot_t *slot )
+{
+    slot->access_count++;
+}
+
+/** Decrement slot access counter.
+ *
+ * This function decrements the slot access counter by one.
+ *
+ * \note To ease the handling of errors in retrieving a key slot
+ *       a NULL input pointer is valid, and the function returns
+ *       successfully without doing anything in that case.
+ *
+ * \param[in] slot  The key slot.
+ * \retval #PSA_SUCCESS
+ *             \p slot is NULL or the key slot access pointer has been
+ *             decremented successfully.
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ *             The access counter was equal to 0.
+ *
+ */
+psa_status_t psa_decrement_key_slot_access_count( psa_key_slot_t *slot );
+
 /** Test whether a lifetime designates a key in an external cryptoprocessor.
  *
  * \param lifetime      The lifetime to test.