Change to use test-hook-based approach

Since we are applying hooks transparently to all tests, we cannot setup
and teardown test hooks in the tests. Instead we must do this in the
test wrappers which are used to pre-poison and unpoison memory.

Signed-off-by: David Horstmann <david.horstmann@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 55adfa7..458affd 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -73,10 +73,6 @@
 #include "mbedtls/sha512.h"
 #include "mbedtls/xtea.h"
 
-#if defined(MBEDTLS_TEST_HOOKS)
-#include "test/memory.h"
-#endif
-
 #define ARRAY_LENGTH(array) (sizeof(array) / sizeof(*(array)))
 
 /****************************************************************/
@@ -5532,6 +5528,14 @@
     return status;
 }
 
+/* Memory copying test hooks */
+#if defined(MBEDTLS_TEST_HOOKS)
+void (*psa_input_pre_copy_hook)(const uint8_t *input, size_t input_len) = NULL;
+void (*psa_input_post_copy_hook)(const uint8_t *input, size_t input_len) = NULL;
+void (*psa_output_pre_copy_hook)(const uint8_t *output, size_t output_len) = NULL;
+void (*psa_output_post_copy_hook)(const uint8_t *output, size_t output_len) = NULL;
+#endif
+
 /** Copy from an input buffer to a local copy.
  *
  * \param[in] input             Pointer to input buffer.
@@ -5553,7 +5557,9 @@
     }
 
 #if defined(MBEDTLS_TEST_HOOKS)
-    MBEDTLS_TEST_MEMORY_UNPOISON(input, input_len);
+    if (psa_input_pre_copy_hook != NULL) {
+        psa_input_pre_copy_hook(input, input_len);
+    }
 #endif
 
     if (input_len > 0) {
@@ -5561,7 +5567,9 @@
     }
 
 #if defined(MBEDTLS_TEST_HOOKS)
-    MBEDTLS_TEST_MEMORY_POISON(input, input_len);
+    if (psa_input_post_copy_hook != NULL) {
+        psa_input_post_copy_hook(input, input_len);
+    }
 #endif
 
     return PSA_SUCCESS;
@@ -5588,7 +5596,9 @@
     }
 
 #if defined(MBEDTLS_TEST_HOOKS)
-    MBEDTLS_TEST_MEMORY_UNPOISON(output, output_len);
+    if (psa_output_pre_copy_hook != NULL) {
+        psa_output_pre_copy_hook(output, output_len);
+    }
 #endif
 
     if (output_copy_len > 0) {
@@ -5596,7 +5606,9 @@
     }
 
 #if defined(MBEDTLS_TEST_HOOKS)
-    MBEDTLS_TEST_MEMORY_POISON(output, output_len);
+    if (psa_output_post_copy_hook != NULL) {
+        psa_output_post_copy_hook(output, output_len);
+    }
 #endif
 
     return PSA_SUCCESS;
diff --git a/library/psa_crypto_invasive.h b/library/psa_crypto_invasive.h
index e7ab9b3..a1281d1 100644
--- a/library/psa_crypto_invasive.h
+++ b/library/psa_crypto_invasive.h
@@ -76,6 +76,14 @@
 psa_status_t psa_crypto_copy_output(const uint8_t *output_copy, size_t output_copy_len,
                                     uint8_t *output, size_t output_len);
 
+/*
+ * Test hooks to use for memory unpoisoning/poisoning in copy functions.
+ */
+extern void (*psa_input_pre_copy_hook)(const uint8_t *input, size_t input_len);
+extern void (*psa_input_post_copy_hook)(const uint8_t *input, size_t input_len);
+extern void (*psa_output_pre_copy_hook)(const uint8_t *output, size_t output_len);
+extern void (*psa_output_post_copy_hook)(const uint8_t *output, size_t output_len);
+
 #endif /* MBEDTLS_TEST_HOOKS && MBEDTLS_PSA_CRYPTO_C */
 
 #endif /* PSA_CRYPTO_INVASIVE_H */