Use thread-local flag to enable memory poisoning

Allow memory poisoning to be enabled and disabled at runtime using a
thread-local flag. This allows poisoning to be disabled whenever a PSA
function is called but not through the test wrappers, removing false
positive use-after-poisons.

Signed-off-by: David Horstmann <david.horstmann@arm.com>
diff --git a/tests/include/test/memory.h b/tests/include/test/memory.h
index 43fbb63..0bf8f46 100644
--- a/tests/include/test/memory.h
+++ b/tests/include/test/memory.h
@@ -21,9 +21,12 @@
  * memory as poisoned, which can be used to enforce some memory access
  * policies.
  *
+ * Support for the C11 thread_local keyword is also required.
+ *
  * Currently, only Asan (Address Sanitizer) is supported.
  */
-#if defined(MBEDTLS_TEST_HAVE_ASAN)
+#if defined(MBEDTLS_TEST_HAVE_ASAN) && \
+    (__STDC_VERSION__ >= 201112L)
 #  define MBEDTLS_TEST_MEMORY_CAN_POISON
 #endif
 
@@ -61,6 +64,12 @@
 
 #if defined(MBEDTLS_TEST_MEMORY_CAN_POISON)
 
+/** Thread-local variable used to enable memory poisoning. This is set and
+ *  unset in the test wrappers so that calls to PSA functions from the library
+ *  do not poison memory.
+ */
+extern _Thread_local int mbedtls_test_memory_poisoning_enabled;
+
 /** Poison a memory area so that any attempt to read or write from it will
  * cause a runtime failure.
  *
@@ -68,7 +77,10 @@
  */
 void mbedtls_test_memory_poison(const unsigned char *ptr, size_t size);
 #define MBEDTLS_TEST_MEMORY_POISON(ptr, size)    \
-    mbedtls_test_memory_poison(ptr, size)
+    do { \
+        mbedtls_test_memory_poisoning_enabled = 1; \
+        mbedtls_test_memory_poison(ptr, size); \
+    } while (0)
 
 /** Undo the effect of mbedtls_test_memory_poison().
  *
@@ -79,7 +91,10 @@
  */
 void mbedtls_test_memory_unpoison(const unsigned char *ptr, size_t size);
 #define MBEDTLS_TEST_MEMORY_UNPOISON(ptr, size)    \
-    mbedtls_test_memory_unpoison(ptr, size)
+    do { \
+        mbedtls_test_memory_unpoison(ptr, size); \
+        mbedtls_test_memory_poisoning_enabled = 0; \
+    } while (0)
 
 #else /* MBEDTLS_TEST_MEMORY_CAN_POISON */
 #define MBEDTLS_TEST_MEMORY_POISON(ptr, size) ((void) (ptr), (void) (size))
diff --git a/tests/src/test_memory.c b/tests/src/test_memory.c
index c277be8..dbb7b20 100644
--- a/tests/src/test_memory.c
+++ b/tests/src/test_memory.c
@@ -13,12 +13,15 @@
 #include <test/macros.h>
 #include <test/memory.h>
 
-#if defined(MBEDTLS_TEST_HAVE_ASAN)
+#if defined(MBEDTLS_TEST_MEMORY_CAN_POISON)
 #include <sanitizer/asan_interface.h>
 #include <stdint.h>
 #endif
 
-#if defined(MBEDTLS_TEST_HAVE_ASAN)
+#if defined(MBEDTLS_TEST_MEMORY_CAN_POISON)
+
+_Thread_local int mbedtls_test_memory_poisoning_enabled = 0;
+
 static void align_for_asan(const unsigned char **p_ptr, size_t *p_size)
 {
     uintptr_t start = (uintptr_t) *p_ptr;
@@ -36,6 +39,9 @@
 
 void mbedtls_test_memory_poison(const unsigned char *ptr, size_t size)
 {
+    if (!mbedtls_test_memory_poisoning_enabled) {
+        return;
+    }
     if (size == 0) {
         return;
     }
@@ -51,4 +57,4 @@
     align_for_asan(&ptr, &size);
     __asan_unpoison_memory_region(ptr, size);
 }
-#endif /* Asan */
+#endif /* Memory poisoning */