From 4124bca309586234748c9c99600ec2f3b2f6d775 Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka@google.com>
Date: Thu, 27 May 2021 10:10:50 -0700
Subject: [PATCH] [scudo] Enabled MTE in tests

Reviewed By: pcc, hctim

Differential Revision: https://reviews.llvm.org/D103305
---
 compiler-rt/lib/scudo/standalone/memtag.h     | 30 ++++++++++++++++++-
 .../standalone/tests/scudo_unit_test_main.cpp |  3 ++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/compiler-rt/lib/scudo/standalone/memtag.h b/compiler-rt/lib/scudo/standalone/memtag.h
index 9773cb97f05d..fec258b11a42 100644
--- a/compiler-rt/lib/scudo/standalone/memtag.h
+++ b/compiler-rt/lib/scudo/standalone/memtag.h
@@ -67,15 +67,27 @@ inline bool systemSupportsMemoryTagging() {
 }
 
 inline bool systemDetectsMemoryTagFaultsTestOnly() {
+#ifndef PR_SET_TAGGED_ADDR_CTRL
+#define PR_SET_TAGGED_ADDR_CTRL 54
+#endif
 #ifndef PR_GET_TAGGED_ADDR_CTRL
 #define PR_GET_TAGGED_ADDR_CTRL 56
 #endif
+#ifndef PR_TAGGED_ADDR_ENABLE
+#define PR_TAGGED_ADDR_ENABLE (1UL << 0)
+#endif
 #ifndef PR_MTE_TCF_SHIFT
 #define PR_MTE_TCF_SHIFT 1
 #endif
+#ifndef PR_MTE_TAG_SHIFT
+#define PR_MTE_TAG_SHIFT 3
+#endif
 #ifndef PR_MTE_TCF_NONE
 #define PR_MTE_TCF_NONE (0UL << PR_MTE_TCF_SHIFT)
 #endif
+#ifndef PR_MTE_TCF_SYNC
+#define PR_MTE_TCF_SYNC (1UL << PR_MTE_TCF_SHIFT)
+#endif
 #ifndef PR_MTE_TCF_MASK
 #define PR_MTE_TCF_MASK (3UL << PR_MTE_TCF_SHIFT)
 #endif
@@ -84,11 +96,23 @@ inline bool systemDetectsMemoryTagFaultsTestOnly() {
           PR_MTE_TCF_MASK) != PR_MTE_TCF_NONE;
 }
 
+inline void enableSystemMemoryTaggingTestOnly() {
+  prctl(PR_SET_TAGGED_ADDR_CTRL,
+        PR_TAGGED_ADDR_ENABLE | PR_MTE_TCF_SYNC | (0xfffe << PR_MTE_TAG_SHIFT),
+        0, 0, 0);
+}
+
 #else // !SCUDO_LINUX
 
 inline bool systemSupportsMemoryTagging() { return false; }
 
-inline bool systemDetectsMemoryTagFaultsTestOnly() { return false; }
+inline bool systemDetectsMemoryTagFaultsTestOnly() {
+  UNREACHABLE("memory tagging not supported");
+}
+
+inline void enableSystemMemoryTaggingTestOnly() {
+  UNREACHABLE("memory tagging not supported");
+}
 
 #endif // SCUDO_LINUX
 
@@ -250,6 +274,10 @@ inline bool systemDetectsMemoryTagFaultsTestOnly() {
   UNREACHABLE("memory tagging not supported");
 }
 
+inline void enableSystemMemoryTaggingTestOnly() {
+  UNREACHABLE("memory tagging not supported");
+}
+
 inline void disableMemoryTagChecksTestOnly() {
   UNREACHABLE("memory tagging not supported");
 }
diff --git a/compiler-rt/lib/scudo/standalone/tests/scudo_unit_test_main.cpp b/compiler-rt/lib/scudo/standalone/tests/scudo_unit_test_main.cpp
index 9bbf6e75a5cd..1983b80ac0ca 100644
--- a/compiler-rt/lib/scudo/standalone/tests/scudo_unit_test_main.cpp
+++ b/compiler-rt/lib/scudo/standalone/tests/scudo_unit_test_main.cpp
@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "memtag.h"
 #include "tests/scudo_unit_test.h"
 
 // Match Android's default configuration, which disables Scudo's mismatch
@@ -33,6 +34,8 @@ __scudo_default_options() {
 // for Fuchsia builds.
 #if !SCUDO_FUCHSIA
 int main(int argc, char **argv) {
+  if (scudo::archSupportsMemoryTagging())
+    scudo::enableSystemMemoryTaggingTestOnly();
   testing::InitGoogleTest(&argc, argv);
   return RUN_ALL_TESTS();
 }