Refactor mbedtls_psa_generate_key_iop_complete()

- Move the checks on the size to the start of the
  function to avaoid costly calls to mbedtls_ecp_gen_privkey()
  in case of invalid size.
- Improve the readability of error checking

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
index 82e8736..4500196 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -625,14 +625,6 @@
 {
     *key_len = 0;
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    status = mbedtls_ecp_gen_privkey(&operation->ecp.grp, &operation->ecp.d,
-                                     mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE);
-
-    if (status) {
-        return mbedtls_to_psa_error(status);
-    }
-
-    operation->num_ops = 1;
 
     *key_len = PSA_BITS_TO_BYTES(operation->ecp.grp.nbits);
 
@@ -640,6 +632,15 @@
         return PSA_ERROR_BUFFER_TOO_SMALL;
     }
 
+    status = mbedtls_ecp_gen_privkey(&operation->ecp.grp, &operation->ecp.d,
+                                     mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE);
+
+    if (status != 0) {
+        return mbedtls_to_psa_error(status);
+    }
+
+    operation->num_ops = 1;
+
     mbedtls_mpi_write_binary(&operation->ecp.d, key_output, key_output_size);
 
     return mbedtls_to_psa_error(status);