Crypto: Implement key storage APIs

Change-Id: I9967c52aaa4d531ec89642b0e8f2bc50b2da5cfe
Signed-off-by: Jamie Fox <jamie.fox@arm.com>
diff --git a/secure_fw/services/crypto/crypto_key.c b/secure_fw/services/crypto/crypto_key.c
index f3fd277..c7405ca 100644
--- a/secure_fw/services/crypto/crypto_key.c
+++ b/secure_fw/services/crypto/crypto_key.c
@@ -71,6 +71,46 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
+psa_status_t tfm_crypto_check_key_storage(uint32_t *index)
+{
+#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
+    return PSA_ERROR_NOT_SUPPORTED;
+#else
+    uint32_t i;
+
+    for (i = 0; i < TFM_CRYPTO_MAX_KEY_HANDLES; i++) {
+        if (handle_owner[i].in_use == TFM_CRYPTO_NOT_IN_USE) {
+            *index = i;
+            return PSA_SUCCESS;
+        }
+    }
+
+    return PSA_ERROR_INSUFFICIENT_MEMORY;
+#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
+}
+
+psa_status_t tfm_crypto_set_key_storage(uint32_t index,
+                                        psa_key_handle_t key_handle)
+{
+#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
+    return PSA_ERROR_NOT_SUPPORTED;
+#else
+    psa_status_t status;
+    int32_t partition_id;
+
+    status = tfm_crypto_get_caller_id(&partition_id);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    handle_owner[index].owner = partition_id;
+    handle_owner[index].handle = key_handle;
+    handle_owner[index].in_use = TFM_CRYPTO_IN_USE;
+
+    return PSA_SUCCESS;
+#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
+}
+
 psa_status_t tfm_crypto_set_key_domain_parameters(psa_invec in_vec[],
                                    size_t in_len,
                                    psa_outvec out_vec[],
diff --git a/secure_fw/services/crypto/crypto_key_derivation.c b/secure_fw/services/crypto/crypto_key_derivation.c
index 80958f5..c186518 100644
--- a/secure_fw/services/crypto/crypto_key_derivation.c
+++ b/secure_fw/services/crypto/crypto_key_derivation.c
@@ -274,8 +274,7 @@
     const psa_key_attributes_t *key_attributes = in_vec[1].base;
     psa_key_derivation_operation_t *operation = NULL;
     psa_key_handle_t *key_handle = out_vec[0].base;
-    //int32_t partition_id = 0;
-    //uint32_t index = 0;
+    uint32_t index;
 
     /* Look up the corresponding operation context */
     status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
@@ -284,18 +283,18 @@
     if (status != PSA_SUCCESS) {
         return status;
     }
-    /*
-    status = tfm_crypto_check_key_storage(&partition_id, &index);
+
+    status = tfm_crypto_check_key_storage(&index);
     if (status != PSA_SUCCESS) {
         return status;
     }
-    */
-    status = psa_key_derivation_output_key(key_attributes, operation, key_handle);
-    /*
+
+    status = psa_key_derivation_output_key(key_attributes, operation,
+                                           key_handle);
     if (status == PSA_SUCCESS) {
-        status = tfm_crypto_set_key_storage(partition_id, index, *key_handle);
+        status = tfm_crypto_set_key_storage(index, *key_handle);
     }
-    */
+
     return status;
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
diff --git a/secure_fw/services/crypto/tfm_crypto_api.h b/secure_fw/services/crypto/tfm_crypto_api.h
index 811e04a..74b0624 100644
--- a/secure_fw/services/crypto/tfm_crypto_api.h
+++ b/secure_fw/services/crypto/tfm_crypto_api.h
@@ -83,31 +83,25 @@
                                            uint32_t *index);
 
 /**
- * \brief Checks that there is enough local storage in RAM to keep
- *        another key, and returns the index of the storage to use
- *        and the ID of the partition to associated to the key
+ * \brief Checks that there is enough local storage in RAM to keep another key,
+ *        and returns the index of the storage to use.
  *
- * \param[out] partition_id ID of the requesting partition
- * \param[out] index        Index of the local storage to use
+ * \param[out] index  Index of the local storage to use
  *
  * \return Return values as described in \ref psa_status_t
  */
-psa_status_t tfm_crypto_check_key_storage(int32_t *partition_id,
-                                          uint32_t *index);
+psa_status_t tfm_crypto_check_key_storage(uint32_t *index);
 
 /**
- * \brief Sets the index of the local storage in use with a key
- *        requested by a partition specified by the input parameter
- *        partition_id, and stores the corresponding key_handle
+ * \brief Sets the index of the local storage in use with a key requested by the
+ *        calling partition, and stores the corresponding key_handle.
  *
- * \param[in] partition_id ID of the requesting partition
- * \param[in] index        Index of the local storage to use
- * \param[in] key_handle   Corresponding key handle to associate
+ * \param[in] index       Index of the local storage to use
+ * \param[in] key_handle  Corresponding key handle to associate
  *
  * \return Return values as described in \ref psa_status_t
  */
-psa_status_t tfm_crypto_set_key_storage(int32_t partition_id,
-                                        uint32_t index,
+psa_status_t tfm_crypto_set_key_storage(uint32_t index,
                                         psa_key_handle_t key_handle);
 /**
  * \brief Allocate an operation context in the backend