From 3b0c371c04cc92e10a6f50b0adb5c0a96f68a5cc Mon Sep 17 00:00:00 2001
From: Gabor Mezei <gabor.mezei@arm.com>
Date: Wed, 24 Jan 2024 13:07:17 +0100
Subject: [PATCH] Add allocate and copy style output buffer handling

Add a new macro `LOCAL_OUTPUT_ALLOC_WITH_COPY` to support the output buffer
handling of the multipart operations like `psa_cipher_update`. This will
allocate a local buffer and copy the content of the original buffer.

Signed-off-by: Gabor Mezei <gabor.mezei@arm.com>
---
 library/psa_crypto.c      | 50 +++++++++++++++++++++++++++++++++++++++
 library/psa_crypto_core.h | 19 +++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 00e4f2aa98..05467bdb7c 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -182,6 +182,23 @@ mbedtls_psa_drbg_context_t *const mbedtls_psa_random_state =
     } \
     output_copy = LOCAL_OUTPUT_COPY_OF_##output.buffer;
 
+/* Allocate a copy of the buffer output and set the pointer output_copy to
+ * point to the start of the copy.
+ *
+ * Assumptions:
+ * - psa_status_t status exists
+ * - An exit label is declared
+ * - output is the name of a pointer to the buffer to be copied
+ * - LOCAL_OUTPUT_DECLARE(output, output_copy) has previously been called
+ */
+#define LOCAL_OUTPUT_ALLOC_WITH_COPY(output, length, output_copy) \
+    status = psa_crypto_local_output_alloc_with_copy(output, length, \
+                                                     &LOCAL_OUTPUT_COPY_OF_##output); \
+    if (status != PSA_SUCCESS) { \
+        goto exit; \
+    } \
+    output_copy = LOCAL_OUTPUT_COPY_OF_##output.buffer;
+
 /* Free the local output copy allocated previously by LOCAL_OUTPUT_ALLOC()
  * after first copying back its contents to the original buffer.
  *
@@ -5783,6 +5800,39 @@ psa_status_t psa_crypto_local_output_alloc(uint8_t *output, size_t output_len,
     return PSA_SUCCESS;
 }
 
+psa_status_t psa_crypto_local_output_alloc_with_copy(uint8_t *output, size_t output_len,
+                                                     psa_crypto_local_output_t *local_output)
+{
+    psa_status_t status;
+    *local_output = PSA_CRYPTO_LOCAL_OUTPUT_INIT;
+
+    if (output_len == 0) {
+        return PSA_SUCCESS;
+    }
+    local_output->buffer = mbedtls_calloc(output_len, 1);
+    if (local_output->buffer == NULL) {
+        /* Since we dealt with the zero-length case above, we know that
+         * a NULL return value means a failure of allocation. */
+        return PSA_ERROR_INSUFFICIENT_MEMORY;
+    }
+    local_output->length = output_len;
+    local_output->original = output;
+
+    status = psa_crypto_copy_input(output, output_len,
+                                   local_output->buffer, local_output->length);
+    if (status != PSA_SUCCESS) {
+        goto error;
+    }
+
+    return PSA_SUCCESS;
+
+error:
+    mbedtls_free(local_output->buffer);
+    local_output->buffer = NULL;
+    local_output->length = 0;
+    return status;
+}
+
 psa_status_t psa_crypto_local_output_free(psa_crypto_local_output_t *local_output)
 {
     psa_status_t status;
diff --git a/library/psa_crypto_core.h b/library/psa_crypto_core.h
index 4731064f85..ae8af09d91 100644
--- a/library/psa_crypto_core.h
+++ b/library/psa_crypto_core.h
@@ -560,6 +560,25 @@ typedef struct psa_crypto_local_output_s {
 psa_status_t psa_crypto_local_output_alloc(uint8_t *output, size_t output_len,
                                            psa_crypto_local_output_t *local_output);
 
+/** Allocate a local copy of an output buffer and copy the contents into it.
+ *
+ * \note                        This allocates and copies a buffer
+ *                              whose contents will be copied back to the
+ *                              original in a future call to
+ *                              psa_crypto_local_output_free().
+ *
+ * \param[in] output            Pointer to output buffer.
+ * \param[in] output_len        Length of the output buffer.
+ * \param[out] local_output     Pointer to a psa_crypto_local_output_t struct to
+ *                              populate with the local output copy.
+ * \return                      #PSA_SUCCESS, if the buffer was successfully
+ *                              copied.
+ * \return                      #PSA_ERROR_INSUFFICIENT_MEMORY, if a copy of
+ *                              the buffer cannot be allocated.
+ */
+psa_status_t psa_crypto_local_output_alloc_with_copy(uint8_t *output, size_t output_len,
+                                                     psa_crypto_local_output_t *local_output);
+
 /** Copy from a local copy of an output buffer back to the original, then
  *  free the local copy.
  *