Improve interruptible key agreement implementation
Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index 840bb6d..1ab2dbb 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -7791,23 +7791,34 @@
size_t key_size = 0;
size_t key_len = 0;
psa_key_attributes_t private_key_attributes;
+ psa_key_type_t private_key_type;
+
+ if (operation->id != 0 || operation->error_occurred) {
+ return PSA_ERROR_BAD_STATE;
+ }
status = psa_get_key_attributes(private_key, &private_key_attributes);
if (status != PSA_SUCCESS) {
+ operation->error_occurred = 1;
return status;
}
- if (!PSA_KEY_TYPE_IS_ECC_KEY_PAIR(private_key_attributes.type) ||
+ private_key_type = psa_get_key_type(&private_key_attributes);
+ if (!PSA_KEY_TYPE_IS_ECC_KEY_PAIR(private_key_type) ||
!PSA_ALG_IS_ECDH(alg)) {
+ operation->error_occurred = 1;
return PSA_ERROR_INVALID_ARGUMENT;
}
- key_size = PSA_BITS_TO_BYTES(psa_get_key_bits(&private_key_attributes));
+ key_size = PSA_EXPORT_KEY_OUTPUT_SIZE(private_key_type,
+ psa_get_key_bits(&private_key_attributes));
if (key_size == 0) {
+ operation->error_occurred = 1;
return PSA_ERROR_INVALID_ARGUMENT;
}
private_key_buffer = mbedtls_calloc(key_size, 1);
if (private_key_buffer == NULL) {
+ operation->error_occurred = 1;
return PSA_ERROR_INSUFFICIENT_MEMORY;
}
@@ -7825,8 +7836,11 @@
exit:
mbedtls_free(private_key_buffer);
+ if (status != PSA_SUCCESS) {
+ operation->error_occurred = 1;
+ }
return status;
- #else
+#else
(void) operation;
(void) private_key;
(void) peer_key;
@@ -7834,7 +7848,7 @@
(void) alg;
(void) attributes;
return PSA_ERROR_NOT_SUPPORTED;
- #endif
+#endif
}
psa_status_t psa_key_agreement_iop_complete(
@@ -7844,6 +7858,10 @@
#if defined(MBEDTLS_ECP_RESTARTABLE) && \
defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH)
+ if (operation->id == 0 || operation->error_occurred) {
+ return PSA_ERROR_BAD_STATE;
+ }
+
psa_status_t status;
uint8_t intermediate_key[PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE];
size_t key_len = 0;
@@ -7859,6 +7877,9 @@
key_len, key);
}
+ if (status != PSA_SUCCESS && status != PSA_OPERATION_INCOMPLETE) {
+ operation->error_occurred = 1;
+ }
mbedtls_platform_zeroize(intermediate_key, PSA_RAW_KEY_AGREEMENT_OUTPUT_MAX_SIZE);
return status;
#else
@@ -7876,14 +7897,12 @@
psa_status_t status;
status = psa_driver_wrapper_key_agreement_abort(operation);
-
+ operation->error_occurred = 0;
return status;
#else
(void) operation;
return PSA_ERROR_NOT_SUPPORTED;
#endif
-
-
}
/****************************************************************/
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 77c8ac2..6556c6e 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -12,7 +12,6 @@
#include <psa/crypto.h>
#include "psa_crypto_core.h"
-#include "psa_crypto_driver_wrappers.h"
#include "psa_crypto_ecp.h"
#include "psa_crypto_random_impl.h"
#include "mbedtls/psa_util.h"
@@ -659,8 +658,8 @@
size_t bits = 0;
status = mbedtls_psa_ecp_load_representation(
- attributes->type,
- attributes->bits,
+ psa_get_key_type(attributes),
+ psa_get_key_bits(attributes),
private_key_buffer,
private_key_buffer_len,
&ecp);
diff --git a/tf-psa-crypto/include/psa/crypto_struct.h b/tf-psa-crypto/include/psa/crypto_struct.h
index b300557..b9e0f4e 100644
--- a/tf-psa-crypto/include/psa/crypto_struct.h
+++ b/tf-psa-crypto/include/psa/crypto_struct.h
@@ -510,13 +510,14 @@
unsigned int MBEDTLS_PRIVATE(id);
psa_driver_key_agreement_interruptible_context_t MBEDTLS_PRIVATE(ctx);
uint32_t MBEDTLS_PRIVATE(num_ops);
+ unsigned int MBEDTLS_PRIVATE(error_occurred) : 1;
#endif
};
#if defined(MBEDTLS_PSA_CRYPTO_CLIENT) && !defined(MBEDTLS_PSA_CRYPTO_C)
#define PSA_KEY_AGREEMENT_IOP_INIT { 0 }
#else
-#define PSA_KEY_AGREEMENT_IOP_INIT { 0, { 0 }, 0 }
+#define PSA_KEY_AGREEMENT_IOP_INIT { 0, { 0 }, 0, 0 }
#endif
static inline struct psa_key_agreement_iop_s