From 372b8bb6c511d25018ed9390f851915c1c6e8c5a Mon Sep 17 00:00:00 2001
From: David Horstmann <david.horstmann@arm.com>
Date: Fri, 24 Nov 2023 16:21:04 +0000
Subject: [PATCH] Add memory poisoning hooks

Signed-off-by: David Horstmann <david.horstmann@arm.com>
---
 library/CMakeLists.txt |  3 ++-
 library/psa_crypto.c   | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/library/CMakeLists.txt b/library/CMakeLists.txt
index 5c297e0a1f..52178e474a 100644
--- a/library/CMakeLists.txt
+++ b/library/CMakeLists.txt
@@ -329,7 +329,8 @@ foreach(target IN LISTS target_libraries)
                $<INSTALL_INTERFACE:include/>
         PRIVATE ${MBEDTLS_DIR}/library/
                 # Needed to include psa_crypto_driver_wrappers.h
-                ${CMAKE_CURRENT_BINARY_DIR})
+                ${CMAKE_CURRENT_BINARY_DIR}
+                ${MBEDTLS_DIR}/tests/include/)
     # Pass-through MBEDTLS_CONFIG_FILE and MBEDTLS_USER_CONFIG_FILE
     if(MBEDTLS_CONFIG_FILE)
         target_compile_definitions(${target}
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index e3187d8d28..ddc834f95c 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -72,6 +72,10 @@
 #include "mbedtls/sha512.h"
 #include "md_psa.h"
 
+#if defined(MBEDTLS_TEST_HOOKS)
+#include "test/memory.h"
+#endif
+
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF) ||          \
     defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT) ||  \
     defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXPAND)
@@ -8451,10 +8455,18 @@ psa_status_t psa_crypto_copy_input(const uint8_t *input, size_t input_len,
         return PSA_ERROR_CORRUPTION_DETECTED;
     }
 
+#if defined(MBEDTLS_TEST_HOOKS)
+    MBEDTLS_TEST_MEMORY_UNPOISON(input, input_len);
+#endif
+
     if (input_len > 0) {
         memcpy(input_copy, input, input_len);
     }
 
+#if defined(MBEDTLS_TEST_HOOKS)
+    MBEDTLS_TEST_MEMORY_POISON(input, input_len);
+#endif
+
     return PSA_SUCCESS;
 }
 
@@ -8478,10 +8490,18 @@ psa_status_t psa_crypto_copy_output(const uint8_t *output_copy, size_t output_co
         return PSA_ERROR_BUFFER_TOO_SMALL;
     }
 
+#if defined(MBEDTLS_TEST_HOOKS)
+    MBEDTLS_TEST_MEMORY_UNPOISON(output, output_len);
+#endif
+
     if (output_copy_len > 0) {
         memcpy(output, output_copy, output_copy_len);
     }
 
+#if defined(MBEDTLS_TEST_HOOKS)
+    MBEDTLS_TEST_MEMORY_POISON(output, output_len);
+#endif
+
     return PSA_SUCCESS;
 }