[Sanitizers] Add support for -sanitize=memtag-stack (#85515)

This sanitizer adds MTE (memory tagging extension) checks to stack
variable accesses. Enablement simply requires setting an attribute on
function definitions, and the instrumentation is added by LLVM.

The corresponding swift-driver change is at:
swiftlang/swift-driver#2016.

rdar://161721201

Author: DanBlackwell

include/swift/Basic/Sanitizers.def

@@ -20,10 +20,11 @@
 
 // SANITIZER(enum_bit, kind, name, file)
 
-SANITIZER(0, Address,   address,    asan)
-SANITIZER(1, Thread,    thread,     tsan)
-SANITIZER(2, Undefined, undefined,  ubsan)
-SANITIZER(3, Fuzzer,    fuzzer,     fuzzer)
-SANITIZER(4, Scudo,     scudo,      scudo_standalone)
+SANITIZER(0, Address,     address,      asan)
+SANITIZER(1, Thread,      thread,       tsan)
+SANITIZER(2, Undefined,   undefined,    ubsan)
+SANITIZER(3, Fuzzer,      fuzzer,       fuzzer)
+SANITIZER(4, Scudo,       scudo,        scudo_standalone)
+SANITIZER(5, MemTagStack, memtag-stack, memtag-stack)
 
 #undef SANITIZER

lib/IRGen/IRGen.cpp

@@ -189,6 +189,11 @@ swift::getIRTargetOptions(const IRGenOptions &Opts, ASTContext &Ctx,
   }
 
   clang::TargetOptions &ClangOpts = Clang->getTargetInfo().getTargetOpts();
+
+  // Add +mte target feature when memtag-stack sanitizer is enabled
+  if (Opts.Sanitizers & SanitizerKind::MemTagStack)
+    ClangOpts.Features.push_back("+mte");
+
   return std::make_tuple(TargetOpts, ClangOpts.CPU, ClangOpts.Features, ClangOpts.Triple);
 }
 

lib/IRGen/IRGenSIL.cpp

@@ -1907,6 +1907,8 @@ IRGenSILFunction::IRGenSILFunction(IRGenModule &IGM, SILFunction *f,
   // being in the external file or via annotations.
   if (IGM.IRGen.Opts.Sanitizers & SanitizerKind::Address)
     CurFn->addFnAttr(llvm::Attribute::SanitizeAddress);
+  if (IGM.IRGen.Opts.Sanitizers & SanitizerKind::MemTagStack)
+    CurFn->addFnAttr(llvm::Attribute::SanitizeMemTag);
   if (IGM.IRGen.Opts.Sanitizers & SanitizerKind::Thread) {
     auto declContext = f->getDeclContext();
     if (isa_and_nonnull<DestructorDecl>(declContext)) {

lib/Option/SanitizerOptions.cpp

@@ -152,6 +152,11 @@ OptionSet<SanitizerKind> swift::parseSanitizerArgValues(
     bool isShared = (kind != SanitizerKind::Fuzzer);
     bool sanitizerSupported = sanitizerRuntimeLibExists(fileName, isShared);
 
+    if (kind == SanitizerKind::MemTagStack)
+      // MemTagStack requires no runtime, so ignore library check above
+      sanitizerSupported =
+          Triple.isOSDarwin() && Triple.getArch() == llvm::Triple::aarch64;
+
     // TSan is explicitly not supported for 32 bits.
     if (kind == SanitizerKind::Thread && !Triple.isArch64Bit())
       sanitizerSupported = false;
@@ -176,6 +181,18 @@ OptionSet<SanitizerKind> swift::parseSanitizerArgValues(
       Triple.getTriple());
   }
 
+  // MemTagStack and ASan can not be enabled concurrently.
+  if ((sanitizerSet & SanitizerKind::MemTagStack)
+        && (sanitizerSet & SanitizerKind::Address)) {
+    SmallString<128> b1;
+    SmallString<128> b2;
+    Diags.diagnose(SourceLoc(), diag::error_argument_not_allowed_with,
+        (A->getOption().getPrefixedName()
+            + toStringRef(SanitizerKind::Address)).toStringRef(b1),
+        (A->getOption().getPrefixedName()
+            + toStringRef(SanitizerKind::MemTagStack)).toStringRef(b2));
+  }
+
   // Address and thread sanitizers can not be enabled concurrently.
   if ((sanitizerSet & SanitizerKind::Thread)
         && (sanitizerSet & SanitizerKind::Address)) {

test/Driver/sanitizers.swift

@@ -39,6 +39,21 @@
 // RUN: %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=undefined -target x86_64-unknown-linux-gnu %s 2>&1 | %FileCheck -check-prefix=UBSAN_LINUX %s
 // RUN: %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=undefined -target x86_64-unknown-windows-msvc %s 2>&1 | %FileCheck -check-prefix=UBSAN_WINDOWS %s
 
+/*
+ * MemTagStack Sanitizer Tests
+ */
+// RUN: %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target arm64-apple-macosx10.9 %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN -check-prefix=MEMTAGSAN_OSX_ARM64 %s
+// RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target x86_64-apple-macosx10.9 %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN_OSX_X86 %s
+// RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target x86-apple-macosx10.9 %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN_OSX_32 %s
+// RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target x86_64-apple-ios7.1-simulator %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN_IOSSIM %s
+// RUN: %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target arm64-apple-ios7.1 %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN_IOS %s
+// RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target x86_64-apple-tvos9.0-simulator %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN_tvOS_SIM %s
+// RUN: %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target arm64-apple-tvos9.0 %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN_tvOS %s
+// RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target i386-apple-watchos2.0-simulator %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN_watchOS_SIM %s
+// RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target armv7k-apple-watchos2.0 %s 2>&1  | %FileCheck -check-prefix=MEMTAGSAN_watchOS %s
+// RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target x86_64-unknown-windows-msvc %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN_WINDOWS %s
+// RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -sanitize=memtag-stack -target x86_64-unknown-linux-gnu %s 2>&1 | %FileCheck -check-prefix=MEMTAGSAN_LINUX %s
+
 /*
  * Multiple Sanitizers At Once
  */
@@ -50,6 +65,7 @@
 // RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -target x86_64-apple-macosx10.9 -sanitize=address,unknown %s 2>&1 | %FileCheck -check-prefix=BADARG %s
 // RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -target x86_64-apple-macosx10.9 -sanitize=address -sanitize=unknown %s 2>&1 | %FileCheck -check-prefix=BADARG %s
 // RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -target x86_64-apple-macosx10.9 -sanitize=address,thread %s 2>&1 | %FileCheck -check-prefix=INCOMPATIBLESANITIZERS %s
+// RUN: not %swiftc_driver -sdk '""' -resource-dir %S/Inputs/fake-resource-dir/lib/swift/ -driver-print-jobs -target arm64-apple-macosx10.9 -sanitize=memtag-stack,address %s 2>&1 | %FileCheck -check-prefix=INCOMPATIBLESANITIZERS_2 %s
 
 /*
  * Make sure we don't accidentally add the sanitizer library path when building libraries or modules
@@ -106,7 +122,23 @@
 
 // UBSAN: -rpath @executable_path
 
+// MEMTAGSAN: swift
+// MEMTAGSAN: -sanitize=memtag-stack
+
+// MEMTAGSAN_OSX_ARM64-NOT: unsupported option '-sanitize=memtag-stack' for target 'arm64-apple-macosx10.9'
+// MEMTAGSAN_OSX_X86: unsupported option '-sanitize=memtag-stack' for target 'x86_64-apple-macosx10.9'
+// MEMTAGSAN_OSX_32: unsupported option '-sanitize=memtag-stack' for target 'x86-apple-macosx10.9'
+// MEMTAGSAN_IOSSIM: unsupported option '-sanitize=memtag-stack' for target 'x86_64-apple-ios7.1-simulator'
+// MEMTAGSAN_IOS-NOT: unsupported option '-sanitize=memtag-stack' for target 'arm64-apple-ios7.1'
+// MEMTAGSAN_tvOS_SIM: unsupported option '-sanitize=memtag-stack' for target 'x86_64-apple-tvos9.0-simulator'
+// MEMTAGSAN_tvOS-NOT: unsupported option '-sanitize=memtag-stack' for target 'arm64-apple-tvos9.0'
+// MEMTAGSAN_watchOS_SIM: unsupported option '-sanitize=memtag-stack' for target 'i386-apple-watchos2.0-simulator'
+// MEMTAGSAN_watchOS: unsupported option '-sanitize=memtag-stack' for target 'armv7k-apple-watchos2.0'
+// MEMTAGSAN_LINUX: unsupported option '-sanitize=memtag-stack' for target 'x86_64-unknown-linux-gnu'
+// MEMTAGSAN_WINDOWS: unsupported option '-sanitize=memtag-stack' for target 'x86_64-unknown-windows-msvc'
+
 // MULTIPLE_SAN_LINUX: -fsanitize=address,undefined,fuzzer
 
 // BADARG: unsupported argument 'unknown' to option '-sanitize='
 // INCOMPATIBLESANITIZERS: argument '-sanitize=address' is not allowed with '-sanitize=thread'
+// INCOMPATIBLESANITIZERS_2: argument '-sanitize=address' is not allowed with '-sanitize=memtag-stack'

test/IRGen/memtag_stack.swift

@@ -0,0 +1,14 @@
+// REQUIRES: CODEGENERATOR=AArch64
+// REQUIRES: OS=macosx
+
+// RUN: %swift -emit-ir -sanitize=memtag-stack -parse-as-library -target arm64-apple-macosx10.9 %s | %FileCheck %s
+
+// Check that functions have the sanitize_memtag attribute when -sanitize=memtag-stack is enabled
+
+// CHECK-LABEL: define {{.*}} @"${{.*}}testFunction{{.*}}"
+// CHECK-SAME: [[ATTRS:#[0-9]+]]
+// CHECK: attributes [[ATTRS]] = {{.*}} sanitize_memtag
+
+func testFunction() -> Int {
+  return 42
+}