[scudo] Enabled MTE in tests

Reviewed By: pcc, hctim

Differential Revision: https://reviews.llvm.org/D103305

GitOrigin-RevId: 4124bca309586234748c9c99600ec2f3b2f6d775
diff --git a/memtag.h b/memtag.h
index 9773cb9..fec258b 100644
--- a/memtag.h
+++ b/memtag.h
@@ -67,15 +67,27 @@
 }
 
 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 @@
           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 @@
   UNREACHABLE("memory tagging not supported");
 }
 
+inline void enableSystemMemoryTaggingTestOnly() {
+  UNREACHABLE("memory tagging not supported");
+}
+
 inline void disableMemoryTagChecksTestOnly() {
   UNREACHABLE("memory tagging not supported");
 }
diff --git a/tests/scudo_unit_test_main.cpp b/tests/scudo_unit_test_main.cpp
index 9bbf6e7..1983b80 100644
--- a/tests/scudo_unit_test_main.cpp
+++ b/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 @@
 // 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();
 }