From 2d28c46055b8006938df8f939d05e1f28e024064 Mon Sep 17 00:00:00 2001
From: Dave Rodgman <dave.rodgman@arm.com>
Date: Fri, 28 Jul 2023 18:22:56 +0100
Subject: [PATCH] Fix asm Memsan workaround

Signed-off-by: Dave Rodgman <dave.rodgman@arm.com>
---
 library/constant_time_impl.h       |  7 ++++---
 tests/include/test/constant_flow.h | 22 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/library/constant_time_impl.h b/library/constant_time_impl.h
index c9e6a83ab..191769b19 100644
--- a/library/constant_time_impl.h
+++ b/library/constant_time_impl.h
@@ -73,17 +73,18 @@ static inline mbedtls_ct_uint_t mbedtls_ct_compiler_opaque(mbedtls_ct_uint_t x)
 #if defined(MBEDTLS_CT_ASM)
     /* Prevent false positives from Memsan - otherwise it will report the asm as
      * accessing secret data. */
-    TEST_CF_PUBLIC(&x, sizeof(x));
+    TEST_CF_SAVE_SECRET(x);
 
     asm volatile ("" : [x] "+r" (x) :);
 
-    /* Mark the return value as secret. This is needed so that code of the form:
+    /* Mark the return value as secret (if it was previously marked secret).
+     * This is needed so that code of the form:
      *
      * if (mbedtls_ct_compiler_opaque(secret)) { ... }
      *
      * will fail const-flow tests.
      */
-    TEST_CF_SECRET(&x, sizeof(x));
+    TEST_CF_RESTORE_SECRET(x);
     return x;
 #else
     return x ^ mbedtls_ct_zero;
diff --git a/tests/include/test/constant_flow.h b/tests/include/test/constant_flow.h
index f3d676e28..572835713 100644
--- a/tests/include/test/constant_flow.h
+++ b/tests/include/test/constant_flow.h
@@ -32,14 +32,27 @@
  *  #define TEST_CF_SECRET(ptr, size)
  *  #define TEST_CF_PUBLIC(ptr, size)
  *
+ * and
+ *
+ * #define TEST_CF_SAVE_SECRET(variable)
+ * #define TEST_CF_RESTORE_SECRET(variable)
+ *
  * that can be used in tests to mark a memory area as secret (no branch or
  * memory access should depend on it) or public (default, only needs to be
  * marked explicitly when it was derived from secret data).
  *
+ * The SAVE/RESTORE forms mark a variable as public, and subsequently restore its
+ * previous secret/not-secret state. This is used where library code is generating
+ * false positives and needs to temporarily disable Memsan checks for a particular
+ * variable, and then restore it's original state afterwards so it doesn't interfere
+ * with other checks.
+ *
  * Arguments:
  * - ptr: a pointer to the memory area to be marked
  * - size: the size in bytes of the memory area
  *
+ * - variable: a variable name
+ *
  * Implementation:
  * The basic idea is that of ctgrind <https://github.com/agl/ctgrind>: we can
  * re-use tools that were designed for checking use of uninitialized memory.
@@ -63,6 +76,9 @@
 #define TEST_CF_PUBLIC  __msan_unpoison
 // void __msan_unpoison(const volatile void *a, size_t size);
 
+#define TEST_CF_SAVE_SECRET(_x)    int _test_cf_is_public_ ## _x = __msan_test_shadow(&(_x), sizeof(_x)) == -1; TEST_CF_PUBLIC(&(_x), sizeof(_x));
+#define TEST_CF_RESTORE_SECRET(_x) do { if (!_test_cf_is_public_ ## _x) TEST_CF_SECRET(&(_x), sizeof(_x)); } while(0)
+
 #elif defined(MBEDTLS_TEST_CONSTANT_FLOW_VALGRIND)
 #include <valgrind/memcheck.h>
 
@@ -71,12 +87,18 @@
 #define TEST_CF_PUBLIC  VALGRIND_MAKE_MEM_DEFINED
 // VALGRIND_MAKE_MEM_DEFINED(_qzz_addr, _qzz_len)
 
+#define TEST_CF_SAVE_SECRET(_x)
+#define TEST_CF_RESTORE_SECRET(_x)
+
 #else /* MBEDTLS_TEST_CONSTANT_FLOW_MEMSAN ||
          MBEDTLS_TEST_CONSTANT_FLOW_VALGRIND */
 
 #define TEST_CF_SECRET(ptr, size)
 #define TEST_CF_PUBLIC(ptr, size)
 
+#define TEST_CF_SAVE_SECRET(_x)
+#define TEST_CF_RESTORE_SECRET(_x)
+
 #endif /* MBEDTLS_TEST_CONSTANT_FLOW_MEMSAN ||
           MBEDTLS_TEST_CONSTANT_FLOW_VALGRIND */