Crypto: Fixes for PSA Crypto API compliance

Details:
 - Require that setup functions are called with handles initialised to
   0. Make 0 the invalid handle and number handles from 1.
 - Return PSA_ERROR_BUFFER_TOO_SMALL if the output buffer is too small
   in psa_asymmetric_encrypt.
 - Release generator operations if setup or abort operations fail.

Change-Id: Ic1da83e056a00364ef8eba3ea23cbf191f8886e4
Signed-off-by: Jamie Fox <jamie.fox@arm.com>
diff --git a/interface/include/tfm_crypto_defs.h b/interface/include/tfm_crypto_defs.h
index 436cc81..adddf4c 100644
--- a/interface/include/tfm_crypto_defs.h
+++ b/interface/include/tfm_crypto_defs.h
@@ -121,7 +121,7 @@
  * \brief This value is used to mark an handle as invalid.
  *
  */
-#define TFM_CRYPTO_INVALID_HANDLE (~0x0u)
+#define TFM_CRYPTO_INVALID_HANDLE (0x0u)
 
 /**
  * \brief Define miscellaneous literal constants that are used in the service
diff --git a/secure_fw/services/crypto/crypto_alloc.c b/secure_fw/services/crypto/crypto_alloc.c
index db9c6af..e8b95c9 100644
--- a/secure_fw/services/crypto/crypto_alloc.c
+++ b/secure_fw/services/crypto/crypto_alloc.c
@@ -94,6 +94,11 @@
 {
     uint32_t i = 0;
 
+    /* Handle must be initialised before calling a setup function */
+    if (*handle != TFM_CRYPTO_INVALID_HANDLE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
     /* Init to invalid values */
     if (ctx == NULL) {
         return PSA_ERROR_INVALID_ARGUMENT;
@@ -104,13 +109,12 @@
         if (operation[i].in_use == TFM_CRYPTO_NOT_IN_USE) {
             operation[i].in_use = TFM_CRYPTO_IN_USE;
             operation[i].type = type;
-            *handle = i;
+            *handle = i + 1;
             *ctx = (void *) &(operation[i].operation);
             return PSA_SUCCESS;
         }
     }
 
-    *handle = TFM_CRYPTO_INVALID_HANDLE;
     return PSA_ERROR_NOT_PERMITTED;
 }
 
@@ -119,12 +123,12 @@
     uint32_t h_val = *handle;
 
     if ( (h_val != TFM_CRYPTO_INVALID_HANDLE) &&
-         (h_val < TFM_CRYPTO_CONC_OPER_NUM) &&
-         (operation[h_val].in_use == TFM_CRYPTO_IN_USE) ) {
+         (h_val <= TFM_CRYPTO_CONC_OPER_NUM) &&
+         (operation[h_val - 1].in_use == TFM_CRYPTO_IN_USE) ) {
 
-        memset_operation_context(h_val);
-        operation[h_val].in_use = TFM_CRYPTO_NOT_IN_USE;
-        operation[h_val].type = TFM_CRYPTO_OPERATION_NONE;
+        memset_operation_context(h_val - 1);
+        operation[h_val - 1].in_use = TFM_CRYPTO_NOT_IN_USE;
+        operation[h_val - 1].type = TFM_CRYPTO_OPERATION_NONE;
         *handle = TFM_CRYPTO_INVALID_HANDLE;
         return PSA_SUCCESS;
     }
@@ -137,11 +141,11 @@
                                          void **ctx)
 {
     if ( (handle != TFM_CRYPTO_INVALID_HANDLE) &&
-         (handle < TFM_CRYPTO_CONC_OPER_NUM) &&
-         (operation[handle].in_use == TFM_CRYPTO_IN_USE) &&
-         (operation[handle].type == type) ) {
+         (handle <= TFM_CRYPTO_CONC_OPER_NUM) &&
+         (operation[handle - 1].in_use == TFM_CRYPTO_IN_USE) &&
+         (operation[handle - 1].type == type) ) {
 
-        *ctx = (void *) &(operation[handle].operation);
+        *ctx = (void *) &(operation[handle - 1].operation);
         return PSA_SUCCESS;
     }
 
diff --git a/secure_fw/services/crypto/crypto_asymmetric.c b/secure_fw/services/crypto/crypto_asymmetric.c
index 3db7938..cd58b11 100644
--- a/secure_fw/services/crypto/crypto_asymmetric.c
+++ b/secure_fw/services/crypto/crypto_asymmetric.c
@@ -79,6 +79,8 @@
                                            psa_outvec out_vec[],
                                            size_t out_len)
 {
+    psa_status_t status;
+
     if (!((in_len == 2) || (in_len == 3)) || (out_len != 1)) {
         return PSA_CONNECTION_REFUSED;
     }
@@ -96,12 +98,24 @@
     size_t salt_length = 0;
     uint8_t *output = out_vec[0].base;
     size_t output_size = out_vec[0].len;
+    psa_key_type_t type;
+    size_t key_bits;
 
     if (in_len == 3) {
         salt = in_vec[2].base;
         salt_length = in_vec[2].len;
     }
 
+    status = psa_get_key_information(handle, &type, &key_bits);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    /* Check that the output buffer is large enough */
+    if (output_size < PSA_ASYMMETRIC_ENCRYPT_OUTPUT_SIZE(type, key_bits, alg)) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+
     return psa_asymmetric_encrypt(handle, alg, input, input_length,
                                   salt, salt_length,
                                   output, output_size, &(out_vec[0].len));
diff --git a/secure_fw/services/crypto/crypto_generator.c b/secure_fw/services/crypto/crypto_generator.c
index 1cf138d..4328ca8 100644
--- a/secure_fw/services/crypto/crypto_generator.c
+++ b/secure_fw/services/crypto/crypto_generator.c
@@ -156,6 +156,8 @@
 
     status = psa_generator_abort(generator);
     if (status != PSA_SUCCESS) {
+        /* Release the operation context, ignore if the operation fails. */
+        (void)tfm_crypto_operation_release(handle_out);
         return status;
     }
 
@@ -211,8 +213,14 @@
 
     *handle_out = handle;
 
-    return psa_key_derivation(generator, key_handle, alg, salt, salt_length,
-                              label, label_length, capacity);
+    status = psa_key_derivation(generator, key_handle, alg, salt, salt_length,
+                                label, label_length, capacity);
+    if (status != PSA_SUCCESS) {
+        /* Release the operation context, ignore if the operation fails. */
+        (void)tfm_crypto_operation_release(handle_out);
+    }
+
+    return status;
 }
 
 psa_status_t tfm_crypto_key_agreement(psa_invec in_vec[],
@@ -249,8 +257,14 @@
 
     *handle_out = handle;
 
-    return psa_key_agreement(generator, private_key,
-                             peer_key, peer_key_length, alg);
+    status = psa_key_agreement(generator, private_key,
+                               peer_key, peer_key_length, alg);
+    if (status != PSA_SUCCESS) {
+        /* Release the operation context, ignore if the operation fails. */
+        (void)tfm_crypto_operation_release(handle_out);
+    }
+
+    return status;
 }
 
 psa_status_t tfm_crypto_generate_random(psa_invec in_vec[],