Merge pull request #8283 from daverodgman/more-aes-checks

More AES guards testing and some fixes
diff --git a/ChangeLog.d/add-psa-example-program-hash.txt b/ChangeLog.d/add-psa-example-program-hash.txt
new file mode 100644
index 0000000..ba4da20
--- /dev/null
+++ b/ChangeLog.d/add-psa-example-program-hash.txt
@@ -0,0 +1,2 @@
+Features
+   * Added an example program showing how to hash with the PSA API.
diff --git a/ChangeLog.d/pkwrite-pem-use-heap.txt b/ChangeLog.d/pkwrite-pem-use-heap.txt
new file mode 100644
index 0000000..11db7b6
--- /dev/null
+++ b/ChangeLog.d/pkwrite-pem-use-heap.txt
@@ -0,0 +1,4 @@
+Changes
+   * Use heap memory to allocate DER encoded public/private key.
+     This reduces stack usage significantly for writing a public/private
+     key to a PEM string.
diff --git a/library/pkwrite.c b/library/pkwrite.c
index 03db145..e38bc27 100644
--- a/library/pkwrite.c
+++ b/library/pkwrite.c
@@ -756,27 +756,38 @@
 int mbedtls_pk_write_pubkey_pem(const mbedtls_pk_context *key, unsigned char *buf, size_t size)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    unsigned char output_buf[PUB_DER_MAX_BYTES];
+    unsigned char *output_buf = NULL;
+    output_buf = mbedtls_calloc(1, PUB_DER_MAX_BYTES);
+    if (output_buf == NULL) {
+        return MBEDTLS_ERR_PK_ALLOC_FAILED;
+    }
     size_t olen = 0;
 
     if ((ret = mbedtls_pk_write_pubkey_der(key, output_buf,
-                                           sizeof(output_buf))) < 0) {
-        return ret;
+                                           PUB_DER_MAX_BYTES)) < 0) {
+        goto cleanup;
     }
 
     if ((ret = mbedtls_pem_write_buffer(PEM_BEGIN_PUBLIC_KEY, PEM_END_PUBLIC_KEY,
-                                        output_buf + sizeof(output_buf) - ret,
+                                        output_buf + PUB_DER_MAX_BYTES - ret,
                                         ret, buf, size, &olen)) != 0) {
-        return ret;
+        goto cleanup;
     }
 
-    return 0;
+    ret = 0;
+cleanup:
+    mbedtls_free(output_buf);
+    return ret;
 }
 
 int mbedtls_pk_write_key_pem(const mbedtls_pk_context *key, unsigned char *buf, size_t size)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    unsigned char output_buf[PRV_DER_MAX_BYTES];
+    unsigned char *output_buf = NULL;
+    output_buf = mbedtls_calloc(1, PRV_DER_MAX_BYTES);
+    if (output_buf == NULL) {
+        return MBEDTLS_ERR_PK_ALLOC_FAILED;
+    }
     const char *begin, *end;
     size_t olen = 0;
 #if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
@@ -789,8 +800,8 @@
     int is_rsa_opaque = 0;
 #endif
 
-    if ((ret = mbedtls_pk_write_key_der(key, output_buf, sizeof(output_buf))) < 0) {
-        return ret;
+    if ((ret = mbedtls_pk_write_key_der(key, output_buf, PRV_DER_MAX_BYTES)) < 0) {
+        goto cleanup;
     }
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
@@ -833,15 +844,21 @@
         }
     } else
 #endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
-    return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
-
-    if ((ret = mbedtls_pem_write_buffer(begin, end,
-                                        output_buf + sizeof(output_buf) - ret,
-                                        ret, buf, size, &olen)) != 0) {
-        return ret;
+    {
+        ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
+        goto cleanup;
     }
 
-    return 0;
+    if ((ret = mbedtls_pem_write_buffer(begin, end,
+                                        output_buf + PRV_DER_MAX_BYTES - ret,
+                                        ret, buf, size, &olen)) != 0) {
+        goto cleanup;
+    }
+
+    ret = 0;
+cleanup:
+    mbedtls_zeroize_and_free(output_buf, PRV_DER_MAX_BYTES);
+    return ret;
 }
 #endif /* MBEDTLS_PEM_WRITE_C */
 
diff --git a/library/ssl_tls13_keys.c b/library/ssl_tls13_keys.c
index 6905d92..5ae6210 100644
--- a/library/ssl_tls13_keys.c
+++ b/library/ssl_tls13_keys.c
@@ -1019,14 +1019,14 @@
 
 #if !defined(MBEDTLS_USE_PSA_CRYPTO)
     if ((ret = mbedtls_cipher_setkey(&transform->cipher_ctx_enc,
-                                     key_enc, mbedtls_cipher_info_get_key_bitlen(cipher_info),
+                                     key_enc, (int) mbedtls_cipher_info_get_key_bitlen(cipher_info),
                                      MBEDTLS_ENCRYPT)) != 0) {
         MBEDTLS_SSL_DEBUG_RET(1, "mbedtls_cipher_setkey", ret);
         return ret;
     }
 
     if ((ret = mbedtls_cipher_setkey(&transform->cipher_ctx_dec,
-                                     key_dec, mbedtls_cipher_info_get_key_bitlen(cipher_info),
+                                     key_dec, (int) mbedtls_cipher_info_get_key_bitlen(cipher_info),
                                      MBEDTLS_DECRYPT)) != 0) {
         MBEDTLS_SSL_DEBUG_RET(1, "mbedtls_cipher_setkey", ret);
         return ret;
diff --git a/programs/Makefile b/programs/Makefile
index 80637e9..116883b 100644
--- a/programs/Makefile
+++ b/programs/Makefile
@@ -109,6 +109,7 @@
 	psa/hmac_demo \
 	psa/key_ladder_demo \
 	psa/psa_constant_names \
+	psa/psa_hash \
 	random/gen_entropy \
 	random/gen_random_ctr_drbg \
 	ssl/dtls_client \
@@ -324,6 +325,10 @@
 	echo "  CC    psa/psa_constant_names.c"
 	$(CC) $(LOCAL_CFLAGS) $(CFLAGS) psa/psa_constant_names.c    $(LOCAL_LDFLAGS) $(LDFLAGS) -o $@
 
+psa/psa_hash$(EXEXT): psa/psa_hash.c $(DEP)
+	echo "  CC    psa/psa_hash.c"
+	$(CC) $(LOCAL_CFLAGS) $(CFLAGS) psa/psa_hash.c    $(LOCAL_LDFLAGS) $(LDFLAGS) -o $@
+
 random/gen_entropy$(EXEXT): random/gen_entropy.c $(DEP)
 	echo "  CC    random/gen_entropy.c"
 	$(CC) $(LOCAL_CFLAGS) $(CFLAGS) random/gen_entropy.c $(LOCAL_LDFLAGS) $(LDFLAGS) -o $@
diff --git a/programs/psa/CMakeLists.txt b/programs/psa/CMakeLists.txt
index 7ba4af6..c8ee626 100644
--- a/programs/psa/CMakeLists.txt
+++ b/programs/psa/CMakeLists.txt
@@ -4,6 +4,7 @@
     hmac_demo
     key_ladder_demo
     psa_constant_names
+    psa_hash
 )
 
 if(GEN_FILES)
diff --git a/programs/psa/psa_hash.c b/programs/psa/psa_hash.c
new file mode 100644
index 0000000..d3a6bf8
--- /dev/null
+++ b/programs/psa/psa_hash.c
@@ -0,0 +1,171 @@
+/*
+ *  Example computing a SHA-256 hash using the PSA Crypto API
+ *
+ *  The example computes the SHA-256 hash of a test string using the
+ *  one-shot API call psa_hash_compute() and the using multi-part
+ *  operation, which requires psa_hash_setup(), psa_hash_update() and
+ *  psa_hash_finish(). The multi-part operation is popular on embedded
+ *  devices where a rolling hash needs to be computed.
+ *
+ *
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include "psa/crypto.h"
+#include <string.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "mbedtls/build_info.h"
+#include "mbedtls/platform.h"
+
+/* Information about hashing with the PSA API can be
+ * found here:
+ * https://arm-software.github.io/psa-api/crypto/1.1/api/ops/hashes.html
+ *
+ * The algorithm used by this demo is SHA 256.
+ * Please see include/psa/crypto_values.h to see the other
+ * algorithms that are supported by Mbed TLS.
+ * If you switch to a different algorithm you will need to update
+ * the hash data in the EXAMPLE_HASH_VALUE macro below. */
+
+#if !defined(MBEDTLS_PSA_CRYPTO_C) || !defined(PSA_WANT_ALG_SHA_256)
+int main(void)
+{
+    mbedtls_printf("MBEDTLS_PSA_CRYPTO_C and PSA_WANT_ALG_SHA_256"
+                   "not defined.\r\n");
+    return EXIT_SUCCESS;
+}
+#else
+
+#define HASH_ALG PSA_ALG_SHA_256
+
+const uint8_t sample_message[] = "Hello World!";
+/* sample_message is terminated with a null byte which is not part of
+ * the message itself so we make sure to subtract it in order to get
+ * the message length. */
+const size_t sample_message_length = sizeof(sample_message) - 1;
+
+#define EXPECTED_HASH_VALUE {                                                    \
+        0x7f, 0x83, 0xb1, 0x65, 0x7f, 0xf1, 0xfc, 0x53, 0xb9, 0x2d, 0xc1, 0x81, \
+        0x48, 0xa1, 0xd6, 0x5d, 0xfc, 0x2d, 0x4b, 0x1f, 0xa3, 0xd6, 0x77, 0x28, \
+        0x4a, 0xdd, 0xd2, 0x00, 0x12, 0x6d, 0x90, 0x69 \
+}
+
+const uint8_t expected_hash[] = EXPECTED_HASH_VALUE;
+const size_t expected_hash_len = sizeof(expected_hash);
+
+int main(void)
+{
+    psa_status_t status;
+    uint8_t hash[PSA_HASH_LENGTH(HASH_ALG)];
+    size_t hash_length;
+    psa_hash_operation_t hash_operation = PSA_HASH_OPERATION_INIT;
+    psa_hash_operation_t cloned_hash_operation = PSA_HASH_OPERATION_INIT;
+
+    mbedtls_printf("PSA Crypto API: SHA-256 example\n\n");
+
+    status = psa_crypto_init();
+    if (status != PSA_SUCCESS) {
+        mbedtls_printf("psa_crypto_init failed\n");
+        return EXIT_FAILURE;
+    }
+
+    /* Compute hash using multi-part operation */
+    status = psa_hash_setup(&hash_operation, HASH_ALG);
+    if (status == PSA_ERROR_NOT_SUPPORTED) {
+        mbedtls_printf("unknown hash algorithm supplied\n");
+        return EXIT_FAILURE;
+    } else if (status != PSA_SUCCESS) {
+        mbedtls_printf("psa_hash_setup failed\n");
+        return EXIT_FAILURE;
+    }
+
+    status = psa_hash_update(&hash_operation, sample_message, sample_message_length);
+    if (status != PSA_SUCCESS) {
+        mbedtls_printf("psa_hash_update failed\n");
+        goto cleanup;
+    }
+
+    status = psa_hash_clone(&hash_operation, &cloned_hash_operation);
+    if (status != PSA_SUCCESS) {
+        mbedtls_printf("PSA hash clone failed\n");
+        goto cleanup;
+    }
+
+    status = psa_hash_finish(&hash_operation, hash, sizeof(hash), &hash_length);
+    if (status != PSA_SUCCESS) {
+        mbedtls_printf("psa_hash_finish failed\n");
+        goto cleanup;
+    }
+
+    /* Check the result of the operation against the sample */
+    if (hash_length != expected_hash_len ||
+        (memcmp(hash, expected_hash, expected_hash_len) != 0)) {
+        mbedtls_printf("Multi-part hash operation gave the wrong result!\n\n");
+        goto cleanup;
+    }
+
+    status =
+        psa_hash_verify(&cloned_hash_operation, expected_hash,
+                        expected_hash_len);
+    if (status != PSA_SUCCESS) {
+        mbedtls_printf("psa_hash_verify failed\n");
+        goto cleanup;
+    } else {
+        mbedtls_printf("Multi-part hash operation successful!\n");
+    }
+
+    /* Clear local variables prior to one-shot hash demo */
+    memset(hash, 0, sizeof(hash));
+    hash_length = 0;
+
+    /* Compute hash using one-shot function call */
+    status = psa_hash_compute(HASH_ALG,
+                              sample_message, sample_message_length,
+                              hash, sizeof(hash),
+                              &hash_length);
+    if (status != PSA_SUCCESS) {
+        mbedtls_printf("psa_hash_compute failed\n");
+        goto cleanup;
+    }
+
+    if (hash_length != expected_hash_len ||
+        (memcmp(hash, expected_hash, expected_hash_len) != 0)) {
+        mbedtls_printf("One-shot hash operation gave the wrong result!\n\n");
+        goto cleanup;
+    }
+
+    mbedtls_printf("One-shot hash operation successful!\n\n");
+
+    /* Print out result */
+    mbedtls_printf("The SHA-256( '%s' ) is: ", sample_message);
+
+    for (size_t j = 0; j < expected_hash_len; j++) {
+        mbedtls_printf("%02x", hash[j]);
+    }
+
+    mbedtls_printf("\n");
+
+    mbedtls_psa_crypto_free();
+    return EXIT_SUCCESS;
+
+cleanup:
+    psa_hash_abort(&hash_operation);
+    psa_hash_abort(&cloned_hash_operation);
+    return EXIT_FAILURE;
+}
+#endif /* !MBEDTLS_PSA_CRYPTO_C || !PSA_WANT_ALG_SHA_256 */