From 902e94cad45ad9f009a4c61ac31d3a210ec69668 Mon Sep 17 00:00:00 2001
From: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Date: Fri, 22 Sep 2023 17:54:59 +0100
Subject: [PATCH 1/7] docs(changelog): add scope for MbedTLS PSA Crypto

Added scope for MbedTLS PSA Crypto API implementation.

Change-Id: I0a7c6242841dfb94f6ca5e9174ffd520c5fe00fa
Signed-off-by: Manish V Badarkhe <Manish.Badarkhe@arm.com>
---
 changelog.yaml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/changelog.yaml b/changelog.yaml
index 3f7979e40..2d6986b77 100644
--- a/changelog.yaml
+++ b/changelog.yaml
@@ -809,6 +809,9 @@ subsections:
           - title: mbedTLS
             scope: mbedtls
 
+          - title: mbedTLS-PSA
+            scope: mbedtls-psa
+
       - title: Console
         scope: console
 

From 5782b890d29646924d8bd3f46acdc73a6e02feb2 Mon Sep 17 00:00:00 2001
From: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Date: Wed, 6 Sep 2023 09:08:28 +0100
Subject: [PATCH 2/7] feat(mbedtls-psa): introduce PSA_CRYPTO build option

This is a preparatory patch to provide MbedTLS PSA Crypto
API support, with below changes -

1. Added a build macro PSA_CRYPTO to enable the MbedTLS PSA
   Crypto API support in the subsequent patches.
2. Compile necessary PSA crypto files from MbedTLS source code
   when PSA_CRYPTO=1.

Also, marked PSA_CRYPTO as an experimental feature.

Change-Id: I45188f56c5c98b169b2e21e365150b1825c6c450
Signed-off-by: Manish V Badarkhe <Manish.Badarkhe@arm.com>
---
 Makefile                                       |  6 ++++++
 docs/getting_started/build-options.rst         |  7 +++++++
 drivers/auth/mbedtls/mbedtls_common.mk         | 18 +++++++++++++++++-
 .../drivers/auth/mbedtls/psa_mbedtls_config.h  | 14 ++++++++++++++
 make_helpers/defaults.mk                       |  3 +++
 5 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 include/drivers/auth/mbedtls/psa_mbedtls_config.h

diff --git a/Makefile b/Makefile
index 8e8fba901..1e6dd9071 100644
--- a/Makefile
+++ b/Makefile
@@ -1036,6 +1036,10 @@ ENABLE_FEAT_RNG		=	$(if $(findstring rng,${arch-features}),1,0)
 # Determine if FEAT_SB is supported
 ENABLE_FEAT_SB		=	$(if $(findstring sb,${arch-features}),1,0)
 
+ifeq ($(PSA_CRYPTO),1)
+        $(info PSA_CRYPTO is an experimental feature)
+endif
+
 ################################################################################
 # Process platform overrideable behaviour
 ################################################################################
@@ -1217,6 +1221,7 @@ $(eval $(call assert_booleans,\
 	ERRATA_NON_ARM_INTERCONNECT \
 	CONDITIONAL_CMO \
 	RAS_FFH_SUPPORT \
+	PSA_CRYPTO	\
 )))
 
 # Numeric_Flags
@@ -1407,6 +1412,7 @@ $(eval $(call add_defines,\
 	IMPDEF_SYSREG_TRAP \
 	SVE_VECTOR_LEN \
 	ENABLE_SPMD_LP \
+	PSA_CRYPTO	\
 )))
 
 ifeq (${SANITIZE_UB},trap)
diff --git a/docs/getting_started/build-options.rst b/docs/getting_started/build-options.rst
index 7c84ef163..34d83f255 100644
--- a/docs/getting_started/build-options.rst
+++ b/docs/getting_started/build-options.rst
@@ -1185,6 +1185,12 @@ Common build options
   errata mitigation for platforms with a non-arm interconnect using the errata
   ABI. By default its disabled (``0``).
 
+- ``PSA_CRYPTO``: Boolean option for enabling MbedTLS PSA crypto APIs support.
+  The platform will use PSA compliant Crypto APIs during authentication and
+  image measurement process by enabling this option. It uses APIs defined as
+  per the `PSA Crypto API specification`_. This feature is only supported if
+  using MbedTLS 3.x version. By default it is disabled (``0``).
+
 GICv3 driver options
 --------------------
 
@@ -1306,3 +1312,4 @@ Firmware update options
 .. _GCC: https://gcc.gnu.org/onlinedocs/gcc/Warning-Options.html
 .. _Clang: https://clang.llvm.org/docs/DiagnosticsReference.html
 .. _Firmware Handoff specification: https://github.com/FirmwareHandoff/firmware_handoff/releases/tag/v0.9
+.. _PSA Crypto API specification: https://armmbed.github.io/mbed-crypto/html/
diff --git a/drivers/auth/mbedtls/mbedtls_common.mk b/drivers/auth/mbedtls/mbedtls_common.mk
index 79c45126a..376b6b7ca 100644
--- a/drivers/auth/mbedtls/mbedtls_common.mk
+++ b/drivers/auth/mbedtls/mbedtls_common.mk
@@ -23,7 +23,11 @@ $(info MBEDTLS_VERSION_MAJOR is [${MBEDTLS_MAJOR}] MBEDTLS_VERSION_MINOR is [${M
 ifeq (${MBEDTLS_MAJOR}, 2)
 	MBEDTLS_CONFIG_FILE	?=	"<drivers/auth/mbedtls/mbedtls_config-2.h>"
 else ifeq (${MBEDTLS_MAJOR}, 3)
-	MBEDTLS_CONFIG_FILE	?=	"<drivers/auth/mbedtls/mbedtls_config-3.h>"
+	ifeq (${PSA_CRYPTO},1)
+		MBEDTLS_CONFIG_FILE     ?=      "<drivers/auth/mbedtls/psa_mbedtls_config.h>"
+	else
+		MBEDTLS_CONFIG_FILE	?=	"<drivers/auth/mbedtls/mbedtls_config-3.h>"
+	endif
 endif
 
 $(eval $(call add_define,MBEDTLS_CONFIG_FILE))
@@ -77,6 +81,18 @@ else ifeq (${MBEDTLS_MAJOR}, 3)
 	LIBMBEDTLS_CFLAGS += -Wno-error=redundant-decls
 endif
 
+ifeq (${PSA_CRYPTO},1)
+LIBMBEDTLS_SRCS         += $(addprefix ${MBEDTLS_DIR}/library/,    	\
+					psa_crypto.c                   	\
+					psa_crypto_client.c            	\
+					psa_crypto_driver_wrappers.c   	\
+					psa_crypto_hash.c              	\
+					psa_crypto_rsa.c               	\
+					psa_crypto_ecp.c               	\
+					psa_crypto_slot_management.c   	\
+					)
+endif
+
 # The platform may define the variable 'TF_MBEDTLS_KEY_ALG' to select the key
 # algorithm to use. If the variable is not defined, select it based on
 # algorithm used for key generation `KEY_ALG`. If `KEY_ALG` is not defined,
diff --git a/include/drivers/auth/mbedtls/psa_mbedtls_config.h b/include/drivers/auth/mbedtls/psa_mbedtls_config.h
new file mode 100644
index 000000000..65df2d4da
--- /dev/null
+++ b/include/drivers/auth/mbedtls/psa_mbedtls_config.h
@@ -0,0 +1,14 @@
+/*
+ * Copyright (c) 2023, Arm Ltd. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#ifndef PSA_MBEDTLS_CONFIG_H
+#define PSA_MBEDTLS_CONFIG_H
+
+#include "mbedtls_config-3.h"
+
+#define MBEDTLS_PSA_CRYPTO_C
+
+#endif /* PSA_MBEDTLS_CONFIG_H */
diff --git a/make_helpers/defaults.mk b/make_helpers/defaults.mk
index b7e6f9917..321ae9fd1 100644
--- a/make_helpers/defaults.mk
+++ b/make_helpers/defaults.mk
@@ -359,3 +359,6 @@ CONDITIONAL_CMO			:= 0
 
 # By default, disable SPMD Logical partitions
 ENABLE_SPMD_LP			:= 0
+
+# By default, disable PSA crypto (use MbedTLS legacy crypto API).
+PSA_CRYPTO			:= 0

From 38f893692ad9b8edb5413f4b2b9cd15a9b485685 Mon Sep 17 00:00:00 2001
From: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Date: Wed, 6 Sep 2023 10:22:19 +0100
Subject: [PATCH 3/7] feat(mbedtls-psa): register an ad-hoc PSA crypto driver

An ad-hoc MbedTLS PSA crypto driver is registered by compiling
a new driver file, namely mbedtls_psa_crypto.c when PSA_CRYPTO=1.
As of now, this file is the same as mbedtls_crypto.c, but subsequent
patches will update crypto functions in this file to trigger
PSA crypto APIs.

Change-Id: I404c347990661d87dcf5d0501d238e36914ec3ee
Signed-off-by: Manish V Badarkhe <Manish.Badarkhe@arm.com>
---
 drivers/auth/mbedtls/mbedtls_crypto.mk    |  13 +-
 drivers/auth/mbedtls/mbedtls_psa_crypto.c | 423 ++++++++++++++++++++++
 2 files changed, 432 insertions(+), 4 deletions(-)
 create mode 100644 drivers/auth/mbedtls/mbedtls_psa_crypto.c

diff --git a/drivers/auth/mbedtls/mbedtls_crypto.mk b/drivers/auth/mbedtls/mbedtls_crypto.mk
index 2a9fbbf99..bd367300b 100644
--- a/drivers/auth/mbedtls/mbedtls_crypto.mk
+++ b/drivers/auth/mbedtls/mbedtls_crypto.mk
@@ -1,11 +1,16 @@
 #
-# Copyright (c) 2015-2017, ARM Limited and Contributors. All rights reserved.
+# Copyright (c) 2015-2023, Arm Limited and Contributors. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
 include drivers/auth/mbedtls/mbedtls_common.mk
 
-MBEDTLS_SOURCES	+=		drivers/auth/mbedtls/mbedtls_crypto.c
-
-
+ifeq (${PSA_CRYPTO},1)
+	# Some of the PSA functions are declared in multiple header files
+	# that triggers this warning.
+	TF_CFLAGS      	+=     	-Wno-error=redundant-decls
+	MBEDTLS_SOURCES +=	drivers/auth/mbedtls/mbedtls_psa_crypto.c
+else
+	MBEDTLS_SOURCES +=	drivers/auth/mbedtls/mbedtls_crypto.c
+endif
diff --git a/drivers/auth/mbedtls/mbedtls_psa_crypto.c b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
new file mode 100644
index 000000000..eaeca79c7
--- /dev/null
+++ b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
@@ -0,0 +1,423 @@
+/*
+ * Copyright (c) 2023, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <assert.h>
+#include <stddef.h>
+#include <string.h>
+
+/* mbed TLS headers */
+#include <mbedtls/gcm.h>
+#include <mbedtls/md.h>
+#include <mbedtls/memory_buffer_alloc.h>
+#include <mbedtls/oid.h>
+#include <mbedtls/platform.h>
+#include <mbedtls/version.h>
+#include <mbedtls/x509.h>
+
+#include <common/debug.h>
+#include <drivers/auth/crypto_mod.h>
+#include <drivers/auth/mbedtls/mbedtls_common.h>
+
+#include <plat/common/platform.h>
+
+#define LIB_NAME		"mbed TLS PSA"
+
+#if CRYPTO_SUPPORT == CRYPTO_HASH_CALC_ONLY || \
+CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC
+/*
+ * CRYPTO_MD_MAX_SIZE value is as per current stronger algorithm available
+ * so make sure that mbed TLS MD maximum size must be lesser than this.
+ */
+CASSERT(CRYPTO_MD_MAX_SIZE >= MBEDTLS_MD_MAX_SIZE,
+	assert_mbedtls_md_size_overflow);
+
+#endif /*
+	* CRYPTO_SUPPORT == CRYPTO_HASH_CALC_ONLY || \
+	* CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC
+	*/
+
+/*
+ * AlgorithmIdentifier  ::=  SEQUENCE  {
+ *     algorithm               OBJECT IDENTIFIER,
+ *     parameters              ANY DEFINED BY algorithm OPTIONAL
+ * }
+ *
+ * SubjectPublicKeyInfo  ::=  SEQUENCE  {
+ *     algorithm            AlgorithmIdentifier,
+ *     subjectPublicKey     BIT STRING
+ * }
+ *
+ * DigestInfo ::= SEQUENCE {
+ *     digestAlgorithm AlgorithmIdentifier,
+ *     digest OCTET STRING
+ * }
+ */
+
+/*
+ * Initialize the library and export the descriptor
+ */
+static void init(void)
+{
+	/* Initialize mbed TLS */
+	mbedtls_init();
+}
+
+#if CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_ONLY || \
+CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC
+/*
+ * Verify a signature.
+ *
+ * Parameters are passed using the DER encoding format following the ASN.1
+ * structures detailed above.
+ */
+static int verify_signature(void *data_ptr, unsigned int data_len,
+			    void *sig_ptr, unsigned int sig_len,
+			    void *sig_alg, unsigned int sig_alg_len,
+			    void *pk_ptr, unsigned int pk_len)
+{
+	mbedtls_asn1_buf sig_oid, sig_params;
+	mbedtls_asn1_buf signature;
+	mbedtls_md_type_t md_alg;
+	mbedtls_pk_type_t pk_alg;
+	mbedtls_pk_context pk = {0};
+	int rc;
+	void *sig_opts = NULL;
+	const mbedtls_md_info_t *md_info;
+	unsigned char *p, *end;
+	unsigned char hash[MBEDTLS_MD_MAX_SIZE];
+
+	/* Get pointers to signature OID and parameters */
+	p = (unsigned char *)sig_alg;
+	end = (unsigned char *)(p + sig_alg_len);
+	rc = mbedtls_asn1_get_alg(&p, end, &sig_oid, &sig_params);
+	if (rc != 0) {
+		return CRYPTO_ERR_SIGNATURE;
+	}
+
+	/* Get the actual signature algorithm (MD + PK) */
+	rc = mbedtls_x509_get_sig_alg(&sig_oid, &sig_params, &md_alg, &pk_alg, &sig_opts);
+	if (rc != 0) {
+		return CRYPTO_ERR_SIGNATURE;
+	}
+
+	/* Parse the public key */
+	mbedtls_pk_init(&pk);
+	p = (unsigned char *)pk_ptr;
+	end = (unsigned char *)(p + pk_len);
+	rc = mbedtls_pk_parse_subpubkey(&p, end, &pk);
+	if (rc != 0) {
+		rc = CRYPTO_ERR_SIGNATURE;
+		goto end2;
+	}
+
+	/* Get the signature (bitstring) */
+	p = (unsigned char *)sig_ptr;
+	end = (unsigned char *)(p + sig_len);
+	signature.tag = *p;
+	rc = mbedtls_asn1_get_bitstring_null(&p, end, &signature.len);
+	if ((rc != 0) || ((size_t)(end - p) != signature.len)) {
+		rc = CRYPTO_ERR_SIGNATURE;
+		goto end1;
+	}
+	signature.p = p;
+
+	/* Calculate the hash of the data */
+	md_info = mbedtls_md_info_from_type(md_alg);
+	if (md_info == NULL) {
+		rc = CRYPTO_ERR_SIGNATURE;
+		goto end1;
+	}
+	p = (unsigned char *)data_ptr;
+	rc = mbedtls_md(md_info, p, data_len, hash);
+	if (rc != 0) {
+		rc = CRYPTO_ERR_SIGNATURE;
+		goto end1;
+	}
+
+	/* Verify the signature */
+	rc = mbedtls_pk_verify_ext(pk_alg, sig_opts, &pk, md_alg, hash,
+			mbedtls_md_get_size(md_info),
+			signature.p, signature.len);
+	if (rc != 0) {
+		rc = CRYPTO_ERR_SIGNATURE;
+		goto end1;
+	}
+
+	/* Signature verification success */
+	rc = CRYPTO_SUCCESS;
+
+end1:
+	mbedtls_pk_free(&pk);
+end2:
+	mbedtls_free(sig_opts);
+	return rc;
+}
+
+/*
+ * Match a hash
+ *
+ * Digest info is passed in DER format following the ASN.1 structure detailed
+ * above.
+ */
+static int verify_hash(void *data_ptr, unsigned int data_len,
+		       void *digest_info_ptr, unsigned int digest_info_len)
+{
+	mbedtls_asn1_buf hash_oid, params;
+	mbedtls_md_type_t md_alg;
+	const mbedtls_md_info_t *md_info;
+	unsigned char *p, *end, *hash;
+	unsigned char data_hash[MBEDTLS_MD_MAX_SIZE];
+	size_t len;
+	int rc;
+
+	/*
+	 * Digest info should be an MBEDTLS_ASN1_SEQUENCE, but padding after
+	 * it is allowed.  This is necessary to support multiple hash
+	 * algorithms.
+	 */
+	p = (unsigned char *)digest_info_ptr;
+	end = p + digest_info_len;
+	rc = mbedtls_asn1_get_tag(&p, end, &len, MBEDTLS_ASN1_CONSTRUCTED |
+				  MBEDTLS_ASN1_SEQUENCE);
+	if (rc != 0) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	end = p + len;
+
+	/* Get the hash algorithm */
+	rc = mbedtls_asn1_get_alg(&p, end, &hash_oid, &params);
+	if (rc != 0) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	rc = mbedtls_oid_get_md_alg(&hash_oid, &md_alg);
+	if (rc != 0) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	md_info = mbedtls_md_info_from_type(md_alg);
+	if (md_info == NULL) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	/* Hash should be octet string type and consume all bytes */
+	rc = mbedtls_asn1_get_tag(&p, end, &len, MBEDTLS_ASN1_OCTET_STRING);
+	if ((rc != 0) || ((size_t)(end - p) != len)) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	/* Length of hash must match the algorithm's size */
+	if (len != mbedtls_md_get_size(md_info)) {
+		return CRYPTO_ERR_HASH;
+	}
+	hash = p;
+
+	/* Calculate the hash of the data */
+	p = (unsigned char *)data_ptr;
+	rc = mbedtls_md(md_info, p, data_len, data_hash);
+	if (rc != 0) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	/* Compare values */
+	rc = memcmp(data_hash, hash, mbedtls_md_get_size(md_info));
+	if (rc != 0) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	return CRYPTO_SUCCESS;
+}
+#endif /*
+	* CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_ONLY || \
+	* CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC
+	*/
+
+#if CRYPTO_SUPPORT == CRYPTO_HASH_CALC_ONLY || \
+CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC
+/*
+ * Map a generic crypto message digest algorithm to the corresponding macro used
+ * by Mbed TLS.
+ */
+static inline mbedtls_md_type_t md_type(enum crypto_md_algo algo)
+{
+	switch (algo) {
+	case CRYPTO_MD_SHA512:
+		return MBEDTLS_MD_SHA512;
+	case CRYPTO_MD_SHA384:
+		return MBEDTLS_MD_SHA384;
+	case CRYPTO_MD_SHA256:
+		return MBEDTLS_MD_SHA256;
+	default:
+		/* Invalid hash algorithm. */
+		return MBEDTLS_MD_NONE;
+	}
+}
+
+/*
+ * Calculate a hash
+ *
+ * output points to the computed hash
+ */
+static int calc_hash(enum crypto_md_algo md_algo, void *data_ptr,
+		     unsigned int data_len,
+		     unsigned char output[CRYPTO_MD_MAX_SIZE])
+{
+	const mbedtls_md_info_t *md_info;
+
+	md_info = mbedtls_md_info_from_type(md_type(md_algo));
+	if (md_info == NULL) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	/*
+	 * Calculate the hash of the data, it is safe to pass the
+	 * 'output' hash buffer pointer considering its size is always
+	 * bigger than or equal to MBEDTLS_MD_MAX_SIZE.
+	 */
+	return mbedtls_md(md_info, data_ptr, data_len, output);
+}
+#endif /*
+	* CRYPTO_SUPPORT == CRYPTO_HASH_CALC_ONLY || \
+	* CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC
+	*/
+
+#if TF_MBEDTLS_USE_AES_GCM
+/*
+ * Stack based buffer allocation for decryption operation. It could
+ * be configured to balance stack usage vs execution speed.
+ */
+#define DEC_OP_BUF_SIZE		128
+
+static int aes_gcm_decrypt(void *data_ptr, size_t len, const void *key,
+			   unsigned int key_len, const void *iv,
+			   unsigned int iv_len, const void *tag,
+			   unsigned int tag_len)
+{
+	mbedtls_gcm_context ctx;
+	mbedtls_cipher_id_t cipher = MBEDTLS_CIPHER_ID_AES;
+	unsigned char buf[DEC_OP_BUF_SIZE];
+	unsigned char tag_buf[CRYPTO_MAX_TAG_SIZE];
+	unsigned char *pt = data_ptr;
+	size_t dec_len;
+	int diff, i, rc;
+	size_t output_length __unused;
+
+	mbedtls_gcm_init(&ctx);
+
+	rc = mbedtls_gcm_setkey(&ctx, cipher, key, key_len * 8);
+	if (rc != 0) {
+		rc = CRYPTO_ERR_DECRYPTION;
+		goto exit_gcm;
+	}
+
+#if (MBEDTLS_VERSION_MAJOR < 3)
+	rc = mbedtls_gcm_starts(&ctx, MBEDTLS_GCM_DECRYPT, iv, iv_len, NULL, 0);
+#else
+	rc = mbedtls_gcm_starts(&ctx, MBEDTLS_GCM_DECRYPT, iv, iv_len);
+#endif
+	if (rc != 0) {
+		rc = CRYPTO_ERR_DECRYPTION;
+		goto exit_gcm;
+	}
+
+	while (len > 0) {
+		dec_len = MIN(sizeof(buf), len);
+
+#if (MBEDTLS_VERSION_MAJOR < 3)
+		rc = mbedtls_gcm_update(&ctx, dec_len, pt, buf);
+#else
+		rc = mbedtls_gcm_update(&ctx, pt, dec_len, buf, sizeof(buf), &output_length);
+#endif
+
+		if (rc != 0) {
+			rc = CRYPTO_ERR_DECRYPTION;
+			goto exit_gcm;
+		}
+
+		memcpy(pt, buf, dec_len);
+		pt += dec_len;
+		len -= dec_len;
+	}
+
+#if (MBEDTLS_VERSION_MAJOR < 3)
+	rc = mbedtls_gcm_finish(&ctx, tag_buf, sizeof(tag_buf));
+#else
+	rc = mbedtls_gcm_finish(&ctx, NULL, 0, &output_length, tag_buf, sizeof(tag_buf));
+#endif
+
+	if (rc != 0) {
+		rc = CRYPTO_ERR_DECRYPTION;
+		goto exit_gcm;
+	}
+
+	/* Check tag in "constant-time" */
+	for (diff = 0, i = 0; i < tag_len; i++)
+		diff |= ((const unsigned char *)tag)[i] ^ tag_buf[i];
+
+	if (diff != 0) {
+		rc = CRYPTO_ERR_DECRYPTION;
+		goto exit_gcm;
+	}
+
+	/* GCM decryption success */
+	rc = CRYPTO_SUCCESS;
+
+exit_gcm:
+	mbedtls_gcm_free(&ctx);
+	return rc;
+}
+
+/*
+ * Authenticated decryption of an image
+ */
+static int auth_decrypt(enum crypto_dec_algo dec_algo, void *data_ptr,
+			size_t len, const void *key, unsigned int key_len,
+			unsigned int key_flags, const void *iv,
+			unsigned int iv_len, const void *tag,
+			unsigned int tag_len)
+{
+	int rc;
+
+	assert((key_flags & ENC_KEY_IS_IDENTIFIER) == 0);
+
+	switch (dec_algo) {
+	case CRYPTO_GCM_DECRYPT:
+		rc = aes_gcm_decrypt(data_ptr, len, key, key_len, iv, iv_len,
+				     tag, tag_len);
+		if (rc != 0)
+			return rc;
+		break;
+	default:
+		return CRYPTO_ERR_DECRYPTION;
+	}
+
+	return CRYPTO_SUCCESS;
+}
+#endif /* TF_MBEDTLS_USE_AES_GCM */
+
+/*
+ * Register crypto library descriptor
+ */
+#if CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC
+#if TF_MBEDTLS_USE_AES_GCM
+REGISTER_CRYPTO_LIB(LIB_NAME, init, verify_signature, verify_hash, calc_hash,
+		    auth_decrypt, NULL);
+#else
+REGISTER_CRYPTO_LIB(LIB_NAME, init, verify_signature, verify_hash, calc_hash,
+		    NULL, NULL);
+#endif
+#elif CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_ONLY
+#if TF_MBEDTLS_USE_AES_GCM
+REGISTER_CRYPTO_LIB(LIB_NAME, init, verify_signature, verify_hash, NULL,
+		    auth_decrypt, NULL);
+#else
+REGISTER_CRYPTO_LIB(LIB_NAME, init, verify_signature, verify_hash, NULL,
+		    NULL, NULL);
+#endif
+#elif CRYPTO_SUPPORT == CRYPTO_HASH_CALC_ONLY
+REGISTER_CRYPTO_LIB(LIB_NAME, init, NULL, NULL, calc_hash, NULL, NULL);
+#endif /* CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC */

From 4eaaaa19299040cfee0585d7daa744dee716d398 Mon Sep 17 00:00:00 2001
From: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Date: Wed, 6 Sep 2023 11:01:37 +0100
Subject: [PATCH 4/7] feat(mbedtls-psa): initialise mbedtls psa crypto

Initialised Mbedtls PSA cryto during Crypto init using
function call 'psa_crypto_init'.

MbedTLS currently requires a Random Number Generator (RNG) once
PSA Crypto support is enabled. However, TF-A itself doesn't engage
in cryptographic operations that demand randomness. Consequently,
we simulate the presence of an external TRNG (through the configuration
option 'MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG) while, in reality, we offer
a dummy implementation of mbedtls_psa_external_get_random() that always
returns an error.

Change-Id: Ife6d03909c0e6081438d2b2519ef500e5dcdb88f
Signed-off-by: Manish V Badarkhe <Manish.Badarkhe@arm.com>
---
 drivers/auth/mbedtls/mbedtls_psa_crypto.c     | 30 ++++++++++++++++++-
 .../drivers/auth/mbedtls/psa_mbedtls_config.h | 14 +++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/drivers/auth/mbedtls/mbedtls_psa_crypto.c b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
index eaeca79c7..64778e8da 100644
--- a/drivers/auth/mbedtls/mbedtls_psa_crypto.c
+++ b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
@@ -16,11 +16,14 @@
 #include <mbedtls/platform.h>
 #include <mbedtls/version.h>
 #include <mbedtls/x509.h>
+#include <psa/crypto.h>
+#include <psa/crypto_platform.h>
+#include <psa/crypto_types.h>
+#include <psa/crypto_values.h>
 
 #include <common/debug.h>
 #include <drivers/auth/crypto_mod.h>
 #include <drivers/auth/mbedtls/mbedtls_common.h>
-
 #include <plat/common/platform.h>
 
 #define LIB_NAME		"mbed TLS PSA"
@@ -56,6 +59,21 @@ CASSERT(CRYPTO_MD_MAX_SIZE >= MBEDTLS_MD_MAX_SIZE,
  * }
  */
 
+/*
+ * We pretend using an external RNG (through MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG
+ * mbedTLS config option) so we need to provide an implementation of
+ * mbedtls_psa_external_get_random(). Provide a fake one, since we do not
+ * actually have any external RNG and TF-A itself doesn't engage in
+ * cryptographic operations that demands randomness.
+ */
+psa_status_t mbedtls_psa_external_get_random(
+			mbedtls_psa_external_random_context_t *context,
+			uint8_t *output, size_t output_size,
+			size_t *output_length)
+{
+	return PSA_ERROR_INSUFFICIENT_ENTROPY;
+}
+
 /*
  * Initialize the library and export the descriptor
  */
@@ -63,6 +81,16 @@ static void init(void)
 {
 	/* Initialize mbed TLS */
 	mbedtls_init();
+
+	/* Initialise PSA mbedTLS */
+	psa_status_t status = psa_crypto_init();
+
+	if (status != PSA_SUCCESS) {
+		ERROR("Failed to initialize %s crypto (%d).\n", LIB_NAME, status);
+		panic();
+	}
+
+	INFO("PSA crypto initialized successfully!\n");
 }
 
 #if CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_ONLY || \
diff --git a/include/drivers/auth/mbedtls/psa_mbedtls_config.h b/include/drivers/auth/mbedtls/psa_mbedtls_config.h
index 65df2d4da..ad825f0aa 100644
--- a/include/drivers/auth/mbedtls/psa_mbedtls_config.h
+++ b/include/drivers/auth/mbedtls/psa_mbedtls_config.h
@@ -11,4 +11,18 @@
 
 #define MBEDTLS_PSA_CRYPTO_C
 
+/*
+ * Using PSA crypto API requires an RNG right now. If we don't define the macro
+ * below then we get build errors.
+ *
+ * This is a functionality gap in mbedTLS. The technical limitation is that
+ * psa_crypto_init() is all-or-nothing, and fixing that would require separate
+ * initialization of the keystore, the RNG, etc.
+ *
+ * By defining MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG, we pretend using an external
+ * RNG. As a result, the PSA crypto init code does nothing when it comes to
+ * initializing the RNG, as we are supposed to take care of that ourselves.
+ */
+#define MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG
+
 #endif /* PSA_MBEDTLS_CONFIG_H */

From 2ed061c43525b8a9cd82b38d31277a8df594edd5 Mon Sep 17 00:00:00 2001
From: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Date: Wed, 6 Sep 2023 12:00:20 +0100
Subject: [PATCH 5/7] feat(mbedtls-psa): use PSA crypto API for hash
 verification

A PSA crypto API is used for hash verification, namely
psa_hash_verify, which internally calculates and verifies the hash
of the given data payload.

Signed-off-by: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Change-Id: Ib48aa6b74b59aea6036333ff00f6ca566b910e60
---
 drivers/auth/mbedtls/mbedtls_psa_crypto.c | 53 +++++++++++++----------
 1 file changed, 30 insertions(+), 23 deletions(-)

diff --git a/drivers/auth/mbedtls/mbedtls_psa_crypto.c b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
index 64778e8da..b0b3d8606 100644
--- a/drivers/auth/mbedtls/mbedtls_psa_crypto.c
+++ b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
@@ -42,6 +42,16 @@ CASSERT(CRYPTO_MD_MAX_SIZE >= MBEDTLS_MD_MAX_SIZE,
 	* CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC
 	*/
 
+static inline psa_algorithm_t mbedtls_md_psa_alg_from_type(
+						mbedtls_md_type_t md_type)
+{
+	assert((md_type == MBEDTLS_MD_SHA256) ||
+	       (md_type == MBEDTLS_MD_SHA384) ||
+	       (md_type == MBEDTLS_MD_SHA512));
+
+	return PSA_ALG_CATEGORY_HASH | (psa_algorithm_t) (md_type + 0x5);
+}
+
 /*
  * AlgorithmIdentifier  ::=  SEQUENCE  {
  *     algorithm               OBJECT IDENTIFIER,
@@ -195,11 +205,11 @@ static int verify_hash(void *data_ptr, unsigned int data_len,
 {
 	mbedtls_asn1_buf hash_oid, params;
 	mbedtls_md_type_t md_alg;
-	const mbedtls_md_info_t *md_info;
 	unsigned char *p, *end, *hash;
-	unsigned char data_hash[MBEDTLS_MD_MAX_SIZE];
 	size_t len;
 	int rc;
+	psa_status_t status;
+	psa_algorithm_t psa_md_alg;
 
 	/*
 	 * Digest info should be an MBEDTLS_ASN1_SEQUENCE, but padding after
@@ -222,38 +232,35 @@ static int verify_hash(void *data_ptr, unsigned int data_len,
 		return CRYPTO_ERR_HASH;
 	}
 
-	rc = mbedtls_oid_get_md_alg(&hash_oid, &md_alg);
-	if (rc != 0) {
-		return CRYPTO_ERR_HASH;
-	}
-
-	md_info = mbedtls_md_info_from_type(md_alg);
-	if (md_info == NULL) {
-		return CRYPTO_ERR_HASH;
-	}
-
 	/* Hash should be octet string type and consume all bytes */
 	rc = mbedtls_asn1_get_tag(&p, end, &len, MBEDTLS_ASN1_OCTET_STRING);
 	if ((rc != 0) || ((size_t)(end - p) != len)) {
 		return CRYPTO_ERR_HASH;
 	}
-
-	/* Length of hash must match the algorithm's size */
-	if (len != mbedtls_md_get_size(md_info)) {
-		return CRYPTO_ERR_HASH;
-	}
 	hash = p;
 
-	/* Calculate the hash of the data */
-	p = (unsigned char *)data_ptr;
-	rc = mbedtls_md(md_info, p, data_len, data_hash);
+	rc = mbedtls_oid_get_md_alg(&hash_oid, &md_alg);
 	if (rc != 0) {
 		return CRYPTO_ERR_HASH;
 	}
 
-	/* Compare values */
-	rc = memcmp(data_hash, hash, mbedtls_md_get_size(md_info));
-	if (rc != 0) {
+	/* convert the md_alg to psa_algo */
+	psa_md_alg = mbedtls_md_psa_alg_from_type(md_alg);
+
+	/* Length of hash must match the algorithm's size */
+	if (len != PSA_HASH_LENGTH(psa_md_alg)) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	/*
+	 * Calculate Hash and compare it against the retrieved hash from
+	 * the certificate (one shot API).
+	 */
+	status = psa_hash_compare(psa_md_alg,
+				  data_ptr, (size_t)data_len,
+				  (const uint8_t *)hash, len);
+
+	if (status != PSA_SUCCESS) {
 		return CRYPTO_ERR_HASH;
 	}
 

From 484b58696d627c68869d86e2c401a9088392659e Mon Sep 17 00:00:00 2001
From: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Date: Wed, 6 Sep 2023 12:05:05 +0100
Subject: [PATCH 6/7] feat(mbedtls-psa): use PSA crypto API during hash
 calculation

A PSA crypto API is used for hash verification, namely
'psa_hash_compute', which internally calculates hash of the given
data payload.

Signed-off-by: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Change-Id: Ic90e43e68b836ee2add4b868663cfca2340c8108
---
 drivers/auth/mbedtls/mbedtls_psa_crypto.c | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/drivers/auth/mbedtls/mbedtls_psa_crypto.c b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
index b0b3d8606..b70d9c2c1 100644
--- a/drivers/auth/mbedtls/mbedtls_psa_crypto.c
+++ b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
@@ -301,19 +301,26 @@ static int calc_hash(enum crypto_md_algo md_algo, void *data_ptr,
 		     unsigned int data_len,
 		     unsigned char output[CRYPTO_MD_MAX_SIZE])
 {
-	const mbedtls_md_info_t *md_info;
+	size_t hash_length;
+	psa_status_t status;
+	psa_algorithm_t psa_md_alg;
 
-	md_info = mbedtls_md_info_from_type(md_type(md_algo));
-	if (md_info == NULL) {
-		return CRYPTO_ERR_HASH;
-	}
+	/* convert the md_alg to psa_algo */
+	psa_md_alg = mbedtls_md_psa_alg_from_type(md_type(md_algo));
 
 	/*
 	 * Calculate the hash of the data, it is safe to pass the
 	 * 'output' hash buffer pointer considering its size is always
 	 * bigger than or equal to MBEDTLS_MD_MAX_SIZE.
 	 */
-	return mbedtls_md(md_info, data_ptr, data_len, output);
+	status = psa_hash_compute(psa_md_alg, data_ptr, (size_t)data_len,
+				  (uint8_t *)output, CRYPTO_MD_MAX_SIZE,
+				  &hash_length);
+	if (status != PSA_SUCCESS) {
+		return CRYPTO_ERR_HASH;
+	}
+
+	return CRYPTO_SUCCESS;
 }
 #endif /*
 	* CRYPTO_SUPPORT == CRYPTO_HASH_CALC_ONLY || \

From eaa62e825e31fb22a6245d9a5ab9cf5c9f8c0e46 Mon Sep 17 00:00:00 2001
From: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Date: Fri, 22 Sep 2023 18:06:05 +0100
Subject: [PATCH 7/7] feat(mbedtls-psa): use PSA crypto API during signature
 verification

A PSA crypto API is used for signature verification, namely
psa_verify_message, which calculates the signature and verify it
for the given signed message.

Note, this functionality is as of now tested for RSA based Key
algorithm, subsequent patches will provide the signature verification
support for all the key algorithms available in TF-A (provided by KEY_ALG
build option).

Signed-off-by: Manish V Badarkhe <Manish.Badarkhe@arm.com>
Change-Id: I16214e3cac6a7af445530a6dbd66cfc6ceedcfc6
---
 drivers/auth/mbedtls/mbedtls_psa_crypto.c | 94 +++++++++++++++--------
 1 file changed, 64 insertions(+), 30 deletions(-)

diff --git a/drivers/auth/mbedtls/mbedtls_psa_crypto.c b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
index b70d9c2c1..2fa8e63ee 100644
--- a/drivers/auth/mbedtls/mbedtls_psa_crypto.c
+++ b/drivers/auth/mbedtls/mbedtls_psa_crypto.c
@@ -105,6 +105,26 @@ static void init(void)
 
 #if CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_ONLY || \
 CRYPTO_SUPPORT == CRYPTO_AUTH_VERIFY_AND_HASH_CALC
+
+static void construct_psa_key_alg_and_type(mbedtls_pk_type_t pk_alg,
+					   mbedtls_md_type_t md_alg,
+					   psa_algorithm_t *psa_alg,
+					   psa_key_type_t *psa_key_type)
+{
+	psa_algorithm_t psa_md_alg = mbedtls_md_psa_alg_from_type(md_alg);
+
+	switch (pk_alg) {
+	case MBEDTLS_PK_RSASSA_PSS:
+		*psa_alg = PSA_ALG_RSA_PSS(psa_md_alg);
+		*psa_key_type = PSA_KEY_TYPE_RSA_PUBLIC_KEY;
+		break;
+	default:
+		*psa_alg = PSA_ALG_NONE;
+		*psa_key_type = PSA_KEY_TYPE_NONE;
+		break;
+	}
+}
+
 /*
  * Verify a signature.
  *
@@ -120,12 +140,16 @@ static int verify_signature(void *data_ptr, unsigned int data_len,
 	mbedtls_asn1_buf signature;
 	mbedtls_md_type_t md_alg;
 	mbedtls_pk_type_t pk_alg;
-	mbedtls_pk_context pk = {0};
 	int rc;
 	void *sig_opts = NULL;
-	const mbedtls_md_info_t *md_info;
 	unsigned char *p, *end;
-	unsigned char hash[MBEDTLS_MD_MAX_SIZE];
+
+	/* construct PSA key algo and type */
+	psa_status_t status = PSA_SUCCESS;
+	psa_key_attributes_t psa_key_attr = PSA_KEY_ATTRIBUTES_INIT;
+	psa_key_id_t psa_key_id = PSA_KEY_ID_NULL;
+	psa_key_type_t psa_key_type;
+	psa_algorithm_t psa_alg;
 
 	/* Get pointers to signature OID and parameters */
 	p = (unsigned char *)sig_alg;
@@ -141,16 +165,6 @@ static int verify_signature(void *data_ptr, unsigned int data_len,
 		return CRYPTO_ERR_SIGNATURE;
 	}
 
-	/* Parse the public key */
-	mbedtls_pk_init(&pk);
-	p = (unsigned char *)pk_ptr;
-	end = (unsigned char *)(p + pk_len);
-	rc = mbedtls_pk_parse_subpubkey(&p, end, &pk);
-	if (rc != 0) {
-		rc = CRYPTO_ERR_SIGNATURE;
-		goto end2;
-	}
-
 	/* Get the signature (bitstring) */
 	p = (unsigned char *)sig_ptr;
 	end = (unsigned char *)(p + sig_len);
@@ -158,28 +172,45 @@ static int verify_signature(void *data_ptr, unsigned int data_len,
 	rc = mbedtls_asn1_get_bitstring_null(&p, end, &signature.len);
 	if ((rc != 0) || ((size_t)(end - p) != signature.len)) {
 		rc = CRYPTO_ERR_SIGNATURE;
-		goto end1;
+		goto end2;
 	}
 	signature.p = p;
 
-	/* Calculate the hash of the data */
-	md_info = mbedtls_md_info_from_type(md_alg);
-	if (md_info == NULL) {
+	/* Convert this pk_alg and md_alg to PSA key type and key algorithm */
+	construct_psa_key_alg_and_type(pk_alg, md_alg,
+				       &psa_alg, &psa_key_type);
+
+
+	if ((psa_alg == PSA_ALG_NONE) || (psa_key_type == PSA_KEY_TYPE_NONE)) {
 		rc = CRYPTO_ERR_SIGNATURE;
-		goto end1;
-	}
-	p = (unsigned char *)data_ptr;
-	rc = mbedtls_md(md_info, p, data_len, hash);
-	if (rc != 0) {
-		rc = CRYPTO_ERR_SIGNATURE;
-		goto end1;
+		goto end2;
 	}
 
-	/* Verify the signature */
-	rc = mbedtls_pk_verify_ext(pk_alg, sig_opts, &pk, md_alg, hash,
-			mbedtls_md_get_size(md_info),
-			signature.p, signature.len);
-	if (rc != 0) {
+	/* filled-in key_attributes */
+	psa_set_key_algorithm(&psa_key_attr, psa_alg);
+	psa_set_key_type(&psa_key_attr, psa_key_type);
+	psa_set_key_usage_flags(&psa_key_attr, PSA_KEY_USAGE_VERIFY_MESSAGE);
+
+	/* Get the key_id using import API */
+	status = psa_import_key(&psa_key_attr,
+				pk_ptr,
+				(size_t)pk_len,
+				&psa_key_id);
+
+	if (status != PSA_SUCCESS) {
+		rc = CRYPTO_ERR_SIGNATURE;
+		goto end2;
+	}
+
+	/*
+	 * Hash calculation and Signature verification of the given data payload
+	 * is wrapped under the psa_verify_message function.
+	 */
+	status = psa_verify_message(psa_key_id, psa_alg,
+				    data_ptr, data_len,
+				    signature.p, signature.len);
+
+	if (status != PSA_SUCCESS) {
 		rc = CRYPTO_ERR_SIGNATURE;
 		goto end1;
 	}
@@ -188,7 +219,10 @@ static int verify_signature(void *data_ptr, unsigned int data_len,
 	rc = CRYPTO_SUCCESS;
 
 end1:
-	mbedtls_pk_free(&pk);
+	/*
+	 * Destroy the key if it is created successfully
+	 */
+	psa_destroy_key(psa_key_id);
 end2:
 	mbedtls_free(sig_opts);
 	return rc;